Merge 4cf456a65f into 98acc01e0d

litellm/main.py

@@ -189,6 +189,7 @@ from litellm.utils import (
     TextCompletionResponse,
     TranscriptionResponse,
     Usage,
+    _get_retry_after_from_exception_header,  # pyright: ignore[reportPrivateUsage]  # shared internal retry parser
     add_provider_specific_params_to_optional_params,
     async_mock_completion_streaming_obj,
     convert_to_model_response_object,
@@ -6634,7 +6635,47 @@ def completion(
         )
 
 
-def completion_with_retries(*args, **kwargs):
+class _CompletionRetryOutcome(Protocol):
+    def exception(self) -> BaseException | None: ...
+
+
+class _CompletionRetryState(Protocol):
+    @property
+    def outcome(self) -> _CompletionRetryOutcome | None: ...
+
+
+def _completion_retry_after(error: Exception) -> float | None:
+    headers: Final = (
+        getattr(getattr(error, "response", None), "headers", None)
+        or getattr(error, "litellm_response_headers", None)
+        or getattr(error, "headers", None)
+    )
+    if headers is None:
+        return None
+    retry_after: Final = _get_retry_after_from_exception_header(
+        headers  # pyright: ignore[reportAny]  # provider exceptions expose untyped headers
+    )
+    return retry_after if 0 < retry_after <= 60 else None
+
+
+def _completion_retry_wait(
+    retry_state: _CompletionRetryState,
+    fallback_wait: Callable[[_CompletionRetryState], float],
+) -> float:
+    outcome: Final = retry_state.outcome
+    error: Final = outcome.exception() if outcome is not None else None
+    if isinstance(error, Exception):
+        retry_after: Final = _completion_retry_after(error)
+        if retry_after is not None:
+            return retry_after
+    return fallback_wait(retry_state)
+
+
+def _completion_constant_wait(_: _CompletionRetryState) -> float:
+    return 0.0
+
+
+def completion_with_retries(*args, _initial_retry_exception: Exception | None = None, **kwargs):
     """
     Executes a litellm.completion() with 3 retries
     """
@@ -6644,19 +6685,24 @@ def completion_with_retries(*args, **kwargs):
         raise Exception(f"tenacity import failed please run `pip install tenacity`. Error{e}")
 
     num_retries: Final = kwargs.pop("num_retries", 3)
-    # reset retries in .completion()
     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.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[Callable[[_CompletionRetryState], float]] = (
+        tenacity.wait_exponential(multiplier=1, max=10)
+        if retry_strategy == "exponential_backoff_retry"
+        else _completion_constant_wait
+    )
+    retryer: Final = tenacity.Retrying(
+        wait=partial(_completion_retry_wait, fallback_wait=fallback_wait),
+        stop=tenacity.stop_after_attempt(num_retries),
+        reraise=True,
+    )
+    if _initial_retry_exception is not None:
+        initial_retry_after: Final = _completion_retry_after(_initial_retry_exception)
+        if initial_retry_after is not None:
+            time.sleep(initial_retry_after)
     return retryer(original_function, *args, **kwargs)
 
 

litellm/utils.py

@@ -16,6 +16,7 @@ import io
 import itertools
 import json
 import logging
+import math
 import os
 import random
 import re
@@ -2316,7 +2317,7 @@ def client(original_function):
                         or isinstance(e, openai.APIConnectionError)
                     ):
                         kwargs["num_retries"] = num_retries
