Source code for langchain_community.document_loaders.s3_file

from __future__ import annotations

import os
import tempfile
from typing import TYPE_CHECKING, Any, Callable, List, Optional, Union

from langchain_community.document_loaders.unstructured import UnstructuredBaseLoader

if TYPE_CHECKING:
    import botocore


class S3FileLoader(UnstructuredBaseLoader):
    """从`Amazon AWS S3`加载文件。"""

    def __init__(
        self,
        bucket: str,
        key: str,
        *,
        region_name: Optional[str] = None,
        api_version: Optional[str] = None,
        use_ssl: Optional[bool] = True,
        verify: Union[str, bool, None] = None,
        endpoint_url: Optional[str] = None,
        aws_access_key_id: Optional[str] = None,
        aws_secret_access_key: Optional[str] = None,
        aws_session_token: Optional[str] = None,
        boto_config: Optional[botocore.client.Config] = None,
        mode: str = "single",
        post_processors: Optional[List[Callable]] = None,
        **unstructured_kwargs: Any,
    ):
        """初始化使用存储桶和键名。

:param bucket: S3存储桶的名称。
:param key: S3对象的键。

:param region_name: 与客户端关联的区域的名称。
    客户端与单个区域关联。

:param api_version: 要使用的API版本。默认情况下，botocore在创建客户端时将使用最新的API版本。只有在要使用客户端的先前API版本时才需要指定此参数。

:param use_ssl: 是否使用SSL。默认情况下，使用SSL。
    请注意，并非所有服务都支持非SSL连接。

:param verify: 是否验证SSL证书。
    默认情况下，会验证SSL证书。您可以提供以下值：

    * False - 不验证SSL证书。仍将使用SSL（除非use_ssl为False），但不会验证SSL证书。
    * path/to/cert/bundle.pem - CA证书捆绑包的文件名，用于使用。如果要使用与botocore使用的不同CA证书捆绑包，则可以指定此参数。

:param endpoint_url: 用于构建客户端的完整URL。
    通常，botocore在与服务通信时会自动构建适当的URL。您可以指定一个完整的URL（包括“http/https”方案）来覆盖此行为。如果提供了此值，则将忽略``use_ssl``。

:param aws_access_key_id: 创建客户端时要使用的访问密钥。
    这是完全可选的，如果未提供，则会自动使用会话配置的凭据。只有在要覆盖用于此特定客户端的凭据时才需要提供此参数。

:param aws_secret_access_key: 创建客户端时要使用的秘密密钥。与上面的aws_access_key_id具有相同的语义。

:param aws_session_token: 创建客户端时要使用的会话令牌。与上面的aws_access_key_id具有相同的语义。

:type boto_config: botocore.client.Config
:param boto_config: 高级boto3客户端配置选项。如果在客户端配置中指定了一个值，其值将优先于环境变量和配置值，但不会优先于显式传递给方法的值。如果在会话上设置了默认配置对象，则在创建客户端时使用的配置对象将是使用提供给此调用的配置调用``merge()``的结果。
:param mode: 读取文件的模式。有效选项为：single、paged和elements。
:param post_processors: 要应用于提取的元素的后处理函数。
:param **unstructured_kwargs: 在调用`partition`时传递的任意额外kwargs。
"""
        super().__init__(mode, post_processors, **unstructured_kwargs)
        self.bucket = bucket
        self.key = key
        self.region_name = region_name
        self.api_version = api_version
        self.use_ssl = use_ssl
        self.verify = verify
        self.endpoint_url = endpoint_url
        self.aws_access_key_id = aws_access_key_id
        self.aws_secret_access_key = aws_secret_access_key
        self.aws_session_token = aws_session_token
        self.boto_config = boto_config

    def _get_elements(self) -> List:
        """获取元素。"""
        from unstructured.partition.auto import partition

        try:
            import boto3
        except ImportError:
            raise ImportError(
                "Could not import `boto3` python package. "
                "Please install it with `pip install boto3`."
            )
        s3 = boto3.client(
            "s3",
            region_name=self.region_name,
            api_version=self.api_version,
            use_ssl=self.use_ssl,
            verify=self.verify,
            endpoint_url=self.endpoint_url,
            aws_access_key_id=self.aws_access_key_id,
            aws_secret_access_key=self.aws_secret_access_key,
            aws_session_token=self.aws_session_token,
            config=self.boto_config,
        )
        with tempfile.TemporaryDirectory() as temp_dir:
            file_path = f"{temp_dir}/{self.key}"
            os.makedirs(os.path.dirname(file_path), exist_ok=True)
            s3.download_file(self.bucket, self.key, file_path)
            return partition(filename=file_path, **self.unstructured_kwargs)

    def _get_metadata(self) -> dict:
        return {"source": f"s3://{self.bucket}/{self.key}"}