Source code for langchain_community.document_loaders.s3_directory

from __future__ import annotations

from typing import TYPE_CHECKING, List, Optional, Union

from langchain_core.documents import Document

from langchain_community.document_loaders.base import BaseLoader
from langchain_community.document_loaders.s3_file import S3FileLoader

if TYPE_CHECKING:
    import botocore


class S3DirectoryLoader(BaseLoader):
    """从`Amazon AWS S3`目录加载。"""

    def __init__(
        self,
        bucket: str,
        prefix: 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,
    ):
        """初始化使用存储桶和键名。

:param bucket: S3存储桶的名称。
:param prefix: 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()``后的结果。
"""
        self.bucket = bucket
        self.prefix = prefix
        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 load(self) -> List[Document]:
        """加载文档。"""
        try:
            import boto3
        except ImportError:
            raise ImportError(
                "Could not import boto3 python package. "
                "Please install it with `pip install boto3`."
            )
        s3 = boto3.resource(
            "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,
        )
        bucket = s3.Bucket(self.bucket)
        docs = []
        for obj in bucket.objects.filter(Prefix=self.prefix):
            # Skip directories
            if obj.size == 0 and obj.key.endswith("/"):
                continue
            loader = S3FileLoader(
                self.bucket,
                obj.key,
                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,
                boto_config=self.boto_config,
            )
            docs.extend(loader.load())
        return docs