refactor(sdk): simplify Decisions retry calls

litellm/main.py

@@ -6680,21 +6680,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),
@@ -6703,24 +6706,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),
@@ -6729,7 +6736,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

@@ -307,7 +307,7 @@ except (ImportError, AttributeError, TypeError):
 # Convert to str (if necessary)
 claude_json_str = json.dumps(json_data)
 import importlib.metadata
-from collections.abc import AsyncIterator, Awaitable, Callable, Collection, Iterable, Iterator, Mapping, Sequence, Sized
+from collections.abc import AsyncIterator, Callable, Collection, Iterable, Iterator, Mapping, Sequence, Sized
 from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable
 
 from typing_extensions import assert_never
@@ -2286,20 +2286,15 @@ 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):
-                decision_num_retries, decision_kwargs = _get_wrapper_num_retries(
-                    kwargs=cast("dict[str, object]", kwargs),  # cast-ok: decorator kwargs always have string keys
-                    exception=e,
-                )
+                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(decision_kwargs, is_async=False)
+                    and not _is_litellm_router_call(retry_call.keywords, is_async=False)
                     and isinstance(e, openai.APIError)
                 ):
-                    return cast(  # cast-ok: legacy retry helper forwards arbitrary SDK arguments
-                        "Callable[..., object]", litellm.completion_with_retries
-                    )(
-                        *args,
-                        **{**decision_kwargs, "num_retries": decision_num_retries, "original_function": wrapper},
+                    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()
@@ -2630,18 +2625,13 @@ def client(original_function):
                     and not _is_litellm_router_call(kwargs, is_async=True)
                     and isinstance(e, openai.APIError)
                 ):
-                    return await cast(  # cast-ok: legacy retry helper forwards arbitrary SDK arguments
-                        "Callable[..., Awaitable[object]]", litellm.acompletion_with_retries
-                    )(
-                        *args,
-                        **{
-                            **kwargs,
-                            "num_retries": num_retries,
-                            "original_function": wrapper_async,
-                            "retry_strategy": "exponential_backoff_retry"
-                            if isinstance(e, openai.RateLimitError)
-                            else "constant_retry",
-                        },
+                    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: