|
|
@ -1,7 +1,7 @@
|
|
|
|
"""Load documents from a GCS file."""
|
|
|
|
"""Load documents from a GCS file."""
|
|
|
|
import os
|
|
|
|
import os
|
|
|
|
import tempfile
|
|
|
|
import tempfile
|
|
|
|
from typing import List
|
|
|
|
from typing import Callable, List, Optional
|
|
|
|
|
|
|
|
|
|
|
|
from langchain.docstore.document import Document
|
|
|
|
from langchain.docstore.document import Document
|
|
|
|
from langchain.document_loaders.base import BaseLoader
|
|
|
|
from langchain.document_loaders.base import BaseLoader
|
|
|
@ -11,18 +11,42 @@ from langchain.document_loaders.unstructured import UnstructuredFileLoader
|
|
|
|
class GCSFileLoader(BaseLoader):
|
|
|
|
class GCSFileLoader(BaseLoader):
|
|
|
|
"""Load Documents from a GCS file."""
|
|
|
|
"""Load Documents from a GCS file."""
|
|
|
|
|
|
|
|
|
|
|
|
def __init__(self, project_name: str, bucket: str, blob: str):
|
|
|
|
def __init__(
|
|
|
|
|
|
|
|
self,
|
|
|
|
|
|
|
|
project_name: str,
|
|
|
|
|
|
|
|
bucket: str,
|
|
|
|
|
|
|
|
blob: str,
|
|
|
|
|
|
|
|
loader_func: Optional[Callable[[str], BaseLoader]] = None,
|
|
|
|
|
|
|
|
):
|
|
|
|
"""Initialize with bucket and key name.
|
|
|
|
"""Initialize with bucket and key name.
|
|
|
|
|
|
|
|
|
|
|
|
Args:
|
|
|
|
Args:
|
|
|
|
project_name: The name of the project to load
|
|
|
|
project_name: The name of the project to load
|
|
|
|
bucket: The name of the GCS bucket.
|
|
|
|
bucket: The name of the GCS bucket.
|
|
|
|
blob: The name of the GCS blob to load.
|
|
|
|
blob: The name of the GCS blob to load.
|
|
|
|
|
|
|
|
loader_func: A loader function that instatiates a loader based on a
|
|
|
|
|
|
|
|
file_path argument. If nothing is provided, the
|
|
|
|
|
|
|
|
UnstructuredFileLoader is used.
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
Examples:
|
|
|
|
|
|
|
|
To use an alternative PDF loader:
|
|
|
|
|
|
|
|
>> from from langchain.document_loaders import PyPDFLoader
|
|
|
|
|
|
|
|
>> loader = GCSFileLoader(..., loader_func=PyPDFLoader)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
To use UnstructuredFileLoader with additional arguments:
|
|
|
|
|
|
|
|
>> loader = GCSFileLoader(...,
|
|
|
|
|
|
|
|
>> loader_func=lambda x: UnstructuredFileLoader(x, mode="elements"))
|
|
|
|
|
|
|
|
|
|
|
|
"""
|
|
|
|
"""
|
|
|
|
self.bucket = bucket
|
|
|
|
self.bucket = bucket
|
|
|
|
self.blob = blob
|
|
|
|
self.blob = blob
|
|
|
|
self.project_name = project_name
|
|
|
|
self.project_name = project_name
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def default_loader_func(file_path: str) -> BaseLoader:
|
|
|
|
|
|
|
|
return UnstructuredFileLoader(file_path)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
self._loader_func = loader_func if loader_func else default_loader_func
|
|
|
|
|
|
|
|
|
|
|
|
def load(self) -> List[Document]:
|
|
|
|
def load(self) -> List[Document]:
|
|
|
|
"""Load documents."""
|
|
|
|
"""Load documents."""
|
|
|
|
try:
|
|
|
|
try:
|
|
|
@ -44,5 +68,5 @@ class GCSFileLoader(BaseLoader):
|
|
|
|
os.makedirs(os.path.dirname(file_path), exist_ok=True)
|
|
|
|
os.makedirs(os.path.dirname(file_path), exist_ok=True)
|
|
|
|
# Download the file to a destination
|
|
|
|
# Download the file to a destination
|
|
|
|
blob.download_to_filename(file_path)
|
|
|
|
blob.download_to_filename(file_path)
|
|
|
|
loader = UnstructuredFileLoader(file_path)
|
|
|
|
loader = self._loader_func(file_path)
|
|
|
|
return loader.load()
|
|
|
|
return loader.load()
|
|
|
|