cleveragents-core/benchmarks/providers_cost_table_bench.py

"""ASV benchmarks for ProviderCostTable construction and iteration.

Measures the performance of:
- ProviderCostTable construction with default entries
- ProviderCostTable construction with custom entries merged in
- Iterating all providers and their models (throughput)
- CostEntry.estimate_cost across the full default table

Note: Basic lookup benchmarks (get_cost_entry, list_providers, list_models)
are already covered by cost_controls_bench.py and are not duplicated here.
"""

from __future__ import annotations

import importlib
import sys
from pathlib import Path

_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.providers.cost_table import CostEntry, ProviderCostTable  # noqa: E402

# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------

_CUSTOM_ENTRIES: dict[str, dict[str, CostEntry]] = {
    "custom_provider": {
        "custom-model-v1": CostEntry(
            input_cost_per_token=0.000002,
            output_cost_per_token=0.000008,
        ),
        "custom-model-v2": CostEntry(
            input_cost_per_token=0.000001,
            output_cost_per_token=0.000004,
        ),
    },
    "openai": {
        "gpt-5": CostEntry(
            input_cost_per_token=0.00001,
            output_cost_per_token=0.00003,
        ),
    },
}


class TimeCostTableConstruction:
    """Benchmark ProviderCostTable construction cost."""

    def time_construct_default(self) -> None:
        """Time construction of a cost table with only default entries."""
        ProviderCostTable()

    def time_construct_with_custom_entries(self) -> None:
        """Time construction of a cost table with custom entries merged in."""
        ProviderCostTable(custom_entries=_CUSTOM_ENTRIES)

    def time_construct_with_empty_custom(self) -> None:
        """Time construction with an empty custom entries dict (no-op merge)."""
        ProviderCostTable(custom_entries={})


class TimeCostTableIteration:
    """Benchmark iterating all providers and models in the cost table."""

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

    def time_iterate_all_providers_and_models(self) -> None:
        """Time iterating every provider and every model in the default table."""
        for provider in self.table.list_providers():
            for model in self.table.list_models(provider):
                self.table.get_cost_entry(provider, model)

    def time_iterate_all_providers_and_models_custom(self) -> None:
        """Time iterating every provider and model in a table with custom entries."""
        for provider in self.table_with_custom.list_providers():
            for model in self.table_with_custom.list_models(provider):
                self.table_with_custom.get_cost_entry(provider, model)

    def time_estimate_cost_all_entries(self) -> None:
        """Time calling estimate_cost for every known provider/model pair."""
        for provider in self.table.list_providers():
            for model in self.table.list_models(provider):
                entry = self.table.get_cost_entry(provider, model)
                entry.estimate_cost(input_tokens=1000, output_tokens=500)


class TimeCostTableFallback:
    """Benchmark the fallback path for unknown provider/model combinations."""

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

    def time_lookup_unknown_provider(self) -> None:
        """Time lookup for a completely unknown provider (returns DEFAULT_COST)."""
        self.table.get_cost_entry("nonexistent_provider", "nonexistent_model")

    def time_lookup_known_provider_unknown_model(self) -> None:
        """Time lookup for a known provider but unknown model (returns DEFAULT_COST)."""
        self.table.get_cost_entry("anthropic", "claude-99-ultra")

    def time_default_cost_estimate(self) -> None:
        """Time estimating cost using the DEFAULT_COST fallback entry."""
        ProviderCostTable.DEFAULT_COST.estimate_cost(
            input_tokens=10_000, output_tokens=5_000
        )