LLM NetworkX

graph_utils.py

from collections import defaultdict
from typing import Generic, TypeVar

T = TypeVar("T")


class SimpleGraph(Generic[T]):
    """
    Lightweight graph implementation, only for the needs of the health module.
    """

    def __init__(self):
        self.graph: dict[T, list[T]] = defaultdict(list)
        self.reverse_graph: dict[T, list[T]] = defaultdict(list)
        self.nodes: set[T] = set()

    def add_node(self, node: T):
        """
        Adds a node to the graph
        """

        self.nodes.add(node)

    def add_edge(self, from_node: T, to_node: T):
        """
        Adds an edge between two nodes
        """

        self.graph[from_node].append(to_node)
        self.reverse_graph[to_node].append(from_node)
        self.nodes.add(from_node)
        self.nodes.add(to_node)

    def topological_sort(self) -> list[T]:
        """
        Returns a topological sort of the graph. This is a linear algorithm
        that runs in O(V + E) time and O(V) space.
        """

        visited = set()
        stack = []

        def dfs(n):
            visited.add(n)
            for neighbor in self.graph[n]:
                if neighbor not in visited:
                    dfs(neighbor)
            stack.append(n)

        for node in self.nodes:
            if node not in visited:
                dfs(node)

        return list(reversed(stack))

    def predecessors(self, node: T) -> list[T]:
        """
        Returns a list of all nodes that have edges pointing to the given node.
        """

        return self.reverse_graph[node]