forked from cleveragents/cleveragents-core
106 lines
2.8 KiB
Python
106 lines
2.8 KiB
Python
"""Airspeed Velocity benchmarks for plan generation workflows.
|
|
|
|
Measures invoke and streaming performance of PlanGenerationGraph using
|
|
built-in FakeListLLM responses and minimal domain models.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import tiktoken
|
|
|
|
from cleveragents.agents.graphs.plan_generation import PlanGenerationGraph
|
|
from cleveragents.domain.models.core import (
|
|
Context,
|
|
ContextType,
|
|
Plan,
|
|
PlanStatus,
|
|
Project,
|
|
)
|
|
from cleveragents.domain.models.core.project import ProjectSettings
|
|
|
|
|
|
def _sample_project() -> Project:
|
|
return Project(
|
|
id=1,
|
|
name="bench-project",
|
|
path=Path(".").resolve(),
|
|
settings=ProjectSettings(
|
|
auto_build=False,
|
|
auto_apply=False,
|
|
confirm_apply=True,
|
|
max_context_size=52_428_800,
|
|
default_model="mock-gpt",
|
|
),
|
|
current_plan_id=1,
|
|
)
|
|
|
|
|
|
def _sample_plan() -> Plan:
|
|
return Plan(
|
|
id=1,
|
|
project_id=1,
|
|
name="bench-plan",
|
|
prompt="Add error handling",
|
|
status=PlanStatus.PENDING,
|
|
current=True,
|
|
build=None,
|
|
build_started_at=None,
|
|
build_completed_at=None,
|
|
model_used=None,
|
|
token_count=None,
|
|
result=None,
|
|
applied_at=None,
|
|
files_created=None,
|
|
files_modified=None,
|
|
files_deleted=None,
|
|
)
|
|
|
|
|
|
def _sample_contexts() -> list[Context]:
|
|
return [
|
|
Context(
|
|
id=1,
|
|
plan_id=1,
|
|
type=ContextType.FILE,
|
|
path="src/example.py",
|
|
content="print('hello')",
|
|
file_hash=None,
|
|
size=18,
|
|
)
|
|
]
|
|
|
|
|
|
class PlanGenerationSuite:
|
|
"""Benchmark PlanGenerationGraph invoke and stream paths."""
|
|
|
|
def setup(self) -> None:
|
|
self.graph = PlanGenerationGraph()
|
|
self.project = _sample_project()
|
|
self.plan = _sample_plan()
|
|
self.contexts = _sample_contexts()
|
|
|
|
enc = tiktoken.get_encoding("cl100k_base")
|
|
self.prompt_tokens = {
|
|
"analyze": len(enc.encode(self.graph.analyze_prompt.template)),
|
|
"generate": len(enc.encode(self.graph.generate_prompt.template)),
|
|
"validate": len(enc.encode(self.graph.validate_prompt.template)),
|
|
}
|
|
|
|
def time_invoke(self) -> None:
|
|
self.graph.invoke(self.project, self.plan, self.contexts, thread_id="bench")
|
|
|
|
def time_stream(self) -> None:
|
|
for _ in self.graph.stream(self.project, self.plan, self.contexts):
|
|
pass
|
|
|
|
def track_analyze_prompt_tokens(self) -> int:
|
|
return self.prompt_tokens["analyze"]
|
|
|
|
def track_generate_prompt_tokens(self) -> int:
|
|
return self.prompt_tokens["generate"]
|
|
|
|
def track_validate_prompt_tokens(self) -> int:
|
|
return self.prompt_tokens["validate"]
|