codekingpro/portable-devtools

1from __future__ import annotations2 3from typing import Any, Iterator, Optional, Sequence4 5from langchain_core.documents import Document6 7from langchain_community.document_loaders.base import BaseLoader8from langchain_community.utilities.max_compute import MaxComputeAPIWrapper9 10 11class MaxComputeLoader(BaseLoader):12    """Load from `Alibaba Cloud MaxCompute` table."""13 14    def __init__(15        self,16        query: str,17        api_wrapper: MaxComputeAPIWrapper,18        *,19        page_content_columns: Optional[Sequence[str]] = None,20        metadata_columns: Optional[Sequence[str]] = None,21    ):22        """Initialize Alibaba Cloud MaxCompute document loader.23 24        Args:25            query: SQL query to execute.26            api_wrapper: MaxCompute API wrapper.27            page_content_columns: The columns to write into the `page_content` of the28                Document. If unspecified, all columns will be written to `page_content`.29            metadata_columns: The columns to write into the `metadata` of the Document.30                If unspecified, all columns not added to `page_content` will be written.31        """32        self.query = query33        self.api_wrapper = api_wrapper34        self.page_content_columns = page_content_columns35        self.metadata_columns = metadata_columns36 37    @classmethod38    def from_params(39        cls,40        query: str,41        endpoint: str,42        project: str,43        *,44        access_id: Optional[str] = None,45        secret_access_key: Optional[str] = None,46        **kwargs: Any,47    ) -> MaxComputeLoader:48        """Convenience constructor that builds the MaxCompute API wrapper from49            given parameters.50 51        Args:52            query: SQL query to execute.53            endpoint: MaxCompute endpoint.54            project: A project is a basic organizational unit of MaxCompute, which is55                similar to a database.56            access_id: MaxCompute access ID. Should be passed in directly or set as the57                environment variable `MAX_COMPUTE_ACCESS_ID`.58            secret_access_key: MaxCompute secret access key. Should be passed in59                directly or set as the environment variable60                `MAX_COMPUTE_SECRET_ACCESS_KEY`.61        """62        api_wrapper = MaxComputeAPIWrapper.from_params(63            endpoint, project, access_id=access_id, secret_access_key=secret_access_key64        )65        return cls(query, api_wrapper, **kwargs)66 67    def lazy_load(self) -> Iterator[Document]:68        for row in self.api_wrapper.query(self.query):69            if self.page_content_columns:70                page_content_data = {71                    k: v for k, v in row.items() if k in self.page_content_columns72                }73            else:74                page_content_data = row75            page_content = "\n".join(f"{k}: {v}" for k, v in page_content_data.items())76            if self.metadata_columns:77                metadata = {k: v for k, v in row.items() if k in self.metadata_columns}78            else:79                metadata = {k: v for k, v in row.items() if k not in page_content_data}80            yield Document(page_content=page_content, metadata=metadata)81