cleveragents-core/benchmarks/llm_trace_bench.py

"""ASV benchmarks for LLM trace observability overhead.

Measures the time to create LLMTrace models, record traces via
TraceService, compute metrics, and exercise lifecycle hooks.
"""

from __future__ import annotations

import importlib
import sys
from pathlib import Path
from unittest.mock import MagicMock

_SRC = str(Path(__file__).resolve().parents[1] / "src")
if _SRC not in sys.path:
    sys.path.insert(0, _SRC)

import cleveragents  # noqa: E402

importlib.reload(cleveragents)

from cleveragents.application.services.trace_service import TraceService  # noqa: E402
from cleveragents.domain.models.observability.llm_trace import LLMTrace  # noqa: E402
from cleveragents.domain.models.observability.metrics import (  # noqa: E402
    MetricCollector,
    OperationalMetricKey,
)

VALID_PLAN_ID = "01HX0000000000PPPPPPPPPPPP"


class InMemoryRepo:
    """Minimal in-memory repo for benchmarking."""

    def __init__(self) -> None:
        self._store: dict[str, LLMTrace] = {}

    def save(self, trace: LLMTrace) -> None:
        self._store[trace.trace_id] = trace

    def get(self, trace_id: str) -> LLMTrace | None:
        return self._store.get(trace_id)

    def list_by_plan(self, plan_id: str) -> list[LLMTrace]:
        return [t for t in self._store.values() if t.plan_id == plan_id]

    def list_by_decision(self, decision_id: str) -> list[LLMTrace]:
        return [t for t in self._store.values() if t.decision_id == decision_id]


def _make_trace(idx: int) -> LLMTrace:
    idx_str = f"{idx:08d}"
    ulid = f"01HX{'A' * (22 - len(idx_str))}{idx_str}"
    return LLMTrace(
        trace_id=ulid,
        plan_id=VALID_PLAN_ID,
        actor="planner",
        provider="openai",
        model="gpt-4o",
        prompt_tokens=500 + idx,
        completion_tokens=200 + idx,
        cost_usd=0.01 + idx * 0.001,
        latency_ms=300.0 + idx,
    )


class LLMTraceModelCreationSuite:
    """Benchmark LLMTrace model creation."""

    def time_create_single_trace(self) -> None:
        """Create a single LLMTrace model."""
        _make_trace(0)

    def time_create_100_traces(self) -> None:
        """Create 100 LLMTrace models."""
        for i in range(100):
            _make_trace(i)


class LLMTraceServiceSuite:
    """Benchmark TraceService operations."""

    def setup(self) -> None:
        self._svc = TraceService(
            settings=MagicMock(),
            repository=InMemoryRepo(),  # type: ignore[arg-type]
        )
        self._traces = [_make_trace(i) for i in range(100)]

    def time_record_100_traces(self) -> None:
        """Record 100 traces."""
        for t in self._traces:
            self._svc.record_trace(t)

    def time_compute_metrics_100(self) -> None:
        """Compute metrics over 100 traces."""
        for t in self._traces:
            self._svc.record_trace(t)
        self._svc.compute_metrics(VALID_PLAN_ID)


class MetricCollectorSuite:
    """Benchmark MetricCollector operations."""

    def time_record_metric(self) -> None:
        """Create a single metric entry."""
        MetricCollector.record(
            OperationalMetricKey.PLAN_DURATION_MS,
            1500.0,
            VALID_PLAN_ID,
        )

    def time_plan_duration_convenience(self) -> None:
        """Use plan_duration convenience method."""
        MetricCollector.plan_duration(VALID_PLAN_ID, 1500.0)

    def time_lifecycle_hooks(self) -> None:
        """Exercise all lifecycle hooks."""
        svc = TraceService(
            settings=MagicMock(),
            repository=InMemoryRepo(),  # type: ignore[arg-type]
        )
        svc.on_plan_start(VALID_PLAN_ID)
        svc.on_actor_invocation(VALID_PLAN_ID, "planner", 250.0)
        svc.on_tool_execution(VALID_PLAN_ID, "file_write", errored=False)