From 6cf555ebd58a69ecbde528f79bc90feeeb754604 Mon Sep 17 00:00:00 2001
From: yangdx <yangdx@znipower.com>
Date: Mon, 24 Feb 2025 03:33:22 +0800
Subject: [PATCH] Fix linting

---
 lightrag/kg/networkx_impl.py | 44 ++++++++++++++++++++----------------
 1 file changed, 24 insertions(+), 20 deletions(-)

diff --git a/lightrag/kg/networkx_impl.py b/lightrag/kg/networkx_impl.py
index e68d2887..e8180dda 100644
--- a/lightrag/kg/networkx_impl.py
+++ b/lightrag/kg/networkx_impl.py
@@ -184,7 +184,7 @@ class NetworkXStorage(BaseGraphStorage):
             #     else:
             #         labels.add(node_data["entity_type"])
             labels.add(str(node))  # Add node id as a label
-        
+
         # Return sorted list
         return sorted(list(labels))
 
@@ -193,52 +193,58 @@ class NetworkXStorage(BaseGraphStorage):
     ) -> KnowledgeGraph:
         """
         Get complete connected subgraph for specified node (including the starting node itself)
-        
+
         Args:
             node_label: Label of the starting node
             max_depth: Maximum depth of the subgraph
-            
+
         Returns:
             KnowledgeGraph object containing nodes and edges
         """
         result = KnowledgeGraph()
         seen_nodes = set()
         seen_edges = set()
-        
+
         # Handle special case for "*" label
         if node_label == "*":
             # For "*", return the entire graph including all nodes and edges
-            subgraph = self._graph.copy()  # Create a copy to avoid modifying the original graph
+            subgraph = (
+                self._graph.copy()
+            )  # Create a copy to avoid modifying the original graph
         else:
             # Find nodes with matching node id (partial match)
             nodes_to_explore = []
             for n, attr in self._graph.nodes(data=True):
                 if node_label in str(n):  # Use partial matching
                     nodes_to_explore.append(n)
-            
+
             if not nodes_to_explore:
                 logger.warning(f"No nodes found with label {node_label}")
                 return result
 
             # Get subgraph using ego_graph
             subgraph = nx.ego_graph(self._graph, nodes_to_explore[0], radius=max_depth)
-        
+
         # Check if number of nodes exceeds max_graph_nodes
-        max_graph_nodes=500
+        max_graph_nodes = 500
         if len(subgraph.nodes()) > max_graph_nodes:
-            origin_nodes=len(subgraph.nodes())
+            origin_nodes = len(subgraph.nodes())
             node_degrees = dict(subgraph.degree())
-            top_nodes = sorted(node_degrees.items(), key=lambda x: x[1], reverse=True)[:max_graph_nodes]
+            top_nodes = sorted(node_degrees.items(), key=lambda x: x[1], reverse=True)[
+                :max_graph_nodes
+            ]
             top_node_ids = [node[0] for node in top_nodes]
             # Create new subgraph with only top nodes
             subgraph = subgraph.subgraph(top_node_ids)
-            logger.info(f"Reduced graph from {origin_nodes} nodes to {max_graph_nodes} nodes by degree (depth={max_depth})")
+            logger.info(
+                f"Reduced graph from {origin_nodes} nodes to {max_graph_nodes} nodes by degree (depth={max_depth})"
+            )
 
         # Add nodes to result
         for node in subgraph.nodes():
             if str(node) in seen_nodes:
                 continue
-                
+
             node_data = dict(subgraph.nodes[node])
             # Get entity_type as labels
             labels = []
@@ -247,28 +253,26 @@ class NetworkXStorage(BaseGraphStorage):
                     labels.extend(node_data["entity_type"])
                 else:
                     labels.append(node_data["entity_type"])
-            
+
             # Create node with properties
             node_properties = {k: v for k, v in node_data.items()}
 
             result.nodes.append(
                 KnowledgeGraphNode(
-                    id=str(node),
-                    labels=[str(node)],
-                    properties=node_properties
+                    id=str(node), labels=[str(node)], properties=node_properties
                 )
             )
             seen_nodes.add(str(node))
-        
+
         # Add edges to result
         for edge in subgraph.edges():
             source, target = edge
             edge_id = f"{source}-{target}"
             if edge_id in seen_edges:
                 continue
-                
+
             edge_data = dict(subgraph.edges[edge])
-            
+
             # Create edge with complete information
             result.edges.append(
                 KnowledgeGraphEdge(
@@ -280,7 +284,7 @@ class NetworkXStorage(BaseGraphStorage):
                 )
             )
             seen_edges.add(edge_id)
-        
+
         # logger.info(result.edges)
 
         logger.info(