diff --git a/src/openrouter/utils/retries.py b/src/openrouter/utils/retries.py index 88a91b1..a8c8255 100644 --- a/src/openrouter/utils/retries.py +++ b/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 diff --git a/tests/test_retries.py b/tests/test_retries.py new file mode 100644 index 0000000..b6b5652 --- /dev/null +++ b/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()