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"]
class KeyPool:
 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())
KeyPool( keys: list[str] | None = None, max_retries: int = 3, cooldown_seconds: int = 60, strategy: str = 'round_robin')
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()
max_retries: int
cooldown_seconds: int
strategy: str
@classmethod
def from_doppler( cls, token: str, project: str, config: str, key_prefix: str | None = None, max_retries: int = 3, cooldown_seconds: int = 60, strategy: str = 'round_robin', force_refresh: bool = False) -> KeyPool:
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_",
... )
@classmethod
def from_env( cls, suffix: str, env_file: str | None = None, max_retries: int = 3, cooldown_seconds: int = 60, strategy: str = 'round_robin') -> KeyPool:
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")
@classmethod
def from_json( cls, path: str, suffix: str | None = None, max_retries: int = 3, cooldown_seconds: int = 60, strategy: str = 'round_robin') -> KeyPool:
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")
def add_key(self, key: str) -> None:
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")
def remove_key(self, key: str) -> None:
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")
def get_key(self) -> str:
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)
def handle_response( self, key: str, status_code: int, headers: dict[str, str] | None = None, body: dict | str | None = None) -> KeyState:
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 using Retry-After if present, otherwise the pool's cooldown_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()
def mark_rate_limited( self, key: str, retry_after: float | None = None, error_code: str | None = None, error_message: str | None = None) -> None:
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")
def mark_invalid( self, key: str, error_code: str | None = None, error_message: str | None = None) -> None:
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")
def mark_success(self, key: str) -> None:
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)
def status(self) -> dict[str, dict]:
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}}
class AllKeysExhaustedError(builtins.Exception):
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.

class KeyState(enum.Enum):
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.

ACTIVE = <KeyState.ACTIVE: 'active'>
COOLDOWN = <KeyState.COOLDOWN: 'cooldown'>
DISABLED = <KeyState.DISABLED: 'disabled'>