gmakstutis/LightRAG
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add chunk tracking cleanup to entity/relation deletion and creation

Browse source 
• Clean up chunk storage on delete
• Track chunks in create operations
• Normalize relation keys consistently
This commit is contained in:
yangdx 2025-10-26 17:06:16 +08:00
parent bf1897a67e
commit a3370b024d
1 changed files with 153 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

188
lightrag/utils_graph.py
View file

@ -12,15 +12,24 @@ from .base import StorageNameSpace
async def adelete_by_entity( async def adelete_by_entity(
chunk_entity_relation_graph, entities_vdb, relationships_vdb, entity_name: str chunk_entity_relation_graph,
entities_vdb,
relationships_vdb,
entity_name: str,
entity_chunks_storage=None,
relation_chunks_storage=None,
) -> DeletionResult: ) -> DeletionResult:
"""Asynchronously delete an entity and all its relationships. """Asynchronously delete an entity and all its relationships.
Also cleans up entity_chunks_storage and relation_chunks_storage to remove chunk tracking.
Args: Args:
chunk_entity_relation_graph: Graph storage instance chunk_entity_relation_graph: Graph storage instance
entities_vdb: Vector database storage for entities entities_vdb: Vector database storage for entities
relationships_vdb: Vector database storage for relationships relationships_vdb: Vector database storage for relationships
entity_name: Name of the entity to delete entity_name: Name of the entity to delete
entity_chunks_storage: Optional KV storage for tracking chunks that reference this entity
relation_chunks_storage: Optional KV storage for tracking chunks that reference relations
""" """
graph_db_lock = get_graph_db_lock(enable_logging=False) graph_db_lock = get_graph_db_lock(enable_logging=False)
# Use graph database lock to ensure atomic graph and vector db operations # Use graph database lock to ensure atomic graph and vector db operations
@ -39,14 +48,43 @@ async def adelete_by_entity(
edges = await chunk_entity_relation_graph.get_node_edges(entity_name) edges = await chunk_entity_relation_graph.get_node_edges(entity_name)
related_relations_count = len(edges) if edges else 0 related_relations_count = len(edges) if edges else 0
# Clean up chunk tracking storages before deletion
if entity_chunks_storage is not None:
# Delete entity's entry from entity_chunks_storage
await entity_chunks_storage.delete([entity_name])
logger.info(
f"Entity Delete: removed chunk tracking for `{entity_name}`"
)
if relation_chunks_storage is not None and edges:
# Delete all related relationships from relation_chunks_storage
from .utils import make_relation_chunk_key
relation_keys_to_delete = []
for src, tgt in edges:
# Normalize entity order for consistent key generation
normalized_src, normalized_tgt = sorted([src, tgt])
storage_key = make_relation_chunk_key(normalized_src, normalized_tgt)
relation_keys_to_delete.append(storage_key)
if relation_keys_to_delete:
await relation_chunks_storage.delete(relation_keys_to_delete)
logger.info(
f"Entity Delete: removed chunk tracking for {len(relation_keys_to_delete)} relations"
)
await entities_vdb.delete_entity(entity_name) await entities_vdb.delete_entity(entity_name)
await relationships_vdb.delete_entity_relation(entity_name) await relationships_vdb.delete_entity_relation(entity_name)
await chunk_entity_relation_graph.delete_node(entity_name) await chunk_entity_relation_graph.delete_node(entity_name)
message = f"Entity '{entity_name}' and its {related_relations_count} relationships have been deleted." message = f"Entity Delete: remove '{entity_name}' and its {related_relations_count} relations"
logger.info(message) logger.info(message)
await _delete_by_entity_done( await _delete_by_entity_done(
entities_vdb, relationships_vdb, chunk_entity_relation_graph entities_vdb,
relationships_vdb,
chunk_entity_relation_graph,
entity_chunks_storage,
relation_chunks_storage,
) )
return DeletionResult( return DeletionResult(
status="success", status="success",
@ -66,17 +104,23 @@ async def adelete_by_entity(
async def _delete_by_entity_done( async def _delete_by_entity_done(
entities_vdb, relationships_vdb, chunk_entity_relation_graph entities_vdb,
relationships_vdb,
chunk_entity_relation_graph,
entity_chunks_storage=None,
relation_chunks_storage=None,
) -> None: ) -> None:
"""Callback after entity deletion is complete, ensures updates are persisted""" """Callback after entity deletion is complete, ensures updates are persisted"""
storages = [entities_vdb, relationships_vdb, chunk_entity_relation_graph]
if entity_chunks_storage is not None:
storages.append(entity_chunks_storage)
if relation_chunks_storage is not None:
storages.append(relation_chunks_storage)
await asyncio.gather( await asyncio.gather(
*[ *[
cast(StorageNameSpace, storage_inst).index_done_callback() cast(StorageNameSpace, storage_inst).index_done_callback()
for storage_inst in [ # type: ignore for storage_inst in storages # type: ignore
entities_vdb,
relationships_vdb,
chunk_entity_relation_graph,
]
] ]
) )
@ -86,14 +130,18 @@ async def adelete_by_relation(
relationships_vdb, relationships_vdb,
source_entity: str, source_entity: str,
target_entity: str, target_entity: str,
relation_chunks_storage=None,
) -> DeletionResult: ) -> DeletionResult:
"""Asynchronously delete a relation between two entities. """Asynchronously delete a relation between two entities.
Also cleans up relation_chunks_storage to remove chunk tracking.
Args: Args:
chunk_entity_relation_graph: Graph storage instance chunk_entity_relation_graph: Graph storage instance
relationships_vdb: Vector database storage for relationships relationships_vdb: Vector database storage for relationships
source_entity: Name of the source entity source_entity: Name of the source entity
target_entity: Name of the target entity target_entity: Name of the target entity
relation_chunks_storage: Optional KV storage for tracking chunks that reference this relation
""" """
relation_str = f"{source_entity} -> {target_entity}" relation_str = f"{source_entity} -> {target_entity}"
graph_db_lock = get_graph_db_lock(enable_logging=False) graph_db_lock = get_graph_db_lock(enable_logging=False)
@ -118,6 +166,19 @@ async def adelete_by_relation(
status_code=404, status_code=404,
) )
# Clean up chunk tracking storage before deletion
if relation_chunks_storage is not None:
from .utils import make_relation_chunk_key
# Normalize entity order for consistent key generation
normalized_src, normalized_tgt = sorted([source_entity, target_entity])
storage_key = make_relation_chunk_key(normalized_src, normalized_tgt)
await relation_chunks_storage.delete([storage_key])
logger.info(
f"Relation Delete: removed chunk tracking for `{source_entity}`~`{target_entity}`"
)
# Delete relation from vector database # Delete relation from vector database
rel_ids_to_delete = [ rel_ids_to_delete = [
compute_mdhash_id(source_entity + target_entity, prefix="rel-"), compute_mdhash_id(source_entity + target_entity, prefix="rel-"),
@ -131,9 +192,11 @@ async def adelete_by_relation(
[(source_entity, target_entity)] [(source_entity, target_entity)]
) )
message = f"Successfully deleted relation from '{source_entity}' to '{target_entity}'" message = f"Relation Delete: `{source_entity}`~`{target_entity}` deleted successfully"
logger.info(message) logger.info(message)
await _delete_relation_done(relationships_vdb, chunk_entity_relation_graph) await _delete_relation_done(
relationships_vdb, chunk_entity_relation_graph, relation_chunks_storage
)
return DeletionResult( return DeletionResult(
status="success", status="success",
doc_id=relation_str, doc_id=relation_str,
@ -151,15 +214,18 @@ async def adelete_by_relation(
) )
async def _delete_relation_done(relationships_vdb, chunk_entity_relation_graph) -> None: async def _delete_relation_done(
relationships_vdb, chunk_entity_relation_graph, relation_chunks_storage=None
) -> None:
"""Callback after relation deletion is complete, ensures updates are persisted""" """Callback after relation deletion is complete, ensures updates are persisted"""
storages = [relationships_vdb, chunk_entity_relation_graph]
if relation_chunks_storage is not None:
storages.append(relation_chunks_storage)
await asyncio.gather( await asyncio.gather(
*[ *[
cast(StorageNameSpace, storage_inst).index_done_callback() cast(StorageNameSpace, storage_inst).index_done_callback()
for storage_inst in [ # type: ignore for storage_inst in storages # type: ignore
relationships_vdb,
chunk_entity_relation_graph,
]
] ]
) )
@ -506,7 +572,7 @@ async def aedit_entity(
relation_chunks_storage, relation_chunks_storage,
) )
logger.info(f"Entity '{entity_name}' successfully updated") logger.info(f"Entity Edit: `{entity_name}` successfully updated")
return await get_entity_info( return await get_entity_info(
chunk_entity_relation_graph, chunk_entity_relation_graph,
entities_vdb, entities_vdb,
@ -586,12 +652,14 @@ async def aedit_relation(
source_entity, target_entity source_entity, target_entity
) )
# Important: First delete the old relation record from the vector database # Important: First delete the old relation record from the vector database
old_relation_id = compute_mdhash_id( # Delete both permutations to handle relationships created before normalization
source_entity + target_entity, prefix="rel-" rel_ids_to_delete = [
) compute_mdhash_id(source_entity + target_entity, prefix="rel-"),
await relationships_vdb.delete([old_relation_id]) compute_mdhash_id(target_entity + source_entity, prefix="rel-"),
logger.info( ]
f"Deleted old relation record from vector database for relation {source_entity} -> {target_entity}" await relationships_vdb.delete(rel_ids_to_delete)
logger.debug(
f"Relation Delete: delete vdb for `{source_entity}`~`{target_entity}`"
) )
# 2. Update relation information in the graph # 2. Update relation information in the graph
@ -690,14 +758,8 @@ async def aedit_relation(
} }
) )
reason = (
"source_id changed"
if source_id_changed
else "initialized from graph"
)
logger.info( logger.info(
f"Updated relation_chunks_storage for '{source_entity}' -> '{target_entity}': " f"Relation Delete: update chunk tracking for `{source_entity}`~`{target_entity}`"
f"{len(updated_chunk_ids)} chunks ({reason})"
) )
# 5. Save changes # 5. Save changes
@ -706,7 +768,7 @@ async def aedit_relation(
) )
logger.info( logger.info(
f"Relation from '{source_entity}' to '{target_entity}' successfully updated" f"Relation Delete: `{source_entity}`~`{target_entity}`' successfully updated"
) )
return await get_relation_info( return await get_relation_info(
chunk_entity_relation_graph, chunk_entity_relation_graph,
@ -744,10 +806,13 @@ async def acreate_entity(
relationships_vdb, relationships_vdb,
entity_name: str, entity_name: str,
entity_data: dict[str, Any], entity_data: dict[str, Any],
entity_chunks_storage=None,
relation_chunks_storage=None,
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Asynchronously create a new entity. """Asynchronously create a new entity.
Creates a new entity in the knowledge graph and adds it to the vector database. Creates a new entity in the knowledge graph and adds it to the vector database.
Also synchronizes entity_chunks_storage to track chunk references.
Args: Args:
chunk_entity_relation_graph: Graph storage instance chunk_entity_relation_graph: Graph storage instance
@ -755,6 +820,8 @@ async def acreate_entity(
relationships_vdb: Vector database storage for relationships relationships_vdb: Vector database storage for relationships
entity_name: Name of the new entity entity_name: Name of the new entity
entity_data: Dictionary containing entity attributes, e.g. {"description": "description", "entity_type": "type"} entity_data: Dictionary containing entity attributes, e.g. {"description": "description", "entity_type": "type"}
entity_chunks_storage: Optional KV storage for tracking chunks that reference this entity
relation_chunks_storage: Optional KV storage for tracking chunks that reference relations
Returns: Returns:
Dictionary containing created entity information Dictionary containing created entity information
@ -805,12 +872,34 @@ async def acreate_entity(
# Update vector database # Update vector database
await entities_vdb.upsert(entity_data_for_vdb) await entities_vdb.upsert(entity_data_for_vdb)
# Update entity_chunks_storage to track chunk references
if entity_chunks_storage is not None:
source_id = node_data.get("source_id", "")
chunk_ids = [cid for cid in source_id.split(GRAPH_FIELD_SEP) if cid]
if chunk_ids:
await entity_chunks_storage.upsert(
{
entity_name: {
"chunk_ids": chunk_ids,
"count": len(chunk_ids),
}
}
)
logger.info(
f"Entity Create: tracked {len(chunk_ids)} chunks for `{entity_name}`"
)
# Save changes # Save changes
await _edit_entity_done( await _edit_entity_done(
entities_vdb, relationships_vdb, chunk_entity_relation_graph entities_vdb,
relationships_vdb,
chunk_entity_relation_graph,
entity_chunks_storage,
relation_chunks_storage,
) )
logger.info(f"Entity '{entity_name}' successfully created") logger.info(f"Entity Create: '{entity_name}' successfully created")
return await get_entity_info( return await get_entity_info(
chunk_entity_relation_graph, chunk_entity_relation_graph,
entities_vdb, entities_vdb,
@ -829,10 +918,12 @@ async def acreate_relation(
source_entity: str, source_entity: str,
target_entity: str, target_entity: str,
relation_data: dict[str, Any], relation_data: dict[str, Any],
relation_chunks_storage=None,
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Asynchronously create a new relation between entities. """Asynchronously create a new relation between entities.
Creates a new relation (edge) in the knowledge graph and adds it to the vector database. Creates a new relation (edge) in the knowledge graph and adds it to the vector database.
Also synchronizes relation_chunks_storage to track chunk references.
Args: Args:
chunk_entity_relation_graph: Graph storage instance chunk_entity_relation_graph: Graph storage instance
@ -841,6 +932,7 @@ async def acreate_relation(
source_entity: Name of the source entity source_entity: Name of the source entity
target_entity: Name of the target entity target_entity: Name of the target entity
relation_data: Dictionary containing relation attributes, e.g. {"description": "description", "keywords": "keywords"} relation_data: Dictionary containing relation attributes, e.g. {"description": "description", "keywords": "keywords"}
relation_chunks_storage: Optional KV storage for tracking chunks that reference this relation
Returns: Returns:
Dictionary containing created relation information Dictionary containing created relation information
@ -917,11 +1009,37 @@ async def acreate_relation(
# Update vector database # Update vector database
await relationships_vdb.upsert(relation_data_for_vdb) await relationships_vdb.upsert(relation_data_for_vdb)
# Update relation_chunks_storage to track chunk references
if relation_chunks_storage is not None:
from .utils import make_relation_chunk_key
# Normalize entity order for consistent key generation
normalized_src, normalized_tgt = sorted([source_entity, target_entity])
storage_key = make_relation_chunk_key(normalized_src, normalized_tgt)
source_id = edge_data.get("source_id", "")
chunk_ids = [cid for cid in source_id.split(GRAPH_FIELD_SEP) if cid]
if chunk_ids:
await relation_chunks_storage.upsert(
{
storage_key: {
"chunk_ids": chunk_ids,
"count": len(chunk_ids),
}
}
)
logger.info(
f"Relation Create: tracked {len(chunk_ids)} chunks for `{source_entity}`~`{target_entity}`"
)
# Save changes # Save changes
await _edit_relation_done(relationships_vdb, chunk_entity_relation_graph) await _edit_relation_done(
relationships_vdb, chunk_entity_relation_graph, relation_chunks_storage
)
logger.info( logger.info(
f"Relation from '{source_entity}' to '{target_entity}' successfully created" f"Relation Create: `{source_entity}`~`{target_entity}` successfully created"
) )
return await get_relation_info( return await get_relation_info(
chunk_entity_relation_graph, chunk_entity_relation_graph,