cleveragents-core/benchmarks/langgraph_execution_bench.py

"""ASV benchmarks for LangGraph execution hotspots.

Measures performance of:
- GraphExecutor.step execution
- RouterNode.route decision making
- Checkpoint read/write operations
- Graph state transitions
- Node execution throughput
"""

from __future__ import annotations

import asyncio
import sys
from pathlib import Path
from typing import Any

try:
    from cleveragents.langgraph.graph import GraphConfig, LangGraph
    from cleveragents.langgraph.nodes import Edge, Node, NodeConfig, NodeType
    from cleveragents.langgraph.state import GraphState
except ModuleNotFoundError:
    sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))
    from cleveragents.langgraph.graph import GraphConfig, LangGraph
    from cleveragents.langgraph.nodes import Edge, Node, NodeConfig, NodeType
    from cleveragents.langgraph.state import GraphState


def _create_simple_graph() -> LangGraph:
    """Create a simple representative plan graph for benchmarking."""
    config = GraphConfig(
        name="bench-graph",
        nodes={
            "start": NodeConfig(name="start", type=NodeType.START),
            "process": NodeConfig(name="process", type=NodeType.FUNCTION),
            "route": NodeConfig(name="route", type=NodeType.CONDITIONAL),
            "end": NodeConfig(name="end", type=NodeType.END),
        },
        edges=[
            Edge(source="start", target="process"),
            Edge(source="process", target="route"),
            Edge(source="route", target="end"),
        ],
        entry_point="start",
        checkpointing=False,
    )
    return LangGraph(config)


def _create_checkpoint_graph() -> LangGraph:
    """Create a graph with checkpointing enabled for persistence benchmarks."""
    checkpoint_dir = Path("/tmp/langgraph_bench_checkpoints")
    checkpoint_dir.mkdir(exist_ok=True)

    config = GraphConfig(
        name="bench-graph-checkpoint",
        nodes={
            "start": NodeConfig(name="start", type=NodeType.START),
            "process": NodeConfig(name="process", type=NodeType.FUNCTION),
            "end": NodeConfig(name="end", type=NodeType.END),
        },
        edges=[
            Edge(source="start", target="process"),
            Edge(source="process", target="end"),
        ],
        entry_point="start",
        checkpointing=True,
        checkpoint_dir=checkpoint_dir,
    )
    return LangGraph(config)


def _create_complex_graph() -> LangGraph:
    """Create a more complex graph with multiple nodes for throughput testing."""
    config = GraphConfig(
        name="bench-graph-complex",
        nodes={
            "start": NodeConfig(name="start", type=NodeType.START),
            "node1": NodeConfig(name="node1", type=NodeType.FUNCTION),
            "node2": NodeConfig(name="node2", type=NodeType.FUNCTION),
            "node3": NodeConfig(name="node3", type=NodeType.FUNCTION),
            "route": NodeConfig(name="route", type=NodeType.CONDITIONAL),
            "end": NodeConfig(name="end", type=NodeType.END),
        },
        edges=[
            Edge(source="start", target="node1"),
            Edge(source="node1", target="node2"),
            Edge(source="node2", target="node3"),
            Edge(source="node3", target="route"),
            Edge(source="route", target="end"),
        ],
        entry_point="start",
        checkpointing=False,
        parallel_execution=True,
    )
    return LangGraph(config)


class LangGraphExecutionSuite:
    """Benchmark LangGraph execution hotspots."""

    def setup(self) -> None:
        """Initialize graphs for benchmarking."""
        self.simple_graph = _create_simple_graph()
        self.checkpoint_graph = _create_checkpoint_graph()
        self.complex_graph = _create_complex_graph()

        # Sample input state
        self.sample_state = GraphState()

    def time_simple_graph_execute(self) -> None:
        """Measure simple graph execution time."""
        asyncio.run(self.simple_graph.execute(self.sample_state))

    def time_complex_graph_execute(self) -> None:
        """Measure complex graph execution time."""
        asyncio.run(self.complex_graph.execute(self.sample_state))

    def time_checkpoint_graph_execute(self) -> None:
        """Measure graph execution with checkpointing enabled."""
        asyncio.run(self.checkpoint_graph.execute(self.sample_state))

    def time_state_manager_get_state(self) -> None:
        """Measure state manager state retrieval."""
        self.simple_graph.state_manager.get_state()

    def time_state_manager_update_state(self) -> None:
        """Measure state manager state update."""
        new_state = GraphState()
        self.simple_graph.state_manager.state = new_state

    def track_graph_nodes_count(self) -> int:
        """Track number of nodes in simple graph."""
        return len(self.simple_graph.nodes)

    def track_graph_edges_count(self) -> int:
        """Track number of edges in simple graph."""
        return len(self.simple_graph.config.edges)

    def track_execution_history_length(self) -> int:
        """Track execution history length after execution."""
        return len(self.simple_graph.get_execution_history())


class LangGraphRouterNodeSuite:
    """Benchmark RouterNode routing performance."""

    def setup(self) -> None:
        """Initialize router node for benchmarking."""
        self.router_config = NodeConfig(
            name="router",
            type=NodeType.CONDITIONAL,
            condition={"type": "simple_condition"},
        )
        self.router_node = Node(self.router_config)

    def time_router_node_creation(self) -> None:
        """Measure router node instantiation time."""
        Node(self.router_config)

    def time_router_node_config_access(self) -> None:
        """Measure router node config access."""
        _ = self.router_node.config.condition


class LangGraphStateTransitionSuite:
    """Benchmark graph state transitions and updates."""

    def setup(self) -> None:
        """Initialize state for benchmarking."""
        self.state = GraphState()
        self.graph = _create_simple_graph()

    def time_state_creation(self) -> None:
        """Measure GraphState creation time."""
        GraphState()

    def time_state_dict_conversion(self) -> None:
        """Measure state to dict conversion."""
        self.state.to_dict()

    def time_state_from_dict(self) -> None:
        """Measure state from dict creation."""
        GraphState.from_dict({"messages": []})

    def time_state_copy(self) -> None:
        """Measure state copy operation."""
        self.state.copy()