forked from Archives/langchain
Harrison/gdrive pdf (#1433)
Co-authored-by: LM <93918064+LuisMalhadas@users.noreply.github.com> Co-authored-by: Luis Malhadas <luis@sia.so>
This commit is contained in:
parent
64c6435545
commit
519f0187b6
@ -26,16 +26,26 @@ class GoogleDriveLoader(BaseLoader, BaseModel):
|
||||
token_path: Path = Path.home() / ".credentials" / "token.json"
|
||||
folder_id: Optional[str] = None
|
||||
document_ids: Optional[List[str]] = None
|
||||
file_ids: Optional[List[str]] = None
|
||||
|
||||
@root_validator
|
||||
def validate_folder_id_or_document_ids(
|
||||
cls, values: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
"""Validate that either folder_id or document_ids is set, but not both."""
|
||||
if values.get("folder_id") and values.get("document_ids"):
|
||||
raise ValueError("Cannot specify both folder_id and document_ids")
|
||||
if not values.get("folder_id") and not values.get("document_ids"):
|
||||
raise ValueError("Must specify either folder_id or document_ids")
|
||||
if values.get("folder_id") and (
|
||||
values.get("document_ids") or values.get("file_ids")
|
||||
):
|
||||
raise ValueError(
|
||||
"Cannot specify both folder_id and document_ids nor "
|
||||
"folder_id and file_ids"
|
||||
)
|
||||
if (
|
||||
not values.get("folder_id")
|
||||
and not values.get("document_ids")
|
||||
and not values.get("file_ids")
|
||||
):
|
||||
raise ValueError("Must specify either folder_id, document_ids, or file_ids")
|
||||
return values
|
||||
|
||||
@validator("credentials_path")
|
||||
@ -115,13 +125,16 @@ class GoogleDriveLoader(BaseLoader, BaseModel):
|
||||
.execute()
|
||||
)
|
||||
items = results.get("files", [])
|
||||
returns = []
|
||||
for item in items:
|
||||
if item["mimeType"] == "application/vnd.google-apps.document":
|
||||
returns.append(self._load_document_from_id(item["id"]))
|
||||
elif item["mimeType"] == "application/pdf":
|
||||
returns.extend(self._load_file_from_id(item["id"]))
|
||||
else:
|
||||
pass
|
||||
|
||||
return [
|
||||
self._load_document_from_id(item["id"])
|
||||
for item in items
|
||||
# Only support Google Docs for now
|
||||
if item["mimeType"] == "application/vnd.google-apps.document"
|
||||
]
|
||||
return returns
|
||||
|
||||
def _load_documents_from_ids(self) -> List[Document]:
|
||||
"""Load documents from a list of IDs."""
|
||||
@ -130,9 +143,53 @@ class GoogleDriveLoader(BaseLoader, BaseModel):
|
||||
|
||||
return [self._load_document_from_id(doc_id) for doc_id in self.document_ids]
|
||||
|
||||
def _load_file_from_id(self, id: str) -> List[Document]:
|
||||
"""Load a file from an ID."""
|
||||
from io import BytesIO
|
||||
|
||||
from googleapiclient.discovery import build
|
||||
from googleapiclient.http import MediaIoBaseDownload
|
||||
|
||||
creds = self._load_credentials()
|
||||
service = build("drive", "v3", credentials=creds)
|
||||
|
||||
request = service.files().get_media(fileId=id)
|
||||
fh = BytesIO()
|
||||
downloader = MediaIoBaseDownload(fh, request)
|
||||
done = False
|
||||
while done is False:
|
||||
status, done = downloader.next_chunk()
|
||||
content = fh.getvalue()
|
||||
|
||||
from PyPDF2 import PdfReader
|
||||
|
||||
pdf_reader = PdfReader(BytesIO(content))
|
||||
|
||||
return [
|
||||
Document(
|
||||
page_content=page.extract_text(),
|
||||
metadata={
|
||||
"source": f"https://drive.google.com/file/d/{id}/view",
|
||||
"page": i,
|
||||
},
|
||||
)
|
||||
for i, page in enumerate(pdf_reader.pages)
|
||||
]
|
||||
|
||||
def _load_file_from_ids(self) -> List[Document]:
|
||||
"""Load files from a list of IDs."""
|
||||
if not self.file_ids:
|
||||
raise ValueError("file_ids must be set")
|
||||
docs = []
|
||||
for file_id in self.file_ids:
|
||||
docs.extend(self._load_file_from_id(file_id))
|
||||
return docs
|
||||
|
||||
def load(self) -> List[Document]:
|
||||
"""Load documents."""
|
||||
if self.folder_id:
|
||||
return self._load_documents_from_folder()
|
||||
else:
|
||||
elif self.document_ids:
|
||||
return self._load_documents_from_ids()
|
||||
else:
|
||||
return self._load_file_from_ids()
|
||||
|
Loading…
Reference in New Issue
Block a user