From 4a3d9107d4f345267cfdcabf07d201a4863644d0 Mon Sep 17 00:00:00 2001 From: Andrew Peabody Date: Sun, 18 Oct 2026 10:15:54 -0700 Subject: [PATCH 01/20] fix(retries): honor Retry-After headers, apply exponential backoff with jitter, and filter transient errors --- litellm/main.py | 131 ++++- litellm/utils.py | 31 +- .../unit/test_completion_transient_retries.py | 466 ++++++++++++++++++ 3 files changed, 591 insertions(+), 37 deletions(-) create mode 100644 tests/unit/test_completion_transient_retries.py diff --git a/litellm/main.py b/litellm/main.py index 38591441e01..250c07b28fc 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -167,6 +167,7 @@ from litellm.utils import ( Usage, add_provider_specific_params_to_optional_params, async_mock_completion_streaming_obj, + calculate_retry_after, convert_to_model_response_object, create_pretrained_tokenizer, create_tokenizer, @@ -183,6 +184,7 @@ from litellm.utils import ( mock_completion_streaming_obj, pre_process_non_default_params, read_config_args, + should_retry, should_run_mock_completion, supports_httpx_timeout, validate_and_fix_openai_messages, @@ -6334,55 +6336,142 @@ def completion( ) +class _RetryOutcome(Protocol): + def exception(self) -> BaseException | None: ... + + +class _RetryState(Protocol): + @property + def attempt_number(self) -> int: ... + + @property + def outcome(self) -> _RetryOutcome | None: ... + + +def _compute_retry_wait( + exc: BaseException | None, + attempt_number: int, + num_retries: int, + retry_strategy: str, +) -> float: + if retry_strategy != "exponential_backoff_retry": + return 0.0 + remaining: Final[int] = max(0, num_retries - attempt_number) + headers: Final[httpx.Headers | None] = ( + httpx.Headers(exc.response.headers) + if isinstance(exc, openai.APIStatusError) and getattr(exc.response, "headers", None) + else None + ) + return float( + calculate_retry_after( + remaining_retries=remaining, + max_retries=num_retries, + response_headers=headers, + ) + ) + + +def _build_retry_predicate(retry_policy_override: bool) -> Callable[[BaseException], bool]: + """Check if an exception is transient or explicitly enabled via RetryPolicy.""" + + def _predicate(e: BaseException) -> bool: + if retry_policy_override and isinstance(e, openai.APIError): + return True + status_code: Final[object] = getattr(e, "status_code", None) + if isinstance(status_code, int): + return should_retry(status_code) + return isinstance(e, (openai.APITimeoutError, openai.APIConnectionError)) + + return _predicate + + +def _build_retry_wait_fn(num_retries: int, retry_strategy: str) -> Callable[[_RetryState], float]: + """Tenacity wait callback honoring provider Retry-After headers and jitter via calculate_retry_after.""" + + def _wait_fn(retry_state: _RetryState) -> float: + outcome: Final[_RetryOutcome | None] = retry_state.outcome + exc: Final[BaseException | None] = outcome.exception() if outcome is not None else None + return _compute_retry_wait( + exc=exc, + attempt_number=retry_state.attempt_number, + num_retries=num_retries, + retry_strategy=retry_strategy, + ) + + return _wait_fn + + def completion_with_retries(*args, **kwargs): """ - Executes a litellm.completion() with 3 retries + Executes a litellm.completion() with retries """ try: import tenacity except Exception as e: raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") - num_retries: Final = kwargs.pop("num_retries", 3) + raw_retries: Final[object] = kwargs.pop("num_retries", 3) + num_retries: Final[int] = raw_retries if isinstance(raw_retries, int) else 3 # reset retries in .completion() kwargs["max_retries"] = 0 kwargs["num_retries"] = 0 - retry_strategy: Final[RetryStrategy] = kwargs.pop("retry_strategy", "constant_retry") + retry_strategy: Final[str] = ( + "constant_retry" if kwargs.pop("retry_strategy", None) == "constant_retry" else "exponential_backoff_retry" + ) original_function: Final = kwargs.pop("original_function", completion) - if retry_strategy == "exponential_backoff_retry": - retryer = tenacity.Retrying( - wait=tenacity.wait_exponential(multiplier=1, max=10), - stop=tenacity.stop_after_attempt(num_retries), - reraise=True, + initial_exc: Final[object] = kwargs.pop("_initial_exception", None) + retry_policy_override: Final[bool] = kwargs.pop("retry_policy", None) is not None + if isinstance(initial_exc, BaseException): + wait_seconds: Final[float] = _compute_retry_wait( + exc=initial_exc, + attempt_number=0, + num_retries=num_retries, + retry_strategy=retry_strategy, ) - else: - retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + time.sleep(wait_seconds) + retryer: Final = tenacity.Retrying( + wait=_build_retry_wait_fn(num_retries, retry_strategy), + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_override)), + reraise=True, + ) return retryer(original_function, *args, **kwargs) async def acompletion_with_retries(*args, **kwargs): """ [DEPRECATED]. Use 'acompletion' or router.acompletion instead! - Executes a litellm.completion() with 3 retries + Executes a litellm.acompletion() with retries """ try: import tenacity except Exception as e: raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") - num_retries: Final = kwargs.pop("num_retries", 3) + raw_retries: Final[object] = kwargs.pop("num_retries", 3) + num_retries: Final[int] = raw_retries if isinstance(raw_retries, int) else 3 kwargs["max_retries"] = 0 kwargs["num_retries"] = 0 - retry_strategy: Final[RetryStrategy] = kwargs.pop("retry_strategy", "constant_retry") - original_function: Final = kwargs.pop("original_function", completion) - if retry_strategy == "exponential_backoff_retry": - retryer = tenacity.AsyncRetrying( - wait=tenacity.wait_exponential(multiplier=1, max=10), - stop=tenacity.stop_after_attempt(num_retries), - reraise=True, + retry_strategy: Final[str] = ( + "constant_retry" if kwargs.pop("retry_strategy", None) == "constant_retry" else "exponential_backoff_retry" + ) + original_function: Final = kwargs.pop("original_function", acompletion) + initial_exc: Final[object] = kwargs.pop("_initial_exception", None) + retry_policy_override: Final[bool] = kwargs.pop("retry_policy", None) is not None + if isinstance(initial_exc, BaseException): + wait_seconds: Final[float] = _compute_retry_wait( + exc=initial_exc, + attempt_number=0, + num_retries=num_retries, + retry_strategy=retry_strategy, ) - else: - retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + await asyncio.sleep(wait_seconds) + retryer: Final = tenacity.AsyncRetrying( + wait=_build_retry_wait_fn(num_retries, retry_strategy), + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_override)), + reraise=True, + ) return await retryer(original_function, *args, **kwargs) diff --git a/litellm/utils.py b/litellm/utils.py index ab97b69897f..715eb2938d9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1475,9 +1475,11 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu exception=exception, retry_policy=kwargs.get("retry_policy"), ) - kwargs["retry_policy"] = reset_retry_policy() if retry_policy_num_retries is not None: num_retries = retry_policy_num_retries + kwargs["retry_policy"] = reset_retry_policy() + else: + kwargs.pop("retry_policy", None) return num_retries, kwargs @@ -1990,13 +1992,12 @@ def client(original_function): if ( num_retries and not is_completion_litellm_router_call ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying + status_code: Final[object] = getattr(e, "status_code", None) if ( - isinstance(e, openai.APIError) - or isinstance(e, openai.Timeout) - or isinstance(e, openai.APIConnectionError) - ): + "retry_policy" in kwargs or not isinstance(status_code, int) or should_retry(status_code) + ) and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)): kwargs["num_retries"] = num_retries - return litellm.completion_with_retries(*args, **kwargs) + return litellm.completion_with_retries(*args, _initial_exception=e, **kwargs) elif ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict @@ -2305,19 +2306,17 @@ def client(original_function): if ( num_retries and not is_acompletion_litellm_router_call and not isinstance(e, ImportError) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying - try: + status_code: Final[object] = getattr(e, "status_code", None) + if ( + "retry_policy" in kwargs or not isinstance(status_code, int) or should_retry(status_code) + ) and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)): litellm.num_retries = None # set retries to None to prevent infinite loops kwargs["num_retries"] = num_retries kwargs["original_function"] = original_function - if isinstance(e, openai.RateLimitError): # rate limiting specific error - kwargs["retry_strategy"] = "exponential_backoff_retry" - elif isinstance(e, openai.APIError): # generic api error - kwargs["retry_strategy"] = "constant_retry" - result = await litellm.acompletion_with_retries(*args, **kwargs) - except Exception: - pass - else: - return result + try: + return await litellm.acompletion_with_retries(*args, _initial_exception=e, **kwargs) + except openai.APIError as retry_err: + e = retry_err elif ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict diff --git a/tests/unit/test_completion_transient_retries.py b/tests/unit/test_completion_transient_retries.py new file mode 100644 index 00000000000..957a26ef956 --- /dev/null +++ b/tests/unit/test_completion_transient_retries.py @@ -0,0 +1,466 @@ +from unittest.mock import AsyncMock, MagicMock, patch + +import httpx +import openai +import pytest + +import litellm +from litellm import ( + AuthenticationError, + BadRequestError, + InternalServerError, + RateLimitError, + ServiceUnavailableError, + acompletion_with_retries, + completion_with_retries, +) + + +def _make_rate_limit_error(retry_after: str | None = None) -> RateLimitError: + headers = {"retry-after": retry_after} if retry_after else {} + resp = MagicMock() + resp.headers = headers + resp.status_code = 429 + err = RateLimitError( + message="Rate limit reached", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 429 + err.response = resp + return err + + +def _make_auth_error() -> AuthenticationError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 401 + err = AuthenticationError( + message="Invalid API Key", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 401 + return err + + +def _make_bad_request_error() -> BadRequestError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 400 + err = BadRequestError( + message="Invalid argument", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 400 + return err + + +def _make_service_unavailable_error() -> ServiceUnavailableError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 503 + err = ServiceUnavailableError( + message="The server is temporarily overloaded", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 503 + return err + + +def _make_internal_server_error() -> InternalServerError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 500 + err = InternalServerError( + message="Internal server error", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 500 + return err + + +def test_completion_with_retries_honors_retry_after(): + """Verify completion_with_retries reads Retry-After and applies jitter.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise _make_rate_limit_error(retry_after="2") + return {"status": "ok"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert res == {"status": "ok"} + assert len(calls) == 2 + mock_sleep.assert_called_once() + sleep_arg = mock_sleep.call_args[0][0] + assert 2.0 <= sleep_arg <= 2.8 + + +def test_completion_with_retries_non_transient_auth_error_fails_immediately(): + """Non-transient errors like 401 should fail on attempt 1 without retrying.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with patch("time.sleep") as mock_sleep: + with pytest.raises(AuthenticationError): + completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +def test_completion_with_retries_bad_request_fails_immediately(): + """Non-transient errors like 400 Bad Request should fail on attempt 1 without retrying.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + raise _make_bad_request_error() + + with patch("time.sleep") as mock_sleep: + with pytest.raises(BadRequestError): + completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +def test_completion_with_retries_service_unavailable_503(): + """503 Service Unavailable should be retried with exponential backoff.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) < 3: + raise _make_service_unavailable_error() + return {"status": "recovered"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + assert len(calls) == 3 + assert mock_sleep.call_count == 2 + + +def test_completion_with_retries_internal_server_error_500(): + """500 Internal Server Error should be retried.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) < 2: + raise _make_internal_server_error() + return {"status": "recovered"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + assert len(calls) == 2 + mock_sleep.assert_called_once() + + +@pytest.mark.asyncio +async def test_acompletion_with_retries_honors_retry_after(): + """Verify acompletion_with_retries reads Retry-After and applies jitter.""" + calls = [] + + async def mock_acompletion(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise _make_rate_limit_error(retry_after="3") + return {"status": "ok"} + + with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + res = await acompletion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_acompletion, + ) + assert res == {"status": "ok"} + assert len(calls) == 2 + mock_sleep.assert_called_once() + sleep_arg = mock_sleep.call_args[0][0] + assert 3.0 <= sleep_arg <= 3.8 + + +@pytest.mark.asyncio +async def test_acompletion_with_retries_non_transient_auth_error_fails_immediately(): + """Async 401 should fail on attempt 1 without retrying.""" + calls = [] + + async def mock_acompletion(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + with pytest.raises(AuthenticationError): + await acompletion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_acompletion, + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +def test_completion_decorator_retries_transient_error(): + """Verify litellm.completion(..., num_retries=2) retries on RateLimitError and sleeps on first retry.""" + calls = [] + + def mock_chat_fn(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise _make_rate_limit_error(retry_after="1") + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + return mock_response + + with ( + patch("time.sleep") as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), + ): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + ) + assert len(calls) == 2 + mock_sleep.assert_called_once() + assert 1.0 <= mock_sleep.call_args[0][0] <= 1.8 + + +def test_completion_decorator_fails_fast_on_auth_error(): + """Verify litellm.completion(..., num_retries=2) fails immediately on AuthenticationError without retry.""" + calls = [] + + def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with ( + patch("time.sleep") as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), + ): + with pytest.raises(AuthenticationError): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +@pytest.mark.asyncio +async def test_acompletion_decorator_preserves_retried_count_on_failure(): + """Verify litellm.acompletion(..., num_retries=3) attaches num_retries to the raised exception.""" + calls = [] + + async def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_internal_server_error() + + with ( + patch("asyncio.sleep", new_callable=AsyncMock), + patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), + ): + with pytest.raises(InternalServerError, match="LiteLLM Retried: 3 times"): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + api_key="test-key", + num_retries=3, + ) + assert len(calls) == 4 + + +def test_completion_with_retries_connection_error_with_constant_retry(): + """Verify APIConnectionError (no status_code) retries and constant_retry waits 0.0s.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise openai.APIConnectionError(request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions")) + return {"status": "recovered"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + retry_strategy="constant_retry", + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + assert len(calls) == 2 + mock_sleep.assert_called_once_with(0.0) + + +@pytest.mark.asyncio +async def test_acompletion_decorator_retries_transient_error(): + """Verify litellm.acompletion(..., num_retries=2) retries on RateLimitError and sleeps on first retry.""" + calls = [] + + async def mock_chat_fn(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise _make_rate_limit_error(retry_after="1") + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + return mock_response + + with ( + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), + ): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + api_key="test-key", + num_retries=2, + ) + assert len(calls) == 2 + mock_sleep.assert_called_once() + assert 1.0 <= mock_sleep.call_args[0][0] <= 1.8 + + +def test_completion_retry_policy_bad_request_honors_full_count(): + """Explicit RetryPolicy(BadRequestErrorRetries=3) should execute all 3 retries (4 total attempts).""" + from litellm.types.router import RetryPolicy + + calls = [] + + def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_bad_request_error() + + with ( + patch("time.sleep") as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), + ): + with pytest.raises(BadRequestError): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + retry_policy=RetryPolicy(BadRequestErrorRetries=3), + ) + assert len(calls) == 4 + assert mock_sleep.call_count == 3 + + +@pytest.mark.asyncio +async def test_acompletion_retry_policy_bad_request_honors_full_count(): + """Explicit RetryPolicy(BadRequestErrorRetries=3) on async should execute all 3 retries (4 total attempts).""" + from litellm.types.router import RetryPolicy + + calls = [] + + async def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_bad_request_error() + + with ( + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), + ): + with pytest.raises(BadRequestError): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + api_key="test-key", + retry_policy=RetryPolicy(BadRequestErrorRetries=3), + ) + assert len(calls) == 4 + assert mock_sleep.call_count == 3 + + +def test_completion_unconfigured_retry_policy_fails_fast_on_auth_error(): + """An unconfigured RetryPolicy() alongside num_retries=2 on sync should not retry AuthenticationError.""" + from litellm.types.router import RetryPolicy + + calls = [] + + def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with ( + patch("time.sleep") as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), + ): + with pytest.raises(AuthenticationError): + litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + retry_policy=RetryPolicy(), + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +@pytest.mark.asyncio +async def test_acompletion_unconfigured_retry_policy_fails_fast_on_auth_error(): + """An unconfigured RetryPolicy() alongside num_retries=2 should not retry AuthenticationError.""" + from litellm.types.router import RetryPolicy + + calls = [] + + async def mock_chat_fn(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with ( + patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, + patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), + ): + with pytest.raises(AuthenticationError): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + api_key="test-key", + num_retries=2, + retry_policy=RetryPolicy(), + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() From 3f7ec221005b0e7b02113f6de7d502ea85124ec9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 22:53:24 -0700 Subject: [PATCH 02/20] fix(retries): read Retry-After from mapped provider errors, honor it under constant_retry, and keep the provider error when async retries cannot run The retry wait read headers only from openai.APIStatusError.response, which litellm's mapped exceptions do not carry, so a real provider's Retry-After was never seen. Read them through the shared response-header extractor instead, honor a valid Retry-After under constant_retry too, and keep the original provider error when acompletion retries cannot start (tenacity missing). Move the retry tests into the mapped test_main.py and drive them through respx upstreams instead of patching our own handlers. --- litellm/main.py | 28 +- litellm/utils.py | 2 + .../unit/test_completion_transient_retries.py | 466 ------------------ tests/unit/test_main.py | 406 ++++++++++++++- 4 files changed, 421 insertions(+), 481 deletions(-) delete mode 100644 tests/unit/test_completion_transient_retries.py diff --git a/litellm/main.py b/litellm/main.py index 86dcfb30a2d..d5d9db8b0ff 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -79,6 +79,9 @@ from litellm.litellm_core_utils.chat_completion_agentic_loop import ( ) from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.dd_tracing import tracer +from litellm.litellm_core_utils.exception_mapping_utils import ( + _get_response_headers, # pyright: ignore[reportPrivateUsage] # shared provider header extractor +) from litellm.litellm_core_utils.get_litellm_params import ( AWS_CREDENTIAL_KWARGS_KEYS, OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS, @@ -171,6 +174,7 @@ from litellm.utils import ( TextCompletionStreamWrapper, TranscriptionResponse, Usage, + _get_retry_after_from_exception_header, # pyright: ignore[reportPrivateUsage] # shared Retry-After parser add_provider_specific_params_to_optional_params, async_mock_completion_streaming_obj, calculate_retry_after, @@ -6386,21 +6390,17 @@ def _compute_retry_wait( num_retries: int, retry_strategy: str, ) -> float: - if retry_strategy != "exponential_backoff_retry": - return 0.0 - remaining: Final[int] = max(0, num_retries - attempt_number) - headers: Final[httpx.Headers | None] = ( - httpx.Headers(exc.response.headers) - if isinstance(exc, openai.APIStatusError) and getattr(exc.response, "headers", None) - else None - ) - return float( - calculate_retry_after( - remaining_retries=remaining, - max_retries=num_retries, - response_headers=headers, + headers: Final[httpx.Headers | None] = _get_response_headers(exc) if isinstance(exc, Exception) else None + if retry_strategy == "exponential_backoff_retry": + return float( + calculate_retry_after( + remaining_retries=max(0, num_retries - attempt_number), + max_retries=num_retries, + response_headers=headers, + ) ) - ) + retry_after: Final[int | None] = _get_retry_after_from_exception_header(headers) + return float(retry_after) if retry_after is not None and 0 < retry_after <= 60 else 0.0 def _build_retry_predicate(retry_policy_override: bool) -> Callable[[BaseException], bool]: diff --git a/litellm/utils.py b/litellm/utils.py index f0c1e844002..764a2b2abca 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2401,6 +2401,8 @@ def client(original_function): return await litellm.acompletion_with_retries(*args, _initial_exception=e, **kwargs) except openai.APIError as retry_err: e = retry_err + except Exception as retry_unavailable: + verbose_logger.debug("acompletion retries did not run: %s", retry_unavailable) elif ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict diff --git a/tests/unit/test_completion_transient_retries.py b/tests/unit/test_completion_transient_retries.py deleted file mode 100644 index 957a26ef956..00000000000 --- a/tests/unit/test_completion_transient_retries.py +++ /dev/null @@ -1,466 +0,0 @@ -from unittest.mock import AsyncMock, MagicMock, patch - -import httpx -import openai -import pytest - -import litellm -from litellm import ( - AuthenticationError, - BadRequestError, - InternalServerError, - RateLimitError, - ServiceUnavailableError, - acompletion_with_retries, - completion_with_retries, -) - - -def _make_rate_limit_error(retry_after: str | None = None) -> RateLimitError: - headers = {"retry-after": retry_after} if retry_after else {} - resp = MagicMock() - resp.headers = headers - resp.status_code = 429 - err = RateLimitError( - message="Rate limit reached", - response=resp, - llm_provider="openai", - model="gpt-4o", - ) - err.status_code = 429 - err.response = resp - return err - - -def _make_auth_error() -> AuthenticationError: - resp = MagicMock() - resp.headers = {} - resp.status_code = 401 - err = AuthenticationError( - message="Invalid API Key", - response=resp, - llm_provider="openai", - model="gpt-4o", - ) - err.status_code = 401 - return err - - -def _make_bad_request_error() -> BadRequestError: - resp = MagicMock() - resp.headers = {} - resp.status_code = 400 - err = BadRequestError( - message="Invalid argument", - response=resp, - llm_provider="openai", - model="gpt-4o", - ) - err.status_code = 400 - return err - - -def _make_service_unavailable_error() -> ServiceUnavailableError: - resp = MagicMock() - resp.headers = {} - resp.status_code = 503 - err = ServiceUnavailableError( - message="The server is temporarily overloaded", - response=resp, - llm_provider="openai", - model="gpt-4o", - ) - err.status_code = 503 - return err - - -def _make_internal_server_error() -> InternalServerError: - resp = MagicMock() - resp.headers = {} - resp.status_code = 500 - err = InternalServerError( - message="Internal server error", - response=resp, - llm_provider="openai", - model="gpt-4o", - ) - err.status_code = 500 - return err - - -def test_completion_with_retries_honors_retry_after(): - """Verify completion_with_retries reads Retry-After and applies jitter.""" - calls = [] - - def mock_completion(*args, **kwargs): - calls.append(1) - if len(calls) == 1: - raise _make_rate_limit_error(retry_after="2") - return {"status": "ok"} - - with patch("time.sleep") as mock_sleep: - res = completion_with_retries( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - num_retries=3, - original_function=mock_completion, - ) - assert res == {"status": "ok"} - assert len(calls) == 2 - mock_sleep.assert_called_once() - sleep_arg = mock_sleep.call_args[0][0] - assert 2.0 <= sleep_arg <= 2.8 - - -def test_completion_with_retries_non_transient_auth_error_fails_immediately(): - """Non-transient errors like 401 should fail on attempt 1 without retrying.""" - calls = [] - - def mock_completion(*args, **kwargs): - calls.append(1) - raise _make_auth_error() - - with patch("time.sleep") as mock_sleep: - with pytest.raises(AuthenticationError): - completion_with_retries( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - num_retries=3, - original_function=mock_completion, - ) - assert len(calls) == 1 - mock_sleep.assert_not_called() - - -def test_completion_with_retries_bad_request_fails_immediately(): - """Non-transient errors like 400 Bad Request should fail on attempt 1 without retrying.""" - calls = [] - - def mock_completion(*args, **kwargs): - calls.append(1) - raise _make_bad_request_error() - - with patch("time.sleep") as mock_sleep: - with pytest.raises(BadRequestError): - completion_with_retries( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - num_retries=3, - original_function=mock_completion, - ) - assert len(calls) == 1 - mock_sleep.assert_not_called() - - -def test_completion_with_retries_service_unavailable_503(): - """503 Service Unavailable should be retried with exponential backoff.""" - calls = [] - - def mock_completion(*args, **kwargs): - calls.append(1) - if len(calls) < 3: - raise _make_service_unavailable_error() - return {"status": "recovered"} - - with patch("time.sleep") as mock_sleep: - res = completion_with_retries( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - num_retries=3, - original_function=mock_completion, - ) - assert res == {"status": "recovered"} - assert len(calls) == 3 - assert mock_sleep.call_count == 2 - - -def test_completion_with_retries_internal_server_error_500(): - """500 Internal Server Error should be retried.""" - calls = [] - - def mock_completion(*args, **kwargs): - calls.append(1) - if len(calls) < 2: - raise _make_internal_server_error() - return {"status": "recovered"} - - with patch("time.sleep") as mock_sleep: - res = completion_with_retries( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - num_retries=2, - original_function=mock_completion, - ) - assert res == {"status": "recovered"} - assert len(calls) == 2 - mock_sleep.assert_called_once() - - -@pytest.mark.asyncio -async def test_acompletion_with_retries_honors_retry_after(): - """Verify acompletion_with_retries reads Retry-After and applies jitter.""" - calls = [] - - async def mock_acompletion(*args, **kwargs): - calls.append(1) - if len(calls) == 1: - raise _make_rate_limit_error(retry_after="3") - return {"status": "ok"} - - with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: - res = await acompletion_with_retries( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - num_retries=3, - original_function=mock_acompletion, - ) - assert res == {"status": "ok"} - assert len(calls) == 2 - mock_sleep.assert_called_once() - sleep_arg = mock_sleep.call_args[0][0] - assert 3.0 <= sleep_arg <= 3.8 - - -@pytest.mark.asyncio -async def test_acompletion_with_retries_non_transient_auth_error_fails_immediately(): - """Async 401 should fail on attempt 1 without retrying.""" - calls = [] - - async def mock_acompletion(*args, **kwargs): - calls.append(1) - raise _make_auth_error() - - with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: - with pytest.raises(AuthenticationError): - await acompletion_with_retries( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - num_retries=3, - original_function=mock_acompletion, - ) - assert len(calls) == 1 - mock_sleep.assert_not_called() - - -def test_completion_decorator_retries_transient_error(): - """Verify litellm.completion(..., num_retries=2) retries on RateLimitError and sleeps on first retry.""" - calls = [] - - def mock_chat_fn(*args, **kwargs): - calls.append(1) - if len(calls) == 1: - raise _make_rate_limit_error(retry_after="1") - mock_response = MagicMock() - mock_response.choices = [MagicMock()] - return mock_response - - with ( - patch("time.sleep") as mock_sleep, - patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), - ): - litellm.completion( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - num_retries=2, - ) - assert len(calls) == 2 - mock_sleep.assert_called_once() - assert 1.0 <= mock_sleep.call_args[0][0] <= 1.8 - - -def test_completion_decorator_fails_fast_on_auth_error(): - """Verify litellm.completion(..., num_retries=2) fails immediately on AuthenticationError without retry.""" - calls = [] - - def mock_chat_fn(*args, **kwargs): - calls.append(1) - raise _make_auth_error() - - with ( - patch("time.sleep") as mock_sleep, - patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), - ): - with pytest.raises(AuthenticationError): - litellm.completion( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - num_retries=2, - ) - assert len(calls) == 1 - mock_sleep.assert_not_called() - - -@pytest.mark.asyncio -async def test_acompletion_decorator_preserves_retried_count_on_failure(): - """Verify litellm.acompletion(..., num_retries=3) attaches num_retries to the raised exception.""" - calls = [] - - async def mock_chat_fn(*args, **kwargs): - calls.append(1) - raise _make_internal_server_error() - - with ( - patch("asyncio.sleep", new_callable=AsyncMock), - patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), - ): - with pytest.raises(InternalServerError, match="LiteLLM Retried: 3 times"): - await litellm.acompletion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "hi"}], - api_key="test-key", - num_retries=3, - ) - assert len(calls) == 4 - - -def test_completion_with_retries_connection_error_with_constant_retry(): - """Verify APIConnectionError (no status_code) retries and constant_retry waits 0.0s.""" - calls = [] - - def mock_completion(*args, **kwargs): - calls.append(1) - if len(calls) == 1: - raise openai.APIConnectionError(request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions")) - return {"status": "recovered"} - - with patch("time.sleep") as mock_sleep: - res = completion_with_retries( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - num_retries=2, - retry_strategy="constant_retry", - original_function=mock_completion, - ) - assert res == {"status": "recovered"} - assert len(calls) == 2 - mock_sleep.assert_called_once_with(0.0) - - -@pytest.mark.asyncio -async def test_acompletion_decorator_retries_transient_error(): - """Verify litellm.acompletion(..., num_retries=2) retries on RateLimitError and sleeps on first retry.""" - calls = [] - - async def mock_chat_fn(*args, **kwargs): - calls.append(1) - if len(calls) == 1: - raise _make_rate_limit_error(retry_after="1") - mock_response = MagicMock() - mock_response.choices = [MagicMock()] - return mock_response - - with ( - patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, - patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), - ): - await litellm.acompletion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "hi"}], - api_key="test-key", - num_retries=2, - ) - assert len(calls) == 2 - mock_sleep.assert_called_once() - assert 1.0 <= mock_sleep.call_args[0][0] <= 1.8 - - -def test_completion_retry_policy_bad_request_honors_full_count(): - """Explicit RetryPolicy(BadRequestErrorRetries=3) should execute all 3 retries (4 total attempts).""" - from litellm.types.router import RetryPolicy - - calls = [] - - def mock_chat_fn(*args, **kwargs): - calls.append(1) - raise _make_bad_request_error() - - with ( - patch("time.sleep") as mock_sleep, - patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), - ): - with pytest.raises(BadRequestError): - litellm.completion( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - retry_policy=RetryPolicy(BadRequestErrorRetries=3), - ) - assert len(calls) == 4 - assert mock_sleep.call_count == 3 - - -@pytest.mark.asyncio -async def test_acompletion_retry_policy_bad_request_honors_full_count(): - """Explicit RetryPolicy(BadRequestErrorRetries=3) on async should execute all 3 retries (4 total attempts).""" - from litellm.types.router import RetryPolicy - - calls = [] - - async def mock_chat_fn(*args, **kwargs): - calls.append(1) - raise _make_bad_request_error() - - with ( - patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, - patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), - ): - with pytest.raises(BadRequestError): - await litellm.acompletion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "hi"}], - api_key="test-key", - retry_policy=RetryPolicy(BadRequestErrorRetries=3), - ) - assert len(calls) == 4 - assert mock_sleep.call_count == 3 - - -def test_completion_unconfigured_retry_policy_fails_fast_on_auth_error(): - """An unconfigured RetryPolicy() alongside num_retries=2 on sync should not retry AuthenticationError.""" - from litellm.types.router import RetryPolicy - - calls = [] - - def mock_chat_fn(*args, **kwargs): - calls.append(1) - raise _make_auth_error() - - with ( - patch("time.sleep") as mock_sleep, - patch.object(litellm.main.openai_chat_completions, "completion", side_effect=mock_chat_fn), - ): - with pytest.raises(AuthenticationError): - litellm.completion( - model="gpt-4o", - messages=[{"role": "user", "content": "hi"}], - num_retries=2, - retry_policy=RetryPolicy(), - ) - assert len(calls) == 1 - mock_sleep.assert_not_called() - - -@pytest.mark.asyncio -async def test_acompletion_unconfigured_retry_policy_fails_fast_on_auth_error(): - """An unconfigured RetryPolicy() alongside num_retries=2 should not retry AuthenticationError.""" - from litellm.types.router import RetryPolicy - - calls = [] - - async def mock_chat_fn(*args, **kwargs): - calls.append(1) - raise _make_auth_error() - - with ( - patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep, - patch.object(litellm.main.openai_chat_completions, "acompletion", side_effect=mock_chat_fn), - ): - with pytest.raises(AuthenticationError): - await litellm.acompletion( - model="gpt-4o-mini", - messages=[{"role": "user", "content": "hi"}], - api_key="test-key", - num_retries=2, - retry_policy=RetryPolicy(), - ) - assert len(calls) == 1 - mock_sleep.assert_not_called() diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index c3794ea7544..888877ce603 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -20,7 +20,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from openai import APITimeoutError, AsyncOpenAI +from openai import APIConnectionError, APITimeoutError, AsyncOpenAI from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion @@ -47,6 +47,7 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.litellm_params import ControlOptions from litellm.types.llms.openai import AllMessageValues, HttpxBinaryResponseContent from litellm.types.prompts.init_prompts import PromptSpec +from litellm.types.router import RetryPolicy from litellm.types.utils import Delta, ModelResponseStream, StandardCallbackDynamicParams, StreamingChoices, Usage from tests.unit.llms.bedrock.slow_upstream import ( STREAM_TIMEOUT_SECONDS, @@ -9626,3 +9627,406 @@ async def test_atranscription_ignores_a_non_string_base_url_when_api_base_is_set assert route.called assert response.text == "four score" + + +def _make_rate_limit_error(retry_after: str | None = None) -> litellm.RateLimitError: + headers = {"retry-after": retry_after} if retry_after else {} + resp = MagicMock() + resp.headers = headers + resp.status_code = 429 + err = litellm.RateLimitError( + message="Rate limit reached", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 429 + err.response = resp + return err + + +def _make_auth_error() -> litellm.AuthenticationError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 401 + err = litellm.AuthenticationError( + message="Invalid API Key", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 401 + return err + + +def _make_bad_request_error() -> litellm.BadRequestError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 400 + err = litellm.BadRequestError( + message="Invalid argument", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 400 + return err + + +def _make_service_unavailable_error() -> litellm.ServiceUnavailableError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 503 + err = litellm.ServiceUnavailableError( + message="The server is temporarily overloaded", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 503 + return err + + +def _make_internal_server_error() -> litellm.InternalServerError: + resp = MagicMock() + resp.headers = {} + resp.status_code = 500 + err = litellm.InternalServerError( + message="Internal server error", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 500 + return err + + +def test_completion_with_retries_honors_retry_after(): + """Verify completion_with_retries reads Retry-After and applies jitter.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise _make_rate_limit_error(retry_after="2") + return {"status": "ok"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert res == {"status": "ok"} + assert len(calls) == 2 + mock_sleep.assert_called_once() + sleep_arg = mock_sleep.call_args[0][0] + assert 2.0 <= sleep_arg <= 2.8 + + +def test_completion_with_retries_non_transient_auth_error_fails_immediately(): + """Non-transient errors like 401 should fail on attempt 1 without retrying.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with patch("time.sleep") as mock_sleep: + with pytest.raises(litellm.AuthenticationError): + completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +def test_completion_with_retries_bad_request_fails_immediately(): + """Non-transient errors like 400 Bad Request should fail on attempt 1 without retrying.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + raise _make_bad_request_error() + + with patch("time.sleep") as mock_sleep: + with pytest.raises(litellm.BadRequestError): + completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +def test_completion_with_retries_service_unavailable_503(): + """503 Service Unavailable should be retried with exponential backoff.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) < 3: + raise _make_service_unavailable_error() + return {"status": "recovered"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + assert len(calls) == 3 + assert mock_sleep.call_count == 2 + + +def test_completion_with_retries_internal_server_error_500(): + """500 Internal Server Error should be retried.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) < 2: + raise _make_internal_server_error() + return {"status": "recovered"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + assert len(calls) == 2 + mock_sleep.assert_called_once() + + +@pytest.mark.asyncio +async def test_acompletion_with_retries_honors_retry_after(): + """Verify acompletion_with_retries reads Retry-After and applies jitter.""" + calls = [] + + async def mock_acompletion(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise _make_rate_limit_error(retry_after="3") + return {"status": "ok"} + + with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + res = await acompletion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_acompletion, + ) + assert res == {"status": "ok"} + assert len(calls) == 2 + mock_sleep.assert_called_once() + sleep_arg = mock_sleep.call_args[0][0] + assert 3.0 <= sleep_arg <= 3.8 + + +@pytest.mark.asyncio +async def test_acompletion_with_retries_non_transient_auth_error_fails_immediately(): + """Async 401 should fail on attempt 1 without retrying.""" + calls = [] + + async def mock_acompletion(*args, **kwargs): + calls.append(1) + raise _make_auth_error() + + with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: + with pytest.raises(litellm.AuthenticationError): + await acompletion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_acompletion, + ) + assert len(calls) == 1 + mock_sleep.assert_not_called() + + +def test_completion_with_retries_connection_error_with_constant_retry(): + """Verify APIConnectionError (no status_code) retries and constant_retry waits 0.0s.""" + calls = [] + + def mock_completion(*args, **kwargs): + calls.append(1) + if len(calls) == 1: + raise APIConnectionError(request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions")) + return {"status": "recovered"} + + with patch("time.sleep") as mock_sleep: + res = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + retry_strategy="constant_retry", + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + assert len(calls) == 2 + mock_sleep.assert_called_once_with(0.0) + + +_RETRY_HOST: Final = "https://retry-upstream.invalid" +_RETRY_RESPONSES_OK: Final = { + "id": "resp-retry", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "test-model", + "output": [ + { + "type": "message", + "id": "msg-retry", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "ok", "annotations": []}], + } + ], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, +} +_RETRY_ANTHROPIC_OK: Final = { + "id": "msg-retry", + "type": "message", + "role": "assistant", + "model": "test-model", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} + + +@dataclass(frozen=True, slots=True) +class _RetryUpstream: + model: str + api_base: str + url: str + ok_body: Mapping[str, object] + + +_RESPONSES_BRIDGE: Final = _RetryUpstream( + "openai/responses/test-model", f"{_RETRY_HOST}/v1", f"{_RETRY_HOST}/v1/responses", _RETRY_RESPONSES_OK +) +_ANTHROPIC: Final = _RetryUpstream( + "anthropic/test-model", _RETRY_HOST, f"{_RETRY_HOST}/v1/messages", _RETRY_ANTHROPIC_OK +) + + +@pytest.fixture +def retry_upstream(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "num_retries", None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + with respx.mock(assert_all_called=False) as upstream: + yield upstream + + +async def _completion_with_retries_kwargs(sync_mode: bool, upstream: _RetryUpstream, **kwargs: object) -> object: + request: Final = { + "model": upstream.model, + "messages": [{"role": "user", "content": "hi"}], + "api_base": upstream.api_base, + "api_key": "test-key", + **kwargs, + } + if sync_mode: + return litellm.completion(**request) + return await litellm.acompletion(**request) + + +def _provider_error(status_code: int, headers: Mapping[str, str] | None = None) -> httpx.Response: + return httpx.Response(status_code, json={"error": {"message": "upstream says no", "type": "x"}}, headers=headers) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("retry_strategy", [None, "constant_retry"]) +@pytest.mark.parametrize("upstream", [_RESPONSES_BRIDGE, _ANTHROPIC], ids=["responses-bridge", "anthropic"]) +async def test_completion_num_retries_waits_for_provider_retry_after( + retry_upstream: respx.MockRouter, sync_mode: bool, retry_strategy: str | None, upstream: _RetryUpstream +): + route: Final = retry_upstream.post(upstream.url).mock( + side_effect=[_provider_error(429, {"retry-after": "7"}), httpx.Response(200, json=upstream.ok_body)] + ) + strategy: Final = {} if retry_strategy is None else {"retry_strategy": retry_strategy} + + with patch("time.sleep") as sync_sleep, patch("asyncio.sleep", new_callable=AsyncMock) as async_sleep: + await _completion_with_retries_kwargs(sync_mode, upstream, num_retries=1, **strategy) + + waits: Final = [call.args[0] for call in (*sync_sleep.call_args_list, *async_sleep.call_args_list) if call.args[0]] + assert route.call_count == 2 + assert len(waits) == 1 + assert 7 <= waits[0] < 8 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + "retry_kwargs", + [{"num_retries": 2}, {"num_retries": 2, "retry_policy": RetryPolicy()}], + ids=["num-retries", "unconfigured-retry-policy"], +) +async def test_completion_num_retries_does_not_retry_auth_error( + retry_upstream: respx.MockRouter, sync_mode: bool, retry_kwargs: Mapping[str, object] +): + route: Final = retry_upstream.post(_RESPONSES_BRIDGE.url).mock(return_value=_provider_error(401)) + + with pytest.raises(litellm.AuthenticationError): + await _completion_with_retries_kwargs(sync_mode, _RESPONSES_BRIDGE, **retry_kwargs) + + assert route.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + ("status_code", "retry_kwargs", "expected_error"), + [ + (500, {"num_retries": 3}, litellm.InternalServerError), + (400, {"retry_policy": RetryPolicy(BadRequestErrorRetries=3)}, litellm.BadRequestError), + ], + ids=["transient-500", "retry-policy-400"], +) +async def test_completion_num_retries_exhausts_configured_attempts( + retry_upstream: respx.MockRouter, + sync_mode: bool, + status_code: int, + retry_kwargs: Mapping[str, object], + expected_error: type[Exception], +): + route: Final = retry_upstream.post(_RESPONSES_BRIDGE.url).mock(return_value=_provider_error(status_code)) + + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(expected_error): + await _completion_with_retries_kwargs(sync_mode, _RESPONSES_BRIDGE, **retry_kwargs) + + assert route.call_count == 4 + + +@pytest.mark.asyncio +async def test_acompletion_keeps_provider_error_when_tenacity_is_missing( + retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + import sys + + monkeypatch.setitem(sys.modules, "tenacity", None) + route: Final = retry_upstream.post(_RESPONSES_BRIDGE.url).mock( + return_value=_provider_error(429, {"retry-after": "7"}) + ) + + with pytest.raises(litellm.RateLimitError): + await _completion_with_retries_kwargs(False, _RESPONSES_BRIDGE, num_retries=1) + + assert route.call_count == 1 From 77a06ff96b0710d4d091dac30ab90118b6d338a0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 00:15:30 -0700 Subject: [PATCH 03/20] fix(retries): drop empty retry policies, keep num_retries when the policy has no count for the error, scope policy retries to the first error's class, and read whole-number float num_retries --- litellm/main.py | 29 +++++++++++++++--------- litellm/utils.py | 49 +++++++++++++++++------------------------ tests/unit/test_main.py | 42 ++++++++++++++++++++++++++++++++--- 3 files changed, 78 insertions(+), 42 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index d5d9db8b0ff..bf6ec4d6f7f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6403,11 +6403,16 @@ def _compute_retry_wait( return float(retry_after) if retry_after is not None and 0 < retry_after <= 60 else 0.0 -def _build_retry_predicate(retry_policy_override: bool) -> Callable[[BaseException], bool]: +_NUM_RETRIES_ADAPTER: Final = TypeAdapter(int) + + +def _build_retry_predicate( + retry_policy_exception_type: type[BaseException] | None, +) -> Callable[[BaseException], bool]: """Check if an exception is transient or explicitly enabled via RetryPolicy.""" def _predicate(e: BaseException) -> bool: - if retry_policy_override and isinstance(e, openai.APIError): + if retry_policy_exception_type is not None and isinstance(e, retry_policy_exception_type): return True status_code: Final[object] = getattr(e, "status_code", None) if isinstance(status_code, int): @@ -6442,8 +6447,7 @@ def completion_with_retries(*args, **kwargs): except Exception as e: raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") - raw_retries: Final[object] = kwargs.pop("num_retries", 3) - num_retries: Final[int] = raw_retries if isinstance(raw_retries, int) else 3 + num_retries: Final[int] = _NUM_RETRIES_ADAPTER.validate_python(kwargs.pop("num_retries", 3)) # reset retries in .completion() kwargs["max_retries"] = 0 kwargs["num_retries"] = 0 @@ -6452,7 +6456,10 @@ def completion_with_retries(*args, **kwargs): ) original_function: Final = kwargs.pop("original_function", completion) initial_exc: Final[object] = kwargs.pop("_initial_exception", None) - retry_policy_override: Final[bool] = kwargs.pop("retry_policy", None) is not None + retry_policy: Final[object] = kwargs.pop("retry_policy", None) + retry_policy_exception_type: Final[type[BaseException] | None] = ( + type(initial_exc) if retry_policy is not None and isinstance(initial_exc, BaseException) else None + ) if isinstance(initial_exc, BaseException): wait_seconds: Final[float] = _compute_retry_wait( exc=initial_exc, @@ -6464,7 +6471,7 @@ def completion_with_retries(*args, **kwargs): retryer: Final = tenacity.Retrying( wait=_build_retry_wait_fn(num_retries, retry_strategy), stop=tenacity.stop_after_attempt(num_retries), - retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_override)), + retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_exception_type)), reraise=True, ) return retryer(original_function, *args, **kwargs) @@ -6480,8 +6487,7 @@ async def acompletion_with_retries(*args, **kwargs): except Exception as e: raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") - raw_retries: Final[object] = kwargs.pop("num_retries", 3) - num_retries: Final[int] = raw_retries if isinstance(raw_retries, int) else 3 + num_retries: Final[int] = _NUM_RETRIES_ADAPTER.validate_python(kwargs.pop("num_retries", 3)) kwargs["max_retries"] = 0 kwargs["num_retries"] = 0 retry_strategy: Final[str] = ( @@ -6489,7 +6495,10 @@ async def acompletion_with_retries(*args, **kwargs): ) original_function: Final = kwargs.pop("original_function", acompletion) initial_exc: Final[object] = kwargs.pop("_initial_exception", None) - retry_policy_override: Final[bool] = kwargs.pop("retry_policy", None) is not None + retry_policy: Final[object] = kwargs.pop("retry_policy", None) + retry_policy_exception_type: Final[type[BaseException] | None] = ( + type(initial_exc) if retry_policy is not None and isinstance(initial_exc, BaseException) else None + ) if isinstance(initial_exc, BaseException): wait_seconds: Final[float] = _compute_retry_wait( exc=initial_exc, @@ -6501,7 +6510,7 @@ async def acompletion_with_retries(*args, **kwargs): retryer: Final = tenacity.AsyncRetrying( wait=_build_retry_wait_fn(num_retries, retry_strategy), stop=tenacity.stop_after_attempt(num_retries), - retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_override)), + retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_exception_type)), reraise=True, ) return await retryer(original_function, *args, **kwargs) diff --git a/litellm/utils.py b/litellm/utils.py index 764a2b2abca..108b4fc07ba 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1550,22 +1550,21 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu num_retries = kwargs.get("num_retries", None) if num_retries is None: num_retries = litellm.num_retries - if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy: Final[Callable[..., int | None]] = ( - litellm_utils.get_num_retries_from_retry_policy - ) - reset_retry_policy: Final = litellm_utils.reset_retry_policy - retry_policy_num_retries: Final[int | None] = get_num_retries_from_retry_policy( - exception=exception, - retry_policy=kwargs.get("retry_policy"), - ) - if retry_policy_num_retries is not None: - num_retries = retry_policy_num_retries - kwargs["retry_policy"] = reset_retry_policy() - else: - kwargs.pop("retry_policy", None) - - return num_retries, kwargs + retry_policy: Final[object] = kwargs.pop("retry_policy", None) + if not retry_policy: + return num_retries, kwargs + get_num_retries_from_retry_policy: Final[Callable[..., int | None]] = ( + litellm_utils.get_num_retries_from_retry_policy + ) + retry_policy_num_retries: Final[int | None] = get_num_retries_from_retry_policy( + exception=exception, + retry_policy=retry_policy, + ) + if retry_policy_num_retries is None: + return num_retries, kwargs + reset_retry_policy: Final = litellm_utils.reset_retry_policy + kwargs["retry_policy"] = reset_retry_policy() + return retry_policy_num_retries, kwargs def _get_wrapper_timeout(kwargs: dict[str, object]) -> object: @@ -2058,17 +2057,7 @@ def client(original_function): return result except Exception as e: if call_type == CallTypes.completion.value: - num_retries = kwargs.get("num_retries", None) or litellm.num_retries or None - if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy: Callable[..., int | None] = ( - litellm_utils.get_num_retries_from_retry_policy - ) - reset_retry_policy = litellm_utils.reset_retry_policy - num_retries = get_num_retries_from_retry_policy( - exception=e, - retry_policy=kwargs.get("retry_policy"), - ) - kwargs["retry_policy"] = reset_retry_policy() # prevent infinite loops + num_retries, kwargs = _get_wrapper_num_retries(kwargs=kwargs, exception=e) litellm.num_retries = None # set retries to None to prevent infinite loops context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {}) @@ -2096,9 +2085,11 @@ def client(original_function): elif call_type == CallTypes.responses.value: num_retries = kwargs.get("num_retries", None) or litellm.num_retries or None if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy = litellm_utils.get_num_retries_from_retry_policy + get_responses_num_retries: Final[Callable[..., int | None]] = ( + litellm_utils.get_num_retries_from_retry_policy + ) reset_retry_policy = litellm_utils.reset_retry_policy - num_retries = get_num_retries_from_retry_policy( + num_retries = get_responses_num_retries( exception=e, retry_policy=kwargs.get("retry_policy"), ) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 888877ce603..b5a04c9ce54 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -9975,8 +9975,13 @@ async def test_completion_num_retries_waits_for_provider_retry_after( @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.parametrize( "retry_kwargs", - [{"num_retries": 2}, {"num_retries": 2, "retry_policy": RetryPolicy()}], - ids=["num-retries", "unconfigured-retry-policy"], + [ + {"num_retries": 2}, + {"num_retries": 2, "retry_policy": RetryPolicy()}, + {"num_retries": 2, "retry_policy": None}, + {"num_retries": 2, "retry_policy": {}}, + ], + ids=["num-retries", "unconfigured-retry-policy", "none-retry-policy", "empty-retry-policy"], ) async def test_completion_num_retries_does_not_retry_auth_error( retry_upstream: respx.MockRouter, sync_mode: bool, retry_kwargs: Mapping[str, object] @@ -9996,8 +10001,9 @@ async def test_completion_num_retries_does_not_retry_auth_error( [ (500, {"num_retries": 3}, litellm.InternalServerError), (400, {"retry_policy": RetryPolicy(BadRequestErrorRetries=3)}, litellm.BadRequestError), + (429, {"num_retries": 3, "retry_policy": RetryPolicy(BadRequestErrorRetries=1)}, litellm.RateLimitError), ], - ids=["transient-500", "retry-policy-400"], + ids=["transient-500", "retry-policy-400", "retry-policy-for-another-error-keeps-num-retries"], ) async def test_completion_num_retries_exhausts_configured_attempts( retry_upstream: respx.MockRouter, @@ -10015,6 +10021,36 @@ async def test_completion_num_retries_exhausts_configured_attempts( assert route.call_count == 4 +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_completion_whole_float_num_retries_retries_that_many_times(retry_upstream: respx.MockRouter, sync_mode: bool): + route: Final = retry_upstream.post(_RESPONSES_BRIDGE.url).mock(return_value=_provider_error(500)) + + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(litellm.InternalServerError): + await _completion_with_retries_kwargs(sync_mode, _RESPONSES_BRIDGE, num_retries=2.0) + + assert route.call_count == 3 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_completion_retry_policy_count_only_retries_its_own_error_class( + retry_upstream: respx.MockRouter, sync_mode: bool +): + route: Final = retry_upstream.post(_RESPONSES_BRIDGE.url).mock( + side_effect=[_provider_error(429), _provider_error(401), httpx.Response(200, json=_RETRY_RESPONSES_OK)] + ) + + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(litellm.AuthenticationError): + await _completion_with_retries_kwargs( + sync_mode, _RESPONSES_BRIDGE, retry_policy=RetryPolicy(RateLimitErrorRetries=2) + ) + + assert route.call_count == 2 + + @pytest.mark.asyncio async def test_acompletion_keeps_provider_error_when_tenacity_is_missing( retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch From 0552a336a25bb53104d834a43c5a3ab783d1263b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 01:44:40 -0700 Subject: [PATCH 04/20] fix(retries): honor a retry_policy passed straight to completion_with_retries, and reset a global num_retries before the async status check as sync does --- litellm/main.py | 43 +++++++++++++++++++++--- litellm/utils.py | 2 +- tests/unit/test_main.py | 74 +++++++++++++++++++++++++++++++++++++++-- 3 files changed, 111 insertions(+), 8 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 79c78181bce..7d91d6d494b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -139,6 +139,7 @@ from litellm.llms.vertex_ai.common_utils import ( ) from litellm.realtime_api.main import realtime_health_check from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway +from litellm.router_utils.get_retry_from_policy import get_num_retries_from_retry_policy from litellm.router_utils.reasoning_effort_capability import ( model_group_from_metadata, reasoning_effort_needs_allowing, @@ -151,7 +152,7 @@ from litellm.types.completion import ( ) from litellm.types.litellm_params import ControlOptions, RetryStrategy from litellm.types.llms.openai import AllMessageValues -from litellm.types.router import GenericLiteLLMParams +from litellm.types.router import GenericLiteLLMParams, RetryPolicy from litellm.types.utils import ( BaseUrlRoutingHint, CustomPricingLiteLLMParams, @@ -6405,16 +6406,26 @@ def _compute_retry_wait( _NUM_RETRIES_ADAPTER: Final = TypeAdapter(int) +_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy) + + +def _policy_retries_for(exc: BaseException | None, retry_policy: RetryPolicy | None) -> int | None: + if retry_policy is None or not isinstance(exc, Exception): + return None + return get_num_retries_from_retry_policy(exception=exc, retry_policy=retry_policy) def _build_retry_predicate( retry_policy_exception_type: type[BaseException] | None, + direct_retry_policy: RetryPolicy | None, ) -> Callable[[BaseException], bool]: """Check if an exception is transient or explicitly enabled via RetryPolicy.""" def _predicate(e: BaseException) -> bool: if retry_policy_exception_type is not None and isinstance(e, retry_policy_exception_type): return True + if (_policy_retries_for(e, direct_retry_policy) or 0) > 0: + return True status_code: Final[object] = getattr(e, "status_code", None) if isinstance(status_code, int): return should_retry(status_code) @@ -6423,6 +6434,18 @@ def _build_retry_predicate( return _predicate +def _build_retry_stop(num_retries: int, direct_retry_policy: RetryPolicy | None) -> Callable[[_RetryState], bool]: + def _stop(retry_state: _RetryState) -> bool: + outcome: Final[_RetryOutcome | None] = retry_state.outcome + policy_retries: Final[int | None] = _policy_retries_for( + outcome.exception() if outcome is not None else None, direct_retry_policy + ) + max_attempts: Final[int] = num_retries if policy_retries is None else policy_retries + 1 + return retry_state.attempt_number >= max_attempts + + return _stop + + def _build_retry_wait_fn(num_retries: int, retry_strategy: str) -> Callable[[_RetryState], float]: """Tenacity wait callback honoring provider Retry-After headers and jitter via calculate_retry_after.""" @@ -6461,6 +6484,11 @@ def completion_with_retries(*args, **kwargs): retry_policy_exception_type: Final[type[BaseException] | None] = ( type(initial_exc) if retry_policy is not None and isinstance(initial_exc, BaseException) else None ) + direct_retry_policy: Final[RetryPolicy | None] = ( + None + if isinstance(initial_exc, BaseException) or not retry_policy + else _RETRY_POLICY_ADAPTER.validate_python(retry_policy) + ) if isinstance(initial_exc, BaseException): wait_seconds: Final[float] = _compute_retry_wait( exc=initial_exc, @@ -6471,8 +6499,8 @@ def completion_with_retries(*args, **kwargs): time.sleep(wait_seconds) retryer: Final = tenacity.Retrying( wait=_build_retry_wait_fn(num_retries, retry_strategy), - stop=tenacity.stop_after_attempt(num_retries), - retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_exception_type)), + stop=_build_retry_stop(num_retries, direct_retry_policy), + retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_exception_type, direct_retry_policy)), reraise=True, ) return retryer(original_function, *args, **kwargs) @@ -6500,6 +6528,11 @@ async def acompletion_with_retries(*args, **kwargs): retry_policy_exception_type: Final[type[BaseException] | None] = ( type(initial_exc) if retry_policy is not None and isinstance(initial_exc, BaseException) else None ) + direct_retry_policy: Final[RetryPolicy | None] = ( + None + if isinstance(initial_exc, BaseException) or not retry_policy + else _RETRY_POLICY_ADAPTER.validate_python(retry_policy) + ) if isinstance(initial_exc, BaseException): wait_seconds: Final[float] = _compute_retry_wait( exc=initial_exc, @@ -6510,8 +6543,8 @@ async def acompletion_with_retries(*args, **kwargs): await asyncio.sleep(wait_seconds) retryer: Final = tenacity.AsyncRetrying( wait=_build_retry_wait_fn(num_retries, retry_strategy), - stop=tenacity.stop_after_attempt(num_retries), - retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_exception_type)), + stop=_build_retry_stop(num_retries, direct_retry_policy), + retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_exception_type, direct_retry_policy)), reraise=True, ) return await retryer(original_function, *args, **kwargs) diff --git a/litellm/utils.py b/litellm/utils.py index cea3b41ae92..b1f689fda98 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2381,11 +2381,11 @@ def client(original_function): if ( num_retries and not is_acompletion_litellm_router_call and not isinstance(e, ImportError) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying + litellm.num_retries = None # set retries to None to prevent infinite loops status_code: Final[object] = getattr(e, "status_code", None) if ( "retry_policy" in kwargs or not isinstance(status_code, int) or should_retry(status_code) ) and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)): - litellm.num_retries = None # set retries to None to prevent infinite loops kwargs["num_retries"] = num_retries kwargs["original_function"] = original_function try: diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index b598784d4e1..aaa7d8222ae 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -14,7 +14,7 @@ from dataclasses import dataclass from datetime import datetime from importlib import import_module from pathlib import Path -from typing import Final +from typing import Final, NoReturn from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -9855,6 +9855,53 @@ async def test_acompletion_with_retries_non_transient_auth_error_fails_immediate mock_sleep.assert_not_called() +@pytest.mark.parametrize( + "retry_policy, make_error, expected_calls", + [ + pytest.param(RetryPolicy(BadRequestErrorRetries=3), _make_bad_request_error, 4, id="count-retries-its-class"), + pytest.param({"BadRequestErrorRetries": 1}, _make_bad_request_error, 2, id="dict-policy"), + pytest.param(RetryPolicy(DefaultRetries=2), _make_auth_error, 3, id="default-retries-any-class"), + pytest.param(RetryPolicy(AuthenticationErrorRetries=2), _make_bad_request_error, 1, id="count-for-another-class"), + pytest.param( + RetryPolicy(RateLimitErrorRetries=0), + lambda: _make_rate_limit_error(retry_after="1"), + 1, + id="zero-count-stops-a-transient-error", + ), + ], +) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_completion_with_retries_honors_a_directly_passed_retry_policy( + retry_policy: RetryPolicy | Mapping[str, int], make_error: Callable[[], Exception], expected_calls: int, sync_mode: bool +): + def raise_error(*args: object, **kwargs: object) -> NoReturn: + raise make_error() + + mock_completion: Final = MagicMock(side_effect=raise_error) + + async def mock_acompletion(*args: object, **kwargs: object) -> None: + mock_completion(*args, **kwargs) + + request: Final = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "num_retries": 3, + "retry_policy": retry_policy, + } + + async def call_with_retries() -> None: + if sync_mode: + completion_with_retries(original_function=mock_completion, **request) + else: + await acompletion_with_retries(original_function=mock_acompletion, **request) + + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(type(make_error())): + await call_with_retries() + assert mock_completion.call_count == expected_calls + + def test_completion_with_retries_connection_error_with_constant_retry(): """Verify APIConnectionError (no status_code) retries and constant_retry waits 0.0s.""" calls = [] @@ -10023,7 +10070,9 @@ async def test_completion_num_retries_exhausts_configured_attempts( @pytest.mark.asyncio @pytest.mark.parametrize("sync_mode", [True, False]) -async def test_completion_whole_float_num_retries_retries_that_many_times(retry_upstream: respx.MockRouter, sync_mode: bool): +async def test_completion_whole_float_num_retries_retries_that_many_times( + retry_upstream: respx.MockRouter, sync_mode: bool +): route: Final = retry_upstream.post(_RESPONSES_BRIDGE.url).mock(return_value=_provider_error(500)) with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): @@ -10051,6 +10100,27 @@ async def test_completion_retry_policy_count_only_retries_its_own_error_class( assert route.call_count == 2 +@pytest.mark.asyncio +async def test_global_num_retries_is_spent_the_same_way_by_sync_and_async_calls( + retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + async def requests_sent(sync_mode: bool) -> tuple[int, int]: + monkeypatch.setattr(litellm, "num_retries", 2) + route: Final = retry_upstream.post(_ANTHROPIC.url).mock( + side_effect=[_provider_error(401), *(_provider_error(500) for _ in range(3))] + ) + route.calls.clear() + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(litellm.AuthenticationError): + await _completion_with_retries_kwargs(sync_mode, _ANTHROPIC) + auth_requests: Final = route.call_count + with pytest.raises(litellm.InternalServerError): + await _completion_with_retries_kwargs(sync_mode, _ANTHROPIC) + return auth_requests, route.call_count - auth_requests + + assert await requests_sent(sync_mode=True) == await requests_sent(sync_mode=False) == (1, 1) + + @pytest.mark.asyncio async def test_acompletion_keeps_provider_error_when_tenacity_is_missing( retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch From ceb3a6f2d560fab63f14cb8619ca0c63c8d47bdd Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 04:05:15 -0700 Subject: [PATCH 05/20] fix(retries): keep the context window fallback and a policy's later-error counts, log async retry successes, and round up fractional num_retries --- litellm/main.py | 207 +++++++++++++++++++++++++--------------- litellm/utils.py | 37 +++---- tests/unit/test_main.py | 187 ++++++++++++++++++++++++++++-------- 3 files changed, 296 insertions(+), 135 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 7d91d6d494b..d3f598d07f0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -14,15 +14,17 @@ import contextvars import datetime import inspect import json +import math import os import random import sys import time import traceback -from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence from concurrent import futures from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from copy import deepcopy +from dataclasses import dataclass from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args @@ -6405,26 +6407,24 @@ def _compute_retry_wait( return float(retry_after) if retry_after is not None and 0 < retry_after <= 60 else 0.0 -_NUM_RETRIES_ADAPTER: Final = TypeAdapter(int) +_NUM_RETRIES_ADAPTER: Final = TypeAdapter(float) _RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy) -def _policy_retries_for(exc: BaseException | None, retry_policy: RetryPolicy | None) -> int | None: - if retry_policy is None or not isinstance(exc, Exception): - return None - return get_num_retries_from_retry_policy(exception=exc, retry_policy=retry_policy) - - def _build_retry_predicate( retry_policy_exception_type: type[BaseException] | None, - direct_retry_policy: RetryPolicy | None, + retry_policy: RetryPolicy | None, ) -> Callable[[BaseException], bool]: """Check if an exception is transient or explicitly enabled via RetryPolicy.""" def _predicate(e: BaseException) -> bool: if retry_policy_exception_type is not None and isinstance(e, retry_policy_exception_type): return True - if (_policy_retries_for(e, direct_retry_policy) or 0) > 0: + if ( + retry_policy is not None + and isinstance(e, Exception) + and (get_num_retries_from_retry_policy(exception=e, retry_policy=retry_policy) or 0) > 0 + ): return True status_code: Final[object] = getattr(e, "status_code", None) if isinstance(status_code, int): @@ -6434,16 +6434,93 @@ def _build_retry_predicate( return _predicate -def _build_retry_stop(num_retries: int, direct_retry_policy: RetryPolicy | None) -> Callable[[_RetryState], bool]: - def _stop(retry_state: _RetryState) -> bool: - outcome: Final[_RetryOutcome | None] = retry_state.outcome - policy_retries: Final[int | None] = _policy_retries_for( - outcome.exception() if outcome is not None else None, direct_retry_policy - ) - max_attempts: Final[int] = num_retries if policy_retries is None else policy_retries + 1 - return retry_state.attempt_number >= max_attempts +@dataclass(frozen=True, slots=True) +class _RetryPlan: + num_retries: int + retry_policy: RetryPolicy | None + initial_exception: BaseException | None - return _stop + @property + def retry_policy_exception_type(self) -> type[BaseException] | None: + if self.retry_policy is None or self.initial_exception is None: + return None + return type(self.initial_exception) + + +def _retry_plan_after_first_error( + first_error: Exception, num_retries: int, retry_policy: RetryPolicy +) -> _RetryPlan | None: + """The retries a direct retry_policy grants after the first error, kept for the whole call like the router's.""" + policy_retries: Final[int | None] = get_num_retries_from_retry_policy( + exception=first_error, retry_policy=retry_policy + ) + if policy_retries is not None: + if policy_retries <= 0: + return None + return _RetryPlan(num_retries=policy_retries, retry_policy=retry_policy, initial_exception=first_error) + if num_retries <= 1 or not _build_retry_predicate(None, None)(first_error): + return None + return _RetryPlan(num_retries=num_retries - 1, retry_policy=None, initial_exception=first_error) + + +def _run_completion_retries( + original_function: Callable[..., object], + args: tuple[object, ...], + kwargs: Mapping[str, object], + plan: _RetryPlan, + retry_strategy: str, +) -> object: + try: + import tenacity + except Exception as e: + raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") + + if plan.initial_exception is not None: + time.sleep( + _compute_retry_wait( + exc=plan.initial_exception, + attempt_number=0, + num_retries=plan.num_retries, + retry_strategy=retry_strategy, + ) + ) + retryer: Final = tenacity.Retrying( + wait=_build_retry_wait_fn(plan.num_retries, retry_strategy), + stop=tenacity.stop_after_attempt(plan.num_retries), + retry=tenacity.retry_if_exception(_build_retry_predicate(plan.retry_policy_exception_type, plan.retry_policy)), + reraise=True, + ) + return retryer(original_function, *args, **kwargs) + + +async def _run_acompletion_retries( + original_function: Callable[..., Awaitable[object]], + args: tuple[object, ...], + kwargs: Mapping[str, object], + plan: _RetryPlan, + retry_strategy: str, +) -> object: + try: + import tenacity + except Exception as e: + raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") + + if plan.initial_exception is not None: + await asyncio.sleep( + _compute_retry_wait( + exc=plan.initial_exception, + attempt_number=0, + num_retries=plan.num_retries, + retry_strategy=retry_strategy, + ) + ) + retryer: Final = tenacity.AsyncRetrying( + wait=_build_retry_wait_fn(plan.num_retries, retry_strategy), + stop=tenacity.stop_after_attempt(plan.num_retries), + retry=tenacity.retry_if_exception(_build_retry_predicate(plan.retry_policy_exception_type, plan.retry_policy)), + reraise=True, + ) + return await retryer(original_function, *args, **kwargs) def _build_retry_wait_fn(num_retries: int, retry_strategy: str) -> Callable[[_RetryState], float]: @@ -6462,92 +6539,72 @@ def _build_retry_wait_fn(num_retries: int, retry_strategy: str) -> Callable[[_Re return _wait_fn -def completion_with_retries(*args, **kwargs): +def completion_with_retries( + *args: object, original_function: Callable[..., object] = completion, **kwargs: object +) -> object: """ Executes a litellm.completion() with retries """ - try: - import tenacity - except Exception as e: - raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") - - num_retries: Final[int] = _NUM_RETRIES_ADAPTER.validate_python(kwargs.pop("num_retries", 3)) + num_retries: Final[int] = math.ceil(_NUM_RETRIES_ADAPTER.validate_python(kwargs.pop("num_retries", 3))) # reset retries in .completion() kwargs["max_retries"] = 0 kwargs["num_retries"] = 0 retry_strategy: Final[str] = ( "constant_retry" if kwargs.pop("retry_strategy", None) == "constant_retry" else "exponential_backoff_retry" ) - original_function: Final = kwargs.pop("original_function", completion) initial_exc: Final[object] = kwargs.pop("_initial_exception", None) retry_policy: Final[object] = kwargs.pop("retry_policy", None) - retry_policy_exception_type: Final[type[BaseException] | None] = ( - type(initial_exc) if retry_policy is not None and isinstance(initial_exc, BaseException) else None + validated_retry_policy: Final[RetryPolicy | None] = ( + _RETRY_POLICY_ADAPTER.validate_python(retry_policy) if retry_policy else None ) - direct_retry_policy: Final[RetryPolicy | None] = ( - None - if isinstance(initial_exc, BaseException) or not retry_policy - else _RETRY_POLICY_ADAPTER.validate_python(retry_policy) - ) - if isinstance(initial_exc, BaseException): - wait_seconds: Final[float] = _compute_retry_wait( - exc=initial_exc, - attempt_number=0, + if isinstance(initial_exc, BaseException) or validated_retry_policy is None: + initial_plan: Final = _RetryPlan( num_retries=num_retries, - retry_strategy=retry_strategy, + retry_policy=validated_retry_policy, + initial_exception=initial_exc if isinstance(initial_exc, BaseException) else None, ) - time.sleep(wait_seconds) - retryer: Final = tenacity.Retrying( - wait=_build_retry_wait_fn(num_retries, retry_strategy), - stop=_build_retry_stop(num_retries, direct_retry_policy), - retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_exception_type, direct_retry_policy)), - reraise=True, - ) - return retryer(original_function, *args, **kwargs) + return _run_completion_retries(original_function, args, kwargs, initial_plan, retry_strategy) + try: + return original_function(*args, **kwargs) + except Exception as first_error: + plan: Final = _retry_plan_after_first_error(first_error, num_retries, validated_retry_policy) + if plan is None: + raise + return _run_completion_retries(original_function, args, kwargs, plan, retry_strategy) -async def acompletion_with_retries(*args, **kwargs): +async def acompletion_with_retries( + *args: object, original_function: Callable[..., Awaitable[object]] = acompletion, **kwargs: object +) -> object: """ [DEPRECATED]. Use 'acompletion' or router.acompletion instead! Executes a litellm.acompletion() with retries """ - try: - import tenacity - except Exception as e: - raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}") - - num_retries: Final[int] = _NUM_RETRIES_ADAPTER.validate_python(kwargs.pop("num_retries", 3)) + num_retries: Final[int] = math.ceil(_NUM_RETRIES_ADAPTER.validate_python(kwargs.pop("num_retries", 3))) kwargs["max_retries"] = 0 kwargs["num_retries"] = 0 retry_strategy: Final[str] = ( "constant_retry" if kwargs.pop("retry_strategy", None) == "constant_retry" else "exponential_backoff_retry" ) - original_function: Final = kwargs.pop("original_function", acompletion) initial_exc: Final[object] = kwargs.pop("_initial_exception", None) retry_policy: Final[object] = kwargs.pop("retry_policy", None) - retry_policy_exception_type: Final[type[BaseException] | None] = ( - type(initial_exc) if retry_policy is not None and isinstance(initial_exc, BaseException) else None + validated_retry_policy: Final[RetryPolicy | None] = ( + _RETRY_POLICY_ADAPTER.validate_python(retry_policy) if retry_policy else None ) - direct_retry_policy: Final[RetryPolicy | None] = ( - None - if isinstance(initial_exc, BaseException) or not retry_policy - else _RETRY_POLICY_ADAPTER.validate_python(retry_policy) - ) - if isinstance(initial_exc, BaseException): - wait_seconds: Final[float] = _compute_retry_wait( - exc=initial_exc, - attempt_number=0, + if isinstance(initial_exc, BaseException) or validated_retry_policy is None: + initial_plan: Final = _RetryPlan( num_retries=num_retries, - retry_strategy=retry_strategy, + retry_policy=validated_retry_policy, + initial_exception=initial_exc if isinstance(initial_exc, BaseException) else None, ) - await asyncio.sleep(wait_seconds) - retryer: Final = tenacity.AsyncRetrying( - wait=_build_retry_wait_fn(num_retries, retry_strategy), - stop=_build_retry_stop(num_retries, direct_retry_policy), - retry=tenacity.retry_if_exception(_build_retry_predicate(retry_policy_exception_type, direct_retry_policy)), - reraise=True, - ) - return await retryer(original_function, *args, **kwargs) + return await _run_acompletion_retries(original_function, args, kwargs, initial_plan, retry_strategy) + try: + return await original_function(*args, **kwargs) + except Exception as first_error: + plan: Final = _retry_plan_after_first_error(first_error, num_retries, validated_retry_policy) + if plan is None: + raise + return await _run_acompletion_retries(original_function, args, kwargs, plan, retry_strategy) def responses_with_retries(*args, **kwargs): diff --git a/litellm/utils.py b/litellm/utils.py index b1f689fda98..bb1b5221c50 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1562,8 +1562,7 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu ) if retry_policy_num_retries is None: return num_retries, kwargs - reset_retry_policy: Final = litellm_utils.reset_retry_policy - kwargs["retry_policy"] = reset_retry_policy() + kwargs["retry_policy"] = retry_policy return retry_policy_num_retries, kwargs @@ -2071,7 +2070,7 @@ def client(original_function): ) and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)): kwargs["num_retries"] = num_retries return litellm.completion_with_retries(*args, _initial_exception=e, **kwargs) - elif ( + if ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict and model in context_window_fallback_dict @@ -2378,22 +2377,24 @@ def client(original_function): is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) - if ( - num_retries and not is_acompletion_litellm_router_call and not isinstance(e, ImportError) - ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying + retries_apply: Final = ( + bool(num_retries) and not is_acompletion_litellm_router_call and not isinstance(e, ImportError) + ) # router has its own logic for retrying + if retries_apply: litellm.num_retries = None # set retries to None to prevent infinite loops - status_code: Final[object] = getattr(e, "status_code", None) - if ( - "retry_policy" in kwargs or not isinstance(status_code, int) or should_retry(status_code) - ) and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)): - kwargs["num_retries"] = num_retries - kwargs["original_function"] = original_function - try: - return await litellm.acompletion_with_retries(*args, _initial_exception=e, **kwargs) - except openai.APIError as retry_err: - e = retry_err - except Exception as retry_unavailable: - verbose_logger.debug("acompletion retries did not run: %s", retry_unavailable) + status_code: Final[object] = getattr(e, "status_code", None) + if ( + retries_apply + and ("retry_policy" in kwargs or not isinstance(status_code, int) or should_retry(status_code)) + and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)) + ): + kwargs["num_retries"] = num_retries + try: + return await litellm.acompletion_with_retries(*args, _initial_exception=e, **kwargs) + except openai.APIError as retry_err: + e = retry_err + except Exception as retry_unavailable: + verbose_logger.debug("acompletion retries did not run: %s", retry_unavailable) elif ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index aaa7d8222ae..28e595f7775 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -9630,11 +9630,11 @@ async def test_atranscription_ignores_a_non_string_base_url_when_api_base_is_set def _make_rate_limit_error(retry_after: str | None = None) -> litellm.RateLimitError: - headers = {"retry-after": retry_after} if retry_after else {} - resp = MagicMock() + headers: Final = {"retry-after": retry_after} if retry_after else {} + resp: Final = MagicMock() resp.headers = headers resp.status_code = 429 - err = litellm.RateLimitError( + err: Final = litellm.RateLimitError( message="Rate limit reached", response=resp, llm_provider="openai", @@ -9646,10 +9646,10 @@ def _make_rate_limit_error(retry_after: str | None = None) -> litellm.RateLimitE def _make_auth_error() -> litellm.AuthenticationError: - resp = MagicMock() + resp: Final = MagicMock() resp.headers = {} resp.status_code = 401 - err = litellm.AuthenticationError( + err: Final = litellm.AuthenticationError( message="Invalid API Key", response=resp, llm_provider="openai", @@ -9660,10 +9660,10 @@ def _make_auth_error() -> litellm.AuthenticationError: def _make_bad_request_error() -> litellm.BadRequestError: - resp = MagicMock() + resp: Final = MagicMock() resp.headers = {} resp.status_code = 400 - err = litellm.BadRequestError( + err: Final = litellm.BadRequestError( message="Invalid argument", response=resp, llm_provider="openai", @@ -9674,10 +9674,10 @@ def _make_bad_request_error() -> litellm.BadRequestError: def _make_service_unavailable_error() -> litellm.ServiceUnavailableError: - resp = MagicMock() + resp: Final = MagicMock() resp.headers = {} resp.status_code = 503 - err = litellm.ServiceUnavailableError( + err: Final = litellm.ServiceUnavailableError( message="The server is temporarily overloaded", response=resp, llm_provider="openai", @@ -9688,10 +9688,10 @@ def _make_service_unavailable_error() -> litellm.ServiceUnavailableError: def _make_internal_server_error() -> litellm.InternalServerError: - resp = MagicMock() + resp: Final = MagicMock() resp.headers = {} resp.status_code = 500 - err = litellm.InternalServerError( + err: Final = litellm.InternalServerError( message="Internal server error", response=resp, llm_provider="openai", @@ -9703,16 +9703,16 @@ def _make_internal_server_error() -> litellm.InternalServerError: def test_completion_with_retries_honors_retry_after(): """Verify completion_with_retries reads Retry-After and applies jitter.""" - calls = [] + calls: Final[list[int]] = [] - def mock_completion(*args, **kwargs): + def mock_completion(*args: object, **kwargs: object) -> dict[str, str]: calls.append(1) if len(calls) == 1: raise _make_rate_limit_error(retry_after="2") return {"status": "ok"} with patch("time.sleep") as mock_sleep: - res = completion_with_retries( + res: Final = completion_with_retries( model="gpt-4o", messages=[{"role": "user", "content": "hi"}], num_retries=3, @@ -9721,15 +9721,15 @@ def test_completion_with_retries_honors_retry_after(): assert res == {"status": "ok"} assert len(calls) == 2 mock_sleep.assert_called_once() - sleep_arg = mock_sleep.call_args[0][0] + sleep_arg: Final = mock_sleep.call_args[0][0] assert 2.0 <= sleep_arg <= 2.8 def test_completion_with_retries_non_transient_auth_error_fails_immediately(): """Non-transient errors like 401 should fail on attempt 1 without retrying.""" - calls = [] + calls: Final[list[int]] = [] - def mock_completion(*args, **kwargs): + def mock_completion(*args: object, **kwargs: object) -> NoReturn: calls.append(1) raise _make_auth_error() @@ -9747,9 +9747,9 @@ def test_completion_with_retries_non_transient_auth_error_fails_immediately(): def test_completion_with_retries_bad_request_fails_immediately(): """Non-transient errors like 400 Bad Request should fail on attempt 1 without retrying.""" - calls = [] + calls: Final[list[int]] = [] - def mock_completion(*args, **kwargs): + def mock_completion(*args: object, **kwargs: object) -> NoReturn: calls.append(1) raise _make_bad_request_error() @@ -9767,16 +9767,16 @@ def test_completion_with_retries_bad_request_fails_immediately(): def test_completion_with_retries_service_unavailable_503(): """503 Service Unavailable should be retried with exponential backoff.""" - calls = [] + calls: Final[list[int]] = [] - def mock_completion(*args, **kwargs): + def mock_completion(*args: object, **kwargs: object) -> dict[str, str]: calls.append(1) if len(calls) < 3: raise _make_service_unavailable_error() return {"status": "recovered"} with patch("time.sleep") as mock_sleep: - res = completion_with_retries( + res: Final = completion_with_retries( model="gpt-4o", messages=[{"role": "user", "content": "hi"}], num_retries=3, @@ -9789,16 +9789,16 @@ def test_completion_with_retries_service_unavailable_503(): def test_completion_with_retries_internal_server_error_500(): """500 Internal Server Error should be retried.""" - calls = [] + calls: Final[list[int]] = [] - def mock_completion(*args, **kwargs): + def mock_completion(*args: object, **kwargs: object) -> dict[str, str]: calls.append(1) if len(calls) < 2: raise _make_internal_server_error() return {"status": "recovered"} with patch("time.sleep") as mock_sleep: - res = completion_with_retries( + res: Final = completion_with_retries( model="gpt-4o", messages=[{"role": "user", "content": "hi"}], num_retries=2, @@ -9812,16 +9812,16 @@ def test_completion_with_retries_internal_server_error_500(): @pytest.mark.asyncio async def test_acompletion_with_retries_honors_retry_after(): """Verify acompletion_with_retries reads Retry-After and applies jitter.""" - calls = [] + calls: Final[list[int]] = [] - async def mock_acompletion(*args, **kwargs): + async def mock_acompletion(*args: object, **kwargs: object) -> dict[str, str]: calls.append(1) if len(calls) == 1: raise _make_rate_limit_error(retry_after="3") return {"status": "ok"} with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: - res = await acompletion_with_retries( + res: Final = await acompletion_with_retries( model="gpt-4o", messages=[{"role": "user", "content": "hi"}], num_retries=3, @@ -9830,16 +9830,16 @@ async def test_acompletion_with_retries_honors_retry_after(): assert res == {"status": "ok"} assert len(calls) == 2 mock_sleep.assert_called_once() - sleep_arg = mock_sleep.call_args[0][0] + sleep_arg: Final = mock_sleep.call_args[0][0] assert 3.0 <= sleep_arg <= 3.8 @pytest.mark.asyncio async def test_acompletion_with_retries_non_transient_auth_error_fails_immediately(): """Async 401 should fail on attempt 1 without retrying.""" - calls = [] + calls: Final[list[int]] = [] - async def mock_acompletion(*args, **kwargs): + async def mock_acompletion(*args: object, **kwargs: object) -> NoReturn: calls.append(1) raise _make_auth_error() @@ -9902,18 +9902,50 @@ async def test_completion_with_retries_honors_a_directly_passed_retry_policy( assert mock_completion.call_count == expected_calls +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_completion_with_retries_keeps_the_first_errors_retry_policy_budget(sync_mode: bool): + errors: Final = iter((_make_bad_request_error(), *(_make_service_unavailable_error() for _ in range(4)))) + + def raise_next_error(*args: object, **kwargs: object) -> NoReturn: + raise next(errors) + + mock_completion: Final = MagicMock(side_effect=raise_next_error) + + async def mock_acompletion(*args: object, **kwargs: object) -> None: + mock_completion(*args, **kwargs) + + request: Final = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "num_retries": 1, + "retry_policy": RetryPolicy(BadRequestErrorRetries=3), + } + + async def call_with_retries() -> None: + if sync_mode: + completion_with_retries(original_function=mock_completion, **request) + else: + await acompletion_with_retries(original_function=mock_acompletion, **request) + + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(litellm.ServiceUnavailableError): + await call_with_retries() + assert mock_completion.call_count == 4 + + def test_completion_with_retries_connection_error_with_constant_retry(): """Verify APIConnectionError (no status_code) retries and constant_retry waits 0.0s.""" - calls = [] + calls: Final[list[int]] = [] - def mock_completion(*args, **kwargs): + def mock_completion(*args: object, **kwargs: object) -> dict[str, str]: calls.append(1) if len(calls) == 1: raise APIConnectionError(request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions")) return {"status": "recovered"} with patch("time.sleep") as mock_sleep: - res = completion_with_retries( + res: Final = completion_with_retries( model="gpt-4o", messages=[{"role": "user", "content": "hi"}], num_retries=2, @@ -9926,7 +9958,7 @@ def test_completion_with_retries_connection_error_with_constant_retry(): _RETRY_HOST: Final = "https://retry-upstream.invalid" -_RETRY_RESPONSES_OK: Final = { +_RETRY_RESPONSES_OK: Final[Mapping[str, object]] = { "id": "resp-retry", "object": "response", "created_at": 1, @@ -9943,7 +9975,7 @@ _RETRY_RESPONSES_OK: Final = { ], "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, } -_RETRY_ANTHROPIC_OK: Final = { +_RETRY_ANTHROPIC_OK: Final[Mapping[str, object]] = { "id": "msg-retry", "type": "message", "role": "assistant", @@ -10070,16 +10102,17 @@ async def test_completion_num_retries_exhausts_configured_attempts( @pytest.mark.asyncio @pytest.mark.parametrize("sync_mode", [True, False]) -async def test_completion_whole_float_num_retries_retries_that_many_times( - retry_upstream: respx.MockRouter, sync_mode: bool +@pytest.mark.parametrize(("num_retries", "expected_requests"), [(2.0, 3), (2.5, 4)], ids=["whole", "fractional"]) +async def test_completion_float_num_retries_retries_the_rounded_up_count( + retry_upstream: respx.MockRouter, sync_mode: bool, num_retries: float, expected_requests: int ): - route: Final = retry_upstream.post(_RESPONSES_BRIDGE.url).mock(return_value=_provider_error(500)) + route: Final = retry_upstream.post(_ANTHROPIC.url).mock(return_value=_provider_error(500)) with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): with pytest.raises(litellm.InternalServerError): - await _completion_with_retries_kwargs(sync_mode, _RESPONSES_BRIDGE, num_retries=2.0) + await _completion_with_retries_kwargs(sync_mode, _ANTHROPIC, num_retries=num_retries) - assert route.call_count == 3 + assert route.call_count == expected_requests @pytest.mark.asyncio @@ -10100,6 +10133,74 @@ async def test_completion_retry_policy_count_only_retries_its_own_error_class( assert route.call_count == 2 +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_completion_num_retries_keeps_the_context_window_fallback( + retry_upstream: respx.MockRouter, sync_mode: bool +): + def prompt_too_long_for_the_small_model(request: httpx.Request) -> httpx.Response: + if json.loads(request.content)["model"] == "test-model-large": + return httpx.Response(200, json=_ANTHROPIC.ok_body) + return httpx.Response( + 400, + json={ + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "prompt is too long: 250000 tokens > 200000 maximum", + }, + }, + ) + + route: Final = retry_upstream.post(_ANTHROPIC.url).mock(side_effect=prompt_too_long_for_the_small_model) + + response: Final = await _completion_with_retries_kwargs( + sync_mode, + _ANTHROPIC, + num_retries=2, + context_window_fallback_dict={_ANTHROPIC.model: "anthropic/test-model-large"}, + ) + + assert isinstance(response, litellm.ModelResponse) + assert response.choices[0].message.content == "ok" + assert [json.loads(call.request.content)["model"] for call in route.calls] == ["test-model", "test-model-large"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_completion_retry_policy_retries_a_later_error_its_policy_names( + retry_upstream: respx.MockRouter, sync_mode: bool +): + route: Final = retry_upstream.post(_ANTHROPIC.url).mock( + side_effect=[_provider_error(400), _provider_error(401), httpx.Response(200, json=_ANTHROPIC.ok_body)] + ) + + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + response: Final = await _completion_with_retries_kwargs( + sync_mode, + _ANTHROPIC, + retry_policy=RetryPolicy(BadRequestErrorRetries=3, AuthenticationErrorRetries=3), + ) + + assert isinstance(response, litellm.ModelResponse) + assert route.call_count == 3 + + +@pytest.mark.asyncio +async def test_acompletion_logs_success_when_a_retry_succeeds( + retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + recorder: Final = _SuccessEventRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + retry_upstream.post(_ANTHROPIC.url).mock( + side_effect=[_provider_error(500), httpx.Response(200, json=_ANTHROPIC.ok_body)] + ) + + await _completion_with_retries_kwargs(False, _ANTHROPIC, num_retries=1, retry_strategy="constant_retry") + + assert (await _wait_for_success_event(recorder, "acompletion")).call_type == "acompletion" + + @pytest.mark.asyncio async def test_global_num_retries_is_spent_the_same_way_by_sync_and_async_calls( retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch @@ -10118,7 +10219,9 @@ async def test_global_num_retries_is_spent_the_same_way_by_sync_and_async_calls( await _completion_with_retries_kwargs(sync_mode, _ANTHROPIC) return auth_requests, route.call_count - auth_requests - assert await requests_sent(sync_mode=True) == await requests_sent(sync_mode=False) == (1, 1) + sync_requests: Final = await requests_sent(sync_mode=True) + assert sync_requests[0] == 1 + assert sync_requests == await requests_sent(sync_mode=False) @pytest.mark.asyncio From 6f4dac7e835b58f7a5217d1bc0d87dea2886790a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 06:02:26 -0700 Subject: [PATCH 06/20] fix(retries): run deployment hooks per async retry, log only the final outcome, fix positional context window fallback Async retries now re-enter acompletion with the caller's original request, so deployment hooks such as code interpreter interception run on each attempt instead of seeing kwargs they already rewrote. A retry error is raised with the same num_retries and timeout context as before and skips the outer failure logging, so callbacks see only the outcome the caller gets. The context window fallback no longer assigns into the args tuple, so a positional model works in sync and async calls. A retry_policy count now decides every error class it covers, so a zero count on a subclass stops it after its parent class was retried. --- litellm/main.py | 36 ++--- litellm/utils.py | 121 ++++++++++------- tests/unit/test_main.py | 290 +++++++++++++++++++++++++++++++--------- 3 files changed, 308 insertions(+), 139 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index d3f598d07f0..c11e6f6349e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6411,21 +6411,17 @@ _NUM_RETRIES_ADAPTER: Final = TypeAdapter(float) _RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy) -def _build_retry_predicate( - retry_policy_exception_type: type[BaseException] | None, - retry_policy: RetryPolicy | None, -) -> Callable[[BaseException], bool]: - """Check if an exception is transient or explicitly enabled via RetryPolicy.""" +def _build_retry_predicate(retry_policy: RetryPolicy | None) -> Callable[[BaseException], bool]: + """A RetryPolicy count decides for the errors it covers, and other errors retry only when transient.""" def _predicate(e: BaseException) -> bool: - if retry_policy_exception_type is not None and isinstance(e, retry_policy_exception_type): - return True - if ( - retry_policy is not None - and isinstance(e, Exception) - and (get_num_retries_from_retry_policy(exception=e, retry_policy=retry_policy) or 0) > 0 - ): - return True + policy_retries: Final[int | None] = ( + get_num_retries_from_retry_policy(exception=e, retry_policy=retry_policy) + if retry_policy is not None and isinstance(e, Exception) + else None + ) + if policy_retries is not None: + return policy_retries > 0 status_code: Final[object] = getattr(e, "status_code", None) if isinstance(status_code, int): return should_retry(status_code) @@ -6440,12 +6436,6 @@ class _RetryPlan: retry_policy: RetryPolicy | None initial_exception: BaseException | None - @property - def retry_policy_exception_type(self) -> type[BaseException] | None: - if self.retry_policy is None or self.initial_exception is None: - return None - return type(self.initial_exception) - def _retry_plan_after_first_error( first_error: Exception, num_retries: int, retry_policy: RetryPolicy @@ -6458,9 +6448,9 @@ def _retry_plan_after_first_error( if policy_retries <= 0: return None return _RetryPlan(num_retries=policy_retries, retry_policy=retry_policy, initial_exception=first_error) - if num_retries <= 1 or not _build_retry_predicate(None, None)(first_error): + if num_retries <= 1 or not _build_retry_predicate(None)(first_error): return None - return _RetryPlan(num_retries=num_retries - 1, retry_policy=None, initial_exception=first_error) + return _RetryPlan(num_retries=num_retries - 1, retry_policy=retry_policy, initial_exception=first_error) def _run_completion_retries( @@ -6487,7 +6477,7 @@ def _run_completion_retries( retryer: Final = tenacity.Retrying( wait=_build_retry_wait_fn(plan.num_retries, retry_strategy), stop=tenacity.stop_after_attempt(plan.num_retries), - retry=tenacity.retry_if_exception(_build_retry_predicate(plan.retry_policy_exception_type, plan.retry_policy)), + retry=tenacity.retry_if_exception(_build_retry_predicate(plan.retry_policy)), reraise=True, ) return retryer(original_function, *args, **kwargs) @@ -6517,7 +6507,7 @@ async def _run_acompletion_retries( retryer: Final = tenacity.AsyncRetrying( wait=_build_retry_wait_fn(plan.num_retries, retry_strategy), stop=tenacity.stop_after_attempt(plan.num_retries), - retry=tenacity.retry_if_exception(_build_retry_predicate(plan.retry_policy_exception_type, plan.retry_policy)), + retry=tenacity.retry_if_exception(_build_retry_predicate(plan.retry_policy)), reraise=True, ) return await retryer(original_function, *args, **kwargs) diff --git a/litellm/utils.py b/litellm/utils.py index bb1b5221c50..c02dfec9069 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1553,20 +1553,22 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu retry_policy: Final[object] = kwargs.pop("retry_policy", None) if not retry_policy: return num_retries, kwargs - get_num_retries_from_retry_policy: Final[Callable[..., int | None]] = ( - litellm_utils.get_num_retries_from_retry_policy - ) - retry_policy_num_retries: Final[int | None] = get_num_retries_from_retry_policy( - exception=exception, - retry_policy=retry_policy, - ) + kwargs["retry_policy"] = retry_policy + retry_policy_num_retries: Final[int | None] = _get_retry_policy_num_retries(kwargs, exception) if retry_policy_num_retries is None: return num_retries, kwargs - kwargs["retry_policy"] = retry_policy return retry_policy_num_retries, kwargs -def _get_wrapper_timeout(kwargs: dict[str, object]) -> object: +def _get_retry_policy_num_retries(kwargs: Mapping[str, object], exception: Exception) -> int | None: + get_num_retries_from_retry_policy: Final[Callable[..., int | None]] = ( + litellm_utils.get_num_retries_from_retry_policy + ) + retry_policy: Final[object] = kwargs.get("retry_policy") + return get_num_retries_from_retry_policy(exception=exception, retry_policy=retry_policy) if retry_policy else None + + +def _get_wrapper_timeout(kwargs: Mapping[str, object]) -> object: from litellm.litellm_core_utils.litellm_logging import Logging if kwargs.get("stream") is not True: @@ -1578,6 +1580,13 @@ def _get_wrapper_timeout(kwargs: dict[str, object]) -> object: return kwargs.get("timeout") if stream_timeout is None else stream_timeout +def _add_retry_context(error: Exception, kwargs: Mapping[str, object]) -> None: + deployment_num_retries: Final = kwargs.get("num_retries") + if deployment_num_retries is not None: + setattr(error, "num_retries", deployment_num_retries) + setattr(error, "timeout", _get_wrapper_timeout(kwargs=kwargs)) + + def check_coroutine(value) -> bool: get_coroutine_checker: Final = litellm_utils.get_coroutine_checker return get_coroutine_checker().is_async_callable(value) @@ -2066,7 +2075,9 @@ def client(original_function): ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying status_code: Final[object] = getattr(e, "status_code", None) if ( - "retry_policy" in kwargs or not isinstance(status_code, int) or should_retry(status_code) + _get_retry_policy_num_retries(kwargs, e) is not None + or not isinstance(status_code, int) + or should_retry(status_code) ) and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)): kwargs["num_retries"] = num_retries return litellm.completion_with_retries(*args, _initial_exception=e, **kwargs) @@ -2076,11 +2087,10 @@ def client(original_function): and model in context_window_fallback_dict and not is_completion_litellm_router_call ): + fallback_model: Final = context_window_fallback_dict[model] if len(args) > 0: - args[0] = context_window_fallback_dict[model] - else: - kwargs["model"] = context_window_fallback_dict[model] - return original_function(*args, **kwargs) + return original_function(fallback_model, *args[1:], **kwargs) + return original_function(*args, **{**kwargs, "model": fallback_model}) elif call_type == CallTypes.responses.value: num_retries = kwargs.get("num_retries", None) or litellm.num_retries or None if kwargs.get("retry_policy", None): @@ -2139,6 +2149,7 @@ def client(original_function): kwargs.pop("_is_litellm_internal_call", None) # discard if injected _is_litellm_internal_call: Final = is_internal_call.get() _deployment_call_end_time: datetime.datetime | None = None + request_kwargs: Final = {**kwargs} try: if logging_obj is None: @@ -2353,6 +2364,41 @@ def client(original_function): return result except Exception as e: + num_retries, kwargs = _get_wrapper_num_retries(kwargs=kwargs, exception=e) + is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) + retries_apply: Final = ( + call_type == CallTypes.acompletion.value + and bool(num_retries) + and not is_acompletion_litellm_router_call + and not isinstance(e, ImportError) + ) # router has its own logic for retrying + if retries_apply: + litellm.num_retries = None # set retries to None to prevent infinite loops + status_code: Final[object] = getattr(e, "status_code", None) + acompletion_retried: Final = ( + retries_apply + and ( + _get_retry_policy_num_retries(kwargs, e) is not None + or not isinstance(status_code, int) + or should_retry(status_code) + ) + and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)) + ) + if acompletion_retried: + kwargs["num_retries"] = num_retries + try: + return await litellm.acompletion_with_retries( + *args, + original_function=litellm.acompletion, + _initial_exception=e, + **{**request_kwargs, "num_retries": num_retries}, + ) + except openai.APIError as retry_err: + _add_retry_context(retry_err, kwargs) + raise + except Exception as retry_unavailable: + verbose_logger.debug("acompletion retries did not run: %s", retry_unavailable) + traceback_exception: Final = traceback.format_exc() # Reuse the timestamp taken right when the deployment call itself failed, before # the failure hook ran, so a slow callback doesn't inflate the reported duration. @@ -2370,44 +2416,20 @@ def client(original_function): except Exception as e: raise e - retry_call_type: Final = _get_call_type_from_public_name(original_function.__name__) - num_retries, kwargs = _get_wrapper_num_retries(kwargs=kwargs, exception=e) - if retry_call_type == CallTypes.acompletion.value: + if call_type == CallTypes.acompletion.value: context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {}) - - is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) - - retries_apply: Final = ( - bool(num_retries) and not is_acompletion_litellm_router_call and not isinstance(e, ImportError) - ) # router has its own logic for retrying - if retries_apply: - litellm.num_retries = None # set retries to None to prevent infinite loops - status_code: Final[object] = getattr(e, "status_code", None) if ( - retries_apply - and ("retry_policy" in kwargs or not isinstance(status_code, int) or should_retry(status_code)) - and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)) - ): - kwargs["num_retries"] = num_retries - try: - return await litellm.acompletion_with_retries(*args, _initial_exception=e, **kwargs) - except openai.APIError as retry_err: - e = retry_err - except Exception as retry_unavailable: - verbose_logger.debug("acompletion retries did not run: %s", retry_unavailable) - elif ( - isinstance(e, litellm.exceptions.ContextWindowExceededError) + not acompletion_retried + and isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict and model in context_window_fallback_dict and not is_acompletion_litellm_router_call ): + fallback_model: Final = context_window_fallback_dict[model] if len(args) > 0: - args[0] = context_window_fallback_dict[model] - else: - kwargs["model"] = context_window_fallback_dict[model] - result = await original_function(*args, **kwargs) - return result - elif retry_call_type == CallTypes.aresponses.value: + return await original_function(fallback_model, *args[1:], **kwargs) + return await original_function(*args, **{**kwargs, "model": fallback_model}) + elif call_type == CallTypes.aresponses.value: is_aresponses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) if ( @@ -2427,13 +2449,8 @@ def client(original_function): else: return result - deployment_num_retries: Final = kwargs.get("num_retries") - if deployment_num_retries is not None: - setattr(e, "num_retries", deployment_num_retries) - - timeout: Final = _get_wrapper_timeout(kwargs=kwargs) - setattr(e, "timeout", timeout) - raise e + _add_retry_context(e, kwargs) + raise finally: # Restore trace_id/session_id contextvars to their pre-call value once diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 28e595f7775..d82c1551d75 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -9673,6 +9673,12 @@ def _make_bad_request_error() -> litellm.BadRequestError: return err +def _make_content_policy_error() -> litellm.ContentPolicyViolationError: + return litellm.ContentPolicyViolationError( + message="Output blocked by content filtering policy", model="gpt-4o", llm_provider="openai" + ) + + def _make_service_unavailable_error() -> litellm.ServiceUnavailableError: resp: Final = MagicMock() resp.headers = {} @@ -9703,13 +9709,7 @@ def _make_internal_server_error() -> litellm.InternalServerError: def test_completion_with_retries_honors_retry_after(): """Verify completion_with_retries reads Retry-After and applies jitter.""" - calls: Final[list[int]] = [] - - def mock_completion(*args: object, **kwargs: object) -> dict[str, str]: - calls.append(1) - if len(calls) == 1: - raise _make_rate_limit_error(retry_after="2") - return {"status": "ok"} + mock_completion: Final = MagicMock(side_effect=[_make_rate_limit_error(retry_after="2"), {"status": "ok"}]) with patch("time.sleep") as mock_sleep: res: Final = completion_with_retries( @@ -9719,7 +9719,7 @@ def test_completion_with_retries_honors_retry_after(): original_function=mock_completion, ) assert res == {"status": "ok"} - assert len(calls) == 2 + assert mock_completion.call_count == 2 mock_sleep.assert_called_once() sleep_arg: Final = mock_sleep.call_args[0][0] assert 2.0 <= sleep_arg <= 2.8 @@ -9727,11 +9727,7 @@ def test_completion_with_retries_honors_retry_after(): def test_completion_with_retries_non_transient_auth_error_fails_immediately(): """Non-transient errors like 401 should fail on attempt 1 without retrying.""" - calls: Final[list[int]] = [] - - def mock_completion(*args: object, **kwargs: object) -> NoReturn: - calls.append(1) - raise _make_auth_error() + mock_completion: Final = MagicMock(side_effect=_make_auth_error()) with patch("time.sleep") as mock_sleep: with pytest.raises(litellm.AuthenticationError): @@ -9741,17 +9737,13 @@ def test_completion_with_retries_non_transient_auth_error_fails_immediately(): num_retries=3, original_function=mock_completion, ) - assert len(calls) == 1 + assert mock_completion.call_count == 1 mock_sleep.assert_not_called() def test_completion_with_retries_bad_request_fails_immediately(): """Non-transient errors like 400 Bad Request should fail on attempt 1 without retrying.""" - calls: Final[list[int]] = [] - - def mock_completion(*args: object, **kwargs: object) -> NoReturn: - calls.append(1) - raise _make_bad_request_error() + mock_completion: Final = MagicMock(side_effect=_make_bad_request_error()) with patch("time.sleep") as mock_sleep: with pytest.raises(litellm.BadRequestError): @@ -9761,19 +9753,15 @@ def test_completion_with_retries_bad_request_fails_immediately(): num_retries=3, original_function=mock_completion, ) - assert len(calls) == 1 + assert mock_completion.call_count == 1 mock_sleep.assert_not_called() def test_completion_with_retries_service_unavailable_503(): """503 Service Unavailable should be retried with exponential backoff.""" - calls: Final[list[int]] = [] - - def mock_completion(*args: object, **kwargs: object) -> dict[str, str]: - calls.append(1) - if len(calls) < 3: - raise _make_service_unavailable_error() - return {"status": "recovered"} + mock_completion: Final = MagicMock( + side_effect=[_make_service_unavailable_error(), _make_service_unavailable_error(), {"status": "recovered"}] + ) with patch("time.sleep") as mock_sleep: res: Final = completion_with_retries( @@ -9783,19 +9771,13 @@ def test_completion_with_retries_service_unavailable_503(): original_function=mock_completion, ) assert res == {"status": "recovered"} - assert len(calls) == 3 + assert mock_completion.call_count == 3 assert mock_sleep.call_count == 2 def test_completion_with_retries_internal_server_error_500(): """500 Internal Server Error should be retried.""" - calls: Final[list[int]] = [] - - def mock_completion(*args: object, **kwargs: object) -> dict[str, str]: - calls.append(1) - if len(calls) < 2: - raise _make_internal_server_error() - return {"status": "recovered"} + mock_completion: Final = MagicMock(side_effect=[_make_internal_server_error(), {"status": "recovered"}]) with patch("time.sleep") as mock_sleep: res: Final = completion_with_retries( @@ -9805,20 +9787,14 @@ def test_completion_with_retries_internal_server_error_500(): original_function=mock_completion, ) assert res == {"status": "recovered"} - assert len(calls) == 2 + assert mock_completion.call_count == 2 mock_sleep.assert_called_once() @pytest.mark.asyncio async def test_acompletion_with_retries_honors_retry_after(): """Verify acompletion_with_retries reads Retry-After and applies jitter.""" - calls: Final[list[int]] = [] - - async def mock_acompletion(*args: object, **kwargs: object) -> dict[str, str]: - calls.append(1) - if len(calls) == 1: - raise _make_rate_limit_error(retry_after="3") - return {"status": "ok"} + mock_acompletion: Final = AsyncMock(side_effect=[_make_rate_limit_error(retry_after="3"), {"status": "ok"}]) with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: res: Final = await acompletion_with_retries( @@ -9828,7 +9804,7 @@ async def test_acompletion_with_retries_honors_retry_after(): original_function=mock_acompletion, ) assert res == {"status": "ok"} - assert len(calls) == 2 + assert mock_acompletion.call_count == 2 mock_sleep.assert_called_once() sleep_arg: Final = mock_sleep.call_args[0][0] assert 3.0 <= sleep_arg <= 3.8 @@ -9837,11 +9813,7 @@ async def test_acompletion_with_retries_honors_retry_after(): @pytest.mark.asyncio async def test_acompletion_with_retries_non_transient_auth_error_fails_immediately(): """Async 401 should fail on attempt 1 without retrying.""" - calls: Final[list[int]] = [] - - async def mock_acompletion(*args: object, **kwargs: object) -> NoReturn: - calls.append(1) - raise _make_auth_error() + mock_acompletion: Final = AsyncMock(side_effect=_make_auth_error()) with patch("asyncio.sleep", new_callable=AsyncMock) as mock_sleep: with pytest.raises(litellm.AuthenticationError): @@ -9851,7 +9823,7 @@ async def test_acompletion_with_retries_non_transient_auth_error_fails_immediate num_retries=3, original_function=mock_acompletion, ) - assert len(calls) == 1 + assert mock_acompletion.call_count == 1 mock_sleep.assert_not_called() @@ -9934,15 +9906,73 @@ async def test_completion_with_retries_keeps_the_first_errors_retry_policy_budge assert mock_completion.call_count == 4 +@pytest.mark.parametrize( + "retry_policy, outcomes, expected_error, expected_calls", + [ + pytest.param( + RetryPolicy(AuthenticationErrorRetries=2), + (_make_service_unavailable_error, _make_auth_error, _make_auth_error), + litellm.AuthenticationError, + 3, + id="retries-a-later-error-its-class-names", + ), + pytest.param( + RetryPolicy(RateLimitErrorRetries=0), + (_make_service_unavailable_error, lambda: _make_rate_limit_error(retry_after="1"), lambda: {"ok": True}), + litellm.RateLimitError, + 2, + id="zero-count-stops-a-later-error", + ), + pytest.param( + RetryPolicy(BadRequestErrorRetries=3, ContentPolicyViolationErrorRetries=0), + (_make_bad_request_error, _make_content_policy_error, lambda: {"ok": True}), + litellm.ContentPolicyViolationError, + 2, + id="zero-count-subclass-stops-after-its-parent-class", + ), + ], +) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_completion_with_retries_applies_the_retry_policy_to_later_errors( + retry_policy: RetryPolicy, + outcomes: tuple[Callable[[], object], ...], + expected_error: type[Exception], + expected_calls: int, + sync_mode: bool, +): + mock_completion: Final = MagicMock(side_effect=[make() for make in outcomes]) + + async def mock_acompletion(*args: object, **kwargs: object) -> object: + return mock_completion(*args, **kwargs) + + request: Final = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "hi"}], + "num_retries": 3, + "retry_policy": retry_policy, + } + + async def call_with_retries() -> None: + if sync_mode: + completion_with_retries(original_function=mock_completion, **request) + else: + await acompletion_with_retries(original_function=mock_acompletion, **request) + + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(expected_error): + await call_with_retries() + assert mock_completion.call_count == expected_calls + + def test_completion_with_retries_connection_error_with_constant_retry(): """Verify APIConnectionError (no status_code) retries and constant_retry waits 0.0s.""" - calls: Final[list[int]] = [] - - def mock_completion(*args: object, **kwargs: object) -> dict[str, str]: - calls.append(1) - if len(calls) == 1: - raise APIConnectionError(request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions")) - return {"status": "recovered"} + mock_completion: Final = MagicMock( + side_effect=[ + APIConnectionError(request=httpx.Request("POST", "https://api.openai.com/v1/chat/completions")), + {"status": "recovered"}, + ] + ) with patch("time.sleep") as mock_sleep: res: Final = completion_with_retries( @@ -9953,7 +9983,7 @@ def test_completion_with_retries_connection_error_with_constant_retry(): original_function=mock_completion, ) assert res == {"status": "recovered"} - assert len(calls) == 2 + assert mock_completion.call_count == 2 mock_sleep.assert_called_once_with(0.0) @@ -10135,8 +10165,9 @@ async def test_completion_retry_policy_count_only_retries_its_own_error_class( @pytest.mark.asyncio @pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("positional_model", [False, True], ids=["model-keyword", "model-positional"]) async def test_completion_num_retries_keeps_the_context_window_fallback( - retry_upstream: respx.MockRouter, sync_mode: bool + retry_upstream: respx.MockRouter, sync_mode: bool, positional_model: bool ): def prompt_too_long_for_the_small_model(request: httpx.Request) -> httpx.Response: if json.loads(request.content)["model"] == "test-model-large": @@ -10154,11 +10185,20 @@ async def test_completion_num_retries_keeps_the_context_window_fallback( route: Final = retry_upstream.post(_ANTHROPIC.url).mock(side_effect=prompt_too_long_for_the_small_model) - response: Final = await _completion_with_retries_kwargs( - sync_mode, - _ANTHROPIC, - num_retries=2, - context_window_fallback_dict={_ANTHROPIC.model: "anthropic/test-model-large"}, + request: Final = { + "messages": [{"role": "user", "content": "hi"}], + "api_base": _ANTHROPIC.api_base, + "api_key": "test-key", + "num_retries": 2, + "context_window_fallback_dict": {_ANTHROPIC.model: "anthropic/test-model-large"}, + } + model_args: Final = (_ANTHROPIC.model,) if positional_model else () + model_kwargs: Final = {} if positional_model else {"model": _ANTHROPIC.model} + + response: Final = ( + litellm.completion(*model_args, **model_kwargs, **request) + if sync_mode + else await litellm.acompletion(*model_args, **model_kwargs, **request) ) assert isinstance(response, litellm.ModelResponse) @@ -10201,6 +10241,128 @@ async def test_acompletion_logs_success_when_a_retry_succeeds( assert (await _wait_for_success_event(recorder, "acompletion")).call_type == "acompletion" +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_completion_retry_policy_count_retries_a_later_error_of_its_class( + retry_upstream: respx.MockRouter, sync_mode: bool +): + route: Final = retry_upstream.post(_RESPONSES_BRIDGE.url).mock( + side_effect=[_provider_error(503), *(_provider_error(401) for _ in range(3))] + ) + + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + with pytest.raises(litellm.AuthenticationError): + await _completion_with_retries_kwargs( + sync_mode, _RESPONSES_BRIDGE, num_retries=3, retry_policy=RetryPolicy(AuthenticationErrorRetries=2) + ) + + assert route.call_count == 4 + + +class _CodeInterpreterRewriteHook(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.tool_types_seen: list[tuple[str, ...]] = [] # mutable-ok: test recorder of deployment hook calls + + async def async_pre_call_deployment_hook( + self, kwargs: dict[str, object], call_type: object + ) -> dict[str, object] | None: + tools: Final = kwargs.get("tools") + tool_types: Final = ( + tuple(str(tool.get("type")) for tool in tools if isinstance(tool, dict)) if isinstance(tools, list) else () + ) + self.tool_types_seen.append(tool_types) + if "code_interpreter" not in tool_types: + return None + return { + **kwargs, + "tools": [{"type": "function", "function": {"name": "run_code", "parameters": {"type": "object"}}}], + } + + +@pytest.mark.asyncio +async def test_acompletion_retry_runs_deployment_hooks_on_the_request_as_sent( + retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + hook: Final = _CodeInterpreterRewriteHook() + monkeypatch.setattr(litellm, "callbacks", [hook]) + route: Final = retry_upstream.post(_ANTHROPIC.url).mock( + side_effect=[_provider_error(429, {"retry-after": "1"}), httpx.Response(200, json=_ANTHROPIC.ok_body)] + ) + + with patch("asyncio.sleep", new_callable=AsyncMock): + await _completion_with_retries_kwargs(False, _ANTHROPIC, num_retries=1, tools=[{"type": "code_interpreter"}]) + + assert hook.tool_types_seen == [("code_interpreter",), ("code_interpreter",)] + assert [[tool["name"] for tool in json.loads(call.request.content)["tools"]] for call in route.calls] == [ + ["run_code"], + ["run_code"], + ] + + +class _OutcomeRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.outcomes: list[str] = [] # mutable-ok: test recorder of success and failure callback events + + def _record_failure(self, kwargs: dict[str, object]) -> None: + self.outcomes.append(f"failure:{type(kwargs.get('exception')).__name__}") + + def log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.outcomes.append("success") + + async def async_log_success_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self.outcomes.append("success") + + def log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record_failure(kwargs) + + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record_failure(kwargs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize( + "second_response, expected_outcomes", + [ + pytest.param(lambda: _provider_error(401), ("failure:AuthenticationError",), id="429-then-401"), + pytest.param(lambda: httpx.Response(200, json=_ANTHROPIC.ok_body), ("success",), id="429-then-200"), + ], +) +async def test_retried_completion_logs_only_the_outcome_the_caller_gets( + retry_upstream: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sync_mode: bool, + second_response: Callable[[], httpx.Response], + expected_outcomes: tuple[str, ...], +): + recorder: Final = _OutcomeRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + retry_upstream.post(_ANTHROPIC.url).mock(side_effect=[_provider_error(429, {"retry-after": "1"}), second_response()]) + + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + try: + await _completion_with_retries_kwargs(sync_mode, _ANTHROPIC, num_retries=1) + except litellm.AuthenticationError: + pass + for _ in range(100): + if recorder.outcomes: + break + await asyncio.sleep(0.05) + await asyncio.sleep(0.2) + + assert tuple(recorder.outcomes) == expected_outcomes + + @pytest.mark.asyncio async def test_global_num_retries_is_spent_the_same_way_by_sync_and_async_calls( retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch From 0e762225602d34c8872b70c93e9ca7e6f5194443 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 07:00:56 -0700 Subject: [PATCH 07/20] fix(retries): keep per-request async success callbacks on async retries function_setup no longer pops async callbacks out of the caller's success_callback list, so the async retry replays the caller's list intact and its success log runs them. The retry logging test now waits on the logging futures instead of real sleeps. --- litellm/utils.py | 29 ++++++++++---------- tests/unit/test_main.py | 60 ++++++++++++++++++++++++++++++++++++----- 2 files changed, 67 insertions(+), 22 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index c02dfec9069..6b934700ab2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1087,6 +1087,10 @@ def _ocr_document_summary(document: object) -> str: return kind +def _runs_as_async_success_callback(callback: str | Callable[..., object] | CustomLogger) -> bool: + return check_coroutine(callback) or callback in ("dynamodb", "s3") + + # Runs once per call to check if the user wants to send their data anywhere - PostHog/Sentry/Slack/etc. @@ -1202,20 +1206,15 @@ def function_setup( dynamic_failure_callbacks: list[str | Callable | CustomLogger] | None = None dynamic_async_failure_callbacks: Final[list[str | Callable | CustomLogger] | None] = None if kwargs.get("success_callback", None) is not None and isinstance(kwargs["success_callback"], list): - removed_async_items = [] - for index, callback in enumerate(kwargs["success_callback"]): - if coroutine_checker.is_async_callable(callback) or callback == "dynamodb" or callback == "s3": - if dynamic_async_success_callbacks is not None and isinstance( - dynamic_async_success_callbacks, list - ): - dynamic_async_success_callbacks.append(callback) - else: - dynamic_async_success_callbacks = [callback] - removed_async_items.append(index) - # Pop the async items from success_callback in reverse order to avoid index issues - for index in reversed(removed_async_items): - kwargs["success_callback"].pop(index) - dynamic_success_callbacks = kwargs.pop("success_callback") + requested_success_callbacks: Final[list[str | Callable[..., object] | CustomLogger]] = kwargs.pop( + "success_callback" + ) + dynamic_async_success_callbacks = [ + callback for callback in requested_success_callbacks if _runs_as_async_success_callback(callback) + ] or None + dynamic_success_callbacks = [ + callback for callback in requested_success_callbacks if not _runs_as_async_success_callback(callback) + ] if kwargs.get("failure_callback", None) is not None and isinstance(kwargs["failure_callback"], list): dynamic_failure_callbacks = kwargs.pop("failure_callback") @@ -2371,7 +2370,7 @@ def client(original_function): and bool(num_retries) and not is_acompletion_litellm_router_call and not isinstance(e, ImportError) - ) # router has its own logic for retrying + ) if retries_apply: litellm.num_retries = None # set retries to None to prevent infinite loops status_code: Final[object] = getattr(e, "status_code", None) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index d82c1551d75..5d3678dac87 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -6,10 +6,11 @@ import io import json import logging import os +import queue import threading import urllib.parse -from collections.abc import Callable, Mapping -from concurrent.futures import ThreadPoolExecutor +from collections.abc import Callable, Iterator, Mapping +from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass from datetime import datetime from importlib import import_module @@ -42,7 +43,9 @@ from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs from litellm.litellm_core_utils.get_litellm_params import stored_control_options from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt +from litellm.litellm_core_utils.thread_pool_executor import executor as logging_executor from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.types.litellm_params import ControlOptions from litellm.types.llms.openai import AllMessageValues, HttpxBinaryResponseContent @@ -10042,6 +10045,30 @@ def retry_upstream(monkeypatch: pytest.MonkeyPatch): yield upstream +@pytest.fixture +def logging_futures() -> Iterator[queue.SimpleQueue[Future[object]]]: + futures: Final = queue.SimpleQueue[Future[object]]() + real_submit: Final = logging_executor.submit + + def submit_and_track(fn: Callable[..., object], *args: object) -> Future[object]: + future: Final = real_submit(fn, *args) + futures.put(future) + return future + + with patch.object( # test-quality-ok: wraps the real submit only to collect the futures to join, the pool still runs + logging_executor, "submit", side_effect=submit_and_track + ): + yield futures + + +async def _finish_logging(logging_futures: queue.SimpleQueue[Future[object]]) -> None: + await asyncio.sleep(0) + await GLOBAL_LOGGING_WORKER.flush() + while not logging_futures.empty(): + await asyncio.wrap_future(logging_futures.get()) + await GLOBAL_LOGGING_WORKER.flush() + + async def _completion_with_retries_kwargs(sync_mode: bool, upstream: _RetryUpstream, **kwargs: object) -> object: request: Final = { "model": upstream.model, @@ -10340,6 +10367,7 @@ class _OutcomeRecorder(CustomLogger): ) async def test_retried_completion_logs_only_the_outcome_the_caller_gets( retry_upstream: respx.MockRouter, + logging_futures: queue.SimpleQueue[Future[object]], monkeypatch: pytest.MonkeyPatch, sync_mode: bool, second_response: Callable[[], httpx.Response], @@ -10354,15 +10382,33 @@ async def test_retried_completion_logs_only_the_outcome_the_caller_gets( await _completion_with_retries_kwargs(sync_mode, _ANTHROPIC, num_retries=1) except litellm.AuthenticationError: pass - for _ in range(100): - if recorder.outcomes: - break - await asyncio.sleep(0.05) - await asyncio.sleep(0.2) + await _finish_logging(logging_futures) assert tuple(recorder.outcomes) == expected_outcomes +@pytest.mark.asyncio +async def test_async_retry_runs_the_callers_async_success_callback_without_emptying_its_list( + retry_upstream: respx.MockRouter, logging_futures: queue.SimpleQueue[Future[object]] +): + successes: Final = queue.SimpleQueue[object]() + + async def record_success(kwargs: object, response: object, start_time: object, end_time: object) -> None: + successes.put(response) + + success_callback: Final = [record_success] + retry_upstream.post(_ANTHROPIC.url).mock( + side_effect=[_provider_error(429, {"retry-after": "1"}), httpx.Response(200, json=_ANTHROPIC.ok_body)] + ) + + with patch("asyncio.sleep", new_callable=AsyncMock): + await _completion_with_retries_kwargs(False, _ANTHROPIC, num_retries=1, success_callback=success_callback) + await _finish_logging(logging_futures) + + assert successes.qsize() == 1 + assert success_callback == [record_success] + + @pytest.mark.asyncio async def test_global_num_retries_is_spent_the_same_way_by_sync_and_async_calls( retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch From ee8435be22e490fac7bab911ca56d0077d58c4aa Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 10:55:40 -0700 Subject: [PATCH 08/20] fix(retries): log one outcome per retried async call and keep retry_policy parity on responses and malformed policies --- litellm/utils.py | 91 ++++++++++++++++++++++++++++++----------- tests/unit/test_main.py | 90 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 158 insertions(+), 23 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 6b934700ab2..a13a994cd09 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1540,10 +1540,12 @@ async def _client_async_logging_helper( ) -def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tuple[int | None, dict[str, Any]]: +def _get_wrapper_num_retries( + kwargs: dict[str, Any], exception: Exception +) -> tuple[int | None | Exception, dict[str, Any]]: """ Get the number of retries from the kwargs and the retry policy. - Used for the wrapper functions. + Used for the wrapper functions. A malformed retry policy comes back as its parse error in place of the count. """ num_retries = kwargs.get("num_retries", None) @@ -1553,7 +1555,10 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu if not retry_policy: return num_retries, kwargs kwargs["retry_policy"] = retry_policy - retry_policy_num_retries: Final[int | None] = _get_retry_policy_num_retries(kwargs, exception) + try: + retry_policy_num_retries: Final[int | None] = _get_retry_policy_num_retries(kwargs, exception) + except (AttributeError, TypeError, ValueError) as retry_policy_error: + return retry_policy_error, kwargs if retry_policy_num_retries is None: return num_retries, kwargs return retry_policy_num_retries, kwargs @@ -1586,6 +1591,32 @@ def _add_retry_context(error: Exception, kwargs: Mapping[str, object]) -> None: setattr(error, "timeout", _get_wrapper_timeout(kwargs=kwargs)) +_RETRY_ATTEMPT_KWARG: Final = "_litellm_wrapper_retry_attempt" + + +def _with_retry_policy_spent(kwargs: Mapping[str, object]) -> Mapping[str, object]: + if not kwargs.get("retry_policy"): + return {**kwargs} + return {**kwargs, "retry_policy": litellm_utils.reset_retry_policy()} + + +async def _end_async_wrapper_failure( + logging_obj: LiteLLMLoggingObject, + error: Exception, + traceback_exception: str, + start_time: datetime.datetime, + end_time: datetime.datetime, + log_failure: bool, +) -> None: + unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) + if not log_failure: + return + logging_obj.failure_handler( + error, traceback_exception, start_time, end_time + ) # DO NOT MAKE THREADED - router retry fallback relies on this! + await logging_obj.async_failure_handler(error, traceback_exception, start_time, end_time) + + def check_coroutine(value) -> bool: get_coroutine_checker: Final = litellm_utils.get_coroutine_checker return get_coroutine_checker().is_async_callable(value) @@ -2065,6 +2096,8 @@ def client(original_function): except Exception as e: if call_type == CallTypes.completion.value: num_retries, kwargs = _get_wrapper_num_retries(kwargs=kwargs, exception=e) + if isinstance(num_retries, Exception): + raise num_retries litellm.num_retries = None # set retries to None to prevent infinite loops context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {}) @@ -2087,9 +2120,10 @@ def client(original_function): and not is_completion_litellm_router_call ): fallback_model: Final = context_window_fallback_dict[model] - if len(args) > 0: - return original_function(fallback_model, *args[1:], **kwargs) - return original_function(*args, **{**kwargs, "model": fallback_model}) + fallback_kwargs: Final = _with_retry_policy_spent(kwargs) + if args: + return original_function(fallback_model, *args[1:], **fallback_kwargs) + return original_function(*args, **{**fallback_kwargs, "model": fallback_model}) elif call_type == CallTypes.responses.value: num_retries = kwargs.get("num_retries", None) or litellm.num_retries or None if kwargs.get("retry_policy", None): @@ -2146,6 +2180,7 @@ def client(original_function): model: Final[str | None] = args[0] if len(args) > 0 else kwargs.get("model", None) is_completion_with_fallbacks: Final = kwargs.get("fallbacks") is not None kwargs.pop("_is_litellm_internal_call", None) # discard if injected + is_wrapper_retry_attempt: Final = kwargs.pop(_RETRY_ATTEMPT_KWARG, None) is True _is_litellm_internal_call: Final = is_internal_call.get() _deployment_call_end_time: datetime.datetime | None = None request_kwargs: Final = {**kwargs} @@ -2363,7 +2398,8 @@ def client(original_function): return result except Exception as e: - num_retries, kwargs = _get_wrapper_num_retries(kwargs=kwargs, exception=e) + wrapper_num_retries, kwargs = _get_wrapper_num_retries(kwargs=kwargs, exception=e) + num_retries: Final = None if isinstance(wrapper_num_retries, Exception) else wrapper_num_retries is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) retries_apply: Final = ( call_type == CallTypes.acompletion.value @@ -2390,10 +2426,19 @@ def client(original_function): *args, original_function=litellm.acompletion, _initial_exception=e, - **{**request_kwargs, "num_retries": num_retries}, + **{**request_kwargs, "num_retries": num_retries, _RETRY_ATTEMPT_KWARG: True}, ) except openai.APIError as retry_err: _add_retry_context(retry_err, kwargs) + if logging_obj and not _is_litellm_internal_call: + await _end_async_wrapper_failure( + logging_obj, + retry_err, + traceback.format_exc(), + start_time, + datetime.datetime.now(), # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with + log_failure=True, + ) raise except Exception as retry_unavailable: verbose_logger.debug("acompletion retries did not run: %s", retry_unavailable) @@ -2403,17 +2448,16 @@ def client(original_function): # the failure hook ran, so a slow callback doesn't inflate the reported duration. end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with if logging_obj and not _is_litellm_internal_call: - unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) - try: - logging_obj.failure_handler( - e, traceback_exception, start_time, end_time - ) # DO NOT MAKE THREADED - router retry fallback relies on this! - except Exception as e: - raise e - try: - await logging_obj.async_failure_handler(e, traceback_exception, start_time, end_time) - except Exception as e: - raise e + await _end_async_wrapper_failure( + logging_obj, + e, + traceback_exception, + start_time, + end_time, + log_failure=not is_wrapper_retry_attempt, + ) + if isinstance(wrapper_num_retries, Exception): + raise wrapper_num_retries if call_type == CallTypes.acompletion.value: context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {}) @@ -2425,9 +2469,10 @@ def client(original_function): and not is_acompletion_litellm_router_call ): fallback_model: Final = context_window_fallback_dict[model] - if len(args) > 0: - return await original_function(fallback_model, *args[1:], **kwargs) - return await original_function(*args, **{**kwargs, "model": fallback_model}) + fallback_kwargs: Final = _with_retry_policy_spent(kwargs) + if args: + return await original_function(fallback_model, *args[1:], **fallback_kwargs) + return await original_function(*args, **{**fallback_kwargs, "model": fallback_model}) elif call_type == CallTypes.aresponses.value: is_aresponses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) @@ -2442,7 +2487,7 @@ def client(original_function): kwargs["retry_strategy"] = "exponential_backoff_retry" elif isinstance(e, openai.APIError): # generic api error kwargs["retry_strategy"] = "constant_retry" - result = await litellm.aresponses_with_retries(*args, **kwargs) + result = await litellm.aresponses_with_retries(*args, **_with_retry_policy_spent(kwargs)) except Exception: pass else: diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 5d3678dac87..6fda46b66fd 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -24,6 +24,7 @@ import respx from openai import APIConnectionError, APITimeoutError, AsyncOpenAI from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion +from pydantic import ValidationError import litellm from litellm import ( @@ -10387,6 +10388,95 @@ async def test_retried_completion_logs_only_the_outcome_the_caller_gets( assert tuple(recorder.outcomes) == expected_outcomes +@pytest.mark.asyncio +@pytest.mark.parametrize( + "later_responses, expected_outcomes", + [ + pytest.param( + lambda: [_provider_error(429), httpx.Response(200, json=_ANTHROPIC.ok_body)], + ("success",), + id="429-429-200", + ), + pytest.param( + lambda: [_provider_error(429), _provider_error(401)], + ("failure:AuthenticationError",), + id="429-429-401", + ), + pytest.param( + lambda: [_provider_error(429), _provider_error(429)], + ("failure:RateLimitError",), + id="429-429-429", + ), + ], +) +async def test_async_completion_retried_twice_logs_only_the_outcome_the_caller_gets( + retry_upstream: respx.MockRouter, + logging_futures: queue.SimpleQueue[Future[object]], + monkeypatch: pytest.MonkeyPatch, + later_responses: Callable[[], list[httpx.Response]], + expected_outcomes: tuple[str, ...], +): + recorder: Final = _OutcomeRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + route: Final = retry_upstream.post(_ANTHROPIC.url).mock(side_effect=[_provider_error(429), *later_responses()]) + + with patch("asyncio.sleep", new_callable=AsyncMock): + try: + await _completion_with_retries_kwargs(False, _ANTHROPIC, num_retries=2) + except (litellm.AuthenticationError, litellm.RateLimitError): + pass + await _finish_logging(logging_futures) + + assert route.call_count == 3 + assert tuple(recorder.outcomes) == expected_outcomes + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("retry_policy", "policy_error"), + [("not-a-policy", AttributeError), ({"AuthenticationErrorRetries": "bad"}, ValidationError)], + ids=["str", "bad-count"], +) +async def test_async_completion_with_a_malformed_retry_policy_still_logs_the_failure( + retry_upstream: respx.MockRouter, + logging_futures: queue.SimpleQueue[Future[object]], + monkeypatch: pytest.MonkeyPatch, + retry_policy: object, + policy_error: type[Exception], +): + recorder: Final = _OutcomeRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + route: Final = retry_upstream.post(_ANTHROPIC.url).mock(return_value=_provider_error(401)) + + with pytest.raises(policy_error): + await _completion_with_retries_kwargs(False, _ANTHROPIC, num_retries=1, retry_policy=retry_policy) + await _finish_logging(logging_futures) + + assert route.call_count == 1 + assert tuple(recorder.outcomes) == ("failure:AuthenticationError",) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("num_retries", [None, 2], ids=["policy-only", "policy-and-num-retries"]) +async def test_aresponses_bridged_to_chat_applies_its_retry_policy_only_on_the_first_attempt( + retry_upstream: respx.MockRouter, num_retries: int | None +): + route: Final = retry_upstream.post(_ANTHROPIC.url).mock(return_value=_provider_error(401)) + retry_kwargs: Final = {} if num_retries is None else {"num_retries": num_retries} + + with patch("asyncio.sleep", new_callable=AsyncMock), pytest.raises(litellm.AuthenticationError): + await litellm.aresponses( + model=_ANTHROPIC.model, + input="hi", + api_base=_ANTHROPIC.api_base, + api_key="test-key", + retry_policy=RetryPolicy(AuthenticationErrorRetries=2), + **retry_kwargs, + ) + + assert route.call_count == 5 + + @pytest.mark.asyncio async def test_async_retry_runs_the_callers_async_success_callback_without_emptying_its_list( retry_upstream: respx.MockRouter, logging_futures: queue.SimpleQueue[Future[object]] From 48bb03923c4f627a52141be6d3c7604fba04371b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 11:16:16 -0700 Subject: [PATCH 09/20] fix(retries): mark wrapper retry attempts with a private sentinel so a request body cannot skip failure logging --- litellm/utils.py | 5 +++-- tests/unit/test_main.py | 18 ++++++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index a13a994cd09..8d179029f78 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1592,6 +1592,7 @@ def _add_retry_context(error: Exception, kwargs: Mapping[str, object]) -> None: _RETRY_ATTEMPT_KWARG: Final = "_litellm_wrapper_retry_attempt" +_RETRY_ATTEMPT_MARKER: Final = object() def _with_retry_policy_spent(kwargs: Mapping[str, object]) -> Mapping[str, object]: @@ -2180,7 +2181,7 @@ def client(original_function): model: Final[str | None] = args[0] if len(args) > 0 else kwargs.get("model", None) is_completion_with_fallbacks: Final = kwargs.get("fallbacks") is not None kwargs.pop("_is_litellm_internal_call", None) # discard if injected - is_wrapper_retry_attempt: Final = kwargs.pop(_RETRY_ATTEMPT_KWARG, None) is True + is_wrapper_retry_attempt: Final = kwargs.pop(_RETRY_ATTEMPT_KWARG, None) is _RETRY_ATTEMPT_MARKER _is_litellm_internal_call: Final = is_internal_call.get() _deployment_call_end_time: datetime.datetime | None = None request_kwargs: Final = {**kwargs} @@ -2426,7 +2427,7 @@ def client(original_function): *args, original_function=litellm.acompletion, _initial_exception=e, - **{**request_kwargs, "num_retries": num_retries, _RETRY_ATTEMPT_KWARG: True}, + **{**request_kwargs, "num_retries": num_retries, _RETRY_ATTEMPT_KWARG: _RETRY_ATTEMPT_MARKER}, ) except openai.APIError as retry_err: _add_retry_context(retry_err, kwargs) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 6fda46b66fd..8803dcb4e20 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -10456,6 +10456,24 @@ async def test_async_completion_with_a_malformed_retry_policy_still_logs_the_fai assert tuple(recorder.outcomes) == ("failure:AuthenticationError",) +@pytest.mark.asyncio +async def test_async_completion_with_a_caller_sent_retry_marker_still_logs_the_failure( + retry_upstream: respx.MockRouter, + logging_futures: queue.SimpleQueue[Future[object]], + monkeypatch: pytest.MonkeyPatch, +): + recorder: Final = _OutcomeRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + route: Final = retry_upstream.post(_ANTHROPIC.url).mock(return_value=_provider_error(401)) + + with pytest.raises(litellm.AuthenticationError): + await _completion_with_retries_kwargs(False, _ANTHROPIC, _litellm_wrapper_retry_attempt=True) + await _finish_logging(logging_futures) + + assert route.call_count == 1 + assert tuple(recorder.outcomes) == ("failure:AuthenticationError",) + + @pytest.mark.asyncio @pytest.mark.parametrize("num_retries", [None, 2], ids=["policy-only", "policy-and-num-retries"]) async def test_aresponses_bridged_to_chat_applies_its_retry_policy_only_on_the_first_attempt( From 45f63afc98b9379daf4942ae87f1b70a40f07c5d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 12:39:56 -0700 Subject: [PATCH 10/20] fix(retries): keep a retried async stream's request ids until the caller finishes reading it --- litellm/utils.py | 4 +++- tests/unit/test_main.py | 44 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/litellm/utils.py b/litellm/utils.py index edf195f8cce..f561dc46a02 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2423,7 +2423,7 @@ def client(original_function): if acompletion_retried: kwargs["num_retries"] = num_retries try: - return await litellm.acompletion_with_retries( + result = await litellm.acompletion_with_retries( *args, original_function=litellm.acompletion, _initial_exception=e, @@ -2443,6 +2443,8 @@ def client(original_function): raise except Exception as retry_unavailable: verbose_logger.debug("acompletion retries did not run: %s", retry_unavailable) + else: + return result traceback_exception: Final = traceback.format_exc() # Reuse the timestamp taken right when the deployment call itself failed, before diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 9c819f58aaf..8e6e1288913 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -36,6 +36,7 @@ from litellm import ( responses_with_retries, ) from litellm import main as litellm_main +from litellm._logging import session_id_var, trace_id_var from litellm.caching.base_cache import BaseCache from litellm.caching.caching import Cache from litellm.constants import CONTROL_OPTIONS_KEY @@ -10568,6 +10569,49 @@ async def test_async_completion_with_a_caller_sent_retry_marker_still_logs_the_f assert tuple(recorder.outcomes) == ("failure:AuthenticationError",) +_ANTHROPIC_STREAM_EVENTS: Final = ( + { + "type": "message_start", + "message": {**_RETRY_ANTHROPIC_OK, "content": [], "stop_reason": None}, + }, + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "ok"}}, + {"type": "content_block_stop", "index": 0}, + {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 1}}, + {"type": "message_stop"}, +) + + +@pytest.mark.asyncio +async def test_async_streaming_retry_keeps_the_call_ids_while_the_caller_reads_the_stream( + retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + sse_body: Final = "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in _ANTHROPIC_STREAM_EVENTS) + route: Final = retry_upstream.post(_ANTHROPIC.url).mock( + side_effect=[ + _provider_error(429, {"retry-after": "0"}), + httpx.Response(200, content=sse_body.encode(), headers={"content-type": "text/event-stream"}), + ] + ) + trace_id_var.set("outer-trace") + session_id_var.set("outer-session") + try: + stream: Final = await _completion_with_retries_kwargs( + False, _ANTHROPIC, num_retries=1, stream=True, litellm_session_id="retried-stream-session" + ) + ids_while_reading: Final = (trace_id_var.get(), session_id_var.get()) + text: Final = "".join([chunk.choices[0].delta.content or "" async for chunk in stream]) + finally: + trace_id_var.set("") + session_id_var.set("") + + assert route.call_count == 2 + assert text == "ok" + assert ids_while_reading[1] == "retried-stream-session" + assert ids_while_reading[0] != "outer-trace" + + @pytest.mark.asyncio @pytest.mark.parametrize("num_retries", [None, 2], ids=["policy-only", "policy-and-num-retries"]) async def test_aresponses_bridged_to_chat_applies_its_retry_policy_only_on_the_first_attempt( From d5a1f0d6059a1b0317999e97ab1c5db3928b7d77 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:05:24 -0700 Subject: [PATCH 11/20] test(retries): skip real backoff sleeps in the retried stream ids test --- tests/unit/test_main.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 8e6e1288913..a008c62a002 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -10597,9 +10597,10 @@ async def test_async_streaming_retry_keeps_the_call_ids_while_the_caller_reads_t trace_id_var.set("outer-trace") session_id_var.set("outer-session") try: - stream: Final = await _completion_with_retries_kwargs( - False, _ANTHROPIC, num_retries=1, stream=True, litellm_session_id="retried-stream-session" - ) + with patch("asyncio.sleep", new_callable=AsyncMock): + stream: Final = await _completion_with_retries_kwargs( + False, _ANTHROPIC, num_retries=1, stream=True, litellm_session_id="retried-stream-session" + ) ids_while_reading: Final = (trace_id_var.get(), session_id_var.get()) text: Final = "".join([chunk.choices[0].delta.content or "" async for chunk in stream]) finally: From d66cef757c2302ec080b3ee03d04d2bd6bde0ee9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 13:31:37 -0700 Subject: [PATCH 12/20] fix(retries): restore the caller's request ids when a retried stream ends and keep retrying when a request has an _initial_exception key --- litellm/utils.py | 13 ++++++++++--- tests/unit/test_main.py | 31 +++++++++++++++++++++++++++++-- 2 files changed, 39 insertions(+), 5 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index f561dc46a02..7f7e1f6f7b7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2113,7 +2113,9 @@ def client(original_function): or should_retry(status_code) ) and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)): kwargs["num_retries"] = num_retries - return litellm.completion_with_retries(*args, _initial_exception=e, **kwargs) + return litellm.completion_with_retries( + *args, original_function=litellm.completion, **{**kwargs, "_initial_exception": e} + ) if ( isinstance(e, litellm.exceptions.ContextWindowExceededError) and context_window_fallback_dict @@ -2422,12 +2424,17 @@ def client(original_function): ) if acompletion_retried: kwargs["num_retries"] = num_retries + _restore_correlation_context_if_supported(logging_obj) try: result = await litellm.acompletion_with_retries( *args, original_function=litellm.acompletion, - _initial_exception=e, - **{**request_kwargs, "num_retries": num_retries, _RETRY_ATTEMPT_KWARG: _RETRY_ATTEMPT_MARKER}, + **{ + **request_kwargs, + "num_retries": num_retries, + "_initial_exception": e, + _RETRY_ATTEMPT_KWARG: _RETRY_ATTEMPT_MARKER, + }, ) except openai.APIError as retry_err: _add_retry_context(retry_err, kwargs) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index a008c62a002..ecf4ece9a68 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -53,7 +53,14 @@ from litellm.types.litellm_params import ControlOptions from litellm.types.llms.openai import AllMessageValues, HttpxBinaryResponseContent from litellm.types.prompts.init_prompts import PromptSpec from litellm.types.router import RetryPolicy -from litellm.types.utils import Delta, ModelResponseStream, StandardCallbackDynamicParams, StreamingChoices, Usage +from litellm.types.utils import ( + Delta, + ModelResponse, + ModelResponseStream, + StandardCallbackDynamicParams, + StreamingChoices, + Usage, +) from tests.unit.llms.bedrock.slow_upstream import ( STREAM_TIMEOUT_SECONDS, UPSTREAM_FIRST_BYTE_SECONDS, @@ -10203,6 +10210,24 @@ async def test_completion_num_retries_waits_for_provider_retry_after( assert 7 <= waits[0] < 8 +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_completion_num_retries_still_retries_when_the_caller_sends_an_initial_exception_key( + retry_upstream: respx.MockRouter, sync_mode: bool +): + route: Final = retry_upstream.post(_ANTHROPIC.url).mock( + side_effect=[_provider_error(429, {"retry-after": "0"}), httpx.Response(200, json=_ANTHROPIC.ok_body)] + ) + + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + response: Final = await _completion_with_retries_kwargs( + sync_mode, _ANTHROPIC, num_retries=1, _initial_exception="x" + ) + + assert route.call_count == 2 + assert isinstance(response, ModelResponse) + + @pytest.mark.asyncio @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.parametrize( @@ -10583,7 +10608,7 @@ _ANTHROPIC_STREAM_EVENTS: Final = ( @pytest.mark.asyncio -async def test_async_streaming_retry_keeps_the_call_ids_while_the_caller_reads_the_stream( +async def test_async_streaming_retry_keeps_the_call_ids_while_read_and_restores_the_callers_after( retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ): monkeypatch.setattr(litellm, "request_correlation_in_logs", True) @@ -10603,6 +10628,7 @@ async def test_async_streaming_retry_keeps_the_call_ids_while_the_caller_reads_t ) ids_while_reading: Final = (trace_id_var.get(), session_id_var.get()) text: Final = "".join([chunk.choices[0].delta.content or "" async for chunk in stream]) + ids_after_reading: Final = (trace_id_var.get(), session_id_var.get()) finally: trace_id_var.set("") session_id_var.set("") @@ -10611,6 +10637,7 @@ async def test_async_streaming_retry_keeps_the_call_ids_while_the_caller_reads_t assert text == "ok" assert ids_while_reading[1] == "retried-stream-session" assert ids_while_reading[0] != "outer-trace" + assert ids_after_reading == ("outer-trace", "outer-session") @pytest.mark.asyncio From 3a629c8955ca3661c340d101697cfffeff7aa81c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 14:23:35 -0700 Subject: [PATCH 13/20] fix(retries): start a direct completion_with_retries call's backoff at the initial retry delay --- litellm/main.py | 22 ++++++++++++++-------- tests/unit/test_main.py | 23 ++++++++++++++++++++++- 2 files changed, 36 insertions(+), 9 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 4b801a30e2d..8dbf2a07c81 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6399,7 +6399,7 @@ class _RetryState(Protocol): def _compute_retry_wait( exc: BaseException | None, - attempt_number: int, + retries_sent: int, num_retries: int, retry_strategy: str, ) -> float: @@ -6407,7 +6407,7 @@ def _compute_retry_wait( if retry_strategy == "exponential_backoff_retry": return float( calculate_retry_after( - remaining_retries=max(0, num_retries - attempt_number), + remaining_retries=max(0, num_retries - retries_sent), max_retries=num_retries, response_headers=headers, ) @@ -6478,13 +6478,15 @@ def _run_completion_retries( time.sleep( _compute_retry_wait( exc=plan.initial_exception, - attempt_number=0, + retries_sent=0, num_retries=plan.num_retries, retry_strategy=retry_strategy, ) ) retryer: Final = tenacity.Retrying( - wait=_build_retry_wait_fn(plan.num_retries, retry_strategy), + wait=_build_retry_wait_fn( + plan.num_retries, retry_strategy, first_attempt_is_a_retry=plan.initial_exception is not None + ), stop=tenacity.stop_after_attempt(plan.num_retries), retry=tenacity.retry_if_exception(_build_retry_predicate(plan.retry_policy)), reraise=True, @@ -6508,13 +6510,15 @@ async def _run_acompletion_retries( await asyncio.sleep( _compute_retry_wait( exc=plan.initial_exception, - attempt_number=0, + retries_sent=0, num_retries=plan.num_retries, retry_strategy=retry_strategy, ) ) retryer: Final = tenacity.AsyncRetrying( - wait=_build_retry_wait_fn(plan.num_retries, retry_strategy), + wait=_build_retry_wait_fn( + plan.num_retries, retry_strategy, first_attempt_is_a_retry=plan.initial_exception is not None + ), stop=tenacity.stop_after_attempt(plan.num_retries), retry=tenacity.retry_if_exception(_build_retry_predicate(plan.retry_policy)), reraise=True, @@ -6522,7 +6526,9 @@ async def _run_acompletion_retries( return await retryer(original_function, *args, **kwargs) -def _build_retry_wait_fn(num_retries: int, retry_strategy: str) -> Callable[[_RetryState], float]: +def _build_retry_wait_fn( + num_retries: int, retry_strategy: str, first_attempt_is_a_retry: bool +) -> Callable[[_RetryState], float]: """Tenacity wait callback honoring provider Retry-After headers and jitter via calculate_retry_after.""" def _wait_fn(retry_state: _RetryState) -> float: @@ -6530,7 +6536,7 @@ def _build_retry_wait_fn(num_retries: int, retry_strategy: str) -> Callable[[_Re exc: Final[BaseException | None] = outcome.exception() if outcome is not None else None return _compute_retry_wait( exc=exc, - attempt_number=retry_state.attempt_number, + retries_sent=retry_state.attempt_number if first_attempt_is_a_retry else retry_state.attempt_number - 1, num_retries=num_retries, retry_strategy=retry_strategy, ) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index ecf4ece9a68..ef85fcb1d42 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -39,7 +39,7 @@ from litellm import main as litellm_main from litellm._logging import session_id_var, trace_id_var from litellm.caching.base_cache import BaseCache from litellm.caching.caching import Cache -from litellm.constants import CONTROL_OPTIONS_KEY +from litellm.constants import CONTROL_OPTIONS_KEY, INITIAL_RETRY_DELAY, JITTER from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.custom_prompt_management import CustomPromptManagement from litellm.litellm_core_utils.core_helpers import get_litellm_metadata_from_kwargs @@ -9897,6 +9897,27 @@ def test_completion_with_retries_internal_server_error_500(): mock_sleep.assert_called_once() +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_completion_with_retries_backoff_starts_at_the_initial_retry_delay(sync_mode: bool): + outcomes: Final = [_make_internal_server_error(), _make_internal_server_error(), {"status": "recovered"}] + messages: Final = [{"role": "user", "content": "hi"}] + + with patch("time.sleep") as sync_sleep, patch("asyncio.sleep", new_callable=AsyncMock) as async_sleep: + if sync_mode: + completion_with_retries( + model="gpt-4o", messages=messages, num_retries=3, original_function=MagicMock(side_effect=outcomes) + ) + else: + await acompletion_with_retries( + model="gpt-4o", messages=messages, num_retries=3, original_function=AsyncMock(side_effect=outcomes) + ) + + first_wait, second_wait = [call.args[0] for call in (*sync_sleep.call_args_list, *async_sleep.call_args_list)] + assert INITIAL_RETRY_DELAY <= first_wait < INITIAL_RETRY_DELAY + JITTER + assert 2 * INITIAL_RETRY_DELAY <= second_wait < 2 * INITIAL_RETRY_DELAY + JITTER + + @pytest.mark.asyncio async def test_acompletion_with_retries_honors_retry_after(): """Verify acompletion_with_retries reads Retry-After and applies jitter.""" From 2f6259fe5c82754e3d6b398b57308f4dadcd2166 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 15:41:25 -0700 Subject: [PATCH 14/20] fix(retries): keep a retried call's own ids on every attempt and its failure log, and run an app's wrapper on litellm.completion once --- litellm/litellm_core_utils/litellm_logging.py | 8 +- litellm/utils.py | 24 +++++- tests/unit/test_main.py | 86 +++++++++++++++++++ 3 files changed, 115 insertions(+), 3 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 78ff073fdcb..ad0277e0e7d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -703,7 +703,8 @@ class Logging(LiteLLMLoggingBaseClass): # ids leaking into a different, later request on the same thread. Sync # support is deferred to a follow-up PR with its own safe-restore # mechanism; async calls (the proxy's only call path) are unaffected. - if supports_correlation_logging and litellm.request_correlation_in_logs: + self._stamps_correlation_context: bool = supports_correlation_logging and litellm.request_correlation_in_logs + if self._stamps_correlation_context: set_trace_id(self.litellm_trace_id) set_session_id(self.litellm_session_id) # set_trace_id()/set_session_id() sanitize (strip control chars, bound @@ -3030,6 +3031,11 @@ class Logging(LiteLLMLoggingBaseClass): trace_id_var.set(self._pre_call_trace_id) session_id_var.set(self._pre_call_session_id) + def _stamp_correlation_context(self) -> None: + if self._stamps_correlation_context: + trace_id_var.set(self._own_trace_id) + session_id_var.set(self._own_session_id) + def _restore_correlation_context_if_unclaimed(self) -> None: """Guarded variant for __del__-triggered cleanup only. diff --git a/litellm/utils.py b/litellm/utils.py index 7f7e1f6f7b7..8b94c80e032 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -954,6 +954,12 @@ def _restore_correlation_context_if_supported(logging_obj: object) -> None: restore() +def _stamp_correlation_context_if_supported(logging_obj: object) -> None: + stamp: Final = getattr(logging_obj, "_stamp_correlation_context", None) + if stamp is not None: + stamp() + + def _is_streaming_response_for_correlation(result: object) -> bool: """True if `result` is a lazy stream wrapper rather than an already-complete response. @@ -1612,6 +1618,7 @@ async def _end_async_wrapper_failure( unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) if not log_failure: return + _stamp_correlation_context_if_supported(logging_obj) logging_obj.failure_handler( error, traceback_exception, start_time, end_time ) # DO NOT MAKE THREADED - router retry fallback relies on this! @@ -2113,8 +2120,12 @@ def client(original_function): or should_retry(status_code) ) and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)): kwargs["num_retries"] = num_retries + from litellm.chat_completions.dispatch import completion as public_completion + return litellm.completion_with_retries( - *args, original_function=litellm.completion, **{**kwargs, "_initial_exception": e} + *args, + original_function=public_completion, + **{**kwargs, "_initial_exception": e}, ) if ( isinstance(e, litellm.exceptions.ContextWindowExceededError) @@ -2194,6 +2205,8 @@ def client(original_function): # Type assertion: logging_obj is guaranteed to be non-None after function_setup assert logging_obj is not None, "logging_obj should not be None after function_setup" + if is_wrapper_retry_attempt: + _stamp_correlation_context_if_supported(logging_obj) if not _is_litellm_internal_call: bind_budget_reservation_to_callbacks(logging_obj.litellm_params) @@ -2425,12 +2438,19 @@ def client(original_function): if acompletion_retried: kwargs["num_retries"] = num_retries _restore_correlation_context_if_supported(logging_obj) + from litellm.chat_completions.dispatch import acompletion as public_acompletion + try: result = await litellm.acompletion_with_retries( *args, - original_function=litellm.acompletion, + original_function=public_acompletion, **{ **request_kwargs, + "litellm_trace_id": ( + request_kwargs.get("litellm_trace_id") + if logging_obj is None + else logging_obj.litellm_trace_id + ), "num_retries": num_retries, "_initial_exception": e, _RETRY_ATTEMPT_KWARG: _RETRY_ATTEMPT_MARKER, diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index ef85fcb1d42..f8c6383a9f1 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -10661,6 +10661,92 @@ async def test_async_streaming_retry_keeps_the_call_ids_while_read_and_restores_ assert ids_after_reading == ("outer-trace", "outer-session") +def _drained(ids: queue.SimpleQueue[tuple[str, str]]) -> tuple[tuple[str, str], ...]: + return tuple(ids.get() for _ in range(ids.qsize())) + + +def _with_a_proxy_style_logging_obj(request: Mapping[str, object]) -> Mapping[str, object]: + logging_obj, data = litellm.utils.function_setup( + original_function="acompletion", + rules_obj=litellm.utils.Rules(), + start_time=datetime.now(), + is_async_call=True, + **{**request, "litellm_call_id": "proxy-call-id"}, + ) + return {**data, "litellm_logging_obj": logging_obj} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("caller_logging_obj", [False, True], ids=["sdk", "proxy-style-logging-obj"]) +async def test_async_retried_failure_keeps_the_call_ids_on_every_attempt_and_its_failure_log( + retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, caller_logging_obj: bool +): + monkeypatch.setattr(litellm, "request_correlation_in_logs", True) + attempt_ids: Final = queue.SimpleQueue[tuple[str, str]]() + failure_log_ids: Final = queue.SimpleQueue[tuple[str, str]]() + + def provider_says_retry_later(request: httpx.Request) -> httpx.Response: + attempt_ids.put((trace_id_var.get(), session_id_var.get())) + return _provider_error(429, {"retry-after": "0"}) + + def record_failure_log_ids(kwargs: object, response: object, start_time: object, end_time: object) -> None: + failure_log_ids.put((trace_id_var.get(), session_id_var.get())) + + retry_upstream.post(_ANTHROPIC.url).mock(side_effect=provider_says_retry_later) + trace_id_var.set("outer-trace") + session_id_var.set("outer-session") + try: + request: Final = { + "model": _ANTHROPIC.model, + "messages": [{"role": "user", "content": "hi"}], + "api_base": _ANTHROPIC.api_base, + "api_key": "test-key", + "num_retries": 2, + "litellm_session_id": "call-session", + "failure_callback": [record_failure_log_ids], + } + call_request: Final = _with_a_proxy_style_logging_obj(request) if caller_logging_obj else request + with patch("asyncio.sleep", new_callable=AsyncMock), pytest.raises(litellm.RateLimitError): + await litellm.acompletion(**call_request) + ids_after_the_call: Final = (trace_id_var.get(), session_id_var.get()) + finally: + trace_id_var.set("") + session_id_var.set("") + + attempts: Final = _drained(attempt_ids) + assert len(attempts) == 3 + assert len(set(attempts)) == 1 + assert attempts[0][0] not in ("", "outer-trace") + assert attempts[0][1] == "call-session" + assert _drained(failure_log_ids) == (attempts[0],) + assert ids_after_the_call == ("outer-trace", "outer-session") + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_a_wrapper_rebound_onto_the_public_completion_runs_once_for_a_retried_call( + retry_upstream: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool +): + route: Final = retry_upstream.post(_ANTHROPIC.url).mock( + side_effect=[_provider_error(429, {"retry-after": "0"}), httpx.Response(200, json=_ANTHROPIC.ok_body)] + ) + public_name: Final = "completion" if sync_mode else "acompletion" + public_function: Final = getattr(litellm, public_name) + wrapper_calls: Final = queue.SimpleQueue[object]() + + def instrumented(*args: object, **kwargs: object) -> object: + wrapper_calls.put(kwargs) + return public_function(*args, **kwargs) + + monkeypatch.setattr(litellm, public_name, instrumented) + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + response: Final = await _completion_with_retries_kwargs(sync_mode, _ANTHROPIC, num_retries=1) + + assert route.call_count == 2 + assert isinstance(response, ModelResponse) + assert wrapper_calls.qsize() == 1 + + @pytest.mark.asyncio @pytest.mark.parametrize("num_retries", [None, 2], ids=["policy-only", "policy-and-num-retries"]) async def test_aresponses_bridged_to_chat_applies_its_retry_policy_only_on_the_first_attempt( From 4daf86d0668b656ad58564e18139f73b2f121527 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 16:44:39 -0700 Subject: [PATCH 15/20] fix(retries): run a configured context window fallback before any retry, as the Router does --- litellm/utils.py | 54 ++++++++++++++++++++++++----------------- tests/unit/test_main.py | 22 +++++++++++------ 2 files changed, 46 insertions(+), 30 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 14de7f5b9bf..d8946163e99 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1865,6 +1865,17 @@ def _get_call_type_from_public_name(function_name: str) -> str: return function_name +def _context_window_fallback_applies( + error: Exception, context_window_fallback_dict: object, model: str, *, is_router_call: bool +) -> bool: + return ( + not is_router_call + and isinstance(error, litellm.exceptions.ContextWindowExceededError) + and isinstance(context_window_fallback_dict, dict) + and model in context_window_fallback_dict + ) + + def _is_litellm_router_call(kwargs: Mapping[str, object], *, is_async: bool) -> bool: """Router completion uses metadata. Async generic calls retry with litellm_metadata; sync generic calls need SDK retries.""" metadata_buckets: Final = ( @@ -2111,8 +2122,11 @@ def client(original_function): context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {}) is_completion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False) + context_window_fallback_applies: Final = isinstance(model, str) and _context_window_fallback_applies( + e, context_window_fallback_dict, model, is_router_call=is_completion_litellm_router_call + ) if ( - num_retries and not is_completion_litellm_router_call + num_retries and not is_completion_litellm_router_call and not context_window_fallback_applies ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying status_code: Final[object] = getattr(e, "status_code", None) if ( @@ -2128,12 +2142,7 @@ def client(original_function): original_function=public_completion, **{**kwargs, "_initial_exception": e}, ) - if ( - isinstance(e, litellm.exceptions.ContextWindowExceededError) - and context_window_fallback_dict - and model in context_window_fallback_dict - and not is_completion_litellm_router_call - ): + if context_window_fallback_applies: fallback_model: Final = context_window_fallback_dict[model] fallback_kwargs: Final = _with_retry_policy_spent(kwargs) if args: @@ -2427,8 +2436,17 @@ def client(original_function): if retries_apply: litellm.num_retries = None # set retries to None to prevent infinite loops status_code: Final[object] = getattr(e, "status_code", None) + context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {}) + context_window_fallback_applies: Final = ( + call_type == CallTypes.acompletion.value + and isinstance(model, str) + and _context_window_fallback_applies( + e, context_window_fallback_dict, model, is_router_call=is_acompletion_litellm_router_call + ) + ) acompletion_retried: Final = ( retries_apply + and not context_window_fallback_applies and ( _get_retry_policy_num_retries(kwargs, e) is not None or not isinstance(status_code, int) @@ -2490,21 +2508,13 @@ def client(original_function): if isinstance(wrapper_num_retries, Exception): raise wrapper_num_retries - if call_type == CallTypes.acompletion.value: - context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {}) - if ( - not acompletion_retried - and isinstance(e, litellm.exceptions.ContextWindowExceededError) - and context_window_fallback_dict - and model in context_window_fallback_dict - and not is_acompletion_litellm_router_call - ): - fallback_model: Final = context_window_fallback_dict[model] - fallback_kwargs: Final = _with_retry_policy_spent(kwargs) - if args: - return await original_function(fallback_model, *args[1:], **fallback_kwargs) - return await original_function(*args, **{**fallback_kwargs, "model": fallback_model}) - elif call_type == CallTypes.aresponses.value: + if context_window_fallback_applies: + fallback_model: Final = context_window_fallback_dict[model] + fallback_kwargs: Final = _with_retry_policy_spent(kwargs) + if args: + return await original_function(fallback_model, *args[1:], **fallback_kwargs) + return await original_function(*args, **{**fallback_kwargs, "model": fallback_model}) + if call_type == CallTypes.aresponses.value: is_aresponses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) if ( diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index f8c6383a9f1..3d5f6ac0142 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -10335,8 +10335,13 @@ async def test_completion_retry_policy_count_only_retries_its_own_error_class( @pytest.mark.asyncio @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.parametrize("positional_model", [False, True], ids=["model-keyword", "model-positional"]) -async def test_completion_num_retries_keeps_the_context_window_fallback( - retry_upstream: respx.MockRouter, sync_mode: bool, positional_model: bool +@pytest.mark.parametrize( + "retry_settings", + [{"num_retries": 2}, {"retry_policy": RetryPolicy(BadRequestErrorRetries=3)}], + ids=["num-retries", "retry-policy-for-400s"], +) +async def test_completion_retries_keep_the_context_window_fallback( + retry_upstream: respx.MockRouter, sync_mode: bool, positional_model: bool, retry_settings: Mapping[str, object] ): def prompt_too_long_for_the_small_model(request: httpx.Request) -> httpx.Response: if json.loads(request.content)["model"] == "test-model-large": @@ -10358,17 +10363,18 @@ async def test_completion_num_retries_keeps_the_context_window_fallback( "messages": [{"role": "user", "content": "hi"}], "api_base": _ANTHROPIC.api_base, "api_key": "test-key", - "num_retries": 2, + **retry_settings, "context_window_fallback_dict": {_ANTHROPIC.model: "anthropic/test-model-large"}, } model_args: Final = (_ANTHROPIC.model,) if positional_model else () model_kwargs: Final = {} if positional_model else {"model": _ANTHROPIC.model} - response: Final = ( - litellm.completion(*model_args, **model_kwargs, **request) - if sync_mode - else await litellm.acompletion(*model_args, **model_kwargs, **request) - ) + with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): + response: Final = ( + litellm.completion(*model_args, **model_kwargs, **request) + if sync_mode + else await litellm.acompletion(*model_args, **model_kwargs, **request) + ) assert isinstance(response, litellm.ModelResponse) assert response.choices[0].message.content == "ok" From 9c1672690d6bd49d2ce7422fa7ea931eb028b32c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 17:04:48 -0700 Subject: [PATCH 16/20] fix(retries): accept any mapping as the context window fallback map --- litellm/utils.py | 2 +- tests/unit/test_main.py | 10 ++++++++-- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index d8946163e99..b2e4cb102c2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1871,7 +1871,7 @@ def _context_window_fallback_applies( return ( not is_router_call and isinstance(error, litellm.exceptions.ContextWindowExceededError) - and isinstance(context_window_fallback_dict, dict) + and isinstance(context_window_fallback_dict, Mapping) and model in context_window_fallback_dict ) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 3d5f6ac0142..8a07bc30529 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -15,6 +15,7 @@ from dataclasses import dataclass from datetime import datetime from importlib import import_module from pathlib import Path +from types import MappingProxyType from typing import Final, NoReturn from unittest.mock import AsyncMock, MagicMock, patch @@ -10340,8 +10341,13 @@ async def test_completion_retry_policy_count_only_retries_its_own_error_class( [{"num_retries": 2}, {"retry_policy": RetryPolicy(BadRequestErrorRetries=3)}], ids=["num-retries", "retry-policy-for-400s"], ) +@pytest.mark.parametrize("fallback_map_type", [dict, MappingProxyType], ids=["dict", "read-only-mapping"]) async def test_completion_retries_keep_the_context_window_fallback( - retry_upstream: respx.MockRouter, sync_mode: bool, positional_model: bool, retry_settings: Mapping[str, object] + retry_upstream: respx.MockRouter, + sync_mode: bool, + positional_model: bool, + retry_settings: Mapping[str, object], + fallback_map_type: Callable[[Mapping[str, str]], Mapping[str, str]], ): def prompt_too_long_for_the_small_model(request: httpx.Request) -> httpx.Response: if json.loads(request.content)["model"] == "test-model-large": @@ -10364,7 +10370,7 @@ async def test_completion_retries_keep_the_context_window_fallback( "api_base": _ANTHROPIC.api_base, "api_key": "test-key", **retry_settings, - "context_window_fallback_dict": {_ANTHROPIC.model: "anthropic/test-model-large"}, + "context_window_fallback_dict": fallback_map_type({_ANTHROPIC.model: "anthropic/test-model-large"}), } model_args: Final = (_ANTHROPIC.model,) if positional_model else () model_kwargs: Final = {} if positional_model else {"model": _ANTHROPIC.model} From 4b085928fb28455311d87449c1495830f375c772 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:45:08 -0700 Subject: [PATCH 17/20] fix(retries): hand the budget reservation back before waiting to retry A client that disconnects during the retry wait left the reservation bound to callbacks that never run, so the request-end release skipped it. Unbind before the wait, as main did, and each retried attempt binds it again. --- litellm/main.py | 2 -- litellm/utils.py | 2 ++ tests/unit/test_main.py | 43 +++++++++++++++++++++++++++++++++++++++++ 3 files changed, 45 insertions(+), 2 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 91e8b863291..6e73f5a670a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6530,8 +6530,6 @@ async def _run_acompletion_retries( def _build_retry_wait_fn( num_retries: int, retry_strategy: str, first_attempt_is_a_retry: bool ) -> Callable[[_RetryState], float]: - """Tenacity wait callback honoring provider Retry-After headers and jitter via calculate_retry_after.""" - def _wait_fn(retry_state: _RetryState) -> float: outcome: Final[_RetryOutcome | None] = retry_state.outcome exc: Final[BaseException | None] = outcome.exception() if outcome is not None else None diff --git a/litellm/utils.py b/litellm/utils.py index 40e5d9feffb..3df763db67f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2474,6 +2474,8 @@ def client(original_function): if acompletion_retried: kwargs["num_retries"] = num_retries _restore_correlation_context_if_supported(logging_obj) + if logging_obj and not _is_litellm_internal_call: + unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) from litellm.chat_completions.dispatch import acompletion as public_acompletion try: diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 8a07bc30529..d6c2f4f0ea9 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -10842,6 +10842,49 @@ async def test_acompletion_keeps_provider_error_when_tenacity_is_missing( assert route.call_count == 1 +async def _until(condition: Callable[[], bool]) -> None: + while not condition(): + await asyncio.sleep(0.01) + + +@pytest.mark.asyncio +async def test_acompletion_hands_the_budget_reservation_back_while_it_waits_to_retry( + retry_upstream: respx.MockRouter, +) -> None: + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + route: Final = retry_upstream.post(_ANTHROPIC.url).mock(return_value=_provider_error(429, {"retry-after": "30"})) + call: Final = asyncio.create_task( + _completion_with_retries_kwargs( + False, _ANTHROPIC, num_retries=1, metadata={"user_api_key_budget_reservation": reservation} + ) + ) + + await asyncio.wait_for(_until(lambda: route.call_count == 1), 10) + await asyncio.wait_for(_until(lambda: reservation["callback_bound"] is False), 5) + call.cancel() + with pytest.raises(asyncio.CancelledError): + await call + + assert route.call_count == 1 + assert reservation["callback_bound"] is False + + +@pytest.mark.asyncio +async def test_acompletion_retry_that_succeeds_leaves_the_budget_reservation_to_the_cost_callback( + retry_upstream: respx.MockRouter, +) -> None: + reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + retry_upstream.post(_ANTHROPIC.url).mock( + side_effect=[_provider_error(429, {"retry-after": "0"}), httpx.Response(200, json=_ANTHROPIC.ok_body)] + ) + + await _completion_with_retries_kwargs( + False, _ANTHROPIC, num_retries=1, metadata={"user_api_key_budget_reservation": reservation} + ) + + assert reservation["callback_bound"] is True + + @pytest.mark.parametrize("stream", (False, True)) @pytest.mark.parametrize("asynchronous", (False, True)) @pytest.mark.parametrize("profile,custom_price", ((None, False), ("us.openai.gpt-6-sol", False), (None, True))) From 3ce23b558c0b2b9f40e0c40ca8390a4351bf734a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:05:16 -0700 Subject: [PATCH 18/20] test(retries): wait on the reservation hand-back event instead of polling with deadlines --- tests/unit/test_main.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index d6c2f4f0ea9..ac80ca59a92 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -10842,16 +10842,22 @@ async def test_acompletion_keeps_provider_error_when_tenacity_is_missing( assert route.call_count == 1 -async def _until(condition: Callable[[], bool]) -> None: - while not condition(): - await asyncio.sleep(0.01) +class _WatchedReservation(dict[str, object]): + def __init__(self) -> None: + super().__init__(reserved_cost=0.5, entries=[], finalized=False, callback_bound=False) + self.handed_back: Final = asyncio.Event() + + def __setitem__(self, key: str, value: object) -> None: + super().__setitem__(key, value) + if key == "callback_bound" and value is False: + self.handed_back.set() @pytest.mark.asyncio async def test_acompletion_hands_the_budget_reservation_back_while_it_waits_to_retry( retry_upstream: respx.MockRouter, ) -> None: - reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} + reservation: Final = _WatchedReservation() route: Final = retry_upstream.post(_ANTHROPIC.url).mock(return_value=_provider_error(429, {"retry-after": "30"})) call: Final = asyncio.create_task( _completion_with_retries_kwargs( @@ -10859,8 +10865,7 @@ async def test_acompletion_hands_the_budget_reservation_back_while_it_waits_to_r ) ) - await asyncio.wait_for(_until(lambda: route.call_count == 1), 10) - await asyncio.wait_for(_until(lambda: reservation["callback_bound"] is False), 5) + await reservation.handed_back.wait() call.cancel() with pytest.raises(asyncio.CancelledError): await call From 86c4b204d1cba765d5f66c08c82124c433c41a2c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 19:23:43 -0700 Subject: [PATCH 19/20] test(retries): retry the budget success case with no backoff wait --- tests/unit/test_main.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index ac80ca59a92..3b5a59138fd 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -10880,11 +10880,15 @@ async def test_acompletion_retry_that_succeeds_leaves_the_budget_reservation_to_ ) -> None: reservation: Final = {"reserved_cost": 0.5, "entries": [], "finalized": False, "callback_bound": False} retry_upstream.post(_ANTHROPIC.url).mock( - side_effect=[_provider_error(429, {"retry-after": "0"}), httpx.Response(200, json=_ANTHROPIC.ok_body)] + side_effect=[_provider_error(500), httpx.Response(200, json=_ANTHROPIC.ok_body)] ) await _completion_with_retries_kwargs( - False, _ANTHROPIC, num_retries=1, metadata={"user_api_key_budget_reservation": reservation} + False, + _ANTHROPIC, + num_retries=1, + retry_strategy="constant_retry", + metadata={"user_api_key_budget_reservation": reservation}, ) assert reservation["callback_bound"] is True From ac45a34f2b84e587c10421e560e51c3fe02be1f5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 09:12:05 -0800 Subject: [PATCH 20/20] fix(main): do not retry a failed agentic follow-up under a direct retry_policy --- litellm/main.py | 2 ++ tests/unit/test_main.py | 32 +++++++++++++++++++++++++------- 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 8b2f046592b..c878da7cd0b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6680,6 +6680,8 @@ def _retry_plan_after_first_error( first_error: Exception, num_retries: int, retry_policy: RetryPolicy ) -> _RetryPlan | None: """The retries a direct retry_policy grants after the first error, kept for the whole call like the router's.""" + if is_agentic_followup_failure(first_error): + return None policy_retries: Final[int | None] = get_num_retries_from_retry_policy( exception=first_error, retry_policy=retry_policy ) diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 7427c6bec41..3afc337a524 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -11143,31 +11143,49 @@ def test_completion_with_retries_internal_server_error_500(): mock_sleep.assert_called_once() -async def _with_three_retries(sync_mode: bool, original_function: MagicMock) -> object: +async def _with_three_retries( + sync_mode: bool, original_function: MagicMock, retry_policy: RetryPolicy | None +) -> object: messages: Final = [{"role": "user", "content": "hi"}] if sync_mode: return completion_with_retries( - model="gpt-4o", messages=messages, num_retries=3, original_function=original_function + model="gpt-4o", + messages=messages, + num_retries=3, + retry_policy=retry_policy, + original_function=original_function, ) return await acompletion_with_retries( - model="gpt-4o", messages=messages, num_retries=3, original_function=original_function + model="gpt-4o", + messages=messages, + num_retries=3, + retry_policy=retry_policy, + original_function=original_function, ) @pytest.mark.asyncio @pytest.mark.parametrize("sync_mode", [True, False]) -async def test_completion_with_retries_stops_at_a_failed_agentic_follow_up(sync_mode: bool): +@pytest.mark.parametrize("retry_policy", [None, RetryPolicy(DefaultRetries=2)], ids=["no-policy", "policy"]) +@pytest.mark.parametrize("failures_before_follow_up", [0, 1]) +async def test_completion_with_retries_stops_at_a_failed_agentic_follow_up( + sync_mode: bool, retry_policy: RetryPolicy | None, failures_before_follow_up: int +): follow_up_failure: Final = _make_internal_server_error() mark_agentic_followup_failure(follow_up_failure, ()) - outcomes: Final = [_make_internal_server_error(), follow_up_failure, {"status": "recovered"}] + outcomes: Final = [ + *(_make_internal_server_error() for _ in range(failures_before_follow_up)), + follow_up_failure, + {"status": "recovered"}, + ] mock_completion: Final = MagicMock(side_effect=outcomes) if sync_mode else AsyncMock(side_effect=outcomes) with patch("time.sleep"), patch("asyncio.sleep", new_callable=AsyncMock): with pytest.raises(litellm.InternalServerError) as raised: - await _with_three_retries(sync_mode, mock_completion) + await _with_three_retries(sync_mode, mock_completion, retry_policy) assert raised.value is follow_up_failure - assert mock_completion.call_count == 2 + assert mock_completion.call_count == failures_before_follow_up + 1 @pytest.mark.asyncio