rmi-backend/app/databus/provider_core.py

"""
DataBus Provider Core - Infrastructure & Base Classes
======================================================

Circuit breakers, rate limiters, quota tracking, and core provider classes.
This module contains NO external API implementations.
"""

import logging
import os
import time
from collections.abc import Callable
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

logger = logging.getLogger("databus.providers.core")


class ProviderTier(Enum):
    """Provider tiers: LOCAL > FREE_API > FREEMIUM > PAID"""

    LOCAL = "local"  # Our own data - instant, free, unlimited
    FREE_API = "free_api"  # Free external API - no key needed
    FREEMIUM = "freemium"  # Free tier with key - limited credits
    PAID = "paid"  # Paid API - precious credits


@dataclass
class Provider:
    """A single data source in a fallback chain."""

    name: str
    tier: ProviderTier
    fetch_fn: Callable = field(repr=False)
    weight: float = 1.0  # Higher = preferred within tier
    rate_limit_rps: float = 1.0
    monthly_quota: int = 0  # 0 = unlimited
    requires_key: bool = False
    key_env: str = ""
    timeout: float = 15.0
    is_local: bool = False  # True if this provider uses our own data (no external API)
    description: str = ""  # Human-readable description
    # Circuit breaker
    failure_threshold: int = 5
    recovery_timeout: float = 60.0


@dataclass
class ProviderChain:
    """A fallback chain for a specific data type."""

    data_type: str
    providers: list[Provider]
    description: str = ""

    async def fetch(self, vault: Any = None, cache: Any = None, **kwargs: Any) -> Any | None:
        """Try each provider in order until one succeeds.

        Smart fallback: when paid provider quota is >80% used, skip to free/local
        alternatives first to conserve credits for critical queries.
        """
        providers_sorted = sorted(self.providers, key=lambda p: (-p.weight, p.tier.value))

        # ── Credit pressure: if paid providers are near quota, bump free providers up ──
        credit_pressure = False
        for p in providers_sorted:
            if p.monthly_quota > 0 and p.tier.value in ("paid", "freemium"):
                used = _quota_usage.get(p.name, 0)
                if used > p.monthly_quota * 0.8:  # 80% threshold
                    credit_pressure = True
                    logger.info(
                        f"Credit pressure: {p.name} at {used}/{p.monthly_quota} ({used * 100 // p.monthly_quota}%)"
                    )

        if credit_pressure:
            # Re-sort: push free/local providers above paid/freemium near quota
            providers_sorted.sort(key=lambda p: (0 if p.tier.value in ("local", "free_api") else 1, -p.weight))

        for provider in providers_sorted:
            # Check circuit breaker
            if not _circuit_breakers.get(provider.name, _CircuitBreaker()).can_call():
                logger.debug(f"Circuit breaker open for {provider.name}")
                continue

            # Check rate limit
            if not _rate_limiters.get(provider.name, _RateLimiter()).can_call():
                logger.debug(f"Rate limit exceeded for {provider.name}")
                continue

            # Check quota
            if provider.monthly_quota > 0:
                used = _quota_usage.get(provider.name, 0)
                if used >= provider.monthly_quota:
                    logger.debug(f"Monthly quota exceeded for {provider.name}")
                    continue

            try:
                # Get API key from env (vault is pool manager, use os.getenv for direct keys)
                api_key = None
                if provider.requires_key and provider.key_env:
                    api_key = os.getenv(provider.key_env, "")

                result = await provider.fetch_fn(api_key=api_key, **kwargs)

                if result is not None:
                    _rate_limiters[provider.name].record_call()
                    if provider.monthly_quota > 0:
                        _quota_usage[provider.name] = _quota_usage.get(provider.name, 0) + 1
                    _circuit_breakers[provider.name].record_success()
                    return result

            except Exception as e:
                logger.warning(f"Provider {provider.name} failed: {e}")
                _circuit_breakers[provider.name].record_failure()
                continue

        return None


# ── Circuit Breaker ────────────────────────────────────────────


class _CircuitBreaker:
    """Circuit breaker to prevent cascading failures."""

    def __init__(self, threshold: int = 5, timeout: float = 60.0):
        self.threshold = threshold
        self.timeout = timeout
        self.failures = 0
        self.last_failure = 0.0
        self.open = False

    def can_call(self) -> bool:
        if self.open:
            if time.time() - self.last_failure > self.timeout:
                self.open = False
                self.failures = 0
                return True
            return False
        return True

    def record_failure(self) -> None:
        self.failures += 1
        self.last_failure = time.time()
        if self.failures >= self.threshold:
            self.open = True

    def record_success(self) -> None:
        self.failures = 0
        self.open = False


class _RateLimiter:
    """Simple rate limiter based on time intervals."""

    def __init__(self, rps: float = 1.0):
        self.rps = rps
        self.min_interval = 1.0 / rps
        self.last_call = 0.0

    def can_call(self) -> bool:
        return time.time() - self.last_call >= self.min_interval

    def record_call(self) -> None:
        self.last_call = time.time()


# ── Shared State ───────────────────────────────────────────────

_circuit_breakers: dict[str, _CircuitBreaker] = {}
_rate_limiters: dict[str, _RateLimiter] = {}
_quota_usage: dict[str, int] = {}


def reset_state() -> None:
    """Reset all circuit breakers, rate limiters, and quota tracking.

    Useful for testing and debugging.
    """
    global _circuit_breakers, _rate_limiters, _quota_usage
    _circuit_breakers = {}
    _rate_limiters = {}
    _quota_usage = {}