Retry transient OpenRouter errors

Author: wassname

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 do_request() -> httpx.Response:
             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 do_request() -> httpx.Response:
             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

tests/test_retries.py

@@ -0,0 +1,75 @@
+import unittest
+
+import httpx
+
+from openrouter.utils.retries import (
+    _is_retryable_response,
+    _is_retryable_response_async,
+)
+
+
+def response(status_code, json_body=None):
+    request = httpx.Request("POST", "https://openrouter.ai/api/v1/chat/completions")
+    if json_body is None:
+        return httpx.Response(status_code, request=request)
+    return httpx.Response(status_code, json=json_body, request=request)
+
+
+class RetryResponseTests(unittest.TestCase):
+    def test_retries_request_timeout_and_rate_limit_by_default(self):
+        self.assertTrue(_is_retryable_response(response(408), []))
+        self.assertTrue(_is_retryable_response(response(429), []))
+
+    def test_retries_configured_status_code_patterns(self):
+        self.assertTrue(_is_retryable_response(response(502), ["5XX"]))
+        self.assertTrue(_is_retryable_response(response(529), ["5XX"]))
+        self.assertFalse(_is_retryable_response(response(400), ["5XX"]))
+
+    def test_retries_400_with_transient_inner_error_code(self):
+        res = response(
+            400,
+            {
+                "error": {
+                    "message": "Provider returned error",
+                    "code": 502,
+                    "metadata": {"provider_name": "example"},
+                }
+            },
+        )
+
+        self.assertTrue(_is_retryable_response(res, ["5XX"]))
+        self.assertEqual(res.json()["error"]["code"], 502)
+
+    def test_does_not_retry_400_with_non_transient_inner_error_code(self):
+        res = response(
+            400,
+            {
+                "error": {
+                    "message": "logprobs must be between 0 and 5",
+                    "code": 400,
+                }
+            },
+        )
+
+        self.assertFalse(_is_retryable_response(res, ["5XX"]))
+
+
+class AsyncRetryResponseTests(unittest.IsolatedAsyncioTestCase):
+    async def test_retries_400_with_transient_inner_error_code_async(self):
+        res = response(
+            400,
+            {
+                "error": {
+                    "message": "Provider returned error",
+                    "code": 502,
+                    "metadata": {"provider_name": "example"},
+                }
+            },
+        )
+
+        self.assertTrue(await _is_retryable_response_async(res, ["5XX"]))
+        self.assertEqual(res.json()["error"]["code"], 502)
+
+
+if __name__ == "__main__":
+    unittest.main()