chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,403 @@
|
||||
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
|
||||
Reference in New Issue
Block a user