polish Memgraph implementation

lightrag/kg/memgraph_impl.py

@@ -89,19 +89,50 @@ class MemgraphStorage(BaseGraphStorage):
         pass
 
     async def has_node(self, node_id: str) -> bool:
+        """
+        Check if a node exists in the graph.
+
+        Args:
+            node_id: The ID of the node to check.
+
+        Returns:
+            bool: True if the node exists, False otherwise.
+
+        Raises:
+            Exception: If there is an error checking the node existence.
+        """
         async with self._driver.session(
             database=self._DATABASE, default_access_mode="READ"
         ) as session:
+            try:
                 query = "MATCH (n:base {entity_id: $entity_id}) RETURN count(n) > 0 AS node_exists"
                 result = await session.run(query, entity_id=node_id)
                 single_result = await result.single()
-            await result.consume()
+                await result.consume()  # Ensure result is fully consumed
                 return single_result["node_exists"]
+            except Exception as e:
+                logger.error(f"Error checking node existence for {node_id}: {str(e)}")
+                await result.consume()  # Ensure the result is consumed even on error
+                raise
 
     async def has_edge(self, source_node_id: str, target_node_id: str) -> bool:
+        """
+        Check if an edge exists between two nodes in the graph.
+
+        Args:
+            source_node_id: The ID of the source node.
+            target_node_id: The ID of the target node.
+
+        Returns:
+            bool: True if the edge exists, False otherwise.
+
+        Raises:
+            Exception: If there is an error checking the edge existence.
+        """
         async with self._driver.session(
             database=self._DATABASE, default_access_mode="READ"
         ) as session:
+            try:
                 query = (
                     "MATCH (a:base {entity_id: $source_entity_id})-[r]-(b:base {entity_id: $target_entity_id}) "
                     "RETURN COUNT(r) > 0 AS edgeExists"
@@ -112,31 +143,113 @@ class MemgraphStorage(BaseGraphStorage):
                     target_entity_id=target_node_id,
                 )
                 single_result = await result.single()
-            await result.consume()
+                await result.consume()  # Ensure result is fully consumed
                 return single_result["edgeExists"]
+            except Exception as e:
+                logger.error(
+                    f"Error checking edge existence between {source_node_id} and {target_node_id}: {str(e)}"
+                )
+                await result.consume()  # Ensure the result is consumed even on error
+                raise
 
     async def get_node(self, node_id: str) -> dict[str, str] | None:
+        """Get node by its label identifier, return only node properties
+
+        Args:
+            node_id: The node label to look up
+
+        Returns:
+            dict: Node properties if found
+            None: If node not found
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
         async with self._driver.session(
             database=self._DATABASE, default_access_mode="READ"
         ) as session:
+            try:
                 query = "MATCH (n:base {entity_id: $entity_id}) RETURN n"
                 result = await session.run(query, entity_id=node_id)
-            records = await result.fetch(2)
+                try:
-            await result.consume()
+                    records = await result.fetch(
+                        2
+                    )  # Get 2 records for duplication check
+
+                    if len(records) > 1:
+                        logger.warning(
+                            f"Multiple nodes found with label '{node_id}'. Using first node."
+                        )
                     if records:
                         node = records[0]["n"]
                         node_dict = dict(node)
+                        # Remove base label from labels list if it exists
                         if "labels" in node_dict:
                             node_dict["labels"] = [
-                        label for label in node_dict["labels"] if label != "base"
+                                label
+                                for label in node_dict["labels"]
+                                if label != "base"
                             ]
                         return node_dict
                     return None
+                finally:
+                    await result.consume()  # Ensure result is fully consumed
+            except Exception as e:
+                logger.error(f"Error getting node for {node_id}: {str(e)}")
+                raise
 
-    async def get_all_labels(self) -> list[str]:
+    async def node_degree(self, node_id: str) -> int:
+        """Get the degree (number of relationships) of a node with the given label.
+        If multiple nodes have the same label, returns the degree of the first node.
+        If no node is found, returns 0.
+
+        Args:
+            node_id: The label of the node
+
+        Returns:
+            int: The number of relationships the node has, or 0 if no node found
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
         async with self._driver.session(
             database=self._DATABASE, default_access_mode="READ"
         ) as session:
