nvidia--skillspector/contrib/batch_scan/api_pool.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""API Key Pool — multi-key load-balancer with per-key concurrency slots.

Each key has a configurable number of concurrent slots (default 5).  The pool
distributes requests across keys using least-loaded scheduling — it *never*
blocks unless every non-rate-limited key is at capacity.  A single key can
serve multiple callers simultaneously; rate-limit (HTTP 429) is the only
signal that removes a key from rotation.

Contrast with the previous mutex-per-key design where :meth:`acquire` blocked
as soon as every key had *one* active request, coupling worker count to key
count.  In the new design, throughput scales with workers independently of
how many keys are configured — keys just need enough aggregate slots.

Integration point
-----------------
Wrap a LangChain ``BaseChatModel`` with :class:`PooledChatModel` to give
it transparent access to the key pool.  The wrapper is API-compatible with
the models returned by :func:`skillspector.llm_utils.get_chat_model` and
can be used wherever a standard ``BaseChatModel`` is expected.

Configuration
-------------
Multi-key mode (recommended for batch scans)::

    export SKILLSPECTOR_API_KEYS="
      sk-or-xxx1|https://api.openai.com/v1|gpt-5.4
      sk-or-xxx2|https://api.openai.com/v1|gpt-5.4
    "

Single-key mode (backward-compatible — no pool needed)::

    export OPENAI_API_KEY=sk-or-xxx1

