Merge 1b7b35cf68 into c58fc0c996

litellm/main.py

@@ -25,7 +25,7 @@ from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
 from copy import deepcopy
 from functools import partial
 from types import MappingProxyType
-from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args
+from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, Union, cast, get_args
 
 from litellm._logging import redact_string
 from litellm._uuid import uuid
@@ -47,6 +47,7 @@ from litellm.types.llms.base import LiteLLMBaseModel
 
 # Other utils are imported directly to avoid circular imports
 from litellm.utils import (
+    _get_retry_after_from_exception_header,  # pyright: ignore[reportPrivateUsage]  # shared retry header parser
     exception_type,
     filter_out_litellm_params,
     get_litellm_params,
@@ -79,6 +80,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 exception header extractor
+)
 from litellm.litellm_core_utils.get_litellm_params import (
     AWS_CREDENTIAL_KWARGS_KEYS,
     OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS,
@@ -6260,6 +6264,32 @@ def completion(
         )
 
 
+class _RetryState(Protocol):
+    @property
+    def outcome(self) -> futures.Future[object] | None: ...
+
+
+_RetryStateT: Final = TypeVar("_RetryStateT", bound=_RetryState)
+
+
+def _retry_after_wait(retry_state: _RetryStateT, fallback_wait: Callable[[_RetryStateT], float]) -> float:
+    """Use a provider retry hint when available, otherwise keep the configured backoff."""
+    outcome: Final = retry_state.outcome
+    if outcome is None:
+        return fallback_wait(retry_state)
+
+    exception: Final = outcome.exception()
+    if not isinstance(exception, Exception):
+        return fallback_wait(retry_state)
+
+    response_headers: Final = _get_response_headers(exception)
+    retry_after: Final = _get_retry_after_from_exception_header(response_headers)
+    if retry_after > 0:
+        return min(retry_after, 60)
+
+    return fallback_wait(retry_state)
+
+
 def completion_with_retries(*args, **kwargs):
     """
     Executes a litellm.completion() with 3 retries
@@ -6275,14 +6305,16 @@ def completion_with_retries(*args, **kwargs):
     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.Retrying(
-            wait=tenacity.wait_exponential(multiplier=1, max=10),
-            stop=tenacity.stop_after_attempt(num_retries),
-            reraise=True,
-        )
-    else:
-        retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True)
+    fallback_wait: Final = (
+        tenacity.wait_exponential(multiplier=1, max=10)
+        if retry_strategy == "exponential_backoff_retry"
+        else tenacity.wait_none()
+    )
+    retryer = tenacity.Retrying(
+        wait=partial(_retry_after_wait, fallback_wait=fallback_wait),
+        stop=tenacity.stop_after_attempt(num_retries),
+        reraise=True,
+    )
     return retryer(original_function, *args, **kwargs)
 
 
@@ -6301,14 +6333,16 @@ async def acompletion_with_retries(*args, **kwargs):
     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,
-        )
-    else:
-        retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True)
+    fallback_wait: Final = (
+        tenacity.wait_exponential(multiplier=1, max=10)
+        if retry_strategy == "exponential_backoff_retry"
+        else tenacity.wait_none()
+    )
+    retryer = tenacity.AsyncRetrying(
+        wait=partial(_retry_after_wait, fallback_wait=fallback_wait),
+        stop=tenacity.stop_after_attempt(num_retries),
+        reraise=True,
+    )
     return await retryer(original_function, *args, **kwargs)
 
 

litellm/utils.py

@@ -7534,36 +7534,40 @@ _should_retry = should_retry
 
 def _get_retry_after_from_exception_header(
     response_headers: Mapping[str, str] | None = None,
-):
+) -> float | int:
     """
     Reimplementation of openai's calculate retry after, since that one can't be imported.
     https://github.com/openai/openai-python/blob/af67cfab4210d8e497c05390ce14f39105c77519/src/openai/_base_client.py#L631
     """
+    if response_headers is None:
+        return -1
+
     try:
-        import email  # openai import
+        import email.utils
 
-        # About the Retry-After header: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After
-        #
-        # <http-date>". See https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After#syntax for
-        # details.
-        if response_headers is not None:
-            retry_header: Final = response_headers.get("retry-after", "")
+        retry_after_ms_header: Final = response_headers.get("retry-after-ms")
+        if retry_after_ms_header is not None:
             try:
-                retry_after = int(retry_header)
-            except Exception:
-                retry_date_tuple: Final = email.utils.parsedate_tz(retry_header)
-                if retry_date_tuple is None:
-                    retry_after = -1
-                else:
-                    retry_date: Final = email.utils.mktime_tz(retry_date_tuple)
-                    retry_after = int(retry_date - time.time())
-        else:
-            retry_after = -1
+                return float(retry_after_ms_header) / 1000
+            except ValueError:
+                pass
 
-        return retry_after
+        retry_header: Final = response_headers.get("retry-after")
+        if retry_header is None:
+            return -1
 
+        try:
+            return int(retry_header)
+        except ValueError:
+            pass
+
+        retry_date_tuple: Final = email.utils.parsedate_tz(retry_header)
+        if retry_date_tuple is None:
+            return -1
+
+        return int(email.utils.mktime_tz(retry_date_tuple) - time.time())
     except Exception:
-        retry_after = -1
+        return -1
 
 
 def calculate_retry_after(
@@ -7578,7 +7582,7 @@ def calculate_retry_after(
     jitter: Final = JITTER * random.random()
 
     # If the API asks us to wait a certain amount of time (and it's a reasonable amount), just do what it says.
-    if retry_after is not None and 0 < retry_after <= 60:
+    if 0 < retry_after <= 60:
         return retry_after + jitter
 
     # Calculate exponential backoff

tests/unit/test_main.py

@@ -2,6 +2,7 @@ import asyncio
 import base64
 import contextlib
 import copy
+import email.utils
 import io
 import json
 import logging
@@ -15,7 +16,7 @@ from datetime import datetime
 from importlib import import_module
 from pathlib import Path
 from typing import Final
-from unittest.mock import AsyncMock, MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
 import httpx
 import pytest
@@ -23,6 +24,7 @@ import respx
 from openai import APITimeoutError, AsyncOpenAI
 from openai._legacy_response import LegacyAPIResponse
 from openai.types.chat.chat_completion import ChatCompletion
+from tenacity import RetryCallState, Retrying
 
 import litellm
 from litellm import (
@@ -3083,6 +3085,182 @@ def throw_retryable_error(*_, **__):
     raise RuntimeError("BOOM")
 
 
+def _retry_state(exception: BaseException | None = None) -> RetryCallState:
+    state: Final = RetryCallState(retry_object=Retrying(), fn=None, args=(), kwargs={})
+    if exception is not None:
+        state.set_exception((type(exception), exception, exception.__traceback__))
+    return state
+
+
+def _rate_limit_error(headers: dict[str, str]) -> litellm.RateLimitError:
+    response = httpx.Response(
+        status_code=429,
+        headers=headers,
+        request=httpx.Request("POST", "https://example.invalid/v1"),
+    )
+    return litellm.RateLimitError(
+        message="rate limited",
+        model="test-model",
+        llm_provider="test-provider",
+        response=response,
+    )
+
+
+@pytest.mark.parametrize(
+    ("headers", "expected_wait"),
+    [
+        ({"retry-after": "20"}, 20),
+        ({"retry-after-ms": "2500"}, 2.5),
+        ({"retry-after": "120"}, 60),
+    ],
+)
+def test_retry_after_wait_uses_and_caps_provider_hint(headers: dict[str, str], expected_wait: float) -> None:
+    fallback_wait = Mock(return_value=1.0)
+
+    wait = litellm_main._retry_after_wait(_retry_state(_rate_limit_error(headers)), fallback_wait)
+
+    assert wait == expected_wait
+    fallback_wait.assert_not_called()
+
+
+def test_retry_after_wait_preserves_fallback_strategy_without_hint() -> None:
+    fallback_wait = Mock(return_value=3.0)
+
+    wait = litellm_main._retry_after_wait(_retry_state(RuntimeError("request failed")), fallback_wait)
+
+    assert wait == 3.0
+    fallback_wait.assert_called_once()
+
+
+def test_retry_after_wait_falls_back_for_invalid_provider_hint() -> None:
+    fallback_wait = Mock(return_value=3.0)
+
+    wait = litellm_main._retry_after_wait(
+        _retry_state(_rate_limit_error({"retry-after-ms": "invalid"})), fallback_wait
+    )
+
+    assert wait == 3.0
+    fallback_wait.assert_called_once()
+
+
+def test_retry_after_wait_falls_back_without_outcome_or_for_base_exception() -> None:
+    fallback_wait = Mock(return_value=3.0)
+
+    no_outcome_wait = litellm_main._retry_after_wait(_retry_state(), fallback_wait)
+    base_exception_wait = litellm_main._retry_after_wait(_retry_state(KeyboardInterrupt()), fallback_wait)
+
+    assert no_outcome_wait == 3.0
+    assert base_exception_wait == 3.0
+    assert fallback_wait.call_count == 2
+
+
+@pytest.mark.parametrize("milliseconds", [None, "invalid", ""])
+def test_retry_after_wait_parses_http_date_header(milliseconds: str | None) -> None:
+    retry_after: Final = email.utils.formatdate(1_000_030, usegmt=True)
+    headers: Final = (
+        {"retry-after": retry_after}
+        if milliseconds is None
+        else {"retry-after": retry_after, "retry-after-ms": milliseconds}
+    )
+    fallback_wait: Final = Mock(return_value=1.0)
+
+    with patch.object(litellm.utils.time, "time", return_value=1_000_000):
+        wait: Final = litellm_main._retry_after_wait(_retry_state(_rate_limit_error(headers)), fallback_wait)
+
+    assert wait == 30
+    fallback_wait.assert_not_called()
+
+
+def test_retry_after_wait_falls_back_for_invalid_retry_after_header() -> None:
+    fallback_wait = Mock(return_value=3.0)
+
+    wait = litellm_main._retry_after_wait(
+        _retry_state(_rate_limit_error({"retry-after": "invalid"})), fallback_wait
+    )
+
+    assert wait == 3.0
+    fallback_wait.assert_called_once()
+
+
+def test_retry_after_wait_falls_back_without_retry_header() -> None:
+    fallback_wait = Mock(return_value=3.0)
+
+    wait = litellm_main._retry_after_wait(_retry_state(_rate_limit_error({})), fallback_wait)
+
+    assert wait == 3.0
+    fallback_wait.assert_called_once()
+
+
+def test_retry_after_parser_handles_empty_headers() -> None:
+    retry_after = litellm.utils._get_retry_after_from_exception_header(httpx.Headers())
+
+    assert retry_after == -1
+
+
+_RETRY_AFTER_CASES: Final = (
+    pytest.param({"retry-after": "30", "retry-after-ms": "1500"}, 1.5, id="milliseconds-first"),
+    pytest.param({"retry-after": "3", "retry-after-ms": "invalid"}, 3.0, id="invalid-milliseconds"),
+    pytest.param({"retry-after": "3", "retry-after-ms": ""}, 3.0, id="empty-milliseconds"),
+    pytest.param({"retry-after": "invalid", "retry-after-ms": "2500"}, 2.5, id="invalid-seconds"),
+    pytest.param({"retry-after-ms": "120000"}, 60.0, id="milliseconds-cap"),
+    pytest.param({"retry-after": "120"}, 60.0, id="seconds-cap"),
+    pytest.param({}, None, id="missing-headers"),
+    pytest.param({"retry-after": "invalid", "retry-after-ms": "invalid"}, None, id="invalid-headers"),
+    pytest.param({"retry-after-ms": "invalid"}, None, id="invalid-milliseconds-only"),
+    pytest.param({"retry-after": "0"}, None, id="zero-delay"),
+    pytest.param({"retry-after-ms": "-1000"}, None, id="negative-delay"),
+)
+
+
+@pytest.mark.parametrize("retry_strategy", ["constant_retry", "exponential_backoff_retry"])
+@pytest.mark.parametrize(("headers", "provider_wait"), _RETRY_AFTER_CASES)
+def test_completion_with_retries_retries_after_provider_error(
+    retry_strategy: str, headers: dict[str, str], provider_wait: float | None
+) -> None:
+    original_function: Final = Mock(side_effect=[_rate_limit_error(headers), "ok"])
+
+    with patch("tenacity.nap.time.sleep") as sleep:
+        result: Final = litellm_main.completion_with_retries(
+            original_function=original_function,
+            num_retries=2,
+            retry_strategy=retry_strategy,
+        )
+
+    assert result == "ok"
+    assert original_function.call_count == 2
+    expected_wait: Final = (
+        provider_wait
+        if provider_wait is not None
+        else (1.0 if retry_strategy == "exponential_backoff_retry" else 0.0)
+    )
+    sleep.assert_called_once_with(expected_wait)
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("retry_strategy", ["constant_retry", "exponential_backoff_retry"])
+@pytest.mark.parametrize(("headers", "provider_wait"), _RETRY_AFTER_CASES)
+async def test_acompletion_with_retries_retries_after_provider_error(
+    retry_strategy: str, headers: dict[str, str], provider_wait: float | None
+) -> None:
+    original_function: Final = AsyncMock(side_effect=[_rate_limit_error(headers), "ok"])
+
+    with patch("asyncio.sleep", new_callable=AsyncMock) as sleep:
+        result: Final = await litellm_main.acompletion_with_retries(
+            original_function=original_function,
+            num_retries=2,
+            retry_strategy=retry_strategy,
+        )
+
+    assert result == "ok"
+    assert original_function.await_count == 2
+    expected_wait: Final = (
+        provider_wait
+        if provider_wait is not None
+        else (1.0 if retry_strategy == "exponential_backoff_retry" else 0.0)
+    )
+    sleep.assert_awaited_once_with(expected_wait)
+
+
 @pytest.mark.asyncio
 async def test_retrying() -> None:
     litellm.num_retries = 10

tests/unit/test_utils.py

@@ -10075,3 +10075,32 @@ def test_vertex_llama3_model_info_resolves_without_the_meta_prefix():
         prefixed_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai")
         assert litellm.get_model_info(model=short_name, custom_llm_provider="vertex_ai") == prefixed_info, model
         assert prefixed_info["key"] == f"vertex_ai/{model}", model
+
+
+@pytest.mark.parametrize(
+    ("headers", "provider_wait"),
+    [
+        pytest.param({"retry-after": "30", "retry-after-ms": "1500"}, 1.5, id="milliseconds-first"),
+        pytest.param({"retry-after": "3", "retry-after-ms": "invalid"}, 3.0, id="invalid-milliseconds"),
+        pytest.param({"retry-after": "3", "retry-after-ms": ""}, 3.0, id="empty-milliseconds"),
+        pytest.param({"retry-after-ms": "2500"}, 2.5, id="milliseconds-only"),
+        pytest.param({"retry-after": "invalid", "retry-after-ms": "invalid"}, None, id="invalid-headers"),
+        pytest.param(
+            {"retry-after": "Thu, 01 Jan 999999999999999999999999 00:00:00 GMT"}, None, id="overflowing-date"
+        ),
+        pytest.param({}, None, id="missing-headers"),
+    ],
+)
+def test_calculate_retry_after_preserves_provider_hint_and_jitter(
+    headers: dict[str, str], provider_wait: float | None
+) -> None:
+    wait: Final = litellm.utils.calculate_retry_after(
+        remaining_retries=1, max_retries=3, response_headers=httpx.Headers(headers)
+    )
+    expected_wait: Final = (
+        provider_wait
+        if provider_wait is not None
+        else min(litellm.utils.INITIAL_RETRY_DELAY * 2**2, litellm.utils.MAX_RETRY_DELAY)
+    )
+
+    assert expected_wait <= wait <= expected_wait + litellm.utils.JITTER