+            try:
+                query = """
+                    MATCH (n:base {entity_id: $entity_id})
+                    OPTIONAL MATCH (n)-[r]-()
+                    RETURN COUNT(r) AS degree
+                """
+                result = await session.run(query, entity_id=node_id)
+                try:
+                    record = await result.single()
+
+                    if not record:
+                        logger.warning(f"No node found with label '{node_id}'")
+                        return 0
+
+                    degree = record["degree"]
+                    return degree
+                finally:
+                    await result.consume()  # Ensure result is fully consumed
+            except Exception as e:
+                logger.error(f"Error getting node degree for {node_id}: {str(e)}")
+                raise
+
+    async def get_all_labels(self) -> list[str]:
+        """
+        Get all existing node labels in the database
+        Returns:
+            ["Person", "Company", ...]  # Alphabetically sorted label list
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
+        async with self._driver.session(
+            database=self._DATABASE, default_access_mode="READ"
+        ) as session:
+            try:
                 query = """
                 MATCH (n:base)
                 WHERE n.entity_id IS NOT NULL
@@ -149,37 +262,89 @@ class MemgraphStorage(BaseGraphStorage):
                     labels.append(record["label"])
                 await result.consume()
                 return labels
+            except Exception as e:
+                logger.error(f"Error getting all labels: {str(e)}")
+                await result.consume()  # Ensure the result is consumed even on error
+                raise
 
     async def get_node_edges(self, source_node_id: str) -> list[tuple[str, str]] | None:
+        """Retrieves all edges (relationships) for a particular node identified by its label.
+
+        Args:
+            source_node_id: Label of the node to get edges for
+
+        Returns:
+            list[tuple[str, str]]: List of (source_label, target_label) tuples representing edges
+            None: If no edges found
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
+        try:
             async with self._driver.session(
                 database=self._DATABASE, default_access_mode="READ"
             ) as session:
-            query = """
+                try:
-            MATCH (n:base {entity_id: $entity_id})
+                    query = """MATCH (n:base {entity_id: $entity_id})
                             OPTIONAL MATCH (n)-[r]-(connected:base)
                             WHERE connected.entity_id IS NOT NULL
-            RETURN n, r, connected
+                            RETURN n, r, connected"""
-            """
                     results = await session.run(query, entity_id=source_node_id)
+
                     edges = []
                     async for record in results:
                         source_node = record["n"]
                         connected_node = record["connected"]
+
+                        # Skip if either node is None
                         if not source_node or not connected_node:
                             continue
-                source_label = source_node.get("entity_id")
+
-                target_label = connected_node.get("entity_id")
+                        source_label = (
+                            source_node.get("entity_id")
+                            if source_node.get("entity_id")
+                            else None
+                        )
+                        target_label = (
+                            connected_node.get("entity_id")
+                            if connected_node.get("entity_id")
+                            else None
+                        )
+
                         if source_label and target_label:
                             edges.append((source_label, target_label))
-            await results.consume()
+
+                    await results.consume()  # Ensure results are consumed
                     return edges
+                except Exception as e:
+                    logger.error(
+                        f"Error getting edges for node {source_node_id}: {str(e)}"
+                    )
+                    await results.consume()  # Ensure results are consumed even on error
+                    raise
+        except Exception as e:
+            logger.error(f"Error in get_node_edges for {source_node_id}: {str(e)}")
+            raise
 
     async def get_edge(
         self, source_node_id: str, target_node_id: str
     ) -> dict[str, str] | None:
+        """Get edge properties between two nodes.
+
+        Args:
+            source_node_id: Label of the source node
+            target_node_id: Label of the target node
+
+        Returns:
+            dict: Edge properties if found, default properties if not found or on error
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
         async with self._driver.session(
             database=self._DATABASE, default_access_mode="READ"
         ) as session:
