From 9a081c8593533fccce50600378a68f433019fc7a Mon Sep 17 00:00:00 2001 From: Timothy Jaeryang Baek Date: Thu, 12 Dec 2024 22:32:28 -0800 Subject: [PATCH] refac --- backend/open_webui/main.py | 2 ++ backend/open_webui/routers/files.py | 14 +++++++++----- backend/open_webui/routers/knowledge.py | 10 ++++++++-- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/backend/open_webui/main.py b/backend/open_webui/main.py index ce9d47959..31604984f 100644 --- a/backend/open_webui/main.py +++ b/backend/open_webui/main.py @@ -812,6 +812,8 @@ async def chat_completion( user=Depends(get_verified_user), bypass_filter: bool = False, ): + if not request.app.state.MODELS: + await get_all_models(request) try: model_id = form_data.get("model", None) diff --git a/backend/open_webui/routers/files.py b/backend/open_webui/routers/files.py index bc553ab26..e56eef273 100644 --- a/backend/open_webui/routers/files.py +++ b/backend/open_webui/routers/files.py @@ -21,7 +21,7 @@ from open_webui.env import SRC_LOG_LEVELS from open_webui.constants import ERROR_MESSAGES -from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, status +from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, status, Request from fastapi.responses import FileResponse, StreamingResponse @@ -39,7 +39,9 @@ router = APIRouter() @router.post("/", response_model=FileModelResponse) -def upload_file(file: UploadFile = File(...), user=Depends(get_verified_user)): +def upload_file( + request: Request, file: UploadFile = File(...), user=Depends(get_verified_user) +): log.info(f"file.content_type: {file.content_type}") try: unsanitized_filename = file.filename @@ -68,7 +70,7 @@ def upload_file(file: UploadFile = File(...), user=Depends(get_verified_user)): ) try: - process_file(ProcessFileForm(file_id=id)) + process_file(request, ProcessFileForm(file_id=id)) file_item = Files.get_file_by_id(id=id) except Exception as e: log.exception(e) @@ -183,13 +185,15 @@ class ContentForm(BaseModel): @router.post("/{id}/data/content/update") async def update_file_data_content_by_id( - id: str, form_data: ContentForm, user=Depends(get_verified_user) + request: Request, id: str, form_data: ContentForm, user=Depends(get_verified_user) ): file = Files.get_file_by_id(id) if file and (file.user_id == user.id or user.role == "admin"): try: - process_file(ProcessFileForm(file_id=id, content=form_data.content)) + process_file( + request, ProcessFileForm(file_id=id, content=form_data.content) + ) file = Files.get_file_by_id(id=id) except Exception as e: log.exception(e) diff --git a/backend/open_webui/routers/knowledge.py b/backend/open_webui/routers/knowledge.py index 0f4dd9283..7f9947d7a 100644 --- a/backend/open_webui/routers/knowledge.py +++ b/backend/open_webui/routers/knowledge.py @@ -242,6 +242,7 @@ class KnowledgeFileIdForm(BaseModel): @router.post("/{id}/file/add", response_model=Optional[KnowledgeFilesResponse]) def add_file_to_knowledge_by_id( + request: Request, id: str, form_data: KnowledgeFileIdForm, user=Depends(get_verified_user), @@ -274,7 +275,9 @@ def add_file_to_knowledge_by_id( # Add content to the vector database try: - process_file(ProcessFileForm(file_id=form_data.file_id, collection_name=id)) + process_file( + request, ProcessFileForm(file_id=form_data.file_id, collection_name=id) + ) except Exception as e: log.debug(e) raise HTTPException( @@ -318,6 +321,7 @@ def add_file_to_knowledge_by_id( @router.post("/{id}/file/update", response_model=Optional[KnowledgeFilesResponse]) def update_file_from_knowledge_by_id( + request: Request, id: str, form_data: KnowledgeFileIdForm, user=Depends(get_verified_user), @@ -349,7 +353,9 @@ def update_file_from_knowledge_by_id( # Add content to the vector database try: - process_file(ProcessFileForm(file_id=form_data.file_id, collection_name=id)) + process_file( + request, ProcessFileForm(file_id=form_data.file_id, collection_name=id) + ) except Exception as e: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST,