From b2fce411c805a114511ce69761ce5e244e872a0b Mon Sep 17 00:00:00 2001 From: JunghwanNA <70629228+shaun0927@users.noreply.github.com> Date: Thu, 16 Apr 2026 21:20:34 +0900 Subject: [PATCH] Prevent silent KB vector loss when file reindexing fails A file content update should not leave knowledge collections emptier than before. The previous flow deleted old vectors first and then attempted the rebuild, so any downstream processing failure could leave the KB without embeddings even though the route returned success.\n\nThis change rebuilds the file's knowledge entries first and only deletes the stale vector ids after the new insert succeeds. The fix is kept narrow to the file-update propagation path and adds a focused regression test for the helper that preserves old ids when reindexing raises.\n\nConstraint: Keep the fix scoped to one retrieval correctness bug without refactoring broader embedding flows\nRejected: Fail the route after deleting old vectors | still leaves the KB empty and does not preserve data\nRejected: Broad temporary-collection refactor | larger than needed for a first bugfix PR\nConfidence: high\nScope-risk: narrow\nReversibility: clean\nDirective: Preserve old vector ids until rebuild success remains observable in tests; do not revert to delete-first ordering\nTested: py_compile on modified modules; PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 PYTHONPATH=backend pytest -q backend/open_webui/test/util/test_knowledge_collections.py; local source-faithful harness reproducing old failure mode\nNot-tested: Full application boot with live vector database backends --- backend/open_webui/routers/files.py | 16 ++-- .../test/util/test_knowledge_collections.py | 86 +++++++++++++++++++ .../open_webui/utils/knowledge_collections.py | 32 +++++++ 3 files changed, 128 insertions(+), 6 deletions(-) create mode 100644 backend/open_webui/test/util/test_knowledge_collections.py create mode 100644 backend/open_webui/utils/knowledge_collections.py 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, + )