Update PathRAG/storage.py

PathRAG/storage.py

@@ -154,7 +154,7 @@ class NanoVectorDBStorage(BaseVectorStorage):
             relations = [
                 dp
                 for dp in self.client_storage["data"]
-                if dp["src_id"] == entity_name or dp["tgt_id"] == entity_name
+                if dp.get("src_id") == entity_name or dp.get("tgt_id") == entity_name
             ]
             ids_to_delete = [relation["__id__"] for relation in relations]
 
@@ -181,10 +181,6 @@ class NetworkXStorage(BaseGraphStorage):
         if os.path.exists(file_name):
             return nx.read_graphml(file_name)
         return None
-    # def load_nx_graph(file_name) -> nx.Graph:
-    #     if os.path.exists(file_name):
-    #         return nx.read_graphml(file_name)
-    #     return None
 
     @staticmethod
     def write_nx_graph(graph: nx.DiGraph, file_name):
@@ -195,49 +191,27 @@ class NetworkXStorage(BaseGraphStorage):
 
     @staticmethod
     def stable_largest_connected_component(graph: nx.Graph) -> nx.Graph:
-        """Refer to https://github.com/microsoft/graphrag/index/graph/utils/stable_lcc.py
-        Return the largest connected component of the graph, with nodes and edges sorted in a stable way.
-        """
         from graspologic.utils import largest_connected_component
-
         graph = graph.copy()
         graph = cast(nx.Graph, largest_connected_component(graph))
         node_mapping = {
             node: html.unescape(node.upper().strip()) for node in graph.nodes()
-        }  # type: ignore
+        }
         graph = nx.relabel_nodes(graph, node_mapping)
         return NetworkXStorage._stabilize_graph(graph)
 
     @staticmethod
     def _stabilize_graph(graph: nx.Graph) -> nx.Graph:
-        """Refer to https://github.com/microsoft/graphrag/index/graph/utils/stable_lcc.py
-        Ensure an undirected graph with the same relationships will always be read the same way.
-        """
         fixed_graph = nx.DiGraph() if graph.is_directed() else nx.Graph()
-
-        sorted_nodes = graph.nodes(data=True)
-        sorted_nodes = sorted(sorted_nodes, key=lambda x: x[0])
-
+        sorted_nodes = sorted(list(graph.nodes(data=True)), key=lambda x: x[0])
         fixed_graph.add_nodes_from(sorted_nodes)
+        
         edges = list(graph.edges(data=True))
-
         if not graph.is_directed():
-
-            def _sort_source_target(edge):
-                source, target, edge_data = edge
-                if source > target:
-                    temp = source
-                    source = target
-                    target = temp
-                return source, target, edge_data
-
-            edges = [_sort_source_target(edge) for edge in edges]
-
-        def _get_edge_key(source: Any, target: Any) -> str:
-            return f"{source} -> {target}"
-
-        edges = sorted(edges, key=lambda x: _get_edge_key(x[0], x[1]))
-
+            edges = sorted(edges, key=lambda x: (min(x[0], x[1]), max(x[0], x[1])))
+        else:
+            edges = sorted(edges, key=lambda x: (x[0], x[1]))
+            
         fixed_graph.add_edges_from(edges)
         return fixed_graph
 
@@ -282,21 +256,20 @@ class NetworkXStorage(BaseGraphStorage):
         if self._graph.has_node(source_node_id):
             return list(self._graph.edges(source_node_id))
         return None
+        
     async def get_node_in_edges(self, source_node_id: str):
         if self._graph.has_node(source_node_id):
             return list(self._graph.in_edges(source_node_id))
         return None
+        
     async def get_node_out_edges(self, source_node_id: str):
         if self._graph.has_node(source_node_id):
             return list(self._graph.out_edges(source_node_id))
         return None
     
-    async def get_pagerank(self,source_node_id:str):
-        pagerank_list=nx.pagerank(self._graph)
-        if source_node_id in pagerank_list:
-            return pagerank_list[source_node_id]
-        else:
-            print("pagerank failed")
+    async def get_pagerank(self, source_node_id: str):
+        pagerank_list = nx.pagerank(self._graph)
+        return pagerank_list.get(source_node_id)
 
     async def upsert_node(self, node_id: str, node_data: dict[str, str]):
         self._graph.add_node(node_id, **node_data)
@@ -307,11 +280,6 @@ class NetworkXStorage(BaseGraphStorage):
         self._graph.add_edge(source_node_id, target_node_id, **edge_data)
 
     async def delete_node(self, node_id: str):
-        """
-        Delete a node from the graph based on the specified node_id.
-
-        :param node_id: The node_id to delete
-        """
         if self._graph.has_node(node_id):
             self._graph.remove_node(node_id)
             logger.info(f"Node {node_id} deleted from the graph.")
@@ -323,19 +291,21 @@ class NetworkXStorage(BaseGraphStorage):
             raise ValueError(f"Node embedding algorithm {algorithm} not supported")
         return await self._node_embed_algorithms[algorithm]()
 
-    # @TODO: NOT USED
     async def _node2vec_embed(self):
         from graspologic import embed
-
         embeddings, nodes = embed.node2vec_embed(
             self._graph,
             **self.global_config["node2vec_params"],
         )
-
         nodes_ids = [self._graph.nodes[node_id]["id"] for node_id in nodes]
         return embeddings, nodes_ids
     
-    async def edges(self):
-        return self._graph.edges()
+    # --- CHANGE: Added missing methods ---
+    # These methods are required by operate.py for graph traversal.
     async def nodes(self):
+        """Returns all nodes in the graph."""
         return self._graph.nodes()
+
+    async def edges(self):
+        """Returns all edges in the graph."""
+        return self._graph.edges()