When ``SKILLSPECTOR_API_KEYS`` is not set, :func:`create_api_key_pool_from_env`
returns ``None`` and the caller should fall back to the single-key provider path.
"""

from __future__ import annotations

import os
import threading
import time
from dataclasses import dataclass

from skillspector.logging_config import get_logger

logger = get_logger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_API_KEYS_ENV = "SKILLSPECTOR_API_KEYS"
_DEFAULT_MAX_CONCURRENT_PER_KEY = 5
_MAX_RATE_LIMIT_RETRIES = 5
_BACKOFF_BASE_S = 30.0
_BACKOFF_CAP_S = 300.0


# ---------------------------------------------------------------------------
# ApiKey — single key tracked by the pool
# ---------------------------------------------------------------------------


@dataclass
class ApiKey:
    """A single API key with concurrency and rate-limit metadata.

    Attributes
    ----------
    key :
        API key string (e.g. ``"sk-or-xxx"``).
    base_url :
        Optional base URL override for the provider endpoint.
    model :
        Model label to use with this key.
    rate_limited :
        ``True`` when this key is cooling down after a 429 response.
    rate_limited_until :
        Monotonic timestamp when this key becomes eligible again after a
        429.  Only meaningful when *rate_limited* is ``True``.
    consecutive_429 :
        Count of consecutive rate-limit hits.  Used to compute the next
        backoff duration via :math:`30 \\times 2^n` seconds, capped at 300.
    total_requests :
        Cumulative request count served by this key.  Used for
        least-loaded scheduling.
    active_requests :
        Number of callers currently using this key.
    max_concurrent :
        Maximum number of simultaneous callers allowed on this key
        (default 5).  One key serves up to this many concurrent LLM calls.
    """

    key: str
    base_url: str | None
    model: str
    rate_limited: bool = False
    rate_limited_until: float = 0.0
    consecutive_429: int = 0
    total_requests: int = 0
    active_requests: int = 0
    max_concurrent: int = _DEFAULT_MAX_CONCURRENT_PER_KEY

    @property
    def available(self) -> bool:
        """``True`` when this key can accept at least one more caller."""
        return not self.rate_limited and self.active_requests < self.max_concurrent


# ---------------------------------------------------------------------------
# ApiKeyPool — multi-key load-balancer
# ---------------------------------------------------------------------------


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
        ------
        RuntimeError
            If *timeout* expires before a slot becomes available.
        """
        deadline = time.monotonic() + timeout if timeout is not None else None

        with self._condition:
            while True:
                now = time.monotonic()

                # Step 1: recover rate-limited keys whose backoff has expired
                self._recover_expired_keys(now)

                # Step 2: find available keys (not rate-limited, slots open)
                available = [k for k in self._keys if k.available]
                if available:
                    key = min(available, key=lambda k: k.active_requests)
                    key.active_requests += 1
                    key.total_requests += 1
                    self._total_requests_served += 1
                    _now_active = sum(k.active_requests for k in self._keys)
                    if _now_active > self._peak_active_requests:
                        self._peak_active_requests = _now_active
                    logger.debug(
                        "Pool: slot on key …%s (%d/%d active)",
                        key.key[-8:],
                        key.active_requests,
                        key.max_concurrent,
                    )
                    return key

                # Step 3: no capacity — compute wait time
                wait_for = self._next_available_in(now)
                remaining = self._remaining_timeout(deadline)
                if remaining is not None and remaining <= 0:
                    raise RuntimeError(
                        "ApiKeyPool: timed out waiting for available slot "
                        f"({self._capacity_summary()})"
                    )

                if wait_for is None:
                    self._condition.wait(timeout=remaining)
                else:
                    wait = min(wait_for, remaining or wait_for)
                    logger.debug(
                        "Pool: at capacity, waiting %.1fs (%s)",
                        wait,
                        self._capacity_summary(),
                    )
                    self._condition.wait(timeout=wait)

    def try_acquire(self) -> ApiKey | None:
        """Non-blocking acquire — returns a key immediately or ``None``.

        Unlike :meth:`acquire`, this never blocks.  If a slot is available
        right now, return the least-loaded key; otherwise return ``None``.
        Useful in async contexts where blocking would stall the event loop.
        """
        with self._lock:
            self._recover_expired_keys(time.monotonic())
            available = [k for k in self._keys if k.available]
            if not available:
                return None
            key = min(available, key=lambda k: k.active_requests)
            key.active_requests += 1
            key.total_requests += 1
            self._total_requests_served += 1
            _now_active = sum(k.active_requests for k in self._keys)
            if _now_active > self._peak_active_requests:
                self._peak_active_requests = _now_active
            return key

    def release(self, key: ApiKey, *, success: bool = True) -> None:
        """Release a slot on *key* back to the pool.

        Parameters
        ----------
        key :
            The key previously obtained from :meth:`acquire`.
        success :
            ``True`` if the API call succeeded; ``False`` if it failed with
            a rate-limit error (HTTP 429).  On failure the key is marked
            rate-limited with exponential backoff.
        """
        with self._condition:
            key.active_requests = max(0, key.active_requests - 1)

            if success:
                key.consecutive_429 = 0
                logger.debug(
                    "Pool: released slot on key …%s (%d/%d active)",
                    key.key[-8:],
                    key.active_requests,
                    key.max_concurrent,
                )
            else:
                key.consecutive_429 += 1
                backoff = min(
                    _BACKOFF_BASE_S * (2 ** (key.consecutive_429 - 1)),
                    _BACKOFF_CAP_S,
                )
                key.rate_limited_until = time.monotonic() + backoff
                key.rate_limited = True
                self._rate_limits_hit += 1
                logger.warning(
                    "Pool: key …%s rate-limited for %.0fs "
                    "(consecutive=%d)",
                    key.key[-8:],
                    backoff,
                    key.consecutive_429,
                )

            self._condition.notify_all()

    def record_retry_success(self) -> None:
        """Increment the retry-success counter for reporting.

        Only call this when a retry (after a key switch due to 429)
        actually succeeds, not on every attempt.
        """
        with self._lock:
            self._retry_successes += 1

    @property
    def rate_limits_hit(self) -> int:
        """Total number of 429 responses encountered across all keys."""
        with self._lock:
            return self._rate_limits_hit

    @property
    def retry_successes(self) -> int:
        """Total number of successful retries after a key switch."""
        with self._lock:
            return self._retry_successes

    @property
    def keys_configured(self) -> int:
        """Total number of keys in the pool."""
        return len(self._keys)

    @property
    def total_capacity(self) -> int:
        """Sum of ``max_concurrent`` across all keys."""
        return sum(k.max_concurrent for k in self._keys)

    @property
    def active_requests(self) -> int:
        """Total active requests across all keys."""
        with self._lock:
            return sum(k.active_requests for k in self._keys)

    def snapshot(self) -> dict[str, object]:
        """Return a snapshot dict suitable for report metadata."""
        with self._lock:
            rate_limited = sum(1 for k in self._keys if k.rate_limited)
            active = sum(k.active_requests for k in self._keys)
            return {
                "keys_configured": len(self._keys),
                "total_capacity": sum(k.max_concurrent for k in self._keys),
                "active_requests": active,
                "peak_active_requests": self._peak_active_requests,
                "total_requests_served": self._total_requests_served,
                "keys_rate_limited": rate_limited,
                "keys_available": len(self._keys) - rate_limited,
                "rate_limits_hit": self._rate_limits_hit,
                "retry_successes": self._retry_successes,
            }

    # -- Internal -------------------------------------------------------------

    def _recover_expired_keys(self, now: float) -> None:
        """Promote rate-limited keys whose backoff has expired."""
        for k in self._keys:
            if k.rate_limited and now >= k.rate_limited_until:
                k.rate_limited = False
                k.consecutive_429 = 0
                logger.info(
                    "Pool: key …%s recovered (backoff expired)", k.key[-8:]
                )

    def _next_available_in(self, now: float) -> float | None:
        """Seconds until the earliest rate-limited key recovers, or ``None``."""
        rate_limited = [k for k in self._keys if k.rate_limited]
        if not rate_limited:
            return None
        earliest = min(k.rate_limited_until for k in rate_limited)
        return max(0.0, earliest - now)

    def _capacity_summary(self) -> str:
        active = sum(k.active_requests for k in self._keys)
        total = sum(k.max_concurrent for k in self._keys)
        rate_limited = sum(1 for k in self._keys if k.rate_limited)
        return (
            f"{active}/{total} slots active, "
            f"{rate_limited} key(s) rate-limited"
        )

    @staticmethod
    def _remaining_timeout(deadline: float | None) -> float | None:
        if deadline is None:
            return None
        return max(0.0, deadline - time.monotonic())