-                        return litellm.completion_with_retries(*args, **kwargs)
+                        return litellm.completion_with_retries(*args, _initial_retry_exception=e, **kwargs)
                 elif (
                     isinstance(e, litellm.exceptions.ContextWindowExceededError)
                     and context_window_fallback_dict
@@ -8512,7 +8513,7 @@ _should_retry = should_retry
 
 def _get_retry_after_from_exception_header(
     response_headers: Mapping[str, str] | None = None,
-):
+) -> float:
     """
     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
@@ -8524,24 +8525,34 @@ def _get_retry_after_from_exception_header(
         #
         # <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", "")
-            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
+        if response_headers is None:
+            return -1.0
 
-        return retry_after
+        retry_ms_header: Final = response_headers.get("retry-after-ms")
+        if retry_ms_header is not None:
+            try:
+                retry_after_ms: Final = float(retry_ms_header) / 1000
+            except (TypeError, ValueError):
+                pass
+            else:
+                if math.isfinite(retry_after_ms):
+                    return retry_after_ms
+
+        retry_header: Final[str | None] = response_headers.get("retry-after")
+        if retry_header is None:
+            return -1.0
+        try:
+            retry_after: Final = float(retry_header)
+        except (TypeError, ValueError):
+            retry_date_tuple: Final = email.utils.parsedate_tz(retry_header)
+            if retry_date_tuple is None:
+                return -1.0
+            retry_date: Final = email.utils.mktime_tz(retry_date_tuple)
+            return float(retry_date - time.time())
+        return retry_after if math.isfinite(retry_after) else -1.0
 
     except Exception:
-        retry_after = -1
+        return -1.0
 
 
 def calculate_retry_after(
@@ -8556,7 +8567,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

@@ -4097,6 +4097,84 @@ def throw_retryable_error(*_, **__):
     raise RuntimeError("BOOM")
 
 
+@pytest.mark.parametrize(
+    ("headers", "retry_strategy", "expected_wait"),
+    [
+        ({"retry-after": "20"}, "constant_retry", 20),
+        ({"retry-after-ms": "1500"}, "exponential_backoff_retry", 1.5),
+        ({"retry-after-ms": "1500", "retry-after": "20"}, "constant_retry", 1.5),
+        ({"retry-after-ms": "invalid", "retry-after": "2.5"}, "constant_retry", 2.5),
+        ({"retry-after": "60"}, "constant_retry", 60),
+        ({"retry-after": "61"}, "exponential_backoff_retry", 1),
+        ({"retry-after": "invalid"}, "constant_retry", 0),
+        ({}, "constant_retry", 0),
+    ],
+)
+def test_completion_retries_respect_provider_wait(
+    headers: dict[str, str], retry_strategy: str, expected_wait: float
+) -> None:
+    response: Final = httpx.Response(
+        status_code=429,
+        headers=headers,
+        request=httpx.Request("POST", "https://example.com/v1/chat/completions"),
+    )
+    error: Final = litellm.RateLimitError(
+        message="rate limited",
+        llm_provider="openai",
+        model="test-model",
+        response=response,
+        headers={"x-request-id": "request-1"},
+    )
+    outcomes: Final = iter((error, "recovered"))
+    sleep: Final = MagicMock()
+
+    def completion_stub(**kwargs: object) -> str:
+        outcome: Final = next(outcomes)
+        if isinstance(outcome, Exception):
+            raise outcome
+        return outcome
+
+    with patch("tenacity.nap.time.sleep", sleep):
+        result: Final = litellm_main.completion_with_retries(
+            num_retries=2,
+            retry_strategy=retry_strategy,
+            original_function=completion_stub,
+        )
+
+    assert result == "recovered"
+    sleep.assert_called_once_with(expected_wait)
+
+
+def test_completion_first_retry_respects_provider_wait() -> None:
+    from litellm.utils import client
+
+    response: Final = httpx.Response(
+        status_code=429,
+        headers={"retry-after": "2"},
+        request=httpx.Request("POST", "https://example.com/v1/chat/completions"),
+    )
+    error: Final = litellm.RateLimitError(
+        message="rate limited", llm_provider="openai", model="gpt-4o-mini", response=response
+    )
+    errors: Final = iter((error, error))
+    sleep: Final = MagicMock()
+
+    def completion(**kwargs: object) -> None:
+        raise next(errors)
+
+    with patch("tenacity.nap.time.sleep", sleep), pytest.raises(litellm.RateLimitError):
+        client(completion)(
+            model="gpt-4o-mini",
+            messages=[{"role": "user", "content": "hi"}],
+            num_retries=1,
+            original_function=completion,
+            litellm_logging_obj=_completion_logging_obj("first-retry"),
+        )
+
+    assert next(errors, None) is None
+    sleep.assert_called_once_with(2.0)
+
+
 @pytest.mark.asyncio
 async def test_retrying() -> None:
     litellm.num_retries = 10

tests/unit/test_utils.py

@@ -140,6 +140,23 @@ def test_get_base_model_from_metadata_ignores_a_non_string_root_value(root_base_
     )
 
 
+@pytest.mark.parametrize(
+    ("headers", "expected"),
+    (
+        ((("retry-after", "inf"),), -1.0),
+        ((("retry-after", "nan"),), -1.0),
+        ((("retry-after-ms", "inf"),), -1.0),
+        ((("retry-after-ms", "nan"), ("retry-after", "2")), 2.0),
+        ((("retry-after-ms", "1e309"), ("retry-after", "2")), 2.0),
+        ((("retry-after-ms", "1500"),), 1.5),
+    ),
+)
+def test_retry_after_parser_rejects_nonfinite_values(headers: tuple[tuple[str, str], ...], expected: float) -> None:
+    from litellm.utils import _get_retry_after_from_exception_header
+
+    assert _get_retry_after_from_exception_header(httpx.Headers(dict(headers))) == expected
+
+
 def _reject_short_post_call_response(input: str, model: str | None = None) -> dict[str, object]:
     if len(input) < 200:
         return {