feat: Implement robust API key rotation, retry, and logging

This commit introduces significant enhancements to the API key management system, focusing on improved reliability, intelligent key rotation, and comprehensive failure logging.

Key changes include:

- **Enhanced Key Rotation Logic:**
- `src/rotator_library/client.py`: Refactored `acompletion` to use a more robust key selection mechanism, including a 5-second wait when all keys are on cooldown, and removed the `excluded_keys` parameter for `get_next_smart_key`.
- `src/rotator_library/usage_manager.py`: Implemented daily usage resets and global statistics accumulation. The `get_next_smart_key` now intelligently filters out keys on cooldown and selects the least-used active key.

- **Comprehensive Failure Logging:**
- **NEW `src/rotator_library/failure_logger.py`:** Introduced a dedicated module for structured logging of failed API calls, capturing details like API key ending, model, attempt number, error type, message, raw response, and request data.
- `src/rotator_library/client.py`: Integrated `log_failure` to centralize error reporting for all API call exceptions.

- **Improved Retry Mechanism:**
- `src/rotator_library/client.py`: Simplified the retry loop, ensuring that retriable server errors lead to a retry with the same key, while other permanent errors (authentication, rate limit) trigger immediate key rotation.

- **Refined Usage Tracking:**
- `src/rotator_library/usage_manager.py`: Updated `record_success` to correctly apply usage to daily and global statistics. `record_rotation_error` now sets a cooldown period for failed keys, optionally parsing `retry_delay` from error messages.

- **Minor Improvements:**
- `src/proxy_app/main.py`: Streamlined Gemini API key loading to correctly handle both `GEMINI_API_KEY` and `GEMINI_API_KEY_n` formats. Removed redundant comments.

src/proxy_app/main.py

@@ -7,6 +7,9 @@ import logging
 from pathlib import Path
 import sys
 
+# This is necessary for the app to find the rotator_library module
+sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
+
 from src.rotator_library.client import RotatingClient
 
 # Configure logging
@@ -24,22 +27,20 @@ if not PROXY_API_KEY:
 gemini_keys = []
 i = 1
 while True:
+    # Start with GEMINI_API_KEY_1, then GEMINI_API_KEY_2, etc.
     key = os.getenv(f"GEMINI_API_KEY_{i}")
+    if not key and i == 1:
+        # Fallback for a single key named just GEMINI_API_KEY
+        key = os.getenv("GEMINI_API_KEY")
+    
     if key:
         gemini_keys.append(key)
         i += 1
     else:
-        # Also check for the key without a number for the first one
-        if i == 1:
-            key = os.getenv("GEMINI_API_KEY")
-            if key:
-                gemini_keys.append(key)
-                i += 1
-                continue
         break
 
 if not gemini_keys:
-    raise ValueError("No GEMINI_API_KEY environment variables found.")
+    raise ValueError("No GEMINI_API_KEY or GEMINI_API_KEY_n environment variables found.")
 
 # Initialize the rotating client
 rotating_client = RotatingClient(api_keys=gemini_keys)
@@ -67,11 +68,8 @@ async def chat_completions(request: Request, _=Depends(verify_api_key)):
         response = await rotating_client.acompletion(**data)
 
         if is_streaming:
-            # For streaming responses, we return a StreamingResponse.
-            # The client's wrapper ensures usage is logged upon completion.
             return StreamingResponse(response, media_type="text/event-stream")
         else:
-            # For non-streaming, the response is a regular JSON object.
             return response
 
     except Exception as e:

src/rotator_library/client.py

@@ -5,6 +5,7 @@ import logging
 from typing import List, Dict, Any, AsyncGenerator
 
 from src.rotator_library.usage_manager import UsageManager
