codekingpro/portable-devtools

1from __future__ import annotations2 3import os4import tempfile5from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union6 7from langchain_community.document_loaders.unstructured import UnstructuredBaseLoader8 9if TYPE_CHECKING:10    import botocore11 12 13class S3FileLoader(UnstructuredBaseLoader):14    """Load from `Amazon AWS S3` file."""15 16    def __init__(17        self,18        bucket: str,19        key: str,20        *,21        region_name: Optional[str] = None,22        api_version: Optional[str] = None,23        use_ssl: Optional[bool] = True,24        verify: Union[str, bool, None] = None,25        endpoint_url: Optional[str] = None,26        aws_access_key_id: Optional[str] = None,27        aws_secret_access_key: Optional[str] = None,28        aws_session_token: Optional[str] = None,29        boto_config: Optional[botocore.client.Config] = None,30        mode: str = "single",31        post_processors: Optional[List[Callable]] = None,32        **unstructured_kwargs: Any,33    ):34        """Initialize with bucket and key name.35 36        :param bucket: The name of the S3 bucket.37        :param key: The key of the S3 object.38 39        :param region_name: The name of the region associated with the client.40            A client is associated with a single region.41 42        :param api_version: The API version to use.  By default, botocore will43            use the latest API version when creating a client.  You only need44            to specify this parameter if you want to use a previous API version45            of the client.46 47        :param use_ssl: Whether or not to use SSL.  By default, SSL is used.48            Note that not all services support non-ssl connections.49 50        :param verify: Whether or not to verify SSL certificates.51            By default SSL certificates are verified.  You can provide the52            following values:53 54            * False - do not validate SSL certificates.  SSL will still be55              used (unless use_ssl is False), but SSL certificates56              will not be verified.57            * path/to/cert/bundle.pem - A filename of the CA cert bundle to58              uses.  You can specify this argument if you want to use a59              different CA cert bundle than the one used by botocore.60 61        :param endpoint_url: The complete URL to use for the constructed62            client.  Normally, botocore will automatically construct the63            appropriate URL to use when communicating with a service.  You can64            specify a complete URL (including the "http/https" scheme) to65            override this behavior.  If this value is provided, then66            ``use_ssl`` is ignored.67 68        :param aws_access_key_id: The access key to use when creating69            the client.  This is entirely optional, and if not provided,70            the credentials configured for the session will automatically71            be used.  You only need to provide this argument if you want72            to override the credentials used for this specific client.73 74        :param aws_secret_access_key: The secret key to use when creating75            the client.  Same semantics as aws_access_key_id above.76 77        :param aws_session_token: The session token to use when creating78            the client.  Same semantics as aws_access_key_id above.79 80        :type boto_config: botocore.client.Config81        :param boto_config: Advanced boto3 client configuration options. If a value82            is specified in the client config, its value will take precedence83            over environment variables and configuration values, but not over84            a value passed explicitly to the method. If a default config85            object is set on the session, the config object used when creating86            the client will be the result of calling ``merge()`` on the87            default config with the config provided to this call.88        :param mode: Mode in which to read the file. Valid options are: single,89            paged and elements.90        :param post_processors: Post processing functions to be applied to91            extracted elements.92        :param **unstructured_kwargs: Arbitrary additional kwargs to pass in when93            calling `partition`94        """95        super().__init__(mode, post_processors, **unstructured_kwargs)96        self.bucket = bucket97        self.key = key98        self.region_name = region_name99        self.api_version = api_version100        self.use_ssl = use_ssl101        self.verify = verify102        self.endpoint_url = endpoint_url103        self.aws_access_key_id = aws_access_key_id104        self.aws_secret_access_key = aws_secret_access_key105        self.aws_session_token = aws_session_token106        self.boto_config = boto_config107 108    def _get_elements(self) -> List:109        """Get elements."""110        from unstructured.partition.auto import partition111 112        try:113            import boto3114        except ImportError:115            raise ImportError(116                "Could not import `boto3` python package. "117                "Please install it with `pip install boto3`."118            )119        s3 = boto3.client(120            "s3",121            region_name=self.region_name,122            api_version=self.api_version,123            use_ssl=self.use_ssl,124            verify=self.verify,125            endpoint_url=self.endpoint_url,126            aws_access_key_id=self.aws_access_key_id,127            aws_secret_access_key=self.aws_secret_access_key,128            aws_session_token=self.aws_session_token,129            config=self.boto_config,130        )131        with tempfile.TemporaryDirectory() as temp_dir:132            file_path = f"{temp_dir}/{self.key}"133            os.makedirs(os.path.dirname(file_path), exist_ok=True)134            s3.download_file(self.bucket, self.key, file_path)135            return partition(filename=file_path, **self.unstructured_kwargs)136 137    def _get_metadata(self) -> dict:138        return {"source": f"s3://{self.bucket}/{self.key}"}139