Timothy Jaeryang Baek 4 months ago
parent
commit
0f6d302760

+ 9 - 6
backend/open_webui/config.py

@@ -1203,6 +1203,15 @@ if VECTOR_DB == "pgvector" and not PGVECTOR_DB_URL.startswith("postgres"):
 # Information Retrieval (RAG)
 ####################################
 
+
+# If configured, Google Drive will be available as an upload option.
+ENABLE_GOOGLE_DRIVE_INTEGRATION = PersistentConfig(
+    "ENABLE_GOOGLE_DRIVE_INTEGRATION",
+    "google_drive.enable",
+    os.getenv("ENABLE_GOOGLE_DRIVE_INTEGRATION", "False").lower() == "true",
+)
+
+
 # RAG Content Extraction
 CONTENT_EXTRACTION_ENGINE = PersistentConfig(
     "CONTENT_EXTRACTION_ENGINE",
@@ -1438,12 +1447,6 @@ RAG_WEB_SEARCH_DOMAIN_FILTER_LIST = PersistentConfig(
     ],
 )
 
-# If configured, Google Drive will be available as an upload option.
-ENABLE_GOOGLE_DRIVE = PersistentConfig(
-    "ENABLE_GOOGLE_DRIVE",
-    "rag.drive.enable",
-    os.getenv("ENABLE_GOOGLE_DRIVE", "False").lower() == "true",
-)
 
 SEARXNG_QUERY_URL = PersistentConfig(
     "SEARXNG_QUERY_URL",

+ 16 - 5
backend/open_webui/main.py

@@ -183,7 +183,7 @@ from open_webui.config import (
     ENABLE_RAG_LOCAL_WEB_FETCH,
     ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
     ENABLE_RAG_WEB_SEARCH,
-    ENABLE_GOOGLE_DRIVE,
+    ENABLE_GOOGLE_DRIVE_INTEGRATION,
     UPLOAD_DIR,
     # WebUI
     WEBUI_AUTH,
@@ -486,7 +486,7 @@ app.state.config.ENABLE_RAG_WEB_SEARCH = ENABLE_RAG_WEB_SEARCH
 app.state.config.RAG_WEB_SEARCH_ENGINE = RAG_WEB_SEARCH_ENGINE
 app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST = RAG_WEB_SEARCH_DOMAIN_FILTER_LIST
 
-app.state.config.ENABLE_GOOGLE_DRIVE = ENABLE_GOOGLE_DRIVE
+app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION = ENABLE_GOOGLE_DRIVE_INTEGRATION
 app.state.config.SEARXNG_QUERY_URL = SEARXNG_QUERY_URL
 app.state.config.GOOGLE_PSE_API_KEY = GOOGLE_PSE_API_KEY
 app.state.config.GOOGLE_PSE_ENGINE_ID = GOOGLE_PSE_ENGINE_ID
@@ -839,7 +839,18 @@ async def chat_completion(
             except Exception as e:
                 raise e
 
-        form_data, events = await process_chat_payload(request, form_data, user, model)
+        metadata = {
+            "chat_id": form_data.pop("chat_id", None),
+            "message_id": form_data.pop("id", None),
+            "session_id": form_data.pop("session_id", None),
+            "tool_ids": form_data.get("tool_ids", None),
+            "files": form_data.get("files", None),
+        }
+        form_data["metadata"] = metadata
+
+        form_data, events = await process_chat_payload(
+            request, form_data, metadata, user, model
+        )
     except Exception as e:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -850,7 +861,7 @@ async def chat_completion(
         response = await chat_completion_handler(
             request, form_data, user, bypass_filter
         )
-        return await process_chat_response(response, events)
+        return await process_chat_response(response, events, metadata)
     except Exception as e:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -939,7 +950,7 @@ async def get_app_config(request: Request):
             **(
                 {
                     "enable_web_search": app.state.config.ENABLE_RAG_WEB_SEARCH,
-                    "enable_google_drive": app.state.config.ENABLE_GOOGLE_DRIVE,
+                    "enable_google_drive_integration": app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION,
                     "enable_image_generation": app.state.config.ENABLE_IMAGE_GENERATION,
                     "enable_community_sharing": app.state.config.ENABLE_COMMUNITY_SHARING,
                     "enable_message_rating": app.state.config.ENABLE_MESSAGE_RATING,

+ 7 - 7
backend/open_webui/routers/retrieval.py

@@ -347,7 +347,7 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
     return {
         "status": True,
         "pdf_extract_images": request.app.state.config.PDF_EXTRACT_IMAGES,
-        "enable_google_drive": request.app.state.config.ENABLE_GOOGLE_DRIVE,
+        "ENABLE_GOOGLE_DRIVE_INTEGRATION": request.app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION,
         "content_extraction": {
             "engine": request.app.state.config.CONTENT_EXTRACTION_ENGINE,
             "tika_server_url": request.app.state.config.TIKA_SERVER_URL,
@@ -370,7 +370,7 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
             "web_loader_ssl_verification": request.app.state.config.ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
             "search": {
                 "enabled": request.app.state.config.ENABLE_RAG_WEB_SEARCH,
-                "drive": request.app.state.config.ENABLE_GOOGLE_DRIVE,
+                "drive": request.app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION,
                 "engine": request.app.state.config.RAG_WEB_SEARCH_ENGINE,
                 "searxng_query_url": request.app.state.config.SEARXNG_QUERY_URL,
                 "google_pse_api_key": request.app.state.config.GOOGLE_PSE_API_KEY,
@@ -447,7 +447,7 @@ class WebConfig(BaseModel):
 
 class ConfigUpdateForm(BaseModel):
     pdf_extract_images: Optional[bool] = None
-    enable_google_drive: Optional[bool] = None
+    enable_google_drive_integration: Optional[bool] = None
     file: Optional[FileConfig] = None
     content_extraction: Optional[ContentExtractionConfig] = None
     chunk: Optional[ChunkParamUpdateForm] = None
@@ -465,10 +465,10 @@ async def update_rag_config(
         else request.app.state.config.PDF_EXTRACT_IMAGES
     )
 
-    request.app.state.config.ENABLE_GOOGLE_DRIVE = (
-        form_data.enable_google_drive
-        if form_data.enable_google_drive is not None
-        else request.app.state.config.ENABLE_GOOGLE_DRIVE
+    request.app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION = (
+        form_data.ENABLE_GOOGLE_DRIVE_INTEGRATION
+        if form_data.ENABLE_GOOGLE_DRIVE_INTEGRATION is not None
+        else request.app.state.config.ENABLE_GOOGLE_DRIVE_INTEGRATION
     )
 
     if form_data.file is not None:

+ 2 - 11
backend/open_webui/utils/middleware.py

@@ -391,19 +391,10 @@ def apply_params_to_form_data(form_data, model):
     return form_data
 
 
-async def process_chat_payload(request, form_data, user, model):
+async def process_chat_payload(request, form_data, metadata, user, model):
     form_data = apply_params_to_form_data(form_data, model)
     log.debug(f"form_data: {form_data}")
 
-    metadata = {
-        "chat_id": form_data.pop("chat_id", None),
-        "message_id": form_data.pop("id", None),
-        "session_id": form_data.pop("session_id", None),
-        "tool_ids": form_data.get("tool_ids", None),
-        "files": form_data.get("files", None),
-    }
-    form_data["metadata"] = metadata
-
     extra_params = {
         "__event_emitter__": get_event_emitter(metadata),
         "__event_call__": get_event_call(metadata),
@@ -513,7 +504,7 @@ async def process_chat_payload(request, form_data, user, model):
     return form_data, events
 
 
-async def process_chat_response(response, events):
+async def process_chat_response(response, events, metadata):
     if not isinstance(response, StreamingResponse):
         return response
 

+ 1 - 1
src/lib/apis/retrieval/index.ts

@@ -45,7 +45,7 @@ type YoutubeConfigForm = {
 
 type RAGConfigForm = {
 	pdf_extract_images?: boolean;
-	enable_google_drive?: boolean;
+	enable_google_drive_integration?: boolean;
 	chunk?: ChunkConfigForm;
 	content_extraction?: ContentExtractConfigForm;
 	web_loader_ssl_verification?: boolean;

+ 4 - 4
src/lib/components/admin/Settings/Documents.svelte

@@ -56,7 +56,7 @@
 	let chunkOverlap = 0;
 	let pdfExtractImages = true;
 
-	let enableGoogleDrive = false;
+	let enableGoogleDriveIntegration = false;
 
 	let OpenAIUrl = '';
 	let OpenAIKey = '';
@@ -177,7 +177,7 @@
 		}
 		const res = await updateRAGConfig(localStorage.token, {
 			pdf_extract_images: pdfExtractImages,
-			enable_google_drive: enableGoogleDrive,
+			enable_google_drive_integration: enableGoogleDriveIntegration,
 			file: {
 				max_size: fileMaxSize === '' ? null : fileMaxSize,
 				max_count: fileMaxCount === '' ? null : fileMaxCount
@@ -249,7 +249,7 @@
 			fileMaxSize = res?.file.max_size ?? '';
 			fileMaxCount = res?.file.max_count ?? '';
 
-			enableGoogleDrive = res.enable_google_drive;
+			enableGoogleDrive = res.enable_google_drive_integration;
 		}
 	});
 </script>
@@ -598,7 +598,7 @@
 				<div class="flex justify-between items-center text-xs">
 					<div class="text-xs font-medium">{$i18n.t('Enable Google Drive')}</div>
 					<div>
-						<Switch bind:state={enableGoogleDrive} />
+						<Switch bind:state={enableGoogleDriveIntegration} />
 					</div>
 				</div>
 			</div>

+ 1 - 1
src/lib/components/chat/MessageInput/InputMenu.svelte

@@ -154,7 +154,7 @@
 				<div class="line-clamp-1">{$i18n.t('Upload Files')}</div>
 			</DropdownMenu.Item>
 
-			{#if $config?.features?.enable_google_drive}
+			{#if $config?.features?.enable_google_drive_integration}
 				<DropdownMenu.Item
 					class="flex gap-2 items-center px-3 py-2 text-sm font-medium cursor-pointer hover:bg-gray-50 dark:hover:bg-gray-800 rounded-xl"
 					on:click={() => {

+ 1 - 1
src/lib/stores/index.ts

@@ -176,7 +176,7 @@ type Config = {
 		enable_signup: boolean;
 		enable_login_form: boolean;
 		enable_web_search?: boolean;
-		enable_google_drive: boolean;
+		enable_google_drive_integration: boolean;
 		enable_image_generation: boolean;
 		enable_admin_export: boolean;
 		enable_admin_chat_access: boolean;