cleveragents-core/robot/langgraph/pure_graph_lib.py

"""Robot Framework library for PureGraph integration testing."""

import contextlib


class PureGraphLibrary:
    """Custom keywords for exercising PureGraph functionality."""

    ROBOT_AUTO_KEYWORDS = True

    def __init__(self):
        self._graph = None
        self._fn_registry = {}
        self._topo_order = None
        self._exec_result = None

    # ------------------------------------------------------------------
    # Graph construction helpers
    # ------------------------------------------------------------------

    def create_linear_graph_with_function_nodes(self, node_names, function_names):
        """Create a linear graph with the specified node and function names."""
        from cleveragents.langgraph.nodes import Edge, NodeConfig, NodeType
        from cleveragents.langgraph.pure_graph import PureGraph

        nodes = {
            "start": NodeConfig(name="start", type=NodeType.START),
            "end": NodeConfig(name="end", type=NodeType.END),
        }
        edges = []
        previous = "start"
        for node_name, function_name in zip(node_names, function_names, strict=False):
            nodes[node_name] = NodeConfig(
                name=node_name, type=NodeType.FUNCTION, function=function_name
            )
            edges.append(Edge(source=previous, target=node_name))
            previous = node_name
        edges.append(Edge(source=previous, target="end"))

        self._graph = PureGraph(name="robot-test-graph", nodes=nodes, edges=edges)

    def create_linear_graph_with_non_functional_nodes(self, node_names):
        """Create a linear graph with non-functional (inert) nodes."""
        from cleveragents.langgraph.nodes import Edge, NodeConfig, NodeType
        from cleveragents.langgraph.pure_graph import PureGraph

        nodes = {
            "start": NodeConfig(name="start", type=NodeType.START),
            "end": NodeConfig(name="end", type=NodeType.END),
        }
        edges = []
        previous = "start"
        for node_name in node_names:
            nodes[node_name] = NodeConfig(
                name=node_name, type=NodeType.FUNCTION, function=None
            )
            edges.append(Edge(source=previous, target=node_name))
            previous = node_name
        edges.append(Edge(source=previous, target="end"))

        self._graph = PureGraph(name="robot-test-inert", nodes=nodes, edges=edges)

    def create_graph_with_missing_function(self, unregistered_fn_name):
        """Create a graph where a node references a function not in the registry."""
        from cleveragents.langgraph.nodes import Edge, NodeConfig, NodeType
        from cleveragents.langgraph.pure_graph import PureGraph

        nodes = {
            "start": NodeConfig(name="start", type=NodeType.START),
            "end": NodeConfig(name="end", type=NodeType.END),
        }
        edges = []
        previous = "start"
        node_name = f"skipper_{unregistered_fn_name}"
        nodes[node_name] = NodeConfig(
            name=node_name, type=NodeType.FUNCTION, function=unregistered_fn_name
        )
        edges.append(Edge(source=previous, target=node_name))
        previous = node_name
        edges.append(Edge(source=previous, target="end"))

        self._graph = PureGraph(name="robot-test-missing-fn", nodes=nodes, edges=edges)
        self._fn_registry = {}

    def set_function_registry(self, registry_dict):
        """Set or update the function registry used by execute()."""
        self._fn_registry = {}
        for key, value in registry_dict.items():
            if callable(value):
                self._fn_registry[key] = value
            elif isinstance(value, str):
                with contextlib.suppress(Exception):
                    self._fn_registry[key] = eval(value)

    # ------------------------------------------------------------------
    # Topological order queries
    # ------------------------------------------------------------------

    def compute_topological_order(self):
        """Compute and return the topological ordering of the current graph."""
        self._topo_order = self._graph.topological_order()

    def get_topological_order(self):
        """Return the last computed topological order as a list."""
        return self._topo_order or []

    # ------------------------------------------------------------------
    # Execution queries
    # ------------------------------------------------------------------

    def execute_graph(self, initial_value):
        """Execute the current graph starting with the given initial value."""
        self._exec_result = self._graph.execute(
            self._fn_registry, initial=initial_value
        )
        return self._exec_result

    def get_execution_result(self):
        """Return the last execution result."""
        return self._exec_result

    # ------------------------------------------------------------------
    # Assertion helpers
    # ------------------------------------------------------------------

    def assert_topo_order_equals(self, *expected_nodes):
        """Assert topological order matches expected: [start, nodes..., end]."""
        actual = self._topo_order or []
        expected = ["start", *expected_nodes, "end"]
        if actual != expected:
            raise AssertionError(
                f"Topological order mismatch: expected {expected}, got {actual}"
            )

    def assert_execution_result_equals(self, expected):
        """Assert the last execution result equals expected value."""
        actual = self._exec_result
        if actual != expected:
            raise AssertionError(f"Expected {expected}, got {actual}")