138 lines
4.9 KiB
Python
138 lines
4.9 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
|
|
from yuxi.knowledge.graphs.extractors import normalize_extraction_result
|
|
from yuxi.knowledge.graphs.milvus_graph_service import MilvusGraphService
|
|
|
|
|
|
def test_normalize_extraction_result_defaults_and_validates_refs():
|
|
result = normalize_extraction_result(
|
|
{
|
|
"entities": [{"text": "张三"}, {"text": "公司"}],
|
|
"relations": [{"source": "张三", "target": "公司", "text": "任职于"}],
|
|
},
|
|
"llm",
|
|
)
|
|
|
|
assert result["entities"][0]["label"] == "Entity"
|
|
assert result["relations"][0]["label"] == "RELATED_TO"
|
|
assert result["relations"][0]["source"] == {"text": "张三", "label": "Entity", "attributes": []}
|
|
assert result["metadata"] == {"extractor_type": "llm", "schema_version": 1}
|
|
|
|
|
|
def test_normalize_extraction_result_accepts_llm_nested_relation_entities():
|
|
result = normalize_extraction_result(
|
|
{
|
|
"relations": [
|
|
{
|
|
"source": {
|
|
"text": "张三",
|
|
"label": "Person",
|
|
"attributes": [{"text": "工程师", "label": "Occupation"}],
|
|
},
|
|
"target": {"text": "公司", "label": "Organization"},
|
|
"text": "任职于",
|
|
"label": "WORKS_AT",
|
|
}
|
|
]
|
|
},
|
|
"llm",
|
|
)
|
|
|
|
assert result["entities"] == [
|
|
{"text": "张三", "label": "Person", "attributes": [{"text": "工程师", "label": "Occupation"}]},
|
|
{"text": "公司", "label": "Organization", "attributes": []},
|
|
]
|
|
assert result["relations"][0]["source"]["attributes"] == [{"text": "工程师", "label": "Occupation"}]
|
|
assert result["relations"][0]["target"] == {"text": "公司", "label": "Organization", "attributes": []}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"payload",
|
|
[
|
|
{"entities": [{"text": "张三"}], "relations": [{"source": "张三", "target": "不存在", "text": "关系"}]},
|
|
{"entities": [{"text": ""}], "relations": []},
|
|
],
|
|
)
|
|
def test_normalize_extraction_result_rejects_invalid_payload(payload):
|
|
with pytest.raises(ValueError):
|
|
normalize_extraction_result(payload, "llm")
|
|
|
|
|
|
def test_milvus_graph_service_writes_chunk_entity_and_relation():
|
|
tx = MagicMock()
|
|
session = MagicMock()
|
|
session.__enter__.return_value = session
|
|
session.execute_write.side_effect = lambda func: func(tx)
|
|
driver = MagicMock()
|
|
driver.session.return_value = session
|
|
connection = SimpleNamespace(driver=driver)
|
|
service = MilvusGraphService(neo4j_connection=connection)
|
|
chunk = SimpleNamespace(
|
|
chunk_id="chunk_1",
|
|
file_id="file_1",
|
|
db_id="kb_test",
|
|
chunk_index=1,
|
|
content="张三任职于公司",
|
|
start_char_pos=0,
|
|
end_char_pos=8,
|
|
)
|
|
|
|
ent_ids, tags = service.write_chunk_graph(
|
|
"kb_test",
|
|
chunk,
|
|
normalize_extraction_result(
|
|
{
|
|
"relations": [
|
|
{
|
|
"source": {
|
|
"text": "张三",
|
|
"label": "Person",
|
|
"attributes": [{"text": "工程师", "label": "Occupation"}],
|
|
},
|
|
"target": {"text": "公司", "label": "Organization"},
|
|
"text": "任职于",
|
|
"label": "WORKS_AT",
|
|
}
|
|
],
|
|
},
|
|
"llm",
|
|
),
|
|
)
|
|
|
|
assert ent_ids == ["张三", "公司"]
|
|
assert tags == ["Organization", "Person"]
|
|
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)
|
|
assert any("MERGE (source)-[r:RELATION" in query for query in queries)
|
|
entity_call = next(call for call in tx.run.call_args_list if "MERGE (e:Entity" in call.args[0])
|
|
assert entity_call.kwargs["attributes"] == '[{"text": "工程师", "label": "Occupation"}]'
|
|
|
|
|
|
def test_milvus_graph_service_query_nodes_empty_db_id():
|
|
service = MilvusGraphService()
|
|
import asyncio
|
|
|
|
result = asyncio.get_event_loop().run_until_complete(service.query_nodes(db_id=None, keyword="test"))
|
|
assert result == {"nodes": [], "edges": []}
|
|
|
|
|
|
def test_milvus_graph_service_get_labels_empty_db_id():
|
|
service = MilvusGraphService()
|
|
import asyncio
|
|
|
|
result = asyncio.get_event_loop().run_until_complete(service.get_labels(db_id=None))
|
|
assert result == []
|
|
|
|
|
|
def test_milvus_graph_service_get_stats_empty_db_id():
|
|
service = MilvusGraphService()
|
|
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": []} |