tracer-cloud--opensre/tests/benchmarks/toolcall_model_benchmark/pipeline_benchmark.py

"""Investigation pipeline benchmark: baseline vs split tool LLM routing.

Uses the same ``run_investigation`` entry point as production, injects synthetic
Grafana fixtures (no Tracer JWT required), and tracks token usage by
instrumenting API creates.
"""

from __future__ import annotations

import time
from collections.abc import Callable
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any

from core.llm import factory
from core.llm.factory import LLMRole, get_llm, reset_llm_clients
from core.llm.transports.sdk.llm_clients import LLMClient, OpenAILLMClient
from core.state import AgentState
from tests.benchmarks.toolcall_model_benchmark.pricing import estimate_run_cost_usd
from tests.synthetic.rds_postgres.run_suite import _build_resolved_integrations
from tests.synthetic.rds_postgres.scenario_loader import (
    SUITE_DIR,
    ScenarioFixture,
    load_all_scenarios,
)
from tools.investigation.capability import run_investigation
from tools.investigation.state_factory import make_initial_state

__all__ = [
    "LLMCallRecord",
    "InvestigationBenchmarkRun",
    "TokenTotals",
    "configure_baseline_reasoning_for_tools",
    "configure_split_models",
    "estimate_run_cost_usd",
    "get_fixture_by_id",
    "make_investigation_state",
    "reset_llm_singletons",
    "run_investigation_bench",
]


@dataclass
class TokenTotals:
    """Accumulated LLM token usage (all providers summed into one counter)."""

    input_tokens: int = 0
    output_tokens: int = 0

    @property
    def total(self) -> int:
        return self.input_tokens + self.output_tokens


@dataclass(frozen=True)
class LLMCallRecord:
    """One provider API completion (after retries succeed for that attempt)."""

    call_index: int
    model_id: str
    input_tokens: int
    output_tokens: int
    latency_sec: float

    @property
    def total_tokens(self) -> int:
        return self.input_tokens + self.output_tokens


@dataclass
class InvestigationBenchmarkRun:
    label: str
    wall_seconds: float
    tokens: TokenTotals
    tokens_by_model: dict[str, TokenTotals] = field(default_factory=dict)
    llm_calls: list[LLMCallRecord] = field(default_factory=list)
    final_state: dict[str, Any] = field(default_factory=dict)


def reset_llm_singletons() -> None:
    """Reset the cached LLM clients; delegates to core.llm.factory.reset_llm_clients."""
    reset_llm_clients()


def configure_split_models() -> None:
    """Production-style split: reasoning model + separate toolcall model."""
    reset_llm_singletons()
    get_llm(LLMRole.REASONING)
    get_llm(LLMRole.TOOLCALL)


def configure_baseline_reasoning_for_tools() -> None:
    """Ablation: tool nodes use the same client instance as reasoning."""
    reset_llm_singletons()
    reasoning = get_llm(LLMRole.REASONING)
    factory._cache.store(LLMRole.TOOLCALL, reasoning)


def make_investigation_state(fixture: ScenarioFixture) -> AgentState:
    alert = fixture.alert
    return make_initial_state(raw_alert=alert)


def _add_usage(
    totals: TokenTotals,
    by_model: dict[str, TokenTotals],
    model_key: str,
    inp: int,
    out: int,
) -> None:
    totals.input_tokens += inp
    totals.output_tokens += out
    bucket = by_model.setdefault(model_key, TokenTotals())
    bucket.input_tokens += inp
    bucket.output_tokens += out


