File size: 1,620 Bytes
933c2fa |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
"""
Level Computer Module
Computes dependency levels for graph nodes to enable efficient batching for LLM processing.
"""
import networkx as nx
def compute_node_levels(graph: nx.DiGraph) -> dict:
"""
Compute the level for each node based on successor depth.
Nodes are assigned levels based on their position in the dependency graph:
- Level 0: nodes with no successors (leaf nodes)
- Level N: 1 + max(successor levels)
This function handles cycles by condensing the graph into strongly connected
components before computing levels.
Args:
graph: NetworkX directed graph
Returns:
Dictionary mapping each node to its level (int)
"""
# Condense the graph to handle strongly connected components (cycles)
C_graph = nx.condensation(graph)
scc_map = C_graph.graph['mapping']
levels = {}
def level(node):
"""Recursively compute level with memoization."""
if node in levels:
return levels[node]
succ = list(C_graph.successors(node))
if not succ: # No outgoing edges → level 0
levels[node] = 0
else:
levels[node] = 1 + max(level(s) for s in succ)
return levels[node]
# Compute levels for all nodes in condensed graph
for node in C_graph.nodes():
level(node)
node_to_level = {node: levels[scc_map[node]] for node in graph.nodes()}
level_to_node = {}
for node, level in node_to_level.items():
level_to_node.setdefault(level, []).append(node)
return level_to_node
|