+            try:
                 query = """
                 MATCH (start:base {entity_id: $source_entity_id})-[r]-(end:base {entity_id: $target_entity_id})
                 RETURN properties(r) as edge_properties
@@ -201,32 +366,70 @@ class MemgraphStorage(BaseGraphStorage):
                     }.items():
                         if key not in edge_result:
                             edge_result[key] = default_value
+                            logger.warning(
+                                f"Edge between {source_node_id} and {target_node_id} is missing property: {key}. Using default value: {default_value}"
+                            )
                     return edge_result
                 return None
+            except Exception as e:
+                logger.error(
+                    f"Error getting edge between {source_node_id} and {target_node_id}: {str(e)}"
+                )
+                await result.consume()  # Ensure the result is consumed even on error
+                raise
 
     async def upsert_node(self, node_id: str, node_data: dict[str, str]) -> None:
+        """
+        Upsert a node in the Neo4j database.
+
+        Args:
+            node_id: The unique identifier for the node (used as label)
+            node_data: Dictionary of node properties
+        """
         properties = node_data
-        entity_type = properties.get("entity_type", "base")
+        entity_type = properties["entity_type"]
         if "entity_id" not in properties:
-            raise ValueError(
+            raise ValueError("Neo4j: node properties must contain an 'entity_id' field")
-                "Memgraph: node properties must contain an 'entity_id' field"
+
-            )
+        try:
             async with self._driver.session(database=self._DATABASE) as session:
 
                 async def execute_upsert(tx: AsyncManagedTransaction):
-                query = f"""
+                    query = (
-                    MERGE (n:base {{entity_id: $entity_id}})
-                    SET n += $properties
-                    SET n:`{entity_type}`
                         """
-                result = await tx.run(query, entity_id=node_id, properties=properties)
+                    MERGE (n:base {entity_id: $entity_id})
-                await result.consume()
+                    SET n += $properties
+                    SET n:`%s`
+                    """
+                        % entity_type
+                    )
+                    result = await tx.run(
+                        query, entity_id=node_id, properties=properties
+                    )
+                    await result.consume()  # Ensure result is fully consumed
 
                 await session.execute_write(execute_upsert)
+        except Exception as e:
+            logger.error(f"Error during upsert: {str(e)}")
+            raise
 
     async def upsert_edge(
         self, source_node_id: str, target_node_id: str, edge_data: dict[str, str]
     ) -> None:
+        """
+        Upsert an edge and its properties between two nodes identified by their labels.
+        Ensures both source and target nodes exist and are unique before creating the edge.
+        Uses entity_id property to uniquely identify nodes.
+
+        Args:
+            source_node_id (str): Label of the source node (used as identifier)
+            target_node_id (str): Label of the target node (used as identifier)
+            edge_data (dict): Dictionary of properties to set on the edge
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
+        try:
             edge_properties = edge_data
             async with self._driver.session(database=self._DATABASE) as session:
 
@@ -245,27 +448,60 @@ class MemgraphStorage(BaseGraphStorage):
                         target_entity_id=target_node_id,
                         properties=edge_properties,
                     )
-                await result.consume()
+                    try:
+                        await result.fetch(2)
+                    finally:
+                        await result.consume()  # Ensure result is consumed
 
                 await session.execute_write(execute_upsert)
+        except Exception as e:
+            logger.error(f"Error during edge upsert: {str(e)}")
+            raise
 
     async def delete_node(self, node_id: str) -> None:
+        """Delete a node with the specified label
+
+        Args:
+            node_id: The label of the node to delete
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
+
         async def _do_delete(tx: AsyncManagedTransaction):
             query = """
             MATCH (n:base {entity_id: $entity_id})
             DETACH DELETE n
             """
             result = await tx.run(query, entity_id=node_id)
+            logger.debug(f"Deleted node with label {node_id}")
             await result.consume()
 
+        try:
             async with self._driver.session(database=self._DATABASE) as session:
                 await session.execute_write(_do_delete)
+        except Exception as e:
+            logger.error(f"Error during node deletion: {str(e)}")
+            raise
 
     async def remove_nodes(self, nodes: list[str]):
