level_computer.py

"""
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