HAL9000
ecb338cbe6
Fix formatting issues found by CI lint check: - Consolidate multi-line NodeConfig and PureGraph constructor calls onto single lines as preferred by ruff formatter This resolves the failing CI / lint (pull_request) check. ISSUES CLOSED: #9531
59 lines
2.1 KiB
Python
59 lines
2.1 KiB
Python
"""Benchmarks for PureGraph execution throughput and ordering."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, ClassVar
|
|
|
|
from cleveragents.langgraph.nodes import Edge, NodeConfig, NodeType
|
|
from cleveragents.langgraph.pure_graph import PureGraph
|
|
|
|
|
|
class PureGraphBench:
|
|
"""Benchmark suite for PureGraph execution performance."""
|
|
|
|
params: ClassVar[list[int]] = [10, 50, 100, 500]
|
|
param_names: ClassVar[list[str]] = ["node_count"]
|
|
|
|
def setup(self, node_count: int) -> None:
|
|
"""Set up benchmark with varying node counts."""
|
|
self.node_count = node_count
|
|
self.nodes: dict[str, NodeConfig] = {
|
|
"start": NodeConfig(name="start", type=NodeType.START),
|
|
"end": NodeConfig(name="end", type=NodeType.END),
|
|
}
|
|
self.edges: list[Edge] = []
|
|
|
|
# Create a linear chain of nodes
|
|
previous = "start"
|
|
for i in range(node_count):
|
|
node_name = f"node_{i}"
|
|
self.nodes[node_name] = NodeConfig(
|
|
name=node_name, type=NodeType.FUNCTION, function=f"fn_{i}"
|
|
)
|
|
self.edges.append(Edge(source=previous, target=node_name))
|
|
previous = node_name
|
|
|
|
self.edges.append(Edge(source=previous, target="end"))
|
|
|
|
# Create function registry with identity functions
|
|
self.fn_registry: dict[str, Any] = {
|
|
f"fn_{i}": lambda x, i=i: x for i in range(node_count)
|
|
}
|
|
|
|
self.graph = PureGraph(
|
|
name=f"bench_graph_{node_count}", nodes=self.nodes, edges=self.edges
|
|
)
|
|
|
|
def time_topological_order(self, node_count: int) -> None:
|
|
"""Benchmark topological ordering."""
|
|
self.graph.topological_order()
|
|
|
|
def time_execute(self, node_count: int) -> None:
|
|
"""Benchmark graph execution."""
|
|
self.graph.execute(self.fn_registry, initial=0)
|
|
|
|
def time_execute_with_result(self, node_count: int) -> None:
|
|
"""Benchmark graph execution with result accumulation."""
|
|
result = self.graph.execute(self.fn_registry, initial=42)
|
|
assert result == 42, f"Expected 42, got {result}"
|