llama-index

1
"""Base reader class."""
2

3
from abc import ABC
4
from typing import TYPE_CHECKING, Any, Dict, Iterable, List
5

6
if TYPE_CHECKING:
7
    from llama_index.legacy.bridge.langchain import Document as LCDocument
8
from llama_index.legacy.bridge.pydantic import Field
9
from llama_index.legacy.schema import BaseComponent, Document
10

11

12
class BaseReader(ABC):
13
    """Utilities for loading data from a directory."""
14

15
    def lazy_load_data(self, *args: Any, **load_kwargs: Any) -> Iterable[Document]:
16
        """Load data from the input directory lazily."""
17
        raise NotImplementedError(
18
            f"{self.__class__.__name__} does not provide lazy_load_data method currently"
19
        )
20

21
    def load_data(self, *args: Any, **load_kwargs: Any) -> List[Document]:
22
        """Load data from the input directory."""
23
        return list(self.lazy_load_data(*args, **load_kwargs))
24

25
    def load_langchain_documents(self, **load_kwargs: Any) -> List["LCDocument"]:
26
        """Load data in LangChain document format."""
27
        docs = self.load_data(**load_kwargs)
28
        return [d.to_langchain_format() for d in docs]
29

30

31
class BasePydanticReader(BaseReader, BaseComponent):
32
    """Serialiable Data Loader with Pydatnic."""
33

34
    is_remote: bool = Field(
35
        default=False,
36
        description="Whether the data is loaded from a remote API or a local file.",
37
    )
38

39
    class Config:
40
        arbitrary_types_allowed = True
41

42

43
class ReaderConfig(BaseComponent):
44
    """Represents a reader and it's input arguments."""
45

46
    reader: BasePydanticReader = Field(..., description="Reader to use.")
47
    reader_args: List[Any] = Field(default_factory=list, description="Reader args.")
48
    reader_kwargs: Dict[str, Any] = Field(
49
        default_factory=dict, description="Reader kwargs."
50
    )
51

52
    class Config:
53
        arbitrary_types_allowed = True
54

55
    @classmethod
56
    def class_name(cls) -> str:
57
        """Get the name identifier of the class."""
58
        return "ReaderConfig"
59

60
    def to_dict(self, **kwargs: Any) -> Dict[str, Any]:
61
        """Convert the class to a dictionary."""
62
        return {
63
            "loader": self.reader.to_dict(**kwargs),
64
            "reader_args": self.reader_args,
65
            "reader_kwargs": self.reader_kwargs,
66
            "class_name": self.class_name(),
67
        }
68

69
    def read(self) -> List[Document]:
70
        """Call the loader with the given arguments."""
71
        return self.reader.load_data(*self.reader_args, **self.reader_kwargs)
72