codekingpro/portable-devtools

1from typing import Dict, List, Optional, cast2 3from langchain_core.documents import Document4 5from langchain_community.document_loaders.base import BaseLoader6 7 8class DuckDBLoader(BaseLoader):9    """Load from `DuckDB`.10 11    Each document represents one row of the result. The `page_content_columns`12    are written into the `page_content` of the document. The `metadata_columns`13    are written into the `metadata` of the document. By default, all columns14    are written into the `page_content` and none into the `metadata`.15    """16 17    def __init__(18        self,19        query: str,20        database: str = ":memory:",21        read_only: bool = False,22        config: Optional[Dict[str, str]] = None,23        page_content_columns: Optional[List[str]] = None,24        metadata_columns: Optional[List[str]] = None,25    ):26        """27 28        Args:29            query: The query to execute.30            database: The database to connect to. Defaults to ":memory:".31            read_only: Whether to open the database in read-only mode.32              Defaults to False.33            config: A dictionary of configuration options to pass to the database.34              Optional.35            page_content_columns: The columns to write into the `page_content`36              of the document. Optional.37            metadata_columns: The columns to write into the `metadata` of the document.38              Optional.39        """40        self.query = query41        self.database = database42        self.read_only = read_only43        self.config = config or {}44        self.page_content_columns = page_content_columns45        self.metadata_columns = metadata_columns46 47    def load(self) -> List[Document]:48        try:49            import duckdb50        except ImportError:51            raise ImportError(52                "Could not import duckdb python package. "53                "Please install it with `pip install duckdb`."54            )55 56        docs = []57        with duckdb.connect(58            database=self.database, read_only=self.read_only, config=self.config59        ) as con:60            query_result = con.execute(self.query)61            results = query_result.fetchall()62            description = cast(list, query_result.description)63            field_names = [c[0] for c in description]64 65            if self.page_content_columns is None:66                page_content_columns = field_names67            else:68                page_content_columns = self.page_content_columns69 70            if self.metadata_columns is None:71                metadata_columns = []72            else:73                metadata_columns = self.metadata_columns74 75            for result in results:76                page_content = "\n".join(77                    f"{column}: {result[field_names.index(column)]}"78                    for column in page_content_columns79                )80 81                metadata = {82                    column: result[field_names.index(column)]83                    for column in metadata_columns84                }85 86                doc = Document(page_content=page_content, metadata=metadata)87                docs.append(doc)88 89        return docs90