Fix linting
This commit is contained in:
@@ -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(
|
||||
|
Reference in New Issue
Block a user