Added info logs when graph is truncated

This commit is contained in:
yangdx
2025-04-02 23:19:41 +08:00
parent 82c4baba70
commit 72132ee1d6

View File

@@ -276,18 +276,21 @@ class NetworkXStorage(BaseGraphStorage):
graph = await self._get_graph() graph = await self._get_graph()
result = KnowledgeGraph() result = KnowledgeGraph()
# Handle special case for "*" label # Handle special case for "*" label
if node_label == "*": if node_label == "*":
# Get degrees of all nodes # Get degrees of all nodes
degrees = dict(graph.degree()) degrees = dict(graph.degree())
# Sort nodes by degree in descending order and take top max_nodes # Sort nodes by degree in descending order and take top max_nodes
sorted_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True) sorted_nodes = sorted(degrees.items(), key=lambda x: x[1], reverse=True)
# Check if graph is truncated # Check if graph is truncated
if len(sorted_nodes) > max_nodes: if len(sorted_nodes) > max_nodes:
result.is_truncated = True result.is_truncated = True
logger.info(
f"Graph truncated: {len(sorted_nodes)} nodes found, limited to {max_nodes}"
)
limited_nodes = [node for node, _ in sorted_nodes[:max_nodes]] limited_nodes = [node for node, _ in sorted_nodes[:max_nodes]]
# Create subgraph with the highest degree nodes # Create subgraph with the highest degree nodes
subgraph = graph.subgraph(limited_nodes) subgraph = graph.subgraph(limited_nodes)
@@ -301,23 +304,26 @@ class NetworkXStorage(BaseGraphStorage):
bfs_nodes = [] bfs_nodes = []
visited = set() visited = set()
queue = [node_label] queue = [node_label]
# Breadth-first search # Breadth-first search
while queue and len(bfs_nodes) < max_nodes: while queue and len(bfs_nodes) < max_nodes:
current = queue.pop(0) current = queue.pop(0)
if current not in visited: if current not in visited:
visited.add(current) visited.add(current)
bfs_nodes.append(current) bfs_nodes.append(current)
# Add neighbor nodes to queue # Add neighbor nodes to queue
neighbors = list(graph.neighbors(current)) neighbors = list(graph.neighbors(current))
queue.extend([n for n in neighbors if n not in visited]) queue.extend([n for n in neighbors if n not in visited])
# Check if graph is truncated - if we still have nodes in the queue # Check if graph is truncated - if we still have nodes in the queue
# and we've reached max_nodes, then the graph is truncated # and we've reached max_nodes, then the graph is truncated
if queue and len(bfs_nodes) >= max_nodes: if queue and len(bfs_nodes) >= max_nodes:
result.is_truncated = True result.is_truncated = True
logger.info(
f"Graph truncated: breadth-first search limited to {max_nodes} nodes"
)
# Create subgraph with BFS discovered nodes # Create subgraph with BFS discovered nodes
subgraph = graph.subgraph(bfs_nodes) subgraph = graph.subgraph(bfs_nodes)