+from src.rotator_library.failure_logger import log_failure
 from src.rotator_library.error_handler import (
     is_authentication_error,
     is_rate_limit_error,
@@ -54,22 +55,22 @@ class RotatingClient:
         Performs a completion call with smart key rotation and retry logic.
         Handles both streaming and non-streaming requests.
         """
-        failed_keys_for_this_request = []
         model = kwargs.get("model")
         is_streaming = kwargs.get("stream", False)
 
         if not model:
             raise ValueError("'model' is a required parameter.")
 
-        while len(failed_keys_for_this_request) < len(self.api_keys):
+        while True: # Loop until a key succeeds or we decide to give up
             current_key = self.usage_manager.get_next_smart_key(
                 available_keys=self.api_keys,
-                model=model,
-                excluded_keys=failed_keys_for_this_request
+                model=model
             )
 
             if not current_key:
-                raise Exception("All available API keys have failed for this request.")
+                print("All keys are currently on cooldown. Waiting...")
+                await asyncio.sleep(5) # Wait 5 seconds before checking for an available key again
+                continue
 
             for attempt in range(self.max_retries):
                 try:
@@ -85,27 +86,25 @@ class RotatingClient:
                         return response
 
                 except Exception as e:
-                    error_message = str(e)
-                    print(f"Key ...{current_key[-4:]} failed with error: {error_message}")
-
-                    if is_authentication_error(e) or is_rate_limit_error(e):
-                        self.usage_manager.record_rotation_error(current_key, model, error_message)
-                        failed_keys_for_this_request.append(current_key)
-                        break
-
-                    elif is_server_error(e):
-                        if attempt == self.max_retries - 1:
-                            self.usage_manager.record_rotation_error(current_key, model, f"Failed after max retries with error: {error_message}")
-                            failed_keys_for_this_request.append(current_key)
-                            break
-                        else:
-                            await asyncio.sleep(1 * (attempt + 1))
-                            continue
-
-                    elif is_unrecoverable_error(e):
-                        raise e
-
-                    else:
+                    log_failure(
+                        api_key=current_key,
+                        model=model,
+                        attempt=attempt + 1,
+                        error=e,
+                        request_data=kwargs
+                    )
+                    
+                    # For any retriable server error, we just continue the attempt loop
+                    if is_server_error(e) and attempt < self.max_retries - 1:
+                        print(f"Key ...{current_key[-4:]} failed with server error. Retrying...")
+                        await asyncio.sleep(1 * (attempt + 1))
+                        continue
+                    
+                    # For unrecoverable errors, fail fast
+                    if is_unrecoverable_error(e):
                         raise e
 
-        raise Exception("All API keys failed after multiple retries.")
+                    # For all other errors (Auth, RateLimit, or final Server error), record it and break to get a new key
+                    print(f"Key ...{current_key[-4:]} failed permanently. Rotating...")
+                    self.usage_manager.record_rotation_error(current_key, model, e)
+                    break

src/rotator_library/failure_logger.py

@@ -0,0 +1,62 @@
+import logging
+import json
+from logging.handlers import RotatingFileHandler
+import os
+
+def setup_failure_logger():
+    """Sets up a dedicated JSON logger for failed API calls."""
+    log_dir = "logs"
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+
+    logger = logging.getLogger('failure_logger')
+    logger.setLevel(logging.ERROR)
+    
+    # Prevent logs from propagating to the root logger
+    logger.propagate = False
+
+    # Use a rotating file handler to keep log files from growing too large
+    handler = RotatingFileHandler(
+        os.path.join(log_dir, 'failures.log'),
+        maxBytes=5*1024*1024,  # 5 MB
+        backupCount=2
+    )
+
+    # Custom JSON formatter
+    class JsonFormatter(logging.Formatter):
+        def format(self, record):
+            log_record = {
+                "timestamp": self.formatTime(record, self.datefmt),
+                "level": record.levelname,
+                "message": record.getMessage()
+            }
+            return json.dumps(log_record)
+
+    handler.setFormatter(JsonFormatter())
+    
+    # Add handler only if it hasn't been added before
+    if not logger.handlers:
+        logger.addHandler(handler)
+
+    return logger
+
+failure_logger = setup_failure_logger()
+
+def log_failure(api_key: str, model: str, attempt: int, error: Exception, request_data: dict):
+    """Logs a structured message for a failed API call."""
+    
+    # Try to get the raw response from the exception if it exists
+    raw_response = None
+    if hasattr(error, 'response') and hasattr(error.response, 'text'):
+        raw_response = error.response.text
+
+    log_data = {
+        "api_key_ending": api_key[-4:],
+        "model": model,
+        "attempt_number": attempt,
+        "error_type": type(error).__name__,
+        "error_message": str(error),
+        "raw_response": raw_response,
+        "request_data": request_data,
+    }
+    failure_logger.error(log_data)

src/rotator_library/usage_manager.py

@@ -1,20 +1,21 @@
 import json
 import os
 import time
-from typing import Dict, List, Optional
+from datetime import date, datetime
+from typing import Dict, List, Optional, Any
 from filelock import FileLock
 
 class UsageManager:
     """
-    Manages detailed usage and failure data for API keys, stored in a JSON file.
+    Manages daily and global usage statistics and cooldowns for API keys.
     """
     def __init__(self, file_path: str = "key_usage.json"):
         self.file_path = file_path
         self.lock = FileLock(f"{self.file_path}.lock")
         self.usage_data = self._load_usage()
+        self._reset_daily_stats_if_needed()
 
     def _load_usage(self) -> Dict:
-        """Loads usage data from the JSON file."""
         with self.lock:
             if not os.path.exists(self.file_path):
                 return {}
@@ -25,63 +26,99 @@ class UsageManager:
                 return {}
 
     def _save_usage(self):
-        """Saves the current usage data to the JSON file."""
         with self.lock:
             with open(self.file_path, 'w') as f:
                 json.dump(self.usage_data, f, indent=2)
 
-    def get_next_smart_key(self, available_keys: List[str], model: str, excluded_keys: List[str]) -> Optional[str]:
+    def _reset_daily_stats_if_needed(self):
+        """Checks if daily stats need to be reset for any key."""
+        today_str = date.today().isoformat()
+        needs_saving = False
+        for key, data in self.usage_data.items():
+            daily_data = data.get("daily", {})
+            last_date_str = daily_data.get("date")
+            if last_date_str != today_str:
+                needs_saving = True
+                # Add yesterday's daily stats to global stats
+                global_data = data.setdefault("global", {"models": {}})
+                for model, stats in daily_data.get("models", {}).items():
+                    global_model_stats = global_data["models"].setdefault(model, {"success_count": 0, "prompt_tokens": 0, "completion_tokens": 0})
+                    global_model_stats["success_count"] += stats.get("success_count", 0)
+                    global_model_stats["prompt_tokens"] += stats.get("prompt_tokens", 0)
+                    global_model_stats["completion_tokens"] += stats.get("completion_tokens", 0)
+                
+                # Reset daily stats
+                data["daily"] = {"date": today_str, "models": {}}
+        
+        if needs_saving:
+            self._save_usage()
+
+    def get_next_smart_key(self, available_keys: List[str], model: str) -> Optional[str]:
         """
-        Finds the best key to use based on the lowest usage count for the given model.
+        Gets the least-used, available key based on daily stats.
         """
         best_key = None
         min_usage = float('inf')
+        
+        # Filter for keys that are not on cooldown
+        active_keys = []
+        for key in available_keys:
+            cooldown_until = self.usage_data.get(key, {}).get("cooldown_until")
+            if not cooldown_until or time.time() > cooldown_until:
+                active_keys.append(key)
 
-        eligible_keys = [k for k in available_keys if k not in excluded_keys]
-
-        if not eligible_keys:
+        if not active_keys:
             return None
 
-        # Initialize all available keys in usage data if they aren't present
-        for key in eligible_keys:
-            self.usage_data.setdefault(key, {"models": {}, "last_rotation_error": None})
-
-        # Find the key with the minimum success_count for the given model
-        for key in eligible_keys:
-            model_usage = self.usage_data[key].get("models", {}).get(model, {})
-            usage_count = model_usage.get("success_count", 0)
+        # Find the key with the minimum daily success_count for the given model
+        for key in active_keys:
+            key_data = self.usage_data.setdefault(key, {"daily": {"date": date.today().isoformat(), "models": {}}, "global": {"models": {}}, "cooldown_until": None})
+            daily_model_usage = key_data.get("daily", {}).get("models", {}).get(model, {})
+            usage_count = daily_model_usage.get("success_count", 0)
 
             if usage_count < min_usage:
                 min_usage = usage_count
                 best_key = key
         
-        # If all have the same usage count, it will pick the first one in the list
-        return best_key if best_key else eligible_keys[0]
-
+        return best_key if best_key else active_keys[0]
 
     def record_success(self, key: str, model: str, usage: Dict):
-        """Records a successful API call and its token usage."""
-        key_data = self.usage_data.setdefault(key, {"models": {}, "last_rotation_error": None})
-        model_data = key_data["models"].setdefault(model, {
-            "success_count": 0,
-            "prompt_tokens": 0,
-            "completion_tokens": 0
-        })
+        key_data = self.usage_data.setdefault(key, {"daily": {"date": date.today().isoformat(), "models": {}}, "global": {"models": {}}, "cooldown_until": None})
         
-        model_data["success_count"] += 1
-        model_data["prompt_tokens"] += usage.get("prompt_tokens", 0)
-        model_data["completion_tokens"] += usage.get("completion_tokens", 0)
+        # Ensure daily stats are for today
+        if key_data["daily"].get("date") != date.today().isoformat():
+            self._reset_daily_stats_if_needed() # Should be rare, but as a safeguard
+            key_data = self.usage_data[key]
+
+        daily_model_data = key_data["daily"]["models"].setdefault(model, {"success_count": 0, "prompt_tokens": 0, "completion_tokens": 0})
+        
+        daily_model_data["success_count"] += 1
+        daily_model_data["prompt_tokens"] += usage.get("prompt_tokens", 0)
+        daily_model_data["completion_tokens"] += usage.get("completion_tokens", 0)
         
         key_data["last_used_ts"] = time.time()
         self._save_usage()
 
-    def record_rotation_error(self, key: str, model: str, error: str):
-        """Records the error that caused a key to be rotated."""
-        key_data = self.usage_data.setdefault(key, {"models": {}, "last_rotation_error": None})
-            
+    def record_rotation_error(self, key: str, model: str, error: Exception):
+        key_data = self.usage_data.setdefault(key, {"daily": {"date": date.today().isoformat(), "models": {}}, "global": {"models": {}}, "cooldown_until": None})
+        
+        # Default cooldown of 24 hours
+        cooldown_seconds = 86400 
+        
+        # Try to parse retry_delay from the error message (very provider-specific)
+        error_str = str(error).lower()
+        if "retry_delay" in error_str:
+            try:
+                # A simple way to parse, might need to be more robust
+                delay_str = error_str.split("retry_delay")[1].split("seconds:")[1].strip().split("}")[0]
+                cooldown_seconds = int(delay_str)
+            except (IndexError, ValueError):
+                pass # Stick to default
+
+        key_data["cooldown_until"] = time.time() + cooldown_seconds
         key_data["last_rotation_error"] = {
             "timestamp": time.time(),
             "model": model,
-            "error": error
+            "error": str(error)
         }
         self._save_usage()