codekingpro/portable-devtools

1from typing import Any, Callable, Iterator, List, Optional, Tuple2 3from langchain_core.documents import Document4 5from langchain_community.document_loaders.base import BaseLoader6 7 8def default_joiner(docs: List[Tuple[str, Any]]) -> str:9    """Default joiner for content columns."""10    return "\n".join([doc[1] for doc in docs])11 12 13class ColumnNotFoundError(Exception):14    """Column not found error."""15 16    def __init__(self, missing_key: str, query: str):17        super().__init__(f'Column "{missing_key}" not selected in query:\n{query}')18 19 20class RocksetLoader(BaseLoader):21    """Load from a `Rockset` database.22 23    To use, you should have the `rockset` python package installed.24 25    Example:26        .. code-block:: python27 28            # This code will load 3 records from the "langchain_demo"29            # collection as Documents, with the `text` column used as30            # the content31 32            from langchain_community.document_loaders import RocksetLoader33            from rockset import RocksetClient, Regions, models34 35            loader = RocksetLoader(36                RocksetClient(Regions.usw2a1, "<api key>"),37                models.QueryRequestSql(38                    query="select * from langchain_demo limit 3"39                ),40                ["text"]41            )42        )43    """44 45    def __init__(46        self,47        client: Any,48        query: Any,49        content_keys: List[str],50        metadata_keys: Optional[List[str]] = None,51        content_columns_joiner: Callable[[List[Tuple[str, Any]]], str] = default_joiner,52    ):53        """Initialize with Rockset client.54 55        Args:56            client: Rockset client object.57            query: Rockset query object.58            content_keys: The collection columns to be written into the `page_content`59                of the Documents.60            metadata_keys: The collection columns to be written into the `metadata` of61                the Documents. By default, this is all the keys in the document.62            content_columns_joiner: Method that joins content_keys and its values into a63                string. It's method that takes in a List[Tuple[str, Any]]],64                representing a list of tuples of (column name, column value).65                By default, this is a method that joins each column value with a new66                line. This method is only relevant if there are multiple content_keys.67        """68        try:69            from rockset import QueryPaginator, RocksetClient70            from rockset.models import QueryRequestSql71        except ImportError:72            raise ImportError(73                "Could not import rockset client python package. "74                "Please install it with `pip install rockset`."75            )76 77        if not isinstance(client, RocksetClient):78            raise ValueError(79                f"client should be an instance of rockset.RocksetClient, "80                f"got {type(client)}"81            )82 83        if not isinstance(query, QueryRequestSql):84            raise ValueError(85                f"query should be an instance of rockset.model.QueryRequestSql, "86                f"got {type(query)}"87            )88 89        self.client = client90        self.query = query91        self.content_keys = content_keys92        self.content_columns_joiner = content_columns_joiner93        self.metadata_keys = metadata_keys94        self.paginator = QueryPaginator95        self.request_model = QueryRequestSql96 97        try:98            self.client.set_application("langchain")99        except AttributeError:100            # ignore101            pass102 103    def lazy_load(self) -> Iterator[Document]:104        query_results = self.client.Queries.query(105            sql=self.query106        ).results  # execute the SQL query107        for doc in query_results:  # for each doc in the response108            try:109                yield Document(110                    page_content=self.content_columns_joiner(111                        [(col, doc[col]) for col in self.content_keys]112                    ),113                    metadata={col: doc[col] for col in self.metadata_keys}114                    if self.metadata_keys is not None115                    else doc,116                )  # try to yield the Document117            except (118                KeyError119            ) as e:  # either content_columns or metadata_columns is invalid120                raise ColumnNotFoundError(121                    e.args[0], self.query122                )  # raise that the column isn't in the db schema123