services/gemini_service/api_key_config.py

"""
API Key Service Configuration

Configures automatic API key selection and rotation via middleware.
"""
from typing import List, Optional
import os
import logging

logger = logging.getLogger(__name__)


class APIKeyServiceConfig:
    """Configuration for API key middleware."""
    
    _rotation_strategy: str = "least_used"  # or "round_robin"
    _cooldown_seconds: int = 60
    _max_requests_per_minute: int = 60
    _retry_on_quota_error: bool = True
    _api_keys: Optional[List[str]] = None
    
    @classmethod
    def register(
        cls,
        rotation_strategy: str = "least_used",
        cooldown_seconds: int = 60,
        max_requests_per_minute: int = 60,
        retry_on_quota_error: bool = True
    ) -> None:
        """
        Register API key service configuration.
        
        Args:
            rotation_strategy: "least_used" or "round_robin"
            cooldown_seconds: Time to wait before reusing a key after quota error
            max_requests_per_minute: Rate limit per key
            retry_on_quota_error: Auto-retry with different key on 429
        
        Example:
            APIKeyServiceConfig.register(
                rotation_strategy="least_used",
                cooldown_seconds=60,
                retry_on_quota_error=True
            )
        """
        cls._rotation_strategy = rotation_strategy
        cls._cooldown_seconds = cooldown_seconds
        cls._max_requests_per_minute = max_requests_per_minute
        cls._retry_on_quota_error = retry_on_quota_error
        
        # Load API keys from env
        cls._load_api_keys()
        
        logger.info(
            f"API Key Service configured: "
            f"keys={len(cls._api_keys or [])}, "
            f"strategy={rotation_strategy}, "
            f"retry={retry_on_quota_error}"
        )
    
    @classmethod
    def _load_api_keys(cls):
        """Load API keys from environment variables."""
        keys_str = os.getenv("GEMINI_API_KEYS", "")
        if not keys_str:
            # Fallback to single key
            single_key = os.getenv("GEMINI_API_KEY", "")
            if single_key:
                cls._api_keys = [single_key]
            else:
                cls._api_keys = []
                logger.warning("No Gemini API keys configured!")
        else:
            cls._api_keys = [k.strip() for k in keys_str.split(",") if k.strip()]
        
        if cls._api_keys:
            logger.info(f"Loaded {len(cls._api_keys)} Gemini API key(s)")
    
    @classmethod
    def get_api_keys(cls) -> List[str]:
        """Get loaded API keys."""
        if cls._api_keys is None:
            cls._load_api_keys()
        return cls._api_keys or []
    
    @classmethod
    def get_key_count(cls) -> int:
        """Get number of available keys."""
        return len(cls.get_api_keys())
    
    @classmethod
    def get_config(cls) -> dict:
        """Get current configuration."""
        return {
            "key_count": cls.get_key_count(),
            "rotation_strategy": cls._rotation_strategy,
            "cooldown_seconds": cls._cooldown_seconds,
            "max_requests_per_minute": cls._max_requests_per_minute,
            "retry_on_quota_error": cls._retry_on_quota_error
        }