+        """Delete multiple nodes
+
+        Args:
+            nodes: List of node labels to be deleted
+        """
         for node in nodes:
             await self.delete_node(node)
 
     async def remove_edges(self, edges: list[tuple[str, str]]):
+        """Delete multiple edges
+
+        Args:
+            edges: List of edges to be deleted, each edge is a (source, target) tuple
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
         for source, target in edges:
 
             async def _do_delete_edge(tx: AsyncManagedTransaction):
@@ -276,15 +512,32 @@ class MemgraphStorage(BaseGraphStorage):
                 result = await tx.run(
                     query, source_entity_id=source, target_entity_id=target
                 )
-                await result.consume()
+                logger.debug(f"Deleted edge from '{source}' to '{target}'")
+                await result.consume()  # Ensure result is fully consumed
 
-            async with self._driver.session(database=self._DATABASE) as session:
-                await session.execute_write(_do_delete_edge)
-
-    async def drop(self) -> dict[str, str]:
             try:
                 async with self._driver.session(database=self._DATABASE) as session:
-                query = "MATCH (n) DETACH DELETE n"
+                    await session.execute_write(_do_delete_edge)
+            except Exception as e:
+                logger.error(f"Error during edge deletion: {str(e)}")
+                raise
+
+    async def drop(self) -> dict[str, str]:
+        """Drop all data from storage and clean up resources
+
+        This method will delete all nodes and relationships in the Neo4j database.
+
+        Returns:
+            dict[str, str]: Operation status and message
+            - On success: {"status": "success", "message": "data dropped"}
+            - On failure: {"status": "error", "message": "<error details>"}
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
+        try:
+            async with self._driver.session(database=self._DATABASE) as session:
+                query = "DROP GRAPH"
                 result = await session.run(query)
                 await result.consume()
                 logger.info(
@@ -295,30 +548,36 @@ class MemgraphStorage(BaseGraphStorage):
             logger.error(f"Error dropping Memgraph database {self._DATABASE}: {e}")
             return {"status": "error", "message": str(e)}
 
-    async def node_degree(self, node_id: str) -> int:
-        async with self._driver.session(
-            database=self._DATABASE, default_access_mode="READ"
-        ) as session:
-            query = """
-                MATCH (n:base {entity_id: $entity_id})
-                OPTIONAL MATCH (n)-[r]-()
-                RETURN COUNT(r) AS degree
-            """
-            result = await session.run(query, entity_id=node_id)
-            record = await result.single()
-            await result.consume()
-            if not record:
-                return 0
-            return record["degree"]
-
     async def edge_degree(self, src_id: str, tgt_id: str) -> int:
+        """Get the total degree (sum of relationships) of two nodes.
+
+        Args:
+            src_id: Label of the source node
+            tgt_id: Label of the target node
+
+        Returns:
+            int: Sum of the degrees of both nodes
+        """
         src_degree = await self.node_degree(src_id)
         trg_degree = await self.node_degree(tgt_id)
+
+        # Convert None to 0 for addition
         src_degree = 0 if src_degree is None else src_degree
         trg_degree = 0 if trg_degree is None else trg_degree
-        return int(src_degree) + int(trg_degree)
+
+        degrees = int(src_degree) + int(trg_degree)
+        return degrees
 
     async def get_nodes_by_chunk_ids(self, chunk_ids: list[str]) -> list[dict]:
+        """Get all nodes that are associated with the given chunk_ids.
+
+        Args:
+            chunk_ids: List of chunk IDs to find associated nodes for
+
+        Returns:
+            list[dict]: A list of nodes, where each node is a dictionary of its properties.
+                        An empty list if no matching nodes are found.
+        """
         async with self._driver.session(
             database=self._DATABASE, default_access_mode="READ"
         ) as session:
@@ -339,6 +598,15 @@ class MemgraphStorage(BaseGraphStorage):
                 return nodes
 
     async def get_edges_by_chunk_ids(self, chunk_ids: list[str]) -> list[dict]:
+        """Get all edges that are associated with the given chunk_ids.
+
+        Args:
+            chunk_ids: List of chunk IDs to find associated edges for
+
+        Returns:
+            list[dict]: A list of edges, where each edge is a dictionary of its properties.
+                        An empty list if no matching edges are found.
+        """
         async with self._driver.session(
             database=self._DATABASE, default_access_mode="READ"
         ) as session:
@@ -364,22 +632,44 @@ class MemgraphStorage(BaseGraphStorage):
         max_depth: int = 3,
         max_nodes: int = MAX_GRAPH_NODES,
     ) -> KnowledgeGraph:
