This commit is contained in:
Timothy J. Baek 2024-10-03 22:35:43 -07:00
parent e9b68524e8
commit d834bd2a18

View File

@ -11,11 +11,13 @@ from open_webui.apps.webui.models.knowledge import (
KnowledgeResponse, KnowledgeResponse,
) )
from open_webui.apps.webui.models.files import Files, FileModel 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.constants import ERROR_MESSAGES
from open_webui.utils.utils import get_admin_user, get_verified_user 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() router = APIRouter()
@ -149,6 +151,9 @@ def add_file_to_knowledge_by_id(
detail=ERROR_MESSAGES.FILE_NOT_PROCESSED, 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: if knowledge:
data = knowledge.data or {} data = knowledge.data or {}
file_ids = data.get("file_ids", []) file_ids = data.get("file_ids", [])
@ -208,11 +213,13 @@ def remove_file_from_knowledge_by_id(
detail=ERROR_MESSAGES.NOT_FOUND, detail=ERROR_MESSAGES.NOT_FOUND,
) )
Files.delete_file_by_id(form_data.file_id) # Remove content from the vector database
VECTOR_DB_CLIENT.delete( VECTOR_DB_CLIENT.delete(
collection_name=knowledge.id, filter={"file_id": form_data.file_id} collection_name=knowledge.id, filter={"file_id": form_data.file_id}
) )
Files.delete_file_by_id(form_data.file_id)
if knowledge: if knowledge:
data = knowledge.data or {} data = knowledge.data or {}
file_ids = data.get("file_ids", []) file_ids = data.get("file_ids", [])