codekingpro/portable-devtools

1import logging2from typing import Iterator, List, Optional3 4from langchain_core.documents import Document5 6from langchain_community.document_loaders.base import BaseLoader7 8logger = logging.getLogger(__name__)9 10 11class CouchbaseLoader(BaseLoader):12    """Load documents from `Couchbase`.13 14    Each document represents one row of the result. The `page_content_fields` are15    written into the `page_content`of the document. The `metadata_fields` are written16    into the `metadata` of the document. By default, all columns are written into17    the `page_content` and none into the `metadata`.18    """19 20    def __init__(21        self,22        connection_string: str,23        db_username: str,24        db_password: str,25        query: str,26        *,27        page_content_fields: Optional[List[str]] = None,28        metadata_fields: Optional[List[str]] = None,29    ) -> None:30        """Initialize Couchbase document loader.31 32        Args:33            connection_string (str): The connection string to the Couchbase cluster.34            db_username (str): The username to connect to the Couchbase cluster.35            db_password (str): The password to connect to the Couchbase cluster.36            query (str): The SQL++ query to execute.37            page_content_fields (Optional[List[str]]): The columns to write into the38                `page_content` field of the document. By default, all columns are39                written.40            metadata_fields (Optional[List[str]]): The columns to write into the41                `metadata` field of the document. By default, no columns are written.42        """43        try:44            from couchbase.auth import PasswordAuthenticator45            from couchbase.cluster import Cluster46            from couchbase.options import ClusterOptions47        except ImportError as e:48            raise ImportError(49                "Could not import couchbase package."50                "Please install couchbase SDK with `pip install couchbase`."51            ) from e52        if not connection_string:53            raise ValueError("connection_string must be provided.")54 55        if not db_username:56            raise ValueError("db_username must be provided.")57 58        if not db_password:59            raise ValueError("db_password must be provided.")60 61        auth = PasswordAuthenticator(62            db_username,63            db_password,64        )65 66        self.cluster: Cluster = Cluster(connection_string, ClusterOptions(auth))67        self.query = query68        self.page_content_fields = page_content_fields69        self.metadata_fields = metadata_fields70 71    def lazy_load(self) -> Iterator[Document]:72        """Load Couchbase data into Document objects lazily."""73        from datetime import timedelta74 75        # Ensure connection to Couchbase cluster76        self.cluster.wait_until_ready(timedelta(seconds=5))77 78        # Run SQL++ Query79        result = self.cluster.query(self.query)80        for row in result:81            metadata_fields = self.metadata_fields82            page_content_fields = self.page_content_fields83 84            if not page_content_fields:85                page_content_fields = list(row.keys())86 87            if not metadata_fields:88                metadata_fields = []89 90            metadata = {field: row[field] for field in metadata_fields}91 92            document = "\n".join(93                f"{k}: {v}" for k, v in row.items() if k in page_content_fields94            )95 96            yield (Document(page_content=document, metadata=metadata))97