scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-06 07:48:06 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 112

refactor(router): delete the python router core, the rust core is the only core (#50178)

Browse source
This commit is contained in:
Nate Armstrong 2026-11-05 10:40:46 -08:00 • committed by GitHub
parent 85bc58085e
commit 067f93340f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
111 changed files with 4758 additions and 42639 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm-rust/crates/python-bridge/src/router/config.rs
View file

@ -188,6 +188,8 @@ struct RequestSpec {
web_search: Option<bool>,
route_handler: Option<String>,
decisions_media: Option<bool>,
batch_create: Option<bool>,
embedding_batch: Option<bool>,
wildcard_request: Option<String>,
unrouted: Option<bool>,
passthrough_on_no_deployment: Option<bool>,
@ -255,6 +257,8 @@ pub(crate) fn parse_request(value: Value, model: Option<&str>) -> Result<RouterR
web_search: spec.web_search.unwrap_or(false),
route_handler: spec.route_handler,
decisions_media: spec.decisions_media.unwrap_or(false),
batch_create: spec.batch_create.unwrap_or(false),
embedding_batch: spec.embedding_batch.unwrap_or(false),
wildcard_request: spec.wildcard_request,
unrouted: spec.unrouted.unwrap_or(false),
passthrough_on_no_deployment: spec.passthrough_on_no_deployment.unwrap_or(false),

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

@ -68,7 +68,12 @@ pub(crate) enum Stage {
/// Pre-call checks and sub-routers: return the deployments that remain.
Plugins,
/// A strategy router-core does not implement: return the candidates with the pick first.
Strategy(String),
/// When candidates share an id, `records` holds each candidate's deployment, in order, and the
/// host answers with the pick's position, so a pick splits over every entry as Python's does.
Strategy {
strategy: String,
records: Vec<Value>,
},
/// Tag routing: return the deployments that remain.
Tags,
/// `enable_pre_call_checks`: return the facts the filter decides by.
@ -79,6 +84,10 @@ pub(crate) enum Stage {
RouteIncapable { handler: String },
/// `_decisions_deployment_refuses_media`: return the candidates that refuse the call's media.
MediaRefusing,
/// `deployments_admitting_embedding_batch`: return the candidates the embedding pick keeps.
EmbeddingBatch,
/// The override strategy's pre-call check on a deployment named by id.
OverrideCheck(String),
/// The routing plugins: run them when `run`, then return the candidate models they left.
RoutingPlugins { run: bool },
/// The host's `CustomLogger` deployment filters: return the ids that remain, or a failure.
@ -163,6 +172,8 @@ pub(crate) enum HostCall {
/// The name the attempt resolved: the requested name or the fallback target.
model: String,
model_group: String,
/// The group a pre-routing hook selected for `model`, when it differs.
routed: Option<String>,
/// The `order` level of a same-group order fallback.
target_order: Option<i64>,
/// Every deployment `model` resolved to for this caller, the picked one included.
@ -279,6 +290,8 @@ pub(crate) enum HostCallResponse {
RefuseFallbacks(bool),
/// The candidate ids that remain, pick first.
Select(Vec<String>),
/// A strategy's pick among candidates sharing an id: positions in the candidates, pick first.
Positions(Vec<usize>),
/// The model group to resolve instead, a custom strategy's deployment, or `None`.
PreRoute(Option<PreRouted>),
/// The facts a `pre_call` select step asked for.
@ -309,9 +322,11 @@ impl HostCallResponse {
Self::Slept => StepKind::Sleep,
Self::AllowFallback(_) => StepKind::AllowFallback,
Self::RefuseFallbacks(_) => StepKind::RefuseFallbacks,
Self::Select(_) | Self::PreCall(_) | Self::Models(_) | Self::Filtered(_) => {
StepKind::Select
}
Self::Select(_)
| Self::Positions(_)
| Self::PreCall(_)
| Self::Models(_)
| Self::Filtered(_) => StepKind::Select,
Self::PreRoute(_) => StepKind::PreRoute,
Self::Cache(_) => StepKind::Cache,
Self::Secrets(_) => StepKind::Secrets,
@ -343,6 +358,13 @@ impl HostCallResponse {
{
return Err(DriveError::UnknownCandidate { id: id.clone() });
}
if let (Self::Positions(positions), HostCall::Select { candidates, .. }) = (self, call)
&& let Some(position) = positions.iter().find(|p| **p >= candidates.len())
{
return Err(DriveError::UnknownCandidate {
id: format!("#{position}"),
});
}
if let (Self::Secrets(values), HostCall::Secrets { refs, .. }) = (self, call)
&& values.len() != refs.len()
{
@ -654,6 +676,28 @@ impl RouterCache for StepCache {
}
/// A deployment as Python spells a `model_list` entry, with its id in `model_info`.
/// Whether two candidates share an id, which a pick answered by id can't tell apart.
fn shares_ids(candidates: &[Arc<Deployment>]) -> bool {
let mut seen = std::collections::HashSet::new();
!candidates.iter().all(|d| seen.insert(d.id.as_str()))
}
/// The candidates a host answer names, in its order. An id names every candidate that carries
/// it, so deployments sharing an id stay together through a filter, as Python's list keeps them.
fn answered_by_id(candidates: &[Arc<Deployment>], ids: &[String]) -> Vec<Arc<Deployment>> {
let mut taken = vec![false; candidates.len()];
let mut kept = Vec::new();
for id in ids {
for (position, deployment) in candidates.iter().enumerate() {
if !taken[position] && deployment.id == *id {
taken[position] = true;
kept.push(deployment.clone());
}
}
}
kept
}
pub(crate) fn deployment_record(deployment: &Deployment) -> Value {
deployment.record()
}
@ -693,6 +737,7 @@ impl Host for StepHost {
deployment: deployment_record(&attempt.deployment),
model: attempt.name.clone(),
model_group: attempt.model_group.clone(),
routed: attempt.routed.clone(),
target_order: attempt.target_order,
candidate_ids: attempt.candidate_ids.clone(),
retry: attempt.retry,
@ -736,13 +781,22 @@ impl Host for StepHost {
let call = HostCall::Select {
stage: match stage {
SelectStage::Plugins => Stage::Plugins,
SelectStage::Strategy(strategy) => Stage::Strategy(strategy.to_string()),
SelectStage::Strategy(strategy) => Stage::Strategy {
strategy: strategy.to_string(),
records: if shares_ids(&candidates) {
candidates.iter().map(|d| deployment_record(d)).collect()
} else {
Vec::new()
},
},
SelectStage::Tags => Stage::Tags,
SelectStage::TextOnly => Stage::TextOnly,
SelectStage::RouteIncapable { handler } => Stage::RouteIncapable {
handler: handler.clone(),
},
SelectStage::MediaRefusing => Stage::MediaRefusing,
SelectStage::EmbeddingBatch => Stage::EmbeddingBatch,
SelectStage::OverrideCheck(strategy) => Stage::OverrideCheck(strategy.to_string()),
},
pool: ctx.pool.iter().map(|d| d.id.clone()).collect(),
model_group: ctx.model_group.to_owned(),
@ -751,13 +805,14 @@ impl Host for StepHost {
group: group.map(str::to_owned),
candidates: candidates.iter().map(|d| d.id.clone()).collect(),
};
let HostCallResponse::Select(ids) = answered(call_host(call).await)? else {
unreachable!("{CHECKED}");
};
Ok(ids
.iter()
.filter_map(|id| candidates.iter().find(|d| d.id == *id).cloned())
.collect())
match answered(call_host(call).await)? {
HostCallResponse::Select(ids) => Ok(answered_by_id(&candidates, &ids)),
HostCallResponse::Positions(positions) => Ok(positions
.iter()
.map(|position| candidates[*position].clone())
.collect()),
_ => unreachable!("{CHECKED}"),
}
}
async fn pre_route(&self, pre: &PreRoute<'_>) -> Result<Option<PreRouted>, Error> {
@ -855,10 +910,7 @@ impl Host for StepHost {
unreachable!("{CHECKED}");
};
let ids = result.map_err(Error::Attempt)?;
Ok(ids
.iter()
.filter_map(|id| candidates.iter().find(|d| d.id == *id).cloned())
.collect())
Ok(answered_by_id(&candidates, &ids))
}
async fn allow_fallback(
@ -922,6 +974,7 @@ impl Host for StepHost {
| RouterEvent::Blocked { .. }
| RouterEvent::CooldownSet { .. }
| RouterEvent::Selected { .. }
| RouterEvent::PickStarted { .. }
) {
return;
}
@ -1461,7 +1514,7 @@ mod tests {
let mut driver = start(&router, "primary");
let step = store.resume(&mut driver, None);
let Step::HostCall(HostCall::Select {
stage: Stage::Strategy(strategy),
stage: Stage::Strategy { strategy, records },
candidates,
..
}) = step
@ -1470,6 +1523,7 @@ mod tests {
};
assert_eq!(strategy, "least-busy");
assert_eq!(candidates, ["a1", "a2"]);
assert!(records.is_empty(), "distinct ids send no records");
let step = store.resume(
&mut driver,
Some(HostCallResponse::Select(vec!["a2".to_owned()])),
@ -1480,6 +1534,40 @@ mod tests {
);
}
#[rstest]
fn a_pick_among_deployments_sharing_an_id_is_answered_by_position() {
let mut second = deployment("primary", "a1");
second["litellm_params"]["api_base"] = json!("https://second.example");
let router = router(
json!({
"model_list": [deployment("primary", "a1"), second],
"routing_strategy": "least-busy",
}),
json!({}),
);
let mut store = Store::default();
let mut driver = start(&router, "primary");
let step = store.resume(&mut driver, None);
let Step::HostCall(HostCall::Select {
stage: Stage::Strategy { records, .. },
..
}) = step
else {
panic!("a host strategy should ask for a selection, got {step:?}");
};
assert_eq!(
records.len(),
2,
"each candidate sharing an id sends its record"
);
let step = store.resume(&mut driver, Some(HostCallResponse::Positions(vec![1])));
assert!(
matches!(&step, Step::HostCall(HostCall::Attempt { deployment, .. })
if deployment["litellm_params"]["api_base"] == "https://second.example"),
"expected the second entry, got {step:?}"
);
}
#[rstest]
fn a_selection_is_driven_through_cache_steps(pair: Router) {
let request = RouterRequest {

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

@ -28,7 +28,7 @@ use litellm_router_core::{
use pyo3::{
exceptions::{PyRuntimeError, PyValueError},
prelude::*,
types::{PyBool, PyDict, PyList, PyString},
types::{PyBool, PyDict, PyList, PyString, PyTuple},
};
use serde_json::Value;
@ -457,7 +457,7 @@ impl RouterCoreHandle {
}
/// Record a success a logging callback saw, driven like `report_failure`; it ends with
/// `done` (`key`: the success counter's key, or None for an unknown deployment).
/// `done` (`key`: the tokens counter's key when usage was added, else None).
/// `usage_model` is the model its usage is keyed by, and `tracks_usage` the host's answer to
/// whether it counts any (see `Output::tracks_usage`). `callbacks`: see `report_failure`.
#[pyo3(signature = (id, total_tokens, requests, usage_model=None, tracks_usage=None, callbacks=None))]
@ -1193,6 +1193,12 @@ fn host_call_response(
}
Ok(HostCallResponse::Filtered(Ok(outcome.extract()?)))
}
HostCall::Select {
stage: Stage::Strategy { records, .. },
..
} if !records.is_empty() => Ok(HostCallResponse::Positions(
outcome.ok_or_else(missing)?.extract()?,
)),
HostCall::Select { .. } => Ok(HostCallResponse::Select(
outcome.ok_or_else(missing)?.extract()?,
)),
@ -1285,6 +1291,15 @@ fn pre_route_outcome(outcome: Option<Bound<'_, PyAny>>) -> PyResult<Option<PreRo
if let Ok(model) = outcome.cast::<PyString>() {
return Ok(Some(PreRouted::Model(model.to_str()?.to_owned())));
}
// `(model, group)`: the model to resolve, served as the group a session binding stamped.
if let Ok(pair) = outcome.cast::<PyTuple>()
&& pair.len() == 2
{
return Ok(Some(PreRouted::Bound {
model: pair.get_item(0)?.extract()?,
group: pair.get_item(1)?.extract()?,
}));
}
let spec: DeploymentSpec = serde_json::from_value(document(&outcome)?)
.map_err(|error| PyValueError::new_err(error.to_string()))?;
let deployment = Deployment::from_spec(spec)
@ -1401,7 +1416,7 @@ fn cause_items(dict: &Bound<'_, PyDict>, cause: RetryCause) -> PyResult<()> {
fn pick_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> {
error_items(dict, error)?;
match error {
Error::NoHealthyDeployments { model } | Error::UnknownModel { model } => {
Error::NoHealthyDeployments { model, .. } | Error::UnknownModel { model } => {
dict.set_item("failed_model", model)
}
_ => Ok(()),
@ -1443,6 +1458,10 @@ fn error_items(dict: &Bound<'_, PyDict>, error: &Error) -> PyResult<()> {
dict.set_item("error_model", model)?;
dict.set_item("error_detail", format!("{until} {timezone}"))
}
Error::NoHealthyDeployments {
request_ids: Some(ids),
..
} => dict.set_item("error_request_ids", PyList::new(dict.py(), ids)?),
Error::NonNumericWeight { metric } => dict.set_item("error_detail", metric),
Error::MissingRoutingGroup { group } => dict.set_item("error_model", group),
Error::ModelInfo { exception, detail } => {
@ -1637,6 +1656,10 @@ fn events_list(py: Python<'_>, events: Vec<RouterEvent>) -> PyResult<Bound<'_, P
dict.set_item("deployment_id", deployment_id)?;
dict.set_item("cooldown_time", cooldown.as_secs_f64())?;
}
RouterEvent::PickStarted { model } => {
dict.set_item("kind", "pick_started")?;
dict.set_item("model", model)?;
}
RouterEvent::Selected {
deployment_id,
pick,
@ -1685,6 +1708,7 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) ->
deployment,
model,
model_group,
routed,
target_order,
candidate_ids,
retry,
@ -1696,6 +1720,7 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) ->
max_retries,
model_group_size,
} => {
dict.set_item("routed", routed)?;
dict.set_item("scope", scope)?;
dict.set_item("max_retries", max_retries)?;
dict.set_item("model_group_size", model_group_size)?;
@ -1729,7 +1754,12 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) ->
} => {
let (stage, strategy) = match stage {
Stage::Plugins => ("plugins", None),
Stage::Strategy(strategy) => ("strategy", Some(strategy)),
Stage::Strategy { strategy, records } => {
if !records.is_empty() {
dict.set_item("records", to_py(py, &Value::Array(records))?)?;
}
("strategy", Some(strategy))
}
Stage::Tags => ("tags", None),
Stage::PreCall => ("pre_call", None),
Stage::TextOnly => ("text_only", None),
@ -1738,6 +1768,8 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) ->
("route_incapable", None)
}
Stage::MediaRefusing => ("media_refusing", None),
Stage::EmbeddingBatch => ("embedding_batch", None),
Stage::OverrideCheck(strategy) => ("override_check", Some(strategy)),
Stage::ProviderRequest => ("provider_request", None),
Stage::CallbackFilters {
excluded_ids,

192
litellm-rust/crates/router-core/src/call.rs
View file

@ -29,7 +29,7 @@ use crate::{
host_index::python_str,
lookup::Lookup,
mock::MockFailure,
outcome::OutcomeWrite,
outcome::{OutcomeCallbacks, OutcomeWrite},
prerouting::{PreRoute, PreRouted},
registry::{ResolveKind, ResolveScope},
retry::{RetryContext, resolve_retry_policy, should_retry},
@ -63,6 +63,12 @@ pub struct RouterResponse {
struct Failed {
error: Box<Error>,
deployment_id: Option<String>,
/// The group a pre-routing hook rewrote the level's name to for the failed attempt (an auto
/// router's tier), which its order fallbacks and fallback lookup key off
/// (`get_pre_routing_selection`).
routed: Option<String>,
/// The group the host stamped that attempt as served for ([`Candidates::served_as`]).
served_as: Option<String>,
}
impl From<Error> for Failed {
@ -70,6 +76,8 @@ impl From<Error> for Failed {
Self {
error: Box::new(error),
deployment_id: None,
routed: None,
served_as: None,
}
}
}
@ -103,6 +111,8 @@ struct Target {
/// `_failover_excluded_ids` in the level's metadata: the deployments weighted failover
/// already excluded on the way here.
failover_excluded: Arc<[String]>,
/// The group a pre-routing hook selected for the level's failed attempt, if it rewrote it.
routed: Option<String>,
}
/// A step of a level's fallback chain: a configured target, or an order level of the group.
@ -166,10 +176,18 @@ fn unattempted_chain(
target: &Target,
at: FallbackHop,
attempted: &[Tried],
served: (Option<&str>, Option<&str>),
) -> bool {
let mut lookup = vec![target.group.as_str()];
if request.model != target.group {
lookup.push(request.model.as_str());
let (routed, served_as) = served;
let mut lookup: Vec<&str> = Vec::new();
for group in routed
.into_iter()
.chain(served_as)
.chain([target.group.as_str(), request.model.as_str()])
{
if !lookup.contains(&group) {
lookup.push(group);
}
}
list.lookup_groups(&lookup, lists.keys)
.is_some_and(|targets| {
@ -285,6 +303,7 @@ impl Core {
params: None,
excluded: Arc::from([]),
failover_excluded: Arc::from([]),
routed: None,
};
let at = FallbackHop::default();
// A fallbacks mock stands in for the requested group's retry loop: no pick, no attempt.
@ -343,6 +362,8 @@ impl Core {
let Failed {
error,
deployment_id: failed_id,
routed,
served_as,
} = failed;
let error = *error;
self.level_failed(at, &target, &error, false);
@ -360,7 +381,14 @@ impl Core {
self.level_failed(at, &target, &raised, true);
return Err(raised);
}
let (targets, order_path) = self.fallback_targets(snapshot, request, &target, &error);
let (targets, order_path) = self.fallback_targets(
snapshot,
request,
&target,
routed.as_deref(),
served_as.as_deref(),
&error,
);
if !order_path {
self.host.observe(&RouterEvent::FallingBack {
level: at.level,
@ -368,6 +396,7 @@ impl Core {
});
}
let mut target = target;
target.routed = routed;
// `_maybe_run_weighted_failover`: its own `run_async_fallback`, before the regular one.
let failover = if order_path {
None
@ -463,7 +492,12 @@ impl Core {
if at.depth >= max_fallbacks {
return Err(None);
}
let failed = Tried::Target(FallbackTarget::Name(target.group.clone()));
let failed = Tried::Target(FallbackTarget::Name(
target
.routed
.clone()
.unwrap_or_else(|| target.group.clone()),
));
if !chain.attempted.contains(&failed) {
chain.attempted.push(failed);
}
@ -488,17 +522,34 @@ impl Core {
let mut unroutable: Option<Value> = None;
for next in targets {
let (tried, group, order) = match &next {
Next::Order(group, order) => (
Tried::Order(group.clone(), *order),
group.clone(),
Some(*order),
),
Next::Order(group, order) => {
// An order level of the group a pre-routing hook selected passes the
// fallback gates as any group the level did not serve does.
if *group != target.group {
match self.host.allow_fallback(request, group).await {
Ok(true) => {}
Ok(false) => continue,
Err(raised) => return Err(Some(Error::Attempt(raised))),
}
}
(
Tried::Order(group.clone(), *order),
group.clone(),
Some(*order),
)
}
Next::Failover(group, excluded) => {
let mut key = excluded.to_vec();
key.sort();
(Tried::Failover(group.clone(), key), group.clone(), None)
}
Next::Fallback(FallbackTarget::Name(name)) if *name == target.group => continue,
// The group that just failed: the pre-routing selection when there was one, so
// a subagent bound to another tier can fall back to the model it asked for.
Next::Fallback(FallbackTarget::Name(name))
if *name == target.routed.as_deref().unwrap_or(&target.group) =>
{
continue;
}
// A provider-scoped request never leaves its group; a dict target with no
// `model` names no group, so it is skipped too.
Next::Fallback(fallback)
@ -573,6 +624,7 @@ impl Core {
params: params.clone(),
excluded,
failover_excluded: target.failover_excluded.clone(),
routed: None,
};
let parent = at;
let at = FallbackHop {
@ -704,6 +756,8 @@ impl Core {
snapshot: &RouterStateSnapshot,
request: &RouterRequest,
target: &Target,
routed: Option<&str>,
served_as: Option<&str>,
error: &Error,
) -> (Option<Vec<Next>>, bool) {
// A pick's error decides by its exception's classes as an attempt's does: a pre-call
@ -711,11 +765,19 @@ impl Core {
let failure = retried_failure(error);
let is_a = |name: &str| failure.as_ref().is_some_and(|f| f.is_a(name));
let lists = Lists::of(snapshot, request);
// `fallback_lookup_groups`: the group being served, then the group the request was for.
let mut lookup = vec![target.group.as_str()];
if request.model != target.group {
lookup.push(request.model.as_str());
// `fallback_lookup_groups`: the group a pre-routing hook selected, the group being served,
// then the group the request was for.
let mut lookup: Vec<&str> = routed.into_iter().collect();
for group in served_as
.into_iter()
.chain([target.group.as_str(), request.model.as_str()])
{
if !lookup.contains(&group) {
lookup.push(group);
}
}
// `order_model_group`: the pre-routing selection's order levels, else the level's own.
let order_group = routed.unwrap_or(target.group.as_str());
let chain = |list: &Fallbacks| {
if list.is_client_side(lists.keys) {
Some(list.as_chain())
@ -734,7 +796,7 @@ impl Core {
if !skip_order {
// A routing group's members, with a priority group's priorities as their order.
let listing_team = request.listing_team_id.as_deref();
let mut orders: Vec<i64> = group_members(snapshot, &target.group, listing_team)
let mut orders: Vec<i64> = group_members(snapshot, order_group, listing_team)
.iter()
.filter(|d| d.listed_for_team(listing_team))
.filter_map(|d| d.order())
@ -742,11 +804,14 @@ impl Core {
orders.sort_unstable();
orders.dedup();
if orders.len() > 1 {
let current = target.order.unwrap_or(orders[0]);
let current = target
.order
.filter(|_| routed.is_none())
.unwrap_or(orders[0]);
let mut targets: Vec<Next> = orders
.into_iter()
.filter(|order| *order > current)
.map(|order| Next::Order(target.group.clone(), order))
.map(|order| Next::Order(order_group.to_owned(), order))
.collect();
if let Some(list) = lists.fallbacks {
targets.extend(fallbacks_of(chain(list).unwrap_or_default()));
@ -879,6 +944,8 @@ impl Core {
return Err(Failed {
error: Box::new(Error::Attempt(AttemptError::host_raised(&["Exception"]))),
deployment_id: None,
routed: None,
served_as: None,
});
}
let Some(first_failure) = retried_failure(&failed.error) else {
@ -965,6 +1032,8 @@ impl Core {
return Err(Failed {
error: Box::new(Error::Attempt(AttemptError::host_raised(&["Exception"]))),
deployment_id: None,
routed: None,
served_as: None,
});
}
if num_retries == 0 {
@ -1219,6 +1288,7 @@ impl Core {
} else {
candidates.model_group.clone()
},
routed: routed_name(target, &candidates),
target_order: target.order,
candidate_ids: candidates.all.iter().map(|d| d.id.clone()).collect(),
retry,
@ -1236,23 +1306,43 @@ impl Core {
// the stream's end.
match self.host.invoke(&attempt).await {
Ok(output) => {
// `_should_raise_content_policy_error`: a blocked response is raised only when
// a content-policy fallback would serve it, and counts as no failure.
let routed = routed_name(target, &candidates);
let served = (routed.as_deref(), candidates.served_as.as_deref());
let raised = match output.blocked {
Some(Blocked::ContentFilter) => self.has_content_policy_fallback(
snapshot, request, target, at, attempted, served,
),
Some(Blocked::Refusal) => self.refusal_fallback_available(
snapshot, request, target, at, attempted, served,
),
None => false,
};
let tracks = output
.tracks_usage
.unwrap_or_else(|| deployment.tracks_usage());
// A shared logging object logs the success of the response that serves the
// request alone, so a raised blocked response counts its request and no success.
let callbacks = match raised && request.outcome_callbacks.shared_logging {
true => OutcomeCallbacks {
success: false,
..request.outcome_callbacks
},
false => request.outcome_callbacks,
};
let recorded = match (output.deferred, tracks) {
(true, true) => {
let tokens = output.total_tokens as f64;
self.health.add_usage(&deployment, None, tokens, 1.0).await
self.health
.add_usage(&deployment, None, tokens, 1.0)
.await
.map(|_| ())
}
(true, false) => Ok(()),
(false, _) => {
self.health
.record_success(
request.outcome_callbacks,
&deployment,
output.total_tokens,
tracks,
)
.record_success(callbacks, &deployment, output.total_tokens, tracks)
.await
}
};
@ -1262,16 +1352,6 @@ impl Core {
self.host.observe(&RouterEvent::Succeeded {
deployment_id: deployment.id.clone(),
});
// `_should_raise_content_policy_error`: a blocked response is raised only when
// a content-policy fallback would serve it, and counts as no failure.
let raised =
match output.blocked {
Some(Blocked::ContentFilter) => self
.has_content_policy_fallback(snapshot, request, target, at, attempted),
Some(Blocked::Refusal) => self
.refusal_fallback_available(snapshot, request, target, at, attempted),
None => false,
};
if let Some(blocked) = output.blocked.filter(|_| raised) {
self.host.observe(&RouterEvent::Blocked {
level: at.level,
@ -1284,6 +1364,8 @@ impl Core {
return Err(Failed {
error: Box::new(Error::Attempt(failure)),
deployment_id: Some(deployment.id.clone()),
routed: routed_name(target, &candidates),
served_as: candidates.served_as.clone(),
});
}
Ok(RouterResponse {
@ -1303,6 +1385,8 @@ impl Core {
Err(Failed {
error: Box::new(Error::Attempt(failure)),
deployment_id: Some(deployment.id.clone()),
routed: routed_name(target, &candidates),
served_as: candidates.served_as.clone(),
})
}
}
@ -1325,6 +1409,9 @@ impl Core {
&snapshot.settings.routing_strategy,
request,
)?;
self.host.observe(&RouterEvent::PickStarted {
model: target.group.clone(),
});
let routed = self
.pre_routed(PreRoute {
op,
@ -1334,22 +1421,25 @@ impl Core {
retry,
})
.await?;
let (routed, served_as) = match routed {
PreRouted::Bound { model, group } => (PreRouted::Model(model), Some(group)),
other => (other, None),
};
match routed {
PreRouted::Model(name) => {
snapshot.missing_routing_group(request, &name)?;
let candidates = self
let mut candidates = self
.candidates(snapshot, op, request, &name, inputs)
.await?;
candidates.served_as = served_as;
if candidates.healthy.is_empty() {
return Err(Error::NoHealthyDeployments {
model: candidates.failed_model(&snapshot.registry),
}
.into());
return Err(candidates.none_healthy(&snapshot.registry).into());
}
let deployment = self.pick(snapshot, op, request, &candidates).await?;
Ok((deployment, candidates))
}
// A custom strategy's pick is served as is, for the requested name.
PreRouted::Bound { .. } => unreachable!("mapped to a model above"),
PreRouted::Deployment(deployment) => {
let candidates = Candidates {
name: target.group.clone(),
@ -1358,6 +1448,7 @@ impl Core {
all: vec![deployment.clone()],
healthy: vec![deployment.clone()],
wildcard_request: None,
served_as: None,
};
Ok((deployment, candidates))
}
@ -1383,6 +1474,7 @@ impl Core {
deployment,
name: target.group.clone(),
model_group: target.group.clone(),
routed: None,
target_order: target.order,
candidate_ids: Vec::new(),
retry,
@ -1417,10 +1509,11 @@ impl Core {
target: &Target,
at: FallbackHop,
attempted: &[Tried],
served: (Option<&str>, Option<&str>),
) -> bool {
let lists = Lists::of(snapshot, request);
match lists.content_policy {
Some(list) => unattempted_chain(list, &lists, request, target, at, attempted),
Some(list) => unattempted_chain(list, &lists, request, target, at, attempted, served),
None => snapshot.settings.dynamic.fallbacks.has_star(),
}
}
@ -1434,20 +1527,22 @@ impl Core {
target: &Target,
at: FallbackHop,
attempted: &[Tried],
served: (Option<&str>, Option<&str>),
) -> bool {
if request.disable_fallbacks {
return false;
}
let lists = Lists::of(snapshot, request);
if lists.content_policy.is_some() {
return self.has_content_policy_fallback(snapshot, request, target, at, attempted);
return self
.has_content_policy_fallback(snapshot, request, target, at, attempted, served);
}
if snapshot.settings.dynamic.fallbacks.has_star() {
return true;
}
lists
.fallbacks
.is_some_and(|list| unattempted_chain(list, &lists, request, target, at, attempted))
lists.fallbacks.is_some_and(|list| {
unattempted_chain(list, &lists, request, target, at, attempted, served)
})
}
/// `async_deployment_callback_on_failure`, when it runs for this failure.
@ -1553,3 +1648,8 @@ impl Core {
Ok(())
}
}
/// The name a pre-routing hook rewrote the level's group to for this pick, if it did.
fn routed_name(target: &Target, candidates: &Candidates) -> Option<String> {
(candidates.name != target.group).then(|| candidates.name.clone())
}

22
litellm-rust/crates/router-core/src/call/stream.rs
View file

@ -37,6 +37,9 @@ pub struct StreamOrigin {
pub attempted_retries: u32,
/// The budget the first stream retry committed to (`committed_retry_budget_for_request`).
pub committed_budget: Option<u32>,
/// The group a pre-routing hook selected for `group`, which the fallback lookup keys on first.
#[serde(default)]
pub routed: Option<String>,
}
/// What may still take over a stream that fails before content.
@ -61,14 +64,22 @@ impl StreamOrigin {
params: None,
excluded: Arc::from([]),
failover_excluded: Arc::from([]),
routed: self.routed.clone(),
}
}
/// `fallback_lookup_groups`: the served group, then the group the request was for.
fn lookup<'a>(&'a self, request: &'a RouterRequest) -> Vec<&'a str> {
let mut lookup = vec![self.group.as_str()];
if request.model != self.group {
lookup.push(request.model.as_str());
let mut lookup: Vec<&str> = Vec::new();
for group in self
.routed
.as_deref()
.into_iter()
.chain([self.group.as_str(), request.model.as_str()])
{
if !lookup.contains(&group) {
lookup.push(group);
}
}
lookup
}
@ -186,7 +197,10 @@ impl Router {
scope: 0,
},
&mut chain,
Failed::from(Error::Attempt(failure)),
Failed {
routed: origin.routed.clone(),
..Failed::from(Error::Attempt(failure))
},
)
.await
}

11
litellm-rust/crates/router-core/src/error.rs
View file

@ -40,7 +40,13 @@ pub enum Error {
#[error("no deployment serves model `{model}`")]
UnknownModel { model: String },
#[error("no deployment of `{model}` is available")]
NoHealthyDeployments { model: String },
NoHealthyDeployments {
model: String,
/// The ids of the deployments the request could route to before the cooldown filter
/// (`request_deployments`), when the candidate list came up empty; `None` when the
/// strategy refused every candidate.
request_ids: Option<Vec<String>>,
},
#[error("the attempt failed")]
Attempt(#[from] AttemptError),
#[error("the routing cache failed")]
@ -201,8 +207,9 @@ impl Error {
Self::UnknownModel { model } => Self::UnknownModel {
model: model.clone(),
},
Self::NoHealthyDeployments { model } => Self::NoHealthyDeployments {
Self::NoHealthyDeployments { model, request_ids } => Self::NoHealthyDeployments {
model: model.clone(),
request_ids: request_ids.clone(),
},
Self::Paused { model } => Self::Paused {
model: model.clone(),

24
litellm-rust/crates/router-core/src/health.rs
View file

@ -554,35 +554,35 @@ impl Health {
return Ok(());
}
callbacks
.record(
OutcomeWrite::AttemptUsage,
self.add_usage(deployment, None, total_tokens as f64, 1.0),
)
.record(OutcomeWrite::AttemptUsage, async {
self.add_usage(deployment, None, total_tokens as f64, 1.0)
.await
.map(|_| ())
})
.await
}
/// `_increment_deployment_usage`, for a deployment that tracks usage: tokens and requests
/// under `model` (metadata's `deployment`, which a wildcard route fills; else the
/// deployment's own), unless there is nothing to add.
/// deployment's own), unless there is nothing to add. Returns the tokens counter's key when
/// it added to it.
pub async fn add_usage(
&self,
deployment: &Deployment,
model: Option<&str>,
total_tokens: f64,
requests: f64,
) -> Result<(), Error> {
) -> Result<Option<String>, Error> {
if total_tokens <= 0.0 && requests <= 0.0 {
return Ok(());
return Ok(None);
}
let minute = current_minute(self.clock.now());
let model = model.unwrap_or(&deployment.params.model);
let tpm_key = usage_key(&deployment.id, model, UsageMetric::Tpm, &minute);
self.cache
.async_increment_usage(
&[
(
usage_key(&deployment.id, model, UsageMetric::Tpm, &minute),
total_tokens,
),
(tpm_key.clone(), total_tokens),
(
usage_key(&deployment.id, model, UsageMetric::Rpm, &minute),
requests,
@ -591,7 +591,7 @@ impl Health {
Some(USAGE_TTL),
)
.await?;
Ok(())
Ok(Some(tpm_key))
}
/// `async_deployment_callback_on_failure`: a failed request counts against the deployment's

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

@ -83,6 +83,13 @@ pub struct RouterRequest {
/// A Decisions call carrying audio or videos: the deployments that refuse them are skipped
/// unless that leaves nothing (`_without_decisions_deployments_refusing_media`).
pub decisions_media: bool,
/// A batch create (`batch_create_routing`): when the caller may use a deployment declared
/// `model_info.mode: batch`, the Azure deployments that are not are skipped
/// (`without_deployments_that_cannot_batch`).
pub batch_create: bool,
/// An embedding pick: the host keeps the candidates whose `max_embedding_batch_size` admits
/// the input (`deployments_admitting_embedding_batch`).
pub embedding_batch: bool,
/// `f"{provider}/{model}"` for a lookup that cannot ask the host
/// (`_core_get_routable_deployments`); a pick asks [`Host::provider_request`].
pub wildcard_request: Option<String>,
@ -291,6 +298,9 @@ pub struct Attempt {
pub name: String,
/// The model group being served, after fallbacks.
pub model_group: String,
/// The group a pre-routing hook selected for `name`, when it differs (the tier stamp the
/// fallback lookup keys on first).
pub routed: Option<String>,
/// The `order` level an order fallback targets in the same group.
pub target_order: Option<i64>,
/// Every deployment `name` resolved to for this caller, the picked one included.
@ -383,11 +393,22 @@ pub enum SelectStage {
/// `_decisions_deployment_refuses_media`: return the candidates that refuse the audio or
/// videos a Decisions call carries.
MediaRefusing,
/// `deployments_admitting_embedding_batch`: return the candidates that admit the embedding
/// input, or what the pick keeps when none does (`deployments_when_none_admits_embedding_batch`).
EmbeddingBatch,
/// `_async_override_selector_pre_call_check`: a deployment named by id under a per-request
/// strategy override still passes that strategy's pre-call check (rpm accounting); return the
/// candidates, or raise its refusal.
OverrideCheck(RoutingStrategy),
}
/// Routing events, delivered through [`Host::observe`].
#[derive(Clone, Debug, PartialEq)]
pub enum RouterEvent {
/// A pick for `model` (the level's group) begins: its route phase opens.
PickStarted {
model: String,
},
Selected {
model_group: String,
deployment_id: String,

36
litellm-rust/crates/router-core/src/lib.rs
View file

@ -366,7 +366,7 @@ impl Router {
})
.await?;
let name = match routed {
PreRouted::Model(name) => name,
PreRouted::Model(name) | PreRouted::Bound { model: name, .. } => name,
PreRouted::Deployment(deployment) => return Ok(deployment),
};
snapshot.missing_routing_group(request, &name)?;
@ -374,6 +374,10 @@ impl Router {
.core
.candidates(&snapshot, op, request, &name, &SelectInputs::default())
.await?;
// An empty candidate list names what the request could route to before the cooldown filter.
if candidates.healthy.is_empty() && !candidates.kind.is_single() {
return Err(candidates.none_healthy(&snapshot.registry));
}
self.core.pick(&snapshot, op, request, &candidates).await
}
@ -398,9 +402,7 @@ impl Router {
.await?;
// `_core_async_get_healthy_deployments` raises when nothing is left.
if candidates.healthy.is_empty() {
return Err(Error::NoHealthyDeployments {
model: candidates.failed_model(&snapshot.registry),
});
return Err(candidates.none_healthy(&snapshot.registry));
}
Ok(candidates)
}
@ -434,7 +436,7 @@ impl Router {
/// A success a logging callback saw (`deployment_callback_on_success`): counted for a
/// configured deployment, with its tokens and requests when it tracks usage. Returns the
/// success counter's key, or `None` for an unknown deployment.
/// tokens counter's key when it added usage (`_increment_deployment_usage`), else `None`.
pub async fn report_success(
&self,
id: &str,
@ -462,22 +464,22 @@ impl Router {
let Some(deployment) = snapshot.registry.deployment(id) else {
return Ok(None);
};
let key = self.core.health.count_success(id).await?;
if report
self.core.health.count_success(id).await?;
if !report
.tracks_usage
.unwrap_or_else(|| deployment.tracks_usage())
{
self.core
.health
.add_usage(
&deployment,
report.usage_model.as_deref(),
total_tokens,
requests,
)
.await?;
return Ok(None);
}
Ok(Some(key))
self.core
.health
.add_usage(
&deployment,
report.usage_model.as_deref(),
total_tokens,
requests,
)
.await
}
/// `sync_deployment_callback_on_success`: count a success against `id`, configured or not,

9
litellm-rust/crates/router-core/src/outcome.rs
View file

@ -28,6 +28,10 @@ pub struct OutcomeCallbacks {
pub failure: bool,
/// The attempt runs through `make_call`, which counts its usage whatever the callbacks do.
pub make_call: bool,
/// The request's attempts log on one logging object the caller passed (`litellm_logging_obj`),
/// which logs one success a request, the response that serves it: an attempt whose blocked
/// response is raised for a fallback counts no success.
pub shared_logging: bool,
}
impl Default for OutcomeCallbacks {
@ -36,6 +40,7 @@ impl Default for OutcomeCallbacks {
success: true,
failure: true,
make_call: true,
shared_logging: false,
}
}
}
@ -78,21 +83,25 @@ mod tests {
success: true,
failure: true,
make_call: true,
shared_logging: false,
};
const UNBOUND: OutcomeCallbacks = OutcomeCallbacks {
success: false,
failure: false,
make_call: true,
shared_logging: false,
};
const UNBOUND_WITHOUT_MAKE_CALL: OutcomeCallbacks = OutcomeCallbacks {
success: false,
failure: false,
make_call: false,
shared_logging: false,
};
const BOUND_SYNC: OutcomeCallbacks = OutcomeCallbacks {
success: false,
failure: true,
make_call: true,
shared_logging: false,
};
#[rstest]

3
litellm-rust/crates/router-core/src/prerouting.rs
View file

@ -33,6 +33,9 @@ pub struct PreRoute<'a> {
pub enum PreRouted {
/// Resolve this model group instead of the requested name.
Model(String),
/// Resolve `model`, served as `group`: the group a session binding stamped on the request,
/// which the level's fallbacks also look up.
Bound { model: String, group: String },
/// A custom strategy's pick: serve this deployment, without resolving or filtering.
Deployment(Arc<Deployment>),
}

17
litellm-rust/crates/router-core/src/registry.rs
View file

@ -234,6 +234,23 @@ impl Deployment {
/// `_deployment_supports_web_search`: `model_info.supports_web_search` when it is set, read
/// for truthiness, else true.
/// `is_batch_capable`: declared `model_info.mode: batch`.
pub(crate) fn is_batch_capable(&self) -> bool {
self.model_info.raw.get("mode").and_then(Value::as_str) == Some("batch")
}
/// `_cannot_run_batches`: an Azure (or Azure AI) deployment not declared batch-capable, by
/// its `custom_llm_provider`, else its model's provider prefix.
pub(crate) fn cannot_run_batches(&self) -> bool {
let provider = self.params.custom_llm_provider.as_deref().or_else(|| {
self.params
.model
.split_once('/')
.map(|(provider, _)| provider)
});
matches!(provider, Some("azure" | "azure_ai")) && !self.is_batch_capable()
}
pub(crate) fn supports_web_search(&self) -> bool {
self.model_info
.raw

80
litellm-rust/crates/router-core/src/select.rs
View file

@ -16,7 +16,7 @@ use crate::{
host::{RouterRequest, RoutingContext, SelectStage},
registry::{Deployment, Registry, ResolveKind},
resolve::{Caller, resolve_request},
strategy::{StrategyScope, routing_context},
strategy::{StrategyScope, request_override, routing_context},
};
/// Randomness for the shuffle and retry jitter, injectable so tests are deterministic.
@ -77,11 +77,23 @@ pub struct Candidates {
/// The provider-prefixed spelling the name resolved under (`get_pattern`), when the host
/// was asked for one.
pub wildcard_request: Option<String>,
/// The group the host stamped the pick as served for when a pre-routing hook rewrote the name
/// (a Claude Code subagent's bound router): a fallback lookup key.
pub served_as: Option<String>,
}
impl Candidates {
/// The `model` a no-deployments error names: what `_common_checks_available_deployment`
/// returned, so the alias-resolved name a wildcard serves and not its pattern.
/// The error an empty candidate list raises: it names the deployments the request could
/// route to before the cooldown filter.
pub(crate) fn none_healthy(&self, registry: &Registry) -> Error {
Error::NoHealthyDeployments {
model: self.failed_model(registry),
request_ids: Some(self.all.iter().map(|d| d.id.clone()).collect()),
}
}
pub(crate) fn failed_model(&self, registry: &Registry) -> String {
let model = match self.kind {
ResolveKind::Wildcard => registry.alias_of(&self.name).unwrap_or(&self.name),
@ -466,6 +478,7 @@ impl Core {
all,
healthy,
wildcard_request,
served_as: None,
});
}
@ -503,6 +516,7 @@ impl Core {
all,
healthy,
wildcard_request,
served_as: None,
});
}
@ -602,6 +616,34 @@ impl Core {
}
}
// `deployments_admitting_embedding_batch`: an embedding pick keeps the deployments whose
// input limit admits the batch; when none does, the host keeps what the pick waits for or
// the 400 names.
if request.embedding_batch && !resolution.kind.is_single() && !healthy.is_empty() {
let mut ctx = self.context(
snapshot,
op,
request,
&resolution.group,
name,
StrategyScope::Default,
);
ctx.pool = &all;
healthy = self
.host
.select(&SelectStage::EmbeddingBatch, &ctx, healthy)
.await?;
}
// `without_deployments_that_cannot_batch`: a batch create on a group the caller may use a
// batch deployment of skips the Azure deployments that cannot run batches.
if request.batch_create
&& !resolution.kind.is_single()
&& healthy.iter().any(|d| d.is_batch_capable())
{
healthy.retain(|d| !d.cannot_run_batches());
}
if let Some(pinned) = request.pinned_deployment_id.as_deref()
&& healthy.iter().any(|d| d.id == pinned)
{
@ -653,6 +695,7 @@ impl Core {
all,
healthy,
wildcard_request,
served_as: None,
})
}
@ -691,14 +734,43 @@ impl Core {
) -> Result<Arc<Deployment>, Error> {
let no_deployment = || Error::NoHealthyDeployments {
model: candidates.failed_model(&snapshot.registry),
request_ids: None,
};
// A deployment named by its id is returned as is, before any strategy runs.
// A deployment named by its id is returned as is, before any strategy runs, past a
// per-request strategy override's pre-call check.
if candidates.kind.is_single() {
return candidates
let named = candidates
.healthy
.first()
.cloned()
.ok_or_else(no_deployment);
.ok_or_else(no_deployment)?;
if let Some(strategy) = request_override(request.routing_strategy.as_deref())
&& !matches!(
strategy,
RoutingStrategy::SimpleShuffle | RoutingStrategy::Unset
)
{
let ctx = self.context(
snapshot,
op,
request,
&candidates.model_group,
&candidates.name,
StrategyScope::Override,
);
let kept = self
.host
.select(
&SelectStage::OverrideCheck(strategy),
&ctx,
vec![named.clone()],
)
.await?;
if kept.is_empty() {
return Err(no_deployment());
}
}
return Ok(named);
}
let (strategy, scope) = routing_context(
&snapshot.registry,

7
litellm/litellm_core_utils/litellm_logging.py
View file

@ -4074,12 +4074,11 @@ class Logging(LiteLLMLoggingBaseClass):
start_time=start_time,
end_time=end_time,
)
if (
callable(callback) is True and is_sync_request and customLogger is not None
): # custom logger functions
if callable(callback) is True and is_sync_request: # custom logger functions
function_logger = customLogger if customLogger is not None else CustomLogger()
print_verbose(f"success callbacks: Running Custom Callback Function - {callback}")
customLogger.log_event(
function_logger.log_event(
kwargs=call_details,
response_obj=result,
start_time=start_time,

3
litellm/llms/anthropic/messages/adapters/handler.py
View file

@ -123,10 +123,9 @@ def _extract_proxy_litellm_metadata(
def _requested_router(kwargs: Mapping[str, object]) -> "Router | None":
from litellm.router import Router
from litellm.router_python_core import PythonCore
requested: Final = kwargs.get("litellm_router")
return requested if isinstance(requested, (Router, PythonCore)) else None # pyright: ignore[reportReturnType] # the router passes its core, which carries Router's members
return requested if isinstance(requested, Router) else None
def request_safeguards_evaluator(

307
litellm/router.py
View file

@ -1,41 +1,208 @@
"""
litellm.Router: load balancing, retries, fallbacks and cooldowns across deployments.
``Router`` holds no routing logic. It builds its configuration (``litellm/router_config.py``), picks
one ``RouterCore`` at construction (``litellm/router_core.py``) and delegates every member to it:
``Router`` holds no routing logic. It builds its configuration (``litellm/router_config.py``), builds
its ``RouterCore`` from it (the Rust core, ``litellm/router_core.py``) and delegates every member to it:
the endpoint methods (acompletion, aembedding, aresponses, aanthropic_messages, ...) and the entry
points (registry, settings, health, select, apply, observe, register_plugin, retire). The members
outside code used before the entry points live in ``DeprecatedRouterMembers`` until they are removed.
"""
import functools
import inspect
import weakref
from collections.abc import Mapping
from collections.abc import Callable, Mapping, Sequence
from types import FunctionType, MappingProxyType
from typing import TYPE_CHECKING, Final
from typing import TYPE_CHECKING, Any, Final, Literal
from litellm import router_python_core
from litellm.router_config import RouterConfig, build_router_config
from litellm.router_core import RouterCore, make_core
from litellm.router_deprecated import DeprecatedRouterMembers
from litellm.router_python_core import (
PythonCore,
as_output_cap,
model_info_is_active_for_environment,
stamped_client_ceiling,
)
from litellm.router_python_host import PythonHost, python_host
from litellm.router_rust.steps.prerouting import as_output_cap, stamped_client_ceiling
from litellm.router_rust.surface.registry import model_info_is_active_for_environment
from litellm.types.litellm_params import RoutingStrategyName
from litellm.types.router import (
AlertingConfig,
AllowedFailsPolicy,
AssistantsTypedDict,
AutoRouterCapabilityLimit,
CustomRoutingStrategyBase,
Deployment,
DeploymentTypedDict,
FallbackAccessCheck,
FallbackBudgetCheck,
GuardrailTypedDict,
LiteLLM_Params,
ModelGroupInfo,
OptionalPreCallChecks,
RetryPolicy,
RouterGeneralSettings,
RouterModelGroupAliasItem,
RoutingGroup,
RoutingPlugin,
SearchToolTypedDict,
)
from litellm.types.utils import GenericBudgetConfigType
_CORE_SIGNATURE: Final = inspect.signature(PythonCore.__init__) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # reflection reads only the parameters' names
if TYPE_CHECKING:
from litellm.router_rust.core import RustCore
# ``Router(...)``'s parameters, their defaults and its docstring: ``router_arguments`` binds a call
# against them and ``Router.__init__`` reports them. The body is never run.
def _router_init(
self,
model_list: list[DeploymentTypedDict] | list[dict[str, Any]] | None = None,
## ASSISTANTS API ##
assistants_config: AssistantsTypedDict | None = None,
## SEARCH API ##
search_tools: list[SearchToolTypedDict] | None = None,
## GUARDRAIL API ##
guardrail_list: list[GuardrailTypedDict] | None = None,
## CACHING ##
redis_url: str | None = None,
redis_host: str | None = None,
redis_port: int | None = None,
redis_password: str | None = None,
redis_db: int | None = None,
cache_responses: bool | None = False,
cache_kwargs: dict = {}, # additional kwargs to pass to RedisCache (see caching.py)
caching_groups: list[tuple] | None = None, # if you want to cache across model groups
client_ttl: int = 3600, # ttl for cached clients - will re-initialize after this time in seconds
## SCHEDULER ##
polling_interval: float | None = None,
default_priority: int | None = None,
## RELIABILITY ##
num_retries: int | None = None,
max_fallbacks: int | None = None, # max fallbacks to try before exiting the call. Defaults to 5.
timeout: float | None = None,
stream_timeout: float | None = None,
default_litellm_params: dict | None = None, # default params for Router.chat.completion.create
default_max_parallel_requests: int | None = None,
set_verbose: bool = False,
debug_level: Literal["DEBUG", "INFO"] = "INFO",
default_fallbacks: list[str] | None = None, # generic fallbacks, works across all deployments
fallbacks: list = [],
context_window_fallbacks: list = [],
content_policy_fallbacks: list = [],
model_group_alias: dict[str, str | RouterModelGroupAliasItem] | None = {},
enable_pre_call_checks: bool = False,
enable_tag_filtering: bool = False,
tag_filtering_match_any: bool = True,
tag_routing_prefix: str = "",
plugins: list[RoutingPlugin] | None = None,
retry_after: int = 0, # min time to wait before retrying a failed request
retry_policy: RetryPolicy | dict | None = None, # set custom retries for different exceptions
model_group_retry_policy: dict[str, RetryPolicy] = {}, # set custom retry policies based on model group
allowed_fails: int | None = None, # Number of times a deployment can failbefore being added to cooldown
allowed_fails_policy: AllowedFailsPolicy | None = None, # set custom allowed fails policy
cooldown_time: float | None = None, # (seconds) time to cooldown a deployment after failure
disable_cooldowns: bool | None = None,
fallback_on_unsupported_params: bool = True,
routing_strategy: RoutingStrategyName = "simple-shuffle",
optional_pre_call_checks: OptionalPreCallChecks | None = None,
routing_strategy_args: dict = {}, # just for latency-based
routing_groups: list[RoutingGroup | dict] | None = None,
provider_budget_config: GenericBudgetConfigType | None = None,
alerting_config: AlertingConfig | None = None,
router_general_settings: RouterGeneralSettings | None = RouterGeneralSettings(),
deployment_affinity_ttl_seconds: int = 3600,
model_group_affinity_config: dict[str, list[str]] | None = None,
ignore_invalid_deployments: bool = False,
enable_health_check_routing: bool = False,
health_check_staleness_threshold: int | None = None,
health_check_ignore_transient_errors: bool = False,
background_health_check_model_groups: Sequence[str] | None = None,
enable_weighted_failover: bool = False,
enable_mid_stream_fallback_continuation: bool = False,
fallback_access_check: FallbackAccessCheck | None = None,
fallback_budget_check: FallbackBudgetCheck | None = None,
auto_router_capability_limit: AutoRouterCapabilityLimit | None = None,
session_affinity_repin_on_failover: bool = False,
_host: PythonHost = python_host,
) -> None:
"""
Initialize the Router class with the given parameters for caching, reliability, and routing strategy.
Args:
model_list (Optional[list]): List of models to be used. Defaults to None.
redis_url (Optional[str]): URL of the Redis server. Defaults to None.
redis_host (Optional[str]): Hostname of the Redis server. Defaults to None.
redis_port (Optional[int]): Port of the Redis server. Defaults to None.
redis_password (Optional[str]): Password of the Redis server. Defaults to None.
cache_responses (Optional[bool]): Flag to enable caching of responses. Defaults to False.
cache_kwargs (dict): Additional kwargs to pass to RedisCache. Defaults to {}.
caching_groups (Optional[List[tuple]]): List of model groups for caching across model groups. Defaults to None.
client_ttl (int): Time-to-live for cached clients in seconds. Defaults to 3600.
polling_interval: (Optional[float]): frequency of polling queue. Only for '.scheduler_acompletion()'. Default is 3ms.
default_priority: (Optional[int]): the priority given to every chat or text completion request that sets no 'priority' of its own; '.acompletion()' and '.atext_completion()' queue it through the scheduler. Default is None.
num_retries (Optional[int]): Number of retries for failed requests. Defaults to 2.
timeout (Optional[float]): Timeout for requests. Defaults to None.
default_litellm_params (dict): Default parameters for Router.chat.completion.create. Defaults to {}.
set_verbose (bool): Flag to set verbose mode. Defaults to False.
debug_level (Literal["DEBUG", "INFO"]): Debug level for logging. Defaults to "INFO".
fallbacks (List): List of fallback options. Defaults to [].
context_window_fallbacks (List): List of context window fallback options. Defaults to [].
enable_pre_call_checks (boolean): Filter out deployments which are outside context window limits for a given prompt
model_group_alias (Optional[dict]): Alias for model groups. Defaults to {}.
retry_after (int): Minimum time to wait before retrying a failed request. Defaults to 0.
allowed_fails (Optional[int]): Number of allowed fails before adding to cooldown. Defaults to None.
cooldown_time (float): Time to cooldown a deployment after failure in seconds. Defaults to 1.
routing_strategy (Literal["simple-shuffle", "least-busy", "usage-based-routing", "latency-based-routing", "cost-based-routing"]): Routing strategy used for the implicit "default" group (any model not claimed by an entry in `routing_groups`). Defaults to "simple-shuffle".
routing_strategy_args (dict): Additional args for the default group's routing strategy (e.g. latency window). Defaults to {}.
routing_groups (Optional[List[RoutingGroup]]): Named subsets of `model_name`s with a group routing strategy.
Priority groups apply only to group calls and may overlap. Other groups supply their members' default
strategy, with at most one such group per model. Unclaimed models use the top-level strategy.
alerting_config (AlertingConfig): Slack alerting configuration. Defaults to None.
provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None.
deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600.
session_affinity_repin_on_failover (bool): After a successful replacement response, move a session pin
whose deployment is in cooldown or removed. Async routes only. Defaults to False. Applied at startup.
ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error.
enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False.
enable_mid_stream_fallback_continuation (bool): When True, a chat-completions stream that breaks after plain assistant text has been delivered continues on a fallback deployment via assistant prefill instead of surfacing the error. Only deployments whose model supports assistant prefill are eligible, so the partial text is continued, not regenerated; if none is, the original error is surfaced. Streams that emitted tool calls, thinking blocks, audio/images, or constrained (JSON / forced tool_choice) output are never continued. Async-only. Defaults to False.
fallback_access_check (Optional[FallbackAccessCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects is skipped. Defaults to None (every configured fallback is attempted).
fallback_budget_check (Optional[FallbackBudgetCheck]): Awaited before each cross-model-group fallback attempt on the async path; a fallback target it rejects as over budget is skipped. Defaults to None (budget is not re-checked on fallback).
Returns:
Router: An instance of the litellm.Router class.
Example Usage:
```python
from litellm import Router
model_list = [
{
"model_name": "azure-gpt-3.5-turbo", # model alias
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/<your-deployment-name-1>",
"api_key": <your-api-key>,
"api_version": <your-api-version>,
"api_base": <your-api-base>
},
},
{
"model_name": "azure-gpt-3.5-turbo", # model alias
"litellm_params": { # params for litellm completion/embedding call
"model": "azure/<your-deployment-name-2>",
"api_key": <your-api-key>,
"api_version": <your-api-version>,
"api_base": <your-api-base>
},
},
{
"model_name": "openai-gpt-3.5-turbo", # model alias
"litellm_params": { # params for litellm completion/embedding call
"model": "gpt-3.5-turbo",
"api_key": <your-api-key>,
},
]
router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": ["openai-gpt-3.5-turbo"]}])
```
"""
_CORE_SIGNATURE: Final = inspect.signature(_router_init)
def router_arguments(
@ -55,55 +222,129 @@ def _class_descriptor(cls: type, name: str) -> bool:
if TYPE_CHECKING:
from typing import overload
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
from litellm.types.llms.openai import AllMessageValues, OpenAIFileObject
from litellm.types.utils import EmbeddingResponse, LiteLLMBatch, ModelResponse
class Router(RustCore, DeprecatedRouterMembers):
"""What type checkers see: the constructor's parameters and the endpoints' return types. The
endpoints are served by the core at run time (litellm/router_rust/invoke/endpoints.py)."""
__init__ = _router_init # pyright: ignore[reportAssignmentType, reportIncompatibleMethodOverride] # Router(...)'s parameters
class Router(PythonCore, DeprecatedRouterMembers):
@classmethod
def from_config(cls, config: RouterConfig, host: PythonHost = python_host) -> "Router": ...
@overload
async def acompletion(
self, model: str, messages: list[AllMessageValues], stream: Literal[True], **kwargs: object
) -> CustomStreamWrapper: ...
@overload
async def acompletion(
self, model: str, messages: list[AllMessageValues], stream: Literal[False] = False, **kwargs: object
) -> ModelResponse: ...
@overload
async def acompletion(
self, model: str, messages: list[AllMessageValues], stream: bool = False, **kwargs: object
) -> CustomStreamWrapper | ModelResponse: ...
async def acompletion( # pyright: ignore[reportIncompatibleMethodOverride] # the endpoint's typed signature
self, model: str, messages: list[AllMessageValues], stream: bool = False, **kwargs: object
) -> CustomStreamWrapper | ModelResponse: ...
def completion( # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion
self, model: str, messages: list[dict[str, str]], **kwargs: object
) -> ModelResponse | CustomStreamWrapper: ...
def embedding( # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion
self, model: str, input: str | list[object], is_async: bool | None = False, **kwargs: object
) -> EmbeddingResponse: ...
async def aembedding( # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion
self, model: str, input: str | list[object], is_async: bool | None = True, **kwargs: object
) -> EmbeddingResponse: ...
async def acreate_file(self, model: str, **kwargs: object) -> OpenAIFileObject: ... # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion
async def acreate_batch(self, model: str, **kwargs: object) -> LiteLLMBatch: ... # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion
async def aretrieve_batch(self, model: str | None = None, **kwargs: object) -> LiteLLMBatch: ... # pyright: ignore[reportIncompatibleMethodOverride] # see acompletion
def __getattr__(self, name: str) -> Any: ... # the endpoints served by name
else:
_MISSING: Final = object()
_replaced_core_members: Final[ # mutable-ok: originals of patched core members, restored on delete
dict[tuple[type, str], list[object]]
] = {}
def _core_class() -> type:
from litellm.router_rust.core import RustCore
return RustCore
class _RouterType(type):
"""
Class-level reads of core members (``Router._helper(...)``) go to ``PythonCore``, and so do
class-level patches, which put the original back when the patch is removed.
``isinstance`` and ``dir`` see ``PythonCore`` too.
Class-level reads of core members (``Router._helper``) go to the core's class, and so do
class-level patches, which put the original back when the patch is removed. ``dir`` lists
the core's members and endpoints too, so ``Mock(spec=Router)`` accepts them.
"""
def __getattr__(cls, name: str) -> object:
return getattr(PythonCore, name)
if name.startswith("__") and name.endswith("__"):
raise AttributeError(name)
core: Final = _core_class()
try:
return getattr(core, name)
except AttributeError:
if not _is_endpoint(name):
raise
return _endpoint_function(name)
def __setattr__(cls, name: str, value: object) -> None:
if any(name in vars(base) for base in cls.__mro__) or not hasattr(PythonCore, name):
core: Final = _core_class()
if any(name in vars(base) for base in cls.__mro__) or not (hasattr(core, name) or _is_endpoint(name)):
super().__setattr__(name, value)
return
_replaced_core_members.setdefault((PythonCore, name), []).append(vars(PythonCore).get(name, _MISSING))
setattr(PythonCore, name, value)
_replaced_core_members.setdefault((core, name), []).append(vars(core).get(name, _MISSING))
setattr(core, name, value)
def __delattr__(cls, name: str) -> None:
replaced: Final = _replaced_core_members.get((PythonCore, name))
core: Final = _core_class()
replaced: Final = _replaced_core_members.get((core, name))
if not replaced or any(name in vars(base) for base in cls.__mro__):
super().__delattr__(name)
return
original: Final = replaced.pop()
if original is _MISSING:
delattr(PythonCore, name)
delattr(core, name)
return
setattr(PythonCore, name, original)
def __instancecheck__(cls, instance: object) -> bool:
return type.__instancecheck__(cls, instance) or isinstance(instance, PythonCore)
setattr(core, name, original)
def __dir__(cls) -> list[str]: # mutable-ok: dir() returns a list
return sorted({*type.__dir__(cls), *dir(PythonCore)})
from litellm.router_rust.invoke.endpoints import endpoint_names
return sorted({*type.__dir__(cls), *dir(_core_class()), *endpoint_names()})
def _is_endpoint(name: str) -> bool:
from litellm.router_rust.invoke.endpoints import endpoint_names
return name in endpoint_names()
@functools.cache
def _endpoint_function(name: str) -> Callable[..., object]:
"""An endpoint the core serves by name, read off the class (``Router.aretrieve_batch``): a
function that calls the instance's endpoint, so a class-level patch of it replaces it."""
def endpoint(
self: object,
*args: object,
**kwargs: object, # kwargs-ok: the endpoint's own arguments
) -> object:
served: Final = _core_class().__getattr__(self, name) # pyright: ignore[reportAny] # the core's endpoint lookup
return served(*args, **kwargs) # pyright: ignore[reportAny] # see above
endpoint.__name__ = endpoint.__qualname__ = name
return endpoint
def _patchable_method(core: object, name: str) -> bool:
"""A method of the core: an instance-level replacement of it is also kept on the facade, so
``mock.patch.object(router, name)`` finds it there and nested patches unwind in order. A core
that serves its members from another core (``__getattr__``) is asked for the bound method."""
that serves its members from ``__getattr__`` is asked for the bound method."""
static: Final = inspect.getattr_static(type(core), name, None)
if static is not None:
return isinstance(static, FunctionType | staticmethod | classmethod)
@ -183,17 +424,13 @@ else:
delattr(self._core, name)
def __dir__(self) -> list[str]: # mutable-ok: dir() returns a list
return sorted({*super().__dir__(), *dir(self.__dict__.get("_core", PythonCore))})
core: Final = self.__dict__.get("_core")
return sorted({*super().__dir__(), *(dir(core) if core is not None else ())})
Router.__init__.__signature__ = _CORE_SIGNATURE.replace( # pyright: ignore[reportFunctionMemberAccess] # inspect reads it
parameters=[parameter for parameter in _CORE_SIGNATURE.parameters.values() if parameter.name != "_host"]
)
Router.__init__.__doc__ = PythonCore.__init__.__doc__
def __getattr__(name: str) -> object:
"""Module-level helpers that lived in this module before the core moved out."""
return getattr(router_python_core, name)
Router.__init__.__doc__ = _router_init.__doc__
__all__ = [

4
litellm/router_api.py
View file

@ -1491,7 +1491,9 @@ class RoutingContext:
return self._router._params_sent_to_deployment(request, deployment)
def default_litellm_params(self) -> Mapping[str, object]:
return self._router._active_default_litellm_params()
"""The router's ``default_litellm_params`` that are set (a None value is no default)."""
configured: Final = self.settings().default_litellm_params
return {key: value for key, value in (configured or {}).items() if value is not None}
def settings(self) -> RouterSettings:
return self._router.settings()

19
litellm/router_attempt.py
View file

@ -3,8 +3,7 @@ One router attempt, as the host runs it: the request for the deployment the rout
The Rust router core picks a deployment and hands the attempt to its host; ``PythonHost`` builds
the keyword arguments for that one deployment here and calls ``litellm.<op>`` with them. Nothing
here selects, retries, falls back or records cooldowns: the core does all of that. Nothing here
reads a Python router either, so this module outlives ``router_python_core.py``.
here selects, retries, falls back or records cooldowns: the core does all of that.
"""
import re
@ -61,6 +60,16 @@ POSITIONAL: Final = MappingProxyType(
)
# The arguments an endpoint takes without a default, which a call that leaves one out raises
# ``TypeError`` for before it routes.
REQUIRED: Final = MappingProxyType(
{
"atext_completion": ("model", "prompt"),
"text_completion": ("model", "prompt"),
}
)
def request_from_call(op: str, args: tuple[object, ...], kwargs: Mapping[str, object]) -> Mapping[str, object]:
"""The call's arguments as one keyword request. Raises ``TypeError`` like the endpoint would."""
names: Final = POSITIONAL.get(op, ())
@ -71,6 +80,12 @@ def request_from_call(op: str, args: tuple[object, ...], kwargs: Mapping[str, ob
if name in request:
raise TypeError(f"{op}() got multiple values for argument {name!r}")
request[name] = value
missing: Final = [name for name in REQUIRED.get(op, ()) if name not in request]
if missing:
listed: Final = " and ".join(repr(name) for name in missing)
raise TypeError(
f"{op}() missing {len(missing)} required positional argument{'s' if len(missing) > 1 else ''}: {listed}"
)
return request

35
litellm/router_core.py
View file

@ -1,16 +1,13 @@
"""
``RouterCore``: the interface ``litellm.Router`` delegates to, and ``make_core``, which picks one.
``RouterCore``: the interface ``litellm.Router`` delegates to, and ``make_core``, which builds it.
``PythonCore`` (``litellm/router_python_core.py``) is today's routing logic. ``RustCore``
(``litellm/router_rust/core.py`` over litellm-rust/crates/router-core) is a separate implementation
that reaches outside itself only through its host; it does not use ``PythonCore``. The core is
picked once, when the router is built, from its ``RouterConfig``.
The core is ``RustCore`` (``litellm/router_rust/core.py`` over litellm-rust/crates/router-core),
which reaches outside itself only through its host.
"""
from collections.abc import Mapping
from typing import TYPE_CHECKING, Final, Protocol, runtime_checkable
from typing import TYPE_CHECKING, Protocol, runtime_checkable
from litellm._logging import verbose_router_logger
from litellm.router_config import RouterConfig
if TYPE_CHECKING:
@ -52,25 +49,7 @@ class RouterCore(Protocol):
def make_core(config: RouterConfig, host: "PythonHost") -> RouterCore:
"""
The Rust core when the route catalog's router rule allows it and the Rust core serves this
config, else the Python core. Decided once, for the router's life, from the config alone: the
Rust core never builds, wraps or falls back to a Python core.
"""
from litellm.router_python_core import PythonCore
from litellm.rust_bridge.catalog import Route, RouteContext, decision
from litellm.rust_bridge.configuration import Decision
"""The router's core: the Rust core, built once for the router's life from its config."""
from litellm.router_rust.core import RustCore
chosen: Final = decision(RouteContext(Route.ROUTER))
if chosen is not Decision.PYTHON:
from litellm.router_rust.core import RustCore
from litellm.rust_bridge._native import RustBridgeDeclined
try:
return RustCore.build(config, host) # pyright: ignore[reportReturnType] # RustCore serves the RouterCore entry points; other members raise RustCoreUnsupported
except RustBridgeDeclined as declined:
# A config the Rust core cannot take keeps the router on Python while both cores exist.
if chosen is Decision.RUST_REQUIRED:
raise
verbose_router_logger.info("Router served by the Python core: %s", declined)
return PythonCore(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the arguments were bound against PythonCore's own signature
return RustCore.build(config, host) # pyright: ignore[reportReturnType] # RustCore serves the RouterCore entry points; other members raise RustCoreUnsupported

33
litellm/router_deprecated.py
View file

@ -92,10 +92,9 @@ def _note(name: str, replacement: str) -> None:
def _as_router(obj: "DeprecatedRouterMembers") -> "Router":
"""The Router this mixin is part of (it is mixed into ``litellm.Router`` only)."""
from litellm.router import Router
from litellm.router_python_core import PythonCore
if isinstance(obj, (Router, PythonCore)):
return obj # pyright: ignore[reportReturnType] # a PythonCore carries every member the mixin reads
if isinstance(obj, Router):
return obj
raise TypeError(f"{type(obj).__name__} is not a litellm.Router; DeprecatedRouterMembers is mixed into Router only")
@ -112,9 +111,12 @@ def _public_attribute_error(obj: object, core_name: str, public_name: str, error
class _DeprecatedAttribute(Generic[_T]):
"""A former Router attribute, read through an entry point. Assignment still replaces it."""
def __init__(self, replacement: str, read: Callable[["Router"], _T]) -> None:
def __init__(self, replacement: str, read: Callable[["Router"], _T], unbound: object = None) -> None:
self.replacement = replacement
self.read = read
# What a class-level read returns when the core keeps no member of this name: a former
# method reads back as a function taking the router, as it did on the Router class.
self.unbound = unbound
self.name = ""
def __set_name__(self, owner: type, name: str) -> None:
@ -129,7 +131,7 @@ class _DeprecatedAttribute(Generic[_T]):
def __get__(self, obj: "DeprecatedRouterMembers | None", owner: type) -> object:
if obj is None:
return getattr(owner, "_core_" + self.name, self)
return getattr(owner, "_core_" + self.name, self if self.unbound is None else self.unbound)
_note(self.name, self.replacement)
router: Final = _as_router(obj)
try:
@ -216,6 +218,23 @@ def _unreserved_picker(router: "Router") -> Callable[..., Awaitable[object]]:
return picker
async def _async_get_available_deployment(
router: "Router",
model: str,
request_kwargs: dict, # mutable-ok: the old parameter type, unchanged
messages: list[dict[str, str]] | None = None, # mutable-ok: the old parameter type, unchanged
input: str | list | None = None, # mutable-ok: the old parameter type, unchanged
specific_deployment: bool | None = False,
) -> object:
"""``Router.async_get_available_deployment`` read on the class: the router's pick, unbound."""
return await _unreserved_picker(router)(model, request_kwargs, messages, input, specific_deployment)
_async_get_available_deployment.__name__ = _async_get_available_deployment.__qualname__ = (
"async_get_available_deployment"
)
def _slack_alerting_logger(router: "Router") -> "SlackAlerting":
"""``observe().slack_alerting_logger``; a router without one has no such attribute, as before."""
logger: Final = router.observe().slack_alerting_logger
@ -282,7 +301,9 @@ class DeprecatedRouterMembers:
cooldown_cache = _attr("health().cooldowns", lambda r: r.health().cooldowns)
health_state_cache = _attr("health().health_checks", lambda r: r.health().health_checks)
deployment_latency_map = _attr("health().latency", lambda r: r.health().latency)
async_get_available_deployment = _attr("select(model, request_kwargs)", lambda r: _unreserved_picker(r))
async_get_available_deployment = _DeprecatedAttribute(
"select(model, request_kwargs)", lambda r: _unreserved_picker(r), unbound=_async_get_available_deployment
)
# -- observers ------------------------------------------------------------------------------
slack_alerting_logger = _attr("observe().slack_alerting_logger", _slack_alerting_logger)

18017
litellm/router_python_core.py
View file

File diff suppressed because it is too large Load diff

752
litellm/router_rust/core.py
View file

File diff suppressed because it is too large Load diff

109
litellm/router_rust/invoke/endpoints.py
View file

@ -73,6 +73,11 @@ class EndpointCore(Protocol):
"""The core itself, for the provider handlers and helpers that take a router."""
...
@property
def public(self) -> object:
"""The ``Router`` that owns the core, for litellm functions that take ``router=``."""
...
@property
def host(self) -> "PythonHost": ...
@ -232,6 +237,9 @@ class Plan:
fail_key: CountKey | None = "deployment"
# One attempt, no retries and no fallbacks.
single_attempt: bool = False
# Whether the attempt's metadata carries the walk's keys (``original_model_group``, the retry
# counts): not for an endpoint whose metadata goes to the provider as the caller sent it.
stamps_walk: bool = True
# The named parameter the endpoint takes beside **kwargs, sent between the deployment params and
# the request (`{**data, name: value, **kwargs}`).
named: tuple[str, ...] = ()
@ -267,8 +275,9 @@ class Plan:
def stamps_failed_deployment(self) -> bool:
"""Whether the attempt stamps ``failed_deployment_id`` on its error, which weighted failover
reads: every endpoint but file and batch creation (whose result belongs to the deployment
that served it), moderation's pass-through and search."""
return self.kind not in ("batch", "cancel_batch", "create_file", "moderation", "search")
that served it), moderation's pass-through, search and vector store creation (whose
handler calls the provider without the attempt wrapper)."""
return self.kind not in ("batch", "cancel_batch", "create_file", "moderation", "search", "vector_store_create")
@property
def metadata_field(self) -> str:
@ -450,6 +459,7 @@ PLANS: Final[Mapping[str, Plan]] = MappingProxyType(
stamp="none",
fail_key="group",
single_attempt=True,
stamps_walk=False,
),
"amoderation": Plan(
"amoderation",
@ -636,6 +646,44 @@ _SYNC_VECTOR_STORE_FILES: Final[Mapping[str, str]] = MappingProxyType(
}
)
# The endpoints ``RustEndpoints`` serves with a method of its own, by the same name.
_DIRECT_ENDPOINTS: Final = frozenset(
{
"avector_store_create",
"acreate_file",
"acreate_batch",
"acancel_batch",
"aretrieve_batch",
"alist_batches",
"aadapter_completion",
"abatch_completion",
"abatch_completion_one_model_multiple_requests",
"abatch_completion_fastest_response",
"schedule_acompletion",
"priority_header_update",
"aguardrail",
"asearch",
}
)
def endpoint_names() -> frozenset[str]:
"""Every endpoint name ``RustEndpoints.member`` may serve: what ``dir()`` of a router lists."""
return frozenset(
{
*PLANS,
*_HOST_GENERIC,
*_VIDEO_DIRECT,
*_SYNC_GENERIC,
*_UNBRANCHED,
*_ASSISTANTS,
*_MANAGED_AGENTS,
*_SYNC_VECTOR_STORE_FILES,
*_DIRECT_ENDPOINTS,
"create_interaction",
}
)
def plan_for(endpoint: str) -> Plan:
plan: Final = generic_plan(endpoint)
@ -968,8 +1016,7 @@ def _search_attempt(
kwargs: dict[str, object], # mutable-ok: the attempt's own kwargs
group: str,
) -> AttemptCall:
"""`SearchAPIRouter.async_search_with_fallbacks_helper`: one of the search tools named
``group`` at random, called with its provider, credentials and params under the request's."""
"""One of the search tools named ``group`` at random, called with its provider, credentials and params under the request's."""
import random
tools: Final = [tool for tool in core.endpoints.search_tools if as_dict(tool).get("search_tool_name") == group]
@ -1359,7 +1406,15 @@ class RustEndpoints:
self._search_tools: list[object] = tools # mutable-ok: the router's live list
self._guardrails: list[object] = list(_as_items(config.get("guardrail_list"))) # mutable-ok: same
parsed: Final = parse_default_priority(config.get("default_priority"))
self._default_priority: Final = None if isinstance(parsed, InvalidPriority) else parsed
if isinstance(parsed, InvalidPriority):
from litellm._logging import verbose_router_logger
verbose_router_logger.warning(
"router_settings.default_priority is ignored, no request is queued by default: %s",
parsed.message,
)
# apply() assigns it after construction, kept as given, as the Python core's attribute is.
self._default_priority: int | None = None if isinstance(parsed, InvalidPriority) else parsed
self._polling_interval: Final = config.get("polling_interval")
self._caching_groups: Final = config.get("caching_groups")
self._default_max_parallel_requests: object = config.get("default_max_parallel_requests")
@ -1375,15 +1430,17 @@ class RustEndpoints:
self._cache_responses = value
def assign_setting(self, name: str, value: object) -> None:
"""``search_tools``, ``guardrail_list`` or ``default_max_parallel_requests`` assigned after
construction, kept as given: later searches, guardrail picks and slot limits built from
then on read it, as the Python core's attributes are read."""
"""``search_tools``, ``guardrail_list``, ``default_max_parallel_requests`` or
``default_priority`` assigned after construction, kept as given: later searches, guardrail
picks, slot limits and queued calls read it, as the Python core's attributes are read."""
if name == "search_tools":
self._search_tools = value # pyright: ignore[reportAttributeAccessIssue] # stored as given, as the Python core stores it
elif name == "guardrail_list":
self._guardrails = value # pyright: ignore[reportAttributeAccessIssue] # same
elif name == "default_max_parallel_requests":
self._default_max_parallel_requests = value
elif name == "default_priority":
self._default_priority = value # pyright: ignore[reportAttributeAccessIssue] # stored as given, as the Python core stores it
else:
raise ValueError(f"not an endpoint setting: {name}")
@ -1490,23 +1547,8 @@ class RustEndpoints:
return self._sync_vector_store_files(name)
if name == "create_interaction":
return self._factory("create_interaction", "interactions.create")
direct: Final = {
"avector_store_create": self.avector_store_create,
"acreate_file": self.acreate_file,
"acreate_batch": self.acreate_batch,
"acancel_batch": self.acancel_batch,
"aretrieve_batch": self.aretrieve_batch,
"alist_batches": self.alist_batches,
"aadapter_completion": self.aadapter_completion,
"abatch_completion": self.abatch_completion,
"abatch_completion_one_model_multiple_requests": self.abatch_completion_one_model_multiple_requests,
"abatch_completion_fastest_response": self.abatch_completion_fastest_response,
"schedule_acompletion": self.schedule_acompletion,
"priority_header_update": self.priority_header_update,
"aguardrail": self.aguardrail,
"asearch": self.asearch,
}.get(name)
if direct is not None:
if name in _DIRECT_ENDPOINTS:
direct: Final[Callable[..., object]] = getattr(self, name) # pyright: ignore[reportAny] # one of this class's own endpoint methods
return direct
if name == "responses":
return self._sync_factory(name, PLANS["responses"])
@ -1538,7 +1580,10 @@ class RustEndpoints:
client: object,
kwargs: dict[str, object], # mutable-ok: the wrapper's own kwargs
) -> object:
plan: Final = generic_plan(name) or _generic(name, host_op)
listed: Final = generic_plan(name)
# The endpoint's own plan, unless the wrapper calls another litellm function for it
# (``avector_store_create`` without a model calls ``vector_stores.main.acreate``).
plan: Final = listed if listed is not None and listed.host_op == host_op else _generic(name, host_op)
if name in _CONTAINER_ENDPOINTS:
return await self._containers(plan, custom_llm_provider, kwargs)
if name in _RESPONSE_ID_ENDPOINTS:
@ -1618,7 +1663,7 @@ class RustEndpoints:
if provider_kwargs.get("model"):
return self._core.call_sync(plan, provider_kwargs, _SINGLE_ATTEMPT)
if name == "vector_store_search":
return self._invoke(plan.host_op)(**{**provider_kwargs, "router": self._core.router})
return self._invoke(plan.host_op)(**{**provider_kwargs, "router": self._core.public})
return self._invoke(plan.host_op)(**provider_kwargs)
sync_wrapper.__name__ = name
@ -1695,7 +1740,7 @@ class RustEndpoints:
if kwargs.get("model"):
return await self._ageneric(plan, kwargs)
if name == "avector_store_search":
return await self._direct(plan.host_op, {**kwargs, "router": self._core.router})
return await self._direct(plan.host_op, {**kwargs, "router": self._core.public})
return await self._direct(plan.host_op, kwargs)
async def _interactions(
@ -1880,7 +1925,12 @@ class RustEndpoints:
assert isinstance(scheduler, Scheduler)
async def healthy_deployments() -> Sequence[object]:
healthy: Final = await self._core.select(model, dict(request), kind="healthy")
try:
healthy: Final = await self._core.select(model, dict(request), kind="healthy")
except Exception: # noqa: BLE001 # `_async_get_healthy_deployments` lists none on any lookup failure
# It lists none while every deployment cools down, or for a group it cannot look up:
# the request keeps its place in the queue, and the call itself raises the error.
return []
return list(healthy) if isinstance(healthy, Sequence) else [healthy]
await scheduler.wait_for_turn(
@ -2440,6 +2490,7 @@ __all__ = [
"attempt_call",
"clean_realtime_session_deployment",
"compacted_attempt",
"endpoint_names",
"failover_fields",
"generic_plan",
"log_refused_call",

110
litellm/router_rust/invoke/streaming.py
View file

@ -108,11 +108,16 @@ class StreamAttempt:
deployment_id: str | None
attempted_retries: int
committed_budget: int | None = None
# The group a pre-routing hook selected for ``group``, which a fallback lookup keys on first.
routed: str | None = None
# The deployment's max_parallel_requests slot the attempt holds; an open chat stream takes it.
slot: contextlib.ExitStack | None = None
# Set by the wrapper that reports the stream's outcome when it ends; the attempt's answer is then
# deferred and counts only the request and the tokens read so far (``counted_tokens``).
reports_outcome: bool = False
# Whether the logging object had logged a failure when the wrapper took the stream: an earlier
# attempt's failure blocks this one's report, the stream's own log of its failure does not.
failure_logged_before: frozenset[str] = frozenset()
counted_tokens: int = 0
outcome_reported: bool = False
# The background recording of the stream's failure, which a same-group retry waits for so it
@ -132,6 +137,7 @@ class StreamAttempt:
"deployment_id": self.deployment_id,
"attempted_retries": self.attempted_retries,
"committed_budget": self.committed_budget,
"routed": self.routed,
}
def policy(self, trigger: BaseException | None = None) -> Mapping[str, object]:
@ -192,6 +198,13 @@ class StreamAttempt:
"""The outcome callbacks that run for this attempt when its stream ends."""
return self.core.outcome_callbacks(self.kwargs, sync=sync)
def take_outcome(self) -> None:
"""The wrapper reports this attempt's outcome when its stream ends."""
self.reports_outcome = True
self.failure_logged_before = frozenset(
event for event in ("async_failure", "sync_failure") if failure_already_logged(self.kwargs, event)
)
def _claim(self) -> bool:
if not self.reports_outcome or self.outcome_reported:
return False
@ -209,9 +222,9 @@ class StreamAttempt:
def report_failure(self, error: BaseException) -> None:
"""The stream failed after it opened: count the failure, cool down by the rules and count one
more request (``async_deployment_callback_on_failure``), in the background. A logging object
shared with an earlier failure blocks it: the stream logs through ``attempt_scoped_copy``, read
here, when the stream fails."""
if self._claim() and not failure_already_logged(self.kwargs, "async_failure"):
shared with an earlier failure blocks it, read when the wrapper took the stream: the stream's own
failure log marks the same object before this runs."""
if self._claim() and "async_failure" not in self.failure_logged_before:
self.failure_recording = _record_in_background(
self.core.report_stream_failure(self._failure_logged(error), self._callbacks(sync=False))
)
@ -235,7 +248,7 @@ class StreamAttempt:
def report_failure_sync(self, error: BaseException) -> None:
"""A sync stream failed after it opened: counted and cooled down, with no request counted."""
if self._claim() and not failure_already_logged(self.kwargs, "sync_failure"):
if self._claim() and "sync_failure" not in self.failure_logged_before:
logged: Final = self._failure_logged(error)
_record_now(lambda: self.core.report_stream_failure_sync(logged, self._callbacks(sync=True)))
@ -284,20 +297,28 @@ async def _areported(
items: AsyncIterable[_ItemT],
total_tokens: Callable[[object], int],
observe: Callable[[object], None] | None = None,
*,
fails_before_first_item: bool = True,
) -> AsyncGenerator[_ItemT, None]:
"""``items``, reporting the attempt's stream once they end: a success with the tokens
``total_tokens`` reads off the last item, or the failure they raise. Closed early, nothing is
reported. The wrapper marks the attempt (``reports_outcome``) when it wraps the stream, before
the attempt answers."""
``total_tokens`` reads off the last item, or the failure they raise, unless they raise it before
their first item and ``fails_before_first_item`` is off. Closed early, nothing is reported. The
wrapper marks the attempt (``reports_outcome``) when it wraps the stream, before the attempt
answers."""
last: object = None # rebind-ok: the newest item
started = False # rebind-ok: whether an item came
try:
async for item in items:
last = item
started = True
if observe is not None:
observe(item)
yield item
except Exception as error:
attempt.report_failure(error)
if started or fails_before_first_item:
attempt.report_failure(error)
else:
attempt.outcome_reported = True # rebind-ok: a stream that never opened counts nothing
raise
attempt.report_success(total_tokens(last), last)
@ -918,7 +939,7 @@ def wrap_chat_stream( # noqa: C901 # one function per Python wrapper it ports
after content, the provider's error is raised. ``held_slot`` (the deployment's max_parallel_requests
slot) is released when the stream is exhausted, closed, or falls back."""
slot: Final = held_slot if held_slot is not None else contextlib.ExitStack()
attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome
attempt.take_outcome()
async def close_model_response() -> None:
if not hasattr(model_response, "aclose"):
@ -1013,7 +1034,7 @@ def wrap_chat_stream_sync(
attempt: StreamAttempt, model_response: CustomStreamWrapper, first_chunk: _FirstChunk[Iterator[object]]
) -> FallbackAwareStreamWrapper:
"""``_completion_streaming_iterator``: the sync twin, which never continues after content."""
attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome
attempt.take_outcome()
def stream_with_fallbacks() -> Generator[object, None, None]:
fallback_response: object = None # rebind-ok: tracked for cleanup
@ -1253,7 +1274,7 @@ def wrap_anthropic_stream(
replays_a_finished_agentic_turn: Final = isinstance(source_iterator, FakeAnthropicMessagesStreamIterator)
model: Final = attempt.group
usage: Final = _AnthropicStreamUsage()
attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome
attempt.take_outcome()
async def close_source() -> None:
with anyio.CancelScope(shield=True), contextlib.suppress(BaseException):
@ -1527,7 +1548,7 @@ def wrap_responses_stream( # noqa: C901 # one function per Python wrapper it p
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator, get_openai_response_types
source_iterator: Final = response
attempt.reports_outcome = True # rebind-ok: the wrapper reports this attempt's outcome
attempt.take_outcome()
openai_types: Final = get_openai_response_types()
terminal_event_types: Final = (
openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
@ -1664,7 +1685,7 @@ def wrap_responses_stream( # noqa: C901 # one function per Python wrapper it p
else cast("MidStreamFallbackError", stream_error) # cast-ok: not a context window error
)
if response_id_exposed:
_raise_original_mid_stream_error(e)
_raise_released_stream_error(attempt, e)
async with contextlib.aclosing(
_responses_fallback_attempt(attempt, e, source_iterator, wrapper.adopt_fallback_headers, held)
) as fallback_stream:
@ -1686,6 +1707,33 @@ def wrap_responses_stream( # noqa: C901 # one function per Python wrapper it p
return wrapper
def _raise_released_stream_error(attempt: StreamAttempt, e: MidStreamFallbackError) -> NoReturn:
"""A background stream whose response id is already out raises its provider error as is; a content
policy refusal before output reads as the router's error would after the walk found no fallback."""
from litellm.litellm_core_utils.routing_common_utils import format_content_policy_no_fallback_note
from litellm.router_utils.fallback_event_handlers import fallback_lookup_groups, fallbacks_disabled_for_request
original: Final = e.original_exception
request: Final = attempt.request
model_group: Final = request.get("model")
if (
attempt.core._host.dynamic().expose_router_debug_in_errors # pyright: ignore[reportPrivateUsage] # the attempt's own core
and isinstance(original, litellm_exceptions.ContentPolicyViolationError)
and isinstance(model_group, str)
and not fallbacks_disabled_for_request(request)
and _mid_stream_error_before_output(e)
):
walked: Final = attempt.core._walked_fallbacks # pyright: ignore[reportPrivateUsage] # the attempt's own core
original.message += format_content_policy_no_fallback_note(
model_group,
fallback_lookup_groups(request, model_group),
request.get("content_policy_fallbacks", walked.get("content_policy_fallbacks")),
request.get("fallbacks", walked.get("fallbacks")),
attempt.depth > 0,
)
_raise_original_mid_stream_error(e)
def _responses_continuation_request(
attempt: StreamAttempt, e: MidStreamFallbackError, source_iterator: "BaseResponsesAPIStreamingIterator"
) -> Mapping[str, object] | None:
@ -1782,6 +1830,18 @@ async def _responses_fallback_attempt(
verbose_router_logger.debug("stream_with_fallbacks(aresponses): error closing fallback: %s", exc)
def _reported_transcription_stream(attempt: StreamAttempt, response: object) -> object:
"""A transcription stream reports its attempt's outcome when it ends: a stream that fails after it
opened cools its deployment down by the rules, as the router's failure callback does."""
from litellm.litellm_core_utils.audio_utils.transcription_streaming import LoggingAsyncTranscriptionStream
if not isinstance(response, LoggingAsyncTranscriptionStream):
return response
attempt.take_outcome()
response._iterator = _areported(attempt, response._iterator, lambda _last: 0) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the stream's own iterator, wrapped
return response
# --- the attempt hook ---------------------------------------------------------------------------
@ -1798,13 +1858,37 @@ async def open_attempt_stream(attempt: StreamAttempt, response: object) -> objec
return wrap_chat_stream(attempt, response, first_chunk, held_slot)
if not attempt.kwargs.get("stream"):
return response
if attempt.op == "atranscription":
return _reported_transcription_stream(attempt, response)
if attempt.op == "atext_completion":
return _reported_text_completion_stream(attempt, response)
if attempt.op == "aresponses" and isinstance(response, BaseResponsesAPIStreamingIterator):
# `_aresponses_streaming_iterator` wraps a call that asked for a stream, not one the
# deployment's own `stream` setting forced.
if not attempt.request.get("stream"):
return response
return wrap_responses_stream(attempt, response)
if attempt.op in ("aanthropic_messages", "anthropic_messages") and hasattr(response, "__aiter__"):
return wrap_anthropic_stream(attempt, cast("AsyncIterator[bytes]", response)) # cast-ok: a byte stream
return response
def _reported_text_completion_stream(attempt: StreamAttempt, response: object) -> object:
"""A text-completion stream reports its attempt's outcome when it ends, as litellm's logging
does for it: a success with the stream's tokens, or a failure after the first chunk. One that
fails before its first chunk counts nothing, as its logging carries no deployment then."""
from litellm.text_completions.streaming_iterator import TextCompletionStreamWrapper
if not isinstance(response, TextCompletionStreamWrapper):
return response
attempt.take_outcome()
chunks: Final = cast("AsyncIterable[object]", response.completion_stream) # cast-ok: an untyped attribute
response.completion_stream = _areported( # rebind-ok: the stream's own iterator, wrapped
attempt, chunks, _chat_total_tokens, fails_before_first_item=False
)
return response
def open_attempt_stream_sync(attempt: StreamAttempt, response: object) -> object:
if attempt.op == "completion" and isinstance(response, CustomStreamWrapper):
if response.completion_stream is None and response.make_call is not None: # pyright: ignore[reportUnknownMemberType] # untyped attributes

44
litellm/router_rust/invoke/walk.py
View file

@ -41,6 +41,17 @@ _SERVED_KEYS: Final = ("model_info", "deployment")
_HOP_KWARGS: Final = ("fallback_depth", "max_fallbacks", "attempted_targets", "include_fallback_errors")
def _reentrant_hop_stamps(request: Mapping[str, object], field: str) -> Mapping[str, object]:
"""A hop ``run_async_fallback`` re-entered with its own ``attempted_targets`` keeps the fallback
stamps it carries in place of the root level's (``async_function_with_fallbacks`` skips them)."""
from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets
if not isinstance(request.get("attempted_targets"), AttemptedFallbackTargets):
return {}
bucket: Final = as_str_mapping(request.get(field)) or {}
return {key: bucket[key] for key in ("attempted_fallbacks", "original_model_group") if key in bucket}
class _Scope:
__slots__ = ("kwargs", "metadata")
@ -146,6 +157,7 @@ class WalkMetadata:
root: Final[dict[str, object]] = {"attempted_fallbacks": 0} # mutable-ok: the root level's metadata
if model is not None:
root["original_model_group"] = model
root.update(_reentrant_hop_stamps(request, field))
self._request: Final = request
self._field: Final = field
self._params_base_model: Final = params_base_model
@ -289,8 +301,11 @@ class WalkMetadata:
"""The level's ``request_retry_count``, as the retry loop has counted it so far."""
return self._scope(step.get("scope")).metadata.get("request_retry_count")
def stamp(self, step: Mapping[str, object], kwargs: Mapping[str, object]) -> Mapping[str, object]:
"""The attempt's kwargs with its level's walk metadata and hop kwargs laid over them."""
def stamp(
self, step: Mapping[str, object], kwargs: Mapping[str, object], deployment: Mapping[str, object] | None = None
) -> Mapping[str, object]:
"""The attempt's kwargs with its level's walk metadata and hop kwargs laid over them.
``deployment``: the deployment the attempt calls, whose own ``mock_timeout`` a hop keeps."""
scope: Final = self._scope(step.get("scope"))
metadata: Final = scope.metadata
metadata["attempted_retries"] = step.get("retry")
@ -299,10 +314,31 @@ class WalkMetadata:
self._write_through(
scope, {"attempted_retries": metadata["attempted_retries"], "max_retries": metadata["max_retries"]}
)
served: Final = dict(as_str_mapping(kwargs.get(self._field)) or {})
# `include_fallback_errors` is the walk's own flag: the fallback loggers see it, the provider never does.
hop: Final = {
key: scope.kwargs[key] for key in _HOP_KWARGS if key in scope.kwargs and key != "include_fallback_errors"
}
depth: Final = hop.get("fallback_depth")
given: Final = kwargs.get(self._field)
# The requested group's attempts stamp the caller's own dict in place, as
# `_update_kwargs_with_deployment` stamps it: the proxy's post-call write-backs land in the dict
# its spend row reads. A fallback hop stamps its own copy (`run_async_fallback`'s kwargs).
served: Final = (
given # pyright: ignore[reportUnknownVariableType] # the caller's metadata dict
if isinstance(given, dict) and not (isinstance(depth, int) and depth > 0)
else dict(as_str_mapping(given) or {})
)
served.update({key: metadata[key] for key in _WALK_KEYS if key in metadata})
metadata.update({key: served.get(key) for key in _SERVED_KEYS})
return {**kwargs, **{key: scope.kwargs[key] for key in _HOP_KWARGS if key in scope.kwargs}, self._field: served}
# `async_function_with_fallbacks` pops the request's `mock_timeout` for the requested group's own
# retries: no hop sees it. A deployment's own `mock_timeout` is merged in per call, so it stays.
own: Final = as_str_mapping((deployment or {}).get("litellm_params")) or {}
sent: Final = (
{key: value for key, value in kwargs.items() if key != "mock_timeout"}
if isinstance(depth, int) and depth > 0 and "mock_timeout" not in own
else kwargs
)
return {**sent, **hop, self._field: served}
async def _log_failed_fallbacks(ctx: "RoutingContext", detail: str, parent_otel_span: "Span | None") -> None:

36
litellm/router_rust/steps/callbacks.py
View file

@ -218,6 +218,12 @@ class OptionalChecks:
self.optional_callbacks.append(callback)
self._host.callbacks.add_litellm_callback(callback)
def adopt(self, callback: object) -> None:
"""Own a check already registered with the host: list it in ``optional_callbacks``."""
if self.optional_callbacks is None:
self.optional_callbacks = []
self.optional_callbacks.append(callback)
def ensure_deployment_affinity_callback(self) -> None:
"""``_ensure_deployment_affinity_callback``: the affinity check with every global flag off."""
if any(isinstance(cb, DeploymentAffinityCheck) for cb in self.optional_callbacks or ()):
@ -323,7 +329,9 @@ class OptionalChecks:
if self.optional_callbacks is None or not any(type(cb) is callback_cls for cb in self.optional_callbacks):
return
self.optional_callbacks = [cb for cb in self.optional_callbacks if type(cb) is not callback_cls]
others: Final = tuple(owner for owner in live_check_owners if id(owner) != id(self._owner))
others: Final = tuple(
owner for owner in live_check_owners if id(owner) != id(self._owner) and _router_alive(owner)
)
if any(any(type(cb) is callback_cls for cb in (owner.optional_callbacks or ())) for owner in others):
return
callbacks: Final = _callbacks(self._host)
@ -552,22 +560,42 @@ class RustOptionalChecks:
def optional_callbacks(self) -> "list[object] | None": # mutable-ok: the router's live list
return self._opt.optional_callbacks
@optional_callbacks.setter
def optional_callbacks(self, value: "list[object] | None") -> None: # mutable-ok: the list a caller assigns
"""Assigning the list replaces it, as on the Python router: the checks read the new one."""
self._opt.optional_callbacks = value
@property
def router_budget_logger(self) -> RouterBudgetLimiting | None:
return self._opt.router_budget_logger
# The affinity settings are plain attributes on the Python router: an assignment is what an
# affinity check added afterwards (`add_optional_pre_call_checks`) is built with.
@property
def deployment_affinity_ttl_seconds(self) -> object:
return self._opt.deployment_affinity_ttl_seconds
@deployment_affinity_ttl_seconds.setter
def deployment_affinity_ttl_seconds(self, value: object) -> None:
self._opt.deployment_affinity_ttl_seconds = value
@property
def session_affinity_repin_on_failover(self) -> object:
return self._opt.session_affinity_repin_on_failover
@session_affinity_repin_on_failover.setter
def session_affinity_repin_on_failover(self, value: object) -> None:
self._opt.session_affinity_repin_on_failover = value
@property
def model_group_affinity_config(self) -> "dict[str, list[str]] | None": # mutable-ok: the argument as passed
return self._opt.model_group_affinity_config
@model_group_affinity_config.setter
def model_group_affinity_config(self, value: "dict[str, list[str]] | None") -> None:
self._opt.model_group_affinity_config = value
def add_optional_pre_call_checks(self, optional_pre_call_checks: "Sequence[object] | None") -> None:
self._opt.add(optional_pre_call_checks)
@ -648,3 +676,9 @@ __all__ = [
"filtered_ids",
"register_router_selector",
]
def _router_alive(owner: object) -> bool:
"""A core whose public ``Router`` was dropped holds no check: it only waits for the collector."""
handle: Final = vars(owner).get("_core_owner") if hasattr(owner, "__dict__") else None
return not callable(handle) or handle() is not None

10
litellm/router_rust/steps/precall.py
View file

@ -27,6 +27,7 @@ from litellm.litellm_core_utils.token_counter import (
single_encoding_openai_model,
split_embedding_input_tokens,
)
from litellm.router_attempt import for_requested_model
from litellm.router_python_host import python_host
from litellm.router_strategy.complexity_router.context_compaction import compaction_pending
from litellm.types.router import LiteLLM_Params
@ -229,10 +230,15 @@ class PreCallFacts:
return self._has_input and any(limit is not None for limit in self._windows.values())
def _tokenizer_model(self) -> str:
named: Final = self._registry._core_deployments_named(self._model) or () # pyright: ignore[reportPrivateUsage] # the host registry's index
# A wildcard deployment serves the requested model (`PatternMatchRouter.route`).
named: Final = tuple(
for_requested_model(record, self._name)
for record in self._registry._core_deployments_named(self._model) or () # pyright: ignore[reportPrivateUsage] # the host registry's index
)
return (
single_encoding_openai_model(
served_openai_models(named) or served_openai_models(self._record(id) for id in self._candidates)
served_openai_models(named)
or served_openai_models(for_requested_model(self._record(id), self._name) for id in self._candidates)
)
or ""
)

25
litellm/router_rust/steps/prerouting.py
View file

@ -65,6 +65,7 @@ from litellm.types.utils import (
from litellm.utils import provider_rejectable_params
if TYPE_CHECKING:
from litellm.caching.redis_cache import RedisCache
from litellm.router_python_host import PythonHost
from litellm.router_rust.surface.registry import HostRegistry
from litellm.router_strategy.auto_router.auto_router import PreRoutingHookResponse
@ -406,6 +407,13 @@ class PreRoutingStrategies:
if router._uses_deployment_pin: # pyright: ignore[reportPrivateUsage] # the Python router reads the same flag
self._ensure_deployment_affinity_callback()
def hand_over_affinity(self, adopt: Callable[[object], None]) -> None:
"""Give the affinity check registered while the deployments were built, before the optional
checks existed, to those checks: they own it from then on, as ``optional_callbacks``."""
if self._affinity is not None:
adopt(self._affinity)
self._affinity = None
def _ensure_deployment_affinity_callback(self) -> None:
"""``_ensure_deployment_affinity_callback``: the affinity check (every global flag off) a
complexity router's session pin needs, registered once."""
@ -568,7 +576,9 @@ def _record_routing_decision(
)
def _as_output_cap(value: object) -> int | None:
def as_output_cap(value: object) -> int | None:
"""A client-sent output cap coerced to an int: ints, floats and numeric strings, never bools
or negatives."""
if isinstance(value, bool) or not isinstance(value, int | float | str):
return None
try:
@ -578,7 +588,8 @@ def _as_output_cap(value: object) -> int | None:
return cap if cap >= 0 else None
def _stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str, int] | None:
def stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str, int] | None:
"""None means unstamped; an empty mapping means the original caller supplied no limit."""
stamped: Final = next(
(
bucket.get(CLIENT_OUTPUT_CEILING_METADATA_KEY)
@ -593,7 +604,7 @@ def _stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str
{
str(carrier): cap
for carrier, value in stamped.items()
if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := _as_output_cap(value)) is not None
if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None
}
)
@ -601,7 +612,7 @@ def _stamped_client_ceiling(request_kwargs: Mapping[str, object]) -> Mapping[str
def _restore_client_output_ceiling(
request_kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs, restored in place
) -> None:
callers_ceiling: Final = _stamped_client_ceiling(request_kwargs)
callers_ceiling: Final = stamped_client_ceiling(request_kwargs)
if callers_ceiling is None:
return
for carrier in OUTPUT_TOKEN_CEILING_PARAMS:
@ -710,6 +721,10 @@ class PreRouter:
self._ctx: Final = ctx
self._session_cache: Final = DualCache(redis_cache=routing_cache.redis_cache, in_memory_cache=InMemoryCache())
def attach_redis_cache(self, redis_cache: "RedisCache") -> None:
"""``update_redis_cache``: the Claude Code session bindings are shared through it from now on."""
self._session_cache.attach_redis_cache(redis_cache)
# -- the step --------------------------------------------------------------------------------
async def pre_route(
@ -1104,4 +1119,4 @@ def _allowlist(params: Mapping[str, object]) -> frozenset[str]:
return frozenset(entry for entry in declared if isinstance(entry, str))
__all__ = ["PreRouter", "PreRoutingStrategies", "is_strategy_marker"]
__all__ = ["PreRouter", "PreRoutingStrategies", "as_output_cap", "is_strategy_marker", "stamped_client_ceiling"]

24
litellm/router_rust/surface/health.py
View file

@ -255,6 +255,17 @@ class RustHealthCallbacks:
deployment: Final = self._deployment_of(deployment_id)
return deployment is not None and _authenticates_per_caller(deployment.litellm_params)
def ran_as_request_deployment(self, logged_id: object, deployment_id: object) -> bool:
"""Whether a call against ``deployment_id`` ran under credentials the request supplied
(``_handle_clientside_credential``): as ``logged_id``, a request-scoped deployment the router
does not hold, so its failure is that id's, never the shared deployment's."""
return (
logged_id is not None
and deployment_id is not None
and str(logged_id) != str(deployment_id)
and self._deployment_of(logged_id) is None
)
def failure_exits(
self, exception: object, logged: Mapping[str, object], deployment_id: object
) -> Mapping[str, bool]:
@ -498,16 +509,17 @@ class RustHealthCallbacks:
completion_response: object,
start_time: object,
end_time: object,
) -> None:
"""Count a success and the call's tokens and request against a tracked deployment."""
await self.report_stream_success(kwargs, completion_response)
) -> str | None:
"""Count a success and the call's tokens and request against a tracked deployment; returns
the counter's key, as the sync callback does."""
return await self.report_stream_success(kwargs, completion_response)
async def report_stream_success(
self,
kwargs: Mapping[str, object],
completion_response: object,
callbacks: Mapping[str, bool] | None = None,
) -> None:
) -> str | None:
"""``deployment_callback_on_success``, for a stream that ended. The re-pin is the host's half
of that callback, so it reads ``callbacks`` here."""
from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
@ -528,7 +540,7 @@ class RustHealthCallbacks:
total: Final = float(tokens) if isinstance(tokens, int | float) else 0.0
counted: Final = get_counted_usage_tokens(as_str_mapping(kwargs.get("litellm_params")) or {})
usage_model: Final = target[1].get("deployment")
await self._drive_cache_steps(
finished: Final = await self._drive_cache_steps(
self._handle.report_success(
target[0],
total if counted is None else max(0.0, total - counted),
@ -538,6 +550,8 @@ class RustHealthCallbacks:
callbacks,
)
)
key: Final = finished.get("key")
return key if isinstance(key, str) else None
async def count_realtime_session_success(self, deployment_id: str) -> None:
"""``increment_deployment_successes_for_current_minute`` for a realtime session that ended

20
litellm/router_rust/surface/model_list.py
View file

@ -187,10 +187,16 @@ class HostedModelList(list[Entry]):
self,
handle: Callable[[], "RouterCoreHandle | None"],
apply: Callable[[Mapping[str, object]], None],
live: Callable[[object], Mapping[str, object]] = lambda _: {},
live_info: Callable[[object], Mapping[str, object]] = lambda _: {},
) -> None:
super().__init__()
self._handle: Final = handle
self._apply: Final = apply
# A deployment's litellm_params values its record carries as text, by id (``live_params``).
self._live: Final = live
# A deployment's model_info values its record reads back changed (non-string keys), by id.
self._live_info: Final = live_info
# The handle and registry version the contents were last read at.
self._read_at: tuple[object, int] | None = None
self._stable: _Stable = {}
@ -225,7 +231,7 @@ class HostedModelList(list[Entry]):
keys, counts = _keys(records, self._id_counts if appended else {})
entries: Final[list[Entry]] = [] # mutable-ok: filled here
for key, record in zip(keys, records, strict=True):
typed = typed_deployment(_record(record))
typed = self._with_live(typed_deployment(_record(record)))
held = self._stable.get(key) if key is not None else None
kept = held is not None and held[0] == typed
entry = held[1] if held is not None and kept else _HostedEntry(typed, self._entry_written)
@ -236,6 +242,18 @@ class HostedModelList(list[Entry]):
self._id_counts = counts
return entries
def _with_live(self, record: Mapping[str, object]) -> Mapping[str, object]:
params: Final = as_str_mapping(record.get("litellm_params"))
info: Final = as_str_mapping(record.get("model_info"))
model_id: Final = (info or {}).get("id")
live: Final = self._live(model_id) if params is not None else {}
live_info: Final = self._live_info(model_id) if info is not None else {}
if live:
record = {**record, "litellm_params": {**(params or {}), **live}} # rebind-ok: given its live values
if live_info:
record = {**record, "model_info": {**(info or {}), **live_info}} # rebind-ok: see above
return record
# -- a build ----------------------------------------------------------------------------------
def begin_build(self, *, replacing: bool) -> None:

120
litellm/router_rust/surface/registry.py
View file

@ -42,6 +42,7 @@ router and every ``registry()`` query dispatches the way it does on the Python c
import copy
import dataclasses
import json
import os
import re
import threading
@ -50,7 +51,7 @@ from functools import lru_cache, partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, Literal, TypeAlias, cast # noqa: TID251 # entries as a caller wrote them
from pydantic import TypeAdapter
from pydantic import BaseModel, TypeAdapter
from litellm import exceptions as litellm_exceptions
from litellm import get_secret_str
@ -93,7 +94,7 @@ from litellm.litellm_core_utils.routing_groups import apply_routing_group_priori
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS
from litellm.responses.bridge_check import keep_registered_mode
from litellm.router_api import RouterRegistry
from litellm.router_api import RouterRegistry, deployment_id
from litellm.router_rust.invoke.shaping import ROUTER_CLIENTS_TARGET
from litellm.router_rust.surface.indexes import HostedIndexes
from litellm.router_rust.surface.model_info import HostModelInfo
@ -579,6 +580,34 @@ def _index_key(value: object) -> str | None:
return value if isinstance(value, str) else None
def _round_trips(value: object) -> bool:
"""Whether JSON reads ``value`` back equal to itself."""
try:
loaded: Final[object] = json.loads(json.dumps(value, default=str)) # pyright: ignore[reportAny] # plain JSON
except (TypeError, ValueError):
return False
return loaded == value
def _has_non_str_keys(value: object) -> bool:
"""Whether ``value`` holds a dict key that is not a string, which a JSON record reads back as one."""
pending: list[object] = [value]
seen: set[int] = set()
while pending:
current = pending.pop()
if id(current) in seen:
continue
seen.add(id(current))
if isinstance(current, Mapping):
for key, item in current.items():
if not isinstance(key, str):
return True
pending.append(item)
elif isinstance(current, list | tuple):
pending.extend(current)
return False
class HostRegistry(HostModelInfo):
"""The live model_list and its indexes, kept the way the Python router core keeps them."""
@ -597,7 +626,13 @@ class HostRegistry(HostModelInfo):
self._deployment_index_lock: Final = threading.RLock()
self._core_model_group_alias: dict[str, str | RouterModelGroupAliasItem] = model_group_alias # mutable-ok: the live alias map # fmt: skip
# The live model_list: router-core's deployments read back, every write a changeset.
self._core_model_list: Final[HostedModelList] = HostedModelList(self._live_handle, self._hand_changes)
self._core_model_list: Final[HostedModelList] = HostedModelList(
self._live_handle, self._hand_changes, self.live_params, self.live_info
)
# The litellm_params values a deployment's JSON cannot carry (an Exception or model object as
# ``mock_response``, a callable ``azure_ad_token_provider``), by deployment id: an attempt is called with them, as Python's list held them.
self._live_params: Final[dict[str, Mapping[str, object]]] = {} # mutable-ok: written as deployments register
self._live_info: Final[dict[str, Mapping[str, object]]] = {} # mutable-ok: written as deployments register
# The indexes the registry hands out: router-core's, read back, every write a changeset.
self._indexes: Final = HostedIndexes(self._live_handle, self._hand_changes)
self._core_deployment_latency_map: dict[str, float] = {} # mutable-ok: the live latency map
@ -854,6 +889,7 @@ class HostRegistry(HostModelInfo):
self.build_runner(entries, build.path, resolve_secrets, install, partial(self._answer_shown, build))
finally:
self._core_model_list.end_build()
self._drop_removed_live_values()
if build.failure is not None:
# router-core served what it built before the entry that raised (`set_model_list`
# cleared the list first), so the router is left as the Python core leaves it.
@ -872,9 +908,11 @@ class HostRegistry(HostModelInfo):
"""The answer to one of router-core's build steps. A ``set_model_list`` deployment that
fails is left out when the router ignores invalid deployments (``_create_deployment``),
and retried on the next price data reload when its provider did not resolve."""
if build.failure is not None:
return {"invalid": True}
kind: Final = step["kind"]
if build.failure is not None:
# The build already failed: every later step is skipped, a secrets step unresolved.
refs_left: Final = step.get("refs")
return [None] * len(refs_left) if kind == "secrets" and isinstance(refs_left, list) else {"invalid": True} # pyright: ignore[reportUnknownArgumentType] # only the count is read
index: Final = int(str(step["index"]))
refs: Final = step.get("refs")
if kind == "secrets" and isinstance(refs, list) and not any(_ref_env_only(ref) for ref in refs):
@ -969,7 +1007,13 @@ class HostRegistry(HostModelInfo):
routed: Final = as_str_mapping(step["entry"]) or {}
routed_info: Final = as_str_mapping(routed.get("model_info")) or {}
if "id" not in entry.model_info:
entry.model_info["id"] = routed_info.get("id")
# router-core hashes the JSON form; a value JSON reads back changed hashes as the
# Python core hashes it (`deployment_id` over the params as given).
entry.model_info["id"] = (
routed_info.get("id")
if _round_trips(entry.litellm_params)
else deployment_id(entry.model_name, entry.litellm_params)
)
if entry.per_organization:
entry.litellm_params["organization"] = (as_str_mapping(routed.get("litellm_params")) or {}).get(
"organization"
@ -1159,6 +1203,40 @@ class HostRegistry(HostModelInfo):
deployment.litellm_params.rpm = getattr(deployment, "rpm") # noqa: B009 # an extra field of the entry # rebind-ok: the core sets it on the deployment it adds
if deployment.litellm_params.tpm is None and getattr(deployment, "tpm", None) is not None:
deployment.litellm_params.tpm = getattr(deployment, "tpm") # noqa: B009 # an extra field of the entry # rebind-ok: see above
live: Final = {
key: value
for key, value in {
**deployment.litellm_params.__dict__,
**(deployment.litellm_params.__pydantic_extra__ or {}),
}.items()
# Values JSON cannot carry: errors, callables, a str subclass that marks the value (an
# admin's secret reference, which only the deployment's own key may resolve), and any
# other value JSON reads back changed (a set, a tuple, a datetime), which the Python
# core keeps as given. A model is read back as its dump, as on the Python core.
if isinstance(value, BaseException)
or callable(value)
or (isinstance(value, str) and type(value) is not str)
or not (isinstance(value, BaseModel) or _round_trips(value))
}
# model_info values JSON reads back changed: a dict with keys that are not strings.
live_info: Final = {
key: value
for key, value in {
**deployment.model_info.__dict__,
**(deployment.model_info.__pydantic_extra__ or {}),
}.items()
if _has_non_str_keys(value)
}
model_id: Final = deployment.model_info.id
if model_id is not None:
if live:
self._live_params[model_id] = live
else:
self._live_params.pop(model_id, None)
if live_info:
self._live_info[model_id] = live_info
else:
self._live_info.pop(model_id, None)
data_sources: Final = deployment.litellm_params.get("dataSources", []) or []
for source, key, value in data_source_values:
data_sources[source]["parameters"][key] = value
@ -1270,6 +1348,17 @@ class HostRegistry(HostModelInfo):
raise RuntimeError("the host registry writes deployments through its Rust core")
self.changes_runner(changes)
self._core_model_list.refresh()
self._drop_removed_live_values()
def _drop_removed_live_values(self) -> None:
"""Forget the live values of ids no deployment carries any more, so a router that rotates
ids does not keep the removed deployments' errors and callables."""
handle: Final = self._live_handle()
if handle is None or not (self._live_params or self._live_info):
return
for held in (self._live_params, self._live_info):
for model_id in [id for id in held if handle.deployment(id) is None]:
del held[model_id]
def refresh(self) -> None:
"""Read router-core's list again after a write (only what was appended, when that is all
@ -1447,6 +1536,14 @@ class HostRegistry(HostModelInfo):
idx: Final[object] = None if handle is None or key is None else handle.id_position(key)
return None if idx is None else self._core_model_list[idx] # pyright: ignore[reportArgumentType, reportCallIssue] # a position, or what a caller wrote, read as Python reads it
def live_params(self, model_id: object) -> Mapping[str, object]:
"""The litellm_params values of deployment ``model_id`` that its JSON record cannot carry."""
return self._live_params.get(model_id, {}) if isinstance(model_id, str) else {}
def live_info(self, model_id: object) -> Mapping[str, object]:
"""The model_info values a deployment's JSON record reads back changed, by id."""
return self._live_info.get(model_id, {}) if isinstance(model_id, str) else {}
def _core_get_deployment(self, model_id: str) -> Deployment | None:
entry: Final = self._core_get_model_info(model_id)
return None if entry is None else Deployment(**entry) # pyright: ignore[reportArgumentType] # a stored entry
@ -1541,8 +1638,14 @@ class HostRegistry(HostModelInfo):
for name in (self._core_routing_groups or {})
if (group := self.get_routing_group(name)) is not None
]
# The provider is inferred only for a name no deployment or alias serves: inferring it for an
# authenticating provider's name would start its sign-in (`_get_wildcard_deployments`).
wildcard: Final = (
[model_name, f"{inferred_provider(model_name)}/{model_name}"] if model_name is not None else []
[]
if model_name is None
else [model_name]
if model_name in alias_map or model_name in self._core_model_names
else [model_name, f"{inferred_provider(model_name)}/{model_name}"]
)
return {"aliases": aliases, "groups": groups, "wildcard": wildcard, "only": only}
@ -1561,7 +1664,8 @@ class HostRegistry(HostModelInfo):
info: Final = _model_info_of(model)
if team_id is not None and info.get("team_id") == team_id and model_name == info.get("team_public_model_name"):
return True
if model_name is not None and model.get("model_name") == model_name: # pyright: ignore[reportUnnecessaryComparison] # the core checks it
# A record with no model_name raises KeyError, as the Python core's lookup does.
if model_name is not None and model["model_name"] == model_name: # pyright: ignore[reportUnnecessaryComparison] # the core checks it
model_team_id: Final = info.get("team_id")
return team_id is None or model_team_id is None or model_team_id == team_id
return False

252
litellm/router_utils/search_api_router.py
View file

@ -1,252 +0,0 @@
"""
Router utilities for Search API integration.
Handles search tool selection, load balancing, and fallback logic for search requests.
"""
import asyncio
import random
import traceback
from collections.abc import Callable
from functools import partial
from types import MappingProxyType
from typing import TYPE_CHECKING, Final
from litellm._logging import verbose_router_logger
from litellm.router_api import ChangeSet
from litellm.router_deprecated import renamed_router_param
from litellm.types.router import SearchToolLiteLLMParams, SearchToolTypedDict
if TYPE_CHECKING:
from litellm.router_api import RoutingContext
class SearchAPIRouter:
"""
Static utility class for routing search API calls through the LiteLLM router.
Provides methods for search tool selection, load balancing, and fallback handling.
"""
@staticmethod
def _resolve_search_provider_credentials(
*,
tool_litellm_params: SearchToolLiteLLMParams,
) -> tuple[str | None, str | None]:
"""
Resolve search provider credentials from tool configuration ONLY.
Credentials are stored only in search_tool.litellm_params, never in team/key metadata.
This ensures secrets are not exposed in team/key API responses.
Args:
tool_litellm_params: Search tool litellm_params with credentials
Returns:
Tuple of (api_key, api_base) from tool configuration
"""
resolved_api_key: Final[str | None] = tool_litellm_params.get("api_key")
resolved_api_base: Final[str | None] = tool_litellm_params.get("api_base")
return resolved_api_key, resolved_api_base
@staticmethod
@renamed_router_param("router_instance")
async def update_router_search_tools(ctx: "RoutingContext", search_tools: list) -> None:
"""
Update the router with search tools from the database.
This method is called by a cron job to sync search tools from DB to router.
Args:
ctx: The routing context of the router to update
search_tools: List of search tool configurations from the database
"""
try:
verbose_router_logger.debug("Adding %s search tools to router", len(search_tools))
# Convert search tools to the format expected by the router
router_search_tools: Final[list] = []
for tool in search_tools:
# Create dict that matches SearchToolTypedDict structure
router_search_tool: SearchToolTypedDict = {
"search_tool_id": tool.get("search_tool_id"),
"search_tool_name": tool.get("search_tool_name"),
"litellm_params": tool.get("litellm_params", {}),
"search_tool_info": tool.get("search_tool_info"),
}
router_search_tools.append(router_search_tool)
# Update the router's search_tools list
ctx.apply(ChangeSet(assign={"search_tools": router_search_tools}))
verbose_router_logger.info("Successfully updated router with %s search tool(s)", len(router_search_tools))
except Exception as e:
verbose_router_logger.exception("Error updating router with search tools: %s", e)
raise e
@staticmethod
@renamed_router_param("router_instance")
def get_matching_search_tools(
ctx: "RoutingContext",
search_tool_name: str,
) -> list:
"""
Get all search tools matching the given name.
Args:
ctx: The routing context of the router
search_tool_name: Name of the search tool to find
Returns:
List of matching search tool configurations
Raises:
ValueError: If no matching search tools are found
"""
matching_tools: Final = [
tool for tool in ctx.settings().search_tools if tool.get("search_tool_name") == search_tool_name
]
if not matching_tools:
raise ValueError(f"Search tool '{search_tool_name}' not found in router.search_tools")
return matching_tools
@staticmethod
@renamed_router_param("router_instance", new="litellm_routing_ctx")
async def async_search_with_fallbacks(
litellm_routing_ctx: "RoutingContext",
original_function: Callable,
**kwargs,
):
"""
Helper function to make a search API call through the router with load balancing and fallbacks.
Reuses the router's retry/fallback infrastructure.
Args:
litellm_routing_ctx: The routing context of the router
original_function: The original litellm.asearch function
**kwargs: Search parameters including search_tool_name, query, etc.
Returns:
SearchResponse from the search API
"""
ctx: Final = litellm_routing_ctx
try:
search_tool_name: Final = kwargs.get("search_tool_name", kwargs.get("model"))
if not search_tool_name:
raise ValueError("search_tool_name or model parameter is required for search")
# Set up kwargs for the fallback system
kwargs["model"] = search_tool_name # Use model field for compatibility with fallback system
kwargs["original_generic_function"] = original_function
kwargs["original_function"] = partial(
SearchAPIRouter.async_search_with_fallbacks_helper,
litellm_routing_ctx=ctx,
)
# Update kwargs before fallbacks (for logging, metadata, etc)
ctx.update_kwargs_before_fallbacks(
model=search_tool_name,
kwargs=kwargs,
metadata_variable_name="litellm_metadata",
)
available_search_tool_names: Final = [tool.get("search_tool_name") for tool in ctx.settings().search_tools]
verbose_router_logger.debug(
"Inside SearchAPIRouter.async_search_with_fallbacks() - search_tool_name: %s, Available Search Tools: %s",
search_tool_name,
available_search_tool_names,
)
# Use the existing retry/fallback infrastructure
response: Final = await ctx.with_fallbacks(**kwargs)
return response
except Exception as e:
from litellm.router_utils.handle_error import send_llm_exception_alert
asyncio.create_task(
send_llm_exception_alert(
ctx=ctx,
request_kwargs=kwargs,
error_traceback_str=traceback.format_exc(),
original_exception=e,
)
)
raise e
@staticmethod
@renamed_router_param("router_instance", new="litellm_routing_ctx")
async def async_search_with_fallbacks_helper(
litellm_routing_ctx: "RoutingContext",
model: str,
original_generic_function: Callable,
**kwargs,
):
"""
Helper function for search API calls - selects a search tool and calls the original function.
Called by async_function_with_fallbacks for each retry attempt.
Args:
litellm_routing_ctx: The routing context of the router
model: The search tool name (passed as model for compatibility)
original_generic_function: The original litellm.asearch function
**kwargs: Search parameters
Returns:
SearchResponse from the selected search provider
"""
ctx: Final = litellm_routing_ctx
search_tool_name: Final = model # model field contains the search_tool_name
try:
# Find matching search tools
matching_tools: Final = SearchAPIRouter.get_matching_search_tools(
ctx=ctx,
search_tool_name=search_tool_name,
)
# Simple random selection for load balancing across multiple providers with same name
# For search tools, we use simple random choice since they don't have TPM/RPM constraints
selected_tool: Final = random.choice(matching_tools)
# Extract search provider and other params from litellm_params
litellm_params: Final = selected_tool.get("litellm_params", {})
search_provider: Final = litellm_params.get("search_provider")
if not search_provider:
raise ValueError(f"search_provider not found in litellm_params for search tool '{search_tool_name}'")
api_key, api_base = SearchAPIRouter._resolve_search_provider_credentials(
tool_litellm_params=litellm_params,
)
protected_params: Final = frozenset(("search_provider", "api_key", "api_base"))
search_params: Final = MappingProxyType(
{
key: value
for params in (litellm_params, kwargs)
for key, value in params.items()
if key not in protected_params and value is not None
}
)
verbose_router_logger.debug("Selected search tool with provider: %s", search_provider)
# Call the original search function with the provider config
response: Final = await original_generic_function(
search_provider=search_provider,
api_key=api_key,
api_base=api_base,
**search_params,
)
return response
except Exception as e:
verbose_router_logger.error(
"Error in SearchAPIRouter.async_search_with_fallbacks_helper for %s: %s", search_tool_name, e
)
raise e

2
litellm/rust_bridge/catalog.py
View file

@ -23,7 +23,6 @@ class Route(str, Enum):
OCR = "ocr"
TOKEN_COUNTER = "token_counter"
TOKENIZER = "tokenizer"
ROUTER = "router"
class Target(str, Enum):
@ -225,7 +224,6 @@ RULES: Final[Rules] = (
RouteRule(Route.RESPONSES, Rollout.PYTHON_ONLY),
RouteRule(Route.TOKEN_COUNTER, Rollout.PYTHON_ONLY),
RouteRule(Route.TOKENIZER, Rollout.PYTHON_ONLY),
RouteRule(Route.ROUTER, Rollout.RUST_OPT_IN),
RouteRule(Route.TRANSCRIPTION, Rollout.RUST_REQUIRED, providers=frozenset({"bedrock"})),
)

1
tests/code_coverage_tests/router_code_coverage.py
View file

@ -76,6 +76,7 @@ ignored_function_names = [
"_acreate_batch",
"_acreate_file",
"__init__",
"_router_init", # Never run: it carries Router.__init__'s signature and docstring
"avector_store_create", # Tested via proxy vector_store_endpoints (files lack "router" in name)
"_override_vector_store_methods_for_router", # No-op placeholder, called during Router init
"_merge_tools_from_deployment", # Tested indirectly via _update_kwargs_with_deployment (test files lack "router" in name)

1
tests/code_coverage_tests/router_litellm_globals_check.py
View file

@ -20,7 +20,6 @@ CHECKED = (
"litellm/router_api.py",
"litellm/router_core.py",
"litellm/router_deprecated.py",
"litellm/router_python_core.py",
"litellm/router_rust/core.py",
)

34
tests/code_coverage_tests/router_surface_check.py
View file

@ -20,7 +20,7 @@ Member uses are found the same way the boundary survey found them: attribute
access on a receiver whose name looks like a router (llm_router,
litellm_router_instance, router, ...) or on the Router class, plus
getattr/setattr/hasattr with a constant name on such a receiver. Members
include the container file endpoints Router.__init__ registers by name.
include the endpoints the Rust core serves by name, container file endpoints included.
Run: python tests/code_coverage_tests/router_surface_check.py [--write]
--write regenerates the allowlist from the current tree (for shrinking it).
@ -41,7 +41,6 @@ ROUTER_FILES = {
"litellm/router_api.py",
"litellm/router_config.py",
"litellm/router_core.py",
"litellm/router_python_core.py",
"litellm/router_python_host.py",
"litellm/router_rust/core.py",
"litellm/router_deprecated.py",
@ -49,8 +48,15 @@ ROUTER_FILES = {
}
# The Rust core's host-side glue modules are part of the router.
ROUTER_FILE_PREFIXES = ("litellm/router_rust/",)
ROUTER_CLASS_FILES = ["litellm/router.py", "litellm/router_python_core.py", "litellm/router_deprecated.py"]
ROUTER_CLASS_NAMES = {"Router", "PythonCore", "DeprecatedRouterMembers"}
ROUTER_CLASS_FILES = ["litellm/router.py", "litellm/router_rust/core.py", "litellm/router_deprecated.py"]
ROUTER_CLASS_NAMES = {"Router", "RustCore", "DeprecatedRouterMembers"}
# The Rust core serves its endpoint methods by name (``RustCore.__getattr__``); these name lists hold them.
ROUTER_ENDPOINT_LISTS = {
"litellm/router_rust/surface/legacy_members.py": {"_ROUTER_ENDPOINTS", "_ASSISTANTS_ENDPOINTS"},
"litellm/router_rust/invoke/endpoints.py": {"_DIRECT_ENDPOINTS"},
}
# ``RustEndpoints.member`` serves these methods by name, from its ``direct`` table.
ROUTER_ENDPOINT_TABLE = "litellm/router_rust/invoke/endpoints.py"
SCAN_ROOTS = ["litellm", "enterprise"]
SKIP_DIRS = {
@ -70,7 +76,7 @@ RECEIVER = re.compile(
r"(^|[._])(llm_router|router|litellm_router|litellm_router_instance|router_instance"
r"|_router|proxy_router|llm_router_instance|router_obj|mock_router|Router)$"
)
# Router.__init__ registers these with setattr(self, name, ...), so they never appear as self.<name>.
# The Rust core serves these by name too, so they never appear as self.<name>.
CONTAINER_FILE_ENDPOINTS = "litellm/containers/endpoints.json"
DYNAMIC_FUNCS = {"getattr", "setattr", "hasattr"}
@ -109,6 +115,24 @@ def router_members() -> Set[str]:
and isinstance(node.ctx, ast.Store)
):
members.add(node.attr)
for rel, names in ROUTER_ENDPOINT_LISTS.items():
with open(os.path.join(ROOT, rel), encoding="utf-8") as f:
tree = ast.parse(f.read())
for node in tree.body:
target = node.target if isinstance(node, ast.AnnAssign) else None
if isinstance(target, ast.Name) and target.id in names and node.value is not None:
members.update(
c.value for c in ast.walk(node.value) if isinstance(c, ast.Constant) and isinstance(c.value, str)
)
with open(os.path.join(ROOT, ROUTER_ENDPOINT_TABLE), encoding="utf-8") as f:
tree = ast.parse(f.read())
for node in ast.walk(tree):
if isinstance(node, ast.AnnAssign) and isinstance(node.target, ast.Name) and node.target.id == "direct":
for table in ast.walk(node.value):
if isinstance(table, ast.Dict):
members.update(
k.value for k in table.keys if isinstance(k, ast.Constant) and isinstance(k.value, str)
)
with open(os.path.join(ROOT, CONTAINER_FILE_ENDPOINTS), encoding="utf-8") as f:
for endpoint in json.load(f)["endpoints"]:
members.update((endpoint["name"], endpoint["async_name"]))

22
tests/unit/caching/test_redis_cluster_cache.py
View file

@ -155,20 +155,15 @@ def test_cache_init_creates_redis_cache_without_cluster_config(
)
def test_router_create_redis_cache_cluster_detection(startup_nodes, env_var, expected_cache_type, monkeypatch):
"""
Test that Router._create_redis_cache() creates RedisClusterCache when
either startup_nodes is in config or REDIS_CLUSTER_NODES env var is set.
Test that a Router given Redis settings builds a RedisClusterCache when
either startup_nodes is in cache_kwargs or REDIS_CLUSTER_NODES env var is set.
Also verifies that explicit startup_nodes take precedence over env var.
Regression test for https://github.com/BerriAI/litellm/issues/22748
"""
from litellm import Router
cache_config = dict(
host="mockhost",
port=6379,
password="mock-password",
startup_nodes=startup_nodes,
)
cache_kwargs = {"startup_nodes": startup_nodes} if startup_nodes is not None else None
if env_var is not None:
monkeypatch.setenv("REDIS_CLUSTER_NODES", env_var)
@ -178,8 +173,15 @@ def test_router_create_redis_cache_cluster_detection(startup_nodes, env_var, exp
def _mock_redis_cache_init(*args, **kwargs): ...
with patch.object(RedisCache, "__init__", _mock_redis_cache_init):
redis_cache = Router._create_redis_cache(cache_config)
assert isinstance(redis_cache, expected_cache_type)
router = Router(
model_list=[{"model_name": "m", "litellm_params": {"model": "gpt-4o", "api_key": "x"}}],
redis_host="mockhost",
redis_port=6379,
redis_password="mock-password",
cache_kwargs=cache_kwargs,
)
redis_cache = router.health().store.redis_cache
assert type(redis_cache) is expected_cache_type
def _isolated_redis_cache(host: str) -> RedisCache:

51
tests/unit/conftest.py
View file

@ -44,20 +44,33 @@ os.environ["PYTHON_DOTENV_DISABLED"] = "1"
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
import litellm # noqa: E402 # litellm reads LITELLM_LOCAL_MODEL_COST_MAP at import
import litellm.router as litellm_router_module # noqa: E402 # same import-time dependency
import litellm.utils as litellm_utils_module # noqa: E402 # same import-time dependency
from litellm._logging import ALL_LOGGERS # noqa: E402 # same import-time dependency
from litellm.anthropic_beta_headers_manager import (
reload_beta_headers_config, # noqa: E402 # same import-time dependency
)
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER # noqa: E402 # same import-time dependency
from litellm.anthropic_beta_headers_manager import reload_beta_headers_config # noqa: E402 # same import-time dependency
from litellm.litellm_core_utils.prompt_templates import factory as prompt_factory_module # noqa: E402 # same import-time dependency
from litellm.litellm_core_utils.prompt_templates import (
factory as prompt_factory_module, # noqa: E402 # same import-time dependency
)
from litellm.litellm_core_utils.prompt_templates import ( # noqa: E402 # same import-time dependency
image_handling as image_handling_module,
)
from litellm.llms.gemini.chat import transformation as gemini_chat_transformation_module # noqa: E402 # same import-time dependency
from litellm.llms.custom_httpx.async_client_cleanup import ( # noqa: E402 # same import-time dependency
close_litellm_async_clients,
)
from litellm.proxy.db import tool_registry_writer as tool_registry_writer_module # noqa: E402 # same import-time dependency
from litellm.llms.gemini.chat import (
transformation as gemini_chat_transformation_module, # noqa: E402 # same import-time dependency
)
from litellm.proxy.db import (
tool_registry_writer as tool_registry_writer_module, # noqa: E402 # same import-time dependency
)
from litellm.router_rust.surface.replay import ( # noqa: E402 # same import-time dependency
_live_registries as live_router_registries, # pyright: ignore[reportPrivateUsage] # the routers the model-cost replay reaches
)
from litellm.router_utils.global_callback_holders import ( # noqa: E402 # same import-time dependency
routers_holding_global_callbacks,
)
LOOPBACK_HOSTS: Final = ["127.0.0.1", "::1", "localhost"]
AMBIENT_AZURE_CREDENTIAL_ENV_VARS: Final = (
@ -252,12 +265,6 @@ def socket_tempdir(socket_dir: Path, monkeypatch: pytest.MonkeyPatch) -> Path:
return socket_dir
@pytest.fixture
def python_router_core(monkeypatch: pytest.MonkeyPatch) -> None:
"""Keep ``Router`` on ``PythonCore``, for a test that patches a member the Rust core never reads."""
monkeypatch.setenv("LITELLM_RUST", "0")
@pytest.fixture(autouse=True, scope="session")
def bundled_tiktoken_cache() -> None:
importlib.import_module("litellm.litellm_core_utils.default_encoding")
@ -293,10 +300,12 @@ def isolate_host_environment(isolated_aws_config_files: tuple[Path, Path]) -> It
@pytest.fixture(autouse=True)
def isolate_litellm_globals() -> Iterator[None]:
original_callbacks: Final = {name: list(getattr(litellm, name) or []) for name in CALLBACK_LISTS}
original_callback_holders: Final = frozenset(litellm_router_module._routers_holding_global_callbacks)
original_callback_holders: Final = frozenset(routers_holding_global_callbacks)
original_reset: Final = {name: getattr(litellm, name) for name in RESET_TO_NONE_GLOBALS}
original_restored: Final = {name: getattr(litellm, name) for name in RESTORED_GLOBALS if hasattr(litellm, name)}
original_clients: Final = {name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__}
original_clients: Final = {
name: litellm.__dict__[name] for name in MODULE_LEVEL_CLIENTS if name in litellm.__dict__
}
original_loggers: Final = {
logger: (logger.level, logger.disabled, logger.propagate, list(logger.handlers), list(logger.filters))
for logger in ALL_LOGGERS
@ -305,7 +314,7 @@ def isolate_litellm_globals() -> Iterator[None]:
_flush_client_caches()
for name in CALLBACK_LISTS:
setattr(litellm, name, [])
litellm_router_module._routers_holding_global_callbacks.clear()
routers_holding_global_callbacks.clear()
for name in RESET_TO_NONE_GLOBALS:
setattr(litellm, name, None)
for name in MODULE_LEVEL_CLIENTS:
@ -320,8 +329,8 @@ def isolate_litellm_globals() -> Iterator[None]:
litellm.__dict__.update(original_clients)
for name, value in (original_callbacks | original_reset | original_restored).items():
setattr(litellm, name, value)
litellm_router_module._routers_holding_global_callbacks.clear()
litellm_router_module._routers_holding_global_callbacks.update(original_callback_holders)
routers_holding_global_callbacks.clear()
routers_holding_global_callbacks.update(original_callback_holders)
for logger, (level, disabled, propagate, handlers, filters) in original_loggers.items():
logger.setLevel(level)
logger.disabled = disabled
@ -333,19 +342,17 @@ def isolate_litellm_globals() -> Iterator[None]:
@pytest.fixture(autouse=True)
def isolate_router_model_cost_state() -> Iterator[None]:
original_live_routers: Final = frozenset(litellm_router_module._live_routers)
original_live_registries: Final = frozenset(live_router_registries)
original_runtime_registered_model_cost: Final = {
model_key: dict(model_value)
for model_key, model_value in litellm_utils_module._runtime_registered_model_cost.items()
}
litellm_utils_module._invalidate_model_cost_lowercase_map()
litellm_router_module._live_routers.clear()
live_router_registries.clear()
litellm_utils_module._runtime_registered_model_cost.clear()
yield
for router in tuple(litellm_router_module._live_routers):
litellm_router_module._live_routers.discard(router)
for router in original_live_routers:
litellm_router_module._live_routers.add(router)
live_router_registries.clear()
live_router_registries.update(original_live_registries)
litellm_utils_module._runtime_registered_model_cost.clear()
litellm_utils_module._runtime_registered_model_cost.update(original_runtime_registered_model_cost)
litellm_utils_module._invalidate_model_cost_lowercase_map()

54
tests/unit/containers/test_azure_container_transformation.py
View file

@ -558,13 +558,14 @@ class TestAzureContainerKnownFailureRegressions:
self, monkeypatch
):
"""Operations without container_id (create, list) must NOT route via
_ageneric_api_call_with_fallbacks using a caller-supplied model_id.
a deployment using a caller-supplied model_id.
Security boundary: only the path that holds a validated container_id
is trusted to fall back to the forwarded model_id. A caller setting
model_id without container_id on POST /v1/containers must not gain
access to an arbitrary deployment UUID.
"""
import litellm.containers
from litellm.router import Router
router = Router(
@ -582,32 +583,27 @@ class TestAzureContainerKnownFailureRegressions:
]
)
fallback_called = {"called": False}
called_with: dict = {}
async def _mock_fallback(original_function, **kwargs):
fallback_called["called"] = True
async def _capture(**kwargs):
called_with.update(kwargs)
return {}
monkeypatch.setattr(router, "_ageneric_api_call_with_fallbacks", _mock_fallback)
monkeypatch.setattr(litellm.containers, "acreate_container", _capture)
original_called = {"called": False}
async def _noop(**kwargs):
original_called["called"] = True
return {}
# No container_id — simulates create/list; caller injects a model_id
await router._init_containers_api_endpoints(
original_function=_noop,
# No container_id — simulates create; caller injects a model_id
await router.acreate_container(
name="c",
model_id="deployment-uuid-123",
custom_llm_provider="azure",
)
assert not fallback_called["called"], (
"_ageneric_api_call_with_fallbacks must NOT be called when "
assert called_with, "the container call must still run directly"
assert called_with.get("api_key") is None, (
"the deployment's credentials must NOT be applied when "
"container_id is absent, even if model_id is supplied"
)
assert original_called["called"], "original_function must be called directly"
assert called_with.get("api_base") is None
def test_regression_httpx_empty_params_strips_query_string(self):
"""httpx erases the URL query-string when params={} (empty dict) is passed.
@ -974,9 +970,10 @@ class TestAzureContainerKnownFailureRegressions:
self, monkeypatch
):
"""Native Azure container IDs (cntr_ + hex, no LiteLLM payload) must
still route through _ageneric_api_call_with_fallbacks using the
model_id forwarded from the proxy ownership check so that deployment
credentials (api_base) are applied."""
still route through the deployment named by the model_id forwarded
from the proxy ownership check so that deployment credentials
(api_base) are applied."""
import litellm.containers
from litellm.router import Router
router = Router(
@ -996,25 +993,22 @@ class TestAzureContainerKnownFailureRegressions:
called_with: dict = {}
async def _mock_fallback(original_function, **kwargs):
async def _capture(**kwargs):
called_with.update(kwargs)
return {}
monkeypatch.setattr(router, "_ageneric_api_call_with_fallbacks", _mock_fallback)
monkeypatch.setattr(litellm.containers, "aretrieve_container", _capture)
native_azure_id = "cntr_6a058b43d24c8190a226cfb1d35405b20115fb7875ff11df"
async def _noop(**kwargs):
return {}
await router._init_containers_api_endpoints(
original_function=_noop,
await router.aretrieve_container(
container_id=native_azure_id,
model_id="deployment-uuid-123",
custom_llm_provider="azure",
)
assert called_with.get("model") == "deployment-uuid-123", (
"_ageneric_api_call_with_fallbacks must be called with the forwarded "
"model_id when the container_id carries no LiteLLM routing payload"
assert called_with.get("api_base") == "https://my-resource.cognitiveservices.azure.com", (
"the forwarded model_id's deployment credentials must be applied "
"when the container_id carries no LiteLLM routing payload"
)
assert called_with.get("container_id") == native_azure_id

6
tests/unit/litellm_core_utils/test_get_llm_provider_logic.py
View file

@ -796,7 +796,6 @@ def test_metadata_callers_never_log_in_for_a_declared_chatgpt_deployment(
from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base
from litellm.proxy.spend_tracking.savings import _resolve_model
from litellm.router import Router
from litellm.types.router import LiteLLM_Params
from litellm.utils import supports_native_streaming
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
@ -805,4 +804,7 @@ def test_metadata_callers_never_log_in_for_a_declared_chatgpt_deployment(
assert get_api_base("chatgpt/gpt-5.4", {"model": "chatgpt/gpt-5.4"}) == "https://chatgpt.com/backend-api/codex"
identity = _resolve_model("chatgpt/gpt-5.4", None)
assert identity is not None and identity.provider == "chatgpt"
assert Router._model_group_llm_provider(LiteLLM_Params(model="chatgpt/gpt-5.4")) == ("gpt-5.4", "chatgpt")
group = Router(model_list=[{"model_name": "codex", "litellm_params": {"model": "chatgpt/gpt-5.4"}}]).registry().group(
"codex"
)
assert group is not None and group.providers == ["chatgpt"]

50
tests/unit/llms/anthropic/messages/test_advisor_orchestration.py
View file

@ -1053,32 +1053,30 @@ async def test_executor_failure_is_not_tagged():
def _router_with_advisor_deployment(
recorder, advisor_model="claude-opus-4-8", deployment_model=None, model_group_alias=None
advisor_model="claude-opus-4-8", deployment_model=None, model_group_alias=None
):
"""Build a Router whose only deployment is the advisor model on Foundry.
The recorder replaces ``litellm.anthropic_messages`` before construction
because Router binds it at init time, so the returned Router exercises the
real deployment-resolution path and records what it dispatched.
The Router resolves ``litellm.anthropic_messages`` per call, so each test
patches it with a recorder around the handler call; the Router then
exercises the real deployment-resolution path and records what it dispatched.
"""
import litellm
from litellm.router import Router
with patch("litellm.anthropic_messages", new=recorder):
return Router(
model_list=[
{
"model_name": advisor_model,
"litellm_params": {
"model": deployment_model or f"azure_ai/{advisor_model}",
"api_base": "http://127.0.0.1:1/foundry",
"api_key": "fake-foundry-key",
},
}
],
model_group_alias=model_group_alias,
num_retries=0,
)
return Router(
model_list=[
{
"model_name": advisor_model,
"litellm_params": {
"model": deployment_model or f"azure_ai/{advisor_model}",
"api_base": "http://127.0.0.1:1/foundry",
"api_key": "fake-foundry-key",
},
}
],
model_group_alias=model_group_alias,
num_retries=0,
)
@pytest.mark.asyncio
@ -1094,7 +1092,7 @@ async def test_advisor_sub_call_routes_through_proxy_router():
router_calls.append(kwargs)
return _make_text_response("Use trial division.", model="claude-opus-4-8")
router = _router_with_advisor_deployment(recorder)
router = _router_with_advisor_deployment()
call_count = 0
@ -1111,6 +1109,7 @@ async def test_advisor_sub_call_routes_through_proxy_router():
side_effect=mock_call,
),
patch.object(proxy_server, "llm_router", router),
patch("litellm.anthropic_messages", new=recorder),
):
h = AdvisorOrchestrationHandler()
result = await h.handle(
@ -1155,7 +1154,7 @@ async def test_advisor_sub_call_routes_through_router_for_alias_and_wildcard(rou
router_calls.append(kwargs)
return _make_text_response("Use trial division.", model="claude-opus-4-8")
router = _router_with_advisor_deployment(recorder, **router_kwargs)
router = _router_with_advisor_deployment(**router_kwargs)
call_count = 0
@ -1172,6 +1171,7 @@ async def test_advisor_sub_call_routes_through_router_for_alias_and_wildcard(rou
side_effect=mock_call,
),
patch.object(proxy_server, "llm_router", router),
patch("litellm.anthropic_messages", new=recorder),
):
h = AdvisorOrchestrationHandler()
await h.handle(
@ -1204,7 +1204,7 @@ async def test_advisor_sub_call_bypasses_router_for_unconfigured_model():
router_calls.append(kwargs)
return _make_text_response("should not be used")
router = _router_with_advisor_deployment(recorder, advisor_model="some-other-model")
router = _router_with_advisor_deployment(advisor_model="some-other-model")
call_count = 0
@ -1223,6 +1223,7 @@ async def test_advisor_sub_call_bypasses_router_for_unconfigured_model():
side_effect=mock_call,
),
patch.object(proxy_server, "llm_router", router),
patch("litellm.anthropic_messages", new=recorder),
):
h = AdvisorOrchestrationHandler()
await h.handle(
@ -1253,7 +1254,7 @@ async def test_advisor_sub_call_client_override_bypasses_router():
router_calls.append(kwargs)
return _make_text_response("should not be used")
router = _router_with_advisor_deployment(recorder)
router = _router_with_advisor_deployment()
sub_calls = []
@ -1278,6 +1279,7 @@ async def test_advisor_sub_call_client_override_bypasses_router():
side_effect=mock_call,
),
patch.object(proxy_server, "llm_router", router),
patch("litellm.anthropic_messages", new=recorder),
patch.dict(proxy_server.general_settings, {"allow_client_side_credentials": True}),
patch.object(litellm, "user_url_validation", False),
):

717
tests/unit/llms/base_llm/responses/test_transformation.py
View file

@ -27,8 +27,8 @@ from litellm.types.llms.openai import(
ResponsesAPIStreamEvents,
)
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
from typing import Any, AsyncIterator, Final, List
from unittest.mock import AsyncMock, MagicMock, Mock, patch
from typing import Any, Final, List
from unittest.mock import Mock, patch
@pytest.mark.asyncio
@ -751,586 +751,6 @@ class TestBaseResponsesAPIStreamingIterator:
mock_logging_obj.async_failure_handler.assert_not_called()
mock_logging_obj.failure_handler.assert_not_called()
@pytest.fixture()
def _vcr_outcome_gate_router_unit(request, vcr):
install_live_call_probe(request, vcr)
yield
record_vcr_outcome(request, vcr)
@pytest.fixture(scope="function")
def setup_and_teardown_router_unit():
"""
This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained.
"""
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue())
importlib.reload(litellm)
try:
if hasattr(litellm, "proxy") and hasattr(litellm.proxy, "proxy_server"):
importlib.reload(litellm.proxy.proxy_server)
except Exception:
pass
loop = asyncio.get_event_loop_policy().new_event_loop()
asyncio.set_event_loop(loop)
yield
loop.close()
asyncio.set_event_loop(None)
def _make_router() -> Router:
return Router(
model_list=[
{
"model_name": "primary",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "sk-test",
},
},
{
"model_name": "fallback",
"litellm_params": {
"model": "openai/gpt-4o",
"api_key": "sk-test",
},
},
]
)
def _make_completed_event(input_tokens: int, output_tokens: int, total_tokens: int) -> ResponseCompletedEvent:
response = ResponsesAPIResponse.model_construct(
usage=ResponseAPIUsage(
input_tokens=input_tokens,
output_tokens=output_tokens,
total_tokens=total_tokens,
)
)
return ResponseCompletedEvent.model_construct(
type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED,
response=response,
)
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
def test_extract_partial_responses_usage_native_completed():
"""Native path: completed_response carries usage → returned as-is."""
completed = _make_completed_event(11, 7, 18)
source = MagicMock()
source.completed_response = completed
usage = Router._extract_partial_responses_usage(source)
assert usage is not None
assert usage.input_tokens == 11
assert usage.output_tokens == 7
assert usage.total_tokens == 18
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
def test_extract_partial_responses_usage_no_completed_response():
"""Native path: no completed_response → returns None."""
source = MagicMock()
source.completed_response = None
usage = Router._extract_partial_responses_usage(source)
assert usage is None
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
def test_extract_partial_responses_usage_bridge_iterator_no_completed_response():
"""
Regression for #35411: the bridge iterator
(LiteLLMCompletionStreamingIterator) overrides __init__ without calling
super().__init__(), so completed_response was never set until the stream
reached RESPONSE_COMPLETED. On a mid-stream provider error (before
completion) the fallback recovery path read source_iterator.completed_response
and raised AttributeError, masking the real provider error and bypassing
fallbacks. The attribute must always exist and default to None.
"""
from litellm.responses.litellm_completion_transformation.streaming_iterator import (
LiteLLMCompletionStreamingIterator,
)
wrapper = MagicMock()
wrapper.logging_obj = MagicMock()
iterator = LiteLLMCompletionStreamingIterator(
model="anthropic/claude-sonnet-4-5",
litellm_custom_stream_wrapper=wrapper,
request_input="hi",
responses_api_request={},
)
assert iterator.completed_response is None
# No chat chunks collected yet and no completed_response → must return
# None instead of raising AttributeError.
assert Router._extract_partial_responses_usage(iterator) is None
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
def test_combine_responses_fallback_usage_sums_completed_event():
"""Partial-stream usage is summed into the fallback event's usage."""
fallback_event = _make_completed_event(5, 3, 8)
partial = ResponseAPIUsage(input_tokens=11, output_tokens=7, total_tokens=18)
Router._combine_responses_fallback_usage(fallback_event, partial)
combined = fallback_event.response.usage
assert combined is not None
assert combined.input_tokens == 16
assert combined.output_tokens == 10
assert combined.total_tokens == 26
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
def test_combine_responses_fallback_usage_passthrough_for_unknown_event():
"""Events that are not completed/failed/incomplete are not mutated."""
other = MagicMock() # not a ResponseCompletedEvent etc. → isinstance false
partial = ResponseAPIUsage(input_tokens=1, output_tokens=1, total_tokens=2)
Router._combine_responses_fallback_usage(other, partial)
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
@pytest.mark.asyncio
async def test_aresponses_streaming_iterator_passthrough():
"""
Without MidStreamFallbackError, the wrapper yields source events
unchanged and returns a BaseResponsesAPIStreamingIterator subclass.
"""
from litellm.responses.streaming_iterator import (
BaseResponsesAPIStreamingIterator,
)
events = [_make_completed_event(1, 1, 2)]
class _FakeSource:
"""Minimal source iterator. Provides every attribute the wrapper
constructor reads from source_iterator."""
def __init__(self) -> None:
self._i = 0
self.completed_response = None
self.response = MagicMock()
self.model = "openai/gpt-4o-mini"
self.logging_obj = MagicMock()
self.responses_api_provider_config = MagicMock()
self.start_time = 0.0
self.litellm_metadata = {}
self.custom_llm_provider = "openai"
self.request_data = {}
self.call_type = "aresponses"
self._hidden_params: dict = {}
def __aiter__(self) -> AsyncIterator[Any]:
return self
async def __anext__(self):
if self._i >= len(events):
raise StopAsyncIteration
ev = events[self._i]
self._i += 1
return ev
async def aclose(self):
return None
router = _make_router()
source = _FakeSource()
wrapper = await router._aresponses_streaming_iterator(source, initial_kwargs={"model": "primary"})
assert isinstance(wrapper, BaseResponsesAPIStreamingIterator)
collected = [ev async for ev in wrapper]
assert len(collected) == 1
assert collected[0].type == ResponsesAPIStreamEvents.RESPONSE_COMPLETED
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
@pytest.mark.asyncio
async def test_aresponses_with_streaming_fallbacks_non_streaming_passthrough():
"""Non-streaming response is returned unchanged, no wrap."""
router = _make_router()
plain_response = MagicMock()
async def fake_original(**_kwargs):
return plain_response
with patch.object(
router,
"_ageneric_api_call_with_fallbacks_helper",
new=AsyncMock(return_value=plain_response),
):
out = await router._aresponses_with_streaming_fallbacks(
original_function=fake_original,
model="primary",
stream=False,
)
assert out is plain_response
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
@pytest.mark.asyncio
async def test_aresponses_with_streaming_fallbacks_wraps_streaming_iterator():
"""Streaming response is wrapped via _aresponses_streaming_iterator."""
from litellm.responses.streaming_iterator import (
BaseResponsesAPIStreamingIterator,
)
router = _make_router()
streaming_iter = MagicMock(spec=BaseResponsesAPIStreamingIterator)
wrapped = MagicMock(spec=BaseResponsesAPIStreamingIterator)
async def fake_original(**_kwargs):
return streaming_iter
with (
patch.object(
router,
"_ageneric_api_call_with_fallbacks_helper",
new=AsyncMock(return_value=streaming_iter),
),
patch.object(
router,
"_aresponses_streaming_iterator",
new=AsyncMock(return_value=wrapped),
) as mock_wrap,
):
out = await router._aresponses_with_streaming_fallbacks(
original_function=fake_original,
model="primary",
stream=True,
)
assert out is wrapped
mock_wrap.assert_awaited_once()
def _make_three_tier_router(**router_kwargs) -> Router:
return Router(
model_list=[
{"model_name": "primary", "litellm_params": {"model": "openai/primary-model", "api_key": "sk-test"}},
{"model_name": "fb1", "litellm_params": {"model": "openai/fb1-model", "api_key": "sk-test"}},
{"model_name": "fb2", "litellm_params": {"model": "openai/fb2-model", "api_key": "sk-test"}},
],
num_retries=0,
**router_kwargs,
)
def _mid_stream_failure(model: str):
import litellm
from litellm.exceptions import MidStreamFallbackError
return MidStreamFallbackError(
message="stream dropped",
model=model,
llm_provider="openai",
original_exception=litellm.InternalServerError(message="stream dropped", llm_provider="openai", model=model),
is_pre_first_chunk=True,
)
def _scripted_responses_stream(events: list, error: Exception | None = None):
from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator
class _ScriptedStream(BaseResponsesAPIStreamingIterator):
def __init__(self) -> None:
self._events = list(events)
self._hidden_params: dict = {}
self.completed_response = None
def __aiter__(self):
return self
async def __anext__(self):
if self._events:
return self._events.pop(0)
if error is not None:
raise error
raise StopAsyncIteration
async def aclose(self) -> None:
return None
return _ScriptedStream()
def _three_tier_original(calls: list, primary_fails_pre_stream: bool):
import litellm
completed_event = _make_completed_event(1, 1, 2)
async def fake_original(**kwargs):
model = kwargs["model"]
calls.append(model)
if model == "openai/primary-model":
if primary_fails_pre_stream:
raise litellm.InternalServerError(message="primary down", llm_provider="openai", model=model)
return _scripted_responses_stream([], _mid_stream_failure(model))
if model == "openai/fb1-model":
return _scripted_responses_stream([], _mid_stream_failure(model))
return _scripted_responses_stream([completed_event])
return fake_original, completed_event
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
@pytest.mark.asyncio
async def test_aresponses_pre_stream_primary_failure_then_hop_stream_failure_reaches_second_entry():
"""Regression: fallbacks=[{"primary": ["fb1", "fb2"]}]. The primary fails before streaming,
fb1 is reached through the regular fallback chain and then fails mid-stream. Only the
primary's stream used to be wrapped, so fb1's mid-stream failure either re-raised or
re-tried fb1 itself; fb2 was unreachable."""
router = _make_three_tier_router(fallbacks=[{"primary": ["fb1", "fb2"]}])
calls: list = []
fake_original, completed_event = _three_tier_original(calls, primary_fails_pre_stream=True)
stream = await router._aresponses_with_streaming_fallbacks(
original_function=fake_original, model="primary", stream=True, input="hi"
)
collected = [event async for event in stream]
assert calls == ["openai/primary-model", "openai/fb1-model", "openai/fb2-model"]
assert collected == [completed_event]
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
@pytest.mark.asyncio
async def test_aresponses_two_consecutive_mid_stream_failures_reach_second_entry():
"""Regression: the primary and fb1 both fail mid-stream; fb2 must still be tried."""
router = _make_three_tier_router(fallbacks=[{"primary": ["fb1", "fb2"]}])
calls: list = []
fake_original, completed_event = _three_tier_original(calls, primary_fails_pre_stream=False)
stream = await router._aresponses_with_streaming_fallbacks(
original_function=fake_original, model="primary", stream=True, input="hi"
)
collected = [event async for event in stream]
assert calls == ["openai/primary-model", "openai/fb1-model", "openai/fb2-model"]
assert collected == [completed_event]
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
@pytest.mark.asyncio
async def test_aresponses_per_request_fallbacks_survive_into_hop_streams():
"""Regression: a request-level fallbacks list (key or team router_settings) is popped
before each attempt runs, so a hop's mid-stream re-entry used to see only the router's
own (empty) list and gave up after fb1."""
router = _make_three_tier_router()
calls: list = []
fake_original, completed_event = _three_tier_original(calls, primary_fails_pre_stream=False)
stream = await router._aresponses_with_streaming_fallbacks(
original_function=fake_original,
model="primary",
stream=True,
input="hi",
fallbacks=[{"primary": ["fb1", "fb2"]}],
)
collected = [event async for event in stream]
assert calls == ["openai/primary-model", "openai/fb1-model", "openai/fb2-model"]
assert collected == [completed_event]
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
@pytest.mark.asyncio
async def test_aresponses_attempt_strips_the_controls_carrier_and_wraps_every_hop_stream():
"""Each attempt of the chain, not only the primary's, comes back wrapped for mid-stream
failover, and the per-request controls carrier rides into the wrapper's re-entry kwargs
without ever reaching the provider call."""
from types import MappingProxyType
from litellm.router_utils.fallback_event_handlers import (
MID_STREAM_FALLBACK_CONTROLS_KEY,
MidStreamFallbackControls,
)
router = _make_three_tier_router()
completed_event = _make_completed_event(1, 1, 2)
hop_stream = _scripted_responses_stream([completed_event])
seen: dict = {}
async def fake_original(**kwargs):
seen.update(kwargs)
return hop_stream
controls = MidStreamFallbackControls(MappingProxyType({"fallbacks": [{"primary": ["fb1", "fb2"]}]}))
stream = await router._ageneric_api_call_with_fallbacks_responses_attempt(
model="fb1",
original_generic_function=fake_original,
stream=True,
input="hi",
**{MID_STREAM_FALLBACK_CONTROLS_KEY: controls},
)
collected = [event async for event in stream]
assert seen["model"] == "openai/fb1-model"
assert MID_STREAM_FALLBACK_CONTROLS_KEY not in seen
assert "fallbacks" not in seen
assert stream is not hop_stream
assert collected == [completed_event]
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
@pytest.mark.asyncio
async def test_aresponses_fallback_on_in_stream_error_event():
"""A retriable in-stream error event (429) must trigger the router's mid-stream
fallback path: the wrapper catches MidStreamFallbackError raised by the source
iterator and yields the fallback stream instead of surfacing the error."""
import json
from unittest.mock import Mock
import litellm
from litellm.exceptions import MidStreamFallbackError
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
from litellm.types.llms.openai import ErrorEvent, ErrorEventError
router = _make_router()
error_payload = {
"type": "error",
"error": {"type": "tokens", "code": "rate_limit_exceeded", "message": "rate limited"},
}
sse_bytes = f"data: {json.dumps(error_payload)}\n\n".encode()
async def mock_aiter_bytes():
yield sse_bytes
mock_response = Mock()
mock_response.headers = {}
mock_response.aiter_bytes = mock_aiter_bytes
mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
mock_logging_obj.model_call_details = {"litellm_params": {}}
mock_logging_obj.completion_start_time = None
mock_config = Mock(spec=BaseResponsesAPIConfig)
mock_config.transform_streaming_response.return_value = ErrorEvent(
type=ResponsesAPIStreamEvents.ERROR,
sequence_number=0,
error=ErrorEventError(type="tokens", code="rate_limit_exceeded", message="rate limited"),
)
source = ResponsesAPIStreamingIterator(
response=mock_response,
model="gpt-5",
responses_api_provider_config=mock_config,
logging_obj=mock_logging_obj,
custom_llm_provider="openai",
)
fallback_event = _make_completed_event(1, 1, 2)
class _FallbackStream:
def __init__(self) -> None:
self._done = False
def __aiter__(self):
return self
async def __anext__(self):
if self._done:
raise StopAsyncIteration
self._done = True
return fallback_event
with patch.object(
router,
"async_function_with_fallbacks_common_utils",
new=AsyncMock(return_value=_FallbackStream()),
) as mock_fallback:
wrapped = await router._aresponses_streaming_iterator(
response=source,
initial_kwargs={"model": "primary", "input": "original question"},
)
collected = [ev async for ev in wrapped]
assert collected == [fallback_event]
mock_fallback.assert_awaited_once()
raised = mock_fallback.await_args.kwargs["e"]
assert isinstance(raised, MidStreamFallbackError)
assert raised.status_code == 429
assert isinstance(raised.original_exception, litellm.RateLimitError)
assert raised.original_exception.status_code == 429
assert mock_fallback.await_args.kwargs["kwargs"]["input"] == "original question"
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
@pytest.mark.asyncio
async def test_aresponses_fallback_uses_continuation_input_after_partial_content():
"""When output text was already streamed before the error, the fallback re-entry
must carry a continuation input with the partial assistant text instead of
retrying the original input from scratch (which would duplicate streamed content)."""
import json
from unittest.mock import Mock
from litellm.exceptions import MidStreamFallbackError
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig
from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator
from litellm.types.llms.openai import ErrorEvent, ErrorEventError
router = _make_router()
events = [
{"type": "response.output_text.delta", "delta": "partial answer"},
{"type": "error", "error": {"type": "server_error", "code": "internal_error", "message": "boom"}},
]
sse_payload = b"".join(f"data: {json.dumps(event)}\n\n".encode() for event in events)
async def mock_aiter_bytes():
yield sse_payload
mock_response = Mock()
mock_response.headers = {}
mock_response.aiter_bytes = mock_aiter_bytes
mock_logging_obj = MagicMock(spec=LiteLLMLoggingObj)
mock_logging_obj.model_call_details = {"litellm_params": {}}
mock_logging_obj.completion_start_time = None
mock_config = Mock(spec=BaseResponsesAPIConfig)
def transform(model, parsed_chunk, logging_obj):
if parsed_chunk.get("type") == "error":
return ErrorEvent(
type=ResponsesAPIStreamEvents.ERROR,
sequence_number=0,
error=ErrorEventError(**parsed_chunk["error"]),
)
delta_event = Mock()
delta_event.type = ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA
delta_event.delta = parsed_chunk["delta"]
return delta_event
mock_config.transform_streaming_response.side_effect = transform
source = ResponsesAPIStreamingIterator(
response=mock_response,
model="gpt-5",
responses_api_provider_config=mock_config,
logging_obj=mock_logging_obj,
custom_llm_provider="openai",
)
fallback_event = _make_completed_event(1, 1, 2)
class _FallbackStream:
def __init__(self) -> None:
self._done = False
def __aiter__(self):
return self
async def __anext__(self):
if self._done:
raise StopAsyncIteration
self._done = True
return fallback_event
with patch.object(
router,
"async_function_with_fallbacks_common_utils",
new=AsyncMock(return_value=_FallbackStream()),
) as mock_fallback:
wrapped = await router._aresponses_streaming_iterator(
response=source,
initial_kwargs={"model": "primary", "input": "original question"},
)
collected = [ev async for ev in wrapped]
assert collected[-1] == fallback_event
raised = mock_fallback.await_args.kwargs["e"]
assert isinstance(raised, MidStreamFallbackError)
assert raised.is_pre_first_chunk is False
assert raised.generated_content == "partial answer"
continuation = mock_fallback.await_args.kwargs["kwargs"]["input"]
assert isinstance(continuation, list)
assert continuation[0]["content"][0]["text"] == "original question"
assert continuation[-2]["role"] == "assistant"
assert continuation[-2]["content"][0]["text"] == "partial answer"
assert continuation[-1]["role"] == "user"
def _sse(*events: Mapping[str, object]) -> httpx.Response:
return httpx.Response(
200,
@ -1894,42 +1314,109 @@ async def test_bridged_responses_fallback_continues_after_a_server_run_tool_sear
router.discard()
@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit")
def _openai_completed(model: str) -> Mapping[str, object]:
return {
"type": "response.completed", "sequence_number": 1,
"response": {
"id": f"resp_{model}", "object": "response", "created_at": 0, "status": "completed", "model": model,
"output": [], "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
},
}
@pytest.mark.asyncio
async def test_aresponses_client_error_event_skips_fallback():
"""A 400-mapped in-stream error (raised as APIError, not MidStreamFallbackError)
must surface to the caller without invoking the router's fallback path."""
import litellm
@pytest.mark.parametrize(
("primary_fails_pre_stream", "request_level_fallbacks"),
[(True, False), (False, False), (False, True)],
ids=["pre_stream_then_hop_stream_failure", "two_mid_stream_failures", "request_level_fallbacks"],
)
async def test_responses_fallback_chain_reaches_the_second_entry_after_a_hop_stream_fails(
primary_fails_pre_stream: bool,
request_level_fallbacks: bool,
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Regression: fallbacks=[{"primary": ["fb1", "fb2"]}]. fb1 is reached (after the primary failed before or
during streaming) and then fails mid-stream; fb2 must still be tried, including when the fallbacks list
came with the request rather than the router."""
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
calls: Final[list[str]] = []
router = _make_router()
def upstream(request: httpx.Request) -> httpx.Response:
model: Final = json.loads(request.content)["model"]
calls.append(model)
if model == "primary-model" and primary_fails_pre_stream:
return httpx.Response(500, json={"error": {"message": "primary down", "type": "server_error"}})
if model in ("primary-model", "fb1-model"):
return _openai_stream_cut_after("")
return _sse(_openai_completed(model))
class _ClientErrorSource:
completed_response = None
def __aiter__(self):
return self
async def __anext__(self):
raise litellm.APIError(
status_code=400,
message="bad request",
llm_provider="openai",
model="gpt-5",
)
wrapped = await router._aresponses_streaming_iterator(
response=_ClientErrorSource(),
initial_kwargs={"model": "primary"},
respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=upstream)
fallbacks: Final = [{"primary": ["fb1", "fb2"]}]
router: Final = Router(
model_list=[
{"model_name": name, "litellm_params": {"model": f"openai/{name}-model", "api_key": "sk-test"}}
for name in ("primary", "fb1", "fb2")
],
num_retries=0,
**({} if request_level_fallbacks else {"fallbacks": fallbacks}),
)
try:
stream: Final = await router.aresponses(
model="primary", input="hi", stream=True, **({"fallbacks": fallbacks} if request_level_fallbacks else {})
)
events: Final = [event async for event in stream]
finally:
router.discard()
with patch.object(
router,
"async_function_with_fallbacks_common_utils",
new=AsyncMock(),
) as mock_fallback:
with pytest.raises(litellm.APIError) as exc_info:
async for _ in wrapped:
pass
assert calls == ["primary-model", "fb1-model", "fb2-model"]
assert [event.type for event in events] == [ResponsesAPIStreamEvents.RESPONSE_COMPLETED]
assert exc_info.value.status_code == 400
mock_fallback.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
("error", "falls_back"),
[
({"type": "tokens", "code": "rate_limit_exceeded", "message": "rate limited"}, True),
({"type": "invalid_request_error", "code": "invalid_request_error", "message": "bad request"}, False),
],
ids=["retriable_429_falls_back", "client_400_surfaces"],
)
async def test_responses_in_stream_error_event_falls_back_only_when_retriable(
error: Mapping[str, object], falls_back: bool, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A retriable in-stream error event (429) takes the router's mid-stream fallback path; a 400-mapped one
surfaces to the caller without invoking a fallback."""
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
calls: Final[list[str]] = []
def upstream(request: httpx.Request) -> httpx.Response:
model: Final = json.loads(request.content)["model"]
calls.append(model)
if model == "primary-model":
return _sse({"type": "error", "sequence_number": 0, "error": error})
return _sse(_openai_completed(model))
respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=upstream)
router: Final = Router(
model_list=[
{"model_name": name, "litellm_params": {"model": f"openai/{name}-model", "api_key": "sk-test"}}
for name in ("primary", "fallback")
],
fallbacks=[{"primary": ["fallback"]}],
num_retries=0,
)
try:
stream: Final = await router.aresponses(model="primary", input="hi", stream=True)
if falls_back:
events: Final = [event async for event in stream]
assert [event.type for event in events] == [ResponsesAPIStreamEvents.RESPONSE_COMPLETED]
assert calls == ["primary-model", "fallback-model"]
else:
with pytest.raises(litellm.BadRequestError) as exc_info:
async for _ in stream:
pass
assert exc_info.value.status_code == 400
assert calls == ["primary-model"]
finally:
router.discard()

31
tests/unit/llms/github_copilot/test_per_user_auth.py
View file

@ -1372,24 +1372,30 @@ def _per_user_router():
)
def _route_with_credential_name(router: litellm.Router, credential_name: str) -> ModelResponse:
return router.completion(
model="copilot",
messages=[{"role": "user", "content": "hi"}],
litellm_credential_name=credential_name,
mock_response="ok",
)
def test_per_user_deployment_rejects_a_caller_credential_override(per_user_credential):
"""Caller kwargs win the litellm_params merge downstream, so the deployment
layer is the last place that still sees both names: a different name on a
per-user deployment must be refused there even if discovery missed the hop."""
router: Final = _per_user_router()
deployment: Final = router.get_model_list()[0]
kwargs: Final = {"litellm_credential_name": "shared-cred", "metadata": {}}
with pytest.raises(litellm.BadRequestError):
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
with pytest.raises(BadRequestError, match="litellm_credential_name cannot be overridden"):
_route_with_credential_name(router, "shared-cred")
def test_per_user_deployment_allows_the_configured_credential_name(per_user_credential):
"""The configured name passes the deployment layer and reaches the per-user login, which asks the
caller to connect GitHub Copilot since no user token is stored."""
router: Final = _per_user_router()
deployment: Final = router.get_model_list()[0]
kwargs: Final = {"litellm_credential_name": "copilot-cred", "metadata": {}}
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
assert kwargs["litellm_credential_name"] == "copilot-cred"
assert "model_info" in kwargs
with pytest.raises(CallerCredentialAuthenticationError, match="copilot-cred"):
_route_with_credential_name(router, "copilot-cred")
def test_shared_deployment_keeps_caller_credential_override():
@ -1400,8 +1406,5 @@ def test_shared_deployment_keeps_caller_credential_override():
patch.object(Authenticator, "get_api_key", return_value="shared-token"),
):
router: Final = _per_user_router()
deployment: Final = router.get_model_list()[0]
kwargs: Final = {"litellm_credential_name": "other-cred", "metadata": {}}
router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs)
assert kwargs["litellm_credential_name"] == "other-cred"
assert "model_info" in kwargs
response: Final = _route_with_credential_name(router, "other-cred")
assert response.choices[0].message.content == "ok"

6
tests/unit/llms/openai/live/test_handler.py
View file

@ -1254,7 +1254,7 @@ def _upstream_failing_a_client_send() -> _FakeUpstream:
ids=["session-closed", "upstream-1011-after-a-frame", "upstream-1006-after-a-frame", "upstream-send-failed"],
)
async def test_live_session_counts_a_router_success_only_when_it_ended_cleanly(
upstream_factory: Callable[[], _FakeUpstream], counted_successes: int
upstream_factory: Callable[[], _FakeUpstream], counted_successes: int, monkeypatch: pytest.MonkeyPatch
) -> None:
router: Final = Router(
model_list=[
@ -1276,10 +1276,10 @@ async def test_live_session_counts_a_router_success_only_when_it_ended_cleanly(
api_key="live-test-key",
)
monkeypatch.setattr(litellm, "_arealtime", _arealtime)
await asyncio.wait_for(
router._ageneric_api_call_with_fallbacks(
router._arealtime( # the routed realtime endpoint the proxy's websocket route calls
model="live-voice",
original_function=_arealtime,
litellm_logging_obj=Logging(
model="gpt-live-1",
messages=[],

3
tests/unit/proxy/auth/test_fallback_model_access.py
View file

@ -55,10 +55,9 @@ def _broken_access_group_lookup(*args: object, **kwargs: object) -> NoReturn:
@pytest.mark.asyncio
@pytest.mark.usefixtures("python_router_core") # the Rust core keeps this lookup on its host registry
async def test_is_model_authorized_for_token_fails_closed_when_the_lookup_breaks(monkeypatch: pytest.MonkeyPatch):
router = _router()
monkeypatch.setattr(router, "_core_get_model_access_groups", _broken_access_group_lookup)
monkeypatch.setattr(type(router.registry()), "access_groups", _broken_access_group_lookup)
assert (
await is_model_authorized_for_token(

13
tests/unit/proxy/auth/test_model_access_group_budgets.py
View file

@ -222,11 +222,15 @@ async def test_group_without_a_budget_is_not_attributed():
@pytest.mark.asyncio
@pytest.mark.usefixtures("python_router_core") # the Rust core keeps this lookup on its host registry
async def test_empty_registry_skips_the_access_group_matching_entirely(monkeypatch: pytest.MonkeyPatch):
router = Router(model_list=MODEL_LIST)
access_group_lookups = MagicMock(wraps=router._core_get_model_access_groups)
monkeypatch.setattr(router, "_core_get_model_access_groups", access_group_lookups)
registry_type = type(router.registry())
access_group_lookups = MagicMock(wraps=registry_type.access_groups)
monkeypatch.setattr(
registry_type,
"access_groups",
lambda registry, *args, **kwargs: access_group_lookups(registry, *args, **kwargs),
)
assert await _matched(key_models=["tier-a"], cache=await _cache(budgeted_groups=()), llm_router=router) == ()
assert access_group_lookups.call_count == 0
@ -259,12 +263,11 @@ def _broken_access_group_lookup(*args: object, **kwargs: object) -> NoReturn:
@pytest.mark.asyncio
@pytest.mark.usefixtures("python_router_core") # the Rust core keeps this lookup on its host registry
async def test_stamp_does_not_break_auth_when_the_access_group_lookup_fails(monkeypatch: pytest.MonkeyPatch):
cache = await _cache()
valid_token = UserAPIKeyAuth(api_key="hashed", models=["tier-a"], team_id=TEAM_ID, user_id=USER_ID)
router = Router(model_list=MODEL_LIST)
monkeypatch.setattr(router, "_core_get_model_access_groups", _broken_access_group_lookup)
monkeypatch.setattr(type(router.registry()), "access_groups", _broken_access_group_lookup)
await stamp_matched_model_access_groups(
model="gpt-4o",

7
tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py
View file

@ -153,7 +153,7 @@ class TestUnmappedModelBudgetEnforcement:
Regression test for the stale-cache bug surfaced in PR review:
upgrading an explicitly free deployment to paid via ``upsert_deployment``
(same deployment count, same router instance) must invalidate the
cached ``_is_model_cost_zero=True`` answer so budget checks resume
cached zero-cost answer so budget checks resume
immediately — not after the next proxy restart.
"""
from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo
@ -178,7 +178,6 @@ class TestUnmappedModelBudgetEnforcement:
)
# Warm the cache as zero-cost.
assert is_model_cost_zero(model="ramping-model", llm_router=router) is True
assert router._zero_cost_cache.get("ramping-model") is True
# In-place pricing update: same deployment count, same router id,
# same model name. The pre-fix cache key was
@ -200,9 +199,7 @@ class TestUnmappedModelBudgetEnforcement:
)
)
# Cache must have been cleared by ``_invalidate_model_group_info_cache``.
assert router._zero_cost_cache == {}
# Subsequent call sees the new pricing and enforces budget.
# The next call sees the new pricing and enforces budget.
assert is_model_cost_zero(model="ramping-model", llm_router=router) is False
def test_strategy_router_alias_with_zero_pricing_enforces_budget(self):

3
tests/unit/proxy/hooks/test_autorouter_baseline_cache.py
View file

@ -20,6 +20,7 @@ from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.llms.anthropic.prompt_cache_prediction import NativePredictionTarget, TokenCounter
from litellm.proxy.hooks.autorouter_baseline_cache import AutoRouterBaselineCache, CapturedBaselineObservation
from litellm.router import Router
from litellm.router_rust.steps.prerouting import _record_routing_decision # pyright: ignore[reportPrivateUsage] # production trusted stamp owner
from litellm.types.router import RetryPolicy
from litellm.types.utils import CallTypes, StandardLoggingRoutingDecision
@ -86,7 +87,7 @@ def _kwargs(logging_obj: Logging, trusted: bool = True, *, explicit_logging: boo
context: Final = _OBJECTS.validate_json(
'{"max_tokens":16,"litellm_metadata":{"user_api_key_hash":"test-caller-hash"}}'
)
Router._record_routing_decision( # pyright: ignore[reportUnknownMemberType, reportPrivateUsage] # production trusted stamp owner
_record_routing_decision(
context,
StandardLoggingRoutingDecision(
router_model_name="test-router",

22
tests/unit/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -2288,8 +2288,8 @@ class TestTeamModelSiblingRouting:
async def test_router_finds_all_sibling_team_deployments(self):
"""
When two team deployments share team_public_model_name="gpt-4.1-mini",
the router's _common_checks_available_deployment must return BOTH as
healthy_deployments (not collapse to one).
routing that public name for the team must reach BOTH deployments
(not collapse to one).
"""
import litellm
@ -2338,14 +2338,16 @@ class TestTeamModelSiblingRouting:
result = router.map_team_model(public_name, team_id)
assert result == public_name
# _common_checks_available_deployment should return both deployments
model, healthy = router._common_checks_available_deployment(
model=public_name,
request_kwargs={"metadata": {"user_api_key_team_id": team_id}},
)
assert isinstance(healthy, list)
assert len(healthy) == 2
api_bases = {d["litellm_params"]["api_base"] for d in healthy}
# routing the public name for the team must reach both siblings
api_bases = set()
for _ in range(30):
response = await router.acompletion(
model=public_name,
messages=[{"role": "user", "content": "hi"}],
mock_response="ok",
metadata={"user_api_key_team_id": team_id},
)
api_bases.add(response._hidden_params["api_base"])
assert api_bases == {
"https://eastus.openai.azure.com",
"https://westus.openai.azure.com",

11
tests/unit/proxy/management_endpoints/test_ptu_model_settings.py
View file

@ -11,6 +11,8 @@ from unittest.mock import patch as patch_ctx
import pytest
from fastapi import HTTPException
import litellm
from litellm.litellm_core_utils.llm_cost_calc.utils import generic_cost_per_token
from litellm.litellm_core_utils.ptu_pricing import ptu_terms
from litellm.llms.gemini.cost_calculator import cost_per_web_search_request
@ -1047,9 +1049,9 @@ class TestPtuDeploymentsAreNotBilledPerToken:
def test_the_cost_map_contributes_no_price_to_a_priced_ptu_deployment(self, backend):
"""The acceptance criterion, read off the entry the router registers for the deployment.
Zeroing only the per-token pair leaves the cache-tier fields unset, which is exactly what
Router._inherit_builtin_cache_pricing back-fills from the public cost map, so a cached
prompt would still be billed at the public rate."""
Zeroing only the per-token pair leaves the cache-tier fields unset, which the router
back-fills from the public cost map, so a cached prompt would still be billed at the
public rate."""
priced = _ptu_priced_deployment(
Deployment(
model_name="ptu-deployment",
@ -1062,7 +1064,8 @@ class TestPtuDeploymentsAreNotBilledPerToken:
),
)
)
registered = Router._deployment_model_cost_payload(priced)
Router(model_list=[priced.to_json(exclude_none=True)])
registered = litellm.model_cost["dep-ptu"]
charged = {
k: v
for k, v in registered.items()

15
tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py
View file

@ -79,10 +79,7 @@ def test_initialize_deployment_for_pass_through_success(reusable_credentials):
)
# Test the initialization
router._initialize_deployment_for_pass_through(
deployment=deployment,
custom_llm_provider="vertex_ai",
)
router.add_deployment(deployment)
# Verify the credentials were properly set
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
@ -114,10 +111,7 @@ def test_initialize_deployment_for_pass_through_missing_params():
ValueError,
match="vertex_project, and vertex_location must be set in litellm_params for pass-through endpoints",
):
router._initialize_deployment_for_pass_through(
deployment=deployment,
custom_llm_provider="vertex_ai",
)
router.add_deployment(deployment)
def test_initialize_deployment_when_pass_through_disabled():
@ -133,10 +127,7 @@ def test_initialize_deployment_when_pass_through_disabled():
)
# This should exit without error, even with missing vertex parameters
router._initialize_deployment_for_pass_through(
deployment=deployment,
custom_llm_provider="vertex_ai",
)
router.add_deployment(deployment)
# If we reach this point, the test passes as the method exited without raising any errors
assert True

15
tests/unit/proxy/test_budget_reservation.py
View file

@ -3057,6 +3057,7 @@ async def test_streaming_cancel_after_only_keepalive_pings_reconciles_to_input_c
@pytest.mark.asyncio
async def test_streaming_cancel_while_holding_back_provider_output_keeps_reservation(
spend_counter_state,
monkeypatch: pytest.MonkeyPatch,
):
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
@ -3086,9 +3087,17 @@ async def test_streaming_cancel_while_holding_back_provider_output_keeps_reserva
}
]
)
response = await router._aanthropic_messages_streaming_iterator(
response=AnthropicMessagesStreamingResponse(completion_stream=held_back, hidden_params={"additional_headers": {}}),
initial_kwargs={"model": "claude-haiku-4-5"},
provider_response = AnthropicMessagesStreamingResponse(
completion_stream=held_back, hidden_params={"additional_headers": {}}
)
monkeypatch.setattr( # test-quality-ok: the provider call the router streams from
litellm, "anthropic_messages", AsyncMock(return_value=provider_response)
)
response = await router.aanthropic_messages(
model="claude-haiku-4-5",
messages=[{"role": "user", "content": "hi"}],
max_tokens=16,
stream=True,
)
async def ping_then_cancel(user_api_key_dict, response, request_data):

10
tests/unit/proxy/test_common_request_processing.py
View file

@ -9027,7 +9027,7 @@ class TestStreamingClientDisconnectBilling:
from litellm.llms.anthropic.messages.adapters.transformation import (
AnthropicAdapter,
)
from litellm.router import FallbackAwareAnthropicMessagesStream
from litellm.router_rust.invoke.streaming import FallbackAwareAnthropicMessagesStream, StreamAttempt
async def _sse_frames() -> AsyncGenerator[bytes, None]:
yield b"event: message_start\n\n"
@ -9046,7 +9046,7 @@ class TestStreamingClientDisconnectBilling:
)
assert isinstance(source_iterator, AnthropicSSEStream)
streamed: Final = prepare_response_for_header_attachment(
FallbackAwareAnthropicMessagesStream(_sse_frames(), source_iterator)
FallbackAwareAnthropicMessagesStream(_sse_frames(), source_iterator, MagicMock(spec=StreamAttempt))
)
billed: Final = await _bill_partial_streamed_spend_on_disconnect(
@ -9586,7 +9586,7 @@ def _apply_stream_usage_tracking(
from litellm.llms.anthropic.messages.responses_adapters.streaming_iterator import (
AnthropicResponsesStreamWrapper,
)
from litellm.router import FallbackAwareAnthropicMessagesStream
from litellm.router_rust.invoke.streaming import FallbackAwareAnthropicMessagesStream, StreamAttempt
from litellm.types.llms.openai import ResponsesAPIResponse
async def _sse_frames() -> AsyncGenerator[bytes, None]:
@ -9595,7 +9595,9 @@ def _apply_stream_usage_tracking(
recorder: Final = _RecordingSuccessLogger()
stream: Final = await self._start_partial_responses_stream(recorder)
wrapper: Final = AnthropicResponsesStreamWrapper(responses_stream=stream, model="gpt-4o-mini")
streamed: Final = FallbackAwareAnthropicMessagesStream(_sse_frames(), wrapper.async_anthropic_sse_wrapper())
streamed: Final = FallbackAwareAnthropicMessagesStream(
_sse_frames(), wrapper.async_anthropic_sse_wrapper(), MagicMock(spec=StreamAttempt)
)
billed: Final = await self._bill_and_wait(recorder, streamed, stream.logging_obj)
assert billed is True

8
tests/unit/proxy/test_litellm_pre_call_utils.py
View file

@ -9213,7 +9213,9 @@ async def test_add_litellm_data_to_request_strips_router_reserved_stamps_on_resp
@pytest.mark.asyncio
async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_strip():
async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_strip(
monkeypatch: pytest.MonkeyPatch,
):
"""Regression for the #38586 break: a client that planted a reserved key in
litellm_metadata made the router hand downstream a scrubbed copy, so the proxy's
post_call write-backs (guardrail telemetry, applied guardrails) landed in a dict the
@ -9247,13 +9249,13 @@ async def test_router_keeps_proxy_metadata_bucket_identity_after_reserved_stamp_
]
)
forwarded_buckets = []
original_acompletion = router._acompletion
original_acompletion = litellm.acompletion
async def _spy(*args, **spy_kwargs):
forwarded_buckets.append(spy_kwargs["metadata"])
return await original_acompletion(*args, **spy_kwargs)
router._acompletion = _spy
monkeypatch.setattr(litellm, "acompletion", _spy)
await router.acompletion(**request_data)

13
tests/unit/proxy/test_proxy_server_endpoints_and_startup.py
View file

@ -5635,6 +5635,13 @@ def _routing_groups_router():
)
def _routing_groups_of(router: litellm.Router) -> dict[str, tuple[list[str], str]]:
return {
group["group_name"]: (group["models"], group["routing_strategy"])
for group in router.settings().as_dict()["routing_groups"] or []
}
@pytest.mark.asyncio
async def test_invalid_db_routing_groups_do_not_abort_other_router_settings():
from unittest.mock import AsyncMock, MagicMock
@ -5656,8 +5663,7 @@ async def test_invalid_db_routing_groups_do_not_abort_other_router_settings():
await ProxyConfig()._add_router_settings_from_db_config(llm_router=router, prisma_client=mock_prisma_client)
assert router.num_retries == 7
assert router._model_to_group == {"m1": "g1"}
assert router._get_routing_context("m1", None)[0] == "latency-based-routing"
assert _routing_groups_of(router) == {"g1": (["m1"], "latency-based-routing")}
@pytest.mark.asyncio
@ -5678,8 +5684,7 @@ async def test_valid_db_routing_groups_still_replace_router_groups():
await ProxyConfig()._add_router_settings_from_db_config(llm_router=router, prisma_client=mock_prisma_client)
assert router.num_retries == 7
assert router._model_to_group == {"m2": "g2"}
assert router._get_routing_context("m2", None)[0] == "least-busy"
assert _routing_groups_of(router) == {"g2": (["m2"], "least-busy")}
def _content_policy_router(**router_kwargs: object) -> litellm.Router:

86
tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py
View file

@ -60,45 +60,31 @@ def test_direct_vector_store_embedding_executor_rejects_invalid_value():
def test_router_vector_store_search_injects_executor_and_request_metadata():
router = litellm.Router(model_list=[])
original = MagicMock(return_value="searched")
wrapped = router.factory_function(original, call_type="vector_store_search")
assert (
wrapped(
vector_store_id="store",
query="query",
custom_llm_provider="valkey",
litellm_metadata={"user_api_key_team_id": "team-a"},
with patch( # test-quality-ok: isolates the SDK search the router dispatches to
"litellm.vector_stores.main.search", return_value="searched"
) as search:
assert (
router.vector_store_search(
vector_store_id="store",
query="query",
custom_llm_provider="valkey",
litellm_metadata={"user_api_key_team_id": "team-a"},
)
== "searched"
)
== "searched"
)
call_kwargs = original.call_args.kwargs
call_kwargs = search.call_args.kwargs
assert call_kwargs["custom_llm_provider"] == "valkey"
executor = call_kwargs["_direct_vector_store_embedding_executor"]
assert isinstance(executor, RouterVectorStoreEmbeddingExecutor)
assert executor.metadata == {"user_api_key_team_id": "team-a"}
assert litellm.Router._vector_store_request_metadata({"metadata": {"user_api_key_team_id": "team-b"}}) == {
"user_api_key_team_id": "team-b"
}
assert litellm.Router._vector_store_request_metadata({}) == {}
with patch.object( # test-quality-ok: fallback dispatch is the boundary this wrapper delegates to
router, "_generic_api_call_with_fallbacks", return_value="routed"
) as fallback:
assert wrapped(model="vector-alias", vector_store_id="store", query="query") == "routed"
assert fallback.call_args.kwargs["model"] == "vector-alias"
assert fallback.call_args.kwargs["original_function"] is original
create_original = MagicMock(return_value="created")
wrapped_create = router.factory_function(create_original, call_type="vector_store_create")
assert wrapped_create(name="store") == "created"
create_original.assert_called_once_with(name="store")
with patch.object( # test-quality-ok: fallback dispatch is the boundary this wrapper delegates to
router, "_generic_api_call_with_fallbacks", return_value="created-through-router"
) as fallback:
assert wrapped_create(model="vector-alias", name="store") == "created-through-router"
fallback.assert_called_once_with(original_function=create_original, model="vector-alias", name="store")
with patch( # test-quality-ok: isolates the SDK create the router dispatches to
"litellm.vector_stores.main.create", return_value="created"
) as create:
assert router.vector_store_create(name="store") == "created"
create.assert_called_once_with(name="store")
@pytest.mark.asyncio
@ -171,15 +157,16 @@ async def test_router_avector_store_search_passes_correct_args():
# Create a router
router = litellm.Router(model_list=[])
# Mock the router's _init_vector_store_api_endpoints method to avoid real API calls
with patch.object(router, "_init_vector_store_api_endpoints") as mock_init:
mock_init.return_value = {
"object": "vector_store.search_results.page",
"search_query": "test query",
"data": [],
}
# Call router's avector_store_search
with patch( # test-quality-ok: isolates the SDK search the router dispatches to
"litellm.vector_stores.main.asearch",
new=AsyncMock(
return_value={
"object": "vector_store.search_results.page",
"search_query": "test query",
"data": [],
}
),
) as asearch:
await router.avector_store_search(
vector_store_id="test_store_id",
query="test query",
@ -187,17 +174,14 @@ async def test_router_avector_store_search_passes_correct_args():
metadata={"user_api_key_team_id": "team-a"},
)
# Verify the internal method was called with correct args
mock_init.assert_called_once()
call_args = mock_init.call_args
# Check that the original function is passed correctly
assert call_args[1]["vector_store_id"] == "test_store_id"
assert call_args[1]["query"] == "test query"
assert call_args[1]["custom_llm_provider"] == "bedrock"
executor = call_args[1]["_direct_vector_store_embedding_executor"]
assert isinstance(executor, RouterVectorStoreEmbeddingExecutor)
assert executor.metadata["user_api_key_team_id"] == "team-a"
asearch.assert_awaited_once()
call_kwargs = asearch.call_args.kwargs
assert call_kwargs["vector_store_id"] == "test_store_id"
assert call_kwargs["query"] == "test query"
assert call_kwargs["custom_llm_provider"] == "bedrock"
executor = call_kwargs["_direct_vector_store_embedding_executor"]
assert isinstance(executor, RouterVectorStoreEmbeddingExecutor)
assert executor.metadata["user_api_key_team_id"] == "team-a"
@pytest.mark.asyncio

89
tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py
View file

@ -4259,75 +4259,6 @@ class TestCompletedResponseLatchedOnStreamEnd:
assert iterator.completed_response.type == "response.completed"
class TestFallbackWrapperStopAsyncIterationFallback:
"""Regression: FallbackResponsesStreamWrapper.__anext__ only sniffed
terminal events off forwarded chunks. When the inner generator raises
StopAsyncIteration without a sniffable chunk (the bridge path ends this
way), the wrapper re-raised without checking source_iterator for a
latched completed_response, leaving its own completed_response None."""
def test_falls_back_to_source_completed_response_on_stop(self):
"""When the inner async generator raises StopAsyncIteration and the
wrapper never sniffed a terminal chunk, it must copy
source_iterator.completed_response so the proxy ownership hook
still sees the terminal event."""
import asyncio
from types import SimpleNamespace
from litellm.router import Router
source = SimpleNamespace(
response=None,
model="deepseek/deepseek-chat",
logging_obj=None,
responses_api_provider_config=None,
start_time=None,
litellm_metadata=None,
custom_llm_provider="deepseek",
request_data={},
call_type="aresponses",
_hidden_params={},
completed_response=SimpleNamespace(
type="response.completed",
response=SimpleNamespace(id="resp_src", output=[], container=None),
),
)
async def empty_gen():
return
yield # pragma: no cover
async def _drive():
router = Router(
model_list=[
{
"model_name": "deepseek/deepseek-chat",
"litellm_params": {"model": "deepseek/deepseek-chat", "api_key": "sk-test"},
}
]
)
wrapper = await router._aresponses_streaming_iterator(
response=source, # type: ignore[arg-type]
initial_kwargs={},
)
wrapper._async_generator = empty_gen()
out = []
try:
async for chunk in wrapper:
out.append(chunk)
except StopAsyncIteration:
pass
return wrapper, out
wrapper, _ = asyncio.run(_drive())
assert wrapper.completed_response is not None, (
"FallbackResponsesStreamWrapper.completed_response is None after "
"StopAsyncIteration even though source_iterator had one — the "
"proxy ownership hook will log a spurious warning"
)
assert wrapper.completed_response.type == "response.completed"
class TestEnsureOutputItemContentPartAdded:
"""Test that _ensure_output_item_for_chunk emits content_part.added after
output_item.added for message items."""
@ -6564,7 +6495,6 @@ def test_an_application_inference_profile_gets_no_item_marks_under_drop_params()
assert AnthropicCacheControlHook.count_request_cache_breakpoints(messages) == 0
def test_a_bedrock_deployment_keeps_item_marks_by_its_own_flag_not_a_siblings_on_the_same_model(
local_model_cost_map: None,
):
@ -6589,6 +6519,7 @@ def test_a_bedrock_deployment_keeps_item_marks_by_its_own_flag_not_a_siblings_on
bridged_system({"id": "opted-in-id", "supports_prompt_cache_breakpoint": True}),
) == ({"role": "system", "content": "House rules"}, marked_system)
def test_the_shared_input_transform_keeps_every_item_mark_for_a_request_it_sends():
messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages(
input=cast(ResponseInputParam, list(_marked_turns(5))), # cast-ok: cache_control is not an input TypedDict key
@ -6901,7 +6832,11 @@ def test_a_new_message_after_a_paused_search_turn_replays_only_that_turn_without
[
{"role": "user", "content": "Plan a Tokyo trip."},
_reasoning_item("rs_tokyo", tokyo_thinking),
{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": "Tokyo is sunny."}]},
{
"type": "message",
"role": "assistant",
"content": [{"type": "output_text", "text": "Tokyo is sunny."}],
},
_hosted_search("ws_srvtoolu_01Tokyo", "completed"),
{"role": "user", "content": "Now Osaka."},
_reasoning_item("rs_osaka", _ANSWERED_THINKING),
@ -7264,7 +7199,11 @@ def test_a_tool_output_with_an_unhashable_call_id_is_a_400_naming_its_input_item
) -> None:
history: Final = cast(
ResponseInputParam,
[_WEATHER_ASK, _weather_call(call_id="call_weather"), {"type": output_type, "call_id": call_id, "output": "sunny"}],
[
_WEATHER_ASK,
_weather_call(call_id="call_weather"),
{"type": output_type, "call_id": call_id, "output": "sunny"},
],
)
with pytest.raises(litellm.BadRequestError) as raised:
@ -7368,7 +7307,11 @@ def test_a_non_string_role_still_converts_for_a_target_that_judges_it_itself() -
id="dict_function_call_id",
),
pytest.param(
[_WEATHER_ASK, _weather_call(call_id=7), {"type": "function_call_output", "call_id": "7", "output": "sunny"}],
[
_WEATHER_ASK,
_weather_call(call_id=7),
{"type": "function_call_output", "call_id": "7", "output": "sunny"},
],
"input[1].call_id",
"int",
id="int_function_call_id",

49
tests/unit/responses/test_responses_router_cooldown.py
View file

@ -11,14 +11,12 @@ from unittest.mock import AsyncMock, patch
import httpx
import pytest
import litellm
from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments
@pytest.mark.asyncio
async def test_responses_api_rate_limit_marks_deployment_for_cooldown():
failing_deployment_id = "deployment-rate-limited"
deployment_ids = {"mock-api-key-1": "deployment-a", "mock-api-key-2": "deployment-b"}
router = litellm.Router(
model_list=[
@ -28,7 +26,7 @@ async def test_responses_api_rate_limit_marks_deployment_for_cooldown():
"model": "openai/gpt-5.1-codex",
"api_key": "mock-api-key-1",
},
"model_info": {"id": failing_deployment_id},
"model_info": {"id": "deployment-a"},
},
{
"model_name": "openai.gpt-5.1-codex",
@ -36,7 +34,7 @@ async def test_responses_api_rate_limit_marks_deployment_for_cooldown():
"model": "openai/gpt-5.1-codex",
"api_key": "mock-api-key-2",
},
"model_info": {"id": "deployment-healthy"},
"model_info": {"id": "deployment-b"},
},
],
num_retries=0,
@ -53,22 +51,16 @@ async def test_responses_api_rate_limit_marks_deployment_for_cooldown():
),
)
def pin_to_failing_deployment(seq):
for d in seq:
if d["model_info"]["id"] == failing_deployment_id:
return d
return seq[0]
called_api_keys = []
with (
patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler",
new_callable=AsyncMock,
side_effect=rate_limit_error,
),
patch(
"litellm.router_strategy.simple_shuffle.random.choice",
side_effect=pin_to_failing_deployment,
),
async def fail_and_record(*args, **kwargs):
called_api_keys.append(kwargs["litellm_params"].api_key)
raise rate_limit_error
with patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler",
new_callable=AsyncMock,
side_effect=fail_and_record,
):
with pytest.raises(litellm.RateLimitError):
await router.aresponses(
@ -76,10 +68,15 @@ async def test_responses_api_rate_limit_marks_deployment_for_cooldown():
input="hi",
)
cooldown_ids = await async_get_cooldown_deployments(
ctx=router._core_ctx, parent_otel_span=None
)
assert failing_deployment_id in cooldown_ids, (
f"Responses API failure callback did not register cooldown for "
f"{failing_deployment_id!r}; cooldown set was {cooldown_ids}"
# num_retries=0: exactly one deployment took the request and failed.
(called_api_key,) = called_api_keys
failed_id = deployment_ids[called_api_key]
cooldown_ids = [
model_id
for model_id, _ in await router.cooldown_cache.async_get_active_cooldowns(
model_ids=list(deployment_ids.values()), parent_otel_span=None
)
]
assert cooldown_ids == [failed_id], (
f"Responses API failure callback did not register cooldown for {failed_id!r}; cooldown set was {cooldown_ids}"
)

108
tests/unit/router_strategy/adaptive_router/test_router_dispatch.py
View file

@ -26,63 +26,6 @@ def _adaptive(r, name):
return r.adaptive_routers[name][0].strategy
def _params(**overrides):
base = {"model": "auto_router/adaptive_router"}
base.update(overrides)
return LiteLLM_Params(**base)
# ---- Fix 1 & 2: opt-in prefix routing -----------------------------------
def test_auto_router_check_excludes_adaptive_router_prefix():
r = Router(model_list=[])
assert (
r._is_auto_router_deployment(
litellm_params=_params(model="auto_router/adaptive_router")
)
is False
)
def test_auto_router_check_excludes_complexity_router_prefix():
r = Router(model_list=[])
assert (
r._is_auto_router_deployment(
litellm_params=_params(model="auto_router/complexity_router")
)
is False
)
def test_auto_router_check_still_matches_plain_auto_router_prefix():
r = Router(model_list=[])
assert (
r._is_auto_router_deployment(
litellm_params=_params(model="auto_router/my-semantic-router")
)
is True
)
def test_adaptive_router_check_recognizes_prefix():
r = Router(model_list=[])
assert (
r._is_adaptive_router_deployment(
litellm_params=_params(model="auto_router/adaptive_router")
)
is True
)
def test_adaptive_router_check_rejects_other_prefixes():
r = Router(model_list=[])
assert (
r._is_adaptive_router_deployment(litellm_params=_params(model="openai/gpt-4o"))
is False
)
# ---- Fix 3: cost field path --------------------------------------------
@ -481,44 +424,6 @@ def test_init_adaptive_router_rejects_duplicate_model_name():
r.init_adaptive_router_deployment(deployment=deployment)
def test_finalize_adaptive_router_if_configured_initializes_and_is_idempotent():
"""`_finalize_adaptive_router_if_configured` walks the model_list, builds an
AdaptiveRouter for each adaptive deployment, and is a safe no-op on
re-entry (models already in self.adaptive_routers are skipped)."""
r = Router(
model_list=[
{
"model_name": "fast",
"litellm_params": {"model": "openai/gpt-4o-mini"},
"model_info": {"input_cost_per_token": 0.00000015},
},
{
"model_name": "smart",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"input_cost_per_token": 0.0000025},
},
{
"model_name": "my-router",
"litellm_params": {
"model": "auto_router/adaptive_router",
"adaptive_router_config": {
"available_models": ["fast", "smart"],
},
},
},
]
)
# Router __init__ already called _finalize_adaptive_router_if_configured.
assert "my-router" in r.adaptive_routers
original = _adaptive(r, "my-router")
# Calling again must be idempotent: the existing AdaptiveRouter instance
# is preserved, not rebuilt.
r._finalize_adaptive_router_if_configured()
assert _adaptive(r, "my-router") is original
def test_finalize_prunes_stale_adaptive_router_hooks_from_callbacks():
"""Replacing the Router (hot-reload path) must not leave stale
AdaptiveRouterPostCallHook instances in `litellm.callbacks` — otherwise
@ -572,16 +477,3 @@ def test_finalize_prunes_stale_adaptive_router_hooks_from_callbacks():
litellm.callbacks.append(cb)
def test_finalize_adaptive_router_if_configured_noop_when_none_configured():
"""With no adaptive deployments in model_list, the finalizer leaves
`adaptive_routers` empty."""
r = Router(
model_list=[
{
"model_name": "fast",
"litellm_params": {"model": "openai/gpt-4o-mini"},
}
]
)
r._finalize_adaptive_router_if_configured()
assert r.adaptive_routers == {}

4
tests/unit/router_strategy/test_baseline_request.py
View file

@ -37,7 +37,7 @@ def test_chat_projection_applies_extra_body_after_top_level_parameters() -> None
def test_baseline_projection_keeps_caller_tools_and_request_parameter_precedence() -> None:
from litellm.router import Router
from litellm.router_rust.invoke.shaping import merge_tools_from_deployment
deployment: Final = {
"tools": [{"type": "function", "function": {"name": "configured"}}],
@ -50,7 +50,7 @@ def test_baseline_projection_keeps_caller_tools_and_request_parameter_precedence
"max_tokens": 128,
}
actual_request: Final = dict(caller)
Router._merge_tools_from_deployment({"litellm_params": deployment}, actual_request)
merge_tools_from_deployment(deployment, actual_request, {})
snapshot: Final = capture_baseline_parameters(caller)
assert snapshot is not None
assert baseline_request({"tools": [{"name": "routed-only"}], "max_tokens": 4}, snapshot, deployment) == {

379
tests/unit/router_strategy/test_complexity_router.py
View file

@ -10,9 +10,8 @@ import logging
import math
import sys
import time
from collections.abc import AsyncIterator, Mapping, Sequence
from collections.abc import Iterator, Mapping, Sequence
from copy import deepcopy
from functools import partial
from types import MappingProxyType
from typing import Dict, Final, List, Literal
from unittest.mock import DEFAULT, AsyncMock, MagicMock, patch
@ -35,7 +34,6 @@ from litellm.constants import (
SESSION_ID_GENERATED_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.router import as_output_cap, stamped_client_ceiling
from litellm.router_api import RoutingContext
from litellm.router_strategy.auto_router.auto_router_model_naming import (
@ -95,7 +93,6 @@ from litellm.types.router import (
LiteLLM_Params,
PreRoutingHookResponse,
RouterErrors,
TaggedPreRoutingStrategy,
)
from litellm.types.utils import StandardLoggingRoutingDecision
@ -1684,36 +1681,6 @@ class TestEdgeCases:
class TestRouterComplexityDeploymentMethods:
"""Tests for Router._is_complexity_router_deployment and Router.init_complexity_router_deployment."""
def test_is_complexity_router_deployment_true(self):
"""_is_complexity_router_deployment returns True for complexity router models."""
router = Router(
model_list=[
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "openai/gpt-4o-mini"},
}
]
)
from litellm.types.router import LiteLLM_Params
params = LiteLLM_Params(model="auto_router/complexity_router/my-router")
assert router._is_complexity_router_deployment(params) is True
def test_is_complexity_router_deployment_false(self):
"""_is_complexity_router_deployment returns False for regular models."""
router = Router(
model_list=[
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "openai/gpt-4o-mini"},
}
]
)
from litellm.types.router import LiteLLM_Params
params = LiteLLM_Params(model="openai/gpt-4o-mini")
assert router._is_complexity_router_deployment(params) is False
def test_init_complexity_router_deployment(self):
"""init_complexity_router_deployment registers a ComplexityRouter."""
router = Router(
@ -2506,106 +2473,6 @@ class TestComplexityRouterTagBasedRouting:
assert us is not None and us.model == "gpt-us"
class TestPreRoutingStrategyRegistry:
"""Directly exercise the tag-scoped registry/selection helpers behind #33655."""
def _router(self) -> Router:
return Router(model_list=[{"model_name": "x", "litellm_params": {"model": "openai/gpt-4o-mini"}}])
@staticmethod
def _deployment(tags: list) -> Deployment:
return Deployment(
model_name="smart",
litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini", tags=tags),
)
def test_deployment_tags_normalizes_to_tuple(self):
router = self._router()
assert router._deployment_tags(self._deployment(["cn", "row"])) == ("cn", "row")
untagged = Deployment(model_name="smart", litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"))
assert router._deployment_tags(untagged) == ()
def test_register_scopes_by_tags_and_rejects_exact_duplicate(self):
router = self._router()
registry: dict = {}
router._register_pre_routing_strategy(
registry=registry, deployment=self._deployment(["cn"]), strategy="CN", strategy_label="Test"
)
router._register_pre_routing_strategy(
registry=registry, deployment=self._deployment(["us"]), strategy="US", strategy_label="Test"
)
assert [entry.tags for entry in registry["smart"]] == [("cn",), ("us",)]
assert router._has_registered_strategy(registry, "smart", ("cn",)) is True
assert router._has_registered_strategy(registry, "smart", ("row",)) is False
with pytest.raises(ValueError, match="already exists"):
router._register_pre_routing_strategy(
registry=registry, deployment=self._deployment(["cn"]), strategy="CN2", strategy_label="Test"
)
def test_select_prefers_request_tag_then_default_then_first(self):
router = self._router()
cn, us, fallback = object(), object(), object()
router.complexity_routers = {
"smart": [
TaggedPreRoutingStrategy(tags=("cn",), strategy=cn),
TaggedPreRoutingStrategy(tags=("us",), strategy=us),
]
}
assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["us"]}}).strategy is us
assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["cn"]}}).strategy is cn
assert router._select_pre_routing_strategy("missing", {"metadata": {"tags": ["cn"]}}) is None
router.complexity_routers = {
"smart": [
TaggedPreRoutingStrategy(tags=("cn",), strategy=cn),
TaggedPreRoutingStrategy(tags=("default",), strategy=fallback),
]
}
assert router._select_pre_routing_strategy("smart", {}).strategy is fallback
router.complexity_routers = {
"smart": [
TaggedPreRoutingStrategy(tags=("cn",), strategy=cn),
TaggedPreRoutingStrategy(tags=("us",), strategy=us),
]
}
assert router._select_pre_routing_strategy("smart", {}).strategy is cn
@staticmethod
def _router_with_plain_smart_deployment(enable_tag_filtering: bool) -> Router:
return Router(
model_list=[{"model_name": "smart", "litellm_params": {"model": "openai/gpt-4o-mini"}}],
enable_tag_filtering=enable_tag_filtering,
)
def test_select_falls_through_to_plain_deployments_when_no_tag_matches_under_tag_filtering(self):
router = self._router_with_plain_smart_deployment(enable_tag_filtering=True)
cn, us = object(), object()
router.complexity_routers = {"smart": [TaggedPreRoutingStrategy(tags=("cn",), strategy=cn)]}
assert router._select_pre_routing_strategy("smart", {}) is None
assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["cn"]}}).strategy is cn
router.complexity_routers = {
"smart": [
TaggedPreRoutingStrategy(tags=("cn",), strategy=cn),
TaggedPreRoutingStrategy(tags=("us",), strategy=us),
]
}
assert router._select_pre_routing_strategy("smart", {}) is None
assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["row"]}}) is None
assert router._select_pre_routing_strategy("smart", {"metadata": {"tags": ["us"]}}).strategy is us
router.complexity_routers["router-only"] = [TaggedPreRoutingStrategy(tags=("cn",), strategy=cn)]
assert router._select_pre_routing_strategy("router-only", {}).strategy is cn
def test_select_keeps_capturing_when_tag_filtering_is_disabled(self):
router = self._router_with_plain_smart_deployment(enable_tag_filtering=False)
cn = object()
router.complexity_routers = {"smart": [TaggedPreRoutingStrategy(tags=("cn",), strategy=cn)]}
assert router._select_pre_routing_strategy("smart", {}).strategy is cn
class TestAsyncPreRoutingHookMultiFormat:
"""Test async_pre_routing_hook with multiple input formats."""
@ -3637,11 +3504,10 @@ def _native_classifier_router(
deployment_model: str = "openai/gpt-6-astra",
failure: Exception | None = None,
native_router: Router | None = None,
http_handler: AsyncHTTPHandler | None = None,
) -> tuple[ComplexityRouter, MagicMock]:
dependency: Final = MagicMock(
aresponses=(
native_router.factory_function(partial(litellm.aresponses, client=http_handler), call_type="aresponses")
native_router.aresponses
if native_router is not None
else AsyncMock(return_value=_native_classifier_response(output), side_effect=failure)
),
@ -3685,15 +3551,14 @@ def _native_classifier_router(
@pytest.fixture
async def native_classifier_http() -> AsyncIterator[tuple[AsyncHTTPHandler, MagicMock]]:
def native_classifier_http(monkeypatch: pytest.MonkeyPatch) -> Iterator[MagicMock]:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
respond: Final = MagicMock(
return_value=httpx.Response(200, json=_native_classifier_response('{"tier":"REASONING"}').model_dump())
)
async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client:
handler: Final = AsyncHTTPHandler()
await handler.client.aclose()
handler.client = client
yield handler, respond
with respx.mock(assert_all_mocked=True, assert_all_called=False) as upstream:
upstream.post(host="classifier.test").mock(side_effect=lambda request: respond(request))
yield respond
class TestEncryptedTaskClassifier:
@ -3871,9 +3736,9 @@ class TestEncryptedTaskClassifier:
["anthropic/test-classifier", "openai/chat_completions/gpt-6-astra", "xai/test-classifier"],
)
async def test_incompatible_classifier_does_not_flatten_encryption(
self, deployment_model: str, native_classifier_http: tuple[AsyncHTTPHandler, MagicMock]
self, deployment_model: str, native_classifier_http: MagicMock
):
handler, respond = native_classifier_http
respond = native_classifier_http
native: Final = Router(
model_list=[
{
@ -3887,7 +3752,7 @@ class TestEncryptedTaskClassifier:
],
num_retries=0,
)
router, _ = _native_classifier_router(native_router=native, http_handler=handler)
router, _ = _native_classifier_router(native_router=native)
result: Final = await router.async_pre_routing_hook(
model="encrypted-router", request_kwargs={"input": [_encrypted_agent_task()]}
@ -3900,9 +3765,9 @@ class TestEncryptedTaskClassifier:
@pytest.mark.asyncio
@pytest.mark.parametrize("blocked", [True, False])
async def test_native_classifier_validates_selected_deployment(
self, blocked: bool, native_classifier_http: tuple[AsyncHTTPHandler, MagicMock]
self, blocked: bool, native_classifier_http: MagicMock
):
handler, respond = native_classifier_http
respond = native_classifier_http
native: Final = Router(
model_list=[
{
@ -3923,7 +3788,7 @@ class TestEncryptedTaskClassifier:
],
num_retries=0,
)
router, _ = _native_classifier_router(native_router=native, http_handler=handler)
router, _ = _native_classifier_router(native_router=native)
task: Final = _encrypted_agent_task()
result: Final = await router.async_pre_routing_hook(model="encrypted-router", request_kwargs={"input": [task]})
@ -3991,7 +3856,6 @@ class TestLLMClassifier:
assert outcome.score is not None
@pytest.mark.asyncio
@pytest.mark.parametrize("redact", (False, True))
@pytest.mark.parametrize(
"override,threshold,tier,model",
(
@ -4007,14 +3871,11 @@ class TestLLMClassifier:
async def test_heuristic_v2_routes_directly_to_predicted_builtin_tier(
self,
mock_router_instance: MagicMock,
redact: bool,
monkeypatch: pytest.MonkeyPatch,
override: Mapping[str, float | None],
threshold: float,
tier: str,
model: str,
) -> None:
monkeypatch.setattr(litellm, "turn_off_message_logging", redact)
artifact: Final = _heuristic_v2_artifact()
router: Final = ComplexityRouter(
model_name="tier-router",
@ -4049,11 +3910,7 @@ class TestLLMClassifier:
"tier-probability:complex=0.892157",
"tier-probability:reasoning=0.980392",
]
redacted: Final = Router._redact_prompt_text_if_needed(
request_kwargs={}, routing_decision=response.routing_decision
)
assert ("signals" in redacted) is not redact
assert redacted["heuristic_v2_forecast"] == {
assert response.routing_decision["heuristic_v2_forecast"] == {
"probabilities": {
"SIMPLE": 11 / 102,
"MEDIUM": 21 / 102,
@ -5065,20 +4922,6 @@ class TestRouterPreRoutingAliasOverrides:
assert "thinking" not in request_kwargs
assert "output_config" not in request_kwargs
def test_drop_client_effort_carriers_helper_edge_shapes(self):
no_pin: Dict = {"thinking": {"type": "adaptive"}}
Router._drop_client_carriers_a_tier_pin_supersedes(no_pin, {"temperature": 0.1})
assert no_pin == {"thinking": {"type": "adaptive"}}
non_dict_carriers: Dict = {"output_config": "max", "reasoning": 3}
Router._drop_client_carriers_a_tier_pin_supersedes(non_dict_carriers, {"reasoning_effort": "low"})
assert non_dict_carriers == {"output_config": "max", "reasoning": 3}
effort_only: Dict = {"output_config": {"effort": "max"}, "reasoning": {"effort": "high"}}
Router._pop_effort_from_nested_carrier(effort_only, "output_config")
Router._pop_effort_from_nested_carrier(effort_only, "reasoning")
assert effort_only == {}
@pytest.mark.asyncio
async def test_client_effort_carriers_survive_when_gate_drops_the_tier_pin(self):
"""The tier-param gate removes a pin the routed target cannot take, and a
@ -5496,15 +5339,6 @@ class TestRouterPreRoutingSharedAliasName:
assert cn_result is not None and cn_result.model == "gpt-cn"
assert "drop_params" not in cn_kwargs
def test_forwardable_alias_marker_params_reads_the_marker_entry_only(self):
router = Router(model_list=[self._plain_entry(), self._marker_entry(), self._tier_entry()])
forwarded = dict(router._forwardable_alias_marker_params(model="gpt4o", strategy_tags=(), request_kwargs={}))
assert forwarded["drop_params"] is True
assert "api_key" not in forwarded and "api_base" not in forwarded
assert router._forwardable_alias_marker_params(model="gemini-flash", strategy_tags=(), request_kwargs={}) == ()
@staticmethod
def _region_marker_entry() -> dict:
return {
@ -5690,27 +5524,6 @@ class TestRouterPreRoutingSharedAliasName:
assert not any(key.startswith("_alias_marker") for key in sent)
def test_forwarded_alias_marker_keys_the_deployment_sets(self):
deployment = {"litellm_params": {"model": "bedrock/x", "aws_region_name": "us-east-1", "timeout": None}}
assert Router._forwarded_alias_marker_keys_the_deployment_sets(
deployment=deployment, forwarded_keys=("aws_region_name", "timeout", "drop_params")
) == ("aws_region_name",)
assert Router._forwarded_alias_marker_keys_the_deployment_sets(deployment=deployment, forwarded_keys=()) == ()
assert Router._forwarded_alias_marker_keys_the_deployment_sets(deployment=deployment, forwarded_keys=None) == ()
assert Router._forwarded_alias_marker_keys_the_deployment_sets(deployment={}, forwarded_keys=("x",)) == ()
def test_deployment_sets_litellm_param(self):
params = {"aws_region_name": "us-east-1", "timeout": None, "use_litellm_proxy": False, "custom_flag": False}
assert Router._deployment_sets_litellm_param(params, "aws_region_name") is True
assert Router._deployment_sets_litellm_param(params, "timeout") is False
assert Router._deployment_sets_litellm_param(params, "missing") is False
assert Router._deployment_sets_litellm_param(params, "use_litellm_proxy") is False
assert Router._deployment_sets_litellm_param({"use_litellm_proxy": True}, "use_litellm_proxy") is True
assert Router._deployment_sets_litellm_param(params, "custom_flag") is True
class TestAdaptiveSoftFloors:
def test_adaptive_defaults_use_cost_weighted_cold_policy(self):
config = ComplexityRouterConfig(
@ -9857,51 +9670,6 @@ class TestRoutingDecisionIsPerAttempt:
assert bucket["routing_decision"]["router_model_name"] == "smart-router"
class TestRecordRoutingDecision:
"""Direct coverage of the single recording point, whose contract is write-or-clear:
the request's metadata must describe the current attempt and nothing else."""
DECISION = {"router_model_name": "smart-router", "router_type": "complexity", "routed_model": "gpt-4o-mini"}
def test_none_clears_a_previous_decision_from_both_buckets(self):
request_kwargs: Dict = {
"metadata": {"routing_decision": self.DECISION, "keep": 1},
"litellm_metadata": {"routing_decision": self.DECISION},
}
Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
assert "routing_decision" not in request_kwargs["metadata"]
assert "routing_decision" not in request_kwargs["litellm_metadata"]
assert request_kwargs["metadata"]["keep"] == 1
def test_none_creates_no_bucket_on_a_request_that_had_none(self):
request_kwargs: Dict = {}
Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
assert request_kwargs == {}
def test_clearing_the_decision_takes_the_savings_facts_with_it(self) -> None:
"""A fallback to a plain model group re-enters the hook with the same
`request_kwargs`. The baseline and the conversation shape ride inside the
decision rather than beside it, so one clear cannot leave either behind and
attribute an auto-router saving to a deployment that never routed."""
from litellm.types.router import BaselineRouteStamp
decision: Final = {
"router_model_name": "smart-router",
"router_type": "complexity",
"routed_model": "gpt-4o-mini",
"savings_baseline_model": "anthropic/claude-opus-5",
"savings_baseline_deployment_id": "opus-deployment",
"conversation_continuing": False,
}
request_kwargs: Final[dict[str, dict[str, object]]] = {"litellm_metadata": {}}
Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=decision)
stamp: Final = request_kwargs["litellm_metadata"]["_autorouter_baseline_route"]
assert isinstance(stamp, BaselineRouteStamp)
assert stamp.baseline_deployment_id == "opus-deployment"
Router._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
assert request_kwargs["litellm_metadata"] == {}
class TestEscalationIsRecordedConsistently:
"""An escalation keyword records two separate facts on every path: that the caller
asked, and whether the tier actually moved. Dropping the ask when there is nowhere
@ -10062,44 +9830,6 @@ class TestRedactedLoggingDropsPromptText:
assert decision["routed_model"] == "gpt-4o"
assert decision["escalated"] is False
def test_only_verbatim_prompt_fields_are_classified_as_prompt_text(self, monkeypatch):
"""The field classification is the whole contract, so pin it directly: anything
that quotes the prompt goes, anything derived from it stays."""
monkeypatch.setattr(litellm, "turn_off_message_logging", True)
full = {
"router_model_name": "smart-router",
"router_type": "complexity",
"routed_model": "gpt-4o",
"cause": "literal_keyword_match",
"tier": "REASONING",
"score": 0.8,
"tier_boundaries": {"simple_medium": 0.15, "medium_complex": 0.35, "complex_reasoning": 0.6},
"classifier_model": "claude-haiku",
"classifier_crux": "deploy the requested service to k8s",
"classifier_primary_rule": "SUP-2",
"classifier_capability_boundary": "supported",
"classifier_p_solve": 0.8,
"classifier_calibrated_p_solve": 0.65,
"classifier_calibration_version": "fitted-v1",
"classifier_threshold": 0.5,
"escalated": True,
"tier_litellm_params": {"reasoning_effort": "xhigh"},
"signals": ["code (python)"],
"matched_keyword": "deploy to k8s",
"escalation_keyword": "LITELLM ESCALATE",
}
kept = Router._redact_prompt_text_if_needed(request_kwargs={}, routing_decision=full)
assert set(full) - set(kept) == {
"signals",
"matched_keyword",
"escalation_keyword",
"classifier_crux",
}
assert kept["classifier_p_solve"] == 0.8
assert kept["classifier_calibrated_p_solve"] == 0.65
assert kept["classifier_calibration_version"] == "fitted-v1"
assert kept["tier_litellm_params"] == {"reasoning_effort": "xhigh"}
@pytest.mark.asyncio
async def test_redaction_via_request_header_is_honored(self):
request_kwargs: Dict = {"metadata": {"headers": {"x-litellm-enable-message-redaction": True}}}
@ -15761,14 +15491,6 @@ class TestHealthFallbackDispatch:
else:
assert decision["tier"] == "REASONING"
assert decision["tier_label"] == "Advanced"
redacted: Final = Router._redact_prompt_text_if_needed(
request_kwargs={"metadata": {"headers": {"x-litellm-enable-message-redaction": True}}},
routing_decision=decision,
)
assert "classifier_crux" not in redacted and "signals" not in redacted
assert {field: value for field, value in redacted.items() if field.startswith("classifier_")} == {
field: value for field, value in expected.items() if field != "classifier_crux"
}
@pytest.mark.asyncio
@pytest.mark.parametrize("peer", (True, False), ids=("peer_failover", "default_fallback"))
@ -16718,13 +16440,15 @@ class TestTierHealthFailover:
("live-big", "openai/gpt-4o-mini", 200000),
)
router_instance.enable_pre_call_checks = True
dead_ids = {d["model_info"]["id"] for d in router_instance.model_list if d["model_name"] == "dead-big"}
async def active_cooldowns(model_ids, parent_otel_span):
return [(i, {"exception_received": "boom"}) for i in model_ids if i in dead_ids]
router_instance.cooldown_cache.async_get_active_cooldowns = active_cooldowns
router_instance.cache = DualCache()
for deployment in router_instance.model_list:
if deployment["model_name"] == "dead-big":
router_instance.cooldown_cache.add_deployment_to_cooldown(
model_id=deployment["model_info"]["id"],
original_exception=RuntimeError("boom"),
exception_status=500,
cooldown_time=600,
)
router = ComplexityRouter(
model_name="health-window-router",
ctx=RoutingContext(router_instance),
@ -17623,24 +17347,6 @@ class TestMaxTokensFromTierModel:
assert sent["max_tokens"] == 64000
@pytest.mark.parametrize(
"tier_params, responses_call, expected",
[
({"max_tokens": 1, "temperature": 0.2}, False, {"max_tokens": 1, "temperature": 0.2}),
({"max_tokens": 1}, True, {"max_output_tokens": 1}),
({"max_completion_tokens": 2}, False, {"max_tokens": 2}),
({"max_completion_tokens": 2}, True, {"max_output_tokens": 2}),
({"max_output_tokens": 3}, False, {"max_tokens": 3}),
({"max_output_tokens": 3}, True, {"max_output_tokens": 3}),
({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 1}),
({"max_tokens": 1, "max_completion_tokens": 2, "max_output_tokens": 3}, True, {"max_output_tokens": 3}),
({"max_completion_tokens": 2, "max_output_tokens": 3}, False, {"max_tokens": 2}),
({"reasoning_effort": "low"}, True, {"reasoning_effort": "low"}),
],
)
def test_every_tier_alias_collapses_onto_the_surface_key(self, tier_params, responses_call, expected):
assert dict(Router._tier_ceiling_under_the_surface_name(tier_params, responses_call=responses_call)) == expected
@pytest.mark.asyncio
async def test_the_default_fallback_exit_carries_the_ceiling(self):
routed: dict = {"max_tokens": 8192}
@ -17813,47 +17519,6 @@ class TestMaxTokensFromTierModel:
assert (None if stamped is None else dict(stamped)) == expected
def test_restoring_the_callers_ceiling_reads_the_stamp_and_replaces_every_carrier(self):
stamped: dict = {"max_output_tokens": 500, "metadata": {"_client_output_ceiling": {"max_tokens": 8192}}}
Router._restore_client_output_ceiling(stamped)
assert {k: v for k, v in stamped.items() if k != "metadata"} == {"max_tokens": 8192}
coerced: dict = {
"max_tokens": 64000,
"metadata": {"_client_output_ceiling": {"max_tokens": "8192", "max_completion_tokens": 100.0}},
}
Router._restore_client_output_ceiling(coerced)
assert {k: v for k, v in coerced.items() if k != "metadata"} == {
"max_tokens": 8192,
"max_completion_tokens": 100,
}
unstamped: dict = {"max_tokens": 64000, "metadata": {}}
Router._restore_client_output_ceiling(unstamped)
assert unstamped["max_tokens"] == 64000
@pytest.mark.asyncio
async def test_pinning_stamps_the_callers_carriers_once(self):
router = self._router()
request_kwargs: dict = {"max_completion_tokens": 8192}
router._stamp_client_output_ceiling(request_kwargs)
assert stamped_client_ceiling(request_kwargs) == {"max_completion_tokens": 8192}
first = router._pin_tier_params_onto_request(
model="big", tier_litellm_params={"max_tokens": 64000}, request_kwargs=request_kwargs, responses_call=False
)
second = router._pin_tier_params_onto_request(
model="big", tier_litellm_params={"max_tokens": 32000}, request_kwargs=request_kwargs, responses_call=False
)
none = router._pin_tier_params_onto_request(
model="big", tier_litellm_params=None, request_kwargs=request_kwargs, responses_call=False
)
assert (first, second, none) == (True, True, False)
assert request_kwargs["max_tokens"] == 32000
assert request_kwargs["metadata"]["_client_output_ceiling"] == {"max_completion_tokens": 8192}
class _OutputCeilingRecorder(CustomLogger):
def __init__(self) -> None:
super().__init__()

33
tests/unit/router_strategy/test_llm_v2.py
View file

@ -389,9 +389,7 @@ async def test_caller_constraints_respect_claude_code_prompt_policy(user_agent:
@pytest.mark.asyncio
@pytest.mark.parametrize("calibrated", (False, True))
async def test_routing_metadata_preserves_exact_forecasts_and_redaction(
calibrated: bool, monkeypatch: pytest.MonkeyPatch
) -> None:
async def test_routing_metadata_preserves_exact_forecasts(calibrated: bool) -> None:
base: Final = _config().llm_v2_config
assert base is not None
calibration: Final = LLMV2Calibration(
@ -410,23 +408,18 @@ async def test_routing_metadata_preserves_exact_forecasts_and_redaction(
assert result.model == ("capable" if calibrated else "efficient")
decision: Final = result.routing_decision
assert decision is not None
monkeypatch.setattr(litellm, "turn_off_message_logging", True)
redacted: Final = Router._redact_prompt_text_if_needed(request_kwargs={}, routing_decision=decision)
assert redacted is not None
assert "signals" not in redacted
for record in (decision, redacted):
assert record["classifier_efficient_p_solve"] == 0.900000123
assert record["classifier_capable_p_solve"] == 0.920000321
assert record["classifier_max_quality_gap"] == 0.05
assert record["classifier_prompt_version"] == LLM_V2_PROMPT_VERSION
if calibrated:
assert record["classifier_calibration_version"] == "test-pair-v1"
assert record["classifier_calibrated_efficient_p_solve"] == calibration.efficient.calibrate(0.900000123)
assert record["classifier_calibrated_capable_p_solve"] == calibration.capable.calibrate(0.920000321)
else:
assert "classifier_calibration_version" not in record
assert "classifier_calibrated_efficient_p_solve" not in record
assert "classifier_calibrated_capable_p_solve" not in record
assert decision["classifier_efficient_p_solve"] == 0.900000123
assert decision["classifier_capable_p_solve"] == 0.920000321
assert decision["classifier_max_quality_gap"] == 0.05
assert decision["classifier_prompt_version"] == LLM_V2_PROMPT_VERSION
if calibrated:
assert decision["classifier_calibration_version"] == "test-pair-v1"
assert decision["classifier_calibrated_efficient_p_solve"] == calibration.efficient.calibrate(0.900000123)
assert decision["classifier_calibrated_capable_p_solve"] == calibration.capable.calibrate(0.920000321)
else:
assert "classifier_calibration_version" not in decision
assert "classifier_calibrated_efficient_p_solve" not in decision
assert "classifier_calibrated_capable_p_solve" not in decision
@pytest.mark.asyncio

34
tests/unit/router_strategy/test_quality_router.py
View file

@ -1024,39 +1024,7 @@ class TestSetResponseHeadersLiftsDecision:
class TestRouterQualityDeploymentMethods:
"""Tests for Router._is_quality_router_deployment and Router.init_quality_router_deployment."""
def test_is_quality_router_deployment_true(self):
"""_is_quality_router_deployment returns True for quality router models."""
from litellm.router import Router
from litellm.types.router import LiteLLM_Params
router = Router(
model_list=[
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "openai/gpt-4o-mini"},
}
]
)
params = LiteLLM_Params(model="auto_router/quality_router/my-router")
assert router._is_quality_router_deployment(params) is True
def test_is_quality_router_deployment_false(self):
"""_is_quality_router_deployment returns False for regular models."""
from litellm.router import Router
from litellm.types.router import LiteLLM_Params
router = Router(
model_list=[
{
"model_name": "gpt-4o-mini",
"litellm_params": {"model": "openai/gpt-4o-mini"},
}
]
)
params = LiteLLM_Params(model="openai/gpt-4o-mini")
assert router._is_quality_router_deployment(params) is False
"""Tests for Router.init_quality_router_deployment."""
def test_init_quality_router_deployment(self):
"""init_quality_router_deployment registers a QualityRouter."""

1218
tests/unit/router_strategy/test_router_routing_groups.py
View file

File diff suppressed because it is too large Load diff

50
tests/unit/router_strategy/test_router_routing_plugins.py
View file

@ -259,56 +259,6 @@ async def test_router_without_plugins_is_unaffected():
assert response.choices[0].message.content == "hi"
@pytest.mark.asyncio
async def test_run_routing_plugins_narrows_candidates_and_records_signals():
"""Unit-level check of _run_routing_plugins in isolation, independent of acompletion."""
router = Router(
model_list=_smart_router_model_list(),
plugins=[LanguageDetector(), DomainClassifier(), TenantPolicy(), BudgetPolicy()],
)
request_kwargs = {"metadata": {"tenant": "acme-corp"}}
context = await router._run_routing_plugins(
model="smart-router",
request_kwargs=request_kwargs,
messages=[{"role": "user", "content": "hi"}],
)
assert context.candidate_models == ["openai/gpt-4o-mini", "anthropic/claude-haiku-4-5"]
assert context.signals["domain-classifier"]["domain"] == "coding"
assert request_kwargs["metadata"]["_routing_plugin_candidate_models"] == context.candidate_models
def test_filter_by_routing_plugin_candidates_narrows_and_raises_when_empty():
"""Unit-level check of _filter_by_routing_plugin_candidates in isolation."""
router = Router(model_list=_smart_router_model_list(), plugins=[TenantPolicy()])
healthy_deployments = router.model_list
narrowed = router._filter_by_routing_plugin_candidates(
healthy_deployments=healthy_deployments,
request_kwargs={"metadata": {"_routing_plugin_candidate_models": ["openai/gpt-4o-mini"]}},
)
assert [d["litellm_params"]["model"] for d in narrowed] == ["openai/gpt-4o-mini"]
with pytest.raises(ValueError, match="No deployments left after routing-plugin filtering"):
router._filter_by_routing_plugin_candidates(
healthy_deployments=healthy_deployments,
request_kwargs={"metadata": {"_routing_plugin_candidate_models": ["nonexistent/model"]}},
)
def test_json_default_stable_id_is_stable_across_instances():
"""generate_model_id's json.dumps `default=` fallback must not embed an object's
memory address (e.g. plain str() on an object with no custom __repr__ falls back
to object.__repr__'s `<module.Class object at 0x...>`) -- that would make the
deployment id churn on every process restart for any deployment whose
litellm_params contain a live plugin instance."""
router = Router(model_list=_smart_router_model_list())
assert router._json_default_stable_id(LanguageDetector()) == router._json_default_stable_id(LanguageDetector())
assert router._json_default_stable_id(LanguageDetector()) != router._json_default_stable_id(TenantPolicy())
def testgenerate_model_id_is_stable_when_litellm_params_contain_a_plugin_instance():
"""End-to-end: a deployment id built from litellm_params containing a routing
plugin instance (e.g. complexity_router_config.plugins) must be identical across

96
tests/unit/router_strategy/test_router_tag_routing.py
View file

@ -4,7 +4,6 @@
import logging
from typing import Final, Literal
from unittest.mock import patch
import pytest
@ -1846,23 +1845,22 @@ async def test_tag_fail_open_survives_deployment_narrowing(
) -> None:
router: Final = _fail_open_router(route=route)
requested_model: Final = "chat" if route == "model" else "chat-test" if route == "wildcard" else "public-chat"
with patch("time.time", return_value=1_700_000_000.0):
if narrowing == "cooldown":
router.cooldown_cache.add_deployment_to_cooldown(
model_id="opted-in",
original_exception=RuntimeError("unavailable"),
exception_status=500,
cooldown_time=60,
)
candidates: Final = await router.async_get_healthy_deployments(
model=requested_model,
request_kwargs={
"litellm_metadata": {"tags": list(tags), "inherited_tags": [], "user_api_key_team_id": "team-a"},
**(
{"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {}
),
},
if narrowing == "cooldown":
router.cooldown_cache.add_deployment_to_cooldown(
model_id="opted-in",
original_exception=RuntimeError("unavailable"),
exception_status=500,
cooldown_time=60,
)
candidates: Final = await router.async_get_healthy_deployments(
model=requested_model,
request_kwargs={
"litellm_metadata": {"tags": list(tags), "inherited_tags": [], "user_api_key_team_id": "team-a"},
**(
{"previous_response_id": _response_id_for_remaining_deployment()} if narrowing == "affinity" else {}
),
},
)
assert isinstance(candidates, (list, tuple))
assert tuple(deployment["model_info"]["id"] for deployment in candidates) == ("remaining",)
@ -2700,6 +2698,13 @@ def _quality_high_cost_low_router():
)
def _cool_down(router: litellm.Router, *model_ids: str) -> None:
for model_id in model_ids:
router.cooldown_cache.add_deployment_to_cooldown(
model_id=model_id, original_exception=RuntimeError("unavailable"), exception_status=500, cooldown_time=60
)
@pytest.mark.asyncio()
async def test_plain_tag_exhaustion_with_universal_default_tag_raises_by_default():
# Every deployment in the group is tagged "default" (a legitimate cross-cutting
@ -2707,21 +2712,16 @@ async def test_plain_tag_exhaustion_with_universal_default_tag_raises_by_default
# the quality:high deployments unhealthy, a request asking for quality:high
# must still hard-fail, not silently get served by a cost:low deployment just
# because it happens to also carry "default".
from unittest.mock import AsyncMock, patch
router = _quality_high_cost_low_router()
_cool_down(router, "quality-high-1", "quality-high-2")
with patch(
"litellm.router_python_core.async_get_cooldown_deployments",
new=AsyncMock(return_value=["quality-high-1", "quality-high-2"]),
):
with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info:
await router.acompletion(
model="gpt-4",
messages=[{"role": "user", "content": "hi"}],
metadata={"tags": ["quality:high"]},
mock_response="hi",
)
with pytest.raises(Exception, match='Not allowed to access model due to tags configuration\\.') as exc_info:
await router.acompletion(
model="gpt-4",
messages=[{"role": "user", "content": "hi"}],
metadata={"tags": ["quality:high"]},
mock_response="hi",
)
from litellm.types.router import RouterErrors
@ -2734,18 +2734,14 @@ async def test_plain_tag_exhaustion_with_universal_default_tag_falls_open_when_a
for deployment in router.model_list:
deployment["model_info"]["allow_fail_open"] = True
from unittest.mock import AsyncMock, patch
_cool_down(router, "quality-high-1", "quality-high-2")
with patch(
"litellm.router_python_core.async_get_cooldown_deployments",
new=AsyncMock(return_value=["quality-high-1", "quality-high-2"]),
):
response = await router.acompletion(
model="gpt-4",
messages=[{"role": "user", "content": "hi"}],
metadata={"tags": ["quality:high"]},
mock_response="hi",
)
response = await router.acompletion(
model="gpt-4",
messages=[{"role": "user", "content": "hi"}],
metadata={"tags": ["quality:high"]},
mock_response="hi",
)
assert response._hidden_params["model_id"] in ("cost-low-1", "cost-low-2")
@ -3095,24 +3091,6 @@ async def test_router_selecting_tag_is_consumed_on_litellm_metadata_shaped_reque
assert deployment["model_info"]["id"] == "tier-gemini-flash"
def test_consumed_request_tags_stamp_names_the_routed_group_and_spent_tags_only_on_a_tag_match():
from litellm.types.router import ConsumedRequestTagsStamp, PreRoutingHookResponse
router = _tagged_marker_router()
strategy = router.auto_routers["gpt4o"][0]
rewrite = PreRoutingHookResponse(model="gemini-flash", messages=None)
consumed = router._consumed_request_tags_stamp(
selected_strategy=strategy, pre_routing_hook_response=rewrite, request_tags=["route"]
)
unmatched = router._consumed_request_tags_stamp(
selected_strategy=strategy, pre_routing_hook_response=rewrite, request_tags=["other"]
)
assert consumed == ConsumedRequestTagsStamp(model_group="gemini-flash", tags=("route",))
assert unmatched is None
@pytest.mark.asyncio()
async def test_tagged_request_direct_to_plain_group_still_rejected():
# Sent straight to the tier, no router selection consumed the tag, so strict

113
tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py
View file

@ -440,41 +440,18 @@ async def test_async_user_parameter_does_not_trigger_deployment_affinity():
model_group = "azure-sdk-test"
user_id = "sdk-user-123"
choice_calls = {"count": 0}
def deterministic_choice(seq):
choice_calls["count"] += 1
if choice_calls["count"] == 1:
return seq[0]
return seq[1] if len(seq) > 1 else seq[0]
with (
patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post,
patch(
"litellm.router_strategy.simple_shuffle.random.choice",
side_effect=deterministic_choice,
),
):
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post:
mock_post.return_value = MockResponse(mock_response_data, 200)
# First call with 'user' parameter (end-user id)
first_response = await router.aresponses(
model=model_group,
input="Hi",
user=user_id,
)
first_model_id = first_response._hidden_params["model_id"]
# Second call with same 'user' parameter should NOT be pinned by affinity
second_response = await router.aresponses(
model=model_group,
input="Follow-up",
user=user_id,
)
assert second_response._hidden_params["model_id"] != first_model_id
# Repeated calls with the same 'user' parameter must keep spreading across the group
model_ids = set()
for _ in range(20):
response = await router.aresponses(model=model_group, input="Hi", user=user_id)
model_ids.add(response._hidden_params["model_id"])
assert len(model_ids) == 2
@pytest.mark.asyncio
@ -839,61 +816,29 @@ async def test_model_group_affinity_config_only_applies_to_configured_group():
)
user_api_key_hash = "test-per-group-key"
choice_calls = {"count": 0}
def deterministic_choice(seq):
choice_calls["count"] += 1
if choice_calls["count"] == 1:
return seq[0]
return seq[1] if len(seq) > 1 else seq[0]
async def served_model_ids(model: str) -> set[str]:
model_ids = set()
for _ in range(20):
response = await router.aresponses(
model=model,
input="Hello",
truncation="auto",
litellm_metadata={"user_api_key_hash": user_api_key_hash},
)
model_ids.add(response._hidden_params["model_id"])
return model_ids
with (
patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post,
patch(
"litellm.router_strategy.simple_shuffle.random.choice",
side_effect=deterministic_choice,
),
):
with patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post:
mock_post.return_value = MockResponse(mock_response_data, 200)
# gpt-4: affinity should work — second request pinned to same deployment
first = await router.aresponses(
model="gpt-4",
input="Hello",
truncation="auto",
litellm_metadata={"user_api_key_hash": user_api_key_hash},
)
first_model_id = first._hidden_params["model_id"]
second = await router.aresponses(
model="gpt-4",
input="Follow-up",
truncation="auto",
litellm_metadata={"user_api_key_hash": user_api_key_hash},
)
assert second._hidden_params["model_id"] == first_model_id
# claude-3: no affinity configured — should NOT be pinned
choice_calls["count"] = 0
first_claude = await router.aresponses(
model="claude-3",
input="Hello",
truncation="auto",
litellm_metadata={"user_api_key_hash": user_api_key_hash},
)
first_claude_id = first_claude._hidden_params["model_id"]
second_claude = await router.aresponses(
model="claude-3",
input="Follow-up",
truncation="auto",
litellm_metadata={"user_api_key_hash": user_api_key_hash},
)
# With deterministic choice and len>1, second call picks seq[1]
assert second_claude._hidden_params["model_id"] != first_claude_id
# gpt-4: affinity pins every request from the key to one deployment
assert len(await served_model_ids("gpt-4")) == 1
# claude-3: no affinity configured, so requests keep spreading across the group
assert len(await served_model_ids("claude-3")) == 2
@pytest.mark.asyncio

52
tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py
View file

@ -849,31 +849,19 @@ async def test_affinity_falls_back_to_same_encryption_boundary_on_model_group_sw
num_retries=0,
)
def first_call_picks_account_a(seq):
for d in seq:
if d["model_info"]["id"] == "gpt-5.3-codex-account-a":
return d
return seq[0]
with (
patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler",
new_callable=AsyncMock,
return_value=first_resp,
),
patch(
"litellm.router_strategy.simple_shuffle.random.choice",
side_effect=first_call_picks_account_a,
),
with patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler",
new_callable=AsyncMock,
return_value=first_resp,
):
r1 = await router.aresponses(model="gpt-5.3-codex", input="hi")
assert r1._hidden_params["model_id"] == "gpt-5.3-codex-account-a"
origin_account = r1._hidden_params["model_id"].removeprefix("gpt-5.3-codex-")
assert origin_account in ("account-a", "account-b")
encoded_id = _extract_encoded_item_id(r1)
assert encoded_id.startswith("encitem_")
# simple_shuffle.random.choice NOT patched: prove affinity narrows the
# candidate pool to a single deployment regardless of which one shuffle picks.
# Affinity narrows the candidate pool to the origin's account regardless of which one shuffle picks.
with patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler",
new_callable=AsyncMock,
@ -890,7 +878,7 @@ async def test_affinity_falls_back_to_same_encryption_boundary_on_model_group_sw
],
)
assert r2._hidden_params["model_id"] == "gpt-5.4-account-a"
assert r2._hidden_params["model_id"] == f"gpt-5.4-{origin_account}"
@pytest.mark.asyncio
@ -977,25 +965,15 @@ async def test_affinity_falls_back_to_same_boundary_on_alias_switch():
num_retries=0,
)
def pick_account_a(seq):
for d in seq:
if d["model_info"]["id"] == "gpt-5.3-codex-account-a":
return d
return seq[0]
with (
patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler",
new_callable=AsyncMock,
return_value=first_resp,
),
patch(
"litellm.router_strategy.simple_shuffle.random.choice",
side_effect=pick_account_a,
),
with patch(
"litellm.llms.custom_httpx.llm_http_handler.BaseLLMHTTPHandler.async_response_api_handler",
new_callable=AsyncMock,
return_value=first_resp,
):
r1 = await router.aresponses(model="gpt-5.3-codex", input="hi")
origin_account = r1._hidden_params["model_id"].removeprefix("gpt-5.3-codex-")
assert origin_account in ("account-a", "account-b")
encoded_id = _extract_encoded_item_id(r1)
assert encoded_id.startswith("encitem_")
@ -1015,7 +993,7 @@ async def test_affinity_falls_back_to_same_boundary_on_alias_switch():
],
)
assert r2._hidden_params["model_id"] == "gpt-5.2-codex-account-a"
assert r2._hidden_params["model_id"] == f"gpt-5.2-codex-{origin_account}"
def test_boundary_fallback_no_router_ref_returns_empty():

85
tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py
View file

@ -146,26 +146,6 @@ async def test_healthy_legacy_string_pin_is_preserved_with_repinning_enabled() -
router.discard()
@pytest.mark.parametrize("route_bucket", ["metadata", "litellm_metadata"])
async def test_repin_session_after_success_reads_split_metadata(route_bucket: str) -> None:
router: Final = _repinning_router()
try:
await router.cache.async_set_cache(_repin_key(), {"model_id": "original"})
_cool_down_original(router)
session_bucket: Final = "litellm_metadata" if route_bucket == "metadata" else "metadata"
await router._repin_session_after_success(
{
"model_info": {"id": "replacement"},
route_bucket: {"deployment_model_name": "repin-group"},
session_bucket: {"session_id": "repin-session", "user_api_key_hash": "caller-1"},
},
{"status": "completed"},
)
assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "replacement"}
finally:
router.discard()
@pytest.mark.parametrize("enabled", [False, True])
async def test_session_selection_observes_another_workers_repin_only_when_enabled(enabled: bool) -> None:
server: Final = FakeServer()
@ -320,49 +300,6 @@ async def test_unsuccessful_or_nested_response_envelopes_never_repin(response: M
router.discard()
class _ReadCountingCache(InMemoryCache):
def __init__(self) -> None:
super().__init__()
self.read_keys: tuple[str, ...] = ()
def get_cache(self, key: str, **kwargs: object) -> object:
self.read_keys += (key,)
return super().get_cache(key, **kwargs)
@pytest.mark.parametrize("case", ["disabled", "no-session", "generated-session", "group-disabled", "active"])
async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: str) -> None:
router: Final = _repinning_router(case != "disabled")
reads: Final = _ReadCountingCache()
try:
router.cooldown_cache.cooldown_store.in_memory_cache = reads
callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck))
if case == "group-disabled":
callback.model_group_affinity_config = {"repin-group": []}
await router._repin_session_after_success(
{
"model_info": {"id": "replacement"},
"metadata": {
"deployment_model_name": "repin-group",
**({"session_id": "repin-session"} if case != "no-session" else {}),
SESSION_ID_GENERATED_METADATA_KEY: case == "generated-session",
},
},
{"status": "completed"},
)
expected: Final = (
{
router.cooldown_cache.get_cooldown_cache_key(model_id)
for model_id in ("original", "replacement", "other")
}
if case == "active"
else set()
)
assert set(reads.read_keys) == expected
finally:
router.discard()
class MockResponse:
def __init__(self, json_data, status_code):
self._json_data = json_data
@ -690,15 +627,19 @@ async def test_auto_router_session_affinity_writes_scoped_pin_and_follows_it():
)
async def test_seeded_session_pin_is_invisible_outside_its_scope(model, key_hash):
"""The pin binds (auto-routed request, api key, session): a direct call to the
group and a different key reusing the session id must both ignore it."""
router = _smart_router()
try:
await router.cache.async_set_cache(
key=_session_pin_key("scoped-session", "key-1"), value={"model_id": "deployment-2"}
)
assert await _one_turn(router, model, "scoped-session", key_hash) == "deployment-1"
finally:
_cleanup_router_callbacks(router)
group and a different key reusing the session id must both ignore it: across fresh
routers each seeded with the pin, the request still reaches the other deployment."""
served = set()
for _ in range(20):
router = _smart_router()
try:
await router.cache.async_set_cache(
key=_session_pin_key("scoped-session", "key-1"), value={"model_id": "deployment-2"}
)
served.add(await _one_turn(router, model, "scoped-session", key_hash))
finally:
_cleanup_router_callbacks(router)
assert "deployment-1" in served
@pytest.mark.asyncio

53
tests/unit/router_utils/test_client_initalization_utils.py
View file

@ -81,14 +81,17 @@ def test_slot_is_released_when_the_held_call_raises():
assert limit.in_flight == 1
def _router_limit(router: Router, model_name: str) -> MaxParallelRequestsLimit:
deployment: Final = router.get_deployment_by_model_group_name(model_group_name=model_name)
assert deployment is not None
client: Final = router._get_client(
deployment=deployment.model_dump(), kwargs={}, client_type="max_parallel_requests"
)
assert isinstance(client, MaxParallelRequestsLimit)
return client
async def _concurrent_outcomes(router: Router, calls: int) -> list[str]:
async def attempt() -> str:
try:
await router.acompletion(
model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], mock_response="ok", mock_delay=0.3
)
except litellm.RateLimitError as e:
return e.message
return "ok"
return await asyncio.gather(*[attempt() for _ in range(calls)])
@pytest.mark.parametrize(
@ -103,23 +106,25 @@ def _router_limit(router: Router, model_name: str) -> MaxParallelRequestsLimit:
@pytest.mark.asyncio
async def test_router_deployment_rejects_past_its_derived_cap(litellm_params: dict[str, int], expected_cap: int):
router: Final = Router(
model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", **litellm_params}}]
model_list=[
{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "k", **litellm_params}}
],
num_retries=0,
)
limit: Final = _router_limit(router, "gpt-5.6")
assert limit.max_parallel_requests == expected_cap
release: Final = asyncio.Event()
holders: Final = [asyncio.create_task(_hold(limit, release)) for _ in range(expected_cap)]
await asyncio.sleep(0)
assert limit.in_flight == expected_cap
assert f"max_parallel_requests={expected_cap}" in _expect_rejection(limit).message
release.set()
assert await asyncio.wait_for(asyncio.gather(*holders), timeout=2) == ["ok"] * expected_cap
outcomes: Final = await _concurrent_outcomes(router, expected_cap + 1)
assert outcomes.count("ok") == expected_cap
rejections: Final = [outcome for outcome in outcomes if outcome != "ok"]
assert len(rejections) == 1
assert f"max_parallel_requests={expected_cap}" in rejections[0]
def test_router_without_any_concurrency_setting_has_no_limit():
router: Final = Router(model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6"}}])
deployment: Final = router.get_deployment_by_model_group_name(model_group_name="gpt-5.6")
assert deployment is not None
assert (
router._get_client(deployment=deployment.model_dump(), kwargs={}, client_type="max_parallel_requests") is None
@pytest.mark.asyncio
async def test_router_without_any_concurrency_setting_has_no_limit():
router: Final = Router(
model_list=[{"model_name": "gpt-5.6", "litellm_params": {"model": "openai/gpt-5.6", "api_key": "k"}}],
num_retries=0,
)
assert await _concurrent_outcomes(router, 20) == ["ok"] * 20

145
tests/unit/router_utils/test_cooldown_handlers.py
View file

@ -1277,80 +1277,6 @@ class TestResolveAllowedFailsFromPolicyFallsThrough:
assert result == 3, "must fall through to BadRequestErrorAllowedFails when the more specific field is unset"
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
class TestDeploymentCallbackOnFailureCooldownTimePrecedence:
def test_model_info_cooldown_time_used_in_primary_sync_path(self):
"""
Router.deployment_callback_on_failure (the primary sync failure-callback path,
as opposed to the fallback path covered by TestFallbackDeploymentCooldown) must
also honor a model_info.cooldown_time, not just litellm_params.cooldown_time.
"""
router = _make_router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "openai/gpt-4"},
"model_info": {"id": "primary", "cooldown_time": 15.0},
},
],
)
exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4")
kwargs = {
"exception": exc,
"litellm_params": {
"model_info": {"id": "primary", "cooldown_time": 15.0},
},
}
with patch("litellm.router_python_core.set_cooldown_deployments") as mock_set_cooldown:
router.deployment_callback_on_failure(
kwargs=kwargs,
completion_response=None,
start_time=0,
end_time=1,
)
mock_set_cooldown.assert_called_once()
call_kwargs = mock_set_cooldown.call_args[1]
assert call_kwargs["time_to_cooldown"] == 15.0, (
"model_info.cooldown_time must be honored in the primary sync failure-callback path"
)
def test_litellm_params_cooldown_time_still_honored_as_fallback(self):
"""cooldown_time has pre-existing litellm_params support on this primary
path; it must keep working when model_info doesn't set it."""
router = _make_router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "openai/gpt-4", "cooldown_time": 20.0},
"model_info": {"id": "primary"},
},
],
)
exc = litellm.RateLimitError("Rate limit", "openai", "gpt-4")
kwargs = {
"exception": exc,
"litellm_params": {
"model_info": {"id": "primary"},
"cooldown_time": 20.0,
},
}
with patch("litellm.router_python_core.set_cooldown_deployments") as mock_set_cooldown:
router.deployment_callback_on_failure(
kwargs=kwargs,
completion_response=None,
start_time=0,
end_time=1,
)
call_kwargs = mock_set_cooldown.call_args[1]
assert call_kwargs["time_to_cooldown"] == 20.0, "litellm_params.cooldown_time must still be honored"
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
class TestCallerScopedOAuthAuthFailureCooldown:
@staticmethod
@ -2465,22 +2391,19 @@ async def test_cooldown_time_zero_uses_zero_not_default():
num_retries=0,
)
with patch.object(router.cooldown_cache, "add_deployment_to_cooldown") as mock_add_cooldown:
try:
await router.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
mock_response="litellm.RateLimitError",
)
except litellm.RateLimitError:
pass
mock_add_cooldown.assert_not_called()
try:
await router.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
mock_response="litellm.RateLimitError",
)
except litellm.RateLimitError:
pass
cooldown_list = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None)
assert len(cooldown_list) == 0
healthy_deployments, _ = await router._async_get_healthy_deployments(model="gpt-3.5-turbo", parent_otel_span=None)
healthy_deployments = await router.select("gpt-3.5-turbo", {}, kind="healthy")
assert len(healthy_deployments) == 1
@ -2559,20 +2482,17 @@ def test_single_deployment_no_cooldowns(num_deployments: int):
router = Router(model_list=model_list, num_retries=0)
with patch.object(router.cooldown_cache, "add_deployment_to_cooldown", new=MagicMock()) as mock_client:
try:
router.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
mock_response="litellm.RateLimitError",
)
except litellm.RateLimitError:
pass
try:
router.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
mock_response="litellm.RateLimitError",
)
except litellm.RateLimitError:
pass
if num_deployments == 1:
mock_client.assert_not_called()
else:
mock_client.assert_called_once()
cooled = get_cooldown_deployments(RoutingContext(router), parent_otel_span=None)
assert len(cooled) == (0 if num_deployments == 1 else 1)
@pytest.mark.asyncio
@ -3608,33 +3528,6 @@ def test_naming_tag_routed_deployments_leaves_the_request_metadata_untouched() -
assert metadata == {"tags": ["tenant-a"]}
@pytest.mark.parametrize(
("has_request_input", "expected"), ((True, (TENANT_A, SHARED)), (False, EVERY_TENANT)), ids=("chat", "no-input")
)
def test_region_routable_deployments_follow_the_pre_call_region_check(
has_request_input: bool, expected: tuple[CooldownDeploymentInfo, ...]
) -> None:
router: Final = _tenant_router((), **REGION_CHECKS_ON)
kept: Final = router._region_routable_deployments(
router.get_model_list(model_name="tenants") or [], {"allowed_model_region": "eu"}, has_request_input
)
assert router._deployment_ids(kept) == frozenset(deployment["model_id"] for deployment in expected)
def test_routable_cooled_deployment_ids_skip_healthy_deployments_and_peel_tag_fallbacks() -> None:
router: Final = _tenant_router((TENANT_A, SHARED), **TAGS_ON)
routable: Final = asyncio.run(
router._async_routable_cooled_deployment_ids(
model="tenants",
request_deployments=router.get_model_list(model_name="tenants") or [],
cooldown_deployments=(TENANT_A["model_id"], SHARED["model_id"]),
request_kwargs={"metadata": {}},
has_request_input=True,
)
)
assert routable == frozenset((TENANT_A["model_id"], SHARED["model_id"]))
REACHABILITY_REQUESTS: Final = (
{"metadata": {"tags": ["tenant-a"]}},
{"metadata": {}},

483
tests/unit/router_utils/test_fallback_event_handlers.py
View file

@ -10,13 +10,11 @@ from unittest.mock import MagicMock, patch
import httpx
import pytest
import respx
from pydantic import TypeAdapter
import litellm
from litellm import Router
from litellm.integrations.custom_logger import CustomLogger
from litellm.llms.anthropic.messages.native.fake_stream_iterator import FakeAnthropicMessagesStreamIterator
from litellm.litellm_core_utils import get_llm_provider_logic
from litellm.litellm_core_utils.litellm_logging import Logging
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
@ -39,7 +37,6 @@ from litellm.router_utils.fallback_event_handlers import (
is_bad_request_without_fallback,
log_failure_fallback_event,
log_success_fallback_event,
mid_stream_fallback_hop_kwargs,
mid_stream_fallback_snapshot_kwargs,
mid_stream_retry_kwargs,
record_pre_routing_selection,
@ -49,9 +46,6 @@ from litellm.router_utils.fallback_event_handlers import (
)
from litellm.types.router import FallbackSkipReason, SkippedFallbackTarget
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
from tests.fake_openai_endpoint import FAKE_OPENAI_API_BASE
from typing import Dict
import os
class StreamingWrapper:
@ -245,65 +239,10 @@ async def test_run_async_fallback_preserves_superseded_usage_when_hop_fails_befo
assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),)
class _RateLimitedHop(Exception):
def __init__(self, deployment_id: str = "fallback-dep-1") -> None:
super().__init__("fallback deployment rate limited")
self.status_code = 429
self.failed_deployment_id = deployment_id
_EARLIER_FAILURE_MARKER: Final = datetime(2026, 1, 1)
_FIRST_HOP_FAILURE_MARKER: Final = datetime(2026, 1, 2)
_FINAL_HOP_FAILURE_MARKER: Final = datetime(2026, 1, 3)
def _two_failure_fallback_router() -> Router:
router: Final = Router(
model_list=[
{
"model_name": "fallback-model",
"litellm_params": {"model": "openai/gpt-5.6", "api_key": "fake-key"},
"model_info": {"id": f"fallback-dep-{index}"},
}
for index in (1, 2)
],
allowed_fails=0,
cooldown_time=60,
num_retries=0,
)
async def fail_both_hops(*args: object, **kwargs: object) -> NoReturn:
logging_obj: Final[object] = kwargs["litellm_logging_obj"]
assert isinstance(logging_obj, Logging)
assert logging_obj.model_call_details.get("has_logged_async_failure", False) is False
logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FIRST_HOP_FAILURE_MARKER
try:
raise _RateLimitedHop("fallback-dep-1")
except _RateLimitedHop:
logging_obj.mark_logging_complete("async_failure")
logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FINAL_HOP_FAILURE_MARKER
raise _RateLimitedHop("fallback-dep-2")
router.async_function_with_fallbacks = fail_both_hops
return router
def _two_deployment_router() -> Router:
return Router(
model_list=[
{
"model_name": "fallback-model",
"litellm_params": {"model": "openai/gpt-5.6", "api_key": "fake-key"},
"model_info": {"id": f"fallback-dep-{index}"},
}
for index in (1, 2)
],
num_retries=0,
)
def _entra_fallback_router() -> Router:
return Router(
model_list=[
@ -361,67 +300,6 @@ async def test_fallback_hop_cools_down_only_a_deployment_entra_rejected_its_own_
assert await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None) == cooled_down
@pytest.mark.asyncio
@pytest.mark.parametrize("failure_already_logged", [False, True])
async def test_run_async_fallback_cools_down_a_failed_hop_once(failure_already_logged: bool) -> None:
"""The manual cooldown stands in for the failure callbacks the logging object blocks once a failure has
been logged. A hop whose own failure was the first one logged already ran them, so only a hop that failed
with the flag already set gets the manual cooldown."""
router: Final = _two_deployment_router()
logging_obj: Final = _request_logging_obj()
if failure_already_logged:
logging_obj.model_call_details["current_attempt_api_call_start_time"] = _EARLIER_FAILURE_MARKER
logging_obj.mark_logging_complete("async_failure")
async def rate_limited_hop(*args: object, **kwargs: object) -> NoReturn:
hop_logging_obj: Final = kwargs["litellm_logging_obj"]
assert isinstance(hop_logging_obj, Logging)
hop_logging_obj.model_call_details["current_attempt_api_call_start_time"] = _FINAL_HOP_FAILURE_MARKER
if hop_logging_obj.model_call_details.get("has_logged_async_failure", False) is not True:
hop_logging_obj.mark_logging_complete("async_failure")
raise _RateLimitedHop()
with pytest.raises(_RateLimitedHop):
await run_async_fallback(
litellm_routing_ctx=RoutingContext(router),
fallback_model_group=["fallback-model"],
original_model_group="primary-model",
original_exception=RuntimeError("stream dropped before content"),
max_fallbacks=3,
fallback_depth=0,
litellm_logging_obj=logging_obj,
original_function=rate_limited_hop,
metadata={},
)
cooled_down: Final = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None)
assert cooled_down == (["fallback-dep-1"] if failure_already_logged else [])
@pytest.mark.asyncio
@pytest.mark.usefixtures("python_router_core") # the Rust core walks fallbacks in Rust, so it has no such member
async def test_run_async_fallback_cools_down_the_final_deployment_after_retry_failure_logging() -> None:
router: Final = _two_failure_fallback_router()
logging_obj: Final = _request_logging_obj()
assert logging_obj.model_call_details.get("has_logged_async_failure", False) is False
with pytest.raises(_RateLimitedHop):
await run_async_fallback(
litellm_routing_ctx=RoutingContext(router),
fallback_model_group=["fallback-model"],
original_model_group="primary-model",
original_exception=RuntimeError("stream dropped before content"),
max_fallbacks=3,
fallback_depth=0,
litellm_logging_obj=logging_obj,
metadata={},
)
assert logging_obj.model_call_details.get("has_logged_async_failure", False) is True
cooled_down: Final = await async_get_cooldown_deployments(ctx=RoutingContext(router), parent_otel_span=None)
assert cooled_down == ["fallback-dep-2"]
@pytest.mark.asyncio
async def test_run_async_fallback_forwards_include_fallback_errors_to_nested_call():
"""A nested fallback (multi-hop) must keep collecting errors, so the opt-in
@ -1362,7 +1240,7 @@ class TestTriggerCooldownForFailedDeployment:
async def _stalled_stream_error() -> Exception:
from litellm.router import _aread_first_chunk
from litellm.router_rust.invoke.streaming import _aread_first_chunk # noqa: PLC2701
async def _stall(**kwargs: object) -> NoReturn:
raise httpx.ReadTimeout("Timeout on reading data from socket")
@ -2244,41 +2122,6 @@ async def test_streaming_refusal_without_fallback_row_passes_frames_through():
assert len(fake.calls) == 1
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
async def test_refused_stream_built_from_a_finished_agentic_turn_is_delivered_without_a_fallback():
"""Websearch, MCP and advisor loops run their tools before streaming the final reply, so a
fallback would run every tool call again for a reply the client then gets from another model."""
fake: Final = FakeAnthropicUpstream()
router: Final = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
with fake.install():
stream: Final = await router._aanthropic_messages_streaming_iterator(
response=FakeAnthropicMessagesStreamIterator(REFUSAL_RESPONSE),
initial_kwargs=mid_stream_fallback_hop_kwargs(
model="fable-tier",
original_generic_function=litellm.anthropic_messages,
controls=None,
kwargs={
"model": "fable-tier",
"max_tokens": 16,
"stream": True,
"messages": [{"role": "user", "content": "hi"}],
},
),
)
events: Final = [
json.loads(line[len(b"data: ") :])
for line in (await _collect(stream)).splitlines()
if line.startswith(b"data: ")
]
assert fake.calls == []
message_delta: Final = next(event for event in events if event["type"] == "message_delta")
assert message_delta["delta"]["stop_reason"] == "refusal"
assert message_delta["delta"]["stop_details"] == REFUSAL_RESPONSE["stop_details"]
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
async def test_streaming_refusal_on_routed_tier_matches_tier_keyed_row_without_inbound_metadata():
@ -2424,28 +2267,6 @@ async def test_configured_content_policy_list_stays_authoritative_over_generic_r
assert len(fake.calls) == 1
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_refusal_fallback_available_arms_on_generic_rows_only_without_content_policy():
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"tier-group": ["opus-target"]}])
stamped = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}}
assert router._refusal_fallback_available("router-group", stamped) is True
assert router._refusal_fallback_available("router-group", {}) is False
assert router._refusal_fallback_available("router-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_chat_content_filter_gate_unchanged_by_generic_rows():
"""The generic-row arming is scoped to /v1/messages safeguard refusals; the chat surface's
content_filter gate keeps its long-standing content-policy-only semantics."""
from litellm.types.utils import Choices, ModelResponse
router = Router(model_list=[FABLE_TIER, OPUS_TARGET], fallbacks=[{"fable-tier": ["opus-target"]}])
response = ModelResponse(choices=[Choices(finish_reason="content_filter")])
assert router._should_raise_content_policy_error(model="fable-tier", response=response, kwargs={}) is False
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"])
@ -2493,74 +2314,6 @@ async def test_disable_fallbacks_beats_a_content_policy_row_too():
assert len(fake.calls) == 1
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_refusal_gate_keys_on_pre_routing_tier_stamp():
router = _router(content_policy_fallbacks=[{"tier-group": ["opus-target"]}])
def anthropic_messages(**kwargs: Any) -> None:
return None
refusal_kwargs = {"litellm_metadata": {PRE_ROUTING_SELECTED_MODEL_KEY: "tier-group"}}
assert (
router._should_raise_anthropic_refusal_error(
model="router-group",
original_generic_function=anthropic_messages,
response=dict(REFUSAL_RESPONSE),
kwargs=refusal_kwargs,
)
is True
)
assert (
router._should_raise_anthropic_refusal_error(
model="router-group",
original_generic_function=anthropic_messages,
response=dict(REFUSAL_RESPONSE),
kwargs={},
)
is False
)
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_has_content_policy_fallback_default_fallbacks_arm():
router = Router(model_list=[OPUS_TARGET], fallbacks=[{"*": ["opus-target"]}])
assert router._has_content_policy_fallback("any-group", {}) is True
assert router._has_content_policy_fallback("any-group", {"content_policy_fallbacks": [{"other": ["x"]}]}) is False
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_get_fallback_model_group_for_lookup_groups_orders_tier_before_requested():
router = _router(content_policy_fallbacks=None)
fallbacks = [{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}]
assert router._get_fallback_model_group_for_lookup_groups(
fallbacks=fallbacks, lookup_groups=("tier1", "smart-router")
) == ["backup-a"]
assert router._get_fallback_model_group_for_lookup_groups(
fallbacks=fallbacks, lookup_groups=("tier9", "smart-router")
) == ["backup-b"]
assert router._get_fallback_model_group_for_lookup_groups(fallbacks=fallbacks, lookup_groups=()) is None
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
def test_refusal_gate_ignores_other_generic_call_types():
router = _router(content_policy_fallbacks=[{"fable-tier": ["opus-target"]}])
def aresponses(**kwargs: Any) -> None:
return None
assert (
router._should_raise_anthropic_refusal_error(
model="fable-tier",
original_generic_function=aresponses,
response=dict(REFUSAL_RESPONSE),
kwargs={},
)
is False
)
class FallbackEventLogger(CustomLogger):
def __init__(self) -> None:
super().__init__()
@ -2600,61 +2353,6 @@ def _create_fallback_event_router() -> Router:
)
@pytest.mark.parametrize(
"function_name",
["_acompletion", "_atext_completion", "_aembedding"],
)
@pytest.mark.asyncio
async def test_run_async_fallback(function_name):
"""
Basic test - given a list of fallback models, run the original function with the fallback models
"""
router = create_test_router()
original_function = getattr(router, function_name)
litellm.set_verbose = True
fallback_model_group = ["gpt-4"]
original_model_group = "gpt-3.5-turbo"
original_exception = litellm.exceptions.InternalServerError(
message="Simulated error",
llm_provider="openai",
model="gpt-3.5-turbo",
)
request_kwargs = {
"mock_response": "hello this is a test for run_async_fallback",
"metadata": {"previous_models": ["gpt-3.5-turbo"]},
}
if function_name == "_aembedding":
request_kwargs["input"] = "hello this is a test for run_async_fallback"
elif function_name == "_atext_completion":
request_kwargs["prompt"] = "hello this is a test for run_async_fallback"
elif function_name == "_acompletion":
request_kwargs["messages"] = [{"role": "user", "content": "Hello, world!"}]
result = await run_async_fallback(
litellm_routing_ctx=RoutingContext(router),
original_function=original_function,
num_retries=1,
fallback_model_group=fallback_model_group,
original_model_group=original_model_group,
original_exception=original_exception,
max_fallbacks=5,
fallback_depth=0,
**request_kwargs,
)
assert result is not None
if function_name == "_acompletion":
assert isinstance(result, litellm.ModelResponse)
elif function_name == "_atext_completion":
assert isinstance(result, litellm.TextCompletionResponse)
elif function_name == "_aembedding":
assert isinstance(result, litellm.EmbeddingResponse)
@pytest.mark.asyncio
async def test_log_success_fallback_event():
"""
@ -2714,71 +2412,6 @@ async def test_log_failure_fallback_event():
)
@pytest.mark.asyncio
@pytest.mark.parametrize("function_name", ["_acompletion", "_atext_completion"])
async def test_failed_fallbacks_raise_most_recent_exception(function_name):
"""
Tests that if all fallbacks fail, the most recent occuring exception is raised
meaning the exception from the last fallback model is raised
"""
router = create_test_router()
original_function = getattr(router, function_name)
fallback_model_group = ["gpt-4"]
original_model_group = "gpt-3.5-turbo"
original_exception = litellm.exceptions.InternalServerError(
message="Simulated error",
llm_provider="openai",
model="gpt-3.5-turbo",
)
request_kwargs: Dict[str, Any] = {"metadata": {"previous_models": ["gpt-3.5-turbo"]}}
if function_name == "_aembedding":
request_kwargs["input"] = "hello this is a test for run_async_fallback"
elif function_name == "_atext_completion":
request_kwargs["prompt"] = "hello this is a test for run_async_fallback"
elif function_name == "_acompletion":
request_kwargs["messages"] = [{"role": "user", "content": "Hello, world!"}]
with pytest.raises(litellm.exceptions.RateLimitError):
await run_async_fallback(
litellm_routing_ctx=RoutingContext(router),
original_function=original_function,
num_retries=1,
fallback_model_group=fallback_model_group,
original_model_group=original_model_group,
original_exception=original_exception,
mock_response="litellm.RateLimitError",
max_fallbacks=5,
fallback_depth=0,
**request_kwargs,
)
def create_test_router():
return Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
{
"model_name": "gpt-4",
"litellm_params": {
"model": "gpt-4",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
],
fallbacks=[{"gpt-3.5-turbo": ["gpt-4"]}],
)
class CustomTestLogger(CustomLogger):
def __init__(self):
super().__init__()
@ -2800,120 +2433,6 @@ class CustomTestLogger(CustomLogger):
self.failure_fallback_events.append((original_model_group, kwargs, original_exception))
def create_test_router_2():
return Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
{
"model_name": "gpt-4",
"litellm_params": {
"model": "gpt-4",
"api_key": "very-fake-key",
},
},
{
"model_name": "fake-openai-endpoint-2",
"litellm_params": {
"model": "openai/fake-openai-endpoint-2",
"api_key": "working-key-since-this-is-fake-endpoint",
"api_base": FAKE_OPENAI_API_BASE,
},
},
],
)
@pytest.mark.asyncio
@pytest.mark.parametrize("function_name", ["_acompletion", "_atext_completion"])
async def test_multiple_fallbacks(function_name, respx_mock: respx.MockRouter, monkeypatch):
"""
Tests that if multiple fallbacks passed:
- fallback 1 = bad configured deployment / failing endpoint
- fallback 2 = working deployment / working endpoint
Assert that:
- a success response is received from the working endpoint (fallback 2)
"""
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
invalid_key = httpx.Response(
401,
json={
"error": {
"message": "Incorrect API key provided",
"type": "invalid_request_error",
"code": "invalid_api_key",
}
},
)
respx_mock.post("https://api.openai.com/v1/chat/completions").mock(return_value=invalid_key)
respx_mock.post("https://api.openai.com/v1/completions").mock(return_value=invalid_key)
respx_mock.post(f"{FAKE_OPENAI_API_BASE}/chat/completions").mock(
return_value=httpx.Response(
200,
json={
"id": "chatcmpl-fake",
"object": "chat.completion",
"created": 1700000000,
"model": "fake-openai-endpoint-2",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
},
)
)
respx_mock.post(f"{FAKE_OPENAI_API_BASE}/completions").mock(
return_value=httpx.Response(
200,
json={
"id": "cmpl-fake",
"object": "text_completion",
"created": 1700000000,
"model": "fake-openai-endpoint-2",
"choices": [{"index": 0, "text": "hi", "logprobs": None, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
},
)
)
router_2 = create_test_router_2()
original_function = getattr(router_2, function_name)
fallback_model_group = ["gpt-4", "fake-openai-endpoint-2"]
original_model_group = "gpt-3.5-turbo"
original_exception = Exception("Simulated error")
request_kwargs: dict[str, object] = {"metadata": {"previous_models": ["gpt-3.5-turbo"]}}
if function_name == "_aembedding":
request_kwargs["input"] = "hello this is a test for run_async_fallback"
elif function_name == "_atext_completion":
request_kwargs["prompt"] = "hello this is a test for run_async_fallback"
elif function_name == "_acompletion":
request_kwargs["messages"] = [{"role": "user", "content": "Hello, world!"}]
result = await run_async_fallback(
litellm_routing_ctx=RoutingContext(router_2),
original_function=original_function,
num_retries=1,
fallback_model_group=fallback_model_group,
original_model_group=original_model_group,
original_exception=original_exception,
max_fallbacks=5,
fallback_depth=0,
**request_kwargs,
)
print(result)
print(result._hidden_params)
assert result._hidden_params["api_base"] == FAKE_OPENAI_API_BASE
def _provider_error(status_code: int) -> httpx.Response:
return httpx.Response(status_code=status_code, request=httpx.Request(method="POST", url="https://provider.test"))

406
tests/unit/router_utils/test_health_check_allowed_fails_integration.py
View file

@ -6,7 +6,8 @@ should increment the same counters used by allowed_fails_policy, using the
actual exception type from the health check error.
"""
from unittest.mock import ANY, patch
import time
from unittest.mock import patch
import pytest
@ -14,9 +15,21 @@ import litellm
from litellm.proxy.health_check import run_with_timeout
from litellm.router import Router
from litellm.router_api import RoutingContext
from litellm.router_utils.cooldown_handlers import get_cooldown_deployments
from litellm.router_utils.router_callbacks.track_deployment_metrics import (
get_deployment_failures_for_current_minute,
)
from litellm.types.router import AllowedFailsPolicy
def _cooled(router: Router) -> list[str]:
return list(get_cooldown_deployments(RoutingContext(router), parent_otel_span=None))
def _failures(router: Router, model_id: str) -> int:
return get_deployment_failures_for_current_minute(RoutingContext(router), model_id)
def _make_model(model_id: str, model_name: str = "gpt-4") -> dict:
return {
"model_name": model_name,
@ -277,43 +290,29 @@ class TestHealthCheckCooldownIntegration:
class TestWriteHealthStateIntegration:
"""Test _write_health_state_to_router_cache integrates with cooldown pipeline."""
"""Test _write_health_state_to_router_cache feeds the router's cooldown pipeline."""
def test_unhealthy_endpoint_triggers_set_cooldown(self):
"""_write_health_state_to_router_cache should call set_cooldown_deployments for unhealthy endpoints."""
def test_unhealthy_endpoint_cools_down_the_deployment(self):
"""An unhealthy endpoint's exception counts against the policy and cools the deployment down."""
import litellm.proxy.proxy_server as proxy_module
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
router = Router(
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5),
model_list=[_make_model("deploy-1"), _make_model("deploy-2")],
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0),
enable_health_check_routing=True,
)
timeout_exc = litellm.Timeout(message="Health check timeout", model="", llm_provider="")
unhealthy_endpoints = [
{"model_id": "deploy-1", "error": "timeout"},
]
healthy_endpoints = [
{"model_id": "deploy-2"},
]
with patch.object(proxy_module, "llm_router", router):
with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown:
_write_health_state_to_router_cache(
healthy_endpoints=healthy_endpoints,
unhealthy_endpoints=unhealthy_endpoints,
exceptions_by_model_id={"deploy-1": timeout_exc},
)
mock_cooldown.assert_called_once_with(
ANY,
original_exception=timeout_exc,
exception_status=408, # Timeout has status_code 408
deployment="deploy-1",
time_to_cooldown=router.settings().cooldown_time,
)
assert mock_cooldown.call_args.args[0].belongs_to(router)
_write_health_state_to_router_cache(
healthy_endpoints=[{"model_id": "deploy-2"}],
unhealthy_endpoints=[{"model_id": "deploy-1", "error": "timeout"}],
exceptions_by_model_id={"deploy-1": timeout_exc},
)
assert _cooled(router) == ["deploy-1"]
def test_unhealthy_endpoint_without_exception_skips_cooldown(self):
"""Unhealthy endpoints without an exception key should not trigger cooldown."""
@ -321,270 +320,72 @@ class TestWriteHealthStateIntegration:
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
router = Router(
model_list=[_make_model("deploy-1")],
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=5),
model_list=[_make_model("deploy-1"), _make_model("deploy-2")],
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=0),
enable_health_check_routing=True,
)
unhealthy_endpoints = [
{"model_id": "deploy-1", "error": "unknown failure"},
]
with patch.object(proxy_module, "llm_router", router):
with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown:
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=unhealthy_endpoints,
# no exceptions_by_model_id → cooldown should not fire
)
mock_cooldown.assert_not_called()
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=[{"model_id": "deploy-1", "error": "unknown failure"}],
)
assert _failures(router, "deploy-1") == 0
assert _cooled(router) == []
def test_unhealthy_endpoint_increments_failure_counter(self):
"""Unhealthy endpoints should call increment_deployment_failures_for_current_minute."""
"""An unhealthy endpoint counts one failure for its deployment."""
import litellm.proxy.proxy_server as proxy_module
from litellm.proxy.proxy_server import _write_health_state_to_router_cache
router = Router(
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
model_list=[_make_model("deploy-1"), _make_model("deploy-2")],
allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=10),
enable_health_check_routing=True,
)
rate_exc = litellm.RateLimitError(message="Rate limited", model="gpt-4", llm_provider="openai")
unhealthy_endpoints = [
{"model_id": "deploy-1", "error": "rate limited"},
]
with patch.object(proxy_module, "llm_router", router):
with patch(
"litellm.router_utils.router_callbacks.track_deployment_metrics.increment_deployment_failures_for_current_minute"
) as mock_increment:
with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments"):
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=unhealthy_endpoints,
exceptions_by_model_id={"deploy-1": rate_exc},
)
mock_increment.assert_called_once_with(ANY, "deploy-1")
assert mock_increment.call_args.args[0].belongs_to(router)
class TestHealthCheckFilterBypassWithPolicy:
"""
When allowed_fails_policy is set, the binary health check filter should be
bypassed so cooldown is the sole routing exclusion mechanism.
"""
def test_filter_bypassed_when_policy_set(self):
"""Binary health check filter is a no-op when allowed_fails_policy is configured."""
import time
from litellm.caching.caching import DualCache
from litellm.router_utils.health_state_cache import DeploymentHealthCache
router = Router(
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=3),
enable_health_check_routing=True,
)
# Mark deploy-1 as unhealthy in the health state cache
cache = DualCache()
health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0)
health_cache.set_deployment_health_states(
{
"deploy-1": {
"is_healthy": False,
"timestamp": time.time(),
"reason": "test",
},
}
)
router.health_state_cache = health_cache
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
# Filter should pass all through because policy is set
result = router._filter_health_check_unhealthy_deployments(deployments)
assert len(result) == 2, "Binary filter should be bypassed when allowed_fails_policy is set"
def test_filter_active_when_no_policy(self):
"""Binary health check filter still works when no allowed_fails_policy is configured."""
import time
from litellm.caching.caching import DualCache
from litellm.router_utils.health_state_cache import DeploymentHealthCache
router = Router(
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
enable_health_check_routing=True,
)
cache = DualCache()
health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0)
health_cache.set_deployment_health_states(
{
"deploy-1": {
"is_healthy": False,
"timestamp": time.time(),
"reason": "test",
},
}
)
router.health_state_cache = health_cache
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
result = router._filter_health_check_unhealthy_deployments(deployments)
assert len(result) == 1
assert result[0]["model_info"]["id"] == "deploy-2"
@pytest.mark.asyncio
async def test_async_filter_bypassed_when_policy_set(self):
"""Async version also bypasses when allowed_fails_policy is set."""
import time
from litellm.caching.caching import DualCache
from litellm.router_utils.health_state_cache import DeploymentHealthCache
router = Router(
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
allowed_fails_policy=AllowedFailsPolicy(TimeoutErrorAllowedFails=2),
enable_health_check_routing=True,
)
cache = DualCache()
health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0)
health_cache.set_deployment_health_states(
{
"deploy-1": {
"is_healthy": False,
"timestamp": time.time(),
"reason": "test",
},
}
)
router.health_state_cache = health_cache
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
result = await router._async_filter_health_check_unhealthy_deployments(deployments)
assert len(result) == 2
def _make_scoped_router_with_unhealthy(self, policy) -> Router:
import time
from litellm.caching.caching import DualCache
from litellm.router_utils.health_state_cache import DeploymentHealthCache
router = Router(
model_list=[
_make_model("bad-listed"),
_make_model("ok-listed"),
_make_model("bad-unlisted", "gpt-5"),
],
allowed_fails_policy=policy,
enable_health_check_routing=True,
background_health_check_model_groups=["gpt-4"],
)
cache = DualCache()
health_cache = DeploymentHealthCache(cache=cache, staleness_threshold=60.0)
health_cache.set_deployment_health_states(
{
model_id: {
"is_healthy": False,
"timestamp": time.time(),
"reason": "test",
}
for model_id in ("bad-listed", "bad-unlisted")
}
)
router.health_state_cache = health_cache
return router
def test_filter_with_policy_still_applies_to_listed_groups(self):
"""A model-group allowlist keeps the filter active for listed groups even with a policy set."""
router = self._make_scoped_router_with_unhealthy(AllowedFailsPolicy(AuthenticationErrorAllowedFails=3))
deployments = [
_make_model("bad-listed"),
_make_model("ok-listed"),
_make_model("bad-unlisted", "gpt-5"),
]
result = router._filter_health_check_unhealthy_deployments(deployments)
assert [d["model_info"]["id"] for d in result] == ["ok-listed", "bad-unlisted"]
@pytest.mark.asyncio
async def test_async_filter_with_policy_still_applies_to_listed_groups(self):
"""Async version: listed groups stay filtered with a policy set, unlisted stay untouched."""
router = self._make_scoped_router_with_unhealthy(AllowedFailsPolicy(TimeoutErrorAllowedFails=2))
deployments = [
_make_model("bad-listed"),
_make_model("ok-listed"),
_make_model("bad-unlisted", "gpt-5"),
]
result = await router._async_filter_health_check_unhealthy_deployments(deployments)
assert [d["model_info"]["id"] for d in result] == ["ok-listed", "bad-unlisted"]
class TestAllDeploymentsInCooldownSafetyNet:
"""
When enable_health_check_routing=True and ALL deployments enter cooldown,
the async routing path should bypass the cooldown filter and return all
deployments rather than blocking all traffic.
"""
def test_raw_cooldown_filter_returns_empty_when_all_cooled(self):
"""The raw _filter_cooldown_deployments has no safety net -- it returns empty."""
router = Router(
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
enable_health_check_routing=True,
)
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
result = router._filter_cooldown_deployments(
healthy_deployments=deployments,
cooldown_deployments=["deploy-1", "deploy-2"],
)
assert result == [] # raw filter has no safety net
@pytest.mark.asyncio
async def test_async_routing_path_bypasses_all_cooldown(self):
"""In the async routing path, all-in-cooldown with enable_health_check_routing
returns the full list instead of empty (safety net)."""
from unittest.mock import AsyncMock
from litellm.router_utils.cooldown_handlers import (
_async_get_cooldown_deployments,
)
router = Router(
model_list=[_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")],
allowed_fails_policy=AllowedFailsPolicy(AuthenticationErrorAllowedFails=0),
enable_health_check_routing=True,
)
deployments = [_make_model("deploy-1"), _make_model("deploy-2", "gpt-5")]
# Simulate all deployments in cooldown
with patch(
"litellm.router.async_get_cooldown_deployments",
new=AsyncMock(return_value=["deploy-1", "deploy-2"]),
):
# The safety net in async_get_available_deployment should restore
# all deployments when the cooldown filter empties the list
_pre = deployments.copy()
filtered = router._filter_cooldown_deployments(
healthy_deployments=deployments,
cooldown_deployments=["deploy-1", "deploy-2"],
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=[{"model_id": "deploy-1", "error": "rate limited"}],
exceptions_by_model_id={"deploy-1": rate_exc},
)
# If filtered is empty and enable_health_check_routing is True,
# the routing path restores _pre_cooldown_deployments
if not filtered and router.enable_health_check_routing:
filtered = _pre
assert len(filtered) == 2, "Safety net should return all deployments when all are in cooldown"
assert _failures(router, "deploy-1") == 1
assert _cooled(router) == []
@pytest.mark.parametrize(
("allowed_fails_policy", "model_groups", "routable"),
[
(None, None, ["deploy-2"]),
(AllowedFailsPolicy(AuthenticationErrorAllowedFails=3), None, ["deploy-1", "deploy-2"]),
(AllowedFailsPolicy(AuthenticationErrorAllowedFails=3), ["gpt-4"], ["deploy-2"]),
],
ids=["no-policy-filters", "policy-bypasses-filter", "policy-keeps-filter-for-listed-groups"],
)
@pytest.mark.asyncio
async def test_health_check_filter_is_bypassed_by_an_allowed_fails_policy(
allowed_fails_policy: AllowedFailsPolicy | None, model_groups: list[str] | None, routable: list[str]
) -> None:
"""With allowed_fails_policy set, cooldown is the sole routing exclusion, except for groups the
health check is scoped to."""
router = Router(
model_list=[_make_model("deploy-1"), _make_model("deploy-2")],
allowed_fails_policy=allowed_fails_policy,
enable_health_check_routing=True,
background_health_check_model_groups=model_groups,
)
router.health().health_checks.set_deployment_health_states(
{"deploy-1": {"is_healthy": False, "timestamp": time.time(), "reason": "test"}}
)
healthy = await router.select("gpt-4", {}, kind="healthy")
assert sorted(deployment["model_info"]["id"] for deployment in healthy) == routable
class TestHealthCheckIgnoreTransientErrors:
@ -614,17 +415,14 @@ class TestHealthCheckIgnoreTransientErrors:
]
with patch.object(proxy_module, "llm_router", router):
with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown:
with patch(
"litellm.router_utils.router_callbacks.track_deployment_metrics.increment_deployment_failures_for_current_minute"
) as mock_increment:
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=unhealthy_endpoints,
exceptions_by_model_id={"deploy-1": rate_exc},
)
mock_cooldown.assert_not_called()
mock_increment.assert_not_called()
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=unhealthy_endpoints,
exceptions_by_model_id={"deploy-1": rate_exc},
)
assert _failures(router, "deploy-1") == 0
assert _cooled(router) == []
def test_408_skipped_when_flag_enabled(self):
"""408 from health check does not trigger cooldown when flag is set."""
@ -645,13 +443,13 @@ class TestHealthCheckIgnoreTransientErrors:
]
with patch.object(proxy_module, "llm_router", router):
with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown:
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=unhealthy_endpoints,
exceptions_by_model_id={"deploy-1": timeout_exc},
)
mock_cooldown.assert_not_called()
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=unhealthy_endpoints,
exceptions_by_model_id={"deploy-1": timeout_exc},
)
assert _cooled(router) == []
def test_401_still_triggers_cooldown_when_flag_enabled(self):
"""Auth errors (401) still trigger cooldown even when flag is set."""
@ -672,13 +470,13 @@ class TestHealthCheckIgnoreTransientErrors:
]
with patch.object(proxy_module, "llm_router", router):
with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown:
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=unhealthy_endpoints,
exceptions_by_model_id={"deploy-1": auth_exc},
)
mock_cooldown.assert_called_once()
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=unhealthy_endpoints,
exceptions_by_model_id={"deploy-1": auth_exc},
)
assert _cooled(router) == ["deploy-1"]
def test_429_not_written_to_health_state_cache_when_flag_enabled(self):
"""429 endpoint is excluded from health state cache when flag is set,
@ -729,13 +527,13 @@ class TestHealthCheckIgnoreTransientErrors:
]
with patch.object(proxy_module, "llm_router", router):
with patch("litellm.router_utils.cooldown_handlers.set_cooldown_deployments") as mock_cooldown:
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=unhealthy_endpoints,
exceptions_by_model_id={"deploy-1": rate_exc},
)
mock_cooldown.assert_called_once()
_write_health_state_to_router_cache(
healthy_endpoints=[],
unhealthy_endpoints=unhealthy_endpoints,
exceptions_by_model_id={"deploy-1": rate_exc},
)
assert _cooled(router) == ["deploy-1"]
class TestSharedCacheTransientErrorFilter:

157
tests/unit/router_utils/test_kubernetes_pod_discovery.py
View file

@ -3,6 +3,7 @@ import copy
import hashlib
import json
import logging
import os
import re
import socket
import threading
@ -17,7 +18,6 @@ import respx
from openai import AsyncOpenAI, OpenAI
import litellm
import litellm.router_python_core
from litellm.constants import (
CALLER_SESSION_ID_METADATA_KEY,
KUBERNETES_POD_ROUTING_KEY,
@ -277,12 +277,11 @@ def _empty_proxy_environment() -> Mapping[str, str]:
return {}
def _patch_router_discovery(
monkeypatch: pytest.MonkeyPatch,
*discoveries: KubernetesPodDiscovery,
) -> None:
discovery_iter: Final = iter(discoveries)
monkeypatch.setattr(litellm.router_python_core, "KubernetesPodDiscovery", lambda: next(discovery_iter))
def _clear_proxy_environment(monkeypatch: pytest.MonkeyPatch) -> None:
"""Clear the proxy environment the router's own pod discovery reads, so it substitutes pod IPs."""
for name in tuple(os.environ):
if name.lower().endswith("_proxy"):
monkeypatch.delenv(name)
def _stub_sync_dns(monkeypatch: pytest.MonkeyPatch, *ips: str) -> None:
@ -1290,7 +1289,6 @@ def test_concurrent_sync_refresh_uses_cached_ips_or_service_url_while_in_flight(
async def test_router_sends_pod_hosts_without_forwarding_discovery_flag(
monkeypatch: pytest.MonkeyPatch,
) -> None:
clock: Final = _clock((0.0, 0.0, 0.0, 0.0, 0.0, 5.0, 5.0))
loop: Final = asyncio.get_running_loop()
async_lookups: Final = count()
sync_lookups: Final = count()
@ -1326,9 +1324,7 @@ async def test_router_sends_pod_hosts_without_forwarding_discovery_flag(
r"http://(?:10\.0\.0\.1|10\.0\.0\.2|vllm-headless\.ns\.svc\.cluster\.local):8000/v1/chat/completions"
)
).mock(return_value=httpx.Response(200, json=_CHAT_RESPONSE))
discovery_enabled: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment)
control: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment)
_patch_router_discovery(monkeypatch, discovery_enabled, control)
_clear_proxy_environment(monkeypatch)
discovery_enabled_router: Final = Router(model_list=[_deployment()])
control_router: Final = Router(model_list=[_deployment(kubernetes_pod_discovery=None)])
cached_async_client: Final = AsyncOpenAI(api_key="fake", base_url=_SERVICE_URL)
@ -1364,8 +1360,8 @@ async def test_router_sends_pod_hosts_without_forwarding_discovery_flag(
"vllm-headless.ns.svc.cluster.local",
)
assert all("kubernetes_pod_discovery" not in body for body in bodies)
assert next(async_lookups) == 1
assert next(sync_lookups) == 1
# The router's discovery resolves the service once; the sync call reuses the cached pod set.
assert (next(async_lookups), next(sync_lookups)) == (1, 0)
@pytest.mark.asyncio
@ -1405,9 +1401,7 @@ async def test_router_session_requests_stick_to_one_pod_for_sync_and_async(
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
)
async_discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
sync_discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
_patch_router_discovery(monkeypatch, async_discovery, sync_discovery)
_clear_proxy_environment(monkeypatch)
async_router: Final = Router(model_list=[_deployment()])
sync_router: Final = Router(model_list=[_deployment()])
async_client: Final = _cache_async_client(async_router)
@ -1451,8 +1445,7 @@ async def test_router_session_mapping_repeats_across_pods(
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
)
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
_patch_router_discovery(monkeypatch, discovery)
_clear_proxy_environment(monkeypatch)
router: Final = Router(model_list=[_deployment()])
client: Final = _cache_async_client(router)
for session_id in session_ids:
@ -1493,8 +1486,7 @@ async def test_generated_session_id_uses_round_robin_pod_selection(
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
)
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
_patch_router_discovery(monkeypatch, discovery)
_clear_proxy_environment(monkeypatch)
router: Final = Router(model_list=[_deployment()])
client: Final = _cache_async_client(router)
for _ in range(6):
@ -1533,8 +1525,7 @@ async def test_empty_session_id_uses_round_robin_pod_selection(
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
)
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
_patch_router_discovery(monkeypatch, discovery)
_clear_proxy_environment(monkeypatch)
router: Final = Router(model_list=[_deployment()])
client: Final = _cache_async_client(router)
for _ in range(6):
@ -1567,8 +1558,7 @@ async def test_session_requests_do_not_advance_round_robin_cursor(
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
)
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
_patch_router_discovery(monkeypatch, discovery)
_clear_proxy_environment(monkeypatch)
router: Final = Router(model_list=[_deployment()])
client: Final = _cache_async_client(router)
await _router_acompletion(router)
@ -1587,14 +1577,13 @@ async def test_session_requests_do_not_advance_round_robin_cursor(
)
@pytest.mark.asyncio
async def test_session_pod_membership_change_remaps_only_affected_sessions(
def test_session_pod_membership_change_remaps_only_affected_sessions(
monkeypatch: pytest.MonkeyPatch,
) -> None:
loop: Final = asyncio.get_running_loop()
lookups: Final = count()
now: Final = [0.0]
async def getaddrinfo(
def getaddrinfo(
host: str | None,
port: str | int | None,
family: int = 0,
@ -1606,44 +1595,34 @@ async def test_session_pod_membership_change_remaps_only_affected_sessions(
ips: Final = ("10.0.0.1", "10.0.0.2", "10.0.0.3") if lookup_number == 0 else ("10.0.0.2", "10.0.0.3")
return _records(*ips, port=int(port or 0))
monkeypatch.setattr(loop, "getaddrinfo", getaddrinfo)
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setattr(socket, "getaddrinfo", getaddrinfo)
discovery: Final = KubernetesPodDiscovery(
refresh_interval_seconds=10,
clock=lambda: now[0],
proxy_environment=_empty_proxy_environment,
)
def host_for(session_id: str) -> str | None:
return httpx.URL(
_api_base(discovery.resolve_deployment(_deployment(), {"metadata": {"session_id": session_id}}))
).host
session_ids: Final = tuple(f"session-{index}" for index in range(60))
with respx.mock(assert_all_called=True) as respx_mock:
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
)
discovery: Final = KubernetesPodDiscovery(
refresh_interval_seconds=10,
clock=_clock((0.0,) * 62 + (10.0,) * 10),
proxy_environment=_empty_proxy_environment,
)
_patch_router_discovery(monkeypatch, discovery)
router: Final = Router(model_list=[_deployment()])
client: Final = _cache_async_client(router)
for session_id in session_ids:
await _router_acompletion(router, metadata={"session_id": session_id})
initial_hosts: Final = tuple(call.request.url.host for call in route.calls)
removed_session_ids: Final = tuple(
session_id for session_id, host in zip(session_ids, initial_hosts) if host == "10.0.0.1"
)
surviving_session_ids: Final = tuple(
session_id for session_id, host in zip(session_ids, initial_hosts) if host == "10.0.0.2"
)
assert removed_session_ids
assert len(surviving_session_ids) >= 2
await _router_acompletion(router)
await _router_acompletion(router, metadata={"session_id": surviving_session_ids[0]})
await _router_acompletion(router, metadata={"session_id": removed_session_ids[0]})
await _router_acompletion(router, metadata={"session_id": surviving_session_ids[1]})
hosts: Final = tuple(call.request.url.host for call in route.calls)
await client.close()
initial_hosts: Final = {session_id: host_for(session_id) for session_id in session_ids}
removed_session_ids: Final = tuple(
session_id for session_id in session_ids if initial_hosts[session_id] == "10.0.0.1"
)
surviving_session_ids: Final = tuple(
session_id for session_id in session_ids if initial_hosts[session_id] in ("10.0.0.2", "10.0.0.3")
)
assert removed_session_ids
assert surviving_session_ids
assert hosts[60] == "10.0.0.1"
assert hosts[61] == "10.0.0.2"
assert hosts[62] in ("10.0.0.2", "10.0.0.3")
assert hosts[63] == "10.0.0.2"
now[0] = 10.0
remapped_hosts: Final = {session_id: host_for(session_id) for session_id in session_ids}
assert all(remapped_hosts[session_id] == initial_hosts[session_id] for session_id in surviving_session_ids)
assert all(remapped_hosts[session_id] in ("10.0.0.2", "10.0.0.3") for session_id in removed_session_ids)
assert next(lookups) == 2
@ -1684,15 +1663,7 @@ async def test_custom_routing_strategy_keeps_sync_and_async_sessions_on_one_pod(
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
)
async_discovery: Final = KubernetesPodDiscovery(
clock=lambda: 0.0,
proxy_environment=_empty_proxy_environment,
)
sync_discovery: Final = KubernetesPodDiscovery(
clock=lambda: 0.0,
proxy_environment=_empty_proxy_environment,
)
_patch_router_discovery(monkeypatch, async_discovery, sync_discovery)
_clear_proxy_environment(monkeypatch)
async_router: Final = Router(model_list=[_deployment()])
sync_router: Final = Router(model_list=[_deployment()])
async_router.set_custom_routing_strategy(_ModelListRoutingStrategy(async_router))
@ -1735,8 +1706,7 @@ async def test_top_level_litellm_session_id_keeps_requests_on_one_pod(
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
)
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
_patch_router_discovery(monkeypatch, discovery)
_clear_proxy_environment(monkeypatch)
router: Final = Router(model_list=[_deployment()])
client: Final = _cache_async_client(router)
for _ in range(6):
@ -1849,8 +1819,7 @@ async def test_generated_metadata_ignores_top_level_litellm_session_id(
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(
return_value=httpx.Response(200, json=_CHAT_RESPONSE)
)
discovery: Final = KubernetesPodDiscovery(clock=lambda: 0.0, proxy_environment=_empty_proxy_environment)
_patch_router_discovery(monkeypatch, discovery)
_clear_proxy_environment(monkeypatch)
router: Final = Router(model_list=[_deployment()])
client: Final = _cache_async_client(router)
for _ in range(6):
@ -1872,7 +1841,6 @@ async def test_generated_metadata_ignores_top_level_litellm_session_id(
async def test_custom_routing_strategy_resolves_pod_hosts_for_sync_and_async_requests(
monkeypatch: pytest.MonkeyPatch,
) -> None:
clock: Final = _clock((0.0,) * 40)
loop: Final = asyncio.get_running_loop()
async_lookups: Final = count()
sync_lookups: Final = count()
@ -1908,15 +1876,7 @@ async def test_custom_routing_strategy_resolves_pod_hosts_for_sync_and_async_req
r"http://(?:10\.0\.0\.1|10\.0\.0\.2|vllm-headless\.ns\.svc\.cluster\.local):8000/v1/chat/completions"
)
).mock(return_value=httpx.Response(200, json=_CHAT_RESPONSE))
async_discovery: Final = KubernetesPodDiscovery(
clock=clock,
proxy_environment=_empty_proxy_environment,
)
sync_discovery: Final = KubernetesPodDiscovery(
clock=clock,
proxy_environment=_empty_proxy_environment,
)
_patch_router_discovery(monkeypatch, async_discovery, sync_discovery)
_clear_proxy_environment(monkeypatch)
async_router: Final = Router(model_list=[_deployment()])
sync_router: Final = Router(model_list=[_deployment()])
async_router.set_custom_routing_strategy(_ModelListRoutingStrategy(async_router))
@ -1932,23 +1892,8 @@ async def test_custom_routing_strategy_resolves_pod_hosts_for_sync_and_async_req
assert next(sync_lookups) == 1
@pytest.mark.asyncio
async def test_unresolved_router_selectors_keep_service_hostname() -> None:
router: Final = Router(model_list=[_deployment()])
sync_selected: Final = router._get_available_deployment_unresolved(model="gpu-model", request_kwargs={})
async_selected: Final = await router._async_get_available_deployment_unresolved(
model="gpu-model",
request_kwargs={},
)
assert _api_base(sync_selected) == _SERVICE_URL
assert _api_base(async_selected) == _SERVICE_URL
@pytest.mark.asyncio
async def test_router_async_retry_uses_next_discovered_pod(monkeypatch: pytest.MonkeyPatch) -> None:
clock: Final = _clock((0.0,) * 8)
loop: Final = asyncio.get_running_loop()
attempts: Final = count()
@ -1974,8 +1919,7 @@ async def test_router_async_retry_uses_next_discovered_pod(monkeypatch: pytest.M
route: Final = respx_mock.post(
url__regex=re.compile(r"http://(?:10\.0\.0\.1|10\.0\.0\.2):8000/v1/chat/completions")
).mock(side_effect=response_for)
discovery: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment)
_patch_router_discovery(monkeypatch, discovery)
_clear_proxy_environment(monkeypatch)
router: Final = Router(model_list=[_deployment()], num_retries=1)
response: Final = await router.acompletion(
@ -1989,7 +1933,6 @@ async def test_router_async_retry_uses_next_discovered_pod(monkeypatch: pytest.M
def test_router_sync_retry_uses_next_discovered_pod(monkeypatch: pytest.MonkeyPatch) -> None:
clock: Final = _clock((0.0,) * 8)
attempts: Final = count()
def getaddrinfo(
@ -2014,8 +1957,7 @@ def test_router_sync_retry_uses_next_discovered_pod(monkeypatch: pytest.MonkeyPa
route: Final = respx_mock.post(
url__regex=re.compile(r"http://(?:10\.0\.0\.1|10\.0\.0\.2):8000/v1/chat/completions")
).mock(side_effect=response_for)
discovery: Final = KubernetesPodDiscovery(clock=clock, proxy_environment=_empty_proxy_environment)
_patch_router_discovery(monkeypatch, discovery)
_clear_proxy_environment(monkeypatch)
router: Final = Router(model_list=[_deployment()], num_retries=1)
response: Final = router.completion(
@ -2088,8 +2030,7 @@ async def test_router_session_retry_uses_next_rendezvous_pod(monkeypatch: pytest
session_id: Final = "retry-session"
with respx.mock(assert_all_called=True) as respx_mock:
route: Final = respx_mock.post(url__regex=_THREE_POD_CHAT_PATTERN).mock(side_effect=response_for)
discovery: Final = KubernetesPodDiscovery(proxy_environment=_empty_proxy_environment)
_patch_router_discovery(monkeypatch, discovery)
_clear_proxy_environment(monkeypatch)
router: Final = Router(model_list=[_deployment()], num_retries=1)
response: Final = await router.acompletion(

266
tests/unit/router_utils/test_router_health_check_routing.py
View file

@ -6,236 +6,70 @@ import time
import pytest
from litellm.caching.caching import DualCache
from litellm.router_utils.health_state_cache import DeploymentHealthCache
from litellm.router import Router
def _make_deployment(model_id: str, model_name: str = "gpt-4") -> dict:
"""Helper to create a deployment dict for testing."""
return {
"model_name": model_name,
"litellm_params": {"model": model_name, "api_key": "fake"},
"litellm_params": {"model": "openai/gpt-4", "api_key": "fake"},
"model_info": {"id": model_id},
}
def _make_health_cache(
unhealthy_ids: set = None, staleness_threshold: float = 60.0
) -> DeploymentHealthCache:
"""Create a health cache pre-populated with unhealthy deployment IDs."""
cache = DualCache()
health_cache = DeploymentHealthCache(
cache=cache, staleness_threshold=staleness_threshold
def _health_routed_router(enable: bool, unhealthy_ids: set[str], model_groups: list[str] | None = None) -> Router:
router = Router(
model_list=[
_make_deployment("deploy-1", model_name="prod"),
_make_deployment("deploy-2", model_name="prod"),
_make_deployment("deploy-3", model_name="prod"),
_make_deployment("other-1", model_name="other"),
_make_deployment("other-2", model_name="other"),
],
enable_health_check_routing=enable,
background_health_check_model_groups=model_groups,
)
if unhealthy_ids:
now = time.time()
states = {}
for uid in unhealthy_ids:
states[uid] = {
"is_healthy": False,
"timestamp": now,
"reason": "test_unhealthy",
}
health_cache.set_deployment_health_states(states)
return health_cache
now = time.time()
router.health().health_checks.set_deployment_health_states(
{uid: {"is_healthy": False, "timestamp": now, "reason": "test_unhealthy"} for uid in unhealthy_ids}
)
return router
class TestFilterHealthCheckUnhealthyDeployments:
"""Test the sync filter method."""
@pytest.mark.parametrize(
("enable", "unhealthy_ids", "model_groups", "routable"),
[
(True, {"deploy-2"}, None, {"prod": ["deploy-1", "deploy-3"], "other": ["other-1", "other-2"]}),
(False, {"deploy-2"}, None, {"prod": ["deploy-1", "deploy-2", "deploy-3"], "other": ["other-1", "other-2"]}),
(True, set(), None, {"prod": ["deploy-1", "deploy-2", "deploy-3"], "other": ["other-1", "other-2"]}),
(
True,
{"deploy-1", "deploy-2", "deploy-3"},
None,
{"prod": ["deploy-1", "deploy-2", "deploy-3"], "other": ["other-1", "other-2"]},
),
(True, {"deploy-2", "other-1"}, ["prod"], {"prod": ["deploy-1", "deploy-3"], "other": ["other-1", "other-2"]}),
(True, {"deploy-2", "other-1"}, None, {"prod": ["deploy-1", "deploy-3"], "other": ["other-2"]}),
],
ids=[
"removes-unhealthy",
"noop-when-disabled",
"all-pass-when-cache-empty",
"safety-net-when-all-unhealthy",
"scoped-to-listed-model-groups",
"unscoped-when-model-groups-unset",
],
)
@pytest.mark.asyncio
async def test_health_check_routing_filters_unhealthy_deployments(
enable: bool, unhealthy_ids: set[str], model_groups: list[str] | None, routable: dict[str, list[str]]
) -> None:
router = _health_routed_router(enable, unhealthy_ids, model_groups)
def _make_router_like(
self,
enable: bool,
health_cache: DeploymentHealthCache,
model_groups: frozenset[str] | None = None,
):
"""Create a minimal object that behaves like Router for filter testing."""
class FakeRouter:
def __init__(self):
self._core_enable_health_check_routing = enable
self._core_health_state_cache = health_cache
self._core_allowed_fails_policy = None
self._core_background_health_check_model_groups = model_groups
# Import the actual method and bind it
from litellm.router import Router
fake = FakeRouter()
# Use the unbound method
fake._filter_health_check_unhealthy_deployments = (
Router._filter_health_check_unhealthy_deployments.__get__(fake, FakeRouter)
)
return fake
def test_filter_removes_unhealthy_deployments(self):
"""Unhealthy deployments should be removed from candidates."""
health_cache = _make_health_cache(unhealthy_ids={"deploy-2"})
router = self._make_router_like(enable=True, health_cache=health_cache)
deployments = [
_make_deployment("deploy-1"),
_make_deployment("deploy-2"),
_make_deployment("deploy-3"),
]
result = router._filter_health_check_unhealthy_deployments(deployments)
assert len(result) == 2
assert all(d["model_info"]["id"] != "deploy-2" for d in result)
def test_filter_noop_when_disabled(self):
"""When enable_health_check_routing=False, filter should be a no-op."""
health_cache = _make_health_cache(unhealthy_ids={"deploy-1"})
router = self._make_router_like(enable=False, health_cache=health_cache)
deployments = [
_make_deployment("deploy-1"),
_make_deployment("deploy-2"),
]
result = router._filter_health_check_unhealthy_deployments(deployments)
assert len(result) == 2 # no filtering
def test_filter_returns_all_when_all_unhealthy(self):
"""Safety net: if ALL deployments are unhealthy, return all (don't cause outage)."""
health_cache = _make_health_cache(
unhealthy_ids={"deploy-1", "deploy-2", "deploy-3"}
)
router = self._make_router_like(enable=True, health_cache=health_cache)
deployments = [
_make_deployment("deploy-1"),
_make_deployment("deploy-2"),
_make_deployment("deploy-3"),
]
result = router._filter_health_check_unhealthy_deployments(deployments)
assert len(result) == 3 # all returned, safety net
def test_filter_returns_all_when_cache_empty(self):
"""When cache is empty, all deployments should pass through."""
health_cache = _make_health_cache() # empty
router = self._make_router_like(enable=True, health_cache=health_cache)
deployments = [
_make_deployment("deploy-1"),
_make_deployment("deploy-2"),
]
result = router._filter_health_check_unhealthy_deployments(deployments)
assert len(result) == 2
def test_filter_scoped_to_listed_model_groups(self):
"""With an allowlist, only deployments in listed groups are filtered on health."""
health_cache = _make_health_cache(unhealthy_ids={"bad-listed", "bad-unlisted"})
router = self._make_router_like(
enable=True, health_cache=health_cache, model_groups=frozenset({"prod"})
)
deployments = [
_make_deployment("bad-listed", model_name="prod"),
_make_deployment("ok-listed", model_name="prod"),
_make_deployment("bad-unlisted", model_name="other"),
_make_deployment("ok-unlisted", model_name="other"),
]
result = router._filter_health_check_unhealthy_deployments(deployments)
assert [d["model_info"]["id"] for d in result] == [
"ok-listed",
"bad-unlisted",
"ok-unlisted",
]
def test_filter_unscoped_when_model_groups_unset(self):
"""Without an allowlist, unhealthy deployments in every group are filtered."""
health_cache = _make_health_cache(unhealthy_ids={"bad-listed", "bad-unlisted"})
router = self._make_router_like(enable=True, health_cache=health_cache)
deployments = [
_make_deployment("bad-listed", model_name="prod"),
_make_deployment("ok-listed", model_name="prod"),
_make_deployment("bad-unlisted", model_name="other"),
_make_deployment("ok-unlisted", model_name="other"),
]
result = router._filter_health_check_unhealthy_deployments(deployments)
assert [d["model_info"]["id"] for d in result] == ["ok-listed", "ok-unlisted"]
class TestAsyncFilterHealthCheckUnhealthyDeployments:
"""Test the async filter method."""
def _make_router_like(
self,
enable: bool,
health_cache: DeploymentHealthCache,
model_groups: frozenset[str] | None = None,
):
from litellm.router import Router
class FakeRouter:
def __init__(self):
self._core_enable_health_check_routing = enable
self._core_health_state_cache = health_cache
self._core_allowed_fails_policy = None
self._core_background_health_check_model_groups = model_groups
fake = FakeRouter()
fake._async_filter_health_check_unhealthy_deployments = (
Router._async_filter_health_check_unhealthy_deployments.__get__(
fake, FakeRouter
)
)
return fake
@pytest.mark.asyncio
async def test_async_filter_removes_unhealthy(self):
"""Async version: unhealthy deployments removed."""
health_cache = _make_health_cache(unhealthy_ids={"deploy-2"})
router = self._make_router_like(enable=True, health_cache=health_cache)
deployments = [
_make_deployment("deploy-1"),
_make_deployment("deploy-2"),
_make_deployment("deploy-3"),
]
result = await router._async_filter_health_check_unhealthy_deployments(
healthy_deployments=deployments
)
assert len(result) == 2
assert all(d["model_info"]["id"] != "deploy-2" for d in result)
@pytest.mark.asyncio
async def test_async_filter_safety_net(self):
"""Async version: safety net when all unhealthy."""
health_cache = _make_health_cache(unhealthy_ids={"deploy-1", "deploy-2"})
router = self._make_router_like(enable=True, health_cache=health_cache)
deployments = [
_make_deployment("deploy-1"),
_make_deployment("deploy-2"),
]
result = await router._async_filter_health_check_unhealthy_deployments(
healthy_deployments=deployments
)
assert len(result) == 2 # safety net
@pytest.mark.asyncio
async def test_async_filter_scoped_to_listed_model_groups(self):
"""Async version: only deployments in listed groups are filtered on health."""
health_cache = _make_health_cache(unhealthy_ids={"bad-listed", "bad-unlisted"})
router = self._make_router_like(
enable=True, health_cache=health_cache, model_groups=frozenset({"prod"})
)
deployments = [
_make_deployment("bad-listed", model_name="prod"),
_make_deployment("ok-listed", model_name="prod"),
_make_deployment("bad-unlisted", model_name="other"),
_make_deployment("ok-unlisted", model_name="other"),
]
result = await router._async_filter_health_check_unhealthy_deployments(
healthy_deployments=deployments
)
assert [d["model_info"]["id"] for d in result] == [
"ok-listed",
"bad-unlisted",
"ok-unlisted",
]
for model_group, expected in routable.items():
healthy = await router.select(model_group, {}, kind="healthy")
assert sorted(d["model_info"]["id"] for d in healthy) == expected, model_group
class TestBuildDeploymentHealthStates:

277
tests/unit/router_utils/test_router_interactions_endpoints.py
View file

@ -2,10 +2,11 @@
Tests for Router interactions API endpoint initialization functions.
"""
from unittest.mock import AsyncMock, MagicMock, patch
from unittest.mock import AsyncMock
import pytest
import litellm
from litellm import Router
@ -39,191 +40,6 @@ class TestInitializeInteractionsEndpoints:
assert callable(router.aget_interaction)
assert callable(router.get_interaction)
def test_initialize_interactions_endpoints_can_be_called_directly(self):
"""Test that _initialize_interactions_endpoints can be called directly to reinitialize endpoints."""
router = Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4"},
}
]
)
# Call _initialize_interactions_endpoints directly
router._initialize_interactions_endpoints()
# Verify the interaction methods still exist after re-initialization
assert hasattr(router, "acreate_interaction")
assert hasattr(router, "create_interaction")
assert callable(router.acreate_interaction)
class TestInitInteractionsApiEndpoints:
"""Test cases for _init_interactions_api_endpoints method"""
@pytest.mark.asyncio
async def test_init_interactions_api_endpoints_passes_custom_llm_provider(self):
"""Test that _init_interactions_api_endpoints passes custom_llm_provider to the original function."""
router = Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4"},
}
]
)
mock_function = AsyncMock(return_value={"result": "success"})
result = await router._init_interactions_api_endpoints(
original_function=mock_function,
custom_llm_provider="gemini",
interaction_id="test-id",
)
mock_function.assert_called_once_with(
custom_llm_provider="gemini",
interaction_id="test-id",
)
assert result == {"result": "success"}
@pytest.mark.asyncio
async def test_init_interactions_api_endpoints_defaults_to_gemini(self):
"""Test that _init_interactions_api_endpoints defaults to gemini when no custom_llm_provider is specified."""
router = Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4"},
}
]
)
mock_function = AsyncMock(return_value={"result": "success"})
result = await router._init_interactions_api_endpoints(
original_function=mock_function,
interaction_id="test-id",
)
mock_function.assert_called_once_with(
custom_llm_provider="gemini",
interaction_id="test-id",
)
assert result == {"result": "success"}
@pytest.mark.asyncio
async def test_init_interactions_api_endpoints_does_not_override_existing_provider(
self,
):
"""Test that _init_interactions_api_endpoints does not override custom_llm_provider if already in kwargs."""
router = Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4"},
}
]
)
mock_function = AsyncMock(return_value={"result": "success"})
# Pass custom_llm_provider in kwargs directly (not as separate param)
result = await router._init_interactions_api_endpoints(
original_function=mock_function,
custom_llm_provider="vertex_ai",
)
# Should use the provided custom_llm_provider
mock_function.assert_called_once_with(
custom_llm_provider="vertex_ai",
)
assert result == {"result": "success"}
@pytest.mark.asyncio
async def test_init_interactions_api_endpoints_clears_model_when_equals_agent(
self,
):
"""Managed agent interactions must not pass agent name as model to the SDK."""
router = Router(model_list=[])
mock_function = AsyncMock(return_value={"result": "success"})
await router._init_interactions_api_endpoints(
original_function=mock_function,
agent="mqy-custom-slides-agent",
model="mqy-custom-slides-agent",
input="hello",
)
mock_function.assert_called_once_with(
custom_llm_provider="gemini",
agent="mqy-custom-slides-agent",
model=None,
input="hello",
)
class TestRouterCreateInteractionRouting:
"""acreate_interaction routing: agent-only vs model + fallbacks."""
@pytest.mark.asyncio
async def test_acreate_interaction_agent_only_uses_init_interactions(self):
"""Agent-only create must not use model-group fallback lookup."""
router = Router(
model_list=[
{
"model_name": "openai/*",
"litellm_params": {"model": "gpt-4"},
}
]
)
with (
patch.object(
router,
"_init_interactions_api_endpoints",
new_callable=AsyncMock,
return_value={"id": "int-1"},
) as mock_init,
patch.object(
router,
"_ageneric_api_call_with_fallbacks",
new_callable=AsyncMock,
) as mock_generic,
):
result = await router.acreate_interaction(
agent="mqy-custom-slides-agent",
input="hello",
custom_llm_provider="gemini",
)
mock_init.assert_called_once()
mock_generic.assert_not_called()
assert result == {"id": "int-1"}
@pytest.mark.asyncio
async def test_init_interactions_model_uses_generic_fallbacks(self):
"""Model-based create uses _ageneric_api_call_with_fallbacks inside _init_interactions."""
router = Router(model_list=[])
with patch.object(
router,
"_ageneric_api_call_with_fallbacks",
new_callable=AsyncMock,
return_value={"id": "int-1"},
) as mock_generic:
result = await router._init_interactions_api_endpoints(
original_function=AsyncMock(),
model="gemini-2.5-flash",
input="hello",
custom_llm_provider="gemini",
)
mock_generic.assert_called_once()
assert result == {"id": "int-1"}
class TestInitializeManagedAgentsEndpoints:
"""Tests for _initialize_managed_agents_endpoints."""
@ -248,51 +64,60 @@ class TestInitializeManagedAgentsEndpoints:
assert hasattr(router, method_name), f"missing {method_name}"
assert callable(getattr(router, method_name)), f"{method_name} not callable"
def test_initialize_managed_agents_endpoints_can_be_called_directly(self):
router = Router(model_list=[])
router._initialize_managed_agents_endpoints()
assert callable(router.acreate_agent)
assert callable(router.alist_agents)
@pytest.mark.parametrize(
("request_kwargs", "provider"),
[({}, "gemini"), ({"custom_llm_provider": "vertex_ai"}, "vertex_ai")],
ids=["defaults-to-gemini", "keeps-the-caller-provider"],
)
@pytest.mark.asyncio
async def test_interaction_endpoints_default_the_provider_to_gemini(
monkeypatch: pytest.MonkeyPatch, request_kwargs: dict[str, str], provider: str
) -> None:
provider_call = AsyncMock(return_value={"result": "success"})
monkeypatch.setattr(litellm.interactions, "aget", provider_call)
router = Router(model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4"}}])
result = await router.aget_interaction(interaction_id="test-id", **request_kwargs)
assert result == {"result": "success"}
assert provider_call.call_args.kwargs["custom_llm_provider"] == provider
assert provider_call.call_args.kwargs["interaction_id"] == "test-id"
class TestInitManagedAgentsApiEndpoints:
"""Tests for _init_managed_agents_api_endpoints."""
@pytest.mark.asyncio
async def test_agent_only_interaction_does_not_send_the_agent_name_as_the_model(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Managed agent interactions must not pass agent name as model to the SDK."""
provider_call = AsyncMock(return_value={"id": "int-1"})
monkeypatch.setattr(litellm.interactions, "acreate", provider_call)
router = Router(model_list=[{"model_name": "openai/*", "litellm_params": {"model": "gpt-4"}}])
@pytest.mark.asyncio
async def test_init_managed_agents_api_endpoints_defaults_to_gemini(self):
router = Router(model_list=[])
mock_fn = AsyncMock(return_value={"agents": []})
result = await router.acreate_interaction(
agent="mqy-custom-slides-agent", model="mqy-custom-slides-agent", input="hello"
)
await router._init_managed_agents_api_endpoints(
original_function=mock_fn,
)
assert result == {"id": "int-1"}
assert provider_call.call_args.kwargs["agent"] == "mqy-custom-slides-agent"
assert provider_call.call_args.kwargs.get("model") is None
assert provider_call.call_args.kwargs["custom_llm_provider"] == "gemini"
call_kwargs = mock_fn.call_args.kwargs
assert call_kwargs["custom_llm_provider"] == "gemini"
@pytest.mark.asyncio
async def test_init_managed_agents_api_endpoints_passes_custom_provider(self):
router = Router(model_list=[])
mock_fn = AsyncMock(return_value={"agents": []})
@pytest.mark.parametrize(
("request_kwargs", "provider"),
[({}, "gemini"), ({"custom_llm_provider": "vertex_ai"}, "vertex_ai")],
ids=["defaults-to-gemini", "keeps-the-caller-provider"],
)
@pytest.mark.asyncio
async def test_managed_agent_endpoints_default_the_provider_to_gemini(
monkeypatch: pytest.MonkeyPatch, request_kwargs: dict[str, str], provider: str
) -> None:
provider_call = AsyncMock(return_value={"agents": []})
monkeypatch.setattr(litellm.interactions.agents, "alist", provider_call)
router = Router(model_list=[])
await router._init_managed_agents_api_endpoints(
original_function=mock_fn,
custom_llm_provider="vertex_ai",
)
result = await router.alist_agents(**request_kwargs)
call_kwargs = mock_fn.call_args.kwargs
assert call_kwargs["custom_llm_provider"] == "vertex_ai"
@pytest.mark.asyncio
async def test_init_managed_agents_api_endpoints_does_not_override_existing_provider(
self,
):
router = Router(model_list=[])
mock_fn = AsyncMock(return_value={"agents": []})
await router._init_managed_agents_api_endpoints(
original_function=mock_fn,
custom_llm_provider="vertex_ai",
)
mock_fn.assert_called_once_with(custom_llm_provider="vertex_ai")
assert result == {"agents": []}
assert provider_call.call_args.kwargs["custom_llm_provider"] == provider

114
tests/unit/router_utils/test_router_utils_common_utils.py
View file

@ -54,16 +54,12 @@ class TestFilterTeamBasedModels:
{"model_info": {"id": "deployment-2"}},
]
def test_filter_team_based_models_none_request_kwargs(
self, sample_deployments_with_teams
):
def test_filter_team_based_models_none_request_kwargs(self, sample_deployments_with_teams):
"""Test that when request_kwargs is None, all deployments are returned unchanged"""
result = filter_team_based_models(sample_deployments_with_teams, None)
assert result == sample_deployments_with_teams
def test_filter_team_based_models_empty_request_kwargs(
self, sample_deployments_with_teams
):
def test_filter_team_based_models_empty_request_kwargs(self, sample_deployments_with_teams):
"""Test with empty request_kwargs"""
result = filter_team_based_models(sample_deployments_with_teams, {})
# Should include all deployments since no team_id in request
@ -76,9 +72,7 @@ class TestFilterTeamBasedModels:
# Should include only non-team based deployments
assert len(result) == 1
def test_filter_team_based_models_team_match_metadata(
self, sample_deployments_with_teams
):
def test_filter_team_based_models_team_match_metadata(self, sample_deployments_with_teams):
"""Test filtering when team_id is in metadata"""
request_kwargs = {"metadata": {"user_api_key_team_id": "team-a"}}
result = filter_team_based_models(sample_deployments_with_teams, request_kwargs)
@ -93,9 +87,7 @@ class TestFilterTeamBasedModels:
result_ids = [d.get("model_info", {}).get("id") for d in result]
assert sorted(result_ids) == sorted(expected_ids)
def test_filter_team_based_models_team_match_litellm_metadata(
self, sample_deployments_with_teams
):
def test_filter_team_based_models_team_match_litellm_metadata(self, sample_deployments_with_teams):
"""Test filtering when team_id is in litellm_metadata"""
request_kwargs = {"litellm_metadata": {"user_api_key_team_id": "team-b"}}
result = filter_team_based_models(sample_deployments_with_teams, request_kwargs)
@ -110,9 +102,7 @@ class TestFilterTeamBasedModels:
result_ids = [d.get("model_info", {}).get("id") for d in result]
assert sorted(result_ids) == sorted(expected_ids)
def test_filter_team_based_models_priority_metadata_over_litellm(
self, sample_deployments_with_teams
):
def test_filter_team_based_models_priority_metadata_over_litellm(self, sample_deployments_with_teams):
"""Test that metadata.user_api_key_team_id takes priority over litellm_metadata.user_api_key_team_id"""
request_kwargs = {
"metadata": {
@ -127,9 +117,7 @@ class TestFilterTeamBasedModels:
result_ids = [d.get("model_info", {}).get("id") for d in result]
assert sorted(result_ids) == sorted(expected_ids)
def test_filter_team_based_models_no_matching_team(
self, sample_deployments_with_teams
):
def test_filter_team_based_models_no_matching_team(self, sample_deployments_with_teams):
"""Test when request team doesn't match any deployment teams"""
request_kwargs = {"metadata": {"user_api_key_team_id": "team-nonexistent"}}
result = filter_team_based_models(sample_deployments_with_teams, request_kwargs)
@ -139,9 +127,7 @@ class TestFilterTeamBasedModels:
result_ids = [d.get("model_info", {}).get("id") for d in result]
assert result_ids == expected_ids
def test_filter_team_based_models_no_team_restrictions(
self, sample_deployments_no_teams
):
def test_filter_team_based_models_no_team_restrictions(self, sample_deployments_no_teams):
"""Test with deployments that have no team restrictions"""
request_kwargs = {"metadata": {"user_api_key_team_id": "any-team"}}
result = filter_team_based_models(sample_deployments_no_teams, request_kwargs)
@ -188,9 +174,7 @@ class TestFilterTeamBasedModels:
def test_filter_team_based_models_empty_deployments(self):
"""Test with empty deployments list"""
result = filter_team_based_models(
[], {"metadata": {"user_api_key_team_id": "team-a"}}
)
result = filter_team_based_models([], {"metadata": {"user_api_key_team_id": "team-a"}})
assert result == []
def test_filter_team_based_models_none_team_id_in_deployment(self):
@ -347,9 +331,7 @@ class TestFilterWebSearchDeployments:
deployments = [
{"model_info": {"id": "d1"}}, # No supports_web_search - defaults to True
{"model_info": {"id": "d2"}}, # No supports_web_search - defaults to True
{
"model_info": {"id": "d3", "supports_web_search": False}
}, # Explicit False
{"model_info": {"id": "d3", "supports_web_search": False}}, # Explicit False
]
request_kwargs = {"tools": [{"type": "web_search"}]}
result = filter_web_search_deployments(deployments, request_kwargs)
@ -375,26 +357,8 @@ class TestFilterWebSearchDeployments:
assert result == deployment
def test_invalidate_model_group_info_cache():
"""Test that _invalidate_model_group_info_cache clears the LRU cache."""
router = Router(
model_list=[
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-4", "api_key": "fake-key"},
}
]
)
# Populate the cache
router._cached_get_model_group_info("gpt-4")
assert router._cached_get_model_group_info.cache_info().currsize > 0
# Invalidate and verify cache is cleared
router._invalidate_model_group_info_cache()
assert router._cached_get_model_group_info.cache_info().currsize == 0
def test_filter_deployments_by_model_access_groups_access_group_only_key():
@pytest.mark.asyncio
async def test_filter_deployments_by_model_access_groups_access_group_only_key():
"""
Access-group-only keys should only route to deployments in allowed groups,
even when multiple deployments share the same public model name.
@ -403,13 +367,13 @@ def test_filter_deployments_by_model_access_groups_access_group_only_key():
model_list=[
{
"model_name": "gpt-5",
"litellm_params": {"model": "openai/gpt-5.1", "api_key": "key-1"},
"model_info": {"access_groups": ["AG1"]},
"litellm_params": {"model": "openai/gpt-5.1", "api_key": "key-1", "mock_response": "ag1"},
"model_info": {"id": "ag1", "access_groups": ["AG1"]},
},
{
"model_name": "gpt-5",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "key-2"},
"model_info": {"access_groups": ["AG2"]},
"litellm_params": {"model": "openai/gpt-4o", "api_key": "key-2", "mock_response": "ag2"},
"model_info": {"id": "ag2", "access_groups": ["AG2"]},
},
]
)
@ -420,21 +384,20 @@ def test_filter_deployments_by_model_access_groups_access_group_only_key():
models=["AG2"],
team_models=["AG2"],
)
metadata = {"user_api_key_team_id": "team-2", "user_api_key_auth": scoped_key}
filtered = router._filter_deployments_by_model_access_groups(
model="gpt-5",
healthy_deployments=router.get_all_deployments(model_name="gpt-5"),
request_kwargs={
"metadata": {
"user_api_key_team_id": "team-2",
"user_api_key_auth": scoped_key,
}
},
request_team_id="team-2",
)
healthy = await router.select("gpt-5", {"metadata": metadata}, kind="healthy")
served = {
(
await router.acompletion(
model="gpt-5", messages=[{"role": "user", "content": "hi"}], metadata=dict(metadata)
)
)._hidden_params["model_id"]
for _ in range(5)
}
assert len(filtered) == 1
assert filtered[0].get("model_info", {}).get("access_groups") == ["AG2"]
assert [deployment["model_info"]["id"] for deployment in healthy] == ["ag2"]
assert served == {"ag2"}
class TestAddModelFileIdMappings:
@ -734,12 +697,7 @@ class TestWarnOnProviderCredentialMismatch:
)
def test_silent_when_no_provider_scoped_credentials_are_set(self):
assert (
warn_on_provider_credential_mismatch(
model_name="gpt-5.5", litellm_params={"model": "gpt-5.5"}
)
is None
)
assert warn_on_provider_credential_mismatch(model_name="gpt-5.5", litellm_params={"model": "gpt-5.5"}) is None
def test_vertex_params_name_vertex_not_bedrock(self):
"""The hint must follow the params that were actually set, otherwise it
@ -789,8 +747,7 @@ class TestWarnOnProviderCredentialMismatch:
mismatch_warnings = [r for r in caplog.records if "resolves to provider" in r.getMessage()]
assert len(mismatch_warnings) == 1, (
"exactly the prefix-less deployment should warn; "
f"got {[r.getMessage() for r in mismatch_warnings]}"
f"exactly the prefix-less deployment should warn; got {[r.getMessage() for r in mismatch_warnings]}"
)
assert "aws_region_name" in mismatch_warnings[0].getMessage()
@ -867,7 +824,14 @@ class TestWarnOnProviderCredentialMismatch:
({"model": "no-provider-knows-this-model"}, None),
({"api_base": "https://my-resource.openai.azure.com"}, None),
],
ids=["declared_wins", "prefix_beats_host_flip", "prefix_beats_cohere_chat_flip", "unprefixed_inferred", "unknown", "no_model"],
ids=[
"declared_wins",
"prefix_beats_host_flip",
"prefix_beats_cohere_chat_flip",
"unprefixed_inferred",
"unknown",
"no_model",
],
)
def test_provider_for_generic_call(litellm_params, expected, monkeypatch):
monkeypatch.setenv("AZURE_AI_API_BASE", "https://unrelated.openai.azure.com")
@ -903,7 +867,9 @@ def test_router_warns_only_for_a_capability_flag_set_under_litellm_params(caplog
def test_router_serves_deployments_without_model_capability_flags():
capability_flags: Final = frozenset(ProviderSpecificModelInfo.__annotations__) | frozenset(
key for key in chain.from_iterable(GetModelCostMap.loaded_model_cost_map().values()) if key.startswith("supports_")
key
for key in chain.from_iterable(GetModelCostMap.loaded_model_cost_map().values())
if key.startswith("supports_")
)
declared_params: Final = {"supports_system_message": False, "supports_system_prompt": True}
router: Final = Router(

13
tests/unit/test_conftest_isolation.py
View file

@ -9,11 +9,10 @@ import pytest
import litellm
from litellm import Router
from litellm import router as litellm_router_module
from litellm import utils as litellm_utils_module
from litellm.caching.caching import DualCache
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from tests.unit import conftest as unit_harness
CANARY_MODEL = "conftest-isolation-canary-model"
@ -117,12 +116,12 @@ def test_live_router_membership_is_scoped_to_this_test():
}
]
)
assert _CanaryRouterHolder.router._core in litellm_router_module._live_routers
assert _CanaryRouterHolder.router._core._hosted in unit_harness.live_router_registries
def test_live_router_membership_was_rolled_back():
assert _CanaryRouterHolder.router is not None
assert _CanaryRouterHolder.router._core not in litellm_router_module._live_routers
assert _CanaryRouterHolder.router._core._hosted not in unit_harness.live_router_registries
def test_aws_cache_reset_sees_new_handlers_and_replaced_caches(monkeypatch: pytest.MonkeyPatch) -> None:
@ -160,7 +159,7 @@ def test_aws_cache_reset_sees_a_replaced_module(monkeypatch: pytest.MonkeyPatch)
@pytest.fixture(scope="module", autouse=True)
def registered_collection_models() -> Iterator[None]:
original_registrations: Final = dict(litellm_utils_module._runtime_registered_model_cost)
original_routers: Final = frozenset(litellm_router_module._live_routers)
original_routers: Final = frozenset(unit_harness.live_router_registries)
litellm.register_model(
{"collection-isolation-ledger": {"litellm_provider": "openai", "input_cost_per_token": 0.001}}
)
@ -175,8 +174,8 @@ def registered_collection_models() -> Iterator[None]:
yield
litellm_utils_module._runtime_registered_model_cost.clear()
litellm_utils_module._runtime_registered_model_cost.update(original_registrations)
litellm_router_module._live_routers.discard(collection_router._core)
litellm_router_module._live_routers.update(original_routers)
unit_harness.live_router_registries.discard(collection_router._core._hosted)
unit_harness.live_router_registries.update(original_routers)
def test_collection_model_metadata_is_not_replayed_into_a_test_catalog() -> None:

1
tests/unit/test_private_usage_aliases.py
View file

@ -1949,7 +1949,6 @@ ALIAS_CASES: Final = (
"is_model_access_group_for_wildcard_route",
False,
),
("litellm.router", "Router", "_replay_model_cost_registrations", "replay_model_cost_registrations", False),
("litellm.router", "Router", "_routing_groups", "routing_groups", True),
("litellm.router", "Router", "_update_redis_cache", "update_redis_cache", False),
(

80
tests/unit/test_redact_string_in_error_paths.py
View file

@ -39,9 +39,7 @@ class TestRedactStringFunction:
text = "This is a normal error message with no secrets"
assert redact_string(text) == text
@pytest.mark.skipif(
not _ENABLE_SECRET_REDACTION, reason="redaction disabled via env var"
)
@pytest.mark.skipif(not _ENABLE_SECRET_REDACTION, reason="redaction disabled via env var")
def test_redaction_enabled_by_default(self):
text = "Bearer sk-9876567890abcdefghij"
result = redact_string(text)
@ -77,9 +75,7 @@ class TestOpenAIRealtimeRedaction:
from litellm.llms.openai.realtime.handler import OpenAIRealtime
handler = OpenAIRealtime()
secret_error = RuntimeError(
"Connection failed for api_key=sk-9876567890abcdefghij"
)
secret_error = RuntimeError("Connection failed for api_key=sk-9876567890abcdefghij")
kwargs = self._call_kwargs()
mock_ws = kwargs["websocket"]
@ -96,9 +92,7 @@ class TestBedrockRealtimeRedaction:
"""Test that redact_string produces safe close reasons for Bedrock-style errors."""
def test_internal_error_message_redacted(self):
secret_error = RuntimeError(
"Failed with aws_secret_access_key=AKIAIOSFODNN7EXAMPLE123456"
)
secret_error = RuntimeError("Failed with aws_secret_access_key=AKIAIOSFODNN7EXAMPLE123456")
reason = redact_string(f"Internal error: {str(secret_error)}")
assert "AKIAIOSFODNN7EXAMPLE123456" not in reason
@ -112,17 +106,13 @@ class TestLLMHTTPHandlerRealtimeRedaction:
def test_internal_server_error_pattern(self):
error_msg = "Connection failed for api_key=sk-secret-key-12345678"
assert "sk-secret-key-12345678" not in redact_string(
f"Internal server error: {error_msg}"
)
assert "sk-secret-key-12345678" not in redact_string(f"Internal server error: {error_msg}")
class TestProxyStreamingDataGeneratorRedaction:
def test_redact_traceback_format_exc(self):
try:
raise RuntimeError(
"Failed connecting to api_key=sk-9876567890abcdefghij at https://api.example.com"
)
raise RuntimeError("Failed connecting to api_key=sk-9876567890abcdefghij at https://api.example.com")
except RuntimeError:
raw_tb = traceback.format_exc()
@ -134,8 +124,7 @@ class TestProxyStreamingDataGeneratorRedaction:
class TestRouterFallbackFailureTracebackRedaction:
"""Test the fallback-failure logs in router.py's
async_function_with_fallbacks_common_utils. Both call sites must redact the
"""Test the router's fallback-failure logs. Both call sites must redact the
traceback at the call site with redact_string() rather than hand a live
exception to exc_info=True. SecretRedactionFilter rewrites record.exc_text,
but record.exc_info stays an exception object no filter can rewrite, so any
@ -144,57 +133,44 @@ class TestRouterFallbackFailureTracebackRedaction:
@pytest.mark.asyncio
async def test_fallback_failure_does_not_leak_secret_via_exc_info(self, caplog):
"""The helper is driven from inside an `except` block because that is the only
way production reaches it, and the entry-point debug log takes its traceback
from the active exception. With no exception in flight sys.exc_info() is empty,
so an exc_info=True regression there would degrade to (None, None, None) and
this test would pass against it.
"""
"""Both the primary group and its fallback fail with the secret in the error, so the
router logs the fallback failure while that exception is in flight."""
import litellm
secret = "sk-testsecretvalue1234567890abcdef"
router = litellm.Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key"},
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": "fake-key",
"mock_response": Exception(f"primary deployment failed api_key={secret}"),
},
},
{
"model_name": "claude-3-haiku",
"litellm_params": {"model": "anthropic/claude-3-haiku-20240307", "api_key": "fake-key"},
"litellm_params": {
"model": "anthropic/claude-3-haiku-20240307",
"api_key": "fake-key",
"mock_response": Exception(f"boom api_key={secret}"),
},
},
],
fallbacks=[{"gpt-3.5-turbo": ["claude-3-haiku"]}],
)
secret = "sk-testsecretvalue1234567890abcdef"
with patch(
"litellm.router.run_async_fallback",
new=AsyncMock(side_effect=RuntimeError(f"boom api_key={secret}")),
):
try:
raise ValueError(f"primary deployment failed api_key={secret}")
except ValueError as original_exception:
with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"):
with pytest.raises(ValueError, match='primary deployment failed api_key=sk-testsecretvalu'):
await router.async_function_with_fallbacks_common_utils(
e=original_exception,
disable_fallbacks=False,
fallbacks=[{"gpt-3.5-turbo": ["claude-3-haiku"]}],
context_window_fallbacks=None,
content_policy_fallbacks=None,
model_group="gpt-3.5-turbo",
args=(),
kwargs={"model": "gpt-3.5-turbo"},
)
with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"):
with pytest.raises(litellm.InternalServerError):
await router.acompletion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "hi"}])
debug_records = [r for r in caplog.records if r.levelno == logging.DEBUG]
assert debug_records, "expected the entry-point debug log, which carries the active traceback"
error_records = [r for r in caplog.records if r.levelno == logging.ERROR]
assert error_records, "expected an error log for the fallback failure"
assert any(
"Cooldown Deployments" in r.getMessage() for r in error_records
), "expected the fallback-failure log, not an unrelated error"
assert any("Error occurred while trying to do fallbacks" in r.getMessage() for r in error_records), (
"expected the fallback-failure log, not an unrelated error"
)
for record in caplog.records:
assert secret not in record.getMessage()
@ -255,9 +231,7 @@ class TestGeminiIngestionHeaders:
mock_client = AsyncMock()
mock_response = MagicMock()
mock_response.status_code = 200
mock_response.headers = {
"x-goog-upload-url": "https://upload.example.com/upload123"
}
mock_response.headers = {"x-goog-upload-url": "https://upload.example.com/upload123"}
mock_client.post.return_value = mock_response
with patch(

240
tests/unit/test_responses_streaming_container_ownership.py
View file

@ -1,9 +1,8 @@
"""
Regression for #30210.
When streaming /v1/responses goes through the proxy + Router, the
streaming iterator is wrapped by ``Router._aresponses_streaming_iterator``
which returns ``FallbackResponsesStreamWrapper``. That wrapper set
When streaming /v1/responses goes through the proxy + Router, the router
returns a ``FallbackResponsesStreamWrapper``. That wrapper set
``self.completed_response = None`` in __init__ and never updated it,
so the proxy's container-ownership hook (which reads
``getattr(stream_response, "completed_response", None)`` via
@ -12,91 +11,81 @@ saw None on every streaming call and silently recorded nothing —
follow-up ``GET /v1/containers/<id>/files`` then 403'd for the very
key that created the container.
Tests below construct the wrapper from a fake async generator that
yields one terminal ``response.completed`` chunk and assert the
wrapper now carries that chunk on ``completed_response`` so the
proxy hook can walk it.
Tests below stream a mocked upstream through ``Router.aresponses`` and
assert the returned stream carries the terminal chunk on
``completed_response`` so the proxy hook can walk it.
"""
import asyncio
import json
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
import httpx
import respx
import litellm
_DELTA = {
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": "msg_1",
"output_index": 0,
"content_index": 0,
"delta": "hello",
"logprobs": [],
}
def _make_wrapper_class():
"""Pull ``FallbackResponsesStreamWrapper`` out by running
``Router._aresponses_streaming_iterator`` long enough to construct
the class then return it. Mirrors how the wrapper is actually
instantiated in production."""
from litellm.router import Router
from litellm.responses.streaming_iterator import (
BaseResponsesAPIStreamingIterator,
)
def _terminal_event(event_type: str, status: str, sequence_number: int = 2, created_at: int = 1761600000):
return {
"type": event_type,
"sequence_number": sequence_number,
"response": {
"id": "resp_test",
"object": "response",
"created_at": created_at,
"status": status,
"model": "gpt-5.5",
"output": [],
"container": {"id": "cntr_test", "type": "code_interpreter"},
"usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2},
},
}
# Minimal source iterator stub with every attribute the wrapper
# copies in __init__ (see router.py:2552-2583).
source = SimpleNamespace(
response=None,
model="openai/gpt-5.5",
logging_obj=None,
responses_api_provider_config=None,
start_time=None,
litellm_metadata=None,
custom_llm_provider="openai",
request_data={},
call_type="aresponses",
_hidden_params={},
)
# The class is defined inside _aresponses_streaming_iterator; capture
# it by patching FallbackResponsesStreamWrapper into a sentinel on
# construction.
captured = {}
real_router_module = __import__("litellm.router", fromlist=["Router"])
def _stream_through_router(*events):
"""Stream ``events`` from a mocked upstream through ``Router.aresponses``;
return the router's stream object after draining it, and the chunks."""
body = b"".join(f"event: {e['type']}\ndata: {json.dumps(e)}\n\n".encode() for e in events)
async def _drive():
async def empty_gen():
if False:
yield # pragma: no cover
return
router = Router(
router = litellm.Router(
model_list=[
{
"model_name": "openai/gpt-5.5",
"litellm_params": {"model": "openai/gpt-5.5", "api_key": "sk-test"},
"model_name": "gpt-5.5",
"litellm_params": {
"model": "openai/gpt-5.5",
"api_key": "sk-test",
"api_base": "https://api.openai.com/v1",
},
}
]
)
wrapped = await router._aresponses_streaming_iterator(
response=source, # type: ignore[arg-type]
initial_kwargs={},
)
captured["wrapper_cls"] = type(wrapped)
captured["instance"] = wrapped
with respx.mock() as mock_router:
mock_router.post("https://api.openai.com/v1/responses").mock(
return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=body)
)
stream = await router.aresponses(model="gpt-5.5", input="hi", stream=True)
chunks = [chunk async for chunk in stream]
return stream, chunks
asyncio.run(_drive())
return captured["wrapper_cls"], captured["instance"]
def _terminal_chunk(event_type: str):
"""A SimpleNamespace shaped like the openai responses-api terminal
event chunks the wrapper inspects (.type attribute)."""
return SimpleNamespace(
type=event_type,
response=SimpleNamespace(
id="resp_test",
output=[],
container={"id": "cntr_test", "type": "code_interpreter"},
),
)
def _non_terminal_chunk(event_type: str = "response.output_text.delta"):
return SimpleNamespace(type=event_type, delta="hello")
with patch.object(litellm, "disable_aiohttp_transport", True):
litellm.in_memory_llm_clients_cache.flush_cache()
try:
return asyncio.run(_drive())
finally:
litellm.in_memory_llm_clients_cache.flush_cache()
class TestStreamWrapperCapturesTerminalEvent:
@ -104,76 +93,29 @@ class TestStreamWrapperCapturesTerminalEvent:
"""The #30210 bug: a forwarded ``response.completed`` chunk used
to leave ``completed_response`` at None on the wrapper. Verify
it now carries the chunk."""
wrapper_cls, _ = _make_wrapper_class()
async def gen():
yield _non_terminal_chunk()
yield _terminal_chunk("response.completed")
wrapper = wrapper_cls(gen())
# Drain the wrapper.
out = asyncio.run(_drain(wrapper))
assert len(out) == 2
assert wrapper.completed_response is not None, (
stream, chunks = _stream_through_router(_DELTA, _terminal_event("response.completed", "completed"))
assert len(chunks) == 2
assert stream.completed_response is not None, (
"FallbackResponsesStreamWrapper.completed_response is still None "
"after a response.completed chunk passed through — the proxy "
"container-ownership hook will 403 follow-up file lookups (#30210)"
)
assert wrapper.completed_response.type == "response.completed"
assert stream.completed_response.type == "response.completed"
def test_terminal_incomplete_event_is_recorded(self):
wrapper_cls, _ = _make_wrapper_class()
async def gen():
yield _terminal_chunk("response.incomplete")
wrapper = wrapper_cls(gen())
asyncio.run(_drain(wrapper))
assert wrapper.completed_response is not None
assert wrapper.completed_response.type == "response.incomplete"
def test_terminal_failed_event_is_recorded(self):
wrapper_cls, _ = _make_wrapper_class()
async def gen():
yield _terminal_chunk("response.failed")
wrapper = wrapper_cls(gen())
asyncio.run(_drain(wrapper))
assert wrapper.completed_response is not None
assert wrapper.completed_response.type == "response.failed"
def test_non_terminal_chunks_do_not_set_completed_response(self):
wrapper_cls, _ = _make_wrapper_class()
async def gen():
yield _non_terminal_chunk("response.output_text.delta")
yield _non_terminal_chunk("response.code_interpreter.in_progress")
wrapper = wrapper_cls(gen())
asyncio.run(_drain(wrapper))
assert (
wrapper.completed_response is None
), "non-terminal chunks must not set completed_response"
stream, _ = _stream_through_router(_terminal_event("response.incomplete", "incomplete", sequence_number=1))
assert stream.completed_response is not None
assert stream.completed_response.type == "response.incomplete"
def test_first_terminal_event_wins(self):
"""Real streams only emit one terminal event, but defend against
future producers emitting more: keep the first one (the inner
source iterator behaves the same way)."""
wrapper_cls, _ = _make_wrapper_class()
first = _terminal_chunk("response.completed")
first.response.id = "resp_first"
second = _terminal_chunk("response.completed")
second.response.id = "resp_second"
async def gen():
yield first
yield second
wrapper = wrapper_cls(gen())
asyncio.run(_drain(wrapper))
assert wrapper.completed_response.response.id == "resp_first"
stream, _ = _stream_through_router(
_terminal_event("response.completed", "completed", sequence_number=1, created_at=1),
_terminal_event("response.completed", "completed", sequence_number=2, created_at=2),
)
assert stream.completed_response.response.created_at == 1
class TestProxyOwnershipHookReadsCompletedResponse:
@ -186,19 +128,10 @@ class TestProxyOwnershipHookReadsCompletedResponse:
ProxyBaseLLMRequestProcessing,
)
wrapper_cls, _ = _make_wrapper_class()
stream, _ = _stream_through_router(_terminal_event("response.completed", "completed", sequence_number=1))
async def gen():
yield _terminal_chunk("response.completed")
wrapper = wrapper_cls(gen())
asyncio.run(_drain(wrapper))
extracted = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(
wrapper
)
extracted = ProxyBaseLLMRequestProcessing.extract_completed_responses_response(stream)
assert extracted is not None
assert extracted.id == "resp_test"
assert extracted.container["id"] == "cntr_test"
@ -209,24 +142,18 @@ class TestSilentSkipNowLogged:
def test_warning_logged_when_completed_response_missing(self, caplog):
import logging
from litellm.proxy.common_request_processing import (
ProxyBaseLLMRequestProcessing,
)
# Wrap a generator that produces NO terminal event so the
# wrapper stays at completed_response=None — same shape as the
# A stream that never saw a terminal event: the same shape as the
# pre-fix bug.
wrapper_cls, _ = _make_wrapper_class()
async def gen():
yield _non_terminal_chunk()
wrapper = wrapper_cls(gen())
stream = SimpleNamespace(completed_response=None)
async def driver():
async def inner_gen():
async for c in wrapper:
yield c
yield SimpleNamespace(type="response.output_text.delta", delta="hello")
# Patch _record_container_owners_from_responses_if_needed to
# a noop async so the warning branch is exercised in
@ -237,7 +164,7 @@ class TestSilentSkipNowLogged:
new=MagicMock(),
):
wrapped = ProxyBaseLLMRequestProcessing._wrap_responses_stream_for_container_ownership(
original_stream_response=wrapper,
original_stream_response=stream,
wrapped_generator=inner_gen(),
user_api_key_dict=MagicMock(),
)
@ -248,14 +175,5 @@ class TestSilentSkipNowLogged:
asyncio.run(driver())
assert any(
"Container ownership recording skipped on streaming /v1/responses"
in r.message
for r in caplog.records
"Container ownership recording skipped on streaming /v1/responses" in r.message for r in caplog.records
), "silent-skip warning never fired despite completed_response=None"
async def _drain(it):
out = []
async for chunk in it:
out.append(chunk)
return out

31
tests/unit/test_router/core/conftest.py
View file

@ -1,11 +1,11 @@
from collections.abc import Callable, Generator
from collections.abc import Callable
from pathlib import Path
from typing import Final
import pytest
from litellm.router_core import RouterCore
from tests.unit.test_router.core.scripted_host import CORE_BUILDERS, CoreDeclined, ScriptedHost
from tests.unit.test_router.core.scripted_host import CORE_BUILDERS, ScriptedHost
@pytest.fixture
@ -16,15 +16,12 @@ def host() -> ScriptedHost:
@pytest.fixture(params=sorted(CORE_BUILDERS))
def make_core(request: pytest.FixtureRequest, host: ScriptedHost) -> Callable[..., RouterCore]:
"""``make_core(**router_arguments)``: the core under test, on the ``host`` fixture. Every test that
uses it runs once per core in ``CORE_BUILDERS`` (``-k python`` or ``-k rust`` picks one)."""
uses it runs once per core in ``CORE_BUILDERS``."""
name: Final[str] = str(request.param) # pyright: ignore[reportAny] # the fixture parameter
builder: Final = CORE_BUILDERS[name]
def build(**arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments
try:
return builder(host, **arguments)
except CoreDeclined as declined:
pytest.skip(f"{name} core declines this configuration: {declined}")
return builder(host, **arguments)
return build
@ -47,23 +44,3 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None:
name: Final = "::".join((path.relative_to(here).as_posix(), *item.nodeid.split("::")[1:]))
if name in _RUST_PARITY_PENDING:
item.add_marker(pytest.mark.xfail(reason="Rust core parity pending", strict=True))
@pytest.hookimpl(hookwrapper=True)
def pytest_runtest_call(item: pytest.Item) -> Generator[None, object, None]:
"""A test that reaches a member or request the Rust core does not serve yet is skipped with
that member named, so the run separates "not served yet" from a wrong answer."""
outcome: Final = yield
excinfo: Final[object] = getattr(outcome, "excinfo", None)
if not isinstance(excinfo, tuple):
return
info: Final[tuple[object, ...]] = excinfo # pyright: ignore[reportUnknownVariableType] # pluggy's (type, value, traceback)
if len(info) < 2:
return
from litellm.router_rust.core import RustCoreUnsupported
error: Final[object] = info[1]
cause: Final[object] = getattr(error, "__cause__", None)
force_exception: Final[object] = getattr(outcome, "force_exception", None)
if (isinstance(error, RustCoreUnsupported) or isinstance(cause, RustCoreUnsupported)) and callable(force_exception):
force_exception(pytest.skip.Exception(f"not served by the Rust core yet: {error}", _use_item_location=True))

35
tests/unit/test_router/core/scripted_host.py
View file

@ -22,7 +22,6 @@ from litellm.litellm_core_utils.agentic_followup_failure import mark_agentic_fol
from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager
from litellm.router_config import RouterDynamicConfig, build_router_config
from litellm.router_core import RouterCore
from litellm.router_python_core import PythonCore
from litellm.router_python_host import CallbackList, Endpoint, PythonHost
from litellm.types.utils import EmbeddingResponse, ModelResponse
@ -482,43 +481,19 @@ def deployment(
}
def python_core(host: ScriptedHost, **arguments: object) -> PythonCore: # kwargs-ok: the router's constructor arguments
"""A ``PythonCore`` built the way ``Router`` builds one, on ``host``."""
config: Final = build_router_config(arguments)
return PythonCore(**config.arguments, _host=host) # pyright: ignore[reportArgumentType] # the same arguments Router binds
class CoreDeclined(Exception):
"""The core does not serve this configuration (the router would stay on the Python core)."""
CoreBuilder = Callable[..., RouterCore]
"""``builder(host, **router_arguments)``: a router core on ``host``."""
def rust_core(host: ScriptedHost, **arguments: object) -> RouterCore: # kwargs-ok: the router's constructor arguments
"""A ``RustCore`` on ``host``, built the way ``make_core`` builds one when the Rust router is on."""
"""A ``RustCore`` on ``host``, built the way ``make_core`` builds one."""
from litellm.router_rust.core import RustCore
from litellm.rust_bridge._native import RustBridgeDeclined
config: Final = build_router_config(arguments) # an invalid config raises here, as for the Python core
try:
return RustCore.build(config, host)
except RustBridgeDeclined as declined: # make_core keeps such a router on the Python core
raise CoreDeclined(str(declined)) from declined
return RustCore.build(build_router_config(arguments), host)
CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"python": python_core, "rust": rust_core}
def records_outcomes_through_callbacks(core: RouterCore) -> bool:
"""Whether ``core`` learns attempt outcomes from callbacks it registers on its host. The Python
core does; the Rust core records every outcome from the attempt itself and registers none."""
return isinstance(core, PythonCore)
"""Every core the suite runs on. Another core adds its builder here; a builder that cannot take a
configuration raises ``CoreDeclined``, and the test is skipped for that core."""
CORE_BUILDERS: Final[dict[str, CoreBuilder]] = {"rust": rust_core}
"""Every core the suite runs on."""
__all__ = [
@ -526,7 +501,6 @@ __all__ = [
"INFERENCE_OPS",
"Call",
"CoreBuilder",
"CoreDeclined",
"HostCallbacks",
"ScriptedHost",
"auth_error",
@ -537,7 +511,6 @@ __all__ = [
"deployment",
"dynamic_config",
"embedding_response",
"python_core",
"rate_limit",
"rust_core",
"server_error",

74
tests/unit/test_router/core/test_construction_precall_characterization.py
View file

@ -7,6 +7,7 @@ creates, alerting, optional pre-call check registration, ``observe``, ``register
Each test asserts what the Python core does today; oddities are named in the test's docstring.
"""
import inspect
import logging
from collections.abc import Callable, Generator, Mapping
from typing import Final
@ -17,9 +18,9 @@ import litellm
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_cache import RedisCache
from litellm.router import Router
from litellm.router_api import ChangeSet, router_init_args
from litellm.router_core import RouterCore
from litellm.router_python_core import PythonCore
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import PromptCachingDeploymentCheck
from litellm.types.router import (
@ -34,7 +35,7 @@ from tests.unit.test_router.core.construction_precall_helpers import (
collect_other_routers,
host_callback_names,
)
from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, records_outcomes_through_callbacks
from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rust_core
MakeCore = Callable[..., RouterCore]
MESSAGES: Final = [{"role": "user", "content": "hi"}]
@ -694,24 +695,6 @@ def test_alerting_config_registers_slack_alerting_and_its_slow_response_callback
)
def test_the_python_core_registers_its_sync_outcome_callback_before_the_slow_response_callback(
host: ScriptedHost,
) -> None:
"""Python core only: it learns sync outcomes from a success callback on its host. The Rust core
records outcomes from the attempt itself and registers no outcome callbacks."""
from tests.unit.test_router.core.scripted_host import python_core
python_core(
host,
model_list=[deployment("g", "a")],
alerting_config=AlertingConfig(webhook_url="https://example.invalid/hook", alerting_threshold=10),
)
assert host_callback_names(host)["success_callback"] == [
"sync_deployment_callback_on_success",
"response_taking_too_long_callback",
]
def test_observe_reports_no_slack_alerting_without_an_alerting_config(make_core: MakeCore) -> None:
assert _one(make_core).observe().slack_alerting_logger is None
@ -884,11 +867,7 @@ def test_runtime_removal_keeps_the_host_callback_while_another_live_router_uses_
host) still has one of that type."""
collect_other_routers()
other_host: Final = CallbackHost()
from tests.unit.test_router.core.scripted_host import python_core
other: Final = python_core(
other_host, model_list=[deployment("g", "o")], optional_pre_call_checks=["prompt_caching"]
)
other: Final = rust_core(other_host, model_list=[deployment("g", "o")], optional_pre_call_checks=["prompt_caching"])
core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"])
core.apply(ChangeSet(settings={"optional_pre_call_checks": []}))
assert core.optional_callbacks == [] # pyright: ignore[reportAttributeAccessIssue]
@ -948,30 +927,13 @@ def test_retiring_one_router_keeps_the_callbacks_of_another_on_the_same_host(
assert host.lists["input_callback"] == [second.leastbusy_logger] # pyright: ignore[reportAttributeAccessIssue]
def test_retiring_one_python_core_keeps_the_outcome_callback_of_another_on_the_same_host(
host: ScriptedHost,
) -> None:
"""Python core only: each one registers an async outcome callback bound to itself; the Rust core
records outcomes from the attempt itself and registers none."""
from tests.unit.test_router.core.scripted_host import python_core
collect_other_routers()
first: Final = python_core(host, model_list=[deployment("g", "a")], routing_strategy="least-busy")
second: Final = python_core(host, model_list=[deployment("g", "a")], routing_strategy="least-busy")
first.retire()
assert host_callback_names(host)["_async_success_callback"] == ["deployment_callback_on_success"]
assert all(callback.__self__ is second for callback in host.lists["_async_success_callback"]) # pyright: ignore[reportAttributeAccessIssue]
def test_retire_hands_a_vacated_check_to_another_live_router_on_a_different_host(
make_core: MakeCore, host: ScriptedHost
) -> None:
"""Oddity: retire re-registers the same-keyed check of any other live router in the process into
this host's list, even one built on a different host."""
collect_other_routers()
from tests.unit.test_router.core.scripted_host import python_core
other: Final = python_core(
other: Final = rust_core(
CallbackHost(), model_list=[deployment("g", "o")], optional_pre_call_checks=["prompt_caching"]
)
core: Final = _one(make_core, optional_pre_call_checks=["prompt_caching"])
@ -997,14 +959,6 @@ class RecordingLogging:
self.calls.append(("add", type(callback).__name__))
def _bound_outcome_callbacks(core: RouterCore) -> list[tuple[str, object]]: # mutable-ok: a fresh list
return (
[("functions", ["async_failure", "async_success", "failure"])]
if records_outcomes_through_callbacks(core)
else []
)
def test_observe_request_scoped_removes_every_callback_the_router_registered_globally(
make_core: MakeCore, host: ScriptedHost
) -> None:
@ -1018,13 +972,12 @@ def test_observe_request_scoped_removes_every_callback_the_router_registered_glo
def test_observe_with_a_logging_object_binds_the_router_outcome_callbacks(
make_core: MakeCore, host: ScriptedHost
) -> None:
"""The Python core binds its three outcome callbacks; the Rust core records every attempt's
outcome itself and binds none, as it registers none on the host."""
"""The core records every attempt's outcome itself and binds none, as it registers none on the host."""
core: Final = _one(make_core)
before: Final = host_callback_names(host)
logging_obj: Final = RecordingLogging()
core.observe(logging_obj) # pyright: ignore[reportArgumentType]
assert logging_obj.calls == _bound_outcome_callbacks(core)
assert logging_obj.calls == []
assert host_callback_names(host) == before
@ -1035,7 +988,6 @@ def test_observe_request_scoped_then_binding_hands_selectors_and_policies_to_the
logging_obj: Final = RecordingLogging()
core.observe(logging_obj) # pyright: ignore[reportArgumentType]
assert logging_obj.calls == [
*_bound_outcome_callbacks(core),
("replace", "LeastBusyLoggingHandler"),
("add", "PromptCachingDeploymentCheck"),
]
@ -1048,7 +1000,7 @@ def test_binding_without_request_scoping_hands_the_request_no_selector_or_policy
core: Final = _one(make_core, routing_strategy="least-busy", optional_pre_call_checks=["prompt_caching"])
logging_obj: Final = RecordingLogging()
core.observe(logging_obj) # pyright: ignore[reportCallIssue] # observe takes a logging object on every core
assert logging_obj.calls == _bound_outcome_callbacks(core)
assert logging_obj.calls == []
def test_binding_skips_a_policy_the_host_already_runs(make_core: MakeCore, host: ScriptedHost) -> None:
@ -1060,7 +1012,7 @@ def test_binding_skips_a_policy_the_host_already_runs(make_core: MakeCore, host:
host.lists["callbacks"].append(PromptCachingDeploymentCheck(cache=DualCache()))
logging_obj: Final = RecordingLogging()
core.observe(logging_obj) # pyright: ignore[reportCallIssue] # observe takes a logging object on every core
assert logging_obj.calls == [*_bound_outcome_callbacks(core), ("replace", "LeastBusyLoggingHandler")]
assert logging_obj.calls == [("replace", "LeastBusyLoggingHandler")]
# --- register_plugin ---------------------------------------------------------------------
@ -1224,7 +1176,7 @@ def test_registry_settings_and_health_views_are_built_once(make_core: MakeCore)
def test_get_valid_args_lists_every_constructor_argument_but_self_and_the_host() -> None:
names: Final = PythonCore.get_valid_args()
names: Final = router_init_args()
assert names[0] == "model_list"
assert "self" not in names
assert "_host" not in names
@ -1232,12 +1184,12 @@ def test_get_valid_args_lists_every_constructor_argument_but_self_and_the_host()
assert len(names) == len(set(names))
def test_get_valid_args_matches_the_router_init_args() -> None:
assert tuple(PythonCore.get_valid_args()) == tuple(router_init_args())
def test_get_valid_args_matches_the_router_signature() -> None:
assert tuple(inspect.signature(Router).parameters) == tuple(router_init_args())
def test_get_valid_args_on_a_built_router(make_core: MakeCore) -> None:
assert _one(make_core).get_valid_args() == PythonCore.get_valid_args() # pyright: ignore[reportAttributeAccessIssue]
assert _one(make_core).get_valid_args() == list(router_init_args()) # pyright: ignore[reportAttributeAccessIssue]
# --- print_deployment ----------------------------------------------------------------------

6
tests/unit/test_router/core/test_construction_precall_characterization_precall.py
View file

@ -20,7 +20,7 @@ import litellm
from litellm.router_core import RouterCore
from litellm.types.router import RouterRateLimitError
from tests.unit.test_router.core.construction_precall_helpers import CallbackHost, restoring_model_cost, windowed
from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, python_core
from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rust_core
MakeCore = Callable[..., RouterCore]
PICKS: Final = 60 # three-way shuffles miss a deployment in 20 picks about once in a thousand runs
@ -326,7 +326,7 @@ async def test_a_deployment_context_window_is_written_into_the_process_cost_map(
model_list=[deployment("g", "declares-five", model=UNMAPPED_MODEL, model_info={"max_input_tokens": 5})]
)
first.retire()
other: Final = python_core(
other: Final = rust_core(
CallbackHost(),
model_list=[deployment("g", "declares-nothing", model=UNMAPPED_MODEL)],
enable_pre_call_checks=True,
@ -343,7 +343,7 @@ async def test_a_deployment_context_window_stays_off_a_mapped_models_process_ent
model_list=[deployment("g", "declares-five", model=NO_JSON_MODE_MODEL, model_info={"max_input_tokens": 5})]
)
first.retire()
other: Final = python_core(
other: Final = rust_core(
CallbackHost(),
model_list=[deployment("g", "declares-nothing", model=NO_JSON_MODE_MODEL)],
enable_pre_call_checks=True,

20
tests/unit/test_router/core/test_endpoints.py
View file

@ -22,7 +22,6 @@ from tests.unit.test_router.core.scripted_host import (
completion_response,
deployment,
embedding_response,
records_outcomes_through_callbacks,
)
MESSAGES: Final = [{"role": "user", "content": "hi"}]
@ -279,21 +278,18 @@ async def test_request_level_api_key_replaces_the_deployment_key_on_the_host_cal
)
async def test_construction_registers_one_success_and_one_failure_callback_per_mode(
make_core: MakeCore, host: ScriptedHost
):
async def test_construction_registers_no_outcome_callbacks(make_core: MakeCore, host: ScriptedHost):
"""The core records every attempt's outcome itself, so it registers no callbacks on its host."""
core: Final = _one_group(make_core)
registered: Final = {name: len(callbacks) for name, callbacks in host.lists.items()}
one: Final = 1 if records_outcomes_through_callbacks(core) else 0
assert registered == {
assert core.registry() is not None
assert {name: len(callbacks) for name, callbacks in host.lists.items()} == {
"callbacks": 0,
"success_callback": one,
"_async_success_callback": one,
"failure_callback": one,
"_async_failure_callback": one,
"success_callback": 0,
"_async_success_callback": 0,
"failure_callback": 0,
"_async_failure_callback": 0,
"input_callback": 0,
}
assert all(callable(callbacks[0]) for callbacks in host.lists.values() if callbacks)
async def test_retire_removes_the_callbacks_the_core_registered(make_core: MakeCore, host: ScriptedHost):

31
tests/unit/test_router/core/test_legacy_members.py
View file

@ -1,11 +1,10 @@
"""
The old public Router members outside the entry points, on each core: registry lookups, the
The old public Router members outside the entry points: registry lookups, the
routing-state lifecycle, the routing callbacks' checks, one provider call's accounting and the
pre-routing strategy builders (``RustLegacyMembers``, litellm/router_rust/surface/legacy_members.py).
"""
import asyncio
import inspect
from collections.abc import Mapping
from typing import Final
@ -13,18 +12,10 @@ import pytest
import litellm
from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
from litellm.containers.endpoint_factory import (
_generated_endpoints, # pyright: ignore[reportPrivateUsage] # the names the routers register
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.token_counter import served_openai_models
from litellm.router_api import RoutingContext
from litellm.router_core import RouterCore
from litellm.router_python_core import PythonCore
from litellm.router_rust.surface.legacy_members import (
_ASSISTANTS_ENDPOINTS, # pyright: ignore[reportPrivateUsage] # pinned against the Python core below
_ROUTER_ENDPOINTS, # pyright: ignore[reportPrivateUsage] # see above
)
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
from litellm.router_strategy.auto_router.auto_router_model_naming import HEURISTIC_V2_CAPABILITY
from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy
@ -46,7 +37,6 @@ from tests.unit.test_router.core.scripted_host import (
ScriptedHost,
completion_response,
deployment,
python_core,
rate_limit,
rust_core,
server_error,
@ -179,12 +169,6 @@ def test_served_openai_models_reads_the_group_and_follows_a_model_list_change(ma
) # pyright: ignore[reportAttributeAccessIssue] # see above
def test_the_tier_params_never_dropped_constant_matches(make_core: MakeCore) -> None:
core: Final = make_core(model_list=list(GROUP))
assert core.TIER_PARAMS_NEVER_DROPPED == PythonCore.TIER_PARAMS_NEVER_DROPPED # pyright: ignore[reportAttributeAccessIssue] # a class constant
def test_auto_router_capability_violation_resolves_the_limit_on_every_call(make_core: MakeCore) -> None:
limits: Final[list[int | None]] = [None] # mutable-ok: what the resolver answers
@ -326,19 +310,6 @@ def test_initialize_router_endpoints_replaces_only_what_it_binds(make_core: Make
assert core.acreate_file is assigned # pyright: ignore[reportAttributeAccessIssue] # see above
def test_the_endpoint_names_are_the_ones_the_python_core_binds() -> None:
bound: Final = vars(python_core(ScriptedHost(), model_list=[]))
listed: Final = {*_ROUTER_ENDPOINTS, *_ASSISTANTS_ENDPOINTS, *_generated_endpoints}
assert listed <= set(bound)
wrappers: Final = {
name
for name, value in bound.items()
if inspect.isfunction(value) and value.__qualname__.startswith("PythonCore.factory_function.")
}
assert wrappers <= listed
# --- pre-routing strategies ----------------------------------------------------------------------

36
tests/unit/test_router/test_io_token_rate_limits.py
View file

@ -1476,7 +1476,6 @@ class TestRouterIOTokenIntegration:
assert await _io_counters_since(router, first_minute) == (0, 0)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("call", "original_function", "mock_error"),
@ -1549,41 +1548,6 @@ class TestContextSlotRetention:
set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False, metadata_channel="metadata")
assert get_io_token_rate_limit_request_kwargs() is None
@pytest.mark.asyncio
async def test_router_does_not_pin_kwargs_without_io_limits(self):
router = Router(
model_list=[
{
"model_name": "plain",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test"},
}
]
)
set_io_token_rate_limit_request_kwargs(None, metadata_channel="metadata")
kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
deployment = router.get_deployment_by_model_group_name("plain")
assert deployment is not None
router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs)
assert get_io_token_rate_limit_request_kwargs() is None
@pytest.mark.asyncio
async def test_router_pins_kwargs_for_io_limited_deployment(self):
router = Router(
model_list=[
{
"model_name": "limited",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-test", "itpm": 100},
}
],
optional_pre_call_checks=["enforce_model_rate_limits"],
)
set_io_token_rate_limit_request_kwargs(None, metadata_channel="metadata")
kwargs = {"messages": [{"role": "user", "content": "hello"}], "metadata": {}}
deployment = router.get_deployment_by_model_group_name("limited")
assert deployment is not None
router._update_kwargs_with_deployment(deployment=deployment.model_dump(), kwargs=kwargs)
assert get_io_token_rate_limit_request_kwargs() is kwargs
@pytest.mark.asyncio
async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_token_count() -> None:

14340
tests/unit/test_router/test_router.py
View file

File diff suppressed because it is too large Load diff

569
tests/unit/test_router/test_router_endpoints.py
View file

@ -7,6 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, Mock, patch
import httpx
import litellm
import pytest
import respx
from litellm import Router
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
@ -129,7 +130,7 @@ async def test_aspeech_sets_deployment_metadata():
mock_response = MagicMock()
with patch("litellm.aspeech", return_value=mock_response) as mock_aspeech:
response = await router._aspeech(
response = await router.aspeech(
model="tts",
input="the quick brown fox jumped over the lazy dogs",
voice="alloy",
@ -165,9 +166,7 @@ async def test_moderation_endpoint_with_api_base():
)
# Mock the OpenAI client to verify api_base is passed
with patch(
"litellm.main.openai_chat_completions.get_openai_client"
) as mock_get_client:
with patch("litellm.main.openai_chat_completions.get_openai_client") as mock_get_client:
mock_client = AsyncMock()
mock_response = MagicMock()
mock_response.model_dump.return_value = {
@ -185,16 +184,14 @@ async def test_moderation_endpoint_with_api_base():
mock_client.moderations.create = AsyncMock(return_value=mock_response)
mock_get_client.return_value = mock_client
response = await router.amoderation(
model="openai/omni-moderation-latest", input="hello this is a test"
)
response = await router.amoderation(model="openai/omni-moderation-latest", input="hello this is a test")
# Verify that get_openai_client was called with the custom api_base
mock_get_client.assert_called()
call_kwargs = mock_get_client.call_args.kwargs
assert (
call_kwargs.get("api_base") == custom_api_base
), f"Expected api_base to be {custom_api_base}, but got {call_kwargs.get('api_base')}"
assert call_kwargs.get("api_base") == custom_api_base, (
f"Expected api_base to be {custom_api_base}, but got {call_kwargs.get('api_base')}"
)
print(f"✓ Moderation endpoint correctly uses api_base: {custom_api_base}")
@ -217,13 +214,6 @@ async def test_aaaaatext_completion_endpoint(model_list, sync_mode):
mock_response="I'm fine, thank you!",
)
response_2 = await router._atext_completion(
model="gpt-5-mini",
prompt="Hello, how are you?",
mock_response="I'm fine, thank you!",
)
assert response_2.choices[0].text == "I'm fine, thank you!"
assert response.choices[0].text == "I'm fine, thank you!"
@ -252,147 +242,9 @@ async def test_router_with_empty_choices(model_list):
assert response is not None
@pytest.mark.parametrize("sync_mode", [True, False])
def test_generic_api_call_with_fallbacks_basic(sync_mode):
"""
Test both the sync and async versions of generic_api_call_with_fallbacks with a basic successful call
"""
# Create a mock function that will be passed to generic_api_call_with_fallbacks
if sync_mode:
from unittest.mock import Mock
mock_function = Mock()
mock_function.__name__ = "test_function"
else:
mock_function = AsyncMock()
mock_function.__name__ = "test_function"
# Create a mock response
mock_response = {
"id": "resp_123456",
"role": "assistant",
"content": "This is a test response",
"model": "test-model",
"usage": {"input_tokens": 10, "output_tokens": 20},
}
mock_function.return_value = mock_response
# Create a router with a test model
router = Router(
model_list=[
{
"model_name": "test-model-alias",
"litellm_params": {
"model": "anthropic/test-model",
"api_key": "fake-api-key",
},
}
]
)
# Call the appropriate generic_api_call_with_fallbacks method
if sync_mode:
response = router._generic_api_call_with_fallbacks(
model="test-model-alias",
original_function=mock_function,
messages=[{"role": "user", "content": "Hello"}],
max_tokens=100,
)
else:
response = asyncio.run(
router._ageneric_api_call_with_fallbacks(
model="test-model-alias",
original_function=mock_function,
messages=[{"role": "user", "content": "Hello"}],
max_tokens=100,
)
)
# Verify the mock function was called
mock_function.assert_called_once()
# Verify the response
assert response == mock_response
@pytest.mark.asyncio
async def test_aadapter_completion():
"""
Test the aadapter_completion method which uses async_function_with_fallbacks
"""
mock_response = {
"id": "adapter_resp_123",
"object": "adapter.completion",
"created": 1677858242,
"model": "test-model-with-adapter",
"choices": [
{
"text": "This is a test adapter response",
"index": 0,
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
}
with patch.object(Router, "_aadapter_completion", new_callable=AsyncMock) as mock_method:
mock_method.return_value = mock_response
router = Router(
model_list=[
{
"model_name": "test-adapter-model",
"litellm_params": {
"model": "anthropic/test-model",
"api_key": "fake-api-key",
},
}
]
)
router.async_function_with_fallbacks = AsyncMock(return_value=mock_response)
response = await router.aadapter_completion(
adapter_id="test-adapter-id",
model="test-adapter-model",
prompt="This is a test prompt",
max_tokens=100,
)
assert response == mock_response
router.async_function_with_fallbacks.assert_called_once()
call_kwargs = router.async_function_with_fallbacks.call_args.kwargs
assert call_kwargs["adapter_id"] == "test-adapter-id"
assert call_kwargs["model"] == "test-adapter-model"
assert call_kwargs["prompt"] == "This is a test prompt"
assert call_kwargs["max_tokens"] == 100
assert call_kwargs["original_function"] == router._aadapter_completion
assert "metadata" in call_kwargs
assert call_kwargs["metadata"]["model_group"] == "test-adapter-model"
@pytest.mark.asyncio
async def test__aadapter_completion():
"""
Test the _aadapter_completion method directly
"""
mock_response = {
"id": "adapter_resp_123",
"object": "adapter.completion",
"created": 1677858242,
"model": "test-model-with-adapter",
"choices": [
{
"text": "This is a test adapter response",
"index": 0,
"finish_reason": "stop",
}
],
"usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
}
async def test_aadapter_completion_calls_the_adapter_with_the_deployments_params():
mock_response = {"id": "adapter_resp_123", "object": "adapter.completion"}
with patch("litellm.aadapter_completion", new_callable=AsyncMock) as mock_adapter_completion:
mock_adapter_completion.return_value = mock_response
@ -409,43 +261,22 @@ async def test__aadapter_completion():
]
)
router.async_get_available_deployment = AsyncMock(
return_value={
"model_name": "test-adapter-model",
"litellm_params": {
"model": "test-model",
"api_key": "fake-api-key",
},
"model_info": {
"id": "test-unique-id",
},
}
)
router.async_routing_strategy_pre_call_checks = AsyncMock()
response = await router._aadapter_completion(
response = await router.aadapter_completion(
adapter_id="test-adapter-id",
model="test-adapter-model",
prompt="This is a test prompt",
max_tokens=100,
)
assert response == mock_response
mock_adapter_completion.assert_called_once()
call_kwargs = mock_adapter_completion.call_args.kwargs
assert call_kwargs["adapter_id"] == "test-adapter-id"
assert call_kwargs["model"] == "test-model"
assert call_kwargs["prompt"] == "This is a test prompt"
assert call_kwargs["max_tokens"] == 100
assert call_kwargs["api_key"] == "fake-api-key"
assert call_kwargs["caching"] == router.cache_responses
assert router.success_calls["test-model"] == 1
assert router.total_calls["test-model"] == 1
router.async_routing_strategy_pre_call_checks.assert_called_once()
assert response["id"] == "adapter_resp_123"
mock_adapter_completion.assert_called_once()
call_kwargs = mock_adapter_completion.call_args.kwargs
assert call_kwargs["adapter_id"] == "test-adapter-id"
assert call_kwargs["model"] == "anthropic/test-model"
assert call_kwargs["prompt"] == "This is a test prompt"
assert call_kwargs["max_tokens"] == 100
assert call_kwargs["api_key"] == "fake-api-key"
assert call_kwargs["metadata"]["model_group"] == "test-adapter-model"
def test_initialize_router_endpoints():
@ -482,101 +313,6 @@ def test_initialize_router_endpoints():
assert callable(router.adelete_responses)
@pytest.mark.asyncio
async def test_init_responses_api_endpoints():
"""
A simpler test for _init_responses_api_endpoints that focuses on the basic functionality
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
# Create a router with a basic model
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/test-model",
"api_key": "fake-api-key",
},
}
]
)
# Just mock the _ageneric_api_call_with_fallbacks method
router._ageneric_api_call_with_fallbacks = AsyncMock()
# Add a mock implementation of _get_model_id_from_response_id to the Router instance
ResponsesAPIRequestUtils.get_model_id_from_response_id = MagicMock(
return_value=None
)
# Call without a response_id (no model extraction should happen)
await router._init_responses_api_endpoints(
original_function=AsyncMock(), thread_id="thread_xyz"
)
# Verify _ageneric_api_call_with_fallbacks was called but model wasn't changed
first_call_kwargs = router._ageneric_api_call_with_fallbacks.call_args.kwargs
assert "model" not in first_call_kwargs
assert first_call_kwargs["thread_id"] == "thread_xyz"
# Reset the mock
router._ageneric_api_call_with_fallbacks.reset_mock()
# Change the return value for the second call
ResponsesAPIRequestUtils.get_model_id_from_response_id.return_value = (
"claude-3-sonnet"
)
# Call with a response_id
await router._init_responses_api_endpoints(
original_function=AsyncMock(), response_id="resp_claude_123"
)
# Verify model was updated in the kwargs
second_call_kwargs = router._ageneric_api_call_with_fallbacks.call_args.kwargs
assert second_call_kwargs["model"] == "claude-3-sonnet"
assert second_call_kwargs["response_id"] == "resp_claude_123"
@pytest.mark.asyncio
async def test_init_vector_store_api_endpoints():
"""
Test that _init_vector_store_api_endpoints correctly passes custom_llm_provider to kwargs
"""
router = Router(
model_list=[
{
"model_name": "test-model",
"litellm_params": {
"model": "openai/test-model",
"api_key": "fake-api-key",
},
}
]
)
mock_original_function = AsyncMock(return_value={"status": "success"})
result = await router._init_vector_store_api_endpoints(
original_function=mock_original_function, vector_store_id="test-store"
)
mock_original_function.assert_called_once_with(vector_store_id="test-store")
assert result == {"status": "success"}
mock_original_function.reset_mock()
await router._init_vector_store_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
vector_store_id="test-store",
)
mock_original_function.assert_called_once_with(vector_store_id="test-store", custom_llm_provider="openai")
@pytest.mark.asyncio
async def test_avector_store_create_sends_only_the_callers_metadata_to_the_provider():
def echo_vector_store(request: httpx.Request) -> httpx.Response:
@ -612,10 +348,8 @@ async def test_avector_store_create_sends_only_the_callers_metadata_to_the_provi
def test_apply_default_settings():
"""
Test the apply_default_settings method.
This test verifies that apply_default_settings correctly initializes
default pre-call checks and doesn't modify existing router state.
apply_default_settings initializes the default pre-call checks and leaves the optional
callbacks the router already has in place.
"""
router = Router()
@ -637,16 +371,6 @@ def test_apply_default_settings():
assert mock_callback in router_with_callbacks.optional_callbacks
with patch.object(Router, "apply_default_settings") as mock_apply:
Router()
mock_apply.assert_called_once()
router_test = Router()
with patch.object(router_test, "add_optional_pre_call_checks") as mock_add_checks:
router_test.apply_default_settings()
mock_add_checks.assert_called_once_with([])
def test_initialize_core_endpoints():
"""
@ -664,8 +388,6 @@ def test_initialize_core_endpoints():
]
)
router._initialize_core_endpoints()
core_endpoints = [
"amoderation",
"aanthropic_messages",
@ -679,7 +401,6 @@ def test_initialize_core_endpoints():
"acancel_responses",
"adelete_responses",
"alist_input_items",
"_arealtime",
"acreate_fine_tuning_job",
"acancel_fine_tuning_job",
"alist_fine_tuning_jobs",
@ -710,8 +431,6 @@ def test_initialize_specialized_endpoints():
]
)
router._initialize_specialized_endpoints()
specialized_endpoints = [
"avector_store_search",
"avector_store_create",
@ -770,8 +489,6 @@ def test_initialize_vector_store_endpoints():
]
)
router._initialize_vector_store_endpoints()
vector_store_endpoints = [
"avector_store_search",
"avector_store_create",
@ -800,8 +517,6 @@ def test_initialize_vector_store_file_endpoints():
]
)
router._initialize_vector_store_file_endpoints()
vector_store_file_endpoints = [
"avector_store_file_create",
"vector_store_file_create",
@ -838,8 +553,6 @@ def test_initialize_google_genai_endpoints():
]
)
router._initialize_google_genai_endpoints()
google_genai_endpoints = [
"agenerate_content",
"generate_content",
@ -868,8 +581,6 @@ def test_initialize_ocr_search_endpoints():
]
)
router._initialize_ocr_search_endpoints()
ocr_search_endpoints = [
"aocr",
"ocr",
@ -898,8 +609,6 @@ def test_initialize_video_endpoints():
]
)
router._initialize_video_endpoints()
video_endpoints = [
"avideo_generation",
"video_generation",
@ -934,8 +643,6 @@ def test_initialize_container_endpoints():
]
)
router._initialize_container_endpoints()
container_endpoints = [
"acreate_container",
"create_container",
@ -968,8 +675,6 @@ def test_initialize_skills_endpoints():
]
)
router._initialize_skills_endpoints()
skills_endpoints = [
"acreate_skill",
"alist_skills",
@ -982,35 +687,34 @@ def test_initialize_skills_endpoints():
assert callable(getattr(router, endpoint))
@pytest.mark.asyncio
async def test_init_containers_api_endpoints():
"""
Test that _init_containers_api_endpoints calls the original function
directly when there is no managed container ID (no embedded model_id).
"""
router = Router(model_list=[])
_CONTAINER_FILE: Final = {
"id": "cfile_xyz",
"object": "container.file",
"container_id": "cfile_upstream_abc",
"created_at": 0,
"bytes": 1,
"path": "/a",
"source": "user",
}
mock_response = {"id": "cntr_test", "name": "Test Container"}
mock_original_function = AsyncMock(return_value=mock_response)
result = await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
name="Test Container",
)
mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container")
assert result == mock_response
async def _container_file_request(router: Router, container_id: str) -> httpx.Request:
with respx.mock(assert_all_called=False) as upstream:
route = upstream.route().mock(return_value=httpx.Response(200, json=_CONTAINER_FILE))
await router.aretrieve_container_file(container_id=container_id, file_id="cfile_xyz")
assert route.call_count == 1
return route.calls.last.request
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_managed_id_routes_via_generic_fallbacks():
async def test_containers_managed_id_routes_to_its_deployment_with_its_credentials(monkeypatch):
"""
Managed ``cntr_`` IDs embed ``model_id``; router should decode and use
``_ageneric_api_call_with_fallbacks`` so deployment credentials apply.
Managed ``cntr_`` IDs embed ``model_id``; the router decodes it so the deployment's
credentials and api_base apply, and the upstream gets the unwrapped container ID.
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
router = Router(
model_list=[
{
@ -1023,31 +727,21 @@ async def test_init_containers_api_endpoints_managed_id_routes_via_generic_fallb
}
]
)
router._ageneric_api_call_with_fallbacks = AsyncMock()
managed_id = ResponsesAPIRequestUtils.build_container_id(
custom_llm_provider="azure",
model_id="azure-router-model",
container_id="cfile_upstream_abc",
)
await router._init_containers_api_endpoints(
original_function=AsyncMock(),
custom_llm_provider="openai",
container_id=managed_id,
file_id="cfile_xyz",
)
request = await _container_file_request(router, managed_id)
router._ageneric_api_call_with_fallbacks.assert_called_once()
call_kw = router._ageneric_api_call_with_fallbacks.call_args.kwargs
assert call_kw["model"] == "azure-router-model"
assert call_kw["container_id"] == "cfile_upstream_abc"
assert call_kw["file_id"] == "cfile_xyz"
assert call_kw["custom_llm_provider"] == "azure"
assert request.url.host == "westus.api.cognitive.microsoft.com"
assert request.url.path.endswith("/containers/cfile_upstream_abc/files/cfile_xyz")
assert request.headers["api-key"] == "fake-key"
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_managed_id_without_model_id_unwraps():
async def test_containers_managed_id_without_model_id_unwraps(monkeypatch):
"""
Managed ``cntr_`` IDs may be encoded with an empty ``model_id`` (e.g. when a
streaming response had no router metadata). The router must still unwrap the
@ -1056,64 +750,22 @@ async def test_init_containers_api_endpoints_managed_id_without_model_id_unwraps
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
monkeypatch.setenv("OPENAI_API_KEY", "sk-global")
router = Router(model_list=[])
mock_original_function = AsyncMock(return_value={"ok": True})
managed_id = ResponsesAPIRequestUtils.build_container_id(
custom_llm_provider="openai",
model_id=None,
container_id="cfile_upstream_abc",
)
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
container_id=managed_id,
file_id="cfile_xyz",
)
request = await _container_file_request(router, managed_id)
mock_original_function.assert_called_once()
call_kw = mock_original_function.call_args.kwargs
assert call_kw["container_id"] == "cfile_upstream_abc"
assert call_kw["file_id"] == "cfile_xyz"
assert call_kw["custom_llm_provider"] == "openai"
assert request.url.path == "/v1/containers/cfile_upstream_abc/files/cfile_xyz"
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_managed_id_without_model_id_applies_decoded_provider():
"""
A managed ``cntr_`` ID can encode a non-OpenAI provider (e.g. ``azure``) with
an empty ``model_id`` (streaming events without router ``model_info.id``).
The router must still apply the decoded provider so the request routes to
the correct upstream — not stay on the default ``openai``.
"""
from litellm.responses.utils import ResponsesAPIRequestUtils
router = Router(model_list=[])
mock_original_function = AsyncMock(return_value={"ok": True})
managed_id = ResponsesAPIRequestUtils.build_container_id(
custom_llm_provider="azure",
model_id=None,
container_id="cfile_upstream_abc",
)
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
container_id=managed_id,
file_id="cfile_xyz",
)
mock_original_function.assert_called_once()
call_kw = mock_original_function.call_args.kwargs
assert call_kw["container_id"] == "cfile_upstream_abc"
assert call_kw["file_id"] == "cfile_xyz"
assert call_kw["custom_llm_provider"] == "azure"
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_create_with_model_uses_deployment_credentials(monkeypatch):
async def test_containers_create_with_model_uses_deployment_credentials(monkeypatch):
"""
``POST /v1/containers`` carries no container ID, so a ``model`` in the request
body is the only way to pick a deployment. The upstream call must receive that
@ -1121,29 +773,24 @@ async def test_init_containers_api_endpoints_create_with_model_uses_deployment_c
``OPENAI_API_KEY`` (which may be unset on the proxy).
"""
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
router = Router(
model_list=[
{
"model_name": "gpt-5.4",
"litellm_params": {
"model": "openai/gpt-5.4",
"api_key": "sk-model-list-key",
"api_base": "https://custom.openai.example/v1",
},
}
]
)
mock_original_function = AsyncMock(return_value={"id": "cntr_test", "name": "Test Container"})
with patch("litellm.containers.acreate_container", new_callable=AsyncMock) as create:
create.return_value = {"id": "cntr_test", "name": "Test Container"}
router = Router(
model_list=[
{
"model_name": "gpt-5.4",
"litellm_params": {
"model": "openai/gpt-5.4",
"api_key": "sk-model-list-key",
"api_base": "https://custom.openai.example/v1",
},
}
]
)
await router.acreate_container(name="Test Container", model="gpt-5.4")
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
name="Test Container",
model="gpt-5.4",
)
mock_original_function.assert_called_once()
call_kw = mock_original_function.call_args.kwargs
create.assert_called_once()
call_kw = create.call_args.kwargs
assert call_kw["api_key"] == "sk-model-list-key"
assert call_kw["api_base"] == "https://custom.openai.example/v1"
assert call_kw["model"] == "openai/gpt-5.4"
@ -1151,53 +798,40 @@ async def test_init_containers_api_endpoints_create_with_model_uses_deployment_c
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_create_without_model_calls_directly():
async def test_containers_create_without_model_calls_directly():
"""
Without ``model`` (or with ``model=None`` as the proxy forwards it), create/list
must keep calling the handler directly with global provider credentials.
Without ``model``, create must keep calling the handler directly with global provider credentials.
"""
router = Router(model_list=[])
router._ageneric_api_call_with_fallbacks = AsyncMock()
mock_original_function = AsyncMock(return_value={"id": "cntr_test"})
with patch("litellm.containers.acreate_container", new_callable=AsyncMock) as create:
create.return_value = {"id": "cntr_test"}
router = Router(model_list=[])
await router.acreate_container(name="Test Container")
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
name="Test Container",
model=None,
)
router._ageneric_api_call_with_fallbacks.assert_not_called()
mock_original_function.assert_called_once_with(custom_llm_provider="openai", name="Test Container", model=None)
create.assert_called_once_with(name="Test Container")
@pytest.mark.asyncio
async def test_init_containers_api_endpoints_create_with_unknown_model_passes_through(monkeypatch):
async def test_containers_create_with_unknown_model_passes_through(monkeypatch):
"""
A ``model`` that names no configured deployment must not turn into a 400. The call
falls through to the handler with the caller's model and no injected deployment
credentials, matching the behaviour before model-based routing existed.
"""
monkeypatch.delenv("OPENAI_API_KEY", raising=False)
router = Router(
model_list=[
{
"model_name": "gpt-5.4",
"litellm_params": {"model": "openai/gpt-5.4", "api_key": "sk-model-list-key"},
}
]
)
mock_original_function = AsyncMock(return_value={"id": "cntr_test"})
with patch("litellm.containers.acreate_container", new_callable=AsyncMock) as create:
create.return_value = {"id": "cntr_test"}
router = Router(
model_list=[
{
"model_name": "gpt-5.4",
"litellm_params": {"model": "openai/gpt-5.4", "api_key": "sk-model-list-key"},
}
]
)
await router.acreate_container(name="Test Container", model="does-not-exist")
await router._init_containers_api_endpoints(
original_function=mock_original_function,
custom_llm_provider="openai",
name="Test Container",
model="does-not-exist",
)
mock_original_function.assert_called_once()
call_kw = mock_original_function.call_args.kwargs
create.assert_called_once()
call_kw = create.call_args.kwargs
assert call_kw["model"] == "does-not-exist"
assert call_kw["name"] == "Test Container"
assert "api_key" not in call_kw
@ -1232,30 +866,11 @@ def test_router_model_group_encrypted_content_affinity_callback_registration():
try:
callbacks = router.optional_callbacks or []
encrypted_content_callbacks = [
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
]
deployment_callback = next(
cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)
)
encrypted_content_callbacks = [cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)]
deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck))
assert len(encrypted_content_callbacks) == 1
assert encrypted_content_callbacks[0].enable_global_affinity is False
assert (
encrypted_content_callbacks[0].model_group_affinity_config
== model_group_affinity_config
)
assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(
deployment_callback
)
router._add_encrypted_content_affinity_check(enable_global_affinity=True)
callbacks = router.optional_callbacks or []
encrypted_content_callbacks = [
cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)
]
assert len(encrypted_content_callbacks) == 1
assert encrypted_content_callbacks[0].enable_global_affinity is True
assert encrypted_content_callbacks[0].ctx is router._core_ctx
assert encrypted_content_callbacks[0].model_group_affinity_config == model_group_affinity_config
assert callbacks.index(encrypted_content_callbacks[0]) < callbacks.index(deployment_callback)
finally:
router.discard()

250
tests/unit/test_router/test_router_entry_points.py
View file

@ -332,15 +332,6 @@ def test_settings_snapshot_is_reused_until_a_setting_is_written(router):
warnings.simplefilter("ignore")
router.fallbacks = [{"gpt": ["team-gpt"]}]
assert router.settings().fallbacks == [{"gpt": ["team-gpt"]}]
router._core_fallbacks = [{"gpt": ["other"]}] # the router's own writes invalidate it too
assert router.settings().fallbacks == [{"gpt": ["other"]}]
router.settings().fallbacks.append(
{"x": ["y"]}
) # the snapshot's lists are copies: settings change only via apply()
assert router._core_fallbacks == [{"gpt": ["other"]}]
unrelated = router.settings()
router._core_model_list = list(router._core_model_list)
assert router.settings() is unrelated
def test_deleting_a_setting_shows_as_unset(router):
@ -358,40 +349,6 @@ def test_settings_snapshot_is_frozen(router):
assert set(SETTING_FIELDS) <= set(RouterSettings.__dataclass_fields__)
def test_unset_setting_raises_attribute_error_through_old_name():
r = _router()
try:
del r._core.__dict__["_core_kubernetes_pod_discovery"]
with warnings.catch_warnings():
warnings.simplefilter("ignore")
assert getattr(r, "kubernetes_pod_discovery", "missing") == "missing"
assert r.settings().is_set("kubernetes_pod_discovery") is False
finally:
with warnings.catch_warnings():
warnings.simplefilter("ignore")
r.retire()
def test_unset_attribute_error_names_the_public_attribute(router):
# Reading (or deleting) an old attribute the router never set raises what reading it raised before.
router._core.__dict__.pop("_core_slack_alerting_logger", None)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
with pytest.raises(AttributeError) as read:
_ = router.slack_alerting_logger
with pytest.raises(AttributeError) as deleted:
del router.slack_alerting_logger
del router._core.__dict__["_core_kubernetes_pod_discovery"]
with pytest.raises(AttributeError) as setting:
_ = router.kubernetes_pod_discovery
for caught, name in ((read, "slack_alerting_logger"), (setting, "kubernetes_pod_discovery")):
assert str(caught.value) == f"'Router' object has no attribute '{name}'"
assert caught.value.name == name
assert caught.value.obj is router
# A failed del sets no name/obj in Python either.
assert str(deleted.value) == "'Router' object has no attribute 'slack_alerting_logger'"
def test_bind_to_request_logging_none_raises_as_before(router):
# The old member bound unconditionally, so None fails on the first callback registration.
with warnings.catch_warnings():
@ -456,17 +413,7 @@ def test_router_cannot_be_subclassed():
def test_core_classes_can_still_be_subclassed():
from litellm.router_python_core import PythonCore
class PythonSub(PythonCore):
pass
assert isinstance(object.__new__(PythonSub), Router)
try:
from litellm.router_rust.core import RustCore
except ImportError:
pytest.skip("the native lib does not load")
from litellm.router_rust.core import RustCore
class RustSub(RustCore):
pass
@ -490,15 +437,6 @@ def test_alias_attribute_writes_the_same_storage(router):
assert router.registry().routing_groups() is replacement
def test_core_members_keep_their_old_names_in_errors():
from litellm.router import _name_core_members_after_old_names
_name_core_members_after_old_names()
assert Router._core_get_deployment_model_info.__qualname__ == "Router.get_deployment_model_info"
assert Router._core_get_deployment.__qualname__ == "Router.get_deployment"
def test_wrong_arity_errors_still_name_router(router):
with warnings.catch_warnings():
warnings.simplefilter("ignore")
@ -510,177 +448,10 @@ def test_wrong_arity_errors_still_name_router(router):
router.has_model_id(**None) # pyright: ignore[reportCallIssue, reportArgumentType]
with pytest.raises(TypeError, match=r"^Router\.is_deployment_blocked\(\) missing 1 required"):
router.is_deployment_blocked() # pyright: ignore[reportCallIssue]
with pytest.raises(TypeError, match=r"^Router\.get_deployment_model_info\(\) missing 2 required"):
router.cached_deployment_model_info() # pyright: ignore[reportCallIssue]
_SENTINEL = object()
# (entry point call, Router core member it forwards to, the arguments the core member must get)
FORWARDS = [
(
lambda r: r.registry().deployment_model_info("dep-a", "gpt", cached=True),
"_core_cached_deployment_model_info",
(("dep-a", "gpt"), {}),
),
(
lambda r: r.registry().deployment_model_info("dep-a", "gpt"),
"_core_get_deployment_model_info",
((), {"model_id": "dep-a", "model_name": "gpt"}),
),
(
lambda r: r.registry().router_model_info(None, "gpt", id="dep-a"),
"_core_get_router_model_info",
((), {"deployment": None, "received_model_name": "gpt", "id": "dep-a"}),
),
(lambda r: r.registry().discovered_model_info("dep-a"), "_core_get_discovered_model_info", (("dep-a",), {})),
(
lambda r: r.registry().deployments("gpt", source="routing_groups"),
"_core_get_model_list_from_routing_groups",
((), {"model_name": "gpt"}),
),
(
lambda r: r.registry().deployments("gpt", source="routed_group", team_id="t"),
"_core_get_model_list_of_routed_group",
(("gpt", "t"), {}),
),
(
lambda r: r.registry().deployments("gpt", source="wildcard", team_id="t"),
"_core_wildcard_route_deployments",
(("gpt", "t"), {}),
),
(
lambda r: r.registry().deployments("gpt", source="routable", request={"a": 1}, call_type="aembedding"),
"_core_get_routable_deployments",
(("gpt", {"a": 1}, "aembedding"), {}),
),
(
lambda r: r.registry().deployments("gpt", source="for_request"),
"_core_deployments_for_request",
(("gpt", {}), {}),
),
(
lambda r: r.registry().includes({"model_name": "gpt"}, model_name="gpt", team_id="t"),
"_core_should_include_deployment",
((), {"model_name": "gpt", "model": {"model_name": "gpt"}, "team_id": "t"}),
),
(
lambda r: r.registry().names("gpt", kind="litellm_models", team_id="t"),
"_core_resolved_litellm_models",
(("gpt",), {"team_id": "t"}),
),
(
lambda r: r.registry().names("gpt", kind="litellm_base_models"),
"_core_resolved_litellm_models",
(("gpt",), {"team_id": None, "use_base_model": True}),
),
(
lambda r: r.registry().names("gpt", kind="wildcard", team_id="t"),
"_core_wildcard_routed_models",
(("gpt", "t"), {}),
),
(
lambda r: r.registry().names("gpt", kind="route_candidates"),
"_core_get_candidate_model_ids_for_route",
(("gpt",), {"team_id": None}),
),
(
lambda r: r.registry().resolve("dep-a", kind="from_id", custom_llm_provider="openai"),
"_core_resolve_model_name_from_model_id",
(("dep-a",), {"custom_llm_provider": "openai"}),
),
(
lambda r: r.registry().resolve("rg", kind="routing_group_alternatives"),
"_core_routing_group_has_alternatives",
(("rg",), {}),
),
(
lambda r: r.registry().resolve("gpt", kind="served", served_model="gpt-4o"),
"_core_resolve_served_model_name",
(("gpt", "gpt-4o"), {}),
),
(
lambda r: r.registry().resolve("gpt", kind="upstream", team_id="t"),
"_core_get_routable_upstream_model",
(("gpt",), {"team_id": "t"}),
),
(
lambda r: r.registry().resolve("gpt", kind="search_tool_support", team_id="t"),
"_core_get_search_tool_support",
(("gpt",), {"team_id": "t"}),
),
(
lambda r: r.registry().resolve("grp", kind="wildcard_access_group"),
"_core_is_model_access_group_for_wildcard_route",
((), {"model_access_group": "grp"}),
),
(
lambda r: r.registry().resolve("dep-a", kind="team_model_alternatives"),
"_core_team_model_has_alternatives",
(("dep-a",), {}),
),
(
lambda r: r.registry().group("gpt", include_hidden=True, team_id="t"),
"_core_get_model_group_info",
(("gpt",), {"include_hidden": True, "team_id": "t"}),
),
(lambda r: r.registry().group("gpt", field="cached_info"), "_core_cached_model_group_info", (("gpt",), {})),
(lambda r: r.registry().group("gpt", field="display_name"), "_core_get_configured_display_name", (("gpt",), {})),
(
lambda r: r.registry().group("gpt", field="service_tiers", team_id="t"),
"_core_get_configured_service_tiers",
(("gpt",), {"team_id": "t"}),
),
(lambda r: r.registry().group("gpt", field="listing_info"), "_core_get_model_listing_info", (("gpt",), {})),
(lambda r: r.registry().plugins("budget_limiter"), "_core_get_router_deployment_budget_limiter", ((), {})),
(lambda r: r.settings().allowed_fails_for(ValueError("x")), "_core_get_allowed_fails_from_policy", None),
(lambda r: r.settings().guardrail("g"), "_core_get_available_guardrail", (("g",), {})),
(lambda r: r.health().team_ptu_ceiling("t", "gpt"), "_core_team_ptu_ceiling", (("t", "gpt"), {})),
(
lambda r: r.health().split_priority("gpt", {"priority": 1}),
"_core_split_request_priority",
(("gpt", {"priority": 1}), {}),
),
]
@pytest.mark.parametrize("call, core, expected", FORWARDS, ids=[f[1] for f in FORWARDS])
def test_entry_point_forwards_to_its_core_member(router, monkeypatch, call, core, expected):
from unittest.mock import MagicMock
mock = MagicMock(return_value=_SENTINEL)
monkeypatch.setattr(router, core, mock)
assert call(router) is _SENTINEL
if expected is not None:
assert mock.call_args == expected
else:
mock.assert_called_once()
ASYNC_FORWARDS = [
(lambda r: r.health().fully_unhealthy_model_names(), "_core_async_get_fully_unhealthy_model_names", ((), {})),
(lambda r: r.health().usage("gpt"), "_core_get_model_group_usage", ((), {"model_group": "gpt"})),
(
lambda r: r.health().usage("gpt", kind="remaining"),
"_core_get_remaining_model_group_usage",
((), {"model_group": "gpt"}),
),
]
@pytest.mark.parametrize("call, core, expected", ASYNC_FORWARDS, ids=[f[1] for f in ASYNC_FORWARDS])
def test_async_entry_point_forwards_to_its_core_member(router, monkeypatch, call, core, expected):
from unittest.mock import AsyncMock
mock = AsyncMock(return_value=_SENTINEL)
monkeypatch.setattr(router, core, mock)
assert asyncio.run(call(router)) is _SENTINEL
assert mock.call_args == expected
def test_registry_reads_without_a_core_forward(router):
registry = router.registry()
assert registry.zero_cost() is getattr(router, "_zero_cost_cache", None)
assert registry.includes(router.registry().deployment("dep-a").model_dump(), team_id=None) is True
@ -727,30 +498,13 @@ def test_old_members_refuse_objects_that_are_not_routers():
_as_router(object()) # pyright: ignore[reportArgumentType]
def test_routing_context_forwards_to_the_router(router, monkeypatch):
from unittest.mock import AsyncMock, MagicMock
def test_routing_context_forwards_to_the_router(router):
from litellm.router_api import RoutingContext
ctx = RoutingContext(router)
assert ctx.registry().version == router.registry().version
assert ctx.settings() is router.settings()
assert isinstance(ctx.health(), type(router.health()))
for name in (
"_pre_call_checks",
"_count_pre_call_check_tokens",
"_update_kwargs_before_fallbacks",
"_core_arm_routing_read_prefetch",
"_core_log_retry",
):
monkeypatch.setattr(router, name, MagicMock(return_value=_SENTINEL))
monkeypatch.setattr(router, "_core_async_pre_routing_hook", AsyncMock(return_value=_SENTINEL))
assert ctx.pre_call_checks("gpt", []) is _SENTINEL
assert ctx.count_pre_call_check_tokens(None, "hi") is _SENTINEL
assert ctx.update_kwargs_before_fallbacks("gpt", {}) is _SENTINEL
assert ctx.arm_read_prefetch("gpt") is _SENTINEL
assert asyncio.run(ctx.pre_routing_hook("gpt", {})) is _SENTINEL
assert ctx.log_retry({}, ValueError("x")) is _SENTINEL
def test_attribute_errors_about_other_objects_pass_through_unchanged(router):

13
tests/unit/test_router/test_router_facade.py
View file

@ -1,12 +1,10 @@
"""
The ``litellm.Router`` facade's own helpers in litellm/router.py: how it binds the constructor
arguments, which members it leaves to the facade, and the module attributes it forwards to the core.
arguments and which members it leaves to the facade.
"""
import pytest
import litellm.router as router_module
from litellm import router_python_core
from litellm.router import Router, _class_descriptor, router_arguments
@ -41,12 +39,3 @@ def test_class_descriptor_is_true_for_a_data_descriptor_on_the_class():
def test_class_descriptor_is_false_for_a_plain_method_or_a_missing_name():
assert not _class_descriptor(Router, "from_config")
assert not _class_descriptor(Router, "no_such_member")
def test_module_getattr_forwards_old_helpers_to_the_core_module():
assert router_module.__getattr__("_stream_pre_content_error") is router_python_core._stream_pre_content_error
def test_module_getattr_raises_attribute_error_for_an_unknown_name():
with pytest.raises(AttributeError):
router_module.__getattr__("no_such_router_helper")

4
tests/unit/test_router/test_router_fallbacks.py
View file

@ -1176,7 +1176,7 @@ def test_context_stream_fallback_obeys_all_output_receipts(
wrapped: bool, has_output: bool, generated_content: str, disqualifying: bool, recoverable: bool
) -> None:
from litellm.exceptions import MidStreamFallbackError
from litellm.router import _context_window_stream_fallback_error
from litellm.router_rust.invoke.streaming import _context_window_stream_fallback_error # noqa: PLC2701
original: Final = litellm.ContextWindowExceededError(message="too long", model="primary", llm_provider="openai")
error: Final = (
@ -1201,7 +1201,7 @@ def test_context_stream_fallback_obeys_all_output_receipts(
def test_context_stream_fallback_preserves_other_errors_for_existing_continuation_policy() -> None:
from litellm.exceptions import MidStreamFallbackError
from litellm.router import _context_window_stream_fallback_error
from litellm.router_rust.invoke.streaming import _context_window_stream_fallback_error # noqa: PLC2701
error: Final = MidStreamFallbackError(
message="provider unavailable", model="primary", llm_provider="openai", generated_content="already visible"

515
tests/unit/test_router/test_router_fuzz_regressions_rust_only.py Normal file
View file

@ -0,0 +1,515 @@
"""
Router behavior the router-code differential fuzz run found changed once the Rust core became the
only core: count settings below zero, an allowed_fails_policy write that is no policy, a model
response object as a deployment's mock_response, and endpoints read or patched on the Router class.
"""
import asyncio
import gc
import json
import socket
import threading
import weakref
from collections.abc import Iterator, Mapping
from datetime import datetime
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Final
from unittest import mock
import pytest
import litellm
from litellm import Router
from litellm.router import RetryPolicy
from litellm.types.router import Deployment
from litellm.types.utils import ModelResponse
MESSAGES: Final = [{"role": "user", "content": "hi"}]
def _deployment(name: str, reply: object) -> dict[str, object]:
return {"model_name": name, "litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "mock_response": reply}}
@pytest.mark.asyncio
@pytest.mark.parametrize("setting", ["num_retries", "max_fallbacks", "allowed_fails"])
async def test_a_negative_count_setting_builds_and_serves(setting: str) -> None:
router: Final = Router(model_list=[_deployment("g", "ok")], **{setting: -1})
response: Final = await router.acompletion(model="g", messages=MESSAGES)
assert response.choices[0].message.content == "ok"
@pytest.mark.asyncio
async def test_a_negative_max_fallbacks_takes_no_fallback() -> None:
router: Final = Router(
model_list=[_deployment("g", Exception("boom")), _deployment("fb", "fb")],
fallbacks=[{"g": ["fb"]}],
max_fallbacks=-1,
num_retries=0,
)
with pytest.raises(litellm.InternalServerError):
await router.acompletion(model="g", messages=MESSAGES)
@pytest.mark.asyncio
async def test_a_negative_num_retries_still_falls_back() -> None:
router: Final = Router(
model_list=[_deployment("g", Exception("boom")), _deployment("fb", "fb")],
fallbacks=[{"g": ["fb"]}],
num_retries=-1,
)
response: Final = await router.acompletion(model="g", messages=MESSAGES)
assert response.choices[0].message.content == "fb"
@pytest.mark.asyncio
async def test_an_allowed_fails_policy_written_as_a_list_reads_as_no_policy() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")])
with pytest.warns(DeprecationWarning, match="allowed_fails_policy"):
router.allowed_fails_policy = [] # pyright: ignore[reportAttributeAccessIssue] # the old attribute write
response: Final = await router.acompletion(model="g", messages=MESSAGES)
assert response.choices[0].message.content == "ok"
assert router.settings().allowed_fails_policy == []
def test_a_model_response_mock_reads_back_from_the_model_list_as_a_dict() -> None:
reply: Final = ModelResponse(id="chatcmpl-fixed", choices=[{"message": {"role": "assistant", "content": "fb"}}])
router: Final = Router(model_list=[_deployment("g", reply)])
params: Final = router.get_model_list()[0]["litellm_params"]
assert isinstance(params["mock_response"], dict)
assert params["mock_response"]["id"] == "chatcmpl-fixed"
def test_an_endpoint_the_core_serves_by_name_is_readable_on_the_class() -> None:
endpoint: Final = Router.aretrieve_batch
assert callable(endpoint)
assert endpoint is Router.aretrieve_batch
assert endpoint.__name__ == "aretrieve_batch"
@pytest.mark.asyncio
async def test_a_class_level_patch_of_an_endpoint_reaches_the_instance_and_unwinds() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")])
with mock.patch.object(Router, "asearch", new=mock.AsyncMock(return_value="patched")) as patched:
assert await router.asearch(query="q", search_tool_name="t") == "patched"
patched.assert_awaited_once()
assert router.asearch is not patched
def test_an_assigned_affinity_ttl_builds_the_affinity_check_added_afterwards() -> None:
from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck
router: Final = Router(model_list=[_deployment("g", "ok")])
router.deployment_affinity_ttl_seconds = 7
router.add_optional_pre_call_checks(["deployment_affinity"])
checks: Final = [c for c in litellm.callbacks if isinstance(c, DeploymentAffinityCheck)]
assert [check.ttl_seconds for check in checks][-1:] == [7]
@pytest.mark.asyncio
async def test_a_deleted_timeout_leaves_get_settings_and_still_serves() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")], num_retries=0)
del router.timeout
assert "timeout" not in router.get_settings()
response: Final = await router.acompletion(model="g", messages=MESSAGES)
assert response.choices[0].message.content == "ok"
def test_litellm_params_json_cannot_carry_read_back_as_given_and_hash_like_python() -> None:
from datetime import datetime
from litellm.router_api import deployment_id
params: Final = {"model": "openai/gpt-4o", "api_key": "k", "cfg": {"s": {"a"}}, "when": datetime(2026, 1, 1, 12)}
router: Final = Router(model_list=[{"model_name": "g", "litellm_params": dict(params)}])
served: Final = router.get_model_list()[0]
assert served["litellm_params"]["cfg"] == {"s": {"a"}}
assert served["litellm_params"]["when"] == datetime(2026, 1, 1, 12)
assert served["model_info"]["id"] == deployment_id("g", params)
def test_a_bad_deployment_before_a_secret_reference_raises_its_own_validation_error() -> None:
from pydantic import ValidationError
with pytest.raises(ValidationError):
Router(
model_list=[
{"model_name": "t", "litellm_params": {"api_key": "k"}},
{"model_name": "u", "litellm_params": {"model": "openai/gpt-4o", "api_key": "os.environ/NOPE_KEY"}},
]
)
@pytest.mark.asyncio
async def test_a_failed_vector_store_create_names_no_failed_deployment() -> None:
router: Final = Router(
model_list=[
{
"model_name": "vs",
"litellm_params": {"model": "openai/gpt-4o", "api_key": "k", "api_base": "http://127.0.0.1:9/v1"},
}
],
num_retries=0,
)
with pytest.raises(Exception) as raised: # noqa: PT011 # the provider's connection error, whatever class maps it
await router.avector_store_create(model="vs", name="x")
assert getattr(raised.value, "failed_deployment_id", None) is None
class _Upstream(BaseHTTPRequestHandler):
"""An OpenAI upstream whose path's first segment names how it answers: ``cut`` sends two
Responses stream events then cuts the chunked body, ``filtered`` answers a chat completion the
provider blocked, ``down`` answers 502, and anything else a plain chat completion."""
protocol_version = "HTTP/1.1"
def log_message(self, format: str, *args: object) -> None: # noqa: A002 # the base class's name
del format, args
def _json(self, status: int, body: Mapping[str, object]) -> None:
data: Final = json.dumps(body).encode()
self.send_response(status)
self.send_header("content-type", "application/json")
self.send_header("content-length", str(len(data)))
self.end_headers()
_ = self.wfile.write(data)
def _chunk(self, event: Mapping[str, object]) -> None:
data: Final = f"event: {event['type']}\ndata: {json.dumps(event)}\n\n".encode()
_ = self.wfile.write(f"{len(data):x}\r\n".encode() + data + b"\r\n")
self.wfile.flush()
def _cut(self) -> None:
self.send_response(200)
self.send_header("content-type", "text/event-stream")
self.send_header("transfer-encoding", "chunked")
self.end_headers()
response: Final = {"id": "r1", "object": "response", "created_at": 1, "status": "in_progress", "output": []}
self._chunk({"type": "response.created", "sequence_number": 0, "response": {**response, "model": "gpt-4o"}})
self._chunk(
{
"type": "response.output_text.delta",
"sequence_number": 1,
"item_id": "m1",
"output_index": 0,
"content_index": 0,
"delta": "hello ",
}
)
self.connection.shutdown(socket.SHUT_RDWR)
self.close_connection = True
def do_POST(self) -> None:
_ = self.rfile.read(int(self.headers.get("content-length") or 0))
mode: Final = self.path.split("/")[1]
if mode == "cut":
self._cut()
return
if mode == "down":
self._json(502, {"error": {"message": "provider error", "type": "api_error"}})
return
finish: Final = "content_filter" if mode == "filtered" else "stop"
message: Final = {"role": "assistant", "content": mode}
choice: Final = {"index": 0, "finish_reason": finish, "message": message}
usage: Final = {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}
self._json(200, {"id": "c1", "object": "chat.completion", "created": 1, "choices": [choice], "usage": usage})
@pytest.fixture
def upstream() -> Iterator[str]:
server: Final = ThreadingHTTPServer(("127.0.0.1", 0), _Upstream)
thread: Final = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
yield f"http://127.0.0.1:{server.server_address[1]}"
server.shutdown()
thread.join()
def _served(name: str, deployment_id: str, model: str, api_base: str) -> dict[str, object]:
return {
"model_name": name,
"litellm_params": {"model": model, "api_key": "k", "api_base": api_base},
"model_info": {"id": deployment_id},
}
def _request_logging(call_type: str, *, stream: bool) -> object:
from litellm.litellm_core_utils.litellm_logging import Logging
return Logging(
model="g",
messages=MESSAGES,
stream=stream,
call_type=call_type,
start_time=datetime.now(), # noqa: DTZ005 # litellm's own naive start time
litellm_call_id="call-1",
function_id="f",
)
async def _recorded() -> None:
"""Let the background recordings of a stream's outcome run."""
for _ in range(50):
await asyncio.sleep(0.02)
@pytest.mark.asyncio
async def test_a_responses_stream_cut_after_it_opened_counts_one_fail_with_a_shared_logging_object(
upstream: str,
) -> None:
router: Final = Router(
model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/cut/v1")], num_retries=0, disable_cooldowns=True
)
stream: Final = await router.aresponses(
model="g", input="hi", stream=True, litellm_logging_obj=_request_logging("aresponses", stream=True)
)
with pytest.raises(Exception): # noqa: B017, PT011 # the transport's read error, whatever class maps it
async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True
pass
await _recorded()
assert router.health().store.get_cache(key="d1:fails") == 1
@pytest.mark.asyncio
async def test_a_blocked_response_raised_for_a_fallback_counts_no_success_with_a_shared_logging_object(
upstream: str,
) -> None:
router: Final = Router(
model_list=[
_served("a", "a1", "openai/gpt-4o", f"{upstream}/filtered/v1"),
_served("safe", "s1", "openai/gpt-4o", f"{upstream}/ok/v1"),
],
content_policy_fallbacks=[{"a": ["safe"]}],
num_retries=0,
disable_cooldowns=True,
)
response: Final = await router.acompletion(
model="a", messages=MESSAGES, litellm_logging_obj=_request_logging("acompletion", stream=False)
)
await _recorded()
assert response.choices[0].message.content == "ok"
assert router.health().store.get_cache(key="a1:successes") is None
assert router.health().store.get_cache(key="s1:successes") == 1
@pytest.mark.asyncio
async def test_a_text_completion_stream_that_fails_before_its_first_chunk_counts_nothing(upstream: str) -> None:
router: Final = Router(
model_list=[_served("g", "d1", "openai/gpt-3.5-turbo-instruct", f"{upstream}/down/v1")],
num_retries=0,
disable_cooldowns=True,
)
stream: Final = await router.atext_completion(model="g", prompt="hi", stream=True)
with pytest.raises(litellm.BadGatewayError):
async for _ in stream: # pyright: ignore[reportGeneralTypeIssues] # a stream when stream=True
pass
await _recorded()
assert router.health().store.get_cache(key="d1:successes") is None
assert router.health().store.get_cache(key="d1:fails") is None
@pytest.mark.asyncio
async def test_a_text_completion_without_a_prompt_raises_type_error_before_routing() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")], num_retries=0)
with pytest.raises(TypeError, match="'prompt'"):
await router.atext_completion(model="g") # pyright: ignore[reportCallIssue] # the missing argument
with pytest.raises(TypeError, match="'prompt'"):
router.text_completion(model="g") # pyright: ignore[reportCallIssue] # the missing argument
@pytest.mark.asyncio
async def test_a_default_priority_assigned_after_construction_reads_back_and_queues_calls() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")], num_retries=0)
router.default_priority = 5
response: Final = await router.acompletion(model="g", messages=MESSAGES)
assert router.default_priority == 5
assert response._hidden_params["additional_headers"]["x-litellm-request-prioritization-used"] is True
@pytest.mark.asyncio
async def test_a_prioritized_call_on_an_unknown_group_raises_the_calls_own_error() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")], num_retries=0)
with pytest.raises(litellm.InternalServerError):
await router.schedule_acompletion(
model="nope", messages=MESSAGES, priority=0, mock_testing_fallbacks=True, num_retries=0
)
def test_a_registry_lookup_of_a_record_with_no_model_name_raises_key_error() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")])
with pytest.raises(KeyError):
router.registry().includes({"litellm_params": {}}, model_name="g") # pyright: ignore[reportArgumentType] # the record with no name
@pytest.mark.asyncio
async def test_a_refused_pick_by_deployment_id_names_the_deployments_group() -> None:
router: Final = Router(
model_list=[
{
"model_name": "t2-pub",
"litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "k", "mock_response": "p2"},
"model_info": {"team_id": "T2", "team_public_model_name": "pub", "id": "p2"},
}
],
num_retries=0,
)
with pytest.raises(Exception, match="model=t2-pub"): # noqa: PT011 # the refusal, whatever class names it
await router.async_get_available_deployment(model="p2", messages=MESSAGES, request_kwargs={})
@pytest.mark.asyncio
async def test_a_request_refused_before_the_provider_counts_against_no_deployment(upstream: str) -> None:
router: Final = Router(
model_list=[_served("g", "d1", "openai/gpt-4o", f"{upstream}/ok/v1")], num_retries=0, disable_cooldowns=True
)
with pytest.raises(litellm.UnsupportedParamsError):
await router.acompletion(model="g", messages=MESSAGES, reasoning_effort="low")
await _recorded()
assert router.health().store.get_cache(key="d1:fails") is None
assert not [key for key in router.health().store.in_memory_cache.cache_dict if ":rpm" in key]
@pytest.mark.asyncio
async def test_the_async_pick_read_on_the_class_is_a_function_taking_the_router() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")])
pick: Final = Router.async_get_available_deployment
assert pick.__name__ == "async_get_available_deployment"
assert (await pick(router, "g", {}))["model_name"] == "g"
def _sharing_ids(*ids: str) -> Router:
return Router(
model_list=[{**_deployment("g", f"from-{index}"), "model_info": {"id": id}} for index, id in enumerate(ids)]
)
@pytest.mark.asyncio
async def test_async_picks_split_over_deployments_sharing_an_id() -> None:
router: Final = _sharing_ids("d1", "d1", "b")
served: Final = {
(await router.acompletion(model="g", messages=MESSAGES)).choices[0].message.content for _ in range(60)
}
assert served == {"from-0", "from-1", "from-2"}
def test_sync_picks_split_over_deployments_sharing_an_id() -> None:
router: Final = _sharing_ids("d1", "d1")
served: Final = {router.completion(model="g", messages=MESSAGES).choices[0].message.content for _ in range(40)}
assert served == {"from-0", "from-1"}
@pytest.mark.asyncio
async def test_no_model_is_served_by_the_first_default_fallback() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")], default_fallbacks=["g"])
async_served: Final = await router.acompletion(model=None, messages=MESSAGES)
sync_served: Final = router.completion(model=None, messages=MESSAGES)
assert (async_served.choices[0].message.content, sync_served.choices[0].message.content) == ("ok", "ok")
@pytest.mark.asyncio
async def test_no_model_is_served_by_a_catch_all_wildcard() -> None:
router: Final = Router(
model_list=[
_deployment("g", "ok"),
{
**_deployment("*", "star"),
"litellm_params": {"model": "openai/*", "api_key": "k", "mock_response": "star"},
},
]
)
served: Final = await router.acompletion(model=None, messages=MESSAGES)
assert served.choices[0].message.content == "star"
@pytest.mark.asyncio
async def test_no_model_is_retried_by_the_retry_policy_and_raised_without_a_fallback_note() -> None:
router: Final = Router(
model_list=[_deployment("g", "ok")],
default_fallbacks=["h"],
retry_after=0,
retry_policy=RetryPolicy(BadRequestErrorRetries=1),
)
with pytest.raises(litellm.BadRequestError) as raised:
await router.acompletion(model=None, messages=MESSAGES)
assert raised.value.num_retries == 1
assert "You passed in model=h. There are no healthy deployments for this model" in str(raised.value)
assert "failed with the error above" not in str(raised.value)
class _TokenProvider:
def __call__(self) -> str:
return "token"
def test_a_deleted_deployment_releases_the_callable_it_was_built_with() -> None:
router: Final = Router(model_list=[_deployment("g", "ok")])
provider = _TokenProvider()
released: Final = weakref.ref(provider)
router.add_deployment(
Deployment(
model_name="g",
litellm_params={"model": "openai/gpt-4o", "api_key": "k", "azure_ad_token_provider": provider},
model_info={"id": "rotated"},
)
)
del provider
router.delete_deployment("rotated")
gc.collect()
assert released() is None
@pytest.mark.asyncio
async def test_a_pick_keeps_model_info_keys_that_are_not_strings() -> None:
router: Final = Router(model_list=[{**_deployment("g", "ok"), "model_info": {"id": "d1", "tiers": {1: "a"}}}])
picked: Final = await router.async_get_available_deployment(model="g", messages=MESSAGES, request_kwargs={})
assert picked["model_info"]["tiers"] == {1: "a"}

60
tests/unit/test_router/test_router_get_deployments.py
View file

@ -213,66 +213,6 @@ async def test_async_get_available_deployment_for_pass_through():
pytest.fail(f"Error occurred: {e}")
def test_filter_pass_through_deployments():
"""
Test _filter_pass_through_deployments function
- Tests that it correctly filters deployments with use_in_pass_through=True
"""
try:
litellm.set_verbose = False
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
"use_in_pass_through": True,
},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"use_in_pass_through": False,
},
},
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-35-turbo",
"api_key": os.getenv("AZURE_API_KEY"),
"api_base": os.getenv("AZURE_API_BASE"),
"api_version": os.getenv("AZURE_API_VERSION"),
"use_in_pass_through": True,
},
},
]
router = Router(
model_list=model_list,
)
# Get all healthy deployments
healthy_deployments = router.get_model_list()
# Filter pass-through deployments
pass_through_deployments = router._filter_pass_through_deployments(healthy_deployments)
# Should only have 2 deployments with use_in_pass_through=True
assert len(pass_through_deployments) == 2
# Verify all returned deployments have use_in_pass_through=True
for deployment in pass_through_deployments:
assert deployment["litellm_params"]["use_in_pass_through"] is True
router.reset()
except Exception as e:
traceback.print_exc()
pytest.fail(f"Error occurred: {e}")
async def _weighted_async_deployment_ids(router: Router) -> tuple[str, ...]:
random.seed(2025)
deployments: Final = tuple(

2090
tests/unit/test_router/test_router_helper_utils.py
View file

File diff suppressed because it is too large Load diff

48
tests/unit/test_router/test_router_max_parallel_requests.py
View file

@ -4,7 +4,6 @@ from typing import Final
import litellm
import pytest
from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit
from litellm.utils import calculate_max_parallel_requests
@ -36,53 +35,6 @@ def test_scenario(
assert calculated == expected
@pytest.mark.parametrize(
("max_parallel_requests", "tpm", "rpm", "default_max_parallel_requests", "expected"),
[
pytest.param(10, 300_000, 30, 40, 10, id="explicit-limit"),
pytest.param(None, 300_000, 30, 40, 30, id="rpm"),
pytest.param(None, 300_000, None, 40, 1_800, id="tpm"),
pytest.param(None, 20, None, 40, 1, id="minimum-tpm-limit"),
pytest.param(None, None, None, 40, 40, id="router-default"),
pytest.param(None, None, None, None, None, id="unset"),
],
)
def test_setting_mpr_limits_per_model(
max_parallel_requests: int | None,
tpm: int | None,
rpm: int | None,
default_max_parallel_requests: int | None,
expected: int | None,
) -> None:
deployment: Final = {
"model_name": "limited",
"litellm_params": {
"model": "openai/gpt-4o-mini",
"api_key": "test-key",
"max_parallel_requests": max_parallel_requests,
"tpm": tpm,
"rpm": rpm,
},
"model_info": {"id": "limited-deployment"},
}
router: Final = litellm.Router(
model_list=[deployment],
default_max_parallel_requests=default_max_parallel_requests,
)
limit: Final = router._get_client(
deployment=deployment,
kwargs={},
client_type="max_parallel_requests",
)
if expected is None:
assert limit is None
return
assert isinstance(limit, MaxParallelRequestsLimit)
assert limit.max_parallel_requests == expected
@pytest.mark.asyncio
async def test_max_parallel_requests_rpm_rate_limiting() -> None:
router: Final = litellm.Router(

101
tests/unit/test_router/test_router_provider_endpoints.py
View file

@ -2,6 +2,7 @@ import asyncio
import contextlib
import io
import json
import time
import uuid
from collections.abc import Callable, Iterator
from datetime import datetime, timezone
@ -51,9 +52,11 @@ def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
@pytest.fixture
def router_minute_pinned(monkeypatch: pytest.MonkeyPatch) -> None:
pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc)
monkeypatch.setattr("litellm.router_python_core.get_utc_datetime", lambda: pinned)
def router_minute_pinned() -> None:
"""The usage counters these tests read are per minute: start well inside one so a test never spans two."""
seconds_left: Final = 60 - datetime.now(timezone.utc).second
if seconds_left < 10:
time.sleep(seconds_left + 0.1)
class _RouterLoggingCapture(CustomLogger):
@ -145,14 +148,14 @@ async def test_router_transcription_reuses_router_level_client_for_each_deployme
model="whisper", file=("speech.wav", io.BytesIO(b"offline audio"), "audio/wav")
)
public_client, public_logging = await capture.next_event()
internal_response: Final = await router._atranscription(
second_response: Final = await router.atranscription(
model="whisper", file=("speech.wav", io.BytesIO(b"offline audio"), "audio/wav")
)
internal_client, _ = await capture.next_event()
second_client, _ = await capture.next_event()
upstream_requests: Final = tuple(call.request for call in route.calls)
assert public_client == str(router_level_client)
assert internal_client == str(router_level_client)
assert second_client == str(router_level_client)
assert tuple(str(request.url) for request in upstream_requests) == (route_url, route_url)
assert all(
request.headers.get("content-type", "").startswith("multipart/form-data")
@ -163,7 +166,7 @@ async def test_router_transcription_reuses_router_level_client_for_each_deployme
assert isinstance(public_logging, dict)
assert public_logging.get("model_group") == "whisper"
assert response.text == "hello"
assert internal_response.text == "hello"
assert second_response.text == "hello"
@pytest.mark.asyncio
@ -197,7 +200,7 @@ async def test_router_speech_returns_binary_content_and_logs_model_group(
@pytest.mark.asyncio
async def test_router_rerank_returns_valid_response_from_public_and_underlying_calls(
async def test_router_rerank_returns_valid_response_on_every_call(
respx_mock: respx.MockRouter,
) -> None:
router: Final = Router(
@ -223,7 +226,7 @@ async def test_router_rerank_returns_valid_response_from_public_and_underlying_c
documents=["hello", "world"],
top_n=1,
)
underlying_response: Final = await router._arerank(
second_response: Final = await router.arerank(
model="cohere-rerank",
query="hello",
documents=["hello", "world"],
@ -242,9 +245,9 @@ async def test_router_rerank_returns_valid_response_from_public_and_underlying_c
public_validated: Final = RerankResponse.model_validate(public_response)
assert public_validated.id == "rerank-1"
assert public_validated.results[0]["relevance_score"] == 0.9
underlying_validated: Final = RerankResponse.model_validate(underlying_response)
assert underlying_validated.id == "rerank-1"
assert underlying_validated.results[0]["relevance_score"] == 0.9
second_validated: Final = RerankResponse.model_validate(second_response)
assert second_validated.id == "rerank-1"
assert second_validated.results[0]["relevance_score"] == 0.9
@pytest.mark.asyncio
@ -317,9 +320,9 @@ async def test_router_image_generation_returns_valid_image_response(
)
response: Final = (
router._image_generation(model="gpt-image-1", prompt="a cat")
router.image_generation(model="gpt-image-1", prompt="a cat")
if sync_mode
else await router._aimage_generation(model="gpt-image-1", prompt="a cat")
else await router.aimage_generation(model="gpt-image-1", prompt="a cat")
)
assert route.call_count == 1
@ -516,24 +519,6 @@ def test_router_routing_strategy_init_accepts_provider_budget_strategy(strategy:
assert router.get_settings()["routing_strategy"] == "provider-budget-routing"
def test_router_track_deployment_metrics_updates_observable_usage() -> None:
router: Final = Router(
model_list=[
{
"model_name": "gpt-5-mini",
"litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake"},
"model_info": {"id": "metrics-deployment"},
}
]
)
deployment: Final = router.model_list[0]
router._track_deployment_metrics(deployment=deployment, parent_otel_span=None)
router._track_deployment_metrics(deployment=deployment, parent_otel_span=None)
assert router.cache.get_cache(key="metrics-deployment", local_only=True) == 2
class _GatedIncrementCache(DualCache):
def __init__(self) -> None:
super().__init__(in_memory_cache=InMemoryCache())
@ -548,6 +533,12 @@ class _GatedIncrementCache(DualCache):
self.deployment_success_incremented.set()
return result
async def async_increment_cache(self, key: str, value: float, *args: object, **kwargs: object) -> float | None:
result: Final = await super().async_increment_cache(key, value, *args, **kwargs)
if key.endswith(":successes"):
self.deployment_success_incremented.set()
return result
async def async_increment_cache_pipeline(
self,
increment_list: list[RedisPipelineIncrementOperation],
@ -581,6 +572,12 @@ class _UnavailableIncrementCache(DualCache):
self.deployment_success_incremented.set()
return result
async def async_increment_cache(self, key: str, value: float, *args: object, **kwargs: object) -> float | None:
result: Final = await super().async_increment_cache(key, value, *args, **kwargs)
if key.endswith(":successes"):
self.deployment_success_incremented.set()
return result
async def async_increment_cache_pipeline(
self,
increment_list: list[RedisPipelineIncrementOperation],
@ -637,7 +634,6 @@ async def test_router_failed_pre_header_increment_clears_counted_tokens_stamp()
await asyncio.wait_for(cache.first_increment_started.wait(), timeout=EVENT_TIMEOUT_SECONDS)
await asyncio.wait_for(cache.deployment_success_incremented.wait(), timeout=EVENT_TIMEOUT_SECONDS)
assert metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] == 30
assert get_deployment_successes_for_current_minute(router, "lit-3058-fail") == 1
assert cache.increment_calls == 1
cache.release_first_increment.set()
@ -662,7 +658,7 @@ ASSISTANT_RESPONSE: Final = {
@pytest.mark.asyncio
async def test_router_assistants_endpoint_factory_invokes_provider(
async def test_router_create_assistants_invokes_provider(
respx_mock: respx.MockRouter,
) -> None:
router: Final = Router(model_list=[])
@ -670,8 +666,7 @@ async def test_router_assistants_endpoint_factory_invokes_provider(
200, json={**ASSISTANT_RESPONSE, "id": "asst-1"}
)
response: Final = await router._pass_through_assistants_endpoint_factory(
original_function=litellm.acreate_assistants,
response: Final = await router.acreate_assistants(
custom_llm_provider="openai",
model="gpt-4o-mini",
api_key="sk-fake",
@ -684,29 +679,7 @@ async def test_router_assistants_endpoint_factory_invokes_provider(
@pytest.mark.asyncio
async def test_router_factory_function_returns_invokable_assistants_wrapper(
respx_mock: respx.MockRouter,
) -> None:
router: Final = Router(model_list=[])
route: Final = respx_mock.post("https://api.openai.com/v1/assistants").respond(
200, json={**ASSISTANT_RESPONSE, "id": "asst-2"}
)
wrapper: Final = router.factory_function(litellm.acreate_assistants, call_type="assistants")
response: Final = await wrapper(
custom_llm_provider="openai",
model="gpt-4o-mini",
api_key="sk-fake",
name="offline",
)
assert route.call_count == 1
assert json.loads(route.calls[0].request.content) == {"model": "gpt-4o-mini", "name": "offline"}
assert response.id == "asst-2"
@pytest.mark.asyncio
async def test_router_moderation_endpoint_factory_invokes_default_model(
async def test_router_moderation_without_a_model_invokes_the_default_model(
respx_mock: respx.MockRouter,
monkeypatch: pytest.MonkeyPatch,
) -> None:
@ -721,9 +694,7 @@ async def test_router_moderation_endpoint_factory_invokes_default_model(
},
)
response: Final = await router._pass_through_moderation_endpoint_factory(
original_function=litellm.amoderation,
custom_llm_provider="openai",
response: Final = await router.amoderation(
input="hello",
model=None,
api_key="sk-fake",
@ -801,7 +772,11 @@ async def test_video_status_on_a_wildcard_polls_the_deployment_that_created_the_
) -> None:
router: Final = Router(
model_list=[
{"model_name": "openai/*", "litellm_params": {"model": "openai/*", "api_key": key}, "model_info": {"id": dep}}
{
"model_name": "openai/*",
"litellm_params": {"model": "openai/*", "api_key": key},
"model_info": {"id": dep},
}
for dep, key in (("wc-a", "key-a"), ("wc-b", "key-b"))
]
)

284
tests/unit/test_router/test_router_retries.py
View file

@ -169,9 +169,7 @@ async def test_router_retries_errors(sync_mode: bool, error_type: str, monkeypat
async def test_router_retry_policy(error_type):
from litellm.router import AllowedFailsPolicy, RetryPolicy
retry_policy = RetryPolicy(
ContentPolicyViolationErrorRetries=3, AuthenticationErrorRetries=0
)
retry_policy = RetryPolicy(ContentPolicyViolationErrorRetries=3, AuthenticationErrorRetries=0)
allowed_fails_policy = AllowedFailsPolicy(
ContentPolicyViolationErrorAllowedFails=1000,
@ -361,9 +359,7 @@ def test_retry_rate_limit_error_with_healthy_deployments():
# Act & Assert
try:
response = router.should_retry_this_error(
error=rate_limit_error, healthy_deployments=healthy_deployments
)
response = router.should_retry_this_error(error=rate_limit_error, healthy_deployments=healthy_deployments)
print("response from should_retry_this_error: ", response)
except Exception as e:
pytest.fail(
@ -394,9 +390,7 @@ def test_do_retry_rate_limit_error_with_no_fallbacks_and_no_healthy_deployments(
# Act & Assert
try:
response = router.should_retry_this_error(
error=rate_limit_error, healthy_deployments=healthy_deployments
)
response = router.should_retry_this_error(error=rate_limit_error, healthy_deployments=healthy_deployments)
pytest.fail("Should have raised an error")
except Exception as e:
print("got an exception", e)
@ -438,9 +432,7 @@ def test_raise_context_window_exceeded_error():
healthy_deployments=None,
context_window_fallbacks=context_window_fallbacks,
)
pytest.fail(
"Expected to raise context window exceeded error -> trigger fallback"
)
pytest.fail("Expected to raise context window exceeded error -> trigger fallback")
except Exception as e:
pass
@ -480,116 +472,11 @@ def test_raise_context_window_exceeded_error_no_retry():
healthy_deployments=None,
context_window_fallbacks=context_window_fallbacks,
)
assert (
response == True
), "Should not have raised exception since we do not have context window fallbacks"
assert response == True, "Should not have raised exception since we do not have context window fallbacks"
except litellm.ContextWindowExceededError:
pass
@pytest.mark.parametrize("num_deployments, expected_timeout", [(1, 60), (2, 0.0)])
def test_timeout_for_rate_limit_error_with_healthy_deployments(
num_deployments, expected_timeout
):
"""
Test 1. Timeout is 0.0 when RateLimit Error and healthy deployments are > 0
"""
cooldown_time = 60
rate_limit_error = litellm.RateLimitError(
message="{RouterErrors.no_deployments_available.value}. 12345 Passed model={model_group}. Deployments={deployment_dict}",
llm_provider="",
model="gpt-3.5-turbo",
response=httpx.Response(
status_code=429,
content="",
headers={"retry-after": str(cooldown_time)}, # type: ignore
request=httpx.Request(method="tpm_rpm_limits", url="https://github.com/BerriAI/litellm"), # type: ignore
),
)
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
},
}
]
if num_deployments == 2:
model_list.append(
{
"model_name": "gpt-4",
"litellm_params": {"model": "gpt-3.5-turbo"},
}
)
router = litellm.Router(model_list=model_list)
_timeout = router._time_to_sleep_before_retry(
e=rate_limit_error,
remaining_retries=2,
num_retries=2,
healthy_deployments=[
{
"model_name": "gpt-4",
"litellm_params": {
"api_key": "my-key",
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com",
"model": "azure/gpt-4.1-mini",
},
"model_info": {
"id": "0e30bc8a63fa91ae4415d4234e231b3f9e6dd900cac57d118ce13a720d95e9d6",
"db_model": False,
},
}
],
all_deployments=model_list,
)
if expected_timeout == 0.0:
assert _timeout == expected_timeout
else:
assert _timeout > 0.0
def test_timeout_for_rate_limit_error_with_no_healthy_deployments():
"""
Test 2. Timeout is > 0.0 when RateLimit Error and healthy deployments == 0
"""
healthy_deployments = []
model_list = [
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "azure/gpt-4.1-mini",
"api_key": os.getenv("AZURE_API_KEY"),
"api_version": os.getenv("AZURE_API_VERSION"),
"api_base": os.getenv("AZURE_API_BASE"),
},
}
]
router = litellm.Router(model_list=model_list)
_timeout = router._time_to_sleep_before_retry(
e=rate_limit_error,
remaining_retries=4,
num_retries=4,
healthy_deployments=healthy_deployments,
all_deployments=model_list,
)
print(
"timeout=",
_timeout,
"error is rate_limit_error and there are no healthy deployments",
)
assert _timeout > 0.0
def test_no_retry_for_not_found_error_404():
healthy_deployments = []
@ -614,9 +501,7 @@ def test_no_retry_for_not_found_error_404():
llm_provider="azure",
)
try:
response = router.should_retry_this_error(
error=error, healthy_deployments=healthy_deployments
)
response = router.should_retry_this_error(error=error, healthy_deployments=healthy_deployments)
pytest.fail(
"Should have raised an exception 404 NotFoundError should never be retried, it's typically model_not_found error"
)
@ -652,12 +537,8 @@ def test_no_retry_for_bad_request_error_400():
llm_provider="azure",
)
try:
response = router.should_retry_this_error(
error=error, healthy_deployments=healthy_deployments
)
pytest.fail(
"Should have raised BadRequestError - 400 errors should never be retried"
)
response = router.should_retry_this_error(error=error, healthy_deployments=healthy_deployments)
pytest.fail("Should have raised BadRequestError - 400 errors should never be retried")
except litellm.BadRequestError as e:
print("Correctly raised BadRequestError without retry:", e)
@ -693,12 +574,8 @@ def test_no_retry_for_unprocessable_entity_error_422():
),
)
try:
response = router.should_retry_this_error(
error=error, healthy_deployments=healthy_deployments
)
pytest.fail(
"Should have raised UnprocessableEntityError - 422 errors should never be retried"
)
response = router.should_retry_this_error(error=error, healthy_deployments=healthy_deployments)
pytest.fail("Should have raised UnprocessableEntityError - 422 errors should never be retried")
except litellm.UnprocessableEntityError as e:
print("Correctly raised UnprocessableEntityError without retry:", e)
@ -727,9 +604,7 @@ def test_no_retry_when_no_healthy_deployments():
timeout_error,
]:
try:
response = router.should_retry_this_error(
error=error, healthy_deployments=healthy_deployments
)
response = router.should_retry_this_error(error=error, healthy_deployments=healthy_deployments)
pytest.fail(
"Should have raised an exception, there's no point retrying an error when there are 0 healthy deployments"
)
@ -739,8 +614,6 @@ def test_no_retry_when_no_healthy_deployments():
@pytest.mark.asyncio
async def test_router_retries_model_specific_and_global():
from unittest.mock import MagicMock, patch
litellm.num_retries = 0
router = Router(
model_list=[
@ -748,28 +621,21 @@ async def test_router_retries_model_specific_and_global():
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
"api_key": "sk-test",
"num_retries": 1,
},
}
]
)
with patch.object(
router, "_time_to_sleep_before_retry"
) as mock_async_function_with_retries:
try:
await router.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello, how are you?"}],
mock_response="litellm.RateLimitError",
)
except Exception as e:
print("got exception", e)
with pytest.raises(litellm.RateLimitError) as exc_info:
await router.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello, how are you?"}],
mock_response="litellm.RateLimitError",
)
mock_async_function_with_retries.assert_called_once()
assert mock_async_function_with_retries.call_args.kwargs["num_retries"] == 1
assert exc_info.value.num_retries == 1
@pytest.mark.usefixtures("fake_provider_credentials")
@ -817,68 +683,28 @@ async def test_router_retry_num_retries_tracking():
This verifies the fix for the bug where num_retries was incorrectly set to current_attempt
(0-indexed) instead of the actual number of retries attempted.
"""
from unittest.mock import AsyncMock, patch
router = Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"},
}
],
num_retries=3, # Set at router level to ensure it's used
num_retries=2,
)
# Mock make_call to always raise a RateLimitError
async def mock_make_call(*args, **kwargs):
raise litellm.RateLimitError(
message="Rate limit exceeded",
with pytest.raises(litellm.RateLimitError) as exc_info:
await router.acompletion(
model="gpt-3.5-turbo",
llm_provider="openai",
messages=[{"role": "user", "content": "Hello"}],
mock_response="litellm.RateLimitError",
)
with patch.object(router, "make_call", side_effect=mock_make_call):
with patch.object(
router,
"_async_get_healthy_deployments",
return_value=(
[{"model_info": {"id": "test-id"}}],
[{"model_info": {"id": "test-id"}}],
),
):
with patch.object(
router, "_time_to_sleep_before_retry", return_value=0.01
): # Fast retries for testing
with pytest.raises(litellm.RateLimitError) as exc_info:
await router.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello"}],
)
e = exc_info.value
assert hasattr(
e, "num_retries"
), "Exception should have num_retries attribute"
assert hasattr(
e, "max_retries"
), "Exception should have max_retries attribute"
assert (
e.num_retries == 3
), f"Expected num_retries to be 3, got {e.num_retries}"
assert (
e.max_retries == 3
), f"Expected max_retries to be 3, got {e.max_retries}"
# Verify the error message includes correct retry information
error_str = str(e)
assert (
"LiteLLM Retried: 3 times" in error_str
), f"Error message should indicate 3 retries: {error_str}"
assert (
"LiteLLM Max Retries: 3" in error_str
), f"Error message should show max retries: {error_str}"
e = exc_info.value
assert e.num_retries == 2
assert e.max_retries == 2
assert "LiteLLM Retried: 2 times" in str(e)
assert "LiteLLM Max Retries: 2" in str(e)
@pytest.mark.asyncio
@ -886,51 +712,25 @@ async def test_router_retry_num_retries_single_retry():
"""
Test num_retries tracking with a single retry to verify edge case handling.
"""
from unittest.mock import patch
router = Router(
model_list=[
{
"model_name": "gpt-3.5-turbo",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
},
"litellm_params": {"model": "gpt-3.5-turbo", "api_key": "sk-test"},
}
],
num_retries=1, # Set at router level - single retry
num_retries=1,
)
# Mock make_call to always raise a Timeout error
async def mock_make_call(*args, **kwargs):
raise litellm.Timeout(
message="Request timed out",
with pytest.raises(litellm.InternalServerError) as exc_info:
await router.acompletion(
model="gpt-3.5-turbo",
llm_provider="openai",
messages=[{"role": "user", "content": "Hello"}],
mock_response="litellm.InternalServerError",
)
with patch.object(router, "make_call", side_effect=mock_make_call):
with patch.object(
router,
"_async_get_healthy_deployments",
return_value=(
[{"model_info": {"id": "test-id"}}],
[{"model_info": {"id": "test-id"}}],
),
):
with patch.object(router, "_time_to_sleep_before_retry", return_value=0.01):
with pytest.raises(litellm.Timeout) as exc_info:
await router.acompletion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hello"}],
)
e = exc_info.value
assert (
e.num_retries == 1
), f"Expected num_retries to be 1, got {e.num_retries}"
assert (
e.max_retries == 1
), f"Expected max_retries to be 1, got {e.max_retries}"
assert exc_info.value.num_retries == 1
assert exc_info.value.max_retries == 1
class MyCustomHandler(CustomLogger):
@ -940,18 +740,14 @@ class MyCustomHandler(CustomLogger):
def log_pre_api_call(self, model, messages, kwargs):
print(f"Pre-API Call")
print(
f"previous_models: {kwargs['litellm_params']['metadata'].get('previous_models', None)}"
)
print(f"previous_models: {kwargs['litellm_params']['metadata'].get('previous_models', None)}")
self.previous_models = len(
kwargs["litellm_params"]["metadata"].get("previous_models", [])
) # {"previous_models": [{"model": litellm_model_name, "exception_type": AuthenticationError, "exception_string": <complete_traceback>}]}
print(f"self.previous_models: {self.previous_models}")
def log_post_api_call(self, kwargs, response_obj, start_time, end_time):
print(
f"Post-API Call - response object: {response_obj}; model: {kwargs['model']}"
)
print(f"Post-API Call - response object: {response_obj}; model: {kwargs['model']}")
def log_stream_event(self, kwargs, response_obj, start_time, end_time):
print(f"On Stream")

141
tests/unit/test_router/test_router_timeout.py
View file

@ -84,73 +84,26 @@ def test_router_timeout_with_retries_anthropic_model(num_retries, expected_call_
assert mock_client.call_count == expected_call_count
@pytest.mark.parametrize(
"stream",
[
True,
False,
],
)
def test_unit_test_streaming_timeout(stream):
import os
from dotenv import load_dotenv
import litellm
from litellm.router import Router, RetryPolicy, AllowedFailsPolicy
async def _timeout_sent_to_the_provider_call(router: Router, request_kwargs: Mapping[str, object]) -> object:
sent: Final[list[object]] = []
real_acompletion: Final = litellm.acompletion
litellm.set_verbose = True
async def spy(*args: object, **kwargs: object) -> object:
# A stream is bounded by the stream timeout litellm.acompletion is handed, any other call by its timeout
stream_timeout = kwargs.get("stream_timeout")
sent.append(
stream_timeout if kwargs.get("stream") is True and stream_timeout is not None else kwargs.get("timeout")
)
return await real_acompletion(*args, **kwargs)
model_list = [
{
"model_name": "llama3",
"litellm_params": {
"model": "watsonx/meta-llama/llama-3-1-8b-instruct",
"api_base": os.getenv("WATSONX_URL_US_SOUTH"),
"api_key": os.getenv("WATSONX_API_KEY"),
"project_id": os.getenv("WATSONX_PROJECT_ID_US_SOUTH"),
"timeout": 0.01,
"stream_timeout": 0.0000001,
},
},
{
"model_name": "bedrock-anthropic",
"litellm_params": {
"model": "bedrock/anthropic.claude-3-5-haiku-20241022-v1:0",
"timeout": 0.01,
"stream_timeout": 0.0000001,
},
},
{
"model_name": "llama3-fallback",
"litellm_params": {
"model": "gpt-3.5-turbo",
"api_key": os.getenv("OPENAI_API_KEY"),
},
},
]
router = Router(model_list=model_list)
stream_timeout = 0.0000001
normal_timeout = 0.01
args = {
"kwargs": {"stream": stream},
"data": {"timeout": normal_timeout, "stream_timeout": stream_timeout},
}
assert router._get_stream_timeout(**args) == stream_timeout
assert router._get_non_stream_timeout(**args) == normal_timeout
stream_timeout_val = router._get_timeout(
kwargs={"stream": stream},
data={"timeout": normal_timeout, "stream_timeout": stream_timeout},
)
if stream:
assert stream_timeout_val == stream_timeout
else:
assert stream_timeout_val == normal_timeout
with patch.object(litellm, "acompletion", spy):
response = await router.acompletion(
model="openai-gpt-5.1", messages=[{"role": "user", "content": "hi"}], mock_response="ok", **request_kwargs
)
if request_kwargs.get("stream"):
_ = [chunk async for chunk in response]
assert len(sent) == 1
return sent[0]
@pytest.mark.parametrize(
@ -165,7 +118,8 @@ def test_unit_test_streaming_timeout(stream):
({"service_tier": "flex", "stream": True, "stream_timeout": 30}, 30),
],
)
def test_the_deployment_timeout_for_the_requests_service_tier_wins_over_its_general_timeout(
@pytest.mark.asyncio
async def test_the_deployment_timeout_for_the_requests_service_tier_wins_over_its_general_timeout(
request_kwargs: dict[str, object], expected_timeout: float
) -> None:
router: Final = Router(
@ -183,12 +137,12 @@ def test_the_deployment_timeout_for_the_requests_service_tier_wins_over_its_gene
}
]
)
deployment_params: Final = router.model_list[0]["litellm_params"]
assert router._get_timeout(kwargs=request_kwargs, data=deployment_params) == expected_timeout
assert await _timeout_sent_to_the_provider_call(router, request_kwargs) == expected_timeout
def test_a_stream_keeps_the_deployment_stream_timeout_when_only_the_request_timeout_map_names_its_tier() -> None:
@pytest.mark.asyncio
async def test_a_stream_keeps_the_deployment_stream_timeout_when_only_the_request_timeout_map_names_its_tier() -> None:
router: Final = Router(
model_list=[
{
@ -203,10 +157,9 @@ def test_a_stream_keeps_the_deployment_stream_timeout_when_only_the_request_time
}
]
)
deployment_params: Final = router.model_list[0]["litellm_params"]
assert router._get_timeout(kwargs={"service_tier": "flex", "stream": True}, data=deployment_params) == 30.0
assert router._get_timeout(kwargs={"service_tier": "flex"}, data=deployment_params) == 900.0
assert await _timeout_sent_to_the_provider_call(router, {"service_tier": "flex", "stream": True}) == 30.0
assert await _timeout_sent_to_the_provider_call(router, {"service_tier": "flex"}) == 900.0
async def _drain_responses_stream(router: Router) -> None:
@ -274,7 +227,6 @@ async def test_prompt_managed_stream_keeps_the_target_deployments_own_stream_tim
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("agent_params", "target_params", "request_params"),
@ -350,7 +302,9 @@ _INVOKE_ANSWER: Final = json.dumps(
async def _streamed_chat_text(router: Router) -> str:
response: Final = await router.acompletion(model="claude", messages=[{"role": "user", "content": "hi"}], stream=True)
response: Final = await router.acompletion(
model="claude", messages=[{"role": "user", "content": "hi"}], stream=True
)
return "".join([chunk.choices[0].delta.content or "" async for chunk in response])
@ -422,7 +376,11 @@ class _FailureDetails(CustomLogger):
async def _drain_messages_stream(router: Router) -> None:
response: Final = await router.aanthropic_messages(
model="native", messages=[{"role": "user", "content": "hi"}], max_tokens=10, stream=True, client_side_timeout=True
model="native",
messages=[{"role": "user", "content": "hi"}],
max_tokens=10,
stream=True,
client_side_timeout=True,
)
_ = [chunk async for chunk in response]
@ -477,7 +435,10 @@ async def test_native_live_stream_cut_at_the_stream_timeout_is_not_blamed_on_the
started: Final = failure_details.details.get("api_call_start_time")
assert isinstance(started, datetime)
assert is_caller_timeout_408(failure_details.details, 408, ended=started + timedelta(seconds=STREAM_TIMEOUT_SECONDS))
assert is_caller_timeout_408(
failure_details.details, 408, ended=started + timedelta(seconds=STREAM_TIMEOUT_SECONDS)
)
@pytest.mark.asyncio
async def test_bedrock_timeout_preserves_the_provider_timeout_error(
@ -489,9 +450,9 @@ async def test_bedrock_timeout_preserves_the_provider_timeout_error(
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "test-secret-key")
monkeypatch.setenv("AWS_DEFAULT_REGION", "us-east-1")
model_id: Final = "us.anthropic.claude-haiku-4-5-20251001-v1:0"
route: Final = respx_mock.post(
url__regex=r"^https://bedrock-runtime\.us-east-1\.amazonaws\.com/.*"
).mock(side_effect=httpx.ReadTimeout("upstream timeout"))
route: Final = respx_mock.post(url__regex=r"^https://bedrock-runtime\.us-east-1\.amazonaws\.com/.*").mock(
side_effect=httpx.ReadTimeout("upstream timeout")
)
router: Final = Router(
model_list=[
{
@ -556,9 +517,7 @@ def test_stream_timeout_uses_the_configured_fallback(
)
chunks: Final = tuple(response)
assert tuple(chunk.choices[0].delta.content for chunk in chunks if chunk.choices[0].delta.content) == (
"fallback",
)
assert tuple(chunk.choices[0].delta.content for chunk in chunks if chunk.choices[0].delta.content) == ("fallback",)
request_bodies: Final = tuple(json.loads(call.request.content) for call in route.calls)
assert route.call_count == 2
assert tuple(body["model"] for body in request_bodies) == ("primary", "fallback")
@ -598,9 +557,9 @@ async def test_anthropic_timeout_raises_for_sync_and_streaming(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post(
"https://api.anthropic.com/v1/messages"
).mock(side_effect=httpx.ReadTimeout("scripted Anthropic timeout"))
route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(
side_effect=httpx.ReadTimeout("scripted Anthropic timeout")
)
request_kwargs: Final = {
"model": "anthropic/claude-haiku-4-5-20251001",
"messages": [{"role": "user", "content": "timeout contract"}],
@ -634,9 +593,9 @@ def test_openai_router_timeout_raises_after_httpx_read_timeout(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post(
"https://api.openai.com/v1/chat/completions"
).mock(side_effect=httpx.ReadTimeout("scripted OpenAI timeout"))
route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(
side_effect=httpx.ReadTimeout("scripted OpenAI timeout")
)
router: Final = Router(
model_list=[
{
@ -665,9 +624,9 @@ def test_azure_router_timeout_raises_after_httpx_read_timeout(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
route: Final = respx_mock.post(
"https://azure.example.com/openai/deployments/gpt-4o-mini/chat/completions"
).mock(side_effect=httpx.ReadTimeout("scripted Azure timeout"))
route: Final = respx_mock.post("https://azure.example.com/openai/deployments/gpt-4o-mini/chat/completions").mock(
side_effect=httpx.ReadTimeout("scripted Azure timeout")
)
router: Final = Router(
model_list=[
{

Some files were not shown because too many files have changed in this diff Show more