From 0f6d302760506e695328381641bead397b6fd020 Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Wed, 18 Dec 2024 18:04:56 -0800 Subject: [PATCH] refac --- backend/open_webui/config.py | 15 +++++++------ backend/open_webui/main.py | 21 ++++++++++++++----- backend/open_webui/routers/retrieval.py | 14 ++++++------- backend/open_webui/utils/middleware.py | 13 ++---------- src/lib/apis/retrieval/index.ts | 2 +- .../admin/Settings/Documents.svelte | 8 +++---- .../chat/MessageInput/InputMenu.svelte | 2 +- src/lib/stores/index.ts | 2 +- 8 files changed, 41 insertions(+), 36 deletions(-) diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index 76b211fb1..8e922dff7 100644 --- a/backend/open_webui/config.py +++ b/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", diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 3a47702f2..f30ca2676 100644 --- a/backend/open_webui/main.py +++ b/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, diff --git a/backend/open_webui/routers/retrieval.py b/backend/open_webui/routers/retrieval.py index d19940197..02ab3ca59 100644 --- a/backend/open_webui/routers/retrieval.py +++ b/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: diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index f697339db..4115c7c2b 100644 --- a/backend/open_webui/utils/middleware.py +++ b/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 diff --git a/src/lib/apis/retrieval/index.ts b/src/lib/apis/retrieval/index.ts index 7b8d836ce..c35c37847 100644 --- a/src/lib/apis/retrieval/index.ts +++ b/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; diff --git a/src/lib/components/admin/Settings/Documents.svelte b/src/lib/components/admin/Settings/Documents.svelte index ff580a972..eca6513dd 100644 --- a/src/lib/components/admin/Settings/Documents.svelte +++ b/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; } }); @@ -598,7 +598,7 @@
{$i18n.t('Enable Google Drive')}
- +
diff --git a/src/lib/components/chat/MessageInput/InputMenu.svelte b/src/lib/components/chat/MessageInput/InputMenu.svelte index 7bded7c15..03280b0f6 100644 --- a/src/lib/components/chat/MessageInput/InputMenu.svelte +++ b/src/lib/components/chat/MessageInput/InputMenu.svelte @@ -154,7 +154,7 @@
{$i18n.t('Upload Files')}
- {#if $config?.features?.enable_google_drive} + {#if $config?.features?.enable_google_drive_integration} { diff --git a/src/lib/stores/index.ts b/src/lib/stores/index.ts index 4e82ba966..0319a5b25 100644 --- a/src/lib/stores/index.ts +++ b/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;