Spaces:
Running
Running
File size: 2,331 Bytes
729a1f7 0fee802 729a1f7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
# ────────────────────────────── utils/rotator.py ──────────────────────────────
import os
import itertools
from ..logger import get_logger
from typing import Optional
import httpx
logger = get_logger("ROTATOR", __name__)
class APIKeyRotator:
"""
Round-robin API key rotator.
- Loads keys from env vars with given prefix (e.g., GEMINI_API_1..5)
- get_key() returns current key
- rotate() moves to next key
- on HTTP 401/429/5xx you should call rotate() and retry (bounded)
"""
def __init__(self, prefix: str, max_slots: int = 5):
self.keys = []
for i in range(1, max_slots + 1):
v = os.getenv(f"{prefix}{i}")
if v:
self.keys.append(v.strip())
if not self.keys:
logger.warning(f"No API keys found for prefix {prefix}. Calls will likely fail.")
self._cycle = itertools.cycle([""])
else:
self._cycle = itertools.cycle(self.keys)
self.current = next(self._cycle)
def get_key(self) -> Optional[str]:
return self.current
def rotate(self) -> Optional[str]:
self.current = next(self._cycle)
logger.info("Rotated API key.")
return self.current
async def robust_post_json(url: str, headers: dict, payload: dict, rotator: APIKeyRotator, max_retries: int = 5):
"""
POST JSON with simple retry+rotate on 401/403/429/5xx.
Returns json response.
"""
for attempt in range(max_retries):
try:
async with httpx.AsyncClient(timeout=60) as client:
r = await client.post(url, headers=headers, json=payload)
if r.status_code in (401, 403, 429) or (500 <= r.status_code < 600):
logger.warning(f"HTTP {r.status_code} from provider. Rotating key and retrying ({attempt+1}/{max_retries})")
rotator.rotate()
continue
r.raise_for_status()
return r.json()
except Exception as e:
logger.warning(f"Request error: {e}. Rotating and retrying ({attempt+1}/{max_retries})")
rotator.rotate()
raise RuntimeError("Provider request failed after retries.") |