This commit is contained in:
Timothy Jaeryang Baek 2024-12-18 18:04:56 -08:00
parent 5871df02ac
commit 0f6d302760
8 changed files with 41 additions and 36 deletions

View File

@ -1203,6 +1203,15 @@ if VECTOR_DB == "pgvector" and not PGVECTOR_DB_URL.startswith("postgres"):
# Information Retrieval (RAG) # 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 # RAG Content Extraction
CONTENT_EXTRACTION_ENGINE = PersistentConfig( CONTENT_EXTRACTION_ENGINE = PersistentConfig(
"CONTENT_EXTRACTION_ENGINE", "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 = PersistentConfig(
"SEARXNG_QUERY_URL", "SEARXNG_QUERY_URL",

View File

@ -183,7 +183,7 @@ from open_webui.config import (
ENABLE_RAG_LOCAL_WEB_FETCH, ENABLE_RAG_LOCAL_WEB_FETCH,
ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION, ENABLE_RAG_WEB_LOADER_SSL_VERIFICATION,
ENABLE_RAG_WEB_SEARCH, ENABLE_RAG_WEB_SEARCH,
ENABLE_GOOGLE_DRIVE, ENABLE_GOOGLE_DRIVE_INTEGRATION,
UPLOAD_DIR, UPLOAD_DIR,
# WebUI # WebUI
WEBUI_AUTH, 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_ENGINE = RAG_WEB_SEARCH_ENGINE
app.state.config.RAG_WEB_SEARCH_DOMAIN_FILTER_LIST = RAG_WEB_SEARCH_DOMAIN_FILTER_LIST 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.SEARXNG_QUERY_URL = SEARXNG_QUERY_URL
app.state.config.GOOGLE_PSE_API_KEY = GOOGLE_PSE_API_KEY app.state.config.GOOGLE_PSE_API_KEY = GOOGLE_PSE_API_KEY
app.state.config.GOOGLE_PSE_ENGINE_ID = GOOGLE_PSE_ENGINE_ID app.state.config.GOOGLE_PSE_ENGINE_ID = GOOGLE_PSE_ENGINE_ID
@ -839,7 +839,18 @@ async def chat_completion(
except Exception as e: except Exception as e:
raise 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: except Exception as e:
raise HTTPException( raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST, status_code=status.HTTP_400_BAD_REQUEST,
@ -850,7 +861,7 @@ async def chat_completion(
response = await chat_completion_handler( response = await chat_completion_handler(
request, form_data, user, bypass_filter 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: except Exception as e:
raise HTTPException( raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST, 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_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_image_generation": app.state.config.ENABLE_IMAGE_GENERATION,
"enable_community_sharing": app.state.config.ENABLE_COMMUNITY_SHARING, "enable_community_sharing": app.state.config.ENABLE_COMMUNITY_SHARING,
"enable_message_rating": app.state.config.ENABLE_MESSAGE_RATING, "enable_message_rating": app.state.config.ENABLE_MESSAGE_RATING,

View File

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

View File

@ -391,19 +391,10 @@ def apply_params_to_form_data(form_data, model):
return form_data 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) form_data = apply_params_to_form_data(form_data, model)
log.debug(f"form_data: {form_data}") 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 = { extra_params = {
"__event_emitter__": get_event_emitter(metadata), "__event_emitter__": get_event_emitter(metadata),
"__event_call__": get_event_call(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 return form_data, events
async def process_chat_response(response, events): async def process_chat_response(response, events, metadata):
if not isinstance(response, StreamingResponse): if not isinstance(response, StreamingResponse):
return response return response

View File

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

View File

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

View File

@ -154,7 +154,7 @@
<div class="line-clamp-1">{$i18n.t('Upload Files')}</div> <div class="line-clamp-1">{$i18n.t('Upload Files')}</div>
</DropdownMenu.Item> </DropdownMenu.Item>
{#if $config?.features?.enable_google_drive} {#if $config?.features?.enable_google_drive_integration}
<DropdownMenu.Item <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" 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={() => { on:click={() => {

View File

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