# ---------------------------------------------------------------------------
# PooledChatModel — transparent key-switching wrapper
# ---------------------------------------------------------------------------


class PooledChatModel:
    """LangChain-compatible chat model wrapper with transparent key switching.

    Each :meth:`invoke` / :meth:`ainvoke` call acquires a key from the pool,
    builds a :class:`~langchain_openai.ChatOpenAI` instance on the fly, and
    releases the key when done.  On rate-limit errors the wrapper releases
    the key with ``success=False``, picks a different key, and retries.

    Parameters
    ----------
    pool :
        An :class:`ApiKeyPool` with at least one configured key.
    max_tokens :
        ``max_completion_tokens`` passed to each ``ChatOpenAI`` instance.
    timeout :
        Request timeout in seconds passed to each ``ChatOpenAI`` instance.
    max_retries :
        Maximum number of key-switch retries on rate-limit errors before
        giving up.
    """

    def __init__(
        self,
        pool: ApiKeyPool,
        *,
        max_tokens: int = 4096,
        timeout: float = 30.0,
        max_retries: int = _MAX_RATE_LIMIT_RETRIES,
    ) -> None:
        self._pool = pool
        self._max_tokens = max_tokens
        self._timeout = timeout
        self._max_retries = max_retries

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

    def invoke(self, prompt: str) -> object:
        """Synchronous invoke with automatic key switching on rate-limit."""
        return self._invoke_with_retry(prompt)

    async def ainvoke(self, prompt: str) -> object:
        """Async invoke with automatic key switching on rate-limit."""
        return await self._ainvoke_with_retry(prompt)

    # -- Internal -------------------------------------------------------------

    def _invoke_with_retry(self, prompt: str) -> object:
        """Sync retry loop — acquire slot, call LLM, release, retry on 429."""
        last_exception: Exception | None = None

        for attempt in range(self._max_retries + 1):
            key = self._pool.acquire()
            llm = self._build_llm(key)
            try:
                result = llm.invoke(prompt)
                self._pool.release(key, success=True)
                if attempt > 0:
                    self._pool.record_retry_success()
                return result
            except Exception as exc:
                if self._is_rate_limit(exc) and attempt < self._max_retries:
                    self._pool.release(key, success=False)
                    logger.debug(
                        "PooledChatModel: rate-limited, retrying "
                        "(attempt %d/%d)",
                        attempt + 1,
                        self._max_retries,
                    )
                    continue
                self._pool.release(key, success=True)
                last_exception = exc
                raise

        raise RuntimeError(
            f"PooledChatModel: exhausted {self._max_retries} retries "
            "due to rate-limit errors"
        ) from last_exception

    async def _ainvoke_with_retry(self, prompt: str) -> object:
        """Async retry loop — non-blocking acquire first, block only if full."""
        import asyncio
        last_exception: Exception | None = None

        for attempt in range(self._max_retries + 1):
            key = self._pool.try_acquire()
            if key is None:
                key = await asyncio.to_thread(self._pool.acquire)
            llm = self._build_llm(key)
            try:
                result = await llm.ainvoke(prompt)
                self._pool.release(key, success=True)
                if attempt > 0:
                    self._pool.record_retry_success()
                return result
            except Exception as exc:
                if self._is_rate_limit(exc) and attempt < self._max_retries:
                    self._pool.release(key, success=False)
                    logger.debug(
                        "PooledChatModel: rate-limited, retrying "
                        "(attempt %d/%d)",
                        attempt + 1,
                        self._max_retries,
                    )
                    continue
                self._pool.release(key, success=True)
                last_exception = exc
                raise

        raise RuntimeError(
            f"PooledChatModel: exhausted {self._max_retries} retries "
            "due to rate-limit errors"
        ) from last_exception

    def _build_llm(self, key: ApiKey):
        """Build a fresh :class:`~langchain_openai.ChatOpenAI` for *key*."""
        from langchain_openai import ChatOpenAI
        from pydantic import SecretStr

        try:
            import httpx
            _timeout = httpx.Timeout(self._timeout, connect=8.0)
        except ImportError:
            _timeout = self._timeout

        return ChatOpenAI(
            model=key.model,
            base_url=key.base_url,
            api_key=SecretStr(key.key),
            max_completion_tokens=self._max_tokens,
            timeout=_timeout,
        )

    @staticmethod
    def _is_rate_limit(exc: Exception) -> bool:
        """Detect rate-limit errors from common LLM provider SDKs."""
        try:
            import openai
            if isinstance(exc, openai.RateLimitError):
                return True
        except ImportError:
            pass

        message = str(exc).lower()
        for marker in ("429", "rate limit", "rate_limit", "too many requests"):
            if marker in message:
                return True

        return False


