Merge b2d1ec6822 into b1b236826a

litellm/main.py

@@ -6608,21 +6608,24 @@ def completion(
         )
 
 
-def completion_with_retries(*args, **kwargs):
+def completion_with_retries(
+    *args: object,
+    num_retries: int = 3,
+    retry_strategy: RetryStrategy = "constant_retry",
+    original_function: Callable[..., object] | None = None,
+    **kwargs: object,
+) -> object:
     """
-    Executes a litellm.completion() with 3 retries
+    Retry completion or a supplied SDK call.
     """
     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)
     # 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),
@@ -6631,24 +6634,28 @@ def completion_with_retries(*args, **kwargs):
         )
     else:
         retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True)
-    return retryer(original_function, *args, **kwargs)
+    retry_target: Final[Callable[..., object]] = completion if original_function is None else original_function
+    return retryer(retry_target, *args, **kwargs)
 
 
-async def acompletion_with_retries(*args, **kwargs):
+async def acompletion_with_retries(
+    *args: object,
+    num_retries: int = 3,
+    retry_strategy: RetryStrategy = "constant_retry",
+    original_function: Callable[..., object] | None = None,
+    **kwargs: object,
+) -> object:
     """
     [DEPRECATED]. Use 'acompletion' or router.acompletion instead!
-    Executes a litellm.completion() with 3 retries
+    Retry completion or a supplied SDK call.
     """
     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)
     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),
@@ -6657,7 +6664,8 @@ async def acompletion_with_retries(*args, **kwargs):
         )
     else:
         retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True)
-    return await retryer(original_function, *args, **kwargs)
+    retry_target: Final[Callable[..., object]] = completion if original_function is None else original_function
+    return await retryer(retry_target, *args, **kwargs)
 
 
 def responses_with_retries(*args, **kwargs):

litellm/utils.py

@@ -2349,6 +2349,17 @@ def client(original_function):
                     ):
                         kwargs["num_retries"] = num_retries
                         return litellm.responses_with_retries(*args, **kwargs)
+            elif call_type in (CallTypes.decisions.value, CallTypes.systemone.value):
+                retry_call: Final[partial[object]] = partial(wrapper, *args, **kwargs)
+                decision_num_retries, _ = _get_wrapper_num_retries(kwargs=retry_call.keywords, exception=e)
+                if (
+                    decision_num_retries
+                    and not _is_litellm_router_call(retry_call.keywords, is_async=False)
+                    and isinstance(e, openai.APIError)
+                ):
+                    return litellm.completion_with_retries(
+                        original_function=retry_call, num_retries=decision_num_retries
+                    )
             traceback_exception: Final = traceback.format_exc()
             end_time = datetime.datetime.now()
 
@@ -2692,6 +2703,20 @@ def client(original_function):
                     else:
                         return result
 
+            elif retry_call_type in (CallTypes.adecisions.value, CallTypes.asystemone.value):
+                if (
+                    num_retries
+                    and not _is_litellm_router_call(kwargs, is_async=True)
+                    and isinstance(e, openai.APIError)
+                ):
+                    retry_call: Final[partial[object]] = partial(wrapper_async, *args, **kwargs)
+                    return await litellm.acompletion_with_retries(
+                        num_retries=num_retries,
+                        original_function=retry_call,
+                        retry_strategy="exponential_backoff_retry"
+                        if isinstance(e, openai.RateLimitError)
+                        else "constant_retry",
+                    )
             deployment_num_retries: Final = kwargs.get("num_retries")
             if deployment_num_retries is not None:
                 setattr(e, "num_retries", deployment_num_retries)

tests/unit/test_utils.py

