diff --git a/src/langbot/pkg/api/http/service/knowledge.py b/src/langbot/pkg/api/http/service/knowledge.py index e85d16fd8..93d753021 100644 --- a/src/langbot/pkg/api/http/service/knowledge.py +++ b/src/langbot/pkg/api/http/service/knowledge.py @@ -303,7 +303,7 @@ class KnowledgeService: kb_uuid: str, file_id: str, parser_plugin_id: str | None = None, - ) -> str: + ) -> int: """存储文件""" execution_context = self._execution_context(context) runtime_kb = await self.ap.rag_mgr.get_knowledge_base_by_uuid(execution_context, kb_uuid) diff --git a/src/langbot/pkg/rag/knowledge/kbmgr.py b/src/langbot/pkg/rag/knowledge/kbmgr.py index 4ef1c9d0c..c9aa64edb 100644 --- a/src/langbot/pkg/rag/knowledge/kbmgr.py +++ b/src/langbot/pkg/rag/knowledge/kbmgr.py @@ -336,7 +336,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): execution_context: ExecutionContext, file_id: str, parser_plugin_id: str | None = None, - ) -> str: + ) -> int: await self._assert_execution_context(execution_context) self._require_upload_object_key(execution_context, file_id) # pre checking @@ -412,7 +412,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): execution_context: ExecutionContext, zip_file_id: str, parser_plugin_id: str | None = None, - ) -> str: + ) -> int: """Handle ZIP file by extracting each document and storing them separately.""" await self._assert_execution_context(execution_context) self._require_upload_object_key(execution_context, zip_file_id) @@ -516,7 +516,7 @@ class RuntimeKnowledgeBase(KnowledgeBaseInterface): self.ap.logger.info( f'Successfully processed ZIP file {zip_file_id}, extracted {len(stored_file_tasks)} files' ) - return stored_file_tasks[0] if stored_file_tasks else '' + return stored_file_tasks[0] finally: try: await self._assert_execution_context(execution_context) diff --git a/tests/integration/api/test_application_resources.py b/tests/integration/api/test_application_resources.py index 580ef592b..a86eba40c 100644 --- a/tests/integration/api/test_application_resources.py +++ b/tests/integration/api/test_application_resources.py @@ -255,3 +255,107 @@ async def test_chat_management_uses_scoped_history_and_real_connection_queues(re == [] ) assert len(adapter.websocket_group_session.message_lists[pipeline]) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize('extension', ['txt', 'zip']) +async def test_knowledge_upload_returns_real_integer_task_id(resources, tmp_path, monkeypatch, extension): + import asyncio + import io + import zipfile + + from langbot.pkg.api.context import ExecutionContext + from langbot.pkg.core.taskmgr import AsyncTaskManager + from langbot.pkg.entity.persistence.rag import File + from langbot.pkg.entity.persistence.workspace import Workspace, WorkspaceExecutionState + from langbot.pkg.storage.mgr import StorageMgr + from langbot.pkg.storage.providers import localstorage + + api, ap, engine = resources + workspace = '00000000-0000-0000-0000-00000000000a' + request = context(workspace) + execution = ExecutionContext.from_request(request) + kb_data = { + 'uuid': 'ingestion-kb', + 'workspace_uuid': workspace, + 'name': 'Ingestion', + 'knowledge_engine_plugin_id': 'test/engine', + 'collection_id': 'ingestion-kb', + 'creation_settings': {}, + 'retrieval_settings': {}, + 'initialized': True, + } + async with engine.begin() as connection: + await connection.execute( + sa.insert(Workspace).values( + uuid=workspace, + instance_uuid=request.instance_uuid, + name='Ingestion', + slug='ingestion', + source='cloud_projection', + ) + ) + await connection.execute( + sa.insert(WorkspaceExecutionState).values( + workspace_uuid=workspace, instance_uuid=request.instance_uuid, active_generation=1, source='cloud' + ) + ) + await connection.execute(sa.insert(KnowledgeBase).values(**kb_data)) + ap.event_loop = asyncio.get_running_loop() + ap.task_mgr = AsyncTaskManager(ap) + monkeypatch.setattr(localstorage, 'LOCAL_STORAGE_PATH', str(tmp_path / 'uploads')) + ap.storage_mgr = StorageMgr(ap) + ap.storage_mgr.storage_provider = localstorage.LocalStorageProvider(ap) + entered, release = asyncio.Event(), asyncio.Event() + + async def ingest(_plugin_id, _payload): + entered.set() + await release.wait() + return {'document_id': 'engine-document', 'status': 'completed'} + + ap.plugin_connector = SimpleNamespace( + is_enable_plugin=True, + require_workspace_context=AsyncMock(side_effect=lambda scope: scope), + list_knowledge_engines=AsyncMock( + return_value=[{'plugin_id': 'test/engine', 'capabilities': ['doc_ingestion']}] + ), + call_rag_ingest=AsyncMock(side_effect=ingest), + ) + await ap.rag_mgr.load_knowledge_base(execution, KnowledgeBase(**kb_data)) + content = b'A document for ingestion.' + if extension == 'zip': + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, 'w') as archive: + archive.writestr('document.txt', content) + content = buffer.getvalue() + upload = await ap.storage_mgr.save_scoped( + execution, owner_type='upload_document', owner='ingestion-test', key='document.' + extension, value=content + ) + try: + task_id = await api.execute('knowledge.file.store', request, {'kb_uuid': kb_data['uuid'], 'file_id': upload}) + assert type(task_id) is int + task = ap.task_mgr.get_task_by_id( + task_id, instance_uuid=request.instance_uuid, workspace_uuid=workspace, placement_generation=1 + ) + assert task is not None and task.kind == 'knowledge-operation' + await asyncio.wait_for(entered.wait(), 5) + async with engine.connect() as connection: + pending = ( + (await connection.execute(sa.select(File).where(File.workspace_uuid == workspace))).mappings().one() + ) + assert pending['kb_id'] == kb_data['uuid'] and pending['status'] == 'processing' + release.set() + await asyncio.wait_for(task.task, 5) + async with engine.connect() as connection: + stored = (await connection.execute(sa.select(File).where(File.uuid == pending['uuid']))).mappings().one() + assert stored['status'] == 'completed' and stored['engine_document_id'] == 'engine-document' + assert not await ap.storage_mgr.exists_scoped_object_key( + execution, upload, expected_owner_type='upload_document' + ) + finally: + release.set() + tasks = [wrapper.task for wrapper in ap.task_mgr.get_all_tasks()] + for task in tasks: + if not task.done(): + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) diff --git a/tests/integration/api/test_knowledge.py b/tests/integration/api/test_knowledge.py index a046da6d4..7c7beb5d2 100644 --- a/tests/integration/api/test_knowledge.py +++ b/tests/integration/api/test_knowledge.py @@ -123,7 +123,7 @@ def fake_knowledge_app(): app.knowledge_service.get_files_by_knowledge_base = AsyncMock( return_value=[{'uuid': 'test-file-uuid', 'filename': 'test.pdf'}] ) - app.knowledge_service.store_file = AsyncMock(return_value='test-task-id') + app.knowledge_service.store_file = AsyncMock(return_value=17) app.knowledge_service.delete_file = AsyncMock() app.knowledge_service.retrieve_knowledge_base = AsyncMock(return_value=[{'content': 'test result', 'score': 0.95}]) @@ -304,7 +304,7 @@ class TestKnowledgeBaseFilesEndpoints: assert called.await_args.args[0].workspace_uuid == '00000000-0000-0000-0000-00000000000a' assert called.await_args.args[1:] == ('test-kb-uuid', 'test-file-id') assert called.await_args.kwargs == {'parser_plugin_id': 'test/parser'} - assert data['data'] == {'task_id': 'test-task-id'} + assert data['data'] == {'task_id': 17} @pytest.mark.asyncio async def test_delete_file_from_knowledge_base(self, quart_test_client, fake_knowledge_app): diff --git a/tests/unit_tests/api/service/test_knowledge_service.py b/tests/unit_tests/api/service/test_knowledge_service.py index 1dc3496ca..663bfa0d7 100644 --- a/tests/unit_tests/api/service/test_knowledge_service.py +++ b/tests/unit_tests/api/service/test_knowledge_service.py @@ -265,14 +265,14 @@ async def test_file_listing_checks_parent_knowledge_base_first(): async def test_store_and_delete_file_require_runtime_parent_and_capability(): app = _app() runtime_kb = SimpleNamespace( - store_file=AsyncMock(return_value='task-a'), + store_file=AsyncMock(return_value=17), delete_file=AsyncMock(), ) app.rag_mgr.get_knowledge_base_by_uuid.return_value = runtime_kb app.rag_mgr.get_knowledge_base_details.return_value = {'knowledge_engine': {'capabilities': ['doc_ingestion']}} service = KnowledgeService(app) - assert await service.store_file(CONTEXT, 'kb-a', 'upload.pdf', 'author/parser') == 'task-a' + assert await service.store_file(CONTEXT, 'kb-a', 'upload.pdf', 'author/parser') == 17 runtime_kb.store_file.assert_awaited_once_with(CONTEXT, 'upload.pdf', parser_plugin_id='author/parser') await service.delete_file(CONTEXT, 'kb-a', 'file-a') runtime_kb.delete_file.assert_awaited_once_with(CONTEXT, 'file-a') diff --git a/tests/unit_tests/rag/test_file_storage.py b/tests/unit_tests/rag/test_file_storage.py index f658218fe..7ba7fdc95 100644 --- a/tests/unit_tests/rag/test_file_storage.py +++ b/tests/unit_tests/rag/test_file_storage.py @@ -90,14 +90,14 @@ class TestStoreFile: def create_user_task(coro, **kwargs): coro.close() - return SimpleNamespace(id='task-1', kwargs=kwargs, task=Mock()) + return SimpleNamespace(id=17, kwargs=kwargs, task=Mock()) kb.ap.task_mgr.create_user_task = Mock(side_effect=create_user_task) object_key = _upload_key('documents/test.pdf') task_id = await kb.store_file(CONTEXT, object_key) - assert task_id == 'task-1' + assert task_id == 17 kb.ap.storage_mgr.storage_provider.exists.assert_awaited_once_with(object_key) kb.ap.persistence_mgr.execute_async.assert_awaited_once() call_kwargs = kb.ap.task_mgr.create_user_task.call_args.kwargs @@ -189,12 +189,12 @@ class TestStoreZipFile: } ) ) - kb.store_file = AsyncMock(side_effect=['task-pdf', 'task-txt', 'task-md', 'task-html']) + kb.store_file = AsyncMock(side_effect=[17, 18, 19, 20]) zip_key = _upload_key('archive.zip') task_id = await kb._store_zip_file(CONTEXT, zip_key, parser_plugin_id='parser/plugin') - assert task_id == 'task-pdf' + assert task_id == 17 assert kb.ap.storage_mgr.storage_provider.save.await_count == 4 saved_names = [call.args[0] for call in kb.ap.storage_mgr.storage_provider.save.await_args_list] assert {name.rsplit('.', 1)[-1] for name in saved_names} == {'pdf', 'txt', 'md', 'html'}