+        """
+        Retrieve a connected subgraph of nodes where the label includes the specified `node_label`.
+
+        Args:
+            node_label: Label of the starting node, * means all nodes
+            max_depth: Maximum depth of the subgraph, Defaults to 3
+            max_nodes: Maxiumu nodes to return by BFS, Defaults to 1000
+
+        Returns:
+            KnowledgeGraph object containing nodes and edges, with an is_truncated flag
+            indicating whether the graph was truncated due to max_nodes limit
+
+        Raises:
+            Exception: If there is an error executing the query
+        """
         result = KnowledgeGraph()
         seen_nodes = set()
         seen_edges = set()
+        try:
             async with self._driver.session(
                 database=self._DATABASE, default_access_mode="READ"
             ) as session:
                 if node_label == "*":
                     count_query = "MATCH (n) RETURN count(n) as total"
+                    count_result = None
+                    try:
                         count_result = await session.run(count_query)
                         count_record = await count_result.single()
-                await count_result.consume()
                         if count_record and count_record["total"] > max_nodes:
                             result.is_truncated = True
                             logger.info(
                                 f"Graph truncated: {count_record['total']} nodes found, limited to {max_nodes}"
                             )
+                    finally:
+                        if count_result:
+                            await count_result.consume()
+
+                    # Run the main query to get nodes with highest degree
                     main_query = """
                     MATCH (n)
                     OPTIONAL MATCH (n)-[r]-()
@@ -394,88 +684,97 @@ class MemgraphStorage(BaseGraphStorage):
                     RETURN filtered_nodes AS node_info,
                             collect(DISTINCT r) AS relationships
                     """
