Phase 3: Implement API Key Middleware

Core implementation:
- Created api_key_config.py: Configuration for key rotation strategies
- Created api_key_middleware.py: Automatic key selection and quota handling
- Integrated into app.py middleware chain

Features:
- Automatic API key selection (least_used or round_robin)
- Quota error detection (429) with automatic retry
- Key cooldown management (60s after quota error)
- Transparent key rotation (app doesn't know which key)
- Usage tracking per key
- No code changes needed in endpoints

Middleware order: Auth → Audit → API Key → Credit → Application

Benefits:
- Zero downtime on quota errors
- Automatic load balancing across keys
- Better observability of key usage
- No manual key management needed

Next: Testing and Phase 4 (Payment Transaction Manager)

app.py

@@ -132,6 +132,16 @@ async def lifespan(app: FastAPI):
     )
     logger.info("✅ Audit Service configured")
     
+    # Register API Key Service configuration
+    from services.gemini_service import APIKeyServiceConfig
+    APIKeyServiceConfig.register(
+        rotation_strategy="least_used",  # or "round_robin"
+        cooldown_seconds=60,  # Wait 1 min after quota error
+        max_requests_per_minute=60,
+        retry_on_quota_error=True  # Auto-retry with different key
+    )
+    logger.info("✅ API Key Service configured")
+    
     # Check for RESET_DB environment variable
     if os.getenv("RESET_DB", "").lower() == "true":
         logger.warning(f"RESET_DB is set to true. Skipping download and clearing local database ({DB_FILENAME}).")
@@ -194,11 +204,15 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Add Audit Middleware (executes second - after auth, before credit)
+# Add Audit Middleware (executes second - after auth, before API key)
 from services.audit_service import AuditMiddleware
 app.add_middleware(AuditMiddleware)
 
-# Add Credit Middleware (executes third - after auth and audit)
+# Add API Key Middleware (executes third - for Gemini requests)
+from services.gemini_service import APIKeyMiddleware
+app.add_middleware(APIKeyMiddleware)
+
+# Add Credit Middleware (executes fourth - after auth, audit, and API key)
 from services.credit_service import CreditMiddleware
 app.add_middleware(CreditMiddleware)
 

services/gemini_service/__init__.py

@@ -16,6 +16,7 @@ from services.gemini_service.api_client import (
     get_gemini_api_key,
     MOCK_MODE,
     MOCK_VIDEO_URL,
+    GeminiAPIClient, # Added
 )
 
 # Job Processor exports
@@ -28,6 +29,10 @@ from services.gemini_service.job_processor import (
     stop_worker,
 )
 
