From 7d92a756d9378e6c9a00eb69ba5bced83508502c Mon Sep 17 00:00:00 2001
From: yangdx <gzdaniel@me.com>
Date: Sun, 26 Oct 2025 01:56:08 +0800
Subject: [PATCH] Fix entity/relation chunk cleanup and self-loop handling
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

• Delete empty chunk entries
• Clean up self-loop chunk data
• Move storage key calculation
• Add proper logging messages
• Prevent stale chunk references
---
 lightrag/utils_graph.py | 21 ++++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/lightrag/utils_graph.py b/lightrag/utils_graph.py
index 1ea60600..068efda3 100644
--- a/lightrag/utils_graph.py
+++ b/lightrag/utils_graph.py
@@ -417,6 +417,11 @@ async def aedit_entity(
                     logger.info(
                         f"Updated entity_chunks for '{entity_name}' with {len(chunk_ids)} chunk IDs"
                     )
+                else:
+                    await entity_chunks_storage.delete([entity_name])
+                    logger.info(
+                        f"Removed entity_chunks entry for '{entity_name}' due to empty source_id"
+                    )
 
             # 4. Save changes
             await _edit_entity_done(
@@ -540,8 +545,8 @@ async def aedit_relation(
             if relation_chunks_storage and "source_id" in updated_data:
                 source_ids = source_id.split(GRAPH_FIELD_SEP) if source_id else []
                 chunk_ids = [cid for cid in source_ids if cid]
+                storage_key = make_relation_chunk_key(source_entity, target_entity)
                 if chunk_ids:
-                    storage_key = make_relation_chunk_key(source_entity, target_entity)
                     await relation_chunks_storage.upsert(
                         {
                             storage_key: {
@@ -554,6 +559,11 @@ async def aedit_relation(
                     logger.info(
                         f"Updated relation_chunks for '{source_entity}' -> '{target_entity}' with {len(chunk_ids)} chunk IDs"
                     )
+                else:
+                    await relation_chunks_storage.delete([storage_key])
+                    logger.info(
+                        f"Removed relation_chunks entry for '{source_entity}' -> '{target_entity}' due to empty source_id"
+                    )
 
             # 4. Save changes
             await _edit_relation_done(relationships_vdb, chunk_entity_relation_graph)
@@ -1084,13 +1094,14 @@ async def amerge_entities(
                     new_src = target_entity if src in source_entities else src
                     new_tgt = target_entity if tgt in source_entities else tgt
 
-                    # Skip self-loops
-                    if new_src == new_tgt:
-                        continue
-
                     old_key = make_relation_chunk_key(src, tgt)
                     new_key = make_relation_chunk_key(new_src, new_tgt)
 
+                    # Skip self-loops but ensure stale chunk entries are removed
+                    if new_src == new_tgt:
+                        old_relation_keys_to_delete.append(old_key)
+                        continue
+
                     if old_key != new_key:
                         old_chunks_data = await relation_chunks_storage.get_by_id(
                             old_key