@@ -5805,6 +5805,126 @@ def test_responses_uses_sdk_retries_without_router(monkeypatch: pytest.MonkeyPat
         litellm.in_memory_llm_clients_cache.flush_cache()
 
 
+@pytest.mark.asyncio
+@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
+@pytest.mark.parametrize("systemone_format", [False, True], ids=["decisions", "systemone"])
+@pytest.mark.parametrize(
+    "num_retries,global_retries,upstream_failures,expected_attempts,router_call",
+    [
+        pytest.param(2, None, 1, 2, False, id="recovers"),
+        pytest.param(2, None, 3, 3, False, id="exhausts"),
+        pytest.param(0, None, 1, 1, False, id="disabled"),
+        pytest.param(None, None, 1, 1, False, id="omitted"),
+        pytest.param(2, None, 1, 1, True, id="router-owns-retries"),
+        pytest.param(None, 2, 1, 2, False, id="global-default"),
+        pytest.param(0, 2, 1, 1, False, id="disabled-overrides-global"),
+    ],
+)
+async def test_decisions_sdk_retries(
+    monkeypatch: pytest.MonkeyPatch,
+    is_async: bool,
+    systemone_format: bool,
+    num_retries: int | None,
+    global_retries: int | None,
+    upstream_failures: int,
+    expected_attempts: int,
+    router_call: bool,
+) -> None:
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    monkeypatch.setattr(litellm, "num_retries", global_retries)
+    success: Final = threading.Event()
+
+    def record_success(
+        kwargs: dict[str, object], completion_response: object, start_time: datetime, end_time: datetime
+    ) -> None:
+        success.set()
+
+    monkeypatch.setattr(litellm, "success_callback", [record_success])
+    retry_kwargs: Final = {} if num_retries is None else {"num_retries": num_retries}
+    call: Final = (
+        partial(
+            litellm.asystemone if is_async else litellm.systemone,
+            model="typesafe/jev-latest",
+            state="I was charged twice.",
+            questions={"needs_refund": {"type": "noul", "instructions": "Is a refund needed?"}},
+        )
+        if systemone_format
+        else partial(
+            litellm.adecisions if is_async else litellm.decisions,
+            model="typesafe/jev-latest",
+            input="I was charged twice.",
+            questions=[{"name": "needs_refund", "type": "predicate", "instructions": "Is a refund needed?"}],
+        )
+    )
+    request: Final = partial(
+        call,
+        api_base="https://decisions-retry.local",
+        api_key="test-key",
+        metadata={"model_group": "decisions-retry"} if router_call else None,
+        **retry_kwargs,
+    )
+    responses: Final = tuple(
+        httpx.Response(503, json={"error": {"message": "temporary outage", "type": "server_error"}})
+        for _ in range(upstream_failures)
+    ) + (
+        httpx.Response(
+            200,
+            json={"model": "jev-latest", "answers": {"needs_refund": {"type": "noul", "noul": 0.9}}},
+        ),
+    )
+    with respx.mock(assert_all_called=True) as upstream:
+        route: Final = upstream.post("https://decisions-retry.local/v1/systemone").mock(side_effect=responses)
+        if upstream_failures < expected_attempts:
+            pending: Final = request()
+            result: Final = await pending if asyncio.iscoroutine(pending) else pending
+            assert result.model_dump()["answers"] == (
+                {"needs_refund": {"type": "noul", "noul": 0.9}}
+                if systemone_format
+                else ({"name": "needs_refund", "type": "predicate", "probability": 0.9},)
+            )
+            assert await asyncio.to_thread(success.wait, 5), "Recovered requests must run success callbacks"
+        elif is_async:
+            with pytest.raises(litellm.ServiceUnavailableError, match="temporary outage"):
+                await request()
+        else:
+            with pytest.raises(litellm.ServiceUnavailableError, match="temporary outage"):
+                request()
+
+    assert route.call_count == expected_attempts
+    bodies: Final = tuple(json.loads(request.content) for request, _ in route.calls)
+    assert all(body == bodies[0] for body in bodies)
+    assert "num_retries" not in bodies[0]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
+async def test_decisions_retries_do_not_bypass_an_exhausted_budget(
+    monkeypatch: pytest.MonkeyPatch, is_async: bool
+) -> None:
+    monkeypatch.setattr(litellm, "max_budget", 1)
+    monkeypatch.setattr(litellm, "_current_cost", 2)
+    request: Final = partial(
+        litellm.adecisions if is_async else litellm.decisions,
+        model="typesafe/jev-latest",
+        input="hello",
+        questions=[{"name": "answer", "type": "predicate", "instructions": "Answer?"}],
+        api_base="https://decisions-budget.local",
+        api_key="test-key",
+        num_retries=2,
+    )
+    with respx.mock(assert_all_called=False) as upstream:
+        upstream.post("https://decisions-budget.local/v1/systemone").respond(
+            503, json={"error": {"message": "temporary outage"}}
+        )
+        if is_async:
+            with pytest.raises(litellm.BudgetExceededError):
+                await request()
+        else:
+            with pytest.raises(litellm.BudgetExceededError):
+                request()
+        assert len(upstream.calls) == 0
+
+
 @pytest.mark.asyncio
 async def test_wrapper_async_replays_cached_converted_chat_stream_as_stream(
     monkeypatch: pytest.MonkeyPatch,