From 6e69da36f7fa4aa5a8b70cfe20900157cc60042b Mon Sep 17 00:00:00 2001 From: Farrizal Alchudry Mutaqien Date: Sun, 27 Jul 2025 23:01:13 +0700 Subject: [PATCH] feat: google drive folder sync --- backend/open_webui/config.py | 13 + backend/open_webui/main.py | 33 ++ backend/open_webui/routers/knowledge.py | 305 ++++++++++ backend/open_webui/routers/retrieval.py | 37 ++ backend/open_webui/services/__init__.py | 1 + backend/open_webui/services/google_drive.py | 523 ++++++++++++++++++ .../services/google_drive_scheduler.py | 238 ++++++++ src/lib/apis/knowledge/index.ts | 76 +++ .../admin/Settings/Documents.svelte | 138 ++++- .../workspace/Knowledge/KnowledgeBase.svelte | 38 +- .../KnowledgeBase/AddContentMenu.svelte | 23 + .../KnowledgeBase/GoogleDriveSyncModal.svelte | 255 +++++++++ src/lib/i18n/locales/en-US/translation.json | 44 +- src/lib/stores/index.ts | 2 + src/lib/types/google-drive.ts | 291 ++++++++++ src/lib/types/index.ts | 3 + src/lib/utils/google-drive-picker.ts | 88 +-- uv.lock | 466 ++++++++-------- 18 files changed, 2297 insertions(+), 277 deletions(-) create mode 100644 backend/open_webui/services/__init__.py create mode 100644 backend/open_webui/services/google_drive.py create mode 100644 backend/open_webui/services/google_drive_scheduler.py create mode 100644 src/lib/components/workspace/Knowledge/KnowledgeBase/GoogleDriveSyncModal.svelte create mode 100644 src/lib/types/google-drive.ts diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index 58810c9e3e..20bbb701fe 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -2503,6 +2503,19 @@ GOOGLE_DRIVE_API_KEY = PersistentConfig( os.environ.get('GOOGLE_DRIVE_API_KEY', ''), ) +GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON = PersistentConfig( + "GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON", + "google_drive.service_account_json", + os.environ.get("GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON", ""), +) + +# If configured, Google Drive folder sync will be available for knowledge bases. +ENABLE_GOOGLE_DRIVE_FOLDER_SYNC = PersistentConfig( + "ENABLE_GOOGLE_DRIVE_FOLDER_SYNC", + "google_drive.enable_folder_sync", + os.getenv("ENABLE_GOOGLE_DRIVE_FOLDER_SYNC", "False").lower() == "true", +) + ENABLE_ONEDRIVE_INTEGRATION = PersistentConfig( 'ENABLE_ONEDRIVE_INTEGRATION', 'onedrive.enable', diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 0081788068..642d6c3568 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -344,6 +344,8 @@ from open_webui.config import ( GOOGLE_PSE_ENGINE_ID, GOOGLE_DRIVE_CLIENT_ID, GOOGLE_DRIVE_API_KEY, + GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON, + ENABLE_GOOGLE_DRIVE_FOLDER_SYNC, ENABLE_ONEDRIVE_INTEGRATION, ONEDRIVE_CLIENT_ID_PERSONAL, ONEDRIVE_CLIENT_ID_BUSINESS, @@ -472,6 +474,7 @@ from open_webui.env import ( AUDIT_INCLUDED_PATHS, AUDIT_LOG_LEVEL, CHANGELOG, + ENV, REDIS_URL, REDIS_CLUSTER, REDIS_KEY_PREFIX, @@ -561,6 +564,8 @@ from open_webui.tasks import ( list_tasks, ) # Import from tasks.py +from open_webui.services.google_drive_scheduler import google_drive_scheduler + from open_webui.utils.redis import get_sentinels_from_env @@ -702,11 +707,31 @@ async def lifespan(app: FastAPI): # Mark application as ready to accept traffic from a startup perspective. app.state.startup_complete = True + # Start Google Drive sync scheduler if enabled + if app.state.config.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC: + try: + await google_drive_scheduler.start() + log.info("Google Drive sync scheduler started successfully") + except Exception as e: + log.error(f"Failed to start Google Drive sync scheduler: {e}") + # Non-critical failure - continue application startup + yield if hasattr(app.state, 'redis_task_command_listener'): app.state.redis_task_command_listener.cancel() + # Stop Google Drive sync scheduler + if ( + hasattr(app.state, "config") + and app.state.config.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC + ): + try: + await google_drive_scheduler.stop() + log.info("Google Drive sync scheduler stopped") + except Exception as e: + log.error(f"Error stopping Google Drive sync scheduler: {e}") + app = FastAPI( title='Open WebUI', @@ -1046,6 +1071,10 @@ app.state.config.BYPASS_WEB_SEARCH_EMBEDDING_AND_RETRIEVAL = BYPASS_WEB_SEARCH_E app.state.config.BYPASS_WEB_SEARCH_WEB_LOADER = BYPASS_WEB_SEARCH_WEB_LOADER app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION = ENABLE_GOOGLE_DRIVE_INTEGRATION +app.state.config.GOOGLE_DRIVE_API_KEY = GOOGLE_DRIVE_API_KEY +app.state.config.GOOGLE_DRIVE_CLIENT_ID = GOOGLE_DRIVE_CLIENT_ID +app.state.config.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON = GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON +app.state.config.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC = ENABLE_GOOGLE_DRIVE_FOLDER_SYNC app.state.config.ENABLE_ONEDRIVE_INTEGRATION = ENABLE_ONEDRIVE_INTEGRATION app.state.config.OLLAMA_CLOUD_WEB_SEARCH_API_KEY = OLLAMA_CLOUD_WEB_SEARCH_API_KEY @@ -2076,6 +2105,7 @@ async def get_app_config(request: Request): 'enable_admin_chat_access': ENABLE_ADMIN_CHAT_ACCESS, 'enable_admin_analytics': ENABLE_ADMIN_ANALYTICS, 'enable_google_drive_integration': app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION, + 'enable_google_drive_folder_sync': app.state.config.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC, 'enable_onedrive_integration': app.state.config.ENABLE_ONEDRIVE_INTEGRATION, 'enable_memories': app.state.config.ENABLE_MEMORIES, **( @@ -2124,6 +2154,9 @@ async def get_app_config(request: Request): 'client_id': GOOGLE_DRIVE_CLIENT_ID.value, 'api_key': GOOGLE_DRIVE_API_KEY.value, }, + 'google_drive_folder_sync': { + 'service_account_json': GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON.value, + }, 'onedrive': { 'client_id_personal': ONEDRIVE_CLIENT_ID_PERSONAL, 'client_id_business': ONEDRIVE_CLIENT_ID_BUSINESS, diff --git a/backend/open_webui/routers/knowledge.py b/backend/open_webui/routers/knowledge.py index ead782cdbf..a19490b5d0 100644 --- a/backend/open_webui/routers/knowledge.py +++ b/backend/open_webui/routers/knowledge.py @@ -4,6 +4,7 @@ from fastapi import APIRouter, Depends, HTTPException, status, Request, Query from fastapi.responses import StreamingResponse from fastapi.concurrency import run_in_threadpool import logging +import time import io import zipfile from urllib.parse import quote @@ -19,6 +20,7 @@ from open_webui.models.knowledge import ( KnowledgeUserResponse, ) from open_webui.models.files import Files, FileModel, FileMetadataResponse +from open_webui.models.users import UserModel from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT from open_webui.routers.retrieval import ( process_file, @@ -26,6 +28,7 @@ from open_webui.routers.retrieval import ( process_files_batch, BatchProcessFilesForm, ) +from open_webui.services.google_drive import google_drive_service from open_webui.storage.provider import Storage from open_webui.constants import ERROR_MESSAGES @@ -1103,3 +1106,305 @@ async def export_knowledge_by_id(id: str, user=Depends(get_admin_user), db: Sess media_type='application/zip', headers={'Content-Disposition': content_disposition}, ) + + +############################ +# Google Drive Sync +############################ + + +class GoogleDriveSyncForm(BaseModel): + folder_id: str + include_nested: bool = True + sync_interval_days: float = 1.0 # Float to support fractional days + + +class GoogleDriveServiceAccountResponse(BaseModel): + email: str + + +@router.get( + "/google-drive/service-account-email", + response_model=GoogleDriveServiceAccountResponse, +) +async def get_google_drive_service_account_email( + user: UserModel = Depends(get_verified_user), +) -> GoogleDriveServiceAccountResponse: + """Get the Google Drive service account email for sharing folders.""" + if not google_drive_service.is_configured(): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Google Drive service account not configured", + ) + + email = google_drive_service.get_service_account_email() + if not email: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Failed to get service account email", + ) + + return {"email": email} + + +@router.post("/{id}/google-drive/sync", response_model=Optional[KnowledgeFilesResponse]) +async def sync_google_drive_folder( + id: str, + form_data: GoogleDriveSyncForm, + request: Request, + user: UserModel = Depends(get_verified_user), +) -> Optional[KnowledgeFilesResponse]: + """Sync a Google Drive folder with the knowledge base.""" + knowledge = Knowledges.get_knowledge_by_id(id=id) + if not knowledge: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=ERROR_MESSAGES.NOT_FOUND, + ) + + if ( + knowledge.user_id != user.id + and not has_access(user.id, "write", knowledge.access_control) + and user.role != "admin" + ): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=ERROR_MESSAGES.ACCESS_PROHIBITED, + ) + + if not google_drive_service.is_configured(): + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Google Drive service account not configured", + ) + + try: + # Get current files in knowledge base + current_file_ids = knowledge.data.get("file_ids", []) if knowledge.data else [] + current_files = Files.get_files_by_ids(current_file_ids) + + # Create a map of Google Drive file IDs to local file IDs + gdrive_file_map = {} + for file in current_files: + if file.meta and file.meta.get("google_drive_id"): + gdrive_file_map[file.meta["google_drive_id"]] = file.id + + # Get files from Google Drive folder + gdrive_files = google_drive_service.list_folder_files( + form_data.folder_id, form_data.include_nested + ) + + log.info( + f"Google Drive sync: Found {len(gdrive_files)} files in folder {form_data.folder_id}" + ) + for gdrive_file in gdrive_files: + log.info( + f"Google Drive file: {gdrive_file.get('name', 'Unknown')} (ID: {gdrive_file.get('id', 'Unknown')}, Type: {gdrive_file.get('mimeType', 'Unknown')})" + ) + + # Track files to keep and files to add + files_to_keep = set() + files_to_add = [] + + for gdrive_file in gdrive_files: + gdrive_id = gdrive_file["id"] + + if gdrive_id in gdrive_file_map: + # File exists, check if it needs updating + local_file_id = gdrive_file_map[gdrive_id] + local_file = Files.get_file_by_id(local_file_id) + + if local_file and local_file.meta: + local_modified = local_file.meta.get("google_drive_modified") + gdrive_modified = gdrive_file["modifiedTime"] + + if local_modified != gdrive_modified: + # File was modified, re-download and update + files_to_add.append(gdrive_file) + # Remove old file + Files.delete_file_by_id(local_file_id) + else: + # File is up to date, keep it + files_to_keep.add(local_file_id) + else: + # Local file metadata is missing, re-download + files_to_add.append(gdrive_file) + else: + # New file, add it + files_to_add.append(gdrive_file) + + # Remove files that are no longer in Google Drive + files_to_remove = set(current_file_ids) - files_to_keep + for file_id in files_to_remove: + Files.delete_file_by_id(file_id) + + # Download and add new/updated files + new_file_ids = list(files_to_keep) + + log.info( + f"Google Drive sync: Processing {len(files_to_add)} files to add/update" + ) + + for gdrive_file in files_to_add: + try: + log.info( + f"Google Drive sync: Processing file '{gdrive_file.get('name', 'Unknown')}' (ID: {gdrive_file.get('id', 'Unknown')})" + ) + + # Download file from Google Drive + file_content, filename = google_drive_service.download_file( + gdrive_file["id"], gdrive_file + ) + + log.info( + f"Google Drive sync: Downloaded file '{filename}', size: {len(file_content)} bytes" + ) + + # Debug: Check if content is actually text + try: + content_preview = ( + file_content.decode("utf-8")[:200] + if isinstance(file_content, bytes) + else str(file_content)[:200] + ) + log.info( + f"Google Drive sync: File content preview: {content_preview}..." + ) + except Exception as e: + log.error(f"Google Drive sync: Could not decode file content: {e}") + + # Create file object + file_obj = io.BytesIO(file_content) + file_obj.name = filename + + # Upload to storage + import uuid + + file_id = str(uuid.uuid4()) + tags = { + "OpenWebUI-User-Email": user.email, + "OpenWebUI-User-Id": user.id, + "OpenWebUI-User-Name": user.name, + "OpenWebUI-File-Id": file_id, + } + + contents, file_path = Storage.upload_file( + file_obj, f"{file_id}_{filename}", tags + ) + log.info( + f"Google Drive sync: Uploaded to storage, contents size: {len(contents)} bytes" + ) + + # Create file record with content in data field + from open_webui.models.files import FileForm + + # Convert bytes to string for text content + text_content = ( + file_content.decode("utf-8") + if isinstance(file_content, bytes) + else str(file_content) + ) + + file_item = Files.insert_new_file( + user.id, + FileForm( + id=file_id, + filename=filename, + path=file_path, + data={ + "content": text_content, # Store the actual text content + }, + meta={ + "name": filename, + "content_type": "text/plain", # Changed from application/octet-stream + "size": len(contents), + "google_drive_id": gdrive_file["id"], + "google_drive_modified": gdrive_file["modifiedTime"], + "google_drive_path": gdrive_file["path"], + "collection_name": id, + }, + ), + ) + + if file_item: + new_file_ids.append(file_id) + log.info( + f"Google Drive sync: Created file record for '{filename}' (File ID: {file_id})" + ) + + # Process file for vector storage + try: + from open_webui.routers.retrieval import process_file + + process_file( + request, + ProcessFileForm( + file_id=file_id, + collection_name=id, + ), + user, + ) + log.info( + f"Google Drive sync: Successfully processed file '{filename}' for vector storage" + ) + except Exception as e: + log.error( + f"Failed to process file {file_id} ('{filename}') for vector storage: {e}" + ) + else: + log.error( + f"Google Drive sync: Failed to create file record for '{filename}'" + ) + + except Exception as e: + log.error(f"Failed to sync file {gdrive_file['name']}: {e}") + continue + + # Update knowledge base data + sync_data = { + "google_drive_folder_id": form_data.folder_id, + "google_drive_include_nested": form_data.include_nested, + "google_drive_sync_interval_days": form_data.sync_interval_days, + "google_drive_last_sync": int(time.time()), + } + + updated_data = knowledge.data.copy() if knowledge.data else {} + updated_data.update(sync_data) + updated_data["file_ids"] = new_file_ids + + knowledge = Knowledges.update_knowledge_data_by_id(id=id, data=updated_data) + + log.info( + f"Google Drive sync completed: {len(new_file_ids)} total files in knowledge base" + ) + log.info(f"Google Drive sync: Knowledge object type: {type(knowledge)}") + log.info(f"Google Drive sync: Knowledge object: {knowledge}") + + if knowledge: + files = Files.get_file_metadatas_by_ids(new_file_ids) + log.info(f"Google Drive sync: Returning {len(files)} file metadata records") + + # Handle both dict and model cases + if hasattr(knowledge, "model_dump"): + knowledge_dict = knowledge.model_dump() + else: + knowledge_dict = ( + knowledge if isinstance(knowledge, dict) else knowledge.__dict__ + ) + + return KnowledgeFilesResponse( + **knowledge_dict, + files=files, + ) + else: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail="Failed to update knowledge base", + ) + + except Exception as e: + log.error(f"Google Drive sync error: {e}") + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail=f"Failed to sync Google Drive folder: {str(e)}", + ) diff --git a/backend/open_webui/routers/retrieval.py b/backend/open_webui/routers/retrieval.py index 6c9e988dd6..054c02cae9 100644 --- a/backend/open_webui/routers/retrieval.py +++ b/backend/open_webui/routers/retrieval.py @@ -97,6 +97,7 @@ from open_webui.utils.misc import ( ) from open_webui.utils.auth import get_admin_user, get_verified_user from open_webui.utils.access_control import has_permission +from open_webui.services.google_drive import google_drive_service from open_webui.config import ( ENV, @@ -503,6 +504,10 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)): 'ALLOWED_FILE_EXTENSIONS': request.app.state.config.ALLOWED_FILE_EXTENSIONS, # Integration settings 'ENABLE_GOOGLE_DRIVE_INTEGRATION': request.app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION, + 'GOOGLE_DRIVE_API_KEY': request.app.state.config.GOOGLE_DRIVE_API_KEY, + 'GOOGLE_DRIVE_CLIENT_ID': request.app.state.config.GOOGLE_DRIVE_CLIENT_ID, + 'GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON': request.app.state.config.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON, + 'ENABLE_GOOGLE_DRIVE_FOLDER_SYNC': request.app.state.config.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC, 'ENABLE_ONEDRIVE_INTEGRATION': request.app.state.config.ENABLE_ONEDRIVE_INTEGRATION, # Web search settings 'web': { @@ -713,6 +718,10 @@ class ConfigForm(BaseModel): # Integration settings ENABLE_GOOGLE_DRIVE_INTEGRATION: Optional[bool] = None + GOOGLE_DRIVE_API_KEY: Optional[str] = None + GOOGLE_DRIVE_CLIENT_ID: Optional[str] = None + GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON: Optional[str] = None + ENABLE_GOOGLE_DRIVE_FOLDER_SYNC: Optional[bool] = None ENABLE_ONEDRIVE_INTEGRATION: Optional[bool] = None # Web search settings @@ -1026,12 +1035,36 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend if form_data.ENABLE_GOOGLE_DRIVE_INTEGRATION is not None else request.app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION ) + request.app.state.config.GOOGLE_DRIVE_API_KEY = ( + form_data.GOOGLE_DRIVE_API_KEY + if form_data.GOOGLE_DRIVE_API_KEY is not None + else request.app.state.config.GOOGLE_DRIVE_API_KEY + ) + request.app.state.config.GOOGLE_DRIVE_CLIENT_ID = ( + form_data.GOOGLE_DRIVE_CLIENT_ID + if form_data.GOOGLE_DRIVE_CLIENT_ID is not None + else request.app.state.config.GOOGLE_DRIVE_CLIENT_ID + ) + request.app.state.config.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON = ( + form_data.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON + if form_data.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON is not None + else request.app.state.config.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON + ) + request.app.state.config.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC = ( + form_data.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC + if form_data.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC is not None + else request.app.state.config.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC + ) request.app.state.config.ENABLE_ONEDRIVE_INTEGRATION = ( form_data.ENABLE_ONEDRIVE_INTEGRATION if form_data.ENABLE_ONEDRIVE_INTEGRATION is not None else request.app.state.config.ENABLE_ONEDRIVE_INTEGRATION ) + # Refresh Google Drive service if service account JSON was updated + if form_data.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON is not None: + google_drive_service.refresh_configuration() + if form_data.web is not None: # Web search settings request.app.state.config.ENABLE_WEB_SEARCH = form_data.web.ENABLE_WEB_SEARCH @@ -1166,6 +1199,10 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend 'ALLOWED_FILE_EXTENSIONS': request.app.state.config.ALLOWED_FILE_EXTENSIONS, # Integration settings 'ENABLE_GOOGLE_DRIVE_INTEGRATION': request.app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION, + 'GOOGLE_DRIVE_API_KEY': request.app.state.config.GOOGLE_DRIVE_API_KEY, + 'GOOGLE_DRIVE_CLIENT_ID': request.app.state.config.GOOGLE_DRIVE_CLIENT_ID, + 'GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON': request.app.state.config.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON, + 'ENABLE_GOOGLE_DRIVE_FOLDER_SYNC': request.app.state.config.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC, 'ENABLE_ONEDRIVE_INTEGRATION': request.app.state.config.ENABLE_ONEDRIVE_INTEGRATION, # Web search settings 'web': { diff --git a/backend/open_webui/services/__init__.py b/backend/open_webui/services/__init__.py new file mode 100644 index 0000000000..a15012dfa8 --- /dev/null +++ b/backend/open_webui/services/__init__.py @@ -0,0 +1 @@ +# Services module for Open WebUI diff --git a/backend/open_webui/services/google_drive.py b/backend/open_webui/services/google_drive.py new file mode 100644 index 0000000000..8e03ecf1e6 --- /dev/null +++ b/backend/open_webui/services/google_drive.py @@ -0,0 +1,523 @@ +import json +import logging +from typing import List, Dict, Optional, Tuple, TypedDict, NotRequired +from datetime import datetime +import io + +from google.oauth2 import service_account # type: ignore +from googleapiclient.discovery import build, Resource # type: ignore +from googleapiclient.errors import HttpError # type: ignore +from googleapiclient.http import MediaIoBaseDownload # type: ignore + +from open_webui.config import GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON + + +# Type definitions for Google Drive API responses +class GoogleDriveFile(TypedDict): + """Type definition for Google Drive file information.""" + + id: str + name: str + mimeType: str + modifiedTime: str + size: NotRequired[str] # Optional field, not present for some files + parents: NotRequired[List[str]] # Optional, may not be present + path: str # Custom field added by our service + + +class GoogleDriveFolder(TypedDict): + """Type definition for Google Drive folder information.""" + + id: str + name: str + + +class GoogleDriveAPIResponse(TypedDict): + """Type definition for Google Drive API list response.""" + + files: List[Dict[str, str]] # Raw API response uses generic dict + nextPageToken: NotRequired[str] + + +class FileDownloadResult(TypedDict): + """Type definition for file download result.""" + + content: bytes + filename: str + + +class GoogleDriveSyncConfig(TypedDict): + """Type definition for Google Drive sync configuration stored in knowledge base data.""" + + google_drive_folder_id: str + google_drive_include_nested: bool + google_drive_sync_interval_days: float # Float to support fractional days + google_drive_last_sync: int # Unix timestamp + file_ids: List[str] # List of file IDs in the knowledge base + + +class GoogleDriveFileMetadata(TypedDict): + """Type definition for Google Drive file metadata stored in file records.""" + + name: str + content_type: str + size: int + google_drive_id: str + google_drive_modified: str # ISO timestamp from Google Drive API + google_drive_path: str + collection_name: str + # Additional standard file metadata fields that may be present + source: NotRequired[str] # Usually "google_drive" + original_filename: NotRequired[str] # May differ from processed filename + + +log = logging.getLogger(__name__) + + +class GoogleDriveService: + """Service for interacting with Google Drive using service account authentication.""" + + def __init__(self) -> None: + self.service: Optional[Resource] = None + self._initialize_service() + + def _initialize_service(self) -> None: + """Initialize Google Drive service with service account credentials.""" + try: + if not GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON.value: + log.warning("Google Drive service account JSON not configured") + return + + # Parse service account JSON + service_account_info = json.loads(GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON.value) + + # Create credentials + credentials = service_account.Credentials.from_service_account_info( + service_account_info, + scopes=["https://www.googleapis.com/auth/drive.readonly"], + ) + + # Build service + self.service = build("drive", "v3", credentials=credentials) + log.info("Google Drive service initialized successfully") + + except json.JSONDecodeError as e: + log.error(f"Invalid service account JSON: {e}") + except Exception as e: + log.error(f"Failed to initialize Google Drive service: {e}") + + def is_configured(self) -> bool: + """Check if Google Drive service is properly configured.""" + return self.service is not None + + def refresh_configuration(self) -> None: + """Refresh the service configuration. Call this when config is updated.""" + self._initialize_service() + + def get_service_account_email(self) -> Optional[str]: + """Get the service account email address.""" + try: + if not GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON.value: + return None + + service_account_info = json.loads(GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON.value) + return service_account_info.get("client_email") + except Exception as e: + log.error(f"Failed to get service account email: {e}") + return None + + def list_folder_files( + self, folder_id: str, include_nested: bool = True + ) -> List[GoogleDriveFile]: + """ + List all files in a Google Drive folder. + + Args: + folder_id: Google Drive folder ID + include_nested: Whether to include files from nested folders + + Returns: + List of file information dictionaries + """ + if not self.service: + raise Exception("Google Drive service not configured") + + try: + files = [] + + if include_nested: + # Get all files recursively + files = self._get_files_recursive(folder_id) + else: + # Get files only from the specified folder + files = self._get_files_in_folder(folder_id) + + # Filter supported file types + log.info(f"Google Drive: Found {len(files)} total files before filtering") + for file_info in files: + log.info( + f"Google Drive file found: '{file_info.get('name', 'Unknown')}' (Type: {file_info.get('mimeType', 'Unknown')})" + ) + + supported_files = [] + for file_info in files: + if self._is_supported_file(file_info): + supported_files.append(file_info) + log.info( + f"Google Drive: File '{file_info.get('name', 'Unknown')}' is SUPPORTED" + ) + else: + log.info( + f"Google Drive: File '{file_info.get('name', 'Unknown')}' is NOT SUPPORTED (Type: {file_info.get('mimeType', 'Unknown')})" + ) + + log.info(f"Google Drive: {len(supported_files)} files passed filtering") + return supported_files + + except HttpError as e: + log.error(f"HTTP error listing folder files: {e}") + raise Exception(f"Failed to access Google Drive folder: {e}") + except Exception as e: + log.error(f"Error listing folder files: {e}") + raise + + def _get_files_recursive( + self, folder_id: str, path: str = "" + ) -> List[GoogleDriveFile]: + """Recursively get all files from a folder and its subfolders.""" + files = [] + + # Get files in current folder + files.extend(self._get_files_in_folder(folder_id, path)) + + # Get subfolders and process them recursively + subfolders = self._get_subfolders(folder_id) + for subfolder in subfolders: + subfolder_path = ( + f"{path}/{subfolder['name']}" if path else subfolder["name"] + ) + files.extend(self._get_files_recursive(subfolder["id"], subfolder_path)) + + return files + + def _get_files_in_folder( + self, folder_id: str, path: str = "" + ) -> List[GoogleDriveFile]: + """Get files directly in a specific folder.""" + files = [] + page_token = None + + log.info(f"Google Drive: Querying folder {folder_id} for files") + + # Debug: Test if we can access any files at all + try: + assert self.service is not None # Type guard for mypy + test_response = ( + self.service.files() + .list( + q="", spaces="drive", pageSize=5, fields="files(id, name, mimeType)" + ) + .execute() + ) + log.info( + f"Google Drive: Test query found {len(test_response.get('files', []))} total accessible files" + ) + except Exception as e: + log.error(f"Google Drive: Test query failed: {e}") + + while True: + # Use official Google Drive API query syntax + query = f"'{folder_id}' in parents" + log.info(f"Google Drive API query: {query}") + + try: + # Follow official Google Drive API example pattern with Shared Drive support + assert self.service is not None # Type guard for mypy + response = ( + self.service.files() + .list( + q=query, + spaces="drive", + corpora="allDrives", + includeItemsFromAllDrives=True, + supportsAllDrives=True, + pageSize=100, + pageToken=page_token, + fields="nextPageToken, files(id, name, mimeType, modifiedTime, size, parents)", + ) + .execute() + ) + + log.info( + f"Google Drive API response: {len(response.get('files', []))} items returned" + ) + + # Log all items found (including folders) + all_items = response.get("files", []) + for item in all_items: + log.info( + f"Google Drive item: '{item.get('name', 'Unknown')}' (Type: {item.get('mimeType', 'Unknown')})" + ) + + results = response # Keep compatibility with existing code + + except Exception as e: + log.error(f"Google Drive API error: {e}") + return [] + + items = results.get("files", []) + + # Filter out folders and only keep files + for item in items: + if item.get("mimeType") != "application/vnd.google-apps.folder": + file_info: GoogleDriveFile = { + "id": item["id"], + "name": item["name"], + "mimeType": item["mimeType"], + "modifiedTime": item["modifiedTime"], + "path": f"{path}/{item['name']}" if path else item["name"], + } + # Add optional fields if present + if "size" in item: + file_info["size"] = item["size"] + if "parents" in item: + file_info["parents"] = item["parents"] + files.append(file_info) + log.info(f"Google Drive: Added file '{item['name']}' to results") + + page_token = results.get("nextPageToken") + if not page_token: + break + + return files + + def _get_subfolders(self, folder_id: str) -> List[Dict[str, str]]: + """Get subfolders in a specific folder.""" + subfolders = [] + page_token = None + + while True: + query = f"'{folder_id}' in parents and mimeType='application/vnd.google-apps.folder' and trashed=false" + + assert self.service is not None # Type guard for mypy + results = ( + self.service.files() + .list( + q=query, + spaces="drive", + corpora="allDrives", + includeItemsFromAllDrives=True, + supportsAllDrives=True, + pageSize=100, + pageToken=page_token, + fields="nextPageToken, files(id, name)", + ) + .execute() + ) + + items = results.get("files", []) + subfolders.extend(items) + + page_token = results.get("nextPageToken") + if not page_token: + break + + return subfolders + + def _is_supported_file(self, file_info: GoogleDriveFile) -> bool: + """Check if a file type is supported by the knowledge base.""" + mime_type = file_info.get("mimeType", "") + file_name = file_info.get("name", "") + + # Supported document file extensions for knowledge base + # Focus on actual document types rather than code files + supported_extensions = [ + # Text and markdown documents + "txt", + "md", + "markdown", + "rst", + + # Microsoft Office documents + "doc", + "docx", + "xls", + "xlsx", + "ppt", + "pptx", + "msg", + + # OpenDocument formats + "odt", + "ods", + "odp", + + # PDF documents + "pdf", + + # Data files + "csv", + "tsv", + "json", + "xml", + + # Web documents + "html", + "htm", + + # Rich text + "rtf", + + # eBooks + "epub", + "mobi", + + # Other document formats + "pages", # Apple Pages + "numbers", # Apple Numbers + "key", # Apple Keynote + ] + + # Check file extension + if "." in file_name: + extension = file_name.split(".")[-1].lower() + if extension in supported_extensions: + return True + + # Check Google Workspace files + google_workspace_types = [ + "application/vnd.google-apps.document", + "application/vnd.google-apps.spreadsheet", + "application/vnd.google-apps.presentation", + ] + + if mime_type in google_workspace_types: + return True + + # Check other supported MIME types for documents + supported_mime_types = [ + # Text documents + "text/plain", + "text/markdown", + "text/x-markdown", + "text/csv", + "text/tab-separated-values", + "text/xml", + "text/html", + "text/rtf", + + # Microsoft Office documents + "application/msword", # .doc + "application/vnd.openxmlformats-officedocument.wordprocessingml.document", # .docx + "application/vnd.ms-excel", # .xls + "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", # .xlsx + "application/vnd.ms-powerpoint", # .ppt + "application/vnd.openxmlformats-officedocument.presentationml.presentation", # .pptx + "application/vnd.ms-outlook", # .msg + + # OpenDocument formats + "application/vnd.oasis.opendocument.text", # .odt + "application/vnd.oasis.opendocument.spreadsheet", # .ods + "application/vnd.oasis.opendocument.presentation", # .odp + + # PDF + "application/pdf", + + # eBooks + "application/epub+zip", + "application/x-mobipocket-ebook", + + # Data formats + "application/json", + + # Apple iWork formats (when exported) + "application/x-iwork-pages-sffpages", # .pages + "application/x-iwork-numbers-sffnumbers", # .numbers + "application/x-iwork-keynote-sffkey", # .key + ] + + return mime_type in supported_mime_types + + def download_file( + self, file_id: str, file_info: GoogleDriveFile + ) -> Tuple[bytes, str]: + """ + Download a file from Google Drive. + + Args: + file_id: Google Drive file ID + file_info: File information dictionary + + Returns: + Tuple of (file_content_bytes, filename) + """ + if not self.service: + raise Exception("Google Drive service not configured") + + try: + mime_type = file_info.get("mimeType", "") + file_name = file_info.get("name", "") + + # Handle Google Workspace files (export) + if mime_type.startswith("application/vnd.google-apps"): + return self._export_google_workspace_file(file_id, mime_type, file_name) + else: + # Handle regular files (download) + return self._download_regular_file(file_id, file_name) + + except HttpError as e: + log.error(f"HTTP error downloading file {file_id}: {e}") + raise Exception(f"Failed to download file: {e}") + except Exception as e: + log.error(f"Error downloading file {file_id}: {e}") + raise + + def _export_google_workspace_file( + self, file_id: str, mime_type: str, file_name: str + ) -> Tuple[bytes, str]: + """Export Google Workspace files to supported formats.""" + export_mime_type = None + export_extension = None + + if "document" in mime_type: + export_mime_type = "text/plain" + export_extension = ".txt" + elif "spreadsheet" in mime_type: + export_mime_type = "text/csv" + export_extension = ".csv" + elif "presentation" in mime_type: + export_mime_type = "text/plain" + export_extension = ".txt" + else: + export_mime_type = "application/pdf" + export_extension = ".pdf" + + assert self.service is not None # Type guard for mypy + request = self.service.files().export_media( + fileId=file_id, mimeType=export_mime_type + ) + file_content = request.execute() + + # Update filename with appropriate extension + if not file_name.endswith(export_extension): + file_name = f"{file_name}{export_extension}" + + return file_content, file_name + + def _download_regular_file(self, file_id: str, file_name: str) -> Tuple[bytes, str]: + """Download regular files.""" + assert self.service is not None # Type guard for mypy + request = self.service.files().get_media(fileId=file_id) + file_io = io.BytesIO() + downloader = MediaIoBaseDownload(file_io, request) + + done = False + while done is False: + _, done = downloader.next_chunk() + + file_content = file_io.getvalue() + return file_content, file_name + + +# Global instance +google_drive_service = GoogleDriveService() diff --git a/backend/open_webui/services/google_drive_scheduler.py b/backend/open_webui/services/google_drive_scheduler.py new file mode 100644 index 0000000000..573d82ce72 --- /dev/null +++ b/backend/open_webui/services/google_drive_scheduler.py @@ -0,0 +1,238 @@ +import asyncio +import logging +import time +from typing import Dict, List, Any, Optional +from datetime import datetime, timedelta +import uuid +import io + +from open_webui.env import ENV +from open_webui.models.knowledge import Knowledges, KnowledgeModel +from open_webui.services.google_drive import google_drive_service, GoogleDriveFile +from open_webui.models.files import Files, FileModel, FileForm +from open_webui.routers.retrieval import process_file, ProcessFileForm +from open_webui.storage.provider import Storage + +log = logging.getLogger(__name__) + + +class GoogleDriveSyncScheduler: + """Background scheduler for automatic Google Drive folder sync.""" + + def __init__(self) -> None: + self.running: bool = False + self.sync_tasks: Dict[str, Any] = {} + # In dev environment, check every minute for faster testing + # In production, check every hour + self.check_interval: int = 60 if ENV == "dev" else 3600 + + async def start(self) -> None: + """Start the background sync scheduler.""" + if self.running: + return + + self.running = True + log.info(f"Google Drive sync scheduler started (ENV: {ENV}, check_interval: {self.check_interval}s)") + + # Start the background task + asyncio.create_task(self._sync_loop()) + + async def stop(self) -> None: + """Stop the background sync scheduler.""" + self.running = False + log.info("Google Drive sync scheduler stopped") + + async def _sync_loop(self) -> None: + """Main sync loop that runs in the background.""" + while self.running: + try: + await self._check_and_sync_knowledge_bases() + await asyncio.sleep(self.check_interval) + except Exception as e: + log.error(f"Error in sync loop: {e}") + await asyncio.sleep(self.check_interval) + + async def _check_and_sync_knowledge_bases(self) -> None: + """Check all knowledge bases for Google Drive sync requirements.""" + if not google_drive_service.is_configured(): + return + + try: + # Get all knowledge bases + knowledge_bases = Knowledges.get_knowledge_bases() + + for kb in knowledge_bases: + if not kb.data: + continue + + # Check if this knowledge base has Google Drive sync configured + folder_id = kb.data.get("google_drive_folder_id") + sync_interval_days = kb.data.get("google_drive_sync_interval_days", 1) + last_sync = kb.data.get("google_drive_last_sync", 0) + + if not folder_id: + continue + + # Check if sync is due + current_time = int(time.time()) + sync_interval_seconds = sync_interval_days * 24 * 3600 + + if current_time - last_sync >= sync_interval_seconds: + log.info( + f"Auto-syncing knowledge base {kb.id} with Google Drive folder {folder_id}" + ) + await self._sync_knowledge_base(kb) + + except Exception as e: + log.error(f"Error checking knowledge bases for sync: {e}") + + async def _sync_knowledge_base(self, knowledge_base: KnowledgeModel) -> None: + """Sync a specific knowledge base with its Google Drive folder.""" + try: + if not knowledge_base.data: + return + + # Type guard for mypy + assert knowledge_base.data is not None + + folder_id = knowledge_base.data.get("google_drive_folder_id") + include_nested = knowledge_base.data.get( + "google_drive_include_nested", True + ) + + if not folder_id: + return + + # Get current files in knowledge base + current_file_ids = knowledge_base.data.get("file_ids", []) + current_files = Files.get_files_by_ids(current_file_ids) + + # Create a map of Google Drive file IDs to local file IDs + gdrive_file_map = {} + for file in current_files: + if file.meta and file.meta.get("google_drive_id"): + gdrive_file_map[file.meta["google_drive_id"]] = file.id + + # Get files from Google Drive folder + gdrive_files = google_drive_service.list_folder_files( + folder_id, include_nested + ) + + # Track files to keep and files to add + files_to_keep = set() + files_to_add = [] + + for gdrive_file in gdrive_files: + gdrive_id = gdrive_file["id"] + + if gdrive_id in gdrive_file_map: + # File exists, check if it needs updating + local_file_id = gdrive_file_map[gdrive_id] + local_file = Files.get_file_by_id(local_file_id) + + if local_file and local_file.meta: + local_modified = local_file.meta.get("google_drive_modified") + gdrive_modified = gdrive_file["modifiedTime"] + + if local_modified != gdrive_modified: + # File was modified, re-download and update + files_to_add.append(gdrive_file) + # Remove old file + Files.delete_file_by_id(local_file_id) + else: + # File is up to date, keep it + files_to_keep.add(local_file_id) + else: + # Local file metadata is missing, re-download + files_to_add.append(gdrive_file) + else: + # New file, add it + files_to_add.append(gdrive_file) + + # Remove files that are no longer in Google Drive + files_to_remove = set(current_file_ids) - files_to_keep + for file_id in files_to_remove: + Files.delete_file_by_id(file_id) + + # Download and add new/updated files + new_file_ids = list(files_to_keep) + + for gdrive_file in files_to_add: + try: + # Download file from Google Drive + file_content, filename = google_drive_service.download_file( + gdrive_file["id"], gdrive_file + ) + + # Create file object + file_obj = io.BytesIO(file_content) + file_obj.name = filename + + # Upload to storage + file_id = str(uuid.uuid4()) + tags = { + "OpenWebUI-User-Email": "system", + "OpenWebUI-User-Id": "system", + "OpenWebUI-User-Name": "Google Drive Sync", + "OpenWebUI-File-Id": file_id, + } + + contents, file_path = Storage.upload_file( + file_obj, f"{file_id}_{filename}", tags + ) + + # Create file record + file_form = FileForm( + id=file_id, + filename=filename, + path=file_path, + meta={ + "name": filename, + "content_type": "application/octet-stream", + "size": len(contents), + "google_drive_id": gdrive_file["id"], + "google_drive_modified": gdrive_file["modifiedTime"], + "google_drive_path": gdrive_file["path"], + "collection_name": knowledge_base.id, + }, + ) + file_item = Files.insert_new_file(knowledge_base.user_id, file_form) + + if file_item: + new_file_ids.append(file_id) + + # Process file for vector storage + try: + process_file( + ProcessFileForm( + file_id=file_id, + collection_name=knowledge_base.id, + ) + ) + except Exception as e: + log.error(f"Failed to process file {file_id}: {e}") + + except Exception as e: + log.error(f"Failed to sync file {gdrive_file['name']}: {e}") + continue + + # Update knowledge base data + assert knowledge_base.data is not None # Type guard for mypy + updated_data = knowledge_base.data.copy() + updated_data["file_ids"] = new_file_ids + updated_data["google_drive_last_sync"] = int(time.time()) + + Knowledges.update_knowledge_data_by_id( + id=knowledge_base.id, data=updated_data + ) + + log.info( + f"Successfully synced knowledge base {knowledge_base.id} with {len(files_to_add)} new/updated files" + ) + + except Exception as e: + log.error(f"Error syncing knowledge base {knowledge_base.id}: {e}") + + +# Global scheduler instance +google_drive_scheduler = GoogleDriveSyncScheduler() diff --git a/src/lib/apis/knowledge/index.ts b/src/lib/apis/knowledge/index.ts index f314bae634..3576b25232 100644 --- a/src/lib/apis/knowledge/index.ts +++ b/src/lib/apis/knowledge/index.ts @@ -1,4 +1,5 @@ import { WEBUI_API_BASE_URL } from '$lib/constants'; +import type { GoogleDriveServiceAccount, GoogleDriveSyncResponse } from '$lib/types/google-drive'; export const createNewKnowledge = async ( token: string, @@ -491,6 +492,81 @@ export const deleteKnowledgeById = async (token: string, id: string) => { return res; }; +export const getGoogleDriveServiceAccountEmail = async ( + token: string +): Promise => { + let error = null; + + const res = await fetch(`${WEBUI_API_BASE_URL}/knowledge/google-drive/service-account-email`, { + method: 'GET', + headers: { + Accept: 'application/json', + 'Content-Type': 'application/json', + authorization: `Bearer ${token}` + } + }) + .then(async (res) => { + if (!res.ok) throw await res.json(); + return res.json(); + }) + .then((json) => { + return json; + }) + .catch((err) => { + error = err.detail; + console.error(err); + return null; + }); + + if (error) { + throw error; + } + + return res; +}; + +export const syncGoogleDriveFolder = async ( + token: string, + knowledgeId: string, + folderId: string, + includeNested: boolean = true, + syncIntervalDays: number = 1 +): Promise => { + let error = null; + + const res = await fetch(`${WEBUI_API_BASE_URL}/knowledge/${knowledgeId}/google-drive/sync`, { + method: 'POST', + headers: { + Accept: 'application/json', + 'Content-Type': 'application/json', + authorization: `Bearer ${token}` + }, + body: JSON.stringify({ + folder_id: folderId, + include_nested: includeNested, + sync_interval_days: syncIntervalDays + }) + }) + .then(async (res) => { + if (!res.ok) throw await res.json(); + return res.json(); + }) + .then((json) => { + return json; + }) + .catch((err) => { + error = err.detail; + console.error(err); + return null; + }); + + if (error) { + throw error; + } + + return res; +}; + export const reindexKnowledgeFiles = async (token: string) => { let error = null; diff --git a/src/lib/components/admin/Settings/Documents.svelte b/src/lib/components/admin/Settings/Documents.svelte index ece64afd54..e93c302b14 100644 --- a/src/lib/components/admin/Settings/Documents.svelte +++ b/src/lib/components/admin/Settings/Documents.svelte @@ -194,6 +194,34 @@ return; } + // Google Drive validation + if (RAGConfig.ENABLE_GOOGLE_DRIVE_INTEGRATION && !RAGConfig.GOOGLE_DRIVE_API_KEY) { + toast.error($i18n.t('Google Drive API Key is required when integration is enabled.')); + return; + } + + // Folder sync specific validation + if (RAGConfig.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC) { + if (!RAGConfig.GOOGLE_DRIVE_CLIENT_ID) { + toast.error($i18n.t('Google Drive Client ID is required when folder sync is enabled.')); + return; + } + + if (!RAGConfig.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON) { + toast.error($i18n.t('Service Account JSON is required when folder sync is enabled.')); + return; + } + + // Validate Service Account JSON format + try { + JSON.parse(RAGConfig.GOOGLE_DRIVE_SERVICE_ACCOUNT_JSON); + } catch (e) { + toast.error($i18n.t('Invalid Service Account JSON format.')); + return; + } + } + + if (!RAGConfig.BYPASS_EMBEDDING_AND_RETRIEVAL) { await embeddingModelUpdateHandler(); } @@ -1460,7 +1488,115 @@ -
+ {#if RAGConfig.ENABLE_GOOGLE_DRIVE_INTEGRATION} +
+
+
+ + {$i18n.t('API Key')} * + +
+ +
+ {$i18n.t('Required for basic Google Drive file picker functionality')} +
+
+ +
+
+ + {$i18n.t('Client ID')} + +
+ +
+ {$i18n.t('Required for OAuth authentication when using the file picker')} +
+
+ +
+ +
+
+ + {$i18n.t('Enable Folder Sync')} + +
+
+ +
+
+ + {#if RAGConfig.ENABLE_GOOGLE_DRIVE_FOLDER_SYNC} +
+
+ {$i18n.t('Folder Sync Configuration')} +
+ +
+
+ + {$i18n.t('Service Account JSON')} * + +
+