cleveragents-core/benchmarks/providers_fallback_selector_bench.py

"""ASV benchmarks for FallbackSelector with varied configurations.

Measures the performance of:
- FallbackSelector construction with custom fallback order
- FallbackSelector construction with an injected CostTracker
- Selection when all providers are unconfigured (exhausts full list)
- Selection with a custom provider order
- Selection with a cost tracker attached

Note: Basic select() calls (no requirements, with tool_calls, with all
requirements) against a default registry 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
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.providers.cost_tracker import CostTracker  # noqa: E402
from cleveragents.providers.fallback_selector import FallbackSelector  # noqa: E402
from cleveragents.providers.registry import (  # noqa: E402
    ProviderRegistry,
    reset_provider_registry,
)


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_settings(**overrides: object) -> MagicMock:
    """Build a minimal mock Settings object for registry construction."""
    settings = MagicMock()
    settings.openai_api_key = None
    settings.anthropic_api_key = None
    settings.google_api_key = None
    settings.gemini_api_key = None
    settings.azure_api_key = None
    settings.openrouter_api_key = None
    settings.cohere_api_key = None
    settings.groq_api_key = None
    settings.together_api_key = None
    settings.default_provider = None
    settings.default_model = None
    settings.azure_openai_endpoint = None
    settings.azure_openai_api_version = None
    settings.azure_openai_deployment = None
    settings.openrouter_organization = None
    settings.mock_providers = False
    for key, value in overrides.items():
        setattr(settings, key, value)
    return settings


class TimeFallbackSelectorConstruction:
    """Benchmark FallbackSelector construction with various configurations."""

    def setup(self) -> None:
        reset_provider_registry()
        self.registry = ProviderRegistry(settings=_make_settings())
        self.tracker = CostTracker(budget_per_plan=10.0, budget_per_day=50.0)

    def time_construct_default_order(self) -> None:
        """Time construction with the default fallback order."""
        FallbackSelector(registry=self.registry)

    def time_construct_custom_order(self) -> None:
        """Time construction with a custom provider fallback order."""
        FallbackSelector(
            registry=self.registry,
            fallback_providers=["anthropic", "openai", "groq"],
        )

    def time_construct_with_cost_tracker(self) -> None:
        """Time construction with a CostTracker injected."""
        FallbackSelector(registry=self.registry, cost_tracker=self.tracker)

    def time_construct_single_provider_order(self) -> None:
        """Time construction with a single-provider fallback list."""
        FallbackSelector(
            registry=self.registry,
            fallback_providers=["openai"],
        )


class TimeFallbackSelectorNoProviders:
    """Benchmark selection when no providers are configured (exhausts full list)."""

    def setup(self) -> None:
        reset_provider_registry()
        # No API keys set — all providers unconfigured
        self.registry = ProviderRegistry(settings=_make_settings())
        self.selector = FallbackSelector(registry=self.registry)
        self.selector_short = FallbackSelector(
            registry=self.registry,
            fallback_providers=["openai", "anthropic"],
        )

    def time_select_exhausts_all_providers(self) -> None:
        """Time selection that must iterate through all providers and find none."""
        self.selector.select()

    def time_select_exhausts_short_list(self) -> None:
        """Time selection that exhausts a short 2-provider list."""
        self.selector_short.select()

    def time_select_with_requirements_no_providers(self) -> None:
        """Time selection with capability requirements when no providers exist."""
        self.selector.select(
            require_tool_calls=True,
            require_streaming=True,
        )


class TimeFallbackSelectorWithConfiguredProvider:
    """Benchmark selection when one provider is configured."""

    def setup(self) -> None:
        reset_provider_registry()
        # OpenAI configured, all others not
        self.registry = ProviderRegistry(
            settings=_make_settings(openai_api_key="sk-bench-key")
        )
        self.selector_default = FallbackSelector(registry=self.registry)
        # Custom order: openai is last — must skip all others first
        self.selector_openai_last = FallbackSelector(
            registry=self.registry,
            fallback_providers=["anthropic", "google", "groq", "openai"],
        )
        self.tracker = CostTracker(budget_per_plan=100.0)
        self.selector_with_tracker = FallbackSelector(
            registry=self.registry,
            cost_tracker=self.tracker,
        )

    def time_select_first_provider_matches(self) -> None:
        """Time selection where the first provider in the list is configured."""
        self.selector_default.select()

    def time_select_last_provider_matches(self) -> None:
        """Time selection where the configured provider is last in the list."""
        self.selector_openai_last.select()

    def time_select_with_cost_tracker_under_budget(self) -> None:
        """Time selection with a cost tracker when budget is not exceeded."""
        self.selector_with_tracker.select()

    def time_select_vision_requirement_openai(self) -> None:
        """Time selection requiring vision support (OpenAI supports it)."""
        self.selector_default.select(require_vision=True)

    def time_select_json_mode_requirement_openai(self) -> None:
        """Time selection requiring JSON mode (OpenAI supports it)."""
        self.selector_default.select(require_json_mode=True)