Merge pull request HKUDS#2265 from danielaskdd/edit-kg-new

Refactor: Enhance KG Editing with Chunk Tracking

Author: danielaskdd

lightrag/api/routers/graph_routes.py

@@ -299,6 +299,7 @@ async def create_entity(request: EntityCreateRequest):
             entity_data (dict): Entity properties including:
                 - description (str): Textual description of the entity
                 - entity_type (str): Category/type of the entity (e.g., PERSON, ORGANIZATION, LOCATION)
+                - source_id (str): Related chunk_id from which the description originates
                 - Additional custom properties as needed
 
         Response Schema:
@@ -309,6 +310,7 @@ async def create_entity(request: EntityCreateRequest):
                     "entity_name": "Tesla",
                     "description": "Electric vehicle manufacturer",
                     "entity_type": "ORGANIZATION",
+                    "source_id": "chunk-123<SEP>chunk-456"
                     ... (other entity properties)
                 }
             }
@@ -361,10 +363,11 @@ async def create_relation(request: RelationCreateRequest):
         """
         Create a new relationship between two entities in the knowledge graph
 
-        This endpoint establishes a directed relationship between two existing entities.
-        Both the source and target entities must already exist in the knowledge graph.
-        The system automatically generates vector embeddings for the relationship to
-        enable semantic search and graph traversal.
+        This endpoint establishes an undirected relationship between two existing entities.
+        The provided source/target order is accepted for convenience, but the backend
+        stored edge is undirected and may be returned with the entities swapped.
+        Both entities must already exist in the knowledge graph. The system automatically
+        generates vector embeddings for the relationship to enable semantic search and graph traversal.
 
         Prerequisites:
             - Both source_entity and target_entity must exist in the knowledge graph
@@ -376,6 +379,7 @@ async def create_relation(request: RelationCreateRequest):
             relation_data (dict): Relationship properties including:
                 - description (str): Textual description of the relationship
                 - keywords (str): Comma-separated keywords describing the relationship type
+                - source_id (str): Related chunk_id from which the description originates
                 - weight (float): Relationship strength/importance (default: 1.0)
                 - Additional custom properties as needed
 
@@ -388,6 +392,7 @@ async def create_relation(request: RelationCreateRequest):
                     "tgt_id": "Tesla",
                     "description": "Elon Musk is the CEO of Tesla",
                     "keywords": "CEO, founder",
+                    "source_id": "chunk-123<SEP>chunk-456"
                     "weight": 1.0,
                     ... (other relationship properties)
                 }

lightrag/lightrag.py

@@ -3582,6 +3582,7 @@ async def aedit_entity(
         """Asynchronously edit entity information.
 
         Updates entity information in the knowledge graph and re-embeds the entity in the vector database.
+        Also synchronizes entity_chunks_storage and relation_chunks_storage to track chunk references.
 
         Args:
             entity_name: Name of the entity to edit
@@ -3600,6 +3601,8 @@ async def aedit_entity(
             entity_name,
             updated_data,
             allow_rename,
+            self.entity_chunks,
+            self.relation_chunks,
         )
 
     def edit_entity(
@@ -3616,6 +3619,7 @@ async def aedit_relation(
         """Asynchronously edit relation information.
 
         Updates relation (edge) information in the knowledge graph and re-embeds the relation in the vector database.
+        Also synchronizes the relation_chunks_storage to track which chunks reference this relation.
 
         Args:
             source_entity: Name of the source entity
@@ -3634,6 +3638,7 @@ async def aedit_relation(
             source_entity,
             target_entity,
             updated_data,
+            self.relation_chunks,
         )
 
     def edit_relation(

lightrag/utils.py

@@ -2551,6 +2551,52 @@ def apply_source_ids_limit(
     return truncated
 
 
+def compute_incremental_chunk_ids(
+    existing_full_chunk_ids: list[str],
+    old_chunk_ids: list[str],
+    new_chunk_ids: list[str],
+) -> list[str]:
+    """
+    Compute incrementally updated chunk IDs based on changes.
+
+    This function applies delta changes (additions and removals) to an existing
+    list of chunk IDs while maintaining order and ensuring deduplication.
+    Delta additions from new_chunk_ids are placed at the end.
+
+    Args:
+        existing_full_chunk_ids: Complete list of existing chunk IDs from storage
+        old_chunk_ids: Previous chunk IDs from source_id (chunks being replaced)
+        new_chunk_ids: New chunk IDs from updated source_id (chunks being added)
+
+    Returns:
+        Updated list of chunk IDs with deduplication
+
+    Example:
+        >>> existing = ['chunk-1', 'chunk-2', 'chunk-3']
+        >>> old = ['chunk-1', 'chunk-2']
+        >>> new = ['chunk-2', 'chunk-4']
+        >>> compute_incremental_chunk_ids(existing, old, new)
+        ['chunk-3', 'chunk-2', 'chunk-4']
+    """
+    # Calculate changes
+    chunks_to_remove = set(old_chunk_ids) - set(new_chunk_ids)
+    chunks_to_add = set(new_chunk_ids) - set(old_chunk_ids)
+
+    # Apply changes to full chunk_ids
+    # Step 1: Remove chunks that are no longer needed
+    updated_chunk_ids = [
+        cid for cid in existing_full_chunk_ids if cid not in chunks_to_remove
+    ]
+
+    # Step 2: Add new chunks (preserving order from new_chunk_ids)
+    # Note: 'cid not in updated_chunk_ids' check ensures deduplication
+    for cid in new_chunk_ids:
+        if cid in chunks_to_add and cid not in updated_chunk_ids:
+            updated_chunk_ids.append(cid)
+
+    return updated_chunk_ids
+
+
 def subtract_source_ids(
     source_ids: Iterable[str],
     ids_to_remove: Collection[str],