keycycle 0.3.2__py3-none-any.whl

keycycle/adapters/openai_adapter.py

@@ -0,0 +1,381 @@
+import asyncio
+import logging
+import time
+from typing import List, Any, Optional, Callable, Generator, AsyncGenerator
+
+from ..config.dataclasses import KeyUsage, RateLimits
+from ..config.constants import (
+    TEMP_RATE_LIMIT_MAX_RETRIES,
+    TEMP_RATE_LIMIT_INITIAL_DELAY,
+    TEMP_RATE_LIMIT_MAX_DELAY,
+    TEMP_RATE_LIMIT_MULTIPLIER,
+    KEY_ROTATION_DELAY_SECONDS,
+)
+from ..core.utils import (
+    is_rate_limit_error,
+    is_temporary_rate_limit_error,
+    is_payment_required_error,
+    get_key_suffix,
+)
+from ..core.backoff import ExponentialBackoff, BackoffConfig
+from ..key_rotation.rotation_manager import RotatingKeyManager
+
+logger = logging.getLogger(__name__)
+
+try:
+    import openai
+    from openai import OpenAI, AsyncOpenAI, RateLimitError
+    HAS_OPENAI = True
+except ImportError:
+    HAS_OPENAI = False
+    # Generic placeholders
+    class OpenAI:
+        pass
+    class AsyncOpenAI:
+        pass
+    RateLimitError = Exception
+
+PROVIDER_BASE_URLS = {
+    "openai": "https://api.openai.com/v1",
+    "openrouter": "https://openrouter.ai/api/v1",
+    "gemini": "https://generativelanguage.googleapis.com/v1beta/openai/",
+    "cerebras": "https://api.cerebras.ai/v1",
+    "groq": "https://api.groq.com/openai/v1",
+    "cohere": "https://api.cohere.ai/compatibility/v1",
+    "moonshot": "https://api.moonshot.ai/v1",
+    "mistral": "https://api.mistral.ai/v1",
+}
+
+class BaseRotatingClient:
+    def __init__(self,
+        manager: RotatingKeyManager,
+        limit_resolver: Callable[[str, Optional[str]], RateLimits],
+        default_model: str,
+        estimated_tokens: int = 1000,
+        max_retries: int = 5,
+        base_url: Optional[str] = None,
+        provider: Optional[str] = None,
+        client_kwargs: dict = None
+    ):
+        """
+        Initialize the rotating client.
+        
+        Args:
+            manager: The key rotation manager
+            limit_resolver: Function to resolve rate limits for a model
+            default_model: Default model to use
+            estimated_tokens: Estimated tokens per request
+            max_retries: Maximum number of retries on rate limit
+            base_url: Base URL for the API (takes precedence over provider)
+            provider: Provider name (openai, openrouter, gemini, cerebras, groq)
+            client_kwargs: Additional kwargs to pass to OpenAI client
+        """
+        
+        if not HAS_OPENAI:
+            raise ImportError("The 'openai' library is required. Install with `pip install openai`.")
+            
+        self.manager = manager
+        self.limit_resolver = limit_resolver
+        self.default_model = default_model
+        self.estimated_tokens = estimated_tokens
+        self.max_retries = max_retries
+        self.client_kwargs = client_kwargs or {}
+
+        if base_url:
+            self.base_url = base_url
+        elif provider:
+            provider_lower = provider.lower()
+            if provider_lower not in PROVIDER_BASE_URLS:
+                raise ValueError(
+                    f"Unknown provider: {provider}. "
+                    f"Valid providers: {', '.join(PROVIDER_BASE_URLS.keys())}"
+                )
+            self.base_url = PROVIDER_BASE_URLS[provider_lower]
+        else:
+            self.base_url = None  # Will use OpenAI's default
+
+        if self.base_url:
+            self.client_kwargs['base_url'] = self.base_url
+
+    def _record_usage(self, key_usage: KeyUsage, model_id: str, actual_tokens: int) -> None:
+        self.manager.record_usage(
+            key_obj=key_usage,
+            model_id=model_id,
+            actual_tokens=actual_tokens,
+            estimated_tokens=self.estimated_tokens
+        )
+
+    def _extract_usage(self, response: Any) -> int:
+        try:
+            if hasattr(response, 'usage') and response.usage:
+                return response.usage.total_tokens
+        except (AttributeError, TypeError):
+            pass
+        return 0
+
+# --- SYNC IMPLEMENTATION ---
+
+class RotatingOpenAIClient(BaseRotatingClient):
+    def _get_fresh_client(self, api_key: str) -> OpenAI:
+        return OpenAI(api_key=api_key, **self.client_kwargs)
+
+    def __getattr__(self, name):
+        return SyncProxyHelper(self, [name])
+
+    def _execute(self, path: List[str], args, kwargs):
+        model_id = kwargs.get('model', self.default_model)
+        if 'model' not in kwargs:
+            kwargs['model'] = model_id
+        limits = self.limit_resolver(model_id, None)
+
+        if kwargs.get('stream', False) and 'stream_options' not in kwargs:
+            kwargs['stream_options'] = {"include_usage": True}
+
+        for attempt in range(self.max_retries + 1):
+            key_usage = self.manager.get_key(model_id, limits, self.estimated_tokens)
+            if not key_usage:
+                raise RuntimeError(f"No available keys for {model_id}")
+
+            # Create backoff for temporary rate limits
+            temp_backoff = ExponentialBackoff(BackoffConfig(
+                initial_interval=TEMP_RATE_LIMIT_INITIAL_DELAY,
+                max_interval=TEMP_RATE_LIMIT_MAX_DELAY,
+                multiplier=TEMP_RATE_LIMIT_MULTIPLIER,
+            ))
+
+            for temp_attempt in range(TEMP_RATE_LIMIT_MAX_RETRIES + 1):
+                try:
+                    real_client = self._get_fresh_client(key_usage.api_key)
+
+                    target = real_client
+                    for p in path:
+                        target = getattr(target, p)
+
+                    result = target(*args, **kwargs)
+
+                    if kwargs.get('stream', False):
+                        return self._wrap_stream(result, key_usage, model_id)
+
+                    self._record_usage(key_usage, model_id, self._extract_usage(result))
+                    return result
+
+                except Exception as e:
+                    # Check for temporary rate limit first - retry with SAME key
+                    if is_temporary_rate_limit_error(e) and temp_attempt < TEMP_RATE_LIMIT_MAX_RETRIES:
+                        delay = temp_backoff.get_next_interval()
+                        logger.info(
+                            "Temporary rate limit on key ...%s for %s. Waiting %.1fs (%d/%d).",
+                            get_key_suffix(key_usage.api_key), model_id, delay,
+                            temp_attempt + 1, TEMP_RATE_LIMIT_MAX_RETRIES
+                        )
+                        time.sleep(delay)
+                        continue  # Retry with SAME key
+
+                    # Payment required (quota/credits exhausted) - key is dead
+                    # for this process, rotate to next key
+                    if is_payment_required_error(e) and attempt < self.max_retries:
+                        logger.warning(
+                            "402/PaymentRequired hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
+                            model_id, get_key_suffix(key_usage.api_key), attempt + 1, self.max_retries + 1
+                        )
+                        key_usage.mark_dead()
+                        self.manager.force_rotate_index()
+                        time.sleep(KEY_ROTATION_DELAY_SECONDS)
+                        break  # Break inner loop, continue outer loop with new key
+
+                    # Hard rate limit - rotate to next key
+                    if is_rate_limit_error(e) and attempt < self.max_retries:
+                        logger.warning(
+                            "429/RateLimit hit for %s on key ...%s. Rotating. (Attempt %d/%d)",
+                            model_id, get_key_suffix(key_usage.api_key), attempt + 1, self.max_retries + 1
+                        )
+                        key_usage.trigger_cooldown()
+                        self.manager.force_rotate_index()
+                        time.sleep(KEY_ROTATION_DELAY_SECONDS)
+                        break  # Break inner loop, continue outer loop with new key
+
+                    self._record_usage(key_usage, model_id, 0)
+                    raise
+            else:
+                # Inner loop exhausted without success - continue to next key
+                if attempt < self.max_retries:
+                    logger.warning(
+                        "Temporary rate limit retries exhausted for key ...%s. Rotating.",
+                        get_key_suffix(key_usage.api_key)
+                    )
+                    key_usage.trigger_cooldown()
+                    self.manager.force_rotate_index()
+                    continue
+
+        raise RuntimeError(f"All retry attempts exhausted for {model_id}")
+
+    def _wrap_stream(self, generator: Generator, key_usage: KeyUsage, model_id: str):
+        final_tokens = 0
+        try:
+            for chunk in generator:
+                if hasattr(chunk, 'usage') and chunk.usage:
+                    final_tokens = chunk.usage.total_tokens
+                yield chunk
+        except Exception as e:
+            if is_payment_required_error(e):
+                logger.warning(
+                    "402/PaymentRequired hit during streaming for %s on key ...%s. Marking key dead.",
+                    model_id, get_key_suffix(key_usage.api_key)
+                )
+                key_usage.mark_dead()
+                self.manager.force_rotate_index()
+            elif is_rate_limit_error(e):
+                logger.warning(
+                    "Rate limit hit during streaming for %s on key ...%s.",
+                    model_id, get_key_suffix(key_usage.api_key)
+                )
+                key_usage.trigger_cooldown()
+                self.manager.force_rotate_index()
+            raise
+        finally:
+            self._record_usage(key_usage, model_id, final_tokens)
+
+class SyncProxyHelper:
+    def __init__(self, client: RotatingOpenAIClient, path: List[str]):
+        self.client = client
+        self.path = path
+
+    def __getattr__(self, name):
+        return SyncProxyHelper(self.client, self.path + [name])
+
+    def __call__(self, *args, **kwargs):
+        return self.client._execute(self.path, args, kwargs)
+
+
+# --- ASYNC IMPLEMENTATION ---
+
+class RotatingAsyncOpenAIClient(BaseRotatingClient):
+    def _get_fresh_client(self, api_key: str) -> AsyncOpenAI:
+        return AsyncOpenAI(api_key=api_key, **self.client_kwargs)
+
+    def __getattr__(self, name):
+        return AsyncProxyHelper(self, [name])
+
+    async def _execute(self, path: List[str], args, kwargs: dict):
+        model_id = kwargs.get('model', self.default_model)
+        if 'model' not in kwargs:
+            kwargs['model'] = model_id
+        limits = self.limit_resolver(model_id, None)
+
+        if kwargs.get('stream', False) and 'stream_options' not in kwargs:
+            kwargs['stream_options'] = {"include_usage": True}
+
+        for attempt in range(self.max_retries + 1):
+            key_usage = self.manager.get_key(model_id, limits, self.estimated_tokens)
+            if not key_usage:
+                raise RuntimeError(f"No available keys for {model_id}")
+
+            # Create backoff for temporary rate limits
+            temp_backoff = ExponentialBackoff(BackoffConfig(
+                initial_interval=TEMP_RATE_LIMIT_INITIAL_DELAY,
+                max_interval=TEMP_RATE_LIMIT_MAX_DELAY,
+                multiplier=TEMP_RATE_LIMIT_MULTIPLIER,
+            ))
+
+            for temp_attempt in range(TEMP_RATE_LIMIT_MAX_RETRIES + 1):
+                try:
+                    real_client = self._get_fresh_client(key_usage.api_key)
+
+                    target = real_client
+                    for p in path:
+                        target = getattr(target, p)
+
+                    result = await target(*args, **kwargs)
+
+                    if kwargs.get('stream', False):
+                        return self._wrap_stream(result, key_usage, model_id)
+
+                    self._record_usage(key_usage, model_id, self._extract_usage(result))
+                    return result
+
+                except Exception as e:
+                    # Check for temporary rate limit first - retry with SAME key
+                    if is_temporary_rate_limit_error(e) and temp_attempt < TEMP_RATE_LIMIT_MAX_RETRIES:
+                        delay = temp_backoff.get_next_interval()
+                        logger.info(
+                            "Temporary rate limit on key ...%s for %s. Waiting %.1fs (%d/%d).",
+                            get_key_suffix(key_usage.api_key), model_id, delay,
+                            temp_attempt + 1, TEMP_RATE_LIMIT_MAX_RETRIES
+                        )
+                        await asyncio.sleep(delay)
+                        continue  # Retry with SAME key
+
+                    # Payment required (quota/credits exhausted) - key is dead
+                    # for this process, rotate to next key
+                    if is_payment_required_error(e) and attempt < self.max_retries:
+                        logger.warning(
+                            "402/PaymentRequired hit for %s on key ...%s. Marking key dead and rotating. (Attempt %d/%d)",
+                            model_id, get_key_suffix(key_usage.api_key), attempt + 1, self.max_retries + 1
+                        )
+                        key_usage.mark_dead()
+                        self.manager.force_rotate_index()
+                        await asyncio.sleep(KEY_ROTATION_DELAY_SECONDS)
+                        break  # Break inner loop, continue outer loop with new key
+
+                    # Hard rate limit - rotate to next key
+                    if is_rate_limit_error(e) and attempt < self.max_retries:
+                        logger.warning(
+                            "429/RateLimit hit for %s on key ...%s. Rotating. (Attempt %d/%d)",
+                            model_id, get_key_suffix(key_usage.api_key), attempt + 1, self.max_retries + 1
+                        )
+                        key_usage.trigger_cooldown()
+                        self.manager.force_rotate_index()
+                        await asyncio.sleep(KEY_ROTATION_DELAY_SECONDS)
+                        break  # Break inner loop, continue outer loop with new key
+
+                    self._record_usage(key_usage, model_id, 0)
+                    raise
+            else:
+                # Inner loop exhausted without success - continue to next key
+                if attempt < self.max_retries:
+                    logger.warning(
+                        "Temporary rate limit retries exhausted for key ...%s. Rotating.",
+                        get_key_suffix(key_usage.api_key)
+                    )
+                    key_usage.trigger_cooldown()
+                    self.manager.force_rotate_index()
+                    continue
+
+        raise RuntimeError(f"All retry attempts exhausted for {model_id}")
+
+    async def _wrap_stream(self, generator: AsyncGenerator, key_usage: KeyUsage, model_id: str):
+        final_tokens = 0
+        try:
+            async for chunk in generator:
+                if hasattr(chunk, 'usage') and chunk.usage:
+                    final_tokens = chunk.usage.total_tokens
+                yield chunk
+        except Exception as e:
+            if is_payment_required_error(e):
+                logger.warning(
+                    "402/PaymentRequired hit during streaming for %s on key ...%s. Marking key dead.",
+                    model_id, get_key_suffix(key_usage.api_key)
+                )
+                key_usage.mark_dead()
+                self.manager.force_rotate_index()
+            elif is_rate_limit_error(e):
+                logger.warning(
+                    "Rate limit hit during streaming for %s on key ...%s.",
+                    model_id, get_key_suffix(key_usage.api_key)
+                )
+                key_usage.trigger_cooldown()
+                self.manager.force_rotate_index()
+            raise
+        finally:
+            self._record_usage(key_usage, model_id, final_tokens)
+
+class AsyncProxyHelper:
+    def __init__(self, client: RotatingAsyncOpenAIClient, path: List[str]):
+        self.client = client
+        self.path = path
+
+    def __getattr__(self, name):
+        return AsyncProxyHelper(self.client, self.path + [name])
+
+    async def __call__(self, *args, **kwargs):
+        return await self.client._execute(self.path, args, kwargs)

