diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 7f6d560d1db..93d344aa3ea 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -73,7 +73,11 @@ class Feature(str, Enum): MESSAGES_INTERCEPTOR = "messages_interceptor" """A registered MessagesInterceptor that takes over the call.""" AGENTIC_LOOP_HOOK = "agentic_loop_hook" - """A logger overriding async_should_run_agentic_loop, which may run follow-up calls.""" + """A logger overriding async_should_run_agentic_loop, which may run follow-up calls. + + Every interception that turns a stream into a converted stream in its pre-call deployment hook + (websearch, code interpreter, headroom) also overrides this method, so this feature keeps those + calls on Python too.""" TLS_POLICY = "tls_policy" """An OpenSSL security level or ECDH curve that native TLS cannot apply.""" MANAGED_RESOURCE_IDS = "managed_resource_ids" diff --git a/litellm/rust_bridge/dispatch.py b/litellm/rust_bridge/dispatch.py index 42c1a96c03d..aea1519701d 100644 --- a/litellm/rust_bridge/dispatch.py +++ b/litellm/rust_bridge/dispatch.py @@ -25,6 +25,16 @@ def call_hook( return hook(request, args, kwargs) +async def _stamped(start: Callable[[], Awaitable[ResultT]], kwargs: Mapping[str, object]) -> ResultT: + from litellm.utils import stamp_async_failure + + try: + return await start() + except Exception as error: + stamp_async_failure(error, kwargs) + raise + + @dataclass(frozen=True, slots=True) class PublicDispatch(Generic[RequestT]): route: Route @@ -100,7 +110,7 @@ class PublicDispatch(Generic[RequestT]): return await runtime.arun( self.context(native_request), binding=binding, - native=lambda hook: native(hook, native_request, args, kwargs), + native=lambda hook: _stamped(lambda: native(hook, native_request, args, kwargs), kwargs), python=python, rules=selected_rules, ) @@ -115,7 +125,7 @@ class PublicDispatch(Generic[RequestT]): return await runtime.arun( context, binding=binding, - native=lambda hook: native(hook, request, args, kwargs), + native=lambda hook: _stamped(lambda: native(hook, request, args, kwargs), kwargs), python=lambda: python(*args, **kwargs), rules=selected_rules, ) diff --git a/litellm/utils.py b/litellm/utils.py index ab97b69897f..4a1b189c261 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1482,7 +1482,14 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu return num_retries, kwargs -def _get_wrapper_timeout(kwargs: dict[str, object]) -> object: +def stamp_async_failure(error: Exception, kwargs: Mapping[str, object]) -> None: + num_retries: Final = kwargs.get("num_retries") + if num_retries is not None: + setattr(error, "num_retries", num_retries) + setattr(error, "timeout", _get_wrapper_timeout(kwargs=kwargs)) + + +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: @@ -2350,12 +2357,7 @@ def client(original_function): 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) + stamp_async_failure(e, kwargs) raise e finally: diff --git a/tests/test_litellm_rust/test_public_entries.py b/tests/test_litellm_rust/test_public_entries.py index 2f5c33e3ced..47558d4923a 100644 --- a/tests/test_litellm_rust/test_public_entries.py +++ b/tests/test_litellm_rust/test_public_entries.py @@ -285,3 +285,22 @@ async def test_sync_entry_with_flag_true_raises_setup_errors_at_await(entry: Ent assert inspect.iscoroutine(pending) with pytest.raises(RuntimeError, match="is unavailable"): await pending + + +@pytest.mark.parametrize("entry", ENTRIES, ids=entry_id) +async def test_async_native_failures_carry_the_retry_count_and_timeout_the_router_reads(entry: Entry) -> None: + failure: Final = litellm.RateLimitError(message="limited", llm_provider="openai", model="gpt-4o") + + async def failing(_: NativeCall) -> object: + raise failure + + entry.async_binding.override(failing) + try: + with scenario_rules(entry, SCENARIOS["required"]), pytest.raises(litellm.RateLimitError) as caught: + await entry.asynchronous(*entry.args(), **entry.kwargs(num_retries=3, timeout=7.0)) + finally: + entry.async_binding.reset() + + assert caught.value is failure + assert getattr(caught.value, "num_retries", None) == 3 + assert getattr(caught.value, "timeout", None) == 7.0