Retry transient OpenRouter errors

src/openrouter/utils/retries.py

@@ -1,11 +1,12 @@
 """Code generated by Speakeasy (https://speakeasy.com). DO NOT EDIT."""
 
 import asyncio
+import json
 import random
 import time
 from datetime import datetime
 from email.utils import parsedate_to_datetime
-from typing import List, Optional
+from typing import Any, List, Optional
 
 import httpx
 
@@ -93,6 +94,79 @@ def _parse_retry_after_header(response: httpx.Response) -> Optional[int]:
     return None
 
 
+TRANSIENT_STATUS_CODES = {408, 429}
+TRANSIENT_ERROR_CODES = {408, 429, 500, 502, 503, 504, 524, 529}
+
+
+def _matches_retry_status_code(response: httpx.Response, status_code: str) -> bool:
+    if "X" in status_code.upper():
+        code_range = int(status_code[0])
+        status_major = response.status_code / 100
+        return code_range <= status_major < code_range + 1
+
+    return response.status_code == int(status_code)
+
+
+def _error_code_from_json_body(response: httpx.Response) -> Optional[int]:
+    try:
+        body: Any = response.json()
+    except json.JSONDecodeError:
+        return None
+
+    if not isinstance(body, dict):
+        return None
+
+    error = body.get("error")
+    if not isinstance(error, dict):
+        return None
+
+    code = error.get("code")
+    if isinstance(code, int):
+        return code
+    if isinstance(code, str) and code.isdigit():
+        return int(code)
+
+    return None
+
+
+def _is_retryable_response(response: httpx.Response, status_codes: List[str]) -> bool:
+    if response.status_code in TRANSIENT_STATUS_CODES:
+        return True
+
+    for status_code in status_codes:
+        if _matches_retry_status_code(response, status_code):
+            return True
+
+    if response.status_code != 400:
+        return False
+
+    if not response.is_closed:
+        response.read()
+
+    inner_code = _error_code_from_json_body(response)
+    return inner_code in TRANSIENT_ERROR_CODES
+
+
+async def _is_retryable_response_async(
+    response: httpx.Response, status_codes: List[str]
+) -> bool:
+    if response.status_code in TRANSIENT_STATUS_CODES:
+        return True
+
+    for status_code in status_codes:
+        if _matches_retry_status_code(response, status_code):
+            return True
+
+    if response.status_code != 400:
+        return False
+
+    if not response.is_closed:
+        await response.aread()
+
+    inner_code = _error_code_from_json_body(response)
+    return inner_code in TRANSIENT_ERROR_CODES
+
+
 def _get_sleep_interval(
     exception: Exception,
     initial_interval: int,
@@ -131,19 +205,8 @@ def retry(func, retries: Retries):
             try:
                 res = func()
 
-                for code in retries.status_codes:
-                    if "X" in code.upper():
-                        code_range = int(code[0])
-
-                        status_major = res.status_code / 100
-
-                        if code_range <= status_major < code_range + 1:
-                            raise TemporaryError(res)
-                    else:
-                        parsed_code = int(code)
-
-                        if res.status_code == parsed_code:
-                            raise TemporaryError(res)
+                if _is_retryable_response(res, retries.status_codes):
+                    raise TemporaryError(res)
             except httpx.ConnectError as exception:
                 if retries.config.retry_connection_errors:
                     raise
@@ -180,19 +243,8 @@ async def retry_async(func, retries: Retries):
             try:
                 res = await func()
 
-                for code in retries.status_codes:
-                    if "X" in code.upper():
-                        code_range = int(code[0])
-
-                        status_major = res.status_code / 100
-
-                        if code_range <= status_major < code_range + 1:
-                            raise TemporaryError(res)
-                    else:
-                        parsed_code = int(code)
-
-                        if res.status_code == parsed_code:
-                            raise TemporaryError(res)
+                if await _is_retryable_response_async(res, retries.status_codes):
+                    raise TemporaryError(res)
             except httpx.ConnectError as exception:
                 if retries.config.retry_connection_errors:
                     raise