Add retrun status to entity and relation delete operations

lightrag/lightrag.py

@@ -1908,11 +1908,14 @@ Rebuilt: {len(entities_to_rebuild)} entities, {len(relationships_to_rebuild)} re
                 status_code=500,
             )
 
-    async def adelete_by_entity(self, entity_name: str) -> None:
+    async def adelete_by_entity(self, entity_name: str) -> DeletionResult:
         """Asynchronously delete an entity and all its relationships.
 
         Args:
-            entity_name: Name of the entity to delete
+            entity_name: Name of the entity to delete.
+
+        Returns:
+            DeletionResult: An object containing the outcome of the deletion process.
         """
         from .utils_graph import adelete_by_entity
 
@@ -1923,16 +1926,29 @@ Rebuilt: {len(entities_to_rebuild)} entities, {len(relationships_to_rebuild)} re
             entity_name,
         )
 
-    def delete_by_entity(self, entity_name: str) -> None:
+    def delete_by_entity(self, entity_name: str) -> DeletionResult:
+        """Synchronously delete an entity and all its relationships.
+
+        Args:
+            entity_name: Name of the entity to delete.
+
+        Returns:
+            DeletionResult: An object containing the outcome of the deletion process.
+        """
         loop = always_get_an_event_loop()
         return loop.run_until_complete(self.adelete_by_entity(entity_name))
 
-    async def adelete_by_relation(self, source_entity: str, target_entity: str) -> None:
+    async def adelete_by_relation(
+        self, source_entity: str, target_entity: str
+    ) -> DeletionResult:
         """Asynchronously delete a relation between two entities.
 
         Args:
-            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.
+
+        Returns:
+            DeletionResult: An object containing the outcome of the deletion process.
         """
         from .utils_graph import adelete_by_relation
 
@@ -1943,7 +1959,18 @@ Rebuilt: {len(entities_to_rebuild)} entities, {len(relationships_to_rebuild)} re
             target_entity,
         )
 
-    def delete_by_relation(self, source_entity: str, target_entity: str) -> None:
+    def delete_by_relation(
+        self, source_entity: str, target_entity: str
+    ) -> DeletionResult:
+        """Synchronously delete a relation between two entities.
+
+        Args:
+            source_entity: Name of the source entity.
+            target_entity: Name of the target entity.
+
+        Returns:
+            DeletionResult: An object containing the outcome of the deletion process.
+        """
         loop = always_get_an_event_loop()
         return loop.run_until_complete(
             self.adelete_by_relation(source_entity, target_entity)

lightrag/utils_graph.py

@@ -4,6 +4,7 @@ import time
 import asyncio
 from typing import Any, cast
 
+from .base import DeletionResult
 from .kg.shared_storage import get_graph_db_lock
 from .prompt import GRAPH_FIELD_SEP
 from .utils import compute_mdhash_id, logger
@@ -12,7 +13,7 @@ from .base import StorageNameSpace
 
 async def adelete_by_entity(
     chunk_entity_relation_graph, entities_vdb, relationships_vdb, entity_name: str
-) -> None:
+) -> DeletionResult:
     """Asynchronously delete an entity and all its relationships.
 
     Args:
@@ -25,18 +26,43 @@ async def adelete_by_entity(
     # Use graph database lock to ensure atomic graph and vector db operations
     async with graph_db_lock:
         try:
+            # Check if the entity exists
+            if not await chunk_entity_relation_graph.has_node(entity_name):
+                logger.warning(f"Entity '{entity_name}' not found.")
+                return DeletionResult(
+                    status="not_found",
+                    doc_id=entity_name,
+                    message=f"Entity '{entity_name}' not found.",
+                    status_code=404,
+                )
+            # Retrieve related relationships before deleting the node
+            edges = await chunk_entity_relation_graph.get_node_edges(entity_name)
+            related_relations_count = len(edges) if edges else 0
+
             await entities_vdb.delete_entity(entity_name)
             await relationships_vdb.delete_entity_relation(entity_name)
             await chunk_entity_relation_graph.delete_node(entity_name)
 
-            logger.info(
+            message = f"Entity '{entity_name}' and its {related_relations_count} relationships have been deleted."
-                f"Entity '{entity_name}' and its relationships have been deleted."
+            logger.info(message)
-            )
             await _delete_by_entity_done(
                 entities_vdb, relationships_vdb, chunk_entity_relation_graph
             )
+            return DeletionResult(
+                status="success",
+                doc_id=entity_name,
+                message=message,
+                status_code=200,
+            )
         except Exception as e:
-            logger.error(f"Error while deleting entity '{entity_name}': {e}")
+            error_message = f"Error while deleting entity '{entity_name}': {e}"
+            logger.error(error_message)
+            return DeletionResult(
+                status="fail",
+                doc_id=entity_name,
+                message=error_message,
+                status_code=500,
+            )
 
 
 async def _delete_by_entity_done(
@@ -60,7 +86,7 @@ async def adelete_by_relation(
     relationships_vdb,
     source_entity: str,
     target_entity: str,
-) -> None:
+) -> DeletionResult:
     """Asynchronously delete a relation between two entities.
 
     Args:
@@ -69,6 +95,7 @@ async def adelete_by_relation(
         source_entity: Name of the source entity
         target_entity: Name of the target entity
     """
+    relation_str = f"{source_entity} -> {target_entity}"
     graph_db_lock = get_graph_db_lock(enable_logging=False)
     # Use graph database lock to ensure atomic graph and vector db operations
     async with graph_db_lock:
@@ -78,10 +105,14 @@ async def adelete_by_relation(
                 source_entity, target_entity
             )
             if not edge_exists:
-                logger.warning(
+                message = f"Relation from '{source_entity}' to '{target_entity}' does not exist"
-                    f"Relation from '{source_entity}' to '{target_entity}' does not exist"
+                logger.warning(message)
+                return DeletionResult(
+                    status="not_found",
+                    doc_id=relation_str,
+                    message=message,
+                    status_code=404,
                 )
-                return
 
             # Delete relation from vector database
             rel_ids_to_delete = [
@@ -96,13 +127,23 @@ async def adelete_by_relation(
                 [(source_entity, target_entity)]
             )
 
-            logger.info(
+            message = f"Successfully deleted relation from '{source_entity}' to '{target_entity}'"
-                f"Successfully deleted relation from '{source_entity}' to '{target_entity}'"
+            logger.info(message)
-            )
             await _delete_relation_done(relationships_vdb, chunk_entity_relation_graph)
+            return DeletionResult(
+                status="success",
+                doc_id=relation_str,
+                message=message,
+                status_code=200,
+            )
         except Exception as e:
-            logger.error(
+            error_message = f"Error while deleting relation from '{source_entity}' to '{target_entity}': {e}"
-                f"Error while deleting relation from '{source_entity}' to '{target_entity}': {e}"
+            logger.error(error_message)
+            return DeletionResult(
+                status="fail",
+                doc_id=relation_str,
+                message=error_message,
+                status_code=500,
             )