feat: Enhance error handling in UsageManager to extract retry delay from multiple error message formats

src/rotator_library/client.py

@@ -19,7 +19,7 @@ if not lib_logger.handlers:
 
 from .usage_manager import UsageManager
 from .failure_logger import log_failure
-from .error_handler import is_server_error, is_unrecoverable_error
+from .error_handler import is_rate_limit_error, is_server_error, is_unrecoverable_error
 from .providers import PROVIDER_PLUGINS
 
 class RotatingClient:
@@ -136,8 +136,8 @@ class RotatingClient:
                     )
                     
                     # 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...")
+                    if is_server_error(e) or is_rate_limit_error(e) and attempt < self.max_retries - 1:
+                        print(f"Key ...{current_key[-4:]} failed with server error or rate limit error. Retrying...")
                         await asyncio.sleep(1 * (attempt + 1))
                         continue
                     

src/rotator_library/usage_manager.py

@@ -6,6 +6,7 @@ from datetime import date, datetime
 from typing import Dict, List, Optional, Any
 from filelock import FileLock
 import litellm
+import re
 
 lib_logger = logging.getLogger('rotator_library')
 lib_logger.propagate = False # Ensure this logger doesn't propagate to root
@@ -128,7 +129,44 @@ class UsageManager:
         error_str = str(error).lower()
         if "retry_delay" in error_str:
             try:
-                delay_str = error_str.split("retry_delay")[1].split("seconds:")[1].strip().split("}")[0]
+                # Try multiple patterns to extract delay from error message
+                delay_str = None
+                
+                # Pattern 1: retry_delay...seconds format
+                if "retry_delay" in error_str and "seconds:" in error_str:
+                    try:
+                        delay_str = error_str.split("retry_delay")[1].split("seconds:")[1].strip().split("}")[0]
+                    except (IndexError, AttributeError):
+                        pass
+                
+                # Pattern 2: retryDelay with 's' suffix (Gemini format)
+                if not delay_str and "retrydelay" in error_str:
+                    try:
+                        match = re.search(r'"retrydelay":\s*"(\d+)s"', error_str)
+                        if match:
+                            delay_str = match.group(1)
+                    except Exception:
+                        pass
+                
+                # Pattern 3: Generic numeric extraction for retry/delay contexts
+                if not delay_str:
+                    try:
+                        # Look for numbers followed by 's' or 'seconds' in retry/delay context
+                        patterns = [
+                            r'retry.*?(\d+)s',
+                            r'delay.*?(\d+)s', 
+                            r'wait.*?(\d+)\s*seconds?'
+                        ]
+                        for pattern in patterns:
+                            match = re.search(pattern, error_str, re.IGNORECASE)
+                            if match:
+                                delay_str = match.group(1)
+                                break
+                    except Exception:
+                        pass
+                
+                if delay_str:
+                    cooldown_seconds = int(delay_str)
                 cooldown_seconds = int(delay_str)
             except (IndexError, ValueError):
                 pass