keycycle/config/__init__.py

@@ -0,0 +1,25 @@
+from .dataclasses import (
+    RateLimits,
+    UsageSnapshot,
+    UsageBucket,
+    KeyUsage,
+    KeyDetailedStats,
+    KeySummary,
+    GlobalStats,
+    ModelAggregatedStats,
+)
+from .enums import RateLimitStrategy
+from .log_config import configure_logging
+
+__all__ = [
+    "RateLimits",
+    "UsageSnapshot",
+    "UsageBucket",
+    "KeyUsage",
+    "KeyDetailedStats",
+    "KeySummary",
+    "GlobalStats",
+    "ModelAggregatedStats",
+    "RateLimitStrategy",
+    "configure_logging",
+]

keycycle/config/constants.py

@@ -0,0 +1,35 @@
+"""Timing constants for key rotation and rate limiting."""
+
+# Time windows (in seconds)
+SECONDS_PER_MINUTE = 60
+SECONDS_PER_HOUR = 3600
+SECONDS_PER_DAY = 86400
+
+# Cooldown configuration
+DEFAULT_COOLDOWN_SECONDS = 30
+
+# Cleanup intervals
+CLEANUP_INTERVAL_SECONDS = 55
+
+# Key wait/polling configuration
+DEFAULT_POLL_INTERVAL = 0.5
+MIN_POLL_INTERVAL = 0.1
+MAX_POLL_INTERVAL = 5.0
+
+# History lookback
+HISTORY_LOOKBACK_SECONDS = 86400  # 24 hours
+
+# API key suffix length for logging
+KEY_SUFFIX_LENGTH = 8
+
+# Temporary rate limit retry configuration
+TEMP_RATE_LIMIT_MAX_RETRIES = 3
+TEMP_RATE_LIMIT_INITIAL_DELAY = 1.0  # seconds
+TEMP_RATE_LIMIT_MAX_DELAY = 10.0  # seconds
+TEMP_RATE_LIMIT_MULTIPLIER = 2.0
+
+# Key rotation delay (seconds to wait after rotating to a new key)
+KEY_ROTATION_DELAY_SECONDS = 0.5
+
+# Usage logging batch size
+USAGE_LOG_BATCH_SIZE = 50

