Improve error handling

Author: Aculeasis

constants.py

@@ -9,9 +9,6 @@
 # Rate limit error code
 RATE_LIMIT_ERROR_CODE = 429
 
-# Rate limit error message
-RATE_LIMIT_ERROR_MESSAGE = "Rate limit exceeded: free-models-per-day"
-
 # Public endpoints that don't require authentication
 PUBLIC_ENDPOINTS = ["/api/v1/models"]
 

routes.py

@@ -16,7 +16,7 @@
 from key_manager import KeyManager
 from utils import (
     verify_access_key,
-    check_rate_limit_openai,
+    check_rate_limit_chat,
     check_rate_limit
 )
 
@@ -128,8 +128,10 @@ async def proxy_endpoint(
                 request, path, api_key, is_stream, is_completion
             )
 
-    except Exception as e:
+    except (Exception, HTTPException) as e:
         logger.error("Error proxying request: %s", str(e))
+        if isinstance(e, HTTPException):
+            raise e
         raise HTTPException(status_code=500, detail=f"Proxy error: {str(e)}") from e
 
 
@@ -143,25 +145,27 @@ async def handle_completions(
     """Handle chat completions using the OpenAI client."""
     try:
         # Extract headers to forward
-        forward_headers = {}
-        for k, v in request.headers.items():
-            if k.lower() in ["http-referer", "x-title"]:
-                forward_headers[k] = v
+        forward_headers = {
+            k: v
+            for k, v in request.headers.items()
+            if k.lower()
+               not in ["host", "content-length", "connection", "authorization"]
+        }
 
         # Create a copy of the request body to modify
         completion_args = request_body.copy()
 
+        # Ensure we don't pass 'stream' twice
+        if "stream" in completion_args:
+            del completion_args["stream"]
+
         # Move non-standard parameters that OpenAI SDK doesn't support directly to extra_body
         extra_body = {}
         openai_unsupported_params = ["include_reasoning", "transforms", "route", "provider"]
         for param in openai_unsupported_params:
             if param in completion_args:
                 extra_body[param] = completion_args.pop(param)
 
-        # Ensure we don't pass 'stream' twice
-        if "stream" in completion_args:
-            del completion_args["stream"]
-
         # Create a properly formatted request to the OpenAI API
         if is_stream:
             logger.info("Making streaming chat completion request")
@@ -186,7 +190,7 @@ async def stream_response() -> AsyncGenerator[bytes, None]:
                     logger.error("Error in streaming response: %s", err)
                     # Check if this is a rate limit error
                     if api_key:
-                        has_rate_limit_error_, reset_time_ms_ = check_rate_limit_openai(err)
+                        has_rate_limit_error_, reset_time_ms_ = check_rate_limit_chat(err)
                         if has_rate_limit_error_:
                             logger.warning("Rate limit detected in stream. Disabling key.")
                             await key_manager.disable_key(
@@ -221,26 +225,30 @@ async def stream_response() -> AsyncGenerator[bytes, None]:
         )
     except (APIError, Exception) as e:
         logger.error("Error in chat completions: %s", str(e))
-        # Check if this is a rate limit error
-        if api_key and isinstance(e, APIError):
-            has_rate_limit_error, reset_time_ms = check_rate_limit_openai(e)
-            if has_rate_limit_error:
-                logger.warning("Rate limit detected in stream. Disabling key.")
-                await key_manager.disable_key(api_key, reset_time_ms)
-
-                # Try again with a new key
-                new_api_key = await key_manager.get_next_key()
-                if new_api_key:
-                    new_client = await get_openai_client(new_api_key)
-                    return await handle_completions(
-                        new_client, request, request_body, new_api_key, is_stream
-                    )
-
+        code = 500
+        detail = f"Error processing chat completion: {str(e)}"
+        if isinstance(e, APIError):
+            # Check if this is a rate limit error
+            if api_key:
+                has_rate_limit_error, reset_time_ms = check_rate_limit_chat(e)
+                if has_rate_limit_error:
+                    logger.warning("Rate limit detected in stream. Disabling key.")
+                    await key_manager.disable_key(api_key, reset_time_ms)
+
+                    # Try again with a new key
+                    new_api_key = await key_manager.get_next_key()
+                    if new_api_key:
+                        new_client = await get_openai_client(new_api_key)
+                        return await handle_completions(
+                            new_client, request, request_body, new_api_key, is_stream
+                        )
+            code = e.code or code
+            detail = e.body or detail
         # Raise the exception
-        raise HTTPException(500, f"Error processing chat completion: {str(e)}") from e
+        raise HTTPException(code, detail) from e
 
 
-async def _check_httpx_err(body: str or bytes, api_key: str or None):
+async def _check_httpx_err(body: str | bytes, api_key: str | None):
     if api_key and (isinstance(body, str) and body.startswith("data: ") or (
             isinstance(body, bytes) and body.startswith(b"data: "))):
         body = body[6:]

utils.py

@@ -4,14 +4,15 @@
 """
 
 import socket
+import time
 import json
 from typing import Optional, Tuple
 
 from fastapi import Header, HTTPException
 from openai import APIError
 
 from config import config, logger
-from constants import RATE_LIMIT_ERROR_MESSAGE, RATE_LIMIT_ERROR_CODE
+from constants import RATE_LIMIT_ERROR_CODE
 
 
 def get_local_ip() -> str:
@@ -55,7 +56,58 @@ async def verify_access_key(
 
     return True
 
-def check_rate_limit_openai(err: APIError) -> Tuple[bool, Optional[int]]:
+
+def parse_google_rate_error(data: str) -> Optional[int]:
+    # data = {
+    #     'error': {
+    #         'code': 429,
+    #         'message': 'You exceeded your current quota, please check your plan and billing details.',
+    #         'status': 'RESOURCE_EXHAUSTED',
+    #         'details': [
+    #             {'@type': 'type.googleapis.com/google.rpc.QuotaFailure', 'violations': [
+    #                 {'quotaMetric': 'generativelanguage.googleapis.com/generate_content_paid_tier_input_token_count',
+    #                  'quotaId': 'GenerateContentPaidTierInputTokensPerModelPerMinute',
+    #                  'quotaDimensions': {'model': 'gemini-2.0-pro-exp', 'location': 'global'},
+    #                  'quotaValue': '10000000'}
+    #             ]},
+    #             {'@type': 'type.googleapis.com/google.rpc.Help', 'links': [
+    #                 {'description': 'Learn more about Gemini API quotas',
+    #                  'url': 'https://ai.google.dev/gemini-api/docs/rate-limits'}
+    #             ]},
+    #             {'@type': 'type.googleapis.com/google.rpc.RetryInfo', 'retryDelay': '5s'}
+    #         ]
+    #     }
+    # }
+    reset_time_ms = None
+    time_units = {'s': 1000, 'm': 60000, 'h': 3600000}
+    try:
+        data = json.loads(data)
+    except Exception as e:
+        logger.info("Json.loads error %s", e)
+    else:
+        retry_delay_ms = None
+        try:
+            message = data["error"].get("message", "")
+
+            retry_info = next((item for item in data['error']['details'] if
+                               item.get('@type') == 'type.googleapis.com/google.rpc.RetryInfo'), {})
+            retry_delay = retry_info.get('retryDelay', '0s')
+
+            num_part = ''.join(c for c in retry_delay if c.isdigit())
+            unit_part = ''.join(c for c in retry_delay if c.isalpha())
+
+            retry_delay_ms = int(num_part) * time_units.get(unit_part, 1000) if num_part else 0
+        except (TypeError, KeyError) as err:
+            logger.info("google reply parsing error %s", err)
+        else:
+            logger.info("google rate limit %s, retry: %s", message, retry_delay)
+
+        if retry_delay_ms:
+            reset_time_ms = int(time.time() * 1000) + retry_delay_ms
+
+    return reset_time_ms
+
+def check_rate_limit_chat(err: APIError) -> Tuple[bool, Optional[int]]:
     """
     Check for rate limit error.
 
@@ -68,15 +120,15 @@ def check_rate_limit_openai(err: APIError) -> Tuple[bool, Optional[int]]:
     has_rate_limit_error = False
     reset_time_ms = None
 
-    if err.code == RATE_LIMIT_ERROR_CODE and isinstance(err.body, dict):
-        try:
-            reset_time_ms = int(err.body["metadata"]["headers"]["X-RateLimit-Reset"])
-            has_rate_limit_error = True
-        except (TypeError, KeyError):
-            pass
-
-    if reset_time_ms is None and RATE_LIMIT_ERROR_MESSAGE in err.message:
+    if err.code == RATE_LIMIT_ERROR_CODE:
         has_rate_limit_error = True
+        if isinstance(err.body, dict):
+            try:
+                reset_time_ms = int(err.body["metadata"]["headers"]["X-RateLimit-Reset"])
+            except (TypeError, KeyError):
+                raw = err.body.get("metadata", {}).get("raw", "")
+                if raw and has_rate_limit_error:
+                    reset_time_ms = parse_google_rate_error(raw)
 
     return has_rate_limit_error, reset_time_ms
 
@@ -100,16 +152,19 @@ def check_rate_limit(data: str or bytes) -> Tuple[bool, Optional[int]]:
     else:
         if isinstance(err, dict) and "error" in err:
             code = err["error"].get("code", 0)
-            msg = err["error"].get("message", 0)
             try:
                 x_rate_limit = int(err["error"]["metadata"]["headers"]["X-RateLimit-Reset"])
             except (TypeError, KeyError):
-                x_rate_limit = 0
+                raw = err["error"].get("metadata", {}).get("raw", "")
+                if raw and code == RATE_LIMIT_ERROR_CODE:
+                    x_rate_limit = parse_google_rate_error(raw)
+                else:
+                    x_rate_limit = 0
 
             if x_rate_limit > 0:
                 has_rate_limit_error = True
                 reset_time_ms = x_rate_limit
-            elif code == RATE_LIMIT_ERROR_CODE and msg == RATE_LIMIT_ERROR_MESSAGE:
+            elif code == RATE_LIMIT_ERROR_CODE:
                 has_rate_limit_error = True
 
     return has_rate_limit_error, reset_time_ms