scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-12 03:48:56 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 57

feat(router): serve the Responses API and Anthropic messages from the Rust backend

Browse source 
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 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Aw1xktpUnix5tWhvSzey8t
This commit is contained in:
Claude 2026-10-07 03:55:48 +00:00
parent 5c464c45e9
commit 6f25f627a6
No known key found for this signature in database
19 changed files with 1085 additions and 324 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm-rust/crates/python-bridge/src/routes/router/host.rs
View file

@ -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<RouterProtocol>);
@ -85,8 +85,8 @@ impl PythonBinding for RouterBinding {
&mut self,
_py: Python<'_>,
arguments: &Bound<'_, PyDict>,
) -> Result<RouterCall, InvokeError<BridgeError>> {
python::router_call(arguments).map_err(InvokeError::Python)
) -> Result<RouterRequest, InvokeError<BridgeError>> {
python::router_request(arguments).map_err(InvokeError::Python)
}
fn encode_response(

19
litellm-rust/crates/python-bridge/src/routes/router/mod.rs
View file

@ -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<Resume<PyObj>>,
}
pub(crate) struct RouterProtocol;
impl Protocol for RouterProtocol {
type Request = RouterCall;
type Request = RouterRequest;
type Response = Routed<PyObj, PyObj>;
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::<RouterProtocol, _, _>(
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)))

39
litellm-rust/crates/python-bridge/src/routes/router/python.rs
View file

@ -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<RouterCall> {
pub(super) fn router_request(arguments: &Bound<'_, PyDict>) -> PyResult<RouterRequest> {
let resume = match arguments.get_item("resume")? {
Some(resume) if !resume.is_none() => Some(self::resume(resume.cast::<PyDict>()?)?),
_ => 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<Resume<PyObj>> {
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>(&classified)?,
deployment_id: optional(resume, "deployment_id")?,
})
}
fn router_call(arguments: &Bound<'_, PyDict>) -> PyResult<RouterCall> {
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<PyObj>) -> 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)?;

1
litellm-rust/crates/router/POC-behavior-spec.md
View file

@ -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

4
litellm-rust/crates/router/POC.md
View file

@ -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)

85
litellm-rust/crates/router/src/engine.rs
View file

@ -108,6 +108,18 @@ pub struct Failed<E> {
pub ops: Vec<Op<E>>,
}
/// 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<E> {
pub group: String,
pub depth: u32,
pub original_group: String,
pub attempted_targets: Vec<String>,
pub error: E,
pub classified: Classified,
pub deployment_id: Option<String>,
}
pub struct Engine {
registry: Registry,
store: Store,
@ -135,6 +147,27 @@ impl Engine {
&self,
host: &H,
call: RouterCall,
) -> Result<Routed<H::Response, H::Error>, RouteError<H::Error, H::Fault>> {
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<H: RouterHost>(
&self,
host: &H,
call: RouterCall,
resume: Resume<H::Error>,
) -> Result<Routed<H::Response, H::Error>, RouteError<H::Error, H::Fault>> {
self.run(host, call, Some(resume)).await
}
async fn run<H: RouterHost>(
&self,
host: &H,
call: RouterCall,
resume: Option<Resume<H::Error>>,
) -> Result<Routed<H::Response, H::Error>, RouteError<H::Error, H::Fault>> {
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<H::Error>,
) -> Attempted<H::Response, H::Error, H::Fault> {
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<String> {
let mut targets: Vec<String> = 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 {

4
litellm-rust/crates/router/src/host.rs
View file

@ -92,8 +92,12 @@ pub struct Attempt<E> {
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<String>,
pub retry: RetryStamp,
pub ops: Vec<Op<E>>,
}

41
litellm-rust/crates/router/tests/engine.rs
View file

@ -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)
);
}

128
litellm/router_backends/python_router.py
View file

@ -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/<id>/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(

303
litellm/router_backends/rust_call.py
View file

@ -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":

101
litellm/router_backends/rust_router.py
View file

@ -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

204
litellm/router_backends/rust_streams.py Normal file
View file

@ -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

9
litellm/router_utils/add_retry_fallback_headers.py
View file

@ -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

85
litellm/router_utils/responses_fallback_stream.py Normal file
View file

@ -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()

275
tests/test_litellm_rust/router/test_responses_routing.py Normal file
View file

@ -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"

46
tests/test_litellm_rust/router/test_rust_router.py
View file

@ -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)]

50
tests/test_litellm_rust/support/router_backends.py Normal file
View file

@ -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-")}

4
tests/unit/router_backends/test_rust_call.py
View file

@ -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,

3
tests/unit/test_responses_streaming_container_ownership.py
View file

@ -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())