diff --git a/backend/open_webui/routers/files.py b/backend/open_webui/routers/files.py index 7ca1c2e73f..f856d45785 100644 --- a/backend/open_webui/routers/files.py +++ b/backend/open_webui/routers/files.py @@ -50,6 +50,7 @@ from open_webui.storage.provider import Storage from open_webui.config import BYPASS_ADMIN_ACCESS_CONTROL, STORAGE_LOCAL_CACHE, STORAGE_PROVIDER, UPLOAD_DIR from open_webui.utils.auth import get_admin_user, get_verified_user +from open_webui.utils.knowledge_collections import reindex_file_in_collection from open_webui.utils.misc import strict_match_mime_type from pydantic import BaseModel @@ -577,14 +578,17 @@ async def update_file_data_content_by_id( knowledges = await Knowledges.get_knowledges_by_file_id(id, db=db) for knowledge in knowledges: try: - # Remove old embeddings for this file from the KB collection - await ASYNC_VECTOR_DB_CLIENT.delete(collection_name=knowledge.id, filter={'file_id': id}) - # Re-add from the now-updated file-{file_id} collection - await process_file( - request, - ProcessFileForm(file_id=id, collection_name=knowledge.id), + # Rebuild first, then drop the stale vectors by their old ids. + # This preserves the previously indexed chunks if the rebuild fails. + await reindex_file_in_collection( + request=request, + file_id=id, + collection_name=knowledge.id, user=user, db=db, + process_file_form_factory=ProcessFileForm, + process_file_func=process_file, + vector_db_client=ASYNC_VECTOR_DB_CLIENT, ) except Exception as e: log.warning(f'Failed to update knowledge {knowledge.id} after content change for file {id}: {e}') diff --git a/backend/open_webui/test/util/test_knowledge_collections.py b/backend/open_webui/test/util/test_knowledge_collections.py new file mode 100644 index 0000000000..ede7488df8 --- /dev/null +++ b/backend/open_webui/test/util/test_knowledge_collections.py @@ -0,0 +1,86 @@ +import asyncio +from types import SimpleNamespace + +from open_webui.utils.knowledge_collections import reindex_file_in_collection + + +def test_reindex_file_in_collection_deletes_old_ids_only_after_success(): + calls = [] + + class DummyVectorDB: + async def query(self, collection_name, filter, limit=None): + calls.append(('query', collection_name, filter)) + return SimpleNamespace(ids=[['old-1', 'old-2']]) + + async def delete(self, collection_name, ids=None, filter=None): + calls.append(('delete', collection_name, ids, filter)) + + async def fake_process_file(request, form, user, db): + calls.append(('process_file', form.file_id, form.collection_name)) + + class DummyProcessFileForm: + def __init__(self, file_id, collection_name=None): + self.file_id = file_id + self.collection_name = collection_name + + asyncio.run( + reindex_file_in_collection( + request=object(), + file_id='file-1', + collection_name='kb-1', + user=object(), + db=None, + process_file_form_factory=DummyProcessFileForm, + process_file_func=fake_process_file, + vector_db_client=DummyVectorDB(), + ) + ) + + assert calls == [ + ('query', 'kb-1', {'file_id': 'file-1'}), + ('process_file', 'file-1', 'kb-1'), + ('delete', 'kb-1', ['old-1', 'old-2'], None), + ] + + +def test_reindex_file_in_collection_preserves_old_ids_when_rebuild_fails(): + calls = [] + + class DummyVectorDB: + async def query(self, collection_name, filter, limit=None): + calls.append(('query', collection_name, filter)) + return SimpleNamespace(ids=[['old-1']]) + + async def delete(self, collection_name, ids=None, filter=None): + calls.append(('delete', collection_name, ids, filter)) + + async def fake_process_file(request, form, user, db): + calls.append(('process_file', form.file_id, form.collection_name)) + raise RuntimeError('reindex failed') + + class DummyProcessFileForm: + def __init__(self, file_id, collection_name=None): + self.file_id = file_id + self.collection_name = collection_name + + try: + asyncio.run( + reindex_file_in_collection( + request=object(), + file_id='file-1', + collection_name='kb-1', + user=object(), + db=None, + process_file_form_factory=DummyProcessFileForm, + process_file_func=fake_process_file, + vector_db_client=DummyVectorDB(), + ) + ) + raise AssertionError('Expected reindex_file_in_collection to raise RuntimeError') + except RuntimeError as exc: + assert str(exc) == 'reindex failed' + + assert calls == [ + ('query', 'kb-1', {'file_id': 'file-1'}), + ('process_file', 'file-1', 'kb-1'), + ] diff --git a/backend/open_webui/utils/knowledge_collections.py b/backend/open_webui/utils/knowledge_collections.py new file mode 100644 index 0000000000..f2764386b1 --- /dev/null +++ b/backend/open_webui/utils/knowledge_collections.py @@ -0,0 +1,32 @@ +async def reindex_file_in_collection( + *, + request, + file_id: str, + collection_name: str, + user, + db, + process_file_form_factory, + process_file_func, + vector_db_client, +) -> None: + result = await vector_db_client.query( + collection_name=collection_name, + filter={'file_id': file_id}, + ) + + existing_ids = [] + if result is not None and result.ids and result.ids[0]: + existing_ids = list(result.ids[0]) + + await process_file_func( + request, + process_file_form_factory(file_id=file_id, collection_name=collection_name), + user=user, + db=db, + ) + + if existing_ids: + await vector_db_client.delete( + collection_name=collection_name, + ids=existing_ids, + )