From d834bd2a189ee660780bd1d091a2d187200423d9 Mon Sep 17 00:00:00 2001 From: "Timothy J. Baek" Date: Thu, 3 Oct 2024 22:35:43 -0700 Subject: [PATCH] refac --- backend/open_webui/apps/webui/routers/knowledge.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/backend/open_webui/apps/webui/routers/knowledge.py b/backend/open_webui/apps/webui/routers/knowledge.py index 9fb8f5a8d..c55d84d22 100644 --- a/backend/open_webui/apps/webui/routers/knowledge.py +++ b/backend/open_webui/apps/webui/routers/knowledge.py @@ -11,11 +11,13 @@ from open_webui.apps.webui.models.knowledge import ( KnowledgeResponse, ) from open_webui.apps.webui.models.files import Files, FileModel +from open_webui.apps.retrieval.vector.connector import VECTOR_DB_CLIENT +from open_webui.apps.retrieval.main import process_file, ProcessFileForm + from open_webui.constants import ERROR_MESSAGES from open_webui.utils.utils import get_admin_user, get_verified_user -from open_webui.apps.retrieval.vector.connector import VECTOR_DB_CLIENT router = APIRouter() @@ -149,6 +151,9 @@ def add_file_to_knowledge_by_id( detail=ERROR_MESSAGES.FILE_NOT_PROCESSED, ) + # Add content to the vector database + process_file(ProcessFileForm(file_id=form_data.file_id, collection_name=id)) + if knowledge: data = knowledge.data or {} file_ids = data.get("file_ids", []) @@ -208,11 +213,13 @@ def remove_file_from_knowledge_by_id( detail=ERROR_MESSAGES.NOT_FOUND, ) - Files.delete_file_by_id(form_data.file_id) + # Remove content from the vector database VECTOR_DB_CLIENT.delete( collection_name=knowledge.id, filter={"file_id": form_data.file_id} ) + Files.delete_file_by_id(form_data.file_id) + if knowledge: data = knowledge.data or {} file_ids = data.get("file_ids", [])