codekingpro/portable-devtools

1# coding:utf-82from typing import List, Optional3 4from langchain_core.documents import Document5 6from langchain_community.document_loaders.base import BaseLoader7from langchain_community.document_loaders.obs_file import OBSFileLoader8 9 10class OBSDirectoryLoader(BaseLoader):11    """Load from `Huawei OBS directory`."""12 13    def __init__(14        self,15        bucket: str,16        endpoint: str,17        config: Optional[dict] = None,18        prefix: str = "",19    ):20        """Initialize the OBSDirectoryLoader with the specified settings.21 22        Args:23            bucket (str): The name of the OBS bucket to be used.24            endpoint (str): The endpoint URL of your OBS bucket.25            config (dict): The parameters for connecting to OBS, provided as a dictionary. The dictionary could have the following keys:26                - "ak" (str, optional): Your OBS access key (required if `get_token_from_ecs` is False and bucket policy is not public read).27                - "sk" (str, optional): Your OBS secret key (required if `get_token_from_ecs` is False and bucket policy is not public read).28                - "token" (str, optional): Your security token (required if using temporary credentials).29                - "get_token_from_ecs" (bool, optional): Whether to retrieve the security token from ECS. Defaults to False if not provided. If set to True, `ak`, `sk`, and `token` will be ignored.30            prefix (str, optional): The prefix to be added to the OBS key. Defaults to "".31 32        Note:33            Before using this class, make sure you have registered with OBS and have the necessary credentials. The `ak`, `sk`, and `endpoint` values are mandatory unless `get_token_from_ecs` is True or the bucket policy is public read. `token` is required when using temporary credentials.34        Example:35            To create a new OBSDirectoryLoader:36            ```37            config = {38                "ak": "your-access-key",39                "sk": "your-secret-key"40            }41            ```42            directory_loader = OBSDirectoryLoader("your-bucket-name", "your-end-endpoint", config, "your-prefix")43        """  # noqa: E50144        try:45            from obs import ObsClient46        except ImportError:47            raise ImportError(48                "Could not import esdk-obs-python python package. "49                "Please install it with `pip install esdk-obs-python`."50            )51        if not config:52            config = dict()53        if config.get("get_token_from_ecs"):54            self.client = ObsClient(server=endpoint, security_provider_policy="ECS")55        else:56            self.client = ObsClient(57                access_key_id=config.get("ak"),58                secret_access_key=config.get("sk"),59                security_token=config.get("token"),60                server=endpoint,61            )62 63        self.bucket = bucket64        self.prefix = prefix65 66    def load(self) -> List[Document]:67        """Load documents."""68        max_num = 100069        mark = None70        docs = []71        while True:72            resp = self.client.listObjects(73                self.bucket, prefix=self.prefix, marker=mark, max_keys=max_num74            )75            if resp.status < 300:76                for content in resp.body.contents:77                    loader = OBSFileLoader(self.bucket, content.key, client=self.client)78                    docs.extend(loader.load())79                if resp.body.is_truncated is True:80                    mark = resp.body.next_marker81                else:82                    break83        return docs84