mirror of https://github.com/hwchase17/langchain
Bos loader (#11525)
**Description:** Add BaiduCloud BOS document loader. --------- Co-authored-by: chenweixu01 <chenweixu01@baidu.com> Co-authored-by: root <root@icoding-cwx.bcc-szzj.baidu.com> Co-authored-by: Bagatur <baskaryan@gmail.com>pull/11684/head^2
parent
fbb82608cd
commit
2363c02cf3
@ -0,0 +1,52 @@
|
||||
from typing import Any, Iterator, List
|
||||
|
||||
from langchain.docstore.document import Document
|
||||
from langchain.document_loaders.base import BaseLoader
|
||||
|
||||
|
||||
class BaiduBOSDirectoryLoader(BaseLoader):
|
||||
"""Load from `Baidu BOS directory`."""
|
||||
|
||||
def __init__(self, conf: Any, bucket: str, prefix: str = ""):
|
||||
"""Initialize with BOS config, bucket and prefix.
|
||||
:param conf(BosConfig): BOS config.
|
||||
:param bucket(str): BOS bucket.
|
||||
:param prefix(str): prefix.
|
||||
"""
|
||||
self.conf = conf
|
||||
self.bucket = bucket
|
||||
self.prefix = prefix
|
||||
|
||||
def load(self) -> List[Document]:
|
||||
return list(self.lazy_load())
|
||||
|
||||
def lazy_load(self) -> Iterator[Document]:
|
||||
"""Load documents."""
|
||||
try:
|
||||
from baidubce.services.bos.bos_client import BosClient
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Please install bce-python-sdk with `pip install bce-python-sdk`."
|
||||
)
|
||||
client = BosClient(self.conf)
|
||||
contents = []
|
||||
marker = ""
|
||||
while True:
|
||||
response = client.list_objects(
|
||||
bucket_name=self.bucket,
|
||||
prefix=self.prefix,
|
||||
marker=marker,
|
||||
max_keys=1000,
|
||||
)
|
||||
contents_len = len(response.contents)
|
||||
contents.extend(response.contents)
|
||||
if response.is_truncated or contents_len < int(str(response.max_keys)):
|
||||
break
|
||||
marker = response.next_marker
|
||||
from baidu_bos_file import BaiduBOSFileLoader
|
||||
|
||||
for content in contents:
|
||||
if str(content.key).endswith("/"):
|
||||
continue
|
||||
loader = BaiduBOSFileLoader(self.conf, self.bucket, str(content.key))
|
||||
yield loader.load()[0]
|
@ -0,0 +1,53 @@
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
from typing import Any, Iterator, List
|
||||
|
||||
from langchain.docstore.document import Document
|
||||
from langchain.document_loaders.base import BaseLoader
|
||||
from langchain.document_loaders.unstructured import UnstructuredFileLoader
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaiduBOSFileLoader(BaseLoader):
|
||||
"""Load from `Baidu Cloud BOS` file."""
|
||||
|
||||
def __init__(self, conf: Any, bucket: str, key: str):
|
||||
"""Initialize with BOS config, bucket and key name.
|
||||
:param conf(BceClientConfiguration): BOS config.
|
||||
:param bucket(str): BOS bucket.
|
||||
:param key(str): BOS file key.
|
||||
"""
|
||||
self.conf = conf
|
||||
self.bucket = bucket
|
||||
self.key = key
|
||||
|
||||
def load(self) -> List[Document]:
|
||||
return list(self.lazy_load())
|
||||
|
||||
def lazy_load(self) -> Iterator[Document]:
|
||||
"""Load documents."""
|
||||
try:
|
||||
from baidubce.services.bos.bos_client import BosClient
|
||||
except ImportError:
|
||||
raise ImportError(
|
||||
"Please using `pip install bce-python-sdk`"
|
||||
+ " before import bos related package."
|
||||
)
|
||||
|
||||
# Initialize BOS Client
|
||||
client = BosClient(self.conf)
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
file_path = f"{temp_dir}/{self.bucket}/{self.key}"
|
||||
os.makedirs(os.path.dirname(file_path), exist_ok=True)
|
||||
# Download the file to a destination
|
||||
logger.debug(f"get object key {self.key} to file {file_path}")
|
||||
client.get_object_to_file(self.bucket, self.key, file_path)
|
||||
try:
|
||||
loader = UnstructuredFileLoader(file_path)
|
||||
documents = loader.load()
|
||||
return iter(documents)
|
||||
except Exception as ex:
|
||||
logger.error(f"load document error = {ex}")
|
||||
return iter([Document(page_content="")])
|
Loading…
Reference in New Issue