core/router.py

"""
Nancy — Provider Router.

Responsible for selecting which provider handles a given request.
Implements:
  - Circuit breaker (trip after N consecutive failures, cooldown period)
  - RPM rate limiting (sliding window per provider)
  - Fallback chains (try next provider when current is unavailable)
  - Model → provider name resolution
"""

from __future__ import annotations

import logging
from typing import Any

from config import settings
from models.provider import ProviderConfig, ProviderState

logger = logging.getLogger("nancy.router")


# ── Model-to-Provider mapping ────────────────────────────────────────────────

# Maps model name aliases to canonical provider names.
# The extension uses the provider name to know which chatbot tab to target.
MODEL_TO_PROVIDER: dict[str, str] = {
    # ChatGPT
    "chatgpt": "chatgpt",
    "gpt-4": "chatgpt",
    "gpt-4o": "chatgpt",
    "gpt-40": "chatgpt",
    "gpt-4o-mini": "chatgpt",
    "gpt-3.5-turbo": "chatgpt",
    # Gemini
    "gemini": "gemini",
    "gemini-pro": "gemini",
    "gemini-2.0-flash": "gemini",
    "gemini-2.5-pro": "gemini",
    # DeepSeek
    "deepseek": "deepseek",
    "deepseek-chat": "deepseek",
    "deepseek-r1": "deepseek",
    # Kimi
    "kimi": "kimi",
    "moonshot": "kimi",
    # Official Paid API / Hybrid models
    "mistral-large": "api-mistral",
    "nvidia-llama3": "api-nvidia-nim",
    "deepseek-api": "api-deepseek",
    "claude-api": "api-anthropic",
    "z-ai-api": "api-z-ai",
    # Browser relay: NIM Portal (free playground)
    "nim": "nim",
    "nim-llama3": "nim",
    "nim-mistral": "nim",
    # Browser relay: z.ai portal
    "zai": "zai",
    "z-ai": "zai",
    # Claude (browser portal)
    "claude": "claude",
    "claude-3": "claude",
}


class ProviderRouter:
    """
    Selects providers for incoming requests with circuit breaking,
    rate limiting, and fallback chain support.

    Usage::

        router = ProviderRouter()
        provider = router.resolve("gpt-4o")        # → "chatgpt"
        available = router.select_provider("chatgpt")  # checks CB + RPM
        router.record_success("chatgpt")
        router.record_failure("chatgpt")
    """

    def __init__(self) -> None:
        self._states: dict[str, ProviderState] = {}
        self._initialize_providers()

    def _initialize_providers(self) -> None:
        """Build ProviderState objects from configuration."""
        for name, raw_config in settings.providers_config.items():
            config = ProviderConfig(**raw_config) if isinstance(raw_config, dict) else ProviderConfig()
            self._states[name] = ProviderState(name, config)
            logger.info(
                "Provider '%s' initialized (rpm=%d, tpm=%d)",
                name,
                config.rpm,
                config.tpm,
            )

        # Ensure all fallback chain providers exist
        for name in settings.fallback_chain:
            if name not in self._states:
                self._states[name] = ProviderState(name, ProviderConfig())
                logger.info("Provider '%s' added from fallback chain with defaults", name)

    # ── Resolution ────────────────────────────────────────────────────

    def resolve(self, model: str) -> str:
        """
        Resolve a model name to a canonical provider name.

        Falls back to ``settings.default_provider`` if the model is unknown.
        """
        provider = MODEL_TO_PROVIDER.get(model.lower(), model.lower())
        # If the resolved name is a known provider, use it
        if provider in self._states:
            return provider
        # Otherwise fall back to default
        logger.debug(
            "Unknown model '%s' → defaulting to '%s'",
            model,
            settings.default_provider,
        )
        return settings.default_provider

    # ── Provider Selection with Circuit Breaker + Rate Limit ──────────

    def select_provider(
        self,
        preferred: str,
        exclude: set[str] | None = None,
    ) -> str | None:
        """
        Select the best available provider.

        1. Try the preferred provider first.
        2. If it's unavailable (circuit open, rate limited), walk the fallback chain.
        3. Return ``None`` if no provider is available.

        Args:
            preferred: The preferred provider name.
            exclude: Set of provider names to skip (already tried and failed).
        """
        exclude = exclude or set()
        candidates = [preferred] + [
            p for p in settings.fallback_chain if p != preferred
        ]

        for name in candidates:
            if name in exclude:
                continue

            state = self._states.get(name)
            if not state:
                continue

            # Check circuit breaker
            if not state.should_allow_request(
                settings.cb_failure_threshold,
                settings.cb_cooldown_seconds,
            ):
                logger.debug("Provider '%s' circuit is OPEN — skipping", name)
                continue

            # Check rate limit
            if not state.check_rate_limit():
                logger.debug("Provider '%s' rate limited — skipping", name)
                continue

            # Record the request
            state.record_request()
            logger.info("Selected provider: '%s'", name)
            return name

        logger.error(
            "No available provider (preferred=%s, exclude=%s)",
            preferred,
            exclude,
        )
        return None

    # ── Feedback ──────────────────────────────────────────────────────

    def record_success(self, provider: str) -> None:
        """Record a successful completion for the given provider."""
        state = self._states.get(provider)
        if state:
            state.record_success()
            logger.debug("Provider '%s' success recorded", provider)

    def record_failure(self, provider: str) -> None:
        """Record a failure for the given provider (may trip circuit)."""
        state = self._states.get(provider)
        if state:
            state.record_failure(settings.cb_failure_threshold)
            logger.warning(
                "Provider '%s' failure recorded (consecutive=%d, circuit=%s)",
                provider,
                state.consecutive_failures,
                state.circuit_state.value,
            )

    # ── Observability ─────────────────────────────────────────────────

    def get_provider_states(self) -> list[dict[str, Any]]:
        """Return status dicts for all providers."""
        return [state.to_dict() for state in self._states.values()]

    def get_available_models(self) -> list[str]:
        """Return list of all recognized model names."""
        return sorted(MODEL_TO_PROVIDER.keys())

    def get_available_providers(self) -> list[str]:
        """Return list of all configured provider names."""
        return sorted(self._states.keys())

    def is_provider_available(self, provider: str) -> bool:
        """Check if a specific provider is currently available."""
        state = self._states.get(provider)
        if not state:
            return False
        return (
            state.should_allow_request(
                settings.cb_failure_threshold,
                settings.cb_cooldown_seconds,
            )
            and state.check_rate_limit()
        )

    def is_provider_healthy(self, provider: str) -> bool:
        """Alias for is_provider_available for admin dashboard compatibility."""
        return self.is_provider_available(provider)

    @property
    def fallback_chain(self) -> list[str]:
        """Return the active fallback chain of providers."""
        return settings.fallback_chain


# Module-level singleton
provider_router = ProviderRouter()