open_keypool
open_keypool — a minimal Python library for pooling and rotating API keys.
Avoids HTTP 429 rate-limit errors by cycling through a pool of keys with cooldown and disablement support. Provide a list of keys (or pull them from Doppler), choose a rotation strategy (round-robin or least-recently-used), and the pool handles cooldown on rate-limit responses and permanent disablement on invalid keys — all thread-safe.
Install
pip install open-keypool
Quickstart — Local keys array
from open_keypool import KeyPool, AllKeysExhaustedError
pool = KeyPool(keys=["sk-key1", "sk-key2", "sk-key3"], strategy="round_robin")
for attempt in range(pool.max_retries):
key = pool.get_key()
response = call_your_api(key)
if response.status_code == 429:
retry_after = float(response.headers.get("Retry-After", 0))
pool.mark_rate_limited(key, retry_after=retry_after or None)
elif response.status_code in (401, 403):
pool.mark_invalid(key)
else:
pool.mark_success(key)
break
Quickstart — Doppler
import os
from open_keypool import KeyPool
DOPPLER_TOKEN = os.getenv("DOPPLER_TOKEN", "dp.st.YOUR_SERVICE_TOKEN")
PROJECT_NAME = "refactor-ai"
CONFIG_NAME = "dev"
pool = KeyPool.from_doppler(
token=DOPPLER_TOKEN,
project=PROJECT_NAME,
config=CONFIG_NAME,
key_prefix="MY_APP_",
strategy="lru",
)
1"""open_keypool — a minimal Python library for pooling and rotating API keys. 2 3Avoids HTTP 429 rate-limit errors by cycling through a pool of keys with 4cooldown and disablement support. Provide a list of keys (or pull them from 5Doppler), choose a rotation strategy (round-robin or least-recently-used), 6and the pool handles cooldown on rate-limit responses and permanent 7disablement on invalid keys — all thread-safe. 8 9Install 10------- 11.. code-block:: bash 12 13 pip install open-keypool 14 15Quickstart — Local keys array 16----------------------------- 17.. code-block:: python 18 19 from open_keypool import KeyPool, AllKeysExhaustedError 20 21 pool = KeyPool(keys=["sk-key1", "sk-key2", "sk-key3"], strategy="round_robin") 22 23 for attempt in range(pool.max_retries): 24 key = pool.get_key() 25 response = call_your_api(key) 26 if response.status_code == 429: 27 retry_after = float(response.headers.get("Retry-After", 0)) 28 pool.mark_rate_limited(key, retry_after=retry_after or None) 29 elif response.status_code in (401, 403): 30 pool.mark_invalid(key) 31 else: 32 pool.mark_success(key) 33 break 34 35Quickstart — Doppler 36-------------------- 37.. code-block:: python 38 39 import os 40 from open_keypool import KeyPool 41 42 DOPPLER_TOKEN = os.getenv("DOPPLER_TOKEN", "dp.st.YOUR_SERVICE_TOKEN") 43 PROJECT_NAME = "refactor-ai" 44 CONFIG_NAME = "dev" 45 46 pool = KeyPool.from_doppler( 47 token=DOPPLER_TOKEN, 48 project=PROJECT_NAME, 49 config=CONFIG_NAME, 50 key_prefix="MY_APP_", 51 strategy="lru", 52 ) 53""" 54 55from open_keypool.core import AllKeysExhaustedError, KeyPool, KeyState 56 57__all__ = ["KeyPool", "AllKeysExhaustedError", "KeyState"]
91class KeyPool: 92 """A thread-safe pool of API keys with cooldown and rotation strategies. 93 94 Manages a collection of API keys, cycling through them using a configurable 95 strategy (round-robin or least-recently-used). Keys that receive HTTP 429 96 responses can be placed on cooldown and will automatically recover after 97 the cooldown period. Permanently invalid keys can be explicitly disabled. 98 99 Parameters 100 ---------- 101 keys : list[str] 102 Initial list of API key strings. At least one key is required. 103 max_retries : int, optional 104 Maximum number of retries the caller should attempt per operation. 105 Stored as ``self.max_retries`` for the caller's reference; the library 106 does not perform automatic retries. Default is 3. 107 cooldown_seconds : int, optional 108 Number of seconds a key stays in COOLDOWN after being rate-limited 109 before it becomes ACTIVE again. Default is 60. 110 strategy : str, optional 111 Rotation strategy. ``"round_robin"`` cycles through keys in insertion 112 order. ``"lru"`` selects the key with the oldest ``last_used`` 113 timestamp. Default is ``"round_robin"``. 114 115 Raises 116 ------ 117 ValueError 118 If *keys* is empty or ``None``. 119 120 Examples 121 -------- 122 >>> pool = KeyPool(["key-a", "key-b", "key-c"], strategy="round_robin") 123 >>> pool.get_key() 124 >>> pool.mark_success(pool.get_key()) 125 """ 126 127 def __init__( 128 self, 129 keys: list[str] | None = None, 130 max_retries: int = 3, 131 cooldown_seconds: int = 60, 132 strategy: str = "round_robin", 133 ) -> None: 134 if not keys: 135 raise ValueError("KeyPool requires at least one key (keys must be a non-empty list).") 136 137 if strategy not in ("round_robin", "lru"): 138 raise ValueError(f"Unknown strategy '{strategy}'. Use 'round_robin' or 'lru'.") 139 140 self.max_retries: int = max_retries 141 self.cooldown_seconds: int = cooldown_seconds 142 self.strategy: str = strategy 143 144 self._records: list[_KeyRecord] = [_KeyRecord(key=k) for k in keys] 145 self._round_robin_index: int = 0 146 self._lock: threading.Lock = threading.Lock() 147 148 # ------------------------------------------------------------------ 149 # Internal helpers 150 # ------------------------------------------------------------------ 151 152 def _find_record(self, key: str) -> _KeyRecord | None: 153 """Return the ``_KeyRecord`` for *key*, or ``None`` if not found.""" 154 for rec in self._records: 155 if rec.key == key: 156 return rec 157 return None 158 159 def _recover_cooldown_keys(self) -> None: 160 """Flip COOLDOWN keys whose cooldown has expired back to ACTIVE.""" 161 now = time.monotonic() 162 for rec in self._records: 163 if rec.state == KeyState.COOLDOWN and rec.cooldown_until is not None and now >= rec.cooldown_until: 164 rec.state = KeyState.ACTIVE 165 rec.cooldown_until = None 166 167 def _active_records(self) -> list[_KeyRecord]: 168 """Return all currently-ACTIVE records.""" 169 return [r for r in self._records if r.state == KeyState.ACTIVE] 170 171 @classmethod 172 def from_doppler( 173 cls, 174 token: str, 175 project: str, 176 config: str, 177 key_prefix: str | None = None, 178 max_retries: int = 3, 179 cooldown_seconds: int = 60, 180 strategy: str = "round_robin", 181 force_refresh: bool = False, 182 ) -> KeyPool: 183 """Create a ``KeyPool`` by fetching API keys from Doppler. 184 185 Calls the Doppler secrets-download REST endpoint and populates the 186 pool with secret values that match *key_prefix* (if given). Results 187 are cached in an in-memory, module-level ``TTLCache`` (1 hour TTL) 188 keyed by ``(project, config, key_prefix)`` so that repeated calls 189 within the same process avoid redundant network requests. 190 191 The cache is purely in-memory and empties naturally on every fresh 192 process start — it is never persisted to disk. 193 194 Parameters 195 ---------- 196 token : str 197 Doppler service-token for bearer authentication 198 (e.g. ``"dp.st.YOUR_SERVICE_TOKEN"``). 199 project : str 200 Doppler project name. 201 config : str 202 Doppler config/environment name (e.g. ``"dev"``, ``"prd"``). 203 key_prefix : str | None, optional 204 If provided, only secrets whose name starts with this string are 205 included as keys. ``None`` (default) includes all secrets. 206 max_retries : int, optional 207 Passed through to the ``KeyPool`` constructor. Default is 3. 208 cooldown_seconds : int, optional 209 Passed through to the ``KeyPool`` constructor. Default is 60. 210 strategy : str, optional 211 Passed through to the ``KeyPool`` constructor. 212 Default is ``"round_robin"``. 213 force_refresh : bool, optional 214 If ``True``, bypass the cache and re-fetch from Doppler even when 215 a valid cache entry exists. Default is ``False``. 216 217 Returns 218 ------- 219 KeyPool 220 A new ``KeyPool`` instance populated with the fetched keys. 221 222 Raises 223 ------ 224 RuntimeError 225 If the Doppler API call fails (non-2xx status) or returns zero 226 keys — the cache is **not** populated on failure. 227 228 Examples 229 -------- 230 >>> import os 231 >>> pool = KeyPool.from_doppler( 232 ... token=os.getenv("DOPPLER_TOKEN", "dp.st.YOUR_SERVICE_TOKEN"), 233 ... project="refactor-ai", 234 ... config="dev", 235 ... key_prefix="API_KEY_", 236 ... ) 237 """ 238 cache_key = (project, config, key_prefix) 239 240 if not force_refresh: 241 cached = _DOPPLER_CACHE.get(cache_key) 242 if cached is not None: 243 return cls( 244 keys=list(cached), 245 max_retries=max_retries, 246 cooldown_seconds=cooldown_seconds, 247 strategy=strategy, 248 ) 249 250 try: 251 response = httpx.get( 252 _DOPPLER_DOWNLOAD_URL, 253 params={"project": project, "config": config}, 254 headers={"Authorization": f"Bearer {token}"}, 255 ) 256 response.raise_for_status() 257 data = response.json() 258 except httpx.HTTPError as exc: 259 raise RuntimeError( 260 f"Doppler API request failed: {exc.__class__.__name__}" 261 ) from exc 262 263 secrets: dict[str, str] = data.get("secrets", {}) 264 fetched_keys: list[str] = [] 265 for name, secret_data in secrets.items(): 266 if isinstance(secret_data, dict): 267 raw = secret_data.get("raw") or secret_data.get("computed", "") 268 else: 269 raw = str(secret_data) 270 if key_prefix is None or name.startswith(key_prefix): 271 fetched_keys.append(raw) 272 273 if not fetched_keys: 274 raise RuntimeError( 275 f"Doppler returned zero keys for project='{project}', " 276 f"config='{config}', key_prefix={key_prefix!r}" 277 ) 278 279 _DOPPLER_CACHE[cache_key] = tuple(fetched_keys) 280 281 return cls( 282 keys=fetched_keys, 283 max_retries=max_retries, 284 cooldown_seconds=cooldown_seconds, 285 strategy=strategy, 286 ) 287 288 @classmethod 289 def from_env( 290 cls, 291 suffix: str, 292 env_file: str | None = None, 293 max_retries: int = 3, 294 cooldown_seconds: int = 60, 295 strategy: str = "round_robin", 296 ) -> KeyPool: 297 """Create a ``KeyPool`` from environment variables matching a suffix. 298 299 Calls ``python-dotenv``'s ``load_dotenv()`` to load a ``.env`` file 300 (if *env_file* is given or a ``.env`` exists in the current directory), 301 then scans **all** environment variables for names ending with 302 *suffix*. The matching variable values are used as API keys. 303 304 Parameters 305 ---------- 306 suffix : str 307 Environment-variable name suffix to match (case-sensitive). 308 For example, ``"GROQ_KEY"`` matches ``TSN_GROQ_KEY``, 309 ``BACKUP_GROQ_KEY``, etc. 310 env_file : str | None, optional 311 Path to a ``.env`` file to load before scanning. ``None`` 312 (default) lets ``load_dotenv()`` find ``.env`` automatically. 313 max_retries : int, optional 314 Passed through to the ``KeyPool`` constructor. Default is 3. 315 cooldown_seconds : int, optional 316 Passed through to the ``KeyPool`` constructor. Default is 60. 317 strategy : str, optional 318 Passed through to the ``KeyPool`` constructor. 319 Default is ``"round_robin"``. 320 321 Returns 322 ------- 323 KeyPool 324 A new ``KeyPool`` instance populated with matching env values. 325 326 Raises 327 ------ 328 ValueError 329 If *suffix* is empty or ``None``. 330 RuntimeError 331 If no environment variables match *suffix*. 332 333 Examples 334 -------- 335 >>> # .env contains: TSN_GROQ_KEY_1=sk-abc TSN_GROQ_KEY_2=sk-def 336 >>> pool = KeyPool.from_env(suffix="GROQ_KEY") 337 """ 338 if not suffix: 339 raise ValueError("suffix must be a non-empty string.") 340 341 load_dotenv(env_file) 342 343 import os 344 345 matched: list[str] = [] 346 for name, value in os.environ.items(): 347 if name.endswith(suffix) and value.strip(): 348 matched.append(value.strip()) 349 350 if not matched: 351 raise RuntimeError( 352 f"No environment variables ending with '{suffix}' found " 353 f"(env_file={env_file!r})." 354 ) 355 356 return cls( 357 keys=matched, 358 max_retries=max_retries, 359 cooldown_seconds=cooldown_seconds, 360 strategy=strategy, 361 ) 362 363 @classmethod 364 def from_json( 365 cls, 366 path: str, 367 suffix: str | None = None, 368 max_retries: int = 3, 369 cooldown_seconds: int = 60, 370 strategy: str = "round_robin", 371 ) -> KeyPool: 372 """Create a ``KeyPool`` from a JSON file. 373 374 The JSON file must contain a **flat object** whose values are the 375 API key strings. Only entries whose *key name* ends with *suffix* 376 are included (if *suffix* is ``None``, all entries are used). 377 378 Expected JSON format:: 379 380 { 381 "TSN_GROQ_KEY_1": "sk-abc123", 382 "TSN_GROQ_KEY_2": "sk-def456", 383 "OTHER_SECRET": "sk-ghi789" 384 } 385 386 Parameters 387 ---------- 388 path : str 389 Path to the JSON file. 390 suffix : str | None, optional 391 If provided, only entries whose key name ends with this string 392 are included. ``None`` (default) includes all entries. 393 max_retries : int, optional 394 Passed through to the ``KeyPool`` constructor. Default is 3. 395 cooldown_seconds : int, optional 396 Passed through to the ``KeyPool`` constructor. Default is 60. 397 strategy : str, optional 398 Passed through to the ``KeyPool`` constructor. 399 Default is ``"round_robin"``. 400 401 Returns 402 ------- 403 KeyPool 404 A new ``KeyPool`` instance populated with matching JSON values. 405 406 Raises 407 ------ 408 FileNotFoundError 409 If *path* does not exist. 410 ValueError 411 If the file is not valid JSON or is not a flat object. 412 RuntimeError 413 If no entries match *suffix* (or the object is empty). 414 415 Examples 416 -------- 417 >>> # keys.json: {"GROQ_1": "sk-abc", "GROQ_2": "sk-def", "OTHER": "sk-ghi"} 418 >>> pool = KeyPool.from_json("keys.json", suffix="GROQ") 419 """ 420 import json as _json 421 import os 422 423 if not os.path.isfile(path): 424 raise FileNotFoundError(f"JSON file not found: {path}") 425 426 with open(path, "r", encoding="utf-8") as fh: 427 try: 428 data = _json.load(fh) 429 except _json.JSONDecodeError as exc: 430 raise ValueError(f"Invalid JSON in {path}: {exc}") from exc 431 432 if not isinstance(data, dict): 433 raise ValueError( 434 f"Expected a JSON object at top level in {path}, got {type(data).__name__}." 435 ) 436 437 if not data: 438 raise RuntimeError(f"JSON object in {path} is empty.") 439 440 matched: list[str] = [] 441 for name, value in data.items(): 442 if not isinstance(value, str): 443 raise ValueError( 444 f"All values must be strings in {path}. " 445 f"Key '{name}' has type {type(value).__name__}." 446 ) 447 if suffix is None or name.endswith(suffix): 448 if value.strip(): 449 matched.append(value.strip()) 450 451 if not matched: 452 raise RuntimeError( 453 f"No entries ending with '{suffix}' found in {path}." 454 ) 455 456 return cls( 457 keys=matched, 458 max_retries=max_retries, 459 cooldown_seconds=cooldown_seconds, 460 strategy=strategy, 461 ) 462 463 # ------------------------------------------------------------------ 464 # Public API 465 # ------------------------------------------------------------------ 466 467 def add_key(self, key: str) -> None: 468 """Add a new ACTIVE key to the pool at runtime. 469 470 If the key is already present in the pool, this is a no-op. 471 472 Parameters 473 ---------- 474 key : str 475 The API key string to add. 476 477 Examples 478 -------- 479 >>> pool = KeyPool(["key-a"]) 480 >>> pool.add_key("key-b") 481 """ 482 with self._lock: 483 if self._find_record(key) is None: 484 self._records.append(_KeyRecord(key=key)) 485 486 def remove_key(self, key: str) -> None: 487 """Remove a key from the pool regardless of its current state. 488 489 If the key is not in the pool, this is a no-op. 490 491 Parameters 492 ---------- 493 key : str 494 The API key string to remove. 495 496 Examples 497 -------- 498 >>> pool = KeyPool(["key-a", "key-b"]) 499 >>> pool.remove_key("key-b") 500 """ 501 with self._lock: 502 rec = self._find_record(key) 503 if rec is not None: 504 self._records.remove(rec) 505 506 def get_key(self) -> str: 507 """Return the next available ACTIVE key according to the pool's strategy. 508 509 Before selecting, any COOLDOWN key whose cooldown period has expired 510 is automatically flipped back to ACTIVE. 511 512 **Round-robin** (``strategy="round_robin"``): iterates through keys 513 in insertion order, maintaining an internal cursor that wraps around. 514 515 **LRU** (``strategy="lru"``): picks the ACTIVE key with the oldest 516 ``last_used`` timestamp and updates it to now upon selection. 517 518 Returns 519 ------- 520 str 521 An ACTIVE API key. 522 523 Raises 524 ------ 525 AllKeysExhaustedError 526 If no ACTIVE key exists in the pool. The message includes the 527 soonest recovery time in seconds when at least one key is in 528 COOLDOWN, otherwise it says all keys are disabled. 529 530 Examples 531 -------- 532 >>> pool = KeyPool(["key-a", "key-b"]) 533 >>> key = pool.get_key() 534 >>> pool.mark_success(key) 535 """ 536 with self._lock: 537 self._recover_cooldown_keys() 538 active = self._active_records() 539 540 if not active: 541 cooldown_records = [r for r in self._records if r.state == KeyState.COOLDOWN] 542 if cooldown_records: 543 now = time.monotonic() 544 soonest = min( 545 (r.cooldown_until - now for r in cooldown_records if r.cooldown_until is not None), 546 default=None, 547 ) 548 if soonest is not None and soonest > 0: 549 raise AllKeysExhaustedError( 550 f"No active keys available. Recovery in {soonest:.1f}s " 551 f"({self._cooldown_summary(cooldown_records, now)})." 552 ) 553 else: 554 raise AllKeysExhaustedError( 555 "No active keys available. All keys are in cooldown or disabled." 556 ) 557 raise AllKeysExhaustedError("No active keys available. All keys are disabled.") 558 559 if self.strategy == "round_robin": 560 self._round_robin_index %= len(active) 561 rec = active[self._round_robin_index] 562 self._round_robin_index += 1 563 else: # lru 564 rec = min(active, key=lambda r: r.last_used) 565 rec.last_used = time.monotonic() 566 567 return rec.key 568 569 def handle_response( 570 self, 571 key: str, 572 status_code: int, 573 headers: dict[str, str] | None = None, 574 body: dict | str | None = None, 575 ) -> KeyState: 576 """Feed an HTTP response to the pool — it decides what to do with the key. 577 578 Introspects the status code and response body and automatically: 579 580 - On **2xx**: marks the key as successful (``mark_success``). 581 - On **429** or **413**, or when the response body contains 582 ``error.code == "rate_limit_exceeded"``: places the key on 583 COOLDOWN using ``Retry-After`` if present, otherwise the pool's 584 ``cooldown_seconds``. 585 - On **401** or **403**: permanently disables the key 586 (``mark_invalid``). 587 - On **5xx**: places the key on COOLDOWN (transient server error). 588 589 All relevant details (status code, error code, error message) are 590 stored on the key record and surfaced in ``status()``. 591 592 Parameters 593 ---------- 594 key : str 595 The API key that was used for the request. 596 status_code : int 597 HTTP status code from the response. 598 headers : dict[str, str] | None, optional 599 Response headers (used to extract ``Retry-After``). 600 body : dict | str | None, optional 601 Parsed JSON body (``dict``) or raw response text (``str``). 602 603 Returns 604 ------- 605 KeyState 606 The new state of the key after processing. 607 608 Examples 609 -------- 610 >>> pool = KeyPool(["key-a", "key-b"]) 611 >>> k = pool.get_key() 612 >>> # Successful call: 613 >>> pool.handle_response(k, 200, body={"choices": [...]}) 614 <KeyState.ACTIVE: 'active'> 615 >>> # Rate-limit (Groq-style in-body): 616 >>> pool.handle_response(k, 200, body={"error": {"code": "rate_limit_exceeded", "message": "TPM limit"}}) 617 <KeyState.COOLDOWN: 'cooldown'> 618 >>> # Re-raise to get a fresh key on cooldown: 619 >>> k2 = pool.get_key() 620 """ 621 headers = headers or {} 622 623 # ── parse body for error details ── 624 error_code = str(status_code) 625 error_message = "" 626 if isinstance(body, dict): 627 err = body.get("error", {}) 628 if isinstance(err, dict): 629 if err.get("code") == "rate_limit_exceeded": 630 error_code = "rate_limit_exceeded" 631 error_message = err.get("message", "") 632 elif isinstance(err, str): 633 error_message = err 634 elif isinstance(body, str): 635 error_message = body[:200] 636 637 with self._lock: 638 rec = self._find_record(key) 639 if rec is None: 640 return KeyState.DISABLED # key doesn't exist, nothing to do 641 642 rec.last_status_code = status_code 643 644 # ── 2xx success ── 645 if 200 <= status_code < 300: 646 rec.failure_count = 0 647 rec.state = KeyState.ACTIVE 648 rec.last_error_code = None 649 rec.last_error_message = None 650 return KeyState.ACTIVE 651 652 # ── rate-limit (429, 413, or rate_limit_exceeded in body) ── 653 if status_code in (429, 413) or error_code == "rate_limit_exceeded": 654 ra = headers.get("Retry-After") 655 try: 656 retry_after = float(ra) if ra else None 657 except (ValueError, TypeError): 658 retry_after = None 659 rec.state = KeyState.COOLDOWN 660 rec.cooldown_until = time.monotonic() + ( 661 retry_after if retry_after is not None else self.cooldown_seconds 662 ) 663 rec.failure_count += 1 664 rec.last_error_code = error_code 665 rec.last_error_message = error_message 666 return KeyState.COOLDOWN 667 668 # ── auth failure (401, 403) ── 669 if status_code in (401, 403): 670 rec.state = KeyState.DISABLED 671 rec.last_error_code = error_code 672 rec.last_error_message = error_message 673 return KeyState.DISABLED 674 675 # ── server error (5xx) — transient, put on cooldown ── 676 if 500 <= status_code < 600: 677 rec.state = KeyState.COOLDOWN 678 rec.cooldown_until = time.monotonic() + self.cooldown_seconds 679 rec.failure_count += 1 680 rec.last_error_code = error_code 681 rec.last_error_message = error_message 682 return KeyState.COOLDOWN 683 684 # ── unknown status — also cooldown ── 685 rec.state = KeyState.COOLDOWN 686 rec.cooldown_until = time.monotonic() + self.cooldown_seconds 687 rec.failure_count += 1 688 rec.last_error_code = error_code 689 rec.last_error_message = error_message 690 return KeyState.COOLDOWN 691 692 def mark_rate_limited( 693 self, 694 key: str, 695 retry_after: float | None = None, 696 error_code: str | None = None, 697 error_message: str | None = None, 698 ) -> None: 699 """Mark a key as rate-limited (COOLDOWN). 700 701 The key will remain in COOLDOWN for *retry_after* seconds (or the 702 pool's ``cooldown_seconds`` if *retry_after* is ``None``). Its 703 ``failure_count`` is incremented. 704 705 Parameters 706 ---------- 707 key : str 708 The API key string. 709 retry_after : float | None, optional 710 Custom cooldown duration in seconds. If ``None``, defaults to 711 ``self.cooldown_seconds``. 712 error_code : str | None, optional 713 Machine-readable code for the last rate-limit error (e.g. 714 ``"rate_limit_exceeded"``, ``"413"``). Stored and surfaced in 715 ``status()``. 716 error_message : str | None, optional 717 Human-readable description of the last rate-limit error. 718 Stored and surfaced in ``status()``. 719 720 Examples 721 -------- 722 >>> pool = KeyPool(["key-a"]) 723 >>> pool.mark_rate_limited("key-a", retry_after=30) 724 >>> pool.mark_rate_limited("key-a", error_code="rate_limit_exceeded", 725 ... error_message="TPM limit 8000 exceeded") 726 """ 727 with self._lock: 728 rec = self._find_record(key) 729 if rec is None: 730 return 731 rec.state = KeyState.COOLDOWN 732 rec.cooldown_until = time.monotonic() + (retry_after if retry_after is not None else self.cooldown_seconds) 733 rec.failure_count += 1 734 rec.last_error_code = error_code 735 rec.last_error_message = error_message 736 737 def mark_invalid( 738 self, 739 key: str, 740 error_code: str | None = None, 741 error_message: str | None = None, 742 ) -> None: 743 """Permanently disable a key (DISABLED). 744 745 Disabled keys never auto-recover. Use this when a key returns an 746 authentication error (e.g. HTTP 401) rather than a rate-limit error. 747 748 Parameters 749 ---------- 750 key : str 751 The API key string. 752 error_code : str | None, optional 753 Machine-readable error code (e.g. ``"401"``, ``"invalid_api_key"``). 754 error_message : str | None, optional 755 Human-readable error description. 756 757 Examples 758 -------- 759 >>> pool = KeyPool(["key-a"]) 760 >>> pool.mark_invalid("key-a") 761 >>> pool.mark_invalid("key-a", error_code="401", error_message="Invalid API key") 762 """ 763 with self._lock: 764 rec = self._find_record(key) 765 if rec is None: 766 return 767 rec.state = KeyState.DISABLED 768 rec.last_error_code = error_code 769 rec.last_error_message = error_message 770 771 def mark_success(self, key: str) -> None: 772 """Reset a key's failure count to 0 and keep it ACTIVE. 773 774 Call this after a successful API response to indicate the key is 775 healthy and reset any transient failure tracking. 776 777 Parameters 778 ---------- 779 key : str 780 The API key string. 781 782 Examples 783 -------- 784 >>> pool = KeyPool(["key-a"]) 785 >>> k = pool.get_key() 786 >>> pool.mark_success(k) 787 """ 788 with self._lock: 789 rec = self._find_record(key) 790 if rec is None: 791 return 792 rec.failure_count = 0 793 rec.state = KeyState.ACTIVE 794 rec.last_status_code = None 795 rec.last_error_code = None 796 rec.last_error_message = None 797 798 def status(self) -> dict[str, dict]: 799 """Return a snapshot of every key's state without exposing raw keys. 800 801 Every key value in the returned dictionary is passed through ``mask()`` 802 so the caller can safely log or print the result. 803 804 Returns 805 ------- 806 dict[str, dict] 807 A mapping of ``{masked_key: {"state": str, "failure_count": int, 808 "cooldown_remaining": float | None, "last_status_code": int | None, 809 "last_error_code": str | None, "last_error_message": str | None}}`` 810 for every key in the pool. 811 812 Examples 813 -------- 814 >>> pool = KeyPool(["sk-abcdef1234567890"]) 815 >>> pool.status() 816 {'sk-abc...7890': {'state': 'active', 'failure_count': 0, 'cooldown_remaining': None, 'last_status_code': None, 'last_error_code': None, 'last_error_message': None}} 817 """ 818 with self._lock: 819 result: dict[str, dict] = {} 820 now = time.monotonic() 821 for rec in self._records: 822 if rec.state == KeyState.COOLDOWN and rec.cooldown_until is not None: 823 remaining = max(0.0, rec.cooldown_until - now) 824 else: 825 remaining = None 826 result[mask(rec.key)] = { 827 "state": rec.state.value, 828 "failure_count": rec.failure_count, 829 "cooldown_remaining": round(remaining, 1) if remaining is not None else None, 830 "last_status_code": rec.last_status_code, 831 "last_error_code": rec.last_error_code, 832 "last_error_message": rec.last_error_message, 833 } 834 return result 835 836 # ------------------------------------------------------------------ 837 # Private helpers 838 # ------------------------------------------------------------------ 839 840 def _cooldown_summary(self, cooldown_records: list[_KeyRecord], now: float) -> str: 841 """Return a brief summary of cooldown keys for error messages.""" 842 parts: list[str] = [] 843 for rec in cooldown_records: 844 if rec.cooldown_until is not None: 845 parts.append(f"{mask(rec.key)} in {max(0, rec.cooldown_until - now):.1f}s") 846 return ", ".join(parts) if parts else "unknown"
A thread-safe pool of API keys with cooldown and rotation strategies.
Manages a collection of API keys, cycling through them using a configurable strategy (round-robin or least-recently-used). Keys that receive HTTP 429 responses can be placed on cooldown and will automatically recover after the cooldown period. Permanently invalid keys can be explicitly disabled.
Parameters
keys : list[str]
Initial list of API key strings. At least one key is required.
max_retries : int, optional
Maximum number of retries the caller should attempt per operation.
Stored as self.max_retries for the caller's reference; the library
does not perform automatic retries. Default is 3.
cooldown_seconds : int, optional
Number of seconds a key stays in COOLDOWN after being rate-limited
before it becomes ACTIVE again. Default is 60.
strategy : str, optional
Rotation strategy. "round_robin" cycles through keys in insertion
order. "lru" selects the key with the oldest last_used
timestamp. Default is "round_robin".
Raises
ValueError
If keys is empty or None.
Examples
>>> pool = KeyPool(["key-a", "key-b", "key-c"], strategy="round_robin")
>>> pool.get_key()
>>> pool.mark_success(pool.get_key())
127 def __init__( 128 self, 129 keys: list[str] | None = None, 130 max_retries: int = 3, 131 cooldown_seconds: int = 60, 132 strategy: str = "round_robin", 133 ) -> None: 134 if not keys: 135 raise ValueError("KeyPool requires at least one key (keys must be a non-empty list).") 136 137 if strategy not in ("round_robin", "lru"): 138 raise ValueError(f"Unknown strategy '{strategy}'. Use 'round_robin' or 'lru'.") 139 140 self.max_retries: int = max_retries 141 self.cooldown_seconds: int = cooldown_seconds 142 self.strategy: str = strategy 143 144 self._records: list[_KeyRecord] = [_KeyRecord(key=k) for k in keys] 145 self._round_robin_index: int = 0 146 self._lock: threading.Lock = threading.Lock()
171 @classmethod 172 def from_doppler( 173 cls, 174 token: str, 175 project: str, 176 config: str, 177 key_prefix: str | None = None, 178 max_retries: int = 3, 179 cooldown_seconds: int = 60, 180 strategy: str = "round_robin", 181 force_refresh: bool = False, 182 ) -> KeyPool: 183 """Create a ``KeyPool`` by fetching API keys from Doppler. 184 185 Calls the Doppler secrets-download REST endpoint and populates the 186 pool with secret values that match *key_prefix* (if given). Results 187 are cached in an in-memory, module-level ``TTLCache`` (1 hour TTL) 188 keyed by ``(project, config, key_prefix)`` so that repeated calls 189 within the same process avoid redundant network requests. 190 191 The cache is purely in-memory and empties naturally on every fresh 192 process start — it is never persisted to disk. 193 194 Parameters 195 ---------- 196 token : str 197 Doppler service-token for bearer authentication 198 (e.g. ``"dp.st.YOUR_SERVICE_TOKEN"``). 199 project : str 200 Doppler project name. 201 config : str 202 Doppler config/environment name (e.g. ``"dev"``, ``"prd"``). 203 key_prefix : str | None, optional 204 If provided, only secrets whose name starts with this string are 205 included as keys. ``None`` (default) includes all secrets. 206 max_retries : int, optional 207 Passed through to the ``KeyPool`` constructor. Default is 3. 208 cooldown_seconds : int, optional 209 Passed through to the ``KeyPool`` constructor. Default is 60. 210 strategy : str, optional 211 Passed through to the ``KeyPool`` constructor. 212 Default is ``"round_robin"``. 213 force_refresh : bool, optional 214 If ``True``, bypass the cache and re-fetch from Doppler even when 215 a valid cache entry exists. Default is ``False``. 216 217 Returns 218 ------- 219 KeyPool 220 A new ``KeyPool`` instance populated with the fetched keys. 221 222 Raises 223 ------ 224 RuntimeError 225 If the Doppler API call fails (non-2xx status) or returns zero 226 keys — the cache is **not** populated on failure. 227 228 Examples 229 -------- 230 >>> import os 231 >>> pool = KeyPool.from_doppler( 232 ... token=os.getenv("DOPPLER_TOKEN", "dp.st.YOUR_SERVICE_TOKEN"), 233 ... project="refactor-ai", 234 ... config="dev", 235 ... key_prefix="API_KEY_", 236 ... ) 237 """ 238 cache_key = (project, config, key_prefix) 239 240 if not force_refresh: 241 cached = _DOPPLER_CACHE.get(cache_key) 242 if cached is not None: 243 return cls( 244 keys=list(cached), 245 max_retries=max_retries, 246 cooldown_seconds=cooldown_seconds, 247 strategy=strategy, 248 ) 249 250 try: 251 response = httpx.get( 252 _DOPPLER_DOWNLOAD_URL, 253 params={"project": project, "config": config}, 254 headers={"Authorization": f"Bearer {token}"}, 255 ) 256 response.raise_for_status() 257 data = response.json() 258 except httpx.HTTPError as exc: 259 raise RuntimeError( 260 f"Doppler API request failed: {exc.__class__.__name__}" 261 ) from exc 262 263 secrets: dict[str, str] = data.get("secrets", {}) 264 fetched_keys: list[str] = [] 265 for name, secret_data in secrets.items(): 266 if isinstance(secret_data, dict): 267 raw = secret_data.get("raw") or secret_data.get("computed", "") 268 else: 269 raw = str(secret_data) 270 if key_prefix is None or name.startswith(key_prefix): 271 fetched_keys.append(raw) 272 273 if not fetched_keys: 274 raise RuntimeError( 275 f"Doppler returned zero keys for project='{project}', " 276 f"config='{config}', key_prefix={key_prefix!r}" 277 ) 278 279 _DOPPLER_CACHE[cache_key] = tuple(fetched_keys) 280 281 return cls( 282 keys=fetched_keys, 283 max_retries=max_retries, 284 cooldown_seconds=cooldown_seconds, 285 strategy=strategy, 286 )
Create a KeyPool by fetching API keys from Doppler.
Calls the Doppler secrets-download REST endpoint and populates the
pool with secret values that match key_prefix (if given). Results
are cached in an in-memory, module-level TTLCache (1 hour TTL)
keyed by (project, config, key_prefix) so that repeated calls
within the same process avoid redundant network requests.
The cache is purely in-memory and empties naturally on every fresh process start — it is never persisted to disk.
Parameters
token : str
Doppler service-token for bearer authentication
(e.g. "dp.st.YOUR_SERVICE_TOKEN").
project : str
Doppler project name.
config : str
Doppler config/environment name (e.g. "dev", "prd").
key_prefix : str | None, optional
If provided, only secrets whose name starts with this string are
included as keys. None (default) includes all secrets.
max_retries : int, optional
Passed through to the KeyPool constructor. Default is 3.
cooldown_seconds : int, optional
Passed through to the KeyPool constructor. Default is 60.
strategy : str, optional
Passed through to the KeyPool constructor.
Default is "round_robin".
force_refresh : bool, optional
If True, bypass the cache and re-fetch from Doppler even when
a valid cache entry exists. Default is False.
Returns
KeyPool
A new KeyPool instance populated with the fetched keys.
Raises
RuntimeError If the Doppler API call fails (non-2xx status) or returns zero keys — the cache is not populated on failure.
Examples
>>> import os
>>> pool = KeyPool.from_doppler(
... token=os.getenv("DOPPLER_TOKEN", "dp.st.YOUR_SERVICE_TOKEN"),
... project="refactor-ai",
... config="dev",
... key_prefix="API_KEY_",
... )
288 @classmethod 289 def from_env( 290 cls, 291 suffix: str, 292 env_file: str | None = None, 293 max_retries: int = 3, 294 cooldown_seconds: int = 60, 295 strategy: str = "round_robin", 296 ) -> KeyPool: 297 """Create a ``KeyPool`` from environment variables matching a suffix. 298 299 Calls ``python-dotenv``'s ``load_dotenv()`` to load a ``.env`` file 300 (if *env_file* is given or a ``.env`` exists in the current directory), 301 then scans **all** environment variables for names ending with 302 *suffix*. The matching variable values are used as API keys. 303 304 Parameters 305 ---------- 306 suffix : str 307 Environment-variable name suffix to match (case-sensitive). 308 For example, ``"GROQ_KEY"`` matches ``TSN_GROQ_KEY``, 309 ``BACKUP_GROQ_KEY``, etc. 310 env_file : str | None, optional 311 Path to a ``.env`` file to load before scanning. ``None`` 312 (default) lets ``load_dotenv()`` find ``.env`` automatically. 313 max_retries : int, optional 314 Passed through to the ``KeyPool`` constructor. Default is 3. 315 cooldown_seconds : int, optional 316 Passed through to the ``KeyPool`` constructor. Default is 60. 317 strategy : str, optional 318 Passed through to the ``KeyPool`` constructor. 319 Default is ``"round_robin"``. 320 321 Returns 322 ------- 323 KeyPool 324 A new ``KeyPool`` instance populated with matching env values. 325 326 Raises 327 ------ 328 ValueError 329 If *suffix* is empty or ``None``. 330 RuntimeError 331 If no environment variables match *suffix*. 332 333 Examples 334 -------- 335 >>> # .env contains: TSN_GROQ_KEY_1=sk-abc TSN_GROQ_KEY_2=sk-def 336 >>> pool = KeyPool.from_env(suffix="GROQ_KEY") 337 """ 338 if not suffix: 339 raise ValueError("suffix must be a non-empty string.") 340 341 load_dotenv(env_file) 342 343 import os 344 345 matched: list[str] = [] 346 for name, value in os.environ.items(): 347 if name.endswith(suffix) and value.strip(): 348 matched.append(value.strip()) 349 350 if not matched: 351 raise RuntimeError( 352 f"No environment variables ending with '{suffix}' found " 353 f"(env_file={env_file!r})." 354 ) 355 356 return cls( 357 keys=matched, 358 max_retries=max_retries, 359 cooldown_seconds=cooldown_seconds, 360 strategy=strategy, 361 )
Create a KeyPool from environment variables matching a suffix.
Calls python-dotenv's load_dotenv() to load a .env file
(if env_file is given or a .env exists in the current directory),
then scans all environment variables for names ending with
suffix. The matching variable values are used as API keys.
Parameters
suffix : str
Environment-variable name suffix to match (case-sensitive).
For example, "GROQ_KEY" matches TSN_GROQ_KEY,
BACKUP_GROQ_KEY, etc.
env_file : str | None, optional
Path to a .env file to load before scanning. None
(default) lets load_dotenv() find .env automatically.
max_retries : int, optional
Passed through to the KeyPool constructor. Default is 3.
cooldown_seconds : int, optional
Passed through to the KeyPool constructor. Default is 60.
strategy : str, optional
Passed through to the KeyPool constructor.
Default is "round_robin".
Returns
KeyPool
A new KeyPool instance populated with matching env values.
Raises
ValueError
If suffix is empty or None.
RuntimeError
If no environment variables match suffix.
Examples
>>> # .env contains: TSN_GROQ_KEY_1=sk-abc TSN_GROQ_KEY_2=sk-def
>>> pool = KeyPool.from_env(suffix="GROQ_KEY")
363 @classmethod 364 def from_json( 365 cls, 366 path: str, 367 suffix: str | None = None, 368 max_retries: int = 3, 369 cooldown_seconds: int = 60, 370 strategy: str = "round_robin", 371 ) -> KeyPool: 372 """Create a ``KeyPool`` from a JSON file. 373 374 The JSON file must contain a **flat object** whose values are the 375 API key strings. Only entries whose *key name* ends with *suffix* 376 are included (if *suffix* is ``None``, all entries are used). 377 378 Expected JSON format:: 379 380 { 381 "TSN_GROQ_KEY_1": "sk-abc123", 382 "TSN_GROQ_KEY_2": "sk-def456", 383 "OTHER_SECRET": "sk-ghi789" 384 } 385 386 Parameters 387 ---------- 388 path : str 389 Path to the JSON file. 390 suffix : str | None, optional 391 If provided, only entries whose key name ends with this string 392 are included. ``None`` (default) includes all entries. 393 max_retries : int, optional 394 Passed through to the ``KeyPool`` constructor. Default is 3. 395 cooldown_seconds : int, optional 396 Passed through to the ``KeyPool`` constructor. Default is 60. 397 strategy : str, optional 398 Passed through to the ``KeyPool`` constructor. 399 Default is ``"round_robin"``. 400 401 Returns 402 ------- 403 KeyPool 404 A new ``KeyPool`` instance populated with matching JSON values. 405 406 Raises 407 ------ 408 FileNotFoundError 409 If *path* does not exist. 410 ValueError 411 If the file is not valid JSON or is not a flat object. 412 RuntimeError 413 If no entries match *suffix* (or the object is empty). 414 415 Examples 416 -------- 417 >>> # keys.json: {"GROQ_1": "sk-abc", "GROQ_2": "sk-def", "OTHER": "sk-ghi"} 418 >>> pool = KeyPool.from_json("keys.json", suffix="GROQ") 419 """ 420 import json as _json 421 import os 422 423 if not os.path.isfile(path): 424 raise FileNotFoundError(f"JSON file not found: {path}") 425 426 with open(path, "r", encoding="utf-8") as fh: 427 try: 428 data = _json.load(fh) 429 except _json.JSONDecodeError as exc: 430 raise ValueError(f"Invalid JSON in {path}: {exc}") from exc 431 432 if not isinstance(data, dict): 433 raise ValueError( 434 f"Expected a JSON object at top level in {path}, got {type(data).__name__}." 435 ) 436 437 if not data: 438 raise RuntimeError(f"JSON object in {path} is empty.") 439 440 matched: list[str] = [] 441 for name, value in data.items(): 442 if not isinstance(value, str): 443 raise ValueError( 444 f"All values must be strings in {path}. " 445 f"Key '{name}' has type {type(value).__name__}." 446 ) 447 if suffix is None or name.endswith(suffix): 448 if value.strip(): 449 matched.append(value.strip()) 450 451 if not matched: 452 raise RuntimeError( 453 f"No entries ending with '{suffix}' found in {path}." 454 ) 455 456 return cls( 457 keys=matched, 458 max_retries=max_retries, 459 cooldown_seconds=cooldown_seconds, 460 strategy=strategy, 461 )
Create a KeyPool from a JSON file.
The JSON file must contain a flat object whose values are the
API key strings. Only entries whose key name ends with suffix
are included (if suffix is None, all entries are used).
Expected JSON format::
{
"TSN_GROQ_KEY_1": "sk-abc123",
"TSN_GROQ_KEY_2": "sk-def456",
"OTHER_SECRET": "sk-ghi789"
}
Parameters
path : str
Path to the JSON file.
suffix : str | None, optional
If provided, only entries whose key name ends with this string
are included. None (default) includes all entries.
max_retries : int, optional
Passed through to the KeyPool constructor. Default is 3.
cooldown_seconds : int, optional
Passed through to the KeyPool constructor. Default is 60.
strategy : str, optional
Passed through to the KeyPool constructor.
Default is "round_robin".
Returns
KeyPool
A new KeyPool instance populated with matching JSON values.
Raises
FileNotFoundError If path does not exist. ValueError If the file is not valid JSON or is not a flat object. RuntimeError If no entries match suffix (or the object is empty).
Examples
>>> # keys.json: {"GROQ_1": "sk-abc", "GROQ_2": "sk-def", "OTHER": "sk-ghi"}
>>> pool = KeyPool.from_json("keys.json", suffix="GROQ")
467 def add_key(self, key: str) -> None: 468 """Add a new ACTIVE key to the pool at runtime. 469 470 If the key is already present in the pool, this is a no-op. 471 472 Parameters 473 ---------- 474 key : str 475 The API key string to add. 476 477 Examples 478 -------- 479 >>> pool = KeyPool(["key-a"]) 480 >>> pool.add_key("key-b") 481 """ 482 with self._lock: 483 if self._find_record(key) is None: 484 self._records.append(_KeyRecord(key=key))
Add a new ACTIVE key to the pool at runtime.
If the key is already present in the pool, this is a no-op.
Parameters
key : str The API key string to add.
Examples
>>> pool = KeyPool(["key-a"])
>>> pool.add_key("key-b")
486 def remove_key(self, key: str) -> None: 487 """Remove a key from the pool regardless of its current state. 488 489 If the key is not in the pool, this is a no-op. 490 491 Parameters 492 ---------- 493 key : str 494 The API key string to remove. 495 496 Examples 497 -------- 498 >>> pool = KeyPool(["key-a", "key-b"]) 499 >>> pool.remove_key("key-b") 500 """ 501 with self._lock: 502 rec = self._find_record(key) 503 if rec is not None: 504 self._records.remove(rec)
Remove a key from the pool regardless of its current state.
If the key is not in the pool, this is a no-op.
Parameters
key : str The API key string to remove.
Examples
>>> pool = KeyPool(["key-a", "key-b"])
>>> pool.remove_key("key-b")
506 def get_key(self) -> str: 507 """Return the next available ACTIVE key according to the pool's strategy. 508 509 Before selecting, any COOLDOWN key whose cooldown period has expired 510 is automatically flipped back to ACTIVE. 511 512 **Round-robin** (``strategy="round_robin"``): iterates through keys 513 in insertion order, maintaining an internal cursor that wraps around. 514 515 **LRU** (``strategy="lru"``): picks the ACTIVE key with the oldest 516 ``last_used`` timestamp and updates it to now upon selection. 517 518 Returns 519 ------- 520 str 521 An ACTIVE API key. 522 523 Raises 524 ------ 525 AllKeysExhaustedError 526 If no ACTIVE key exists in the pool. The message includes the 527 soonest recovery time in seconds when at least one key is in 528 COOLDOWN, otherwise it says all keys are disabled. 529 530 Examples 531 -------- 532 >>> pool = KeyPool(["key-a", "key-b"]) 533 >>> key = pool.get_key() 534 >>> pool.mark_success(key) 535 """ 536 with self._lock: 537 self._recover_cooldown_keys() 538 active = self._active_records() 539 540 if not active: 541 cooldown_records = [r for r in self._records if r.state == KeyState.COOLDOWN] 542 if cooldown_records: 543 now = time.monotonic() 544 soonest = min( 545 (r.cooldown_until - now for r in cooldown_records if r.cooldown_until is not None), 546 default=None, 547 ) 548 if soonest is not None and soonest > 0: 549 raise AllKeysExhaustedError( 550 f"No active keys available. Recovery in {soonest:.1f}s " 551 f"({self._cooldown_summary(cooldown_records, now)})." 552 ) 553 else: 554 raise AllKeysExhaustedError( 555 "No active keys available. All keys are in cooldown or disabled." 556 ) 557 raise AllKeysExhaustedError("No active keys available. All keys are disabled.") 558 559 if self.strategy == "round_robin": 560 self._round_robin_index %= len(active) 561 rec = active[self._round_robin_index] 562 self._round_robin_index += 1 563 else: # lru 564 rec = min(active, key=lambda r: r.last_used) 565 rec.last_used = time.monotonic() 566 567 return rec.key
Return the next available ACTIVE key according to the pool's strategy.
Before selecting, any COOLDOWN key whose cooldown period has expired is automatically flipped back to ACTIVE.
Round-robin (strategy="round_robin"): iterates through keys
in insertion order, maintaining an internal cursor that wraps around.
LRU (strategy="lru"): picks the ACTIVE key with the oldest
last_used timestamp and updates it to now upon selection.
Returns
str An ACTIVE API key.
Raises
AllKeysExhaustedError If no ACTIVE key exists in the pool. The message includes the soonest recovery time in seconds when at least one key is in COOLDOWN, otherwise it says all keys are disabled.
Examples
>>> pool = KeyPool(["key-a", "key-b"])
>>> key = pool.get_key()
>>> pool.mark_success(key)
569 def handle_response( 570 self, 571 key: str, 572 status_code: int, 573 headers: dict[str, str] | None = None, 574 body: dict | str | None = None, 575 ) -> KeyState: 576 """Feed an HTTP response to the pool — it decides what to do with the key. 577 578 Introspects the status code and response body and automatically: 579 580 - On **2xx**: marks the key as successful (``mark_success``). 581 - On **429** or **413**, or when the response body contains 582 ``error.code == "rate_limit_exceeded"``: places the key on 583 COOLDOWN using ``Retry-After`` if present, otherwise the pool's 584 ``cooldown_seconds``. 585 - On **401** or **403**: permanently disables the key 586 (``mark_invalid``). 587 - On **5xx**: places the key on COOLDOWN (transient server error). 588 589 All relevant details (status code, error code, error message) are 590 stored on the key record and surfaced in ``status()``. 591 592 Parameters 593 ---------- 594 key : str 595 The API key that was used for the request. 596 status_code : int 597 HTTP status code from the response. 598 headers : dict[str, str] | None, optional 599 Response headers (used to extract ``Retry-After``). 600 body : dict | str | None, optional 601 Parsed JSON body (``dict``) or raw response text (``str``). 602 603 Returns 604 ------- 605 KeyState 606 The new state of the key after processing. 607 608 Examples 609 -------- 610 >>> pool = KeyPool(["key-a", "key-b"]) 611 >>> k = pool.get_key() 612 >>> # Successful call: 613 >>> pool.handle_response(k, 200, body={"choices": [...]}) 614 <KeyState.ACTIVE: 'active'> 615 >>> # Rate-limit (Groq-style in-body): 616 >>> pool.handle_response(k, 200, body={"error": {"code": "rate_limit_exceeded", "message": "TPM limit"}}) 617 <KeyState.COOLDOWN: 'cooldown'> 618 >>> # Re-raise to get a fresh key on cooldown: 619 >>> k2 = pool.get_key() 620 """ 621 headers = headers or {} 622 623 # ── parse body for error details ── 624 error_code = str(status_code) 625 error_message = "" 626 if isinstance(body, dict): 627 err = body.get("error", {}) 628 if isinstance(err, dict): 629 if err.get("code") == "rate_limit_exceeded": 630 error_code = "rate_limit_exceeded" 631 error_message = err.get("message", "") 632 elif isinstance(err, str): 633 error_message = err 634 elif isinstance(body, str): 635 error_message = body[:200] 636 637 with self._lock: 638 rec = self._find_record(key) 639 if rec is None: 640 return KeyState.DISABLED # key doesn't exist, nothing to do 641 642 rec.last_status_code = status_code 643 644 # ── 2xx success ── 645 if 200 <= status_code < 300: 646 rec.failure_count = 0 647 rec.state = KeyState.ACTIVE 648 rec.last_error_code = None 649 rec.last_error_message = None 650 return KeyState.ACTIVE 651 652 # ── rate-limit (429, 413, or rate_limit_exceeded in body) ── 653 if status_code in (429, 413) or error_code == "rate_limit_exceeded": 654 ra = headers.get("Retry-After") 655 try: 656 retry_after = float(ra) if ra else None 657 except (ValueError, TypeError): 658 retry_after = None 659 rec.state = KeyState.COOLDOWN 660 rec.cooldown_until = time.monotonic() + ( 661 retry_after if retry_after is not None else self.cooldown_seconds 662 ) 663 rec.failure_count += 1 664 rec.last_error_code = error_code 665 rec.last_error_message = error_message 666 return KeyState.COOLDOWN 667 668 # ── auth failure (401, 403) ── 669 if status_code in (401, 403): 670 rec.state = KeyState.DISABLED 671 rec.last_error_code = error_code 672 rec.last_error_message = error_message 673 return KeyState.DISABLED 674 675 # ── server error (5xx) — transient, put on cooldown ── 676 if 500 <= status_code < 600: 677 rec.state = KeyState.COOLDOWN 678 rec.cooldown_until = time.monotonic() + self.cooldown_seconds 679 rec.failure_count += 1 680 rec.last_error_code = error_code 681 rec.last_error_message = error_message 682 return KeyState.COOLDOWN 683 684 # ── unknown status — also cooldown ── 685 rec.state = KeyState.COOLDOWN 686 rec.cooldown_until = time.monotonic() + self.cooldown_seconds 687 rec.failure_count += 1 688 rec.last_error_code = error_code 689 rec.last_error_message = error_message 690 return KeyState.COOLDOWN
Feed an HTTP response to the pool — it decides what to do with the key.
Introspects the status code and response body and automatically:
- On 2xx: marks the key as successful (
mark_success). - On 429 or 413, or when the response body contains
error.code == "rate_limit_exceeded": places the key on COOLDOWN usingRetry-Afterif present, otherwise the pool'scooldown_seconds. - On 401 or 403: permanently disables the key
(
mark_invalid). - On 5xx: places the key on COOLDOWN (transient server error).
All relevant details (status code, error code, error message) are
stored on the key record and surfaced in status().
Parameters
key : str
The API key that was used for the request.
status_code : int
HTTP status code from the response.
headers : dict[str, str] | None, optional
Response headers (used to extract Retry-After).
body : dict | str | None, optional
Parsed JSON body (dict) or raw response text (str).
Returns
KeyState The new state of the key after processing.
Examples
>>> pool = KeyPool(["key-a", "key-b"])
>>> k = pool.get_key()
>>> # Successful call:
>>> pool.handle_response(k, 200, body={"choices": [...]})
<KeyState.ACTIVE: 'active'>
>>> # Rate-limit (Groq-style in-body):
>>> pool.handle_response(k, 200, body={"error": {"code": "rate_limit_exceeded", "message": "TPM limit"}})
<KeyState.COOLDOWN: 'cooldown'>
>>> # Re-raise to get a fresh key on cooldown:
>>> k2 = pool.get_key()
692 def mark_rate_limited( 693 self, 694 key: str, 695 retry_after: float | None = None, 696 error_code: str | None = None, 697 error_message: str | None = None, 698 ) -> None: 699 """Mark a key as rate-limited (COOLDOWN). 700 701 The key will remain in COOLDOWN for *retry_after* seconds (or the 702 pool's ``cooldown_seconds`` if *retry_after* is ``None``). Its 703 ``failure_count`` is incremented. 704 705 Parameters 706 ---------- 707 key : str 708 The API key string. 709 retry_after : float | None, optional 710 Custom cooldown duration in seconds. If ``None``, defaults to 711 ``self.cooldown_seconds``. 712 error_code : str | None, optional 713 Machine-readable code for the last rate-limit error (e.g. 714 ``"rate_limit_exceeded"``, ``"413"``). Stored and surfaced in 715 ``status()``. 716 error_message : str | None, optional 717 Human-readable description of the last rate-limit error. 718 Stored and surfaced in ``status()``. 719 720 Examples 721 -------- 722 >>> pool = KeyPool(["key-a"]) 723 >>> pool.mark_rate_limited("key-a", retry_after=30) 724 >>> pool.mark_rate_limited("key-a", error_code="rate_limit_exceeded", 725 ... error_message="TPM limit 8000 exceeded") 726 """ 727 with self._lock: 728 rec = self._find_record(key) 729 if rec is None: 730 return 731 rec.state = KeyState.COOLDOWN 732 rec.cooldown_until = time.monotonic() + (retry_after if retry_after is not None else self.cooldown_seconds) 733 rec.failure_count += 1 734 rec.last_error_code = error_code 735 rec.last_error_message = error_message
Mark a key as rate-limited (COOLDOWN).
The key will remain in COOLDOWN for retry_after seconds (or the
pool's cooldown_seconds if retry_after is None). Its
failure_count is incremented.
Parameters
key : str
The API key string.
retry_after : float | None, optional
Custom cooldown duration in seconds. If None, defaults to
self.cooldown_seconds.
error_code : str | None, optional
Machine-readable code for the last rate-limit error (e.g.
"rate_limit_exceeded", "413"). Stored and surfaced in
status().
error_message : str | None, optional
Human-readable description of the last rate-limit error.
Stored and surfaced in status().
Examples
>>> pool = KeyPool(["key-a"])
>>> pool.mark_rate_limited("key-a", retry_after=30)
>>> pool.mark_rate_limited("key-a", error_code="rate_limit_exceeded",
... error_message="TPM limit 8000 exceeded")
737 def mark_invalid( 738 self, 739 key: str, 740 error_code: str | None = None, 741 error_message: str | None = None, 742 ) -> None: 743 """Permanently disable a key (DISABLED). 744 745 Disabled keys never auto-recover. Use this when a key returns an 746 authentication error (e.g. HTTP 401) rather than a rate-limit error. 747 748 Parameters 749 ---------- 750 key : str 751 The API key string. 752 error_code : str | None, optional 753 Machine-readable error code (e.g. ``"401"``, ``"invalid_api_key"``). 754 error_message : str | None, optional 755 Human-readable error description. 756 757 Examples 758 -------- 759 >>> pool = KeyPool(["key-a"]) 760 >>> pool.mark_invalid("key-a") 761 >>> pool.mark_invalid("key-a", error_code="401", error_message="Invalid API key") 762 """ 763 with self._lock: 764 rec = self._find_record(key) 765 if rec is None: 766 return 767 rec.state = KeyState.DISABLED 768 rec.last_error_code = error_code 769 rec.last_error_message = error_message
Permanently disable a key (DISABLED).
Disabled keys never auto-recover. Use this when a key returns an authentication error (e.g. HTTP 401) rather than a rate-limit error.
Parameters
key : str
The API key string.
error_code : str | None, optional
Machine-readable error code (e.g. "401", "invalid_api_key").
error_message : str | None, optional
Human-readable error description.
Examples
>>> pool = KeyPool(["key-a"])
>>> pool.mark_invalid("key-a")
>>> pool.mark_invalid("key-a", error_code="401", error_message="Invalid API key")
771 def mark_success(self, key: str) -> None: 772 """Reset a key's failure count to 0 and keep it ACTIVE. 773 774 Call this after a successful API response to indicate the key is 775 healthy and reset any transient failure tracking. 776 777 Parameters 778 ---------- 779 key : str 780 The API key string. 781 782 Examples 783 -------- 784 >>> pool = KeyPool(["key-a"]) 785 >>> k = pool.get_key() 786 >>> pool.mark_success(k) 787 """ 788 with self._lock: 789 rec = self._find_record(key) 790 if rec is None: 791 return 792 rec.failure_count = 0 793 rec.state = KeyState.ACTIVE 794 rec.last_status_code = None 795 rec.last_error_code = None 796 rec.last_error_message = None
Reset a key's failure count to 0 and keep it ACTIVE.
Call this after a successful API response to indicate the key is healthy and reset any transient failure tracking.
Parameters
key : str The API key string.
Examples
>>> pool = KeyPool(["key-a"])
>>> k = pool.get_key()
>>> pool.mark_success(k)
798 def status(self) -> dict[str, dict]: 799 """Return a snapshot of every key's state without exposing raw keys. 800 801 Every key value in the returned dictionary is passed through ``mask()`` 802 so the caller can safely log or print the result. 803 804 Returns 805 ------- 806 dict[str, dict] 807 A mapping of ``{masked_key: {"state": str, "failure_count": int, 808 "cooldown_remaining": float | None, "last_status_code": int | None, 809 "last_error_code": str | None, "last_error_message": str | None}}`` 810 for every key in the pool. 811 812 Examples 813 -------- 814 >>> pool = KeyPool(["sk-abcdef1234567890"]) 815 >>> pool.status() 816 {'sk-abc...7890': {'state': 'active', 'failure_count': 0, 'cooldown_remaining': None, 'last_status_code': None, 'last_error_code': None, 'last_error_message': None}} 817 """ 818 with self._lock: 819 result: dict[str, dict] = {} 820 now = time.monotonic() 821 for rec in self._records: 822 if rec.state == KeyState.COOLDOWN and rec.cooldown_until is not None: 823 remaining = max(0.0, rec.cooldown_until - now) 824 else: 825 remaining = None 826 result[mask(rec.key)] = { 827 "state": rec.state.value, 828 "failure_count": rec.failure_count, 829 "cooldown_remaining": round(remaining, 1) if remaining is not None else None, 830 "last_status_code": rec.last_status_code, 831 "last_error_code": rec.last_error_code, 832 "last_error_message": rec.last_error_message, 833 } 834 return result
Return a snapshot of every key's state without exposing raw keys.
Every key value in the returned dictionary is passed through mask()
so the caller can safely log or print the result.
Returns
dict[str, dict]
A mapping of {masked_key: {"state": str, "failure_count": int,
"cooldown_remaining": float | None, "last_status_code": int | None,
"last_error_code": str | None, "last_error_message": str | None}}
for every key in the pool.
Examples
>>> pool = KeyPool(["sk-abcdef1234567890"])
>>> pool.status()
{'sk-abc...7890': {'state': 'active', 'failure_count': 0, 'cooldown_remaining': None, 'last_status_code': None, 'last_error_code': None, 'last_error_message': None}}
50class AllKeysExhaustedError(Exception): 51 """Raised when no ACTIVE key is available in the pool. 52 53 The exception message includes the soonest recovery time in seconds if any 54 key is in COOLDOWN, otherwise it reports that all keys are disabled. 55 """ 56 57 pass
Raised when no ACTIVE key is available in the pool.
The exception message includes the soonest recovery time in seconds if any key is in COOLDOWN, otherwise it reports that all keys are disabled.
21class KeyState(Enum): 22 """Possible states for a key in the pool. 23 24 Attributes: 25 ACTIVE: The key is healthy and available for use. 26 COOLDOWN: The key is temporarily unavailable (rate-limited) and will 27 automatically recover after *cooldown_seconds* elapses. 28 DISABLED: The key is permanently unusable and will never auto-recover. 29 """ 30 31 ACTIVE = "active" 32 COOLDOWN = "cooldown" 33 DISABLED = "disabled"
Possible states for a key in the pool.
Attributes: ACTIVE: The key is healthy and available for use. COOLDOWN: The key is temporarily unavailable (rate-limited) and will automatically recover after cooldown_seconds elapses. DISABLED: The key is permanently unusable and will never auto-recover.