-                result_set = await session.run(main_query, {"max_nodes": max_nodes})
+                    result_set = None
-                record = await result_set.single()
+                    try:
-                await result_set.consume()
+                        result_set = await session.run(
-            else:
+                            main_query, {"max_nodes": max_nodes}
-                # BFS fallback for Memgraph (no APOC)
-                from collections import deque
-
-                # Get the starting node
-                start_query = "MATCH (n:base {entity_id: $entity_id}) RETURN n"
-                node_result = await session.run(start_query, entity_id=node_label)
-                node_record = await node_result.single()
-                await node_result.consume()
-                if not node_record:
-                    return result
-                start_node = node_record["n"]
-                queue = deque([(start_node, 0)])
-                visited = set()
-                bfs_nodes = []
-                while queue and len(bfs_nodes) < max_nodes:
-                    current_node, depth = queue.popleft()
-                    node_id = current_node.get("entity_id")
-                    if node_id in visited:
-                        continue
-                    visited.add(node_id)
-                    bfs_nodes.append(current_node)
-                    if depth < max_depth:
-                        # Get neighbors
-                        neighbor_query = """
-                        MATCH (n:base {entity_id: $entity_id})-[]-(m:base)
-                        RETURN m
-                        """
-                        neighbors_result = await session.run(
-                            neighbor_query, entity_id=node_id
                         )
-                        neighbors = [
+                        record = await result_set.single()
-                            rec["m"] for rec in await neighbors_result.to_list()
+                    finally:
-                        ]
+                        if result_set:
-                        await neighbors_result.consume()
+                            await result_set.consume()
-                        for neighbor in neighbors:
+
-                            neighbor_id = neighbor.get("entity_id")
+                else:
-                            if neighbor_id not in visited:
+                    bfs_query = """
-                                queue.append((neighbor, depth + 1))
+                    MATCH (start) WHERE start.entity_id = $entity_id
-                # Build subgraph
+                    WITH start
-                subgraph_ids = [n.get("entity_id") for n in bfs_nodes]
+                    CALL {
-                # Nodes
+                        WITH start
-                for n in bfs_nodes:
+                        MATCH path = (start)-[*0..$max_depth]-(node)
-                    node_id = n.get("entity_id")
+                        WITH nodes(path) AS path_nodes, relationships(path) AS path_rels
+                        UNWIND path_nodes AS n
+                        WITH collect(DISTINCT n) AS all_nodes, collect(DISTINCT path_rels) AS all_rel_lists
+                        WITH all_nodes, reduce(r = [], x IN all_rel_lists | r + x) AS all_rels
+                        RETURN all_nodes, all_rels
+                    }
+                    WITH all_nodes AS nodes, all_rels AS relationships, size(all_nodes) AS total_nodes
+
+                    // Apply node limiting here
+                    WITH CASE
+                        WHEN total_nodes <= $max_nodes THEN nodes
+                        ELSE nodes[0..$max_nodes]
+                    END AS limited_nodes,
+                    relationships,
+                    total_nodes,
+                    total_nodes > $max_nodes AS is_truncated
+                    UNWIND limited_nodes AS node
+                    WITH collect({node: node}) AS node_info, relationships, total_nodes, is_truncated
+                    RETURN node_info, relationships, total_nodes, is_truncated
+                    """
+                    result_set = None
+                    try:
+                        result_set = await session.run(
+                            bfs_query,
+                            {
+                                "entity_id": node_label,
+                                "max_depth": max_depth,
+                                "max_nodes": max_nodes,
+                            },
+                        )
+                        record = await result_set.single()
+                        if not record:
+                            logger.debug(f"No record found for node {node_label}")
+                            return result
+
+                        for node_info in record["node_info"]:
+                            node = node_info["node"]
+                            node_id = node.id
                             if node_id not in seen_nodes:
+                                seen_nodes.add(node_id)
                                 result.nodes.append(
                                     KnowledgeGraphNode(
-                                id=node_id,
+                                        id=f"{node_id}",
-                                labels=[node_id],
+                                        labels=[node.get("entity_id")],
-                                properties=dict(n),
+                                        properties=dict(node),
                                     )
                                 )
-                        seen_nodes.add(node_id)
+
-                # Edges
+                        for rel in record["relationships"]:
-                if subgraph_ids:
+                            edge_id = rel.id
-                    edge_query = """
-                    MATCH (a:base)-[r]-(b:base)
-                    WHERE a.entity_id IN $ids AND b.entity_id IN $ids
-                    RETURN DISTINCT r, a, b
-                    """
-                    edge_result = await session.run(edge_query, ids=subgraph_ids)
-                    async for record in edge_result:
-                        r = record["r"]
-                        a = record["a"]
-                        b = record["b"]
-                        edge_id = f"{a.get('entity_id')}-{b.get('entity_id')}"
                             if edge_id not in seen_edges:
+                                seen_edges.add(edge_id)
+                                start = rel.start_node
+                                end = rel.end_node
                                 result.edges.append(
                                     KnowledgeGraphEdge(
-                                    id=edge_id,
+                                        id=f"{edge_id}",
-                                    type="DIRECTED",
+                                        type=rel.type,
-                                    source=a.get("entity_id"),
+                                        source=f"{start.id}",
-                                    target=b.get("entity_id"),
+                                        target=f"{end.id}",
-                                    properties=dict(r),
+                                        properties=dict(rel),
                                     )
                                 )
-                            seen_edges.add(edge_id)
+
-                    await edge_result.consume()
                         logger.info(
                             f"Subgraph query successful | Node count: {len(result.nodes)} | Edge count: {len(result.edges)}"
                         )
+
+                        return result
+
+                    finally:
+                        if result_set:
+                            await result_set.consume()
+
+        except Exception as e:
+            logger.error(f"Error getting knowledge graph: {str(e)}")
             return result