OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-26 18:26:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

refac

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

15
backend/open_webui/config.py
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",

21
backend/open_webui/main.py
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,

14
backend/open_webui/routers/retrieval.py
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:

13
backend/open_webui/utils/middleware.py
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