| """ | |
| Shared LLM + Zep API key slot (same index for both) with rate-limit handling: | |
| try primary → on limit try secondary → on limit sleep 60s → retry from primary once more. | |
| """ | |
| from __future__ import annotations | |
| import contextvars | |
| import time | |
| from typing import Any, Callable, TypeVar | |
| from ..config import Config | |
| from ..utils.logger import get_logger | |
| logger = get_logger("mirofish.api_key_runtime") | |
| T = TypeVar("T") | |
| _key_slot: contextvars.ContextVar[int] = contextvars.ContextVar("miro_api_key_slot", default=0) | |
| def reset_pipeline_api_keys() -> None: | |
| """Call at the start of each full pipeline run (worker thread).""" | |
| _key_slot.set(0) | |
| def current_key_slot() -> int: | |
| return _key_slot.get() | |
| def _max_key_slots() -> int: | |
| lk = Config.get_llm_api_keys() | |
| zk = Config.get_zep_api_keys() | |
| n = max(len(lk or []), len(zk or []), 1) | |
| return n | |
| def set_key_slot(i: int) -> None: | |
| n = _max_key_slots() | |
| _key_slot.set(max(0, min(i, n - 1))) | |
| def get_llm_key() -> str | None: | |
| keys = Config.get_llm_api_keys() | |
| if not keys: | |
| return None | |
| i = min(_key_slot.get(), len(keys) - 1) | |
| return keys[i] | |
| def get_zep_key() -> str | None: | |
| keys = Config.get_zep_api_keys() | |
| if not keys: | |
| return None | |
| i = min(_key_slot.get(), len(keys) - 1) | |
| return keys[i] | |
| def is_rate_limit_error(exc: BaseException) -> bool: | |
| try: | |
| import openai | |
| if isinstance(exc, getattr(openai, "RateLimitError", ())): | |
| return True | |
| if isinstance(exc, getattr(openai, "APIStatusError", ())): | |
| if getattr(exc, "status_code", None) == 429: | |
| return True | |
| except Exception: | |
| pass | |
| st = getattr(exc, "status_code", None) | |
| if st == 429: | |
| return True | |
| msg = str(exc).lower() | |
| for phrase in ( | |
| "rate limit", | |
| "429", | |
| "too many requests", | |
| "quota", | |
| "resource_exhausted", | |
| "resource exhausted", | |
| "throttl", | |
| "over capacity", | |
| "capacity", | |
| ): | |
| if phrase in msg: | |
| return True | |
| return False | |
| def is_auth_or_invalid_key_error(exc: BaseException) -> bool: | |
| """401 / invalid API key — rotate to next paired slot when LLM_2/ZEP_2 exist.""" | |
| try: | |
| import openai | |
| if isinstance(exc, getattr(openai, "AuthenticationError", ())): | |
| return True | |
| if isinstance(exc, getattr(openai, "APIStatusError", ())): | |
| if getattr(exc, "status_code", None) == 401: | |
| return True | |
| except Exception: | |
| pass | |
| if getattr(exc, "status_code", None) == 401: | |
| return True | |
| msg = str(exc).lower() | |
| if "invalid api key" in msg or "invalid_api_key" in msg: | |
| return True | |
| if "401" in msg and ("unauthorized" in msg or "authentication" in msg): | |
| return True | |
| return False | |
| def should_rotate_api_key(exc: BaseException) -> bool: | |
| return is_rate_limit_error(exc) or is_auth_or_invalid_key_error(exc) | |
| def call_with_limit_rotation(op: Callable[[], T]) -> T: | |
| """ | |
| Run op() using keys at the current slot. On rate-limit: | |
| advance to next key slot if available; else sleep 60s, reset to slot 0, try again; | |
| if still rate-limited after that full cycle, re-raise. | |
| """ | |
| cooldown_used = False | |
| last_exc: BaseException | None = None | |
| while True: | |
| try: | |
| return op() | |
| except BaseException as e: | |
| if not should_rotate_api_key(e): | |
| raise | |
| last_exc = e | |
| idx = current_key_slot() | |
| n = _max_key_slots() | |
| if idx + 1 < n: | |
| set_key_slot(idx + 1) | |
| logger.warning( | |
| "Rate limit hit; switching API key slot to %s (LLM+Zep)", current_key_slot() | |
| ) | |
| continue | |
| if not cooldown_used: | |
| logger.warning("Rate limit on all API key slots; sleeping 60s then retrying from slot 0") | |
| time.sleep(60) | |
| cooldown_used = True | |
| set_key_slot(0) | |
| continue | |
| raise last_exc from None | |