Source code for lalandre_core.utils.api_key_pool

"""
API Key Pool Manager
Distributes API calls across multiple keys using round-robin strategy.
"""

import logging
import os
from threading import Lock
from typing import List

logger = logging.getLogger(__name__)




class APIKeyPool:
    """
    Thread-safe container for API keys with round-robin distribution.

    Keys are loaded from environment variables following the pattern:
        BASE_VAR, BASE_VAR_2, BASE_VAR_3, ..., BASE_VAR_{max_keys}
    """

    def __init__(self, keys: List[str]):
        if not keys:
            raise ValueError("At least one API key must be provided")
        self.keys = keys
        self._lock = Lock()
        self._index = 0
        logger.info("Initialized API key pool with %d key(s)", len(keys))



    @classmethod
    def from_env(
        cls,
        base_env_var: str = "MISTRAL_API_KEY",
        max_keys: int = 10,
        start_index: int = 1,
    ) -> "APIKeyPool":
        """
        Load keys from environment variables.

        Looks for:
        - {base_env_var}          (index 1, main key)
        - {base_env_var}_2, ..., {base_env_var}_{max_keys}

        *start_index* and *max_keys* control the range: indices
        [start_index .. max_keys] are loaded.  This allows splitting
        keys between services (e.g. 1-5 for RAG, 6-10 for workers).
        """
        keys: List[str] = []

        for i in range(start_index, max_keys + 1):
            env_name = base_env_var if i == 1 else f"{base_env_var}_{i}"
            key = os.environ.get(env_name, "").strip()
            if key:
                keys.append(key)

        if not keys:
            raise ValueError(f"No API keys found in environment (looking for {base_env_var})")

        logger.info(
            "Loaded %d API key(s) from environment (%s, indices %d-%d)",
            len(keys),
            base_env_var,
            start_index,
            max_keys,
        )
        return cls(keys)




    def next_key(self) -> str:
        """Return the next key in round-robin order (thread-safe)."""
        with self._lock:
            key = self.keys[self._index % len(self.keys)]
            self._index += 1
            return key


    def __len__(self) -> int:
        return len(self.keys)

    def __iter__(self):
        return iter(self.keys)

    def __repr__(self) -> str:
        return f"APIKeyPool({len(self.keys)} keys)"



__all__ = ["APIKeyPool"]