# ---------------------------------------------------------------------------
# Factory — create pool from environment
# ---------------------------------------------------------------------------


def create_api_key_pool_from_env(
    max_concurrent_per_key: int = _DEFAULT_MAX_CONCURRENT_PER_KEY,
) -> ApiKeyPool | None:
    """Build an :class:`ApiKeyPool` from environment variables.

    Reads ``SKILLSPECTOR_API_KEYS`` — a newline- or semicolon-delimited list
    of ``key|base_url|model`` entries.

    Also supports a fallback format where multiple keys are specified via
    sequentially numbered env vars ``OPENAI_API_KEY``, ``OPENAI_API_KEY_2``,
    etc.

    Parameters
    ----------
    max_concurrent_per_key :
        Maximum simultaneous requests allowed per key (default 5).
        With 10 keys this gives 50 aggregate slots.

    Returns
    -------
    ApiKeyPool or None
        ``None`` when no multi-key configuration is detected.
    """
    keys: list[ApiKey] = []

    raw = os.environ.get(_API_KEYS_ENV, "").strip()
    if raw:
        for line in raw.replace(";", "\n").splitlines():
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            parts = line.split("|")
            if len(parts) < 1:
                continue
            key_str = parts[0].strip()
            base_url = parts[1].strip() if len(parts) > 1 else None
            model = parts[2].strip() if len(parts) > 2 else "gpt-5.4"
            keys.append(ApiKey(
                key=key_str, base_url=base_url, model=model,
                max_concurrent=max_concurrent_per_key,
            ))

    if not keys:
        base = os.environ.get("OPENAI_API_KEY", "").strip()
        base_url = os.environ.get("OPENAI_BASE_URL", None)
        if base:
            keys.append(ApiKey(
                key=base, base_url=base_url, model="gpt-5.4",
                max_concurrent=max_concurrent_per_key,
            ))
        for idx in range(2, 10):
            extra = os.environ.get(f"OPENAI_API_KEY_{idx}", "").strip()
            if not extra:
                break
            keys.append(ApiKey(
                key=extra, base_url=base_url, model="gpt-5.4",
                max_concurrent=max_concurrent_per_key,
            ))

    if len(keys) <= 1:
        return None

    total_cap = len(keys) * max_concurrent_per_key
    logger.info(
        "ApiKeyPool: %d keys × %d slots = %d total capacity",
        len(keys), max_concurrent_per_key, total_cap,
    )
    return ApiKeyPool(keys)