@contextmanager
def _track_llm_usage(
    totals: TokenTotals,
    tokens_by_model: dict[str, TokenTotals],
    llm_calls: list[LLMCallRecord],
) -> Any:
    """Count tokens and record one row per successful API completion (incl. latency)."""
    orig_anthropic = LLMClient.invoke
    orig_openai = OpenAILLMClient.invoke

    def anthropic_invoke(self: Any, prompt_or_messages: Any) -> Any:
        self._ensure_client()
        real_create = self._client.messages.create

        def wrapped_create(*args: Any, **kwargs: Any) -> Any:
            t0 = time.perf_counter()
            resp = real_create(*args, **kwargs)
            lat = time.perf_counter() - t0
            u = getattr(resp, "usage", None)
            if u is not None:
                inp = int(getattr(u, "input_tokens", 0) or 0)
                out = int(getattr(u, "output_tokens", 0) or 0)
                model_key = str(kwargs.get("model") or getattr(self, "_model", "") or "unknown")
                _add_usage(totals, tokens_by_model, model_key, inp, out)
                llm_calls.append(
                    LLMCallRecord(
                        call_index=len(llm_calls) + 1,
                        model_id=model_key,
                        input_tokens=inp,
                        output_tokens=out,
                        latency_sec=lat,
                    )
                )
            return resp

        self._client.messages.create = wrapped_create  # type: ignore[method-assign]
        try:
            return orig_anthropic(self, prompt_or_messages)
        finally:
            self._client.messages.create = real_create  # type: ignore[method-assign]

    def openai_invoke(self: Any, prompt_or_messages: Any) -> Any:
        self._ensure_client()
        real_create = self._client.chat.completions.create

        def wrapped_create(*args: Any, **kwargs: Any) -> Any:
            t0 = time.perf_counter()
            resp = real_create(*args, **kwargs)
            lat = time.perf_counter() - t0
            u = resp.usage
            if u is not None:
                inp = int(u.prompt_tokens or 0)
                out = int(u.completion_tokens or 0)
                model_key = str(kwargs.get("model") or getattr(self, "_model", "") or "unknown")
                _add_usage(totals, tokens_by_model, model_key, inp, out)
                llm_calls.append(
                    LLMCallRecord(
                        call_index=len(llm_calls) + 1,
                        model_id=model_key,
                        input_tokens=inp,
                        output_tokens=out,
                        latency_sec=lat,
                    )
                )
            return resp

        self._client.chat.completions.create = wrapped_create  # type: ignore[method-assign]
        try:
            return orig_openai(self, prompt_or_messages)
        finally:
            self._client.chat.completions.create = real_create  # type: ignore[method-assign]

    LLMClient.invoke = anthropic_invoke  # type: ignore[assignment]
    OpenAILLMClient.invoke = openai_invoke  # type: ignore[assignment]
    try:
        yield
    finally:
        LLMClient.invoke = orig_anthropic  # type: ignore[assignment]
        OpenAILLMClient.invoke = orig_openai  # type: ignore[assignment]


def run_investigation_bench(
    fixture: ScenarioFixture,
    *,
    label: str,
    configure_llm: Callable[[], None],
) -> InvestigationBenchmarkRun:
    """Run one full investigation for a synthetic RDS scenario; return timing + tokens."""
    resolved = _build_resolved_integrations(fixture, use_mock_grafana=True)
    assert resolved is not None
    alert = fixture.alert
    totals = TokenTotals()
    tokens_by_model: dict[str, TokenTotals] = {}
    llm_calls: list[LLMCallRecord] = []

    configure_llm()

    with _track_llm_usage(totals, tokens_by_model, llm_calls):
        t0 = time.perf_counter()
        out = run_investigation(alert, resolved_integrations=resolved)
        elapsed = time.perf_counter() - t0

    out_dict = dict(out) if isinstance(out, dict) else {}

    return InvestigationBenchmarkRun(
        label=label,
        wall_seconds=elapsed,
        tokens=totals,
        tokens_by_model=dict(tokens_by_model),
        llm_calls=list(llm_calls),
        final_state=out_dict,
    )


def get_fixture_by_id(scenario_id: str) -> ScenarioFixture:
    fixtures = load_all_scenarios(SUITE_DIR)
    for f in fixtures:
        if f.scenario_id == scenario_id:
            return f
    raise ValueError(
        f"Unknown scenario_id: {scenario_id!r}. Loaded {len(fixtures)} scenarios from {SUITE_DIR}."
    )