diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 9b6ec8d28d6..61042faa399 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -796,7 +796,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 @@ -3578,6 +3579,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/main.py b/litellm/main.py index 6b9964a450f..aaabe237b14 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -14,12 +14,13 @@ 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 @@ -72,6 +73,7 @@ from litellm.constants import ( ) from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.allowed_tools_tool_choice import ( raise_if_allowed_tools_tool_choice_unsupported, ) @@ -86,6 +88,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 ( AGENTIC_LOOP_CONTROL_KWARGS_KEYS, AWS_CREDENTIAL_KWARGS_KEYS, @@ -103,6 +108,7 @@ from litellm.litellm_core_utils.get_llm_provider_logic import ( from litellm.litellm_core_utils.get_provider_specific_headers import ( ProviderSpecificHeaderUtils, ) +from litellm.litellm_core_utils.get_retry_from_policy import get_num_retries_from_retry_policy from litellm.litellm_core_utils.health_check_utils import ( _filter_model_params, create_health_check_response, @@ -166,7 +172,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, validate_litellm_params_strictly +from litellm.types.router import GenericLiteLLMParams, RetryPolicy, validate_litellm_params_strictly from litellm.types.utils import ( AgenticLoopParams, BaseUrlRoutingHint, @@ -188,8 +194,10 @@ from litellm.utils import ( TextCompletionResponse, 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, convert_to_model_response_object, create_pretrained_tokenizer, create_tokenizer, @@ -207,6 +215,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, supports_pdf_input, @@ -6608,56 +6617,233 @@ def completion( ) -def completion_with_retries(*args, **kwargs): - """ - Executes a litellm.completion() with 3 retries - """ +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, + retries_sent: int, + num_retries: int, + retry_strategy: str, +) -> float: + 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 - retries_sent), + 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 + + +_NUM_RETRIES_ADAPTER: Final = TypeAdapter(float) +_RETRY_POLICY_ADAPTER: Final = TypeAdapter(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 is_agentic_followup_failure(e): + return False + 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) + return isinstance(e, (openai.APITimeoutError, openai.APIConnectionError)) + + return _predicate + + +@dataclass(frozen=True, slots=True) +class _RetryPlan: + num_retries: int + retry_policy: RetryPolicy | None + initial_exception: BaseException | None + + +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 + ) + 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)(first_error): + return None + return _RetryPlan(num_retries=num_retries - 1, retry_policy=retry_policy, 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}") - 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, + if plan.initial_exception is not None: + time.sleep( + _compute_retry_wait( + exc=plan.initial_exception, + retries_sent=0, + num_retries=plan.num_retries, + retry_strategy=retry_strategy, + ) ) - else: - retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + retryer: Final = tenacity.Retrying( + 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, + ) 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 - """ +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}") - num_retries: Final = kwargs.pop("num_retries", 3) + if plan.initial_exception is not None: + await asyncio.sleep( + _compute_retry_wait( + exc=plan.initial_exception, + 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, 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, + ) + return await retryer(original_function, *args, **kwargs) + + +def _build_retry_wait_fn( + num_retries: int, retry_strategy: str, first_attempt_is_a_retry: bool +) -> Callable[[_RetryState], float]: + 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, + 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, + ) + + return _wait_fn + + +def completion_with_retries( + *args: object, original_function: Callable[..., object] = completion, **kwargs: object +) -> object: + """ + Executes a litellm.completion() with retries + """ + 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[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" + ) + initial_exc: Final[object] = kwargs.pop("_initial_exception", None) + retry_policy: Final[object] = kwargs.pop("retry_policy", None) + validated_retry_policy: Final[RetryPolicy | None] = ( + _RETRY_POLICY_ADAPTER.validate_python(retry_policy) if retry_policy else None + ) + if isinstance(initial_exc, BaseException) or validated_retry_policy is None: + initial_plan: Final = _RetryPlan( + num_retries=num_retries, + retry_policy=validated_retry_policy, + initial_exception=initial_exc if isinstance(initial_exc, BaseException) else None, ) - else: - retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) - return await 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: object, original_function: Callable[..., Awaitable[object]] = acompletion, **kwargs: object +) -> object: + """ + [DEPRECATED]. Use 'acompletion' or router.acompletion instead! + Executes a litellm.acompletion() with retries + """ + 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" + ) + initial_exc: Final[object] = kwargs.pop("_initial_exception", None) + retry_policy: Final[object] = kwargs.pop("retry_policy", None) + validated_retry_policy: Final[RetryPolicy | None] = ( + _RETRY_POLICY_ADAPTER.validate_python(retry_policy) if retry_policy else None + ) + if isinstance(initial_exc, BaseException) or validated_retry_policy is None: + initial_plan: Final = _RetryPlan( + num_retries=num_retries, + retry_policy=validated_retry_policy, + initial_exception=initial_exc if isinstance(initial_exc, BaseException) else None, + ) + 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 53b02ebddac..6d9f0c95cca 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1090,6 +1090,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. @@ -1377,20 +1383,15 @@ def function_setup( dynamic_failure_callbacks: list[str | Callable[..., object] | CustomLogger] | None = None dynamic_async_failure_callbacks: Final[list[str | Callable[..., object] | 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") @@ -1730,32 +1731,39 @@ 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) 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"), - ) - kwargs["retry_policy"] = reset_retry_policy() - if retry_policy_num_retries is not None: - num_retries = retry_policy_num_retries - - return num_retries, kwargs + retry_policy: Final[object] = kwargs.pop("retry_policy", None) + if not retry_policy: + return num_retries, kwargs + kwargs["retry_policy"] = retry_policy + 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 -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: @@ -1767,6 +1775,41 @@ 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)) + + +_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]: + 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 + _stamp_correlation_context_if_supported(logging_obj) + logging_obj.failure_handler( + error, traceback_exception, start_time, end_time, for_async_request=True + ) # DO NOT MAKE THREADED - router retry fallback relies on this! + await logging_obj.async_failure_handler(error, traceback_exception, start_time, end_time) + + def _runs_as_async_success_callback(callback: str | Callable[..., object] | CustomLogger) -> bool: get_coroutine_checker: Final = litellm_utils.get_coroutine_checker return get_coroutine_checker().is_async_callable(callback) or callback in ("dynamodb", "s3") @@ -2043,6 +2086,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, Mapping) + 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 = ( @@ -2290,48 +2344,50 @@ 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) + 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", {}) 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 and not is_agentic_followup_failure(e) - ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying - if ( - isinstance(e, openai.APIError) - or isinstance(e, openai.Timeout) - or isinstance(e, openai.APIConnectionError) - ): - kwargs["num_retries"] = num_retries - return litellm.completion_with_retries(*args, **kwargs) - elif ( - isinstance(e, litellm.exceptions.ContextWindowExceededError) - and context_window_fallback_dict - and model in context_window_fallback_dict + num_retries and not is_completion_litellm_router_call - ): - if len(args) > 0: - args[0] = context_window_fallback_dict[model] - else: - kwargs["model"] = context_window_fallback_dict[model] - return original_function(*args, **kwargs) + and not context_window_fallback_applies + and not is_agentic_followup_failure(e) + ): # 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 ( + _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 + from litellm.chat_completions.dispatch import completion as public_completion + + return litellm.completion_with_retries( + *args, + original_function=public_completion, + **{**kwargs, "_initial_exception": e}, + ) + if context_window_fallback_applies: + fallback_model: Final = context_window_fallback_dict[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): - 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"), ) @@ -2380,8 +2436,10 @@ 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 _RETRY_ATTEMPT_MARKER _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: @@ -2389,6 +2447,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) @@ -2614,62 +2674,100 @@ def client(original_function): return _replayed_when_run_unstreamed(result, transcription_runs_unstreamed) except Exception as 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 + and bool(num_retries) + and not is_acompletion_litellm_router_call + and not isinstance(e, ImportError) + and not is_agentic_followup_failure(e) + ) + 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) + or should_retry(status_code) + ) + and isinstance(e, (openai.APIError, openai.APITimeoutError, openai.APIConnectionError)) + ) + 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: + result = await litellm.acompletion_with_retries( + *args, + 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, + }, + ) + 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) + else: + return result + 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. 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, for_async_request=True - ) # 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 - 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: - 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) - - if ( - num_retries - and not is_acompletion_litellm_router_call - and not isinstance(e, ImportError) - and not is_agentic_followup_failure(e) - ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying - try: - 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 - elif ( - isinstance(e, litellm.exceptions.ContextWindowExceededError) - and context_window_fallback_dict - and model in context_window_fallback_dict - and not is_acompletion_litellm_router_call - ): - 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: + 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 ( @@ -2686,19 +2784,14 @@ 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: 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 034b76507cf..6b7cb203770 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -6,24 +6,26 @@ 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 from pathlib import Path -from typing import Final, Literal +from types import MappingProxyType +from typing import Final, Literal, NoReturn from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from openai import APITimeoutError, AsyncOpenAI, OpenAI +from openai import APIConnectionError, APITimeoutError, AsyncOpenAI, OpenAI from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion -from pydantic import BaseModel, ConfigDict, TypeAdapter +from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError import litellm from litellm import ( @@ -35,21 +37,27 @@ 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 +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.agentic_followup_failure import mark_agentic_followup_failure 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 AsyncHTTPHandler, 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, + ModelResponse, ModelResponseStream, StandardCallbackDynamicParams, StreamingChoices, @@ -10976,6 +10984,1209 @@ async def test_atranscription_ignores_a_non_string_base_url_when_api_base_is_set assert response.text == "four score" +def _make_rate_limit_error(retry_after: str | None = None) -> litellm.RateLimitError: + headers: Final = {"retry-after": retry_after} if retry_after else {} + resp: Final = MagicMock() + resp.headers = headers + resp.status_code = 429 + err: Final = 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: Final = MagicMock() + resp.headers = {} + resp.status_code = 401 + err: Final = 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: Final = MagicMock() + resp.headers = {} + resp.status_code = 400 + err: Final = litellm.BadRequestError( + message="Invalid argument", + response=resp, + llm_provider="openai", + model="gpt-4o", + ) + err.status_code = 400 + 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 = {} + resp.status_code = 503 + err: Final = 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: Final = MagicMock() + resp.headers = {} + resp.status_code = 500 + err: Final = 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.""" + 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( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert res == {"status": "ok"} + 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 + + +def test_completion_with_retries_non_transient_auth_error_fails_immediately(): + """Non-transient errors like 401 should fail on attempt 1 without retrying.""" + mock_completion: Final = MagicMock(side_effect=_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 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.""" + mock_completion: Final = MagicMock(side_effect=_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 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.""" + 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( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + 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.""" + mock_completion: Final = MagicMock(side_effect=[_make_internal_server_error(), {"status": "recovered"}]) + + with patch("time.sleep") as mock_sleep: + res: Final = completion_with_retries( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + assert mock_completion.call_count == 2 + mock_sleep.assert_called_once() + + +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, + retry_policy=retry_policy, + original_function=original_function, + ) + return await acompletion_with_retries( + 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]) +@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() 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, retry_policy) + + assert raised.value is follow_up_failure + assert mock_completion.call_count == failures_before_follow_up + 1 + + +@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.""" + 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( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=3, + original_function=mock_acompletion, + ) + assert res == {"status": "ok"} + 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 + + +@pytest.mark.asyncio +async def test_acompletion_with_retries_non_transient_auth_error_fails_immediately(): + """Async 401 should fail on attempt 1 without retrying.""" + mock_acompletion: Final = AsyncMock(side_effect=_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 mock_acompletion.call_count == 1 + 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 + + +@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 + + +@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.""" + 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( + model="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + num_retries=2, + retry_strategy="constant_retry", + original_function=mock_completion, + ) + assert res == {"status": "recovered"} + assert mock_completion.call_count == 2 + mock_sleep.assert_called_once_with(0.0) + + +_RETRY_HOST: Final = "https://retry-upstream.invalid" +_RETRY_RESPONSES_OK: Final[Mapping[str, object]] = { + "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[Mapping[str, object]] = { + "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 + + +@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, + "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]) +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( + "retry_kwargs", + [ + {"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] +): + 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), + (429, {"num_retries": 3, "retry_policy": RetryPolicy(BadRequestErrorRetries=1)}, litellm.RateLimitError), + ], + 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, + 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 +@pytest.mark.parametrize("sync_mode", [True, False]) +@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(_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, _ANTHROPIC, num_retries=num_retries) + + assert route.call_count == expected_requests + + +@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 +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("positional_model", [False, True], ids=["model-keyword", "model-positional"]) +@pytest.mark.parametrize( + "retry_settings", + [{"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], + 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": + 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) + + request: Final = { + "messages": [{"role": "user", "content": "hi"}], + "api_base": _ANTHROPIC.api_base, + "api_key": "test-key", + **retry_settings, + "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} + + 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" + 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 +@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, + logging_futures: queue.SimpleQueue[Future[object]], + 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 + await _finish_logging(logging_futures) + + 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 +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",) + + +_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_read_and_restores_the_callers_after( + 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: + 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]) + ids_after_reading: Final = (trace_id_var.get(), session_id_var.get()) + 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" + 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( + 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]] +): + 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 +): + 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 + + 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 +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 + + +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 = _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( + False, _ANTHROPIC, num_retries=1, metadata={"user_api_key_budget_reservation": reservation} + ) + ) + + await reservation.handed_back.wait() + 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(500), httpx.Response(200, json=_ANTHROPIC.ok_body)] + ) + + await _completion_with_retries_kwargs( + False, + _ANTHROPIC, + num_retries=1, + retry_strategy="constant_retry", + metadata={"user_api_key_budget_reservation": reservation}, + ) + + assert reservation["callback_bound"] is True + + _PDF_FILE_PART: Final = { "type": "file", "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"},