From 6f25f627a61a6e6fc69fa28888994bed0cdcb967 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 7 Oct 2026 03:55:48 +0000 Subject: [PATCH] feat(router): serve the Responses API and Anthropic messages from the Rust backend Both go through a mirror of the generic helper after selection, with metadata in litellm_metadata. A Responses stream holds its opening lifecycle events until its first output and falls back from a drop before then on its hop's chain. A drop after output resumes the chain through a new Engine::resume entry point, continuing from the generated text. Python's Responses stream wrapper class moves to router_utils so both backends share it. Streaming Anthropic messages still raises NotImplementedError Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Aw1xktpUnix5tWhvSzey8t --- .../python-bridge/src/routes/router/host.rs | 8 +- .../python-bridge/src/routes/router/mod.rs | 19 +- .../python-bridge/src/routes/router/python.rs | 39 ++- .../crates/router/POC-behavior-spec.md | 1 + litellm-rust/crates/router/POC.md | 4 +- litellm-rust/crates/router/src/engine.rs | 85 ++++- litellm-rust/crates/router/src/host.rs | 4 + litellm-rust/crates/router/tests/engine.rs | 41 ++- litellm/router_backends/python_router.py | 128 +------- litellm/router_backends/rust_call.py | 303 +++++++++++------- litellm/router_backends/rust_router.py | 101 ++++-- litellm/router_backends/rust_streams.py | 204 ++++++++++++ .../add_retry_fallback_headers.py | 9 + .../router_utils/responses_fallback_stream.py | 85 +++++ .../router/test_responses_routing.py | 275 ++++++++++++++++ .../router/test_rust_router.py | 46 +-- .../support/router_backends.py | 50 +++ tests/unit/router_backends/test_rust_call.py | 4 +- ...responses_streaming_container_ownership.py | 3 +- 19 files changed, 1085 insertions(+), 324 deletions(-) create mode 100644 litellm/router_backends/rust_streams.py create mode 100644 litellm/router_utils/responses_fallback_stream.py create mode 100644 tests/test_litellm_rust/router/test_responses_routing.py create mode 100644 tests/test_litellm_rust/support/router_backends.py diff --git a/litellm-rust/crates/python-bridge/src/routes/router/host.rs b/litellm-rust/crates/python-bridge/src/routes/router/host.rs index e0455d4ec5a..ddc7b3fbc64 100644 --- a/litellm-rust/crates/python-bridge/src/routes/router/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/router/host.rs @@ -3,7 +3,7 @@ use litellm_host_python::{ InvokeError, PythonBinding, PythonHostCalls, PythonOwned, missing_state, }; use litellm_router::{ - engine::{Routed, RouterCall}, + engine::Routed, host::{Attempt, Invoked, RouterHost}, }; use pyo3::{ @@ -13,7 +13,7 @@ use pyo3::{ types::PyDict, }; -use super::{BridgeError, PyObj, RouterHostCall, RouterProtocol, python}; +use super::{BridgeError, PyObj, RouterHostCall, RouterProtocol, RouterRequest, python}; use litellm_host::protocol::Reply; pub(super) struct BridgeHost(pub(super) HostServices); @@ -85,8 +85,8 @@ impl PythonBinding for RouterBinding { &mut self, _py: Python<'_>, arguments: &Bound<'_, PyDict>, - ) -> Result> { - python::router_call(arguments).map_err(InvokeError::Python) + ) -> Result> { + python::router_request(arguments).map_err(InvokeError::Python) } fn encode_response( diff --git a/litellm-rust/crates/python-bridge/src/routes/router/mod.rs b/litellm-rust/crates/python-bridge/src/routes/router/mod.rs index 505b9a443cb..253b24d6210 100644 --- a/litellm-rust/crates/python-bridge/src/routes/router/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/router/mod.rs @@ -14,7 +14,7 @@ use litellm_host::{ }; use litellm_host_python::{HookChain, enter_native, from_py_argument}; use litellm_router::{ - engine::{Engine, Failed, RouteError, Routed, RouterCall}, + engine::{Engine, Failed, Resume, RouteError, Routed, RouterCall}, host::{Attempt, Invoked}, random::PythonRandom, settings::Settings, @@ -62,10 +62,16 @@ pub(crate) enum RouterHostCall { }, } +/// A call to route, or to resume from a stream that failed after its content reached the caller. +pub(crate) struct RouterRequest { + call: RouterCall, + resume: Option>, +} + pub(crate) struct RouterProtocol; impl Protocol for RouterProtocol { - type Request = RouterCall; + type Request = RouterRequest; type Response = Routed; type Error = BridgeError; type HostCall = RouterHostCall; @@ -121,12 +127,17 @@ impl NativeRouter { let engine = Arc::clone(&self.engine); litellm_host_python::run_call( py, - move |_, _, request: RouterCall| { + move |_, _, request: RouterRequest| { Ok(hosted_call::( request, None, move |request, services, _, _| async move { - match engine.route(&BridgeHost(services), request).await { + let host = BridgeHost(services); + let routed = match request.resume { + None => engine.route(&host, request.call).await, + Some(resume) => engine.resume(&host, request.call, resume).await, + }; + match routed { Ok(routed) => Ok(CallOutput::Complete(routed)), Err(RouteError::Failed(failed)) => { Err(BridgeError::Failed(Arc::from(failed))) diff --git a/litellm-rust/crates/python-bridge/src/routes/router/python.rs b/litellm-rust/crates/python-bridge/src/routes/router/python.rs index 62b9f25d455..d157c4270fc 100644 --- a/litellm-rust/crates/python-bridge/src/routes/router/python.rs +++ b/litellm-rust/crates/python-bridge/src/routes/router/python.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use litellm_host_python::from_py; use litellm_router::{ - engine::{Outcome, Override, RouterCall}, + engine::{Outcome, Override, Resume, RouterCall}, failure::{Classified, Raised, Rejection}, host::{Attempt, Invoked, MockFailure, Op, Target, TypedFallback}, operation::Operation, @@ -18,9 +18,40 @@ use pyo3::{ types::{PyDict, PyList, PyTuple}, }; -use super::PyObj; +use super::{PyObj, RouterRequest}; -pub(super) fn router_call(arguments: &Bound<'_, PyDict>) -> PyResult { +pub(super) fn router_request(arguments: &Bound<'_, PyDict>) -> PyResult { + let resume = match arguments.get_item("resume")? { + Some(resume) if !resume.is_none() => Some(self::resume(resume.cast::()?)?), + _ => None, + }; + Ok(RouterRequest { + call: router_call(arguments)?, + resume, + }) +} + +/// `{model_group, fallback_depth, original_model_group, attempted_targets, error, classified, +/// deployment_id}` for the stream that failed. +fn resume(resume: &Bound<'_, PyDict>) -> PyResult> { + let error = resume + .get_item("error")? + .ok_or_else(|| PyValueError::new_err("router resume is missing \"error\""))?; + let classified = resume + .get_item("classified")? + .ok_or_else(|| PyValueError::new_err("router resume is missing \"classified\""))?; + Ok(Resume { + group: required(resume, "model_group")?, + depth: required(resume, "fallback_depth")?, + original_group: required(resume, "original_model_group")?, + attempted_targets: optional(resume, "attempted_targets")?.unwrap_or_default(), + error: Arc::new(error.unbind()), + classified: from_py::(&classified)?, + deployment_id: optional(resume, "deployment_id")?, + }) +} + +fn router_call(arguments: &Bound<'_, PyDict>) -> PyResult { let operation: String = required(arguments, "operation")?; let model: String = required(arguments, "model")?; let mut call = RouterCall::new(Operation::from_name(&operation), model); @@ -105,8 +136,10 @@ pub(super) fn attempt<'py>(py: Python<'py>, attempt: &Attempt) -> PyResul )?, } dict.set_item("model_group", &attempt.model_group)?; + dict.set_item("original_model_group", &attempt.original_group)?; dict.set_item("bucket", attempt.bucket)?; dict.set_item("fallback_depth", attempt.fallback_depth)?; + dict.set_item("attempted_targets", &attempt.attempted_targets)?; dict.set_item("model_group_size", attempt.retry.model_group_size)?; dict.set_item("attempted_retries", attempt.retry.attempted_retries)?; dict.set_item("max_retries", attempt.retry.max_retries)?; diff --git a/litellm-rust/crates/router/POC-behavior-spec.md b/litellm-rust/crates/router/POC-behavior-spec.md index 083508f5826..cf545081ec8 100644 --- a/litellm-rust/crates/router/POC-behavior-spec.md +++ b/litellm-rust/crates/router/POC-behavior-spec.md @@ -535,3 +535,4 @@ Response (`set_response_headers`, PR:11756-11784, HDR): 26. Sync chat streams: a stream failing before content reruns `function_with_fallbacks` on the same group instead of walking the fallback chain, so fallbacks are only reached when the rerun's call itself raises, and a group that keeps failing recurses until `RecursionError` (PR:3624-3735). NOT replicated: the Rust backend applies the async rule (section 9) to sync streams too 27. `_combine_fallback_usage` only adds the failed stream's partial usage when `hasattr(usage, "usage")`, which a `Usage` never has, so partial usage is never added to the fallback's (PR:2942-2947). Matching it means the Rust backend adds nothing 28. Python's sync `_completion` updates none of `total_calls`/`success_calls`/`fail_calls`, while `_acompletion` updates all three (PR:2652-2760 vs 3780-3940) +29. The Responses stream wrapper adopts its fallback's headers before reading it (PR:3486-3492), so when that fallback's own stream falls back again the caller's headers still name the first fallback. The chat wrapper re-adopts once the fallback yields (PR:3104-3116). NOT replicated: the Rust backend names the group that served, as chat does diff --git a/litellm-rust/crates/router/POC.md b/litellm-rust/crates/router/POC.md index ca9fabe11a2..2c0d6d9ce6c 100644 --- a/litellm-rust/crates/router/POC.md +++ b/litellm-rust/crates/router/POC.md @@ -177,6 +177,8 @@ Design notes made while implementing: Streaming (chat, both drivers): `RoutedCall` reads the stream until the first chunk with content (Python's `_stream_chunks_have_generated_content`) and returns a `ReplayStream` (a `FallbackAwareStreamWrapper`, sync and async) that replays the buffered chunks. A `MidStreamFallbackError` before content is classified `stream_failure: before_content` and goes to that hop's own fallback chain without in-group retries; any other pre-content stream error is `terminal`. Either way the result settles the call: Python's outer retry and fallback layers had already returned the stream, so the engine's `Stop::Settled` and `Success::settled` keep them from retrying, trying the next target or rewriting the outcome. A stream error reaching the caller is unwrapped to its `original_exception`, as Python's wrappers do. The sync driver uses the async rule on purpose, because Python's sync wrapper reruns the same group and can recurse until `RecursionError` (spec section 11 item 26). Partial usage is not added to the fallback's, matching Python, whose check never passes (item 27). Like Python, the sync path no longer touches `total_calls`/`success_calls`/`fail_calls` (item 28) -Next, following the revised POC scope: finish increment 1 (`aresponses`, `aanthropic_messages`, and more suites opted into `router_backend`, starting with `tests/unit/test_router_order_fallback.py` and the fallback/cooldown tests in `tests/unit/router_utils` that go through `Router(...)`), then increments 2, 3 and 4. Then sketch the options for increment 3 +Responses API (`RustRouter.aresponses`): `RoutedCall._ageneric` mirrors `_ageneric_api_call_with_fallbacks_helper` after selection (metadata in `litellm_metadata`, failures counted by model group, the Anthropic refusal check). `rust_streams.responses_until_output` holds the opening lifecycle events until the first output, as Python's wrapper does; a mid-stream fallback error before then goes to the hop's chain with a content-policy trigger unwrapped. After output the stream still falls back: the `ResponsesReplay` (Python's `FallbackResponsesStreamWrapper`, moved to `litellm/router_utils/responses_fallback_stream.py` so both backends use it) calls `Engine::resume` through a second `route` with the hop's group, depth, original group and attempted targets, and the input continued from the generated text. Partial usage is added to the fallback's terminal events with Python's helpers (not covered by a parity test yet). Headers after a fallback that falls back again differ on purpose (spec section 11 item 29). `aanthropic_messages` serves non-streaming calls; streaming raises `NotImplementedError` until its same-group retry loop lands + +Next, following the revised POC scope: finish increment 1 (streaming `aanthropic_messages`, and more suites opted into `router_backend`, starting with `tests/unit/test_router_order_fallback.py` and the fallback/cooldown tests in `tests/unit/router_utils` that go through `Router(...)`), then increments 2, 3 and 4. Then sketch the options for increment 3 Known gaps to close or decline: CustomLogger fallback-event hooks and `router_cooldown_event_callback` (increment 3), the 200-key `InMemoryCache` eviction, fallback keys containing `/` (provider-prefixed matching needs the cost map) diff --git a/litellm-rust/crates/router/src/engine.rs b/litellm-rust/crates/router/src/engine.rs index 90e0ac4a5fd..a7f4d717803 100644 --- a/litellm-rust/crates/router/src/engine.rs +++ b/litellm-rust/crates/router/src/engine.rs @@ -108,6 +108,18 @@ pub struct Failed { pub ops: Vec>, } +/// A stream that failed after its content reached the caller, which some operations still +/// fall back from: the hop it was served by, as its attempt reported it, and its failure. +pub struct Resume { + pub group: String, + pub depth: u32, + pub original_group: String, + pub attempted_targets: Vec, + pub error: E, + pub classified: Classified, + pub deployment_id: Option, +} + pub struct Engine { registry: Registry, store: Store, @@ -135,6 +147,27 @@ impl Engine { &self, host: &H, call: RouterCall, + ) -> Result, RouteError> { + self.run(host, call, None).await + } + + /// Falls back from a stream that failed after its content reached the caller, as the + /// stream's own wrapper does in Python: the chain of the hop that served it, with the + /// groups already attempted skipped. The host's bucket 0 is that hop's. + pub async fn resume( + &self, + host: &H, + call: RouterCall, + resume: Resume, + ) -> Result, RouteError> { + self.run(host, call, Some(resume)).await + } + + async fn run( + &self, + host: &H, + call: RouterCall, + resume: Option>, ) -> Result, RouteError> { let snapshot = self.registry.load(); let mut run = Run { @@ -148,8 +181,13 @@ impl Engine { next_rejection: 0, ops: Vec::new(), }; - let model = call.model.clone(); - let result = run.hop(model.clone(), 0, 0, model).await; + let result = match resume { + None => { + let model = call.model.clone(); + run.hop(model.clone(), 0, 0, model).await + } + Some(resume) => run.resume(resume).await, + }; let ops = std::mem::take(&mut run.ops); match result { Ok(success) => Ok(Routed { @@ -249,6 +287,39 @@ impl<'a, H: RouterHost> Run<'a, H> { .as_ref()) } + async fn resume( + &mut self, + resume: Resume, + ) -> Attempted { + self.attempted_targets.extend(resume.attempted_targets); + let failure = Failure { + raised: Raised::Host(resume.error), + classified: resume.classified, + deployment_id: resume.deployment_id, + }; + if let Some(deployment) = failure + .deployment_id + .as_deref() + .and_then(|id| self.snapshot.by_id(id)) + && self.run_failure_callbacks(&failure.classified) + { + cooldown::record_failure( + self.snapshot, + &self.engine.store, + deployment, + &failure.classified, + ) + .await; + } + let hop = Hop { + group: resume.group, + depth: resume.depth, + bucket: 0, + original_group: resume.original_group, + }; + self.fallback(&hop, Box::new(failure)).await + } + /// `async_function_with_fallbacks` for one hop. fn hop( &mut self, @@ -285,8 +356,10 @@ impl<'a, H: RouterHost> Run<'a, H> { let attempt = Attempt { target: Target::Mock(mock), model_group: hop.group.clone(), + original_group: hop.original_group.clone(), bucket: hop.bucket, fallback_depth: hop.depth, + attempted_targets: self.attempted_targets(), retry: RetryStamp { model_group_size: 0, attempted_retries: 0, @@ -522,8 +595,10 @@ impl<'a, H: RouterHost> Run<'a, H> { let attempt = Attempt { target: Target::Deployment(deployment.id.clone()), model_group: hop.group.clone(), + original_group: hop.original_group.clone(), bucket: hop.bucket, fallback_depth: hop.depth, + attempted_targets: self.attempted_targets(), retry, ops: std::mem::take(&mut self.ops), }; @@ -556,6 +631,12 @@ impl<'a, H: RouterHost> Run<'a, H> { } } + fn attempted_targets(&self) -> Vec { + let mut targets: Vec = self.attempted_targets.iter().cloned().collect(); + targets.sort(); + targets + } + /// A failure runs litellm's failure callbacks when litellm raised it, and only once per /// request when the caller shares one logging object across attempts. fn run_failure_callbacks(&mut self, classified: &Classified) -> bool { diff --git a/litellm-rust/crates/router/src/host.rs b/litellm-rust/crates/router/src/host.rs index d6958afa41d..a8e69d4f8ae 100644 --- a/litellm-rust/crates/router/src/host.rs +++ b/litellm-rust/crates/router/src/host.rs @@ -92,8 +92,12 @@ pub struct Attempt { pub target: Target, /// The hop's `kwargs["model"]`. pub model_group: String, + /// The group the request was for, which a stream's later fallback looks its chain up by. + pub original_group: String, pub bucket: u32, pub fallback_depth: u32, + /// The groups the fallback chain has tried so far, which a stream's later fallback skips. + pub attempted_targets: Vec, pub retry: RetryStamp, pub ops: Vec>, } diff --git a/litellm-rust/crates/router/tests/engine.rs b/litellm-rust/crates/router/tests/engine.rs index 44a62f200cf..e91f55655da 100644 --- a/litellm-rust/crates/router/tests/engine.rs +++ b/litellm-rust/crates/router/tests/engine.rs @@ -4,7 +4,7 @@ use std::{ }; use litellm_router::{ - engine::{Engine, Override, RouteError, RouterCall}, + engine::{Engine, Override, Resume, RouteError, RouterCall}, failure::{Classified, ExceptionClass, Raised, StreamFailure}, host::{Attempt, Invoked, Op, RouterHost, Target}, operation::Operation, @@ -526,3 +526,42 @@ async fn a_fallback_hop_stream_falls_back_on_its_own_chain_and_reports_its_depth ("i", 2) ); } + +#[rstest] +#[tokio::test] +async fn a_resumed_stream_falls_back_on_its_hops_chain_skipping_attempted_groups() { + let mut settings = settings(); + settings.fallbacks = Some(vec![chain("g", &["h", "i"])]); + let engine = engine( + vec![ + deployment("a", "g"), + deployment("c", "h"), + deployment("d", "i"), + ], + settings, + ); + let host = ScriptedHost::default(); + let resume = Resume { + group: "h".into(), + depth: 1, + original_group: "g".into(), + attempted_targets: vec!["g".into(), "h".into()], + error: "c#1".into(), + classified: Classified { + stream_failure: Some(StreamFailure::BeforeContent), + ..classified(503, ExceptionClass::ServiceUnavailable) + }, + deployment_id: Some("c".into()), + }; + + let routed = engine.resume(&host, call("g"), resume).await.ok().unwrap(); + + assert_eq!(host.attempts(), ["d"]); + assert_eq!( + ( + routed.outcome.model_group.as_str(), + routed.outcome.attempted_fallbacks + ), + ("i", 2) + ); +} diff --git a/litellm/router_backends/python_router.py b/litellm/router_backends/python_router.py index 9263414877d..804689d6415 100644 --- a/litellm/router_backends/python_router.py +++ b/litellm/router_backends/python_router.py @@ -157,6 +157,7 @@ from litellm.router_utils.add_retry_fallback_headers import ( complexity_router_decision_headers, ensure_response_additional_headers, get_hidden_params_dict, + prepare_fallback_hidden_params, prepare_response_for_header_attachment, replace_complexity_router_headers, response_total_token_count, @@ -2954,11 +2955,7 @@ class PythonRouter: def _prepare_fallback_hidden_params( fallback_response: object, ) -> tuple[dict[str, object], dict[str, object]]: - fallback_hidden_params: Final = get_hidden_params_dict(fallback_response) - fallback_headers: Final = fallback_hidden_params.get("additional_headers") - if not isinstance(fallback_headers, dict): - return fallback_hidden_params, {} - return fallback_hidden_params, cast("dict[str, object]", fallback_headers) + return prepare_fallback_hidden_params(fallback_response) @staticmethod def _adopt_fallback_response_headers( @@ -3363,127 +3360,10 @@ class PythonRouter: iterators on terminate. """ from litellm.exceptions import MidStreamFallbackError - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - _get_openai_response_types, - ) + from litellm.router_utils.responses_fallback_stream import FallbackResponsesStreamWrapper source_iterator: Final = response - # Pre-resolve the set of terminal stream event types so the - # per-chunk type check inside FallbackResponsesStreamWrapper - # stays cheap; mirrors the source-iterator filter at - # responses/streaming_iterator.py:243-247. - _openai_types: Final = _get_openai_response_types() - _RESPONSES_TERMINAL_EVENT_TYPES: Final = ( - _openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - _openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, - _openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, - ) - - class FallbackResponsesStreamWrapper(BaseResponsesAPIStreamingIterator): - """ - Subclasses BaseResponsesAPIStreamingIterator only for isinstance - compatibility (proxy + interactions code paths check the type). - Bypasses the parent constructor and delegates iteration to an - async generator. - """ - - fallback_headers_adopted: bool = False - - def __init__(self, async_generator: AsyncGenerator): - import time - from datetime import datetime - - self._async_generator = async_generator - # Mirror every attribute BaseResponsesAPIStreamingIterator.__init__ - # would have set. The wrapper bypasses super().__init__ (it has no - # httpx.Response of its own and no provider config to drive), so - # we copy from source_iterator where applicable and use safe - # defaults elsewhere. This keeps inherited methods (e.g. - # _check_max_streaming_duration, _handle_failure) safe to call. - # - # The bridge path (LiteLLMCompletionStreamingIterator used by - # Anthropic/Bedrock/Vertex) does not call super().__init__ and - # is missing many of these attributes — use getattr fallbacks - # so wrapper construction never raises AttributeError. The - # bridge stores the logging object as `litellm_logging_obj`. - # base class declares non-Optional types for these - # fields but the bridge path (LiteLLMCompletionStreamingIterator) - # can legitimately omit them at runtime — keep the None - # fallback. Same lines passed mypy on the pre-fix file - # because the surrounding function body wasn't fully - # type-narrowed; the new typed terminal-event tuple above - # is what made these surface. - self.response = getattr(source_iterator, "response", None) - self.model = getattr(source_iterator, "model", None) - self.logging_obj = getattr( - source_iterator, - "logging_obj", - getattr(source_iterator, "litellm_logging_obj", None), - ) - self.finished = False - self.responses_api_provider_config = getattr(source_iterator, "responses_api_provider_config", None) - self.completed_response = None - self.start_time = getattr(source_iterator, "start_time", datetime.now()) - self._failure_handled = False - self._yielded_first_chunk = False - self._generated_content = "" - self._completed_response_cached = False - self._completed_response_logged = False - self._completed_response_cache_hit = None - self._persist_completed_response_before_logging = True - self._stream_created_time = time.time() - self.litellm_metadata = getattr(source_iterator, "litellm_metadata", None) - self.custom_llm_provider = getattr(source_iterator, "custom_llm_provider", None) - self.request_data = getattr(source_iterator, "request_data", {}) or {} - self.call_type = getattr(source_iterator, "call_type", None) - # Preserve hidden params so response headers (model_id, - # api_base, additional_headers) keep flowing. - self._hidden_params = dict(getattr(source_iterator, "_hidden_params", None) or {}) - - def adopt_fallback_headers(self, fallback_response: object) -> tuple[dict[str, object], dict[str, object]]: - prepared: Final = PythonRouter._prepare_fallback_hidden_params(fallback_response) - self._hidden_params = {**prepared[0], "additional_headers": prepared[1]} - self.fallback_headers_adopted = True - return prepared - - def __aiter__(self): - return self - - async def __anext__(self): - try: - chunk: Final = await self._async_generator.__anext__() - except StopAsyncIteration: - # The inner generator is exhausted. If we never sniffed a - # terminal event off a chunk (the bridge path emits the - # final response.completed via common_done_event_logic, - # which raises StopAsyncIteration after returning it), - # fall back to whatever the source iterator latched so - # the proxy's container-ownership hook still sees a - # completed_response instead of logging a spurious - # "no completed_response" warning. - if self.completed_response is None: - self.completed_response = getattr(source_iterator, "completed_response", None) - raise - # Sniff the terminal stream event off each forwarded chunk - # so ``self.completed_response`` is populated regardless of - # which inner iterator produced it (source_iterator, - # fallback_iterator, or any future wrapper). Without this - # the proxy's container-ownership hook (which reads - # ``getattr(stream_response, "completed_response", None)`` - # via _extract_completed_responses_response) silently - # records nothing on streaming /v1/responses calls — every - # follow-up /v1/containers//files call then 403s for - # the very key that created the container (#30210). - if self.completed_response is None and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: - self.completed_response = chunk - return chunk - - async def aclose(self): - # async generators always expose aclose — no defensive check needed. - await self._async_generator.aclose() - async def stream_with_fallbacks(): held_lifecycle_events: tuple[object, ...] = () # rebind-ok: flushed at first output, dropped on fallback try: @@ -3520,7 +3400,7 @@ class PythonRouter: exc, ) - wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks()) + wrapper: Final = FallbackResponsesStreamWrapper(stream_with_fallbacks(), source_iterator) return wrapper async def _aresponses_fallback_attempt( diff --git a/litellm/router_backends/rust_call.py b/litellm/router_backends/rust_call.py index c34b3cc2490..b4866895b6a 100644 --- a/litellm/router_backends/rust_call.py +++ b/litellm/router_backends/rust_call.py @@ -10,8 +10,9 @@ readers look for them. from __future__ import annotations import asyncio +import copy import time -from collections.abc import Callable, Iterator, Mapping, MutableMapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, MutableMapping, Sequence from typing import Final, Literal, Protocol, TypeAlias, cast from pydantic import BaseModel, ConfigDict, TypeAdapter @@ -25,11 +26,20 @@ from litellm.litellm_core_utils.exception_mapping_utils import ( from litellm.litellm_core_utils.secret_redaction import redact_string from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.llms.anthropic.pass_through.messages.utils import safeguard_refusal_error +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.router_backends.python_router import ( - FallbackAwareStreamWrapper, - _stream_chunks_have_generated_content, # pyright: ignore[reportPrivateUsage] # the content rule Python's stream wrapper applies + PythonRouter, + _with_router_resolved_session_model, # pyright: ignore[reportPrivateUsage] # the session rewrite the generic helper applies ) -from litellm.router_strategy.complexity_router.context_compaction import compact_to_fit +from litellm.router_backends.rust_streams import ( + StreamFailed, + chat_until_content, + chat_until_content_sync, + responses_until_output, + unwrapped, +) +from litellm.router_strategy.complexity_router.context_compaction import compact_to_fit, surface_for_call from litellm.router_utils.add_retry_fallback_headers import ( add_fallback_headers_to_response, add_retry_headers_to_response, @@ -37,6 +47,7 @@ from litellm.router_utils.add_retry_fallback_headers import ( from litellm.router_utils.common_utils import ( format_fallback_outcome_message, format_no_fallback_group_message, + provider_for_generic_call, truncate_fallback_error_detail, ) from litellm.router_utils.cooldown_handlers import ( @@ -44,14 +55,17 @@ from litellm.router_utils.cooldown_handlers import ( is_background_response_cost_poll_not_found, is_caller_timeout_408, ) +from litellm.types.llms.openai import ResponseAPIUsage from litellm.types.router import RouterErrors, RouterRateLimitError -from litellm.types.utils import ModelResponse, ModelResponseStream +from litellm.types.utils import ModelResponse from litellm.utils import ( _get_retry_after_from_exception_header, # pyright: ignore[reportPrivateUsage] # the Retry-After parser retries and cooldowns use ) Bucket: TypeAlias = MutableMapping[str, object] Kwargs: TypeAlias = dict[str, object] +Operation: TypeAlias = Literal["completion", "responses", "anthropic_messages"] +ResumeRoute: TypeAlias = Callable[[Mapping[str, object], Kwargs], Awaitable[object]] class AttemptRouter(Protocol): @@ -84,7 +98,9 @@ class AttemptRouter(Protocol): self, deployment_params: Mapping[str, object], kwargs: Mapping[str, object] ) -> Kwargs: ... - def _update_kwargs_with_deployment(self, deployment: Kwargs, kwargs: Kwargs) -> None: ... + def _update_kwargs_with_deployment( + self, deployment: Kwargs, kwargs: Kwargs, function_name: str | None = None + ) -> None: ... def _get_async_openai_model_client(self, deployment: Kwargs, kwargs: Kwargs) -> object: ... @@ -92,6 +108,10 @@ class AttemptRouter(Protocol): def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: Kwargs) -> bool: ... + def _should_raise_anthropic_refusal_error( + self, model: str, original_generic_function: Callable[..., object], response: object, kwargs: Kwargs + ) -> bool: ... + def _set_deployment_num_retries_on_exception(self, exception: Exception, deployment: Kwargs) -> None: ... def _stamp_failed_deployment_id_with_effective_model_info( @@ -113,6 +133,16 @@ _ROUTER_ONLY_KWARGS: Final = ( "mock_testing_rate_limit_error", ) +_UNWRAPPED_FALLBACK_TRIGGERS: Final[Mapping[Operation, tuple[type[Exception], ...]]] = { + "completion": (), + "responses": (litellm.ContentPolicyViolationError,), + "anthropic_messages": (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError), +} +_GENERIC_HANDLERS: Final[Mapping[Operation, str]] = { + "responses": "aresponses", + "anthropic_messages": "anthropic_messages", +} + _EXCEPTION_CLASSES: Final[tuple[tuple[type[BaseException], str], ...]] = ( (litellm.ContextWindowExceededError, "ContextWindowExceeded"), (litellm.ContentPolicyViolationError, "ContentPolicyViolation"), @@ -135,8 +165,10 @@ class Attempt(BaseModel): deployment_id: str | None = None mock: Literal["fallbacks", "context_window_fallbacks", "content_policy_fallbacks", "rate_limit"] | None = None model_group: str + original_model_group: str bucket: int fallback_depth: int + attempted_targets: tuple[str, ...] model_group_size: int attempted_retries: int max_retries: int @@ -206,104 +238,10 @@ def classify(error: BaseException, kwargs: Mapping[str, object], callbacks_ran: } -class _StreamFailed(Exception): - """A stream that failed before its first content chunk.""" - - def __init__(self, error: BaseException, kind: Literal["before_content", "terminal"]) -> None: - super().__init__(str(error)) - self.error: Final = error - self.kind: Final = kind - - -class ReplayStream(FallbackAwareStreamWrapper): - """The attempt's stream, replaying the chunks read while waiting for its first content. - - After content has reached the caller the stream cannot fall back, so a mid-stream - fallback error surfaces as the error that caused it, as the Python router's wrapper does. - """ - - def __init__(self, inner: CustomStreamWrapper, buffered: Sequence[ModelResponseStream]) -> None: - super().__init__( # pyright: ignore[reportUnknownMemberType] # CustomStreamWrapper's untyped constructor - completion_stream=cast(object, inner.completion_stream), # cast-ok: untyped attribute - model=cast(object, inner.model), # cast-ok: untyped attribute - custom_llm_provider=inner.custom_llm_provider, - logging_obj=inner.logging_obj, - _response_headers=getattr(inner, "_response_headers", None), - ) - self._inner: Final = inner - self._replay: Final = iter(tuple(buffered)) - self.chunks = cast(list[ModelResponseStream], inner.chunks) # cast-ok: untyped attribute - self._hidden_params = dict(cast(Mapping[str, object], inner._hidden_params)) # pyright: ignore[reportPrivateUsage] # cast-ok: the stream's own untyped metadata - - def __aiter__(self) -> ReplayStream: - return self - - async def __anext__(self) -> ModelResponseStream: - buffered: Final = next(self._replay, None) - if buffered is not None: - return buffered - try: - return await self._inner.__anext__() - except MidStreamFallbackError as error: - raise _unwrapped(error) from error - - def __iter__(self) -> Iterator[ModelResponseStream]: - return self - - def __next__(self) -> ModelResponseStream: - buffered: Final = next(self._replay, None) - if buffered is not None: - return buffered - try: - return self._inner.__next__() - except MidStreamFallbackError as error: - raise _unwrapped(error) from error - - async def aclose(self) -> None: - await self._inner.aclose() - - -def _unwrapped(error: BaseException) -> BaseException: - """The provider error behind a mid-stream fallback error, which Python's stream wrappers raise in its place.""" - if isinstance(error, MidStreamFallbackError) and error.original_exception is not None: - return error.original_exception - return error - - def _not_fetched(stream: CustomStreamWrapper) -> bool: - return ( - cast(object, stream.completion_stream) is None and cast(object, stream.make_call) is not None - ) # cast-ok: untyped attributes - - -async def _replay_from_first_content(stream: CustomStreamWrapper) -> ReplayStream: - """Reads `stream` until a chunk carries content, so a failure before then can still go to - another deployment; the chunks read so far replay ahead of the rest.""" - buffered: Final[list[ModelResponseStream]] = [] # mutable-ok: chunks read before the first content - try: - async for chunk in stream: - buffered.append(chunk) - if _stream_chunks_have_generated_content([chunk]): - break - except MidStreamFallbackError as error: - raise _StreamFailed(error, "before_content") from error - except Exception as error: - raise _StreamFailed(error, "terminal") from error - return ReplayStream(stream, buffered) - - -def _replay_from_first_content_sync(stream: CustomStreamWrapper) -> ReplayStream: - buffered: Final[list[ModelResponseStream]] = [] # mutable-ok: chunks read before the first content - try: - for chunk in stream: - buffered.append(chunk) - if _stream_chunks_have_generated_content([chunk]): - break - except MidStreamFallbackError as error: - raise _StreamFailed(error, "before_content") from error - except Exception as error: - raise _StreamFailed(error, "terminal") from error - return ReplayStream(stream, buffered) + completion_stream: Final = cast(object, stream.completion_stream) # cast-ok: untyped attribute + make_call: Final = cast(object, stream.make_call) # cast-ok: untyped attribute + return completion_stream is None and make_call is not None class _AttemptFailed(Exception): @@ -321,13 +259,18 @@ class RoutedCall: self, normalizer: AttemptRouter, kwargs: Kwargs, - metadata_key: Literal["metadata", "litellm_metadata"], + operation: Operation, fallbacks: Callable[[str], object], + resume: ResumeRoute | None = None, ) -> None: self._normalizer: Final = normalizer self._kwargs: Final = {key: value for key, value in kwargs.items() if key not in _ROUTER_ONLY_KWARGS} + self._operation: Final = operation + metadata_key: Final = metadata_key_for(operation) self._metadata_key: Final = metadata_key self._fallbacks: Final = fallbacks + self._resume: Final = resume + self._partial_usage: Final[list[ResponseAPIUsage]] = [] # mutable-ok: what failed Responses streams used root: Final = kwargs.setdefault(metadata_key, {}) bucket: Final[Bucket] = _bucket(root) self._buckets: Final[dict[int, Bucket]] = {0: bucket} # mutable-ok: grows as fallback hops open buckets @@ -349,13 +292,13 @@ class RoutedCall: if parsed.mock is not None: return self._mock(parsed, kwargs) try: - response: Final = await self._acompletion(parsed, kwargs) - except _StreamFailed as failed: - return ( - "error", - failed.error, - {**classify(failed.error, kwargs, callbacks_ran=True), "stream_failure": failed.kind}, + response: Final = ( + await self._acompletion(parsed, kwargs) + if self._operation == "completion" + else await self._ageneric(parsed, kwargs) ) + except StreamFailed as failed: + return self._stream_failed(failed, kwargs) except _AttemptFailed as failed: return ("error", failed.error, classify(failed.error, kwargs, callbacks_ran=False)) except Exception as error: @@ -369,12 +312,8 @@ class RoutedCall: return self._mock(parsed, kwargs) try: response: Final = self._completion(parsed, kwargs) - except _StreamFailed as failed: - return ( - "error", - failed.error, - {**classify(failed.error, kwargs, callbacks_ran=True), "stream_failure": failed.kind}, - ) + except StreamFailed as failed: + return self._stream_failed(failed, kwargs) except _AttemptFailed as failed: return ("error", failed.error, classify(failed.error, kwargs, callbacks_ran=False)) except Exception as error: @@ -397,7 +336,21 @@ class RoutedCall: def failure(self, error: object, ops: Sequence[Mapping[str, object]]) -> BaseException: self._apply(ops) - return _unwrapped(self._exception(error)) + return unwrapped(self._exception(error)) + + def _stream_failed(self, failed: StreamFailed, kwargs: Mapping[str, object]) -> Invoked: + if failed.partial_usage is not None: + self._partial_usage.append(failed.partial_usage) + trigger: Final = self._fallback_trigger(failed.error) + return ("error", failed.error, {**classify(trigger, kwargs, callbacks_ran=True), "stream_failure": failed.kind}) + + def _fallback_trigger(self, error: BaseException) -> BaseException: + """The error a stream's fallback is judged and annotated by: the mid-stream fallback error, + or the provider error inside it for the classes this operation's wrapper unwraps.""" + if not isinstance(error, MidStreamFallbackError): + return error + original: Final = error.original_exception + return original if isinstance(original, _UNWRAPPED_FALLBACK_TRIGGERS[self._operation]) else error def _prepare(self, attempt: Attempt) -> Kwargs: self._apply(attempt.ops) @@ -407,8 +360,9 @@ class RoutedCall: bucket["attempted_retries"] = attempt.attempted_retries bucket["max_retries"] = attempt.max_retries first_hop_only: Final = () if attempt.bucket == 0 else ("mock_timeout",) + passed_explicitly: Final = ("model", "messages") if self._operation == "completion" else ("model",) return { - **{key: value for key, value in self._kwargs.items() if key not in ("model", "messages", *first_hop_only)}, + **{key: value for key, value in self._kwargs.items() if key not in (*passed_explicitly, *first_hop_only)}, **self._hop_kwargs[attempt.bucket], self._metadata_key: bucket, } @@ -504,9 +458,9 @@ class RoutedCall: await response.fetch_stream() normalizer.success_calls[model_name] += 1 if isinstance(response, CustomStreamWrapper): - return await _replay_from_first_content(response) + return await chat_until_content(response) return response - except _StreamFailed: + except StreamFailed: raise except _AttemptFailed as failed: self._stamp(failed.error, deployment, kwargs, model_name) @@ -521,6 +475,97 @@ class RoutedCall: self._stamp(error, deployment, kwargs, model_name) raise + async def _ageneric(self, attempt: Attempt, kwargs: Kwargs) -> object: + """`_ageneric_api_call_with_fallbacks_helper` after selection, with the Responses API's and + Anthropic messages' own handling of the stream it returns.""" + normalizer: Final = self._normalizer + deployment: Final = self._deployment(attempt) + handler: Final = _handler(self._operation) + hop_kwargs: Final = {**kwargs, self._metadata_key: copy.deepcopy(kwargs[self._metadata_key])} + try: + normalizer._update_kwargs_with_deployment( # pyright: ignore[reportPrivateUsage] # the attempt body PythonRouter runs + deployment=deployment, kwargs=kwargs, function_name="_ageneric_api_call_with_fallbacks" + ) + data: Final = dict(as_mapping(deployment.get("litellm_params"))) + model_name: Final = str(data["model"]) + normalizer.total_calls[model_name] += 1 + provider: Final = provider_for_generic_call(data) + response_kwargs: Final = { + **data, + "caching": normalizer.cache_responses, + **kwargs, + "model": model_name, + **_with_router_resolved_session_model(kwargs.get("session"), model_name), + **({"custom_llm_provider": provider} if provider is not None else {}), + } + compacted: Final = await compact_to_fit( + normalizer, # pyright: ignore[reportArgumentType] # PythonRouter + deployment, # pyright: ignore[reportArgumentType] # router deployment dict + response_kwargs, + surface_for_call(getattr(handler, "__name__", "")), + ) + response: Final = await handler(**compacted) + if normalizer._should_raise_anthropic_refusal_error( # pyright: ignore[reportPrivateUsage] # as above + model=attempt.model_group, original_generic_function=handler, response=response, kwargs=kwargs + ): + refusal: Final = cast(Mapping[str, object], response) # cast-ok: the refusal gate checked its shape + stop_details: Final = cast(dict[str, object], refusal["stop_details"]) # cast-ok: as above + raise _AttemptFailed(safeguard_refusal_error(model=attempt.model_group, stop_details=stop_details)) + normalizer.success_calls[model_name] += 1 + except _AttemptFailed as failed: + self._stamp_generic(failed.error, deployment, kwargs, attempt.model_group) + raise + except Exception as error: + self._stamp_generic(error, deployment, kwargs, attempt.model_group) + raise + if kwargs.get("stream") and isinstance(response, BaseResponsesAPIStreamingIterator): + return await responses_until_output( + response, self._responses_fallback(attempt, hop_kwargs), tuple(self._partial_usage) + ) + return response + + def _responses_fallback( + self, attempt: Attempt, hop_kwargs: Kwargs + ) -> Callable[[MidStreamFallbackError], Awaitable[object]]: + """`_aresponses_fallback_attempt`: the chain of the hop `attempt` ran in, continuing from the + text the failed stream generated.""" + + async def fall_back(error: MidStreamFallbackError) -> object: + if self._resume is None: + raise error + request_input: Final = cast(str, hop_kwargs.get("input")) # cast-ok: the request's own input + continued: Final = ( + { + "input": PythonRouter._build_responses_continuation_input( # pyright: ignore[reportPrivateUsage] # Python's continuation input + request_input, error.generated_content + ) + } + if error.generated_content and not error.is_pre_first_chunk + else {} + ) + trigger: Final = self._fallback_trigger(error) + return await self._resume( + { + "model_group": attempt.model_group, + "fallback_depth": attempt.fallback_depth, + "original_model_group": attempt.original_model_group, + "attempted_targets": attempt.attempted_targets, + "error": error, + "classified": classify(trigger, hop_kwargs, callbacks_ran=True), + "deployment_id": attempt.deployment_id, + }, + {**hop_kwargs, **continued}, + ) + + return fall_back + + def _stamp_generic( + self, error: BaseException, deployment: Mapping[str, object], kwargs: Mapping[str, object], model_group: str + ) -> None: + """The generic helper counts failures by model group, not by provider model.""" + self._normalizer.fail_calls[model_group] += 1 + self._stamp(error, deployment, kwargs, None) + def _completion(self, attempt: Attempt, kwargs: Kwargs) -> object: normalizer: Final = self._normalizer deployment: Final = self._deployment(attempt) @@ -566,9 +611,9 @@ class RoutedCall: if isinstance(response, CustomStreamWrapper) and _not_fetched(response): response.fetch_sync_stream() if isinstance(response, CustomStreamWrapper): - return _replay_from_first_content_sync(response) + return chat_until_content_sync(response) return response - except _StreamFailed: + except StreamFailed: raise except _AttemptFailed as failed: self._stamp(failed.error, deployment, kwargs, None) @@ -645,7 +690,7 @@ class RoutedCall: ) def _as_exception(self, raised: object) -> Exception: - error: Final = self._exception(raised) + error: Final = self._fallback_trigger(self._exception(raised)) return error if isinstance(error, Exception) else Exception(str(error)) def _open_bucket(self, op: Mapping[str, object]) -> None: @@ -676,13 +721,23 @@ class RoutedCall: self._append_message(op.get("error"), f"\n{message}", require_message=False) def _append_message(self, raised: object, text: str, require_message: bool = True) -> None: - error: Final = self._exception(raised) + error: Final = self._fallback_trigger(self._exception(raised)) message: Final[object] = getattr(error, "message", None) if not litellm.expose_router_debug_in_errors or (require_message and not hasattr(error, "message")): return error.message = f"{message}{text}" # pyright: ignore[reportAttributeAccessIssue] # litellm exceptions carry a message +def _handler(operation: Operation) -> Callable[..., Awaitable[object]]: + """Read at call time, so a patched handler is the one called.""" + name: Final = _GENERIC_HANDLERS[operation] + return cast(Callable[..., Awaitable[object]], getattr(litellm, name)) # cast-ok: litellm's async handler + + +def metadata_key_for(operation: Operation) -> Literal["metadata", "litellm_metadata"]: + return "metadata" if operation == "completion" else "litellm_metadata" + + def _rejection(raised: Mapping[str, object]) -> BaseException: model: Final = str(raised.get("model")) if raised.get("kind") == "no_healthy_deployments": diff --git a/litellm/router_backends/rust_router.py b/litellm/router_backends/rust_router.py index 5ac98c24e67..a9493c2ac54 100644 --- a/litellm/router_backends/rust_router.py +++ b/litellm/router_backends/rust_router.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from enum import Enum -from typing import Final, Literal, Protocol, TypeAlias, cast +from typing import Final, Protocol, TypeAlias, cast from urllib.parse import quote from pydantic import BaseModel, ConfigDict @@ -11,7 +11,15 @@ from pydantic import BaseModel, ConfigDict import litellm from litellm import constants from litellm.router_backends.python_router import PythonRouter -from litellm.router_backends.rust_call import AttemptRouter, Kwargs, RoutedCall, as_mapping, as_sequence +from litellm.router_backends.rust_call import ( + AttemptRouter, + Kwargs, + Operation, + RoutedCall, + as_mapping, + as_sequence, + metadata_key_for, +) from litellm.router_backends.rust_support import unsupported_reason from litellm.router_utils import clientside_credential_handler from litellm.router_utils.cooldown_handlers import ( @@ -90,7 +98,9 @@ class NormalizerView(AttemptRouter, Protocol): def discard(self) -> None: ... - def _update_kwargs_before_fallbacks(self, model: str, kwargs: Kwargs) -> None: ... + def _update_kwargs_before_fallbacks( + self, model: str, kwargs: Kwargs, metadata_variable_name: str = "metadata" + ) -> None: ... _IS_CLIENTSIDE_CREDENTIAL: Final = cast( # cast-ok: the helper's dict parameter is unparameterized @@ -247,10 +257,8 @@ class RustRouter: messages: Sequence[Mapping[str, object]], **kwargs: object, # kwargs-ok: forwards litellm.acompletion's kwargs ) -> object: - call: Final = self._call("acompletion", model, messages, kwargs) - native_route: Final = self._native.route(call.spec, call.driver, True) - routed: Final = cast(Awaitable[object], native_route) # cast-ok: an asynchronous route is a coroutine - return await routed + call: Final = self._call("completion", model, {**kwargs, "messages": messages}) + return await self._route(call, True) def completion( self, @@ -258,16 +266,37 @@ class RustRouter: messages: Sequence[Mapping[str, object]], **kwargs: object, # kwargs-ok: forwards litellm.completion's kwargs ) -> object: - call: Final = self._call("completion", model, messages, kwargs) + call: Final = self._call("completion", model, {**kwargs, "messages": messages}) return self._native.route(call.spec, call.driver, False) - def _call( + async def aresponses( self, - operation: Literal["acompletion", "completion"], - model: str, - messages: Sequence[Mapping[str, object]], - kwargs: Kwargs, - ) -> _Call: + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: forwards litellm.aresponses' kwargs + ) -> object: + """`factory_function`'s wrapper drops its own `custom_llm_provider` and `client`.""" + del custom_llm_provider, client + return await self._route(self._call("responses", _model(kwargs), kwargs), True) + + async def aanthropic_messages( + self, + custom_llm_provider: str | None = None, + client: object = None, + **kwargs: object, # kwargs-ok: forwards litellm.anthropic_messages' kwargs + ) -> object: + """`factory_function`'s wrapper drops its own `custom_llm_provider` and `client`.""" + del custom_llm_provider, client + if kwargs.get("stream") is True: + raise NotImplementedError("the Rust router backend does not stream Anthropic messages yet") + return await self._route(self._call("anthropic_messages", _model(kwargs), kwargs), True) + + async def _route(self, call: _Call, asynchronous: bool) -> object: + native_route: Final = self._native.route(call.spec, call.driver, asynchronous) + routed: Final = cast(Awaitable[object], native_route) # cast-ok: an asynchronous route is a coroutine + return await routed + + def _call(self, operation: Operation, model: str, kwargs: Kwargs) -> _Call: unsupported: Final = [name for name in _UNSUPPORTED_REQUEST_KWARGS if name in kwargs] if unsupported: raise NotImplementedError(f"the Rust router backend does not support request option {unsupported[0]!r} yet") @@ -275,22 +304,17 @@ class RustRouter: if clientside: raise NotImplementedError("the Rust router backend does not support client-side credentials yet") kwargs["model"] = model # rebind-ok: PythonRouter stamps the caller's kwargs the same way - kwargs["messages"] = messages # rebind-ok: as above - kwargs.setdefault("stream", False) - self._normalizer._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) # pyright: ignore[reportPrivateUsage] # PythonRouter's entry stamps + if operation == "completion": + kwargs.setdefault("stream", False) + self._normalizer._update_kwargs_before_fallbacks( # pyright: ignore[reportPrivateUsage] # PythonRouter's entry stamps + model=model, kwargs=kwargs, metadata_variable_name=metadata_key_for(operation) + ) overrides: Final = { name: kwargs[name] for name in ("fallbacks", "context_window_fallbacks", "content_policy_fallbacks") if name in kwargs } mock: Final = next((failure for flag, failure in _MOCK_FAILURES if _flag(kwargs.get(flag))), None) - driver: Final = RoutedCall( - self._normalizer, - kwargs, - "metadata", - lambda name: overrides.get(name, _attribute(self._normalizer, name)), - ) - driver.start(model) tools: Final = (as_mapping(tool) for tool in as_sequence(kwargs.get("tools"))) spec: Final = { "operation": operation, @@ -302,8 +326,30 @@ class RustRouter: "web_search": any(tool.get("type") in _WEB_SEARCH_TOOLS for tool in tools), "mock": mock, } + driver: Final = self._driver(operation, kwargs, overrides, spec) + driver.start(model) return _Call(spec, driver) + def _driver( + self, operation: Operation, kwargs: Kwargs, overrides: Mapping[str, object], spec: Mapping[str, object] + ) -> RoutedCall: + async def resume(point: Mapping[str, object], hop_kwargs: Kwargs) -> object: + """A stream that failed after its content reached the caller falls back as its own call.""" + model: Final = str(point["model_group"]) + self._normalizer._update_kwargs_before_fallbacks( # pyright: ignore[reportPrivateUsage] # as the stream wrapper re-enters the chain + model=model, kwargs=hop_kwargs, metadata_variable_name=metadata_key_for(operation) + ) + driver: Final = self._driver(operation, hop_kwargs, overrides, spec) + return await self._route(_Call({**spec, "model": model, "mock": None, "resume": point}, driver), True) + + return RoutedCall( + self._normalizer, + kwargs, + operation, + lambda name: overrides.get(name, _attribute(self._normalizer, name)), + resume, + ) + @dataclass(frozen=True, slots=True) class _Call: @@ -311,6 +357,13 @@ class _Call: driver: RoutedCall +def _model(kwargs: Kwargs) -> str: + model: Final = kwargs.pop("model", None) + if not isinstance(model, str): + raise TypeError("the router needs a model group to route to") + return model + + def _attribute(owner: object, name: str) -> object: return cast(object, getattr(owner, name)) # cast-ok: getattr is typed Any diff --git a/litellm/router_backends/rust_streams.py b/litellm/router_backends/rust_streams.py new file mode 100644 index 00000000000..078503b7ee6 --- /dev/null +++ b/litellm/router_backends/rust_streams.py @@ -0,0 +1,204 @@ +"""The streams a Rust-routed attempt hands back, read until the attempt has committed. + +An attempt's stream is read until its first content (or, for the Responses API, its first +event past the opening lifecycle events) before the attempt reports success, so a failure +before then can still go to another deployment. What was read replays ahead of the rest. +""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Iterator, Sequence +from typing import Final, Literal, cast + +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator +from litellm.router_backends.python_router import ( + FallbackAwareStreamWrapper, + PythonRouter, + _responses_stream_holds_event, # pyright: ignore[reportPrivateUsage] # the hold rule Python's Responses wrapper applies + _stream_chunks_have_generated_content, # pyright: ignore[reportPrivateUsage] # the content rule Python's stream wrapper applies +) +from litellm.router_utils.responses_fallback_stream import FallbackResponsesStreamWrapper +from litellm.types.llms.openai import ResponseAPIUsage +from litellm.types.utils import ModelResponseStream + +StreamFailureKind = Literal["before_content", "terminal"] + + +class StreamFailed(Exception): + """A stream that failed before it committed. `partial_usage` is what the failed Responses + stream had used, which Python adds to whatever falls back from it.""" + + def __init__( + self, error: BaseException, kind: StreamFailureKind, partial_usage: ResponseAPIUsage | None = None + ) -> None: + super().__init__(str(error)) + self.error: Final = error + self.kind: Final = kind + self.partial_usage: Final = partial_usage + + +def unwrapped(error: BaseException) -> BaseException: + """The provider error behind a mid-stream fallback error, which Python's stream wrappers raise in its place.""" + if isinstance(error, MidStreamFallbackError) and error.original_exception is not None: + return error.original_exception + return error + + +class ReplayStream(FallbackAwareStreamWrapper): + """A chat attempt's stream, replaying the chunks read while waiting for its first content. + + After content has reached the caller the stream cannot fall back, so a mid-stream + fallback error surfaces as the error that caused it, as the Python router's wrapper does. + """ + + def __init__(self, inner: CustomStreamWrapper, buffered: Sequence[ModelResponseStream]) -> None: + super().__init__( # pyright: ignore[reportUnknownMemberType] # CustomStreamWrapper's untyped constructor + completion_stream=cast(object, inner.completion_stream), # cast-ok: untyped attribute + model=cast(object, inner.model), # cast-ok: untyped attribute + custom_llm_provider=inner.custom_llm_provider, + logging_obj=inner.logging_obj, + _response_headers=getattr(inner, "_response_headers", None), + ) + self._inner: Final = inner + self._replay: Final = iter(tuple(buffered)) + self.chunks = cast(list[ModelResponseStream], inner.chunks) # cast-ok: untyped attribute + self._hidden_params = dict(cast(dict[str, object], inner._hidden_params)) # pyright: ignore[reportPrivateUsage] # cast-ok: the stream's own untyped metadata + + def __aiter__(self) -> ReplayStream: + return self + + async def __anext__(self) -> ModelResponseStream: + buffered: Final = next(self._replay, None) + if buffered is not None: + return buffered + try: + return await self._inner.__anext__() + except MidStreamFallbackError as error: + raise unwrapped(error) from error + + def __iter__(self) -> Iterator[ModelResponseStream]: + return self + + def __next__(self) -> ModelResponseStream: + buffered: Final = next(self._replay, None) + if buffered is not None: + return buffered + try: + return self._inner.__next__() + except MidStreamFallbackError as error: + raise unwrapped(error) from error + + async def aclose(self) -> None: + await self._inner.aclose() + + +async def chat_until_content(stream: CustomStreamWrapper) -> ReplayStream: + buffered: Final[list[ModelResponseStream]] = [] # mutable-ok: chunks read before the first content + try: + async for chunk in stream: + buffered.append(chunk) + if _stream_chunks_have_generated_content([chunk]): + break + except MidStreamFallbackError as error: + raise StreamFailed(error, "before_content") from error + except Exception as error: + raise StreamFailed(error, "terminal") from error + return ReplayStream(stream, buffered) + + +def chat_until_content_sync(stream: CustomStreamWrapper) -> ReplayStream: + buffered: Final[list[ModelResponseStream]] = [] # mutable-ok: chunks read before the first content + try: + for chunk in stream: + buffered.append(chunk) + if _stream_chunks_have_generated_content([chunk]): + break + except MidStreamFallbackError as error: + raise StreamFailed(error, "before_content") from error + except Exception as error: + raise StreamFailed(error, "terminal") from error + return ReplayStream(stream, buffered) + + +ResponsesFallback = Callable[[MidStreamFallbackError], Awaitable[object]] + + +async def responses_until_output( + stream: BaseResponsesAPIStreamingIterator, + fallback: ResponsesFallback, + prior_usage: Sequence[ResponseAPIUsage], +) -> ResponsesReplay: + """Holds the opening lifecycle events back until the stream's first output, as Python's + Responses wrapper does, so a failure before then falls back without announcing a response.""" + held: tuple[object, ...] = () # rebind-ok: grows until the first output + try: + async for item in _items(stream): + held = (*held, item) + if not _responses_stream_holds_event(item, len(held) - 1): + break + except MidStreamFallbackError as error: + raise StreamFailed( + error, + "before_content", + PythonRouter._extract_partial_responses_usage(stream), # pyright: ignore[reportPrivateUsage] # Python's partial-usage reader + ) from error + except Exception as error: + raise StreamFailed(error, "terminal") from error + return ResponsesReplay(held, stream, fallback, prior_usage) + + +class ResponsesReplay(FallbackResponsesStreamWrapper): + """A Responses attempt's stream, replaying what was held. After output a mid-stream fallback + error still falls back through `fallback`, with the text generated so far to continue from.""" + + def __init__( + self, + held: Sequence[object], + source: BaseResponsesAPIStreamingIterator, + fallback: ResponsesFallback, + prior_usage: Sequence[ResponseAPIUsage], + ) -> None: + super().__init__(self._replay(held, source, fallback, prior_usage), source) + + async def _replay( + self, + held: Sequence[object], + source: BaseResponsesAPIStreamingIterator, + fallback: ResponsesFallback, + prior_usage: Sequence[ResponseAPIUsage], + ) -> AsyncGenerator[object, None]: + for item in held: + yield _with_usage(item, prior_usage) + try: + async for item in _items(source): + yield _with_usage(item, prior_usage) + except MidStreamFallbackError as error: + partial: Final = PythonRouter._extract_partial_responses_usage(source) # pyright: ignore[reportPrivateUsage] # Python's partial-usage reader + usage: Final = (*prior_usage, partial) if partial is not None else tuple(prior_usage) + try: + response: Final = await fallback(error) + except Exception as fallback_error: + raise unwrapped(fallback_error) from fallback_error + prepared: Final = self.adopt_fallback_headers(response) + if not hasattr(response, "__aiter__"): + yield response + return + async for item in cast(AsyncIterator[object], response): # cast-ok: checked above + PythonRouter._apply_fallback_hidden_params_to_item(item, prepared) # pyright: ignore[reportPrivateUsage] # as Python's wrapper + yield _with_usage(item, usage) + finally: + close: Final[object] = getattr(source, "aclose", None) + if callable(close): + await cast(Callable[[], Awaitable[object]], close)() # cast-ok: an async iterator's aclose + + +def _items(stream: BaseResponsesAPIStreamingIterator) -> AsyncIterator[object]: + return cast(AsyncIterator[object], stream) # cast-ok: every Responses stream iterator is async-iterable + + +def _with_usage(item: object, usage: Sequence[ResponseAPIUsage]) -> object: + for partial in usage: + PythonRouter._combine_responses_fallback_usage(item, partial) # pyright: ignore[reportPrivateUsage, reportArgumentType] # Python's usage merge, a no-op on events without usage + return item diff --git a/litellm/router_utils/add_retry_fallback_headers.py b/litellm/router_utils/add_retry_fallback_headers.py index 6e07693b7ea..d404a200f97 100644 --- a/litellm/router_utils/add_retry_fallback_headers.py +++ b/litellm/router_utils/add_retry_fallback_headers.py @@ -209,6 +209,15 @@ def get_hidden_params_dict( return hidden_params +def prepare_fallback_hidden_params(fallback_response: object) -> tuple[dict[str, object], dict[str, object]]: + """A fallback response's hidden params and its additional headers, which a stream wrapper adopts.""" + fallback_hidden_params: Final = get_hidden_params_dict(fallback_response) + fallback_headers: Final = fallback_hidden_params.get("additional_headers") + if not isinstance(fallback_headers, dict): + return fallback_hidden_params, {} + return fallback_hidden_params, cast("dict[str, object]", fallback_headers) + + def _write_hidden_params(response: object, hidden_params: dict[str, object]) -> None: if isinstance(response, dict): response["_hidden_params"] = hidden_params diff --git a/litellm/router_utils/responses_fallback_stream.py b/litellm/router_utils/responses_fallback_stream.py new file mode 100644 index 00000000000..eddfabba4b4 --- /dev/null +++ b/litellm/router_utils/responses_fallback_stream.py @@ -0,0 +1,85 @@ +"""The Responses API stream a router hands back when its attempts' streams can fall back mid-stream.""" + +import time +from collections.abc import AsyncGenerator +from datetime import datetime +from typing import Final + +from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + _get_openai_response_types, # pyright: ignore[reportPrivateUsage] # the lazily imported OpenAI event types +) +from litellm.router_utils.add_retry_fallback_headers import prepare_fallback_hidden_params + +_OPENAI_TYPES: Final = _get_openai_response_types() +_RESPONSES_TERMINAL_EVENT_TYPES: Final = ( + _OPENAI_TYPES.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, + _OPENAI_TYPES.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, + _OPENAI_TYPES.ResponsesAPIStreamEvents.RESPONSE_FAILED, +) + + +class FallbackResponsesStreamWrapper(BaseResponsesAPIStreamingIterator): + """ + Subclasses BaseResponsesAPIStreamingIterator only for isinstance compatibility (proxy and + interactions code paths check the type). Bypasses the parent constructor and delegates + iteration to an async generator. + """ + + fallback_headers_adopted: bool = False + + def __init__(self, async_generator: AsyncGenerator[object, None], source_iterator: object) -> None: + self._async_generator: Final = async_generator + self._source_iterator: Final = source_iterator + # The bridge path (LiteLLMCompletionStreamingIterator) skips super().__init__ and lacks + # many of these attributes, and stores its logging object as `litellm_logging_obj`. + self.response = getattr(source_iterator, "response", None) + self.model = getattr(source_iterator, "model", None) + self.logging_obj = getattr( + source_iterator, + "logging_obj", + getattr(source_iterator, "litellm_logging_obj", None), + ) + self.finished = False + self.responses_api_provider_config = getattr(source_iterator, "responses_api_provider_config", None) + self.completed_response = None + self.start_time = getattr(source_iterator, "start_time", datetime.now()) + self._failure_handled = False + self._yielded_first_chunk = False + self._generated_content = "" + self._completed_response_cached = False + self._completed_response_logged = False + self._completed_response_cache_hit = None + self._persist_completed_response_before_logging = True + self._stream_created_time = time.time() + self.litellm_metadata = getattr(source_iterator, "litellm_metadata", None) + self.custom_llm_provider = getattr(source_iterator, "custom_llm_provider", None) + self.request_data = getattr(source_iterator, "request_data", {}) or {} + self.call_type = getattr(source_iterator, "call_type", None) + self._hidden_params = dict(getattr(source_iterator, "_hidden_params", None) or {}) + + def adopt_fallback_headers(self, fallback_response: object) -> tuple[dict[str, object], dict[str, object]]: + prepared: Final = prepare_fallback_hidden_params(fallback_response) + self._hidden_params = {**prepared[0], "additional_headers": prepared[1]} + self.fallback_headers_adopted = True + return prepared + + def __aiter__(self) -> "FallbackResponsesStreamWrapper": + return self + + async def __anext__(self) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # yields whatever the wrapped stream yields + try: + chunk: Final = await self._async_generator.__anext__() + except StopAsyncIteration: + # The bridge path emits its final response.completed and then stops, so a terminal + # event this wrapper never saw is read off whatever the source iterator latched. + if self.completed_response is None: + self.completed_response = getattr(self._source_iterator, "completed_response", None) + raise + # The proxy's container-ownership hook reads `completed_response` off the stream it got (#30210) + if self.completed_response is None and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: + self.completed_response = chunk # pyright: ignore[reportAttributeAccessIssue] # the terminal event itself + return chunk + + async def aclose(self) -> None: + await self._async_generator.aclose() diff --git a/tests/test_litellm_rust/router/test_responses_routing.py b/tests/test_litellm_rust/router/test_responses_routing.py new file mode 100644 index 00000000000..aa4d0badecf --- /dev/null +++ b/tests/test_litellm_rust/router/test_responses_routing.py @@ -0,0 +1,275 @@ +"""The Rust router backend against the Python one on the Responses API: the same attempts, the +same fallbacks, and the same events reaching the caller.""" + +from __future__ import annotations + +from collections.abc import Callable, Mapping, Sequence +from types import SimpleNamespace +from typing import Final +from unittest.mock import patch + +import pytest + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator +from tests.test_litellm_rust.support.router_backends import BACKENDS, Backend, router_headers + +pytestmark = pytest.mark.requires_rust_extension + +GROUPS: Final = ("primary", "fb1", "fb2") + + +def _arguments(fallbacks: Sequence[Mapping[str, Sequence[str]]], num_retries: int = 1) -> Mapping[str, object]: + return { + "model_list": tuple( + { + "model_name": group, + "litellm_params": {"model": f"openai/{group}-model", "api_key": "k", "mock_response": f"from {group}"}, + } + for group in GROUPS + ), + "fallbacks": [dict(entry) for entry in fallbacks], + "num_retries": num_retries, + } + + +def _event(kind: str, text: str = "") -> SimpleNamespace: + return SimpleNamespace(type=kind, delta=text) + + +class _Stream(BaseResponsesAPIStreamingIterator): + """A provider stream: `events`, then `failure` if there is one.""" + + def __init__(self, model: str, events: Sequence[SimpleNamespace], failure: MidStreamFallbackError | None) -> None: + self.model = model + self.completed_response = None + self._hidden_params = {} + self._events: Final = iter(tuple(events)) + self._failure: Final = failure + + def __aiter__(self) -> _Stream: + return self + + async def __anext__(self) -> SimpleNamespace: + event: Final = next(self._events, None) + if event is not None: + return event + if self._failure is not None: + raise self._failure + raise StopAsyncIteration + + async def aclose(self) -> None: + return None + + +def _mid_stream(model: str, generated: str) -> MidStreamFallbackError: + return MidStreamFallbackError( + message=f"provider 500 from {model}", + model=model, + llm_provider="openai", + generated_content=generated, + is_pre_first_chunk=not generated, + original_exception=litellm.InternalServerError( + message=f"provider 500 from {model}", model=model, llm_provider="openai" + ), + ) + + +Script = Mapping[str, str] + + +def _provider( + script: Script, calls: list[tuple[str, object]] +) -> Callable[..., object]: # mutable-ok: records each call + """`litellm.aresponses` answering each group by its script: `ok`, `raise`, `drop-before-output` + (an announcement, then a mid-stream error) or `drop-after-output` (some text, then the error).""" + + async def aresponses(**kwargs: object) -> object: # kwargs-ok: litellm.aresponses' surface + metadata: Final = kwargs["litellm_metadata"] + assert isinstance(metadata, Mapping) + group: Final = str(metadata["model_group"]) + model: Final = str(kwargs["model"]) + calls.append((group, kwargs.get("input"))) + behavior: Final = script.get(group, "ok") + if behavior == "raise": + raise litellm.InternalServerError(message=f"500 from {group}", model=model, llm_provider="openai") + if behavior == "drop-before-output": + return _Stream(model, (_event("response.created"),), _mid_stream(model, "")) + if behavior == "drop-after-output": + return _Stream( + model, + (_event("response.created"), _event("response.output_text.delta", f"partial from {group}")), + _mid_stream(model, f"partial from {group}"), + ) + return _Stream( + model, + ( + _event("response.created"), + _event("response.output_text.delta", f"from {group}"), + _event("response.completed"), + ), + None, + ) + + return aresponses + + +Observed = tuple[tuple[str, ...], tuple[tuple[str, object], ...], Mapping[str, object]] + + +async def _stream( + build: Callable[[Mapping[str, object], int], Backend], arguments: Mapping[str, object], script: Script +) -> Observed: + """`build` runs under the patch: Python's router binds `litellm.aresponses` when it is built.""" + calls: Final[list[tuple[str, object]]] = [] # mutable-ok: one entry per provider call + events: Final[list[str]] = [] # mutable-ok: the events in the order the caller got them + response: object = None # rebind-ok: set once the call returns + with patch("litellm.aresponses", side_effect=_provider(script, calls)): + backend: Final = build(arguments, 1) + try: + response = await backend.aresponses(model="primary", input="hi", stream=True) # rebind-ok: see above + async for event in response: # pyright: ignore[reportGeneralTypeIssues,reportUnknownVariableType] # an async stream + events.append(f"{event.type}:{event.delta}") # pyright: ignore[reportUnknownMemberType] # the scripted events + except litellm.InternalServerError as error: + events.append(f"raised {error.message}") + return tuple(events), tuple(calls), router_headers(response) + + +@pytest.mark.parametrize( + ("fallbacks", "script", "expected_events", "expected_groups"), + ( + pytest.param( + [{"primary": ["fb1"]}], + {"primary": "drop-before-output"}, + ("response.created:", "response.output_text.delta:from fb1", "response.completed:"), + ("primary", "fb1"), + id="drop-before-output-falls-back-without-retrying", + ), + pytest.param( + [{"primary": ["fb1"]}], + {"primary": "drop-after-output"}, + ( + "response.created:", + "response.output_text.delta:partial from primary", + "response.created:", + "response.output_text.delta:from fb1", + "response.completed:", + ), + ("primary", "fb1"), + id="drop-after-output-continues-on-the-fallback", + ), + pytest.param( + [{"primary": ["fb1"]}, {"fb1": ["fb2"]}], + {"primary": "drop-after-output", "fb1": "drop-before-output"}, + ( + "response.created:", + "response.output_text.delta:partial from primary", + "response.created:", + "response.output_text.delta:from fb2", + "response.completed:", + ), + ("primary", "fb1", "fb2"), + id="a-fallback-that-drops-uses-its-own-chain", + ), + pytest.param( + [{"primary": ["fb1"]}], + {"primary": "drop-after-output", "fb1": "raise"}, + ( + "response.created:", + "response.output_text.delta:partial from primary", + "raised litellm.InternalServerError: provider 500 from openai/primary-model", + ), + ("primary", "fb1", "fb1"), + id="the-provider-error-surfaces-when-the-fallback-fails", + ), + ), +) +async def test_a_dropped_responses_stream_falls_back_the_same_way_on_both_backends( + fallbacks: Sequence[Mapping[str, Sequence[str]]], + script: Script, + expected_events: tuple[str, ...], + expected_groups: tuple[str, ...], +) -> None: + observed: Final = [await _stream(build, _arguments(fallbacks), script) for build in BACKENDS] + + assert observed[1][:2] == observed[0][:2] + events, calls, _ = observed[1] + assert events == expected_events + assert tuple(group for group, _ in calls) == expected_groups + + +async def test_a_fallback_after_output_reports_the_fallback_in_both_backends_headers() -> None: + observed: Final = [ + await _stream(build, _arguments([{"primary": ["fb1"]}]), {"primary": "drop-after-output"}) for build in BACKENDS + ] + + assert observed[1][2] == observed[0][2] + assert observed[1][2] == { + "x-litellm-model-group": "fb1", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + } + + +async def test_headers_name_the_group_that_served_after_a_fallback_falls_back_again() -> None: + """Python's Responses wrapper adopts its fallback's headers before reading it, so when that + fallback falls back again it still names the first one; its chat wrapper re-adopts them once the + stream yields. The Rust backend names the group that served, as chat does.""" + _, _, headers = await _stream( + BACKENDS[1], + _arguments([{"primary": ["fb1"]}, {"fb1": ["fb2"]}]), + {"primary": "drop-after-output", "fb1": "drop-before-output"}, + ) + + assert headers == { + "x-litellm-model-group": "fb2", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 2, + } + + +async def test_a_fallback_after_output_continues_from_the_generated_text() -> None: + calls: Final[list[tuple[str, object]]] = [] # mutable-ok: one entry per provider call + backend: Final = BACKENDS[1](_arguments([{"primary": ["fb1"]}]), 1) + with patch("litellm.aresponses", side_effect=_provider({"primary": "drop-after-output"}, calls)): + response = await backend.aresponses(model="primary", input="hi", stream=True) + _ = [event async for event in response] # pyright: ignore[reportGeneralTypeIssues,reportUnknownVariableType] # an async stream + + fallback_input: Final = calls[1][1] + assert isinstance(fallback_input, list) + assert fallback_input[-1] == { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "partial from primary"}], + } + + +@pytest.mark.parametrize("build", BACKENDS) +async def test_a_failing_group_falls_back_to_its_mock_response( + build: Callable[[Mapping[str, object], int], Backend], +) -> None: + backend: Final = build(_arguments([{"primary": ["fb1"]}], num_retries=0), 1) + + response: Final = await backend.aresponses(model="primary", input="hi", mock_testing_fallbacks=True) + + assert response.output[0].content[0].text == "from fb1" # pyright: ignore[reportAttributeAccessIssue,reportUnknownMemberType] # ResponsesAPIResponse + assert router_headers(response) == { + "x-litellm-model-group": "fb1", + "x-litellm-attempted-retries": 0, + "x-litellm-attempted-fallbacks": 1, + } + + +@pytest.mark.parametrize("build", BACKENDS) +async def test_anthropic_messages_fall_back_to_the_fallback_groups_mock_response( + build: Callable[[Mapping[str, object], int], Backend], +) -> None: + backend: Final = build(_arguments([{"primary": ["fb1"]}], num_retries=0), 1) + + response: Final = await backend.aanthropic_messages( + model="primary", messages=[{"role": "user", "content": "hi"}], max_tokens=10, mock_testing_fallbacks=True + ) + + assert isinstance(response, Mapping) + assert response["content"][0]["text"] == "from fb1" diff --git a/tests/test_litellm_rust/router/test_rust_router.py b/tests/test_litellm_rust/router/test_rust_router.py index 71931d16a1e..f3abc6bb35b 100644 --- a/tests/test_litellm_rust/router/test_rust_router.py +++ b/tests/test_litellm_rust/router/test_rust_router.py @@ -11,8 +11,7 @@ from unittest.mock import MagicMock, patch import pytest import litellm -from litellm.router_backends.python_router import PythonRouter -from litellm.router_backends.rust_router import NATIVE_ROUTER, RustRouter +from litellm.router_backends.rust_router import RustRouter from litellm.router_backends.selection import select_backend from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.configuration import Rollout @@ -20,22 +19,19 @@ from litellm.types.router import RouterRateLimitError from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.utils import ModelResponse, ModelResponseStream +from tests.test_litellm_rust.support.router_backends import ( + BACKENDS, + Backend, + python_backend, + router_headers, + rust_backend, +) pytestmark = pytest.mark.requires_rust_extension MESSAGES: Final = ({"role": "user", "content": "hi"},) -class Backend(Protocol): - async def acompletion( - self, model: str, messages: Sequence[Mapping[str, object]], **kwargs: object - ) -> object: ... # kwargs-ok: Router.acompletion's surface - - def completion( - self, model: str, messages: Sequence[Mapping[str, object]], **kwargs: object - ) -> object: ... # kwargs-ok: Router.completion's surface - - def _deployment( group: str, name: str, response: object = None ) -> dict[str, object]: # mutable-ok: PythonRouter consumes model_list dicts @@ -46,23 +42,9 @@ def _deployment( } -def _python(arguments: Mapping[str, object], seed: int) -> Backend: - random.seed(seed) - router: Final = PythonRouter( - **{key: list(value) if isinstance(value, tuple) else value for key, value in arguments.items()} - ) - return router # pyright: ignore[reportReturnType] # PythonRouter.acompletion's untyped signature - - -def _rust(arguments: Mapping[str, object], seed: int) -> Backend: - native: Final = NATIVE_ROUTER.load() - assert native is not None - return RustRouter( - {key: list(value) if isinstance(value, tuple) else value for key, value in arguments.items()}, native, seed - ) - - -BACKENDS: Final[tuple[Callable[[Mapping[str, object], int], Backend], ...]] = (_python, _rust) +_python: Final = python_backend +_rust: Final = rust_backend +_headers: Final = router_headers def _content(response: object) -> str: @@ -70,12 +52,6 @@ def _content(response: object) -> str: return str(response.choices[0].message.content) # pyright: ignore[reportAttributeAccessIssue] # non-streaming choice -def _headers(response: object) -> Mapping[str, object]: - hidden: Final = getattr(response, "_hidden_params", {}) - headers: Final = hidden.get("additional_headers", {}) if isinstance(hidden, Mapping) else {} - return {key: value for key, value in headers.items() if key.startswith("x-litellm-")} - - async def _picks(backend: Backend, count: int) -> list[str]: return [_content(await backend.acompletion("g", MESSAGES)) for _ in range(count)] diff --git a/tests/test_litellm_rust/support/router_backends.py b/tests/test_litellm_rust/support/router_backends.py new file mode 100644 index 00000000000..013a13c2a01 --- /dev/null +++ b/tests/test_litellm_rust/support/router_backends.py @@ -0,0 +1,50 @@ +"""Both router backends built from the same arguments and seed, for parity tests.""" + +from __future__ import annotations + +import random +from collections.abc import Callable, Mapping, Sequence +from typing import Final, Protocol + +from litellm.router_backends.python_router import PythonRouter +from litellm.router_backends.rust_router import NATIVE_ROUTER, RustRouter + + +class Backend(Protocol): + async def acompletion( + self, model: str, messages: Sequence[Mapping[str, object]], **kwargs: object + ) -> object: ... # kwargs-ok: Router.acompletion's surface + + def completion( + self, model: str, messages: Sequence[Mapping[str, object]], **kwargs: object + ) -> object: ... # kwargs-ok: Router.completion's surface + + async def aresponses(self, **kwargs: object) -> object: ... # kwargs-ok: Router.aresponses' surface + + async def aanthropic_messages( + self, **kwargs: object + ) -> object: ... # kwargs-ok: Router.aanthropic_messages' surface + + +def _lists(arguments: Mapping[str, object]) -> dict[str, object]: # mutable-ok: Router(...) keyword arguments + return {key: list(value) if isinstance(value, tuple) else value for key, value in arguments.items()} + + +def python_backend(arguments: Mapping[str, object], seed: int) -> Backend: + random.seed(seed) + return PythonRouter(**_lists(arguments)) # pyright: ignore[reportReturnType] # PythonRouter's untyped signatures + + +def rust_backend(arguments: Mapping[str, object], seed: int) -> Backend: + native: Final = NATIVE_ROUTER.load() + assert native is not None + return RustRouter(_lists(arguments), native, seed) + + +BACKENDS: Final[tuple[Callable[[Mapping[str, object], int], Backend], ...]] = (python_backend, rust_backend) + + +def router_headers(response: object) -> Mapping[str, object]: + hidden: Final = getattr(response, "_hidden_params", {}) + headers: Final = hidden.get("additional_headers", {}) if isinstance(hidden, Mapping) else {} + return {key: value for key, value in headers.items() if key.startswith("x-litellm-")} diff --git a/tests/unit/router_backends/test_rust_call.py b/tests/unit/router_backends/test_rust_call.py index bf602289e9e..4c9953b65be 100644 --- a/tests/unit/router_backends/test_rust_call.py +++ b/tests/unit/router_backends/test_rust_call.py @@ -28,7 +28,7 @@ def _call(normalizer: PythonRouter, metadata: dict[str, object]) -> RoutedCall: call: Final = RoutedCall( normalizer, # pyright: ignore[reportArgumentType] # PythonRouter provides the AttemptRouter members {"model": "g", "messages": [], "metadata": metadata, "num_retries": 3, "fallbacks": []}, - "metadata", + "completion", lambda name: getattr(normalizer, name), ) call.start("g") @@ -195,8 +195,10 @@ def call_attempt(bucket: int) -> Attempt: return Attempt( deployment_id="a", model_group="g", + original_model_group="g", bucket=bucket, fallback_depth=0, + attempted_targets=(), model_group_size=1, attempted_retries=1, max_retries=2, diff --git a/tests/unit/test_responses_streaming_container_ownership.py b/tests/unit/test_responses_streaming_container_ownership.py index 07cc309798b..03f642a9821 100644 --- a/tests/unit/test_responses_streaming_container_ownership.py +++ b/tests/unit/test_responses_streaming_container_ownership.py @@ -19,6 +19,7 @@ proxy hook can walk it. """ import asyncio +import functools from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -75,7 +76,7 @@ def _make_wrapper_class(): response=source, # type: ignore[arg-type] initial_kwargs={}, ) - captured["wrapper_cls"] = type(wrapped) + captured["wrapper_cls"] = functools.partial(type(wrapped), source_iterator=source) captured["instance"] = wrapped asyncio.run(_drive())