cleveragents-core/benchmarks/cost_controls_bench.py

"""ASV benchmarks for cost controls overhead.

Measures performance of:
- CostEntry creation and estimation
- ProviderCostTable lookups
- CostTracker budget checks and recording
- FallbackSelector provider selection
"""

from __future__ import annotations

import importlib
import sys
from pathlib import Path

# Ensure the local source tree is importable
_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.config.settings import Settings  # noqa: E402
from cleveragents.domain.models.core.cost_metadata import CostMetadata  # noqa: E402
from cleveragents.providers.cost_table import (  # noqa: E402
    CostEntry,
    ProviderCostTable,
)
from cleveragents.providers.cost_tracker import CostTracker  # noqa: E402
from cleveragents.providers.fallback_selector import FallbackSelector  # noqa: E402
from cleveragents.providers.registry import ProviderRegistry  # noqa: E402


class TimeCostEntry:
    """Benchmark CostEntry creation and estimation."""

    def time_create_cost_entry(self) -> None:
        CostEntry(input_cost_per_token=0.00001, output_cost_per_token=0.00003)

    def time_estimate_cost(self) -> None:
        entry = CostEntry(input_cost_per_token=0.00001, output_cost_per_token=0.00003)
        entry.estimate_cost(1000, 500)


class TimeCostTable:
    """Benchmark ProviderCostTable lookups."""

    def setup(self) -> None:
        self.table = ProviderCostTable()

    def time_lookup_known_model(self) -> None:
        self.table.get_cost_entry("openai", "gpt-4o")

    def time_lookup_unknown_model(self) -> None:
        self.table.get_cost_entry("openai", "unknown-model")

    def time_list_providers(self) -> None:
        self.table.list_providers()

    def time_list_models(self) -> None:
        self.table.list_models("openai")


class TimeCostTracker:
    """Benchmark CostTracker budget checks and recording."""

    def setup(self) -> None:
        self.tracker = CostTracker(budget_per_plan=10.0, budget_per_day=50.0)
        self.tracker_unlimited = CostTracker()

    def time_record_usage_with_budget(self) -> None:
        metadata = CostMetadata()
        self.tracker.record_usage(
            metadata,
            provider="openai",
            model="gpt-4o",
            input_tokens=1000,
            output_tokens=500,
        )

    def time_record_usage_unlimited(self) -> None:
        metadata = CostMetadata()
        self.tracker_unlimited.record_usage(
            metadata,
            provider="openai",
            model="gpt-4o",
            input_tokens=1000,
            output_tokens=500,
        )

    def time_check_plan_budget(self) -> None:
        metadata = CostMetadata()
        self.tracker.check_plan_budget(metadata)

    def time_check_daily_budget(self) -> None:
        self.tracker.check_daily_budget()

    def time_estimate_cost(self) -> None:
        self.tracker.estimate_cost("openai", "gpt-4o", 1000, 500)


class TimeFallbackSelector:
    """Benchmark FallbackSelector provider selection."""

    def setup(self) -> None:
        settings = Settings()
        self.registry = ProviderRegistry(settings)
        self.selector = FallbackSelector(registry=self.registry)

    def time_select_no_requirements(self) -> None:
        self.selector.select()

    def time_select_with_tool_calls(self) -> None:
        self.selector.select(require_tool_calls=True)

    def time_select_with_all_requirements(self) -> None:
        self.selector.select(
            require_tool_calls=True,
            require_streaming=True,
            require_vision=True,
            require_json_mode=True,
        )


class TimeCostMetadata:
    """Benchmark CostMetadata operations."""

    def time_create_metadata(self) -> None:
        CostMetadata()

    def time_record_usage(self) -> None:
        meta = CostMetadata()
        meta.record_usage(
            input_tokens=1000,
            output_tokens=500,
            cost=0.025,
            provider="openai",
        )

    def time_display_dict(self) -> None:
        meta = CostMetadata(total_tokens=1500, total_cost=0.025)
        meta.as_display_dict()