zrguo commited on
Commit
5a851ea
·
unverified ·
2 Parent(s): 06b45ea 1906884

Merge pull request #1006 from HKUDS/dev

Browse files

[Bug]Fix the issue where editing entity names does not properly update the corresponding relationship in the vdb.

Files changed (1) hide show
  1. lightrag/lightrag.py +39 -1
lightrag/lightrag.py CHANGED
@@ -1981,6 +1981,9 @@ class LightRAG:
1981
  new_entity_name, new_node_data
1982
  )
1983
 
 
 
 
1984
  # Get all edges related to the original entity
1985
  edges = await self.chunk_entity_relation_graph.get_node_edges(
1986
  entity_name
@@ -1996,10 +1999,16 @@ class LightRAG:
1996
  await self.chunk_entity_relation_graph.upsert_edge(
1997
  new_entity_name, target, edge_data
1998
  )
 
 
 
1999
  else: # target == entity_name
2000
  await self.chunk_entity_relation_graph.upsert_edge(
2001
  source, new_entity_name, edge_data
2002
  )
 
 
 
2003
 
2004
  # Delete old entity
2005
  await self.chunk_entity_relation_graph.delete_node(entity_name)
@@ -2008,6 +2017,35 @@ class LightRAG:
2008
  old_entity_id = compute_mdhash_id(entity_name, prefix="ent-")
2009
  await self.entities_vdb.delete([old_entity_id])
2010
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2011
  # Update working entity name to new name
2012
  entity_name = new_entity_name
2013
  else:
@@ -2118,7 +2156,7 @@ class LightRAG:
2118
  weight = float(new_edge_data.get("weight", 1.0))
2119
 
2120
  # Create content for embedding
2121
- content = f"{keywords}\t{source_entity}\n{target_entity}\n{description}"
2122
 
2123
  # Calculate relation ID
2124
  relation_id = compute_mdhash_id(
 
1981
  new_entity_name, new_node_data
1982
  )
1983
 
1984
+ # Store relationships that need to be updated
1985
+ relations_to_update = []
1986
+
1987
  # Get all edges related to the original entity
1988
  edges = await self.chunk_entity_relation_graph.get_node_edges(
1989
  entity_name
 
1999
  await self.chunk_entity_relation_graph.upsert_edge(
2000
  new_entity_name, target, edge_data
2001
  )
2002
+ relations_to_update.append(
2003
+ (new_entity_name, target, edge_data)
2004
+ )
2005
  else: # target == entity_name
2006
  await self.chunk_entity_relation_graph.upsert_edge(
2007
  source, new_entity_name, edge_data
2008
  )
2009
+ relations_to_update.append(
2010
+ (source, new_entity_name, edge_data)
2011
+ )
2012
 
2013
  # Delete old entity
2014
  await self.chunk_entity_relation_graph.delete_node(entity_name)
 
2017
  old_entity_id = compute_mdhash_id(entity_name, prefix="ent-")
2018
  await self.entities_vdb.delete([old_entity_id])
2019
 
2020
+ # Update relationship vector representations
2021
+ for src, tgt, edge_data in relations_to_update:
2022
+ description = edge_data.get("description", "")
2023
+ keywords = edge_data.get("keywords", "")
2024
+ source_id = edge_data.get("source_id", "")
2025
+ weight = float(edge_data.get("weight", 1.0))
2026
+
2027
+ # Create new content for embedding
2028
+ content = f"{src}\t{tgt}\n{keywords}\n{description}"
2029
+
2030
+ # Calculate relationship ID
2031
+ relation_id = compute_mdhash_id(src + tgt, prefix="rel-")
2032
+
2033
+ # Prepare data for vector database update
2034
+ relation_data = {
2035
+ relation_id: {
2036
+ "content": content,
2037
+ "src_id": src,
2038
+ "tgt_id": tgt,
2039
+ "source_id": source_id,
2040
+ "description": description,
2041
+ "keywords": keywords,
2042
+ "weight": weight,
2043
+ }
2044
+ }
2045
+
2046
+ # Update vector database
2047
+ await self.relationships_vdb.upsert(relation_data)
2048
+
2049
  # Update working entity name to new name
2050
  entity_name = new_entity_name
2051
  else:
 
2156
  weight = float(new_edge_data.get("weight", 1.0))
2157
 
2158
  # Create content for embedding
2159
+ content = f"{source_entity}\t{target_entity}\n{keywords}\n{description}"
2160
 
2161
  # Calculate relation ID
2162
  relation_id = compute_mdhash_id(