wehub-resource-sync
404 lines
12 KiB
Python
404 lines
12 KiB
Python
import pytest
|
|
|
|
from lightrag.constants import SOURCE_IDS_LIMIT_METHOD_KEEP
|
|
from lightrag.constants import GRAPH_FIELD_SEP
|
|
from lightrag.operate import (
|
|
_handle_single_entity_extraction,
|
|
_merge_nodes_then_upsert,
|
|
_normalize_text_extraction_record_attributes,
|
|
_handle_single_relationship_extraction,
|
|
)
|
|
from lightrag import utils_graph
|
|
from lightrag.utils import VectorStorageConsistencyError
|
|
|
|
|
|
class DummyGraphStorage:
|
|
def __init__(self, node=None):
|
|
self.node = node
|
|
self.upserted_nodes = []
|
|
|
|
async def get_node(self, node_id):
|
|
return self.node
|
|
|
|
async def upsert_node(self, node_id, node_data):
|
|
self.upserted_nodes.append((node_id, node_data))
|
|
self.node = dict(node_data)
|
|
|
|
|
|
class DummyVectorStorage:
|
|
def __init__(self):
|
|
self.global_config = {"workspace": "test"}
|
|
self.upserts = []
|
|
self.deletes = []
|
|
|
|
async def upsert(self, data):
|
|
self.upserts.append(data)
|
|
return None
|
|
|
|
async def delete(self, ids):
|
|
self.deletes.append(ids)
|
|
return None
|
|
|
|
async def get_by_id(self, id_):
|
|
return None
|
|
|
|
async def index_done_callback(self):
|
|
return True
|
|
|
|
|
|
class DummyAsyncContext:
|
|
async def __aenter__(self):
|
|
return None
|
|
|
|
async def __aexit__(self, exc_type, exc, tb):
|
|
return False
|
|
|
|
|
|
class DummyMergeGraphStorage:
|
|
def __init__(self):
|
|
self.nodes = {
|
|
"Canonical": {
|
|
"entity_id": "Canonical",
|
|
"description": "canonical desc",
|
|
"entity_type": "ORG",
|
|
"source_id": "chunk-1",
|
|
"file_path": "canonical.md",
|
|
},
|
|
"Alias": {
|
|
"entity_id": "Alias",
|
|
"description": "alias desc",
|
|
"entity_type": "ORG",
|
|
"source_id": "chunk-2",
|
|
"file_path": "alias.md",
|
|
},
|
|
"Neighbor": {
|
|
"entity_id": "Neighbor",
|
|
"description": "neighbor desc",
|
|
"entity_type": "ORG",
|
|
"source_id": "chunk-3",
|
|
"file_path": "neighbor.md",
|
|
},
|
|
}
|
|
self.edges = {
|
|
("Alias", "Neighbor"): {
|
|
"description": "rel desc",
|
|
"keywords": "alias",
|
|
"source_id": "chunk-rel",
|
|
"weight": 1.0,
|
|
"file_path": "rel.md",
|
|
}
|
|
}
|
|
|
|
async def has_node(self, node_id):
|
|
return node_id in self.nodes
|
|
|
|
async def get_node(self, node_id):
|
|
return self.nodes[node_id]
|
|
|
|
async def upsert_node(self, node_id, node_data):
|
|
self.nodes[node_id] = dict(node_data)
|
|
|
|
async def get_node_edges(self, node_id):
|
|
results = []
|
|
for src, tgt in self.edges:
|
|
if src == node_id or tgt == node_id:
|
|
results.append((src, tgt))
|
|
return results
|
|
|
|
async def get_edge(self, src, tgt):
|
|
return self.edges.get((src, tgt)) or self.edges.get((tgt, src))
|
|
|
|
async def upsert_edge(self, src, tgt, edge_data):
|
|
self.edges[(src, tgt)] = dict(edge_data)
|
|
|
|
async def delete_node(self, node_id):
|
|
self.nodes.pop(node_id, None)
|
|
self.edges = {
|
|
(src, tgt): data
|
|
for (src, tgt), data in self.edges.items()
|
|
if src != node_id and tgt != node_id
|
|
}
|
|
|
|
async def index_done_callback(self):
|
|
return True
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_merge_nodes_then_upsert_handles_missing_legacy_description():
|
|
graph = DummyGraphStorage(node={"source_id": "chunk-1"})
|
|
global_config = {
|
|
"source_ids_limit_method": SOURCE_IDS_LIMIT_METHOD_KEEP,
|
|
"max_source_ids_per_entity": 20,
|
|
}
|
|
|
|
result = await _merge_nodes_then_upsert(
|
|
entity_name="LegacyEntity",
|
|
nodes_data=[],
|
|
knowledge_graph_inst=graph,
|
|
entity_vdb=None,
|
|
global_config=global_config,
|
|
)
|
|
|
|
assert result["description"] == "Entity LegacyEntity"
|
|
assert graph.upserted_nodes[-1][1]["description"] == "Entity LegacyEntity"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acreate_entity_rejects_empty_description():
|
|
with pytest.raises(ValueError, match="description cannot be empty"):
|
|
await utils_graph.acreate_entity(
|
|
chunk_entity_relation_graph=None,
|
|
entities_vdb=None,
|
|
relationships_vdb=None,
|
|
entity_name="EntityA",
|
|
entity_data={"description": " "},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_acreate_relation_rejects_empty_description():
|
|
with pytest.raises(ValueError, match="description cannot be empty"):
|
|
await utils_graph.acreate_relation(
|
|
chunk_entity_relation_graph=None,
|
|
entities_vdb=None,
|
|
relationships_vdb=None,
|
|
source_entity="A",
|
|
target_entity="B",
|
|
relation_data={"description": ""},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aedit_entity_rejects_empty_description():
|
|
with pytest.raises(ValueError, match="description cannot be empty"):
|
|
await utils_graph.aedit_entity(
|
|
chunk_entity_relation_graph=None,
|
|
entities_vdb=None,
|
|
relationships_vdb=None,
|
|
entity_name="EntityA",
|
|
updated_data={"description": None},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aedit_relation_rejects_empty_description():
|
|
with pytest.raises(ValueError, match="description cannot be empty"):
|
|
await utils_graph.aedit_relation(
|
|
chunk_entity_relation_graph=None,
|
|
entities_vdb=None,
|
|
relationships_vdb=None,
|
|
source_entity="A",
|
|
target_entity="B",
|
|
updated_data={"description": " "},
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aedit_entity_allows_updates_without_description(monkeypatch):
|
|
async def fake_edit_impl(*args, **kwargs):
|
|
return {"entity_name": "EntityA", "description": "kept", "source_id": "chunk-1"}
|
|
|
|
monkeypatch.setattr(
|
|
utils_graph, "get_storage_keyed_lock", lambda *a, **k: DummyAsyncContext()
|
|
)
|
|
monkeypatch.setattr(utils_graph, "_edit_entity_impl", fake_edit_impl)
|
|
|
|
result = await utils_graph.aedit_entity(
|
|
chunk_entity_relation_graph=None,
|
|
entities_vdb=DummyVectorStorage(),
|
|
relationships_vdb=DummyVectorStorage(),
|
|
entity_name="EntityA",
|
|
updated_data={"entity_type": "ORG"},
|
|
)
|
|
|
|
assert result["operation_summary"]["operation_status"] == "success"
|
|
|
|
|
|
def test_handle_single_relationship_extraction_ignores_empty_description():
|
|
relation = _handle_single_relationship_extraction(
|
|
["relation", "Alice", "Bob", "works_with", " "],
|
|
chunk_key="chunk-1",
|
|
timestamp=1,
|
|
)
|
|
|
|
assert relation is None
|
|
|
|
|
|
def test_mis_prefixed_relation_row_is_recovered():
|
|
record = _normalize_text_extraction_record_attributes(
|
|
["entity", "Alice", "Acme Corp", "founded", "Alice founded Acme Corp."],
|
|
chunk_key="chunk-1",
|
|
)
|
|
|
|
relation = _handle_single_relationship_extraction(
|
|
record,
|
|
chunk_key="chunk-1",
|
|
timestamp=1,
|
|
)
|
|
|
|
assert relation is not None
|
|
assert relation["src_id"] == "Alice"
|
|
assert relation["tgt_id"] == "Acme Corp"
|
|
|
|
|
|
def test_four_part_entity_row_remains_entity():
|
|
record = _normalize_text_extraction_record_attributes(
|
|
["entity", "Alice", "Person", "Alice is the founder of Acme Corp."],
|
|
chunk_key="chunk-1",
|
|
)
|
|
|
|
entity = _handle_single_entity_extraction(
|
|
record,
|
|
chunk_key="chunk-1",
|
|
timestamp=1,
|
|
)
|
|
|
|
assert entity is not None
|
|
assert entity["entity_name"] == "Alice"
|
|
|
|
|
|
def test_malformed_recovered_relation_still_fails():
|
|
record = _normalize_text_extraction_record_attributes(
|
|
["entity", "Alice", "Acme Corp", "founded", " "],
|
|
chunk_key="chunk-1",
|
|
)
|
|
|
|
relation = _handle_single_relationship_extraction(
|
|
record,
|
|
chunk_key="chunk-1",
|
|
timestamp=1,
|
|
)
|
|
|
|
assert relation is None
|
|
|
|
|
|
def test_unrelated_five_part_prefix_remains_invalid():
|
|
record = _normalize_text_extraction_record_attributes(
|
|
["edge", "Alice", "Acme Corp", "founded", "Alice founded Acme Corp."],
|
|
chunk_key="chunk-1",
|
|
)
|
|
|
|
relation = _handle_single_relationship_extraction(
|
|
record,
|
|
chunk_key="chunk-1",
|
|
timestamp=1,
|
|
)
|
|
|
|
assert relation is None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_merge_entities_preserves_file_path_in_vector_updates(monkeypatch):
|
|
graph = DummyMergeGraphStorage()
|
|
entities_vdb = DummyVectorStorage()
|
|
relationships_vdb = DummyVectorStorage()
|
|
|
|
async def fake_get_entity_info(*args, **kwargs):
|
|
return {"entity_name": "Canonical"}
|
|
|
|
monkeypatch.setattr(utils_graph, "get_entity_info", fake_get_entity_info)
|
|
|
|
await utils_graph._merge_entities_impl(
|
|
chunk_entity_relation_graph=graph,
|
|
entities_vdb=entities_vdb,
|
|
relationships_vdb=relationships_vdb,
|
|
source_entities=["Alias", "Canonical"],
|
|
target_entity="Canonical",
|
|
)
|
|
|
|
relationship_payload = relationships_vdb.upserts[-1]
|
|
entity_payload = entities_vdb.upserts[-1]
|
|
|
|
assert next(iter(relationship_payload.values()))["file_path"] == "rel.md"
|
|
assert set(
|
|
next(iter(entity_payload.values()))["file_path"].split(GRAPH_FIELD_SEP)
|
|
) == {
|
|
"alias.md",
|
|
"canonical.md",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_merge_response_is_built_from_graph_without_reading_vdb():
|
|
# The merge response is built from graph data only. Reading the vector
|
|
# store is redundant (the graph is authoritative) and previously leaked a
|
|
# non-JSON-serializable embedding (pgvector -> numpy) into the API
|
|
# response, causing a 500 on /graph/entity/edit with all-PostgreSQL.
|
|
graph = DummyMergeGraphStorage()
|
|
entities_vdb = DummyVectorStorage()
|
|
relationships_vdb = DummyVectorStorage()
|
|
|
|
vdb_reads = []
|
|
inner_get_by_id = entities_vdb.get_by_id
|
|
|
|
async def _tracked_get_by_id(id_):
|
|
vdb_reads.append(id_)
|
|
return await inner_get_by_id(id_)
|
|
|
|
entities_vdb.get_by_id = _tracked_get_by_id
|
|
|
|
result = await utils_graph._merge_entities_impl(
|
|
chunk_entity_relation_graph=graph,
|
|
entities_vdb=entities_vdb,
|
|
relationships_vdb=relationships_vdb,
|
|
source_entities=["Alias", "Canonical"],
|
|
target_entity="Canonical",
|
|
)
|
|
|
|
assert "graph_data" in result
|
|
assert "vector_data" not in result
|
|
# The response did not read the entity vector store.
|
|
assert vdb_reads == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_aedit_entity_merge_propagates_consistency_error(monkeypatch):
|
|
"""rename-with-merge must surface VectorStorageConsistencyError, not swallow it.
|
|
|
|
The graph is updated during the merge, then the (deferred) vector-store
|
|
flush fails. aedit_entity used to fold every merge exception into a
|
|
partial-success summary and return normally, so the /graph/entity/edit
|
|
route reported HTTP 200 "success" and hid the consistency failure that the
|
|
fail-loud merge path exists to surface. It must re-raise instead.
|
|
"""
|
|
graph = DummyMergeGraphStorage()
|
|
entities_vdb = DummyVectorStorage()
|
|
relationships_vdb = DummyVectorStorage()
|
|
|
|
# Deferred-embedding flush failure: upsert buffers, index_done_callback raises.
|
|
async def _boom():
|
|
raise RuntimeError("embedder down")
|
|
|
|
relationships_vdb.index_done_callback = _boom
|
|
|
|
# Single-attempt VDB wrapper (no retry delays) + a no-op storage lock.
|
|
async def _single_attempt(operation, **kwargs):
|
|
await operation()
|
|
|
|
monkeypatch.setattr(
|
|
utils_graph, "safe_vdb_operation_with_exception", _single_attempt
|
|
)
|
|
monkeypatch.setattr(
|
|
utils_graph, "get_storage_keyed_lock", lambda *a, **k: DummyAsyncContext()
|
|
)
|
|
|
|
async def fake_get_entity_info(*args, **kwargs):
|
|
return {"entity_name": "Canonical"}
|
|
|
|
monkeypatch.setattr(utils_graph, "get_entity_info", fake_get_entity_info)
|
|
|
|
with pytest.raises(VectorStorageConsistencyError) as excinfo:
|
|
await utils_graph.aedit_entity(
|
|
chunk_entity_relation_graph=graph,
|
|
entities_vdb=entities_vdb,
|
|
relationships_vdb=relationships_vdb,
|
|
entity_name="Alias",
|
|
updated_data={"entity_name": "Canonical"},
|
|
allow_rename=True,
|
|
allow_merge=True,
|
|
)
|
|
|
|
assert "lightrag-rebuild-vdb" in str(excinfo.value)
|
|
# Fail-loud happened before source deletion: 'Alias' is still in the graph.
|
|
assert "Alias" in graph.nodes
|