Fix linting

This commit is contained in:
yangdx
2025-02-24 03:33:22 +08:00
parent b4543561f6
commit 6cf555ebd5

View File

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