keycycle 0.3.2__py3-none-any.whl

keycycle/key_rotation/rotating_mixin.py

@@ -0,0 +1,392 @@
+import asyncio
+import functools
+import logging
+import time
+from typing import AsyncIterator, Iterator, Optional, Union, TYPE_CHECKING
+
+from ..config.dataclasses import KeyUsage
+from ..config.log_config import default_logger
+from ..config.constants import (
+    TEMP_RATE_LIMIT_MAX_RETRIES,
+    TEMP_RATE_LIMIT_INITIAL_DELAY,
+    TEMP_RATE_LIMIT_MAX_DELAY,
+    TEMP_RATE_LIMIT_MULTIPLIER,
+)
+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
+if TYPE_CHECKING:
+    from agno.models.response import ModelResponse
+
+
+class _UsageResponse:
+    """Minimal response holder for usage recording when agno is not available."""
+    __slots__ = ('response_usage',)
+
+    def __init__(self, usage=None):
+        self.response_usage = usage
+
+
+class RotatingCredentialsMixin:
+    """
+    Mixin that handles key rotation, 429 detection, and 30s cooldown triggers.
+    """
+    
+    def __init__(
+        self, 
+        *args, 
+        model_id: str,
+        wrapper=None, 
+        rotating_wait=True, 
+        rotating_timeout=10.0, 
+        rotating_estimated_tokens=1000,
+        rotating_max_retries=5, 
+        rotating_fixed_key_id: Union[int, str] = None,
+        logger: Optional[logging.Logger] = None,
+        **kwargs):
+        """
+        Initializes rotation parameters and patches the model if necessary.
+        
+        Args:
+            wrapper: The MultiProviderWrapper instance managing keys.
+            rotating_wait: Whether to wait for a key if none are available.
+            rotating_timeout: Max time to wait for a key.
+            rotating_estimated_tokens: Default token estimate for rate limiting.
+            rotating_max_retries: Number of retries on 429 errors.
+        """
+        self.logger = logger or default_logger
+        self.wrapper = wrapper
+        self.model_id = model_id
+        self._rotating_wait = rotating_wait
+        self._rotating_timeout = rotating_timeout
+        self._estimated_tokens = rotating_estimated_tokens
+        self._max_retries = rotating_max_retries
+        self._fixed_key_id = rotating_fixed_key_id
+
+        super().__init__(*args, **kwargs)
+
+    def __init_subclass__(cls, **kwargs):
+        """
+        Hook to automatically copy function signatures and docstrings 
+        from the parent class to this mixin's wrapper methods.
+        """
+        super().__init_subclass__(**kwargs)
+        methods_to_sync = ['invoke', 'ainvoke', 'invoke_stream', 'ainvoke_stream']
+        for method_name in methods_to_sync:
+            if not hasattr(cls, method_name):
+                continue
+            target_method = None
+            for parent in cls.mro():
+                if parent is cls or parent is RotatingCredentialsMixin:
+                    continue
+                # Check if the method is defined in this specific parent's __dict__
+                if hasattr(parent, method_name) and method_name in parent.__dict__:
+                    target_method = getattr(parent, method_name)
+                    break
+            
+            if target_method:
+                mixin_implementation = getattr(cls, method_name)
+                wrapped_method = functools.wraps(target_method)(mixin_implementation)
+
+                if hasattr(target_method, '__signature__'):
+                    wrapped_method.__signature__ = target_method.__signature__
+                setattr(cls, method_name, wrapped_method)
+
+    def _rotate_credentials(self) -> KeyUsage:
+        key_usage: KeyUsage = self.wrapper.get_key_usage(
+            model_id=self.model_id,
+            estimated_tokens=self._estimated_tokens,
+            wait=self._rotating_wait,
+            timeout=self._rotating_timeout,
+            key_id=self._fixed_key_id
+        )
+        self.api_key = key_usage.api_key
+        
+        if hasattr(self, "client"): self.client = None
+        if hasattr(self, "async_client"): self.async_client = None
+        if hasattr(self, "gemini_client"): self.gemini_client = None
+
+        return key_usage
+
+    def _get_retry_limit(self) -> int:
+        return min(self._max_retries, len(self.wrapper.manager.keys) - 1)
+
+    def _record_usage(self, key_obj: KeyUsage, response: Optional["ModelResponse"]):
+        """
+        Extracts usage from the response and reports it to the manager.
+        Falls back to estimated_tokens if response is None or usage is missing.
+        """
+        if not self.wrapper:
+            return
+
+        actual_tokens = 0
+
+        if response and response.response_usage:
+            actual_tokens = response.response_usage.total_tokens
+
+        # Fallback: If no usage found (or streaming), use the estimate
+        if actual_tokens == 0:
+            actual_tokens = self._estimated_tokens
+
+        self.wrapper.manager.record_usage(
+            key_obj=key_obj,
+            model_id=self.model_id,
+            actual_tokens=actual_tokens,
+            estimated_tokens=self._estimated_tokens
+        )
+
+    def _create_temp_backoff(self) -> ExponentialBackoff:
+        """Create a backoff instance for temporary rate limit retries."""
+        return ExponentialBackoff(BackoffConfig(
+            initial_interval=TEMP_RATE_LIMIT_INITIAL_DELAY,
+            max_interval=TEMP_RATE_LIMIT_MAX_DELAY,
+            multiplier=TEMP_RATE_LIMIT_MULTIPLIER,
+        ))
+
+    def invoke(self, *args, **kwargs) -> "ModelResponse":
+        limit = self._get_retry_limit()
+
+        for attempt in range(limit + 1):
+            key_usage = self._rotate_credentials()
+            temp_backoff = self._create_temp_backoff()
+
+            for temp_attempt in range(TEMP_RATE_LIMIT_MAX_RETRIES + 1):
+                try:
+                    response = super().invoke(*args, **kwargs)
+                    self._record_usage(key_usage, response)
+                    return response
+                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()
+                        self.logger.info(
+                            "Temporary rate limit on key %s [%s]. Waiting %.1fs (%d/%d).",
+                            get_key_suffix(self.api_key), self.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 < limit:
+                        self.logger.warning(
+                            "402 Hit on key %s (Sync) [%s]. Marking key dead and rotating (%d/%d).",
+                            get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
+                        )
+                        key_usage.mark_dead()
+                        self.wrapper.manager.force_rotate_index()
+                        break  # Break inner loop, continue outer with new key
+
+                    # Hard rate limit - rotate key
+                    if is_rate_limit_error(e) and attempt < limit:
+                        self.logger.warning(
+                            "429 Hit on key %s (Sync) [%s]. Rotating and retrying (%d/%d).",
+                            get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
+                        )
+                        key_usage.trigger_cooldown()
+                        self.wrapper.manager.force_rotate_index()
+                        break  # Break inner loop, continue outer with new key
+                    raise
+            else:
+                # Temp retries exhausted, move to next key
+                if attempt < limit:
+                    self.logger.warning(
+                        "Temp rate limit retries exhausted for key %s. Rotating.",
+                        get_key_suffix(self.api_key)
+                    )
+                    key_usage.trigger_cooldown()
+                    self.wrapper.manager.force_rotate_index()
+                    continue
+                raise
+
+    async def ainvoke(self, *args, **kwargs) -> "ModelResponse":
+        limit = self._get_retry_limit()
+
+        for attempt in range(limit + 1):
+            key_usage = self._rotate_credentials()
+            temp_backoff = self._create_temp_backoff()
+
+            for temp_attempt in range(TEMP_RATE_LIMIT_MAX_RETRIES + 1):
+                try:
+                    response = await super().ainvoke(*args, **kwargs)
+                    self._record_usage(key_usage, response)
+                    return response
+                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()
+                        self.logger.info(
+                            "Temporary rate limit on key %s [%s]. Waiting %.1fs (%d/%d).",
+                            get_key_suffix(self.api_key), self.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 < limit:
+                        self.logger.warning(
+                            "402 Hit on key %s (Async) [%s]. Marking key dead and rotating (%d/%d).",
+                            get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
+                        )
+                        key_usage.mark_dead()
+                        self.wrapper.manager.force_rotate_index()
+                        break  # Break inner loop, continue outer with new key
+
+                    # Hard rate limit - rotate key
+                    if is_rate_limit_error(e) and attempt < limit:
+                        self.logger.warning(
+                            "429 Hit on key %s (Async) [%s]. Rotating and retrying (%d/%d).",
+                            get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
+                        )
+                        key_usage.trigger_cooldown()
+                        self.wrapper.manager.force_rotate_index()
+                        break  # Break inner loop, continue outer with new key
+                    raise
+            else:
+                # Temp retries exhausted, move to next key
+                if attempt < limit:
+                    self.logger.warning(
+                        "Temp rate limit retries exhausted for key %s. Rotating.",
+                        get_key_suffix(self.api_key)
+                    )
+                    key_usage.trigger_cooldown()
+                    self.wrapper.manager.force_rotate_index()
+                    continue
+                raise
+
+    def invoke_stream(self, *args, **kwargs) -> Iterator["ModelResponse"]:
+        limit = self._get_retry_limit()
+
+        for attempt in range(limit + 1):
+            key_usage = self._rotate_credentials()
+            temp_backoff = self._create_temp_backoff()
+
+            for temp_attempt in range(TEMP_RATE_LIMIT_MAX_RETRIES + 1):
+                chunks_received = False
+                try:
+                    stream = super().invoke_stream(*args, **kwargs)
+                    final_usage = None
+                    for chunk in stream:
+                        chunks_received = True
+                        if chunk.response_usage:
+                            final_usage = chunk.response_usage
+                        yield chunk
+                    if final_usage:
+                        dummy_response = _UsageResponse(final_usage)
+                        self._record_usage(key_usage, dummy_response)
+                    else:
+                        # If the stream didn't return usage data, fallback to estimation
+                        self._record_usage(key_usage, None)
+
+                    return
+                except Exception as e:
+                    # Only retry temp limits if NO chunks received yet
+                    if (is_temporary_rate_limit_error(e)
+                        and temp_attempt < TEMP_RATE_LIMIT_MAX_RETRIES
+                        and not chunks_received):
+                        delay = temp_backoff.get_next_interval()
+                        self.logger.info(
+                            "Temporary rate limit on key %s (Stream) [%s]. Waiting %.1fs (%d/%d).",
+                            get_key_suffix(self.api_key), self.model_id, delay,
+                            temp_attempt + 1, TEMP_RATE_LIMIT_MAX_RETRIES
+                        )
+                        time.sleep(delay)
+                        continue
+
+                    if is_payment_required_error(e) and attempt < limit:
+                        self.logger.warning(
+                            "402 Hit on key %s (Sync Stream) [%s]. Marking key dead and rotating (%d/%d).",
+                            get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
+                        )
+                        key_usage.mark_dead()
+                        self.wrapper.manager.force_rotate_index()
+                        break
+
+                    if is_rate_limit_error(e) and attempt < limit:
+                        self.logger.warning(
+                            "429 Hit on key %s (Sync Stream) [%s]. Rotating and retrying (%d/%d).",
+                            get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
+                        )
+                        key_usage.trigger_cooldown()
+                        self.wrapper.manager.force_rotate_index()
+                        break
+                    raise
+            else:
+                if attempt < limit:
+                    key_usage.trigger_cooldown()
+                    self.wrapper.manager.force_rotate_index()
+                    continue
+                raise
+
+    async def ainvoke_stream(self, *args, **kwargs) -> AsyncIterator["ModelResponse"]:
+        limit = self._get_retry_limit()
+
+        for attempt in range(limit + 1):
+            key_usage = self._rotate_credentials()
+            temp_backoff = self._create_temp_backoff()
+
+            for temp_attempt in range(TEMP_RATE_LIMIT_MAX_RETRIES + 1):
+                chunks_received = False
+                try:
+                    stream = super().ainvoke_stream(*args, **kwargs)
+
+                    final_usage = None
+                    async for chunk in stream:
+                        chunks_received = True
+                        if chunk.response_usage:
+                            final_usage = chunk.response_usage
+                        yield chunk
+
+                    # Stream completed successfully. Record usage.
+                    if final_usage:
+                        dummy_response = _UsageResponse(final_usage)
+                        self._record_usage(key_usage, dummy_response)
+                    else:
+                        self._record_usage(key_usage, None)
+
+                    return
+                except Exception as e:
+                    # Only retry temp limits if NO chunks received yet
+                    if (is_temporary_rate_limit_error(e)
+                        and temp_attempt < TEMP_RATE_LIMIT_MAX_RETRIES
+                        and not chunks_received):
+                        delay = temp_backoff.get_next_interval()
+                        self.logger.info(
+                            "Temporary rate limit on key %s (Async Stream) [%s]. Waiting %.1fs (%d/%d).",
+                            get_key_suffix(self.api_key), self.model_id, delay,
+                            temp_attempt + 1, TEMP_RATE_LIMIT_MAX_RETRIES
+                        )
+                        await asyncio.sleep(delay)
+                        continue
+
+                    if is_payment_required_error(e) and attempt < limit:
+                        self.logger.warning(
+                            "402 Hit on key %s (Async Stream) [%s]. Marking key dead and rotating (%d/%d).",
+                            get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
+                        )
+                        key_usage.mark_dead()
+                        self.wrapper.manager.force_rotate_index()
+                        break
+
+                    if is_rate_limit_error(e) and attempt < limit:
+                        self.logger.warning(
+                            "429 Hit on key %s (Async Stream) [%s]. Rotating and retrying (%d/%d).",
+                            get_key_suffix(self.api_key), self.model_id, attempt + 1, limit
+                        )
+                        key_usage.trigger_cooldown()
+                        self.wrapper.manager.force_rotate_index()
+                        break
+                    raise
+            else:
+                if attempt < limit:
+                    key_usage.trigger_cooldown()
+                    self.wrapper.manager.force_rotate_index()
+                    continue
+                raise

keycycle/key_rotation/rotation_manager.py

@@ -0,0 +1,231 @@
+import time
+import atexit
+import threading
+from threading import Lock, Event
+import logging
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+from ..config.dataclasses import (
+    RateLimits, UsageSnapshot,
+    GlobalStats, KeySummary,
+    KeyDetailedStats, ModelAggregatedStats,
+    KeyUsage
+)
+from ..config.enums import RateLimitStrategy
+from ..config.log_config import default_logger
+from ..config.constants import (
+    CLEANUP_INTERVAL_SECONDS,
+    HISTORY_LOOKBACK_SECONDS,
+    DEFAULT_COOLDOWN_SECONDS,
+)
+from ..core.utils import get_key_suffix, KeyEntry, normalize_key_entries
+from ..usage.usage_logger import AsyncUsageLogger
+from ..usage.db_logic import UsageDatabase
+
+class RotatingKeyManager:
+    """Manages API key rotation with rate limiting"""
+
+    def __init__(
+        self,
+        api_keys: List[KeyEntry],
+        provider_name: str,
+        strategy: RateLimitStrategy,
+        db: UsageDatabase,
+        logger: Optional[logging.Logger] = None,
+        cooldown_seconds: int = DEFAULT_COOLDOWN_SECONDS,
+        limit_resolver: Optional[Callable[[str, Optional[str]], RateLimits]] = None,
+        api_key_param: str = "api_key",
+    ):
+        self.provider_name = provider_name
+        self.logger = logger or default_logger
+        self.strategy = strategy
+        self.cooldown_seconds = cooldown_seconds
+        self.limit_resolver = limit_resolver
+        self.api_key_param = api_key_param
+
+        # Normalize key entries and create KeyUsage objects with params
+        normalized = normalize_key_entries(api_keys, api_key_param)
+        self.keys = [
+            KeyUsage(api_key=primary, strategy=strategy, params=params)
+            for primary, params in normalized
+        ]
+        self.current_index = 0
+        self.lock = Lock()
+
+        self.db = db
+        self.usage_logger = AsyncUsageLogger(self.db)
+        self._hydrate()
+
+        self._stop_event = Event()
+        self._start_cleanup()
+        atexit.register(self.stop)
+
+        self.logger.info("Initialized %d keys for provider %s.", len(self.keys), provider_name)
+    
+    def force_rotate_index(self) -> None:
+        """
+        Force the internal pointer to increment.
+        Useful when a key hits a 429 despite local checks passing.
+        """
+        with self.lock:
+            self.current_index = (self.current_index + 1) % len(self.keys)
+
+    def _hydrate(self) -> None:
+        """Load historical usage from database to restore state."""
+        self.logger.debug("Loading history for provider %s.", self.provider_name)
+        all_history = self.db.load_provider_history(self.provider_name, HISTORY_LOOKBACK_SECONDS)
+        if not all_history:
+            self.logger.info("No history found in DB for %s.", self.provider_name)
+            return
+
+        key_map = {get_key_suffix(k.api_key): k for k in self.keys}
+        count = 0
+        for row in all_history:
+            suffix, model_id, ts, tokens = row
+            
+            if suffix in key_map:
+                key_map[suffix].record_usage(model_id, tokens=tokens, timestamp=ts)
+                count += 1
+        self.logger.info("Hydrated %d records for %s.", 
+                    count, self.provider_name)
+    
+    def _cleanup_loop(self) -> None:
+        """Periodically clean deques to prevent memory bloat."""
+        while not self._stop_event.is_set():
+            try:
+                with self.lock:
+                    for key in self.keys:
+                        for bucket in key.buckets.values():
+                            bucket.clean()
+                        if self.strategy == RateLimitStrategy.GLOBAL:
+                            key.global_bucket.clean()
+            except Exception as e:
+                self.logger.error("Cleanup loop error: %s", e, exc_info=True)
+            time.sleep(CLEANUP_INTERVAL_SECONDS)
+    
+    def _start_cleanup(self) -> None:
+        self._thread = threading.Thread(target=self._cleanup_loop, daemon=True)
+        self._thread.start()
+
+    def stop(self) -> None:
+        """Stop the cleanup thread and flush logs."""
+        self._stop_event.set()
+        self.usage_logger.stop()  # Flush logs
+        if self._thread.is_alive():
+            self._thread.join(timeout=10)
+
+    def get_key(self, model_id: str, default_limits: RateLimits, estimated_tokens: int = 1000) -> Optional[KeyUsage]:
+        """
+        Get an available API key that can handle the request.
+
+        Args:
+            model_id: The model identifier for rate limit lookup
+            default_limits: Default limits to use if no per-key override exists
+            estimated_tokens: Estimated token usage for this request
+
+        Returns:
+            KeyUsage object if a key is available, None otherwise
+        """
+        with self.lock:
+            for offset in range(len(self.keys)):
+                idx = (self.current_index + offset) % len(self.keys)
+                key: KeyUsage = self.keys[idx]
+
+                if key.is_cooling_down(self.cooldown_seconds):
+                    continue
+
+                # Resolve limits for this specific key (supports per-key overrides)
+                if self.limit_resolver:
+                    key_suffix = get_key_suffix(key.api_key)
+                    limits = self.limit_resolver(model_id, key_suffix)
+                else:
+                    limits = default_limits
+
+                if key.can_use_model(model_id, limits, estimated_tokens):
+                    key.reserve(model_id, estimated_tokens)
+                    self.current_index = idx
+                    return key
+            return None
+    
+    def get_specific_key(self, identifier: Union[int, str], model_id: str, estimated_tokens: int = 1000) -> Optional[KeyUsage]:
+        """
+        Get a specific key by index or identifier.
+        Note: This does NOT check rate limits implicitly to allow 'hard' retrieval,
+        but it DOES reserve the estimated tokens to keep tracking accurate.
+        """
+        with self.lock:
+            key, _ = self._find_key(identifier)
+            if key:
+                key.reserve(model_id, estimated_tokens)
+                return key
+            return None
+
+    def record_usage(
+        self, key_obj: KeyUsage, model_id: str, actual_tokens: int, estimated_tokens: int = 1000
+    ) -> None:
+        """Record usage for a specific API key."""
+        with self.lock:
+            key_obj.commit(model_id, actual_tokens, estimated_tokens)
+        self.usage_logger.log(self.provider_name, model_id, key_obj.api_key, actual_tokens)
+
+    # --- STATS HELPERS ---
+
+    def _find_key(self, identifier: Union[int, str]) -> Tuple[Optional[KeyUsage], int]:
+        """Locate key by index (int) or suffix/full-key (str)."""
+        if isinstance(identifier, int):
+            if 0 <= identifier < len(self.keys):
+                return self.keys[identifier], identifier
+        elif isinstance(identifier, str):
+            for i, k in enumerate(self.keys):
+                if k.api_key == identifier or k.api_key.endswith(identifier):
+                    return k, i
+        return None, -1
+
+    def get_global_stats(self) -> GlobalStats:
+        """Aggregates usage across all keys and models."""
+        total = UsageSnapshot()
+        keys_summary = []
+        with self.lock:
+            for i, key in enumerate(self.keys):
+                snap = key.get_total_snapshot()
+                total = total + snap
+                suffix = get_key_suffix(key.api_key)
+                keys_summary.append(KeySummary(index=i, suffix=suffix, snapshot=snap))
+        return GlobalStats(total=total, keys=keys_summary)
+
+    def get_key_stats(self, identifier: Union[int, str]) -> Optional[KeyDetailedStats]:
+        """Stats for a specific key, including per-model breakdown."""
+        with self.lock:
+            key, idx = self._find_key(identifier)
+            if not key:
+                return None
+            total_snap = key.get_total_snapshot()
+            breakdown = {}
+            for model, bucket in key.buckets.items():
+                breakdown[model] = bucket.get_snapshot()
+            suffix = get_key_suffix(key.api_key)
+            return KeyDetailedStats(index=idx, suffix=suffix, total=total_snap, breakdown=breakdown)
+
+    def get_model_stats(self, model_id: str) -> ModelAggregatedStats:
+        """Aggregates stats for ONE model across ALL keys."""
+        total = UsageSnapshot()
+        contributing_keys = []
+        with self.lock:
+            for i, key in enumerate(self.keys):
+                if model_id in key.buckets:
+                    snap = key.buckets[model_id].get_snapshot()
+                    total = total + snap
+                    suffix = get_key_suffix(key.api_key)
+                    contributing_keys.append(KeySummary(index=i, suffix=suffix, snapshot=snap))
+        return ModelAggregatedStats(model_id=model_id, total=total, keys=contributing_keys)
+
+    def get_granular_stats(self, identifier: Union[int, str], model_id: str) -> Optional[KeySummary]:
+        """Specific Key + Specific Model."""
+        with self.lock:
+            key, idx = self._find_key(identifier)
+            if not key:
+                return None
+            suffix = get_key_suffix(key.api_key)
+            snap = key.buckets[model_id].get_snapshot() if model_id in key.buckets else UsageSnapshot()
+            return KeySummary(index=idx, suffix=suffix, snapshot=snap)
+