""" Rate Tracker — lightweight in-memory sliding window rate limiter per provider. Features: - Configurable per-provider limits (requests per minute) for free-tier providers. - Sliding window using time-bucketed counters for efficiency. - Async-safe with per-provider locks. - Exposes: - `record_call(provider_name)` — record an API call timestamp - `can_call(provider_name)` — whether a call is allowed now - `get_available_providers()` — returns providers sorted by availability (most free capacity first) - `get_status()` — debug snapshot of usage counts and limits This is intended as a lightweight in-memory tracker suitable for single-process deployments. For multi-process / multi-host production, migrate to Redis-based counters. """ import time import asyncio from collections import deque import os from typing import Dict, Deque, Tuple, List import logging logger = logging.getLogger(__name__) # Free-tier limits (RPM) — adjust as necessary DEFAULT_LIMITS = { "cerebras": 30, "groq": 30, "google": 15, "openrouter": 5, "xai": 20, } # Threshold at which providers are considered "near capacity" (e.g., 85%) NEAR_CAPACITY_RATIO = 0.85 # Window length in seconds (60 seconds for per-minute limits) WINDOW_SECONDS = 60 class ProviderWindow: def __init__(self, limit: int): self.limit = limit self.timestamps: Deque[float] = deque() self.lock = asyncio.Lock() def _prune(self, now: float): # remove timestamps older than WINDOW_SECONDS while self.timestamps and (now - self.timestamps[0]) > WINDOW_SECONDS: self.timestamps.popleft() async def record(self) -> None: now = time.time() async with self.lock: self._prune(now) self.timestamps.append(now) async def count(self) -> int: now = time.time() async with self.lock: self._prune(now) return len(self.timestamps) async def available_capacity(self) -> int: cnt = await self.count() return max(0, self.limit - cnt) async def is_near_capacity(self) -> bool: cnt = await self.count() return cnt >= int(self.limit * NEAR_CAPACITY_RATIO) async def is_exhausted(self) -> bool: cnt = await self.count() return cnt >= self.limit class RateTracker: def __init__(self, limits: Dict[str, int] = None): self.limits = limits or DEFAULT_LIMITS.copy() self.windows: Dict[str, ProviderWindow] = { name: ProviderWindow(limit) for name, limit in self.limits.items() } def _ensure_provider(self, provider: str): if provider not in self.windows: # default conservative limit self.limits[provider] = 10 self.windows[provider] = ProviderWindow(10) async def record_call(self, provider: str) -> None: """Record a single API call for provider.""" self._ensure_provider(provider) await self.windows[provider].record() async def can_call(self, provider: str) -> bool: """Return True if call is allowed (limit not exhausted).""" self._ensure_provider(provider) return not await self.windows[provider].is_exhausted() async def is_near_capacity(self, provider: str) -> bool: self._ensure_provider(provider) return await self.windows[provider].is_near_capacity() async def get_available_providers(self) -> List[Tuple[str, int]]: """Return list of (provider, available_capacity) sorted desc by capacity.""" snapshot = [] for name, window in self.windows.items(): cap = await window.available_capacity() snapshot.append((name, cap)) snapshot.sort(key=lambda x: x[1], reverse=True) return snapshot async def get_status(self) -> Dict[str, Dict]: """Return debugging status for all providers.""" status = {} for name, window in self.windows.items(): cnt = await window.count() status[name] = { "limit": window.limit, "count_last_min": cnt, "available": max(0, window.limit - cnt), "near_capacity": cnt >= int(window.limit * NEAR_CAPACITY_RATIO), "exhausted": cnt >= window.limit, } return status # Singleton instance rate_tracker = RateTracker() # Convenience async wrappers async def record_call(provider: str) -> None: await rate_tracker.record_call(provider) async def can_call(provider: str) -> bool: return await rate_tracker.can_call(provider) async def get_available_providers() -> List[Tuple[str, int]]: return await rate_tracker.get_available_providers() async def get_status() -> Dict[str, Dict]: return await rate_tracker.get_status() # If REDIS_URL is provided, prefer a Redis-backed tracker for multi-process safety REDIS_URL = os.getenv("REDIS_URL") if REDIS_URL: try: from .redis_rate_tracker import redis_tracker as _redis_tracker if _redis_tracker: # Overwrite the singleton instance, so the wrapper functions # above will automatically delegate to the Redis tracker. rate_tracker = _redis_tracker logger.info("[RateTracker] Using Redis-backed rate tracker") except Exception: logger.exception("Failed to initialize Redis rate tracker; falling back to in-memory")