gardel/lightrag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix edit feature

Browse Source
This commit is contained in:
zrguo
2025-03-20 22:00:11 +08:00
parent 89d4a13cf5
commit d817defbc4
1 changed files with 29 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files

64
lightrag/lightrag.py
View File

@@ -1947,6 +1947,8 @@ class LightRAG:
# 2. Update entity information in the graph # 2. Update entity information in the graph
new_node_data = {**node_data, **updated_data} new_node_data = {**node_data, **updated_data}
new_node_data["entity_id"] = new_entity_name
if "entity_name" in new_node_data: if "entity_name" in new_node_data:
del new_node_data[ del new_node_data[
"entity_name" "entity_name"
@@ -1963,7 +1965,7 @@ class LightRAG:
# Store relationships that need to be updated # Store relationships that need to be updated
relations_to_update = [] relations_to_update = []
relations_to_delete = []
# Get all edges related to the original entity # Get all edges related to the original entity
edges = await self.chunk_entity_relation_graph.get_node_edges( edges = await self.chunk_entity_relation_graph.get_node_edges(
entity_name entity_name
@@ -1975,6 +1977,12 @@ class LightRAG:
source, target source, target
) )
if edge_data: if edge_data:
relations_to_delete.append(
compute_mdhash_id(source + target, prefix="rel-")
)
relations_to_delete.append(
compute_mdhash_id(target + source, prefix="rel-")
)
if source == entity_name: if source == entity_name:
await self.chunk_entity_relation_graph.upsert_edge( await self.chunk_entity_relation_graph.upsert_edge(
new_entity_name, target, edge_data new_entity_name, target, edge_data
@@ -2000,6 +2008,12 @@ class LightRAG:
f"Deleted old entity '{entity_name}' and its vector embedding from database" f"Deleted old entity '{entity_name}' and its vector embedding from database"
) )
# Delete old relation records from vector database
await self.relationships_vdb.delete(relations_to_delete)
logger.info(
f"Deleted {len(relations_to_delete)} relation records for entity '{entity_name}' from vector database"
)
# Update relationship vector representations # Update relationship vector representations
for src, tgt, edge_data in relations_to_update: for src, tgt, edge_data in relations_to_update:
description = edge_data.get("description", "") description = edge_data.get("description", "")
@@ -2498,39 +2512,21 @@ class LightRAG:
# 4. Get all relationships of the source entities # 4. Get all relationships of the source entities
all_relations = [] all_relations = []
for entity_name in source_entities: for entity_name in source_entities:
# Get all relationships where this entity is the source # Get all relationships of the source entities
outgoing_edges = await self.chunk_entity_relation_graph.get_node_edges( edges = await self.chunk_entity_relation_graph.get_node_edges(
entity_name entity_name
) )
if outgoing_edges: if edges:
for src, tgt in outgoing_edges: for src, tgt in edges:
# Ensure src is the current entity # Ensure src is the current entity
if src == entity_name: if src == entity_name:
edge_data = await self.chunk_entity_relation_graph.get_edge( edge_data = await self.chunk_entity_relation_graph.get_edge(
src, tgt src, tgt
) )
all_relations.append(("outgoing", src, tgt, edge_data)) all_relations.append((src, tgt, edge_data))
# Get all relationships where this entity is the target
incoming_edges = []
all_labels = await self.chunk_entity_relation_graph.get_all_labels()
for label in all_labels:
if label == entity_name:
continue
node_edges = await self.chunk_entity_relation_graph.get_node_edges(
label
)
for src, tgt in node_edges or []:
if tgt == entity_name:
incoming_edges.append((src, tgt))
for src, tgt in incoming_edges:
edge_data = await self.chunk_entity_relation_graph.get_edge(
src, tgt
)
all_relations.append(("incoming", src, tgt, edge_data))
# 5. Create or update the target entity # 5. Create or update the target entity
merged_entity_data["entity_id"] = target_entity
if not target_exists: if not target_exists:
await self.chunk_entity_relation_graph.upsert_node( await self.chunk_entity_relation_graph.upsert_node(
target_entity, merged_entity_data target_entity, merged_entity_data
@@ -2544,8 +2540,11 @@ class LightRAG:
# 6. Recreate all relationships, pointing to the target entity # 6. Recreate all relationships, pointing to the target entity
relation_updates = {} # Track relationships that need to be merged relation_updates = {} # Track relationships that need to be merged
relations_to_delete = []
for rel_type, src, tgt, edge_data in all_relations: for src, tgt, edge_data in all_relations:
relations_to_delete.append(compute_mdhash_id(src + tgt, prefix="rel-"))
relations_to_delete.append(compute_mdhash_id(tgt + src, prefix="rel-"))
new_src = target_entity if src in source_entities else src new_src = target_entity if src in source_entities else src
new_tgt = target_entity if tgt in source_entities else tgt new_tgt = target_entity if tgt in source_entities else tgt
@@ -2652,17 +2651,10 @@ class LightRAG:
entity_id = compute_mdhash_id(entity_name, prefix="ent-") entity_id = compute_mdhash_id(entity_name, prefix="ent-")
await self.entities_vdb.delete([entity_id]) await self.entities_vdb.delete([entity_id])
# Also ensure any relationships specific to this entity are deleted from vector DB # Delete relationships records from vector database
# This is a safety check, as these should have been transformed to the target entity already await self.relationships_vdb.delete(relations_to_delete)
entity_relation_prefix = compute_mdhash_id(entity_name, prefix="rel-")
relations_with_entity = await self.relationships_vdb.search_by_prefix(
entity_relation_prefix
)
if relations_with_entity:
relation_ids = [r["id"] for r in relations_with_entity]
await self.relationships_vdb.delete(relation_ids)
logger.info( logger.info(
f"Deleted {len(relation_ids)} relation records for entity '{entity_name}' from vector database" f"Deleted {len(relations_to_delete)} relation records for entity '{entity_name}' from vector database"
) )
logger.info( logger.info(