keycycle/config/dataclasses.py

@@ -0,0 +1,225 @@
+from typing import Any, Dict, List, Optional, Union
+from .enums import RateLimitStrategy
+from .constants import (
+    SECONDS_PER_MINUTE, SECONDS_PER_HOUR, SECONDS_PER_DAY,
+    DEFAULT_COOLDOWN_SECONDS
+)
+import time
+from collections import deque, defaultdict
+from dataclasses import dataclass, field
+
+# --- CONFIGURATION DATA ---
+
+@dataclass
+class RateLimits:
+    """Rate limits for a provider"""
+    requests_per_minute: int
+    requests_per_hour: int
+    requests_per_day: int
+    tokens_per_minute: Optional[int] = None
+    tokens_per_hour: Optional[int] = None
+    tokens_per_day: Optional[int] = None
+
+# Type alias for per-key rate limit overrides
+# Can be a single RateLimits (applies to all models) or a dict mapping model_id -> RateLimits
+KeyLimitOverride = Union[RateLimits, Dict[str, RateLimits]]
+
+@dataclass
+class UsageSnapshot:
+    """Standardized view of usage counters"""
+    rpm: int = 0
+    rph: int = 0
+    rpd: int = 0
+    tpm: int = 0
+    tph: int = 0
+    tpd: int = 0
+    total_requests: int = 0
+    total_tokens: int = 0
+
+    def __add__(self, other):
+        """Allow summing snapshots for aggregation"""
+        if not isinstance(other, UsageSnapshot): return NotImplemented
+        return UsageSnapshot(
+            self.rpm + other.rpm, self.rph + other.rph, self.rpd + other.rpd,
+            self.tpm + other.tpm, self.tph + other.tph, self.tpd + other.tpd,
+            self.total_requests + other.total_requests, self.total_tokens + other.total_tokens
+        )
+
+# --- STATS DATA TRANSFER OBJECTS (DTOs) ---
+
+@dataclass
+class KeySummary:
+    index: int; suffix: str; snapshot: UsageSnapshot
+@dataclass
+class GlobalStats:
+    total: UsageSnapshot; keys: List[KeySummary]
+@dataclass
+class KeyDetailedStats:
+    index: int; suffix: str; total: UsageSnapshot; breakdown: Dict[str, UsageSnapshot]
+@dataclass
+class ModelAggregatedStats:
+    model_id: str; total: UsageSnapshot; keys: List[KeySummary]
+    
+
+# --- USAGE TRACKING ---
+@dataclass
+class UsageBucket:
+    """Tracks counters for a SINGLE model context"""
+    requests_minute: deque[float] = field(default_factory=deque)
+    requests_hour: deque[float] = field(default_factory=deque)
+    requests_day: deque[float] = field(default_factory=deque)
+    
+    tokens_minute: deque[tuple[float, int]] = field(default_factory=deque)
+    tokens_hour: deque[tuple[float, int]] = field(default_factory=deque)
+    tokens_day: deque[tuple[float, int]] = field(default_factory=deque)
+    
+    total_requests: int = 0
+    total_tokens: int = 0
+    
+    pending_tokens: int = 0
+    
+    def clean(self) -> None:
+        """Clean old entries based on current time"""
+        now = time.time()
+        cutoffs = (now - SECONDS_PER_MINUTE, now - SECONDS_PER_HOUR, now - SECONDS_PER_DAY)
+        
+        for d, cut in zip(
+            [self.requests_minute, self.requests_hour, self.requests_day], cutoffs
+        ):
+            while d and d[0] <= cut: d.popleft()
+            
+        for d, cut in zip(
+            [self.tokens_minute, self.tokens_hour, self.tokens_day], cutoffs
+        ):
+            while d and d[0][0] <= cut: d.popleft()
+    
+    def add(self, tokens: int, timestamp: float):
+        self.requests_minute.append(timestamp)
+        self.requests_hour.append(timestamp)
+        self.requests_day.append(timestamp)
+        self.total_requests += 1
+        
+        if tokens > 0:
+            self.tokens_minute.append((timestamp, tokens))
+            self.tokens_hour.append((timestamp, tokens))
+            self.tokens_day.append((timestamp, tokens))
+            self.total_tokens += tokens
+    
+    def check_limits(self, limits: RateLimits, estimated_tokens: int) -> bool:
+        self.clean()
+        if len(self.requests_minute) >= limits.requests_per_minute: return False
+        if len(self.requests_hour) >= limits.requests_per_hour: return False
+        if len(self.requests_day) >= limits.requests_per_day: return False
+        
+        current_tpm = sum(t[1] for t in self.tokens_minute) + self.pending_tokens
+        current_tph = sum(t[1] for t in self.tokens_hour) + self.pending_tokens
+        current_tpd = sum(t[1] for t in self.tokens_day) + self.pending_tokens
+        
+        if limits.tokens_per_minute and (current_tpm + estimated_tokens > limits.tokens_per_minute): return False
+        if limits.tokens_per_hour and (current_tph + estimated_tokens > limits.tokens_per_hour): return False
+        if limits.tokens_per_day and (current_tpd + estimated_tokens > limits.tokens_per_day): return False
+        
+        return True
+    
+    def reserve(self, tokens: int):
+        """Lock in estimated tokens"""
+        self.pending_tokens += tokens
+    
+    def commit(self, actual_tokens: int, reserved_tokens: int, timestamp: float):
+        """Remove reservation and add actual usage"""
+        self.pending_tokens -= reserved_tokens
+        if self.pending_tokens < 0:
+            import logging
+            logging.getLogger(__name__).warning(
+                "Pending tokens went negative (%d), clamping to 0",
+                self.pending_tokens
+            )
+            self.pending_tokens = 0
+        self.add(actual_tokens, timestamp)
+    
+    def get_snapshot(self) -> UsageSnapshot:
+        """Return current counts as a clean snapshot"""
+        self.clean()
+        return UsageSnapshot(
+            rpm=len(self.requests_minute),
+            rph=len(self.requests_hour),
+            rpd=len(self.requests_day),
+            tpm=sum(t[1] for t in self.tokens_minute),
+            tph=sum(t[1] for t in self.tokens_hour),
+            tpd=sum(t[1] for t in self.tokens_day),
+            total_requests=self.total_requests,
+            total_tokens=self.total_tokens
+        )
+    
+
+@dataclass
+class KeyUsage:
+    """Represents an API Key and holds multiple UsageBuckets (one per model)"""
+    api_key: str
+    strategy: RateLimitStrategy
+    params: Dict[str, Any] = field(default_factory=dict)
+    buckets: Dict[str, UsageBucket] = field(default_factory=lambda: defaultdict(UsageBucket))
+    global_bucket: UsageBucket = field(default_factory=UsageBucket)
+    last_429: float = 0.0
+    dead: bool = False
+
+    def get_client_params(self) -> Dict[str, Any]:
+        """Returns all params for client instantiation."""
+        if self.params:
+            return dict(self.params)
+        return {"api_key": self.api_key}
+    
+    def record_usage(self, model_id: str, tokens: int, timestamp: float = None):
+        ts = timestamp if timestamp else time.time()
+        self.buckets[model_id].add(tokens, ts)
+        if self.strategy == RateLimitStrategy.GLOBAL:
+            self.global_bucket.add(tokens, ts)
+        
+    def can_use_model(self, model_id: str, limits: RateLimits, estimated_tokens: int = 1000) -> bool:
+        """Check limits based on the provider's strategy"""
+        if self.strategy == RateLimitStrategy.GLOBAL:
+            return self.global_bucket.check_limits(limits, estimated_tokens)
+        else: # Per-Model Limits
+            return self.buckets[model_id].check_limits(limits, estimated_tokens)
+
+    def get_total_snapshot(self) -> UsageSnapshot:
+        if self.strategy == RateLimitStrategy.GLOBAL:
+            return self.global_bucket.get_snapshot()
+        total = UsageSnapshot()
+        for b in self.buckets.values():
+            total = total + b.get_snapshot()
+        return total
+    
+    def reserve(self, model_id: str, tokens: int):
+        self.buckets[model_id].reserve(tokens)
+        if self.strategy == RateLimitStrategy.GLOBAL:
+            self.global_bucket.reserve(tokens)
+    
+    def commit(self, model_id: str, actual_tokens: int, reserved_tokens: int, timestamp: float = None):
+        ts = timestamp if timestamp else time.time()
+        self.buckets[model_id].commit(actual_tokens, reserved_tokens, ts)
+        if self.strategy == RateLimitStrategy.GLOBAL:
+            self.global_bucket.commit(actual_tokens, reserved_tokens, ts)
+
+
+    def is_cooling_down(self, cooldown_seconds: int = DEFAULT_COOLDOWN_SECONDS) -> bool:
+        """Returns True if the key is still in its cooldown penalty period, or dead."""
+        if self.dead:
+            return True
+        if self.last_429 == 0:
+            return False
+        return (time.time() - self.last_429) < cooldown_seconds
+
+    def trigger_cooldown(self):
+        """Mark this key as rate-limited."""
+        self.last_429 = time.time()
+
+    def mark_dead(self):
+        """
+        Permanently exclude this key from rotation for the lifetime of this process.
+
+        Used for errors a short cooldown can't fix - e.g. HTTP 402 payment_required
+        when a free-trial key's credits/quota are exhausted. Unlike trigger_cooldown(),
+        this never expires; the key stays dead until the process restarts.
+        """
+        self.dead = True

keycycle/config/enums.py

@@ -0,0 +1,5 @@
+from enum import Enum
+
+class RateLimitStrategy(Enum):
+    PER_MODEL = "per_model"  # Cerebras, Groq, Gemini
+    GLOBAL = "global"        # OpenRouter (Shared limits across all models)