From 5634ac3e2729bd4fdf908b535630efcabd68ed1f Mon Sep 17 00:00:00 2001 From: Wenjie Zhang Date: Sun, 17 May 2026 19:55:40 +0800 Subject: [PATCH] =?UTF-8?q?test:=20=E8=A1=A5=E5=85=85=E5=9B=BE=E8=B0=B1?= =?UTF-8?q?=E6=8A=BD=E5=8F=96=E5=99=A8=E3=80=81metadata=20=E6=8C=81?= =?UTF-8?q?=E4=B9=85=E5=8C=96=E4=B8=8E=E6=A3=80=E7=B4=A2=E5=85=83=E6=95=B0?= =?UTF-8?q?=E6=8D=AE=E6=B5=8B=E8=AF=95=EF=BC=9B=E6=9B=B4=E6=96=B0=20roadma?= =?UTF-8?q?p?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - LLM 抽取器测试:禁止自定义 Prompt、Schema 拼接、配置锁定修改 - metadata 持久化测试:验证 _save_metadata 不覆写数据库已有图谱配置 - 检索测试:验证 resource_id 写入 metadata - roadmap 补充图谱抽取器配置优化条目 --- .../unit/graphs/test_milvus_graph_build.py | 77 +++++++++++++++++-- .../unit/knowledge/test_file_size_fallback.py | 75 ++++++++++++++++++ backend/test/unit/toolkits/test_kbs_tools.py | 3 + 3 files changed, 149 insertions(+), 6 deletions(-) diff --git a/backend/test/unit/graphs/test_milvus_graph_build.py b/backend/test/unit/graphs/test_milvus_graph_build.py index 82c8b747..69509b13 100644 --- a/backend/test/unit/graphs/test_milvus_graph_build.py +++ b/backend/test/unit/graphs/test_milvus_graph_build.py @@ -1,11 +1,11 @@ from __future__ import annotations from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import AsyncMock, MagicMock import pytest -from yuxi.knowledge.graphs.extractors import normalize_extraction_result +from yuxi.knowledge.graphs.extractors import LLMGraphExtractor, normalize_extraction_result from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService @@ -63,6 +63,70 @@ def test_normalize_extraction_result_rejects_invalid_payload(payload): normalize_extraction_result(payload, "llm") +def test_llm_graph_extractor_rejects_custom_prompt(): + extractor = LLMGraphExtractor({"model_spec": "test/model", "prompt": "custom"}) + + with pytest.raises(ValueError, match="不支持自定义完整 Prompt"): + extractor.validate_options() + + +def test_llm_graph_extractor_appends_schema_to_fixed_prompt(): + extractor = LLMGraphExtractor( + { + "model_spec": "test/model", + "schema": "实体类型只能是 Person 或 Organization", + "concurrency_count": 5, + "model_params": {"temperature": 0.1}, + } + ) + + prompt = extractor._build_prompt("张三任职于公司") + + assert "请从下面文本中抽取实体和实体关系" in prompt + assert "抽取 Schema 约束" in prompt + assert "实体类型只能是 Person 或 Organization" in prompt + assert "文本:\n张三任职于公司" in prompt + + +@pytest.mark.asyncio +async def test_milvus_graph_service_configure_persists_updated_concurrency(): + kb = SimpleNamespace( + kb_type="milvus", + additional_params={ + "graph_build_config": { + "locked": True, + "extractor_type": "llm", + "extractor_options": {"model_spec": "test/model", "concurrency_count": 5}, + } + }, + ) + + class Repo: + async def get_by_id(self, db_id): + return kb + + async def update(self, db_id, data): + kb.additional_params = data["additional_params"] + return kb + + chunk_repo = SimpleNamespace( + count_by_db_id=AsyncMock(return_value=0), + count_graph_pending_by_db_id=AsyncMock(return_value=0), + count_graph_indexed_by_db_id=AsyncMock(return_value=0), + ) + service = MilvusGraphService(kb_repo=Repo(), chunk_repo=chunk_repo) + + await service.configure( + "kb_test", + extractor_type="llm", + extractor_options={"model_spec": "test/model", "concurrency_count": 9}, + created_by="user_1", + ) + status = await service.get_status("kb_test") + + assert status["config"]["extractor_options"]["concurrency_count"] == 9 + + def test_milvus_graph_service_writes_chunk_entity_and_relation(): tx = MagicMock() session = MagicMock() @@ -82,7 +146,7 @@ def test_milvus_graph_service_writes_chunk_entity_and_relation(): end_char_pos=8, ) - ent_ids, tags = service.write_chunk_graph( + entities, triples = service.write_chunk_graph( "kb_test", chunk, normalize_extraction_result( @@ -104,8 +168,9 @@ def test_milvus_graph_service_writes_chunk_entity_and_relation(): ), ) - assert ent_ids == ["张三", "公司"] - assert tags == ["Organization", "Person"] + assert [entity["name"] for entity in entities] == ["张三", "公司"] + assert {entity["label"] for entity in entities} == {"Person", "Organization"} + assert triples[0]["relation_type"] == "WORKS_AT" queries = [call.args[0] for call in tx.run.call_args_list] assert any("MERGE (c:Chunk:MilvusKB:`kb_test`" in query for query in queries) assert any("MERGE (e:Entity:MilvusKB:`kb_test`" in query for query in queries) @@ -135,4 +200,4 @@ def test_milvus_graph_service_get_stats_empty_db_id(): import asyncio result = asyncio.get_event_loop().run_until_complete(service.get_stats(db_id=None)) - assert result == {"total_nodes": 0, "total_edges": 0, "entity_types": []} \ No newline at end of file + assert result == {"total_nodes": 0, "total_edges": 0, "entity_types": []} diff --git a/backend/test/unit/knowledge/test_file_size_fallback.py b/backend/test/unit/knowledge/test_file_size_fallback.py index c2bc8715..eb32c874 100644 --- a/backend/test/unit/knowledge/test_file_size_fallback.py +++ b/backend/test/unit/knowledge/test_file_size_fallback.py @@ -1,3 +1,4 @@ +from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -164,3 +165,77 @@ class TestAddFileRecordSizeFallback: metadata = await kb.add_file_record("db1", item, params=params) assert metadata.get("size") is None + + +@pytest.mark.unit +@pytest.mark.asyncio +async def test_save_metadata_does_not_overwrite_existing_kb_config(): + from yuxi.knowledge.base import KnowledgeBase + + class TestKB(KnowledgeBase): + @property + def kb_type(self): + return "test" + + async def _create_kb_instance(self, db_id, config): + pass + + async def _initialize_kb_instance(self, instance): + pass + + async def index_file(self, db_id, file_id, operator_id=None): + return {} + + async def update_content(self, db_id, file_ids, params=None): + return [] + + async def aquery(self, query_text, db_id, **kwargs): + return [] + + def get_query_params_config(self, db_id, **kwargs): + return {"type": "test", "options": []} + + async def delete_file(self, db_id, file_id): + pass + + async def get_file_basic_info(self, db_id, file_id): + return {} + + async def get_file_content(self, db_id, file_id): + return {} + + async def get_file_info(self, db_id, file_id): + return {} + + class ExistingKbRepo: + def __init__(self): + self.created = [] + self.updated = [] + + async def get_by_id(self, db_id): + return SimpleNamespace(db_id=db_id) + + async def create(self, payload): + self.created.append(payload) + + async def update(self, db_id, data): + self.updated.append((db_id, data)) + + kb_repo = ExistingKbRepo() + kb = TestKB(work_dir="/tmp/test_kb") + kb.databases_meta["db1"] = { + "name": "Runtime name", + "description": "Runtime description", + "kb_type": "test", + "metadata": {"graph_build_config": {"extractor_options": {"concurrency_count": 5}}}, + } + + with ( + patch("yuxi.repositories.knowledge_base_repository.KnowledgeBaseRepository", return_value=kb_repo), + patch("yuxi.repositories.knowledge_file_repository.KnowledgeFileRepository", return_value=SimpleNamespace()), + patch("yuxi.repositories.evaluation_repository.EvaluationRepository", return_value=SimpleNamespace()), + ): + await kb._save_metadata() + + assert kb_repo.created == [] + assert kb_repo.updated == [] diff --git a/backend/test/unit/toolkits/test_kbs_tools.py b/backend/test/unit/toolkits/test_kbs_tools.py index 8605b2ef..21e12872 100644 --- a/backend/test/unit/toolkits/test_kbs_tools.py +++ b/backend/test/unit/toolkits/test_kbs_tools.py @@ -101,6 +101,7 @@ async def test_query_kb_returns_milvus_chunks_without_sandbox_paths(monkeypatch) "metadata": { "file_id": "file-1", "source": "auth-guide.pdf", + "resource_id": "db-1", }, } ] @@ -150,6 +151,7 @@ async def test_query_kb_allows_dify_knowledge_base(monkeypatch) -> None: "metadata": { "file_id": "dify-doc-1", "source": "Dify Doc", + "resource_id": "db-1", }, } ] @@ -219,6 +221,7 @@ async def test_query_kb_normalizes_file_metadata_for_open(monkeypatch) -> None: assert result[0]["metadata"] == { "file_id": "file-1", + "resource_id": "db-1", "chunk_id": "chunk-1", "chunk_index": 3, }