From 0dc75363aa98c1f44781cc86e652a6534533ae5c Mon Sep 17 00:00:00 2001 From: Taylor Wilsdon Date: Wed, 18 Dec 2024 13:25:57 -0500 Subject: [PATCH] Add configurable Google Drive toggle in the Documents admin section along with necessary config scaffolding --- backend/open_webui/config.py | 7 +++++++ backend/open_webui/main.py | 5 +++-- backend/open_webui/routers/retrieval.py | 9 +++++++++ src/lib/apis/retrieval/index.ts | 1 + src/lib/components/admin/Settings/Documents.svelte | 2 -- src/lib/components/chat/MessageInput/InputMenu.svelte | 2 ++ src/lib/stores/index.ts | 1 + 7 files changed, 23 insertions(+), 4 deletions(-) diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index a69af97ba..76b211fb1 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -1438,6 +1438,13 @@ 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", "rag.web.search.searxng_query_url", diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index 48e6268b5..3a47702f2 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -183,6 +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, UPLOAD_DIR, # WebUI WEBUI_AUTH, @@ -234,8 +235,6 @@ from open_webui.config import ( CORS_ALLOW_ORIGIN, DEFAULT_LOCALE, OAUTH_PROVIDERS, - GOOGLE_DRIVE_CLIENT_ID, - GOOGLE_DRIVE_API_KEY, # Admin ENABLE_ADMIN_CHAT_ACCESS, ENABLE_ADMIN_EXPORT, @@ -487,6 +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.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 @@ -939,6 +939,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_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 c6a3a0cca..d19940197 100644 --- a/backend/open_webui/routers/retrieval.py +++ b/backend/open_webui/routers/retrieval.py @@ -347,6 +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, "content_extraction": { "engine": request.app.state.config.CONTENT_EXTRACTION_ENGINE, "tika_server_url": request.app.state.config.TIKA_SERVER_URL, @@ -369,6 +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, "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, @@ -445,6 +447,7 @@ class WebConfig(BaseModel): class ConfigUpdateForm(BaseModel): pdf_extract_images: Optional[bool] = None + enable_google_drive: Optional[bool] = None file: Optional[FileConfig] = None content_extraction: Optional[ContentExtractionConfig] = None chunk: Optional[ChunkParamUpdateForm] = None @@ -462,6 +465,12 @@ 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 + ) + if form_data.file is not None: request.app.state.config.FILE_MAX_SIZE = form_data.file.max_size request.app.state.config.FILE_MAX_COUNT = form_data.file.max_count diff --git a/src/lib/apis/retrieval/index.ts b/src/lib/apis/retrieval/index.ts index 21ae792fa..7b8d836ce 100644 --- a/src/lib/apis/retrieval/index.ts +++ b/src/lib/apis/retrieval/index.ts @@ -45,6 +45,7 @@ type YoutubeConfigForm = { type RAGConfigForm = { pdf_extract_images?: boolean; + enable_google_drive?: 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 f802d4b7f..5ff48900a 100644 --- a/src/lib/components/admin/Settings/Documents.svelte +++ b/src/lib/components/admin/Settings/Documents.svelte @@ -592,7 +592,6 @@
-
{$i18n.t('Google Drive')}
@@ -604,7 +603,6 @@
-
diff --git a/src/lib/components/chat/MessageInput/InputMenu.svelte b/src/lib/components/chat/MessageInput/InputMenu.svelte index 299ebb6e3..9f861ec78 100644 --- a/src/lib/components/chat/MessageInput/InputMenu.svelte +++ b/src/lib/components/chat/MessageInput/InputMenu.svelte @@ -154,6 +154,7 @@
{$i18n.t('Upload Files')}
+ {#if $config?.features?.enable_google_drive} { @@ -188,6 +189,7 @@
{$i18n.t('Google Drive')}
 + {/if} diff --git a/src/lib/stores/index.ts b/src/lib/stores/index.ts index 2e3976bf9..4e82ba966 100644 --- a/src/lib/stores/index.ts +++ b/src/lib/stores/index.ts @@ -176,6 +176,7 @@ type Config = { enable_signup: boolean; enable_login_form: boolean; enable_web_search?: boolean; + enable_google_drive: boolean; enable_image_generation: boolean; enable_admin_export: boolean; enable_admin_chat_access: boolean;