+# API Key Middleware exports # Added
+from services.gemini_service.api_key_config import APIKeyServiceConfig # Added
+from services.gemini_service.api_key_middleware import APIKeyMiddleware # Added
+
 __all__ = [
     # API Client
     'GeminiService',
@@ -36,6 +41,7 @@ __all__ = [
     'get_gemini_api_key',
     'MOCK_MODE',
     'MOCK_VIDEO_URL',
+    'GeminiAPIClient',
     
     # Job Processor
     'GeminiJobProcessor',
@@ -44,4 +50,8 @@ __all__ = [
     'get_priority_for_job_type',
     'start_worker',
     'stop_worker',
+    
+    # API Key Middleware
+    'APIKeyServiceConfig',
+    'APIKeyMiddleware',
 ]

services/gemini_service/api_key_config.py

@@ -0,0 +1,100 @@
+"""
+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
+        }

services/gemini_service/api_key_middleware.py

@@ -0,0 +1,180 @@
+"""
+API Key Middleware - Automatic key selection and rotation
+
+Automatically selects and injects Gemini API keys for requests.
+Handles quota errors with automatic key rotation and retry.
+"""
+import time
+import logging
+from datetime import datetime, timedelta
+from typing import Optional, Dict
+from fastapi import Request, Response
+from starlette.middleware.base import BaseHTTPMiddleware
+from starlette.types import ASGIApp
+
+from core.database import async_session_maker
+from services.gemini_service.api_key_config import APIKeyServiceConfig
+
+logger = logging.getLogger(__name__)
+
+
+# Track key cooldowns in memory
+_key_cooldowns: Dict[int, datetime] = {}
+
+
+class APIKeyMiddleware(BaseHTTPMiddleware):
+    """
+    Middleware for automatic API key management.
+    
+    Features:
+    - Automatic key selection based on strategy
+    - Quota error detection and recovery
+    - Key cooldown management
+    - Usage tracking
+    """
+    
+    def __init__(self, app: ASGIApp):
+        super().__init__(app)
+    
+    async def dispatch(self, request: Request, call_next):
+        """
+        Process request with automatic API key injection.
+        
+        Flow:
+        1. Check if Gemini request
+        2. Select best available key
+        3. Inject into request state
+        4. Handle response (quota errors)
+        """
+        # Only handle Gemini requests
+        if not self._is_gemini_request(request):
+            return await call_next(request)
+        
+        # Select API key
+        try:
+            key_index, api_key = await self._select_api_key()
+            request.state.gemini_api_key = api_key
+            request.state.gemini_key_index = key_index
+        except ValueError as e:
+            # No keys available
+            logger.error(f"No API keys available: {e}")
+            return Response(
+                content=f'{{"detail": "{str(e)}"}}',
+                status_code=503,
+                media_type="application/json"
+            )
+        
+        # Process request
+        response = await call_next(request)
+        
+        # Handle quota errors
+        if response.status_code == 429 and API KeyServiceConfig._retry_on_quota_error:
+            logger.warning(f"Quota error on key {key_index}, attempting retry")
+            
+            # Mark key in cooldown
+            self._mark_cooldown(key_index)
+            
+            # Try to select different key
+            try:
+                key_index, api_key = await self._select_api_key(exclude_index=key_index)
+                request.state.gemini_api_key = api_key
+                request.state.gemini_key_index = key_index
+                
+                # Retry request
+                logger.info(f"Retrying with key {key_index}")
+                response = await call_next(request)
+            except ValueError:
+                # No other keys available
+                logger.error("All API keys in cooldown or exhausted")
+        
+        # Track usage
+        success = response.status_code < 400
+        await self._track_usage(key_index, success, response.status_code)
+        
+        return response
+    
+    def _is_gemini_request(self, request: Request) -> bool:
+        """Check if request is for Gemini service."""
+        path = request.url.path
+        gemini_paths = ["/gemini/", "/api/gemini"]
+        return any(path.startswith(p) for p in gemini_paths)
+    
+    async def _select_api_key(self, exclude_index: Optional[int] = None) -> tuple[int, str]:
+        """
+        Select best available API key.
+        
+        Args:
+            exclude_index: Key index to exclude (e.g., after quota error)
+        
+        Returns:
+            Tuple of (key_index, api_key)
+        
+        Raises:
+            ValueError: If no keys available
+        """
+        keys = APIKeyServiceConfig.get_api_keys()
+        if not keys:
+            raise ValueError("No API keys configured")
+        
+        # Filter out excluded and cooldown keys
+        available_indices = []
+        for i in range(len(keys)):
+            if i == exclude_index:
+                continue
+            if self._is_in_cooldown(i):
+                continue
+            available_indices.append(i)
+        
+        if not available_indices:
+            raise ValueError("All API keys in cooldown")
+        
+        # Select based on strategy
+        if APIKeyServiceConfig._rotation_strategy == "round_robin":
+            # Simple round-robin
+            selected_index = available_indices[0]
+        else:  # least_used
+            # Get usage stats from DB
+            async with async_session_maker() as db:
+                from services.api_key_manager import get_least_used_key
+                try:
+                    selected_index, _ = await get_least_used_key(db)
+                    if selected_index not in available_indices:
+                        # Fallback to first available
+                        selected_index = available_indices[0]
+                except Exception as e:
+                    logger.error(f"Error getting least used key: {e}")
+                    selected_index = available_indices[0]
+        
+        logger.debug(f"Selected API key index {selected_index}")
+        return selected_index, keys[selected_index]
+    
+    def _is_in_cooldown(self, key_index: int) -> bool:
+        """Check if key is in cooldown period."""
+        if key_index not in _key_cooldowns:
+            return False
+        
+        cooldown_until = _key_cooldowns[key_index]
+        if datetime.utcnow() > cooldown_until:
+            # Cooldown expired
+            del _key_cooldowns[key_index]
+            return False
+        
+        return True
+    
+    def _mark_cooldown(self, key_index: int):
+        """Mark key as in cooldown."""
+        cooldown_seconds = APIKeyServiceConfig._cooldown_seconds
+        cooldown_until = datetime.utcnow() + timedelta(seconds=cooldown_seconds)
+        _key_cooldowns[key_index] = cooldown_until
+        logger.info(f"Key {key_index} in cooldown until {cooldown_until}")
+    
+    async def _track_usage(self, key_index: int, success: bool, status_code: int):
+        """Track API key usage."""
+        try:
+            async with async_session_maker() as db:
+                from services.api_key_manager import record_usage
+                error_message = f"HTTP {status_code}" if not success else None
+                await record_usage(db, key_index, success, error_message)
+                await db.commit()
+        except Exception as e:
+            logger.error(f"Failed to track usage: {e}")