cleveragents-core/benchmarks/providers_cost_tracker_bench.py

"""ASV benchmarks for CostTracker accumulation and daily spend tracking.

Measures the performance of:
- CostTracker construction with various budget configurations
- Accumulating many record_usage calls (throughput)
- Daily spend tracking and retrieval
- get_cost_entry delegation to the underlying cost table
- Budget evaluation across multiple accumulated calls

Note: Single-call record_usage, check_plan_budget, check_daily_budget, and
estimate_cost 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.domain.models.core.cost_metadata import CostMetadata  # noqa: E402
from cleveragents.providers.cost_table import CostEntry, ProviderCostTable  # noqa: E402
from cleveragents.providers.cost_tracker import CostTracker  # noqa: E402


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

_CUSTOM_TABLE = ProviderCostTable(
    custom_entries={
        "bench_provider": {
            "bench-model": CostEntry(
                input_cost_per_token=0.000001,
                output_cost_per_token=0.000002,
            )
        }
    }
)

_PROVIDERS_AND_MODELS = [
    ("openai", "gpt-4o"),
    ("anthropic", "claude-sonnet-4-20250514"),
    ("google", "gemini-2.0-flash"),
    ("groq", "llama-3.1-70b-versatile"),
    ("cohere", "command-r-plus"),
]


class TimeCostTrackerConstruction:
    """Benchmark CostTracker construction with various budget configurations."""

    def time_construct_no_budget(self) -> None:
        """Time construction with no budget limits (unlimited)."""
        CostTracker()

    def time_construct_plan_budget_only(self) -> None:
        """Time construction with only a per-plan budget."""
        CostTracker(budget_per_plan=10.0)

    def time_construct_daily_budget_only(self) -> None:
        """Time construction with only a per-day budget."""
        CostTracker(budget_per_day=50.0)

    def time_construct_both_budgets(self) -> None:
        """Time construction with both plan and daily budgets."""
        CostTracker(budget_per_plan=10.0, budget_per_day=50.0)

    def time_construct_with_custom_cost_table(self) -> None:
        """Time construction with a pre-built custom cost table injected."""
        CostTracker(budget_per_plan=5.0, cost_table=_CUSTOM_TABLE)


class TimeCostTrackerAccumulation:
    """Benchmark CostTracker accumulation throughput across many calls."""

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

    def time_accumulate_10_calls(self) -> None:
        """Time accumulating 10 record_usage calls on a fresh CostMetadata."""
        metadata = CostMetadata()
        for _ in range(10):
            self.tracker.record_usage(
                metadata,
                provider="openai",
                model="gpt-4o",
                input_tokens=500,
                output_tokens=250,
            )

    def time_accumulate_50_calls(self) -> None:
        """Time accumulating 50 record_usage calls on a fresh CostMetadata."""
        metadata = CostMetadata()
        for _ in range(50):
            self.tracker.record_usage(
                metadata,
                provider="openai",
                model="gpt-4o",
                input_tokens=500,
                output_tokens=250,
            )

    def time_accumulate_mixed_providers(self) -> None:
        """Time accumulating calls across multiple providers and models."""
        metadata = CostMetadata()
        for provider, model in _PROVIDERS_AND_MODELS:
            self.tracker.record_usage(
                metadata,
                provider=provider,
                model=model,
                input_tokens=1000,
                output_tokens=500,
            )

    def time_accumulate_unlimited_budget(self) -> None:
        """Time accumulating 10 calls with no budget limits (no threshold checks)."""
        metadata = CostMetadata()
        for _ in range(10):
            self.tracker_unlimited.record_usage(
                metadata,
                provider="openai",
                model="gpt-4o",
                input_tokens=500,
                output_tokens=250,
            )


class TimeCostTrackerDailySpend:
    """Benchmark daily spend tracking and retrieval."""

    def setup(self) -> None:
        self.tracker = CostTracker(budget_per_day=100.0)
        metadata = CostMetadata()
        # Pre-populate with some daily spend
        for _ in range(5):
            self.tracker.record_usage(
                metadata,
                provider="openai",
                model="gpt-4o",
                input_tokens=1000,
                output_tokens=500,
            )

    def time_get_daily_spend(self) -> None:
        """Time retrieving the current day's accumulated spend."""
        self.tracker.get_daily_spend()

    def time_check_daily_budget_with_spend(self) -> None:
        """Time checking daily budget after some spend has been accumulated."""
        self.tracker.check_daily_budget()


class TimeCostTrackerGetCostEntry:
    """Benchmark CostTracker.get_cost_entry delegation."""

    def setup(self) -> None:
        self.tracker = CostTracker()

    def time_get_cost_entry_known(self) -> None:
        """Time get_cost_entry for a known provider/model pair."""
        self.tracker.get_cost_entry("openai", "gpt-4o")

    def time_get_cost_entry_unknown(self) -> None:
        """Time get_cost_entry for an unknown provider/model (returns default)."""
        self.tracker.get_cost_entry("unknown_provider", "unknown_model")

    def time_get_cost_entry_all_known(self) -> None:
        """Time get_cost_entry for all known provider/model pairs."""
        for provider, model in _PROVIDERS_AND_MODELS:
            self.tracker.get_cost_entry(provider, model)