Class ApiKeyPool

contrib/multilingual/api_pool.py:131–389

Thread-safe pool of API keys with per-key concurrency slots. Each key has *max_concurrent* slots (default 5). :meth:`acquire` picks the least-loaded available key — multiple callers can share the same key as long as slots remain. Only rate-limited keys (HTTP 429) are taken out of

class ApiKeyPool:
    """Thread-safe pool of API keys with per-key concurrency slots.

    Each key has *max_concurrent* slots (default 5).  :meth:`acquire` picks
    the least-loaded available key — multiple callers can share the same key
    as long as slots remain.  Only rate-limited keys (HTTP 429) are taken
    out of rotation; the pool only blocks when every non-rate-limited key
    is at capacity.

    Usage::

        pool = ApiKeyPool([ApiKey("sk-a", ...), ApiKey("sk-b", ...)])
        key = pool.acquire()          # blocks only if all keys full
        try:
            llm_call(key)
            pool.release(key, success=True)
        except RateLimitError:
            pool.release(key, success=False)
            key = pool.acquire()
    """

    def __init__(self, keys: list[ApiKey]) -> None:
        if not keys:
            raise ValueError("ApiKeyPool requires at least one key")
        self._keys = list(keys)
        self._lock = threading.Lock()
        self._condition = threading.Condition(self._lock)
        self._rate_limits_hit: int = 0
        self._retry_successes: int = 0
        self._total_requests_served: int = 0
        self._peak_active_requests: int = 0

    # -- Public API -----------------------------------------------------------

    def acquire(self, timeout: float | None = None) -> ApiKey:
        """Acquire a slot on the least-loaded available key.

        Scheduling priority:

        1. **Recovered keys** — rate-limited keys whose backoff has expired
           become available again.
        2. **Least-loaded key** — among available keys, pick the one with
           the fewest ``active_requests``.
        3. **Block** — if every non-rate-limited key is at capacity, wait
           for a slot to free up or a rate-limited key to recover.

        Parameters
        ----------
        timeout :
            Maximum seconds to wait.  ``None`` means wait indefinitely.

        Returns
        -------
        ApiKey
            A key with at least one available slot.

        Raises
        ------