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