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)
####################################
# 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",

View File

@ -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,

View File

@ -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:

View File

@ -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

View File

@ -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;

View File

@ -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>

View File

@ -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={() => {

View File

@ -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;