feat: create google drive ingest component (#3129)
* feat: create google drive ingest component It adds the basic implementation of the Google Drive Loader. * feat: create google drive ingest component Created a basic Google Drive loader component * Updated the Icon * updated the formatting ran make format to update the formatting * Addressed comments
This commit is contained in:
parent
ca008a1346
commit
3e6c863a8b
1 changed files with 87 additions and 0 deletions
87
src/backend/base/langflow/components/data/GoogleDrive.py
Normal file
87
src/backend/base/langflow/components/data/GoogleDrive.py
Normal file
|
|
@ -0,0 +1,87 @@
|
||||||
|
import json
|
||||||
|
from typing import Optional
|
||||||
|
from google.oauth2.credentials import Credentials
|
||||||
|
from google.auth.exceptions import RefreshError
|
||||||
|
from langflow.custom import Component
|
||||||
|
from langflow.inputs import MessageTextInput
|
||||||
|
from langflow.io import SecretStrInput
|
||||||
|
from langflow.template import Output
|
||||||
|
from langflow.schema import Data
|
||||||
|
from langchain_google_community import GoogleDriveLoader
|
||||||
|
from langflow.helpers.data import docs_to_data
|
||||||
|
|
||||||
|
from json.decoder import JSONDecodeError
|
||||||
|
|
||||||
|
|
||||||
|
class GoogleDriveComponent(Component):
|
||||||
|
display_name = "Google Drive Loader"
|
||||||
|
description = "Loads documents from Google Drive using provided credentials."
|
||||||
|
icon = "Google"
|
||||||
|
|
||||||
|
inputs = [
|
||||||
|
SecretStrInput(
|
||||||
|
name="json_string",
|
||||||
|
display_name="JSON String of the Service Account Token",
|
||||||
|
info="JSON string containing OAuth 2.0 access token information for service account access",
|
||||||
|
required=True,
|
||||||
|
),
|
||||||
|
MessageTextInput(
|
||||||
|
name="document_id", display_name="Document ID", info="Single Google Drive document ID", required=True
|
||||||
|
),
|
||||||
|
]
|
||||||
|
|
||||||
|
outputs = [
|
||||||
|
Output(display_name="Loaded Documents", name="docs", method="load_documents"),
|
||||||
|
]
|
||||||
|
|
||||||
|
def load_documents(self) -> Data:
|
||||||
|
class CustomGoogleDriveLoader(GoogleDriveLoader):
|
||||||
|
creds: Optional[Credentials] = None
|
||||||
|
"""Credentials object to be passed directly."""
|
||||||
|
|
||||||
|
def _load_credentials(self):
|
||||||
|
"""Load credentials from the provided creds attribute or fallback to the original method."""
|
||||||
|
if self.creds:
|
||||||
|
return self.creds
|
||||||
|
else:
|
||||||
|
raise ValueError("No credentials provided.")
|
||||||
|
|
||||||
|
class Config:
|
||||||
|
arbitrary_types_allowed = True
|
||||||
|
|
||||||
|
json_string = self.json_string
|
||||||
|
|
||||||
|
document_ids = [self.document_id]
|
||||||
|
if len(document_ids) != 1:
|
||||||
|
raise ValueError("Expected a single document ID")
|
||||||
|
|
||||||
|
# TODO: Add validation to check if the document ID is valid
|
||||||
|
|
||||||
|
# Load the token information from the JSON string
|
||||||
|
try:
|
||||||
|
token_info = json.loads(json_string)
|
||||||
|
except JSONDecodeError as e:
|
||||||
|
raise ValueError("Invalid JSON string") from e
|
||||||
|
|
||||||
|
# Initialize the custom loader with the provided credentials and document IDs
|
||||||
|
loader = CustomGoogleDriveLoader(
|
||||||
|
creds=Credentials.from_authorized_user_info(token_info), document_ids=document_ids
|
||||||
|
)
|
||||||
|
|
||||||
|
# Load the documents
|
||||||
|
try:
|
||||||
|
docs = loader.load()
|
||||||
|
# catch google.auth.exceptions.RefreshError
|
||||||
|
except RefreshError as e:
|
||||||
|
raise ValueError(
|
||||||
|
"Authentication error: Unable to refresh authentication token. Please try to reauthenticate."
|
||||||
|
) from e
|
||||||
|
except Exception as e:
|
||||||
|
raise ValueError(f"Error loading documents: {e}") from e
|
||||||
|
|
||||||
|
assert len(docs) == 1, "Expected a single document to be loaded."
|
||||||
|
|
||||||
|
data = docs_to_data(docs)
|
||||||
|
# Return the loaded documents
|
||||||
|
self.status = data
|
||||||
|
return Data(data={"text": data})
|
||||||
Loading…
Add table
Add a link
Reference in a new issue