fix(rag): retain interrupted ingestion state and engine identity

This commit is contained in:
RockChinQ
2026-09-21 07:47:31 +00:00
parent 52f5699533
commit 381bb3f852
11 changed files with 1162 additions and 83 deletions
+6 -4
View File
@@ -90,7 +90,7 @@ class TestStoreFile:
def create_user_task(coro, **kwargs):
coro.close()
return SimpleNamespace(id='task-1', kwargs=kwargs)
return SimpleNamespace(id='task-1', kwargs=kwargs, task=Mock())
kb.ap.task_mgr.create_user_task = Mock(side_effect=create_user_task)
@@ -279,7 +279,7 @@ class TestStoreFileTask:
kb._assert_execution_context = AsyncMock(side_effect=assert_execution_context)
kb._set_file_status = AsyncMock(side_effect=[True, True])
kb._ingest_document = AsyncMock(return_value={'status': 'completed'})
kb._ingest_document = AsyncMock(return_value={'status': 'completed', 'document_id': 'file-uuid'})
object_key = _upload_key('scoped.pdf')
file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf')
@@ -290,7 +290,7 @@ class TestStoreFileTask:
@pytest.mark.asyncio
async def test_store_file_task_marks_completed_and_cleans_storage(self):
kb = _make_kb()
kb._ingest_document = AsyncMock(return_value={'status': 'completed'})
kb._ingest_document = AsyncMock(return_value={'status': 'completed', 'document_id': 'file-uuid'})
object_key = _upload_key('test.pdf')
file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf')
task_context = Mock()
@@ -306,7 +306,9 @@ class TestStoreFileTask:
@pytest.mark.asyncio
async def test_store_file_task_marks_failed_and_cleans_storage(self):
kb = _make_kb()
kb._ingest_document = AsyncMock(return_value={'status': 'failed', 'error_message': 'parser failed'})
kb._ingest_document = AsyncMock(
return_value={'status': 'failed', 'error_message': 'parser failed', 'document_id': 'file-uuid'}
)
object_key = _upload_key('bad.pdf')
file_obj = SimpleNamespace(uuid='file-uuid', file_name=object_key, extension='pdf')
task_context = Mock()
+11 -8
View File
@@ -268,7 +268,9 @@ async def test_ingestion_payload_uses_host_owned_kb_collection():
async def test_delete_file_checks_workspace_and_parent_before_plugin_call():
app = _app()
runtime = RuntimeKnowledgeBase(app, _entity(), CONTEXT_A)
app.persistence_mgr.execute_async.return_value = _Result(first=('file-a',))
app.persistence_mgr.execute_async.return_value = _Result(
first=SimpleNamespace(uuid='file-a', status='completed', engine_document_id=None)
)
await runtime.delete_file(CONTEXT_A, 'file-a')
app.plugin_connector.call_rag_delete_document.assert_awaited_once_with(
@@ -384,7 +386,8 @@ class TestRAGManagerCreateKnowledgeBase:
)
assert manager.knowledge_bases == {}
assert app.persistence_mgr.execute_async.await_count == 2
# Insert, interrupted-ingestion reconciliation, rollback delete.
assert app.persistence_mgr.execute_async.await_count == 3
@pytest.mark.asyncio
async def test_sets_default_retrieval_settings(self):
@@ -456,7 +459,9 @@ class TestRuntimeKnowledgeBaseDeleteFile:
@pytest.mark.asyncio
async def test_delete_file_calls_plugin_and_db(self):
app = _app()
app.persistence_mgr.execute_async.return_value = _Result(first=('file-uuid',))
app.persistence_mgr.execute_async.return_value = _Result(
first=SimpleNamespace(uuid='file-uuid', status='completed', engine_document_id=None)
)
await RuntimeKnowledgeBase(app, _entity(), CONTEXT_A).delete_file(
CONTEXT_A,
@@ -514,7 +519,7 @@ class TestRAGManagerLoadKnowledgeBasesFromDB:
}
@pytest.mark.asyncio
async def test_cloud_startup_reuses_validated_binding(self):
async def test_cloud_startup_revalidates_binding_before_recovery_write(self):
class TenantUow:
async def __aenter__(self):
return self
@@ -534,15 +539,13 @@ class TestRAGManagerLoadKnowledgeBasesFromDB:
app.persistence_mgr.tenant_uow = lambda _workspace_uuid: TenantUow()
app.persistence_mgr.execute_async.return_value = _Result([_entity()])
app.workspace_service.list_active_execution_bindings = AsyncMock(return_value=[binding])
app.workspace_service.get_execution_binding = AsyncMock(
side_effect=AssertionError('startup RAG loader repeated a validated binding lookup')
)
app.workspace_service.get_execution_binding = AsyncMock(return_value=binding)
manager = RAGManager(app)
await manager.load_knowledge_bases_from_db()
assert set(manager.knowledge_bases) == {('workspace-a', 'kb-a')}
app.workspace_service.get_execution_binding.assert_not_awaited()
app.workspace_service.get_execution_binding.assert_awaited_once_with('workspace-a', expected_generation=5)
@pytest.mark.asyncio
async def test_handles_load_error_gracefully(self):