From 067f93340fc79f442bf83bf4540e0433914021b1 Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Thu, 5 Nov 2026 10:40:46 -0800 Subject: [PATCH] refactor(router): delete the python router core, the rust core is the only core (#50178) --- .../crates/python-bridge/src/router/config.rs | 4 + .../crates/python-bridge/src/router/host.rs | 122 +- .../crates/python-bridge/src/router/mod.rs | 40 +- litellm-rust/crates/router-core/src/call.rs | 192 +- .../crates/router-core/src/call/stream.rs | 22 +- litellm-rust/crates/router-core/src/error.rs | 11 +- litellm-rust/crates/router-core/src/health.rs | 24 +- litellm-rust/crates/router-core/src/host.rs | 21 + litellm-rust/crates/router-core/src/lib.rs | 36 +- .../crates/router-core/src/outcome.rs | 9 + .../crates/router-core/src/prerouting.rs | 3 + .../crates/router-core/src/registry.rs | 17 + litellm-rust/crates/router-core/src/select.rs | 80 +- litellm/litellm_core_utils/litellm_logging.py | 7 +- .../anthropic/messages/adapters/handler.py | 3 +- litellm/router.py | 307 +- litellm/router_api.py | 4 +- litellm/router_attempt.py | 19 +- litellm/router_core.py | 35 +- litellm/router_deprecated.py | 33 +- litellm/router_python_core.py | 18017 ---------------- litellm/router_rust/core.py | 752 +- litellm/router_rust/invoke/endpoints.py | 109 +- litellm/router_rust/invoke/streaming.py | 110 +- litellm/router_rust/invoke/walk.py | 44 +- litellm/router_rust/steps/callbacks.py | 36 +- litellm/router_rust/steps/precall.py | 10 +- litellm/router_rust/steps/prerouting.py | 25 +- litellm/router_rust/surface/health.py | 24 +- litellm/router_rust/surface/model_list.py | 20 +- litellm/router_rust/surface/registry.py | 120 +- litellm/router_utils/search_api_router.py | 252 - litellm/rust_bridge/catalog.py | 2 - .../router_code_coverage.py | 1 + .../router_litellm_globals_check.py | 1 - .../router_surface_check.py | 34 +- .../unit/caching/test_redis_cluster_cache.py | 22 +- tests/unit/conftest.py | 51 +- .../test_azure_container_transformation.py | 54 +- .../test_get_llm_provider_logic.py | 6 +- .../messages/test_advisor_orchestration.py | 50 +- .../base_llm/responses/test_transformation.py | 717 +- .../llms/github_copilot/test_per_user_auth.py | 31 +- tests/unit/llms/openai/live/test_handler.py | 6 +- .../proxy/auth/test_fallback_model_access.py | 3 +- .../auth/test_model_access_group_budgets.py | 13 +- .../test_unmapped_model_budget_enforcement.py | 7 +- .../hooks/test_autorouter_baseline_cache.py | 3 +- .../test_model_management_endpoints.py | 22 +- .../test_ptu_model_settings.py | 11 +- .../test_llm_passthrough_endpoints.py | 15 +- tests/unit/proxy/test_budget_reservation.py | 15 +- .../proxy/test_common_request_processing.py | 10 +- .../unit/proxy/test_litellm_pre_call_utils.py | 8 +- ...test_proxy_server_endpoints_and_startup.py | 13 +- .../test_vector_store_endpoints.py | 86 +- .../test_litellm_completion_responses.py | 89 +- .../test_responses_router_cooldown.py | 49 +- .../adaptive_router/test_router_dispatch.py | 108 - .../router_strategy/test_baseline_request.py | 4 +- .../router_strategy/test_complexity_router.py | 379 +- tests/unit/router_strategy/test_llm_v2.py | 33 +- .../router_strategy/test_quality_router.py | 34 +- .../test_router_routing_groups.py | 1218 +- .../test_router_routing_plugins.py | 50 - .../test_router_tag_routing.py | 96 +- .../test_deployment_affinity_check.py | 113 +- .../test_encrypted_content_affinity_check.py | 52 +- .../test_session_id_affinity.py | 85 +- .../test_client_initalization_utils.py | 53 +- .../router_utils/test_cooldown_handlers.py | 145 +- .../test_fallback_event_handlers.py | 483 +- ..._health_check_allowed_fails_integration.py | 406 +- .../test_kubernetes_pod_discovery.py | 157 +- .../test_router_health_check_routing.py | 266 +- .../test_router_interactions_endpoints.py | 277 +- .../test_router_utils_common_utils.py | 114 +- tests/unit/test_conftest_isolation.py | 13 +- tests/unit/test_private_usage_aliases.py | 1 - .../unit/test_redact_string_in_error_paths.py | 80 +- ...responses_streaming_container_ownership.py | 240 +- tests/unit/test_router/core/conftest.py | 31 +- tests/unit/test_router/core/scripted_host.py | 35 +- ...t_construction_precall_characterization.py | 74 +- ...uction_precall_characterization_precall.py | 6 +- tests/unit/test_router/core/test_endpoints.py | 20 +- .../test_router/core/test_legacy_members.py | 31 +- .../test_router/test_io_token_rate_limits.py | 36 - tests/unit/test_router/test_router.py | 14340 +----------- .../unit/test_router/test_router_endpoints.py | 569 +- .../test_router/test_router_entry_points.py | 250 +- tests/unit/test_router/test_router_facade.py | 13 +- .../unit/test_router/test_router_fallbacks.py | 4 +- .../test_router_fuzz_regressions_rust_only.py | 515 + .../test_router_get_deployments.py | 60 - .../test_router/test_router_helper_utils.py | 2090 +- .../test_router_max_parallel_requests.py | 48 - .../test_router_provider_endpoints.py | 101 +- tests/unit/test_router/test_router_retries.py | 284 +- tests/unit/test_router/test_router_timeout.py | 141 +- tests/unit/test_router_block_helpers.py | 10 - .../unit/test_router_model_cost_isolation.py | 628 +- tests/unit/test_router_order_fallback.py | 46 - .../test_router_per_deployment_num_retries.py | 196 +- .../unit/test_router_retry_backoff_headers.py | 90 +- .../test_router_retry_non_retryable_errors.py | 316 +- tests/unit/test_router_retry_policy_update.py | 22 +- tests/unit/test_router_silent_experiment.py | 263 +- ...test_router_streaming_fallback_metadata.py | 142 - tests/unit/test_router_weighted_failover.py | 483 +- .../unit/test_streaming_connection_cleanup.py | 319 +- 111 files changed, 4758 insertions(+), 42639 deletions(-) delete mode 100644 litellm/router_python_core.py delete mode 100644 litellm/router_utils/search_api_router.py create mode 100644 tests/unit/test_router/test_router_fuzz_regressions_rust_only.py diff --git a/litellm-rust/crates/python-bridge/src/router/config.rs b/litellm-rust/crates/python-bridge/src/router/config.rs index 0b46f2208c5..de383a0ffa1 100644 --- a/litellm-rust/crates/python-bridge/src/router/config.rs +++ b/litellm-rust/crates/python-bridge/src/router/config.rs @@ -188,6 +188,8 @@ struct RequestSpec { web_search: Option, route_handler: Option, decisions_media: Option, + batch_create: Option, + embedding_batch: Option, wildcard_request: Option, unrouted: Option, passthrough_on_no_deployment: Option, @@ -255,6 +257,8 @@ pub(crate) fn parse_request(value: Value, model: Option<&str>) -> Result, + }, /// 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, /// The `order` level of a same-group order fallback. target_order: Option, /// 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), + /// A strategy's pick among candidates sharing an id: positions in the candidates, pick first. + Positions(Vec), /// The model group to resolve instead, a custom strategy's deployment, or `None`. PreRoute(Option), /// 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]) -> 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], ids: &[String]) -> Vec> { + 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, 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 { diff --git a/litellm-rust/crates/python-bridge/src/router/mod.rs b/litellm-rust/crates/python-bridge/src/router/mod.rs index c69259fcf1c..8157d0f5a32 100644 --- a/litellm-rust/crates/python-bridge/src/router/mod.rs +++ b/litellm-rust/crates/python-bridge/src/router/mod.rs @@ -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>) -> PyResult() { 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::() + && 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) -> PyResult { + 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, diff --git a/litellm-rust/crates/router-core/src/call.rs b/litellm-rust/crates/router-core/src/call.rs index 287bb98d8b7..3a30c779958 100644 --- a/litellm-rust/crates/router-core/src/call.rs +++ b/litellm-rust/crates/router-core/src/call.rs @@ -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, deployment_id: Option, + /// 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, + /// The group the host stamped that attempt as served for ([`Candidates::served_as`]). + served_as: Option, } impl From for Failed { @@ -70,6 +76,8 @@ impl From 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, } /// 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 = 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>, 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 = group_members(snapshot, &target.group, listing_team) + let mut orders: Vec = 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 = 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 { + (candidates.name != target.group).then(|| candidates.name.clone()) +} diff --git a/litellm-rust/crates/router-core/src/call/stream.rs b/litellm-rust/crates/router-core/src/call/stream.rs index a0fce4372dd..f1dc71a9b18 100644 --- a/litellm-rust/crates/router-core/src/call/stream.rs +++ b/litellm-rust/crates/router-core/src/call/stream.rs @@ -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, + /// The group a pre-routing hook selected for `group`, which the fallback lookup keys on first. + #[serde(default)] + pub routed: Option, } /// 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 } diff --git a/litellm-rust/crates/router-core/src/error.rs b/litellm-rust/crates/router-core/src/error.rs index b4112501c36..926772f435c 100644 --- a/litellm-rust/crates/router-core/src/error.rs +++ b/litellm-rust/crates/router-core/src/error.rs @@ -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>, + }, #[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(), diff --git a/litellm-rust/crates/router-core/src/health.rs b/litellm-rust/crates/router-core/src/health.rs index 020a1a0f47d..c8055eafe31 100644 --- a/litellm-rust/crates/router-core/src/health.rs +++ b/litellm-rust/crates/router-core/src/health.rs @@ -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, 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 diff --git a/litellm-rust/crates/router-core/src/host.rs b/litellm-rust/crates/router-core/src/host.rs index 0a72be8f980..cfc97480bf1 100644 --- a/litellm-rust/crates/router-core/src/host.rs +++ b/litellm-rust/crates/router-core/src/host.rs @@ -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, @@ -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, /// The `order` level an order fallback targets in the same group. pub target_order: Option, /// 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, diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs index 5f15b6b8df0..80c8f597157 100644 --- a/litellm-rust/crates/router-core/src/lib.rs +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -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, diff --git a/litellm-rust/crates/router-core/src/outcome.rs b/litellm-rust/crates/router-core/src/outcome.rs index 59e0d14b2c2..7403b819d6d 100644 --- a/litellm-rust/crates/router-core/src/outcome.rs +++ b/litellm-rust/crates/router-core/src/outcome.rs @@ -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] diff --git a/litellm-rust/crates/router-core/src/prerouting.rs b/litellm-rust/crates/router-core/src/prerouting.rs index d6e3df05e7a..8ea1377a5aa 100644 --- a/litellm-rust/crates/router-core/src/prerouting.rs +++ b/litellm-rust/crates/router-core/src/prerouting.rs @@ -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), } diff --git a/litellm-rust/crates/router-core/src/registry.rs b/litellm-rust/crates/router-core/src/registry.rs index 78c7d3b67e3..594aff5f6dd 100644 --- a/litellm-rust/crates/router-core/src/registry.rs +++ b/litellm-rust/crates/router-core/src/registry.rs @@ -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 diff --git a/litellm-rust/crates/router-core/src/select.rs b/litellm-rust/crates/router-core/src/select.rs index fefa7f25ad1..1972b2cbd90 100644 --- a/litellm-rust/crates/router-core/src/select.rs +++ b/litellm-rust/crates/router-core/src/select.rs @@ -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, + /// 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, } 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, 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, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index ce57a9bd4a5..9c49efc529f 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -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, diff --git a/litellm/llms/anthropic/messages/adapters/handler.py b/litellm/llms/anthropic/messages/adapters/handler.py index 31a4e9d9d73..da120cad0e1 100644 --- a/litellm/llms/anthropic/messages/adapters/handler.py +++ b/litellm/llms/anthropic/messages/adapters/handler.py @@ -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( diff --git a/litellm/router.py b/litellm/router.py index 6a60d5d958e..a5d72daea1e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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/", + "api_key": , + "api_version": , + "api_base": + }, + }, + { + "model_name": "azure-gpt-3.5-turbo", # model alias + "litellm_params": { # params for litellm completion/embedding call + "model": "azure/", + "api_key": , + "api_version": , + "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": , + }, + ] + + 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__ = [ diff --git a/litellm/router_api.py b/litellm/router_api.py index 82660439ae9..dc95b742d56 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -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() diff --git a/litellm/router_attempt.py b/litellm/router_attempt.py index 33e0a19c530..b7b8d2ac87d 100644 --- a/litellm/router_attempt.py +++ b/litellm/router_attempt.py @@ -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.`` 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 diff --git a/litellm/router_core.py b/litellm/router_core.py index c077b3f0015..a906e4d86fb 100644 --- a/litellm/router_core.py +++ b/litellm/router_core.py @@ -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 diff --git a/litellm/router_deprecated.py b/litellm/router_deprecated.py index 9d915fe0982..c2e1c197259 100644 --- a/litellm/router_deprecated.py +++ b/litellm/router_deprecated.py @@ -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) diff --git a/litellm/router_python_core.py b/litellm/router_python_core.py deleted file mode 100644 index 8ca2afb9a99..00000000000 --- a/litellm/router_python_core.py +++ /dev/null @@ -1,18017 +0,0 @@ -# +-----------------------------------------------+ -# | | -# | Give Feedback / Get Help | -# | https://github.com/BerriAI/litellm/issues/new | -# | | -# +-----------------------------------------------+ -# -# Thank you ! We ❤️ you! - Krrish & Ishaan - -""" -``PythonCore``: the Python implementation of the router core behind ``litellm.Router``. - -It holds all routing logic (registry, selection, retries, fallbacks, cooldowns, usage) and reaches -outside itself only through its ``PythonHost``: inference calls, sleeps, callback registration and -the process settings (``litellm/router_config.py``). The Rust core in litellm-rust/crates/router-core -implements the same interface. ``litellm/router.py`` picks one at construction and delegates to it. -""" - -import asyncio -import contextlib -import copy -import enum -import inspect -import json -import logging -import re -import threading -import time -import traceback -import weakref -from collections import defaultdict -from collections.abc import ( - AsyncGenerator, - AsyncIterable, - AsyncIterator, - Awaitable, - Callable, - Generator, - Iterable, - Iterator, - Mapping, - MutableMapping, - Sequence, -) -from contextvars import ContextVar -from dataclasses import dataclass -from datetime import datetime, timezone -from functools import lru_cache, partial -from itertools import chain -from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Generic, Literal, NoReturn, Optional, TypeAlias, TypeVar, Union, cast - -import anyio -import httpx -import openai -from openai import AsyncOpenAI -from pydantic import BaseModel, TypeAdapter, ValidationError -from typing_extensions import assert_never, overload - -from litellm import exceptions as litellm_exceptions -from litellm import get_secret_str -from litellm import utils as litellm_utils -from litellm._internal_context import in_mirrored_call, mirrored_call, service_target, with_service_target -from litellm._logging import verbose_router_logger -from litellm._uuid import uuid -from litellm.caching.caching import Cache as LitellmCache -from litellm.caching.caching import ( - DualCache, - InMemoryCache, - RedisCache, - RedisClusterCache, -) -from litellm.caching.redis_cache import log_redis_failure -from litellm.constants import ( - BAGGAGE_MAPPED_TAGS_METADATA_KEY, - CLIENT_OUTPUT_CEILING_METADATA_KEY, - CONSUMED_REQUEST_TAGS_METADATA_KEY, - DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, - DEFAULT_HEALTH_CHECK_INTERVAL, - DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER, - DEFAULT_MAX_LRU_CACHE_SIZE, - INTERNAL_CALL_ORIGIN_METADATA_KEY, - KUBERNETES_POD_ROUTING_KEY, - MID_STREAM_FAILURE_LOG_WAIT_SECONDS, - OUTPUT_TOKEN_CEILING_PARAMS, - PROMPT_CACHE_PIN_CACHE_MAX_SIZE, - REALTIME_SESSION_ENDED_ABNORMALLY_KEY, - REALTIME_SESSION_SUCCESS_LOGGED_KEY, - ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, - ROUTING_REQUEST_TAGS_METADATA_KEY, - RUNTIME_UPDATABLE_ROUTER_SETTINGS, - SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, -) -from litellm.integrations.custom_guardrail import CustomGuardrail, is_guardrail_intervention -from litellm.integrations.custom_logger import CustomLogger -from litellm.integrations.otel.routing import routing_decision_attributes -from litellm.integrations.otel.runtime import phase_attributes, phase_event, phase_span -from litellm.litellm_core_utils import exception_mapping_utils -from litellm.litellm_core_utils.access_windows import access_windows_config_error, filter_reserved_deployments -from litellm.litellm_core_utils.add_retry_fallback_headers import ( - add_fallback_headers_to_response, - add_retry_headers_to_response, - apply_quality_router_decision_headers, - apply_remaining_usage_headers, - apply_response_model_id, - complexity_router_decision_headers, - ensure_response_additional_headers, - get_hidden_params_dict, - prepare_response_for_header_attachment, - replace_complexity_router_headers, - response_total_token_count, -) -from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure -from litellm.litellm_core_utils.asyncify import run_async_function -from litellm.litellm_core_utils.batch_utils import ( - batch_create_routing, - deployment_supported_endpoints, - get_router_metadata_variable_name, - is_batch_retrieve_call_type, - replace_model_in_jsonl, - should_replace_model_in_jsonl, - without_deployments_that_cannot_batch, -) -from litellm.litellm_core_utils.clientside_credential_handler import ( - get_dynamic_litellm_params, - is_clientside_credential, -) -from litellm.litellm_core_utils.completion_timeout import CompletionTimeout -from litellm.litellm_core_utils.core_helpers import ( - as_str_mapping, - coerce_token_limit, - get_litellm_metadata_from_kwargs, - get_metadata_variable_name_from_kwargs, - get_or_create_metadata_bucket, - get_parent_otel_span_from_kwargs, -) -from litellm.litellm_core_utils.coroutine_checker import coroutine_checker -from litellm.litellm_core_utils.credential_accessor import CredentialAccessor, reject_api_key_secret_reference -from litellm.litellm_core_utils.dd_tracing import tracer -from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models, deployment_pricing_models -from litellm.litellm_core_utils.get_litellm_params import SUPPLEMENTAL_KWARGS_KEYS -from litellm.litellm_core_utils.get_llm_provider_logic import ( - declared_authenticating_provider, - inferred_provider, - is_registered_custom_provider, - resolve_model_and_provider_for_metadata, -) -from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap -from litellm.litellm_core_utils.get_retry_from_policy import ( - get_num_retries_from_retry_policy as _get_num_retries_from_retry_policy, -) -from litellm.litellm_core_utils.get_retry_from_policy import resolve_retry_policy -from litellm.litellm_core_utils.hidden_params import ( - HIDDEN_PARAMS_ATTR as _HIDDEN_PARAMS_ATTR, -) -from litellm.litellm_core_utils.hidden_params import ( - get_hidden_params, - set_hidden_params, -) -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES -from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter -from litellm.litellm_core_utils.ptu_pricing import ( - PTU_COST_ATTRIBUTION_ENV_VAR, - declares_ptu, - declares_ptu_shares, - is_ptu_cost_attribution_enabled, - ptu_config_error, - ptu_identity_error, - ptu_terms, - zeroed_ptu_pricing, -) -from litellm.litellm_core_utils.ptu_shares import ( - PTUTeamCeiling, - filter_ptu_shared_deployments, - ptu_capacity_warning, - ptu_share_teams, - share_holder_request_kwargs, - team_ptu_ceiling, -) -from litellm.litellm_core_utils.reasoning_effort_capability import ( - deployment_is_catalog_mapped, - intersect_supported_reasoning_efforts, - resolve_supported_reasoning_efforts, -) -from litellm.litellm_core_utils.request_timeout_resolver import ( - get_configured_request_timeout, - get_service_tier_timeout, -) -from litellm.litellm_core_utils.routing_common_utils import ( - MID_STREAM_CONTINUATION_KWARG, - as_retry_skipped_deployment_ids, - deployment_model_id_pin, - deployments_the_caller_may_use, - filter_team_based_models, - filter_web_search_deployments, - format_content_policy_fallback_hint, - format_content_policy_no_fallback_note, - format_fallback_outcome_message, - format_no_fallback_group_message, - get_request_team_id, - is_proxy_admin_request, - provider_for_generic_call, - request_user_api_key_auth, - resolve_model_group_alias, - team_may_use_deployment, - truncate_fallback_error_detail, - warn_on_ignored_model_capability_flags, - warn_on_provider_credential_mismatch, - with_admin_api_key_secret_reference, - without_model_capability_flags, - without_router_only_kwargs, -) -from litellm.litellm_core_utils.routing_groups import ( - apply_routing_group_priority, - parse_routing_groups, - validate_routing_strategy, -) -from litellm.litellm_core_utils.secret_redaction import redact_string -from litellm.litellm_core_utils.sensitive_data_masker import ( - SensitiveDataMasker, - mask_sensitive_structure, -) -from litellm.litellm_core_utils.token_counter import ( - offload_token_count, - served_openai_models, - single_encoding_openai_model, -) -from litellm.litellm_core_utils.tool_fields import TOOL_FIELDS, drop_policy -from litellm.llms.anthropic.common_utils import AnthropicModelInfo, anthropic_error_frame_exception -from litellm.llms.base_llm.passthrough.transformation import replace_path_segment -from litellm.llms.base_llm.vector_store.transformation import ( - RouterVectorStoreEmbeddingExecutor, - vector_store_request_metadata, -) -from litellm.llms.bedrock.common_utils import get_bedrock_base_model -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client -from litellm.llms.openai_like.json_loader import JSONProviderRegistry -from litellm.llms.openai_like.model_info import ( - MODEL_INFO_DISCOVERY_PROVIDERS, - MODEL_INFO_REFRESH_CONCURRENCY, - MODEL_INFO_REFRESH_SECONDS, - get_openai_compatible_model_info, -) -from litellm.main import Chat as LitellmChat -from litellm.responses.bridge_check import keep_registered_mode -from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults -from litellm.router_api import ( - PLAIN_SETTING_ATTRIBUTES, - SETTING_FIELDS, - ApplyResult, - ChangeSet, - RouterHealth, - RouterRegistry, - RouterSettings, - RoutingObservers, - SelectKind, - aapply_changes, - apply_changes, -) -from litellm.router_api import RoutingContext as RouterContext -from litellm.router_api import deployment_id as _deployment_id -from litellm.router_deprecated import DeprecatedRouterMembers -from litellm.router_host import router_host -from litellm.router_python_host import PythonHost, python_host -from litellm.router_strategy.auto_router.auto_router_model_naming import ( - AUTO_ROUTER_MODEL_PREFIX, - STRATEGY_ROUTER_PARAM_FIELDS, - GatedAutoRouterCapability, - capability_limit_violation, - claimed_capabilities, - classify_strategy_router_model, - count_capability_routers, -) -from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy -from litellm.router_strategy.budget_limiter import RouterBudgetLimiting -from litellm.router_strategy.complexity_router.context_compaction import ( - arm_compaction, - compact_to_fit, - compaction_pending, - initialize_compaction_state, - is_native_compaction_call, - reject_recursive_compactor, - surface_for_call, -) -from litellm.router_strategy.least_busy import ( - CALL_IDENTITY, - SELECTION_RESERVATION_MAX_KEYS, - LeastBusyLoggingHandler, - call_id_of, - owning_picks, - pick_reserves, - picks_reserve, -) -from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler -from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler -from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler -from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage -from litellm.router_strategy.simple_shuffle import simple_shuffle -from litellm.router_strategy.tag_based_routing import ( - baggage_stamp_for_serving_deployment, - get_deployments_for_tag, - get_routing_tags_from_request_kwargs, - get_tags_from_request_kwargs, - is_valid_deployment_tag, - tags_with_serving_deployment_tags, -) -from litellm.router_utils.client_initalization_utils import InitalizeCachedClient, MaxParallelRequestsLimit -from litellm.router_utils.cooldown_cache import CooldownCache -from litellm.router_utils.cooldown_handlers import ( - DEFAULT_COOLDOWN_TIME_SECONDS, - _first_present, # pyright: ignore[reportPrivateUsage] - shared internal helper across router_utils submodules, matching the other cooldown_handlers imports on this line - async_get_cooldown_deployments, - async_get_cooldown_deployments_with_debug_info, - get_cooldown_deployment_info, - get_cooldown_deployments, - is_background_response_poll_not_found, - is_caller_entra_credential_failure, - is_caller_scoped_auth_failure, - is_caller_timeout_408, - is_failure_outside_selected_deployment, - is_video_lookup_not_found, - mark_deployment_selection_failure, - set_cooldown_deployments, -) -from litellm.router_utils.deployment_numbers import refuse_unservable_numbers -from litellm.router_utils.embedding_batch_size_check import ( - async_log_rejection_before_provider_call, - clear_rejected_before_provider_call, - deployments_admitting_embedding_batch, - deployments_when_none_admits_embedding_batch, - embedding_batch_too_large_error, - embedding_deployment_selection, - log_rejection_before_provider_call, -) -from litellm.router_utils.endpoint_eligibility import ( - eligible_deployments, - endpoint_selection, - warn_on_unenforced_endpoint_flag, -) -from litellm.router_utils.fallback_event_handlers import ( - MID_STREAM_FALLBACK_CONTROLS_KEY, - AttemptedFallbackTargets, - attempted_fallback_targets_from, - attempted_retries_for_request, - carry_over_pre_routing_selection, - carry_over_routed_deployment, - carry_superseded_attempt_usage, - check_non_standard_fallback_format, - clear_pre_routing_selection, - committed_retry_budget_for_request, - continuation_fallback_targets, - fallback_lookup_groups, - fallback_target_name, - fallback_target_overrides_request_params, - fallbacks_disabled_for_request, - get_fallback_model_group_for_lookup_groups, - get_pre_routing_selection, - has_unattempted_fallback_target, - is_bad_request_without_fallback, - mid_stream_fallback_hop_kwargs, - mid_stream_fallback_snapshot_kwargs, - mid_stream_retry_kwargs, - per_request_fallback_controls, - record_disable_fallbacks, - record_pre_routing_selection, - record_retry_attempt, - routed_deployment_id, - run_async_fallback, -) -from litellm.router_utils.global_callback_holders import ( - live_check_owners, - routers_holding_global_callbacks, - routers_holding_global_callbacks_lock, -) -from litellm.router_utils.handle_error import ( - async_raise_no_deployment_exception, - send_llm_exception_alert, -) -from litellm.router_utils.health_state_cache import DeploymentHealthCache -from litellm.router_utils.kubernetes_pod_discovery import ( - KubernetesPodDiscovery, - async_resolve_pods_after, - async_resolve_pods_after_bound, - resolve_pods_after, - resolve_pods_after_bound, -) -from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( - MID_STREAM_CONTINUATION_MARKER, - RESPONSES_CONTINUATION_KWARG, - ContinuationOutput, - ContinuationPrefillDeploymentCheck, - ContinuedPartial, - ResponsesContinuation, - ResponsesContinuationSeam, - constrains_continuation, - constrains_responses_continuation, - continuation_deployments, - continuation_output_ceilings, - continued_text, - deployment_can_continue, - responses_continuation_input, - responses_continuation_input_for, -) -from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( - DeploymentAffinityCheck, - warn_on_unknown_model_group_affinity_flags, -) -from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, -) -from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( - async_admit_io_token_turn, - async_io_token_refund_failure, - build_io_token_rate_limit_headers, - deployment_has_io_token_limits, - refund_stale_reservation_before_retry, - set_io_token_rate_limit_request_kwargs, - without_io_token_limits, -) -from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( - ModelRateLimitingCheck, -) -from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import ( - PromptCachingDeploymentCheck, -) -from litellm.router_utils.request_priority import ( - InvalidPriority, - parse_default_priority, - priority_error, - priority_header_fields, - request_drops_params, - resolve_request_priority, - unqueueable_priority_message, -) -from litellm.router_utils.router_callbacks.track_deployment_metrics import ( - find_deployment_metadata, - get_counted_usage_tokens, - increment_deployment_failures_for_current_minute, - increment_deployment_successes_for_current_minute, -) -from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingPrefetch, RoutingReadBatch -from litellm.router_utils.stream_usage_injection import router_config_requests_stream_usage -from litellm.scheduler import FlowItem, Scheduler -from litellm.secret_managers.main import str_to_bool -from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD, RoutingStrategyName -from litellm.types.llms.openai import ( - AllMessageValues, - ChatCompletionToolParam, - FileTypes, - OpenAIFileObject, - OpenAIFilesPurpose, -) -from litellm.types.router import ( - CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, - VALID_LITELLM_ENVIRONMENTS, - AlertingConfig, - AllowedFailsPolicy, - AssistantsTypedDict, - AutoRouterCapabilityLimit, - ConsumedRequestTagsStamp, - CredentialLiteLLMParams, - CustomRoutingStrategyBase, - Deployment, - DeploymentModelListingInfo, - DeploymentTypedDict, - DiscoveredDeploymentModelInfo, - FallbackAccessCheck, - FallbackBudgetCheck, - GenericLiteLLMParams, - GuardrailTypedDict, - LiteLLM_Params, - LiteLLMParamsTypedDict, - MockRouterTestingParams, - ModelGroupInfo, - OptionalPreCallChecks, - PreRoutingStrategy, - RetryPolicy, - RouterCacheEnum, - RouterErrors, - RouterGeneralSettings, - RouterModelGroupAliasItem, - RouterRateLimitError, - RouterRateLimitErrorBasic, - RoutingContext, - RoutingGroup, - RoutingPlugin, - RoutingStrategy, - SearchToolTypedDict, - TaggedPreRoutingStrategy, - holds_secret_pointer, - validate_litellm_params_strictly, -) -from litellm.types.services import ServiceTypes -from litellm.types.utils import ( - AUTOROUTER_CLASSIFIER_CALL_ORIGIN, - COST_MAP_LOOKUP_KEY, - DEPLOYMENT_SCOPED_TOKEN_LIMIT_FIELDS, - PROMPT_QUOTING_ROUTING_DECISION_FIELDS, - CustomPricingLiteLLMParams, - GenericBudgetConfigType, - LiteLLMBatch, - LlmProviders, - LlmProvidersSet, - ModelInfo, - ModelResponseStream, - StandardLoggingPayload, - StandardLoggingRoutingDecision, - Usage, - all_litellm_params, - echoed_cost_map_fields, - is_server_derived_pricing_key, - shared_backend_model_info, -) -from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.videos.utils import VIDEO_ID_CALL_TYPES, decode_video_id_with_provider -from litellm.utils import ( - CustomStreamWrapper, - EmbeddingResponse, - ModelInfoLookupMiss, - ModelResponse, - ProviderConfigManager, - Rules, - TextCompletionResponse, - _restore_correlation_context_if_supported, # pyright: ignore[reportPrivateUsage] # same task-local cleanup as the SDK wrapper - aresolve_api_key_secret_reference, - catalog_with_runtime_registrations, - function_setup, - get_llm_provider, - get_non_default_completion_params, - get_potential_model_names, - get_secret, - get_utc_datetime, - is_generalized_model_info, - is_region_allowed, - lookup_model_info_for_registration, - model_info_with_pricing_overrides, - provider_rejectable_params, - set_live_deployment_replay, - token_counter, -) - -if TYPE_CHECKING: - from opentelemetry.trace import Span as _Span - - from litellm.exceptions import MidStreamFallbackError - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - ) - from litellm.router_api import UntypedDict - from litellm.router_strategy.adaptive_router.adaptive_router import ( - AdaptiveRouter, - ) - from litellm.router_strategy.auto_router.auto_router import ( - AutoRouter, - PreRoutingHookResponse, - ) - from litellm.router_strategy.complexity_router.complexity_router import ( - ComplexityRouter, - ) - from litellm.router_strategy.quality_router.quality_router import ( - QualityRouter, - ) - from litellm.types.llms.anthropic_messages.anthropic_response import ( - AnthropicMessagesResponse, - ) - from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject - from litellm.types.llms.openai import ( - ResponseAPIUsage, - ResponseIncompleteEvent, - ResponsesAPIResponse, - ResponsesAPIStreamingResponse, - ) - - Span = _Span -else: - Span = Any - AutoRouter = Any - ComplexityRouter = Any - AdaptiveRouter = Any - QualityRouter = Any - PreRoutingHookResponse = Any - -RouterStrategySelector: TypeAlias = ( - LeastBusyLoggingHandler - | LowestCostLoggingHandler - | LowestLatencyLoggingHandler - | LowestTPMLoggingHandler - | LowestTPMLoggingHandler_v2 -) - - -def _cost_value_as_float(value: str | float | None) -> float | None: - if value is None: - return None - try: - return float(value) - except (TypeError, ValueError): - return None - - -def model_info_is_active_for_environment(model_info: Mapping[str, object] | None) -> bool: - """Single owner of the environment-gating rule: a deployment whose model_info names - `supported_environments` loads only on pods whose LITELLM_ENVIRONMENT is in that list. - `Router.deployment_is_active_for_environment` delegates here, and the model-write - endpoints consult the same rule to tell a deliberately inactive model from one that - was dropped by a failed reload.""" - if model_info is None: - return True - supported_environments: Final = model_info.get("supported_environments") - if supported_environments is None: - return True - if not isinstance(supported_environments, (list, tuple)): - raise ValueError( - f"supported_environments must be a list of {VALID_LITELLM_ENVIRONMENTS}. " - f"but set as: {supported_environments} for model_info: {model_info}" - ) - litellm_environment: Final = get_secret_str(secret_name="LITELLM_ENVIRONMENT") - if litellm_environment is None: - raise ValueError("Set 'supported_environments' for model but not 'LITELLM_ENVIRONMENT' set in .env") - - if litellm_environment not in VALID_LITELLM_ENVIRONMENTS: - raise ValueError( - f"LITELLM_ENVIRONMENT must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {litellm_environment}" - ) - - for _env in supported_environments: - if _env not in VALID_LITELLM_ENVIRONMENTS: - raise ValueError( - f"supported_environments must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {_env} " - f"for model_info: {model_info}" - ) - - if litellm_environment in supported_environments: - return True - return False - - -_PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") -_CallbackT = TypeVar("_CallbackT") -_DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) - - -_ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset( - {"model", "api_base", "api_key", "api_version", "responses_api_version"} -) -_ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" -_FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS: Final = frozenset({"chatgpt", "anthropic"}) -_CLAUDE_CODE_SESSION_ID_RE: Final = re.compile(r"^[a-zA-Z0-9_\-]{8,}$") -_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 -CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding" - -_PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"}) -_STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS: Final = frozenset({"anthropic_messages", "aresponses", "responses"}) -_video_call_deployment_id: Final[ContextVar[str | None]] = ContextVar("video_call_deployment_id", default=None) -_routed_call_type: Final[ContextVar[str | None]] = ContextVar("routed_call_type", default=None) -_DECISIONS_MEDIA_CALL_TYPES: Final = frozenset({"adecisions", "decisions", "asystemone", "systemone"}) -_ROUTE_CAPABILITY_RESOLVERS: Final[Mapping[str, Callable[[str, LlmProviders], object | None]]] = MappingProxyType( - { - "adecisions": ProviderConfigManager.get_provider_decisions_config, - "decisions": ProviderConfigManager.get_provider_decisions_config, - } -) -_RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType( - { - "prompt_caching": PromptCachingDeploymentCheck, - "enforce_model_rate_limits": ModelRateLimitingCheck, - "encrypted_content_affinity": EncryptedContentAffinityCheck, - } -) - - -def _stream_chunks_have_generated_content(chunks: Sequence[ModelResponseStream]) -> bool: - for chunk in chunks: - if not chunk.choices: - continue - delta = chunk.choices[0].delta - if ( - delta.get("content") - or delta.get("tool_calls") - or delta.get("function_call") - or delta.get("reasoning_content") - or delta.get("thinking_blocks") - or delta.get("reasoning_items") - or delta.get("audio") - or delta.get("images") - or delta.get("annotations") - ): - return True - return False - - -_NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType({}) -_NO_DEFAULT_PARAMS: Final[Mapping[str, object]] = MappingProxyType({}) -_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_CALL_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str]) -_RESOLVED_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy | None) -_ROUTING_KWARGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) -_FALLBACK_HOP_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_REQUEST_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object]) -_DEPLOYMENT_PARAMS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | GenericLiteLLMParams]] = TypeAdapter( - Mapping[str, object] | GenericLiteLLMParams -) -_OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object) -_DEPLOYMENT_SELECTED_EVENT: Final = "litellm.request.deployment_selected" -_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]]) -_DEPLOYMENTS_ADAPTER: Final = TypeAdapter(Sequence[Mapping[str, object]]) - - -def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: str) -> Sequence[Mapping[str, object]]: - """Continue the last text block in place so the prefill reads as the text the caller already received.""" - last: Final = blocks[-1] if blocks else None - if last is not None and last.get("type") == "text" and isinstance(last_text := last.get("text"), str): - return [*blocks[:-1], {**last, "text": f"{last_text}{text}"}] - return [*blocks, {"type": "text", "text": text}] - - -def _is_fallback_hop(request_kwargs: Mapping[str, object]) -> bool: - fallback_depth: Final = request_kwargs.get("fallback_depth") - return isinstance(fallback_depth, int) and fallback_depth > 0 - - -def _deployment_that_just_failed(request_metadata: object) -> str | None: - try: - model_info: Final = _FALLBACK_HOP_ADAPTER.validate_python( - _FALLBACK_HOP_ADAPTER.validate_python(request_metadata).get("model_info") - ) - except ValidationError: - return None - model_id: Final = model_info.get("id") - return model_id if isinstance(model_id, str) else None - - -def _drops_malformed_request_tools( - deployment_params: Mapping[str, object], defaults: Mapping[str, object], request: Mapping[str, object] -) -> bool: - if isinstance(request.get("tools") or [], (list, tuple)) or not deployment_params.get("tools"): - return False - set_defaults: Final = {name: value for name, value in defaults.items() if value is not None} - return drop_policy({**deployment_params, **set_defaults, **request}).drops("tools") - - -def _malformed_request_params_the_deployment_sets( - deployment_params: Mapping[str, object], defaults: Mapping[str, object], request: Mapping[str, object] -) -> frozenset[str]: - overrides: Final = { - key: value - for key, value in request.items() - if key in SUPPLEMENTAL_KWARGS_KEYS and deployment_params.get(key) is not None - } - if not overrides: - return frozenset() - try: - validate_litellm_params_strictly(GenericLiteLLMParams, overrides) - except ValidationError as error: - set_defaults: Final = {name: value for name, value in defaults.items() if value is not None} - policy: Final = drop_policy({**deployment_params, **set_defaults, **request}) - malformed: Final = (str(detail["loc"][0]) for detail in error.errors() if detail["loc"]) - return frozenset(key for key in malformed if policy.drops(key)) - return frozenset() - - -def _drop_malformed_request_params_the_deployment_sets( - kwargs: dict[str, object], # mutable-ok: drops the keys from the router's request kwargs in place - deployment_params: Mapping[str, object], - defaults: Mapping[str, object], -) -> None: - for key in _malformed_request_params_the_deployment_sets(deployment_params, defaults, kwargs): - kwargs.pop(key) - - -def _merged_tools(deployment_tools: Sequence[object], request_tools: object) -> list[object] | None: - if not isinstance(request_tools, (list, tuple)): - return None - return [*deployment_tools, *request_tools] if deployment_tools or request_tools else None - - -def _deployment_pick_attributes(model: str, request_kwargs: Mapping[str, object] | None) -> Mapping[str, str | int]: - """Bounded attributes for one deployment pick; attempt is 1-based within the current model group.""" - kwargs: Final = request_kwargs or {} - metadata: Final = kwargs.get("litellm_metadata", kwargs.get("metadata")) - attempted_retries: Final = metadata.get("attempted_retries") if isinstance(metadata, Mapping) else None - retries: Final = attempted_retries if isinstance(attempted_retries, int) else 0 - reason: Final = "retry" if retries > 0 else "fallback" if _is_fallback_hop(kwargs) else "initial" - return MappingProxyType( - { - "litellm.deployment.attempt": retries + 1, - "litellm.deployment.reason": reason, - "litellm.deployment.model_group": model, - } - ) - - -def _configured_model_info(deployment: Mapping[str, object]) -> Mapping[str, object]: - return _MODEL_INFO_ADAPTER.validate_python(deployment.get("model_info") or {}) - - -def _operator_max_input_tokens(deployment: Mapping[str, object]) -> int | None: - model_info: Final = _configured_model_info(deployment) - catalog_key: Final = model_info.get(COST_MAP_LOOKUP_KEY) - catalog_entry: Final = ( - GetModelCostMap.loaded_model_cost_map().get(catalog_key) if isinstance(catalog_key, str) else None - ) - if catalog_entry is not None and "max_input_tokens" in echoed_cost_map_fields(model_info, catalog_entry): - return None - return coerce_token_limit(model_info.get("max_input_tokens")) - - -def _search_tool_support(deployment: DeploymentTypedDict) -> bool | None: - from litellm.responses.litellm_completion_transformation.tool_search import chat_config_builds_anthropic_request - - configured: Final = _configured_model_info(deployment).get("supports_search_tool") - if isinstance(configured, bool): - return configured - params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) - model: Final = params.get("model") - provider: Final = params.get("custom_llm_provider") - if not isinstance(model, str): - return None - try: - resolved_model, resolved_provider = resolve_model_and_provider_for_metadata( - model=model, custom_llm_provider=provider if isinstance(provider, str) else None - ) - chat_config: Final = ProviderConfigManager.get_provider_chat_config( - model=resolved_model, provider=LlmProviders(resolved_provider) - ) - except (litellm_exceptions.BadRequestError, ValueError): - return None - return True if chat_config_builds_anthropic_request(chat_config) else None - - -def _every_deployment_refused( - deployments: Sequence[Mapping[str, object]], refused_deployment_ids: tuple[str, ...] -) -> bool: - return all(_configured_model_info(deployment).get("id") in refused_deployment_ids for deployment in deployments) - - -def _silent_experiment_targets(silent_model: object) -> tuple[str, ...]: - if silent_model is None or in_mirrored_call(): - return () - try: - targets: Final = _SILENT_MODEL_ADAPTER.validate_python(silent_model) - except ValidationError: - verbose_router_logger.warning( - "silent_model must be a model name or a list of model names, got %r; skipping shadow traffic", - silent_model, - ) - return () - return (targets,) if isinstance(targets, str) else tuple(targets) - - -def _silent_experiment_kwargs_snapshot(kwargs: Mapping[str, object]) -> Mapping[str, object]: - metadata: Final = kwargs.get("metadata") - if not isinstance(metadata, Mapping): - return MappingProxyType({**kwargs}) - return MappingProxyType({**kwargs, "metadata": dict(metadata)}) - - -def _with_router_resolved_session_model(session: object, model_name: str) -> Mapping[str, Mapping[str, object]]: - """ - Realtime client-secret requests carry the model inside ``session`` as well, and the caller's copy of it still - holds the pre-routing model group name, so it has to follow the deployment the router just picked. - - Returns kwargs to merge into the downstream call, empty when there is no session model to resolve. - """ - try: - typed_session: Final = _SESSION_ADAPTER.validate_python(session) - except ValidationError: - return _NO_SESSION_KWARGS - if "model" not in typed_session and typed_session.get("type") != "transcription": - return _NO_SESSION_KWARGS - from litellm.realtime_api.main import with_resolved_session_model - - return MappingProxyType({"session": with_resolved_session_model(dict(typed_session), model_name)}) - - -# Router._aanthropic_messages_streaming_iterator buffers lifecycle chunks -# until real content commits the primary stream, and only while a fallback -# can still take over; a hostile or slow-starting upstream that never emits -# content or an error could otherwise grow that buffer without bound, so -# hitting this cap forces an early commit instead. -MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS: Final = 200 - - -def _anthropic_stream_forwards_ping_live(chunk: object, has_generated_content: bool) -> bool: - """A `ping` keepalive reaches the client live whenever the stream has not committed: it carries no - lifecycle, so it cannot create overlapping lifecycles on the wire, and it keeps the connection alive - while lifecycle frames sit buffered for a possible fallback during a long thinking pass.""" - from litellm.llms.anthropic.messages.native.streaming_iterator import is_anthropic_ping_chunk - - return not has_generated_content and is_anthropic_ping_chunk(chunk) - - -def _is_retriable_anthropic_status(status_code: int) -> bool: - return status_code == 429 or status_code >= 500 - - -def _without_line_breaks(value: object) -> str: - return str(value).replace("\r", "").replace("\n", "") - - -def _anthropic_stream_error_is_gateway_verdict(chunk: object) -> bool: - """AgenticAnthropicStreamingIterator's own retrieval-failure frame is the gateway's verdict, not a provider - failure: another deployment would rerun the same failed hook, so it reaches the client instead of falling back.""" - from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import ( - is_server_fulfilled_tool_leak_error, - ) - - return is_server_fulfilled_tool_leak_error(chunk) - - -def _anthropic_stream_should_decline_fallback(has_generated_content: bool, error: "MidStreamFallbackError") -> bool: - """ - A MidStreamFallbackError raised directly by the source iterator (the - completion-bridge path's CustomStreamWrapper, e.g. on a transport drop) - carries its own pre_first_chunk bookkeeping - gated the same way a - detected SSE error event is, so a fallback is never appended after real - content already reached the client on either path. - """ - return has_generated_content or not error.is_pre_first_chunk - - -def _stream_raised_error_status(error: Exception) -> int | None: - raw_status: Final = getattr(error, "status_code", None) - if isinstance(raw_status, int): - return raw_status - if isinstance(raw_status, str) and raw_status.isdigit(): - return int(raw_status) - response_status: Final = getattr(getattr(error, "response", None), "status_code", None) - return response_status if isinstance(response_status, int) else None - - -def _stream_fallback_error_for_raised( - error: Exception, model: str, has_generated_content: bool, llm_provider: str = "anthropic" -) -> "MidStreamFallbackError | None": - """The pre-stream retry rule (408, 409, 429, 5xx); None means the raise propagates unchanged.""" - if has_generated_content: - return None - status_code: Final = _stream_raised_error_status(error) - if status_code is None: - return _stream_pre_content_error(error, model, llm_provider) - retriable: Final = python_host.invoke("should_retry")(status_code) - return _stream_pre_content_error(error, model, llm_provider) if retriable else None - - -def _stream_pre_content_error( - error: Exception, model: str, llm_provider: str = "anthropic" -) -> "MidStreamFallbackError": - """The envelope the fallback chain judges a failure by when the client has received no content yet.""" - from litellm.exceptions import MidStreamFallbackError - - return MidStreamFallbackError( - message=str(error), - model=model, - llm_provider=llm_provider, - original_exception=error, - is_pre_first_chunk=True, - ) - - -def _context_window_stream_fallback_error( - error: "MidStreamFallbackError | litellm_exceptions.ContextWindowExceededError", - model: object, - llm_provider: str, - has_output: bool, -) -> "MidStreamFallbackError": - if isinstance(error, litellm_exceptions.ContextWindowExceededError): - if has_output: - raise error - return _stream_pre_content_error(error, model if isinstance(model, str) else "", llm_provider) - if isinstance(error.original_exception, litellm_exceptions.ContextWindowExceededError) and ( - has_output or error.generated_content or error.emitted_disqualifying_content - ): - raise error.original_exception - return error - - -def _deployment_num_retries(deployment: "Deployment | None") -> int | None: - """The deployment's own num_retries litellm_param, an int or a digit string the way the config loader leaves it.""" - configured: Final = getattr(deployment.litellm_params, "num_retries", None) if deployment is not None else None - if isinstance(configured, bool) or not isinstance(configured, (int, str)): - return None - return int(configured) if str(configured).isdigit() else None - - -def _request_fallback_list( - kwargs: Mapping[str, object], key: str, router_default: "list[object] | None" -) -> "list[object] | None": # mutable-ok: should_retry_this_error's own parameter type - return cast("list[object] | None", kwargs.get(key, router_default)) # cast-ok: same type as the router attribute - - -def _request_model_group(kwargs: Mapping[str, object]) -> str | None: - model_group: Final = kwargs.get("model") - return model_group if isinstance(model_group, str) else None - - -def _mid_stream_error_before_output(error: "MidStreamFallbackError") -> bool: - return not error.generated_content or error.is_pre_first_chunk - - -def _mid_stream_retry_trigger(error: "MidStreamFallbackError") -> Exception: - """The provider's own error, which is what the retry policy and should_retry_this_error classify.""" - return error.original_exception if error.original_exception is not None else error - - -_RestT: Final = TypeVar("_RestT", Iterator[object], AsyncIterator[object]) -_RouteDeploymentsT: Final = TypeVar("_RouteDeploymentsT", bound=Sequence[Mapping[str, object]] | Mapping[str, object]) - - -@dataclass(frozen=True, slots=True) -class _FirstChunk(Generic[_RestT]): - """ - The chunk Router._acompletion / _completion already read off a stream (None: the - provider closed it empty) and the iterator it was read from, which carries on after it. - """ - - chunk: ModelResponseStream | None - rest: _RestT - - -async def _aread_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk[AsyncIterator[object]]: - """ - Read the first chunk while Router._acompletion still holds the request, so a - stream that fails before any token raises there, inside async_function_with_retries - and the fallback chain, instead of surfacing to the caller mid-iteration. - - Iterates the way `async for` does, one __aiter__ call and then __anext__ on what it - returned: a CustomStreamWrapper subclass may hand iteration to another object and - leave completion_stream None (the MCP gateway's MCPStreamWrapper). - """ - from litellm.exceptions import MidStreamFallbackError - - rest: Final = stream.__aiter__() - try: - return _FirstChunk(chunk=await rest.__anext__(), rest=rest) - except StopAsyncIteration: - return _FirstChunk(chunk=None, rest=rest) - except MidStreamFallbackError as e: - raise _mid_stream_retry_trigger(e) from e - - -def _read_first_chunk(stream: CustomStreamWrapper) -> _FirstChunk[Iterator[object]]: - """Sync twin of _aread_first_chunk for Router._completion.""" - from litellm.exceptions import MidStreamFallbackError - - rest: Final = iter(stream) - try: - return _FirstChunk(chunk=next(rest), rest=rest) - except StopIteration: - return _FirstChunk(chunk=None, rest=rest) - except MidStreamFallbackError as e: - raise _mid_stream_retry_trigger(e) from e - - -def _chunks_with_first( - model_response: Iterable[object], taken: _FirstChunk[Iterator[object]] | None -) -> Iterator[object]: - """The stream with the chunk _read_first_chunk already read put back in front (None: the stream as is).""" - if taken is None: - yield from model_response - return - if taken.chunk is None: - return - yield taken.chunk - yield from taken.rest - - -async def _achunks_with_first( - model_response: AsyncIterable[object], taken: _FirstChunk[AsyncIterator[object]] | None -) -> AsyncIterator[object]: - """Async twin of _chunks_with_first.""" - if taken is None: - async for item in model_response: - yield item - return - if taken.chunk is None: - return - yield taken.chunk - async for item in taken.rest: - yield item - - -def _stream_commits_now(chunk: object, has_generated_content: bool, buffered_chunk_count: int) -> bool: - """ - Whether `chunk` should make Router._aanthropic_messages_streaming_iterator - commit to the primary Anthropic stream (real content arrived, or the - pre-content buffer cap was hit) rather than keep buffering lifecycle - frames toward a possible fallback. - """ - from litellm.llms.anthropic.messages.native.streaming_iterator import ( - is_anthropic_content_delta_chunk, - ) - - if has_generated_content: - return False - return is_anthropic_content_delta_chunk(chunk) or buffered_chunk_count >= MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS - - -@dataclass(frozen=True, slots=True) -class _StreamRetryOpened: - response: object - attempted_retries: int - max_retries: int - - -@dataclass(frozen=True, slots=True) -class _StreamRetriesExhausted: - error: "MidStreamFallbackError" - - -_StreamRetryOutcome: TypeAlias = _StreamRetryOpened | _StreamRetriesExhausted - - -MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: Final = 200 - - -def _retry_policy_ceiling(policy: RetryPolicy) -> int: - """The most retries any error class under this policy can be granted.""" - retry_counts: Final[tuple[object, ...]] = tuple(policy.model_dump().values()) - return max((retries for retries in retry_counts if isinstance(retries, int)), default=0) - - -def _responses_stream_holds_event(item: object, held_event_count: int) -> bool: - from litellm.responses.streaming_iterator import PRE_OUTPUT_LIFECYCLE_EVENT_TYPES - - if held_event_count >= MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS: - return False - return getattr(item, "type", None) in PRE_OUTPUT_LIFECYCLE_EVENT_TYPES - - -def _responses_event_announces_background_run(item: object) -> bool: - return getattr(getattr(item, "response", None), "background", None) is True - - -class FallbackAwareAnthropicMessagesStream: - """ - Bare async generators can't carry the `_hidden_params` attribute the - proxy reads response headers off of (see - router_utils.add_retry_fallback_headers.get_hidden_params_dict), so this - thin wrapper carries it through from the source iterator - mirrors - AnthropicMessagesStreamingResponse. Used by - Router._aanthropic_messages_streaming_iterator. - """ - - def __init__(self, async_generator: AsyncGenerator[bytes, None], source_iterator: object) -> None: - self._async_generator = async_generator - self._source_iterator = source_iterator - self.fallback_headers_adopted = False - self._hidden_params = dict(getattr(source_iterator, "_hidden_params", None) or {}) - self._followed_source_params: object = None - - @property - def has_buffered_provider_output(self) -> bool: - return getattr(self._source_iterator, "has_buffered_provider_output", False) is True - - @property - def chunks(self) -> list[ModelResponseStream] | None: - return cast( # cast-ok: chunks is a list of ModelResponseStream on the inner stream - "list[ModelResponseStream] | None", getattr(self._source_iterator, "chunks", None) - ) - - @property - def messages(self) -> list[AllMessageValues] | None: - return cast( # cast-ok: messages is a list of AllMessageValues on the inner stream - "list[AllMessageValues] | None", getattr(self._source_iterator, "messages", None) - ) - - @property - def model(self) -> str | None: - return cast( # cast-ok: model is a str on the inner stream - "str | None", getattr(self._source_iterator, "model", None) - ) - - def adopt_fallback_source(self, fallback_response: object) -> None: - self._source_iterator = fallback_response - self.fallback_headers_adopted = True - - def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": - from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect - - return partial_responses_incomplete_event_on_disconnect(self._source_iterator) - - def follow_source_attribution(self) -> None: - """ - A retry's or a fallback's stream carries a wrapper of its own, so a hop it makes before its - first byte lands on that wrapper while the proxy reads the headers off this one. Mirrors the - source's current attribution onto this wrapper whenever the source has adopted a new one. - """ - source: Final = self._source_iterator - if not getattr(source, "fallback_headers_adopted", False): - return - source_params: Final = getattr(source, "_hidden_params", None) - if source_params is None or source_params is self._followed_source_params: - return - self._followed_source_params = source_params - hidden_params, headers = PythonCore._prepare_fallback_hidden_params(source) # pyright: ignore[reportPrivateUsage] # this wrapper is the Router's own stream type - self.merge_fallback_hidden_params(hidden_params, headers) - - def __aiter__(self) -> "FallbackAwareAnthropicMessagesStream": - return self - - async def __anext__(self) -> bytes: - return await self._async_generator.__anext__() - - async def aclose(self) -> None: - await self._async_generator.aclose() - - def merge_fallback_hidden_params( - self, - fallback_hidden_params: Mapping[str, object], - fallback_headers: Mapping[str, object], - ) -> None: - """ - Raw bytes can't carry their own _hidden_params the way a - ModelResponseStream/ResponsesAPI event can, so a mid-stream - fallback's provider headers (e.g. Bedrock's x-amzn-requestid) are - merged onto the wrapper itself instead - mirrors - Router._apply_fallback_hidden_params_to_item's merge shape. - """ - existing_headers: Final = cast( # cast-ok: additional_headers is always a dict[str, object] when present - "dict[str, object]", self._hidden_params.get("additional_headers") or {} - ) - self._hidden_params = { - **self._hidden_params, - **fallback_hidden_params, - "additional_headers": dict(replace_complexity_router_headers(existing_headers, fallback_headers)), - } - - -class RoutingArgs(enum.Enum): - ttl = 60 # 1min (RPM/TPM expire key) - - -# Routers that are still in use, so a price data reload can rebuild the cost-map -# entries their deployments own. Weak so a router nothing references any more, such -# as the per-request one built from a caller-supplied user_config, drops out on its -# own rather than leaving entries behind that nothing can withdraw. -_live_routers: Final["weakref.WeakSet[PythonCore]"] = weakref.WeakSet() -# Shared with the other router cores of the process (litellm/router_utils/global_callback_holders.py). -_routers_holding_global_callbacks: Final["weakref.WeakSet[PythonCore]"] = routers_holding_global_callbacks # pyright: ignore[reportAssignmentType] # every member serves the holder members -_routers_holding_global_callbacks_lock: Final = routers_holding_global_callbacks_lock - - -def _replay_live_router_model_cost() -> None: - """Re-assert every live router's deployments after the cost map is refreshed.""" - for router in tuple(_live_routers): - router.replay_model_cost_registrations() - - -set_live_deployment_replay(_replay_live_router_model_cost) - - -RETRY_BREADCRUMB_LIMIT: Final = 4 - - -class FallbackAwareStreamWrapper(CustomStreamWrapper): - """Base for the Router's chat-completion stream wrappers, which are built around the - attempt the Router picked first and have to repoint themselves when a fallback takes over.""" - - fallback_headers_adopted: bool = False - - def adopt_fallback_response_headers( - self, - fallback_response: object, - prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]], - ) -> None: - """Repoint this wrapper at the deployment that served the stream. - - Replaces rather than merges, so the failed attempt's `x-request-id`, rate limit - counters, `model_id` and `api_base` cannot reach the proxy's response headers or - its callbacks. - """ - self._response_headers = getattr(fallback_response, "_response_headers", None) - fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params - if fallback_hidden_params: - self._hidden_params = { - **fallback_hidden_params, - # dict() because add_retry_fallback_headers mutates additional_headers in place - "additional_headers": dict(fallback_headers), - } - self._base_hidden_params = { - **self._hidden_params, - "response_cost": None, - } - self.fallback_headers_adopted = True - - -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: - cap: Final = int(float(value)) - except (ValueError, OverflowError): - return None - return cap if cap >= 0 else 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) - for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")) - if isinstance(bucket, dict) and CLIENT_OUTPUT_CEILING_METADATA_KEY in bucket - ), - None, - ) - if not isinstance(stamped, dict): - return None - return MappingProxyType( - { - carrier: cap - for carrier, value in stamped.items() - if carrier in OUTPUT_TOKEN_CEILING_PARAMS and (cap := as_output_cap(value)) is not None - } - ) - - -def _exact_base_model_max_input_tokens(base_model: str) -> int | None: - try: - model_info: Final = python_host.invoke("get_model_info")(model=base_model) - except Exception: # noqa: BLE001 # get_model_info raises ModelNotMappedError or a bare Exception for an unmapped model - return None - if is_generalized_model_info(model_info): - return None - return model_info.get("max_input_tokens") - - -def _provider_model_id(model_name: str, custom_llm_provider: str) -> str: - model_id: Final = model_name.rsplit("/", 1)[-1] - return get_bedrock_base_model(model_id) if custom_llm_provider.startswith("bedrock") else model_id - - -def _names_route_model(base_model: str, route_key: str, custom_llm_provider: str) -> bool: - route_model_id: Final = _provider_model_id(route_key, custom_llm_provider) - base_model_id: Final = _provider_model_id(base_model, custom_llm_provider) - return route_model_id in (base_model_id, base_model_id.split("@", 1)[0]) - - -def _declared_base_model(user_model_info: Mapping[str, object]) -> str | None: - match user_model_info: - case {"base_model": str() as base_model}: - return base_model - case _: - return None - - -def _base_model_widened_window( - route_model_info: ModelMapInfo, base_model: str | None, custom_llm_provider: str -) -> Mapping[str, int]: - route_max_input_tokens: Final = route_model_info.get("max_input_tokens") - if base_model is None or route_max_input_tokens is None: - return MappingProxyType({}) - if not _names_route_model(base_model, route_model_info["key"], custom_llm_provider): - return MappingProxyType({}) - base_max_input_tokens: Final = _exact_base_model_max_input_tokens(base_model) - if base_max_input_tokens is None or base_max_input_tokens <= route_max_input_tokens: - return MappingProxyType({}) - return MappingProxyType({"max_input_tokens": base_max_input_tokens}) - - -_SETTING_ATTRIBUTES: Final = frozenset("_core_" + name for name in SETTING_FIELDS) - - -class PythonCore(DeprecatedRouterMembers): - _core_model_names: set[str] = set() - _core_cache_responses: bool | None = False - _registry_version: int = 0 - # Bumped by every ``apply``: the static arguments plus ``_core_applied`` are the settings. - _settings_version: int = 0 - _core_applied: Mapping[str, object] = MappingProxyType({}) - - # -- entry points: everything outside code needs besides the endpoint methods --------------- - # See litellm/router_api.py. The members outside code used before these live in - # DeprecatedRouterMembers (litellm/router_deprecated.py). - - def registry(self) -> RouterRegistry: - """Deployment lookups: by id, group, alias, wildcard and team, plus configured plugins.""" - cached: Final = self.__dict__.get("_registry_view") - if isinstance(cached, RouterRegistry): - return cached - view: Final = RouterRegistry(self) - self.__dict__["_registry_view"] = view - return view - - def settings(self) -> RouterSettings: - """ - A frozen snapshot of the router's configuration. It is built once and reused until a - setting attribute is rebound or deleted (``__setattr__``/``__delattr__`` drop it), so the - request path can read it cheaply. - """ - cached: Final = self.__dict__.get("_settings_snapshot") - if isinstance(cached, RouterSettings): - return cached - snapshot: Final = RouterSettings.of(self) - self.__dict__["_settings_snapshot"] = snapshot - return snapshot - - def __setattr__(self, name: str, value: object) -> None: - if ( - (name in PLAIN_SETTING_ATTRIBUTES or name in _SETTING_ATTRIBUTES) - and self.__dict__.get("_core_ready") - and not self.__dict__.get("_core_applying") - ): - # Settings change only through a changeset, so every core serving this router sees it. - self.apply(ChangeSet(assign={name.removeprefix("_core_"): value})) - return - if name in _SETTING_ATTRIBUTES: - self.__dict__.pop("_settings_snapshot", None) - elif name == "_core_model_list" and "_core_model_list" in self.__dict__: - # A whole new deployment list (``router.model_list = ...``) is a deployment change too. - self.__dict__["_registry_version"] = self._registry_version + 1 - super().__setattr__(name, value) - - def __delattr__(self, name: str) -> None: - if name in _SETTING_ATTRIBUTES: - self.__dict__.pop("_settings_snapshot", None) - super().__delattr__(name) - - def health(self) -> RouterHealth: - """Cooldowns, blocked deployments, health checks and usage.""" - cached: Final = self.__dict__.get("_health_view") - if isinstance(cached, RouterHealth): - return cached - view: Final = RouterHealth(self) - self.__dict__["_health_view"] = view - return view - - @overload - async def select( - self, - model: str, - request_kwargs: dict[str, object] | None = None, # mutable-ok: the selection cores take a dict - *, - kind: Literal["pick", "pass_through"] = "pick", - messages: list[dict[str, str]] | None = None, # mutable-ok: the selection cores take a list - input: str | list[object] | None = None, # mutable-ok: the selection cores take a list - specific_deployment: bool | None = False, - parent_otel_span: Span | None = None, - health_check_probe: bool = False, - ) -> "UntypedDict": ... # a mutable dict, as before: the cores' untyped deployment dict, unchanged - - @overload - async def select( - self, - model: str, - request_kwargs: dict[str, object] | None = None, # mutable-ok: the selection cores take a dict - *, - kind: Literal["healthy"], - messages: list[dict[str, str]] | None = None, # mutable-ok: the selection cores take a list - input: str | list[object] | None = None, # mutable-ok: the selection cores take a list - specific_deployment: bool | None = False, - parent_otel_span: Span | None = None, - health_check_probe: bool = False, - ) -> "list[UntypedDict] | UntypedDict": ... # mutable-ok: the core's type (list[dict] | dict), unchanged - - async def select( # pyright: ignore[reportUnknownParameterType] # the selection cores return untyped deployment dicts - self, - model: str, - request_kwargs: dict[str, object] | None = None, # mutable-ok: the selection cores take a dict - *, - kind: SelectKind = "pick", - messages: list[dict[str, str]] | None = None, # mutable-ok: the selection cores take a list - input: str | list[object] | None = None, # mutable-ok: the selection cores take a list - specific_deployment: bool | None = False, - parent_otel_span: Span | None = None, - health_check_probe: bool = False, - ): - """ - Pick a deployment for ``model`` without calling it (``pick``), list the deployments that - are healthy for the request (``healthy``), or pick one for a pass-through route - (``pass_through``). - """ - if kind == "pick": - with picks_reserve(False): - return await self._core_async_get_available_deployment( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - request_kwargs=request_kwargs if request_kwargs is not None else {}, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) - if kind == "healthy": - return await self._core_async_get_healthy_deployments( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - request_kwargs=request_kwargs if request_kwargs is not None else {}, - messages=messages, - input=input, - specific_deployment=specific_deployment, - parent_otel_span=parent_otel_span, - health_check_probe=health_check_probe, - ) - if kind == "pass_through": - with picks_reserve(False): - return self._core_get_available_deployment_for_pass_through( # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # the selection cores return untyped deployment dicts - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) - raise ValueError(f"unknown select kind {kind!r}") - - def apply(self, changes: ChangeSet) -> ApplyResult: - """Apply runtime changes: settings, deployments, the routing state's Redis cache.""" - with self._core_applying_changes(changes): - return apply_changes(self, changes) - - async def aapply(self, changes: ChangeSet) -> ApplyResult: - """``apply``, plus the changes that need I/O (``refresh_model_info``).""" - with self._core_applying_changes(changes): - return await aapply_changes(self, changes) - - @contextlib.contextmanager - def _core_applying_changes(self, changes: ChangeSet) -> Generator[None, None, None]: - """Mark a changeset in progress, then record its settings and bump ``_settings_version``.""" - depth: Final = self.__dict__.get("_core_applying", 0) - self.__dict__["_core_applying"] = depth + 1 - try: - yield - finally: - self.__dict__["_core_applying"] = depth - if changes.assign or changes.settings: - self.__dict__["_core_applied"] = MappingProxyType( - {**self._core_applied, **changes.assign, **changes.settings} - ) - self.__dict__["_settings_version"] = self._settings_version + 1 - - def observe(self, logging_obj: "LiteLLMLogging | None" = None, *, request_scoped: bool = False) -> RoutingObservers: - """ - The router's routing observers. ``request_scoped=True`` stops registering them globally, - so each request binds them instead; passing a request's ``logging_obj`` binds them to it. - """ - if request_scoped: - self._core_serve_through_request_logging() - if logging_obj is not None: - self._core_bind_to_request_logging(logging_obj) - return RoutingObservers(slack_alerting_logger=getattr(self, "_core_slack_alerting_logger", None)) - - def register_plugin(self, plugin: CustomRoutingStrategyBase) -> None: - """Route with a custom strategy instead of the configured one.""" - self._core_set_custom_routing_strategy(plugin) - - def retire(self) -> None: - """Release the router's global callback registrations and routing state.""" - self._core_discard() - - def _core_public(self) -> "PythonCore": - """The ``litellm.Router`` this core serves, for objects the router hands out; the core itself if none.""" - owner: Final = self.__dict__.get("_core_owner") - public: Final = owner() if owner is not None else None - return self if public is None else public - - async def call(self, op: str, request: Mapping[str, object]) -> object: - """Run one endpoint (``acompletion``, ``aembedding``, ...) with the request's keyword payload.""" - result: Final = getattr(self, op)(**request) - return await result if inspect.isawaitable(result) else result - - default_cache_time_seconds: int = 1 * 60 * 60 # 1 hour - tenacity = None - leastbusy_logger: LeastBusyLoggingHandler | None = None - lowesttpm_logger: LowestTPMLoggingHandler | None = None - optional_callbacks: list[CustomLogger | Callable | str] | None = None - - def __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/", - "api_key": , - "api_version": , - "api_base": - }, - }, - { - "model_name": "azure-gpt-3.5-turbo", # model alias - "litellm_params": { # params for litellm completion/embedding call - "model": "azure/", - "api_key": , - "api_version": , - "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": , - }, - ] - - router = Router(model_list=model_list, fallbacks=[{"azure-gpt-3.5-turbo": ["openai-gpt-3.5-turbo"]}]) - ``` - """ - self._host: Final = _host - self._core_ctx: Final = RouterContext(self) - - self.set_verbose = set_verbose - self.ignore_invalid_deployments = ignore_invalid_deployments - self.auto_router_capability_limit = auto_router_capability_limit - self._core_fallback_access_check: Final = fallback_access_check - self._core_fallback_budget_check: Final = fallback_budget_check - self.debug_level = debug_level - self._core_enable_pre_call_checks = enable_pre_call_checks - self._core_enable_tag_filtering = enable_tag_filtering - self._core_tag_filtering_match_any = tag_filtering_match_any - self._core_tag_routing_prefix = tag_routing_prefix - from litellm._service_logger import ServiceLogging - - self.service_logger_obj: ServiceLogging = ServiceLogging() - self._host.suppress_debug_info() # prevents 'Give Feedback/Get help' message from being emitted on Router - Relevant Issue: https://github.com/BerriAI/litellm/issues/5942 - if self.set_verbose is True: - if debug_level == "INFO": - verbose_router_logger.setLevel(logging.INFO) - elif debug_level == "DEBUG": - verbose_router_logger.setLevel(logging.DEBUG) - self._core_router_general_settings: RouterGeneralSettings = router_general_settings or RouterGeneralSettings() - - self.assistants_config = assistants_config - self._core_search_tools = search_tools or [] - self._core_guardrail_list = guardrail_list or [] - self._core_deployment_names: list[str] = [] # names of models under litellm_params. ex. azure/chatgpt-v-2 - self._core_deployment_latency_map = {} - ### CACHING ### - cache_type: Literal["local", "redis", "redis-semantic", "s3"] = "local" # default to an in-memory cache - redis_cache = None - cache_config: Final[dict[str, Any]] = {} - - self.client_ttl = client_ttl - self._core_kubernetes_pod_discovery: Final = KubernetesPodDiscovery() - if redis_url is not None or (redis_host is not None and redis_port is not None): - cache_type = "redis" - - if redis_url is not None: - cache_config["url"] = redis_url - - if redis_host is not None: - cache_config["host"] = redis_host - - if redis_port is not None: - cache_config["port"] = str(redis_port) - - if redis_password is not None: - cache_config["password"] = redis_password - - if redis_db is not None: - verbose_router_logger.warning( - "Deprecated 'redis_db' argument used. Please remove 'redis_db' from your config/database and use 'cache_kwargs' instead." - ) - cache_config["db"] = str(redis_db) - - # Add additional key-value pairs from cache_kwargs - cache_config.update(cache_kwargs) - redis_cache = self._create_redis_cache(cache_config) - - if cache_responses: - if self._host.global_cache() is None: - # the cache can be initialized on the proxy server. We should not overwrite it - self._host.set_global_cache(LitellmCache(type=cache_type, **cache_config)) - self._core_cache_responses = cache_responses - self._core_cache = DualCache( - redis_cache=redis_cache, in_memory_cache=InMemoryCache() - ) # use a dual cache (Redis+In-Memory) for tracking cooldowns, usage, etc. - self._core_least_busy_reservations: InMemoryCache = InMemoryCache( - max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS - ) - self._claude_code_session_router_cache: DualCache = DualCache( - redis_cache=redis_cache, - in_memory_cache=InMemoryCache(), - ) - self._core_prompt_caching_pin_cache: DualCache = DualCache( - redis_cache=redis_cache, - in_memory_cache=InMemoryCache(max_size_in_memory=PROMPT_CACHE_PIN_CACHE_MAX_SIZE), - ) - - ### SCHEDULER ### - self.scheduler = Scheduler(polling_interval=polling_interval, redis_cache=redis_cache) - parsed_default_priority: Final = parse_default_priority(default_priority) - if isinstance(parsed_default_priority, InvalidPriority): - verbose_router_logger.warning( - "router_settings.default_priority is ignored, no request is queued by default: %s", - parsed_default_priority.message, - ) - self.default_priority = ( - None if isinstance(parsed_default_priority, InvalidPriority) else parsed_default_priority - ) - self._core_default_deployment = ( - None # use this to track the users default deployment, when they want to use model = * - ) - self._core_default_max_parallel_requests = default_max_parallel_requests - self._core_provider_default_deployment_ids: list[str] = [] - self._core_pattern_router = PatternMatchRouter() - self._core_team_pattern_routers: dict[str, PatternMatchRouter] = {} # {"TEAM_ID": PatternMatchRouter} - self._core_auto_routers: dict[str, list[TaggedPreRoutingStrategy[AutoRouter]]] = {} - self._core_complexity_routers: dict[str, list[TaggedPreRoutingStrategy[ComplexityRouter]]] = {} - self._core_adaptive_routers: dict[str, list[TaggedPreRoutingStrategy[AdaptiveRouter]]] = {} - self._core_quality_routers: dict[str, list[TaggedPreRoutingStrategy[QualityRouter]]] = {} - self._core_routing_plugins: list[RoutingPlugin] = list(plugins) if plugins else [] - - # Initialize model_group_alias early since it's used in set_model_list - self._core_model_group_alias: dict[str, str | RouterModelGroupAliasItem] = ( - model_group_alias or {} - ) # dict to store aliases for router, ex. {"gpt-4": "gpt-3.5-turbo"}, all requests with gpt-4 -> get routed to gpt-3.5-turbo group - - # Initialize model ID to deployment index mapping for O(1) lookups - self.model_id_to_deployment_index_map: dict[str, int] = {} - # Initialize model name to deployment indices mapping for O(1) lookups - # Maps model_name -> list of indices in model_list - self._core_model_name_to_deployment_indices: dict[str, list[int]] = {} - # Maps (team_id, team_public_model_name) -> list of indices in model_list - self._core_team_model_to_deployment_indices: dict[tuple[str, str], list[int]] = {} - self._core_team_public_model_names: frozenset[str] = frozenset() - # Guards model_list and the three index maps above: writers change them, and readers turn an - # index into a deployment, under it, so a lookup never lands between a removal and its re-index - self._deployment_index_lock = threading.RLock() - - # Initialize cache attributes that ``_invalidate_model_group_info_cache`` - # and ``_invalidate_access_groups_cache`` touch *before* the first - # ``set_model_list`` below (which calls those invalidations as part of - # building the model index) and before ``_init_routing_groups(None)`` - # (which calls them on every group rebuild). - self._access_groups_cache: dict[str, list[str]] | None = None - # Per-router cache for the proxy auth-layer "is this model explicitly - # zero-cost?" check. Lives on the router so it is invalidated alongside - # ``_cached_get_model_group_info`` and dies with the router (no - # ``id()``-reuse risk after GC). See - # ``litellm.proxy.auth.auth_checks._is_model_cost_zero``. - self._zero_cost_cache: dict[str, bool] = {} - self._core_cached_deployment_model_info = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)( - self._core_get_deployment_model_info - ) - self._ptu_share_teams = lru_cache(maxsize=1)(self._collect_ptu_share_teams) - self._cached_team_ptu_ceiling = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(self._resolve_team_ptu_ceiling) - self._cached_deployment_tracks_usage = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)( - self._deployment_tracks_usage - ) - self.served_openai_models = lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE)(self._served_openai_models) - self._discovered_model_info_cache: InMemoryCache = InMemoryCache( - max_size_in_memory=max(len(model_list or ()), 1), - default_ttl=2 * MODEL_INFO_REFRESH_SECONDS, - ) - self._routing_group_rows: tuple[DeploymentTypedDict, ...] | None = None - self._init_routing_groups(None) - self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = () - - self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds - self.session_affinity_repin_on_failover = session_affinity_repin_on_failover - self.model_group_affinity_config = model_group_affinity_config - warn_on_unknown_model_group_affinity_flags(model_group_affinity_config) - - if model_list is not None: - # set_model_list will build indices automatically - self.set_model_list(model_list) - self.healthy_deployments: list = self._core_model_list - for m in model_list: - if "model" in m["litellm_params"]: - self._core_deployment_latency_map[m["litellm_params"]["model"]] = 0 - else: - self._core_model_list: list = [] # initialize an empty list - to allow _add_deployment and delete_deployment to work - # Track this router so a price data reload can rebuild its deployments' - # cost-map entries from the list it is serving at that moment, whether - # those deployments came from model_list or were added later from the DB. - _live_routers.add(self) - live_check_owners.add(self) - - if allowed_fails is not None: - self._core_allowed_fails = allowed_fails - else: - self._core_allowed_fails = self._host.dynamic().allowed_fails - self.explicit_allowed_fails = allowed_fails - self._core_cooldown_time = cooldown_time or DEFAULT_COOLDOWN_TIME_SECONDS - self._core_cooldown_cache = CooldownCache( - cache=self._core_cache, default_cooldown_time=self._core_cooldown_time - ) - self._core_disable_cooldowns = disable_cooldowns - self.fallback_on_unsupported_params = fallback_on_unsupported_params - self._core_enable_health_check_routing = enable_health_check_routing - self.enable_weighted_failover = enable_weighted_failover - self.enable_mid_stream_fallback_continuation = self._is_flag_enabled(enable_mid_stream_fallback_continuation) - self._core_health_check_ignore_transient_errors = health_check_ignore_transient_errors - self._core_background_health_check_model_groups: frozenset[str] | None = ( - frozenset(background_health_check_model_groups) - if background_health_check_model_groups is not None - else None - ) - _staleness: Final = health_check_staleness_threshold or ( - DEFAULT_HEALTH_CHECK_INTERVAL * DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER - ) - self._core_health_state_cache = DeploymentHealthCache( - cache=self._core_cache, staleness_threshold=float(_staleness) - ) - - if num_retries is not None: - self.num_retries = num_retries - elif self._host.dynamic().num_retries is not None: - self.num_retries = self._host.dynamic().num_retries - else: - self.num_retries = openai.DEFAULT_MAX_RETRIES - - if max_fallbacks is not None: - self.max_fallbacks = max_fallbacks - elif self._host.dynamic().max_fallbacks is not None: - self.max_fallbacks = self._host.dynamic().max_fallbacks - else: - self.max_fallbacks = self._host.dynamic().router_max_fallbacks - - self._explicit_timeout = timeout # None when user did not pass timeout - self.timeout = timeout or self._host.dynamic().request_timeout - # Per-attempt request_timeout, independent of router_settings.timeout. - # Only stored when a router timeout is also set, since otherwise - # request_timeout already flows through self.timeout above. - self.request_timeout = get_configured_request_timeout() if timeout is not None else None - self.stream_timeout = stream_timeout - - self.retry_after = retry_after - self._core_routing_strategy = self._normalize_strategy(routing_strategy) - self._routing_groups_input: list[RoutingGroup | dict] | None = routing_groups - - ## SETTING FALLBACKS ## - ### validate if it's set + in correct format - _fallbacks = fallbacks or self._host.dynamic().fallbacks - - self.validate_fallbacks(fallback_param=_fallbacks) - ### set fallbacks - self._core_fallbacks = _fallbacks - - if default_fallbacks is not None or self._host.dynamic().default_fallbacks is not None: - _fallbacks = default_fallbacks or self._host.dynamic().default_fallbacks - if isinstance(_fallbacks, str): - raise TypeError( - f"default_fallbacks must be a list of model names, got '{_fallbacks}'. Use ['{_fallbacks}']" - ) - if self._core_fallbacks is not None: - self._core_fallbacks.append({"*": _fallbacks}) - else: - self._core_fallbacks = [{"*": _fallbacks}] - - self._core_context_window_fallbacks = context_window_fallbacks or self._host.dynamic().context_window_fallbacks - - _content_policy_fallbacks: Final = content_policy_fallbacks or self._host.dynamic().content_policy_fallbacks - self.validate_fallbacks(fallback_param=_content_policy_fallbacks) - self._core_content_policy_fallbacks = _content_policy_fallbacks - self.total_calls: defaultdict = defaultdict(int) # dict to store total calls made to each model - self.fail_calls: defaultdict = defaultdict(int) # dict to store fail_calls made to each model - self.success_calls: defaultdict = defaultdict(int) # dict to store success_calls made to each model - - # make Router.chat.completions.create compatible for openai.chat.completions.create - default_litellm_params = default_litellm_params or {} - self.chat = LitellmChat(params=default_litellm_params, router_obj=self._core_public()) - - # default litellm args - self._core_default_litellm_params = default_litellm_params - self._core_default_litellm_params.setdefault("timeout", timeout) - self._core_default_litellm_params.setdefault("max_retries", 0) - self._core_default_litellm_params.setdefault("metadata", {}).update({"caching_groups": caching_groups}) - - self.deployment_stats: dict = {} # used for debugging load balancing - """ - deployment_stats = { - "122999-2828282-277: - { - "model": "gpt-3", - "api_base": "http://localhost:4000", - "num_requests": 20, - "avg_latency": 0.001, - "num_failures": 0, - "num_successes": 20 - } - } - """ - - ### ROUTING SETUP ### - if self._normalize_strategy(routing_strategy) == "lar1": - from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - - apply_lar1_routing_strategy(self._core_ctx, routing_strategy_args) - else: - self.routing_strategy_init( - routing_strategy=routing_strategy, - routing_strategy_args=routing_strategy_args, - ) - self._init_routing_groups(self._routing_groups_input) - self._override_selectors: dict[str, RouterStrategySelector | None] = {} - self._override_selectors_lock = threading.Lock() - self._request_scoped_selectors: tuple[CustomLogger, ...] = () - self._request_scoped_policies: tuple[CustomLogger, ...] = () - self.access_groups = None - with _routers_holding_global_callbacks_lock: - _routers_holding_global_callbacks.add(self) - ## USAGE TRACKING ## - self._host.callbacks.add_litellm_async_success_callback(self.deployment_callback_on_success) - if isinstance(self._host.callback_list("success_callback"), list): - self._host.callbacks.add_litellm_success_callback(self.sync_deployment_callback_on_success) - else: - self._host.set_callback_list("success_callback", [self.sync_deployment_callback_on_success]) - self._host.callbacks.add_litellm_async_failure_callback(self.async_deployment_callback_on_failure) - ## COOLDOWNS ## - if isinstance(self._host.callback_list("failure_callback"), list): - self._host.callbacks.add_litellm_failure_callback(self.deployment_callback_on_failure) - else: - self._host.set_callback_list("failure_callback", [self.deployment_callback_on_failure]) - self.routing_strategy_args = routing_strategy_args - self._core_provider_budget_config = provider_budget_config - self.router_budget_logger: RouterBudgetLimiting | None = None - if RouterBudgetLimiting.should_init_router_budget_limiter( - model_list=model_list, provider_budget_config=self._core_provider_budget_config - ): - if optional_pre_call_checks is not None: - optional_pre_call_checks.append("router_budget_limiting") - else: - optional_pre_call_checks = ["router_budget_limiting"] - self.retry_policy: RetryPolicy | None = None - if retry_policy is not None: - if isinstance(retry_policy, dict): - self.retry_policy = RetryPolicy(**retry_policy) - elif isinstance(retry_policy, RetryPolicy): - self.retry_policy = retry_policy - if self.retry_policy is not None: - verbose_router_logger.info( - "\x1b[32mRouter Custom Retry Policy Set:\n%s\x1b[0m", - self.retry_policy.model_dump(exclude_none=True), - ) - - self.model_group_retry_policy: dict[str, RetryPolicy] | None = model_group_retry_policy - - self._core_allowed_fails_policy: AllowedFailsPolicy | None = None - if allowed_fails_policy is not None: - if isinstance(allowed_fails_policy, dict): - self._core_allowed_fails_policy = AllowedFailsPolicy(**allowed_fails_policy) - elif isinstance(allowed_fails_policy, AllowedFailsPolicy): - self._core_allowed_fails_policy = allowed_fails_policy - - if self._core_allowed_fails_policy is not None: - verbose_router_logger.info( - "\x1b[32mRouter Custom Allowed Fails Policy Set:\n%s\x1b[0m", - self._core_allowed_fails_policy.model_dump(exclude_none=True), - ) - - self.alerting_config: AlertingConfig | None = alerting_config - - if optional_pre_call_checks is not None: - self.add_optional_pre_call_checks(optional_pre_call_checks) - - # If model_group_affinity_config is set but no global affinity checks were - # enabled, we still need the DeploymentAffinityCheck callback (with global - # flags all False) so per-group config can activate affinity per model group. - if self.model_group_affinity_config: - self._ensure_deployment_affinity_callback() - - if self.alerting_config is not None: - self._initialize_alerting() - - self.initialize_assistants_endpoint() - self.initialize_router_endpoints() - self.apply_default_settings() - # From here on a settings write is a changeset (``__setattr__``). - self.__dict__["_core_ready"] = True - - @staticmethod - def get_valid_args() -> list[str]: - """ - Returns a list of valid arguments for the Router.__init__ method. - """ - arg_spec: Final = inspect.getfullargspec(PythonCore.__init__) - return [arg for arg in (*arg_spec.args, *arg_spec.kwonlyargs) if arg not in ("self", "_host")] - - def apply_default_settings(self): - """ - Apply the default settings to the router. - """ - - default_pre_call_checks: Final[OptionalPreCallChecks] = [] - self.add_optional_pre_call_checks(default_pre_call_checks) - - if self.enable_mid_stream_fallback_continuation: - self._register_continuation_prefill_check() - - def _register_continuation_prefill_check(self) -> None: - # Process-global on purpose: Router._core_discard() must not drop a filter another router still needs. - self._host.callbacks.add_litellm_callback(ContinuationPrefillDeploymentCheck()) - - def _core_discard(self) -> None: - """ - Pseudo-destructor to be invoked to clean up global data structures when router is no longer used. - For now, unhook router's callbacks from all lists - """ - # Stop contributing to cost-map rebuilds straight away rather than waiting - # for this router to be collected. - _live_routers.discard(self) - live_check_owners.discard(self) - self.retire_routing_state() - self._release_global_callbacks() - - def _release_global_callbacks(self) -> None: - with _routers_holding_global_callbacks_lock: - _routers_holding_global_callbacks.discard(self) - slack_alerting_loggers: Final = self._slack_alerting_loggers() - adaptive_router_hooks: Final = self._adaptive_router_hooks() - vacated_keys: Final = self._host.callbacks.drop_callbacks( - ( - *self._callback_selectors(), - *self._override_selectors.values(), - *(self.optional_callbacks or ()), - *slack_alerting_loggers, - *adaptive_router_hooks, - ), - method_owners=(self, *slack_alerting_loggers), - ) - with _routers_holding_global_callbacks_lock: - survivors: Final = tuple(_routers_holding_global_callbacks) - candidates: Final = chain.from_iterable(router._global_callback_candidates() for router in survivors) - for candidate in self._host.callbacks.with_custom_logger_keys(candidates, vacated_keys): - self._register_router_selector(candidate) - adaptive_keys: Final = frozenset(map(self._host.callbacks.get_custom_logger_key, adaptive_router_hooks)) - adaptive_survivor: Final = next((router for router in survivors if router._core_adaptive_routers), None) - if adaptive_survivor is not None and vacated_keys & adaptive_keys: - adaptive_survivor._sync_adaptive_router_hooks() - - def _core_serve_through_request_logging(self) -> None: - self._request_scoped_selectors = (*self._callback_selectors(), *self._request_adaptive_router_hooks()) - self._request_scoped_policies = (*self._optional_loggers(), *self._slack_alerting_loggers()) - self._release_global_callbacks() - - def _core_bind_to_request_logging(self, logging_obj: LiteLLMLogging) -> None: - logging_obj.add_dynamic_function_callbacks( - async_success=self.deployment_callback_on_success, - failure=self.deployment_callback_on_failure, - async_failure=self.async_deployment_callback_on_failure, - ) - for selector in self._request_scoped_selectors: - logging_obj.replace_global_callback(selector) - for policy in self._request_policies_not_set_globally(): - logging_obj.add_dynamic_callback(policy) - - def _request_policies_not_set_globally(self) -> tuple[CustomLogger, ...]: - global_keys: Final = frozenset( - map( - self._host.callbacks.get_custom_logger_key, - (callback for callback in self._host.callback_list("callbacks") if isinstance(callback, CustomLogger)), - ) - ) - return tuple( - policy - for policy in self._request_scoped_policies - if self._host.callbacks.get_custom_logger_key(policy) not in global_keys - ) - - def _request_adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: - from litellm.router_strategy.adaptive_router.hooks import ( - AdaptiveRouterPostCallHook, - ) - - return tuple( - AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy) - for tagged in chain.from_iterable(self._core_adaptive_routers.values()) - ) - - def _routing_callbacks(self) -> Iterable[object]: - if not (self._request_scoped_selectors or self._request_scoped_policies): - return self._host.callback_list("callbacks") - replaced_keys: Final = frozenset( - map(self._host.callbacks.get_custom_logger_key, self._request_scoped_selectors) - ) - return chain( - self._host.callbacks.without_custom_logger_keys(self._host.callback_list("callbacks"), replaced_keys), - self._request_scoped_selectors, - self._request_policies_not_set_globally(), - ) - - def retire_routing_state(self) -> None: - for strategy in (*self._global_callback_candidates(), *self._override_selectors.values()): - if isinstance(strategy, BaseRoutingStrategy | RouterBudgetLimiting): - strategy.retire() - - def _slack_alerting_loggers(self) -> tuple[CustomLogger, ...]: - slack_alerting_logger: Final = getattr(self, "_core_slack_alerting_logger", None) - return (slack_alerting_logger,) if isinstance(slack_alerting_logger, CustomLogger) else () - - def _adaptive_router_hooks(self) -> tuple[CustomLogger, ...]: - from litellm.router_strategy.adaptive_router.hooks import ( - AdaptiveRouterPostCallHook, - ) - - own_strategy_ids: Final = frozenset( - id(tagged.strategy) for tagged in chain.from_iterable(self._core_adaptive_routers.values()) - ) - hooks: Final = self._host.callbacks.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) - return tuple( - hook - for hook in hooks - if isinstance(hook, AdaptiveRouterPostCallHook) and id(hook.adaptive_router) in own_strategy_ids - ) - - def _optional_loggers(self) -> tuple[CustomLogger, ...]: - return tuple(callback for callback in self.optional_callbacks or () if isinstance(callback, CustomLogger)) - - def _global_callback_candidates(self) -> tuple[CustomLogger, ...]: - return (*self._callback_selectors(), *self._optional_loggers(), *self._slack_alerting_loggers()) - - @staticmethod - def _create_redis_cache( - cache_config: dict[str, Any], - ) -> RedisCache | RedisClusterCache: - """ - Initializes either a RedisCache or RedisClusterCache based on the cache_config. - """ - startup_nodes = cache_config.get("startup_nodes") - if not startup_nodes: - _env_cluster_nodes: Final = get_secret("REDIS_CLUSTER_NODES") - if _env_cluster_nodes is not None and isinstance(_env_cluster_nodes, str): - startup_nodes = _OPAQUE_VALUE_ADAPTER.validate_python(json.loads(_env_cluster_nodes)) - - if startup_nodes: - return RedisClusterCache(**{**cache_config, "startup_nodes": startup_nodes}) - else: - return RedisCache(**cache_config) - - def _core_update_redis_cache(self, cache: RedisCache) -> None: - """ - Update the redis cache for the router, if none set. - - Allows proxy user to just do - ```yaml - litellm_settings: - cache: true - ``` - and caching to just work. - """ - self._core_cache.attach_redis_cache(cache) - self._claude_code_session_router_cache.attach_redis_cache(cache) - self._core_prompt_caching_pin_cache.attach_redis_cache(cache) - - # Maps a routing strategy string to the attribute on `self` that holds - # the default group's strategy selector for that strategy. (The selectors - # double as `CustomLogger` callbacks, hence the legacy `*_logger` attrs.) - _DEFAULT_SELECTOR_ATTR_BY_STRATEGY: dict[str, str] = { - "least-busy": "leastbusy_logger", - "usage-based-routing": "lowesttpm_logger", - "usage-based-routing-v2": "lowesttpm_logger_v2", - "latency-based-routing": "lowestlatency_logger", - "cost-based-routing": "lowestcost_logger", - } - - @staticmethod - def _normalize_strategy( - strategy: RoutingStrategy | str | None, - ) -> str | None: - if strategy is None: - return None - if isinstance(strategy, RoutingStrategy): - return strategy.value - return strategy - - @staticmethod - def _validate_routing_strategy(routing_strategy: RoutingStrategy | str | None) -> None: - validate_routing_strategy(routing_strategy) - - def _build_strategy_selector( - self, - strategy: RoutingStrategy | str, - routing_strategy_args: dict, - register_callbacks: bool = True, - ) -> RouterStrategySelector | None: - """ - Constructs a strategy selector for a given strategy. - Returns None for `simple-shuffle` (no selector needed) and unknown - strategies. - """ - selector: RouterStrategySelector | None = None - match self._normalize_strategy(strategy): - case RoutingStrategy.LEAST_BUSY.value: - selector = LeastBusyLoggingHandler( - router_cache=self._core_cache, reservation_store=self._core_least_busy_reservations - ) - case RoutingStrategy.USAGE_BASED_ROUTING.value: - selector = LowestTPMLoggingHandler( - router_cache=self._core_cache, - routing_args=routing_strategy_args, - ) - case RoutingStrategy.USAGE_BASED_ROUTING_V2.value: - selector = LowestTPMLoggingHandler_v2( - router_cache=self._core_cache, - routing_args=routing_strategy_args, - ) - case RoutingStrategy.LATENCY_BASED.value: - selector = LowestLatencyLoggingHandler( - router_cache=self._core_cache, - routing_args=routing_strategy_args, - ) - case RoutingStrategy.COST_BASED.value: - selector = LowestCostLoggingHandler( - router_cache=self._core_cache, - routing_args={}, - ) - case _: - pass - - if selector is not None and register_callbacks: - self._register_router_selector(selector) - - return selector - - def _register_router_selector(self, selector: CustomLogger) -> None: - if isinstance(selector, LeastBusyLoggingHandler): - if isinstance(self._host.callback_list("input_callback"), list): - self._host.callbacks.add_litellm_input_callback(selector) - else: - self._host.set_callback_list("input_callback", [selector]) - if isinstance(self._host.callback_list("callbacks"), list): - self._host.callbacks.add_litellm_callback(selector) - - def _callback_selectors(self) -> tuple[RouterStrategySelector, ...]: - return tuple( - selector - for selector in ( - *(getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()), - *chain.from_iterable(selectors.values() for selectors in self._group_selectors.values()), - ) - if selector is not None - ) - - def _unregister_router_selectors(self, selectors: Sequence[object]) -> None: - """ - Drop router-owned strategy selectors from litellm's global callback - lists. Used before re-init (`routing_strategy_init` / - `_init_routing_groups`) and by `discard` so dead selectors stop - receiving callback events. `function_setup` copies `litellm.callbacks` - into the per-event lists on every request, so those are purged too. - """ - for selector in selectors: - if isinstance(selector, BaseRoutingStrategy): - selector.retire() - self._host.callbacks.drop_callbacks(selectors) - - def _apply_updated_routing_strategy_args(self) -> None: - """ - Re-link the default group's selector to the current `routing_strategy_args`. - - Selectors freeze their `RoutingArgs` at construction, so a runtime args - update would otherwise keep serving the boot-time values until restart. - Latency/usage state survives the rebuild: it lives in the shared router - cache, not on the selector. - """ - strategy: Final = self._normalize_strategy(self._core_routing_strategy) - if strategy == "lar1": - from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - - apply_lar1_routing_strategy(self._core_ctx, self.routing_strategy_args) - return - - attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "") - current: Final = getattr(self, attr, None) if attr is not None else None - if attr is None or current is None: - return - - try: - rebuilt: Final = self._build_strategy_selector( - strategy=strategy or "", - routing_strategy_args=self.routing_strategy_args, - ) - except (TypeError, ValidationError): - verbose_router_logger.exception( - "Invalid routing_strategy_args %s for '%s'; keeping the previous ones", - self.routing_strategy_args, - strategy, - ) - return - - self._unregister_router_selectors((current,)) - setattr(self, attr, rebuilt) - - def routing_strategy_init(self, routing_strategy: RoutingStrategy | str, routing_strategy_args: dict): - verbose_router_logger.info("Routing strategy: %s", routing_strategy) - self._validate_routing_strategy(routing_strategy) - self._reset_custom_routing_strategy() - - self._unregister_router_selectors( - [getattr(self, attr, None) for attr in self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.values()] - + list(getattr(self, "_override_selectors", {}).values()) - ) - self._override_selectors = {} - - self.leastbusy_logger: LeastBusyLoggingHandler | None = None - self.lowesttpm_logger: LowestTPMLoggingHandler | None = None - self.lowesttpm_logger_v2: LowestTPMLoggingHandler_v2 | None = None - self.lowestlatency_logger: LowestLatencyLoggingHandler | None = None - self.lowestcost_logger: LowestCostLoggingHandler | None = None - - selector: Final = self._build_strategy_selector( - strategy=routing_strategy, - routing_strategy_args=routing_strategy_args, - ) - attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(self._normalize_strategy(routing_strategy) or "") - # TODO: legacy `self._logger` attributes are read directly by - # `get_settings()` and external callers. Fold the default group into - # `self._group_selectors["default"]` and drop these attribute writes — - # the dual storage is an antipattern preserved only for back-compat. - if attr is not None: - setattr(self, attr, selector) - - def _init_routing_groups( - self, - groups_input: list[RoutingGroup | dict] | None, - ) -> None: - """ - Validates and indexes `routing_groups`. Each `model_name` may belong to - at most one non-priority group. Constructs per-group strategy selectors so - groups with different `routing_strategy_args` track independent state. - - Models not claimed by a non-priority group are served by the implicit - `"default"` group, whose selectors are the `self._logger` - attributes set up in `routing_strategy_init`. - """ - if not groups_input: - self._replace_routing_groups(()) - return - - known_model_names: Final = frozenset( - m["model_name"] for m in (self._core_model_list or ()) if m.get("model_name") - ) - groups: Final = parse_routing_groups(groups_input, known_model_names=known_model_names) - - alias_names: Final = frozenset(self._core_model_group_alias or ()) - for group in groups: - if group.group_name in known_model_names or group.group_name in alias_names: - if group.routing_strategy == "priority": - raise ValueError("Priority routing group names must not shadow a model or alias") - verbose_router_logger.warning( - "routing_groups: group_name '%s' is shadowed by an existing model_name or model_group_alias; " - "the group's strategy still applies to its members, but the name is not callable until renamed.", - group.group_name, - ) - - built: Final = tuple( - ( - group, - self._build_strategy_selector( - strategy="simple-shuffle" if group.routing_strategy == "priority" else group.routing_strategy, - routing_strategy_args=group.routing_strategy_args or {}, - register_callbacks=False, - ), - ) - for group in groups - ) - self._replace_routing_groups(built) - - def _replace_routing_groups( - self, - built: tuple[tuple[RoutingGroup, RouterStrategySelector | None], ...], - ) -> None: - previous_selectors: Final[Mapping[str, Mapping[str, RouterStrategySelector]]] = getattr( - self, "_group_selectors", {} - ) - self._unregister_router_selectors( - tuple(sel for selectors in previous_selectors.values() for sel in selectors.values()) - ) - for _, selector in built: - if selector is not None: - self._register_router_selector(selector) - - self._core_routing_groups: dict[str, RoutingGroup] = {group.group_name: group for group, _ in built} - self._model_to_group: dict[str, str] = { - model_name: group.group_name - for group, _ in built - if group.routing_strategy != "priority" - for model_name in group.models - } - self._group_selectors: dict[str, dict[str, RouterStrategySelector]] = { - group.group_name: ( - {} if selector is None else {self._normalize_strategy(group.routing_strategy) or "": selector} - ) - for group, selector in built - } - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - - def get_routing_group(self, model_name: str) -> RoutingGroup | None: - """ - The routing group callable as `model_name`, or None. A real deployment - `model_name` added after init shadows a same-named group (mirroring - `_try_early_resolve_deployments_for_model_not_in_names`, where concrete - models win over indirection); config-time collisions are rejected by - `_init_routing_groups`. - """ - if not self._core_routing_groups: - return None - group: Final = self._core_routing_groups.get(model_name) - if ( - group is None - or model_name in self._core_model_name_to_deployment_indices - or model_name in (self._core_model_group_alias or {}) - ): - return None - return group - - def _get_routing_group_deployments( - self, model: str, team_id: str | None = None - ) -> list[DeploymentTypedDict] | None: # mutable-ok: list matches _get_all_deployments' contract for callers - """ - The union of member deployments for a routing group called as `model`, - or None when `model` is not a callable group. The requested name stays - the group name so strategy selectors key their state by it. - - `_common_checks_available_deployment` consults this BEFORE its - early-resolve step so a wildcard `default_deployment` or pattern route - cannot hijack a group call. Overall resolution precedence there: - specific deployment > model id > model_group_alias > routing group > - model_name > team/pattern/default fallbacks. - """ - if not self._core_routing_groups: - return None - routing_group: Final = self.get_routing_group(model) - if routing_group is None: - return None - return [ - apply_routing_group_priority(routing_group, member, deployment) - for member in routing_group.models - for deployment in self._core_get_all_deployments(model_name=member, team_id=team_id) - ] - - def _is_priority_routing_group(self, model: str) -> bool: - resolved: Final = self._core_get_model_from_alias(model=model) or model - group: Final = self.get_routing_group(resolved) - return group is not None and group.routing_strategy == "priority" - - def _core_is_recognized_model(self, model: str) -> bool: - """ - Whether `model` names something this router serves directly: a - deployment model_name, a deployment id, a `model_group_alias`, or a - callable routing group. Proxy request gates share this predicate so a - new virtual-model kind cannot be forgotten at one of them; wildcard, - default-deployment, and deployment-name fallbacks stay caller policy. - """ - return ( - model in self._core_model_names - or self._core_has_model_id(model) - or (self._core_model_group_alias is not None and model in self._core_model_group_alias) - or self.get_routing_group(model) is not None - ) - - def _core_routing_group_has_alternatives(self, model_group: str | None) -> bool: - """ - True when `model_group` names a callable routing group whose member - union spans more than one deployment. Cooldown handling passes the - FAILING REQUEST's model group here: a 429 on a group call cools the - member down so selection moves to the group's alternatives, while a - direct call to a single-deployment member keeps the - single-deployment-model-group cooldown exemption. - """ - if model_group is None: - return False - resolved: Final = self._core_get_model_from_alias(model=model_group) or model_group - group: Final = self.get_routing_group(resolved) - if group is None: - return False - return sum(len(self._core_model_name_to_deployment_indices.get(member) or ()) for member in group.models) > 1 - - def _core_team_model_has_alternatives(self, deployment_id: str) -> bool: - deployment: Final = self._core_get_deployment(model_id=deployment_id) - if deployment is None: - return False - team_id: Final = deployment.model_info.team_id - public_model_name: Final = deployment.model_info.team_public_model_name - if team_id is None or public_model_name is None: - return False - siblings: Final = self._team_deployments_named(team_id, public_model_name) or () - routable_siblings: Final = self._filter_blocked_deployments([dict(sibling) for sibling in siblings]) - return len(routable_siblings) > 1 - - _OVERRIDABLE_ROUTING_STRATEGIES: frozenset[str] = frozenset({"simple-shuffle", *_DEFAULT_SELECTOR_ATTR_BY_STRATEGY}) - - def _get_request_routing_strategy_override(self, request_kwargs: dict | None) -> str | None: - """ - Reads a per-request `routing_strategy` override (forwarded by the proxy - from key/team `router_settings`) out of the request kwargs. - - Only strategies with a per-request-capable selector are honored; - anything else (unknown strings, `lar1`, `provider-budget-routing`) is - ignored with a warning so a bad value stored on a key or team can - never take down that caller's traffic. - """ - if not request_kwargs: - return None - raw_strategy: Final = request_kwargs.get("routing_strategy") - if raw_strategy is None: - return None - strategy = self._normalize_strategy(raw_strategy) if isinstance(raw_strategy, (str, RoutingStrategy)) else None - if not isinstance(strategy, str) or strategy not in self._OVERRIDABLE_ROUTING_STRATEGIES: - verbose_router_logger.warning( - "Ignoring per-request routing_strategy override '%s'; supported overrides: %s.", - raw_strategy, - sorted(self._OVERRIDABLE_ROUTING_STRATEGIES), - ) - return None - return strategy - - def _get_override_strategy_selector(self, strategy: str) -> RouterStrategySelector | None: - """ - Returns the selector for a per-request strategy override. - - Reuses the default group's selector when the override matches the - router's configured strategy (so shared state keeps accumulating in - one place); otherwise lazily builds one selector per strategy and - caches it for the router's lifetime so its usage/latency state - persists across requests. - """ - if strategy == self._normalize_strategy(self._core_routing_strategy): - attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy) - return getattr(self, attr, None) if attr is not None else None - with self._override_selectors_lock: - if strategy not in self._override_selectors: - self._override_selectors[strategy] = self._build_strategy_selector( - strategy=strategy, - routing_strategy_args={}, - register_callbacks=False, - ) - return self._override_selectors[strategy] - - def _override_selector_pre_call_check( - self, strategy: str | None, selector: RouterStrategySelector | None, deployment: dict - ) -> None: - """ - Override selectors are not in `litellm.callbacks`, so the pre-call check that - `routing_strategy_pre_call_checks` runs for the router's own selectors (rpm - accounting for `usage-based-routing-v2`) runs here, for the overriding request only. - """ - if selector is None or strategy is None or selector is not self._override_selectors.get(strategy): - return - selector.pre_call_check(deployment) - - async def _async_override_selector_pre_call_check( - self, - strategy: str | None, - selector: RouterStrategySelector | None, - deployment: dict, - parent_otel_span: Span | None, - ) -> None: - if selector is None or strategy is None or selector is not self._override_selectors.get(strategy): - return - await selector.async_pre_call_check(deployment, parent_otel_span) - - def _bind_override_selector_to_request( - self, strategy: str, selector: RouterStrategySelector | None, request_kwargs: Mapping[str, object] | None - ) -> None: - if selector is None or request_kwargs is None or strategy in self._globally_registered_strategies(): - return - logging_obj: Final = request_kwargs.get("litellm_logging_obj") - if isinstance(logging_obj, LiteLLMLogging): - logging_obj.add_dynamic_callback(selector) - - def _globally_registered_strategies(self) -> frozenset[str]: - configured: Final = ( - self._core_routing_strategy, - *(group.routing_strategy for group in self._core_routing_groups.values()), - ) - return frozenset( - normalized for normalized in map(self._normalize_strategy, configured) if normalized is not None - ) - - def _core_arm_routing_read_prefetch(self, model: str, request_kwargs: dict[str, object] | None = None) -> None: - """Declare the cooldown read (and, for usage-based routing, the usage read) that - `async_get_available_deployment` will make for `model` on the request's Redis batch, so admission's - flush carries it. A miss (alias, no batch) costs nothing: routing then reads as it always has.""" - try: - strategy, selector = self._get_routing_context(model, request_kwargs) - usage_selector: Final = ( - selector - if strategy == "usage-based-routing-v2" and isinstance(selector, LowestTPMLoggingHandler_v2) - else None - ) - deployments: Final = self._core_get_model_list(model_name=model) - if deployments: - RoutingPrefetch.arm(self._core_ctx, usage_selector, deployments) - except Exception as e: # noqa: BLE001 # a prefetch is an optimisation, never a reason to fail the request - verbose_router_logger.debug( - "routing read prefetch not armed for %s: %s", _without_line_breaks(model), _without_line_breaks(e) - ) - - def _get_routing_context( - self, model: str, request_kwargs: dict | None = None - ) -> tuple[str | None, RouterStrategySelector | None]: - """ - Resolves the routing strategy and selector to use for the given model. - - A per-request `routing_strategy` in `request_kwargs` (forwarded by the - proxy from key/team `router_settings`) takes precedence over both the - model's routing group and the router's top-level strategy, since it is - the most specific expression of caller intent. - - Otherwise every model belongs to exactly one group: an explicit entry - from `routing_groups` (either because `model` IS a callable group name, - or because it is a member of one), or the implicit `"default"` group - driven by the router's top-level `routing_strategy` / - `routing_strategy_args`. - - `self._core_routing_strategy` may be either a string or a `RoutingStrategy` - enum member (the constructor accepts both), so it is normalized to a - string here. Downstream call sites and `_select_deployment_*` arms - compare against string literals. - """ - override: Final = self._get_request_routing_strategy_override(request_kwargs) - if override is not None: - verbose_router_logger.debug("routing_group=request-override model=%s strategy=%s", model, override) - override_selector: Final = self._get_override_strategy_selector(override) - self._bind_override_selector_to_request(override, override_selector, request_kwargs) - return override, override_selector - - resolved_model: Final = self._core_get_model_from_alias(model=model) or model - group_name: Final = ( - resolved_model - if self.get_routing_group(resolved_model) is not None - else self._model_to_group.get(resolved_model) - ) - if group_name is None: - strategy = self._normalize_strategy(self._core_routing_strategy) - attr: Final = self._DEFAULT_SELECTOR_ATTR_BY_STRATEGY.get(strategy or "") - selector = getattr(self, attr, None) if attr is not None else None - verbose_router_logger.debug("routing_group=default model=%s strategy=%s", model, strategy) - return strategy, selector - - group: Final = self._core_routing_groups[group_name] - if group.routing_strategy == "priority": - return "simple-shuffle", None - strategy = self._normalize_strategy(group.routing_strategy) - selector = self._group_selectors.get(group_name, {}).get(strategy or "") - verbose_router_logger.debug("routing_group=%s model=%s strategy=%s", group_name, model, strategy) - return strategy, selector - - async def _select_deployment_async( - self, - *, - strategy: str | None, - selector: Any | None, - model: str, - healthy_deployments: list, - messages: list[dict[str, str]] | None, - input: str | list | None, - request_kwargs: dict | None, - reserve: bool, - ) -> Any | None: - """ - Asks the strategy selector for a deployment. Caller handles - `simple-shuffle` separately (it does not flow through a selector). - Returns None for unknown strategies or when the selector is missing. - """ - if selector is None: - return None - match strategy: - case "least-busy": - return await selector.async_get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - reserve=reserve, - call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), - ) - case "usage-based-routing": - # `LowestTPMLoggingHandler` (v1) only exposes the sync - # `get_available_deployments`. Mirror the pre-routing-groups - # top-level fallback by calling it inline so groups using v1 - # still work from async callers. - return selector.get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - ) - case "usage-based-routing-v2" | "cost-based-routing": - return await selector.async_get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - ) - case "latency-based-routing": - return await selector.async_get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - ) - case _: - return None - - def _select_deployment_sync( - self, - *, - strategy: str | None, - selector: Any | None, - model: str, - healthy_deployments: list, - messages: list[dict[str, str]] | None, - input: str | list | None, - request_kwargs: dict | None, - reserve: bool, - ) -> Any | None: - """ - Sync sibling of `_select_deployment_async`. Caller handles - `simple-shuffle` separately. - """ - if selector is None: - return None - - # `cost-based-routing` is intentionally omitted — - # `LowestCostLoggingHandler` only implements - # `async_get_available_deployments` - match strategy: - case "least-busy": - return selector.get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - reserve=reserve, - call_id=call_id_of(CALL_IDENTITY.validate_python(request_kwargs or {})), - ) - case "usage-based-routing" | "usage-based-routing-v2": - return selector.get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - ) - case "latency-based-routing": - return selector.get_available_deployments( - model_group=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - ) - case _: - return None - - def initialize_assistants_endpoint(self): - ## INITIALIZE PASS THROUGH ASSISTANTS ENDPOINT ## - self.acreate_assistants = self.factory_function(self._host.invoke("acreate_assistants")) - self.adelete_assistant = self.factory_function(self._host.invoke("adelete_assistant")) - self.aget_assistants = self.factory_function(self._host.invoke("aget_assistants")) - self.acreate_thread = self.factory_function(self._host.invoke("acreate_thread")) - self.aget_thread = self.factory_function(self._host.invoke("aget_thread")) - self.a_add_message = self.factory_function(self._host.invoke("a_add_message")) - self.aget_messages = self.factory_function(self._host.invoke("aget_messages")) - self.arun_thread = self.factory_function(self._host.invoke("arun_thread")) - - def _initialize_core_endpoints(self): - """Helper to initialize core router endpoints.""" - self.amoderation = self.factory_function(self._host.invoke("amoderation"), call_type="moderation") - self.aanthropic_messages = self.factory_function( - self._host.invoke("anthropic_messages"), call_type="anthropic_messages" - ) - self.anthropic_messages = self.factory_function( - self._host.invoke("anthropic_messages"), call_type="anthropic_messages" - ) - self.agenerate_content = self.factory_function( - self._host.invoke("agenerate_content"), call_type="agenerate_content" - ) - self.aadapter_generate_content = self.factory_function( - self._host.invoke("aadapter_generate_content"), call_type="aadapter_generate_content" - ) - self.aresponses = self.factory_function(self._host.invoke("aresponses"), call_type="aresponses") - self.afile_delete = self.factory_function(self._host.invoke("afile_delete"), call_type="afile_delete") - self.afile_content = self.factory_function(self._host.invoke("afile_content"), call_type="afile_content") - self.responses = self.factory_function(self._host.invoke("responses"), call_type="responses") - self.aget_responses = self.factory_function(self._host.invoke("aget_responses"), call_type="aget_responses") - self.acancel_responses = self.factory_function( - self._host.invoke("acancel_responses"), call_type="acancel_responses" - ) - self.acompact_responses = self.factory_function( - self._host.invoke("acompact_responses"), call_type="acompact_responses" - ) - self.adelete_responses = self.factory_function( - self._host.invoke("adelete_responses"), call_type="adelete_responses" - ) - self.alist_input_items = self.factory_function( - self._host.invoke("alist_input_items"), call_type="alist_input_items" - ) - self._arealtime = self.factory_function(self._host.invoke("_arealtime"), call_type="_arealtime") - self.acreate_realtime_client_secret = self.factory_function( - self._host.invoke("acreate_realtime_client_secret"), call_type="acreate_realtime_client_secret" - ) - self.arealtime_calls = self.factory_function(self._host.invoke("arealtime_calls"), call_type="arealtime_calls") - self.acreate_realtime_transcription_session = self.factory_function( - self._host.invoke("acreate_realtime_transcription_session"), - call_type="acreate_realtime_transcription_session", - ) - self.acreate_realtime_translation_client_secret = self.factory_function( - self._host.invoke("acreate_realtime_translation_client_secret"), - call_type="acreate_realtime_translation_client_secret", - ) - self.arealtime_translation_calls = self.factory_function( - self._host.invoke("arealtime_translation_calls"), - call_type="arealtime_translation_calls", - ) - self._aresponses_websocket = self.factory_function( - self._host.invoke("_aresponses_websocket"), call_type="_aresponses_websocket" - ) - self.acreate_fine_tuning_job = self.factory_function( - self._host.invoke("acreate_fine_tuning_job"), call_type="acreate_fine_tuning_job" - ) - self.acancel_fine_tuning_job = self.factory_function( - self._host.invoke("acancel_fine_tuning_job"), call_type="acancel_fine_tuning_job" - ) - self.alist_fine_tuning_jobs = self.factory_function( - self._host.invoke("alist_fine_tuning_jobs"), call_type="alist_fine_tuning_jobs" - ) - self.aretrieve_fine_tuning_job = self.factory_function( - self._host.invoke("aretrieve_fine_tuning_job"), call_type="aretrieve_fine_tuning_job" - ) - self.afile_list = self.factory_function(self._host.invoke("afile_list"), call_type="alist_files") - self.aimage_edit = self.factory_function(self._host.invoke("aimage_edit"), call_type="aimage_edit") - self.allm_passthrough_route = self.factory_function( - self._host.invoke("allm_passthrough_route"), call_type="allm_passthrough_route" - ) - # Note: acancel_batch is defined as a method on the Router class (not using factory_function) - # to properly handle model-to-provider mapping like acreate_batch and aretrieve_batch - - def _initialize_vector_store_endpoints(self): - """Initialize vector store endpoints.""" - from litellm.vector_stores.main import ( - adelete, - alist, - aretrieve, - asearch, - aupdate, - create, - delete, - list, - retrieve, - search, - update, - ) - - self.avector_store_search = self.factory_function(asearch, call_type="avector_store_search") - self.vector_store_search = self.factory_function(search, call_type="vector_store_search") - self.vector_store_create = self.factory_function(create, call_type="vector_store_create") - self.avector_store_retrieve = self.factory_function(aretrieve, call_type="avector_store_retrieve") - self.vector_store_retrieve = self.factory_function(retrieve, call_type="vector_store_retrieve") - self.avector_store_list = self.factory_function(alist, call_type="avector_store_list") - self.vector_store_list = self.factory_function(list, call_type="vector_store_list") - self.avector_store_update = self.factory_function(aupdate, call_type="avector_store_update") - self.vector_store_update = self.factory_function(update, call_type="vector_store_update") - self.avector_store_delete = self.factory_function(adelete, call_type="avector_store_delete") - self.vector_store_delete = self.factory_function(delete, call_type="vector_store_delete") - - def _initialize_vector_store_file_endpoints(self): - """Initialize vector store file endpoints.""" - from litellm.vector_store_files.main import ( - acreate as avector_store_file_create_fn, - ) - from litellm.vector_store_files.main import ( - adelete as avector_store_file_delete_fn, - ) - from litellm.vector_store_files.main import alist as avector_store_file_list_fn - from litellm.vector_store_files.main import ( - aretrieve as avector_store_file_retrieve_fn, - ) - from litellm.vector_store_files.main import ( - aretrieve_content as avector_store_file_content_fn, - ) - from litellm.vector_store_files.main import ( - aupdate as avector_store_file_update_fn, - ) - from litellm.vector_store_files.main import ( - create as vector_store_file_create_fn, - ) - from litellm.vector_store_files.main import ( - delete as vector_store_file_delete_fn, - ) - from litellm.vector_store_files.main import list as vector_store_file_list_fn - from litellm.vector_store_files.main import ( - retrieve as vector_store_file_retrieve_fn, - ) - from litellm.vector_store_files.main import ( - retrieve_content as vector_store_file_content_fn, - ) - from litellm.vector_store_files.main import ( - update as vector_store_file_update_fn, - ) - - self.avector_store_file_create = self.factory_function( - avector_store_file_create_fn, call_type="avector_store_file_create" - ) - self.vector_store_file_create = self.factory_function( - vector_store_file_create_fn, call_type="vector_store_file_create" - ) - self.avector_store_file_list = self.factory_function( - avector_store_file_list_fn, call_type="avector_store_file_list" - ) - self.vector_store_file_list = self.factory_function( - vector_store_file_list_fn, call_type="vector_store_file_list" - ) - self.avector_store_file_retrieve = self.factory_function( - avector_store_file_retrieve_fn, call_type="avector_store_file_retrieve" - ) - self.vector_store_file_retrieve = self.factory_function( - vector_store_file_retrieve_fn, call_type="vector_store_file_retrieve" - ) - self.avector_store_file_content = self.factory_function( - avector_store_file_content_fn, call_type="avector_store_file_content" - ) - self.vector_store_file_content = self.factory_function( - vector_store_file_content_fn, call_type="vector_store_file_content" - ) - self.avector_store_file_update = self.factory_function( - avector_store_file_update_fn, call_type="avector_store_file_update" - ) - self.vector_store_file_update = self.factory_function( - vector_store_file_update_fn, call_type="vector_store_file_update" - ) - self.avector_store_file_delete = self.factory_function( - avector_store_file_delete_fn, call_type="avector_store_file_delete" - ) - self.vector_store_file_delete = self.factory_function( - vector_store_file_delete_fn, call_type="vector_store_file_delete" - ) - - def _initialize_google_genai_endpoints(self): - """Initialize Google GenAI endpoints.""" - from litellm.google_genai import ( - agenerate_content, - agenerate_content_stream, - generate_content, - generate_content_stream, - ) - - self.agenerate_content = self.factory_function(agenerate_content, call_type="agenerate_content") - self.generate_content = self.factory_function(generate_content, call_type="generate_content") - self.agenerate_content_stream = self.factory_function( - agenerate_content_stream, call_type="agenerate_content_stream" - ) - self.generate_content_stream = self.factory_function( - generate_content_stream, call_type="generate_content_stream" - ) - - def _initialize_ocr_search_endpoints(self): - """Initialize OCR and search endpoints.""" - from litellm.ocr import aocr, ocr - - self.aocr = self.factory_function(aocr, call_type="aocr") - self.ocr = self.factory_function(ocr, call_type="ocr") - - from litellm.search import asearch, search - - self.asearch = self.factory_function(asearch, call_type="asearch") - self.search = self.factory_function(search, call_type="search") - - from litellm.decisions import adecisions, decisions - - self.adecisions = self.factory_function(adecisions, call_type="adecisions") - self.decisions = self.factory_function(decisions, call_type="decisions") - - from litellm.systemone import asystemone, systemone - - self.asystemone = self.factory_function(asystemone, call_type="asystemone") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - self.systemone = self.factory_function(systemone, call_type="systemone") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - - from litellm.extraction import aextraction, extraction - - self.aextraction = self.factory_function(aextraction, call_type="aextraction") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - self.extraction = self.factory_function(extraction, call_type="extraction") # pyright: ignore[reportUnknownMemberType] # factory_function takes a bare Callable - - def _initialize_video_endpoints(self): - """Initialize video endpoints.""" - from litellm.videos import ( - avideo_content, - avideo_create_character, - avideo_edit, - avideo_extension, - avideo_generation, - avideo_get_character, - avideo_list, - avideo_remix, - avideo_status, - video_content, - video_create_character, - video_edit, - video_extension, - video_generation, - video_get_character, - video_list, - video_remix, - video_status, - ) - - self.avideo_generation = self.factory_function(avideo_generation, call_type="avideo_generation") - self.video_generation = self.factory_function(video_generation, call_type="video_generation") - self.avideo_list = self.factory_function(avideo_list, call_type="avideo_list") - self.video_list = self.factory_function(video_list, call_type="video_list") - self.avideo_status = self.factory_function(avideo_status, call_type="avideo_status") - self.video_status = self.factory_function(video_status, call_type="video_status") - self.avideo_content = self.factory_function(avideo_content, call_type="avideo_content") - self.video_content = self.factory_function(video_content, call_type="video_content") - self.avideo_remix = self.factory_function(avideo_remix, call_type="avideo_remix") - self.video_remix = self.factory_function(video_remix, call_type="video_remix") - self.avideo_create_character = self.factory_function( - avideo_create_character, call_type="avideo_create_character" - ) - self.video_create_character = self.factory_function(video_create_character, call_type="video_create_character") - self.avideo_get_character = self.factory_function(avideo_get_character, call_type="avideo_get_character") - self.video_get_character = self.factory_function(video_get_character, call_type="video_get_character") - self.avideo_edit = self.factory_function(avideo_edit, call_type="avideo_edit") - self.video_edit = self.factory_function(video_edit, call_type="video_edit") - self.avideo_extension = self.factory_function(avideo_extension, call_type="avideo_extension") - self.video_extension = self.factory_function(video_extension, call_type="video_extension") - - def _initialize_container_endpoints(self): - """Initialize container endpoints.""" - from litellm.containers import ( - acreate_container, - adelete_container, - alist_containers, - aretrieve_container, - create_container, - delete_container, - list_containers, - retrieve_container, - ) - from litellm.containers.endpoint_factory import ( - _generated_endpoints as container_file_endpoints, - ) - - self.acreate_container = self.factory_function(acreate_container, call_type="acreate_container") - self.create_container = self.factory_function(create_container, call_type="create_container") - self.alist_containers = self.factory_function(alist_containers, call_type="alist_containers") - self.list_containers = self.factory_function(list_containers, call_type="list_containers") - self.aretrieve_container = self.factory_function(aretrieve_container, call_type="aretrieve_container") - self.retrieve_container = self.factory_function(retrieve_container, call_type="retrieve_container") - self.adelete_container = self.factory_function(adelete_container, call_type="adelete_container") - self.delete_container = self.factory_function(delete_container, call_type="delete_container") - - # Auto-register JSON-generated container file endpoints - for name, func in container_file_endpoints.items(): - setattr(self, name, self.factory_function(func, call_type=name)) - - def _initialize_skills_endpoints(self): - """Initialize Anthropic Skills API endpoints.""" - self.acreate_skill = self.factory_function(self._host.invoke("acreate_skill"), call_type="acreate_skill") - self.alist_skills = self.factory_function(self._host.invoke("alist_skills"), call_type="alist_skills") - self.aget_skill = self.factory_function(self._host.invoke("aget_skill"), call_type="aget_skill") - self.adelete_skill = self.factory_function(self._host.invoke("adelete_skill"), call_type="adelete_skill") - - def _initialize_interactions_endpoints(self): - """Initialize Google Interactions API endpoints.""" - from litellm.interactions import acancel as acancel_interaction - from litellm.interactions import acreate as acreate_interaction - from litellm.interactions import adelete as adelete_interaction - from litellm.interactions import aget as aget_interaction - from litellm.interactions import cancel as cancel_interaction - from litellm.interactions import create as create_interaction - from litellm.interactions import delete as delete_interaction - from litellm.interactions import get as get_interaction - - self.acreate_interaction = self.factory_function(acreate_interaction, call_type="acreate_interaction") - self.create_interaction = self.factory_function(create_interaction, call_type="create_interaction") - self.aget_interaction = self.factory_function(aget_interaction, call_type="aget_interaction") - self.get_interaction = self.factory_function(get_interaction, call_type="get_interaction") - self.adelete_interaction = self.factory_function(adelete_interaction, call_type="adelete_interaction") - self.delete_interaction = self.factory_function(delete_interaction, call_type="delete_interaction") - self.acancel_interaction = self.factory_function(acancel_interaction, call_type="acancel_interaction") - self.cancel_interaction = self.factory_function(cancel_interaction, call_type="cancel_interaction") - - def _initialize_managed_agents_endpoints(self): - """Initialize Google Managed Agents API endpoints (v1beta/agents).""" - from litellm.interactions.agents import acreate as acreate_agent - from litellm.interactions.agents import adelete as adelete_agent - from litellm.interactions.agents import aget as aget_agent - from litellm.interactions.agents import alist as alist_agents - from litellm.interactions.agents import alist_versions as alist_agent_versions - from litellm.interactions.agents import create as create_agent - from litellm.interactions.agents import delete as delete_agent - from litellm.interactions.agents import get as get_agent - from litellm.interactions.agents import list as list_agents - from litellm.interactions.agents import list_versions as list_agent_versions - - self.acreate_agent = self.factory_function(acreate_agent, call_type="acreate_agent") - self.create_agent = self.factory_function(create_agent, call_type="create_agent") - self.alist_agents = self.factory_function(alist_agents, call_type="alist_agents") - self.list_agents = self.factory_function(list_agents, call_type="list_agents") - self.aget_agent = self.factory_function(aget_agent, call_type="aget_agent") - self.get_agent = self.factory_function(get_agent, call_type="get_agent") - self.adelete_agent = self.factory_function(adelete_agent, call_type="adelete_agent") - self.delete_agent = self.factory_function(delete_agent, call_type="delete_agent") - self.alist_agent_versions = self.factory_function(alist_agent_versions, call_type="alist_agent_versions") - self.list_agent_versions = self.factory_function(list_agent_versions, call_type="list_agent_versions") - - def _initialize_specialized_endpoints(self): - """Helper to initialize specialized router endpoints (vector store, OCR, search, video, container, skills, interactions).""" - self._initialize_vector_store_endpoints() - self._initialize_vector_store_file_endpoints() - self._initialize_google_genai_endpoints() - self._initialize_ocr_search_endpoints() - # Override vector store methods with router-aware implementations - self._override_vector_store_methods_for_router() - self._initialize_video_endpoints() - self._initialize_container_endpoints() - self._initialize_skills_endpoints() - self._initialize_interactions_endpoints() - self._initialize_managed_agents_endpoints() - - def initialize_router_endpoints(self): - self._initialize_core_endpoints() - self._initialize_specialized_endpoints() - - def validate_fallbacks(self, fallback_param: Sequence[Mapping[str, object]] | None): - """ - Validate the fallbacks parameter. - """ - if fallback_param is None: - return - for fallback_dict in fallback_param: - if not isinstance(fallback_dict, dict): - raise ValueError(f"Item '{fallback_dict}' is not a dictionary.") - if len(fallback_dict) != 1: - raise ValueError( - f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys." - ) - ((model_group, targets),) = fallback_dict.items() - if isinstance(targets, list) or check_non_standard_fallback_format(fallbacks=[fallback_dict]): - continue - raise TypeError( - f"Fallbacks for '{model_group}' must be a list of model names, got '{targets}'. " - f"Use {{'{model_group}': ['{targets}']}}" - ) - - def _add_encrypted_content_affinity_check(self, enable_global_affinity: bool) -> None: - def _move_before_deployment_affinity( - callback_list: list[_CallbackT], - callback_to_move: _CallbackT, - ) -> None: - if callback_to_move not in callback_list: - return - callback_list.remove(callback_to_move) - insert_index: Final = next( - (idx for idx, callback in enumerate(callback_list) if isinstance(callback, DeploymentAffinityCheck)), - len(callback_list), - ) - callback_list.insert(insert_index, callback_to_move) - - if enable_global_affinity or EncryptedContentAffinityCheck.has_model_group_affinity_enabled( - self.model_group_affinity_config - ): - if self.optional_callbacks is None: - self.optional_callbacks = [] - - existing_ec_callback: EncryptedContentAffinityCheck | None = None - for cb in self.optional_callbacks: - if isinstance(cb, EncryptedContentAffinityCheck): - existing_ec_callback = cb - break - - if existing_ec_callback is not None: - existing_ec_callback.ctx = self._core_ctx - existing_ec_callback.enable_global_affinity = ( - existing_ec_callback.enable_global_affinity or enable_global_affinity - ) - existing_ec_callback.model_group_affinity_config = self.model_group_affinity_config or {} - ec_callback = existing_ec_callback - else: - ec_callback = EncryptedContentAffinityCheck( - ctx=self._core_ctx, - enable_global_affinity=enable_global_affinity, - model_group_affinity_config=self.model_group_affinity_config, - ) - self.optional_callbacks.append(ec_callback) - self._host.callbacks.add_litellm_callback(ec_callback) - - _move_before_deployment_affinity(self.optional_callbacks, ec_callback) - _move_before_deployment_affinity(self._host.callback_list("callbacks"), ec_callback) - - def _ensure_deployment_affinity_callback(self) -> None: - """Register the DeploymentAffinityCheck callback (global flags all False) if absent. - - Needed when nothing enabled a global affinity flag but affinity can still - activate per request: per-group `model_group_affinity_config` entries, or the - session-affinity marker a complexity router stamps at pre-routing time. - """ - if any(isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or [])): - return - if self.optional_callbacks is None: - self.optional_callbacks = [] - affinity_callback: Final = DeploymentAffinityCheck( - cache=self._core_cache, - ttl_seconds=self.deployment_affinity_ttl_seconds, - enable_user_key_affinity=False, - enable_responses_api_affinity=False, - enable_session_id_affinity=False, - model_group_affinity_config=self.model_group_affinity_config, - is_priority_group=self._is_priority_routing_group, - session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, - ) - self.optional_callbacks.append(affinity_callback) - self._host.callbacks.add_litellm_callback(affinity_callback) - - def add_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None): - if optional_pre_call_checks is None: - return - - # --------------------------------------------------------------------- - # Unified deployment affinity (session stickiness) - # --------------------------------------------------------------------- - enable_user_key_affinity: Final = "deployment_affinity" in optional_pre_call_checks - enable_responses_api_affinity: Final = "responses_api_deployment_check" in optional_pre_call_checks - enable_session_id_affinity: Final = "session_affinity" in optional_pre_call_checks - if enable_user_key_affinity or enable_responses_api_affinity or enable_session_id_affinity: - if self.optional_callbacks is None: - self.optional_callbacks = [] - - existing_affinity_callback: DeploymentAffinityCheck | None = None - for cb in self.optional_callbacks: - if isinstance(cb, DeploymentAffinityCheck): - existing_affinity_callback = cb - break - - if existing_affinity_callback is not None: - existing_affinity_callback.enable_user_key_affinity = ( - existing_affinity_callback.enable_user_key_affinity or enable_user_key_affinity - ) - existing_affinity_callback.enable_responses_api_affinity = ( - existing_affinity_callback.enable_responses_api_affinity or enable_responses_api_affinity - ) - existing_affinity_callback.enable_session_id_affinity = ( - existing_affinity_callback.enable_session_id_affinity or enable_session_id_affinity - ) - existing_affinity_callback.ttl_seconds = self.deployment_affinity_ttl_seconds - if self.model_group_affinity_config: - existing_affinity_callback.model_group_affinity_config = self.model_group_affinity_config - else: - affinity_callback: Final = DeploymentAffinityCheck( - cache=self._core_cache, - ttl_seconds=self.deployment_affinity_ttl_seconds, - enable_user_key_affinity=enable_user_key_affinity, - enable_responses_api_affinity=enable_responses_api_affinity, - enable_session_id_affinity=enable_session_id_affinity, - model_group_affinity_config=self.model_group_affinity_config, - is_priority_group=self._is_priority_routing_group, - session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, - ) - self.optional_callbacks.append(affinity_callback) - self._host.callbacks.add_litellm_callback(affinity_callback) - - # --------------------------------------------------------------------- - # Encrypted content affinity - # --------------------------------------------------------------------- - self._add_encrypted_content_affinity_check( - enable_global_affinity=("encrypted_content_affinity" in optional_pre_call_checks) - ) - - # --------------------------------------------------------------------- - # Remaining optional pre-call checks - # --------------------------------------------------------------------- - for pre_call_check in optional_pre_call_checks: - _callback: CustomLogger | None = None - if pre_call_check in ( - "deployment_affinity", - "responses_api_deployment_check", - "session_affinity", - "encrypted_content_affinity", - ): - continue - if pre_call_check == "prompt_caching": - _callback = PromptCachingDeploymentCheck( - cache=self._core_prompt_caching_pin_cache, is_priority_group=self._is_priority_routing_group - ) - elif pre_call_check == "router_budget_limiting": - if self._core_get_router_deployment_budget_limiter() is not None: - continue - _callback = RouterBudgetLimiting( - dual_cache=self._core_cache, - provider_budget_config=self._core_provider_budget_config, - model_list=self._core_model_list, - ) - self.router_budget_logger = _callback - elif pre_call_check == "enforce_model_rate_limits": - _callback = ModelRateLimitingCheck(dual_cache=self._core_cache) - - if _callback is None: - continue - - if self.optional_callbacks is not None and any( - isinstance(callback, type(_callback)) for callback in self.optional_callbacks - ): - continue - if self.optional_callbacks is None: - self.optional_callbacks = [] - self.optional_callbacks.append(_callback) - self._host.callbacks.add_litellm_callback(_callback) - - def set_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None) -> None: - if optional_pre_call_checks is None: - return - requested: Final = frozenset(optional_pre_call_checks) - for name, callback_cls in _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS.items(): - if name not in requested: - self._remove_optional_callbacks_of_type(callback_cls) - self.add_optional_pre_call_checks(optional_pre_call_checks) - - def _remove_optional_callbacks_of_type(self, callback_cls: type[CustomLogger]) -> None: - 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] - if any( - router is not self and any(type(cb) is callback_cls for cb in (router.optional_callbacks or [])) - for router in (*_live_routers, *(owner for owner in live_check_owners if not isinstance(owner, PythonCore))) - ): - return - for cb in tuple(self._host.callback_list("callbacks")): - if type(cb) is callback_cls: - self._host.callbacks.remove_callback_from_list_by_object( - self._host.callback_list("callbacks"), cb, require_self=False - ) - - def print_deployment(self, deployment: dict): - """ - returns a copy of the deployment with the api key masked - - Only returns 2 characters of the api key and masks the rest with * (10 *). - """ - try: - _deployment_copy: Final = copy.deepcopy(deployment) - litellm_params: Final[dict] = _deployment_copy["litellm_params"] - - if self._host.dynamic().redact_user_api_key_info: - masker: Final = SensitiveDataMasker(visible_prefix=2, visible_suffix=0) - _deployment_copy["litellm_params"] = masker.mask_dict(litellm_params) - elif "api_key" in litellm_params: - litellm_params["api_key"] = litellm_params["api_key"][:2] + "*" * 10 - - return _deployment_copy - except Exception as e: - verbose_router_logger.debug("Error occurred while printing deployment - %s", e) - raise e - - @staticmethod - def _deployment_params_with_request_reasoning_override( - deployment_params: Mapping[str, object], request_kwargs: Mapping[str, object] - ) -> dict[str, object]: # mutable-ok: litellm's request pipeline consumes a mutable kwargs mapping - """Return deployment params whose equivalent effort controls cannot outrank a request override. - - Providers expose the same setting through several native carriers. A request-level - ``reasoning_effort`` is the portable override, so a deployment's ``thinking`` or nested - ``*.effort`` must not remain beside it and either win or trigger a conflicting-params 400. - Every changed mapping is copied so the Router's shared deployment config stays immutable. - """ - sanitized: Final = dict(deployment_params) - if request_kwargs.get("reasoning_effort") is None: - return sanitized - - sanitized.pop("thinking", None) - PythonCore._pop_effort_from_nested_carrier(sanitized, "output_config") - PythonCore._pop_effort_from_nested_carrier(sanitized, "reasoning") - - extra_body: Final = sanitized.get("extra_body") - if isinstance(extra_body, Mapping): - sanitized_extra_body: Final = dict(extra_body) - sanitized_extra_body.pop("reasoning_effort", None) - sanitized_extra_body.pop("thinking", None) - PythonCore._pop_effort_from_nested_carrier(sanitized_extra_body, "output_config") - PythonCore._pop_effort_from_nested_carrier(sanitized_extra_body, "reasoning") - if sanitized_extra_body: - sanitized["extra_body"] = sanitized_extra_body - else: - sanitized.pop("extra_body", None) - return sanitized - - @staticmethod - def _is_classifier_internal_call(kwargs: Mapping[str, object]) -> bool: - metadata: Final = kwargs.get("metadata") - litellm_metadata: Final = kwargs.get("litellm_metadata") - return any( - isinstance(candidate, Mapping) - and candidate.get(INTERNAL_CALL_ORIGIN_METADATA_KEY) == AUTOROUTER_CLASSIFIER_CALL_ORIGIN - for candidate in (metadata, litellm_metadata) - ) - - def _drop_unsupported_classifier_reasoning_effort( - self, - deployment: DeploymentTypedDict, - model: str, - kwargs: dict[str, object], # mutable-ok: fallback must update the active request and its log body together - ) -> None: - """Let a classifier fallback without reasoning support remain a usable fallback. - - The dashboard only offers explicitly advertised levels, but an existing config can outlive - a model change and fallbacks can target a different group. Unknown capability fails open; - only a provider that explicitly rejects the parameter has it removed. - """ - if kwargs.get("reasoning_effort") is None or not self._is_classifier_internal_call(kwargs): - return - if self._deployment_accepts_param(deployment, model, "reasoning_effort"): - return - verbose_router_logger.warning( - "litellm.router.py: dropping classifier reasoning_effort for model=%s because the selected deployment does not support it", - model, - ) - kwargs.pop("reasoning_effort", None) - proxy_server_request: Final = kwargs.get("proxy_server_request") - if not isinstance(proxy_server_request, dict): - return - body: Final = proxy_server_request.get("body") - if isinstance(body, dict): - body.pop("reasoning_effort", None) - - ### COMPLETION, EMBEDDING, IMG GENERATION FUNCTIONS - - def completion(self, model: str, messages: list[dict[str, str]], **kwargs) -> ModelResponse | CustomStreamWrapper: - """ - Example usage: - response = router.completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hey, how's it going?"}] - """ - try: - verbose_router_logger.debug("router.completion(model=%s,..)", model) - kwargs["model"] = model - kwargs["messages"] = messages - kwargs["original_function"] = self._completion - controls: Final = per_request_fallback_controls(kwargs) - kwargs[MID_STREAM_FALLBACK_CONTROLS_KEY] = controls # rebind-ok: forwarded to every hop - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - - response: Final = self.function_with_fallbacks(**kwargs) - return response - except Exception as e: - raise e - - def _completion(self, model: str, messages: list[dict[str, str]], **kwargs) -> ModelResponse | CustomStreamWrapper: - model_name = None - deployment = None - try: - controls: Final = kwargs.pop(MID_STREAM_FALLBACK_CONTROLS_KEY, None) - input_kwargs_for_streaming_fallback: Final = mid_stream_fallback_snapshot_kwargs( - model=model, controls=controls, kwargs=kwargs - ) - - # pick the one that is available (lowest TPM/RPM) - with picks_reserve(True): - deployment = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._drop_unsupported_classifier_reasoning_effort( - deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment - model=model, - kwargs=kwargs, - ) - # Check for silent model experiment - # Make a local copy of litellm_params to avoid mutating the Router's state - litellm_params: Final = self._deployment_params_with_request_reasoning_override( - deployment["litellm_params"], kwargs - ) - silent_model: Final = litellm_params.pop("silent_model", None) - - for silent_target in _silent_experiment_targets(silent_model): - # Mirroring traffic to a secondary model - # Use threading.Thread (not ThreadPoolExecutor) - executor.submit() - # requires pickling args, which fails when kwargs contain unpicklable - # objects (e.g. _thread.RLock from OTEL spans, loggers) in deployment. - threading.Thread( - target=self._silent_experiment_completion, - args=(silent_target, messages), - kwargs=_silent_experiment_kwargs_snapshot(kwargs), - daemon=True, - ).start() - - kwargs.setdefault("messages", messages) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) - kwargs.pop("silent_model", None) # Ensure it's not in kwargs either - model_name = litellm_params["model"] - potential_model_client: Final = self._get_client(deployment=deployment, kwargs=kwargs) - # check if provided keys == client keys # - dynamic_api_key: Final = kwargs.get("api_key", None) - if ( - dynamic_api_key is not None - and potential_model_client is not None - and dynamic_api_key != potential_model_client.api_key - ): - model_client = None - else: - model_client = potential_model_client - - ### DEPLOYMENT-SPECIFIC PRE-CALL CHECKS ### (e.g. update rpm pre-call. Raise error, if deployment over limit) - ## only run if model group given, not model id - if model in self._core_model_names or not self._core_has_model_id(model): - self.routing_strategy_pre_call_checks(deployment=deployment) - - input_kwargs: Final = without_router_only_kwargs( - { - **litellm_params, - "messages": messages, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - response: Final = self._host.invoke("completion")(**input_kwargs) - verbose_router_logger.info("litellm.completion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - ## CHECK CONTENT FILTER ERROR ## - if isinstance(response, ModelResponse): - _should_raise = self._should_raise_content_policy_error(model=model, response=response, kwargs=kwargs) - if _should_raise: - raise litellm_exceptions.ContentPolicyViolationError( - message="Response output was blocked.", - model=model, - llm_provider="", - ) - - if ( - isinstance(response, CustomStreamWrapper) - and response.completion_stream is None - and response.make_call is not None - ): - response.fetch_sync_stream() - - # Wrap streaming responses so MidStreamFallbackError (raised - # during iteration) triggers the Router's fallback chain. - if isinstance(response, CustomStreamWrapper): - return self._completion_streaming_iterator( - model_response=response, - messages=messages, - initial_kwargs=input_kwargs_for_streaming_fallback, - first_chunk=_read_first_chunk(response), - ) - - return response - except Exception as e: - verbose_router_logger.info("litellm.completion(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - # Set per-deployment num_retries on exception for retry logic - if deployment is not None: - self._set_deployment_num_retries_on_exception(e, deployment) - self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) - raise e - - def _get_silent_experiment_kwargs(self, **kwargs) -> dict: - """ - Prepare kwargs for a silent experiment by ensuring isolation from the primary call. - - Guarantee metadata isolation: safe_deep_copy falls back to the original - reference when deepcopy fails (e.g. metadata contains UserAPIKeyAuth with - parent_otel_span — an OTel Span that is not deepcopy-able). Force a shallow - copy of the metadata dict so mutations (model_group, is_silent_experiment) - never corrupt the main call's metadata. - """ - from litellm.litellm_core_utils.core_helpers import safe_deep_copy - - silent_kwargs: Final = safe_deep_copy(kwargs) - - # safe_deep_copy may fall back to the original metadata reference when - # deepcopy fails (UserAPIKeyAuth.parent_otel_span is not deepcopy-able). - # Detect this via identity check and force a shallow copy so that setting - # model_group / is_silent_experiment on the silent dict doesn't corrupt - # the primary call's metadata. - original_metadata: Final = kwargs.get("metadata") - if original_metadata is not None and silent_kwargs.get("metadata") is original_metadata: - silent_kwargs["metadata"] = dict(original_metadata) - - if "metadata" not in silent_kwargs: - silent_kwargs["metadata"] = {} - - # OTel spans are not safe to use across event loops. The silent - # experiment runs in a new event loop, so strip the span to prevent - # cross-loop tracing races or span corruption. - silent_kwargs["metadata"].pop("litellm_parent_otel_span", None) - - silent_kwargs["metadata"]["is_silent_experiment"] = True - - # Pop logging objects and call IDs to ensure a fresh logging context - # This prevents collisions in the Proxy's database (spend_logs) - silent_kwargs.pop("litellm_call_id", None) - silent_kwargs.pop("litellm_logging_obj", None) - silent_kwargs.pop("standard_logging_object", None) - # DON'T pop proxy_server_request — it's needed for spend log metadata - - return silent_kwargs - - async def _run_silent_experiment( - self, silent_model: str, messages: Sequence[Mapping[str, str]], silent_kwargs: Mapping[str, object] - ) -> None: - remaining_kwargs: Final = MappingProxyType( - {key: value for key, value in silent_kwargs.items() if key != "stream"} - ) - with mirrored_call(): - response: Final = await self.acompletion( - model=silent_model, - messages=cast(list[AllMessageValues], messages), - stream=bool(silent_kwargs.get("stream", False)), - **remaining_kwargs, - ) - if not isinstance(response, CustomStreamWrapper): - return - async for _ in response: - pass - - def _silent_experiment_completion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs): - """ - Run a silent experiment in the background (thread). - """ - try: - # Prevent infinite recursion if silent model also has a silent model - if kwargs.get("metadata", {}).get("is_silent_experiment", False): - return - - messages = copy.deepcopy(messages) - - verbose_router_logger.info("Starting silent experiment for model %s", silent_model) - - silent_kwargs: Final = self._get_silent_experiment_kwargs(**kwargs) - - # Override model_group to correctly attribute metrics to the silent model - silent_kwargs["metadata"]["model_group"] = silent_model - - # Create a new event loop for this thread so that async success - # callbacks (e.g. _ProxyDBLogger) can schedule and run DB writes. - loop: Final = asyncio.new_event_loop() - asyncio.set_event_loop(loop) - try: - - async def _run_silent_completion(): - await self._run_silent_experiment(silent_model, messages, silent_kwargs) - # Drain any fire-and-forget tasks (e.g. alerting hooks) - # scheduled via asyncio.create_task during acompletion. - pending: Final = asyncio.all_tasks() - current: Final = asyncio.current_task() - if current is not None: - pending.discard(current) - if pending: - await asyncio.gather(*pending, return_exceptions=True) - - loop.run_until_complete(_run_silent_completion()) - finally: - loop.close() - except Exception as e: - verbose_router_logger.error("Silent experiment failed for model %s: %s", silent_model, e) - - # fmt: off - - @overload - async def acompletion( - self, model: str, messages: list[AllMessageValues], stream: Literal[True], **kwargs - ) -> CustomStreamWrapper: - ... - - @overload - async def acompletion( - self, model: str, messages: list[AllMessageValues], stream: Literal[False] = False, **kwargs - ) -> ModelResponse: - ... - - @overload - async def acompletion( - self, model: str, messages: list[AllMessageValues], stream: Literal[True, False] = False, **kwargs - ) -> CustomStreamWrapper | ModelResponse: - ... - - # fmt: on - - # The actual implementation of the function - async def acompletion( - self, - model: str, - messages: list[AllMessageValues], - stream: bool = False, - **kwargs, - ): - try: - kwargs["model"] = model - kwargs["messages"] = messages - kwargs["stream"] = stream - kwargs["original_function"] = self._acompletion - - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - start_time: Final = time.time() - is_prompt_management_model: Final = self._is_prompt_management_model(model) - if not is_prompt_management_model: - controls: Final = per_request_fallback_controls(kwargs) - kwargs[MID_STREAM_FALLBACK_CONTROLS_KEY] = controls # rebind-ok: forwarded to every hop - request_priority, request_kwargs = self._core_split_request_priority(model=model, kwargs=kwargs) - original_function: Final[Callable[..., Awaitable[object]]] = ( - self._prompt_management_acompletion(model=model, messages=messages) - if is_prompt_management_model - else self.async_function_with_fallbacks - ) - if request_priority is None: - response = await original_function(**request_kwargs) - else: - response = await self._schedule_factory( - model=model, - priority=request_priority, - original_function=original_function, - args=(), - kwargs=request_kwargs, - ) - end_time: Final = time.time() - _duration: Final = end_time - start_time - asyncio.create_task( - self.service_logger_obj.async_service_success_hook( - service=ServiceTypes.ROUTER, - duration=_duration, - call_type="acompletion", - start_time=start_time, - end_time=end_time, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - ) - ) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - @staticmethod - def _combine_fallback_usage( - fallback_item: ModelResponseStream, - complete_response_object_usage: Usage | None, - ) -> None: - """Merge partial-stream usage with fallback-stream usage on the chunk.""" - from litellm.cost_calculator import BaseTokenUsageProcessor - - usage: Final = cast(Usage | None, getattr(fallback_item, "usage", None)) - usage_objects: Final = [usage] if usage is not None else [] - if complete_response_object_usage is not None: - usage_objects.append(complete_response_object_usage) - combined_usage: Final = BaseTokenUsageProcessor.combine_usage_objects(usage_objects=usage_objects) - setattr(fallback_item, "usage", combined_usage) - - @staticmethod - def _price_cut_partial( - cut_stream: CustomStreamWrapper, - partial_response: ModelResponse | TextCompletionResponse | None, - ) -> "ContinuedPartial | None": - """Price the cut stream's partial, plus any partial it inherited, before a fallback can reuse its logging.""" - from litellm.cost_calculator import BaseTokenUsageProcessor - - cut_logging: Final = cut_stream.logging_obj - partial_usage: Final = getattr(partial_response, "usage", None) - if partial_response is None or not isinstance(partial_usage, Usage): - return None - inherited_usage: Final = cut_logging.model_call_details.get("continued_partial_usage") - inherited_cost: Final = cut_logging.model_call_details.get("continued_partial_cost") - inherited_id: Final = cut_logging.model_call_details.get("continued_response_id") - priced_partial: Final = partial_response.model_copy( - update={"model": cut_stream.model} if cut_stream.model else {} - ) - priced_partial._hidden_params = { # pyright: ignore[reportPrivateUsage] # no public accessor - key: value for key, value in get_hidden_params_dict(partial_response).items() if key != "response_cost" - } - return ContinuedPartial( - usage=BaseTokenUsageProcessor.combine_usage_objects( - usage_objects=[partial_usage, inherited_usage] - if isinstance(inherited_usage, Usage) - else [partial_usage] - ), - response_cost=(cut_logging.response_cost_calculator(result=priced_partial) or 0.0) - + (inherited_cost if isinstance(inherited_cost, float) else 0.0), - response_id=inherited_id if isinstance(inherited_id, str) else partial_response.id, - ) - - @staticmethod - def _carry_partial_spend_to_fallback(fallback_response: object, cut_partial: "ContinuedPartial | None") -> None: - """Bill what the cut legs consumed on the continuation's success log.""" - fallback_logging: Final = getattr(fallback_response, "logging_obj", None) - if cut_partial is None or not isinstance(fallback_logging, LiteLLMLogging): - return - fallback_logging.record_continued_partial( - usage=cut_partial.usage, response_cost=cut_partial.response_cost, response_id=cut_partial.response_id - ) - - @staticmethod - def _prepare_fallback_hidden_params( - fallback_response: object, - ) -> tuple[dict[str, object], dict[str, object]]: - fallback_hidden_params: Final = get_hidden_params_dict(fallback_response) - fallback_headers: Final = fallback_hidden_params.get("additional_headers") - if not isinstance(fallback_headers, dict): - return fallback_hidden_params, {} - return fallback_hidden_params, cast("dict[str, object]", fallback_headers) - - @staticmethod - def _adopt_fallback_response_headers( - wrapper_ref: "weakref.ref[FallbackAwareStreamWrapper]", - fallback_response: object, - ) -> tuple[dict[str, object], dict[str, object]]: - """Repoint the wrapper at `fallback_response`, returning its prepared hidden params.""" - prepared: Final = PythonCore._prepare_fallback_hidden_params(fallback_response) - adopting_wrapper: Final = wrapper_ref() - if adopting_wrapper is not None: - adopting_wrapper.adopt_fallback_response_headers(fallback_response, prepared) - return prepared - - @staticmethod - def _apply_fallback_hidden_params_to_item( - fallback_item: object, - prepared_fallback_hidden_params: tuple[dict[str, object], dict[str, object]], - ) -> None: - if fallback_item is None or not hasattr(fallback_item, _HIDDEN_PARAMS_ATTR): - return - - fallback_hidden_params, fallback_headers = prepared_fallback_hidden_params - item_hidden_params: Final = get_hidden_params_dict(fallback_item) - item_headers = item_hidden_params.get("additional_headers") - if not isinstance(item_headers, dict): - item_headers = {} - - set_hidden_params( - fallback_item, - { - **item_hidden_params, - **fallback_hidden_params, - "additional_headers": {**item_headers, **fallback_headers}, - }, - ) - - async def _acompletion_streaming_iterator( - self, - model_response: CustomStreamWrapper, - messages: list[dict[str, str]], - initial_kwargs: dict[str, Any], - deployment_slot: contextlib.AsyncExitStack | None = None, - first_chunk: _FirstChunk[AsyncIterator[object]] | None = None, - ) -> CustomStreamWrapper: - """ - Helper to iterate over a streaming response. - - Catches errors for fallbacks using the router's fallback system - - `deployment_slot` holds the deployment's max_parallel_requests semaphore; it is - released when the stream is exhausted, closed, or falls back to another deployment - - `first_chunk` carries the chunk _acompletion already read (None: iterate from the start) - """ - from litellm.exceptions import MidStreamFallbackError - - held_slot: Final = deployment_slot if deployment_slot is not None else contextlib.AsyncExitStack() - - class FallbackStreamWrapper(FallbackAwareStreamWrapper): - def __init__(self, async_generator: AsyncGenerator): - # Copy attributes from the original model_response - super().__init__( - completion_stream=async_generator, - model=model_response.model, - custom_llm_provider=model_response.custom_llm_provider, - logging_obj=model_response.logging_obj, - _response_headers=getattr(model_response, "_response_headers", None), - ) - self._async_generator = async_generator - inner_chunks: Final[object] = getattr(model_response, "chunks", None) - if isinstance(inner_chunks, list): - self.chunks = inner_chunks - # Preserve hidden params (including litellm_overhead_time_ms) from original response - model_response_hidden_params: Final = get_hidden_params(model_response) - if model_response_hidden_params is not None: - self._hidden_params = dict(model_response_hidden_params) - - def __aiter__(self): - return self - - async def __anext__(self): - return await self._async_generator.__anext__() - - async def close_model_response() -> None: - if not hasattr(model_response, "aclose"): - return - try: - await model_response.aclose() - except BaseException as e: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing model_response: %s", - e, - ) - - async def stream_with_fallbacks(): - fallback_response = None # Track for cleanup in finally - try: - async for item in _achunks_with_first(model_response, first_chunk): - yield item - except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: - e: Final = ( - _context_window_stream_fallback_error( - stream_error, - getattr(model_response, "model", ""), - model_response.custom_llm_provider or "", - bool(getattr(model_response, "response_uptil_now", "")) - or _stream_chunks_have_generated_content(model_response.chunks), - ) - if isinstance(stream_error, litellm_exceptions.ContextWindowExceededError) - or isinstance(stream_error.original_exception, litellm_exceptions.ContextWindowExceededError) - else stream_error - ) - with anyio.CancelScope(shield=True): - await close_model_response() - await held_slot.aclose() - committed: Final = bool( - not e.is_pre_first_chunk - and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) - ) - cut_hidden_params: Final = get_hidden_params(model_response) or {} - continue_after_content: Final = committed and self._mid_stream_continuation_eligible( - e=e, request_kwargs=initial_kwargs, cut_deployment_id=cut_hidden_params.get("model_id") - ) - if committed and not continue_after_content: - self._raise_original_mid_stream_error(e) - - from litellm.main import stream_chunk_builder - - complete_response_object: Final = stream_chunk_builder( - chunks=model_response.chunks, messages=messages if continue_after_content else None - ) - complete_response_object_usage: Final = cast( - Usage | None, - getattr(complete_response_object, "usage", None), - ) - try: - # Use the router's fallback system - model_group: Final = cast(str, initial_kwargs.get("model")) - fallbacks: Final[list | None] = initial_kwargs.get("fallbacks", self._core_fallbacks) - context_window_fallbacks: Final[list | None] = initial_kwargs.get( - "context_window_fallbacks", self._core_context_window_fallbacks - ) - content_policy_fallbacks: Final[list | None] = initial_kwargs.get( - "content_policy_fallbacks", self._core_content_policy_fallbacks - ) - initial_kwargs["original_function"] = self._acompletion - continuation_messages: Final = ( - self._build_completion_continuation_input(messages, e.generated_content) - if continue_after_content - else None - ) - continuation_output: Final = ( - ContinuationOutput( - e.generated_content, - continuation_messages[-1], - model_group, - complete_response_object, - checks_post_call_rules=initial_kwargs.get(MID_STREAM_CONTINUATION_KWARG) - is not MID_STREAM_CONTINUATION_MARKER, - ) - if continuation_messages is not None - else None - ) - cut_partial: Final = ( - PythonCore._price_cut_partial(model_response, complete_response_object) - if continuation_output is not None - else None - ) - if cut_partial is not None: - model_response.logging_obj.mark_stash_billed_by_continuation() - if continue_after_content: - emitted_tokens: Final = int( - getattr(complete_response_object_usage, "completion_tokens", 0) or 0 - ) - reduced_ceilings: Final = continuation_output_ceilings(initial_kwargs, emitted_tokens) - if reduced_ceilings is None or continuation_messages is None: - self._raise_original_mid_stream_error(e) - await model_response.wait_for_failure_logging(timeout=MID_STREAM_FAILURE_LOG_WAIT_SECONDS) - initial_kwargs.update(reduced_ceilings) - initial_kwargs["messages"] = continuation_messages - initial_kwargs[MID_STREAM_CONTINUATION_KWARG] = MID_STREAM_CONTINUATION_MARKER - else: - initial_kwargs["messages"] = messages - self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) - fallback_response = await self.async_function_with_fallbacks_common_utils( - e=e.original_exception - if isinstance(e.original_exception, litellm_exceptions.ContextWindowExceededError) - else e, - disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - model_group=model_group, - args=(), - kwargs=initial_kwargs, - include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, - ) - - # If fallback returns a streaming response, iterate over it - if hasattr(fallback_response, "__aiter__"): - PythonCore._carry_partial_spend_to_fallback(fallback_response, cut_partial) - prepared_fallback_hidden_params = PythonCore._adopt_fallback_response_headers( - wrapper_ref, fallback_response - ) - fallback_headers_are_settled = False - fallback_items: Final = ( - continuation_output.splice(fallback_response) - if continuation_output is not None - else fallback_response - ) - async for fallback_item in fallback_items: - if not fallback_headers_are_settled: - fallback_headers_are_settled = True - # a fallback that failed over again only repoints itself once it yields - prepared_fallback_hidden_params = PythonCore._adopt_fallback_response_headers( - wrapper_ref, fallback_response - ) - PythonCore._apply_fallback_hidden_params_to_item( - fallback_item, prepared_fallback_hidden_params - ) - if ( - fallback_item - and isinstance(fallback_item, ModelResponseStream) - and hasattr(fallback_item, "usage") - ): - self._combine_fallback_usage( - fallback_item, - complete_response_object_usage if continuation_output is not None else None, - ) - yield fallback_item - else: - # If fallback returns a non-streaming response, yield None - yield None - - except Exception as fallback_error: - # If fallback also fails, log and re-raise original error - verbose_router_logger.error("Fallback also failed: %s", fallback_error) - # No fallback handled the mid-stream error, so surface the - # real provider exception (e.g. RateLimitError) instead of - # leaking the internal MidStreamFallbackError to the client - if ( - isinstance(fallback_error, MidStreamFallbackError) - and fallback_error.original_exception is not None - ): - raise fallback_error.original_exception from fallback_error - raise fallback_error - finally: - # Close the underlying streams to release HTTP connections - # back to the connection pool when the generator is closed - # (e.g. on client disconnect). - # Shield from anyio cancellation so the awaits can complete. - with anyio.CancelScope(shield=True): - await close_model_response() - await held_slot.aclose() - if fallback_response is not None and hasattr(fallback_response, "aclose"): - try: - await fallback_response.aclose() - except BaseException as close_error: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing fallback_response: %s", - close_error, - ) - - wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks()) - # weak, so the generator closing over it does not keep the wrapper out of - # refcount teardown and delay the `finally` that releases the deployment slot - wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response) - return wrapped_response - - @staticmethod - def _extract_partial_responses_usage( - source_iterator: "BaseResponsesAPIStreamingIterator", - ) -> Optional["ResponseAPIUsage"]: - """ - Best-effort: pull partial token usage from a Responses-API streaming - iterator that errored mid-stream, normalized to ResponseAPIUsage so - the caller can combine without crossing token-naming conventions. - - Two sources, in priority order: - 1. The bridge path (LiteLLMCompletionStreamingIterator) accumulates - chat-completion chunks while streaming — feed them through - stream_chunk_builder to recover chat Usage, then translate - (prompt_tokens → input_tokens, completion_tokens → output_tokens). - 2. The native path (ResponsesAPIStreamingIterator) only has a - completed_response object if the stream reached - RESPONSE_COMPLETED before erroring — uncommon mid-stream but - worth checking. Already ResponseAPIUsage-shaped. - - Returns None when no partial usage is recoverable. - """ - from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, - ) - from litellm.types.llms.openai import ( - ResponseAPIUsage, - ResponseCompletedEvent, - ResponseFailedEvent, - ResponseIncompleteEvent, - ) - - # Bridge subclass is the only iterator that accumulates chat-completion - # chunks. isinstance narrows the type so we can read the attribute - # directly instead of getattr-ing on the base class. - if isinstance(source_iterator, LiteLLMCompletionStreamingIterator): - chunks: Final = source_iterator.collected_chat_completion_chunks - if chunks: - try: - from litellm.main import stream_chunk_builder - - built: Final = stream_chunk_builder(chunks=chunks) - # stream_chunk_builder returns ModelResponse | - # TextCompletionResponse | None. ModelResponse sets .usage - # in __init__ rather than declaring it as a class field, so - # static narrowing doesn't expose it. Mirror the sync path - # (_completion_streaming_iterator) and pull via getattr. - chat: Final[object | None] = getattr(built, "usage", None) if built is not None else None - if chat is not None: - # getattr-with-default because the test path may - # substitute a SimpleNamespace lacking some fields; - # real Usage instances always have them. - prompt: Final = int(getattr(chat, "prompt_tokens", 0) or 0) - completion: Final = int(getattr(chat, "completion_tokens", 0) or 0) - total: Final = int(getattr(chat, "total_tokens", prompt + completion) or (prompt + completion)) - return ResponseAPIUsage( - input_tokens=prompt, - output_tokens=completion, - total_tokens=total, - ) - except Exception: - # Builder is best-effort — fall through to native path. - pass - - # Native path: completed_response is set only if RESPONSE_COMPLETED - # arrived before the error (uncommon mid-stream but worth checking). - # Already ResponseAPIUsage-shaped — return as-is. - completed: Final = source_iterator.completed_response - if isinstance( - completed, - (ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent), - ): - return completed.response.usage - return None - - @staticmethod - def _combine_responses_fallback_usage( - fallback_item: "BaseLiteLLMOpenAIResponseObject", - partial_usage: "ResponseAPIUsage", - ) -> None: - """ - Merge partial-stream usage with fallback-stream usage on a - Responses-API streaming event. - - Only mutates events that carry a `response` with a `usage` field - (response.completed / response.failed / response.incomplete). Other - events pass through unchanged. - - Both inputs are ResponseAPIUsage-shaped (see - _extract_partial_responses_usage which normalizes the bridge path), - so we can sum input_tokens / output_tokens / total_tokens directly - and produce a clean ResponseAPIUsage — no token-naming split, no - setattr bypass. - """ - from litellm.types.llms.openai import ( - ResponseAPIUsage, - ResponseCompletedEvent, - ResponseFailedEvent, - ResponseIncompleteEvent, - ) - - if not isinstance( - fallback_item, - (ResponseCompletedEvent, ResponseFailedEvent, ResponseIncompleteEvent), - ): - return - response: Final = fallback_item.response - if response.usage is None: - return - - fb: Final = response.usage - response.usage = ResponseAPIUsage( - input_tokens=(partial_usage.input_tokens or 0) + (fb.input_tokens or 0), - output_tokens=(partial_usage.output_tokens or 0) + (fb.output_tokens or 0), - total_tokens=(partial_usage.total_tokens or 0) + (fb.total_tokens or 0), - ) - - @staticmethod - def _raise_original_mid_stream_error(e: "MidStreamFallbackError") -> NoReturn: - """Surface the provider error the stream wrapper carried instead of the internal MidStreamFallbackError.""" - if e.original_exception is not None: - raise e.original_exception from e - raise e - - def _active_default_litellm_params(self) -> Mapping[str, object]: - default_params: Final = _REQUEST_KWARGS_ADAPTER.validate_python(self._core_default_litellm_params) - return {key: value for key, value in default_params.items() if value is not None} - - def _params_sent_to_deployment( - self, request_kwargs: Mapping[str, object], deployment: Deployment | None - ) -> Mapping[str, object]: - """The request as the deployment received it, merged in the order _update_kwargs_with_deployment applies.""" - deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {} - defaults: Final = self._active_default_litellm_params() - deployment_tool_choice: Final = ( - {"tool_choice": deployment_params["tool_choice"]} - if deployment_params.get("tool_choice") is not None - else {} - ) - return MappingProxyType({**deployment_params, **defaults, **deployment_tool_choice, **request_kwargs}) - - def _names_deployment_that_cannot_continue(self, target: object) -> bool: - name: Final = fallback_target_name(target) - if name is None or name in self._core_model_names or not self._core_has_model_id(name): - return False - deployment: Final = self._core_get_deployment(model_id=name) - return deployment is not None and not deployment_can_continue(deployment.model_dump(exclude_none=True)) - - def _mid_stream_continuation_eligible( - self, - e: "MidStreamFallbackError", - request_kwargs: Mapping[str, object], - cut_deployment_id: object = None, - ) -> bool: - """Whether a stream that broke after plain assistant text may continue via prefill.""" - if not self.enable_mid_stream_fallback_continuation: - return False - if not e.generated_content or e.emitted_disqualifying_content: - return False - if isinstance(e.original_exception, litellm_exceptions.ContentPolicyViolationError): - return False - cut_deployment: Final = ( - self._core_get_deployment(model_id=cut_deployment_id) if isinstance(cut_deployment_id, str) else None - ) - if isinstance(cut_deployment_id, str) and cut_deployment is None: - return False - if constrains_continuation(self._params_sent_to_deployment(request_kwargs, cut_deployment)): - return False - requested_group: Final = request_kwargs.get("model") - lookup_groups: Final = fallback_lookup_groups( - request_kwargs, requested_group if isinstance(requested_group, str) else None - ) - targets: Final = continuation_fallback_targets( - request_kwargs.get("fallbacks", self._core_fallbacks), lookup_groups - ) - if any(fallback_target_overrides_request_params(target) for target in targets): - return False - if any(self._names_deployment_that_cannot_continue(target) for target in targets): - return False - return not any( - isinstance(callback, CustomGuardrail) - and callback._event_hook_is_event_type( # pyright: ignore[reportPrivateUsage] # the only check of a guardrail's hooks - GuardrailEventHooks.post_call - ) - for callback in self._host.callback_list("callbacks") - ) - - @staticmethod - def _build_completion_continuation_input( - messages: Sequence[Mapping[str, object]], - generated_content: str, - ) -> Sequence[Mapping[str, object]] | None: - """Append the partial output as an assistant prefill, or extend a trailing `prefix: True` assistant turn in - place so the request never ends in two assistant messages. None when the request ends in a plain assistant - turn, in a prefill turn whose content has a shape this cannot extend, or when the cut came mid-echo of the - prefill.""" - last: Final = messages[-1] if messages else None - if last is None or last.get("role") != "assistant": - return [*messages, {"role": "assistant", "content": generated_content, "prefix": True}] - if last.get("prefix") is not True: - return None - new_text: Final = continued_text(last, generated_content) - if new_text is None: - return None - content: Final = last.get("content") - if content is None or isinstance(content, str): - return [*messages[:-1], {**last, "content": f"{content or ''}{new_text}", "prefix": True}] - try: - blocks: Final = _CONTENT_BLOCKS_ADAPTER.validate_python(content) - except ValidationError: - return None - extended: Final = _extend_trailing_text_block(blocks, new_text) - return [*messages[:-1], {**last, "content": extended, "prefix": True}] - - async def _aresponses_streaming_iterator( - self, - response: "BaseResponsesAPIStreamingIterator", - initial_kwargs: dict[str, object], - ) -> "BaseResponsesAPIStreamingIterator": - """ - Wrap a Responses-API streaming iterator so MidStreamFallbackError - triggers the Router's fallback chain (parity with - _acompletion_streaming_iterator for the chat-completions path). - - The Responses-API streaming path goes through - _ageneric_api_call_with_fallbacks rather than _acompletion, so the - returned iterator is never wrapped by the chat completions - fallback handler. Without this wrapper, MidStreamFallbackError - raised mid-stream from the underlying CustomStreamWrapper (used by - LiteLLMCompletionStreamingIterator when the Responses API is - served via the completion bridge) propagates unhandled and the - configured cross-provider fallback never fires. - - Full parity with the chat-completions path: - - Pre-first-chunk: retry with the original input unchanged. - - Partial content: continue from the generated text, as an - assistant prefill on a fallback that takes one and as assistant - history plus a user request to continue on any other. A request - whose output is constrained (a JSON schema, a forced tool), or a - stream that began a tool call, is not continued, and the primary's - error surfaces instead. - - Usage combining: merge partial-stream usage onto the fallback's - response.completed event so accounting reflects both attempts. - - Stream cleanup: shielded aclose() on both source and fallback - iterators on terminate. - """ - from litellm.exceptions import MidStreamFallbackError - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - get_openai_response_types, - ) - - source_iterator: Final = response - - # Pre-resolve the set of terminal stream event types so the - # per-chunk type check inside FallbackResponsesStreamWrapper - # stays cheap; mirrors the source-iterator filter at - # responses/streaming_iterator.py:243-247. - _openai_types: Final = get_openai_response_types() - _RESPONSES_TERMINAL_EVENT_TYPES: Final = ( - _openai_types.ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - _openai_types.ResponsesAPIStreamEvents.RESPONSE_INCOMPLETE, - _openai_types.ResponsesAPIStreamEvents.RESPONSE_FAILED, - ) - - class FallbackResponsesStreamWrapper(BaseResponsesAPIStreamingIterator): - """ - Subclasses BaseResponsesAPIStreamingIterator only for isinstance - compatibility (proxy + interactions code paths check the type). - Bypasses the parent constructor and delegates iteration to an - async generator. - """ - - fallback_headers_adopted: bool = False - - def __init__(self, async_generator: "AsyncGenerator[ResponsesAPIStreamingResponse, None]"): - import time - from datetime import datetime - - self._async_generator = async_generator - self._billing_source: object = source_iterator - self._last_sequence_number = -1 - # Mirror every attribute BaseResponsesAPIStreamingIterator.__init__ - # would have set. The wrapper bypasses super().__init__ (it has no - # httpx.Response of its own and no provider config to drive), so - # we copy from source_iterator where applicable and use safe - # defaults elsewhere. This keeps inherited methods (e.g. - # _check_max_streaming_duration, _handle_failure) safe to call. - # - # The bridge path (LiteLLMCompletionStreamingIterator used by - # Anthropic/Bedrock/Vertex) does not call super().__init__ and - # is missing many of these attributes — use getattr fallbacks - # so wrapper construction never raises AttributeError. The - # bridge stores the logging object as `litellm_logging_obj`. - # base class declares non-Optional types for these - # fields but the bridge path (LiteLLMCompletionStreamingIterator) - # can legitimately omit them at runtime — keep the None - # fallback. Same lines passed mypy on the pre-fix file - # because the surrounding function body wasn't fully - # type-narrowed; the new typed terminal-event tuple above - # is what made these surface. - self.response = getattr(source_iterator, "response", None) - self.model = getattr(source_iterator, "model", None) - self.logging_obj = getattr( - source_iterator, - "logging_obj", - getattr(source_iterator, "litellm_logging_obj", None), - ) - self.finished = False - self.responses_api_provider_config = getattr(source_iterator, "responses_api_provider_config", None) - self.completed_response = None - self.start_time = getattr(source_iterator, "start_time", datetime.now()) - self._failure_handled = False - self._yielded_first_chunk = False - self._generated_content = "" - self._completed_response_cached = False - self._completed_response_logged = False - self._completed_response_cache_hit = None - self._persist_completed_response_before_logging = True - self._stream_created_time = time.time() - self.litellm_metadata = getattr(source_iterator, "litellm_metadata", None) - self.custom_llm_provider = getattr(source_iterator, "custom_llm_provider", None) - self.request_data = getattr(source_iterator, "request_data", {}) or {} - self.call_type = getattr(source_iterator, "call_type", None) - # Preserve hidden params so response headers (model_id, - # api_base, additional_headers) keep flowing. - self._hidden_params = dict(getattr(source_iterator, "_hidden_params", None) or {}) - - def adopt_fallback_headers(self, fallback_response: object) -> tuple[dict[str, object], dict[str, object]]: - prepared: Final = PythonCore._prepare_fallback_hidden_params(fallback_response) - self._billing_source = fallback_response - self._hidden_params = {**prepared[0], "additional_headers": prepared[1]} - self.fallback_headers_adopted = True - return prepared - - def partial_response_on_disconnect(self) -> "ResponseIncompleteEvent | None": - """ - The event to bill on a client disconnect. After a mid-stream hop the primary - stream's relayed tokens are folded into the fallback's event, the way - _combine_responses_fallback_usage does for a natural end of stream. - """ - from litellm.responses.streaming_iterator import partial_responses_incomplete_event_on_disconnect - - current: Final = partial_responses_incomplete_event_on_disconnect(self._billing_source) - if self._billing_source is source_iterator: - return current - primary: Final = partial_responses_incomplete_event_on_disconnect(source_iterator) - if current is None or primary is None: - return current or primary - if primary.response.usage is not None: - PythonCore._combine_responses_fallback_usage(current, primary.response.usage) - return current - - def __aiter__(self): - return self - - @property - def yielded_first_chunk(self) -> bool: - return self._yielded_first_chunk - - async def __anext__(self): - try: - chunk: Final = await self._async_generator.__anext__() - except StopAsyncIteration: - # The inner generator is exhausted. If we never sniffed a - # terminal event off a chunk (the bridge path emits the - # final response.completed via common_done_event_logic, - # which raises StopAsyncIteration after returning it), - # fall back to whatever the source iterator latched so - # the proxy's container-ownership hook still sees a - # completed_response instead of logging a spurious - # "no completed_response" warning. - if self.completed_response is None: - self.completed_response = getattr(source_iterator, "completed_response", None) - raise - # Sniff the terminal stream event off each forwarded chunk - # so ``self.completed_response`` is populated regardless of - # which inner iterator produced it (source_iterator, - # fallback_iterator, or any future wrapper). Without this - # the proxy's container-ownership hook (which reads - # ``getattr(stream_response, "completed_response", None)`` - # via extract_completed_responses_response) silently - # records nothing on streaming /v1/responses calls — every - # follow-up /v1/containers//files call then 403s for - # the very key that created the container (#30210). - if self.completed_response is None and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: - self.completed_response = chunk - self._yielded_first_chunk = True - sequence_number: Final = getattr(chunk, "sequence_number", None) - if not isinstance(sequence_number, int): - return chunk - if not self.fallback_headers_adopted: - self._last_sequence_number = max(self._last_sequence_number, sequence_number) - return chunk - self._last_sequence_number = max(self._last_sequence_number + 1, sequence_number) - if self._last_sequence_number == sequence_number: - return chunk - return chunk.model_copy(update={"sequence_number": self._last_sequence_number}) - - async def aclose(self): - # async generators always expose aclose — no defensive check needed. - await self._async_generator.aclose() - - async def stream_with_fallbacks(): - held_lifecycle_events: tuple[object, ...] = () # rebind-ok: flushed at first output, dropped on fallback - response_id_exposed = False # rebind-ok: set once, before an early-released lifecycle event is yielded - try: - async for item in cast( # cast-ok: response iterators are async streams - AsyncIterator[object], source_iterator - ): - if not response_id_exposed and self._aresponses_background_stream_can_release( - item, source_iterator, initial_kwargs - ): - response_id_exposed = True - if not response_id_exposed and _responses_stream_holds_event(item, len(held_lifecycle_events)): - held_lifecycle_events = (*held_lifecycle_events, item) - continue - for held_event in held_lifecycle_events: - yield held_event - held_lifecycle_events = () - yield item - for held_event in held_lifecycle_events: - yield held_event - except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: - e: Final = ( - _context_window_stream_fallback_error( - stream_error, - source_iterator.model or "", - source_iterator.custom_llm_provider or "", - wrapper.yielded_first_chunk, - ) - if isinstance(stream_error, litellm_exceptions.ContextWindowExceededError) - or isinstance(stream_error.original_exception, litellm_exceptions.ContextWindowExceededError) - else stream_error - ) - if response_id_exposed: - self._raise_released_stream_error(e, initial_kwargs) - async with contextlib.aclosing( - self._aresponses_fallback_attempt( - e, source_iterator, initial_kwargs, wrapper.adopt_fallback_headers, held_lifecycle_events - ) - ) as fallback_stream: - async for fallback_item in fallback_stream: - yield fallback_item - except Exception: - for held_event in held_lifecycle_events: - yield held_event - raise - finally: - with anyio.CancelScope(shield=True): - if hasattr(source_iterator, "aclose"): - try: - await source_iterator.aclose() - except Exception as exc: - verbose_router_logger.debug( - "stream_with_fallbacks(aresponses): error closing source: %s", - exc, - ) - - wrapper: Final = FallbackResponsesStreamWrapper( - stream_with_fallbacks() # pyright: ignore[reportArgumentType] # relays response events typed as object - ) - return wrapper - - def _aresponses_background_stream_can_release( - self, - item: object, - source_iterator: "BaseResponsesAPIStreamingIterator", - kwargs: Mapping[str, object], - ) -> bool: - from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator - - if ( - type(source_iterator) is not ResponsesAPIStreamingIterator - or source_iterator.custom_llm_provider != "openai" - ): - return False - if not _responses_event_announces_background_run(item): - return False - model_group: Final = kwargs.get("model") - return isinstance(model_group, str) and not self._stream_can_fall_back(model_group, kwargs) - - def _raise_released_stream_error( - self, e: "MidStreamFallbackError", initial_kwargs: Mapping[str, object] - ) -> NoReturn: - original: Final = e.original_exception - model_group: Final = initial_kwargs.get("model") - if ( - self._host.dynamic().expose_router_debug_in_errors - and isinstance(original, litellm_exceptions.ContentPolicyViolationError) - and isinstance(model_group, str) - and not fallbacks_disabled_for_request(initial_kwargs) - and _mid_stream_error_before_output(e) - ): - hop_depth: Final = initial_kwargs.get("fallback_depth") - original.message += format_content_policy_no_fallback_note( - model_group, - fallback_lookup_groups(initial_kwargs, model_group), - _request_fallback_list( - initial_kwargs, - "content_policy_fallbacks", - cast( # cast-ok: untyped router attribute - "list[object] | None", self._core_content_policy_fallbacks - ), - ), - _request_fallback_list( - initial_kwargs, - "fallbacks", - cast("list[object] | None", self._core_fallbacks), # cast-ok: untyped router attribute - ), - isinstance(hop_depth, int) and hop_depth > 0, - ) - PythonCore._raise_original_mid_stream_error(e) - - def _responses_continuation( - self, - e: "MidStreamFallbackError", - source_iterator: "BaseResponsesAPIStreamingIterator", - request_kwargs: Mapping[str, object], - ) -> ResponsesContinuation | None: - """What a fallback continues from once a Responses stream broke after text, raising the primary's error - when the request constrains its output or the stream began a tool call, since a continued JSON document or - tool call would not parse, and when a content filter cut it, so no other model finishes filtered text.""" - if _mid_stream_error_before_output(e): - return None - if source_iterator.started_tool_call or isinstance( - e.original_exception, litellm_exceptions.ContentPolicyViolationError - ): - PythonCore._raise_original_mid_stream_error(e) - cut_deployment_id: Final = (get_hidden_params(source_iterator) or {}).get("model_id") - cut_deployment: Final = ( - self._core_get_deployment(model_id=cut_deployment_id) if isinstance(cut_deployment_id, str) else None - ) - if constrains_responses_continuation(self._params_sent_to_deployment(request_kwargs, cut_deployment)): - PythonCore._raise_original_mid_stream_error(e) - earlier: Final = request_kwargs.get(RESPONSES_CONTINUATION_KWARG) - if isinstance(earlier, ResponsesContinuation): - return earlier.extended(e.generated_content) - return ResponsesContinuation(request_kwargs.get("input"), e.generated_content) - - @staticmethod - def _responses_continuation_ceilings( - e: "MidStreamFallbackError", request_kwargs: Mapping[str, object] - ) -> Mapping[str, object]: - """The caller's output ceilings less what the cut stream already emitted, raising the primary's error once - nothing is left.""" - emitted_tokens: Final = token_counter(model=e.model or "", text=e.generated_content) - reduced: Final = continuation_output_ceilings(request_kwargs, emitted_tokens) - if reduced is None: - PythonCore._raise_original_mid_stream_error(e) - return reduced - - async def _aresponses_fallback_attempt( - self, - e: "MidStreamFallbackError", - source_iterator: "BaseResponsesAPIStreamingIterator", - initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain - adopt_headers: Callable[[object], tuple[dict[str, object], dict[str, object]]], # mutable-ok: hidden params - held_lifecycle_events: tuple[object, ...], - ) -> AsyncGenerator[object, None]: - """ - Re-enters the Router's fallback chain for a mid-stream Responses API error and yields - whatever the fallback attempt produces. The lifecycle events the primary stream held - back reach the client only when no fallback lands, so the client sees exactly one - response announced, the one whose id completes. Split out of - _aresponses_streaming_iterator to keep each function's cyclomatic complexity within - the repo's C901 budget. - """ - from litellm.exceptions import MidStreamFallbackError - - continuation: Final = self._responses_continuation(e, source_iterator, initial_kwargs) - reduced_ceilings: Final[Mapping[str, object]] = ( - PythonCore._responses_continuation_ceilings(e, initial_kwargs) if continuation is not None else {} - ) - seam: Final = ResponsesContinuationSeam(continuation.emitted_text if continuation is not None else "") - partial_usage: Final = PythonCore._extract_partial_responses_usage(source_iterator) - fallback_response = None # rebind-ok: pre-init so finally can close it if a fallback was actually attempted - fallback_yielded = False # rebind-ok: flipped on the first fallback item so a fallback that dies before its first event still replays the primary's held announcement - try: - model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: model group - fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the common_utils list|None param - "fallbacks", self._core_fallbacks - ) - context_window_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below - "context_window_fallbacks", self._core_context_window_fallbacks - ) - content_policy_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below - "content_policy_fallbacks", self._core_content_policy_fallbacks - ) - initial_kwargs["original_function"] = ( # rebind-ok: the fallback chain re-enters on the same kwargs - self._ageneric_api_call_with_fallbacks_responses_attempt - ) - if continuation is not None: - initial_kwargs.update(reduced_ceilings) - initial_kwargs[RESPONSES_CONTINUATION_KWARG] = continuation # rebind-ok: shaped per fallback deployment - initial_kwargs["input"] = responses_continuation_input(continuation) # rebind-ok: fallback hop input - # The Responses-API path stores observability metadata - # under "litellm_metadata" (not the default "metadata") — - # see _ageneric_api_call_with_fallbacks. Mirroring that - # here ensures model_group, model_group_alias, and trace - # ids land in the same key litellm.aresponses reads from. - self._update_kwargs_before_fallbacks( - model=model_group, - kwargs=initial_kwargs, - metadata_variable_name="litellm_metadata", - ) - # The content-policy dispatch branch matches on the trigger's own type, so a refusal's - # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. - fallback_trigger: Final[Exception] = ( - e.original_exception - if isinstance( - e.original_exception, - (litellm_exceptions.ContentPolicyViolationError, litellm_exceptions.ContextWindowExceededError), - ) - else e - ) - fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success - e=fallback_trigger, - disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - model_group=model_group, - args=(), - kwargs=initial_kwargs, - include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, - ) - prepared_fallback_hidden_params: Final = adopt_headers(fallback_response) - if hasattr(fallback_response, "__aiter__"): - async for fallback_item in seam.splice(fallback_response): - PythonCore._apply_fallback_hidden_params_to_item(fallback_item, prepared_fallback_hidden_params) - if partial_usage is not None: - PythonCore._combine_responses_fallback_usage(fallback_item, partial_usage) - fallback_yielded = True - yield fallback_item - else: - fallback_yielded = True # rebind-ok: see the pre-init above - yield fallback_response - except Exception as fallback_error: - verbose_router_logger.error("Responses streaming fallback also failed: %s", fallback_error) - if not fallback_yielded: - for held_event in held_lifecycle_events: - yield held_event - if isinstance(fallback_error, MidStreamFallbackError) and fallback_error.original_exception is not None: - raise fallback_error.original_exception from fallback_error - raise - finally: - if fallback_response is not None and hasattr(fallback_response, "aclose"): - with anyio.CancelScope(shield=True): - try: - await fallback_response.aclose() - except Exception as exc: - verbose_router_logger.debug( - "stream_with_fallbacks(aresponses): error closing fallback: %s", - exc, - ) - - def _completion_streaming_iterator( - self, - model_response: CustomStreamWrapper, - messages: list[dict[str, str]], - initial_kwargs: dict, - first_chunk: _FirstChunk[Iterator[object]] | None = None, - ) -> CustomStreamWrapper: - """ - Sync equivalent of _acompletion_streaming_iterator. - - Wraps a sync streaming response so that MidStreamFallbackError - (raised by CustomStreamWrapper.__next__) triggers the Router's - fallback chain instead of surfacing directly to the caller. - """ - from litellm.exceptions import MidStreamFallbackError - - class SyncFallbackStreamWrapper(FallbackAwareStreamWrapper): - def __init__(self, sync_generator: Generator): - super().__init__( - completion_stream=sync_generator, - model=model_response.model, - custom_llm_provider=model_response.custom_llm_provider, - logging_obj=model_response.logging_obj, - _response_headers=getattr(model_response, "_response_headers", None), - ) - self._sync_generator = sync_generator - model_response_hidden_params: Final = get_hidden_params(model_response) - if model_response_hidden_params is not None: - self._hidden_params = dict(model_response_hidden_params) - - def __iter__(self): - return self - - def __next__(self): - return next(self._sync_generator) - - router_self: Final = self - - def stream_with_fallbacks(): - fallback_response = None - try: - for item in _chunks_with_first(model_response, first_chunk): - yield item - except (MidStreamFallbackError, litellm_exceptions.ContextWindowExceededError) as stream_error: - e: Final = ( - _context_window_stream_fallback_error( - stream_error, - getattr(model_response, "model", ""), - model_response.custom_llm_provider or "", - bool(getattr(model_response, "response_uptil_now", "")) - or _stream_chunks_have_generated_content(model_response.chunks), - ) - if isinstance(stream_error, litellm_exceptions.ContextWindowExceededError) - or isinstance(stream_error.original_exception, litellm_exceptions.ContextWindowExceededError) - else stream_error - ) - if fallbacks_disabled_for_request(initial_kwargs) or ( - not e.is_pre_first_chunk - and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) - ): - if e.original_exception is not None: - raise e.original_exception from e - raise - - try: - model_group: Final = cast(str, initial_kwargs.get("model")) - fallbacks: Final[list | None] = initial_kwargs.get("fallbacks", router_self._core_fallbacks) - context_window_fallbacks: Final[list | None] = initial_kwargs.get( - "context_window_fallbacks", - router_self._core_context_window_fallbacks, - ) - content_policy_fallbacks: Final[list | None] = initial_kwargs.get( - "content_policy_fallbacks", - router_self._core_content_policy_fallbacks, - ) - initial_kwargs["original_function"] = router_self._completion - initial_kwargs["messages"] = messages - router_self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) - fallback_response = run_async_function( - router_self.async_function_with_fallbacks_common_utils, - e=e.original_exception - if isinstance(e.original_exception, litellm_exceptions.ContextWindowExceededError) - else e, - disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - model_group=model_group, - args=(), - kwargs=initial_kwargs, - include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, - ) - - if hasattr(fallback_response, "__iter__"): - prepared_fallback_hidden_params = PythonCore._adopt_fallback_response_headers( - wrapper_ref, fallback_response - ) - fallback_headers_are_settled = False - for fallback_item in cast(Iterable[object], fallback_response): # cast-ok: __iter__ was checked - if not fallback_headers_are_settled: - fallback_headers_are_settled = True - # a fallback that failed over again only repoints itself once it yields - prepared_fallback_hidden_params = PythonCore._adopt_fallback_response_headers( - wrapper_ref, fallback_response - ) - PythonCore._apply_fallback_hidden_params_to_item( - fallback_item, prepared_fallback_hidden_params - ) - if ( - fallback_item - and isinstance(fallback_item, ModelResponseStream) - and hasattr(fallback_item, "usage") - ): - router_self._combine_fallback_usage(fallback_item, None) - yield fallback_item - else: - yield None - - except Exception as fallback_error: - verbose_router_logger.error("Fallback also failed: %s", fallback_error) - if ( - isinstance(fallback_error, MidStreamFallbackError) - and fallback_error.original_exception is not None - ): - raise fallback_error.original_exception from fallback_error - raise fallback_error - finally: - if hasattr(model_response, "close"): - try: - model_response.close() - except BaseException as close_err: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing model_response: %s", - close_err, - ) - if fallback_response is not None and hasattr(fallback_response, "close"): - try: - fallback_response.close() - except BaseException as close_err: - verbose_router_logger.debug( - "stream_with_fallbacks: error closing fallback_response: %s", - close_err, - ) - - wrapped_response: Final = SyncFallbackStreamWrapper(stream_with_fallbacks()) - # weak, for the same reason as the async twin - wrapper_ref: Final[weakref.ref[FallbackAwareStreamWrapper]] = weakref.ref(wrapped_response) - return wrapped_response - - async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs): - """ - Run a silent experiment in the background. - """ - try: - # Prevent infinite recursion if silent model also has a silent model - if kwargs.get("metadata", {}).get("is_silent_experiment", False): - return - - messages = copy.deepcopy(messages) - - verbose_router_logger.info("Starting silent experiment for model %s", silent_model) - - silent_kwargs: Final = self._get_silent_experiment_kwargs(**kwargs) - # Override model_group to correctly attribute metrics to the silent model - silent_kwargs["metadata"]["model_group"] = silent_model - - # Trigger the silent request - await self._run_silent_experiment(silent_model, messages, silent_kwargs) - except Exception as e: - verbose_router_logger.error("Silent experiment failed for model %s: %s", silent_model, e) - - async def _acompletion( - self, model: str, messages: list[dict[str, str]], **kwargs - ) -> ModelResponse | CustomStreamWrapper: - """ - - Get an available deployment - - call it with a semaphore over the call - - semaphore specific to it's rpm - - in the semaphore, make a check against it's local rpm before running - """ - model_name = None - deployment = None - _timeout_debug_deployment_dict = {} # this is a temporary dict to debug timeout issues - try: - controls: Final = kwargs.pop(MID_STREAM_FALLBACK_CONTROLS_KEY, None) - input_kwargs_for_streaming_fallback: Final = mid_stream_fallback_snapshot_kwargs( - model=model, controls=controls, kwargs=kwargs - ) - - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - start_time: Final = time.time() - deployment = await self._core_async_get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._drop_unsupported_classifier_reasoning_effort( - deployment=cast(DeploymentTypedDict, deployment), # cast-ok: selection returns a router deployment - model=model, - kwargs=kwargs, - ) - - _timeout_debug_deployment_dict = deployment - end_time: Final = time.time() - _duration: Final = end_time - start_time - asyncio.create_task( - self.service_logger_obj.async_service_success_hook( - service=ServiceTypes.ROUTER, - duration=_duration, - call_type="async_get_available_deployment", - start_time=start_time, - end_time=end_time, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - ) - ) - - # debug how often this deployment picked - - self._track_deployment_metrics(deployment=deployment, parent_otel_span=parent_otel_span) - - # Check for silent model experiment - # Make a local copy of litellm_params to avoid mutating the Router's state - litellm_params: Final = self._deployment_params_with_request_reasoning_override( - deployment["litellm_params"], kwargs - ) - silent_model: Final = litellm_params.pop("silent_model", None) - - for silent_target in _silent_experiment_targets(silent_model): - # Mirroring traffic to a secondary model - # This is a silent experiment, so we don't want to block the primary request - asyncio.create_task( - self._silent_experiment_acompletion( - silent_model=silent_target, - messages=messages, # Use messages instead of *args - **_silent_experiment_kwargs_snapshot(kwargs), - ) - ) - - kwargs.setdefault("messages", messages) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) - kwargs.pop("silent_model", None) # Ensure it's not in kwargs either - - model_name = litellm_params["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - input_kwargs: Final = without_router_only_kwargs( - { - **litellm_params, - "messages": messages, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - logging_obj: Final[LiteLLMLogging | None] = kwargs.get("litellm_logging_obj", None) - - max_parallel_requests_limit: Final = self._get_client( - deployment=deployment, - kwargs=kwargs, - client_type="max_parallel_requests", - ) - compacted_input: Final = await compact_to_fit(self._core_ctx, deployment, input_kwargs, "chat") - async with contextlib.AsyncExitStack() as deployment_slot: - if isinstance(max_parallel_requests_limit, MaxParallelRequestsLimit): - deployment_slot.enter_context(max_parallel_requests_limit) - await self.async_routing_strategy_pre_call_checks( - deployment=deployment, - logging_obj=logging_obj, - parent_otel_span=parent_otel_span, - ) - response = await self._host.invoke("acompletion")(**compacted_input) - - ## CHECK CONTENT FILTER ERROR ## - if isinstance(response, ModelResponse): - _should_raise = self._should_raise_content_policy_error( - model=model, response=response, kwargs=kwargs - ) - if _should_raise: - raise litellm_exceptions.ContentPolicyViolationError( - message="Response output was blocked.", - model=model, - llm_provider="", - ) - - if ( - isinstance(response, CustomStreamWrapper) - and response.completion_stream is None - and response.make_call is not None - ): - await response.fetch_stream() - first_chunk: Final = ( - await _aread_first_chunk(response) if isinstance(response, CustomStreamWrapper) else None - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - # debug how often this deployment picked - self._track_deployment_metrics( - deployment=deployment, - response=response, - parent_otel_span=parent_otel_span, - ) - - if isinstance(response, CustomStreamWrapper): - return await self._acompletion_streaming_iterator( - model_response=response, - messages=messages, - initial_kwargs=input_kwargs_for_streaming_fallback, - deployment_slot=deployment_slot.pop_all(), - first_chunk=first_chunk, - ) - - return response - except litellm_exceptions.Timeout as e: - deployment_litellm_params: Final[Mapping[str, object]] = _timeout_debug_deployment_dict.get( - "litellm_params", {} - ) - if self._host.dynamic().expose_router_debug_in_errors: - e.message += ( - f"\n\nDeployment Info: request_timeout: {deployment_litellm_params.get('request_timeout')}" - f"\ntimeout: {deployment_litellm_params.get('timeout')}" - f"\nstream_timeout: {deployment_litellm_params.get('stream_timeout')}" - ) - # Set per-deployment num_retries on exception for retry logic - if deployment is not None: - self._set_deployment_num_retries_on_exception(e, deployment) - self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) - raise e - except Exception as e: - verbose_router_logger.info("litellm.acompletion(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - # Set per-deployment num_retries on exception for retry logic - if deployment is not None: - self._set_deployment_num_retries_on_exception(e, deployment) - self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) - raise e - - def _update_kwargs_before_fallbacks( - self, - model: str, - kwargs: dict, - metadata_variable_name: str | None = "metadata", - ) -> None: - """ - Adds/updates to kwargs: - - litellm_trace_id - - metadata - - num_retries is deliberately left as the caller passed it, with no default filled in. - async_function_with_retries resolves the router/global default itself, and it can only - honour the precedence request > deployment litellm_params > litellm_settings while an - absent num_retries still means "the caller did not ask for one". - """ - kwargs.setdefault("litellm_trace_id", str(uuid.uuid4())) - model_group_alias: str | None = None - if self._core_get_model_from_alias(model=model): - model_group_alias = model - kwargs.setdefault(metadata_variable_name, {}).update( - {"model_group": model, "model_group_alias": model_group_alias} - ) - - def _set_deployment_num_retries_on_exception(self, exception: Exception, deployment: dict) -> None: - """ - Set num_retries from deployment litellm_params on the exception. - - This allows the retry logic in async_function_with_retries to use - per-deployment retry settings instead of the global setting. - """ - # Only set if exception doesn't already have num_retries - if hasattr(exception, "num_retries") and exception.num_retries is not None: - return - - litellm_params: Final = deployment.get("litellm_params", {}) - dep_num_retries: Final = litellm_params.get("num_retries") - if dep_num_retries is not None: - try: - exception.num_retries = int(dep_num_retries) # Handle both int and str - except (ValueError, TypeError): - pass # Skip if value can't be converted to int - - def _set_failed_deployment_id_on_exception(self, exception: Exception, deployment: Mapping[str, Any]) -> None: - """ - Stamp the failed deployment's `model_info.id` on the exception so the - fallback layer can exclude it from subsequent re-picks within the same - request (used by weighted-routing failover). - - Idempotent: never overwrites an existing value, so the id of the - deployment that *first* failed in a chain is preserved if multiple - layers re-raise. - """ - if getattr(exception, "failed_deployment_id", None): - return - deployment_id: Final = (deployment.get("model_info") or {}).get("id") - if deployment_id: - try: - exception.failed_deployment_id = deployment_id - except Exception: - pass - - def _stamp_failed_deployment_id_with_effective_model_info( - self, exception: Exception, deployment: Mapping[str, object], kwargs: Mapping[str, object] - ) -> None: - # A client-side-credential call gets a dynamic deployment id generated inside - # _update_kwargs_with_deployment and stamped into kwargs["model_info"]; stamping - # the static shared deployment's id instead would let one tenant's bad credentials - # cool down the deployment every other tenant sharing this config relies on. - if getattr(exception, "failed_deployment_id", None): - return - effective_model_info: Final = kwargs.get("model_info") or deployment.get("model_info") or MappingProxyType({}) - self._set_failed_deployment_id_on_exception(exception, MappingProxyType({"model_info": effective_model_info})) - configured_deployment_id: Final = _MODEL_INFO_ADAPTER.validate_python(effective_model_info).get( - "original_model_id" - ) - if isinstance(configured_deployment_id, str) and configured_deployment_id: - exception.failed_configured_deployment_id = configured_deployment_id # pyright: ignore[reportAttributeAccessIssue] # rebind-ok: the exception carries the failure to _maybe_run_weighted_failover - - def _stamp_failed_attempt(self, exception: Exception, kwargs: Mapping[str, object], *, picked: bool) -> None: - if picked: - self._stamp_failed_deployment_id_with_effective_model_info(exception, MappingProxyType({}), kwargs) - - @staticmethod - def _stamp_retry_skip_deployment_id(exception: Exception, kwargs: Mapping[str, object]) -> None: - effective_model_info: Final = kwargs.get("model_info") - if not isinstance(effective_model_info, Mapping): - return - model_info: Final = _MODEL_INFO_ADAPTER.validate_python(effective_model_info) - configured_deployment_id: Final = ( - model_info.get("original_model_id") - if isinstance(exception, litellm_exceptions.UnsupportedEndpointError) - else None - ) - deployment_id: Final = configured_deployment_id or model_info.get("id") - if isinstance(deployment_id, str) and deployment_id: - exception.retry_skip_deployment_id = deployment_id # pyright: ignore[reportAttributeAccessIssue] # dynamic stamp, read by _deployment_ids_to_skip_on_retry - - def _update_kwargs_with_default_litellm_params( - self, kwargs: dict, metadata_variable_name: str | None = "metadata" - ) -> None: - """ - Adds default litellm params to kwargs, if set. - - Handles inserting this as either "metadata" or "litellm_metadata" depending on the metadata_variable_name - """ - # 1) copy your defaults and pull out metadata - defaults: Final = dict(without_retrieve_stream_defaults(self._core_default_litellm_params, kwargs)) - metadata_defaults: Final = defaults.pop("metadata", {}) or {} - - # 2) add any non-metadata defaults that aren't already in kwargs - for key, value in defaults.items(): - if value is None: - continue - kwargs.setdefault(key, value) - - # 3) merge in metadata, this handles inserting this as either "metadata" or "litellm_metadata" - kwargs.setdefault(metadata_variable_name, {}).update(metadata_defaults) - - def _handle_clientside_credential( - self, deployment: dict, kwargs: dict, function_name: str | None = None - ) -> Deployment: - """ - Build a per-request Deployment carrying the caller-supplied api_key/api_base, - with its own stable id for cooldown, logging, and cost-map identity. - - This deployment is deliberately never registered with the router (no - upsert_deployment/add_deployment call): doing so used to add it to - self._core_model_list under the shared model_name, which made a request-scoped, - caller-supplied provider credential a permanent, load-balanced deployment - that every other caller of that model group could be routed onto. Its - pricing is still registered directly, so a custom price configured on the - underlying deployment still applies to this call. - """ - model_info: Final = deployment.get("model_info", {}).copy() - litellm_params: Final = deployment["litellm_params"].copy() - dynamic_litellm_params: Final = get_dynamic_litellm_params(litellm_params=litellm_params, request_kwargs=kwargs) - # Use deployment model_name as model_group for generating model_id - metadata_variable_name: Final = get_router_metadata_variable_name( - function_name=function_name, - ) - model_group: Final = kwargs.get(metadata_variable_name, {}).get("model_group") - _model_id: Final = self._core_generate_model_id(model_group=model_group, litellm_params=dynamic_litellm_params) - original_model_id: Final = model_info.get("id") - model_info["id"] = _model_id - model_info["original_model_id"] = original_model_id - deployment_pydantic_obj: Final = Deployment( - model_name=model_group, - litellm_params=LiteLLM_Params.model_validate(dynamic_litellm_params), - model_info=model_info, - ) - PythonCore._register_deployment_pricing(deployment=deployment_pydantic_obj) - return deployment_pydantic_obj - - @staticmethod - def _merge_tools_from_deployment( - deployment: dict, kwargs: dict, defaults: Mapping[str, object] = _NO_DEFAULT_PARAMS - ) -> None: - """ - Merge tools from deployment litellm_params with request kwargs. - When both have tools, concatenate them (deployment tools first, then request tools). - tool_choice: use request value if provided, else deployment's. - """ - dep_params_raw: Final = deployment.get("litellm_params", {}) or {} - if isinstance(dep_params_raw, dict): - dep_params = dep_params_raw - else: - dep_params = dep_params_raw.model_dump(exclude_none=True) - if _drops_malformed_request_tools(dep_params, defaults, kwargs): - for name in TOOL_FIELDS: - kwargs.pop(name, None) - merged_tools: Final = _merged_tools(dep_params.get("tools") or [], kwargs.get("tools") or []) - if merged_tools is not None: - kwargs["tools"] = merged_tools - if "tool_choice" not in kwargs and dep_params.get("tool_choice") is not None: - kwargs["tool_choice"] = dep_params["tool_choice"] - - def _update_kwargs_with_deployment( - self, - deployment: dict, - kwargs: dict, - function_name: str | None = None, - splits_stream_timeout: bool = False, - opens_responses_socket: bool = False, - ) -> None: - """ - 5 jobs: - - Adds selected deployment, model_info and api_base to kwargs["metadata"] (used for logging) - - Adds default litellm params to kwargs, if set. - - Merges tools from deployment with request (proxy-configured tools + request tools). - - On a fallback hop, drops the encrypted reasoning this deployment cannot decrypt, keeping its summary. - - On a Responses continuation, sends the partial response as prefill where this deployment takes one. - """ - for key in self._forwarded_alias_marker_keys_the_deployment_sets( - deployment=deployment, forwarded_keys=kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, ()) - ): - kwargs.pop(key, None) - kwargs.pop("daily_reset_time", None) # pyright: ignore[reportUnknownMemberType] # kwargs is an untyped dict - defaults: Final = _LITELLM_PARAMS_ADAPTER.validate_python(self._core_default_litellm_params) - self._merge_tools_from_deployment(deployment=deployment, kwargs=kwargs, defaults=defaults) - - model_info = deployment.get("model_info", {}).copy() - deployment_litellm_model_name = deployment["litellm_params"]["model"] - deployment_api_base = deployment["litellm_params"].get("api_base") - deployment_model_name: Final = deployment["model_name"] - from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils - - request_context: Final = cast( # cast-ok: the router's legacy request dict has string keys - Mapping[str, object], kwargs - ) - raw_deployment_params: Final[object] = deployment["litellm_params"] - deployment_context: Final = _LITELLM_PARAMS_ADAPTER.validate_python( - raw_deployment_params.model_dump(exclude_none=True) - if isinstance(raw_deployment_params, LiteLLM_Params) - else raw_deployment_params - ) - _drop_malformed_request_params_the_deployment_sets( - kwargs, # pyright: ignore[reportUnknownArgumentType] # kwargs is the router's untyped request dict - deployment_context, - defaults, - ) - declared_provider: Final = deployment_context.get("custom_llm_provider") - declared_model: Final = deployment_context.get("model") - deployment_provider: Final = ( - declared_provider - if isinstance(declared_provider, str) - else inferred_provider(declared_model if isinstance(declared_model, str) else None) - ) - forwarded_authorization: Final = ( - ProviderSpecificHeaderUtils.forwarded_authorization( - request_context.get("provider_specific_header"), deployment_provider - ) - if deployment_provider in _FORWARDED_CREDENTIAL_ISOLATION_PROVIDERS - else None - ) - if is_clientside_credential(request_kwargs=kwargs) or forwarded_authorization is not None: - deployment_pydantic_obj: Final = self._handle_clientside_credential( - deployment=deployment, - kwargs={ - **kwargs, - **({"api_key": forwarded_authorization} if forwarded_authorization is not None else {}), - }, - function_name=function_name, - ) - model_info = deployment_pydantic_obj.model_info.model_dump() - deployment_litellm_model_name = deployment_pydantic_obj.litellm_params.model - deployment_api_base = deployment_pydantic_obj.litellm_params.api_base - - metadata_variable_name: Final = get_router_metadata_variable_name( - function_name=function_name, - ) - deployment_that_just_failed: Final = _deployment_that_just_failed( - _FALLBACK_HOP_ADAPTER.validate_python(kwargs).get(metadata_variable_name) - ) - routing: Final = deployment.get(KUBERNETES_POD_ROUTING_KEY) - - kwargs.setdefault(metadata_variable_name, {}).update( - { - "deployment": deployment_litellm_model_name, - "model_info": model_info, - "api_base": deployment_api_base, - "deployment_model_name": deployment_model_name, - KUBERNETES_POD_ROUTING_KEY: routing, - } - ) - other_bucket: Final = "metadata" if metadata_variable_name == "litellm_metadata" else "litellm_metadata" - other_bucket_metadata: Final = kwargs.get(other_bucket) - if isinstance(other_bucket_metadata, Mapping) and KUBERNETES_POD_ROUTING_KEY in other_bucket_metadata: - kwargs[other_bucket] = { # rebind-ok: this method updates the router's request kwargs in place - **other_bucket_metadata, - KUBERNETES_POD_ROUTING_KEY: routing, - } - - # A retry/fallback reuses this same kwargs dict for the next deployment. - # Refund and clear any reservation the previous deployment attempt left - # here before it's wiped below, instead of relying on that attempt's - # (possibly still-pending) failure event to do it. - refund_stale_reservation_before_retry(self._core_cache, kwargs) - set_io_token_rate_limit_request_kwargs( - kwargs, - store_in_context=deployment_has_io_token_limits(deployment), - opens_responses_socket=opens_responses_socket, - metadata_channel=metadata_variable_name, - ) - - kwargs[metadata_variable_name].setdefault( - ROUTING_REQUEST_TAGS_METADATA_KEY, - tuple(get_tags_from_request_kwargs(kwargs, metadata_variable_name=metadata_variable_name)), - ) - - ## DEPLOYMENT-LEVEL TAGS - deployment_tags: Final = deployment.get("litellm_params", {}).get("tags") - merged_tags: Final = tags_with_serving_deployment_tags(kwargs[metadata_variable_name], deployment_tags) - if merged_tags is not None: - kwargs[metadata_variable_name]["tags"] = list(merged_tags) - - baggage_stamp: Final = baggage_stamp_for_serving_deployment(kwargs[metadata_variable_name], deployment_tags) - if baggage_stamp is not None: - kwargs[metadata_variable_name][BAGGAGE_MAPPED_TAGS_METADATA_KEY] = baggage_stamp - - ## CREDENTIAL NAME AS TAG - credential_name: Final = deployment.get("litellm_params", {}).get("litellm_credential_name") - if credential_name: - credential_tag: Final = f"Credential: {credential_name}" - existing_tags = kwargs[metadata_variable_name].get("tags") or [] - if credential_tag not in existing_tags: - existing_tags.append(credential_tag) - kwargs[metadata_variable_name]["tags"] = existing_tags - - from litellm.llms.github_copilot.per_user_auth import ( - github_copilot_per_user_credential_name, - ) - - configured_per_user_name: Final = github_copilot_per_user_credential_name( - cast( # cast-ok: deployment litellm_params is a str-keyed dict - "Mapping[str, object]", deployment.get("litellm_params") - ) - ) - caller_credential_name: Final[object] = cast( # cast-ok: kwargs is the untyped request dict - "object", kwargs.get("litellm_credential_name") - ) - if ( - configured_per_user_name is not None - and isinstance(caller_credential_name, str) - and caller_credential_name != configured_per_user_name - ): - raise litellm_exceptions.BadRequestError( - message="litellm_credential_name cannot be overridden on a deployment that uses per-user GitHub OAuth", - model=deployment_model_name, - llm_provider="", - ) - - kwargs["model_info"] = model_info - - if function_name == "_ageneric_api_call_with_fallbacks": - from litellm.passthrough.timeout_utils import ( - resolve_llm_passthrough_timeout, - ) - - _router_timeout: Final = ( - self.request_timeout - if self.request_timeout is not None - else float(self._explicit_timeout) - if isinstance(self._explicit_timeout, (int, float)) - else None - ) - _router_stream_timeout: Final = ( - self.stream_timeout - if self.stream_timeout is not None - else self.request_timeout - if self.request_timeout is not None - else self._core_default_litellm_params.get("stream_timeout") - ) - passthrough_timeout: Final = partial( - resolve_llm_passthrough_timeout, - litellm_params=deployment["litellm_params"], - router_timeout=_router_timeout, - router_stream_timeout=_router_stream_timeout, - ) - timeout_request: Final = {**kwargs, "stream": True} if requests_streamed_retrieval(kwargs) else kwargs - kwargs.update( - CompletionTimeout.stream_request_kwargs( - passthrough_timeout(kwargs={**kwargs, "stream": False}), passthrough_timeout(kwargs=timeout_request) - ) - if splits_stream_timeout and timeout_request.get("stream") - else {"timeout": passthrough_timeout(kwargs=timeout_request)} - ) - elif splits_stream_timeout: - request_kwargs: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs) - deployment_params: Final = _DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) - kwargs.update( - CompletionTimeout.stream_request_kwargs( - self._get_non_stream_timeout(kwargs=request_kwargs, data=deployment_params), - self._get_stream_timeout(kwargs=request_kwargs, data=deployment_params) - if request_kwargs.get("stream") - else None, - ) - ) - else: - kwargs["timeout"] = self._get_timeout( - kwargs=_REQUEST_KWARGS_ADAPTER.validate_python(kwargs), - data=_DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]), - ) - - deployment_litellm_params: Final = _DEPLOYMENT_PARAMS_ADAPTER.validate_python(deployment["litellm_params"]) - if router_config_requests_stream_usage( - _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), - ( - _LITELLM_PARAMS_ADAPTER.validate_python( - deployment_litellm_params - if isinstance(deployment_litellm_params, Mapping) - else deployment_litellm_params.model_dump(exclude_none=True) - ), - defaults, - ), - ): - kwargs[STREAM_USAGE_INJECTED_FIELD] = True - self._update_kwargs_with_default_litellm_params(kwargs=kwargs, metadata_variable_name=metadata_variable_name) - pinned_model_id: Final = deployment_model_id_pin(deployment_context) - if pinned_model_id is not None: - kwargs["model_id"] = pinned_model_id # rebind-ok: this method's contract is filling the caller's kwargs - hop_kwargs: Final = _FALLBACK_HOP_ADAPTER.validate_python(kwargs) - if _is_fallback_hop(hop_kwargs): - EncryptedContentAffinityCheck.strip_reasoning_the_targets_cannot_decrypt( - self._core_ctx, - hop_kwargs.get("input"), - hop_kwargs.get("messages"), - (_FALLBACK_HOP_ADAPTER.validate_python(deployment),), - unmarked_origin=deployment_that_just_failed, - ) - if RESPONSES_CONTINUATION_KWARG in hop_kwargs: - del kwargs[RESPONSES_CONTINUATION_KWARG] # rebind-ok: this method prepares the call's kwargs in place - kwargs["input"] = responses_continuation_input_for( # rebind-ok: the deployment's continuation input - hop_kwargs, _FALLBACK_HOP_ADAPTER.validate_python(deployment) - ) - - def _get_async_openai_model_client(self, deployment: dict, kwargs: dict): - """ - Helper to get AsyncOpenAI or AsyncAzureOpenAI client that was created for the deployment - - The same OpenAI client is re-used to optimize latency / performance in production - - If dynamic api key is provided: - Do not re-use the client. Pass model_client=None. The OpenAI/ AzureOpenAI client will be recreated in the handler for the llm provider - """ - potential_model_client: Final = self._get_client(deployment=deployment, kwargs=kwargs, client_type="async") - - # check if provided keys == client keys # - dynamic_api_key: Final = kwargs.get("api_key", None) - if ( - dynamic_api_key is not None - and potential_model_client is not None - and dynamic_api_key != potential_model_client.api_key - ): - model_client = None - else: - model_client = potential_model_client - - return model_client - - def _get_stream_timeout( - self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams - ) -> object: - """Helper to get stream timeout from kwargs or deployment params""" - return ( - kwargs.get("stream_timeout", None) # the params dynamically set by user - or get_service_tier_timeout(kwargs.get("service_tier"), data, "stream_timeout_by_service_tier") - or data.get("stream_timeout", None) # timeout set on litellm_params for this deployment - or self.stream_timeout # timeout set on router - or self.request_timeout # litellm_settings.request_timeout (per-attempt) - or self._core_default_litellm_params.get("stream_timeout", None) - ) - - def _get_non_stream_timeout( - self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams - ) -> object: - """Helper to get non-stream timeout from kwargs or deployment params""" - return ( - kwargs.get("timeout", None) # the params dynamically set by user - or kwargs.get("request_timeout", None) # the params dynamically set by user - or get_service_tier_timeout(kwargs.get("service_tier"), data, "timeout_by_service_tier") - or data.get("timeout", None) # timeout set on litellm_params for this deployment - or data.get("request_timeout", None) # timeout set on litellm_params for this deployment - or self.request_timeout # litellm_settings.request_timeout (per-attempt) - or self.timeout # timeout set on router (router_settings.timeout) - or self._core_default_litellm_params.get("timeout", None) - ) - - def _get_timeout(self, kwargs: Mapping[str, object], data: Mapping[str, object] | GenericLiteLLMParams) -> object: - """Helper to get timeout from kwargs or deployment params""" - stream_timeout: Final = self._get_stream_timeout(kwargs=kwargs, data=data) if kwargs.get("stream") else None - return stream_timeout if stream_timeout is not None else self._get_non_stream_timeout(kwargs=kwargs, data=data) - - async def abatch_completion( - self, - models: list[str], - messages: list[dict[str, str]] | list[list[dict[str, str]]], - **kwargs, - ): - """ - Async Batch Completion. Used for 2 scenarios: - 1. Batch Process 1 request to N models on litellm.Router. Pass messages as List[Dict[str, str]] to use this - 2. Batch Process N requests to M models on litellm.Router. Pass messages as List[List[Dict[str, str]]] to use this - - Example Request for 1 request to N models: - ``` - response = await router.abatch_completion( - models=["gpt-3.5-turbo", "groq-llama"], - messages=[ - {"role": "user", "content": "is litellm becoming a better product ?"} - ], - max_tokens=15, - ) - ``` - - - Example Request for N requests to M models: - ``` - response = await router.abatch_completion( - models=["gpt-3.5-turbo", "groq-llama"], - messages=[ - [{"role": "user", "content": "is litellm becoming a better product ?"}], - [{"role": "user", "content": "who is this"}], - ], - ) - ``` - """ - ############## Helpers for async completion ################## - - async def _async_completion_no_exceptions(model: str, messages: list[AllMessageValues], **kwargs): - """ - Wrapper around self.async_completion that catches exceptions and returns them as a result - """ - try: - return await self.acompletion(model=model, messages=messages, **kwargs) - except Exception as e: - return e - - async def _async_completion_no_exceptions_return_idx( - model: str, - messages: list[AllMessageValues], - idx: int, # index of message this response corresponds to - **kwargs, - ): - """ - Wrapper around self.async_completion that catches exceptions and returns them as a result - """ - try: - return ( - await self.acompletion(model=model, messages=messages, **kwargs), - idx, - ) - except Exception as e: - return e, idx - - ############## Helpers for async completion ################## - - if isinstance(messages, list) and all(isinstance(m, dict) for m in messages): - _tasks = [] - for model in models: - # add each task but if the task fails - _tasks.append(_async_completion_no_exceptions(model=model, messages=messages, **kwargs)) - response = await asyncio.gather(*_tasks) - return response - elif isinstance(messages, list) and all(isinstance(m, list) for m in messages): - _tasks = [] - for idx, message in enumerate(messages): - for model in models: - # Request Number X, Model Number Y - _tasks.append( - _async_completion_no_exceptions_return_idx( - model=model, - idx=idx, - messages=message, - **kwargs, - ) - ) - responses: Final = await asyncio.gather(*_tasks) - final_responses: Final[list[list[object]]] = [[] for _ in range(len(messages))] - for response in responses: - if isinstance(response, tuple): - final_responses[response[1]].append(response[0]) - else: - final_responses[0].append(response) - return final_responses - - async def abatch_completion_one_model_multiple_requests( - self, model: str, messages: list[list[AllMessageValues]], **kwargs - ): - """ - Async Batch Completion - Batch Process multiple Messages to one model_group on litellm.Router - - Use this for sending multiple requests to 1 model - - Args: - model (List[str]): model group - messages (List[List[Dict[str, str]]]): list of messages. Each element in the list is one request - **kwargs: additional kwargs - Usage: - response = await self.abatch_completion_one_model_multiple_requests( - model="gpt-3.5-turbo", - messages=[ - [{"role": "user", "content": "hello"}, {"role": "user", "content": "tell me something funny"}], - [{"role": "user", "content": "hello good mornign"}], - ] - ) - """ - - async def _async_completion_no_exceptions(model: str, messages: list[AllMessageValues], **kwargs): - """ - Wrapper around self.async_completion that catches exceptions and returns them as a result - """ - try: - return await self.acompletion(model=model, messages=messages, **kwargs) - except Exception as e: - return e - - _tasks: Final = [] - for message_request in messages: - # add each task but if the task fails - _tasks.append(_async_completion_no_exceptions(model=model, messages=message_request, **kwargs)) - - response: Final = await asyncio.gather(*_tasks) - return response - - # fmt: off - - @overload - async def abatch_completion_fastest_response( - self, model: str, messages: list[dict[str, str]], stream: Literal[True], **kwargs - ) -> CustomStreamWrapper: - ... - - - - @overload - async def abatch_completion_fastest_response( - self, model: str, messages: list[dict[str, str]], stream: Literal[False] = False, **kwargs - ) -> ModelResponse: - ... - - # fmt: on - - async def abatch_completion_fastest_response( - self, - model: str, - messages: list[dict[str, str]], - stream: bool = False, - **kwargs, - ): - """ - model - List of comma-separated model names. E.g. model="gpt-4, gpt-3.5-turbo" - - Returns fastest response from list of model names. OpenAI-compatible endpoint. - """ - models: Final = [m.strip() for m in model.split(",")] - - async def _async_completion_no_exceptions( - model_name: str, messages: list[dict[str, str]], stream: bool, **kwargs: object - ) -> ModelResponse | CustomStreamWrapper | Exception: - """ - Wrapper around self.acompletion that catches exceptions and returns them as a result - """ - try: - result = await self.acompletion(model=model_name, messages=messages, stream=stream, **kwargs) - return result - except asyncio.CancelledError: - verbose_router_logger.debug("Received 'task.cancel'. Cancelling call w/ model=%s.", model_name) - raise - except Exception as e: - return e - - pending_tasks = [] - - async def check_response(task: asyncio.Task): - nonlocal pending_tasks - try: - result: Final = await task - if isinstance(result, (ModelResponse, CustomStreamWrapper)): - verbose_router_logger.debug("Received successful response. Cancelling other LLM API calls.") - # If a desired response is received, cancel all other pending tasks - for t in pending_tasks: - t.cancel() - return result - except Exception: - # Ignore exceptions, let the loop handle them - pass - finally: - # Remove the task from pending tasks if it finishes - try: - pending_tasks.remove(task) - except KeyError: - pass - - for model_name in models: - task = asyncio.create_task( - _async_completion_no_exceptions(model_name=model_name, messages=messages, stream=stream, **kwargs) - ) - pending_tasks.append(task) - - # Await the first task to complete successfully - while pending_tasks: - done, pending_tasks = await asyncio.wait(pending_tasks, return_when=asyncio.FIRST_COMPLETED) - for completed_task in done: - result = await check_response(completed_task) - - if result is not None: - # Return the first successful result - if (result_hidden_params := get_hidden_params(result)) is not None: - result_hidden_params["fastest_response_batch_completion"] = True - return result - - # If we exit the loop without returning, all tasks failed - raise Exception("All tasks failed") - - ### SCHEDULER ### - - # fmt: off - - @overload - async def schedule_acompletion( - self, - model: str, - messages: list[AllMessageValues], - priority: int | None = None, - stream: Literal[False] = False, - **kwargs, - ) -> ModelResponse: - ... - - @overload - async def schedule_acompletion( - self, model: str, messages: list[AllMessageValues], priority: int | None, stream: Literal[True], **kwargs - ) -> CustomStreamWrapper: - ... - - # fmt: on - - async def schedule_acompletion( - self, - model: str, - messages: list[AllMessageValues], - priority: int | None = None, - stream=False, - **kwargs, - ): - queued_priority, _ = self._core_split_request_priority(model=model, kwargs={**kwargs, "priority": priority}) - if queued_priority is None: - raise priority_error(model=model, message=unqueueable_priority_message(priority)) - return await self.acompletion(model=model, messages=messages, stream=stream, priority=queued_priority, **kwargs) - - def _core_split_request_priority( - self, model: str, kwargs: Mapping[str, object] - ) -> tuple[int | None, Mapping[str, object]]: - request_priority: Final = resolve_request_priority( - requested=kwargs.get("priority"), - default_priority=self.default_priority, - drops_params=lambda: self._drops_params(model=model, kwargs=kwargs), - ) - if isinstance(request_priority, InvalidPriority): - raise priority_error(model=model, message=request_priority.message) - return request_priority, {key: value for key, value in kwargs.items() if key != "priority"} - - def priority_header_update( - self, model: str, kwargs: Mapping[str, object], header: str | None - ) -> Mapping[str, int | str]: - return priority_header_fields(header, drops_params=lambda: self._drops_params(model=model, kwargs=kwargs)) - - def _drops_params(self, model: str, kwargs: Mapping[str, object]) -> bool: - return request_drops_params( - kwargs, self._core_default_litellm_params, self._request_deployment_params(model=model, kwargs=kwargs) - ) - - async def _schedule_factory( - self, - model: str, - priority: int, - original_function: Callable[..., Awaitable[object]], - args: tuple[object, ...], - kwargs: Mapping[str, object], - ): - await self._wait_for_scheduler_turn( - model=model, - priority=priority, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - request_kwargs=kwargs, - ) - try: - _response: Final = await original_function(*args, **kwargs) - response_hidden_params: Final = get_hidden_params(_response) - if response_hidden_params is not None: - additional_headers: Final = cast( # cast-ok: router headers are stored as a mutable mapping - dict[str, object], response_hidden_params.setdefault("additional_headers", {}) - ) - additional_headers.update({"x-litellm-request-prioritization-used": True}) - return _response - except Exception as e: - setattr(e, "priority", priority) - raise e - - async def _wait_for_scheduler_turn( - self, model: str, priority: int, parent_otel_span: Span | None, request_kwargs: Mapping[str, object] - ) -> None: - pinned: Final = self._pinned_deployment(model) - - async def healthy_deployments() -> Sequence[object]: - if pinned is not None: - return (pinned,) - deployments, _ = await self._async_get_healthy_deployments( - model=model, - parent_otel_span=parent_otel_span, - request_kwargs=share_holder_request_kwargs(request_kwargs, self._ptu_share_teams()), - ) - return deployments - - await self.scheduler.wait_for_turn( - request=FlowItem(priority=priority, request_id=str(uuid.uuid4()), model_name=model), - timeout=self.timeout, - get_healthy_deployments=healthy_deployments, - ) - - def _is_prompt_management_model(self, model: str) -> bool: - model_list: Final = self._core_get_model_list(model_name=model) - if model_list is None or len(model_list) != 1: - return False - - litellm_model: Final = model_list[0]["litellm_params"].get("model", None) - if litellm_model is None or "/" not in litellm_model: - return False - - split_litellm_model: Final = litellm_model.split("/")[0] - return split_litellm_model in self._host.dynamic().known_custom_logger_compatible_callbacks - - def _pinned_deployment(self, model: str) -> Deployment | None: - return self._core_get_deployment(model_id=model) if model not in self._core_model_names else None - - def _request_deployment_params(self, model: str, kwargs: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: - deployment: Final = self._pinned_deployment(model) - if deployment is not None: - return (deployment.litellm_params.model_dump(exclude_none=True),) - candidates: Final = ( - self._get_deployment_by_litellm_model(model=model) - if kwargs.get("specific_deployment") is True - else self._core_get_model_list(model_name=model, team_id=get_request_team_id(kwargs)) or () - ) - return tuple(candidate["litellm_params"] for candidate in candidates) - - def _prompt_management_acompletion( - self, - model: str, - messages: list[AllMessageValues], # mutable-ok: _prompt_management_factory takes the caller's list - ) -> Callable[..., Awaitable[object]]: - async def acompletion(**kwargs: object) -> object: # kwargs-ok: called with the request's own keyword payload - return await self._prompt_management_factory(model=model, messages=messages, kwargs=dict(kwargs)) - - return acompletion - - async def _prompt_management_factory( - self, - model: str, - messages: list[AllMessageValues], - kwargs: dict[str, Any], - ): - litellm_logging_object = kwargs.get("litellm_logging_obj", None) - if litellm_logging_object is None: - litellm_logging_object, kwargs = function_setup( - **{ - "original_function": "acompletion", - "rules_obj": Rules(), - "start_time": get_utc_datetime(), - **kwargs, - } - ) - litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object) - caller_timeout: Final = kwargs.get("timeout") or kwargs.get("request_timeout") - caller_stream_timeout: Final = kwargs.get("stream_timeout") - specific_deployment: Final = kwargs.pop("specific_deployment", None) - prompt_management_deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=cast(list[dict[str, str]], messages), # cast-ok: selection reads messages structurally - specific_deployment=specific_deployment, - request_kwargs=kwargs, - ) - - self._update_kwargs_with_deployment( - deployment=prompt_management_deployment, kwargs=kwargs, splits_stream_timeout=True - ) - data: Final = prompt_management_deployment["litellm_params"].copy() - - litellm_model: Final = data.get("model", None) - - # litellm_agent/ prefix only strips the model name, no prompt_id needed - is_litellm_agent_model: Final = isinstance(litellm_model, str) and litellm_model.startswith("litellm_agent/") - - prompt_id = kwargs.get("prompt_id") or prompt_management_deployment["litellm_params"].get("prompt_id", None) - prompt_variables: Final = kwargs.get("prompt_variables") or prompt_management_deployment["litellm_params"].get( - "prompt_variables", None - ) - prompt_label: Final = kwargs.get("prompt_label", None) or prompt_management_deployment["litellm_params"].get( - "prompt_label", None - ) - - if not is_litellm_agent_model and (prompt_id is None or not isinstance(prompt_id, str)): - raise ValueError(f"Prompt ID is not set or not a string. Got={prompt_id}, type={type(prompt_id)}") - if prompt_variables is not None and not isinstance(prompt_variables, dict): - raise ValueError( - f"Prompt variables is set but not a dictionary. Got={prompt_variables}, type={type(prompt_variables)}" - ) - - ( - model, - messages, - optional_params, - ) = litellm_logging_object.get_chat_completion_prompt( - model=litellm_model, - messages=messages, - non_default_params=get_non_default_completion_params(kwargs=kwargs), - prompt_id=prompt_id, - prompt_variables=prompt_variables, - prompt_label=prompt_label, - request_kwargs=kwargs, - injected_for_every_deployment=True, - ) - - # Filter out prompt management specific parameters from data before merging - prompt_management_params: Final = { - "bitbucket_config", - "dotprompt_config", - "prompt_id", - "prompt_variables", - "prompt_label", - "prompt_version", - } - filtered_data: Final = {k: v for k, v in data.items() if k not in prompt_management_params} - - kwargs = {**filtered_data, **kwargs, **optional_params} - kwargs["model"] = model - kwargs["messages"] = messages - kwargs["litellm_logging_obj"] = litellm_logging_object - kwargs["prompt_id"] = prompt_id - kwargs["prompt_variables"] = prompt_variables - kwargs["prompt_label"] = prompt_label - - _model_list: Final = self._core_get_model_list(model_name=model) - if _model_list is None or len(_model_list) == 0: # if direct call to model - kwargs.pop("original_function") - return await self._host.invoke("acompletion")(**kwargs) - - target_group_timeouts: Final = { - key: value - for key, value in ( - ("timeout", caller_timeout if caller_timeout is not None else filtered_data.get("timeout")), - ( - "stream_timeout", - caller_stream_timeout if caller_stream_timeout is not None else filtered_data.get("stream_timeout"), - ), - ) - if value is not None - } - return await self.async_function_with_fallbacks( - **{key: value for key, value in kwargs.items() if key not in ("timeout", "stream_timeout")}, - **target_group_timeouts, - ) - - def image_generation(self, prompt: str, model: str, **kwargs): - try: - kwargs["model"] = model - kwargs["prompt"] = prompt - kwargs["original_function"] = self._image_generation - kwargs.setdefault("metadata", {}).update({"model_group": model}) - response: Final = self.function_with_fallbacks(**kwargs) - - return response - except Exception as e: - raise e - - def _image_generation(self, prompt: str, model: str, **kwargs): - model_name: Final = "" - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _image_generation()- model: %s", model) - with picks_reserve(True): - deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - - self.total_calls[model_name] += 1 - - ### DEPLOYMENT-SPECIFIC PRE-CALL CHECKS ### (e.g. update rpm pre-call. Raise error, if deployment over limit) - self.routing_strategy_pre_call_checks(deployment=deployment) - - response: Final = self._host.invoke("image_generation")( - **{ - **data, - "prompt": prompt, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.image_generation(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.image_generation(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def aimage_generation(self, prompt: str, model: str, **kwargs): - try: - kwargs["model"] = model - kwargs["prompt"] = prompt - kwargs["original_function"] = self._aimage_generation - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _aimage_generation(self, prompt: str, model: str, **kwargs): - model_name = model - start_time: Final = datetime.now() # noqa: DTZ005 # SDK logging uses naive local timestamps - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _image_generation()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - try: - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - except Exception as selection_error: - # Admitted attempts keep SDK-owned logging, including deployment callbacks and retry state - request_kwargs: Final = cast( # cast-ok: Python kwargs have string keys and opaque values - Mapping[str, object], kwargs - ) - if request_kwargs.get("litellm_logging_obj") is None: - refusal_kwargs: Final = {**request_kwargs, "model": model, "prompt": prompt} - refusal_kwargs.setdefault("litellm_call_id", str(uuid.uuid4())) - logging_obj, _ = function_setup( - "aimage_generation", - Rules(), - start_time, - **refusal_kwargs, # pyright: ignore[reportArgumentType] # preserve opaque SDK options, including is_async_call - ) - try: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=selection_error, - traceback_exception=traceback.format_exc(), - prefer_async_handlers=True, - ) - ) - finally: - _restore_correlation_context_if_supported(logging_obj) - raise - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - data: Final = deployment["litellm_params"].copy() - model_name = data["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - - self.total_calls[model_name] += 1 - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("aimage_generation")( - **{ - **data, - "prompt": prompt, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.aimage_generation(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.aimage_generation(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def atranscription(self, file: FileTypes, model: str, **kwargs): - """ - Example Usage: - - ``` - from litellm import Router - client = Router(model_list = [ - { - "model_name": "whisper", - "litellm_params": { - "model": "whisper-1", - }, - }, - ]) - - audio_file = open("speech.mp3", "rb") - transcript = await client.atranscription( - model="whisper", - file=audio_file - ) - - ``` - """ - try: - kwargs["model"] = model - kwargs["file"] = file - kwargs["original_function"] = self._atranscription - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _atranscription(self, file: FileTypes, model: str, **kwargs): - model_name: Final = model - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _atranscription()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - - self.total_calls[model_name] += 1 - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("atranscription")( - **{ - **data, - "file": file, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.atranscription(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.atranscription(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): - """ - Example Usage: - - ``` - from litellm import Router - client = Router(model_list = [ - { - "model_name": "tts", - "litellm_params": { - "model": "tts-1", - }, - }, - ]) - - async with client.aspeech( - model="tts", - voice="alloy", - input="the quick brown fox jumped over the lazy dogs", - api_base=None, - api_key=None, - organization=None, - project=None, - max_retries=1, - timeout=600, - client=None, - optional_params={}, - ) as response: - response.stream_to_file(speech_file_path) - - ``` - """ - try: - kwargs["model"] = model - kwargs["input"] = input - kwargs["voice"] = voice - kwargs["original_function"] = self._aspeech - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): - model_name: Final = model - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _aspeech()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "prompt"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - - self.total_calls[model_name] += 1 - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("aspeech")( - **{ - **data, - "input": input, - "voice": data.get("voice") if voice is None else voice, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.aspeech(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.aspeech(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def arerank(self, model: str, **kwargs): - try: - kwargs["model"] = model - kwargs["input"] = input - kwargs["original_function"] = self._arerank - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _arerank(self, model: str, **kwargs): - model_name = None - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _rerank()- model: %s", model) - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - model_name = data["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - response: Final = await self._host.invoke("arerank")( - **{ - **data, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.arerank(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.arerank(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - def text_completion( - self, - model: str, - prompt: str, - is_retry: bool | None = False, - is_fallback: bool | None = False, - is_async: bool | None = False, - **kwargs, - ): - messages: Final = [{"role": "user", "content": prompt}] - with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): - try: - kwargs["model"] = model - kwargs["prompt"] = prompt - kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) - kwargs.setdefault("metadata", {}).update({"model_group": model}) - - # pick the one that is available (lowest TPM/RPM) - with endpoint_selection("text_completion", model=model), picks_reserve(True): - deployment: Final = self.get_available_deployment( - model=model, - messages=messages, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - - data: Final = deployment["litellm_params"].copy() - for k, v in self._core_default_litellm_params.items(): - if k not in kwargs: # prioritize model-specific params > default router params - kwargs[k] = v - elif k == "metadata": - kwargs[k].update(v) - pinned_model_id: Final = deployment_model_id_pin(_LITELLM_PARAMS_ADAPTER.validate_python(data)) - if pinned_model_id is not None: - kwargs["model_id"] = pinned_model_id # rebind-ok: this call's own kwargs, filled like the defaults - - # call via litellm.completion() - return self._host.invoke("text_completion")( - **{**data, "prompt": prompt, "caching": self._core_cache_responses, **kwargs} - ) - except Exception as e: - raise e - - async def atext_completion( - self, - model: str, - prompt: str, - is_retry: bool | None = False, - is_fallback: bool | None = False, - is_async: bool | None = False, - **kwargs, - ): - try: - kwargs["model"] = model - kwargs["prompt"] = prompt - kwargs["original_function"] = self._atext_completion - - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - request_priority, request_kwargs = self._core_split_request_priority(model=model, kwargs=kwargs) - if request_priority is None: - return await self.async_function_with_fallbacks(**request_kwargs) - return await self._schedule_factory( - model=model, - priority=request_priority, - original_function=self.async_function_with_fallbacks, - args=(), - kwargs=request_kwargs, - ) - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _atext_completion(self, model: str, prompt: str, **kwargs): - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _atext_completion()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": prompt}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) - - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("atext_completion")( - **{ - **data, - "prompt": prompt, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.atext_completion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.atext_completion(model=%s)\x1b[31m Exception %s\x1b[0m", model, e) - if model is not None: - self.fail_calls[model] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def aadapter_completion( - self, - adapter_id: str, - model: str, - is_retry: bool | None = False, - is_fallback: bool | None = False, - is_async: bool | None = False, - **kwargs, - ): - try: - kwargs["model"] = model - kwargs["adapter_id"] = adapter_id - kwargs["original_function"] = self._aadapter_completion - kwargs.setdefault("metadata", {}).update({"model_group": model}) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _aadapter_completion(self, adapter_id: str, model: str, **kwargs): - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _aadapter_completion()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=[{"role": "user", "content": "default text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, splits_stream_timeout=True) - - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("aadapter_completion")( - **{ - **data, - "adapter_id": adapter_id, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.aadapter_completion(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.aadapter_completion(model=%s)\x1b[31m Exception %s\x1b[0m", model, e) - if model is not None: - self.fail_calls[model] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def _asearch_with_fallbacks(self, original_function: Callable, **kwargs) -> object: - """ - Helper function to make a search API call through the router with load balancing and fallbacks. - Reuses the router's retry/fallback infrastructure. - """ - from litellm.router_utils.search_api_router import SearchAPIRouter - - return _OPAQUE_VALUE_ADAPTER.validate_python( - await SearchAPIRouter.async_search_with_fallbacks( - litellm_routing_ctx=self._core_ctx, - original_function=original_function, - **kwargs, - ) - ) - - async def _asearch_with_fallbacks_helper(self, 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. - """ - from litellm.router_utils.search_api_router import SearchAPIRouter - - return await SearchAPIRouter.async_search_with_fallbacks_helper( - litellm_routing_ctx=self._core_ctx, - model=model, - original_generic_function=original_generic_function, - **kwargs, - ) - - async def aguardrail( - self, - guardrail_name: str, - original_function: Callable, - **kwargs, - ): - """ - Execute a guardrail with load balancing and fallbacks. - - Args: - guardrail_name: Name of the guardrail to execute - original_function: The guardrail's execution function (e.g., async_pre_call_hook) - **kwargs: Additional arguments passed to the guardrail - - Returns: - Result from the guardrail execution - """ - kwargs["model"] = guardrail_name # For fallback system compatibility - kwargs["original_generic_function"] = original_function - kwargs["original_function"] = self._aguardrail_helper - self._update_kwargs_before_fallbacks( - model=guardrail_name, - kwargs=kwargs, - metadata_variable_name="litellm_metadata", - ) - verbose_router_logger.debug("Inside aguardrail() - guardrail_name: %s", guardrail_name) - response: Final = await self.async_function_with_fallbacks(**kwargs) - return response - - async def _aguardrail_helper( - self, - model: str, - original_generic_function: Callable, - **kwargs, - ): - """ - Helper for aguardrail - selects a guardrail deployment and executes it. - Called by async_function_with_fallbacks for each retry attempt. - - Args: - model: The guardrail_name (named 'model' for fallback system compatibility) - original_generic_function: The guardrail's execution function - **kwargs: Additional arguments - """ - guardrail_name: Final = model - selected_guardrail: Final = self._core_get_available_guardrail( - guardrail_name=guardrail_name, - ) - - verbose_router_logger.debug( - "Selected guardrail deployment: %s", selected_guardrail.get("litellm_params", {}).get("guardrail") - ) - - # Pass the selected guardrail config to the original function - kwargs["selected_guardrail"] = selected_guardrail - response: Final = await original_generic_function(**kwargs) - return response - - def _core_get_available_guardrail( - self, - guardrail_name: str, - ) -> "GuardrailTypedDict": - """ - Select a guardrail deployment using the router's load balancing strategy. - - Args: - guardrail_name: Name of the guardrail to select - - Returns: - Selected guardrail configuration dict - """ - from litellm.router_strategy.simple_shuffle import simple_shuffle - - healthy_deployments: Final = [g for g in self._core_guardrail_list if g.get("guardrail_name") == guardrail_name] - - if not healthy_deployments: - raise ValueError(f"No guardrail found with name: {guardrail_name}") - - if len(healthy_deployments) == 1: - return healthy_deployments[0] - - # Use simple_shuffle for weighted selection - return simple_shuffle( - resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=healthy_deployments, - model=guardrail_name, - request_kwargs=None, - ) - - async def _ageneric_api_call_with_fallbacks( - self, model: str, original_function: Callable, attempt_function: Callable | None = None, **kwargs - ): - """ - Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router - - attempt_function runs every attempt of the chain instead of the plain helper, so a streaming - endpoint can wrap each attempt's stream with its own mid-stream fallback handling. - """ - try: - kwargs["model"] = model - kwargs["original_generic_function"] = original_function - kwargs["original_function"] = attempt_function or self._ageneric_api_call_with_fallbacks_helper - if attempt_function is not None: - controls: Final = per_request_fallback_controls(kwargs) - kwargs[MID_STREAM_FALLBACK_CONTROLS_KEY] = controls # rebind-ok: forwarded to every hop - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs, metadata_variable_name="litellm_metadata") - verbose_router_logger.debug("Inside ageneric_api_call_with_fallbacks() - model: %s", model) - response: Final = await self.async_function_with_fallbacks(**kwargs) - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - def _add_deployment_model_to_endpoint_for_llm_passthrough_route( - self, kwargs: dict[str, Any], model: str, model_name: str - ) -> dict[str, Any]: - """ - Add the deployment model to the endpoint for LLM passthrough route. - - e.g for bedrock invoke users can pass endpoint as /model/special-bedrock-model/invoke - it should be actually sent as /model/us.anthropic.claude-3-5-sonnet-20240620-v1:0/invoke - """ - if "endpoint" in kwargs and kwargs["endpoint"]: - # For provider-specific endpoints, strip the provider prefix from model_name - # e.g., "bedrock/us.anthropic.claude-3-5-sonnet-20240620-v1:0" -> "us.anthropic.claude-3-5-sonnet-20240620-v1:0" - from litellm import get_llm_provider - - try: - # get_llm_provider returns (model_without_prefix, provider, api_key, api_base) - stripped_model_name, _, _, _ = get_llm_provider( - model=model_name, - custom_llm_provider=kwargs.get("custom_llm_provider"), - api_base=kwargs.get("api_base"), - ) - replacement_model_name = stripped_model_name - except Exception: - # If get_llm_provider fails, fall back to using model_name as-is - replacement_model_name = model_name - - kwargs["endpoint"] = replace_path_segment(kwargs["endpoint"], model, replacement_model_name) - return kwargs - - async def _ageneric_api_call_with_fallbacks_helper(self, model: str, original_generic_function: Callable, **kwargs): - """ - Helper function to make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router - """ - - passthrough_on_no_deployment: Final = kwargs.pop("passthrough_on_no_deployment", False) - function_name: Final = "_ageneric_api_call_with_fallbacks" - handler_name: Final[str | None] = getattr(original_generic_function, "__name__", None) - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - try: - with self._video_call_pin(handler_name, kwargs), self._route_capability_pin(handler_name): - deployment = ( # rebind-ok: set on success, see pre-init - await self._core_async_get_available_deployment( - model=model, - request_kwargs=kwargs, - messages=kwargs.get("messages", None), - input=kwargs.get("input", None), - specific_deployment=kwargs.pop("specific_deployment", None), - ) - ) - except Exception as e: - if passthrough_on_no_deployment: - return await original_generic_function(model=model, **kwargs) - raise e - - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name=function_name, - splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, - opens_responses_socket=handler_name == "_aresponses_websocket", - ) - - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - self.total_calls[model_name] += 1 - - self._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs=kwargs, model=model, model_name=model_name - ) - - custom_llm_provider: Final = provider_for_generic_call(data) - generic_function_name: Final[str] = getattr(original_generic_function, "__name__", "") - - response_kwargs: Final = { - **without_retrieve_stream_defaults(data, kwargs), - "caching": self._core_cache_responses, - **kwargs, - "model": model_name, - **_with_router_resolved_session_model(kwargs.get("session"), model_name), - **( - {"turn_admission": partial(self.async_admit_deployment_turn, model, deployment, kwargs)} - if generic_function_name == "_aresponses_websocket" - else {} - ), - } - # Only set custom_llm_provider if it's not None - if custom_llm_provider is not None: - response_kwargs["custom_llm_provider"] = custom_llm_provider - - compacted_input: Final = await compact_to_fit( - self._core_ctx, - deployment, - response_kwargs, - surface_for_call(generic_function_name), - ) - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await original_generic_function(**compacted_input) - - if self._should_raise_anthropic_refusal_error( - model=model, - original_generic_function=original_generic_function, - response=response, - kwargs=kwargs, - ): - from litellm.llms.anthropic.messages.native.utils import ( - safeguard_refusal_error, - ) - - refusal_details: Final = cast(dict, response["stop_details"]) # cast-ok: gate verified the shape - raise safeguard_refusal_error(model=model, stop_details=refusal_details) - - self.success_calls[model_name] += 1 - realtime_session_logging: Final = ( - _CALL_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj") - if generic_function_name == "_arealtime" - else None - ) - realtime_session_deployment_id: Final = ( - _MODEL_INFO_ADAPTER.validate_python(kwargs["model_info"]).get("id") - if isinstance(realtime_session_logging, LiteLLMLogging) - and realtime_session_logging.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True - and realtime_session_logging.model_call_details.get(REALTIME_SESSION_ENDED_ABNORMALLY_KEY) is not True - else None - ) - if realtime_session_deployment_id is not None: - increment_deployment_successes_for_current_minute( - ctx=self._core_ctx, deployment_id=str(realtime_session_deployment_id) - ) - verbose_router_logger.info("ageneric_api_call_with_fallbacks(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - return response - except Exception as e: - verbose_router_logger.info( - "ageneric_api_call_with_fallbacks(model=%s)\x1b[31m Exception %s\x1b[0m", model, e - ) - if model is not None: - self.fail_calls[model] += 1 - if deployment is not None: - self._set_deployment_num_retries_on_exception(e, deployment) - self._stamp_failed_deployment_id_with_effective_model_info(e, deployment, kwargs) - raise e - - async def _aresponses_with_streaming_fallbacks( - self, original_function: Callable, **kwargs: Any - ) -> Union["ResponsesAPIResponse", "BaseResponsesAPIStreamingIterator"]: - """ - _ageneric_api_call_with_fallbacks for the Responses API, with every attempt's stream - carrying its own mid-stream fallback handling - (see _ageneric_api_call_with_fallbacks_responses_attempt). - """ - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - attempt_function=self._ageneric_api_call_with_fallbacks_responses_attempt, - **kwargs, - ) - - async def _ageneric_api_call_with_fallbacks_responses_attempt( - self, - model: str, - original_generic_function: Callable, - **kwargs: object, # kwargs-ok: forwarded verbatim to the per-attempt helper, shape varies per call site - ) -> Union["ResponsesAPIResponse", "BaseResponsesAPIStreamingIterator"]: - """ - One attempt of the Responses API fallback chain. A streaming result is wrapped with - _aresponses_streaming_iterator over this attempt's own kwargs, so a fallback hop that - fails mid-stream resumes the original group's chain instead of re-raising; the name keeps - _get_router_metadata_variable_name resolving to litellm_metadata for every hop. - """ - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - ) - - controls: Final = kwargs.pop(MID_STREAM_FALLBACK_CONTROLS_KEY, None) - hop_kwargs: Final = mid_stream_fallback_hop_kwargs( - model=model, original_generic_function=original_generic_function, controls=controls, kwargs=kwargs - ) - shared_logging: Final = kwargs.get("litellm_logging_obj") - if isinstance(shared_logging, LiteLLMLogging): - shared_logging.carry_failed_attempt_usage() - response: Final = await self._ageneric_api_call_with_fallbacks_helper( - model=model, original_generic_function=original_generic_function, **kwargs - ) - carry_over_pre_routing_selection(live_kwargs=kwargs, snapshot=hop_kwargs) - carry_over_routed_deployment(live_kwargs=kwargs, snapshot=hop_kwargs) - if kwargs.get("stream") and isinstance(response, BaseResponsesAPIStreamingIterator): - return await self._aresponses_streaming_iterator(response=response, initial_kwargs=hop_kwargs) - return response - - async def _aanthropic_messages_streaming_iterator( - self, - response: AsyncIterator[bytes], - initial_kwargs: dict[str, object], # mutable-ok: mutated in-place before re-entering the fallback chain - ) -> AsyncIterator[bytes]: - """ - Wrap an anthropic_messages (/v1/messages) streaming response so a - mid-stream provider error triggers the Router's fallback chain - (parity with _acompletion_streaming_iterator for the - chat-completions path). See #24004. - - anthropic_messages goes through _ageneric_api_call_with_fallbacks - rather than _acompletion, so the returned byte iterator is never - wrapped by the chat-completions fallback handler. Two failure - shapes land here: - - the completion-bridge path (deployments with no native - /v1/messages endpoint, via - LiteLLMMessagesToCompletionTransformationHandler) already - raises MidStreamFallbackError out of its underlying - CustomStreamWrapper; this wrapper only needs to catch it. - - a native Anthropic/Bedrock passthrough never raises anything - for a provider SSE `event: error` frame (e.g. `overloaded_error`, - `internal_server_error`) - it is forwarded to the client as-is - - so this wrapper detects it via parse_anthropic_error_event and - raises MidStreamFallbackError itself. - - Only an error before any real content (a content_block_delta frame) - has reached the caller triggers a fallback attempt, mirroring the - restriction _acompletion_streaming_iterator applies: once generated - output has already reached the caller, retrying would start a - second, overlapping Anthropic message lifecycle on the same SSE - stream, so the error is left to propagate instead of being retried - invisibly. A non-retriable client error (4xx other than 429) is - never worth a fallback attempt either, so it is also left to - propagate. - - Lifecycle/bookkeeping frames (message_start, content_block_start, - ping, ...) do not by themselves disqualify a fallback attempt - - Anthropic routinely sends message_start before an overload error. - When a fallback can still take over they are BUFFERED rather than - forwarded immediately, since forwarding one and then appending a - fallback attempt's own message_start would produce two overlapping - message lifecycles on one SSE stream; a `ping` carries no lifecycle, - so it is forwarded live even while lifecycle frames sit buffered, - keeping the connection alive during a long thinking pass. Buffered - frames are flushed, in order, the moment real content arrives (the - primary attempt has committed by then anyway) or once the stream - ends without ever producing content or an error. When no fallback - can take over the request is already committed, so every frame, - including pings and provider error frames, is forwarded live and - verbatim instead. - """ - from litellm.llms.anthropic.messages.native.fake_stream_iterator import ( - FakeAnthropicMessagesStreamIterator, - ) - from litellm.llms.anthropic.messages.native.streaming_iterator import ( - aclose_if_supported, - aclose_superseded_attempt, - parse_anthropic_error_event, - parse_anthropic_refusal_stop_details, - ) - from litellm.llms.anthropic.messages.native.utils import ( - safeguard_refusal_error, - ) - - source_iterator: Final = response - replays_a_finished_agentic_turn: Final = isinstance(source_iterator, FakeAnthropicMessagesStreamIterator) - - async def close_source() -> None: - with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): - await aclose_if_supported(source_iterator) - - async def close_superseded() -> None: - with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): - await aclose_superseded_attempt(source_iterator) - - async def stream_with_fallbacks() -> AsyncGenerator[bytes, None]: - from litellm.exceptions import MidStreamFallbackError - - # Lifecycle/bookkeeping frames (message_start, content_block_start, - # ...) are held back rather than forwarded immediately, but only - # while a fallback can still take over: Anthropic routinely sends - # message_start before an overload error, and once a byte reaches - # the client a fallback attempt can only append its OWN - # message_start, producing two overlapping message lifecycles on - # one SSE stream. A `ping` keepalive carries no lifecycle, so it - # is forwarded live even behind buffered frames, keeping the - # connection alive through a long thinking pass. Buffered frames - # are flushed the moment real content (content_block_delta) - # arrives - at that point the primary attempt has committed and a - # clean retry is no longer possible anyway - or once the primary - # stream ends without ever producing content. Hitting - # MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS forces the same early - # commit as real content arriving, so a hostile or pathological - # upstream can't grow the buffer forever. With no fallback able - # to take over there is nothing to buffer for, so every frame, - # including pings and provider error frames, is forwarded live. - model: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group - has_generated_content = not ( # rebind-ok: set once real content is seen, the buffer cap is hit, or neither a retry nor a fallback can take over - self._anthropic_messages_stream_can_retry(initial_kwargs) - or self._stream_can_fall_back(model, initial_kwargs) - ) - buffered_lifecycle_chunks: tuple[bytes, ...] = () # rebind-ok: flushed once committed or on decline - # An error this loop raises off a frame, not the provider's. - raised_here: BaseException | None = None # rebind-ok: set when the loop raises off a frame - try: - async for chunk in source_iterator: - if _anthropic_stream_forwards_ping_live(chunk, has_generated_content): - yield chunk - continue - if _stream_commits_now(chunk, has_generated_content, len(buffered_lifecycle_chunks)): - has_generated_content = True - # A transport can split one SSE data line across byte chunks, so pre-content - # detection parses the accumulated buffer plus the current chunk, never the - # chunk alone; the buffer is already capped, which bounds this window too. - parse_window = ( - b"".join(c for c in (*buffered_lifecycle_chunks, chunk) if isinstance(c, (bytes, bytearray))) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime - if not has_generated_content and isinstance(chunk, (bytes, bytearray)) # pyright: ignore[reportUnnecessaryIsInstance] # bridge-path chunks are not always bytes at runtime - else chunk - ) - error_event = parse_anthropic_error_event(parse_window) - recoverable_frame_error = self._anthropic_messages_recoverable_frame_error( - error_event, chunk, has_generated_content, model, initial_kwargs - ) - refusal_stop_details = ( - parse_anthropic_refusal_stop_details(parse_window) - if not has_generated_content and error_event is None - else None - ) - if ( - refusal_stop_details is not None - and not replays_a_finished_agentic_turn - and self._refusal_fallback_available(model, initial_kwargs) - ): - refusal_error = safeguard_refusal_error(model=model, stop_details=refusal_stop_details) - raised_here = MidStreamFallbackError( - message=refusal_error.message, - model=model, - llm_provider="anthropic", - original_exception=refusal_error, - is_pre_first_chunk=True, - ) - raise raised_here - if not has_generated_content and error_event is None: - buffered_lifecycle_chunks = (*buffered_lifecycle_chunks, chunk) - continue - if recoverable_frame_error is not None: - assert error_event is not None - raised_here = MidStreamFallbackError( - message=error_event[1], - model=model, - llm_provider="anthropic", - original_exception=recoverable_frame_error, - is_pre_first_chunk=True, - ) - raise raised_here - for buffered_chunk in buffered_lifecycle_chunks: - yield buffered_chunk - buffered_lifecycle_chunks = () - yield chunk - for buffered_chunk in buffered_lifecycle_chunks: - yield buffered_chunk - except Exception as stream_error: # noqa: BLE001 # any raised provider error must reach the fallback gate - if stream_error is not raised_here: - await self._areport_raised_messages_stream_failure(stream_error, initial_kwargs, source_iterator) - async for item in self._aanthropic_messages_recover_stream_error( - stream_error, - has_generated_content, - buffered_lifecycle_chunks, - model, - initial_kwargs, - wrapper, - close_superseded, - ): - yield item - finally: - await close_source() - - # Referenced by stream_with_fallbacks via closure - assigned here, before - # the generator body ever runs, so the reference resolves fine despite - # being defined textually after the function that captures it. - wrapper: Final = FallbackAwareAnthropicMessagesStream(stream_with_fallbacks(), source_iterator) - return wrapper - - async def _areport_raised_messages_stream_failure( - self, stream_error: Exception, attempt_kwargs: Mapping[str, object], source: object - ) -> None: - """A /v1/messages stream whose provider iterator raised, counted against the deployment - that served it and cooled down by the rules, the same as a non-stream call's failure, so a - retry picks another deployment. A stream that logs itself (litellm's own stream logging - reaches the router's failure callback), or a failure the request's logging object already - logged, is not counted twice.""" - from litellm.exceptions import MidStreamFallbackError - from litellm.litellm_core_utils.litellm_logging import Logging - - if isinstance(getattr(source, "logging_obj", None), Logging): - return - logging_obj: Final = attempt_kwargs.get("litellm_logging_obj") - if isinstance(logging_obj, Logging) and logging_obj.model_call_details.get("has_logged_async_failure") is True: - return - bucket: Final = as_str_mapping(attempt_kwargs.get(get_metadata_variable_name_from_kwargs(attempt_kwargs))) or {} - model_info: Final = as_str_mapping(bucket.get("model_info")) or {} - deployment_id: Final = model_info.get("id") - if deployment_id is None: - return - deployment: Final = self.get_deployment(model_id=str(deployment_id)) - deployment_params: Final = deployment.litellm_params.model_dump(exclude_none=True) if deployment else {} - provider_error: Final = ( - stream_error.original_exception - if isinstance(stream_error, MidStreamFallbackError) and stream_error.original_exception is not None - else stream_error - ) - now: Final = datetime.now() # noqa: DTZ005 # naive, as the logging pipeline's times - logged: Final = { - "call_type": "anthropic_messages", - "exception": provider_error, - "litellm_params": { - **{ - name: deployment_params[name] for name in ("api_base", "cooldown_time") if name in deployment_params - }, - "metadata": dict(bucket), - "model_info": dict(model_info), - }, - } - try: - _ = self.deployment_callback_on_failure(logged, None, now, now) # pyright: ignore[reportUnknownMemberType] # an untyped callback - await self.async_deployment_callback_on_failure(logged, None, now, now) # pyright: ignore[reportUnknownMemberType] # see above - except Exception as error: # noqa: BLE001 # a recording never fails the stream - verbose_router_logger.warning("a stream's failure was not recorded: %s", error) - - async def _aanthropic_messages_recover_stream_error( - self, - stream_error: Exception, - has_generated_content: bool, - buffered_lifecycle_chunks: tuple[bytes, ...], - model: str, - initial_kwargs: dict[str, object], # mutable-ok: _aanthropic_messages_fallback_attempt mutates it - wrapper: "FallbackAwareAnthropicMessagesStream", - close_superseded: Callable[[], Awaitable[None]], - ) -> AsyncGenerator[bytes, None]: - """Turns a source-iterator failure into a fallback attempt or the error reaching the caller.""" - from litellm.exceptions import MidStreamFallbackError - - if isinstance(stream_error, MidStreamFallbackError) and _anthropic_stream_should_decline_fallback( - has_generated_content, stream_error - ): - for buffered_chunk in buffered_lifecycle_chunks: - yield buffered_chunk - if stream_error.original_exception is not None: - raise stream_error.original_exception from stream_error - raise stream_error - fallback_error: Final = ( - stream_error - if isinstance(stream_error, MidStreamFallbackError) - else _stream_fallback_error_for_raised(stream_error, model, has_generated_content) - ) - if fallback_error is None: - raise stream_error - await close_superseded() - outcome: Final = await self._astream_retry_same_group( - fallback_error, - initial_kwargs, - attempt_function=self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt, - ) - match outcome: - case _StreamRetryOpened(response=retried, attempted_retries=attempted, max_retries=budget): - async for item in self._aanthropic_messages_yield_recovered(retried, wrapper, (attempted, budget)): - yield item - case _StreamRetriesExhausted(error=last_error): - async for item in self._aanthropic_messages_fallback_attempt(last_error, initial_kwargs, wrapper): - yield item - case _: - assert_never(outcome) - - def _anthropic_messages_group_retry_policy(self, kwargs: Mapping[str, object]) -> dict[str, RetryPolicy] | None: - configured: Final = kwargs.get("model_group_retry_policy", self.model_group_retry_policy) - return cast("dict[str, RetryPolicy] | None", configured) # cast-ok: same type as the router attribute - - def _anthropic_messages_resolved_retry_policy(self, kwargs: Mapping[str, object]) -> RetryPolicy | None: - """ - The retry policy for this request's model group, unless the request opted out with num_retries=0. - A policy that does not resolve to a RetryPolicy governs nothing here, so the stream runs as it would - with none; the pre-stream retry loop still reports the malformed policy when an attempt fails. - """ - if kwargs.get("num_retries") == 0: - return None - model_group: Final = _request_model_group(kwargs) - try: - return _RESOLVED_RETRY_POLICY_ADAPTER.validate_python( - resolve_retry_policy( - retry_policy=self.retry_policy, - model_group=model_group, - model_group_retry_policy=self._anthropic_messages_group_retry_policy(kwargs), - ) - ) - except (TypeError, ValidationError) as malformed: - verbose_router_logger.warning( - "The retry policy for %s is not a RetryPolicy, streaming without one: %s", model_group, malformed - ) - return None - - def _anthropic_messages_plain_retry_budget(self, kwargs: Mapping[str, object]) -> int: - """ - The same precedence async_function_with_retries resolves for a failure raised before the - stream opened: the request's num_retries, then the routed deployment's, then the router's. - """ - request_num_retries: Final = kwargs.get("num_retries") - if isinstance(request_num_retries, int): - return request_num_retries - deployment_id: Final = routed_deployment_id(kwargs) - deployment: Final = self._core_get_deployment(deployment_id) if deployment_id is not None else None - deployment_num_retries: Final = _deployment_num_retries(deployment) - if deployment_num_retries is not None: - return deployment_num_retries - return self.num_retries if self.num_retries is not None else 0 - - def _anthropic_messages_policy_retries(self, trigger: Exception, kwargs: Mapping[str, object]) -> int | None: - policy: Final = self._anthropic_messages_resolved_retry_policy(kwargs) - if policy is None: - return None - return _get_num_retries_from_retry_policy(exception=trigger, retry_policy=policy) - - def _stream_retry_budget(self, trigger: Exception, kwargs: Mapping[str, object]) -> tuple[int, bool]: - """ - The budget an earlier retry of this request committed to, else the retry policy's grant when one - names this error, else the plain budget, with whether a policy governs the retry: a committed budget - is kept whichever deployment the retry lands on, as async_function_with_retries keeps its own. - """ - policy_retries: Final = self._anthropic_messages_policy_retries(trigger, kwargs) - committed_budget: Final = committed_retry_budget_for_request(kwargs) - if committed_budget is not None: - return committed_budget, policy_retries is not None - if policy_retries is None: - return self._anthropic_messages_plain_retry_budget(kwargs), False - return policy_retries, True - - def _anthropic_messages_stream_can_retry(self, kwargs: Mapping[str, object]) -> bool: - """ - Whether a pre-content failure of this stream would be retried within its own model group, - the other case where holding lifecycle frames back from the client buys a clean restart. - A retry policy names its budget per error class, so the largest budget it names bounds the - hold: holding frames one attempt too long is safe, forwarding them before a retry is not. - """ - attempted: Final = attempted_retries_for_request(kwargs) - committed_budget: Final = committed_retry_budget_for_request(kwargs) - if committed_budget is not None: - return committed_budget > attempted - plain_budget: Final = self._anthropic_messages_plain_retry_budget(kwargs) - policy: Final = self._anthropic_messages_resolved_retry_policy(kwargs) - ceiling: Final = plain_budget if policy is None else max(plain_budget, _retry_policy_ceiling(policy)) - return ceiling > attempted - - def _anthropic_messages_recoverable_frame_error( - self, - error_event: tuple[str, str, int] | None, - chunk: object, - has_generated_content: bool, - model_group: str, - kwargs: Mapping[str, object], - ) -> Exception | None: - """ - The exception a provider `event: error` frame before content recovers through when a retry of its - class or a fallback can still take over. A frame nothing can take over for (content already out, a - gateway verdict, or a class granted no retry with no fallback) reaches the client as the provider - sent it, the way the last exhausted attempt's does. - """ - if has_generated_content or error_event is None: - return None - error_type, message, status_code = error_event - if not _is_retriable_anthropic_status(status_code) or _anthropic_stream_error_is_gateway_verdict(chunk): - return None - frame_error: Final = anthropic_error_frame_exception(error_type, message, status_code, model_group) - budget, _ = self._stream_retry_budget(frame_error, kwargs) - if budget > attempted_retries_for_request(kwargs): - return frame_error - if self._stream_can_fall_back(model_group, kwargs): - return frame_error - return None - - def _stream_should_retry( - self, - trigger: Exception, - healthy_deployments: list[dict], # mutable-ok: should_retry_this_error's own parameter type - all_deployments: list[dict], # mutable-ok: should_retry_this_error's own parameter type - kwargs: Mapping[str, object], - ) -> bool: - try: - self.should_retry_this_error( - error=trigger, - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - context_window_fallbacks=_request_fallback_list( - kwargs, - "context_window_fallbacks", - cast( # cast-ok: untyped router attribute - "list[object] | None", self._core_context_window_fallbacks - ), - ), - content_policy_fallbacks=_request_fallback_list( - kwargs, - "content_policy_fallbacks", - cast( # cast-ok: untyped router attribute - "list[object] | None", self._core_content_policy_fallbacks - ), - ), - regular_fallbacks=_request_fallback_list( - kwargs, - "fallbacks", - cast("list[object] | None", self._core_fallbacks), # cast-ok: untyped router attribute - ), - ) - except Exception: # noqa: BLE001 # should_retry_this_error declines by raising the error it was given - return False - return True - - async def _astream_retry_same_group( - self, - e: "MidStreamFallbackError", - initial_kwargs: Mapping[str, object], - attempt_function: "Callable[..., Awaitable[object]]", - ) -> _StreamRetryOutcome: - """ - Re-runs the attempt within the request's own model group, the way async_function_with_retries - would have for a failure raised before the stream opened, until a retry opens a stream or the - budget runs out. Each retry's stream carries its own wrapper with the remaining budget, so a - retry that drops before content again continues the same count instead of starting over. - `attempt_function` is the request's own attempt (the /v1/messages attempt, or a chat completions one). - """ - model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: kwargs always carries the model group - retry_kwargs: Final = mid_stream_retry_kwargs(initial_kwargs) - handler: Final = retry_kwargs.get("original_generic_function") - healthy_deployments, all_deployments = await self._async_get_healthy_deployments( - model=model_group, - parent_otel_span=get_parent_otel_span_from_kwargs(retry_kwargs), - request_kwargs=share_holder_request_kwargs(initial_kwargs, self._ptu_share_teams()), - ) - budget, policy_applies = self._stream_retry_budget(_mid_stream_retry_trigger(e), initial_kwargs) - last_error = e # rebind-ok: the newest failure is what the fallback chain and the caller see - for attempt in range(attempted_retries_for_request(initial_kwargs), budget): - trigger = _mid_stream_retry_trigger(last_error) - if is_agentic_followup_failure(trigger): - return _StreamRetriesExhausted(last_error) - if not policy_applies and not self._stream_should_retry( - trigger, healthy_deployments, all_deployments, initial_kwargs - ): - return _StreamRetriesExhausted(last_error) - if get_metadata_variable_name_from_kwargs(retry_kwargs) in retry_kwargs: - self._core_log_retry(kwargs=retry_kwargs, e=trigger) - await self._host.sleep( - self._time_to_sleep_before_retry( - e=trigger, - remaining_retries=budget - attempt, - num_retries=budget, - healthy_deployments=healthy_deployments, - all_deployments=all_deployments, - ) - ) - record_retry_attempt(retry_kwargs, attempted_retries=attempt + 1, max_retries=budget) - verbose_router_logger.debug( - "Retrying stream that failed before content, attempt %s of %s", attempt + 1, budget - ) - carry_superseded_attempt_usage(retry_kwargs) - try: - with endpoint_selection( - getattr(handler, "__name__", None), getattr(handler, "__module__", None), model=model_group - ): - response = await attempt_function(**retry_kwargs) - except Exception as retry_error: # noqa: BLE001 # the fallback chain judges the failed retry - wrapped = _stream_fallback_error_for_raised(retry_error, model_group, False, e.llm_provider) - if wrapped is None: - return _StreamRetriesExhausted(_stream_pre_content_error(retry_error, model_group, e.llm_provider)) - last_error = wrapped - continue - return _StreamRetryOpened(response, attempted_retries=attempt + 1, max_retries=budget) - return _StreamRetriesExhausted(last_error) - - async def _aanthropic_messages_yield_recovered( - self, - recovered: object, - wrapper: "FallbackAwareAnthropicMessagesStream", - retry_counters: tuple[int, int] | None = None, - ) -> AsyncGenerator[bytes, None]: - """ - Hands a retry's or a fallback's response to the client through the wrapper, closing it afterwards. - A retry stamps the retry headers async_function_with_retries would have for a pre-stream retry; - a later hop the recovered stream makes replaces them with its own, the way a fallback's do. - """ - from litellm.llms.anthropic.messages.native.streaming_iterator import ( - aclose_if_supported, - anthropic_messages_response_as_sse_events, - ) - - hidden_params, headers = PythonCore._prepare_fallback_hidden_params(recovered) - wrapper.merge_fallback_hidden_params(hidden_params, headers) - wrapper.adopt_fallback_source(recovered) - if retry_counters is not None: - add_retry_headers_to_response( - response=wrapper, attempted_retries=retry_counters[0], max_retries=retry_counters[1] - ) - try: - if hasattr(recovered, "__aiter__"): - async for item in cast("AsyncIterator[bytes]", recovered): # cast-ok: __aiter__ checked above - wrapper.follow_source_attribution() - yield item - return - # A recovery can resolve to a complete AnthropicMessagesResponse - # dict even for a streaming request (e.g. an agentic tool-use - # interception loop) - yielding it as-is would put a raw dict - # into a byte stream, so it's synthesized into the SSE - # lifecycle a real stream would have sent instead. - for event in anthropic_messages_response_as_sse_events( - cast("AnthropicMessagesResponse", recovered) # cast-ok: non-streaming shape by elimination - ): - yield event - finally: - with anyio.CancelScope(shield=True), contextlib.suppress(BaseException): - await aclose_if_supported(recovered) - - async def _aanthropic_messages_fallback_attempt( - self, - e: "MidStreamFallbackError", - initial_kwargs: dict[str, Any], # mutable-ok: mutated in-place before re-entering the fallback chain - wrapper: "FallbackAwareAnthropicMessagesStream", - ) -> AsyncGenerator[bytes, None]: - """ - Re-enters the Router's fallback chain for a mid-stream - anthropic_messages error and yields whatever the fallback attempt - produces. Split out of _aanthropic_messages_streaming_iterator to - keep each function's cyclomatic complexity within the repo's C901 - budget. - """ - from litellm.exceptions import MidStreamFallbackError - - try: - model_group: Final = cast(str, initial_kwargs.get("model")) # cast-ok: model group - fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the common_utils list|None param - "fallbacks", self._core_fallbacks - ) - context_window_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below - "context_window_fallbacks", self._core_context_window_fallbacks - ) - content_policy_fallbacks: Final[list | None] = initial_kwargs.get( # mutable-ok: matches the param below - "content_policy_fallbacks", self._core_content_policy_fallbacks - ) - initial_kwargs["original_function"] = self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt - self._update_kwargs_before_fallbacks( - model=model_group, - kwargs=initial_kwargs, - metadata_variable_name="litellm_metadata", - ) - # The content-policy and context-window dispatch branches match on the trigger's own type, so - # such an error's MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. - fallback_trigger: Final[Exception] = ( - e.original_exception - if isinstance( - e.original_exception, - (litellm_exceptions.ContentPolicyViolationError, litellm_exceptions.ContextWindowExceededError), - ) - else e - ) - fallback_response: Final = _OPAQUE_VALUE_ADAPTER.validate_python( - await self.async_function_with_fallbacks_common_utils( - e=fallback_trigger, - disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - model_group=model_group, - args=(), - kwargs=initial_kwargs, - include_fallback_errors=initial_kwargs.get("include_fallback_errors", False) is True, - ) - ) - async for fallback_item in self._aanthropic_messages_yield_recovered(fallback_response, wrapper): - yield fallback_item - except Exception as fallback_error: - verbose_router_logger.error("Anthropic messages streaming fallback also failed: %s", fallback_error) - if isinstance(fallback_error, MidStreamFallbackError) and fallback_error.original_exception is not None: - raise fallback_error.original_exception from fallback_error - raise - - async def _aanthropic_messages_with_streaming_fallbacks( - self, - original_function: Callable, - **kwargs: object, # kwargs-ok: forwarded verbatim to original_function, shape varies per call site - ) -> Union["AnthropicMessagesResponse", AsyncIterator[bytes]]: - """ - _ageneric_api_call_with_fallbacks for anthropic_messages, with every attempt's stream - carrying its own mid-stream fallback handling - (see _ageneric_api_call_with_fallbacks_anthropic_messages_attempt). Parity with - _aresponses_with_streaming_fallbacks for the Responses API. - """ - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - attempt_function=self._ageneric_api_call_with_fallbacks_anthropic_messages_attempt, - **kwargs, - ) - - async def _ageneric_api_call_with_fallbacks_anthropic_messages_attempt( - self, - model: str, - original_generic_function: Callable[..., Awaitable[object]], - **kwargs: object, # kwargs-ok: forwarded verbatim to the per-attempt helper, shape varies per call site - ) -> Union["AnthropicMessagesResponse", AsyncIterator[bytes]]: - """ - One attempt of the anthropic_messages fallback chain. A streaming result is wrapped with - _aanthropic_messages_streaming_iterator over this attempt's own kwargs, so a fallback hop - that fails mid-stream resumes the original group's chain instead of re-raising; the name - keeps _get_router_metadata_variable_name resolving to litellm_metadata for every hop. - """ - controls: Final = kwargs.pop(MID_STREAM_FALLBACK_CONTROLS_KEY, None) - hop_kwargs: Final = mid_stream_fallback_hop_kwargs( - model=model, original_generic_function=original_generic_function, controls=controls, kwargs=kwargs - ) - response: Final = await self._ageneric_api_call_with_fallbacks_helper( - model=model, original_generic_function=original_generic_function, **kwargs - ) - carry_over_pre_routing_selection(live_kwargs=kwargs, snapshot=hop_kwargs) - carry_over_routed_deployment(live_kwargs=kwargs, snapshot=hop_kwargs) - if kwargs.get("stream") and hasattr(response, "__aiter__"): - return await self._aanthropic_messages_streaming_iterator( - response=cast("AsyncIterator[bytes]", response), # cast-ok: stream=True always returns a byte iterator - initial_kwargs=hop_kwargs, - ) - return response - - async def _dispatch_generic_call_type( - self, - call_type: str, - original_function: Callable, - **kwargs: object, # kwargs-ok: forwarded verbatim to the per-call-type helper, shape varies per call site - ): - """ - factory_function's shared dispatch for call types with no - call-specific handling, except anthropic_messages: kept out of - factory_function's own async_wrapper (already at the repo's C901 - complexity ceiling) so routing its mid-stream fallback handling - (#24004) doesn't add another branch there. - """ - if call_type == "anthropic_messages": - return await self._aanthropic_messages_with_streaming_fallbacks( - original_function=original_function, **kwargs - ) - return await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) - - @staticmethod - @contextlib.contextmanager - def _video_call_pin(handler_name: str | None, request_kwargs: Mapping[str, object]) -> Generator[None]: - """ - Scopes the deployment a video id names to the routing of one call on that video (status, content, remix, - edit, extension). Any other call routes as if no id were sent, whatever its body carries. - """ - video_id: Final = request_kwargs.get("video_id") - deployment_id: Final = ( - decode_video_id_with_provider(video_id).get("deployment_id") - if handler_name in VIDEO_ID_CALL_TYPES and isinstance(video_id, str) - else None - ) - token: Final = _video_call_deployment_id.set(deployment_id) - try: - yield - finally: - _video_call_deployment_id.reset(token) - - @staticmethod - @contextlib.contextmanager - def _route_capability_pin(handler_name: str | None) -> Generator[None]: - token: Final = _routed_call_type.set(handler_name) - try: - yield - finally: - _routed_call_type.reset(token) - - @staticmethod - def _narrow_to_video_deployment(healthy_deployments: Sequence[_DeploymentT]) -> Sequence[_DeploymentT]: - """ - A video lives on the provider account of the deployment that created it, so a call on an id naming a - deployment that every routing filter kept goes to that deployment, else to whichever one routing picks. - """ - deployment_id: Final = _video_call_deployment_id.get() - creating: Final = tuple( - deployment - for deployment in healthy_deployments - if deployment_id in PythonCore._deployment_ids((deployment,)) - ) - return creating or healthy_deployments - - def _generic_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs): - """ - Make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router - Args: - model: The model to use - original_function: The handler function to call (e.g., litellm.completion) - **kwargs: Additional arguments to pass to the handler function - Returns: - The response from the handler function - """ - handler_name: Final = original_function.__name__ - metadata_variable_name: Final = get_router_metadata_variable_name(function_name="generic_api_call") - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): - try: - verbose_router_logger.debug("Inside _generic_api_call() - handler: %s, model: %s", handler_name, model) - self._update_kwargs_before_fallbacks( - model=model, - kwargs=kwargs, - metadata_variable_name=metadata_variable_name, - ) - with ( - self._video_call_pin(handler_name, kwargs), - self._route_capability_pin(handler_name), - picks_reserve(True), - ): - deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - messages=kwargs.get("messages", None), - input=kwargs.get("input", None), - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="generic_api_call", - splits_stream_timeout=handler_name in _STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, - ) - - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - - self.total_calls[model_name] += 1 - - # For passthrough routes, use the actual model from deployment - # and swap model name in endpoint if present - if "endpoint" in kwargs and kwargs["endpoint"]: - kwargs["endpoint"] = kwargs["endpoint"].replace(model, model_name) - kwargs["model"] = model_name - - # Perform pre-call checks for routing strategy - self.routing_strategy_pre_call_checks(deployment=deployment) - - custom_llm_provider: Final = provider_for_generic_call(data) - - response: Final = original_function( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", handler_name, model_name) - return response - except Exception as e: - verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", handler_name, model, e) - if model is not None: - self.fail_calls[model] += 1 - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - def embedding( - self, - model: str, - input: str | list, - is_async: bool | None = False, - **kwargs, - ) -> EmbeddingResponse: - try: - kwargs["model"] = model - kwargs["input"] = input - kwargs["original_function"] = self._embedding - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = self.function_with_fallbacks(**kwargs) - return response - except Exception as e: - raise e - - def _embedding(self, input: str | list[object], model: str, **kwargs): - model_name = None - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside embedding()- model: %s", model) - with embedding_deployment_selection(), picks_reserve(True): - deployment = self.get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - input=input, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - model_name = data["model"] - batch_error: Final = embedding_batch_too_large_error( - deployment=_OPAQUE_VALUE_ADAPTER.validate_python(deployment), - input=input, - request_kwargs=_OPAQUE_VALUE_ADAPTER.validate_python(kwargs), - model=model, - ) - logging_obj: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj") - if batch_error is not None: - if isinstance(logging_obj, LiteLLMLogging): - log_rejection_before_provider_call(logging_obj, batch_error) - raise batch_error - if isinstance(logging_obj, LiteLLMLogging): - clear_rejected_before_provider_call(logging_obj) - - potential_model_client: Final = self._get_client(deployment=deployment, kwargs=kwargs, client_type="sync") - # check if provided keys == client keys # - dynamic_api_key: Final = kwargs.get("api_key", None) - if ( - dynamic_api_key is not None - and potential_model_client is not None - and dynamic_api_key != potential_model_client.api_key - ): - model_client = None - else: - model_client = potential_model_client - - self.total_calls[model_name] += 1 - - ### DEPLOYMENT-SPECIFIC PRE-CALL CHECKS ### (e.g. update rpm pre-call. Raise error, if deployment over limit) - self.routing_strategy_pre_call_checks(deployment=deployment) - - response: Final = self._host.invoke("embedding")( - **{ - **data, - "input": input, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.embedding(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.embedding(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - async def aembedding( - self, - model: str, - input: str | list, - is_async: bool | None = True, - **kwargs, - ) -> EmbeddingResponse: - try: - kwargs["model"] = model - kwargs["input"] = input - kwargs["original_function"] = self._aembedding - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = await self.async_function_with_fallbacks(**kwargs) - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _aembedding(self, input: str | list[object], model: str, **kwargs): - model_name = None - deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked - try: - verbose_router_logger.debug("Inside _aembedding()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - with embedding_deployment_selection(): - deployment = await self._core_async_get_available_deployment( # rebind-ok: set on success, see pre-init - model=model, - input=input, - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - data: Final = deployment["litellm_params"].copy() - model_name = data["model"] - batch_error: Final = embedding_batch_too_large_error( - deployment=_OPAQUE_VALUE_ADAPTER.validate_python(deployment), - input=input, - request_kwargs=_OPAQUE_VALUE_ADAPTER.validate_python(kwargs), - model=model, - ) - logging_obj: Final = _REQUEST_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj") - if batch_error is not None: - if isinstance(logging_obj, LiteLLMLogging): - await async_log_rejection_before_provider_call(logging_obj, batch_error) - raise batch_error - if isinstance(logging_obj, LiteLLMLogging): - clear_rejected_before_provider_call(logging_obj) - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - - self.total_calls[model_name] += 1 - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("aembedding")( - **{ - **data, - "input": input, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.aembedding(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - return response - except Exception as e: - verbose_router_logger.info("litellm.aembedding(model=%s)\x1b[31m Exception %s\x1b[0m", model_name, e) - if model_name is not None: - self.fail_calls[model_name] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - self._stamp_failed_attempt( - e, _REQUEST_KWARGS_ADAPTER.validate_python(kwargs), picked=deployment is not None - ) - raise e - - #### FILES API #### - async def acreate_file( - self, - model: str, - **kwargs, - ) -> OpenAIFileObject: - try: - kwargs["model"] = model - kwargs["original_function"] = self._acreate_file - self._update_kwargs_before_fallbacks(model=model, kwargs=kwargs) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _acreate_file( - self, - model: str, - **kwargs, - ) -> OpenAIFileObject: - try: - from litellm.litellm_core_utils.routing_common_utils import add_model_file_id_mappings - - verbose_router_logger.debug("Inside _atext_completion()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - healthy_deployments: Final = await self._core_async_get_healthy_deployments( - model=model, - messages=[{"role": "user", "content": "files-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - parent_otel_span=parent_otel_span, - ) - - async def create_file_for_deployment(deployment: dict) -> OpenAIFileObject: - from litellm.litellm_core_utils.core_helpers import safe_deep_copy - - kwargs_copy: Final = safe_deep_copy(kwargs) - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs_copy, - function_name="acreate_file", - ) - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs_copy, - ) - self.total_calls[model_name] += 1 - - ## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ## - # For DB/config deployments, use provider from deployment params - custom_llm_provider = data.get("custom_llm_provider") - stripped_model, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - # Preserve explicitly stored provider, fallback to inferred - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - ## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ## - purpose: Final = cast(OpenAIFilesPurpose | None, kwargs.get("purpose")) - file = cast(FileTypes | None, kwargs.get("file")) - if not file or not purpose: - raise Exception("file and file_purpose are required for create_file") - - replace_model_in_jsonl_bool: Final = should_replace_model_in_jsonl( - purpose=purpose, - passthrough=kwargs.get("passthrough") is True, - ) - if replace_model_in_jsonl_bool: - file = replace_model_in_jsonl( - file_content=file, - new_model_name=stripped_model, - keep_line_numbers=custom_llm_provider == LlmProviders.VERTEX_AI.value, - ) - - kwargs_copy["file"] = file - if custom_llm_provider == LlmProviders.LITELLM_PROXY.value: - kwargs_copy["extra_body"] = MappingProxyType( - { - **(kwargs_copy.get("extra_body") or MappingProxyType({})), - "target_model_names": stripped_model, - } - ) - if ( - "gcs_bucket_name" in data - ): # TODO: Remove this once we have a better way to handle GCS bucket name: Problem is that we need to pass the gcs_bucket_name to the router for the create_file call but it doesn't show up there - kwargs_copy.setdefault("litellm_metadata", {})["gcs_bucket_name"] = data["gcs_bucket_name"] - async with self._deployment_slot( - deployment=deployment, kwargs=kwargs_copy, parent_otel_span=parent_otel_span - ): - response = await self._host.invoke("acreate_file")( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs_copy, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.acreate_file(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - return response - - tasks: Final = [] - - if isinstance(healthy_deployments, dict): - tasks.append(create_file_for_deployment(healthy_deployments)) - else: - for deployment in healthy_deployments: - tasks.append(create_file_for_deployment(deployment)) - - responses: Final = await asyncio.gather(*tasks) - - if len(responses) == 0: - raise Exception("No healthy deployments found.") - - model_file_id_mapping: Final = add_model_file_id_mappings( - healthy_deployments=healthy_deployments, responses=responses - ) - returned_response: Final = cast(OpenAIFileObject, responses[0]) - returned_response_hidden_params: Final = get_hidden_params(returned_response) - if returned_response_hidden_params is not None: - returned_response_hidden_params["model_file_id_mapping"] = model_file_id_mapping - return returned_response - except Exception as e: - verbose_router_logger.exception( - "litellm.acreate_file(model=%s, %s)\x1b[31m Exception %s\x1b[0m", model, kwargs, e - ) - if model is not None: - self.fail_calls[model] += 1 - raise e - - #### VECTOR STORES API #### - async def avector_store_create( - self, - model: str | None, - **kwargs, - ): - """ - Create a vector store for a specific model. - - Args: - model: Model name from router config - **kwargs: Vector store creation parameters - - Returns: - VectorStoreCreateResponse - """ - with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): - try: - # If model is None, use the factory function approach (direct SDK call) - if model is None: - from litellm.vector_stores.main import acreate - - # Use the factory function to handle the call - factory_fn: Final = self.factory_function(acreate, call_type="avector_store_create") - return await factory_fn(**kwargs) - - from litellm.vector_stores import acreate as avector_store_create_sdk - - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "vector-store-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="avector_store_create", - ) - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - # Get custom provider from deployment params - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - async with self._deployment_slot( - deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span - ): - response = await avector_store_create_sdk( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.avector_store_create(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - return response - except Exception as e: - verbose_router_logger.exception( - "litellm.avector_store_create(model=%s)\x1b[31m Exception %s\x1b[0m", model, e - ) - if model is not None: - self.fail_calls[model] += 1 - raise e - - def _override_vector_store_methods_for_router(self): - """ - Override factory-generated vector store methods with router-aware implementations. - This is called after _initialize_vector_store_endpoints() to ensure our custom - methods that handle deployment selection and credential injection are used instead - of the generic factory-generated ones. - """ - # Store references to the custom methods defined above - # These methods handle proper routing through deployments - # The methods are already defined as instance methods above - - async def acreate_batch( - self, - model: str, - **kwargs, - ) -> LiteLLMBatch: - try: - kwargs["model"] = model - kwargs["original_function"] = self._acreate_batch - metadata_variable_name: Final = get_router_metadata_variable_name(function_name="_acreate_batch") - self._update_kwargs_before_fallbacks( - model=model, - kwargs=kwargs, - metadata_variable_name=metadata_variable_name, - ) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _acreate_batch( - self, - model: str, - **kwargs, - ) -> LiteLLMBatch: - try: - verbose_router_logger.debug("Inside _acreate_batch()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - with batch_create_routing(): - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "files-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - - data: Final = deployment["litellm_params"].copy() - model_name: Final = data["model"] - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="_acreate_batch") - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - ## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ## - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("acreate_batch")( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.acreate_batch(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - return response - except Exception as e: - verbose_router_logger.exception( - "litellm._acreate_batch(model=%s, %s)\x1b[31m Exception %s\x1b[0m", model, kwargs, e - ) - if model is not None: - self.fail_calls[model] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - raise e - - async def aretrieve_batch( - self, - model: str | None = None, - **kwargs, - ) -> LiteLLMBatch: - """ - Iterate through all models in a model group to check for batch - - Future Improvement - cache the result. - """ - try: - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - requested_model_group: Final = model - metadata_variable_name: Final = get_router_metadata_variable_name(function_name="aretrieve_batch") - if model is not None: - filtered_model_list: ( - list[DeploymentTypedDict] | list[dict] | dict | None - ) = await self._core_async_get_healthy_deployments( - model=model, - messages=[{"role": "user", "content": "retrieve-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - parent_otel_span=parent_otel_span, - ) - else: - filtered_model_list = self._core_get_model_list() - if filtered_model_list is None: - raise Exception("Router not yet initialized.") - - receieved_exceptions: Final = [] - - async def try_retrieve_batch(model_name: DeploymentTypedDict): - try: - from litellm.litellm_core_utils.core_helpers import safe_deep_copy - - model: Final = model_name["litellm_params"].get("model") - data: Final = model_name["litellm_params"].copy() - custom_llm_provider = data.get("custom_llm_provider") - if model is None: - raise Exception(f"Model not found in litellm_params for deployment: {model_name}") - # Update kwargs with the current model name or any other model-specific adjustments - ## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ## - if not custom_llm_provider: - _, custom_llm_provider, _, _ = get_llm_provider(model=model) - new_kwargs: Final = safe_deep_copy(kwargs) - self._update_kwargs_with_deployment( - deployment=cast(dict, model_name), - kwargs=new_kwargs, - function_name="aretrieve_batch", - ) - model_group: Final = requested_model_group or model_name["model_name"] - if not new_kwargs[metadata_variable_name].get("model_group"): - new_kwargs[metadata_variable_name]["model_group"] = model_group - new_kwargs.pop("custom_llm_provider", None) - data.pop("custom_llm_provider", None) - return await self._host.invoke("aretrieve_batch")( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - **new_kwargs, - }, - ) - except Exception as e: - import traceback - - traceback.print_exc() - receieved_exceptions.append(e) - return None - - # Check all models in parallel - if ( - filtered_model_list is not None - and isinstance(filtered_model_list, list) - and len(filtered_model_list) > 0 - ): - results = await asyncio.gather( - *[try_retrieve_batch(cast(DeploymentTypedDict, model)) for model in filtered_model_list], - return_exceptions=True, - ) - elif filtered_model_list is not None and isinstance(filtered_model_list, dict): - results = await try_retrieve_batch(cast(DeploymentTypedDict, filtered_model_list)) - else: - raise Exception("No healthy deployments found.") - - # Check for successful responses and handle exceptions - if results is not None: - if isinstance(results, LiteLLMBatch): - return results - elif isinstance(results, list): - for result in results: - if isinstance(result, LiteLLMBatch): - return result - - # If no valid Batch response was found, raise the first encountered exception - if receieved_exceptions: - raise receieved_exceptions[0] # Raising the first exception encountered - - # If no exceptions were encountered, raise a generic exception - raise Exception(f"Unable to find batch in any model. Received errors - {receieved_exceptions}") - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def acancel_batch( - self, - model: str, - **kwargs, - ) -> LiteLLMBatch: - """ - Cancel a batch through the router with proper model-to-provider mapping. - """ - try: - kwargs["model"] = model - kwargs["original_function"] = self._acancel_batch - metadata_variable_name: Final = get_router_metadata_variable_name(function_name="_acancel_batch") - self._update_kwargs_before_fallbacks( - model=model, - kwargs=kwargs, - metadata_variable_name=metadata_variable_name, - ) - response: Final = await self.async_function_with_fallbacks(**kwargs) - - return response - except Exception as e: - asyncio.create_task( - send_llm_exception_alert( - ctx=self._core_ctx, - request_kwargs=kwargs, - error_traceback_str=traceback.format_exc(), - original_exception=e, - ) - ) - raise e - - async def _acancel_batch( - self, - model: str, - **kwargs, - ) -> LiteLLMBatch: - try: - verbose_router_logger.debug("Inside _acancel_batch()- model: %s", model) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - with picks_reserve(False): - deployment: Final = await self._core_async_get_available_deployment( - model=model, - messages=[{"role": "user", "content": "batch-api-fake-text"}], - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) - - selected_deployment_id: Final = (deployment.get("model_info") or {}).get("id") - data: Final = deployment["litellm_params"].copy() - resolved_credentials: Final = self._core_get_deployment_credentials_with_provider( - model_id=selected_deployment_id or model - ) - if resolved_credentials is not None: - data.update(resolved_credentials) - data.pop("litellm_credential_name", None) - model_name: Final = data["model"] - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="_acancel_batch") - - model_client: Final = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - ## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ## - custom_llm_provider = data.get("custom_llm_provider") - _, inferred_custom_llm_provider, _, _ = get_llm_provider( - model=data["model"], - custom_llm_provider=custom_llm_provider, - ) - custom_llm_provider = custom_llm_provider or inferred_custom_llm_provider - - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): - response = await self._host.invoke("acancel_batch")( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self._core_cache_responses, - "client": model_client, - **kwargs, - } - ) - - self.success_calls[model_name] += 1 - verbose_router_logger.info("litellm.acancel_batch(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) - - return response - except Exception as e: - verbose_router_logger.exception( - "litellm._acancel_batch(model=%s, %s)\x1b[31m Exception %s\x1b[0m", model, kwargs, e - ) - if model is not None: - self.fail_calls[model] += 1 - self._stamp_retry_skip_deployment_id(e, kwargs) - raise e - - async def alist_batches( - self, - model: str, - **kwargs, - ): - """ - Return all the batches across all deployments of a model group. - """ - - filtered_model_list: Final = self._core_get_model_list(model_name=model) - if filtered_model_list is None: - raise Exception("Router not yet initialized.") - - async def try_retrieve_batch(model: DeploymentTypedDict): - try: - # Update kwargs with the current model name or any other model-specific adjustments - return await self._host.invoke("alist_batches")(**{**model["litellm_params"], **kwargs}) - except Exception: - return None - - # Check all models in parallel - results: Final = await asyncio.gather(*[try_retrieve_batch(model) for model in filtered_model_list]) - - final_results: Final[dict] = { - "object": "list", - "data": [], - "first_id": None, - "last_id": None, - "has_more": False, - } - - for result in results: - if result is not None: - ## check batch id - if final_results["first_id"] is None and hasattr(result, "first_id"): - final_results["first_id"] = getattr(result, "first_id") - final_results["last_id"] = getattr(result, "last_id") - final_results["data"].extend(result.data) - - ## check 'has_more' - if getattr(result, "has_more", False) is True: - final_results["has_more"] = True - - return final_results - - #### PASSTHROUGH API #### - - async def _pass_through_moderation_endpoint_factory( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - **kwargs, - ): - # update kwargs with model_group - self._update_kwargs_before_fallbacks( - model=kwargs.get("model", ""), - kwargs=kwargs, - ) - if kwargs.get("model") and self._core_get_model_list(model_name=kwargs["model"]): - deployment: Final = await self._core_async_get_available_deployment( - model=kwargs["model"], - request_kwargs=kwargs, - ) - kwargs["model"] = deployment["litellm_params"]["model"] - data: Final = deployment["litellm_params"].copy() - self._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - ) - kwargs.update(data) - - return await original_function(**kwargs) - - def factory_function( - self, - original_function: Callable, - call_type: Literal[ - "assistants", - "moderation", - "anthropic_messages", - "aresponses", - "acancel_responses", - "acompact_responses", - "responses", - "aget_responses", - "adelete_responses", - "afile_delete", - "afile_content", - "_arealtime", - "acreate_realtime_client_secret", - "arealtime_calls", - "acreate_realtime_transcription_session", - "acreate_realtime_translation_client_secret", - "arealtime_translation_calls", - "_aresponses_websocket", - "acreate_fine_tuning_job", - "acancel_fine_tuning_job", - "alist_fine_tuning_jobs", - "aretrieve_fine_tuning_job", - "alist_files", - "aimage_edit", - "allm_passthrough_route", - "alist_input_items", - "agenerate_content", - "generate_content", - "agenerate_content_stream", - "generate_content_stream", - "avector_store_search", - "avector_store_create", - "avector_store_retrieve", - "avector_store_list", - "avector_store_update", - "avector_store_delete", - "avector_store_file_create", - "avector_store_file_list", - "avector_store_file_retrieve", - "avector_store_file_content", - "avector_store_file_update", - "avector_store_file_delete", - "vector_store_search", - "vector_store_create", - "vector_store_retrieve", - "vector_store_list", - "vector_store_update", - "vector_store_delete", - "vector_store_file_create", - "vector_store_file_list", - "vector_store_file_retrieve", - "vector_store_file_content", - "vector_store_file_update", - "vector_store_file_delete", - "aocr", - "ocr", - "asearch", - "search", - "adecisions", - "decisions", - "asystemone", - "systemone", - "aextraction", - "extraction", - "aadapter_generate_content", - "avideo_generation", - "video_generation", - "avideo_list", - "video_list", - "avideo_status", - "video_status", - "avideo_content", - "video_content", - "avideo_remix", - "video_remix", - "avideo_create_character", - "video_create_character", - "avideo_get_character", - "video_get_character", - "avideo_edit", - "video_edit", - "avideo_extension", - "video_extension", - "acreate_container", - "create_container", - "alist_containers", - "list_containers", - "aretrieve_container", - "retrieve_container", - "adelete_container", - "delete_container", - "aupload_container_file", - "upload_container_file", - "alist_container_files", - "list_container_files", - "aretrieve_container_file", - "retrieve_container_file", - "adelete_container_file", - "delete_container_file", - "acreate_skill", - "alist_skills", - "aget_skill", - "adelete_skill", - "acreate_interaction", - "create_interaction", - "aget_interaction", - "get_interaction", - "adelete_interaction", - "delete_interaction", - "acancel_interaction", - "cancel_interaction", - "acreate_agent", - "create_agent", - "alist_agents", - "list_agents", - "aget_agent", - "get_agent", - "adelete_agent", - "delete_agent", - "alist_agent_versions", - "list_agent_versions", - ] = "assistants", - ): - """ - Creates appropriate wrapper functions for different API call types. - - Returns: - - A synchronous function for synchronous call types - - An asynchronous function for asynchronous call types - """ - # Handle synchronous call types - if call_type in ( - "responses", - "generate_content", - "generate_content_stream", - "ocr", - "search", - "decisions", - "systemone", - "extraction", - "video_generation", - "video_list", - "video_status", - "video_content", - "video_remix", - "create_container", - "list_containers", - "retrieve_container", - "delete_container", - ): - - def sync_wrapper( - custom_llm_provider: str | None = None, - client: object | None = None, - **kwargs, - ): - return self._generic_api_call_with_fallbacks(original_function=original_function, **kwargs) - - return sync_wrapper - - if call_type in ( - "vector_store_search", - "vector_store_create", - "vector_store_retrieve", - "vector_store_list", - "vector_store_update", - "vector_store_delete", - ): - - def vector_store_sync_wrapper( - custom_llm_provider: str | None = None, - client: object | None = None, - **kwargs, - ): - provider_kwargs: Final = ( - MappingProxyType({**kwargs, "custom_llm_provider": custom_llm_provider}) - if custom_llm_provider and "custom_llm_provider" not in kwargs - else MappingProxyType(kwargs) - ) - search_kwargs: Final = ( - MappingProxyType( - { - **provider_kwargs, - "_direct_vector_store_embedding_executor": RouterVectorStoreEmbeddingExecutor( - router=self._core_public(), - metadata=self._vector_store_request_metadata(kwargs), - ), - } - ) - if call_type == "vector_store_search" - else provider_kwargs - ) - if search_kwargs.get("model"): - return self._generic_api_call_with_fallbacks(original_function=original_function, **search_kwargs) - if call_type == "vector_store_search": - return original_function(**MappingProxyType({**search_kwargs, "router": self._core_public()})) - return original_function(**search_kwargs) - - return vector_store_sync_wrapper - - if call_type in ( - "vector_store_file_create", - "vector_store_file_list", - "vector_store_file_retrieve", - "vector_store_file_content", - "vector_store_file_update", - "vector_store_file_delete", - ): - - def vector_store_file_sync_wrapper( - custom_llm_provider: str | None = None, - client: object | None = None, - **kwargs, - ): - return original_function( - custom_llm_provider=custom_llm_provider, - client=client, - **kwargs, - ) - - return vector_store_file_sync_wrapper - - if call_type in ( - "create_agent", - "list_agents", - "get_agent", - "delete_agent", - "list_agent_versions", - ): - - def managed_agents_sync_wrapper( - custom_llm_provider: str | None = None, - client: object | None = None, - **kwargs, - ): - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - if "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = "gemini" - return original_function(**kwargs) - - return managed_agents_sync_wrapper - - # Handle asynchronous call types - async def async_wrapper( - custom_llm_provider: str | None = None, - client: AsyncOpenAI | None = None, - **kwargs, - ): - if call_type == "assistants": - return await self._pass_through_assistants_endpoint_factory( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - client=client, - **kwargs, - ) - elif call_type == "moderation": - return await self._pass_through_moderation_endpoint_factory( - original_function=original_function, **kwargs - ) - elif call_type in ("asearch", "search"): - return await self._asearch_with_fallbacks( - original_function=original_function, - **kwargs, - ) - elif call_type in ( - "avector_store_file_create", - "avector_store_file_list", - "avector_store_file_retrieve", - "avector_store_file_content", - "avector_store_file_update", - "avector_store_file_delete", - ): - return await self._init_vector_store_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - elif call_type == "aresponses": - return await self._aresponses_with_streaming_fallbacks( - original_function=original_function, - **kwargs, - ) - elif call_type in ( - "acreate_realtime_client_secret", - "arealtime_calls", - "acreate_realtime_transcription_session", - "acreate_realtime_translation_client_secret", - "arealtime_translation_calls", - ): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - client=client, - **kwargs, - ) - elif call_type in ( - "anthropic_messages", - "_arealtime", - "_aresponses_websocket", - "acreate_fine_tuning_job", - "acancel_fine_tuning_job", - "alist_fine_tuning_jobs", - "aretrieve_fine_tuning_job", - "alist_files", - "aimage_edit", - "agenerate_content", - "agenerate_content_stream", - "aocr", - "ocr", - "adecisions", - "asystemone", - "aextraction", - "avideo_generation", - "avideo_list", - "avideo_status", - "avideo_content", - "avideo_remix", - "avideo_create_character", - "avideo_get_character", - "avideo_edit", - "avideo_extension", - "acreate_skill", - "alist_skills", - "aget_skill", - "adelete_skill", - ): - return await self._dispatch_generic_call_type( - call_type=call_type, - original_function=original_function, - **kwargs, - ) - elif call_type in ( - "acreate_container", - "alist_containers", - "aretrieve_container", - "adelete_container", - "aupload_container_file", - "alist_container_files", - "aretrieve_container_file", - "adelete_container_file", - "aretrieve_container_file_content", - ): - return await self._init_containers_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - elif call_type == "allm_passthrough_route": - if client: - kwargs["client"] = client - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - passthrough_on_no_deployment=True, - **kwargs, - ) - elif call_type in ( - "aget_responses", - "acancel_responses", - "acompact_responses", - "adelete_responses", - "alist_input_items", - ): - return await self._init_responses_api_endpoints( - original_function=original_function, - **kwargs, - ) - elif call_type in ( - "avector_store_search", - "avector_store_create", - "avector_store_retrieve", - "avector_store_list", - "avector_store_update", - "avector_store_delete", - ): - vector_store_kwargs: Final = ( - { - **kwargs, - "_direct_vector_store_embedding_executor": RouterVectorStoreEmbeddingExecutor( - router=self._core_public(), - metadata=self._vector_store_request_metadata(kwargs), - ), - } - if call_type == "avector_store_search" - else kwargs - ) - return await self._init_vector_store_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - call_type=call_type, - **vector_store_kwargs, - ) - elif call_type in ("afile_delete", "afile_content"): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - client=client, - **kwargs, - ) - elif call_type in ( - "acreate_interaction", - "create_interaction", - "aget_interaction", - "adelete_interaction", - "acancel_interaction", - ): - return await self._init_interactions_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - elif call_type in ( - "acreate_agent", - "alist_agents", - "aget_agent", - "adelete_agent", - "alist_agent_versions", - ): - return await self._init_managed_agents_api_endpoints( - original_function=original_function, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) - - return async_wrapper - - @staticmethod - def _vector_store_request_metadata(kwargs: Mapping[str, object]) -> Mapping[str, object]: - return vector_store_request_metadata(kwargs) - - async def _init_vector_store_api_endpoints( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - call_type: str | None = None, - **kwargs, - ): - """ - Initialize the Vector Store API endpoints on the router. - - If a model is provided in kwargs, use model-based routing to get - the deployment credentials. Otherwise, call the original function directly. - """ - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - - # If model is provided, use generic API call with fallbacks for proper routing - if kwargs.get("model"): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - **kwargs, - ) - - # For search, pass the router so provider transforms can resolve - # router-managed embedding models (e.g. S3 Vectors query embeddings). - # The merge also overrides any client-supplied `router` key. - if call_type == "avector_store_search": - search_kwargs: Final = MappingProxyType({**kwargs, "router": self._core_public()}) - return await original_function(**search_kwargs) - - # Otherwise, call the original function directly - return await original_function(**kwargs) - - async def _init_containers_api_endpoints( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - **kwargs, - ): - """ - Initialize the Containers API endpoints on the router. - - LiteLLM-managed container IDs (``cntr_...``) encode ``model_id`` and provider - metadata. When present, decode the ID, replace ``container_id`` with the - upstream value, and route through ``_ageneric_api_call_with_fallbacks`` so - deployment credentials (e.g. regional ``api_base`` for Azure) match - :meth:`_init_responses_api_endpoints`. Create/list calls carry no container ID, so - they route through the deployment named by ``model`` when the caller passes one, - falling back to the direct call when no deployment matches. Otherwise call the - handler directly with global provider credentials. - """ - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - - from litellm.responses.utils import ResponsesAPIRequestUtils - - container_id: Final = kwargs.get("container_id") - _forwarded_model_id: Final = kwargs.get("model_id") - if isinstance(container_id, str): - decoded: Final = ResponsesAPIRequestUtils.decode_container_id(container_id) - original_id: Final = decoded.get("response_id", container_id) - if original_id != container_id: - kwargs["container_id"] = original_id - decoded_provider: Final = decoded.get("custom_llm_provider") - if decoded_provider and kwargs.get("custom_llm_provider") == "openai": - kwargs["custom_llm_provider"] = decoded_provider - # Fall back to the model_id forwarded by the proxy when the container_id - # is a native upstream ID (e.g. Azure hex cntr_) that carries no LiteLLM - # routing payload, so deployment credentials (api_base, api_key) are applied. - model_id: Final = decoded.get("model_id") or ( - _forwarded_model_id.strip() - if isinstance(_forwarded_model_id, str) and _forwarded_model_id.strip() - else None - ) - if model_id: - kwargs["model"] = model_id - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - **kwargs, - ) - - requested_model: Final = kwargs.get("model") - if isinstance(requested_model, str) and requested_model.strip(): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - passthrough_on_no_deployment=True, - **kwargs, - ) - - return await original_function(**kwargs) - - async def _init_responses_api_endpoints( - self, - original_function: Callable, - **kwargs, - ): - """ - Initialize the Responses API endpoints on the router. - - GET, DELETE, CANCEL Responses API Requests encode the model_id in the response_id, this function decodes the response_id and sets the model to the model_id. - """ - from litellm.responses.utils import ResponsesAPIRequestUtils - - model_id: Final = ResponsesAPIRequestUtils.get_model_id_from_response_id(kwargs.get("response_id")) - if model_id is not None: - kwargs["model"] = model_id - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - **kwargs, - ) - - async def _init_interactions_api_endpoints( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - **kwargs, - ): - """ - Initialize the Interactions API endpoints on the router. - - GET, DELETE, CANCEL Interactions API Requests don't need model-based routing, - so we call the original function directly with the custom_llm_provider. - """ - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - # Default to gemini for interactions API - if "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = "gemini" - # If the proxy accidentally passed agent name as model, clear it - if kwargs.get("agent") and kwargs.get("model") == kwargs.get("agent"): - kwargs["model"] = None - # Model-based interactions use deployment routing + fallbacks; agent-only calls - # must not enter model-group lookup (agent name is not a LiteLLM deployment). - if kwargs.get("model"): - return await self._ageneric_api_call_with_fallbacks( - original_function=original_function, - **kwargs, - ) - return await original_function(**kwargs) - - async def _init_managed_agents_api_endpoints( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - **kwargs, - ): - """ - Initialize the Managed Agents API endpoints on the router (v1beta/agents). - - CRUD operations for Gemini managed agents don't need model-based routing, - so we call the original function directly with the custom_llm_provider. - """ - if custom_llm_provider and "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = custom_llm_provider - if "custom_llm_provider" not in kwargs: - kwargs["custom_llm_provider"] = "gemini" - return await original_function(**kwargs) - - async def _pass_through_assistants_endpoint_factory( - self, - original_function: Callable, - custom_llm_provider: str | None = None, - client: AsyncOpenAI | None = None, - **kwargs, - ): - """Internal helper function to pass through the assistants endpoint""" - if custom_llm_provider is None: - if self.assistants_config is not None: - custom_llm_provider = self.assistants_config["custom_llm_provider"] - kwargs.update(self.assistants_config["litellm_params"]) - else: - raise Exception( - "'custom_llm_provider' must be set. Either via:\n `Router(assistants_config={'custom_llm_provider': ..})` \nor\n `router.arun_thread(custom_llm_provider=..)`" - ) - return await original_function(custom_llm_provider=custom_llm_provider, client=client, **kwargs) - - #### [END] ASSISTANTS API #### - - async def _maybe_run_weighted_failover( - self, - exception: Exception, - original_model_group: str, - all_deployments: Sequence[DeploymentTypedDict], - args: tuple, - kwargs: dict, - input_kwargs: dict, - ) -> Any | None: - """Same-model-group retry after a failed deployment; returns None if not applicable.""" - strategy, _ = self._get_routing_context(original_model_group, kwargs) - if strategy != "simple-shuffle": - return None - - failed_id: Final[str | None] = getattr(exception, "failed_deployment_id", None) - if not failed_id: - return None - - metadata_variable_name: Final = self._get_metadata_variable_name_from_kwargs(kwargs) - meta = kwargs.get(metadata_variable_name) - if meta is None: - meta = {} - kwargs[metadata_variable_name] = meta - if not isinstance(meta, dict): - return None - configured_failed_id: Final[str | None] = getattr(exception, "failed_configured_deployment_id", None) - prev_excluded: Final = set(meta.get("_failover_excluded_ids") or []) - excluded: Final = prev_excluded | {configured_failed_id or failed_id} - - all_ids: Final = { - (d.get("model_info") or {}).get("id") - for d in all_deployments - if (d.get("model_info") or {}).get("id") is not None - } - # Only consider deployments that are currently healthy (not in cooldown). - # Using all_ids here would cause a wasteful run_async_fallback invocation - # that fails with RouterRateLimitError whenever the "remaining" entries - # are all in cooldown — the inner async_get_healthy_deployments call - # would find an empty list and raise immediately. - cooldown_ids = set(await async_get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=None)) - remaining: Final = (all_ids - cooldown_ids) - excluded - if not remaining: - return None - - verbose_router_logger.debug( - "Weighted failover: exclude=%r, remaining=%s for model_group=%r", - excluded, - len(remaining), - original_model_group, - ) - - meta["_failover_excluded_ids"] = list(excluded) - - entry: Final = { - "model": original_model_group, - "_excluded_deployment_ids": list(excluded), - } - # Build a local copy so the weighted-failover keys do not leak back to - # the caller's shared kwargs dict (any downstream fallback path reads - # the same dict and must not inherit our `_excluded_deployment_ids` - # entry). - failover_kwargs: Final = { - **input_kwargs, - "fallback_model_group": [entry], - "original_model_group": original_model_group, - } - try: - return await run_async_fallback(*args, **failover_kwargs) - except (openai.APIError, RouterRateLimitError, RouterRateLimitErrorBasic): - # Expected model-level failure on the retried deployment. All - # litellm provider errors derive from openai.APIError; if every - # remaining deployment in the group is in cooldown the router - # raises RouterRateLimitError (a ValueError, not an APIError). - # In either case defer to the regular fallback path. Programming - # errors (AttributeError, KeyError, TypeError, etc.) intentionally - # propagate so they remain visible. - return None - - async def async_function_with_fallbacks_common_utils( - self, - e: Exception, - disable_fallbacks: bool | None, - fallbacks: list | None, - context_window_fallbacks: list | None, - content_policy_fallbacks: list | None, - model_group: str | None, - args: tuple, - kwargs: dict, - include_fallback_errors: bool = False, - ): - """ - Common utilities for async_function_with_fallbacks - """ - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("Traceback%s", redact_string(traceback.format_exc())) - original_exception: Final = e - fallback_model_group = None - original_model_group: Final[str | None] = kwargs.get("model") - fallback_walk: Final = attempted_fallback_targets_from(kwargs) # pyright: ignore[reportUnknownArgumentType] # keep legacy kwargs untyped - # A pre-routing hook (complexity / auto / adaptive / quality routers) picks a tier - # behind the router name, and fallbacks are configured per tier, not per router. - lookup_groups: Final[tuple[str, ...]] = fallback_lookup_groups(kwargs, model_group) - fallback_failure_exception_str = "" - no_fallback_group_explained = False - hop_depth: Final = kwargs.get("fallback_depth") - nested_fallback_hop: Final = isinstance(hop_depth, int) and hop_depth > 0 - - unsupported_params_without_fallback: Final = is_bad_request_without_fallback( - e, self.fallback_on_unsupported_params - ) - if ( - disable_fallbacks is True - or original_model_group is None - or is_guardrail_intervention(e) - or unsupported_params_without_fallback - ): - raise e - - input_kwargs: Final = { - "litellm_routing_ctx": self._core_ctx, - "original_exception": original_exception, - **kwargs, - "attempted_targets": fallback_walk, - } - - if "max_fallbacks" not in input_kwargs: - input_kwargs["max_fallbacks"] = self.max_fallbacks - if "fallback_depth" not in input_kwargs: - input_kwargs["fallback_depth"] = 0 - if include_fallback_errors: - input_kwargs["include_fallback_errors"] = True - - context_window_fallback_model_group: Final = ( - self._get_fallback_model_group_for_lookup_groups( - fallbacks=context_window_fallbacks, - lookup_groups=lookup_groups, - ) - if isinstance(e, litellm_exceptions.ContextWindowExceededError) and context_window_fallbacks is not None - else None - ) - _skip_order_fallback: Final = ( - isinstance(e, litellm_exceptions.ContentPolicyViolationError) - or context_window_fallback_model_group is not None - ) - _request_team_id: Final[str | None] = (kwargs.get("metadata", {}) or {}).get("user_api_key_team_id") - # Use wildcard-aware lookup so order-based fallback also works for model - # groups resolved via pattern routing (e.g. `openai/*` -> `openai/gpt-4.1-mini`). - order_model_group: Final = get_pre_routing_selection(kwargs) or original_model_group - all_deployments: Final = self._core_get_model_list(model_name=order_model_group, team_id=_request_team_id) or () - _order_set: Final[set] = { - litellm_utils.get_deployment_order(d) - for d in all_deployments - if litellm_utils.get_deployment_order(d) is not None - } - order_values: Final[list] = sorted(_order_set) - if len(order_values) > 1 and not _skip_order_fallback: - # Determine which order levels have already been tried - current_target: Final = kwargs.get("_target_order") - skip_up_to: Final = current_target if current_target is not None else order_values[0] - # Build order-based fallback entries (skip already-tried levels) - order_fallback_entries: Final[list] = [ - {"model": order_model_group, "_target_order": o} for o in order_values if o > skip_up_to - ] - # Get external fallbacks — handle both standard and non-standard formats - external_fallback_group: list | None = None - if fallbacks is not None and lookup_groups: - if check_non_standard_fallback_format(fallbacks=fallbacks): - # Non-standard formats (e.g. ["claude-3-haiku"] or - # [{"model": "...", "messages": [...]}]) are passed through directly - external_fallback_group = fallbacks - else: - external_fallback_group, generic_idx = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, - lookup_groups=lookup_groups, - ) - if external_fallback_group is None and generic_idx is not None: - external_fallback_group = fallbacks[generic_idx]["*"] - - # Combined list: order fallbacks first, then external - combined_fallbacks: Final = order_fallback_entries + (external_fallback_group or []) - - if combined_fallbacks: - input_kwargs.update( - { - "fallback_model_group": combined_fallbacks, - "original_model_group": original_model_group, - } - ) - try: - return await run_async_fallback( - *args, - **input_kwargs, - ) - except Exception: - if isinstance(original_exception, litellm_exceptions.ContextWindowExceededError): - raise original_exception - raise - - # Weighted intra-group failover (simple-shuffle only); see _maybe_run_weighted_failover. - if ( - self.enable_weighted_failover - and not _skip_order_fallback - and not isinstance(e, litellm_exceptions.ContextWindowExceededError) - and original_model_group is not None - ): - response = await self._maybe_run_weighted_failover( - exception=e, - original_model_group=original_model_group, - all_deployments=all_deployments, - args=args, - kwargs=kwargs, - input_kwargs=input_kwargs, - ) - if response is not None: - return response - - try: - verbose_router_logger.info("Trying to fallback b/w models") - - # check if client-side fallbacks are used (e.g. fallbacks = ["gpt-3.5-turbo", "claude-3-haiku"] or fallbacks=[{"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hey, how's it going?"}]}] - is_non_standard_fallback_format: Final = check_non_standard_fallback_format(fallbacks=fallbacks) - - if is_non_standard_fallback_format and context_window_fallback_model_group is None: - input_kwargs.update( - { - "fallback_model_group": fallbacks, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - - return response - - if isinstance(e, litellm_exceptions.ContextWindowExceededError): - if context_window_fallback_model_group is not None: - input_kwargs.update( - { - "fallback_model_group": context_window_fallback_model_group, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - return response - - elif context_window_fallbacks is None: - error_message = f"model={model_group}. context_window_fallbacks={mask_sensitive_structure(context_window_fallbacks)}. fallbacks={mask_sensitive_structure(fallbacks)}.\n\nSet 'context_window_fallback' - https://docs.litellm.ai/docs/routing#fallbacks" - verbose_router_logger.info( - msg=f"Got 'ContextWindowExceededError'. No context_window_fallback set. Defaulting \ - to fallbacks, if available.{error_message}" - ) - - if self._host.dynamic().expose_router_debug_in_errors: - e.message += f"\n{error_message}" - elif isinstance(e, litellm_exceptions.ContentPolicyViolationError): - if content_policy_fallbacks is not None: - content_policy_fallback_model_group: Final[list[str] | None] = ( - self._get_fallback_model_group_for_lookup_groups( - fallbacks=content_policy_fallbacks, - lookup_groups=lookup_groups, - ) - ) - if content_policy_fallback_model_group is None: - raise original_exception - - input_kwargs.update( - { - "fallback_model_group": content_policy_fallback_model_group, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - return response - else: - error_message = format_content_policy_fallback_hint( - model_group, content_policy_fallbacks, fallbacks - ) - verbose_router_logger.info( - msg=f"Got 'ContentPolicyViolationError'. No content_policy_fallback set. Defaulting \ - to fallbacks, if available.{error_message}" - ) - - if self._host.dynamic().expose_router_debug_in_errors: - e.message += f"\n{error_message}" - if fallbacks is not None and lookup_groups: - verbose_router_logger.debug("inside model fallbacks: %s", mask_sensitive_structure(fallbacks)) - ( - fallback_model_group, - generic_fallback_idx, - ) = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, # if fallbacks = [{"gpt-3.5-turbo": ["claude-3-haiku"]}] - lookup_groups=lookup_groups, - ) - ## if none, check for generic fallback - if fallback_model_group is None and generic_fallback_idx is not None: - fallback_model_group = fallbacks[generic_fallback_idx]["*"] - - if fallback_model_group is None: - masked_fallbacks: Final = mask_sensitive_structure(fallbacks) - verbose_router_logger.info( - "No fallback model group found for lookup_groups=%s. Fallbacks=%s", - " -> ".join(lookup_groups), - masked_fallbacks, - ) - if ( - hasattr(original_exception, "message") - and self._host.dynamic().expose_router_debug_in_errors - and not nested_fallback_hop - ): - original_exception.message += format_no_fallback_group_message(lookup_groups, fallbacks) - no_fallback_group_explained = True - raise original_exception - - input_kwargs.update( - { - "fallback_model_group": fallback_model_group, - "original_model_group": original_model_group, - } - ) - - response = await run_async_fallback( - *args, - **input_kwargs, - ) - - return response - except Exception as new_exception: - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - fallback_failure_exception_str = truncate_fallback_error_detail(redact_string(str(new_exception))) - cooldown_info: Final = await async_get_cooldown_deployments_with_debug_info( - ctx=self._core_ctx, - parent_otel_span=parent_otel_span, - ) - verbose_router_logger.error( - "litellm.router.py::async_function_with_fallbacks() - " - "Error occurred while trying to do fallbacks - %s\n" - "Debug Information:\nCooldown Deployments=%s", - fallback_failure_exception_str, - cooldown_info, - ) - - attempted_fallback_group: Final = input_kwargs.get("fallback_model_group") - if ( - hasattr(original_exception, "message") - and self._host.dynamic().expose_router_debug_in_errors - and not no_fallback_group_explained - and not nested_fallback_hop - ): - original_exception.message += format_fallback_outcome_message( - model_group, - attempted_fallback_group, - fallback_failure_exception_str, - skipped=fallback_walk.skipped, - ) - - raise original_exception - - @tracer.wrap() - async def async_function_with_fallbacks(self, *args, **kwargs): - """ - Try calling the function_with_retries - If it fails after num_retries, fall back to another model group - """ - model_group: Final[str | None] = kwargs.get("model") - compaction_surface: Final = surface_for_call( - getattr(kwargs.get("original_generic_function") or kwargs.get("original_function"), "__name__", "") - ) - if compaction_surface is not None: - kwargs["_context_compaction_state"] = initialize_compaction_state(kwargs, compaction_surface) - clear_pre_routing_selection(kwargs) # pyright: ignore[reportUnknownArgumentType] # **kwargs is untyped at this boundary - if not isinstance(kwargs.get("attempted_targets"), AttemptedFallbackTargets): - _fallback_metadata_key: Final = get_router_metadata_variable_name( - function_name=getattr(kwargs.get("original_function"), "__name__", None) - ) - _sibling_metadata_key: Final = ( - "metadata" if _fallback_metadata_key == "litellm_metadata" else "litellm_metadata" - ) - if isinstance(_sibling_metadata := kwargs.get(_sibling_metadata_key), dict): - # In place, like every other router bucket write: downstream resolves the bucket by - # key presence, so rebinding kwargs to a copy detaches the proxy's request_data write-backs - _sibling_metadata.pop("attempted_fallbacks", None) - _sibling_metadata.pop("original_model_group", None) - if isinstance(_fallback_metadata := kwargs.get(_fallback_metadata_key), dict): - _fallback_metadata["attempted_fallbacks"] = 0 - if model_group is not None: - _fallback_metadata["original_model_group"] = model_group - include_fallback_errors: Final = kwargs.get("include_fallback_errors", False) is True - disable_fallbacks: Final[bool | None] = kwargs.pop("disable_fallbacks", False) - record_disable_fallbacks(kwargs, disable_fallbacks is True) - fallbacks: Final[list | None] = kwargs.get("fallbacks", self._core_fallbacks) - context_window_fallbacks: list | None = kwargs.get( - "context_window_fallbacks", self._core_context_window_fallbacks - ) - content_policy_fallbacks: list | None = kwargs.get( - "content_policy_fallbacks", self._core_content_policy_fallbacks - ) - - mock_timeout: Final = kwargs.pop("mock_timeout", None) - - with owning_picks(self._core_least_busy_reservations, call_id_of(CALL_IDENTITY.validate_python(kwargs))): - try: - self._handle_mock_testing_fallbacks( - kwargs=kwargs, - model_group=model_group, - fallbacks=fallbacks, - context_window_fallbacks=context_window_fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) - - if mock_timeout is not None: - response = await self.async_function_with_retries(*args, **kwargs, mock_timeout=mock_timeout) - else: - response = await self.async_function_with_retries(*args, **kwargs) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("Async Response: %s", response) - response = add_fallback_headers_to_response( - response=response, - attempted_fallbacks=0, - ) - return response - except Exception as e: - return await self.async_function_with_fallbacks_common_utils( - e, - disable_fallbacks, - fallbacks, - context_window_fallbacks, - content_policy_fallbacks, - model_group, - args, - kwargs, - include_fallback_errors=include_fallback_errors, - ) - - def _handle_mock_testing_fallbacks( - self, - kwargs: dict, - model_group: str | None = None, - fallbacks: list | None = None, - context_window_fallbacks: list | None = None, - content_policy_fallbacks: list | None = None, - ): - """ - Helper function to raise a litellm Error for mock testing purposes. - - Raises: - litellm.InternalServerError: when `mock_testing_fallbacks=True` passed in request params - litellm_exceptions.ContextWindowExceededError: when `mock_testing_context_fallbacks=True` passed in request params - litellm.ContentPolicyViolationError: when `mock_testing_content_policy_fallbacks=True` passed in request params - """ - mock_testing_params: Final = MockRouterTestingParams.from_kwargs(kwargs) - if ( - mock_testing_params.mock_testing_fallbacks is not None - and mock_testing_params.mock_testing_fallbacks is True - ): - raise litellm_exceptions.InternalServerError( - model=model_group, - llm_provider="", - message=f"This is a mock exception for model={model_group}, to trigger a fallback. Fallbacks={fallbacks}", - ) - elif ( - mock_testing_params.mock_testing_context_fallbacks is not None - and mock_testing_params.mock_testing_context_fallbacks is True - ): - raise litellm_exceptions.ContextWindowExceededError( - model=model_group, - llm_provider="", - message=f"This is a mock exception for model={model_group}, to trigger a fallback. \ - Context_Window_Fallbacks={context_window_fallbacks}", - ) - elif ( - mock_testing_params.mock_testing_content_policy_fallbacks is not None - and mock_testing_params.mock_testing_content_policy_fallbacks is True - ): - raise litellm_exceptions.ContentPolicyViolationError( - model=model_group, - llm_provider="", - message=f"This is a mock exception for model={model_group}, to trigger a fallback. \ - Context_Policy_Fallbacks={content_policy_fallbacks}", - ) - - @staticmethod - def _deployment_ids_to_skip_on_retry(exception: Exception, already_skipped: object) -> tuple[str, ...]: - failed_deployment_id: Final[str | None] = getattr(exception, "retry_skip_deployment_id", None) or getattr( - exception, "failed_deployment_id", None - ) - status_code: Final = getattr(exception, "status_code", None) - if not failed_deployment_id or not isinstance(status_code, int): - return () - if python_host.invoke("should_retry")(status_code): - return () - already_skipped_ids: Final = as_retry_skipped_deployment_ids(already_skipped) - skipped: Final = tuple(sorted(frozenset((*already_skipped_ids, failed_deployment_id)))) - verbose_router_logger.debug( - "Retry skips deployments that already answered %s to this request: %s", status_code, skipped - ) - return skipped - - @tracer.wrap() - async def async_function_with_retries(self, *args, **kwargs): - verbose_router_logger.debug("Inside async function with retries.") - original_function: Final = kwargs.pop("original_function") - fallbacks: Final = kwargs.pop("fallbacks", self._core_fallbacks) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - context_window_fallbacks: Final = kwargs.pop("context_window_fallbacks", self._core_context_window_fallbacks) - content_policy_fallbacks: Final = kwargs.pop("content_policy_fallbacks", self._core_content_policy_fallbacks) - # Support per-request model_group_retry_policy override (from key/team settings) - model_group_retry_policy: Final = kwargs.pop("model_group_retry_policy", self.model_group_retry_policy) - model_group: Final[str | None] = kwargs.get("model") - request_num_retries: Final[int | None] = kwargs.pop("num_retries", None) - num_retries = request_num_retries - if num_retries is None: - # Fall back to the router setting (then 0) so the comparisons below never - # hit `None > int`, which would mask the real upstream error with a TypeError. - num_retries = self.num_retries if self.num_retries is not None else 0 - - ## ADD MODEL GROUP SIZE TO METADATA - used for model_group_rate_limit_error tracking - _metadata: Final[dict] = kwargs.get("litellm_metadata", kwargs.get("metadata")) or {} - if "model_group" in _metadata and isinstance(_metadata["model_group"], str): - model_list: Final = self._core_get_model_list(model_name=_metadata["model_group"]) - if model_list is not None: - _metadata.update({"model_group_size": len(model_list)}) - - verbose_router_logger.debug( - "async function w/ retries: original_function - %s, num_retries - %s", original_function, num_retries - ) - ## ADD RETRY TRACKING TO METADATA - used for spend logs retry tracking - _metadata["attempted_retries"] = 0 - _metadata["max_retries"] = num_retries # Updated after overrides in exception handler - try: - self._handle_mock_testing_rate_limit_error(model_group=model_group, kwargs=kwargs) - # if the function call is successful, no exception will be raised and we'll break out of the loop - response = await self.make_call(original_function, *args, **kwargs) - response = add_retry_headers_to_response(response=response, attempted_retries=0, max_retries=None) - return response - except Exception as e: - if is_guardrail_intervention(e) or is_agentic_followup_failure(e): - raise - current_attempt = None - original_exception = e - deployment_num_retries: Final = getattr(e, "num_retries", None) - - if ( - request_num_retries is None - and deployment_num_retries is not None - and isinstance(deployment_num_retries, int) - ): - num_retries = deployment_num_retries - """ - Retry Logic - """ - ( - _healthy_deployments, - _all_deployments, - ) = await self._async_get_healthy_deployments( - model=kwargs.get("model") or "", - parent_otel_span=parent_otel_span, - request_kwargs=share_holder_request_kwargs( - cast(Mapping[str, object], kwargs), # cast-ok: kwargs keys are param names - self._ptu_share_teams(), - ), - ) - - # Check retry policy FIRST, before should_retry_this_error - # This allows retry policies to override the healthy deployments check - _retry_policy_applies = False - if request_num_retries != 0 and (self.retry_policy is not None or model_group_retry_policy is not None): - # get num_retries from retry policy - # Use the model_group captured at the start of the function, or get it from metadata - # kwargs.get("model") at this point is the deployment model, not the model_group - _model_group_for_retry_policy = model_group or _metadata.get("model_group") or kwargs.get("model") - # Use per-request model_group_retry_policy if provided, otherwise use self - _retry_policy_retries: Final = _get_num_retries_from_retry_policy( - exception=original_exception, - model_group=_model_group_for_retry_policy, - model_group_retry_policy=model_group_retry_policy, - retry_policy=self.retry_policy, - ) - if _retry_policy_retries is not None: - num_retries = _retry_policy_retries - _retry_policy_applies = True - - # raises an exception if this error should not be retries - # Skip this check if retry policy applies (retry policy takes precedence) - if not _retry_policy_applies: - self.should_retry_this_error( - error=e, - healthy_deployments=_healthy_deployments, - all_deployments=_all_deployments, - context_window_fallbacks=context_window_fallbacks, - regular_fallbacks=fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) - # Update max_retries after overrides (deployment_num_retries / retry_policy) - _metadata["max_retries"] = num_retries - - ## LOGGING - if num_retries > 0: - kwargs = self._core_log_retry(kwargs=kwargs, e=original_exception) - first_skipped_ids: Final = self._deployment_ids_to_skip_on_retry( - exception=original_exception, - already_skipped=kwargs.get("_retry_skipped_deployment_ids"), - ) - if first_skipped_ids: - kwargs["_retry_skipped_deployment_ids"] = first_skipped_ids # rebind-ok: the next attempt reads it - if isinstance( - original_exception, litellm_exceptions.UnsupportedEndpointError - ) and _every_deployment_refused( - _DEPLOYMENTS_ADAPTER.validate_python(_healthy_deployments), first_skipped_ids - ): - raise - else: - raise - - verbose_router_logger.debug("Retrying request with num_retries: %s", num_retries) - # decides how long to sleep before retry - retry_after: Final = self._time_to_sleep_before_retry( - e=original_exception, - remaining_retries=num_retries, - num_retries=num_retries, - healthy_deployments=_healthy_deployments, - all_deployments=_all_deployments, - ) - - await self._host.sleep(retry_after) - - for current_attempt in range(num_retries): - try: - # Update retry tracking metadata before each retry attempt - _metadata["attempted_retries"] = current_attempt + 1 - _metadata["max_retries"] = num_retries - carry_superseded_attempt_usage(kwargs) - # if the function call is successful, no exception will be raised and we'll break out of the loop - response = await self.make_call(original_function, *args, **kwargs) - if coroutine_checker.is_async_callable(response): # async errors are often returned as coroutines - response = await response - - response = add_retry_headers_to_response( - response=response, - attempted_retries=current_attempt + 1, - max_retries=num_retries, - ) - return response - - except Exception as e: - if is_agentic_followup_failure(e): - raise - # Always track the latest error so we raise the most - # recent exception instead of the first one. - original_exception = e - - ## LOGGING - kwargs = self._core_log_retry(kwargs=kwargs, e=e) - remaining_retries = num_retries - current_attempt - 1 - _model: str | None = kwargs.get("model") - if _model is not None: - ( - _healthy_deployments, - _, - ) = await self._async_get_healthy_deployments( - model=_model, - parent_otel_span=parent_otel_span, - request_kwargs=share_holder_request_kwargs( - cast(Mapping[str, object], kwargs), # cast-ok: kwargs keys are param names - self._ptu_share_teams(), - ), - ) - else: - _healthy_deployments = [] - - # Check if this error is non-retryable (e.g., 400 context - # window exceeded). If so, raise immediately instead of - # continuing the retry loop. Respect retry policy - # precedence - only check when no retry policy applies. - if not _retry_policy_applies: - try: - self.should_retry_this_error( - error=e, - healthy_deployments=_healthy_deployments, - all_deployments=_all_deployments, - context_window_fallbacks=context_window_fallbacks, - regular_fallbacks=fallbacks, - content_policy_fallbacks=content_policy_fallbacks, - ) - except Exception: - raise e - - skipped_ids = self._deployment_ids_to_skip_on_retry( - exception=e, - already_skipped=kwargs.get("_retry_skipped_deployment_ids"), - ) - if skipped_ids: - kwargs["_retry_skipped_deployment_ids"] = skipped_ids # rebind-ok: the next attempt reads it - if isinstance(e, litellm_exceptions.UnsupportedEndpointError) and _every_deployment_refused( - _DEPLOYMENTS_ADAPTER.validate_python(_healthy_deployments), skipped_ids - ): - raise - _timeout = self._time_to_sleep_before_retry( - e=e, - remaining_retries=remaining_retries, - num_retries=num_retries, - healthy_deployments=_healthy_deployments, - all_deployments=_all_deployments, - ) - await self._host.sleep(_timeout) - - if type(original_exception) in self._host.dynamic().exception_types: - setattr(original_exception, "max_retries", num_retries) - # current_attempt is 0-indexed (0 to num_retries-1), so after loop completion - # it represents the last attempt index. The actual number of retries attempted - # is current_attempt + 1, which equals num_retries when all retries are exhausted. - # We've already verified num_retries > 0 before entering the loop, so current_attempt - # will always be set (never None) when we reach this point. - actual_retries_attempted: Final = current_attempt + 1 if current_attempt is not None else num_retries - setattr(original_exception, "num_retries", actual_retries_attempted) - - raise original_exception - - async def make_call(self, original_function: Any, *args: object, **kwargs: object): - """ - Handler for making a call to the .completion()/.embeddings()/etc. functions. - """ - model_group: Final = kwargs.get("model") - generic_function: Final = _OPAQUE_VALUE_ADAPTER.validate_python(kwargs.get("original_generic_function")) - endpoint_function: Final = ( - generic_function if callable(generic_function) else _OPAQUE_VALUE_ADAPTER.validate_python(original_function) - ) - try: - with endpoint_selection( - getattr(endpoint_function, "__name__", None), - getattr(endpoint_function, "__module__", None), - model=_OPAQUE_VALUE_ADAPTER.validate_python(model_group), - ): - response = original_function(*args, **kwargs) - if coroutine_checker.is_async_callable(response) or inspect.isawaitable(response): - response = await response - except Exception: - with contextlib.suppress(Exception): - await async_io_token_refund_failure( - self._core_cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs) - ) - raise - await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) - ## PROCESS RESPONSE HEADERS - response = await self.set_response_headers( - response=response, model_group=model_group if isinstance(model_group, str) else None, request_kwargs=kwargs - ) - - return response - - def _handle_mock_testing_rate_limit_error(self, kwargs: dict, model_group: str | None = None): - """ - Helper function to raise a mock litellm.RateLimitError error for testing purposes. - - Raises: - litellm.RateLimitError error when `mock_testing_rate_limit_error=True` passed in request params - """ - mock_testing_rate_limit_error: Final[bool | None] = kwargs.pop("mock_testing_rate_limit_error", None) - - available_models: Final = self._core_get_model_list(model_name=model_group) - num_retries: int | None = None - - if available_models is not None and len(available_models) == 1: - num_retries = available_models[0]["litellm_params"].get("num_retries") - - if mock_testing_rate_limit_error is not None and mock_testing_rate_limit_error is True: - verbose_router_logger.info( - "litellm.router.py::_mock_rate_limit_error() - Raising mock RateLimitError for model=%s", model_group - ) - raise litellm_exceptions.RateLimitError( - model=model_group, - llm_provider="", - message=f"This is a mock exception for model={model_group}, to trigger a rate limit error.", - num_retries=num_retries, - ) - - def should_retry_this_error( - self, - error: Exception, - healthy_deployments: list | None = None, - all_deployments: list | None = None, - context_window_fallbacks: list | None = None, - content_policy_fallbacks: list | None = None, - regular_fallbacks: list | None = None, - ): - """ - 1. raise an exception for ContextWindowExceededError if context_window_fallbacks is not None - 2. raise an exception for ContentPolicyViolationError if content_policy_fallbacks is not None - - 2. raise an exception for RateLimitError if - - there are no fallbacks - - there are no healthy deployments in the same model group - """ - _num_healthy_deployments = 0 - if healthy_deployments is not None and isinstance(healthy_deployments, list): - _num_healthy_deployments = len(healthy_deployments) - - _num_all_deployments = 0 - if all_deployments is not None and isinstance(all_deployments, list): - _num_all_deployments = len(all_deployments) - - ### CHECK IF RATE LIMIT / CONTEXT WINDOW ERROR / CONTENT POLICY VIOLATION ERROR w/ fallbacks available / Bad Request Error - if isinstance(error, litellm_exceptions.ContextWindowExceededError) and context_window_fallbacks is not None: - raise error - - if isinstance(error, litellm_exceptions.ContentPolicyViolationError) and content_policy_fallbacks is not None: - raise error - - status_code: Final = getattr(error, "status_code", None) - if status_code is not None and not self._host.invoke("should_retry")(status_code): - # 401/403 are special cases - allow retry if multiple deployments exist (handled below) - if status_code not in (401, 403) and not isinstance(error, litellm_exceptions.UnsupportedEndpointError): - raise error - - if isinstance(error, litellm_exceptions.NotFoundError): - raise error - # Error we should only retry if there are other deployments - if isinstance(error, openai.RateLimitError): - if ( - _num_healthy_deployments <= 0 # if no healthy deployments - and regular_fallbacks is not None # and fallbacks available - and len(regular_fallbacks) > 0 - ): - raise error # then raise the error - - if isinstance(error, (openai.AuthenticationError, openai.PermissionDeniedError)): - """ - - if other deployments available -> retry - - else -> raise error - """ - if _num_all_deployments <= 1: # if there is only 1 deployment for this model group then don't retry - raise error # then raise error - - # Do not retry if there are no healthy deployments - # just raise the error - if _num_healthy_deployments <= 0: # if no healthy deployments - raise error - - return True - - def function_with_fallbacks(self, *args, **kwargs): - """ - Sync wrapper for async_function_with_fallbacks - - Wrapped to reduce code duplication and prevent bugs. - """ - return run_async_function(self.async_function_with_fallbacks, *args, **kwargs) - - def _get_fallback_model_group_for_lookup_groups( - self, - fallbacks: list[dict[str, list[str]]], # mutable-ok: mirrors the shared resolver's contract - lookup_groups: tuple[str, ...], - ) -> list[str] | None: # mutable-ok: mirrors the shared resolver's contract - fallback_model_group, _ = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, lookup_groups=lookup_groups - ) - return fallback_model_group - - def _get_first_default_fallback(self) -> str | None: - """ - Returns the first model from the default_fallbacks list, if it exists. - """ - if self._core_fallbacks is None: - return None - for fallback in self._core_fallbacks: - if isinstance(fallback, dict) and "*" in fallback: - default_list = fallback["*"] - if isinstance(default_list, list) and len(default_list) > 0: - return default_list[0] - return None - - def _time_to_sleep_before_retry( - self, - e: Exception, - remaining_retries: int, - num_retries: int, - healthy_deployments: list | None = None, - all_deployments: list | None = None, - ) -> int | float: - """ - Calculate back-off, then retry - - It should instantly retry only when: - 1. there are healthy deployments in the same model group - 2. there are fallbacks for the completion call - """ - - ## base case - single deployment - if all_deployments is not None and len(all_deployments) == 1: - pass - elif healthy_deployments is not None and isinstance(healthy_deployments, list) and len(healthy_deployments) > 0: - return 0 - - response_headers: httpx.Headers | None = None - if hasattr(e, "response") and hasattr(e.response, "headers"): - response_headers = e.response.headers - if hasattr(e, "litellm_response_headers"): - response_headers = e.litellm_response_headers - - if response_headers is not None: - timeout = self._host.invoke("calculate_retry_after")( - remaining_retries=remaining_retries, - max_retries=num_retries, - response_headers=response_headers, - min_timeout=self.retry_after, - ) - - else: - timeout = self._host.invoke("calculate_retry_after")( - remaining_retries=remaining_retries, - max_retries=num_retries, - min_timeout=self.retry_after, - ) - - return timeout - - ### HELPER FUNCTIONS - - async def deployment_callback_on_success( - self, - kwargs, # kwargs to completion - completion_response, # response from completion - start_time, - end_time, # start/end time - ): - """ - Track remaining tpm/rpm quota for model in model_list - """ - try: - if kwargs.get("call_type") == "_arealtime": - return - if is_batch_retrieve_call_type(kwargs.get("call_type")): - return - standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None) - if standard_logging_object is None: - raise ValueError("standard_logging_object is None") - litellm_params: Final = kwargs["litellm_params"] - metadata: Final = litellm_params.get("metadata") - if metadata is None: - return - model_group: Final = metadata.get("model_group", None) - model_info: Final = litellm_params.get("model_info", {}) or {} - deployment_id: Final = model_info.get("id", None) - if ( - model_group is None - or deployment_id is None - or self._core_get_deployment(model_id=str(deployment_id)) is None - ): - return - - if self.session_affinity_repin_on_failover: - await self._repin_session_after_success( - _LITELLM_PARAMS_ADAPTER.validate_python(litellm_params), - _OPAQUE_VALUE_ADAPTER.validate_python(completion_response), - ) - - # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits - increment_deployment_successes_for_current_minute( - ctx=self._core_ctx, - deployment_id=str(deployment_id), - ) - - total_tokens: Final[float] = standard_logging_object.get("total_tokens", 0) - counted_tokens: Final = get_counted_usage_tokens(litellm_params) - deployment_name: Final = metadata.get("deployment", None) - return await self._increment_deployment_usage( - deployment_id=str(deployment_id), - deployment_name=deployment_name if isinstance(deployment_name, str) else None, - model_group=model_group, - total_tokens=total_tokens if counted_tokens is None else max(0, total_tokens - counted_tokens), - rpm_increment=1 if counted_tokens is None else 0, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - ) - - except Exception as e: - verbose_router_logger.debug( - "litellm.router.Router::deployment_callback_on_success(): Exception occured - %s", e - ) - - async def _repin_session_after_success(self, litellm_params: Mapping[str, object], response: object) -> None: - if not self.session_affinity_repin_on_failover: - return - try: - from litellm.router_utils.pre_call_checks.deployment_affinity_check import completed_affinity_response - - if not completed_affinity_response(response): - return - callback: Final = next( - (logger for logger in self._optional_loggers() if isinstance(logger, DeploymentAffinityCheck)), - None, - ) - if callback is None: - return - model_group: Final = callback.get_repin_model_group(litellm_params) - if model_group is None or self._is_priority_routing_group(model_group): - return - registered_ids: Final = self._core_get_candidate_model_ids_for_route( - model=model_group, team_id=get_request_team_id(litellm_params) - ) - cooldown_ids: Final = frozenset( - model_id - for model_id, _ in await self._core_cooldown_cache.async_get_active_cooldowns( - model_ids=list(registered_ids), parent_otel_span=None - ) - ) - await callback.async_repin_session_after_success(litellm_params, model_group, registered_ids - cooldown_ids) - except Exception as error: # noqa: BLE001 # Optional pin updates must not interrupt success accounting - verbose_router_logger.debug("Session re-pinning failed: %s", error) - - async def increment_deployment_usage_for_response( - self, - response: object, - request_kwargs: dict[str, object], - ) -> None: - if response is None: - return - try: - deployment_metadata: Final = find_deployment_metadata(request_kwargs) - model_group: Final = request_kwargs.get("model") - if deployment_metadata is None or not isinstance(model_group, str): - return - model_info: Final = deployment_metadata["model_info"] - deployment_id: Final = model_info.get("id") if isinstance(model_info, dict) else None - if deployment_id is None: - return - total_tokens: Final = response_total_token_count(response) - deployment_name: Final = deployment_metadata.get("deployment") - deployment_metadata[ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY] = total_tokens - try: - await self._increment_deployment_usage( - deployment_id=str(deployment_id), - deployment_name=deployment_name if isinstance(deployment_name, str) else None, - model_group=model_group, - total_tokens=total_tokens, - rpm_increment=1, - parent_otel_span=get_parent_otel_span_from_kwargs(request_kwargs), - ) - except Exception: - deployment_metadata.pop(ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY, None) - raise - except Exception as e: - verbose_router_logger.debug( - "litellm.router.Router::increment_deployment_usage_for_response(): Exception occured - %s", e - ) - - async def _increment_deployment_usage( - self, - *, - deployment_id: str, - deployment_name: str | None, - model_group: str, - total_tokens: float, - rpm_increment: int, - parent_otel_span: Span | None, - ) -> str | None: - from litellm.types.caching import RedisPipelineIncrementOperation - - if total_tokens <= 0 and rpm_increment <= 0: - return None - if not self._cached_deployment_tracks_usage(deployment_id, model_group): - return None - - current_minute: Final = get_utc_datetime().strftime("%H-%M") # use the same timezone regardless of system clock - tpm_key: Final = RouterCacheEnum.TPM.value.format( - id=deployment_id, current_minute=current_minute, model=deployment_name - ) - rpm_key: Final = RouterCacheEnum.RPM.value.format( - id=deployment_id, current_minute=current_minute, model=deployment_name - ) - pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [ - RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value) - for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment)) - ] - post_increment_values: Final = await self._core_cache.async_increment_cache_pipeline( - increment_list=pipeline_operations, - parent_otel_span=parent_otel_span, - ) - if post_increment_values is not None and self._core_cache.redis_cache is not None: - for operation, value in zip(pipeline_operations, post_increment_values): - await self._core_cache.async_set_cache( - operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value - ) - return tpm_key - - def _deployment_tracks_usage(self, deployment_id: str, model_group: str) -> bool: - deployment_info: Final = self._core_get_deployment(model_id=deployment_id) - if deployment_info is None: - return False - deployment_model_info: Final = self._core_get_router_model_info( - deployment=deployment_info, - received_model_name=model_group, - ) - configured_limits: Final = ( - deployment_info.get("tpm", None), - deployment_info.get("rpm", None), - deployment_info.litellm_params.tpm, - deployment_info.litellm_params.rpm, - deployment_model_info.get("tpm", None), - deployment_model_info.get("rpm", None), - ) - ## Nothing to track only when neither tpm/rpm nor itpm/otpm limits are - ## set. IO deployments still record TPM/RPM usage here so TPM-aware - ## routing strategies see their real load in mixed model groups; their - ## itpm/otpm enforcement runs separately in ModelRateLimitingCheck. - return any(limit is not None for limit in configured_limits) or deployment_has_io_token_limits( - deployment_info.model_dump() - ) - - def sync_deployment_callback_on_success( - self, - kwargs, # kwargs to completion - completion_response, # response from completion - start_time, - end_time, # start/end time - ) -> str | None: - """ - Tracks the number of successes for a deployment in the current minute (using in-memory cache) - - Returns: - - key: str - The key used to increment the cache - - None: if no key is found - """ - if is_batch_retrieve_call_type(kwargs.get("call_type")): - return None - id = None - if kwargs["litellm_params"].get("metadata") is None: - pass - else: - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - model_info: Final = kwargs["litellm_params"].get("model_info", {}) or {} - id = model_info.get("id", None) - if model_group is None or id is None: - return None - elif isinstance(id, int): - id = str(id) - - if id is not None: - key: Final = increment_deployment_successes_for_current_minute( - ctx=self._core_ctx, - deployment_id=id, - ) - return key - - return None - - def deployment_callback_on_failure( - self, - kwargs, # kwargs to completion - completion_response, # response from completion - start_time, - end_time, # start/end time - ) -> bool: - """ - 2 jobs: - - Tracks the number of failures for a deployment in the current minute (using in-memory cache) - - Puts the deployment in cooldown if it exceeds the allowed fails / minute - - Returns: - - True if the deployment should be put in cooldown - - False if the deployment should not be put in cooldown - """ - verbose_router_logger.debug("Router: Entering 'deployment_callback_on_failure'") - try: - exception: Final = kwargs.get("exception", None) - - if is_failure_outside_selected_deployment(exception): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Failure came from advisor orchestration or deployment selection, not the selected deployment." - ) - return False - - # Cache litellm_params to avoid repeated dict lookups - litellm_params: Final = kwargs.get("litellm_params", {}) - _model_info: Final = litellm_params.get("model_info", {}) - - if is_background_response_poll_not_found(exception, litellm_params): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Provider 404 came from the background response poll, not the deployment's health." - ) - return False - - if is_video_lookup_not_found(exception, kwargs): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Provider 404 says the video is not on this deployment's account, not that the deployment is down." - ) - return False - - exception_status: Final = getattr(exception, "status_code", "") - failure_call_details: Final = _CALL_KWARGS_ADAPTER.validate_python(kwargs) - - if is_caller_timeout_408(failure_call_details, exception_status): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "A timeout the caller set caused this 408, not the deployment's health." - ) - return False - - caller_failure_model_info: Final = ( - _MODEL_INFO_ADAPTER.validate_python(_model_info) if isinstance(_model_info, dict) else None - ) - caller_failure_model_id: Final = ( - caller_failure_model_info.get("id") if caller_failure_model_info is not None else None - ) - caller_failure_deployment: Final = ( - self._core_get_deployment(model_id=caller_failure_model_id) - if isinstance(caller_failure_model_id, str) - else None - ) - if is_caller_scoped_auth_failure(caller_failure_deployment, exception_status): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Caller-scoped OAuth authentication failed, not the deployment." - ) - return False - if is_caller_entra_credential_failure( - caller_failure_deployment, - self._active_default_litellm_params(), - failure_call_details, - exception_status, - exception, - ): - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. " - "Entra ID rejected a credential the caller sent, not the deployment's own." - ) - return False - - exception_headers: Final = exception_mapping_utils._get_response_headers(original_exception=exception) - - # Determine cooldown time with priority: deployment config > response header > router default - deployment_cooldown: Final = _first_present( - _model_info if isinstance(_model_info, dict) else None, litellm_params, key="cooldown_time" - ) - - header_cooldown = None - if exception_headers is not None: - header_cooldown = litellm_utils._get_retry_after_from_exception_header( - response_headers=exception_headers - ) - ############################################## - # Logic to determine cooldown time - # 1. Check if a cooldown time is set in the deployment config - # 2. Check if a cooldown time is set in the response header - # 3. If no cooldown time is set, use the router default cooldown time - ############################################## - if deployment_cooldown is not None and deployment_cooldown >= 0: - _time_to_cooldown = deployment_cooldown - elif header_cooldown is not None and header_cooldown >= 0: - _time_to_cooldown = header_cooldown - else: - _time_to_cooldown = self._core_cooldown_time - - if isinstance(_model_info, dict): - deployment_id: Final[str | None] = _model_info.get("id") - if deployment_id is None: - return False - increment_deployment_failures_for_current_minute( - ctx=self._core_ctx, - deployment_id=deployment_id, - ) - result: Final = set_cooldown_deployments( - ctx=self._core_ctx, - exception_status=exception_status, - original_exception=exception, - deployment=deployment_id, - time_to_cooldown=_time_to_cooldown, - requested_model_group=( - get_litellm_metadata_from_kwargs( - cast("dict[str, object]", kwargs) # cast-ok: untyped kwargs dict - ) - or {} - ).get("model_group"), - request_kwargs=cast("dict[str, object]", kwargs), # cast-ok: kwargs is the untyped request dict - ) # setting deployment_id in cooldown deployments - - return result - else: - verbose_router_logger.debug( - "Router: Exiting 'deployment_callback_on_failure' without cooldown. No model_info found." - ) - return False - - except Exception as e: - raise e - - async def async_deployment_callback_on_failure( - self, kwargs, completion_response: object | None, start_time, end_time - ): - """ - Update RPM usage for a deployment - """ - if is_batch_retrieve_call_type(kwargs.get("call_type")): - return - deployment_name: Final = kwargs["litellm_params"]["metadata"].get( - "deployment", None - ) # handles wildcard routes - by giving the original name sent to `litellm.completion` - model_group: Final = kwargs["litellm_params"]["metadata"].get("model_group", None) - model_info: Final = kwargs["litellm_params"].get("model_info", {}) or {} - id = model_info.get("id", None) - if model_group is None or id is None: - return - elif isinstance(id, int): - id = str(id) - parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) - - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - - ## RPM - rpm_key: Final = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name) - with service_target(ROUTER_USAGE_TARGET): - await self._core_cache.async_increment_cache( - key=rpm_key, - value=1, - parent_otel_span=parent_otel_span, - ttl=RoutingArgs.ttl.value, - ) - - def _get_metadata_variable_name_from_kwargs(self, kwargs: dict) -> Literal["metadata", "litellm_metadata"]: - """ - Helper to return what the "metadata" field should be called in the request data - - - New endpoints return `litellm_metadata` - - Old endpoints return `metadata` - - Context: - - LiteLLM used `metadata` as an internal field for storing metadata - - OpenAI then started using this field for their metadata - - LiteLLM is now moving to using `litellm_metadata` for our metadata - """ - return get_metadata_variable_name_from_kwargs(kwargs) - - def _core_log_retry(self, kwargs: dict, e: Exception) -> dict: - """ - When a retry or fallback happens, record which model group, deployment and attempt just failed and why, - and count it toward the request-wide num_retries_per_request cap - """ - from litellm.litellm_core_utils.get_litellm_params import get_base_model_from_litellm_call_metadata - from litellm.types.router import RetryAttemptRecord - - _metadata_var: Final = "litellm_metadata" if "litellm_metadata" in kwargs else "metadata" - request_metadata: Final = cast("Mapping[str, object]", kwargs[_metadata_var]) # cast-ok: untyped kwargs - model_group: Final = kwargs.get("model") - model_info: Final = request_metadata.get("model_info") - deployment_id: Final = model_info.get("id") if isinstance(model_info, Mapping) else None - base_model: Final = get_base_model_from_litellm_call_metadata( - metadata=request_metadata - ) or self._deployment_params_base_model(deployment_id if isinstance(deployment_id, str) else None) - deployment_model: Final = request_metadata.get("deployment") - attempted_retries: Final = request_metadata.get("attempted_retries") - attempt_record: Final[RetryAttemptRecord] = { - "model_group": model_group if isinstance(model_group, str) else None, - "deployment_id": deployment_id if isinstance(deployment_id, str) else None, - "model": deployment_model if isinstance(deployment_model, str) else None, - "base_model": base_model if isinstance(base_model, str) else None, - "exception_type": type(e).__name__, - "exception_string": str(e), - "attempted_retries": attempted_retries if type(attempted_retries) is int else None, - } - earlier_breadcrumbs: Final = request_metadata.get("previous_models") - kept_breadcrumbs: Final[tuple[object, ...]] = ( - tuple(earlier_breadcrumbs)[-(RETRY_BREADCRUMB_LIMIT - 1) :] - if isinstance(earlier_breadcrumbs, (list, tuple)) - else () - ) - breadcrumbs: Final = (*kept_breadcrumbs, attempt_record) - earlier: Final = request_metadata.get("request_retry_count") - request_retry_count: Final = (earlier if type(earlier) is int and 0 <= earlier else 0) + 1 - decision: Final = request_metadata.get("routing_decision") - router_name: Final = decision.get("router_model_name") if isinstance(decision, Mapping) else None - if isinstance(router_name, str): - phase_event( - "litellm.routing.retry", - { - "litellm.routing.router_model_name": router_name, - "litellm.retry.count": request_retry_count, - "error.type": type(e).__name__, - **({"litellm.deployment.model_group": model_group} if isinstance(model_group, str) else {}), - **({"litellm.deployment.id": deployment_id} if isinstance(deployment_id, str) else {}), - }, - ) - kwargs[_metadata_var]["previous_models"] = breadcrumbs # rebind-ok: the logging object already holds this dict - kwargs[_metadata_var]["request_retry_count"] = request_retry_count # rebind-ok: same dict, read by the cap - return kwargs - - def _deployment_params_base_model(self, deployment_id: str | None) -> str | None: - """The ``base_model`` a deployment declares under ``litellm_params`` (the Azure named-deployment - shape), the same key the success path's ``get_base_model_from_metadata`` reads.""" - deployment: Final = self._core_get_deployment(model_id=deployment_id) if deployment_id is not None else None - if deployment is None: - return None - extra_params: Final = cast( # cast-ok: pydantic extras are untyped, base_model is isinstance-checked below - "Mapping[str, object]", deployment.litellm_params.model_extra or {} - ) - base_model: Final = extra_params.get("base_model") - return base_model if isinstance(base_model, str) else None - - def _update_usage(self, deployment_id: str, parent_otel_span: Span | None) -> int: - """ - Update deployment rpm for that minute - - Returns: - - int: request count - """ - rpm_key: Final = deployment_id - - request_count = self._core_cache.get_cache(key=rpm_key, parent_otel_span=parent_otel_span, local_only=True) - if request_count is None: - request_count = 1 - self._core_cache.set_cache(key=rpm_key, value=request_count, local_only=True, ttl=60) # only store for 60s - else: - request_count += 1 - self._core_cache.set_cache(key=rpm_key, value=request_count, local_only=True) # don't change existing ttl - - return request_count - - def _has_default_fallbacks(self) -> bool: - if self._core_fallbacks is None: - return False - for fallback in self._core_fallbacks: - if isinstance(fallback, dict): - if "*" in fallback: - return True - return False - - def _has_content_policy_fallback(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: - """ - Whether a content-policy fallback would resolve for this request, keyed the same way - async_function_with_fallbacks_common_utils resolves it: the tier a pre-routing hook - selected wins over the requested group. Raising without this returning True would turn - a deliverable response into an error the fallback chain cannot recover from. - """ - content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self._core_content_policy_fallbacks) - if content_policy_fallbacks is not None: - return has_unattempted_fallback_target( - self._get_fallback_model_group_for_lookup_groups( - fallbacks=content_policy_fallbacks, - lookup_groups=fallback_lookup_groups(kwargs, model_group), - ), - kwargs, - ) - if self._has_default_fallbacks(): - return True - verbose_router_logger.debug( - "No content-policy fallback available. Returning original response. model=%s, content_policy_fallbacks=%s", - model_group, - content_policy_fallbacks, - ) - return False - - def _refusal_fallback_available(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: - """ - Whether a safeguard refusal can actually be recovered by the dispatcher. A configured - content-policy list is authoritative; with none configured at all, the dispatcher falls - through to the generic fallbacks lookup, so the gate mirrors that reachability and arms - on a resolving generic chain (tier first, then the requested group, then "*"). - """ - if fallbacks_disabled_for_request(kwargs): - return False - content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self._core_content_policy_fallbacks) - if content_policy_fallbacks is not None: - return self._has_content_policy_fallback(model_group, kwargs) - if self._has_default_fallbacks(): - return True - fallbacks: Final = kwargs.get("fallbacks", self._core_fallbacks) - if fallbacks is None: - return False - resolved, _ = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, - lookup_groups=fallback_lookup_groups(kwargs, model_group), - ) - return has_unattempted_fallback_target(resolved, kwargs) - - def _stream_order_levels(self, model_group: str, kwargs: Mapping[str, Any]) -> tuple[int, ...]: - """ - The distinct deployment order levels the fallback dispatcher would see for this request, - computed the same way: the tier a pre-routing hook selected wins over the requested group. - """ - request_team_id: Final[str | None] = (kwargs.get("metadata", {}) or {}).get("user_api_key_team_id") - order_model_group: Final = get_pre_routing_selection(kwargs) or model_group - all_deployments: Final = self._core_get_model_list(model_name=order_model_group, team_id=request_team_id) or () - return tuple( - sorted( - { - litellm_utils.get_deployment_order(d) - for d in all_deployments - if litellm_utils.get_deployment_order(d) is not None - } - ) - ) - - def _stream_can_fall_back(self, model_group: str, kwargs: Mapping[str, Any]) -> bool: - """ - Whether async_function_with_fallbacks_common_utils could still route a - MidStreamFallbackError somewhere for this request (order levels, weighted - failover, content-policy or generic fallbacks), which is the only case where - holding lifecycle frames back from the client buys a clean retry. Errs toward - True whenever a dispatcher path might reach a fallback. - """ - if fallbacks_disabled_for_request(kwargs): - return False - if self.enable_weighted_failover: - return True - order_levels: Final = self._stream_order_levels(model_group, kwargs) - if len(order_levels) > 1: - current_target: Final = kwargs.get("_target_order") - skip_up_to: Final = current_target if current_target is not None else order_levels[0] - if any(o > skip_up_to for o in order_levels): - return True - content_policy_fallbacks: Final = kwargs.get("content_policy_fallbacks", self._core_content_policy_fallbacks) - if content_policy_fallbacks is not None and self._has_content_policy_fallback(model_group, kwargs): - return True - fallbacks: Final = kwargs.get("fallbacks", self._core_fallbacks) - if not fallbacks: - return False - if check_non_standard_fallback_format(fallbacks=fallbacks): - return True - resolved, _ = get_fallback_model_group_for_lookup_groups( - fallbacks=fallbacks, - lookup_groups=fallback_lookup_groups(kwargs, model_group), - ) - return has_unattempted_fallback_target(resolved, kwargs) - - def _should_raise_content_policy_error(self, model: str, response: ModelResponse, kwargs: dict) -> bool: - """ - Determines if a content policy error should be raised. - - Only raised if a fallback is available. - - Else, original response is returned. - """ - if response.choices and len(response.choices) > 0: - if response.choices[0].finish_reason != "content_filter": - return False - - return self._has_content_policy_fallback(model, kwargs) - - def _should_raise_anthropic_refusal_error( - self, model: str, original_generic_function: Callable, response: object, kwargs: Mapping[str, object] - ) -> bool: - """ - The /v1/messages twin of _should_raise_content_policy_error: an Anthropic safeguard - refusal (stop_reason "refusal" carrying stop_details) re-enters the fallback chain only - when a content-policy fallback is configured; a plain refusal without stop_details, or - any response with nothing configured, is returned to the client unchanged. - """ - from litellm.llms.anthropic.messages.native.utils import ( - get_safeguard_refusal_stop_details, - ) - - if getattr(original_generic_function, "__name__", "") != "anthropic_messages": - return False - if get_safeguard_refusal_stop_details(response) is None: - return False - return self._refusal_fallback_available(model, kwargs) - - def _get_healthy_deployments(self, model: str, parent_otel_span: Span | None): - _all_deployments: list = [] - try: - _, _all_deployments = self._common_checks_available_deployment( - model=model, - ) - if isinstance(_all_deployments, dict): - return [] - except Exception: - pass - - unhealthy_deployments: Final = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - unhealthy_set: Final = set(unhealthy_deployments) - healthy_deployments: list = [d for d in _all_deployments if d["model_info"]["id"] not in unhealthy_set] - healthy_deployments = self._filter_blocked_deployments(healthy_deployments) - - return healthy_deployments, _all_deployments - - async def _async_get_healthy_deployments( - self, model: str, parent_otel_span: Span | None, request_kwargs: Mapping[str, object] | None = None - ) -> tuple[list[dict], list[dict]]: - """ - Returns Tuple of: - - Tuple[List[Dict], List[Dict]]: - 1. healthy_deployments: list of healthy deployments the caller of ``request_kwargs`` may use - 2. all_deployments: list of all deployments that caller may use - """ - _all_deployments: list = [] - try: - _, _all_deployments = self._common_checks_available_deployment( - model=model, - request_kwargs=request_kwargs, - ) - if isinstance(_all_deployments, dict): - return [], _all_deployments - except Exception: - pass - - unhealthy_deployments: Final = await async_get_cooldown_deployments( - ctx=self._core_ctx, parent_otel_span=parent_otel_span - ) - # Convert to set for O(1) lookup instead of O(n) - unhealthy_deployments_set: Final = set(unhealthy_deployments) - healthy_deployments: list = [ - d for d in _all_deployments if d["model_info"]["id"] not in unhealthy_deployments_set - ] - healthy_deployments = self._filter_blocked_deployments(healthy_deployments) - return healthy_deployments, _all_deployments - - def routing_strategy_pre_call_checks(self, deployment: dict): - """ - Mimics 'async_routing_strategy_pre_call_checks' - - Ensures consistent update rpm implementation for 'usage-based-routing-v2' - - Returns: - - None - - Raises: - - Rate Limit Exception - If the deployment is over it's tpm/rpm limits - """ - for _callback in self._routing_callbacks(): - if isinstance(_callback, CustomLogger): - _callback.pre_call_check(deployment) - - async def async_routing_strategy_pre_call_checks( - self, - deployment: dict, - parent_otel_span: Span | None, - logging_obj: LiteLLMLogging | None = None, - ): - """ - For usage-based-routing-v2, enables running rpm checks before the call is made, inside the semaphore. - - -> makes the calls concurrency-safe, when rpm limits are set for a deployment - - Returns: - - None - - Raises: - - Rate Limit Exception - If the deployment is over it's tpm/rpm limits - """ - for _callback in self._routing_callbacks(): - if isinstance(_callback, CustomLogger): - try: - await _callback.async_pre_call_check(deployment, parent_otel_span) - except litellm_exceptions.RateLimitError as e: - ## LOG FAILURE EVENT - if logging_obj is not None: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=e, - traceback_exception=traceback.format_exc(), - prefer_async_handlers=True, - ) - ) - set_cooldown_deployments( - ctx=self._core_ctx, - exception_status=e.status_code, - original_exception=e, - deployment=deployment["model_info"]["id"], - time_to_cooldown=self._core_cooldown_time, - ) - raise e - except Exception as e: - ## LOG FAILURE EVENT - if logging_obj is not None: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=e, - traceback_exception=traceback.format_exc(), - prefer_async_handlers=True, - ) - ) - raise e - - def _deployment_as_configured_now(self, deployment: Mapping[str, object]) -> Mapping[str, object]: - model_info: Final = deployment.get("model_info") - model_id: Final = ( - _MODEL_INFO_ADAPTER.validate_python(model_info).get("id") if isinstance(model_info, dict) else None - ) - current: Final = self._core_get_deployment(model_id) if isinstance(model_id, str) else None - return deployment if current is None else current.model_dump(exclude_none=True) - - async def async_admit_deployment_turn( - self, - model_group: str, - opened_on: Mapping[str, object], - request_kwargs: Mapping[str, object], - ) -> Mapping[str, object]: - deployment: Final = self._deployment_as_configured_now(opened_on) - rate_limit_check: Final = next( - (callback for callback in self._routing_callbacks() if isinstance(callback, ModelRateLimitingCheck)), None - ) - io_token_admission: Final[Mapping[str, object]] = ( - MappingProxyType({}) - if rate_limit_check is None - else await async_admit_io_token_turn(rate_limit_check.dual_cache, deployment, get_utc_datetime()) - ) - budget_limiter: Final = self._core_get_router_deployment_budget_limiter() - if budget_limiter is not None: - try: - await budget_limiter.async_filter_deployments( # pyright: ignore[reportUnknownMemberType] # legacy signature takes untyped deployment dicts - model=model_group, - healthy_deployments=[dict(deployment)], - messages=None, - request_kwargs=dict(request_kwargs), - ) - except ValueError as e: - raise litellm_exceptions.RateLimitError( - message=str(e), - llm_provider="", - model=model_group, - category=litellm_exceptions.RateLimitErrorCategory.LITELLM_RATE_LIMIT, - ) from e - await self._admit_turn_through_routing_strategy(model_group, deployment, request_kwargs) - await self.async_routing_strategy_pre_call_checks( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict - deployment=dict(without_io_token_limits(deployment)), parent_otel_span=None - ) - return io_token_admission - - async def _admit_turn_through_routing_strategy( - self, model_group: str, deployment: Mapping[str, object], request_kwargs: Mapping[str, object] - ) -> None: - routing_kwargs: Final = {key: value for key, value in request_kwargs.items() if key != "litellm_logging_obj"} - strategy, selector = self._get_routing_context( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped request dict - model=model_group, request_kwargs=routing_kwargs - ) - if strategy == "simple-shuffle" or selector is None: - return - selected: Final = await self._select_deployment_async( # pyright: ignore[reportUnknownMemberType] # legacy signature takes untyped deployment dicts - strategy=strategy, - selector=selector, - model=model_group, - healthy_deployments=[dict(deployment)], - messages=None, - input=None, - request_kwargs=routing_kwargs, - reserve=False, - ) - if selected is None: - no_deployment: Final = await async_raise_no_deployment_exception( - ctx=self._core_ctx, model=model_group, parent_otel_span=None - ) - raise litellm_exceptions.RateLimitError( - message=str(no_deployment), - llm_provider="", - model=model_group, - category=litellm_exceptions.RateLimitErrorCategory.LITELLM_RATE_LIMIT, - ) from no_deployment - await self._async_override_selector_pre_call_check( # pyright: ignore[reportUnknownMemberType] # legacy signature takes an untyped deployment dict - strategy, selector, dict(deployment), None - ) - - @contextlib.asynccontextmanager - async def _deployment_slot( - self, deployment: dict, kwargs: Mapping[str, object], parent_otel_span: Span | None - ) -> AsyncGenerator[None, None]: - """Holds the deployment's max_parallel_requests slot, if it has one, around the provider call. Routing - strategy pre-call checks run inside the slot so their rpm accounting stays concurrency-safe.""" - max_parallel_requests_limit: Final = self._get_client( - deployment=deployment, - kwargs=kwargs, - client_type="max_parallel_requests", - ) - async with contextlib.AsyncExitStack() as slot: - if isinstance(max_parallel_requests_limit, MaxParallelRequestsLimit): - slot.enter_context(max_parallel_requests_limit) - await self.async_routing_strategy_pre_call_checks(deployment=deployment, parent_otel_span=parent_otel_span) - yield - - async def async_callback_filter_deployments( - self, - model: str, - healthy_deployments: list[dict], - messages: list[AllMessageValues] | None, - parent_otel_span: Span | None, - request_kwargs: dict | None = None, - logging_obj: LiteLLMLogging | None = None, - ): - """ - For usage-based-routing-v2, enables running rpm checks before the call is made, inside the semaphore. - - -> makes the calls concurrency-safe, when rpm limits are set for a deployment - - Returns: - - None - - Raises: - - Rate Limit Exception - If the deployment is over it's tpm/rpm limits - """ - returned_healthy_deployments = healthy_deployments - for _callback in self._routing_callbacks(): - if isinstance(_callback, CustomLogger): - try: - returned_healthy_deployments = await _callback.async_filter_deployments( - model=model, - healthy_deployments=returned_healthy_deployments, - messages=messages, - request_kwargs=request_kwargs, - parent_otel_span=parent_otel_span, - ) - except Exception as e: - ## LOG FAILURE EVENT - if logging_obj is not None: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=e, - traceback_exception=traceback.format_exc(), - prefer_async_handlers=True, - ) - ) - raise e - return returned_healthy_deployments - - @staticmethod - def _json_default_stable_id(value: object) -> str: - """json.dumps default= for generate_model_id: plain str() on an arbitrary - object (e.g. a RoutingPlugin instance) falls back to object.__repr__'s - ``, so the hash -- and deployment id -- would - change every restart. Use the class name instead, stable across restarts.""" - return f"{type(value).__module__}.{type(value).__qualname__}" - - @staticmethod - def _core_generate_model_id(model_group: str, litellm_params: dict) -> str: # mutable-ok: hashed read-only - """The stable id of a deployment without ``model_info.id``; see ``litellm.router_api.deployment_id``.""" - return _deployment_id(model_group, litellm_params) - - @staticmethod - def _inherit_builtin_cache_pricing(model_info: dict, backend_model: str, custom_llm_provider: str | None) -> None: - """Fill missing cache pricing on a custom-priced deployment entry from - the backend model's built-in cost map entry, so a deployment that - only spells out ``input_cost_per_token``/``output_cost_per_token`` - does not silently bill cache_read/cache_creation at 0. - - User-specified cache fields always win; only ``None``/missing entries - are inherited, and a user-set cache_read price also covers explicit - cache hits. No-op when the backend model has no canonical entry. - """ - cache_fields: Final = ( - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr", - "cache_creation_input_token_cost_above_200k_tokens", - "cache_read_input_token_cost", - "cache_read_input_token_cost_above_200k_tokens", - ) - inherited_fields: Final = ( - cache_fields - if model_info.get("cache_read_input_token_cost") is not None - else (*cache_fields, "explicit_cache_read_input_token_cost") - ) - missing_fields: Final = tuple(f for f in inherited_fields if model_info.get(f) is None) - if not missing_fields: - return - try: - backend_info: Final = python_host.invoke("get_model_info")( - model=backend_model, custom_llm_provider=custom_llm_provider - ) - except Exception: - return - for field in missing_fields: - backend_value = backend_info.get(field) - if backend_value is not None: - model_info[field] = backend_value - - @staticmethod - def _cost_map_backend_model(deployment: Deployment | DeploymentTypedDict) -> str: - model_info_base_model: Final = ( - deployment.model_info.base_model - if isinstance(deployment, Deployment) - else _configured_model_info(deployment).get("base_model") - ) - if isinstance(model_info_base_model, str) and model_info_base_model: - return model_info_base_model - params: Final = _LITELLM_PARAMS_ADAPTER.validate_python( - deployment.litellm_params.model_dump(exclude_none=True) - if isinstance(deployment, Deployment) - else deployment["litellm_params"] - ) - params_base_model: Final = params.get("base_model") - if isinstance(params_base_model, str) and params_base_model: - return params_base_model - model: Final = ( - deployment.litellm_params.model - if isinstance(deployment, Deployment) - else deployment["litellm_params"].get("model", "") - ) - invoked_model: Final = params.get("model_id") - if ( - isinstance(invoked_model, str) - and invoked_model - and (params.get("custom_llm_provider") == "bedrock" or model.startswith("bedrock/")) - ): - profile_model: Final = invoked_model if invoked_model.startswith("bedrock/") else f"bedrock/{invoked_model}" - profile_info: Final = lookup_model_info_for_registration(model=profile_model, custom_llm_provider="bedrock") - if not isinstance(profile_info, ModelInfoLookupMiss): - return profile_model - return model - - @staticmethod - def _inherit_builtin_service_tier_pricing( - model_info: dict, # mutable-ok: deployment cost-map entry filled in place - backend_model: str, - custom_llm_provider: str | None, - ) -> None: - """Inherit missing tier rates so a standalone entry does not fall back to custom standard rates.""" - if ptu_terms(model_info) is not None and is_ptu_cost_attribution_enabled(): - return - if all(model_info.get(field) is None for field in ("input_cost_per_token", "output_cost_per_token")): - return - try: - backend_info: Final = python_host.invoke("get_model_info")( - model=backend_model, custom_llm_provider=custom_llm_provider - ) - except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model - return - backend_entry: Final = python_host.dynamic().model_cost.get(backend_info.get("key") or "") - if not isinstance(backend_entry, dict): - return - for field, backend_value in backend_entry.items(): - if not field.endswith(SERVICE_TIER_COST_KEY_SUFFIXES): - continue - if model_info.get(field) is not None or backend_value is None: - continue - model_info[field] = copy.deepcopy(backend_value) - - @staticmethod - def _inherit_builtin_base_rates_for_off_peak( - model_info: dict, # mutable-ok: cost-map entry filled in place - backend_model: str, - custom_llm_provider: str | None, - ) -> None: - """Fill missing pricing fields on a deployment entry that only sets - ``off_peak_pricing``, from the backend model's built-in cost map entry. - - Cost lookup selects the deployment-scoped entry over the shared backend - entry only when the deployment entry carries a base pricing field, and - ``off_peak_pricing`` is deliberately kept off the shared entry, so a - deployment spelling out only its off-peak schedule would otherwise - never receive the discount. The backend model's entire canonical cost - map entry is copied, field by field, so threshold, tiered, - service-tier, cache, character, and per-second rates as well as - companion billing fields like ``web_search_billing_unit`` and the - regional uplift multipliers all carry over, and peak-hour billing - through the deployment entry matches the shared backend entry exactly. - The raw ``litellm.model_cost`` entry is the copy source rather than - ``get_model_info``'s view of it, since that view synthesizes zero flat - token rates for backends without one and storing those would mark a - tiered-only backend explicitly priced free. Values are deep-copied to - keep the builtin entry isolated. User-specified fields always win; - no-op when any base pricing field is already set or the backend model - has no canonical entry. - """ - if not model_info.get("off_peak_pricing"): - return - if any( - model_info.get(field) is not None - for field in ("input_cost_per_token", "input_cost_per_second", "cost_per_second", "tiered_pricing") - ): - return - try: - backend_info: Final = python_host.invoke("get_model_info")( - model=backend_model, custom_llm_provider=custom_llm_provider - ) - except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model - return - backend_entry: Final = python_host.dynamic().model_cost.get(backend_info.get("key") or "") - if not isinstance(backend_entry, dict): - return - for field, backend_value in backend_entry.items(): - if model_info.get(field) is not None or backend_value is None: - continue - model_info[field] = copy.deepcopy(backend_value) - - @staticmethod - def _inherit_builtin_tiered_output_rate( - model_info: dict, backend_model: str, custom_llm_provider: str | None - ) -> None: - """Fill a missing entry-level output rate on a deployment entry whose tier - table omits one, from the backend model's built-in cost map entry. - - A deployment's custom pricing is registered as its own standalone - ``litellm.model_cost`` entry holding only the supplied fields, and the - tiered-cost output fallback reads that same entry, so a tier table that - spells out only input-side rates would bill every completion at 0. - - A user-specified ``output_cost_per_token`` always wins. No-op without a - tier table, when every tier declares its own output rate, or when the - backend model has no canonical entry or no flat output rate: - ``get_model_info`` synthesizes a zero for tiered-only backends, and - storing that zero would mark the deployment as explicitly priced free. - """ - tiers: Final = model_info.get("tiered_pricing") - if not isinstance(tiers, list) or not tiers: - return - if model_info.get("output_cost_per_token") is not None: - return - if all(isinstance(tier, dict) and "output_cost_per_token" in tier for tier in tiers): - return - try: - backend_info: Final = python_host.invoke("get_model_info")( - model=backend_model, custom_llm_provider=custom_llm_provider - ) - except Exception: # noqa: BLE001 # get_model_info raises plain Exception for an unmapped backend model - return - backend_rate: Final = backend_info.get("output_cost_per_token") - if backend_rate: - model_info["output_cost_per_token"] = backend_rate - - def _create_deployment( - self, - deployment_info: dict, - _model_name: str, - _litellm_params: dict, - _model_info: dict, - *, - declared_id: str | None = None, - duplicate_ids: frozenset[str] = frozenset(), - ) -> Deployment | None: - """ - Create a deployment object and add it to the model list - - If the deployment is not active for the current environment, it is ignored - - Returns: - - Deployment: The deployment object - - None: If the deployment is not active for the current environment (if 'supported_environments' is set in litellm_params) - """ - try: - config_sourced: Final = _model_info.get("db_model") is not True - identity_error: Final = ( - ptu_identity_error( - declared_id=declared_id, - taken=declared_id in duplicate_ids, - current_id=_model_info.get("id"), - model_name=_model_name, - ) - if config_sourced and ptu_terms(_model_info) is not None - else None - ) - ptu_error: Final = ( - (ptu_config_error(_model_info, model_name=_model_name) or identity_error) if config_sourced else None - ) - declares_split: Final = declares_ptu_shares(_model_info) - if ptu_error is not None and (declares_split or is_ptu_cost_attribution_enabled()): - raise ValueError(ptu_error) - access_windows_error: Final = access_windows_config_error(_model_info, model_name=_model_name) - if access_windows_error is not None: - raise ValueError(access_windows_error) - capacity_warning: Final = ( - ptu_capacity_warning( - _model_name, - MappingProxyType( - { # pyright: ignore[reportUnknownArgumentType] # router deployment dicts are untyped - "model_info": _model_info, - "litellm_params": _litellm_params, - } - ), - ) - if is_ptu_cost_attribution_enabled() - else None - ) - if capacity_warning is not None: - verbose_router_logger.warning(capacity_warning) - zeroed_pricing: Final = zeroed_ptu_pricing(_model_info, _litellm_params) if config_sourced else None - merged_params: Final[Mapping[str, object]] = ( - _litellm_params if zeroed_pricing is None else MappingProxyType({**_litellm_params, **zeroed_pricing}) - ) - litellm_params: Final[LiteLLM_Params] = validate_litellm_params_strictly( - LiteLLM_Params, dict(**merged_params) - ) - declared_params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(_litellm_params) - warn_on_provider_credential_mismatch(model_name=_model_name, litellm_params=declared_params) - warn_on_ignored_model_capability_flags(model_name=_model_name, litellm_params=declared_params) - warn_on_unenforced_endpoint_flag(model_name=_model_name, model_info=_model_info) - deployment = without_model_capability_flags( - Deployment( - **deployment_info, - model_name=_model_name, - litellm_params=litellm_params, - model_info=_model_info, - ) - ) - refuse_unservable_numbers(deployment.to_json(exclude_none=True, warnings=False)) # pyright: ignore[reportUnknownMemberType] # pydantic's to_json - PythonCore._register_deployment_pricing(deployment) - - ## Check if LLM Deployment is allowed for this deployment - if self.deployment_is_active_for_environment(deployment=deployment) is not True: - verbose_router_logger.warning( - "Ignoring deployment %s as it is not active for environment %s", - deployment.model_name, - deployment.model_info["supported_environments"], - ) - return None - - # Validate tag_regex patterns BEFORE adding the deployment so we never - # have partially-initialised router state if a pattern is invalid. - _tag_regex: Final = deployment.litellm_params.get("tag_regex") or [] - for pattern in _tag_regex: - try: - re.compile(pattern) - except re.error as exc: - raise ValueError( - f"Invalid regex in tag_regex for model '{deployment.model_name}': {pattern!r} — {exc}" - ) from exc - - deployment = self._add_deployment(deployment=with_admin_api_key_secret_reference(deployment)) - - model: Final = deployment.to_json(exclude_none=True) - - self._add_model_to_list_and_index_map(model=model, model_id=deployment.model_info.id) - return deployment - except Exception as e: - if self.ignore_invalid_deployments: - if isinstance(e, litellm_exceptions.BadRequestError): - self._provider_unresolved_deployments = ( - *self._provider_unresolved_deployments, - partial( - self._create_deployment, - deployment_info=deployment_info, - _model_name=_model_name, - _litellm_params=_litellm_params, - _model_info=_model_info, - declared_id=declared_id, - duplicate_ids=duplicate_ids, - ), - ) - verbose_router_logger.exception( - "Error creating deployment: %s, ignoring and continuing with other deployments.", e - ) - return None - else: - raise e - - def _is_auto_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: - """ - Check if the deployment is an auto-router deployment (semantic router). - - Returns True if the litellm_params model starts with "auto_router/" - but NOT "auto_router/complexity_router" or "auto_router/adaptive_router" - (which use the complexity-router and adaptive-router strategies). - """ - return classify_strategy_router_model(litellm_params.model) == "semantic" - - @staticmethod - def _deployment_tags(deployment: Deployment) -> tuple[str, ...]: - """Deployment tags used to disambiguate strategy registries keyed by model_name.""" - return tuple(deployment.litellm_params.tags or ()) - - def init_auto_router_deployment(self, deployment: Deployment): - """ - Initialize the auto-router deployment. - - This will initialize the auto-router and add it to the auto-routers dictionary. - """ - from litellm.router_strategy.auto_router.auto_router import AutoRouter - - auto_router_config_path: Final[str | None] = deployment.litellm_params.auto_router_config_path - auto_router_config: Final[str | None] = deployment.litellm_params.auto_router_config - if auto_router_config_path is None and auto_router_config is None: - raise ValueError( - "auto_router_config_path or auto_router_config is required for auto-router deployments. Please set it in the litellm_params" - ) - - default_model: Final[str | None] = deployment.litellm_params.auto_router_default_model - if default_model is None: - raise ValueError( - "auto_router_default_model is required for auto-router deployments. Please set it in the litellm_params" - ) - - embedding_model: Final[str | None] = deployment.litellm_params.auto_router_embedding_model - if embedding_model is None: - raise ValueError( - "auto_router_embedding_model is required for auto-router deployments. Please set it in the litellm_params" - ) - - max_input_chars: Final[int | None] = deployment.litellm_params.auto_router_max_input_chars - - autor_router: Final[AutoRouter] = AutoRouter( - model_name=deployment.model_name, - auto_router_config_path=auto_router_config_path, - auto_router_config=auto_router_config, - default_model=default_model, - embedding_model=embedding_model, - ctx=self._core_ctx, - max_input_chars=(max_input_chars if max_input_chars is not None else DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS), - ) - self._register_pre_routing_strategy( - registry=self._core_auto_routers, - deployment=deployment, - strategy=autor_router, - strategy_label="Auto-router", - ) - - def _is_complexity_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: - """ - Check if the deployment is a complexity-router deployment. - - Returns True if the litellm_params model starts with "auto_router/complexity_router" - """ - return classify_strategy_router_model(litellm_params.model) == "complexity" - - def _core_config_deployments(self) -> Iterator[Mapping[str, object]]: - """The model_list rows that came from config.yaml rather than the DB (``model_info.db_model`` unset).""" - for deployment in self._core_model_list: - if not isinstance(deployment, Mapping): - continue - model_info = deployment.get("model_info") - if not (isinstance(model_info, Mapping) and model_info.get("db_model")): - yield deployment - - def auto_router_capability_violation(self, capability: GatedAutoRouterCapability) -> str | None: - """ - Why one more router claiming ``capability`` cannot join this router, or None when it can. - - Judged against every deployment currently on the model_list; an upsert pops the row being - edited first, so an edit of an existing gated router keeps its own slot. The limit is - resolved on every call through ``auto_router_capability_limit``; unset means unlimited, - which is the SDK default, and the proxy injects a resolver backed by its license. - """ - limit: Final = self.auto_router_capability_limit() if self.auto_router_capability_limit is not None else None - others: Final = count_capability_routers( - (deployment for deployment in self._core_model_list if isinstance(deployment, Mapping)), - capability=capability, - ) - return capability_limit_violation(capability=capability, held=others + 1, limit=limit) - - def init_complexity_router_deployment(self, deployment: Deployment): - """ - Initialize the complexity-router deployment. - - This will initialize the complexity-router and add it to the complexity-routers dictionary. - """ - # Import here to avoid circular imports — ComplexityRouter is a CustomLogger - # subclass that imports litellm internals which depend on router.py. - # This matches the AutoRouter pattern in init_auto_router_deployment above. - from litellm.router_strategy.complexity_router.complexity_router import ( - ComplexityRouter, - ) - from litellm.router_strategy.complexity_router.config import ( - ComplexityRouterConfig, - ) - - complexity_router_config: Final[dict | None] = deployment.litellm_params.complexity_router_config - for capability in claimed_capabilities(complexity_router_config): - if (limit_violation := self.auto_router_capability_violation(capability)) is not None: - raise ValueError(limit_violation) - - default_model: Final = ( - ComplexityRouterConfig.model_validate(complexity_router_config).resolve_default_model( - deployment.litellm_params.complexity_router_default_model - ) - if complexity_router_config - else deployment.litellm_params.complexity_router_default_model - ) - - if default_model is None: - raise ValueError( - "complexity_router_default_model is required for complexity-router deployments, " - "or configure tiers in complexity_router_config. Please set it in the litellm_params" - ) - - complexity_router: Final[ComplexityRouter] = ComplexityRouter( - model_name=deployment.model_name, - default_model=default_model, - ctx=self._core_ctx, - complexity_router_config=complexity_router_config, - ) - self._register_pre_routing_strategy( - registry=self._core_complexity_routers, - deployment=deployment, - strategy=complexity_router, - strategy_label="Complexity-router", - ) - if complexity_router._uses_deployment_pin: - self._ensure_deployment_affinity_callback() - - def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: - """True when this deployment opts in via the `auto_router/adaptive_router` model prefix.""" - return classify_strategy_router_model(litellm_params.model) == "adaptive" - - def _deployment_participates_in_adaptive_routing(self, litellm_params: LiteLLM_Params) -> bool: - """True when this deployment owns an `adaptive_routers` entry once finalized: - a dedicated adaptive router, or a complexity router whose config enables the - adaptive companion. Mirrors the two arms of - `_finalize_adaptive_router_if_configured`, which is the registry's only writer.""" - if self._is_adaptive_router_deployment(litellm_params=litellm_params): - return True - if not self._is_complexity_router_deployment(litellm_params=litellm_params): - return False - config: Final = litellm_params.complexity_router_config - if not config: - return False - adaptive_flag: Final[object] = config.get("adaptive") - return bool(adaptive_flag) - - @staticmethod - def _has_registered_strategy( - registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]], - model_name: str, - tags: tuple[str, ...], - ) -> bool: - """True when a strategy for this (model_name, tags) pair is already registered.""" - return any(existing.tags == tags for existing in registry.get(model_name, [])) - - def _register_pre_routing_strategy( - self, - registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]], - deployment: Deployment, - strategy: _PreRoutingStrategyT, - strategy_label: str, - ) -> None: - """ - Register `strategy` under `deployment.model_name`, scoped by its tags. - Reusing a `model_name` is allowed when tags differ; a repeat of the same - (model_name, tags) pair is a misconfiguration and is rejected. - """ - tags: Final = self._deployment_tags(deployment) - if self._has_registered_strategy(registry, deployment.model_name, tags): - raise ValueError( - f"{strategy_label} deployment {deployment.model_name} with tags {list(tags)} already exists. " - "Please use a different model name or set different tags." - ) - registry[deployment.model_name] = [ - *registry.get(deployment.model_name, []), - TaggedPreRoutingStrategy( - tags=tags, - strategy=strategy, - definition_fingerprint=self._routing_definition_fingerprint(deployment), - deployment=deployment, - ), - ] - - @staticmethod - def _routing_definition_fingerprint(deployment: Deployment) -> str | None: - """Identify configured routing fields, excluding external files and live strategy state.""" - try: - definition: Final = json.dumps( - deployment.litellm_params.model_dump( - include=STRATEGY_ROUTER_PARAM_FIELDS | {"model", "tags"}, exclude_none=True, warnings=False - ), - sort_keys=True, - default=PythonCore._json_default_stable_id, - ) - except (TypeError, ValueError): - return None - return PythonCore._core_generate_model_id(deployment.model_name, {"routing_definition": definition}) - - @staticmethod - def _unregister_pre_routing_strategy( - registry: dict[str, list[TaggedPreRoutingStrategy[_PreRoutingStrategyT]]], - model_name: str, - tags: tuple[str, ...], - ) -> bool: - """Drop the strategy registered for this exact (model_name, tags) pair, leaving - strategies registered under the same name with different tags in place. Returns - whether anything was actually dropped.""" - existing: Final = registry.get(model_name, []) - remaining: Final = [entry for entry in existing if entry.tags != tags] - if len(remaining) == len(existing): - return False - if remaining: - registry[model_name] = remaining - else: - registry.pop(model_name, None) - return True - - def _unregister_pre_routing_strategy_for_deployment(self, deployment: Deployment) -> None: - """ - Release the pre-routing strategy a deployment holds, so removing it from the - model_list also frees its (model_name, tags) slot. - - Without this, re-adding the deployment (an edit arriving via upsert_deployment, - or a router recreated under a name that was deleted earlier) hits the - "already exists" guard in `_register_pre_routing_strategy`, which - `ignore_invalid_deployments` swallows - the deployment then silently never - makes it back into the model_list. - - Released from every registry rather than the first match, because registration is - one-to-many: a complexity router configured with `adaptive` is also registered in - `adaptive_routers` under the same (model_name, tags) by the deferred finalize pass. - Guarded on the auto_router/ prefix so removing a *regular* deployment can't evict a - router that merely shares its model_name. - """ - if not deployment.litellm_params.model.startswith("auto_router/"): - return - model_name: Final = deployment.model_name - tags: Final = self._deployment_tags(deployment) - for registry in (self._core_auto_routers, self._core_complexity_routers, self._core_quality_routers): - self._unregister_pre_routing_strategy(registry, model_name, tags) - if self._unregister_pre_routing_strategy(self._core_adaptive_routers, model_name, tags): - self._sync_adaptive_router_hooks() - - def _finalize_adaptive_router_if_configured(self) -> None: - """Locate every adaptive-router deployment in the finalized model_list and - build an AdaptiveRouter for each. Safe no-op when none are configured. - Idempotent: skips any deployment whose (model_name, tags) pair is already - initialized, so hot-reloads don't rebuild routers that would lose state.""" - for entry in self._core_model_list or []: - lp = entry.get("litellm_params") if isinstance(entry, dict) else entry.litellm_params - lp_model = (lp.get("model") if isinstance(lp, dict) else lp.model) if lp else None - if not (lp_model and lp_model.startswith("auto_router/adaptive_router")): - continue - model_name = entry.get("model_name") if isinstance(entry, dict) else entry.model_name - if not model_name or not lp: - continue - deployment = Deployment( - model_name=model_name, - litellm_params=(lp if not isinstance(lp, dict) else LiteLLM_Params.model_validate(lp)), - model_info=(entry.get("model_info") if isinstance(entry, dict) else entry.model_info), - ) - if self._has_registered_strategy( - self._core_adaptive_routers, model_name, self._deployment_tags(deployment) - ): - continue - self.init_adaptive_router_deployment(deployment=deployment) - - for model_name, tagged_complexity_routers in self._core_complexity_routers.items(): - for tagged in tagged_complexity_routers: - complexity_router = tagged.strategy - if not complexity_router.config.adaptive: - continue - if self._has_registered_strategy(self._core_adaptive_routers, model_name, tagged.tags): - continue - adaptive_router: AdaptiveRouter | None = complexity_router._ensure_adaptive_router() - if adaptive_router is not None: - self._core_adaptive_routers[model_name] = [ - *self._core_adaptive_routers.get(model_name, []), - TaggedPreRoutingStrategy( - tags=tagged.tags, - strategy=adaptive_router, - definition_fingerprint=tagged.definition_fingerprint, - deployment=tagged.deployment, - ), - ] - - self._sync_adaptive_router_hooks() - - def _sync_adaptive_router_hooks(self) -> None: - """Rebuild the AdaptiveRouterPostCallHook set so it is exactly one hook per - currently registered adaptive router. Run at every point the adaptive registry - changes, otherwise a released router keeps recording turns through its hook.""" - from litellm.router_strategy.adaptive_router.hooks import ( - AdaptiveRouterPostCallHook, - AdaptiveRouterResponseHeaderHook, - ) - - for callback in self._host.callbacks.get_custom_loggers_for_type(AdaptiveRouterPostCallHook): - self._host.callbacks.remove_callback_from_all_lists(callback) - for tagged_adaptive_routers in self._core_adaptive_routers.values(): - for tagged in tagged_adaptive_routers: - self._host.callbacks.add_litellm_callback(AdaptiveRouterPostCallHook(adaptive_router=tagged.strategy)) - if self._core_adaptive_routers: - self._host.callbacks.add_litellm_callback(AdaptiveRouterResponseHeaderHook()) - - def init_adaptive_router_deployment(self, deployment: Deployment) -> None: - """ - Build an AdaptiveRouter instance for this deployment and register its - post-call hook. Multiple adaptive routers can coexist on a single Router, - keyed by `deployment.model_name`. - - `model_to_prefs` and `model_to_cost` are derived from the OTHER models - already registered in `self._core_model_list` whose `model_name` appears in - `available_models`. Models not yet registered fall back to defaults. - """ - # Local import: AdaptiveRouter -> hooks -> classifier all import litellm - # internals which transitively import this module. (AGENTS.md exception clause.) - from litellm.router_strategy.adaptive_router.adaptive_router import ( - AdaptiveRouter, - ) - from litellm.router_strategy.adaptive_router.hooks import ( - AdaptiveRouterPostCallHook, - ) - from litellm.types.router import ( - AdaptiveRouterConfig, - AdaptiveRouterPreferences, - ) - - raw_config: Final = deployment.litellm_params.adaptive_router_config - if raw_config is None: - raise ValueError("adaptive_router_config is required for adaptive-router deployments.") - - config: Final = AdaptiveRouterConfig(**raw_config) - - model_to_prefs: Final[dict[str, AdaptiveRouterPreferences]] = {} - model_to_cost: Final[dict[str, float]] = {} - # O(k) via the name→indices map: only touch deployments whose name - # is listed in `available_models`, instead of scanning model_list. - for name in config.available_models: - deployments = self._core_deployments_named(name) - if not deployments: - continue - d = deployments[0] - mi_dict: dict[str, Any] = d.get("model_info") or {} - prefs_raw = mi_dict.get("adaptive_router_preferences") - if prefs_raw is not None: - model_to_prefs[name] = AdaptiveRouterPreferences(**prefs_raw) - - # model_info is the conventional pricing location elsewhere in LiteLLM; litellm_params wins if set. - cost = d.get("litellm_params", {}).get("input_cost_per_token") - if cost is None: - cost = mi_dict.get("input_cost_per_token") - if cost is not None: - model_to_cost[name] = float(cost) - - adaptive_router: Final = AdaptiveRouter( - router_name=deployment.model_name, - config=config, - model_to_prefs=model_to_prefs, - model_to_cost=model_to_cost, - ) - self._register_pre_routing_strategy( - registry=self._core_adaptive_routers, - deployment=deployment, - strategy=adaptive_router, - strategy_label="Adaptive-router", - ) - self._host.callbacks.add_litellm_callback(AdaptiveRouterPostCallHook(adaptive_router=adaptive_router)) - verbose_router_logger.info( - "AdaptiveRouter[%s] initialized with %d models", - deployment.model_name, - len(config.available_models), - ) - - def _is_quality_router_deployment(self, litellm_params: LiteLLM_Params) -> bool: - """ - Check if the deployment is a quality-router deployment. - - Returns True if the litellm_params model starts with "auto_router/quality_router". - """ - return classify_strategy_router_model(litellm_params.model) == "quality" - - def init_quality_router_deployment(self, deployment: Deployment): - """ - Initialize the quality-router deployment. - - Resolves the default model from either `quality_router_default_model` or - `quality_router_config["default_model"]`, then instantiates the - QualityRouter and stores it in `self._core_quality_routers`. - """ - # Import here to mirror the AutoRouter / ComplexityRouter init pattern - # and avoid circular imports. - from litellm.router_strategy.quality_router.quality_router import ( - QualityRouter, - ) - - quality_router_config: Final[dict | None] = deployment.litellm_params.quality_router_config - - default_model: str | None = deployment.litellm_params.quality_router_default_model - if default_model is None and quality_router_config: - default_model = quality_router_config.get("default_model") - - if default_model is None: - raise ValueError( - "quality_router_default_model is required for quality-router deployments, " - "or set default_model in quality_router_config. Please configure it in the litellm_params" - ) - - quality_router: Final[QualityRouter] = QualityRouter( - model_name=deployment.model_name, - default_model=default_model, - ctx=self._core_ctx, - quality_router_config=quality_router_config, - ) - self._register_pre_routing_strategy( - registry=self._core_quality_routers, - deployment=deployment, - strategy=quality_router, - strategy_label="Quality-router", - ) - - def deployment_is_active_for_environment(self, deployment: Deployment) -> bool: - """ - Function to check if a llm deployment is active for a given environment. Allows using the same config.yaml across multople environments - - Requires `LITELLM_ENVIRONMENT` to be set in .env. Valid values for environment: - - development - - staging - - production - - Raises: - - ValueError: If LITELLM_ENVIRONMENT is not set in .env or not one of the valid values - - ValueError: If supported_environments is not set in model_info or not one of the valid values - """ - return model_info_is_active_for_environment(model_info=deployment.model_info) - - def set_model_list(self, model_list: list): - """Replace every deployment. All or nothing: when an entry raises, the deployments served - before the call are put back and the error is raised.""" - previous: Final = copy.deepcopy(list(self.__dict__.get("_core_model_list") or ())) - try: - self._replace_model_list(model_list) # pyright: ignore[reportUnknownMemberType] # the legacy untyped model list - except Exception: - self._replace_model_list(previous) # pyright: ignore[reportUnknownMemberType] # see above - raise - - def _replace_model_list(self, model_list: list) -> None: # pyright: ignore[reportUnknownParameterType, reportMissingTypeArgument] # the legacy untyped model list - original_model_list: Final = copy.deepcopy(model_list) - self._discovered_model_info_cache.flush_cache() - with self._deployment_index_lock: - self._core_model_list = [] - self.model_id_to_deployment_index_map = {} # Reset the index - self._core_model_name_to_deployment_indices = {} # Reset the model_name index - self._core_team_model_to_deployment_indices = {} # Reset the team_model index - self._core_team_pattern_routers = {} - self._core_team_public_model_names = frozenset() - # Reset per-strategy router registries so hot-reload doesn't leave - # stale routers pointing at the old model_list. - self._core_quality_routers = {} - self._core_complexity_routers = {} - self._core_auto_routers = {} - self._provider_unresolved_deployments = () - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works - - declared_ids: Final = tuple( - str(entry["model_info"]["id"]) - for entry in original_model_list - if isinstance(entry.get("model_info"), dict) and entry["model_info"].get("id") is not None - ) - duplicate_ids: Final = frozenset(model_id for model_id in declared_ids if declared_ids.count(model_id) > 1) - - ptu_declared: Final = tuple( - str(entry.get("model_name")) - for entry in original_model_list - if isinstance(entry.get("model_info"), dict) - and entry["model_info"].get("db_model") is not True - and declares_ptu(entry["model_info"]) - ) - if ptu_declared and not is_ptu_cost_attribution_enabled(): - verbose_router_logger.warning( - "PTU fields are set on config.yaml deployment(s) %s, but PTU cost attribution is disabled, so no " - "flat cost accrues and this traffic is billed per token. Set %s=True to enable it", - ", ".join(ptu_declared), - PTU_COST_ATTRIBUTION_ENV_VAR, - ) - - for model in original_model_list: - _model_name = model.pop("model_name") - _litellm_params = model.pop("litellm_params") - ## check if litellm params in os.environ - if isinstance(_litellm_params, dict): - for k, v in _litellm_params.items(): - if isinstance(v, str) and v.startswith("os.environ/") and not holds_secret_pointer(k): - _litellm_params[k] = get_secret(v) - - _model_info: dict = model.pop("model_info", {}) - - declared_id = None if _model_info.get("id") is None else str(_model_info["id"]) - - # check if model info has id - if "id" not in _model_info: - _id = self._core_generate_model_id(_model_name, _litellm_params) - _model_info["id"] = _id - - if _litellm_params.get("organization", None) is not None and isinstance( - _litellm_params["organization"], list - ): # Addresses https://github.com/BerriAI/litellm/issues/3949 - for org in _litellm_params["organization"]: - _litellm_params["organization"] = org - self._create_deployment( - deployment_info=model, - _model_name=_model_name, - _litellm_params=_litellm_params, - _model_info=_model_info, - declared_id=declared_id, - duplicate_ids=duplicate_ids, - ) - else: - self._create_deployment( - deployment_info=model, - _model_name=_model_name, - _litellm_params=_litellm_params, - _model_info=_model_info, - declared_id=declared_id, - duplicate_ids=duplicate_ids, - ) - - verbose_router_logger.debug("\nInitialized Model List %s", self._core_get_model_names()) - self._core_model_names = {m["model_name"] for m in model_list} - - # Note: model_name_to_deployment_indices is already built incrementally - # by _create_deployment -> _add_model_to_list_and_index_map - - # Deferred: build the AdaptiveRouter strategy now that all underlying - # deployments have been registered. - self._finalize_adaptive_router_if_configured() - - def _add_deployment(self, deployment: Deployment) -> Deployment: - import os - - #### VALIDATE MODEL ######## - # Check if this is a prompt management model before validating as LLM provider - litellm_model: Final = deployment.litellm_params.model - if isinstance(deployment.litellm_params.drop_params, str): - verbose_router_logger.warning( - "model=%s drop_params=%r is not a flag value, treating it as unset", - deployment.model_name, - deployment.litellm_params.drop_params, - ) - is_prompt_management_model = False - - if "/" in litellm_model: - split_litellm_model: Final = litellm_model.split("/")[0] - if split_litellm_model in self._host.dynamic().known_custom_logger_compatible_callbacks: - is_prompt_management_model = True - - if is_prompt_management_model: - # For prompt management models, skip LLM provider validation - # The actual model will be resolved at runtime from the prompt file - _model = litellm_model - custom_llm_provider = None - dynamic_api_key = None - api_base = None - elif declared_authenticating_provider(litellm_model, deployment.litellm_params.custom_llm_provider) is not None: - _model, custom_llm_provider = ( # rebind-ok: each branch resolves the deployment provider - self._model_group_llm_provider(deployment.litellm_params) - ) - else: - # check if model provider in supported providers - ( - _model, - custom_llm_provider, - dynamic_api_key, - api_base, - ) = self._host.invoke("get_llm_provider")( - model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.get("custom_llm_provider", None), - api_base=deployment.litellm_params.api_base, - litellm_params=deployment.litellm_params, - ) - # done reading model["litellm_params"] - # Check if provider is supported: either in enum or JSON-configured - if ( - custom_llm_provider not in self._host.dynamic().provider_list - and not JSONProviderRegistry.exists(custom_llm_provider) - and not is_registered_custom_provider(custom_llm_provider) - ): - raise Exception(f"Unsupported provider - {custom_llm_provider}") - - #### DEPLOYMENT NAMES INIT ######## - self._core_deployment_names.append(deployment.litellm_params.model) - ############ Users can either pass tpm/rpm as a litellm_param or a router param ########### - # for get_available_deployment, we use the litellm_param["rpm"] - # in this snippet we also set rpm to be a litellm_param - if deployment.litellm_params.rpm is None and getattr(deployment, "rpm", None) is not None: - deployment.litellm_params.rpm = getattr(deployment, "rpm") - - if deployment.litellm_params.tpm is None and getattr(deployment, "tpm", None) is not None: - deployment.litellm_params.tpm = getattr(deployment, "tpm") - - # Check if user is trying to use model_name == "*" - # this is a catch all model for their specific api key - # if deployment.model_name == "*": - # if deployment.litellm_params.model == "*": - # # user wants to pass through all requests to litellm.acompletion for unknown deployments - # self._core_router_general_settings.pass_through_all_models = True - # else: - # self._core_default_deployment = deployment.to_json(exclude_none=True) - # Check if user is using provider specific wildcard routing - # example model_name = "databricks/*" or model_name = "anthropic/*" - if "*" in deployment.model_name: - # store this as a regex pattern - all deployments matching this pattern will be sent to this deployment - # Store deployment.model_name as a regex pattern - self._core_pattern_router.add_pattern(deployment.model_name, deployment.to_json(exclude_none=True)) - if deployment.model_info.id: - self._core_provider_default_deployment_ids.append(deployment.model_info.id) - - _team_id: Final = deployment.model_info.get("team_id") - _team_public_model_name: Final[str | None] = deployment.model_info.get("team_public_model_name") - if _team_id is not None and _team_public_model_name is not None and "*" in _team_public_model_name: - if _team_id not in self._core_team_pattern_routers: - self._core_team_pattern_routers[_team_id] = PatternMatchRouter() - self._core_team_pattern_routers[_team_id].add_pattern( - _team_public_model_name, deployment.to_json(exclude_none=True) - ) - - # Azure GPT-Vision Enhancements, users can pass os.environ/ - data_sources: Final = deployment.litellm_params.get("dataSources", []) or [] - - for data_source in data_sources: - params = data_source.get("parameters", {}) - for param_key in ["endpoint", "key"]: - # if endpoint or key set for Azure GPT Vision Enhancements, check if it's an env var - if param_key in params and params[param_key].startswith("os.environ/"): - env_name = params[param_key].replace("os.environ/", "") - params[param_key] = os.environ.get(env_name, "") - - # # init OpenAI, Azure clients - # InitalizeOpenAISDKClient.set_client( - # litellm_router_instance=self, model=deployment.to_json(exclude_none=True) - # ) - - if custom_llm_provider is not None: - self._initialize_deployment_for_pass_through( - deployment=deployment, - custom_llm_provider=custom_llm_provider, - ) - - ######################################################### - # Check if this is an auto-router deployment - ######################################################### - if self._is_auto_router_deployment(litellm_params=deployment.litellm_params): - self.init_auto_router_deployment(deployment=deployment) - - ######################################################### - # Check if this is a complexity-router deployment - ######################################################### - if self._is_complexity_router_deployment(litellm_params=deployment.litellm_params): - self.init_complexity_router_deployment(deployment=deployment) - - # NOTE: adaptive-router deployments are deferred to the end of - # set_model_list() because their init needs visibility into the OTHER - # deployments listed in `available_models` (which may not yet have - # been processed when this one is created). - ######################################################### - # Check if this is a quality-router deployment - ######################################################### - if self._is_quality_router_deployment(litellm_params=deployment.litellm_params): - self.init_quality_router_deployment(deployment=deployment) - - return deployment - - def _initialize_deployment_for_pass_through(self, deployment: Deployment, custom_llm_provider: str): - """ - Optional: Register vertex credentials for pass-through endpoints if `deployment.litellm_params.use_in_pass_through` is True - - Other providers need no registration here: PassthroughEndpointRouter.get_credentials resolves their credentials per-request from the live router deployments - """ - if deployment.litellm_params.use_in_pass_through is not True: - return - if custom_llm_provider != "vertex_ai": - return - credential_name: Final = deployment.litellm_params.litellm_credential_name - credential_values: Final = ( - CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else {} - ) - from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES - - vertex_project: Final = ( - cast("str | None", credential_values.get("vertex_project")) # cast-ok: vertex credential values are str - or deployment.litellm_params.vertex_project - ) - vertex_location: Final = ( - cast("str | None", credential_values.get("vertex_location")) # cast-ok: vertex credential values are str - or deployment.litellm_params.vertex_location - ) - vertex_credentials: Final = ( - cast( # cast-ok: vertex_credentials holds the typed credential union - "VERTEX_CREDENTIALS_TYPES | None", credential_values.get("vertex_credentials") - ) - or deployment.litellm_params.vertex_credentials - ) - - if vertex_project is None or vertex_location is None: - raise ValueError( - "vertex_project, and vertex_location must be set in litellm_params for pass-through endpoints." - ) - router_host().add_vertex_pass_through_credentials( - project_id=vertex_project, - location=vertex_location, - vertex_credentials=vertex_credentials, - ) - - def add_deployment(self, deployment: Deployment) -> Deployment | None: - """ - Parameters: - - deployment: Deployment - the deployment to be added to the Router - - Returns: - - The added deployment - - OR None (if deployment already exists) - """ - # check if deployment already exists - - _deployment_model_id: Final = deployment.model_info.id - if _deployment_model_id and self._core_has_model_id(_deployment_model_id): - return None - - declared_params: Final = deployment.litellm_params.model_dump(exclude_none=True) - warn_on_provider_credential_mismatch(model_name=deployment.model_name, litellm_params=declared_params) - warn_on_ignored_model_capability_flags(model_name=deployment.model_name, litellm_params=declared_params) - warn_on_unenforced_endpoint_flag( - model_name=deployment.model_name, model_info=deployment.model_info.model_dump() - ) - served: Final = with_admin_api_key_secret_reference(without_model_capability_flags(deployment)) - - # add to model list - _deployment: Final = served.to_json(exclude_none=True) - refuse_unservable_numbers(_deployment) - # initialize client - self._add_deployment(deployment=served) - - PythonCore._register_deployment_pricing(deployment=served) - - # add to model names - self._add_model_to_list_and_index_map(model=_deployment, model_id=served.model_info.id) - self._core_model_names.add(served.model_name) - self._sync_deployment_budget_config(deployment=served) - return served - - def _update_deployment_indices_after_removal(self, model_id: str, removal_idx: int) -> None: - """ - Helper method to update deployment indices after a deployment has been removed from model_list. - - Parameters: - - model_id: str - the id of the deployment that was removed - - removal_idx: int - the index where the deployment was removed from model_list - """ - self._discovered_model_info_cache.delete_cache(model_id) - # Update indices for all models after the removed one - for deployment_id, idx in self.model_id_to_deployment_index_map.items(): - if idx > removal_idx: - self.model_id_to_deployment_index_map[deployment_id] = idx - 1 - # Remove the deleted model from index - if model_id in self.model_id_to_deployment_index_map: - del self.model_id_to_deployment_index_map[model_id] - - # Update model_name_to_deployment_indices - for model_name, indices in list(self._core_model_name_to_deployment_indices.items()): - # Build new list without mutating the original - updated_indices = [] - for idx in indices: - if idx == removal_idx: - # Skip the removed index - continue - elif idx > removal_idx: - # Decrement indices after removal - updated_indices.append(idx - 1) - else: - # Keep indices before removal unchanged - updated_indices.append(idx) - - # Update or remove the entry - if len(updated_indices) > 0: - self._core_model_name_to_deployment_indices[model_name] = updated_indices - else: - del self._core_model_name_to_deployment_indices[model_name] - self._core_model_names.discard(model_name) - - # Update team_model_to_deployment_indices - for key, indices in list(self._core_team_model_to_deployment_indices.items()): - # Build new list without mutating the original - updated_indices = [] - for idx in indices: - if idx == removal_idx: - # Skip the removed index - continue - elif idx > removal_idx: - # Decrement indices after removal - updated_indices.append(idx - 1) - else: - # Keep indices before removal unchanged - updated_indices.append(idx) - - # Update or remove the entry - if len(updated_indices) > 0: - self._core_team_model_to_deployment_indices[key] = updated_indices - else: - del self._core_team_model_to_deployment_indices[key] - self._core_team_public_model_names = frozenset( - public_model_name for _, public_model_name in self._core_team_model_to_deployment_indices - ) - - self._core_pattern_router.remove_deployment(model_id) - for team_id in list(self._core_team_pattern_routers.keys()): - team_pattern_router = self._core_team_pattern_routers[team_id] - team_pattern_router.remove_deployment(model_id) - if not team_pattern_router.patterns: - del self._core_team_pattern_routers[team_id] - - def _update_team_model_index(self, model: dict, idx: int) -> None: - """ - Helper to update team_model_to_deployment_indices for a single deployment. - - Parameters: - - model: dict - the deployment to index - - idx: int - the index in model_list - """ - team_id: Final = (model.get("model_info") or {}).get("team_id") - team_public_model_name: Final = (model.get("model_info") or {}).get("team_public_model_name") - if team_id and team_public_model_name: - key: Final = (team_id, team_public_model_name) - self._core_team_public_model_names = self._core_team_public_model_names | frozenset( - {team_public_model_name} - ) - if key not in self._core_team_model_to_deployment_indices: - self._core_team_model_to_deployment_indices[key] = [] - if idx not in self._core_team_model_to_deployment_indices[key]: - self._core_team_model_to_deployment_indices[key].append(idx) - - def _add_model_to_list_and_index_map(self, model: dict, model_id: str | None = None) -> None: - """ - Helper method to add a model to the model_list and update both indices. - - Parameters: - - model: dict - the model to add to the list - - model_id: Optional[str] - the model ID to use for indexing. If None, will try to get from model["model_info"]["id"] - """ - with self._deployment_index_lock: - idx: Final = len(self._core_model_list) - self._core_model_list.append(model) - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - - # Update model_id index for O(1) lookup - if model_id is not None: - self.model_id_to_deployment_index_map[model_id] = idx - elif model.get("model_info", {}).get("id") is not None: - self.model_id_to_deployment_index_map[model["model_info"]["id"]] = idx - - # Update model_name index for O(1) lookup - model_name: Final = model.get("model_name") - if model_name: - if model_name not in self._core_model_name_to_deployment_indices: - self._core_model_name_to_deployment_indices[model_name] = [] - self._core_model_name_to_deployment_indices[model_name].append(idx) - - # Update team_model index for O(1) team-scoped lookup - self._update_team_model_index(model, idx) - - def _core_upsert_deployment(self, deployment: Deployment) -> Deployment | None: - """ - Add or update deployment - Parameters: - - deployment: Deployment - the deployment to be added to the Router - - Returns: - - The added/updated deployment - """ - _deployment_model_id: Final = deployment.model_info.id or "" - _deployment_on_router: Final[Deployment | None] = self._core_get_deployment(model_id=_deployment_model_id) - served: Final = without_model_capability_flags(deployment) - # `_create_deployment` lists the model name before it can still fail; a dropped deployment - # takes its name back off the list. - names_before: Final = len(self._core_deployment_names) - try: - if _deployment_on_router is not None: - # deployment with this model_id exists on the router - if ( - deployment.model_name == _deployment_on_router.model_name - and served.litellm_params == _deployment_on_router.litellm_params - and deployment.model_info == _deployment_on_router.model_info - ): - # No need to update - return None - - # if there is a new litellm param -> then update the deployment - # remove the previous deployment - with self._deployment_index_lock: - removal_idx: Final = self.model_id_to_deployment_index_map.get(_deployment_model_id) - if removal_idx is not None: - self._core_model_list.pop(removal_idx) - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - self._update_deployment_indices_after_removal( - model_id=_deployment_model_id, removal_idx=removal_idx - ) - - # Free the outgoing deployment's pre-routing strategy slot (keyed by the - # OLD model_name/tags) before the re-add below re-registers it. - self._unregister_pre_routing_strategy_for_deployment(deployment=_deployment_on_router) - - # if the model_id is not in router - self.add_deployment(deployment=deployment) - # add_deployment() builds every strategy EXCEPT the adaptive one, which - # set_model_list() defers until the whole model_list is visible. Re-run that - # deferred pass so an adaptive router whose slot was just released above is - # rebuilt rather than left unregistered. - if self._deployment_participates_in_adaptive_routing(litellm_params=deployment.litellm_params) or ( - _deployment_on_router is not None - and self._deployment_participates_in_adaptive_routing( - litellm_params=_deployment_on_router.litellm_params - ) - ): - self._finalize_adaptive_router_if_configured() - return deployment - except Exception as e: - if self.ignore_invalid_deployments: - verbose_router_logger.warning( - "Error upserting deployment %s (id=%s): %s. Dropping it and continuing with other deployments.", - deployment.model_name, - deployment.model_info.id, - e, - ) - del self._core_deployment_names[names_before:] - self._restore_deployment_after_failed_upsert( - previous_deployment=_deployment_on_router, model_id=_deployment_model_id - ) - return None - else: - raise e - - def _restore_deployment_after_failed_upsert(self, previous_deployment: Deployment | None, model_id: str) -> None: - """Put a deployment back the way it was before a failed upsert popped it. - - A rollback re-admits state that was already serving, so it does not go through the - capability ceiling a newcomer gets: with the ceiling tightened since the deployment first - registered, judging the rollback would drop a serving router over an unrelated failed edit. - """ - if previous_deployment is None or self._core_has_model_id(model_id): - return - limit_resolver: Final = self.auto_router_capability_limit - self.auto_router_capability_limit = None - try: - self.add_deployment(deployment=previous_deployment) - verbose_router_logger.info( - "Restored deployment %s (id=%s); it keeps serving its previous configuration.", - previous_deployment.model_name, - model_id, - ) - except Exception as restore_error: # noqa: BLE001 # best-effort restore: a second failure must not abort the reload - verbose_router_logger.warning( - "Could not restore previously served deployment %s (id=%s) after the failed upsert: %s", - previous_deployment.model_name, - model_id, - restore_error, - ) - finally: - self.auto_router_capability_limit = limit_resolver - - @staticmethod - def _backend_cost_map_keys(model: str, custom_llm_provider: str | None) -> tuple[str, ...]: - """The ``litellm.model_cost`` keys a deployment's shared backend info is registered under.""" - backend_key: Final = model if custom_llm_provider is None else f"{custom_llm_provider}/{model}" - if "responses/" in backend_key: - return (backend_key, backend_key.replace("responses/", "")) - return (backend_key,) - - @staticmethod - def _catalog_entry_has_token_limit(catalog_entry: Mapping[str, object], field: str) -> bool: - return catalog_entry.get(field) is not None or ( - field == "max_tokens" and catalog_entry.get("max_output_tokens") is not None - ) - - @staticmethod - def _shared_backend_model_info_for_keys( - catalog: Mapping[str, Mapping[str, object]], - backend_keys: tuple[str, ...], - shared_model_info: Mapping[str, object], - ) -> Mapping[str, object]: - catalog_token_limit_fields: Final = frozenset( - field - for field in DEPLOYMENT_SCOPED_TOKEN_LIMIT_FIELDS - if any(PythonCore._catalog_entry_has_token_limit(catalog.get(key) or {}, field) for key in backend_keys) - ) - return {field: value for field, value in shared_model_info.items() if field not in catalog_token_limit_fields} - - @staticmethod - def _deployment_model_cost_payload(deployment: Deployment) -> dict[str, object]: # mutable-ok: cost-map entry - """The ``model_info`` a deployment contributes to ``litellm.model_cost``. - - Custom pricing lives on ``litellm_params`` rather than ``model_info``, and - the built-in cache-pricing inheritance is derived rather than stored, so - both are folded back in here. That keeps this reproducible from a - deployment alone, which is what lets a refresh rebuild the same entries. - """ - model_info: Final = dict( - _MODEL_INFO_ADAPTER.validate_python(deployment.model_info.model_dump(exclude_none=True)) - ) - for field in CustomPricingLiteLLMParams.model_fields: - field_value = deployment.litellm_params.get(field) - if field_value is not None: - model_info[field] = field_value - backend_model: Final = PythonCore._cost_map_backend_model(deployment) - pricing_provider: Final = ( - deployment.model_info.pricing_provider or deployment.litellm_params.custom_llm_provider - ) - params: Final = _LITELLM_PARAMS_ADAPTER.validate_python(deployment.litellm_params.model_dump(exclude_none=True)) - if ( - params.get("model_id") - and backend_model != deployment.litellm_params.model - and ( - deployment.litellm_params.custom_llm_provider == "bedrock" - or deployment.litellm_params.model.startswith("bedrock/") - ) - ): - profile_info: Final = lookup_model_info_for_registration(model=backend_model, custom_llm_provider="bedrock") - if not isinstance(profile_info, ModelInfoLookupMiss): - explicit_info: Final = _MODEL_INFO_ADAPTER.validate_python(model_info) - priced_info: Final = model_info_with_pricing_overrides(profile_info, explicit_info) - model_info.update( - {key: value for key, value in priced_info.items() if is_server_derived_pricing_key(key)} - ) - PythonCore._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider=pricing_provider, - ) - if model_info.get("input_cost_per_token") is not None: - PythonCore._inherit_builtin_cache_pricing( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider=pricing_provider, - ) - PythonCore._inherit_builtin_service_tier_pricing( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider=pricing_provider, - ) - PythonCore._inherit_builtin_tiered_output_rate( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider=pricing_provider, - ) - return model_info - - @staticmethod - def _register_deployment_pricing(deployment: Deployment) -> None: - """Register a deployment's custom/inherited pricing in ``litellm.model_cost``. - - Takes only a ``Deployment``, so it registers pricing for a deployment that - is never added to ``self._core_model_list`` (a per-request client-side-credential - deployment) just as readily as one that is. - """ - PythonCore._register_deployment_in_model_cost( - model_id=deployment.model_info.id, - model_info=PythonCore._deployment_model_cost_payload(deployment), - model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - pricing_provider=deployment.model_info.pricing_provider, - ) - - @staticmethod - def _register_deployment_in_model_cost( - *, - model_id: str | None, - model_info: dict, # mutable-ok: cost-map entry - model: str, - custom_llm_provider: str | None, - pricing_provider: str | None = None, - ) -> None: - """Write a deployment's metadata into ``litellm.model_cost``. - - Runs when a deployment is added and again after a price data reload, so - the entries a refresh rebuilds are the ones a fresh boot would produce. - Nothing is recorded for replay: a refresh walks the live routers instead, - so a deleted, repointed or never-added deployment, and a discarded router, - drop out of the rebuild on their own. - - A strategy-router alias is never the deployment actually called or - billed, so custom pricing configured on it must not become a cost-map - price: an explicit zero would let ``_is_cost_explicitly_configured`` - treat the alias as a genuinely free model and waive budget checks for - requests that route to (and bill as) a real deployment. - """ - if classify_strategy_router_model(model) is not None: - model_info = {k: v for k, v in model_info.items() if k not in CustomPricingLiteLLMParams.model_fields} - - if model_id is not None: - python_host.invoke("register_model")( - model_cost={model_id: model_info}, - persist_across_reloads=False, - warning_display_name=model, - custom_llm_provider=pricing_provider or custom_llm_provider, - add_to_provider_model_sets=False, - ) - - ## OLD MODEL REGISTRATION ## Kept to prevent breaking changes - backend_keys: Final = PythonCore._backend_cost_map_keys(model=model, custom_llm_provider=custom_llm_provider) - model_provider_prefix: Final = model.partition("/")[0] - catalog_provider: Final = ( - custom_llm_provider - if custom_llm_provider is not None - else model_provider_prefix - if model_provider_prefix in LlmProvidersSet - else None - ) - potential_names: Final = ( - get_potential_model_names(model, catalog_provider) if catalog_provider is not None else None - ) - catalog_backend_keys: Final = ( - ( - *backend_keys, - model, - potential_names["split_model"], - potential_names["stripped_model_name"], - potential_names["combined_model_name"], - potential_names["combined_stripped_model_name"], - potential_names["provider_prefixed_model_name"], - potential_names["region_free_combined_model_name"], - ) - if potential_names is not None - else (*backend_keys, model) - ) - - # For the shared backend key, keep only cost-map schema fields - # (minus custom pricing) so that one deployment's pricing overrides - # or custom metadata (id, access_via_team_ids, arbitrary keys) - # don't pollute another deployment sharing the same backend model - # name. Each deployment's full model_info is already stored under - # its unique model_id above. - catalog: Final = catalog_with_runtime_registrations() - shared_model_info: Final = PythonCore._shared_backend_model_info_for_keys( - catalog=catalog, - backend_keys=catalog_backend_keys, - shared_model_info=shared_backend_model_info(model_info), - ) - python_host.invoke("register_model")( - model_cost={_key: keep_registered_mode(_key, shared_model_info) for _key in backend_keys}, - persist_across_reloads=False, - ) - - def replay_model_cost_registrations(self) -> None: - """Re-assert this router's deployments onto a freshly fetched catalog. - - Reads ``model_list`` at call time, so only deployments the router still - serves are restored, plus any config deployment the fresh catalog now resolves. - """ - provider_unresolved: Final = self._provider_unresolved_deployments - self._provider_unresolved_deployments = () - for create_deployment in provider_unresolved: - create_deployment() - for entry in tuple(self._core_model_list): - try: - deployment = entry if isinstance(entry, Deployment) else Deployment(**entry) - except Exception: # noqa: BLE001 # a malformed entry must not abort the rest of the rebuild - verbose_router_logger.exception( - "Router: could not rebuild cost-map entry for a deployment during a price data reload" - ) - continue - PythonCore._register_deployment_in_model_cost( - model_id=deployment.model_info.id, - model_info=PythonCore._deployment_model_cost_payload(deployment), - model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - pricing_provider=deployment.model_info.pricing_provider, - ) - self._invalidate_model_group_info_cache() - - _replay_model_cost_registrations = replay_model_cost_registrations - - def _core_delete_deployment(self, id: str) -> Deployment | None: - """ - Parameters: - - id: str - the id of the deployment to be deleted - - Returns: - - The deleted deployment - - OR None (if deleted deployment not found) - """ - try: - with self._deployment_index_lock: - deployment_idx: Final = self.model_id_to_deployment_index_map.get(id) - if deployment_idx is None: - return None - # Pop the item from the list first - item: Final = self._core_model_list.pop(deployment_idx) - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - self._update_deployment_indices_after_removal(model_id=id, removal_idx=deployment_idx) - _budget_limiter: Final = self._core_get_router_deployment_budget_limiter() - if _budget_limiter is not None: - _budget_limiter.unregister_deployment_budget(model_id=id) - try: - self._unregister_pre_routing_strategy_for_deployment( - deployment=item if isinstance(item, Deployment) else Deployment(**item) - ) - except Exception: - verbose_router_logger.exception( - "delete_deployment: could not release pre-routing strategies for model_id=%s; " - "the deployment is out of the model_list and its indices are repaired", - id, - ) - except Exception: - return None - return item - - def _core_get_router_deployment_budget_limiter( - self, - ) -> RouterBudgetLimiting | None: - """ - Return the router's deployment-budget callback. - - Uses exact-type matching so proxy subclasses (e.g. virtual-key model budgets) - registered on litellm.callbacks are not mistaken for router deployment budgets. - """ - if self.router_budget_logger is not None: - return self.router_budget_logger - - if self.optional_callbacks: - for _cb in self.optional_callbacks: - if isinstance(_cb, RouterBudgetLimiting) and type(_cb) is RouterBudgetLimiting: - self.router_budget_logger = _cb - return _cb - return None - - def _deployment_has_budget_limits(self, deployment: Deployment) -> bool: - return ( - deployment.litellm_params.get("max_budget") is not None - and deployment.litellm_params.get("budget_duration") is not None - and deployment.model_info.id is not None - ) - - def _sync_deployment_budget_config(self, deployment: Deployment) -> None: - model_id: Final = deployment.model_info.id - if model_id is None: - return - - _budget_limiter = self._core_get_router_deployment_budget_limiter() - - if not self._deployment_has_budget_limits(deployment=deployment): - if _budget_limiter is not None: - _budget_limiter.unregister_deployment_budget(model_id=model_id) - return - - if _budget_limiter is None: - self.add_optional_pre_call_checks(optional_pre_call_checks=["router_budget_limiting"]) - _budget_limiter = self._core_get_router_deployment_budget_limiter() - - if _budget_limiter is not None: - _budget_limiter.register_deployment_budget(deployment=deployment.to_json(exclude_none=True)) - - def _core_deployments_named(self, model_name: str) -> tuple[DeploymentTypedDict, ...] | None: - """The deployments indexed under `model_name`, read under the index lock, or None when the name is not indexed.""" - with self._deployment_index_lock: - indices: Final = self._core_model_name_to_deployment_indices.get(model_name) - return None if indices is None else tuple(self._core_model_list[idx] for idx in indices) - - def _served_openai_models(self, model_name: str) -> tuple[str | None, ...]: - return served_openai_models(self._core_deployments_named(model_name) or ()) - - def team_public_model_names_of(self, team_id: str) -> frozenset[str]: - """The public names `team_id` serves its own deployments under.""" - with self._deployment_index_lock: - return frozenset(name for owner, name in self._core_team_model_to_deployment_indices if owner == team_id) - - def _team_deployments_named(self, team_id: str, public_model_name: str) -> tuple[DeploymentTypedDict, ...] | None: - """The deployments `team_id` serves under `public_model_name`, or None when that pair is not indexed.""" - with self._deployment_index_lock: - indices: Final = self._core_team_model_to_deployment_indices.get((team_id, public_model_name)) - return None if indices is None else tuple(self._core_model_list[idx] for idx in indices) - - def _core_get_deployment(self, model_id: str) -> Deployment | None: - """ - Returns -> Deployment or None - - Raise Exception -> if model found in invalid format - """ - # Use O(1) lookup via model_id_to_deployment_index_map only - with self._deployment_index_lock: - idx: Final = self.model_id_to_deployment_index_map.get(model_id) - model: Final = None if idx is None else self._core_model_list[idx] - if model is not None: - if isinstance(model, dict): - return Deployment(**model) - elif isinstance(model, Deployment): - return model - else: - raise Exception(f"Model invalid format - {type(model)}") - - return None - - def _core_get_deployment_credentials(self, model_id: str) -> dict | None: - """ - Returns -> dict of credentials for a given model id. - - Returns None if the deployment is paused via `LiteLLM_ProxyModelTable.blocked`, - so file/batch/passthrough callers that resolve credentials directly cannot keep - using a paused deployment. - """ - deployment: Final = self._core_get_deployment(model_id=model_id) - if deployment is None or self._core_is_deployment_blocked(deployment): - return None - return CredentialLiteLLMParams.model_validate( - deployment.litellm_params.model_dump(exclude_none=True) - ).model_dump(exclude_none=True) - - def _core_get_deployment_by_model_group_name(self, model_group_name: str) -> Deployment | None: - """ - Returns -> Deployment or None - - Raise Exception -> if model found in invalid format - - Optimized with O(1) index lookup instead of O(n) linear scan. - """ - # O(1) lookup in model_name index - with self._deployment_index_lock: - indices: Final = self._core_model_name_to_deployment_indices.get(model_group_name) - # Return first deployment for this model_name - model: Final = self._core_model_list[indices[0]] if indices else None - if model is not None: - if isinstance(model, dict): - return Deployment(**model) - elif isinstance(model, Deployment): - return model - else: - raise Exception(f"Model Name invalid - {type(model)}") - return None - - @staticmethod - def _core_deployment_usable_by_team(model: Mapping[str, object] | Deployment, team_id: str | None) -> bool: - """ - A team-scoped deployment (``model_info.team_id`` set) is only usable by - callers from that same team; deployments without a team owner are shared. - """ - model_info: Final = model.get("model_info") if isinstance(model, dict) else model.model_info - owner_team_id: Final = cast(Mapping[str, object], model_info).get("team_id") if model_info is not None else None - return team_may_use_deployment(owner_team_id, team_id) - - def _get_model_group_deployment_usable_by_team( - self, model_group_name: str, team_id: str | None - ) -> Deployment | None: - """ - Like ``get_deployment_by_model_group_name``, but skips deployments owned - by other teams so a shared model name never resolves another team's - credentials. - """ - with self._deployment_index_lock: - indices: Final = self._core_model_name_to_deployment_indices.get(model_group_name) or () - usable: Final = ( - self._core_model_list[idx] - for idx in indices - if self._core_deployment_usable_by_team(self._core_model_list[idx], team_id) - ) - first_usable: Final = next(usable, None) - if first_usable is None: - return None - return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable - - async def _core_arefresh_model_info(self, *, client: AsyncHTTPHandler | None = None) -> None: - """Refresh token limits advertised by configured OpenAI-compatible deployments.""" - deployments: Final = iter(tuple(self._core_model_list)) - - async def refresh_worker() -> None: - for raw_deployment in deployments: - try: - await self._arefresh_deployment_model_info(raw_deployment, client=client) - except Exception: # noqa: BLE001 # one invalid deployment must not prevent refreshing the others - verbose_router_logger.debug("Could not refresh deployment model info") - - await asyncio.gather(*(refresh_worker() for _ in range(MODEL_INFO_REFRESH_CONCURRENCY))) - self._invalidate_model_group_info_cache() - - async def _arefresh_deployment_model_info( - self, raw_deployment: Mapping[str, object], *, client: AsyncHTTPHandler | None - ) -> None: - deployment: Final = Deployment.model_validate(raw_deployment) - params: Final = LiteLLM_Params.model_validate( - MappingProxyType( - { - **deployment.litellm_params.model_dump(exclude_none=True), - **( - self._core_get_deployment_credentials_with_provider(deployment.model_info.id or "") - or MappingProxyType({}) - ), - } - ) - ) - model, provider, dynamic_api_key, api_base = self._host.invoke("get_llm_provider")( - model=params.model, litellm_params=params - ) - if provider not in MODEL_INFO_DISCOVERY_PROVIDERS: - return - if api_base is None or "*" in model or params.get("use_clientside_credentials"): - return - api_key: Final = await aresolve_api_key_secret_reference( - params.api_key or dynamic_api_key, - model=model, - llm_provider=provider, - upstream=MappingProxyType({"api_base": params.api_base}), - ) - headers: Final = TypeAdapter(Mapping[str, str]).validate_python( - params.get("extra_headers") or params.get("headers") or MappingProxyType({}) - ) - auth_headers: Final = ( - MappingProxyType({"authorization": f"Bearer {api_key}"}) if api_key else MappingProxyType({}) - ) - limits: Final = await get_openai_compatible_model_info( - model=model, - api_base=api_base, - headers=MappingProxyType( - { - **auth_headers, - **MappingProxyType({key.lower(): value for key, value in headers.items()}), - } - ), - client=client or get_async_httpx_client(llm_provider=LlmProviders.OPENAI), - cache=self._core_cache.in_memory_cache, - ) - model_id: Final = deployment.model_info.id - if not limits or model_id is None or self._core_get_model_info(model_id) is not raw_deployment: - return - self._discovered_model_info_cache.max_size_in_memory = max(len(self._core_model_list), 1) - self._discovered_model_info_cache.delete_cache(model_id) - self._discovered_model_info_cache.set_cache( - model_id, DiscoveredDeploymentModelInfo(deployment=raw_deployment, limits=limits) - ) - self._invalidate_model_group_info_cache() - - def _core_get_discovered_model_info(self, model_id: str | None) -> Mapping[str, int]: - cached: Final[object] = self._discovered_model_info_cache.get_cache(model_id) - if ( - model_id is not None - and isinstance(cached, DiscoveredDeploymentModelInfo) - and cached.deployment is self._core_get_model_info(model_id) - ): - configured: Final = TypeAdapter(Mapping[str, object]).validate_python(cached.deployment["model_info"]) - return MappingProxyType({key: value for key, value in cached.limits.items() if configured.get(key) is None}) - return MappingProxyType({}) - - def _core_get_model_listing_info(self, model_name: str) -> DeploymentModelListingInfo | None: - """ - Return what the concrete deployments behind model_name contribute to its - /v1/models entry: the cost-map keys for their underlying models, plus the widest - configured or discovered token limits. Resolved via O(1) index - lookup. - - Returns None for wildcard-expanded or unknown names, where the listed name is the - real model name and no deployment-specific information exists, and treats a - malformed configured limit as absent rather than failing the listing. - - The whole group is read rather than just its first deployment, so a group that - mixes models does not advertise a window that depends on config order; the widest - one is reported, which is what get_model_group_info already shows the Admin UI. - Keys are deduplicated, so the ordinary group of interchangeable deployments of one - model still costs the caller a single cost-map lookup. Unlike get_model_group_info, - this never triggers pattern matching or deep copies, so it is safe to call per - listed model on the /v1/models hot path. - """ - deployments: Final = self._core_deployments_named(model_name) - if not deployments: - return None - - model_infos: Final = tuple( - MappingProxyType( - { - **self._core_get_discovered_model_info( - (deployment.get("model_info") or MappingProxyType({})).get("id") - ), - **MappingProxyType( - { - k: v - for k, v in (deployment.get("model_info") or MappingProxyType({})).items() - if v is not None - } - ), - } - ) - for deployment in deployments - ) - raw_deployments: Final[tuple[object, ...]] = deployments - pricing_keys: Final = (deployment_pricing_models(deployment) for deployment in raw_deployments) - cost_map_keys: Final = tuple(dict.fromkeys(keys[0] for keys in pricing_keys if keys)) - return DeploymentModelListingInfo( - cost_map_keys=cost_map_keys, - max_input_tokens=self._widest_configured_limit(model_infos, "max_input_tokens"), - max_output_tokens=self._widest_configured_limit(model_infos, "max_output_tokens"), - ) - - @staticmethod - def _widest_configured_limit(model_infos: Sequence[Mapping[str, object]], field: str) -> int | None: - """The largest usable value of ``field`` across a group's configured model_info blocks.""" - limits: Final = tuple( - limit - for limit in (coerce_token_limit(model_info.get(field)) for model_info in model_infos) - if limit is not None - ) - return max(limits) if limits else None - - def _core_get_configured_token_limits(self, model_name: str) -> "tuple[int | None, int | None]": - """ - Return (max_input_tokens, max_output_tokens) configured or discovered for a concrete - deployment of model_name, via O(1) index lookup. - - Returns (None, None) for wildcard-expanded or unknown names, and treats a - malformed configured value as absent rather than failing the caller. - - Deliberately reads one deployment rather than aggregating the group the way - get_model_listing_info does: its caller truncates an embedding input to this - value, so the widest window in a mixed group would be the wrong answer there. - """ - deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name) - if deployment is None: - return (None, None) - - model_info: Final = MappingProxyType( - { - **self._core_get_discovered_model_info(deployment.model_info.id), - **deployment.model_info.model_dump(exclude_none=True), - } - ) - return ( - coerce_token_limit(model_info.get("max_input_tokens")), - coerce_token_limit(model_info.get("max_output_tokens")), - ) - - def _core_get_configured_mode(self, model_name: str) -> "str | None": - """Return the mode explicitly configured for a concrete deployment.""" - deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name) - if deployment is None: - return None - - mode: Final = deployment.model_info.get("mode") - if isinstance(mode, str) and mode.strip(): - return mode - return None - - def _core_get_configured_display_name(self, model_name: str) -> "str | None": - """ - Return the display_name explicitly configured in a concrete deployment's - model_info for model_name, via O(1) index lookup. - - Returns None for wildcard-expanded or unknown names, and treats a - non-string or empty configured value as absent rather than failing the - listing. Like get_configured_token_limits, this never triggers pattern - matching or deep copies, so it is safe to call per listed model on the - /v1/models hot path. - """ - deployment: Final = self._core_get_deployment_by_model_group_name(model_group_name=model_name) - if deployment is None: - return None - - display_name: Final = deployment.model_info.get("display_name") - if isinstance(display_name, str) and display_name.strip(): - return display_name - return None - - def _core_routable_model_group(self, model_name: str) -> str: - """ - The model group a request to model_name routes to: its target when - model_name is a `model_group_alias`, else model_name itself. - """ - target: Final = self._core_get_model_from_alias(model_name) - return target if target is not None else model_name - - def _routable_deployments(self, model_name: str, team_id: str | None) -> tuple[DeploymentTypedDict, ...]: - """ - The deployments a request from team_id to model_name can route to, in - model_list order and via the same O(1) index lookup as - get_configured_display_name, selected the way routing selects them: a - `model_group_alias` reads its target's deployments, another team's - deployment of the name is left out, a caller with no team reads the - deployments no team owns (every deployment of the name when a team owns - each one, which is what an admin's request routes to), and one an admin - paused via `LiteLLM_ProxyModelTable.blocked` is left out. - - Returns an empty tuple for wildcard-expanded or unknown names. - """ - named: Final = self._core_get_all_deployments( - model_name=self._core_routable_model_group(model_name), team_id=team_id - ) - usable: Final = tuple( - deployment for deployment in named if self._core_deployment_usable_by_team(deployment, team_id) - ) - return tuple( - deployment - for deployment in (usable or named) - if _configured_model_info(deployment).get("blocked") is not True - ) - - def _core_get_configured_service_tiers(self, model_name: str, team_id: str | None = None) -> tuple[object, ...]: - """ - Return the service_tiers value each deployment a request from team_id - to model_name can route to (see _routable_deployments) configures in - its model_info, unvalidated and in model_list order, None for a - deployment that sets none; the caller validates the shape it expects - and decides how the deployments combine. - - Returns an empty tuple for wildcard-expanded or unknown names. - """ - return tuple( - _configured_model_info(deployment).get("service_tiers") - for deployment in self._routable_deployments(model_name, team_id) - ) - - def _core_get_configured_max_input_tokens( - self, model_name: str, team_id: str | None = None - ) -> tuple[int | None, ...]: - """ - Return the max_input_tokens each deployment a request from team_id to - model_name can route to (see _routable_deployments) sets in its own - model_info, in model_list order, None for a deployment that sets none - or a malformed one. Unlike get_model_listing_info, neither the cost map - nor a discovered limit is read, so only an operator's own value counts, - and a value a saved /model/info response copied from the cost map is - skipped too. - - Returns an empty tuple for wildcard-expanded or unknown names. - """ - return tuple( - _operator_max_input_tokens(deployment) for deployment in self._routable_deployments(model_name, team_id) - ) - - def _core_get_search_tool_support(self, model_name: str, team_id: str | None = None) -> tuple[bool | None, ...]: - """ - Whether each deployment a request from team_id to model_name can route - to (see _routable_deployments) answers Codex's client `tool_search`: - its model_info.supports_search_tool when that is a bool, else True when - its route builds an Anthropic Messages request, else None (unknown). - """ - return tuple(_search_tool_support(deployment) for deployment in self._routable_deployments(model_name, team_id)) - - def _core_get_routable_upstream_model(self, model_name: str, team_id: str | None = None) -> str | None: - """ - Return the `litellm_params.model` of the first deployment a request - from team_id to model_name can route to (see _routable_deployments). - - Returns None for wildcard-expanded or unknown names, and when the team - can route to no deployment of the name. - """ - deployment: Final = next(iter(self._routable_deployments(model_name, team_id)), None) - return deployment["litellm_params"].get("model") if deployment is not None else None - - def _core_get_credential_deployment(self, model_id: str, team_id: str | None = None) -> Deployment | None: - """ - The deployment a passthrough endpoint (files, batches, etc.) resolves for a - model id or model name: by deployment id first, then by model_name, then by - the team's exact public model name, then by wildcard pattern (team wildcards - before global ones, so a global "openai/*" never shadows the team's own - entry). Name and wildcard lookups never resolve another team's deployment. - - Returns None when nothing matches or the match is paused via - `LiteLLM_ProxyModelTable.blocked`, so callers cannot bypass an admin pause - by resolving the deployment directly. - """ - deployment: Final = ( - self._core_get_deployment(model_id=model_id) - or self._get_model_group_deployment_usable_by_team(model_group_name=model_id, team_id=team_id) - or self._get_team_public_name_deployment(model_id=model_id, team_id=team_id) - or self._get_wildcard_deployment_usable_by_team(model_id=model_id, team_id=team_id) - ) - if deployment is None or self._core_is_deployment_blocked(deployment): - return None - return deployment - - def _get_team_public_name_deployment(self, model_id: str, team_id: str | None) -> Deployment | None: - if team_id is None: - return None - with self._deployment_index_lock: - team_indices: Final = self._core_team_model_to_deployment_indices.get((team_id, model_id)) - team_model: Final = self._core_model_list[team_indices[0]] if team_indices else None - if team_model is None: - return None - return Deployment(**team_model) if isinstance(team_model, dict) else team_model - - def _get_wildcard_deployment_usable_by_team(self, model_id: str, team_id: str | None) -> Deployment | None: - team_pattern_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - team_wildcard_models: Final = team_pattern_router.route(model_id) if team_pattern_router else None - global_wildcard_models: Final = tuple( - wildcard_model - for wildcard_model in (self._core_pattern_router.route(model_id) or ()) - if self._core_deployment_usable_by_team(wildcard_model, team_id) - ) - potential_wildcard_models: Final = team_wildcard_models or global_wildcard_models - if not potential_wildcard_models: - return None - wildcard_deployment: Final = potential_wildcard_models[0] - if isinstance(wildcard_deployment, dict): - return Deployment(**wildcard_deployment) - if isinstance(wildcard_deployment, Deployment): - return wildcard_deployment - return None - - def _core_get_deployment_credentials_with_provider( - self, model_id: str, team_id: str | None = None - ) -> dict[str, Any] | None: - """ - Get API credentials and provider info from a model name in model_list. - Useful for passthrough endpoints (files, batches, etc.) that need credentials. - - Resolves the deployment with `get_credential_deployment` (by deployment id, - then model_name, team public model name, and wildcard pattern). - - Args: - model_id: Model ID or model name from model_list (e.g., "gpt-4o-litellm") - team_id: Optional team id of the caller. When set, team-scoped - deployments (indexed by team public model name, including team - wildcard models like "openai/*") are also considered. Name and - wildcard lookups never resolve a deployment owned by a - different team, so shared model names can't leak another - team's credentials. - - Returns: - Dictionary containing api_key, api_base, custom_llm_provider, etc. - Returns None if model not found, or if the resolved deployment is - paused via `LiteLLM_ProxyModelTable.blocked` (so passthrough callers - cannot bypass an admin pause by resolving credentials directly). - - Example: - credentials = router.get_deployment_credentials_with_provider("gpt-4o-litellm") - # Returns: {"api_key": "sk-...", "custom_llm_provider": "openai", "model": "gpt-4o", ...} - """ - deployment: Final = self._core_get_credential_deployment(model_id=model_id, team_id=team_id) - if deployment is None: - return None - - # Get basic credentials - credentials: Final = CredentialLiteLLMParams.model_validate( - deployment.litellm_params.model_dump(exclude_none=True) - ).model_dump(exclude_none=True) - - # Resolve litellm_credential_name to actual credentials - if deployment.litellm_params.litellm_credential_name is not None: - credential_values: Final = CredentialAccessor.get_credential_values( - deployment.litellm_params.litellm_credential_name - ) - if not credential_values: - verbose_router_logger.warning( - "Credential '%s' not found in credential_list", deployment.litellm_params.litellm_credential_name - ) - credentials.update(credential_values) - # Remove the credential name since we've resolved it - credentials.pop("litellm_credential_name", None) - - credentials["model"] = deployment.litellm_params.model - - # Add custom_llm_provider - if deployment.litellm_params.custom_llm_provider: - credentials["custom_llm_provider"] = deployment.litellm_params.custom_llm_provider - elif "/" in deployment.litellm_params.model: - # Extract provider from "provider/model" format - credentials["custom_llm_provider"] = deployment.litellm_params.model.split("/")[0] - else: - credentials["custom_llm_provider"] = "openai" # default - - return credentials - - @overload - def _core_get_router_model_info( - self, - deployment: Union[dict, "Deployment"], - received_model_name: str, - id: None = None, - ) -> ModelMapInfo: - pass - - @overload - def _core_get_router_model_info(self, deployment: None, received_model_name: str, id: str) -> ModelMapInfo: - pass - - def _core_get_router_model_info( - self, - deployment: Union[dict, "Deployment"] | None, - received_model_name: str, - id: str | None = None, - ) -> ModelMapInfo: - """ - For a given model id, return the model info (max tokens, input cost, output cost, etc.). - - Augment litellm info with additional params set in `model_info`. - - For azure models, ignore the `model:`. Only set max tokens, cost values if base_model is set. - For every other provider, a base_model naming the same model as `model:` can only widen its context window. - - Returns - - ModelInfo - If found -> typed dict with max tokens, input cost, etc. - - Raises: - - ValueError -> If model is not mapped yet - """ - if id is not None: - _deployment: Final = self._core_get_deployment(model_id=id) - if _deployment is not None: - deployment = _deployment - - if deployment is None: - raise ValueError("Deployment not found") - - ## GET BASE MODEL - base_model = (deployment.get("model_info") or {}).get("base_model", None) - if base_model is None: - base_model = (deployment.get("litellm_params") or {}).get("base_model", None) - - model = base_model - - ## GET PROVIDER - reuse LiteLLM_Params if already constructed - litellm_params_data: Final = deployment.get("litellm_params") - litellm_params: LiteLLM_Params - if isinstance(litellm_params_data, LiteLLM_Params): - litellm_params = litellm_params_data - elif isinstance(litellm_params_data, dict) and "model" in litellm_params_data: - litellm_params = LiteLLM_Params.model_validate(litellm_params_data) - else: - raise ValueError( - f"Deployment missing valid litellm_params. " - f"Got: {type(litellm_params_data).__name__}, " - f"deployment_id: {(deployment.get('model_info') or {}).get('id', 'unknown')}" - ) - _model, custom_llm_provider, _, _ = self._host.invoke("get_llm_provider")( - model=litellm_params.model, - litellm_params=litellm_params, - ) - - ## SET MODEL TO 'model=' - if base_model is None + not azure - if custom_llm_provider == "azure" and base_model is None: - # Router init auto-registers every deployment name into - # litellm.model_cost as a zeroed stub, so membership alone can't - # tell a resolvable name apart; require usable limits/costs. - _azure_fallback_key = _model if _model.startswith("azure/") else f"azure/{_model}" - _fallback_entry = self._host.dynamic().model_cost.get(_azure_fallback_key) - _fallback_resolves = isinstance(_fallback_entry, Mapping) and ( - (_fallback_entry.get("max_input_tokens") or 0) > 0 - or (_fallback_entry.get("input_cost_per_token") or 0) > 0 - ) - if _fallback_resolves: - verbose_router_logger.debug( - "Azure deployment '%s' has no base_model set; using '%s' from the model cost map for max tokens, cost tracking, etc.", - _model, - _azure_fallback_key, - ) - else: - verbose_router_logger.error( - "Could not identify azure model '%s'. Set azure 'base_model' for accurate max tokens, cost tracking, etc.- https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models", - _model, - ) - elif custom_llm_provider != "azure": - model = _model - - if "*" in model: # only call pattern_router for wildcard models - potential_models: Final = self._core_pattern_router.route(received_model_name) - if potential_models is not None: - for potential_model in potential_models: - try: - if (potential_model.get("model_info") or {}).get("id") == ( - deployment.get("model_info") or {} - ).get("id"): - model = (potential_model.get("litellm_params") or {}).get("model") - break - except Exception: - pass - - ## GET LITELLM MODEL INFO - raises exception, if model is not mapped - if model is None: - # Handle case where base_model is None (e.g., Azure models without base_model set) - # Use the original model from litellm_params - model = _model - - if not model.startswith(f"{custom_llm_provider}/"): - model_info_name = f"{custom_llm_provider}/{model}" - else: - model_info_name = model - - model_info: Final = self._host.invoke("get_model_info")(model=model_info_name) - if model_info is None: - return model_info - - ## CHECK USER SET MODEL INFO - raw_user_model_info: Final = deployment.get("model_info") - user_model_info: Final = _MODEL_INFO_ADAPTER.validate_python( - raw_user_model_info.model_dump(exclude_none=True) - if isinstance(raw_user_model_info, BaseModel) - else raw_user_model_info or {} - ) - - # get_model_info() hands back an lru_cache'd dict, so merge into a copy; unset - # values are skipped or Deployment's None pricing defaults would erase the map's - merged_model_info: Final[ModelMapInfo] = { - **copy.deepcopy(model_info), - **_base_model_widened_window( - model_info, - _declared_base_model(user_model_info) or _declared_base_model(litellm_params.model_extra or {}), - custom_llm_provider, - ), - **self._core_get_discovered_model_info((deployment.get("model_info") or {}).get("id")), - **MappingProxyType({key: value for key, value in user_model_info.items() if value is not None}), - } - - return merged_model_info - - def _core_get_model_info(self, id: str) -> dict | None: - """ - For a given model id, return the model info - - Returns - - dict: the model in list with 'model_name', 'litellm_params', Optional['model_info'] - - None: could not find deployment in list - - Optimized with O(1) index lookup instead of O(n) linear scan. - """ - # O(1) lookup via model_id_to_deployment_index_map - with self._deployment_index_lock: - idx: Final = self.model_id_to_deployment_index_map.get(id) - return None if idx is None else self._core_model_list[idx] - - def _core_get_model_group(self, id: str) -> list | None: - """ - Return list of all models in the same model group as that model id - """ - - model_info: Final = self._core_get_model_info(id=id) - if model_info is None: - return None - - model_name: Final = model_info["model_name"] - return self._core_get_model_list(model_name=model_name) - - def _core_get_deployment_model_info(self, model_id: str, model_name: str) -> ModelInfo | None: - """ - For a given model id, return the model info - - 1. Check if model_id is in model info - 2. If not, check if litellm model name is in model info - 3. If not, return None - """ - from litellm.utils import cost_map_omits_token_price, update_dictionary - - model_info: ModelInfo | None = None - custom_model_info: dict | None = None - litellm_model_name_model_info: ModelInfo | None = None - base_model_key: str | None = None - - try: - if model_id in self._host.dynamic().model_cost: - cost_entry: Final = self._host.dynamic().model_cost.get(model_id) or MappingProxyType({}) - if not isinstance(cost_entry, Mapping): - raise TypeError(f"model_cost[{model_id!r}] is not a mapping") - custom_model_info = { - **copy.deepcopy(cost_entry), - **self._core_get_discovered_model_info(model_id), - } - except Exception: - pass - - try: - litellm_model_name_model_info = self._host.invoke("get_model_info")(model=model_name) - except Exception: - pass - - ## check for base model - try: - if custom_model_info is not None: - base_model: Final = custom_model_info.get("base_model", None) - if base_model is not None: - ## update litellm model info with base model info - base_model_info: Final = copy.deepcopy(self._host.invoke("get_model_info")(model=base_model)) - if base_model_info is not None: - base_model_key = base_model_info.get("key") - # Base model provides defaults, custom model info overrides - custom_model_info = update_dictionary( - cast(dict, base_model_info), - custom_model_info, - ) - except Exception: - pass - - # Three mutually exclusive scenarios for the model's metadata: - if custom_model_info is not None and litellm_model_name_model_info is not None: - # (1) It has both custom model_info set and exists in the built-in map - # merge with custom overriding built-in - model_info = cast( - ModelInfo, - update_dictionary( - copy.deepcopy(cast(dict, litellm_model_name_model_info)), - custom_model_info, - ), - ) - elif litellm_model_name_model_info is not None: - # (2) Built-in only — no custom pricing to merge - model_info = copy.deepcopy(litellm_model_name_model_info) - elif custom_model_info is not None: - # (3) Custom only — model not in built-in cost map yet - # custom_model_info already includes base_model defaults at this point, if applicable - model_info = cast(ModelInfo, custom_model_info) - - if model_info is None: - return None - builtin_key: Final = ( - litellm_model_name_model_info.get("key") if litellm_model_name_model_info is not None else None - ) - if cost_map_omits_token_price(model_id, builtin_key, base_model_key): - return cast( # cast-ok: TypedDict spread with overridden keys loses its type - ModelInfo, {**model_info, "input_cost_per_token": None, "output_cost_per_token": None} - ) - return model_info - - @staticmethod - def _model_group_llm_provider(litellm_params: LiteLLM_Params) -> tuple[str, str]: - try: - return resolve_model_and_provider_for_metadata( - model=litellm_params.model, - custom_llm_provider=litellm_params.custom_llm_provider, - ) - except litellm_exceptions.BadRequestError as e: - verbose_router_logger.error("litellm.router.py::get_model_group_info() - %s", e) - return "", "" - - def _set_model_group_info( - self, model_group: str, user_facing_model_group_name: str, team_id: str | None = None - ) -> ModelGroupInfo | None: - """ - For a given model group name, return the combined model info - - Returns: - - ModelGroupInfo if able to construct a model group - - None if error constructing model group info - """ - model_group_info: ModelGroupInfo | None = None - - total_tpm: int | None = None - total_rpm: int | None = None - total_itpm: int | None = None - total_otpm: int | None = None - configurable_clientside_auth_params: CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS = None - reasoning_efforts_initialized = False - reasoning_efforts_unknown = False - serves_team_deployments: Final = team_id is not None and bool( - self._team_deployments_named(team_id, model_group) - ) - model_list: Final = self._core_get_model_list_of_routed_group(model_group, team_id) - model_group_description: Final[str | None] = next( - (d for d in (_configured_model_info(m).get("description") for m in model_list) if isinstance(d, str)), - None, - ) - for model in model_list: - is_match = False - if ( - serves_team_deployments - or "model_name" in model - and model["model_name"] == model_group - or "model_name" in model - and self._core_pattern_router.route(model_group) is not None - ): # exact match - is_match = True - - if not is_match: - continue - # model in model group found # - litellm_params = LiteLLM_Params(**model["litellm_params"]) - # get configurable clientside auth params - configurable_clientside_auth_params = litellm_params.configurable_clientside_auth_params - - # Cache nested dict access to avoid repeated temporary dict allocations - model_litellm_params = model.get("litellm_params", {}) - model_info_dict = model.get("model_info", {}) - - # get model tpm - _deployment_tpm: int | None = None - if _deployment_tpm is None: - _deployment_tpm = model.get("tpm", None) - if _deployment_tpm is None: - _deployment_tpm = model_litellm_params.get("tpm", None) - if _deployment_tpm is None: - _deployment_tpm = model_info_dict.get("tpm", None) - - # get model rpm - _deployment_rpm: int | None = None - if _deployment_rpm is None: - _deployment_rpm = model.get("rpm", None) - if _deployment_rpm is None: - _deployment_rpm = model_litellm_params.get("rpm", None) - if _deployment_rpm is None: - _deployment_rpm = model_info_dict.get("rpm", None) - - _deployment_itpm: int | None = model.get("itpm") - if _deployment_itpm is None: - _deployment_itpm = model_litellm_params.get("itpm", None) - if _deployment_itpm is None: - _deployment_itpm = model_info_dict.get("itpm", None) - - _deployment_otpm: int | None = model.get("otpm") - if _deployment_otpm is None: - _deployment_otpm = model_litellm_params.get("otpm", None) - if _deployment_otpm is None: - _deployment_otpm = model_info_dict.get("otpm", None) - - # get model info - try: - model_id = model_info_dict.get("id", None) - if model_id is not None: - model_info = self._core_get_deployment_model_info( - model_id=model_id, model_name=litellm_params.model - ) - else: - model_info = None - except Exception: - model_info = None - - deployment_is_mapped = deployment_is_catalog_mapped(model_info, model_info_dict) - - # get llm provider - litellm_model, llm_provider = self._model_group_llm_provider(litellm_params) - - if model_info is None: - supported_openai_params = self._host.invoke("get_supported_openai_params")( - model=litellm_model, custom_llm_provider=llm_provider - ) - if supported_openai_params is None: - supported_openai_params = [] - - # Get mode from database model_info if available, otherwise default to "chat" - db_model_info = model.get("model_info", {}) - mode = db_model_info.get("mode", "chat") - input_cost_per_token = _cost_value_as_float(db_model_info.get("input_cost_per_token")) - output_cost_per_token = _cost_value_as_float(db_model_info.get("output_cost_per_token")) - - model_info = ModelMapInfo( - key=model_group, - max_tokens=None, - max_input_tokens=None, - max_output_tokens=None, - input_cost_per_token=input_cost_per_token, - output_cost_per_token=output_cost_per_token, - litellm_provider=llm_provider, - mode=mode, - supported_openai_params=supported_openai_params, - supports_system_messages=None, - ) - - if model_group_info is None: - model_group_info = ModelGroupInfo( - **{ - "model_group": user_facing_model_group_name, - "providers": [llm_provider], - "description": model_group_description, - **model_info, - "supports_continuation": model_info.get("supports_continuation") is True, - "supports_fast_mode": True, - "supported_reasoning_efforts": None, - "supported_endpoints": None, - } - ) - else: - # if max_input_tokens > curr - # if max_output_tokens > curr - # if input_cost_per_token > curr - # if output_cost_per_token > curr - # supports_parallel_function_calling == True - # supports_vision == True - # supports_function_calling == True - if llm_provider not in model_group_info.providers: - model_group_info.providers.append(llm_provider) - if ( - model_info.get("max_input_tokens", None) is not None - and model_info["max_input_tokens"] is not None - and ( - model_group_info.max_input_tokens is None - or model_info["max_input_tokens"] > model_group_info.max_input_tokens - ) - ): - model_group_info.max_input_tokens = model_info["max_input_tokens"] - if ( - model_info.get("max_output_tokens", None) is not None - and model_info["max_output_tokens"] is not None - and ( - model_group_info.max_output_tokens is None - or model_info["max_output_tokens"] > model_group_info.max_output_tokens - ) - ): - model_group_info.max_output_tokens = model_info["max_output_tokens"] - _input_cost_per_token = _cost_value_as_float(model_info.get("input_cost_per_token")) - if _input_cost_per_token is not None and ( - model_group_info.input_cost_per_token is None - or _input_cost_per_token > (model_group_info.input_cost_per_token or 0.0) - ): - model_group_info.input_cost_per_token = _input_cost_per_token - _output_cost_per_token = _cost_value_as_float(model_info.get("output_cost_per_token")) - if _output_cost_per_token is not None and ( - model_group_info.output_cost_per_token is None - or _output_cost_per_token > (model_group_info.output_cost_per_token or 0.0) - ): - model_group_info.output_cost_per_token = _output_cost_per_token - if ( - model_info.get("supports_parallel_function_calling", None) is not None - and model_info["supports_parallel_function_calling"] is True - ): - model_group_info.supports_parallel_function_calling = True - if model_info.get("supports_vision", None) is not None and model_info["supports_vision"] is True: - model_group_info.supports_vision = True - if ( - model_info.get("supports_function_calling", None) is not None - and model_info["supports_function_calling"] is True - ): - model_group_info.supports_function_calling = True - if ( - model_info.get("supports_web_search", None) is not None - and model_info["supports_web_search"] is True - ): - model_group_info.supports_web_search = True - if ( - model_info.get("supports_url_context", None) is not None - and model_info["supports_url_context"] is True - ): - model_group_info.supports_url_context = True - if model_info.get("supports_continuation") is True: - model_group_info.supports_continuation = True - - if model_info.get("supports_reasoning", None) is not None and model_info["supports_reasoning"] is True: - model_group_info.supports_reasoning = True - if ( - model_info.get("supported_openai_params", None) is not None - and model_info["supported_openai_params"] is not None - ): - model_group_info.supported_openai_params = model_info["supported_openai_params"] - if model_info.get("tpm", None) is not None and _deployment_tpm is None: - _deployment_tpm = model_info.get("tpm") - if model_info.get("rpm", None) is not None and _deployment_rpm is None: - _deployment_rpm = model_info.get("rpm") - - model_group_info.supported_endpoints = ( - tuple( - dict.fromkeys( - ( - *(model_group_info.supported_endpoints or ()), - *( - deployment_supported_endpoints( - model_info.get("supported_endpoints"), model, llm_provider - ) - if team_may_use_deployment(_configured_model_info(model).get("team_id"), team_id) - else () - ), - ) - ) - ) - or None - ) - model_group_info.supports_fast_mode = model_group_info.supports_fast_mode and ( - AnthropicModelInfo.supports_fast_mode(litellm_model, llm_provider) - ) - deployment_reasoning_efforts = resolve_supported_reasoning_efforts( - model_info, deployment_is_mapped=deployment_is_mapped - ) - if deployment_reasoning_efforts is None: - reasoning_efforts_unknown = True - model_group_info.supported_reasoning_efforts = None - elif not reasoning_efforts_initialized: - reasoning_efforts_initialized = True - if not reasoning_efforts_unknown: - model_group_info.supported_reasoning_efforts = deployment_reasoning_efforts - elif not reasoning_efforts_unknown: - model_group_info.supported_reasoning_efforts = intersect_supported_reasoning_efforts( - model_group_info.supported_reasoning_efforts, - deployment_reasoning_efforts, - ) - - if _deployment_tpm is not None: - if total_tpm is None: - total_tpm = 0 - total_tpm += _deployment_tpm - - if _deployment_rpm is not None: - if total_rpm is None: - total_rpm = 0 - total_rpm += _deployment_rpm - - if _deployment_itpm is not None: - if total_itpm is None: - total_itpm = 0 - total_itpm += _deployment_itpm - - if _deployment_otpm is not None: - if total_otpm is None: - total_otpm = 0 - total_otpm += _deployment_otpm - if model_group_info is not None: - ## UPDATE WITH TOTAL TPM/RPM FOR MODEL GROUP - if total_tpm is not None: - model_group_info.tpm = total_tpm - - if total_rpm is not None: - model_group_info.rpm = total_rpm - - if total_itpm is not None: - model_group_info.itpm = total_itpm - - if total_otpm is not None: - model_group_info.otpm = total_otpm - - ## UPDATE WITH CONFIGURABLE CLIENTSIDE AUTH PARAMS FOR MODEL GROUP - if configurable_clientside_auth_params is not None: - model_group_info.configurable_clientside_auth_params = configurable_clientside_auth_params - - return model_group_info - - def _core_get_model_group_info( - self, model_group: str, *, include_hidden: bool = False, team_id: str | None = None - ) -> ModelGroupInfo | None: - """ - For a given model group name, return the combined model info - - Returns: - - ModelGroupInfo if able to construct a model group - - None if error constructing model group info or hidden model group (unless include_hidden) - """ - ## Check if model group alias - if model_group in self._core_model_group_alias: - item: Final = self._core_model_group_alias[model_group] - if isinstance(item, str): - _router_model_group = item - elif isinstance(item, dict): - if item["hidden"] is True and not include_hidden: - return None - else: - _router_model_group = item["model"] - else: - return None - - return self._set_model_group_info( - model_group=_router_model_group, - user_facing_model_group_name=model_group, - team_id=team_id, - ) - - ## Check if actual model - return self._set_model_group_info( - model_group=model_group, user_facing_model_group_name=model_group, team_id=team_id - ) - - async def _core_get_model_group_usage(self, model_group: str) -> tuple[int | None, int | None]: - """ - Returns current tpm/rpm usage for model group - - Parameters: - - model_group: str - the received model name from the user (can be a wildcard route). - - Returns: - - usage: Tuple[tpm, rpm] - """ - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - tpm_keys: Final[list[str]] = [] - rpm_keys: Final[list[str]] = [] - - model_list: Final = self._core_get_model_list(model_name=model_group) - if model_list is None: # no matching deployments - return None, None - - for model in model_list: - id: str | None = model.get("model_info", {}).get("id") - litellm_model: str | None = model["litellm_params"].get( - "model" - ) # USE THE MODEL SENT TO litellm.completion() - consistent with how global_router cache is written. - if id is None or litellm_model is None: - continue - tpm_keys.append( - RouterCacheEnum.TPM.value.format( - id=id, - model=litellm_model, - current_minute=current_minute, - ) - ) - rpm_keys.append( - RouterCacheEnum.RPM.value.format( - id=id, - model=litellm_model, - current_minute=current_minute, - ) - ) - combined_tpm_rpm_keys: Final = tpm_keys + rpm_keys - - combined_tpm_rpm_values: Final = await self._core_cache.async_batch_get_cache(keys=combined_tpm_rpm_keys) - if combined_tpm_rpm_values is None: - return None, None - - tpm_usage_list: Final[list | None] = combined_tpm_rpm_values[: len(tpm_keys)] - rpm_usage_list: Final[list | None] = combined_tpm_rpm_values[len(tpm_keys) :] - - ## TPM - tpm_usage: int | None = None - if tpm_usage_list is not None: - for t in tpm_usage_list: - if isinstance(t, int): - if tpm_usage is None: - tpm_usage = 0 - tpm_usage += t - ## RPM - rpm_usage: int | None = None - if rpm_usage_list is not None: - for t in rpm_usage_list: - if isinstance(t, int): - if rpm_usage is None: - rpm_usage = 0 - rpm_usage += t - return tpm_usage, rpm_usage - - async def get_model_group_io_token_usage(self, model_group: str) -> tuple[int | None, int | None]: - """ - Returns current ITPM/OTPM usage for a model group (sum across deployments). - """ - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") - itpm_keys: Final[list[str]] = [] - otpm_keys: Final[list[str]] = [] - - model_list: Final = self._core_get_model_list(model_name=model_group) - if model_list is None: - return None, None - - for model in model_list: - model_id: str | None = model.get("model_info", {}).get("id") - litellm_model: str | None = model["litellm_params"].get("model") - if model_id is None or litellm_model is None: - continue - itpm_keys.append( - RouterCacheEnum.ITPM.value.format( - id=model_id, - model=litellm_model, - current_minute=current_minute, - ) - ) - otpm_keys.append( - RouterCacheEnum.OTPM.value.format( - id=model_id, - model=litellm_model, - current_minute=current_minute, - ) - ) - - combined_values: Final = await self._core_cache.async_batch_get_cache(keys=itpm_keys + otpm_keys) - if combined_values is None: - return None, None - - itpm_values: Final = combined_values[: len(itpm_keys)] - otpm_values: Final = combined_values[len(itpm_keys) :] - - total_itpm: int | None = None - for value in itpm_values: - if isinstance(value, int): - total_itpm = (total_itpm or 0) + value - - total_otpm: int | None = None - for value in otpm_values: - if isinstance(value, int): - total_otpm = (total_otpm or 0) + value - - return total_itpm, total_otpm - - @lru_cache(maxsize=DEFAULT_MAX_LRU_CACHE_SIZE) - def _cached_get_model_group_info(self, model_group: str) -> ModelGroupInfo | None: - """ - Cached version of get_model_group_info, uses @lru_cache wrapper - - This is a speed optimization, since set_response_headers makes a call to get_model_group_info on every request - """ - return self._core_get_model_group_info(model_group) - - def _core_cached_model_group_info(self, model_group: str) -> ModelGroupInfo | None: - return self._cached_get_model_group_info(model_group) - - async def _core_get_remaining_model_group_usage(self, model_group: str) -> dict[str, int]: - model_group_info: Final = self._cached_get_model_group_info(model_group) - - returned_dict: Final[dict[str, int]] = {} - - # ITPM/OTPM groups emit input/output token headers, but they may also set - # tpm/rpm, so build both sets rather than returning early - clients and - # prometheus gauges that read the standard headers still get data. - if model_group_info is not None and (model_group_info.itpm is not None or model_group_info.otpm is not None): - current_itpm, current_otpm = await self.get_model_group_io_token_usage(model_group) - returned_dict.update( - build_io_token_rate_limit_headers( - itpm_limit=model_group_info.itpm, - otpm_limit=model_group_info.otpm, - current_itpm=current_itpm, - current_otpm=current_otpm, - ) - ) - - tpm_limit: Final = model_group_info.tpm if model_group_info is not None else None - rpm_limit: Final = model_group_info.rpm if model_group_info is not None else None - - if tpm_limit is None and rpm_limit is None: - return returned_dict - - current_tpm, current_rpm = await self._core_get_model_group_usage(model_group) - - if tpm_limit is not None: - returned_dict["x-ratelimit-remaining-tokens"] = tpm_limit - (current_tpm or 0) - returned_dict["x-ratelimit-limit-tokens"] = tpm_limit - if rpm_limit is not None: - returned_dict["x-ratelimit-remaining-requests"] = rpm_limit - (current_rpm or 0) - returned_dict["x-ratelimit-limit-requests"] = rpm_limit - - return returned_dict - - async def set_response_headers( - self, - response: object, - model_group: str | None = None, - request_kwargs: dict[str, object] | None = None, - ) -> Any: - """ - Add the most accurate rate limit headers for a given model response. - - ## TODO: add model group rate limit headers - # - if healthy_deployments > 1, return model group rate limit headers - # - else return the model's rate limit headers - """ - response = prepare_response_for_header_attachment(response) - if response is None: - return response - - additional_headers: Final = ensure_response_additional_headers(response) - apply_response_model_id( - response, find_deployment_metadata(request_kwargs) if request_kwargs is not None else None - ) - additional_headers["x-litellm-model-group"] = model_group - apply_quality_router_decision_headers(additional_headers, request_kwargs) - additional_headers.update(complexity_router_decision_headers(request_kwargs)) - - if model_group is not None: - remaining_usage: Final = await self._core_get_remaining_model_group_usage(model_group) - apply_remaining_usage_headers(additional_headers, remaining_usage) - return response - - def _build_model_name_index(self, model_list: list) -> None: - """ - Build model_name -> deployment indices mapping for O(1) lookups. - - This index allows us to find all deployments for a given model_name in O(1) time - instead of O(n) linear scan through the entire model_list. - """ - self._core_model_name_to_deployment_indices.clear() - self._core_team_model_to_deployment_indices.clear() - self._core_team_public_model_names = frozenset() - - for idx, model in enumerate(model_list): - model_name = model.get("model_name") - if model_name: - if model_name not in self._core_model_name_to_deployment_indices: - self._core_model_name_to_deployment_indices[model_name] = [] - self._core_model_name_to_deployment_indices[model_name].append(idx) - - self._update_team_model_index(model, idx) - - def _build_model_id_to_deployment_index_map(self, model_list: list): - """ - Build model index from model list to enable O(1) lookups immediately. - This is called during initialization to avoid the race condition where - requests arrive before model_id_to_deployment_index_map is populated. - """ - # First populate the model_list - self._core_model_list = [] - self._invalidate_model_group_info_cache() - self._invalidate_access_groups_cache() - for _, model in enumerate(model_list): - # Extract model_info from the model dict - model_info = model.get("model_info", {}) - model_id = model_info.get("id") - - # If no ID exists, generate one using the same logic as set_model_list - if model_id is None: - model_name = model.get("model_name", "") - litellm_params = model.get("litellm_params", {}) - model_id = self._core_generate_model_id(model_name, litellm_params) - # Update the model_info in the original list - if "model_info" not in model: - model["model_info"] = {} - model["model_info"]["id"] = model_id - - self._add_model_to_list_and_index_map(model=model, model_id=model_id) - - def _core_get_model_ids(self, model_name: str | None = None, exclude_team_models: bool = False) -> list[str]: - """ - if 'model_name' is none, returns all. - - Returns list of model id's. - - Optimized with O(1) or O(k) index lookup when model_name provided, - instead of O(n) linear scan. - """ - ids: Final = [] - - if model_name is not None: - # O(1) lookup in model_name index, then O(k) iteration where k = deployments for this model_name - for model in self._core_deployments_named(model_name) or (): - if "model_info" in model and "id" in model["model_info"]: - if exclude_team_models and model["model_info"].get("team_id"): - continue - ids.append(model["model_info"]["id"]) - else: - # When model_name is None, return all model IDs - # Use the index map keys for O(n) where n = total deployments - with self._deployment_index_lock: - id_indices: Final = tuple(self.model_id_to_deployment_index_map.items()) - models: Final = tuple(self._core_model_list) - for model_id, idx in id_indices: - model = models[idx] - if "model_info" in model and "id" in model["model_info"]: - if exclude_team_models and model["model_info"].get("team_id"): - continue - ids.append(model_id) - - return ids - - def _core_get_candidate_model_ids_for_route(self, model: str, team_id: str | None = None) -> frozenset[str]: - """ - Deployment ids that could serve ``model`` for ``team_id``, following the same - precedence ``_common_checks_available_deployment`` uses to build a candidate pool: - ``model_group_alias``, then a routing group, then the first matching early-resolve - path for a name that is not a ``model_name`` (team route, wildcard pattern via - ``get_deployments_by_pattern``, team pattern router, default deployment), then the - ``model_name`` and team indexes. Delegating to the router's own resolvers keeps this - aligned with how a route actually resolves rather than re-deriving it, and unlike - ``_common_checks_available_deployment`` it is read-only: it does not apply request - fallbacks and (with ``include_team_models`` left off) does not raise. Lets a pre-call - check tell a genuine cross-group route from same-group unavailability without leaking - deployment ids into request kwargs bound for the provider. - """ - resolved: Final = self._core_get_model_from_alias(model=model) or model - routing_group_members: Final = self._get_routing_group_deployments(model=resolved, team_id=team_id) - if routing_group_members is not None: - return self._deployment_ids(routing_group_members) - early: Final = self._try_early_resolve_deployments_for_model_not_in_names( - model=resolved, request_team_id=team_id - ) - if early is not None: - early_deployments: Final = early[1] - return self._deployment_ids( - (early_deployments,) if isinstance(early_deployments, Mapping) else early_deployments - ) - return self._deployment_ids(self._core_get_all_deployments(model_name=resolved, team_id=team_id)) - - @staticmethod - def _deployment_ids(deployments: Sequence[Mapping[str, object]]) -> frozenset[str]: - return frozenset( - str(model_info["id"]) - for deployment in deployments - for model_info in (deployment.get("model_info"),) - if isinstance(model_info, Mapping) and model_info.get("id") is not None - ) - - def _region_routable_deployments( - self, - deployments: Sequence[DeploymentTypedDict], - request_kwargs: Mapping[str, object] | None, - has_request_input: bool, - ) -> tuple[Mapping[str, Any], ...]: - """ - `deployments` without the ones `_pre_call_checks` keeps this request away from for being outside its allowed - region, typed for the tag routing filter - """ - allowed_model_region: Final = None if request_kwargs is None else request_kwargs.get("allowed_model_region") - if allowed_model_region is None or not self._core_enable_pre_call_checks or not has_request_input: - return tuple(deployments) - if not isinstance(allowed_model_region, str): - return () - return tuple( - deployment - for deployment in deployments - if is_region_allowed( - litellm_params=LiteLLM_Params.model_validate(deployment["litellm_params"]), - allowed_model_region=allowed_model_region, - ) - ) - - async def _async_routable_cooled_deployment_ids( - self, - model: str, - request_deployments: Sequence[DeploymentTypedDict], - cooldown_deployments: Iterable[str], - request_kwargs: Mapping[str, object], - has_request_input: bool, - ) -> frozenset[str]: - """ - Ids of the cooling deployments in `request_deployments` that region and tag routing would send this request - to, each tag routing pick counted along with the picks it falls back to once that one is unavailable, judged - without recording a tag match on the request - """ - cooled_ids: Final = frozenset(cooldown_deployments) - metadata_variable_name: Final = get_metadata_variable_name_from_kwargs(request_kwargs) - metadata: Final = request_kwargs.get(metadata_variable_name) - if metadata_variable_name in request_kwargs and not isinstance(metadata, Mapping): - return frozenset() - tag_request: Final[Mapping[str, Any]] = ( - {**request_kwargs, metadata_variable_name: {**metadata}} - if isinstance(metadata, Mapping) - else request_kwargs - ) - routable_ids = frozenset[str]() # rebind-ok: gains each pass's tag routing picks - remaining = tuple( # rebind-ok: each pass drops what tag routing just picked - deployment - for deployment in self._region_routable_deployments(request_deployments, request_kwargs, has_request_input) - if self._deployment_ids((deployment,)) <= cooled_ids - ) - while remaining: - try: - picked = await get_deployments_for_tag( - ctx=self._core_ctx, - model=model, - healthy_deployments=remaining, - request_kwargs=tag_request, - metadata_variable_name=metadata_variable_name, - request_deployments=request_deployments, - ) - except ValueError: # tag routing sends this request to none of the rest - break - picked_identities = frozenset(id(deployment) for deployment in picked) - rest = tuple(deployment for deployment in remaining if id(deployment) not in picked_identities) - if len(rest) == len(remaining): - break - routable_ids |= self._deployment_ids( - tuple(deployment for deployment in remaining if id(deployment) in picked_identities) - ) - remaining = rest - return routable_ids - - def _core_has_model_id(self, candidate_id: str) -> bool: - """ - O(1) membership check for a deployment ID without allocating large lists. - - Note: Call sites may pass a variable named `model` when it actually - contains a deployment ID. This helper expects the deployment ID string. - - Uses the existing `model_id_to_deployment_index_map` which is kept - in sync by `_build_model_id_to_deployment_index_map` and model-list - mutation helpers. - """ - return candidate_id in self.model_id_to_deployment_index_map - - def _core_resolve_model_name_from_model_id( - self, model_id: str | None, custom_llm_provider: str | None = None - ) -> str | None: - """ - Resolve model_name from model_id. - - This method attempts to find the correct model_name to use with the router - so that litellm_params can be automatically injected from the model config. - - Strategy: - 1. First, check if model_id directly matches a model_name or deployment ID - 2. If not, search through router's model_list to find a match by litellm_params.model - 3. If not, keep model_id when the router routes that name itself (an alias, a routing group, or a wildcard - whose deployment serves `custom_llm_provider`), else name it `custom_llm_provider/model_id` when only that - form routes to such a wildcard - 4. Return the model_name if found, None otherwise - - Args: - model_id: The model_id extracted from decoded video_id - (could be model_name or litellm_params.model value) - custom_llm_provider: The provider extracted from the same id - - Returns: - model_name if found, None otherwise. If None, the request will fall through - to normal flow using environment variables. - """ - if not model_id: - return None - - # Strategy 1: Check if model_id directly matches a model_name or deployment ID - if model_id in self._core_model_names: - return model_id - if self._core_has_model_id(model_id): - deployment = self._core_get_deployment(model_id=model_id) - if deployment is not None and deployment.model_name: - return deployment.model_name - return model_id - - # Strategy 2: Search through router's model_list to find by litellm_params.model - all_models: Final = self._core_get_model_list(model_name=None) - if not all_models: - return None - - for deployment in all_models: - litellm_params = deployment.get("litellm_params", {}) - actual_model = litellm_params.get("model") - - # Match by exact match or by checking if actual_model ends with /model_id or :model_id - # e.g., model_id="veo-2.0-generate-001" matches actual_model="vertex_ai/veo-2.0-generate-001" - matches = ( - actual_model == model_id - or (actual_model and actual_model.endswith(f"/{model_id}")) - or (actual_model and actual_model.endswith(f":{model_id}")) - ) - - if matches: - model_name = deployment.get("model_name") - if model_name: - return model_name - - if self._core_is_recognized_model(model_id): - return model_id - if custom_llm_provider is None: - return None - if self._core_pattern_router.routes_to_provider(model_id, custom_llm_provider): - return model_id - provider_prefixed: Final = f"{custom_llm_provider}/{model_id}" - if not model_id.startswith(f"{custom_llm_provider}/") and self._core_pattern_router.routes_to_provider( - provider_prefixed, custom_llm_provider - ): - return provider_prefixed - return None - - def _core_resolve_served_model_name(self, model_name: str, served_model: str) -> str: - """ - `served_model` when `model_name` routes to it (the alias target, the routing group member, or the name - itself), so a request can return to the group that served an earlier one, else `model_name`. - """ - target: Final = self._core_get_model_from_alias(model_name) or model_name - routing_group: Final = self.get_routing_group(target) - members: Final = routing_group.models if routing_group is not None else () - return served_model if served_model == target or served_model in members else model_name - - def _core_map_team_model( - self, team_model_name: str | None, team_id: str, custom_llm_provider: str | None = None - ) -> str | None: - """ - Resolve explicit team models or select a deployment for model-less calls. - - Model-less calls with no provider use the first team deployment. With a provider, they - use the first team deployment of that provider, pinned by ID when its public name also - covers another provider. - """ - models: Final = self._core_get_model_list(model_name=team_model_name, team_id=team_id) - if not models: - return None - if team_model_name is None and custom_llm_provider is not None: - - def team_deployment_public_name(deployment: DeploymentTypedDict) -> str | None: - model_info: Final = _configured_model_info(deployment) - if model_info.get("team_id") != team_id: - return None - team_public_model_name: Final = model_info.get("team_public_model_name") - return ( - team_public_model_name - if isinstance(team_public_model_name, str) and team_public_model_name - else deployment["model_name"] - ) - - matching_deployment: Final = next( - ( - model - for model in models - if team_deployment_public_name(model) is not None - and provider_for_generic_call(model["litellm_params"]) == custom_llm_provider - ), - None, - ) - if matching_deployment is None: - return None - public_name: Final = team_deployment_public_name(matching_deployment) - if public_name is None: - return None - same_public_name_deployments: Final = tuple( - candidate for candidate in models if team_deployment_public_name(candidate) == public_name - ) - if all( - provider_for_generic_call(candidate["litellm_params"]) == custom_llm_provider - for candidate in same_public_name_deployments - ): - return public_name - deployment_id: Final = _configured_model_info(matching_deployment).get("id") - return deployment_id if isinstance(deployment_id, str) else None - for model in models: - if model.get("model_info", {}).get("team_id") == team_id: - if team_model_name is None: - # No model was specified (e.g. vector store endpoints). - # Return the deployment's public model name so the router - # can route to it and inject the BYOK API key. - return model.get("model_info", {}).get("team_public_model_name") or model.get("model_name") - return team_model_name - - # No team-scoped deployment found; wildcard/pattern routes are - # handled downstream by the pattern_router in _common_checks_available_deployment. - return None - - def _core_should_include_deployment( - self, model_name: str, model: DeploymentTypedDict | dict, team_id: str | None = None - ) -> bool: - """ - Get the team-specific model name if team_id matches the deployment. - """ - if ( - team_id is not None - and (model.get("model_info") or {}).get("team_id") == team_id - and model_name == (model.get("model_info") or {}).get("team_public_model_name") - ): - return True - elif model_name is not None and model["model_name"] == model_name: - # Fallback: check by internal model_name for non-team deployments - # or deployments that haven't been migrated to team_public_model_name yet - model_team_id: Final = (model.get("model_info") or {}).get("team_id") - if ( - team_id is None # requester has no team constraint - or model_team_id is None # global deployment - accessible to all teams - or model_team_id == team_id # deployment belongs to requester's team - ): - return True - # No match: deployment is for a different team or doesn't match the requested model - return False - - def _core_get_all_deployments( - self, - model_name: str, - model_alias: str | None = None, - team_id: str | None = None, - ) -> list[DeploymentTypedDict]: - """ - Return all deployments of a model name - - Used for accurate 'get_model_list'. - - if team_id specified, only return team-specific models - - Optimized with O(1) index lookup instead of O(n) linear scan. - - Note: when team_id is provided, O(1) lookup in - `team_model_to_deployment_indices` only applies when `model_name` is the - team public model name. If a caller passes an internal deployment model - name (for example, `model_name__`), this method falls back - to the standard model-name index / scan path. - """ - returned_models: Final[list[DeploymentTypedDict]] = [] - - # O(1) lookup in team_model index when team_id is provided - if team_id is not None: - team_models: Final = self._team_deployments_named(team_id, model_name) - if team_models is not None: - # O(k) where k = team deployments for this model_name (typically 1-10) - for model in team_models: - if not self._core_should_include_deployment(model_name=model_name, model=model, team_id=team_id): - continue - if model_alias is not None: - alias_model = model.copy() - alias_model["model_name"] = model_alias - returned_models.append(alias_model) - else: - returned_models.append(model) - if returned_models: - return returned_models - - # O(1) lookup in model_name index - named_models: Final = self._core_deployments_named(model_name) - if named_models is not None: - # O(k) where k = deployments for this model_name (typically 1-10) - for model in named_models: - if self._core_should_include_deployment(model_name=model_name, model=model, team_id=team_id): - if model_alias is not None: - # Optimized: Use shallow copy since we only modify top-level model_name - # This is much faster than deepcopy for nested dict structures - alias_model = model.copy() - alias_model["model_name"] = model_alias - returned_models.append(alias_model) - else: - returned_models.append(model) - elif team_id is not None: - # Fallback: if team_id is provided and model_name not in index, - # check if model_name matches any team_public_model_name - # O(n) scan but only when team_id lookup fails - for _idx, model in enumerate(self._core_model_list): - if self._core_should_include_deployment(model_name=model_name, model=model, team_id=team_id): - if model_alias is not None: - # Optimized: Use shallow copy since we only modify top-level model_name - alias_model = model.copy() - alias_model["model_name"] = model_alias - returned_models.append(alias_model) - else: - returned_models.append(model) - - return returned_models - - def _core_get_model_names(self, team_id: str | None = None) -> list[str]: - """ - Returns all possible model names for the router, including models defined via model_group_alias. - - If a team_id is provided, only deployments configured with that team_id (i.e. team‐specific models) - will yield their team public name. - """ - deployments: Final = self._core_get_model_list() or [] - model_names: Final = [] - - for deployment in deployments: - model_info = deployment.get("model_info") - if self._is_team_specific_model(model_info): - team_model_name = self._get_team_specific_model(deployment=deployment, team_id=team_id) - if team_model_name: - model_names.append(team_model_name) - else: - model_names.append(deployment.get("model_name", "")) - - return model_names - - def _core_get_fully_blocked_model_names(self) -> set[str]: - """ - Returns the set of model_names where every backing deployment has `blocked=True`. - - Used by `/v1/models` to hide paused models from client listings while still - surfacing them on admin endpoints (e.g. `/model/info`). A model with at least - one non-blocked deployment is still serviceable and remains visible. - """ - deployments: Final = self._core_get_model_list() or [] - blocked_by_name: Final[dict[str, bool]] = {} - for deployment in deployments: - name = deployment.get("model_name") or "" - if not name: - continue - is_blocked = (deployment.get("model_info") or {}).get("blocked") is True - if name in blocked_by_name: - blocked_by_name[name] = blocked_by_name[name] and is_blocked - else: - blocked_by_name[name] = is_blocked - return {name for name, fully_blocked in blocked_by_name.items() if fully_blocked} - - @staticmethod - def _core_are_all_deployments_blocked( - deployments: Sequence[DeploymentTypedDict], - ) -> bool: - return len(deployments) > 0 and all( - (deployment.get("model_info") or {}).get("blocked") is True for deployment in deployments - ) - - def _is_model_fully_blocked(self, model: str) -> bool: - deployments: Final = self._core_get_model_list(model_name=model) or [] - return self._core_are_all_deployments_blocked(deployments=deployments) - - async def _core_async_get_fully_unhealthy_model_names(self) -> set[str]: - """ - Returns the set of model names where every backing deployment is currently - marked unhealthy by background health checks (and the health state is not stale). - - Used by `/v1/models?healthy_only=true` to hide models that cannot serve any - request. A model with at least one healthy (or unknown-health) deployment - remains visible. Returns an empty set when no health state is available, so - callers fail open to the unfiltered listing. - - Notes: - - Mirrors `_async_filter_health_check_unhealthy_deployments`: when - `allowed_fails_policy` is set, cooldown is the sole routing exclusion - mechanism, so nothing is hidden here either. - - Team-specific public model names (`team_public_model_name`) are - aggregated alongside `model_name`, so team aliases of fully-unhealthy - deployments are hidden too (unlike `get_fully_blocked_model_names`, - which matches `model_name` only). - - Wildcard routes (e.g. `openai/*`) are matched by their literal - deployment name only; models expanded from a wildcard route are not - hidden (fail open). - - Intentionally diverges from the routing-time safety net (which - bypasses the health filter when every candidate is unhealthy and - still attempts the request): hiding here is presentation-only — - it answers "should this model be advertised?", not "should a - request for it still be attempted?". A hidden model can still be - called directly. - """ - if self._core_allowed_fails_policy is not None: - return set() - unhealthy_ids: Final = await self._core_health_state_cache.async_get_unhealthy_deployment_ids() - if not unhealthy_ids: - return set() - deployments: Final = self._core_get_model_list() or [] - unhealthy_by_name: Final[dict[str, bool]] = {} - for deployment in deployments: - model_info = deployment.get("model_info") or {} - names = [deployment.get("model_name") or ""] - team_public_model_name = model_info.get("team_public_model_name") - if team_public_model_name: - names.append(team_public_model_name) - is_unhealthy = model_info.get("id") in unhealthy_ids - for name in names: - if not name: - continue - if name in unhealthy_by_name: - unhealthy_by_name[name] = unhealthy_by_name[name] and is_unhealthy - else: - unhealthy_by_name[name] = is_unhealthy - return {name for name, fully_unhealthy in unhealthy_by_name.items() if fully_unhealthy} - - def _get_team_specific_model(self, deployment: DeploymentTypedDict, team_id: str | None = None) -> str | None: - """ - Get the team-specific model name if team_id matches the deployment. - - Args: - deployment: DeploymentTypedDict - The model deployment - team_id: Optional[str] - If passed, will return router models set with a `team_id` matching the passed `team_id`. - - Returns: - str: The `team_public_model_name` if team_id matches - None: If team_id doesn't match or no team info exists - """ - model_info: Final[dict | None] = deployment.get("model_info") or {} - if model_info is None: - return None - if team_id == model_info.get("team_id"): - return model_info.get("team_public_model_name") - return None - - def _is_team_specific_model(self, model_info: dict | None) -> bool: - """ - Check if model info contains team-specific configuration. - - Args: - model_info: Model information dictionary - - Returns: - bool: True if model has team-specific configuration - """ - return bool(model_info and model_info.get("team_id")) - - def _core_get_model_list_from_model_alias(self, model_name: str | None = None) -> list[DeploymentTypedDict]: - """ - Helper function to get model list from model alias. - - Used by `.get_model_list` to get model list from model alias. - """ - returned_models: Final[list[DeploymentTypedDict]] = [] - - if model_name is not None: - # Fast path: direct dict lookup avoids scanning all aliases for non-alias model names. - if model_name not in self._core_model_group_alias: - return returned_models - alias_items = [(model_name, self._core_model_group_alias[model_name])] - else: - alias_items = list(self._core_model_group_alias.items()) - - for model_alias, model_value in alias_items: - if isinstance(model_value, str): - _router_model_name: str = model_value - elif isinstance(model_value, dict): - _model_value = RouterModelGroupAliasItem(**model_value) - if _model_value["hidden"] is True and model_name is None: - continue - _router_model_name = _model_value["model"] - else: - continue - - if (alias_group := self.get_routing_group(_router_model_name)) is not None: - returned_models.extend( - {**row, "model_name": model_alias} for row in self._materialize_routing_group_rows((alias_group,)) - ) - else: - returned_models.extend( - self._core_get_all_deployments(model_name=_router_model_name, model_alias=model_alias) - ) - - return returned_models - - def _core_get_model_list_from_routing_groups(self, model_name: str | None = None) -> Sequence[DeploymentTypedDict]: - """ - Callable routing groups materialized as model-list rows, mirroring - `get_model_list_from_model_alias`: each member deployment is emitted - under the group's name (via `get_all_deployments`' `model_alias` - rewrite), which is what surfaces groups in `get_model_names`, - `/v1/models` discovery, `get_model_group_usage`, and the - blocked/unhealthy hiding that all read `get_model_list`. - """ - if model_name is not None: - group: Final = self.get_routing_group(model_name) - return self._materialize_routing_group_rows((group,)) if group is not None else () - cached: Final = self._routing_group_rows - if cached is not None: - return cached - rows: Final = self._materialize_routing_group_rows( - tuple( - callable_group - for name in self._core_routing_groups - if (callable_group := self.get_routing_group(name)) is not None - ) - ) - self._routing_group_rows = rows - return rows - - def _materialize_routing_group_rows(self, groups: tuple[RoutingGroup, ...]) -> tuple[DeploymentTypedDict, ...]: - return tuple( - self._as_routing_group_row(apply_routing_group_priority(group, member, deployment)) - for group in groups - for member in group.models - for deployment in self._core_get_all_deployments(model_name=member, model_alias=group.group_name) - ) - - @staticmethod - def _as_routing_group_row(deployment: DeploymentTypedDict) -> DeploymentTypedDict: - """ - A member deployment re-emitted under its group's name must not carry - the member's `access_groups`: access groups grant member names, never - the group, so inheriting them here would let a key holding a member's - access group list and call the whole group. - """ - model_info: Final = {k: v for k, v in (deployment.get("model_info") or {}).items() if k != "access_groups"} - return {**deployment, "model_info": model_info} - - TIER_PARAMS_NEVER_DROPPED: Final = frozenset(all_litellm_params) | frozenset( - { - "additional_drop_params", - "drop_params", - "messages", - "model", - "extra_headers", - "max_tokens", - "max_completion_tokens", - } - ) - - @staticmethod - def _declared_param_allowlist(params: Mapping[str, object]) -> frozenset[str]: - declared: Final = params.get("allowed_openai_params") - if not isinstance(declared, (list, tuple, set, frozenset)): - return frozenset() - return frozenset(entry for entry in declared if isinstance(entry, str)) - - @staticmethod - def _deployment_accepts_param(deployment: DeploymentTypedDict, group: str, param: str) -> bool: - deployment_params: Final = deployment.get("litellm_params") - if not deployment_params: - return True - if param in PythonCore._declared_param_allowlist(deployment_params): - return True - if declared_authenticating_provider( - str(deployment_params.get("model") or ""), deployment_params.get("custom_llm_provider") - ): - return True - return PythonCore._deployment_supports_param(deployment, group, param) - - @staticmethod - def _capability_model_and_provider(deployment_params: LiteLLMParamsTypedDict, group: str) -> tuple[str, str]: - model_label: Final = str(deployment_params.get("model") or group) - authenticating_provider: Final = declared_authenticating_provider( - model_label, deployment_params.get("custom_llm_provider") - ) - if authenticating_provider is not None: - return model_label.removeprefix(f"{authenticating_provider}/"), authenticating_provider - model, custom_llm_provider, _, _ = python_host.invoke("get_llm_provider")( - model=model_label, - custom_llm_provider=deployment_params.get("custom_llm_provider"), - litellm_params=LiteLLM_Params.model_validate(deployment_params), - ) - return model, custom_llm_provider - - @staticmethod - def _deployment_supports_param(deployment: DeploymentTypedDict, group: str, param: str) -> bool: - deployment_params: Final = deployment.get("litellm_params") - if not deployment_params: - return True - deployment_model_info: Final = deployment.get("model_info") - base_model: Final = ( - deployment_model_info.get("base_model") if deployment_model_info else None - ) or deployment_params.get("base_model") - try: - model, custom_llm_provider = PythonCore._capability_model_and_provider(deployment_params, group) - supported: Final = python_host.invoke("get_supported_openai_params")( - model=model, - custom_llm_provider=custom_llm_provider, - base_model=base_model if isinstance(base_model, str) else None, - ) - except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request - verbose_router_logger.debug( - "litellm.router.py::_deployment_supports_param: keeping %s for model=%s. Got - %s", param, group, e - ) - return True - return supported is None or param in supported - - def _tier_params_the_target_accepts( - self, model: str, tier_params: Mapping[str, object], request_kwargs: Mapping[str, object] - ) -> Mapping[str, object]: - """Drop an OpenAI param that no deployment behind ``model`` declares. - - A tier's litellm_params are an operator override applied to every request the tier routes, - so one the target cannot take turns that whole tier into a 400 raised before the request - leaves the proxy. The candidates are exactly what get_optional_params can reject, asked of - the module that raises, so credentials and endpoint controls are never at risk. - - TIER_PARAMS_NEVER_DROPPED is excluded on top of that, for two reasons. No provider lists a - litellm control among its supported params, so "no deployment declares it" means litellm - consumes it rather than that the target refuses it, and dropping one changes litellm's own - behavior: dropping drop_params or additional_drop_params silently disables the sanitization - the operator configured. Providers do list extra_headers, but it carries auth, tenancy and - routing information, so sending fewer headers than configured is worse than today's error. - Token ceilings stay for the same reason: a tier's max_tokens or max_completion_tokens is a - cost bound, and dropping it would let a caller's own larger value through where today the - mismatch fails loudly. - - The trade this filter makes is a param for a working request, which is right for one that - only shapes how the model answers and wrong for anything else. - - A param survives if ANY deployment could take it, because routing has not chosen one yet, - and it survives both an unresolvable provider and a group with no deployments, because a - best-effort filter must never narrow what the request already did. - - A github_copilot or chatgpt deployment counts as accepting everything, decided before any - lookup: resolving either provider runs its OAuth device flow, so a capability question - asked from the routing path can freeze the event loop for minutes waiting on a human. - - allowed_openai_params is the documented escape hatch for an outdated or incomplete - supported-params list: request-time validation extends the supported list with it before - comparing. The filter asks the same question, so a param named by the allowlist on the tier - overlay, the request, or a deployment's own litellm_params is never a drop candidate. - """ - deployments: Final = self._core_get_model_list(model_name=model) or () - if not deployments: - return tier_params - allowlisted: Final = self._declared_param_allowlist(tier_params) | self._declared_param_allowlist( - request_kwargs - ) - candidates: Final = provider_rejectable_params(tier_params) - self.TIER_PARAMS_NEVER_DROPPED - allowlisted - unsupported: Final = frozenset( - param - for param in candidates - if not any(self._deployment_accepts_param(deployment, model, param) for deployment in deployments) - ) - if not unsupported: - return tier_params - verbose_router_logger.warning( - "litellm.router.py: dropping tier params %s for model=%s, no deployment behind it declares them", - ", ".join(sorted(unsupported)), - model, - ) - return MappingProxyType({key: value for key, value in tier_params.items() if key not in unsupported}) - - def _core_get_model_list( - self, model_name: str | None = None, team_id: str | None = None - ) -> list[DeploymentTypedDict] | None: - """ - Includes router model_group_alias'es as well - - if team_id specified, returns matching team-specific models - """ - # Note: model_list and model_group_alias are always initialized in __init__ - # so hasattr checks are unnecessary - returned_models: list[DeploymentTypedDict] = [] - - if model_name is not None: - returned_models.extend(self._core_get_all_deployments(model_name=model_name, team_id=team_id)) - - returned_models.extend(self._core_get_model_list_from_model_alias(model_name=model_name)) - returned_models.extend(self._core_get_model_list_from_routing_groups(model_name=model_name)) - - if len(returned_models) == 0 and model_name is not None: - returned_models.extend(self._get_wildcard_deployments(model_name=model_name, team_id=team_id)) - - if model_name is None: - returned_models += self._core_model_list - - return returned_models - - def _get_wildcard_deployments(self, model_name: str, team_id: str | None = None) -> list[DeploymentTypedDict]: - """ - The deployments of the wildcard routes matching model_name (the proxy-wide - ones, plus team_id's own when given), each emitted under model_name. - """ - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - matches: Final = [ - *self._core_pattern_router.get_deployments_by_pattern(model=model_name), - *(team_router.get_deployments_by_pattern(model=model_name) if team_router is not None else ()), - ] - return [{**DeploymentTypedDict(**m), "model_name": model_name} for m in matches] - - def _core_wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: - """The deployments routing serves ``model`` from through a wildcard: ``team_id``'s own pattern - matching it, else the most specific proxy-wide pattern, never both.""" - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - team_owned: Final[Sequence[Mapping[str, object]]] = ( - team_router.get_deployments_by_pattern(model=model) if team_router is not None else () - ) - return team_owned or self._core_pattern_router.get_deployments_by_pattern(model=model) - - def _core_wildcard_routed_models(self, model: str, team_id: str | None) -> frozenset[str]: - """The models wildcard_route_deployments would route ``model`` to, without copying the deployments.""" - team_router: Final = self._core_team_pattern_routers.get(team_id) if team_id is not None else None - team_owned: Final = ( - team_router.get_models_by_pattern(model=model) if team_router is not None else frozenset[str]() - ) - return team_owned or self._core_pattern_router.get_models_by_pattern(model=model) - - def _core_team_ptu_ceiling(self, team_id: str, model_group: str) -> PTUTeamCeiling | None: - """The per-minute ceiling ``team_id``'s PTU shares set on the group serving ``model_group``, - else None, kept until the model list next changes.""" - if team_id not in self._ptu_share_teams(): - return None - return self._cached_team_ptu_ceiling(team_id, model_group) - - def _collect_ptu_share_teams(self) -> frozenset[str]: - return ptu_share_teams(self._core_get_model_list() or ()) - - def _resolve_team_ptu_ceiling(self, team_id: str, model_group: str) -> PTUTeamCeiling | None: - return team_ptu_ceiling( - self._core_get_model_list() or (), - self._core_model_list, - team_id, - model_group, - lambda name: self._core_wildcard_route_deployments(name, team_id), - ) - - def _core_get_model_list_of_routed_group( - self, model_group: str, team_id: str | None = None - ) -> list[DeploymentTypedDict]: - """ - The deployments a request the router has already resolved to model_group is - served from: team_id's own deployments under that public name when the team has - any, else the ones named model_group, the routing group of that name, or the - wildcard route matching it when neither exists. - - Unlike get_model_list, model_group's own model_group_alias entry is not - followed. The router resolves an alias exactly once, so a group reached as - an alias target is served by its own deployments, never by a second hop: - in the chain X -> T -> U a request to X is served from T's deployments. - """ - if team_id is not None and self._team_deployments_named(team_id, model_group): - return self._core_get_all_deployments(model_name=model_group, team_id=team_id) - own: Final = self._get_all_deployments(model_name=model_group) - named: Final = [*own, *self._core_get_model_list_from_routing_groups(model_name=model_group)] - return named or self._get_wildcard_deployments(model_name=model_group) - - def _core_resolved_litellm_models( - self, model_name: str, team_id: str | None = None, *, use_base_model: bool = False - ) -> tuple[str, ...]: - """The provider model strings `model_name` can actually be served by on this proxy. - - `get_model_list` composes every channel the request path itself uses (exact name, - model_group_alias, routing groups, wildcards), so this answers "which models will - answer a call to this name" rather than "what did the admin call it": the deployment - name is admin-arbitrary, and two names over one provider model are one model. - - Empty when the name resolves to no deployment. That is not the same fact as "the - call will fail" - a provider-qualified public name is served by the SDK with no - deployment behind it - so the fallback for an empty result is the caller's policy, - never this function's. - """ - if use_base_model: - return tuple( - backend_model - for deployment in self._core_get_model_list(model_name=model_name, team_id=team_id) or () - if (backend_model := self._cost_map_backend_model(deployment)) - ) - return tuple( - litellm_model - for deployment in self._core_get_model_list(model_name=model_name, team_id=team_id) or () - if isinstance(litellm_model := deployment.get("litellm_params", {}).get("model"), str) and litellm_model - ) - - def _invalidate_model_group_info_cache(self) -> None: - """Invalidate the cached model group info. - - Call this whenever self._core_model_list is modified to ensure the cache is rebuilt. - Also clears the auth-layer zero-cost cache, which depends on the same - ``ModelGroupInfo`` data — without this, an in-place pricing update on - an existing deployment (same model count) would keep a stale ``True`` - result and bypass budget enforcement. - """ - self._cached_get_model_group_info.cache_clear() - self._core_cached_deployment_model_info.cache_clear() - self._ptu_share_teams.cache_clear() - self._cached_team_ptu_ceiling.cache_clear() - self._cached_deployment_tracks_usage.cache_clear() - self.served_openai_models.cache_clear() - self._zero_cost_cache.clear() - self._routing_group_rows = None - - def _invalidate_access_groups_cache(self) -> None: - """Invalidate the cached access groups. - - Call this whenever self._core_model_list is modified to ensure the cache is rebuilt. It also - advances ``registry().version``, so every write path (``apply``, the old ``set_model_list`` - and ``add_deployment``) moves it. - """ - self._access_groups_cache = None - self._registry_version += 1 - - def _core_get_model_access_groups( - self, - model_name: str | None = None, - model_access_group: str | None = None, - team_id: str | None = None, - ) -> dict[str, list[str]]: - """ - If model_name is provided, only return access groups for that model. - - Parameters: - - model_name: Optional[str] - the received model name from the user (can be a wildcard route). If set, will only return access groups for that model. - - model_access_group: Optional[str] - the received model access group from the user. If set, will only return models for that access group. - - team_id: Optional[str] - the team id, to resolve team-specific models - """ - # Check if this is the no-args hot path (cacheable) - _use_cache: Final = model_name is None and model_access_group is None and team_id is None - - # Return cached result for the no-args hot path - if _use_cache and self._access_groups_cache is not None: - return self._access_groups_cache - - from collections import defaultdict - - access_groups: Final = defaultdict(list) - - model_list: Final = self._core_get_model_list(model_name=model_name, team_id=team_id) - if model_list: - for m in model_list: - _model_info = m.get("model_info") - if _model_info: - for group in _model_info.get("access_groups", []) or []: - if model_access_group is not None: - if group == model_access_group: - model_name = m["model_name"] - access_groups[group].append(model_name) - else: - model_name = m["model_name"] - access_groups[group].append(model_name) - - # Cache the result for the no-args hot path - if _use_cache: - self._access_groups_cache = dict(access_groups) - return self._access_groups_cache - - return access_groups - - def _core_is_model_access_group_for_wildcard_route(self, model_access_group: str) -> bool: - """ - Return True if model access group is a wildcard route - """ - # GET ACCESS GROUPS - access_groups: Final = self._core_get_model_access_groups(model_access_group=model_access_group) - - if len(access_groups) == 0: - return False - - models: Final = access_groups.get(model_access_group, []) - - for model in models: - # CHECK IF MODEL ACCESS GROUP IS A WILDCARD ROUTE - if self._core_pattern_router.route(request=model) is not None: - return True - - return False - - def _core_get_settings(self) -> dict[str, object]: # mutable-ok: the settings snapshot is built fresh per call - """ - Get router settings method, returns a dictionary of the settings and their values. - For example get the set values for routing_strategy_args, routing_strategy, allowed_fails, cooldown_time, num_retries, timeout, max_retries, retry_after - """ - _all_vars: Final = vars(self) - _settings_to_return: Final = {} - vars_to_include: Final = [ - "routing_strategy_args", - "routing_strategy", - "allowed_fails", - "cooldown_time", - "num_retries", - "timeout", - "max_retries", - "retry_after", - "fallbacks", - "context_window_fallbacks", - "model_group_retry_policy", - "retry_policy", - "model_group_alias", - "enable_weighted_failover", - "enable_mid_stream_fallback_continuation", - "enable_tag_filtering", - "tag_routing_prefix", - "fallback_on_unsupported_params", - ] - - for var in vars_to_include: - stored = var if var in _all_vars else "_core_" + var - if stored in _all_vars: - _settings_to_return[var] = _all_vars[stored] - if ( - var == "routing_strategy_args" - and self._core_routing_strategy == "latency-based-routing" - and self.lowestlatency_logger is not None - ): - _settings_to_return[var] = self.lowestlatency_logger.routing_args.json() - - _settings_to_return["routing_groups"] = [ - group.model_dump(exclude=frozenset(("model_priorities",)) if group.model_priorities is None else None) - for group in self._core_routing_groups.values() - ] - return _settings_to_return - - def _switch_routing_strategy(self, routing_strategy: str | None, kwargs: Mapping[str, object]) -> None: - if routing_strategy == "lar1": - from litellm.router_strategy.lar1_routing import apply_lar1_routing_strategy - - apply_lar1_routing_strategy(self._core_ctx, kwargs.get("routing_strategy_args")) - return - self.routing_strategy_init( - routing_strategy=routing_strategy, routing_strategy_args=kwargs.get("routing_strategy_args", {}) - ) - - def _core_update_settings(self, **kwargs) -> None: - """ - Update the router settings. - """ - _int_settings: Final = [ - "timeout", - "num_retries", - "retry_after", - "allowed_fails", - "cooldown_time", - ] - - _existing_router_settings: Final = self._core_get_settings() - model_group_alias_before: Final = self._core_model_group_alias - rebuild_routing_groups = False - routing_args_updated = False - for var in kwargs: - if var in RUNTIME_UPDATABLE_ROUTER_SETTINGS: - if var in _int_settings: - _casted_value = int(kwargs[var]) - setattr(self, var, _casted_value) - elif var == "routing_groups": - rebuild_routing_groups = True - elif var == "optional_pre_call_checks": - self.set_optional_pre_call_checks(kwargs[var]) - elif var == "enable_mid_stream_fallback_continuation": - self._set_mid_stream_fallback_continuation(kwargs[var]) - elif var == "retry_policy": - self._set_retry_policy(kwargs[var]) - else: - value = kwargs[var] - # only run routing strategy init if it has changed - if var == "routing_strategy": - value = self._normalize_strategy(value) - if _existing_router_settings["routing_strategy"] != value: - self._switch_routing_strategy(value, kwargs) - rebuild_routing_groups = True - elif var == "routing_strategy_args": - routing_args_updated = value != self.routing_strategy_args - setattr(self, var, value) - else: - verbose_router_logger.debug("Setting %s is not allowed", var) - - if routing_args_updated: - self._apply_updated_routing_strategy_args() - - if self._core_model_group_alias != model_group_alias_before: - self._invalidate_model_group_info_cache() - - if rebuild_routing_groups: - routing_groups_input: Final = kwargs.get("routing_groups", self._routing_groups_input) - self._init_routing_groups(routing_groups_input) - self._routing_groups_input = routing_groups_input - verbose_router_logger.debug("Updated Router settings: %s", self._core_get_settings()) - - @staticmethod - def _is_flag_enabled(value: object) -> bool: - return value is True or (isinstance(value, str) and str_to_bool(value) is True) - - def _set_mid_stream_fallback_continuation(self, enabled: object) -> None: - self.enable_mid_stream_fallback_continuation = self._is_flag_enabled(enabled) - if self.enable_mid_stream_fallback_continuation: - self._register_continuation_prefill_check() - - def _set_retry_policy(self, value: object) -> None: - policy: Final = RetryPolicy.model_validate(value) if isinstance(value, dict) else value - if policy is None or isinstance(policy, RetryPolicy): - self.retry_policy = policy - - def _get_client(self, deployment, kwargs, client_type=None): - """ - Returns the appropriate client based on the given deployment, kwargs, and client_type. - - Parameters: - deployment (dict): The deployment dictionary containing the clients. - kwargs (dict): The keyword arguments passed to the function. - client_type (str): The type of client to return. - - Returns: - The appropriate client based on the given client_type and kwargs. - """ - model_id: Final = deployment["model_info"]["id"] - parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(kwargs) - if client_type == "max_parallel_requests": - cache_key = f"{model_id}_max_parallel_requests_client" - client = self._core_cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) - if client is None: - InitalizeCachedClient.set_max_parallel_requests_client(ctx=self._core_ctx, model=deployment) - client = self._core_cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) - return client - elif client_type == "async": - if ( - deployment["litellm_params"].get( # pyright: ignore[reportUnknownMemberType] # legacy mapping - "kubernetes_pod_discovery" - ) - is True - ): - return None - if kwargs.get("stream") is True: - cache_key = f"{model_id}_stream_async_client" - client = self._core_cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) - return client - else: - cache_key = f"{model_id}_async_client" - client = self._core_cache.get_cache(key=cache_key, local_only=True, parent_otel_span=parent_otel_span) - return client - else: - if ( - deployment["litellm_params"].get( # pyright: ignore[reportUnknownMemberType] # legacy mapping - "kubernetes_pod_discovery" - ) - is True - ): - return None - if kwargs.get("stream") is True: - cache_key = f"{model_id}_stream_client" - client = self._core_cache.get_cache(key=cache_key, parent_otel_span=parent_otel_span) - return client - else: - cache_key = f"{model_id}_client" - client = self._core_cache.get_cache(key=cache_key, parent_otel_span=parent_otel_span) - return client - - def _count_pre_call_check_tokens( - self, - messages: Sequence[Mapping[str, object]] | None, - input: str | list[object] | None, - request_kwargs: Mapping[str, object] | None = None, - tokenizer_model: str = "", - deployments: Sequence[Mapping[str, object]] = (), - ) -> int: - """ - Count input tokens for context-window pre-call checks. - - Chat Completions send `messages`; the Responses API sends `input` (a string or - a list of Responses input items) plus an optional `instructions` system prompt. - The Responses payload is normalized to chat messages via the shared - LiteLLMCompletionResponsesConfig transform so the same token_counter path covers - both API surfaces and `instructions` tokens are included in the count. - - Prompt content the message list never carries is read from `request_kwargs`: - `tools` (Chat Completions, Responses and Anthropic Messages shapes) and the - Anthropic Messages top-level `system` block. - """ - from litellm.llms.anthropic.messages.native.utils import ( - anthropic_system_to_openai_message, - ) - - extras: Final = request_kwargs if request_kwargs is not None else MappingProxyType({}) - raw_instructions: Final = extras.get("instructions") - instructions: Final = raw_instructions if isinstance(raw_instructions, str) else None - raw_tools: Final = extras.get("tools") - tools: Final = ( - cast(list[ChatCompletionToolParam], raw_tools) # cast-ok: token_counter formats any tool dict shape - if isinstance(raw_tools, list) and raw_tools - else None - ) - system_message: Final = anthropic_system_to_openai_message(extras.get("system")) - if messages is not None: - counted_messages: Final = (system_message, *messages) if system_message is not None else messages - return self._host.invoke("token_counter")(model=tokenizer_model, messages=counted_messages, tools=tools) - if input is not None: - from openai.types.responses.response_create_params import ResponseInputParam - - from litellm.litellm_core_utils.token_counter import split_embedding_input_tokens - from litellm.responses.litellm_completion_transformation.transformation import ( - LiteLLMCompletionResponsesConfig, - ) - - text_input, _ = split_embedding_input_tokens(input, partial(deployment_backend_models, deployments)) - typed_input: Final = cast( # cast-ok: str | list matches transform input - str | ResponseInputParam, text_input - ) - input_messages: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( - input=typed_input, - responses_api_request={"instructions": instructions} if instructions is not None else {}, - only_since_compaction=True, - ) - return self._host.invoke("token_counter")( - model=tokenizer_model, - messages=cast(list, input_messages), # cast-ok: transformed chat messages - tools=tools, - ) - raise ValueError("Either messages or input must be provided to count tokens") - - def _deployment_max_input_tokens(self, model: str, deployment: Mapping[str, object]) -> int | None: - """The deployment's declared context window, or None when it declares none or cannot be resolved.""" - try: - model_info: Final = self._core_get_router_model_info( - deployment=cast(dict, deployment), # cast-ok: router deployments are plain dicts - received_model_name=model, - ) - except Exception as e: # noqa: BLE001 # best-effort: an unmappable deployment must not hide the others - verbose_router_logger.debug( - "litellm.router.py::_deployment_max_input_tokens: skipping deployment. Got - %s", e - ) - return None - max_input_tokens: Final = model_info.get("max_input_tokens") - return max_input_tokens if isinstance(max_input_tokens, int) else None - - def _pre_call_checks_need_token_count( - self, model: str, healthy_deployments: Sequence[Mapping[str, object]] - ) -> bool: - """Whether any healthy deployment declares a context window that a token count could exceed. - - Resolves each deployment the way ``_pre_call_checks`` does, so one unmappable deployment - cannot hide a later one that does declare a limit. - """ - return any( - self._deployment_max_input_tokens(model, deployment) is not None for deployment in healthy_deployments - ) - - async def _acount_pre_call_check_tokens( - self, - model: str, - healthy_deployments: Sequence[Mapping[str, object]], - messages: Sequence[Mapping[str, str]] | None, - input: str | Sequence[object] | None, - request_kwargs: Mapping[str, object] | None, - ) -> int | None: - """Count input tokens off the event loop, so a multi-MB prompt cannot stall the proxy. - - Returns None when no deployment limits its context window, and when counting fails. The - caller pairs this with ``skip_inline_token_count`` so neither case puts the count back on - the loop: a failed count only skips the context-window check - RPM, region and - invalid-params checks still run for every deployment. - """ - if messages is None and input is None: - return None - try: - if not self._pre_call_checks_need_token_count(model, healthy_deployments): - return None - return await offload_token_count(self._count_pre_call_check_tokens)( - messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter - input=cast(str | list | None, input), # cast-ok: forwarded to the sync counter - request_kwargs=request_kwargs, - tokenizer_model=single_encoding_openai_model( - self.served_openai_models(model) or served_openai_models(healthy_deployments) - ) - or "", - deployments=healthy_deployments, - ) - except Exception as e: # noqa: BLE001 # best-effort: an uncountable prompt must not fail the request - verbose_router_logger.error( - "litellm.router.py::_acount_pre_call_check_tokens: failed to count tokens. Got - %s", e - ) - return None - - def _resolve_pre_call_check_input_tokens( - self, - messages: Sequence[Mapping[str, str]] | None, - input: str | Sequence[object] | None, - request_kwargs: Mapping[str, object] | None, - input_token_count: int | None, - count_inline: bool, - model: str, - deployments: Sequence[Mapping[str, object]], - ) -> int | None: - if input_token_count is not None: - return input_token_count - if not count_inline: - return None - try: - return self._count_pre_call_check_tokens( - messages=cast(list[dict[str, str]] | None, messages), # cast-ok: forwarded to the sync counter - input=cast(str | list[object] | None, input), # cast-ok: forwarded to the sync counter - request_kwargs=request_kwargs, - tokenizer_model=single_encoding_openai_model( - self.served_openai_models(model) or served_openai_models(deployments) - ) - or "", - deployments=deployments, - ) - except Exception as e: - verbose_router_logger.error( - "litellm.router.py::_pre_call_checks: failed to count tokens. " - "Skipping the context-window check only; RPM/region/param checks still apply. Got - %s", - e, - ) - return None - - @staticmethod - def _params_neither_dropped_nor_allowed(params: Mapping[str, object], candidates: frozenset[str]) -> frozenset[str]: - if params.get("drop_params") is True: - return frozenset() - dropped: Final = params.get("additional_drop_params") - dropped_names: Final = ( - frozenset(entry for entry in dropped if isinstance(entry, str)) - if isinstance(dropped, list) - else frozenset() - ) - return candidates - PythonCore._declared_param_allowlist(params) - dropped_names - - @staticmethod - def _deployment_lacks_params(deployment: DeploymentTypedDict, group: str, params: frozenset[str]) -> bool: - return any(not PythonCore._deployment_supports_param(deployment, group, param) for param in params) - - @staticmethod - def _deployment_rejects_params( - deployment: DeploymentTypedDict, group: str, params: frozenset[str], request_kwargs: Mapping[str, object] - ) -> bool: - call_settings: Final[Mapping[str, object]] = {**(deployment.get("litellm_params") or {}), **request_kwargs} - return PythonCore._deployment_lacks_params( - deployment, group, PythonCore._params_neither_dropped_nor_allowed(call_settings, params) - ) - - @staticmethod - def _deployment_declares_no_vision(deployment: DeploymentTypedDict, group: str) -> bool: - plain_deployment: Final = cast(Mapping[str, object], deployment) # cast-ok: router deployments are plain dicts - model_info: Final = cast(Mapping[str, object], plain_deployment.get("model_info") or {}) # cast-ok: plain dict - declared: Final = model_info.get("supports_vision") - if declared is not None: - return declared is False - raw_params: Final = plain_deployment.get("litellm_params") or {} - params: Final = cast(Mapping[str, object], raw_params) # cast-ok: router deployments are plain dicts - if not params: - return False - base_model: Final = model_info.get("base_model") or params.get("base_model") - try: - model, custom_llm_provider = PythonCore._capability_model_and_provider(deployment["litellm_params"], group) - return litellm_utils.is_vision_explicitly_disabled( - base_model if isinstance(base_model, str) else model, custom_llm_provider - ) - except Exception as e: # noqa: BLE001 # best-effort filter: an unresolvable provider must not narrow the request - verbose_router_logger.debug( - "litellm.router.py::_deployment_declares_no_vision: keeping deployment for model=%s. Got - %s", group, e - ) - return False - - @staticmethod - def _deployment_serves_route( - deployment: DeploymentTypedDict, - group: str, - resolver: Callable[[str, LlmProviders], object | None], - ) -> bool: - if "litellm_params" not in deployment: - return True - - try: - model, resolved_custom_llm_provider = PythonCore._capability_model_and_provider( - deployment["litellm_params"], group - ) - provider: Final = next( - (member for member in LlmProviders if member.value == resolved_custom_llm_provider), - None, - ) - return provider is not None and resolver(model, provider) is not None - except Exception as e: # noqa: BLE001 # best-effort capability filter - verbose_router_logger.debug( - "litellm.router.py::_deployment_serves_route: keeping deployment for model=%s. Got - %s", group, e - ) - return True - - @staticmethod - def _without_deployments_that_cannot_serve_route( - healthy_deployments: _RouteDeploymentsT, group: str - ) -> _RouteDeploymentsT: - handler_name: Final = _routed_call_type.get() - resolver: Final = _ROUTE_CAPABILITY_RESOLVERS.get(handler_name) if handler_name is not None else None - if resolver is None or isinstance(healthy_deployments, Mapping): - return healthy_deployments - - filtered_deployments: Final = [ - deployment - for deployment in healthy_deployments - if PythonCore._deployment_serves_route( - cast(DeploymentTypedDict, deployment), # cast-ok: router candidates use the deployment shape - group, - resolver, - ) - ] - if not filtered_deployments: - return healthy_deployments - return cast(_RouteDeploymentsT, filtered_deployments) # cast-ok: same list shape as the input candidates - - @staticmethod - def _decisions_deployment_refuses_media( - deployment: DeploymentTypedDict, group: str, request_kwargs: Mapping[str, object], *, audio: bool, videos: bool - ) -> bool: - from litellm.decisions.call import deployment_refuses_media - - plain_deployment: Final = cast(Mapping[str, object], deployment) # cast-ok: router deployments are plain dicts - model_info: Final = cast(Mapping[str, object], plain_deployment.get("model_info") or {}) # cast-ok: plain dict - deployment_params: Final = cast(Mapping[str, object], deployment["litellm_params"]) # cast-ok: plain dict - try: - model, custom_llm_provider = PythonCore._capability_model_and_provider(deployment["litellm_params"], group) - return deployment_refuses_media( - model, - custom_llm_provider, - model_info, - {**deployment_params, **request_kwargs}, - audio=audio, - videos=videos, - ) - except Exception as e: # noqa: BLE001 # best-effort filter: a deployment it cannot resolve answers for itself - verbose_router_logger.debug( - "litellm.router.py::_decisions_deployment_refuses_media: keeping a deployment of model=%s. Got - %s", - group, - e, - ) - return False - - def _without_decisions_deployments_refusing_media( - self, healthy_deployments: _RouteDeploymentsT, group: str, request_kwargs: Mapping[str, object] | None - ) -> _RouteDeploymentsT: - call_kwargs: Final = request_kwargs or {} - audio: Final = bool(call_kwargs.get("audio")) - videos: Final = bool(call_kwargs.get("videos")) - if ( - _routed_call_type.get() not in _DECISIONS_MEDIA_CALL_TYPES - or not (audio or videos) - or isinstance(healthy_deployments, Mapping) - ): - return healthy_deployments - - sent_kwargs: Final = {**self._active_default_litellm_params(), **call_kwargs} - deployments_taking_media: Final = [ - deployment - for deployment in healthy_deployments - if not PythonCore._decisions_deployment_refuses_media( - cast(DeploymentTypedDict, deployment), # cast-ok: router candidates use the deployment shape - group, - sent_kwargs, - audio=audio, - videos=videos, - ) - ] - if not deployments_taking_media: - return healthy_deployments - return cast(_RouteDeploymentsT, deployments_taking_media) # cast-ok: same list shape as the input candidates - - @staticmethod - def _without_declared_text_only_deployments( - healthy_deployments: _RouteDeploymentsT, - group: str, - messages: list[dict[str, str]] | None, # mutable-ok: handed to _request_carries_image as-is - request_kwargs: Mapping[str, object] | None, - ) -> _RouteDeploymentsT: - if isinstance(healthy_deployments, Mapping) or not PythonCore._request_carries_image( - messages, request_kwargs or {} - ): - return healthy_deployments - - vision_capable_deployments: Final = [ - deployment - for deployment in healthy_deployments - if not PythonCore._deployment_declares_no_vision( - cast(DeploymentTypedDict, deployment), # cast-ok: router candidates use the deployment shape - group, - ) - ] - if not vision_capable_deployments: - return healthy_deployments - return cast(_RouteDeploymentsT, vision_capable_deployments) # cast-ok: same list shape as the input candidates - - @staticmethod - def _request_carries_image( - messages: object, - request_kwargs: Mapping[str, object], - ) -> bool: - from litellm.litellm_core_utils.prompt_templates.common_utils import request_contains_image_content - - return request_contains_image_content(messages or request_kwargs.get("messages") or request_kwargs.get("input")) - - @staticmethod - def _pre_call_check_requested_chat_params(request_kwargs: Mapping[str, object]) -> frozenset[str]: - requested: Final = litellm_utils.get_non_default_params(passed_params=dict(request_kwargs)) - return frozenset( - param - for param, value in requested.items() - if param in _PRE_CALL_CHECK_FILTERED_PARAMS and not isinstance(value, str) - ) - - def _pre_call_checks( - self, - model: str, - healthy_deployments: list, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - request_kwargs: dict | None = None, - input_token_count: int | None = None, - skip_inline_token_count: bool = False, - request_deployments: Sequence[DeploymentTypedDict] | None = None, - ): - """ - Filter out model in model group, if: - - - model context window < message length. For azure openai models, requires 'base_model' is set. - https://docs.litellm.ai/docs/proxy/cost_tracking#spend-tracking-for-azure-openai-models - - if region given, filter out models not in that region / unknown region - - [TODO] function call and model doesn't support function calling - """ - - verbose_router_logger.debug("Starting Pre-call checks for deployments in model=%s", model) - - # Optimized: Use list() shallow copy instead of deepcopy - # We only pop from the list, not modify deployment dicts - 100x+ faster on hot path (every request) - _returned_deployments = list(healthy_deployments) - - invalid_model_indices: Final = set[int]() # Use set for O(1) membership checks - - _context_window_error = False - _potential_error_str = "" - - has_countable_input: Final = (messages is not None or input is not None) and not compaction_pending( - request_kwargs - ) - deployment_max_input_tokens: Final = tuple( - self._deployment_max_input_tokens(model, deployment) for deployment in _returned_deployments - ) - input_tokens: Final = self._resolve_pre_call_check_input_tokens( - messages=messages, - input=cast(str | list[object] | None, input), # cast-ok: forwarded to the resolver - request_kwargs=cast(Mapping[str, object] | None, request_kwargs), # cast-ok: forwarded to the resolver - input_token_count=input_token_count, - count_inline=has_countable_input - and not skip_inline_token_count - and any(limit is not None for limit in deployment_max_input_tokens), - model=model, - deployments=_returned_deployments, - ) - request_settings: Final[Mapping[str, object]] = request_kwargs if request_kwargs is not None else {} - requested_params: Final[frozenset[str]] = ( - self._pre_call_check_requested_chat_params(request_settings) - if request_kwargs is not None and self._host.dynamic().drop_params is False - else frozenset() - ) - - for idx, deployment in enumerate(_returned_deployments): - # Cache nested dict access to avoid repeated temporary dict allocations - _litellm_params = deployment.get("litellm_params", {}) - _model_info = deployment.get("model_info", {}) - - max_input_tokens = deployment_max_input_tokens[idx] - if ( - max_input_tokens is not None - and has_countable_input - and input_tokens is not None - and input_tokens > max_input_tokens - ): - invalid_model_indices.add(idx) - _context_window_error = True - deployment_model = ( - _model_info.get("base_model") or _litellm_params.get("base_model") or _litellm_params.get("model") - ) - _potential_error_str += ( - f"Model={deployment_model}, Max Input Tokens={max_input_tokens}, Got={input_tokens}" - ) - continue - - ## REGION CHECK ## - if request_kwargs is not None and request_kwargs.get("allowed_model_region") is not None: - allowed_model_region = request_kwargs.get("allowed_model_region") - - if allowed_model_region is not None: - if not is_region_allowed( - litellm_params=LiteLLM_Params.model_validate(_litellm_params), - allowed_model_region=allowed_model_region, - ): - invalid_model_indices.add(idx) - continue - - typed_deployments: Final = tuple( - cast(DeploymentTypedDict, deployment) # cast-ok: router deployments are plain dicts - for deployment in _returned_deployments - ) - param_candidates: Final = tuple( - (idx, deployment) for idx, deployment in enumerate(typed_deployments) if idx not in invalid_model_indices - ) - candidates_lacking_params: Final = tuple( - self._deployment_lacks_params(deployment, model, requested_params) for _, deployment in param_candidates - ) - no_deployment_takes_params: Final = bool(candidates_lacking_params) and all(candidates_lacking_params) - candidates_failing_params: Final = ( - tuple( - self._deployment_rejects_params(deployment, model, requested_params, request_settings) - for _, deployment in param_candidates - ) - if no_deployment_takes_params - else candidates_lacking_params - ) - excluded_indices: Final = frozenset(invalid_model_indices) | frozenset( - idx for (idx, _), fails in zip(param_candidates, candidates_failing_params, strict=True) if fails - ) - - if len(excluded_indices) == len(_returned_deployments): - group_deployments: Final = ( - tuple(request_deployments) - if request_deployments is not None - else (tuple(self._core_get_all_deployments(model_name=model)) or typed_deployments) - ) - if ( - not _context_window_error - and group_deployments - and all( - self._deployment_rejects_params(deployment, model, requested_params, request_settings) - for deployment in group_deployments - ) - ): - raise litellm_exceptions.UnsupportedParamsError( - message=( - f"no deployment of model={model} supports {sorted(requested_params)}. " - "To drop these, set `litellm.drop_params=True` or for proxy `litellm_settings: drop_params: true`, " - f"or send allowed_openai_params={sorted(requested_params)} to pass them through" - ), - model=model, - ) - if _context_window_error is True: - raise litellm_exceptions.ContextWindowExceededError( - message=f"litellm._pre_call_checks: Context Window exceeded for given call. No models have context window large enough for this call.\n{_potential_error_str}", - model=model, - llm_provider="", - ) - if len(excluded_indices) > 0: - # Single-pass filter using set for O(1) lookups (avoids O(n^2) from repeated pops) - _returned_deployments = [d for i, d in enumerate(_returned_deployments) if i not in excluded_indices] - - return _returned_deployments - - def _core_get_model_from_alias(self, model: str) -> str | None: - """ - Get the model from the alias. - - Returns: - - str, the litellm model name - - None, if model is not in model group alias - """ - return resolve_model_group_alias(self._core_model_group_alias, model) - - def _get_deployment_by_litellm_model(self, model: str) -> list[DeploymentTypedDict]: - """ - Get the deployment by litellm model. - """ - return [m for m in self._core_model_list if m["litellm_params"]["model"] == model] - - def _core_get_routable_deployments( - self, - model: str, - request_kwargs: Mapping[str, object] = MappingProxyType({}), - call_type: str | None = None, - ) -> tuple[DeploymentTypedDict, ...]: - try: - with endpoint_selection(call_type, model=model): - _, candidates = self._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # its signature leaves input and the result untyped - except (litellm_exceptions.BadRequestError, ValueError): - return () - if isinstance(candidates, dict): - return (cast(DeploymentTypedDict, candidates),) # cast-ok: the router returns one deployment as a dict - return tuple(cast(list[DeploymentTypedDict], candidates)) # cast-ok: the router returns deployments as a list - - def _try_early_resolve_deployments_for_model_not_in_names( - self, - model: str, - request_team_id: str | None, - include_team_models: bool = False, - ) -> tuple[str, list | dict] | None: - """ - When ``model`` is not in ``self._core_model_names``, try team routes, pattern routes, - team pattern routers, then default deployment. Returns None if none apply. - """ - if model in self._core_model_names: - return None - # Check for team-specific deployments by team_public_model_name. - # This intentionally takes priority over team pattern routers below, - # so that named team deployments shadow wildcard/pattern routes. - if request_team_id is not None: - team_deployments = self._core_get_all_deployments(model_name=model, team_id=request_team_id) - if team_deployments: - return model, team_deployments - elif include_team_models: - team_deployments = self._team_deployments_across_teams(model) - if team_deployments: - return model, team_deployments - - wildcard_deployments: Final = self._core_wildcard_route_deployments(model, request_team_id) - if wildcard_deployments: - return model, list(wildcard_deployments) - - if self._core_default_deployment is not None: - # Shallow copy with nested litellm_params copy (100x+ faster than deepcopy) - updated_deployment: Final = self._core_default_deployment.copy() - updated_deployment["litellm_params"] = self._core_default_deployment["litellm_params"].copy() - updated_deployment["litellm_params"]["model"] = model - return model, updated_deployment - - return None - - def _team_deployments_across_teams(self, model: str) -> list[DeploymentTypedDict]: - """Every team's deployments under public name `model`, for a proxy admin calling without a team.""" - with self._deployment_index_lock: - team_deployments: Final = [ - self._core_model_list[index] - for (_, public_model_name), indices in self._core_team_model_to_deployment_indices.items() - if public_model_name == model - for index in indices - ] - team_ids: Final = { - team_id - for deployment in team_deployments - for team_id in [(deployment.get("model_info") or {}).get("team_id")] - if team_id is not None - } - if len(team_ids) > 1: - raise litellm_exceptions.BadRequestError( - message=( - f"Model name '{model}' matches deployments from multiple teams. " - "Specify the deployment ID directly to disambiguate." - ), - model=model, - llm_provider="", - ) - return team_deployments - - def _core_deployments_for_request( - self, model: str, request_kwargs: Mapping[str, object] - ) -> Sequence[DeploymentTypedDict]: - """The deployments `model` names for this caller, through the same alias, then team-first, then - global, then admin-across-teams resolution `_common_checks_available_deployment` applies, so - strategy selection and compression policy can never disagree with deployment selection about - which marker a name means.""" - registered_name: Final = self._core_get_model_from_alias(model=model) or model - team_id: Final = get_request_team_id(request_kwargs) - deployments: Final = self._core_get_all_deployments(model_name=registered_name, team_id=team_id) - if deployments or team_id is not None or not is_proxy_admin_request(request_kwargs): - return deployments - return self._team_deployments_across_teams(registered_name) - - @staticmethod - def _is_strategy_marker_deployment(deployment: Mapping[str, object]) -> bool: - litellm_params: Final = deployment.get("litellm_params") - if not isinstance(litellm_params, Mapping): - return False - deployment_model: Final = litellm_params.get("model") - return isinstance(deployment_model, str) and classify_strategy_router_model(deployment_model) is not None - - def _common_checks_available_deployment( - self, - model: str, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - request_kwargs: Mapping[str, object] | None = None, - ) -> tuple[str, list | dict]: - """ - Common checks for 'get_available_deployment' across sync + async call. - - If 'healthy_deployments' returned is None, this means the user chose a specific deployment - - Returns - - str, the litellm model name - - List, if multiple models chosen - - Dict, if specific model chosen - """ - - request_team_id: Final = get_request_team_id(request_kwargs) - # check if aliases set on litellm model alias map - if specific_deployment is True: - return model, self._drop_strategy_markers( - model, - self._filter_reserved_deployments( - model=model, - healthy_deployments=self._get_deployment_by_litellm_model(model=model), - request_team_id=request_team_id, - ), - ) - elif model not in self._core_model_names and self._core_has_model_id(model): - deployment: Final = self._core_get_deployment(model_id=model) - if deployment is not None: - deployment_model: Final = deployment.litellm_params.model - return deployment_model, cast( # cast-ok: contract requires a plain dict for a single deployment - dict, - self._filter_reserved_deployments( - model=deployment_model, - healthy_deployments=( - cast( # cast-ok: model_dump of a router deployment - DeploymentTypedDict, deployment.model_dump(exclude_none=True) - ), - ), - request_team_id=request_team_id, - )[0], - ) - raise ValueError( - f"LiteLLM Router: Trying to call specific deployment, but Model ID :{model} does not exist in Model ID map" - ) - - _model_from_alias: Final = self._core_get_model_from_alias(model=model) - if _model_from_alias is not None: - model = _model_from_alias - - _routing_group_deployments: Final = self._get_routing_group_deployments(model=model, team_id=request_team_id) - if _routing_group_deployments is None: - early: Final = self._try_early_resolve_deployments_for_model_not_in_names( - model=model, - request_team_id=request_team_id, - include_team_models=is_proxy_admin_request(request_kwargs), - ) - if early is not None: - if not isinstance(early[1], list): - return early[0], cast( # cast-ok: contract requires a plain dict for a single deployment - dict, - self._filter_reserved_deployments( - model=early[0], - healthy_deployments=( - cast( # cast-ok: early resolve returns a router deployment - DeploymentTypedDict, early[1] - ), - ), - request_team_id=request_team_id, - )[0], - ) - return early[0], self._drop_strategy_markers( - early[0], - self._filter_reserved_deployments( - model=early[0], - healthy_deployments=early[1], - request_team_id=request_team_id, - ), - ) - - ## get healthy deployments - ### get all deployments - healthy_deployments = ( - _routing_group_deployments - if _routing_group_deployments is not None - else self._core_get_all_deployments(model_name=model, team_id=request_team_id) - ) - _pre_model_access_group_filter_len: Final = len(healthy_deployments) - healthy_deployments = self._filter_reserved_deployments( - model=model, - healthy_deployments=self._filter_deployments_by_model_access_groups( - model=model, - healthy_deployments=healthy_deployments, - request_kwargs=request_kwargs, - request_team_id=request_team_id, - ), - request_team_id=request_team_id, - ) - _access_group_filter_emptied_candidates = ( - _pre_model_access_group_filter_len > 0 and len(healthy_deployments) == 0 - ) - - if len(healthy_deployments) == 0: - # check if the user sent in a deployment name instead - # Do not fall back when access-group filtering removed every candidate; - # _get_deployment_by_litellm_model does not re-apply that filter. - if _pre_model_access_group_filter_len == 0: - _litellm_model_deployments: Final = self._get_deployment_by_litellm_model(model=model) - healthy_deployments = self._filter_reserved_deployments( - model=model, - healthy_deployments=self._filter_deployments_by_model_access_groups( - model=model, - healthy_deployments=_litellm_model_deployments, - request_kwargs=request_kwargs, - request_team_id=request_team_id, - ), - request_team_id=request_team_id, - ) - # If the litellm-model lookup produced candidates that access-group - # filtering then removed, treat this the same as the by-name path - # being emptied: prevent default-model fallback from bypassing the - # restriction (the fallback model may have no access_groups and - # would short-circuit the filter). - if len(_litellm_model_deployments) > 0 and len(healthy_deployments) == 0: - _access_group_filter_emptied_candidates = True - - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("initial list of deployments: %s", healthy_deployments) - - if len(healthy_deployments) == 0: - # Check for default fallbacks if no deployments are found for the requested model - # Do not fall back to another model when access-group filtering removed every - # candidate for the requested name: re-filtering the fallback model can be a - # no-op when it has no access_groups, incorrectly serving a different model. - if self._has_default_fallbacks() and not _access_group_filter_emptied_candidates: - fallback_model: Final = self._get_first_default_fallback() - if fallback_model: - verbose_router_logger.info( - "Model '%s' not found. Attempting to use default fallback model '%s'.", model, fallback_model - ) - # Re-assign model to the fallback and try to get deployments again - model = fallback_model - healthy_deployments = self._core_get_all_deployments(model_name=model, team_id=request_team_id) - healthy_deployments = self._filter_reserved_deployments( - model=model, - healthy_deployments=self._filter_deployments_by_model_access_groups( - model=model, - healthy_deployments=healthy_deployments, - request_kwargs=request_kwargs, - request_team_id=request_team_id, - ), - request_team_id=request_team_id, - ) - - # If still no deployments after checking for fallbacks, raise an error - if len(healthy_deployments) == 0: - raise litellm_exceptions.BadRequestError( - message=f"You passed in model={model}. {RouterErrors.no_healthy_deployments.value}", - model=model, - llm_provider="", - ) - - if self._host.dynamic().model_alias_map and model in self._host.dynamic().model_alias_map: - model = self._host.dynamic().model_alias_map[ - model - ] # update the model to the actual value if an alias has been passed in - - return model, self._drop_strategy_markers(model, healthy_deployments) - - def _drop_strategy_markers( - self, model: str, deployments: Sequence[DeploymentTypedDict] - ) -> list[DeploymentTypedDict]: - """A strategy marker is never a callable deployment, whichever resolution arm produced it.""" - selectable: Final = [d for d in deployments if not self._is_strategy_marker_deployment(d)] - if deployments and not selectable: - raise litellm_exceptions.BadRequestError( - message=f"You passed in model={model}. {RouterErrors.only_strategy_marker_deployments.value}", - model=model, - llm_provider="", - ) - eligible: Final = eligible_deployments(model, selectable) - if isinstance(eligible, tuple): - return list(eligible) - raise eligible - - def _filter_reserved_deployments( - self, - model: str, - healthy_deployments: Sequence[DeploymentTypedDict], - request_team_id: str | None, - ) -> tuple[DeploymentTypedDict, ...]: - result: Final = filter_reserved_deployments( - self._drop_strategy_markers(model, healthy_deployments), request_team_id, now=datetime.now(timezone.utc) - ) - if result.blocking_window is not None and len(result.deployments) == 0: - raise litellm_exceptions.BadRequestError( - message=f"Deployment {model} is reserved for another team until " - f"{result.blocking_window.end:%H:%M} {result.blocking_window.timezone}", - model=model, - llm_provider="", - ) - shared: Final = filter_ptu_shared_deployments(result.deployments, request_team_id) - if shared.withheld and len(shared.deployments) == 0: - raise litellm_exceptions.BadRequestError( - message=f"Deployment {model} is reserved for the teams holding a PTU share of it", - model=model, - llm_provider="", - ) - return shared.deployments - - def _filter_deployments_by_model_access_groups( - self, - model: str, - healthy_deployments: list[DeploymentTypedDict], - request_kwargs: Mapping[str, object] | None, - request_team_id: str | None, - ) -> list[DeploymentTypedDict]: - """ - Restrict candidate deployments to caller-authorized model access groups. - - This is only applied when: - - request metadata includes `user_api_key_auth`, and - - caller permissions for this model are access-group-only - (no explicit model, wildcard, or all-proxy grants). - """ - if not healthy_deployments or request_kwargs is None: - return healthy_deployments - - user_api_key_auth: Final = request_user_api_key_auth(request_kwargs) - if user_api_key_auth is None: - return healthy_deployments - - object_models: Final = set(getattr(user_api_key_auth, "models", []) or []) - object_team_models: Final = set(getattr(user_api_key_auth, "team_models", []) or []) - allowed_models: Final = object_models | object_team_models - if not allowed_models: - return healthy_deployments - - # If caller has direct model/wildcard/all-proxy access, do not constrain - # deployment choice by access group. - if model in allowed_models or "*" in allowed_models or "all-proxy-models" in allowed_models: - return healthy_deployments - - access_groups_for_model: Final = self._core_get_model_access_groups(model_name=model, team_id=request_team_id) - if len(access_groups_for_model) == 0: - return healthy_deployments - - allowed_access_groups: Final = set(access_groups_for_model.keys()) & allowed_models - if not allowed_access_groups: - # No overlap means this request was not authorized via model access - # group membership for this model, so do not force group filtering. - return healthy_deployments - - filtered_deployments: Final = [] - for deployment in healthy_deployments: - deployment_model_info = deployment.get("model_info") or {} - deployment_access_groups = set(deployment_model_info.get("access_groups", []) or []) - if deployment_access_groups & allowed_access_groups: - filtered_deployments.append(deployment) - - return filtered_deployments - - async def _core_async_get_healthy_deployments( - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - parent_otel_span: Span | None = None, - health_check_probe: bool = False, - ) -> list[dict] | dict: - """ - Get the healthy deployments for a model. - - Returns: - - List[Dict], if multiple models chosen - *OR* - - Dict, if specific model chosen - """ - - model, healthy_deployments = self._common_checks_available_deployment( - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) - - # IF TEAM ID SPECIFIED ON MODEL, AND REQUEST CONTAINS USER_API_KEY_TEAM_ID, FILTER OUT MODELS THAT ARE NOT IN THE TEAM - ## THIS PREVENTS WRITING FILES OF OTHER TEAMS TO MODELS THAT ARE TEAM-ONLY MODELS - by_id: Final = healthy_deployments if isinstance(healthy_deployments, dict) else None # pyright: ignore[reportUnknownVariableType] # the legacy deployment dict - healthy_deployments = filter_team_based_models( - healthy_deployments=healthy_deployments, - request_kwargs=request_kwargs, - ) - if by_id is not None and not healthy_deployments: - # Another team's deployment asked for by id: refused under its group's name. - group_name: Final = str(by_id.get("model_name")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above - model = group_name # rebind-ok: the error names the group - - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("healthy_deployments after team filter: %s", healthy_deployments) - - healthy_deployments = filter_web_search_deployments( - healthy_deployments=healthy_deployments, - request_kwargs=request_kwargs, - ) - - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("healthy_deployments after web search filter: %s", healthy_deployments) - - healthy_deployments = ( # pyright: ignore[reportUnknownVariableType] # rebind-ok: route filter narrows - self._without_deployments_that_cannot_serve_route( - healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type - model, - ) - ) - - if isinstance(healthy_deployments, dict): - if (healthy_deployments.get("model_info") or {}).get("blocked") is True: - raise litellm_exceptions.ServiceUnavailableError( - message=f"Model '{model}' is currently paused and cannot accept requests.", - model=model, - llm_provider="", - ) - return healthy_deployments - - request_deployments: Final = cast( # cast-ok: router deployments are plain dicts - Sequence[DeploymentTypedDict], healthy_deployments - ) - - # Health-check-based filtering (before cooldown) - healthy_deployments = ( # rebind-ok: filtered candidates enter health checks - await self._async_filter_health_check_unhealthy_deployments( - healthy_deployments=healthy_deployments, - parent_otel_span=parent_otel_span, - health_check_probe=health_check_probe, - ) - ) - - routing_read_batch: Final = RoutingReadBatch.active() - cooldown_deployments: Final = ( - await async_get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - if routing_read_batch is None - else await routing_read_batch.async_get_cooldown_deployments( - ctx=self._core_ctx, - healthy_deployments=healthy_deployments, - parent_otel_span=parent_otel_span, - ) - ) - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("cooldown deployments: %s", cooldown_deployments) - _pre_cooldown_deployments: Final = healthy_deployments - healthy_deployments = self._filter_cooldown_deployments( - healthy_deployments=healthy_deployments, - cooldown_deployments=cooldown_deployments, - ) - # Safety net: only bypass cooldown filter when health-check routing is - # driving cooldown (i.e. allowed_fails_policy is set). Without a policy, - # cooldowns are from real request failures and must not be bypassed. - if ( - not healthy_deployments - and self._core_enable_health_check_routing - and self._core_allowed_fails_policy is not None - ): - verbose_router_logger.warning( - "All deployments in cooldown via health-check routing, bypassing cooldown filter" - ) - healthy_deployments = _pre_cooldown_deployments - - healthy_deployments = self._filter_blocked_deployments(healthy_deployments) - - healthy_deployments = await self.async_callback_filter_deployments( - model=model, - healthy_deployments=healthy_deployments, - messages=(cast(list[AllMessageValues], messages) if messages is not None else None), - request_kwargs=request_kwargs, - parent_otel_span=parent_otel_span, - ) - - if self._core_enable_pre_call_checks and (messages is not None or input is not None): - deployments_to_check: Final = healthy_deployments - healthy_deployments = self._pre_call_checks( - model=model, - healthy_deployments=deployments_to_check, - messages=messages, - input=input, - request_kwargs=request_kwargs, - input_token_count=await self._acount_pre_call_check_tokens( - model=model, - healthy_deployments=deployments_to_check, - messages=messages, - input=input, - request_kwargs=request_kwargs, - ), - skip_inline_token_count=True, - request_deployments=request_deployments, - ) - # check if user wants to do tag based routing - healthy_deployments = await get_deployments_for_tag( - ctx=self._core_ctx, - model=model, - request_kwargs=request_kwargs, - healthy_deployments=healthy_deployments, - metadata_variable_name=self._get_metadata_variable_name_from_kwargs(request_kwargs), - request_deployments=request_deployments, - ) - - # narrow to whatever `self._core_routing_plugins` left in candidate_models - healthy_deployments = self._filter_by_routing_plugin_candidates( - healthy_deployments=healthy_deployments, - request_kwargs=request_kwargs, - ) - - ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) - _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = self._without_decisions_deployments_refusing_media( - self._without_declared_text_only_deployments( - cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts - model, - messages, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict - ), - model, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict - ) - embedding_batch_admitting: Final = deployments_admitting_embedding_batch( - batch_candidates, - _OPAQUE_VALUE_ADAPTER.validate_python(input), - _OPAQUE_VALUE_ADAPTER.validate_python(request_kwargs), - ) - embedding_batch_routed: Final = ( - embedding_batch_admitting - if embedding_batch_admitting is not None - else deployments_when_none_admits_embedding_batch( - batch_candidates, - cast(list[dict[str, object]], _pre_cooldown_deployments), # cast-ok: router deployments are plain dicts - await self._async_routable_cooled_deployment_ids( - model=model, - request_deployments=request_deployments, - cooldown_deployments=cooldown_deployments, - request_kwargs=cast(Mapping[str, object], request_kwargs), # cast-ok: str-keyed request kwargs - has_request_input=True, - ), - _OPAQUE_VALUE_ADAPTER.validate_python(input), - _OPAQUE_VALUE_ADAPTER.validate_python(request_kwargs), - ) - ) - healthy_deployments = litellm_utils.get_order_filtered_deployments( - list( - self._narrow_to_video_deployment( - without_deployments_that_cannot_batch( - embedding_batch_routed, caller_deployments=lambda: batch_candidates - ) - ) - ), - target_order=_target_order, - ) - - ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in - ## this request via weighted-failover. Always honored, regardless of the - ## router-level flag, so a stale exclusion key on kwargs cannot escape. - _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm_utils.get_excluded_filtered_deployments( - cast(list[dict], healthy_deployments), - excluded_deployment_ids=_excluded_deployment_ids, - ) - - ## RETRY SKIP ## -> drop deployments that already refused this request with a - ## non-retryable status, unless that leaves nothing, so the caller still gets - ## the provider's own error instead of a no-deployments error. - _retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids( - request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None - ) - healthy_deployments = ( - litellm_utils.get_excluded_filtered_deployments( - healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids - ) - or healthy_deployments - ) - - if len(healthy_deployments) == 0: - exception: Final = await async_raise_no_deployment_exception( - ctx=self._core_ctx, - model=model, - parent_otel_span=parent_otel_span, - route_model_ids=await self._async_routable_cooled_deployment_ids( - model=model, - request_deployments=request_deployments, - cooldown_deployments=cooldown_deployments, - request_kwargs=cast(Mapping[str, object], request_kwargs), # cast-ok: str-keyed request kwargs - has_request_input=messages is not None or input is not None, - ), - ) - raise exception - - return healthy_deployments - - @staticmethod - def _pop_effort_from_nested_carrier(request_kwargs: dict[str, object], carrier: str) -> None: - nested: Final = request_kwargs.get(carrier) - if not isinstance(nested, Mapping): - return - sanitized: Final = {key: value for key, value in nested.items() if key != "effort"} - if sanitized: - request_kwargs[carrier] = sanitized # rebind-ok: copy-on-write, so a shared nested carrier is never edited - else: - request_kwargs.pop(carrier, None) - - @staticmethod - def _tier_ceiling_under_the_surface_name( - tier_litellm_params: Mapping[str, object], responses_call: bool - ) -> Mapping[str, object]: - """``max_tokens``, ``max_completion_tokens`` and ``max_output_tokens`` are one - ceiling under three names, and each surface reads exactly one of them: the - Responses bridge builds its internal ``max_tokens`` from ``max_output_tokens`` - and would overwrite the tier's, chat and /v1/messages never read - ``max_output_tokens``, and litellm already renames ``max_tokens`` to - ``max_completion_tokens`` for the OpenAI models that require it. Collapse - whatever the tier carries onto the surface's own name, preferring a value the - operator already wrote under that name.""" - surface_key: Final = "max_output_tokens" if responses_call else "max_tokens" - carried: Final = tuple( - key - for key in (surface_key, "max_tokens", "max_completion_tokens", "max_output_tokens") - if key in tier_litellm_params - ) - if not carried: - return tier_litellm_params - return MappingProxyType( - { - **{k: v for k, v in tier_litellm_params.items() if k not in OUTPUT_TOKEN_CEILING_PARAMS}, - surface_key: tier_litellm_params[carried[0]], - } - ) - - def _pin_tier_params_onto_request( - self, - model: str, - tier_litellm_params: Mapping[str, object] | None, - request_kwargs: dict, - responses_call: bool, - ) -> bool: - """Apply per-tier litellm_params and report whether they pinned an output ceiling.""" - if not tier_litellm_params: - return False - reject_api_key_secret_reference(tier_litellm_params) - accepted_tier_params: Final = self._tier_params_the_target_accepts(model, tier_litellm_params, request_kwargs) - surface_tier_params: Final = self._tier_ceiling_under_the_surface_name( - accepted_tier_params, responses_call=responses_call - ) - self._drop_client_carriers_a_tier_pin_supersedes(request_kwargs, surface_tier_params) - request_kwargs.update(surface_tier_params) - return not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(surface_tier_params) - - @staticmethod - def _restore_client_output_ceiling(request_kwargs: MutableMapping[str, object]) -> None: - """Restore caller intent before routing, so neither a previous tier pin nor alias defaults - become the next attempt's budget. The proxy strips caller-written stamps at ingestion.""" - callers_ceiling: Final = stamped_client_ceiling(request_kwargs) - if callers_ceiling is None: - return - for carrier in OUTPUT_TOKEN_CEILING_PARAMS: - request_kwargs.pop(carrier, None) - request_kwargs.update(callers_ceiling) - - @staticmethod - def _stamp_client_output_ceiling( - request_kwargs: dict[str, object], # mutable-ok: preserve request-owned metadata across routing attempts - ) -> None: - _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) - metadata_bucket.setdefault( - CLIENT_OUTPUT_CEILING_METADATA_KEY, - {carrier: request_kwargs[carrier] for carrier in OUTPUT_TOKEN_CEILING_PARAMS if carrier in request_kwargs}, - ) - - @staticmethod - def _drop_client_carriers_a_tier_pin_supersedes( - request_kwargs: dict[str, object], - tier_litellm_params: Mapping[str, object], - ) -> None: - """Tier litellm_params are deliberate operator overrides, but provider - translations let a caller-supplied carrier of the same setting - (``thinking``, ``output_config.effort``, ``reasoning.effort``) outrank - the ``reasoning_effort`` alias, so a pinned effort only reaches the wire - if the client's other encodings are removed before the merge. Non-effort - fields a carrier also holds (``output_config.format``, - ``reasoning.summary``) are kept. An output ceiling has the same shape: - ``max_tokens``, ``max_completion_tokens`` and ``max_output_tokens`` are - one setting under three names, and a provider handed two of them either - rejects the request or picks one by iteration order.""" - if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_litellm_params): - PythonCore._stamp_client_output_ceiling(request_kwargs) - for carrier in OUTPUT_TOKEN_CEILING_PARAMS: - request_kwargs.pop(carrier, None) - if "reasoning_effort" not in tier_litellm_params: - return - request_kwargs.pop("thinking", None) - PythonCore._pop_effort_from_nested_carrier(request_kwargs, "output_config") - PythonCore._pop_effort_from_nested_carrier(request_kwargs, "reasoning") - - async def _async_get_available_deployment_unresolved( # noqa: ANN202 # legacy selector result remains inferred - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - ): - """ - Async implementation of 'get_available_deployments'. - - Allows all cache calls to be made async => 10x perf impact (8rps -> 100 rps). - """ - with phase_span(f"route {model}"): - return await self._async_get_available_deployment( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - reserve=pick_reserves(default=True), - ) - - async def _async_get_available_deployment( - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, str]] | None, - input: str | list | None, - specific_deployment: bool | None, - reserve: bool, - ): - if ( - self._core_routing_strategy != "usage-based-routing-v2" - and self._core_routing_strategy != "simple-shuffle" - and self._core_routing_strategy != "cost-based-routing" - and self._core_routing_strategy != "latency-based-routing" - and self._core_routing_strategy != "least-busy" - ): # prevent regressions for other routing strategies, that don't have async get available deployments implemented. - return self._get_available_deployment_unresolved( - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) - try: - parent_otel_span: Final = get_parent_otel_span_from_kwargs(request_kwargs) - - ######################################################### - # Execute Pre-Routing Hooks - # this hook can modify the model, messages before the routing decision is made - ######################################################### - responses_call: Final = input is not None and messages is None - pre_routing_hook_response: Final = await self._core_async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) - if pre_routing_hook_response is not None: - model = pre_routing_hook_response.model - messages = pre_routing_hook_response.messages - record_pre_routing_selection(request_kwargs, model) - self._pin_tier_params_onto_request( - model=model, - tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, - request_kwargs=request_kwargs, - responses_call=responses_call, - ) - ######################################################### - - # Resolve the strategy and logger AFTER the pre-routing hook, since - # the hook can replace `model` and routing-group lookup must key - # off the final model name. - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) - pick_attributes: Final = _deployment_pick_attributes( - model, _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) - ) - routing_read_batch: Final = RoutingReadBatch.for_strategy(strategy, strategy_selector) - - with RoutingReadBatch.scoped(routing_read_batch): - healthy_deployments: Final = await self._core_async_get_healthy_deployments( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - parent_otel_span=parent_otel_span, - ) - if isinstance(healthy_deployments, dict): - await self._async_override_selector_pre_call_check( - strategy, strategy_selector, healthy_deployments, parent_otel_span - ) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return healthy_deployments - - # When encrypted content affinity pins to a specific deployment, - if request_kwargs.get("_encrypted_content_affinity_pinned") and len(healthy_deployments) == 1: - await self._async_override_selector_pre_call_check( - strategy, strategy_selector, healthy_deployments[0], parent_otel_span - ) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return healthy_deployments[0] - - start_time: Final = time.time() - if strategy == "simple-shuffle": - shuffled: Final = simple_shuffle( - resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=healthy_deployments, - model=model, - request_kwargs=request_kwargs, - ) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return shuffled - with PrefetchedUsage.scoped( - routing_read_batch.prefetched_usage if routing_read_batch is not None else None - ): - deployment: Final = await self._select_deployment_async( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - reserve=reserve and request_kwargs not in (None, {}), - ) - if deployment is None: - exception: Final = await async_raise_no_deployment_exception( - ctx=self._core_ctx, - model=model, - parent_otel_span=parent_otel_span, - ) - raise exception - await self._async_override_selector_pre_call_check( - strategy, strategy_selector, deployment, parent_otel_span - ) - verbose_router_logger.info( - "get_available_deployment for model: %s, Selected deployment: %s for model: %s", - model, - self.print_deployment(deployment), - model, - ) - - end_time: Final = time.time() - _duration: Final = end_time - start_time - asyncio.create_task( - self.service_logger_obj.async_service_success_hook( - service=ServiceTypes.ROUTER, - duration=_duration, - call_type=".async_get_available_deployments", - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - ) - ) - - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return deployment - except Exception as e: - mark_deployment_selection_failure(e) - traceback_exception: Final = traceback.format_exc() - # if router rejects call -> log to langfuse/otel/etc. - if request_kwargs is not None: - logging_obj: Final = request_kwargs.get("litellm_logging_obj", None) - - if logging_obj is not None: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=e, - traceback_exception=traceback_exception, - prefer_async_handlers=True, - ) - ) - raise e - - _core_async_get_available_deployment = ( # pyright: ignore[reportUnknownVariableType] # legacy selector return - async_resolve_pods_after( - _async_get_available_deployment_unresolved # pyright: ignore[reportUnknownArgumentType] # legacy type - ) - ) - - async def async_get_available_deployment_for_pass_through( - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - ): - """ - Async version of get_available_deployment_for_pass_through - - Only returns deployments configured with use_in_pass_through=True - """ - with phase_span(f"route {model}"): - return await self._async_get_available_deployment_for_pass_through( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) - - async def _async_get_available_deployment_for_pass_through( - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, str]] | None, - input: str | list | None, - specific_deployment: bool | None, - ): - try: - parent_otel_span: Final = get_parent_otel_span_from_kwargs(request_kwargs) - - # 1. Execute pre-routing hook - responses_call: Final = input is not None and messages is None - pre_routing_hook_response: Final = await self._core_async_pre_routing_hook( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - ) - if pre_routing_hook_response is not None: - model = pre_routing_hook_response.model - messages = pre_routing_hook_response.messages - record_pre_routing_selection(request_kwargs, model) - self._pin_tier_params_onto_request( - model=model, - tier_litellm_params=pre_routing_hook_response.litellm_params if pre_routing_hook_response else None, - request_kwargs=request_kwargs, - responses_call=responses_call, - ) - - # 2. Get healthy deployments - healthy_deployments: Final = await self._core_async_get_healthy_deployments( - model=model, - request_kwargs=request_kwargs, - messages=messages, - input=input, - specific_deployment=specific_deployment, - parent_otel_span=parent_otel_span, - ) - - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) - - # 3. If specific deployment returned, verify if it supports pass-through - if isinstance(healthy_deployments, dict): - if (healthy_deployments.get("model_info") or {}).get("blocked") is True: - raise litellm_exceptions.ServiceUnavailableError( - message=f"Model '{model}' is currently paused and cannot accept requests.", - model=model, - llm_provider="", - ) - litellm_params: Final = healthy_deployments.get("litellm_params", {}) - if litellm_params.get("use_in_pass_through"): - await self._async_override_selector_pre_call_check( - strategy, strategy_selector, healthy_deployments, parent_otel_span - ) - return healthy_deployments - else: - raise litellm_exceptions.BadRequestError( - message=f"Deployment {healthy_deployments.get('model_info', {}).get('id')} does not support pass-through endpoint (use_in_pass_through=False)", - model=model, - llm_provider="", - ) - - # 4. Filter deployments that support pass-through - pass_through_deployments = self._filter_pass_through_deployments(healthy_deployments=healthy_deployments) - - if len(pass_through_deployments) == 0: - raise litellm_exceptions.BadRequestError( - message=f"Model {model} has no deployments configured with use_in_pass_through=True. Please add use_in_pass_through: true to the deployment configuration", - model=model, - llm_provider="", - ) - - # 5. Apply load balancing strategy - start_time: Final = time.perf_counter() - if strategy == "simple-shuffle": - return simple_shuffle( - resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=pass_through_deployments, - model=model, - request_kwargs=request_kwargs, - ) - deployment: Final = await self._select_deployment_async( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=pass_through_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), - ) - - if deployment is None: - exception: Final = await async_raise_no_deployment_exception( - ctx=self._core_ctx, - model=model, - parent_otel_span=parent_otel_span, - ) - raise exception - await self._async_override_selector_pre_call_check( - strategy, strategy_selector, deployment, parent_otel_span - ) - - verbose_router_logger.info( - "async_get_available_deployment_for_pass_through model: %s, selected deployment: %s", - model, - self.print_deployment(deployment), - ) - - end_time: Final = time.perf_counter() - _duration: Final = end_time - start_time - asyncio.create_task( - self.service_logger_obj.async_service_success_hook( - service=ServiceTypes.ROUTER, - duration=_duration, - call_type=".async_get_available_deployments", - parent_otel_span=parent_otel_span, - start_time=start_time, - end_time=end_time, - ) - ) - - return deployment - except Exception as e: - mark_deployment_selection_failure(e) - traceback_exception: Final = traceback.format_exc() - if request_kwargs is not None: - logging_obj: Final = request_kwargs.get("litellm_logging_obj", None) - if logging_obj is not None: - asyncio.create_task( - logging_obj.dispatch_failure_handlers( - exception=e, - traceback_exception=traceback_exception, - prefer_async_handlers=True, - ) - ) - raise e - - async def _run_routing_plugins( - self, - model: str, - request_kwargs: dict, - messages: list[dict[str, object]] | None, - ) -> RoutingContext: - """ - Build a RoutingContext for `model`, run it through `self._core_routing_plugins` - in order, then stash the narrowed candidate list and accumulated signals - onto `request_kwargs["metadata"]` so `_filter_by_routing_plugin_candidates` - (called later, during healthy-deployment filtering) can consume them. - """ - from litellm.litellm_core_utils.prompt_templates.factory import ( - resolve_structured_messages, - ) - - candidate_models: Final = list(self._core_resolved_litellm_models(model)) - - metadata_key: Final = self._get_metadata_variable_name_from_kwargs(request_kwargs) - metadata: Final = request_kwargs.setdefault(metadata_key, {}) - - context = RoutingContext( - raw_messages=messages or [], - structured_messages=resolve_structured_messages(messages=messages, request_kwargs=request_kwargs) or [], - candidate_models=candidate_models, - metadata=metadata, - ) - - for plugin in self._core_routing_plugins: - context = await plugin.run(context) - - metadata["routing_plugin_signals"] = context.signals - if len(context.candidate_models) < len(candidate_models): - metadata["_routing_plugin_candidate_models"] = context.candidate_models - - return context - - def _filter_by_routing_plugin_candidates( - self, - healthy_deployments: list[dict] | dict, - request_kwargs: dict, - ) -> list[dict] | dict: - """ - Narrow `healthy_deployments` to whatever `self._core_routing_plugins` left in - `context.candidate_models`. Raises rather than silently falling back to - the unfiltered pool -- a plugin narrowing to nothing is a policy decision - (e.g. no model this tenant's budget allows), not something to bypass. - """ - if not self._core_routing_plugins or not isinstance(healthy_deployments, list): - return healthy_deployments - - metadata_key: Final = self._get_metadata_variable_name_from_kwargs(request_kwargs) - candidate_models: Final = (request_kwargs.get(metadata_key) or {}).get("_routing_plugin_candidate_models") - # `is None` (not falsy-check): a plugin narrowing to an empty list must - # still hit the "no deployments left" raise below, not be treated the - # same as "no plugin ever set this key". - if candidate_models is None: - return healthy_deployments - - candidate_set: Final = set(candidate_models) - filtered: Final = [d for d in healthy_deployments if d.get("litellm_params", {}).get("model") in candidate_set] - - if not filtered: - raise ValueError(f"No deployments left after routing-plugin filtering. candidate_models={candidate_models}") - - return filtered - - def _select_pre_routing_strategy( - self, model: str, request_kwargs: Mapping[str, object] - ) -> "TaggedPreRoutingStrategy[PreRoutingStrategy] | None": - """ - Resolve the pre-routing strategy for `model`, disambiguating deployments - that share a `model_name` by matching the request's tags against each - registered strategy's tags before falling back to the first registered. - Returns the tagged registry entry so the caller can tell whether the - request's tags were what selected it, and can locate the marker - deployment the strategy was registered from via its (model_name, tags) - pair. - - The registries are keyed by the marker deployment's own `model_name`, which - for a team-scoped router is the internal `model_name_{team}_{uuid}` while - the caller sends the team's public name. So the names looked up are the - `model_name`s of whatever deployments this caller's request resolves `model` - to, and `model` itself when it resolves to none. - - With tag filtering enabled, router-wide or by the request's - enable_tag_filtering (which the proxy sets from key/team - router_settings), strategies that all carry real tags matching none of - the request's do not capture it when the name also has plain - deployments: returning None hands the request to ordinary tag-aware - deployment selection. - """ - registries: Final = ( - self._core_auto_routers, - self._core_complexity_routers, - self._core_adaptive_routers, - self._core_quality_routers, - ) - if not any(registries): - return None - deployments: Final = self._core_deployments_for_request(model, request_kwargs) - registered_names: Final = tuple(dict.fromkeys(str(d["model_name"]) for d in deployments)) or (model,) - candidates: Final = tuple( - tagged for registry in registries for name in registered_names for tagged in registry.get(name, []) - ) - if not candidates: - return None - - request_tags: Final = get_routing_tags_from_request_kwargs(request_kwargs) - if request_tags: - for tagged in candidates: - if tagged.tags and is_valid_deployment_tag( - list(tagged.tags), request_tags, self._core_tag_filtering_match_any - ): - return tagged - for tagged in candidates: - if "default" in tagged.tags: - return tagged - request_scoped_filtering: Final = request_kwargs.get("enable_tag_filtering") is True - if ( - (self._core_enable_tag_filtering or request_scoped_filtering) - and all(tagged.tags for tagged in candidates) - and any(not self._is_strategy_marker_deployment(d) for d in deployments) - ): - return None - return candidates[0] - - @staticmethod - def _request_header(request_kwargs: Mapping[str, object], header_name: str) -> str | None: - proxy_server_request: Final = request_kwargs.get("proxy_server_request") - if not isinstance(proxy_server_request, Mapping): - return None - headers: Final = proxy_server_request.get("headers") - if not isinstance(headers, Mapping): - return None - return next( - ( - value - for key, value in headers.items() - if isinstance(key, str) and key.lower() == header_name and isinstance(value, str) - ), - None, - ) - - def _claude_code_session_router_cache_key(self, request_kwargs: Mapping[str, object]) -> str | None: - session_id: Final = self._request_header(request_kwargs, "x-claude-code-session-id") - if session_id is None or _CLAUDE_CODE_SESSION_ID_RE.fullmatch(session_id) is None: - return None - metadata_name: Final = "litellm_metadata" if "litellm_metadata" in request_kwargs else "metadata" - metadata: Final = request_kwargs.get(metadata_name) - if not isinstance(metadata, Mapping): - return None - caller_scope: Final = metadata.get("user_api_key_hash") - if not isinstance(caller_scope, str) or not caller_scope: - return None - return f"claude_code_session_router:v1:{caller_scope}:{session_id}" - - @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) - async def _delete_claude_code_session_router_binding(self, cache_key: str) -> None: - try: - await self._claude_code_session_router_cache.async_delete_cache(key=cache_key) - except Exception as e: # noqa: BLE001 # cache cleanup must not fail an otherwise routable request - verbose_router_logger.warning( - "Failed to delete Claude Code session router binding; the binding may remain until its TTL expires: %s", - e, - ) - - @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) - async def _get_claude_code_session_router_binding(self, cache_key: str) -> object: - session_cache: Final = self._claude_code_session_router_cache - try: - if session_cache.redis_cache is None: - return await session_cache.async_get_cache(key=cache_key) - return await session_cache.redis_cache.async_get_cache(key=cache_key) - except Exception as e: # noqa: BLE001 # an optional binding must not make routing depend on Redis - log_redis_failure( - verbose_router_logger, - logging.WARNING, - "Failed to read Claude Code session router binding; using the requested model", - e, - ) - return None - - @with_service_target(CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET) - async def _resolve_claude_code_session_router( - self, - model: str, - registered_model_name: str, - request_kwargs: Mapping[str, object], - ) -> str: - if is_native_compaction_call(): - return registered_model_name - if not any( - ( - self._core_auto_routers, - self._core_complexity_routers, - self._core_adaptive_routers, - self._core_quality_routers, - ) - ): - return registered_model_name - cache_key: Final = self._claude_code_session_router_cache_key(request_kwargs) - if cache_key is None or not isinstance(request_kwargs, dict): - return registered_model_name - if request_kwargs.get("fallback_depth") not in (None, 0): - return registered_model_name - - agent_id: Final = self._request_header(request_kwargs, "x-claude-code-agent-id") - if agent_id is not None: - bound_model: Final = await self._get_claude_code_session_router_binding(cache_key) - if not isinstance(bound_model, str): - return registered_model_name - bound_registered_model: Final = self._core_get_model_from_alias(model=bound_model) or bound_model - if self._select_pre_routing_strategy(bound_registered_model, request_kwargs) is None: - await self._delete_claude_code_session_router_binding(cache_key) - return registered_model_name - await self._claude_code_session_router_cache.async_set_cache( - key=cache_key, - value=bound_model, - ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, - ) - self._stamp_or_clear_metadata_key(request_kwargs, "model_group", bound_model) - return bound_registered_model - - if self._select_pre_routing_strategy(registered_model_name, request_kwargs) is None: - return registered_model_name - await self._claude_code_session_router_cache.async_set_cache( - key=cache_key, - value=model, - ttl=_CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS, - ) - return registered_model_name - - async def _core_async_pre_routing_hook( - self, - model: str, - request_kwargs: dict[str, object], - messages: list[dict[str, Any]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - ) -> PreRoutingHookResponse | None: - """ - This hook is called before the routing decision is made. - - Used for the litellm auto-router to modify the request before the routing decision is made. - - `model` is whatever the caller asked for, which may be a `model_group_alias` key or a team's - public model name, while the strategy registries and the marker deployment are keyed by the - marker's own `model_name`, so every lookup below resolves the alias first and the team name - through the deployment path. Only the lookups: the caller-facing name stays the alias, since - spend metadata is stamped before routing and the response carries the tier group the - strategy picked. - """ - self._restore_client_output_ceiling(request_kwargs) - self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None) - requested_registered_model_name: Final = self._core_get_model_from_alias(model=model) or model - registered_model_name: Final = await self._resolve_claude_code_session_router( - model=model, - registered_model_name=requested_registered_model_name, - request_kwargs=request_kwargs, - ) - - ######################################################### - # Run the routing-plugin pipeline, if any plugins are configured. - # Plugins narrow the candidate deployment pool (consumed later by - # `_filter_by_routing_plugin_candidates`) and may attach signals for - # downstream strategies (auto-router, complexity-router, ...) to read. - ######################################################### - if self._core_routing_plugins: - await self._run_routing_plugins( - model=registered_model_name, request_kwargs=request_kwargs, messages=messages - ) - - selected_strategy: Final = self._select_pre_routing_strategy( - model=registered_model_name, request_kwargs=request_kwargs - ) - if selected_strategy is None: - await arm_compaction(request_kwargs, None) - self._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, value=None - ) - self._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, key=CONSUMED_REQUEST_TAGS_METADATA_KEY, value=None - ) - return None - - from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter - - marker: Final = self._selected_strategy_marker_deployment( - model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs - ) - definition: Final = selected_strategy.deployment - config_id: Final = definition.model_info.id if definition else None - updated_at: Final = definition.model_info.updated_at if definition else None - router_kind: Final = classify_strategy_router_model(definition.litellm_params.model) if definition else None - provenance: Final[StandardLoggingRoutingDecision] = { - "router_model_name": definition.model_name if definition else registered_model_name, - **({"router_type": router_kind} if router_kind is not None else {}), - **({"router_config_id": config_id} if config_id is not None else {}), - **({"router_config_updated_at": updated_at.isoformat()} if updated_at is not None else {}), - **( - {"router_config_fingerprint": selected_strategy.definition_fingerprint} - if selected_strategy.definition_fingerprint is not None - else {} - ), - } - phase_attributes(routing_decision_attributes(provenance)) - - reject_recursive_compactor(registered_model_name) - await arm_compaction( - request_kwargs, - selected_strategy.strategy.config.context_compaction - if isinstance(selected_strategy.strategy, ComplexityRouter) - else None, - tuple( - dict.fromkeys( - member - for pool in selected_strategy.strategy.config.tiers.values() - for member in ((pool,) if isinstance(pool, str) else pool) - ) - ) - if isinstance(selected_strategy.strategy, ComplexityRouter) - else (), - parent_model=model, - ctx=self._core_ctx, - allow_escalation=isinstance(selected_strategy.strategy, ComplexityRouter) - and selected_strategy.strategy.config.enable_context_window_escalation, - messages=messages, - ) - - await router_host().authorize_strategy_inference( - deployment=marker, - request_kwargs=request_kwargs, - llm_router=self._core_public(), - ) - - # Same tag-aware lookup the proxy's pre-call arming used, so an alias with - # several tag-scoped markers cannot suppress under one and route under another. - compression_policy: Final = router_host().compression_policy( - llm_router=self._core_public(), - model_alias=registered_model_name, - request_kwargs=request_kwargs, - request_tags=get_routing_tags_from_request_kwargs(request_kwargs), - ) - # Shared compression already ran in the pre-call hook, so reuse it rather than - # compressing twice. Conditional on arming having actually happened: only the - # proxy arms, and on the SDK path the shortcut would skip both hops entirely. - needs_independent_routing_compression: Final = compression_policy is not None and not ( - compression_policy.is_same - and compression_policy.model is not None - and router_host().model_hop_compression_armed() - ) - routing_messages: Final = ( - await router_host().messages_for_routing( - policy=compression_policy, messages=messages, request_kwargs=request_kwargs - ) - if needs_independent_routing_compression - else None - ) - - routed: Final = await selected_strategy.strategy.async_pre_routing_hook( - model=registered_model_name, - request_kwargs=request_kwargs, - messages=routing_messages if routing_messages is not None else messages, - input=input, - specific_deployment=specific_deployment, - ) - traced_routed: Final = ( - routed.model_copy(update={"routing_decision": {**provenance, **routed.routing_decision}}) - if routed is not None and routed.routing_decision is not None - else routed - ) - # Routing-only compression must not leak into the response: the model call and - # deployment-context filtering key off this field. Compared by value, since - # pydantic rebuilds the list rather than keeping the object passed in. - pre_routing_hook_response: Final = ( - traced_routed.model_copy(update={"messages": messages}) - if traced_routed is not None and routing_messages is not None and traced_routed.messages == routing_messages - else traced_routed - ) - self._record_routing_decision( - request_kwargs=request_kwargs, - routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None), - ) - self._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, - key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, - value=(pre_routing_hook_response.session_affinity_ttl_seconds if pre_routing_hook_response else None), - ) - self._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, - key=CONSUMED_REQUEST_TAGS_METADATA_KEY, - value=self._consumed_request_tags_stamp( - selected_strategy=selected_strategy, - pre_routing_hook_response=pre_routing_hook_response, - request_tags=get_routing_tags_from_request_kwargs(request_kwargs), - ), - ) - - # `model` (the alias, e.g. "smart-router") is never the deployment actually - # called - apply the router marker's own litellm_params to the request, - # since the tier/route deployment the hook selected won't have them. The - # marker entry is looked up by its `auto_router/` model prefix and the - # selected strategy's tags, never by list position: plain deployments may - # share the alias `model_name` and must not leak their params (`api_base`, - # `api_key`, ...) onto the routed call. Router-only fields (tpm, rpm, - # weight, complexity_router_config, ...) are excluded from the actual - # outbound LLM call downstream by litellm.types.utils.all_litellm_params, - # not here. Custom pricing fields ARE call params, so they must be - # excluded here: they price the alias, not the deployment the hook - # selected, and forwarding them re-registers the routed deployment at - # the alias's price (an explicit 0 makes every alias request bill $0). - # Forwarded params only fill gaps: the keys inserted here ride along on the - # request (top level, so sibling requests sharing a `metadata` dict never see - # them) until `_update_kwargs_with_deployment` drops any the selected - # deployment sets itself (its own `aws_region_name` beats the marker's). - # Per-tier `litellm_params` on the hook response are deliberate overrides - # the caller applies on top, so those keys and equivalent output-limit - # names are never tracked as alias defaults here. - marker_params: Final = ( - self._forwardable_alias_marker_params( - model=registered_model_name, strategy_tags=selected_strategy.tags, request_kwargs=request_kwargs - ) - if pre_routing_hook_response is not None - else () - ) - tier_param_keys: Final = ( - tuple(pre_routing_hook_response.litellm_params or ()) if pre_routing_hook_response is not None else () - ) - overridden_marker_keys: Final = frozenset(tier_param_keys) | ( - OUTPUT_TOKEN_CEILING_PARAMS if not OUTPUT_TOKEN_CEILING_PARAMS.isdisjoint(tier_param_keys) else frozenset() - ) - newly_forwarded: Final = tuple( - (key, value) - for key, value in marker_params - if key not in request_kwargs and key not in overridden_marker_keys - ) - if pre_routing_hook_response is not None: - self._stamp_client_output_ceiling(request_kwargs) - request_kwargs.pop(_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, None) - request_kwargs.update(newly_forwarded) - if newly_forwarded: - request_kwargs.update(((_ALIAS_MARKER_FORWARDED_PARAMS_KWARG, tuple(key for key, _ in newly_forwarded)),)) - - return pre_routing_hook_response - - def _selected_strategy_marker_deployment( - self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] - ) -> DeploymentTypedDict | None: - markers: Final = tuple( - deployment - for deployment in self._core_deployments_for_request(model, request_kwargs) - if "model" in deployment["litellm_params"] - and str(deployment["litellm_params"]["model"]).startswith(AUTO_ROUTER_MODEL_PREFIX) - ) - tag_matched: Final = tuple( - deployment - for deployment in markers - if (tuple(deployment["litellm_params"]["tags"] or ()) if "tags" in deployment["litellm_params"] else ()) - == strategy_tags - ) - return tag_matched[0] if tag_matched else (markers[0] if markers else None) - - def _forwardable_alias_marker_params( - self, model: str, strategy_tags: tuple[str, ...], request_kwargs: Mapping[str, object] - ) -> tuple[tuple[str, object], ...]: - marker: Final = self._selected_strategy_marker_deployment( - model=model, strategy_tags=strategy_tags, request_kwargs=request_kwargs - ) - if marker is None: - return () - return tuple( - (key, value) - for key, value in marker["litellm_params"].items() - if key not in _ALIAS_PARAMS_NEVER_FORWARDED - and key not in CustomPricingLiteLLMParams.model_fields - and value is not None - ) - - @staticmethod - def _forwarded_alias_marker_keys_the_deployment_sets( - deployment: Mapping[str, object], forwarded_keys: object - ) -> tuple[str, ...]: - deployment_litellm_params: Final = deployment.get("litellm_params") - if not isinstance(deployment_litellm_params, Mapping) or not isinstance(forwarded_keys, tuple): - return () - return tuple( - key - for key in forwarded_keys - if isinstance(key, str) and PythonCore._deployment_sets_litellm_param(deployment_litellm_params, key) - ) - - @staticmethod - def _deployment_sets_litellm_param(deployment_litellm_params: Mapping[str, object], key: str) -> bool: - value: Final = deployment_litellm_params.get(key) - if value is None: - return False - field: Final = LiteLLM_Params.model_fields.get(key) - return field is None or value != field.default - - def _consumed_request_tags_stamp( - self, - selected_strategy: "TaggedPreRoutingStrategy[PreRoutingStrategy]", - pre_routing_hook_response: PreRoutingHookResponse | None, - request_tags: Sequence[str], - ) -> ConsumedRequestTagsStamp | None: - """Record which tags picked the router and which model group it rewrote to, or None. - - A request whose tags matched the selected strategy's tags has already spent those - tags on picking the router; re-applying them to the routed tier's model group would - empty the pool unless every tier deployment repeats the marker's tag. Only the - strategy's own tags are spent: the request's other tags keep constraining - deployment selection inside the routed group, and key/team policy tags are - untouched because tag filtering separately re-applies whatever - `metadata.inherited_tags` carries for the stamped group. - """ - if pre_routing_hook_response is None or not selected_strategy.tags or not request_tags: - return None - if not is_valid_deployment_tag(selected_strategy.tags, request_tags, self._core_tag_filtering_match_any): - return None - return ConsumedRequestTagsStamp(model_group=pre_routing_hook_response.model, tags=selected_strategy.tags) - - @staticmethod - def _record_routing_decision( - request_kwargs: dict, - routing_decision: StandardLoggingRoutingDecision | None, - ) -> None: - """Make the request's metadata describe THIS routing attempt, and only this one. - - Fallbacks re-enter the hook with the same `request_kwargs`, so an attempt that - picks a plain model group after an auto-router group failed must clear the - earlier decision; leaving it would attribute the first router's tier and cause - to the deployment that actually served the request. Every attempt therefore - writes or clears, never just writes. - """ - from litellm.router_strategy.auto_router.baseline_request import capture_baseline_parameters - from litellm.types.router import BaselineRouteStamp - - phase_attributes(routing_decision_attributes(routing_decision)) - baseline_model: Final = routing_decision.get("savings_baseline_model") if routing_decision else None - baseline_id: Final = routing_decision.get("savings_baseline_deployment_id") if routing_decision else None - router_name: Final = routing_decision.get("router_model_name") if routing_decision else None - caller_parameters: Final = ( - capture_baseline_parameters(request_kwargs) if router_name and baseline_model else None - ) - PythonCore._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, - key="_autorouter_baseline_route", - value=( - BaselineRouteStamp(router_name, baseline_model, baseline_id, caller_parameters) - if router_name and baseline_model and baseline_id - else None - ), - ) - PythonCore._stamp_or_clear_metadata_key( - request_kwargs=request_kwargs, - key="routing_decision", - value=( - None - if routing_decision is None - else PythonCore._redact_prompt_text_if_needed( - request_kwargs=request_kwargs, routing_decision=routing_decision - ) - ), - ) - - @staticmethod - def _stamp_or_clear_metadata_key(request_kwargs: dict, key: str, value: object | None) -> None: - """Write a proxy-internal metadata key for THIS routing attempt, or clear it. - - Fallbacks and retries re-enter the pre-routing hook with the same - `request_kwargs`, so every attempt must write or clear, never just write; - a value left behind by an earlier attempt would be attributed to this one. - `get_or_create_metadata_bucket` is the single owner of "which dict holds - proxy-internal metadata": it picks `litellm_metadata` when present (so the - value never lands in the `metadata` dict that routes like /v1/messages - forward to the provider) and replaces a non-dict value rather than silently - skipping the write. Clearing pops from BOTH buckets so a request whose - bucket resolution changed between attempts cannot resurrect a stale value. - """ - if value is None: - for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")): - if isinstance(bucket, dict): - bucket.pop(key, None) - return - - _, metadata_bucket = get_or_create_metadata_bucket(request_kwargs) - metadata_bucket[key] = value - - @staticmethod - def _redact_prompt_text_if_needed( - request_kwargs: Mapping[str, object], - routing_decision: StandardLoggingRoutingDecision, - ) -> StandardLoggingRoutingDecision: - """Drop verbatim prompt text from the record when message logging is redacted. - - An operator who turns message logging off has said prompt content must not reach - the logs, so the fields that quote the prompt (the matched keywords, and the - signals that name them) are omitted. Derived values are kept, because a tier, a - cause, a score or an escalation flag aggregates the prompt rather than - reproducing any of it, and dropping them would leave the row unexplainable for - no privacy gain. Applied here rather than in each strategy so a strategy added - later cannot bypass it. - """ - from litellm.litellm_core_utils.redact_messages import ( - should_redact_message_logging, - ) - - if not should_redact_message_logging( - { - "litellm_params": request_kwargs, - "standard_callback_dynamic_params": request_kwargs.get("standard_callback_dynamic_params"), - } - ): - return routing_decision - kept: Final = { - field: value - for field, value in routing_decision.items() - if field not in PROMPT_QUOTING_ROUTING_DECISION_FIELDS - } - return cast(StandardLoggingRoutingDecision, kept) # cast-ok: dropping optional keys preserves the type - - def _get_available_deployment_unresolved( # noqa: ANN202 # legacy selector result remains inferred - self, - model: str, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - request_kwargs: dict | None = None, - ): - """ - Returns the deployment based on routing strategy - """ - if self._core_routing_plugins: - raise ValueError( - "Router(plugins=[...]) is configured, but this call resolved to the synchronous " - "deployment-selection path, which never runs the routing-plugin pipeline. This " - "happens for sync Router methods (e.g. Router.completion()) and for async calls " - "with a routing_strategy that has no async-native selector (e.g. legacy " - "'usage-based-routing', v1). Silently skipping " - "configured plugins would let a policy plugin (e.g. a deny-all rule) be bypassed. " - "Use an async Router method with a supported routing_strategy (simple-shuffle, " - "usage-based-routing-v2, cost-based-routing, latency-based-routing, least-busy), " - "or remove `plugins` from the Router config." - ) - # users need to explicitly call a specific deployment, by setting `specific_deployment = True` as completion()/embedding() kwarg - # When this was no explicit we had several issues with fallbacks timing out - - model, healthy_deployments = self._common_checks_available_deployment( - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) - healthy_deployments = ( # pyright: ignore[reportUnknownVariableType] # rebind-ok: route filter narrows - self._without_deployments_that_cannot_serve_route( - healthy_deployments, # pyright: ignore[reportUnknownArgumentType] # legacy selector type - model, - ) - ) - # As the async pick: a team-owned deployment serves only its own team's callers, also when - # asked for by id (then refused under its group's name). - by_id: Final = healthy_deployments if isinstance(healthy_deployments, dict) else None # pyright: ignore[reportUnknownVariableType] # the legacy deployment dict - healthy_deployments = filter_team_based_models( # rebind-ok: the team filter narrows # pyright: ignore[reportUnknownVariableType] # legacy selector type - healthy_deployments=healthy_deployments, # pyright: ignore[reportArgumentType] # legacy selector type - request_kwargs=request_kwargs, - ) - if by_id is not None and not healthy_deployments: - group_name: Final = str(by_id.get("model_name")) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # see above - model = group_name # rebind-ok: the error names the group - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) - routing_kwargs: Final = _ROUTING_KWARGS_ADAPTER.validate_python(request_kwargs) - pick_attributes: Final = _deployment_pick_attributes(model, routing_kwargs) - - if isinstance(healthy_deployments, dict): - if (healthy_deployments.get("model_info") or {}).get("blocked") is True: - raise litellm_exceptions.ServiceUnavailableError( - message=f"Model '{model}' is currently paused and cannot accept requests.", - model=model, - llm_provider="", - ) - self._override_selector_pre_call_check(strategy, strategy_selector, healthy_deployments) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return healthy_deployments - - parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(request_kwargs) - request_deployments: Final = cast( # cast-ok: router deployments are plain dicts - Sequence[DeploymentTypedDict], healthy_deployments - ) - - # Health-check-based filtering (before cooldown) - healthy_deployments = self._filter_health_check_unhealthy_deployments( - healthy_deployments=healthy_deployments, - parent_otel_span=parent_otel_span, - ) - - cooldown_deployments: Final = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - _pre_cooldown_deployments: Final = healthy_deployments - healthy_deployments = self._filter_cooldown_deployments( - healthy_deployments=healthy_deployments, - cooldown_deployments=cooldown_deployments, - ) - if ( - not healthy_deployments - and self._core_enable_health_check_routing - and self._core_allowed_fails_policy is not None - ): - verbose_router_logger.warning( - "All deployments in cooldown via health-check routing, bypassing cooldown filter" - ) - healthy_deployments = _pre_cooldown_deployments - - healthy_deployments = continuation_deployments( - self._filter_blocked_deployments(healthy_deployments), routing_kwargs - ) - - # filter pre-call checks - if self._core_enable_pre_call_checks and (messages is not None or input is not None): - healthy_deployments = self._pre_call_checks( - model=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - request_deployments=request_deployments, - ) - - ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) - _target_order: Final = (request_kwargs or {}).pop("_target_order", None) - batch_candidates: Final = self._without_decisions_deployments_refusing_media( - self._without_declared_text_only_deployments( - cast(list[dict[str, object]], healthy_deployments), # cast-ok: router deployments are plain dicts - model, - messages, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict - ), - model, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: legacy signature uses a bare dict - ) - embedding_batch_admitting: Final = deployments_admitting_embedding_batch( - batch_candidates, - _OPAQUE_VALUE_ADAPTER.validate_python(input), - routing_kwargs, - ) - embedding_batch_routed: Final = ( - embedding_batch_admitting - if embedding_batch_admitting is not None - else deployments_when_none_admits_embedding_batch( - batch_candidates, - cast(list[dict[str, object]], _pre_cooldown_deployments), # cast-ok: router deployments are plain dicts - self._deployment_ids( - self._region_routable_deployments( - request_deployments, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: str-keyed request kwargs - True, - ) - ) - & frozenset(cooldown_deployments), - _OPAQUE_VALUE_ADAPTER.validate_python(input), - routing_kwargs, - ) - ) - healthy_deployments = litellm_utils.get_order_filtered_deployments( - list( - self._narrow_to_video_deployment( - without_deployments_that_cannot_batch( - embedding_batch_routed, - caller_deployments=lambda: deployments_the_caller_may_use(batch_candidates, routing_kwargs), - ) - ) - ), - target_order=_target_order, - ) - - ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in - ## this request via weighted-failover. See async counterpart in - ## async_get_healthy_deployments for details. - _excluded_deployment_ids: Final = (request_kwargs or {}).pop("_excluded_deployment_ids", None) - healthy_deployments = litellm_utils.get_excluded_filtered_deployments( - healthy_deployments, - excluded_deployment_ids=_excluded_deployment_ids, - ) - - ## RETRY SKIP ## -> see async counterpart in async_get_healthy_deployments. - _retry_skipped_deployment_ids: Final = as_retry_skipped_deployment_ids( - request_kwargs.pop("_retry_skipped_deployment_ids", None) if request_kwargs else None - ) - healthy_deployments = ( - litellm_utils.get_excluded_filtered_deployments( - healthy_deployments, excluded_deployment_ids=_retry_skipped_deployment_ids - ) - or healthy_deployments - ) - - if len(healthy_deployments) == 0: - model_ids = self._core_get_model_ids(model_name=model) - _cooldown_time = self._core_cooldown_cache.get_min_cooldown( - model_ids=model_ids, parent_otel_span=parent_otel_span - ) - _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - raise RouterRateLimitError( - model=model, - cooldown_time=_cooldown_time, - enable_pre_call_checks=self._core_enable_pre_call_checks, - cooldown_list=_cooldown_list, - model_ids=model_ids, - cooldown_deployments=get_cooldown_deployment_info( - self._core_ctx, - self._deployment_ids( - self._region_routable_deployments( - request_deployments, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: str-keyed request kwargs - messages is not None or input is not None, - ) - ), - _cooldown_list, - ), - ) - - if strategy == "simple-shuffle": - # if users pass rpm or tpm, we do a random weighted pick - based on rpm/tpm - ############## Check 'weight' param set for weighted pick ################# - shuffled: Final = simple_shuffle( - resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=healthy_deployments, - model=model, - request_kwargs=request_kwargs, - ) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return shuffled - deployment: Final = self._select_deployment_sync( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=healthy_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), - ) - - if deployment is None: - verbose_router_logger.info("get_available_deployment for model: %s, No deployment available", model) - model_ids = self._core_get_model_ids(model_name=model) - _cooldown_time = self._core_cooldown_cache.get_min_cooldown( - model_ids=model_ids, parent_otel_span=parent_otel_span - ) - _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - raise RouterRateLimitError( - model=model, - cooldown_time=_cooldown_time, - enable_pre_call_checks=self._core_enable_pre_call_checks, - cooldown_list=_cooldown_list, - model_ids=model_ids, - ) - self._override_selector_pre_call_check(strategy, strategy_selector, deployment) - verbose_router_logger.info( - "get_available_deployment for model: %s, Selected deployment: %s for model: %s", - model, - self.print_deployment(deployment), - model, - ) - phase_event(_DEPLOYMENT_SELECTED_EVENT, pick_attributes) - return deployment - - get_available_deployment = ( # pyright: ignore[reportUnknownVariableType] # legacy selector return - resolve_pods_after( - _get_available_deployment_unresolved # pyright: ignore[reportUnknownArgumentType] # legacy annotations - ) - ) - - def _core_get_available_deployment_for_pass_through( # noqa: ANN202 # unannotated before the rename too; its inferred return mixes untyped deployment dicts - self, - model: str, - messages: list[dict[str, str]] | None = None, - input: str | list | None = None, - specific_deployment: bool | None = False, - request_kwargs: dict | None = None, - ): - """ - Returns deployments available for pass-through endpoints (based on load balancing strategy) - - Similar to get_available_deployment, but only returns deployments with use_in_pass_through=True - - Args: - model: Model name - messages: Optional list of messages - input: Optional input data - specific_deployment: Whether to find a specific deployment - request_kwargs: Optional request parameters - - Returns: - Dict: Selected deployment configuration - - Raises: - BadRequestError: If no deployment is configured with use_in_pass_through=True - RouterRateLimitError: If no pass-through deployments are available - """ - # 1. Perform common checks to get healthy deployments list - model, healthy_deployments = self._common_checks_available_deployment( - model=model, - messages=messages, - input=input, - specific_deployment=specific_deployment, - request_kwargs=request_kwargs, - ) - - strategy, strategy_selector = self._get_routing_context(model, request_kwargs) - - # 2. If the returned is a specific deployment (Dict), verify and return directly - if isinstance(healthy_deployments, dict): - if (healthy_deployments.get("model_info") or {}).get("blocked") is True: - raise litellm_exceptions.ServiceUnavailableError( - message=f"Model '{model}' is currently paused and cannot accept requests.", - model=model, - llm_provider="", - ) - litellm_params: Final = healthy_deployments.get("litellm_params", {}) - if litellm_params.get("use_in_pass_through"): - self._override_selector_pre_call_check(strategy, strategy_selector, healthy_deployments) - return healthy_deployments - else: - # Specific deployment does not support pass-through - raise litellm_exceptions.BadRequestError( - message=f"Deployment {healthy_deployments.get('model_info', {}).get('id')} does not support pass-through endpoint (use_in_pass_through=False)", - model=model, - llm_provider="", - ) - - # 3. Filter deployments that support pass-through - pass_through_deployments = self._filter_pass_through_deployments(healthy_deployments=healthy_deployments) - - if len(pass_through_deployments) == 0: - # No deployments support pass-through - raise litellm_exceptions.BadRequestError( - message=f"Model {model} has no deployment configured with use_in_pass_through=True. Please add use_in_pass_through: true in the deployment configuration", - model=model, - llm_provider="", - ) - pass_through_model_ids: Final = tuple( - deployment["model_info"]["id"] - for deployment in pass_through_deployments - if "id" in deployment.get("model_info", {}) - ) - - # 4. Apply health-check and cooldown filtering - parent_otel_span: Final[Span | None] = get_parent_otel_span_from_kwargs(request_kwargs) - request_deployments: Final = cast( # cast-ok: router deployments are plain dicts - Sequence[DeploymentTypedDict], pass_through_deployments - ) - pass_through_deployments = self._filter_health_check_unhealthy_deployments( - healthy_deployments=pass_through_deployments, - parent_otel_span=parent_otel_span, - ) - cooldown_deployments: Final = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - pass_through_deployments = self._filter_cooldown_deployments( - healthy_deployments=pass_through_deployments, - cooldown_deployments=cooldown_deployments, - ) - pass_through_deployments = self._filter_blocked_deployments(pass_through_deployments) - - # 5. Apply pre-call checks (if enabled) - if self._core_enable_pre_call_checks and (messages is not None or input is not None): - pass_through_deployments = self._pre_call_checks( - model=model, - healthy_deployments=pass_through_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - request_deployments=request_deployments, - ) - - if self._is_priority_routing_group(model): - pass_through_deployments = litellm_utils.get_order_filtered_deployments( - pass_through_deployments, - target_order=request_kwargs.pop("_target_order", None) if request_kwargs is not None else None, - ) - - if len(pass_through_deployments) == 0: - model_ids = self._core_get_model_ids(model_name=model) - _cooldown_time = self._core_cooldown_cache.get_min_cooldown( - model_ids=model_ids, parent_otel_span=parent_otel_span - ) - _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - raise RouterRateLimitError( - model=model, - cooldown_time=_cooldown_time, - enable_pre_call_checks=self._core_enable_pre_call_checks, - cooldown_list=_cooldown_list, - model_ids=pass_through_model_ids, - cooldown_deployments=get_cooldown_deployment_info( - self._core_ctx, - self._deployment_ids( - self._region_routable_deployments( - request_deployments, - cast(Mapping[str, object] | None, request_kwargs), # cast-ok: str-keyed request kwargs - messages is not None or input is not None, - ) - ), - _cooldown_list, - ), - ) - - # 6. Apply load balancing strategy - if strategy == "simple-shuffle": - return simple_shuffle( - resolve_model_alias=self._core_get_model_from_alias, - healthy_deployments=pass_through_deployments, - model=model, - request_kwargs=request_kwargs, - ) - deployment: Final = self._select_deployment_sync( - strategy=strategy, - selector=strategy_selector, - model=model, - healthy_deployments=pass_through_deployments, - messages=messages, - input=input, - request_kwargs=request_kwargs, - reserve=pick_reserves(default=False) and request_kwargs not in (None, {}), - ) - - if deployment is None: - verbose_router_logger.info( - "get_available_deployment_for_pass_through model: %s, no available deployments", model - ) - model_ids = self._core_get_model_ids(model_name=model) - _cooldown_time = self._core_cooldown_cache.get_min_cooldown( - model_ids=model_ids, parent_otel_span=parent_otel_span - ) - _cooldown_list = get_cooldown_deployments(ctx=self._core_ctx, parent_otel_span=parent_otel_span) - raise RouterRateLimitError( - model=model, - cooldown_time=_cooldown_time, - enable_pre_call_checks=self._core_enable_pre_call_checks, - cooldown_list=_cooldown_list, - model_ids=model_ids, - ) - self._override_selector_pre_call_check(strategy, strategy_selector, deployment) - - verbose_router_logger.info( - "get_available_deployment_for_pass_through model: %s, selected deployment: %s", - model, - self.print_deployment(deployment), - ) - return deployment - - def _filter_cooldown_deployments( - self, healthy_deployments: list[dict], cooldown_deployments: list[str] - ) -> list[dict]: - """ - Filters out the deployments currently cooling down from the list of healthy deployments - - Args: - healthy_deployments: List of healthy deployments - cooldown_deployments: List of model_ids cooling down. cooldown_deployments is a list of model_id's cooling down, cooldown_deployments = ["16700539-b3cd-42f4-b426-6a12a1bb706a", "16700539-b3cd-42f4-b426-7899"] - - Returns: - List of healthy deployments - """ - if verbose_router_logger.isEnabledFor(logging.DEBUG): - verbose_router_logger.debug("cooldown deployments: %s", cooldown_deployments) - # Convert to set for O(1) lookup and use list comprehension for O(n) filtering - cooldown_set: Final = set(cooldown_deployments) - return [deployment for deployment in healthy_deployments if deployment["model_info"]["id"] not in cooldown_set] - - def _filter_blocked_deployments(self, healthy_deployments: list[dict]) -> list[dict]: - """ - Filters out deployments that an admin has paused via `LiteLLM_ProxyModelTable.blocked`. - - Applied alongside the cooldown filter on every routing entry point that calls - `_common_checks_available_deployment` directly — the primary sync/async path, - the sync pass-through path, and the retry / health-check helpers — so paused - deployments never serve a request. The async pass-through path inherits this - filter through its delegation to `async_get_healthy_deployments`. - """ - return [ - deployment - for deployment in healthy_deployments - if (deployment.get("model_info") or {}).get("blocked") is not True - ] - - @staticmethod - def _core_is_deployment_blocked(deployment: "Deployment") -> bool: - """ - Returns True when a `Deployment` Pydantic instance carries the admin-paused - flag. Used by credential-lookup helpers so passthrough file / batch endpoints - cannot bypass the pause by resolving credentials directly. - """ - model_info: Final[object | None] = getattr(deployment, "model_info", None) - if model_info is None: - return False - return getattr(model_info, "blocked", None) is True - - async def _async_filter_health_check_unhealthy_deployments( - self, - healthy_deployments: list[dict], - parent_otel_span: Span | None = None, - health_check_probe: bool = False, - ) -> list[dict]: - """ - Filter out deployments marked unhealthy by background health checks. - No-op when enable_health_check_routing is False. - When background_health_check_model_groups is set, only deployments in the - listed model groups are filtered; every other group keeps its configured - routing strategy untouched, and a router-level allowed_fails_policy no - longer disables the filter for the listed groups. - Returns all deployments if health state is unavailable, stale, or would - exclude every candidate (safety net). - """ - if not self._core_enable_health_check_routing: - return healthy_deployments - - # When allowed_fails_policy is set, cooldown is the sole routing exclusion - # mechanism -- skip the binary health check filter so the policy threshold - # is respected before any deployment is excluded. With a model-group - # allowlist the filter is already scoped, so listed groups keep it. - scoped_groups: Final = self._core_background_health_check_model_groups - if self._core_allowed_fails_policy is not None and scoped_groups is None: - return healthy_deployments - - unhealthy_ids: Final = await self._core_health_state_cache.async_get_unhealthy_deployment_ids( - parent_otel_span=parent_otel_span - ) - if not unhealthy_ids: - return healthy_deployments - - filtered: Final = [ - d - for d in healthy_deployments - if d["model_info"]["id"] not in unhealthy_ids - or (scoped_groups is not None and d["model_name"] not in scoped_groups) - ] - - if not filtered: - return [] if health_check_probe else healthy_deployments - - return filtered - - def _filter_health_check_unhealthy_deployments( - self, - healthy_deployments: list[dict], - parent_otel_span: Span | None = None, - ) -> list[dict]: - """Sync version of _async_filter_health_check_unhealthy_deployments.""" - if not self._core_enable_health_check_routing: - return healthy_deployments - - scoped_groups: Final = self._core_background_health_check_model_groups - if self._core_allowed_fails_policy is not None and scoped_groups is None: - return healthy_deployments - - unhealthy_ids: Final = self._core_health_state_cache.get_unhealthy_deployment_ids( - parent_otel_span=parent_otel_span - ) - if not unhealthy_ids: - return healthy_deployments - - filtered: Final = [ - d - for d in healthy_deployments - if d["model_info"]["id"] not in unhealthy_ids - or (scoped_groups is not None and d["model_name"] not in scoped_groups) - ] - - if not filtered: - verbose_router_logger.warning("All deployments marked unhealthy by health checks, bypassing health filter") - return healthy_deployments - - return filtered - - def _filter_pass_through_deployments(self, healthy_deployments: list[dict]) -> list[dict]: - """ - Filter out deployments configured with use_in_pass_through=True - - Args: - healthy_deployments: List of healthy deployments - - Returns: - List[Dict]: Only includes a list of deployments that support pass-through - """ - verbose_router_logger.debug( - "Filter pass-through deployments from %s healthy deployments", len(healthy_deployments) - ) - - pass_through_deployments: Final = [ - deployment - for deployment in healthy_deployments - if deployment.get("litellm_params", {}).get("use_in_pass_through", False) - ] - - verbose_router_logger.debug("Found %s deployments with pass-through enabled", len(pass_through_deployments)) - - return pass_through_deployments - - def _track_deployment_metrics(self, deployment, parent_otel_span: Span | None, response=None): - """ - Tracks successful requests rpm usage. - """ - try: - model_id: Final = deployment.get("model_info", {}).get("id", None) - if response is None: - # update self.deployment_stats - if model_id is not None: - self._update_usage(model_id, parent_otel_span) # update in-memory cache for tracking - except Exception as e: - verbose_router_logger.error("Error in _track_deployment_metrics: %s", e) - - def get_num_retries_from_retry_policy(self, exception: Exception, model_group: str | None = None): - return _get_num_retries_from_retry_policy( - exception=exception, - model_group=model_group, - model_group_retry_policy=self.model_group_retry_policy, - retry_policy=self.retry_policy, - ) - - def _core_get_allowed_fails_from_policy(self, exception: Exception) -> int | None: - """ - BadRequestErrorRetries: Optional[int] = None - AuthenticationErrorRetries: Optional[int] = None - TimeoutErrorRetries: Optional[int] = None - RateLimitErrorRetries: Optional[int] = None - ContentPolicyViolationErrorRetries: Optional[int] = None - """ - # if we can find the exception then in the retry policy -> return the number of retries - allowed_fails_policy: Final[AllowedFailsPolicy | None] = self._core_allowed_fails_policy - - if allowed_fails_policy is None: - return None - - if ( - isinstance(exception, litellm_exceptions.AuthenticationError) - and allowed_fails_policy.AuthenticationErrorAllowedFails is not None - ): - return allowed_fails_policy.AuthenticationErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.Timeout) - and allowed_fails_policy.TimeoutErrorAllowedFails is not None - ): - return allowed_fails_policy.TimeoutErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.RateLimitError) - and allowed_fails_policy.RateLimitErrorAllowedFails is not None - ): - return allowed_fails_policy.RateLimitErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.ContentPolicyViolationError) - and allowed_fails_policy.ContentPolicyViolationErrorAllowedFails is not None - ): - return allowed_fails_policy.ContentPolicyViolationErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.BadRequestError) - and allowed_fails_policy.BadRequestErrorAllowedFails is not None - ): - return allowed_fails_policy.BadRequestErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.InternalServerError) - and allowed_fails_policy.InternalServerErrorAllowedFails is not None - ): - return allowed_fails_policy.InternalServerErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.ServiceUnavailableError) - and allowed_fails_policy.ServiceUnavailableErrorAllowedFails is not None - ): - return allowed_fails_policy.ServiceUnavailableErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.BadGatewayError) - and allowed_fails_policy.BadGatewayErrorAllowedFails is not None - ): - return allowed_fails_policy.BadGatewayErrorAllowedFails - if ( - isinstance(exception, litellm_exceptions.NotFoundError) - and allowed_fails_policy.NotFoundErrorAllowedFails is not None - ): - return allowed_fails_policy.NotFoundErrorAllowedFails - - def _initialize_alerting(self): - from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting - - if self.alerting_config is None: - return - - router_alerting_config: Final[AlertingConfig] = self.alerting_config - - _slack_alerting_logger: Final = SlackAlerting( - alerting_threshold=router_alerting_config.alerting_threshold, - alerting=["slack"], - default_webhook_url=router_alerting_config.webhook_url, - ) - - self._core_slack_alerting_logger = _slack_alerting_logger - - self._host.callbacks.add_litellm_callback(_slack_alerting_logger) - self._host.callbacks.add_litellm_success_callback(_slack_alerting_logger.response_taking_too_long_callback) - verbose_router_logger.info("\033[94m\nInitialized Alerting for litellm.Router\033[0m\n") - - def _core_set_custom_routing_strategy(self, CustomRoutingStrategy: CustomRoutingStrategyBase) -> None: - """ - Sets get_available_deployment and async_get_available_deployment on an instanced of litellm.Router - - Use this to set your custom routing strategy - - Args: - CustomRoutingStrategy: litellm.router.CustomRoutingStrategyBase - """ - - strategy: Final = CustomRoutingStrategy - setattr( - self, - "get_available_deployment", - resolve_pods_after_bound( - strategy.get_available_deployment, # pyright: ignore[reportUnknownArgumentType] # legacy stub - self._core_kubernetes_pod_discovery, - ), - ) - setattr( - self, - "_core_async_get_available_deployment", - async_resolve_pods_after_bound( - strategy.async_get_available_deployment, # pyright: ignore[reportUnknownArgumentType] # legacy stub - self._core_kubernetes_pod_discovery, - ), - ) - - def _reset_custom_routing_strategy(self) -> None: - for attr in ("get_available_deployment", "_core_async_get_available_deployment"): - if attr in self.__dict__: - delattr(self, attr) - - def flush_cache(self): - self._host.set_global_cache(None) - self._core_cache.flush_cache() - session_in_memory_cache: Final = self._claude_code_session_router_cache.in_memory_cache - if session_in_memory_cache is not None: - session_in_memory_cache.flush_cache() - self._core_prompt_caching_pin_cache.in_memory_cache.flush_cache() - - def reset(self): - ## clean up on close - self._host.set_callback_list("success_callback", []) - self._host.set_callback_list("_async_success_callback", []) - self._host.set_callback_list("failure_callback", []) - self._host.set_callback_list("_async_failure_callback", []) - self.retry_policy = None - self.flush_cache() - - -def _name_core_members_after_old_names() -> None: - # An old member that moved to ``_core_`` can still be reached as the function itself - # (``cached_deployment_model_info`` is an lru_cache over one), so errors that name the method - # keep the name it had before. - for name, member in vars(PythonCore).items(): - old_name = name.removeprefix("_core_") - if old_name != name and inspect.isfunction(member) and hasattr(DeprecatedRouterMembers, old_name): - member.__qualname__ = f"Router.{old_name}" - - -_name_core_members_after_old_names() diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 395b580d9b2..c6ca4a1f399 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -11,11 +11,9 @@ itself only through its host, and the host never calls back into a Python router Python router would have running for the request reach (``OptionalChecks.outcome_callbacks``) - sleeps, fallback access and budget checks, and cache reads and writes are host steps -``RustCore`` does not build, wrap or read a ``PythonCore``. Settings are resolved from the -``RouterConfig`` and the host's process settings; the views (``registry()``, ``settings()``, -``health()``) are read from the Rust handle. A member the Rust core does not serve raises -``RustCoreUnsupported`` (an ``AttributeError``) instead of reaching a Python core, so deleting -``router_python_core.py`` cannot change what this core does. +Settings are resolved from the ``RouterConfig`` and the host's process settings; the views +(``registry()``, ``settings()``, ``health()``) are read from the Rust handle. A member the Rust +core does not serve raises ``RustCoreUnsupported`` (an ``AttributeError``). """ import contextlib @@ -23,13 +21,14 @@ import dataclasses import inspect import itertools import json -from collections.abc import Awaitable, Callable, Coroutine, Iterable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Coroutine, Generator, Iterable, Mapping, Sequence +from contextvars import ContextVar from datetime import date, datetime, time from functools import partial from time import perf_counter from time import time as unix_time from types import MappingProxyType -from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, final +from typing import TYPE_CHECKING, Final, Literal, Protocol, TypeAlias, TypeVar, final, overload from litellm._internal_context import service_target from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS, HEALTH_CHECKS_TARGET @@ -46,7 +45,12 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import ( ) from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type -from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping, without +from litellm.litellm_core_utils.core_helpers import ( + as_object_list, + as_str_mapping, + get_parent_otel_span_from_kwargs, + without, +) from litellm.litellm_core_utils.routing_common_utils import ( format_content_policy_fallback_hint, format_fallback_outcome_message, @@ -79,6 +83,7 @@ from litellm.router_rust.invoke.endpoints import ( attempt_call, clean_realtime_session_deployment, compacted_attempt, + endpoint_names, failover_fields, generic_plan, log_refused_call, @@ -208,9 +213,11 @@ from litellm.router_utils.kubernetes_pod_discovery import ( resolve_pods_after_bound, ) from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_io_token_refund_failure from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata from litellm.router_utils.routing_read_batch import ROUTER_USAGE_TARGET, RoutingReadBatch from litellm.types.router import ( + AllowedFailsPolicy, Deployment, RouterErrors, RouterGeneralSettings, @@ -218,7 +225,7 @@ from litellm.types.router import ( RoutingGroup, RoutingStrategy, ) -from litellm.types.utils import LlmProviders +from litellm.types.utils import CooldownDeploymentInfo, LlmProviders if TYPE_CHECKING: from opentelemetry.trace import Span @@ -444,6 +451,23 @@ _NUMBER_SETTINGS: Final = frozenset( _UPDATE_INT_SETTINGS: Final = frozenset({"timeout", "num_retries", "retry_after", "allowed_fails", "cooldown_time"}) +def _raised_before_provider(error: BaseException) -> bool: + """A request litellm refused before it reached the provider (a parameter the provider does not + take, a malformed priority): neither a failure nor a request of the deployment it picked.""" + from litellm.exceptions import UnsupportedParamsError + + body: Final = getattr(error, "body", None) + return isinstance(error, UnsupportedParamsError) or (isinstance(body, Mapping) and body.get("param") == "priority") + + +def _shares_logging(request: Mapping[str, object]) -> bool: + """The caller passed the logging object every attempt of the request logs on, which logs one + success a request (``should_run_logging``).""" + from litellm.litellm_core_utils.litellm_logging import Logging + + return isinstance(request.get("litellm_logging_obj"), Logging) + + class RustCoreUnsupported(AttributeError): """A member or request the Rust core does not serve. It is never forwarded to a Python core.""" @@ -541,6 +565,14 @@ def _strategy_name(value: object) -> str | None: return None if value is None else normalize_strategy(value) or "simple-shuffle" +_U32_MAX: Final = 2**32 - 1 + + +def _crate_count(value: int | None) -> int | None: + """A count setting as the crate's u32 reads it: a negative one is 0, a huge one the most it holds.""" + return None if value is None else min(max(value, 0), _U32_MAX) + + def _lenient_int(value: object, default: int) -> int: """A count setting kept as given, as the crate reads it: an int, or a number spelled another way (``1.5``, ``"3"``) truncated to one; ``default`` when it is no number. The Python core stores the @@ -763,9 +795,9 @@ class Resolved: if name in {"model_group_alias", "default_litellm_params"}: updates[name] = MappingProxyType(dict(as_str_mapping(value) or {})) elif name in _NUMBER_SETTINGS: + # `allowed_fails` saved later is no explicit one: only the constructor's argument is + # (`explicit_allowed_fails`), so a value equal to the default keeps the rate-based default. updates[name] = self._number_setting(name, value) - if name == "allowed_fails": - updates["explicit_allowed_fails"] = updates[name] elif name in _SELECTION_SETTINGS: updates[name] = self._selection_setting(name, value) elif name == "disable_cooldowns": @@ -823,7 +855,9 @@ class Resolved: "routing_strategy": self.routing_strategy, "cooldown_time": self.cooldown_time, "disable_cooldowns": self.disable_cooldowns, - "allowed_fails_policy": _plain(self.allowed_fails_policy), + # A policy that is no policy or mapping (an attribute write of a list) is none, as the + # Python core's `get_allowed_fails_from_policy` reads no field off it. + "allowed_fails_policy": _plain(self.allowed_fails_policy) if self._has_fails_policy() else None, "retry_after": self.retry_after, "timeout": self.timeout, "routing_groups": group_documents(self.routing_groups) if self.routing_groups else None, @@ -842,6 +876,9 @@ class Resolved: # An explicit `routing_strategy=None` reaches the Rust core as "None": no strategy picks. return kept if self.routing_strategy is not None else {**kept, "routing_strategy": "None"} + def _has_fails_policy(self) -> bool: + return isinstance(self.allowed_fails_policy, AllowedFailsPolicy | Mapping) + def dynamic_document(self, version: int, model_names: Sequence[str]) -> Mapping[str, object]: values: Final = { "version": version, @@ -849,9 +886,11 @@ class Resolved: "fallbacks": _plain(self.fallbacks), "context_window_fallbacks": _plain(self.context_window_fallbacks), "content_policy_fallbacks": _plain(self.content_policy_fallbacks), - "num_retries": self.num_retries, - "max_fallbacks": self.max_fallbacks, - "allowed_fails": self.allowed_fails, + # A negative count reads as none at all in the Python core (`range(num_retries)`, + # `fallback_depth >= max_fallbacks`, `fails > allowed_fails`); the crate reads a u32. + "num_retries": _crate_count(self.num_retries), + "max_fallbacks": _crate_count(self.max_fallbacks), + "allowed_fails": _crate_count(self.allowed_fails), "allowed_fails_is_default": True if self.allowed_fails_is_default() else None, "fallback_keys": self._fallback_keys(model_names), "default_deployment": _plain(served_default_deployment(self.default_deployment)[0]), @@ -878,7 +917,9 @@ class Resolved: for entry in _as_list(entries) if (mapping := as_str_mapping(entry)) ) - names: Final = {*model_names, *self.model_group_alias} + # A bare name a wildcard serves is in no list of names: the keys' own bare spellings stand for it. + bare: Final = (key.split("/", 1)[1] for key in keys if "/" in key) + names: Final = {*model_names, *self.model_group_alias, *bare} prefixed: Final = { name: spelled for name in names if (spelled := _provider_prefixed_model_group(name, keys)) is not None } @@ -1094,6 +1135,14 @@ class RustHealth(RouterHealth): core._resolved.allowed_fails_policy, # pyright: ignore[reportPrivateUsage] # see store ) + @overload + async def usage( + self, model_group: str, *, kind: Literal["current"] = "current" + ) -> tuple[int | None, int | None]: ... + + @overload + async def usage(self, model_group: str, *, kind: Literal["remaining"]) -> Mapping[str, int]: ... + async def usage( # pyright: ignore[reportIncompatibleMethodOverride] # the same two kinds self, model_group: str, *, kind: Literal["current", "remaining"] = "current" ) -> tuple[int | None, int | None] | Mapping[str, int]: @@ -1167,6 +1216,21 @@ _EndpointImpl: TypeAlias = Callable[..., Awaitable[object]] _StrategyInputs: TypeAlias = tuple[str, object, str, list[dict[str, object]]] +def _pick_answer( + step: Mapping[str, object], candidates: Sequence[object], picked: object +) -> list[object]: # mutable-ok: a JSON list + """A strategy step's answer: the picked id, or its position when the candidates share ids (the + step then carries ``records``), so the pick serves the entry it chose as Python's list does.""" + found: Final = picked_id(picked) + if step.get("records") is None or found is None: + return [found] if found is not None else [] + # The entry the selector returned, else (a selector that returned a copy) the first with its id. + position: Final = next((index for index, record in enumerate(candidates) if record is picked), None) + if position is not None: + return [position] + return [next((index for index, record in enumerate(candidates) if picked_id(record) == found), 0)] + + @final class _PickCounts(Protocol): """The router cache as ``track_deployment_pick`` uses it: each deployment id holds its pick count.""" @@ -1193,6 +1257,14 @@ class _EndpointPort: def router(self) -> "RustCore": return self._core + @property + def public(self) -> object: + """The ``Router`` that owns the core, which litellm functions taking ``router=`` get; the core + itself when built alone.""" + owner: Final[object] = self._core.__dict__.get("_core_owner") + public: Final[object] = owner() if callable(owner) else None + return self._core if public is None else public + @property def host(self) -> PythonHost: return self._core._host # pyright: ignore[reportPrivateUsage] # see the class docstring @@ -1343,6 +1415,7 @@ class RustCore( ) self._hosted._core_cache = self._cache # pyright: ignore[reportPrivateUsage] # usage reads come from the cache the Rust core writes ctx: Final = RoutingContext(self) # pyright: ignore[reportArgumentType] # the strategies read registry(), settings(), health() and select() + self._ctx: Final = ctx strategies: Final = PreRoutingStrategies( self._hosted, host, ctx, config.get("auto_router_capability_limit"), self._cache ) @@ -1373,6 +1446,7 @@ class RustCore( self._checks: Final = OptionalChecks(self, host, self._cache, self._hosted, ctx, config.arguments) self._hosted.checks = self._checks strategies.ensure_affinity = self._checks.ensure_deployment_affinity_callback + strategies.hand_over_affinity(self._checks.adopt) self._checks.build(config.get("optional_pre_call_checks"), config.get("model_list")) self._slack_alerting_logger: SlackAlerting | None = initialize_alerting(host, config.get("alerting_config")) # pyright: ignore[reportArgumentType] # the argument as passed if self._slack_alerting_logger is not None: @@ -1384,6 +1458,11 @@ class RustCore( apply_lar1_routing_strategy(ctx, resolved.routing_strategy_args) # pyright: ignore[reportArgumentType] # the core's RoutingContext + @property + def _core_ctx(self) -> RoutingContext: + """The router's ``RoutingContext``: what a routing helper gets where it once took the router.""" + return self._ctx + @property def _registry_state(self) -> HostRegistry: return self._hosted @@ -1417,6 +1496,9 @@ class RustCore( if enabled: register_continuation_prefill_check(self._host) + def __dir__(self) -> list[str]: # mutable-ok: dir() returns a list + return sorted({*super().__dir__(), *endpoint_names()}) + def __getattr__(self, name: str) -> object: if name in self.__dict__.get("_deleted", ()): raise self._no_attribute(name) @@ -1550,6 +1632,27 @@ class RustCore( str(e).replace("\r", "").replace("\n", ""), ) + @contextlib.contextmanager + def _usage_read_scope(self, model: str, request: Mapping[str, object]) -> Generator[None, None, None]: + """``RoutingReadBatch.for_strategy``: a usage-based-routing-v2 pick reads its tpm/rpm counters + in the cooldown read's MGET (see ``_batch_read``), and the strategy reads them from there.""" + from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 + + try: + selector: Final = self._prefetch_usage_selector(model, request) + deployments: Final = self._hosted._core_get_model_list(model_name=model) or [] # pyright: ignore[reportPrivateUsage] # see __init__ + except Exception: # noqa: BLE001 # the batched read is an optimisation: the reads stay apart + yield + return + if not isinstance(selector, LowestTPMLoggingHandler_v2) or not deployments: + yield + return + token: Final = _USAGE_READ.set((RoutingReadBatch(usage_selector=selector), list(deployments))) + try: + yield + finally: + _USAGE_READ.reset(token) + def _prefetch_usage_selector(self, model: str, request_kwargs: Mapping[str, object] | None) -> object: """``_get_routing_context``'s selector when it is a usage-based-routing-v2 one: the request's strategy override, else the routing group ``model`` resolves into, else the router's own.""" @@ -1913,6 +2016,7 @@ class RustCore( # `_core_update_redis_cache`: the routing state's cache gains the Redis cache. self._cache.attach_redis_cache(changes.redis_cache) self._checks.prompt_caching_pin_cache.attach_redis_cache(changes.redis_cache) + self._prerouter.attach_redis_cache(changes.redis_cache) return ApplyResult( version=self._hosted._registry_version, # pyright: ignore[reportPrivateUsage] # see __init__ upserted=upserted, @@ -1967,6 +2071,8 @@ class RustCore( self._settings = None # the snapshot holds settings fields, not the plain attributes if "model_group_alias" in values: self._hosted._core_model_group_alias = dict(resolved.model_group_alias) # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # see __init__ + # A group's info reads through the alias map. + self._hosted._invalidate_model_group_info_cache() # pyright: ignore[reportPrivateUsage] # see __init__ self._apply_strategy(previous, resolved, given, settled=settled) if "routing_groups" in values: self._settings = None @@ -2162,11 +2268,15 @@ class RustCore( def __delattr__(self, name: str) -> None: # `del router.` unsets it: the settings view reports it unset from then on. - from litellm.router_api import SETTING_FIELDS + from litellm.router_api import PLAIN_SETTING_ATTRIBUTES, SETTING_FIELDS # The deprecated attribute deletes the core's own name (``_core_``). setting: Final = name.removeprefix("_core_") - if setting not in SETTING_FIELDS and setting not in _DELETABLE_ATTRIBUTES: + if ( + setting not in SETTING_FIELDS + and setting not in PLAIN_SETTING_ATTRIBUTES + and setting not in _DELETABLE_ATTRIBUTES + ): if name in self.__dict__ or name in self._unset_attributes: object.__delattr__(self, name) return @@ -2220,6 +2330,10 @@ class RustCore( elif name in PLAIN_DEFAULTS or name == "max_retries": # Plain attributes on the Python core, outside the setting lists: stored as given. self._store_plain({name: value}) + elif name == "default_priority": + # Queued calls read it from the next call on, kept as given, as the Python core's attribute is. + self._endpoints.assign_setting(name, value) + self._deleted = self._deleted - {name} elif name == "scheduler": self._endpoints.set_scheduler(value) self._deleted = self._deleted - {name} @@ -2392,7 +2506,7 @@ class RustCore( return await self._endpoints.generic(plan_for("aanthropic_messages"), kwargs) async def anthropic_messages(self, **kwargs: object) -> object: # kwargs-ok: endpoint payload - """The proxy's ``/v1/messages`` route name for ``aanthropic_messages``; async, as in PythonCore.""" + """The proxy's ``/v1/messages`` route name for ``aanthropic_messages``; async.""" return await self._endpoints.generic(plan_for("anthropic_messages"), kwargs) async def aimage_generation(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload @@ -2446,18 +2560,24 @@ class RustCore( self._refuse_deleted_reads({**request, **(overrides or {})}) self._sync_routing_state() model: Final = request.get("model") - if model is None and "model" in request: - raise self._no_model_error(plan, request) if "model" not in request: raise TypeError(f"{plan.endpoint}() needs a model name") - name: Final = model if isinstance(model, str) else self.routed_name(model) + name: Final = ( + model if isinstance(model, str) else self._no_model_name() if model is None else self.routed_name(model) + ) field: Final = plan.prepare_field or "metadata" if request.get(field) is not None and not isinstance(request.get(field), Mapping): raise TypeError(f"{field} must be a mapping") prepared: Final = dict(request) if prepared_already else self._endpoints.prepare(plan, request, model) # pyright: ignore[reportArgumentType] # the name as given, which the metadata keeps - return self._handle.start( - plan.endpoint, self._routing_inputs(plan, request, prepared, name, overrides) - ), prepared + if not prepared_already: + # `async_function_with_fallbacks`: a selection the caller sent is not one the router made, + # so it never exempts a fallback target from the access gate. + from litellm.router_utils.fallback_event_handlers import clear_pre_routing_selection + + clear_pre_routing_selection(prepared) + # `async_function_with_fallbacks` raises a pick for no model group as it is, with no fallback. + routed: Final = {**(overrides or {}), "disable_fallbacks": True} if model is None else overrides + return self._handle.start(plan.endpoint, self._routing_inputs(plan, request, prepared, name, routed)), prepared def service_logged(self, call_type: str, start_time: float, end_time: float, request: Mapping[str, object]) -> None: """The ``ServiceTypes.ROUTER`` success event of a timed step, as a task of the running loop.""" @@ -2523,28 +2643,17 @@ class RustCore( _ = inferred_provider(model) # pyright: ignore[reportArgumentType] # see above return str(model) - def _no_model_error(self, plan: Plan, request: Mapping[str, object]) -> BaseException: - """``model=None``: the pick fails in ``_common_checks_available_deployment``, naming the - first default fallback when there is one, and with no group the fallback wrapper raises - that error as it is (``async_function_with_fallbacks_common_utils``).""" - import litellm.exceptions as litellm_exceptions - - named: Final = next( + def _no_model_name(self) -> str: + """``model=None``: ``_common_checks_available_deployment`` picks for the first default + fallback when there is one, and else for the name ``None``, which only a wildcard matches.""" + return next( ( - targets[0] + str(targets[0]) for entry in _as_list(self._resolved.fallbacks) if (targets := _as_list((as_str_mapping(entry) or {}).get("*"))) ), - None, + "None", ) - error: Final = litellm_exceptions.BadRequestError( - message=f"You passed in model={named}. {RouterErrors.no_healthy_deployments.value}", - model=named, # pyright: ignore[reportArgumentType] # None, as the Python router passes it - llm_provider="", - ) - if not plan.sync: - _Driver(self, plan, request)._log_pick_failure(error) # pyright: ignore[reportPrivateUsage] # the async pick's own logging - return error def _routing_inputs( self, @@ -2579,7 +2688,10 @@ class RustCore( ``make_call``, and a recovered stream is async.""" sync: Final = plan is not None and plan.sync make_call: Final = plan is None or not plan.single_attempt - callbacks: Final = self.outcome_callbacks(request, sync=sync, make_call=make_call) + callbacks: Final = { + **self.outcome_callbacks(request, sync=sync, make_call=make_call), + **({"shared_logging": True} if _shares_logging(request) else {}), + } return {"outcome_callbacks": callbacks} if callbacks else {} def _selection_fields(self, plan: Plan, request: Mapping[str, object]) -> Mapping[str, object]: @@ -2600,6 +2712,10 @@ class RustCore( **({"web_search": True} if web_search else {}), **({"route_handler": handler} if serves_route_capability(handler) else {}), **({"decisions_media": True} if decisions_media else {}), + # `batch_create_routing`: a batch create keeps to the deployments that can run batches. + **({"batch_create": True} if plan.kind == "batch" else {}), + # `deployments_admitting_embedding_batch`: an embedding keeps to the deployments admitting its input. + **({"embedding_batch": True} if plan.kind == "embedding" else {}), } def _picks_async(self) -> bool: @@ -2607,12 +2723,14 @@ class RustCore( def _process_renamed(self, model: str, group: str) -> str: """``model`` as the error names it (``_common_checks_available_deployment``): the group a name - resolved to by name, alias or routing group, renamed through the process ``model_alias_map``.""" + resolved to by name, alias or routing group, or the group of the deployment an id names, + renamed through the process ``model_alias_map``.""" resolution: Final = self._handle.resolve(model) - if resolution is None or resolution["kind"] not in ("Group", "Alias", "RoutingGroup"): + if resolution is None or resolution["kind"] not in ("Group", "Alias", "RoutingGroup", "Id"): return model + named: Final = str(resolution["model_group"]) if resolution["kind"] == "Id" else group renames: Final = self._host.dynamic().model_alias_map - return renames[group] if group in renames else group # noqa: SIM401 # `renames` is any process mapping, read as Python reads it + return renames[named] if named in renames else named # noqa: SIM401 # `renames` is any process mapping, read as Python reads it def _max_parallel_requests_client(self, deployment: Mapping[str, object], kwargs: Mapping[str, object]) -> object: """``_get_client(client_type="max_parallel_requests")``: the deployment's slot limit.""" @@ -2719,7 +2837,10 @@ class RustCore( with owning_picks(self._selectors.reservations, call_id_of(CALL_IDENTITY.validate_python(request))): call, prepared = self._start(plan, request, routed, prepared_already=prepared_already) arm_compaction_state(self._port, plan, prepared) # pyright: ignore[reportArgumentType] # the call's own dict - with picks_reserve(plan.endpoint not in _UNRESERVED_OPS): + with ( + picks_reserve(plan.endpoint not in _UNRESERVED_OPS), + self._usage_read_scope(str(prepared.get("model")), prepared), + ): driver, step = await self._run(plan, call, prepared, fresh=True, passthrough=passthrough) if step["kind"] == "failed": raise driver.failed(step) @@ -2744,6 +2865,16 @@ class RustCore( driver: Final = _Driver( self, plan, prepared, cut_deployment, stream_failure, fresh=fresh, passthrough=passthrough ) + try: + return await self._run_steps(driver, call) + except BaseException: + driver.close_route_phase() + raise + + async def _run_steps( + self, driver: "_Driver", call: "RouterCall[_CallSteps]" + ) -> "tuple[_Driver, Mapping[str, object]]": + """``_run``'s loop over the call's steps.""" outcome: object = None # rebind-ok: the answer to the last step while True: step = call.next(outcome) @@ -2769,6 +2900,7 @@ class RustCore( case "pre_route": outcome = await driver.pre_route(step) case "done" | "failed": + driver.close_route_phase() return driver, step case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list raise RuntimeError(f"router-core returned an unknown step {other!r}") @@ -3003,10 +3135,17 @@ class RustCore( if not isinstance(request, dict): return None depth: Final = step.get("fallback_depth") + if isinstance(depth, int) and depth > 0: + # A fallback hop routes on its own copy of the metadata buckets (`run_async_fallback`'s + # kwargs): what its hook clears and stamps leaves the caller's metadata as the first hop left it. + for name in ("metadata", "litellm_metadata"): + bucket = request.get(name) + if isinstance(bucket, dict): + request[name] = dict(bucket) # pyright: ignore[reportUnknownArgumentType] # the request's metadata dict r: Final = self._resolved owner: Final[object] = self.__dict__.get("_core_owner") public: Final[object] = owner() if callable(owner) else None - return await self._prerouter.pre_route( + routed: Final[object] = await self._prerouter.pre_route( str(step.get("model")), request, # pyright: ignore[reportUnknownArgumentType] # the request kwargs the hook stamps messages=request.get("messages") if messages is _NO_INPUT else messages, @@ -3023,6 +3162,7 @@ class RustCore( if self._hosted._core_routing_plugins # pyright: ignore[reportPrivateUsage] # the live plugin list else None, ) + return _served_as(routed, request) async def _run_plugins( self, @@ -3075,9 +3215,15 @@ class RustCore( selector: Final = self._selectors.selector_for( strategy, str(step.get("scope")), str(group) if group is not None else None, request ) - candidates: Final = [ - dict(as_str_mapping(self._handle.deployment(str(id))) or {}) for id in _as_list(step.get("candidates")) - ] + # Candidates sharing an id come with their own records, so the pick can land on any of them. + records: Final = step.get("records") + candidates: Final = ( + [dict(as_str_mapping(record) or {}) for record in _as_list(records)] + if records is not None + else [ + dict(as_str_mapping(self._handle.deployment(str(id))) or {}) for id in _as_list(step.get("candidates")) + ] + ) return strategy, selector, str(step.get("model")), candidates def _records_of(self, ids: object) -> list[dict[str, object]]: # mutable-ok: the deployment dicts filters take @@ -3144,6 +3290,45 @@ class RustCore( ) ] + def _override_selector(self, step: Mapping[str, object], request: Mapping[str, object]) -> object: + """The request's own strategy selector when it overrides the router's (``_override_selectors``), + whose pre-call check is in no callback list.""" + strategy: Final = str(step.get("strategy")) + selector: Final = self._selectors.selector_for(strategy, "override", None, request) + return selector if selector is not None and self._selectors.is_override(strategy, selector) else None + + def _embedding_batch( + self, step: Mapping[str, object], request: Mapping[str, object], input: object, *, tags: bool + ) -> list[str]: # mutable-ok: the step's answer + """``deployments_admitting_embedding_batch`` over the candidates; when none admits the input, + nothing while a routable, unpaused cooling deployment does (the pick waits out its cooldown), + else the candidates with the largest limit, so the 400 names the most the group accepts.""" + from litellm.router_utils.embedding_batch_size_check import ( + deployments_admitting_embedding_batch, + deployments_when_none_admits_embedding_batch, + embedding_deployment_selection, + ) + + sent: Final = request.get("input", input) + candidates: Final = self._records_of(step.get("candidates")) + with embedding_deployment_selection(): + admitting: Final = deployments_admitting_embedding_batch(candidates, sent, dict(request)) + if admitting is not None: + return [found for record in admitting if (found := picked_id(record)) is not None] + pool: Final = self._records_of(step.get("pool")) + pool_ids: Final = [found for record in pool if (found := picked_id(record)) is not None] + cooling: Final = [id for id, _ in self._cooldowns.get_active_cooldowns(pool_ids, None)] # pyright: ignore[reportArgumentType] # a list of ids + routable: Final = _region_routable(pool, request, True, self._resolved.enable_pre_call_checks) + routable_cooling: Final = ( + self._tag_routable_ids(str(step.get("model")), pool, routable, frozenset(cooling), request) + if tags + else frozenset(id for record in routable if (id := picked_id(record)) in cooling) + ) + kept: Final = deployments_when_none_admits_embedding_batch( + candidates, pool, routable_cooling, sent, dict(request) + ) + return [found for record in kept if (found := picked_id(record)) is not None] + async def _answer_select( self, step: Mapping[str, object], @@ -3167,6 +3352,14 @@ class RustCore( return self._route_incapable(step) if step.get("stage") == "media_refusing": return self._media_refusing(step, request) + if step.get("stage") == "embedding_batch": + return self._embedding_batch(step, request, input, tags=True) + if step.get("stage") == "override_check": + selector = self._override_selector(step, request) + if selector is not None: + for record in self._records_of(step.get("candidates")): + await selector.async_pre_call_check(dict(record), None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger + return step["candidates"] if step.get("stage") == "provider_request": return _provider_request(step) inputs: Final = self._strategy_inputs(step, request) @@ -3185,7 +3378,7 @@ class RustCore( picked: Final = ( candidates[0] if pinned - else await pick_async( + else await _pick_with_prefetched_usage( strategy, selector, # pyright: ignore[reportArgumentType] # a selector from selector_for model, @@ -3202,8 +3395,7 @@ class RustCore( await selector.async_pre_call_check(dict(as_str_mapping(picked) or {}), None) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger if picked is not None and timed: self.service_logged(".async_get_available_deployments", start_time, clock(), request) - found: Final = picked_id(picked) - return [found] if found is not None else [] + return _pick_answer(step, candidates, picked) def _answer_select_sync( self, step: Mapping[str, object], request: Mapping[str, object], messages: object, input: object @@ -3221,6 +3413,14 @@ class RustCore( return self._route_incapable(step) if step.get("stage") == "media_refusing": return self._media_refusing(step, request) + if step.get("stage") == "embedding_batch": + return self._embedding_batch(step, request, input, tags=False) + if step.get("stage") == "override_check": + selector = self._override_selector(step, request) + if selector is not None: + for record in self._records_of(step.get("candidates")): + selector.pre_call_check(dict(record)) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # every selector is a CustomLogger + return step["candidates"] if step.get("stage") == "provider_request": return _provider_request(step) inputs: Final = self._strategy_inputs(step, request) @@ -3242,8 +3442,7 @@ class RustCore( ) if picked is not None and self._selectors.is_override(strategy, selector): selector.pre_call_check(dict(as_str_mapping(picked) or {})) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # see above - found: Final = picked_id(picked) - return [found] if found is not None else [] + return _pick_answer(step, candidates, picked) # -- cache steps ----------------------------------------------------------------------------- @@ -3259,6 +3458,9 @@ class RustCore( read: Final[object] = await self._cache.async_batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above store: Final = self._cooldowns.cooldown_store + usage_read: Final = _USAGE_READ.get() + if usage_read is not None: + return self._corrected_cooldowns(keys, await self._batch_read_with_usage(keys, *usage_read)) with service_target(ROUTER_COOLDOWNS_TARGET): prefetched: Final = await RoutingReadBatch._read_prefetched(((store, keys),)) # pyright: ignore[reportPrivateUsage] # the batch's own prefetch read results: Final[object] = ( @@ -3266,6 +3468,37 @@ class RustCore( ) return self._corrected_cooldowns(keys, results) + async def _batch_read_with_usage( + self, + keys: list[str], + batch: RoutingReadBatch, + deployments: list[Mapping[str, object]], # mutable-ok: the read's key and deployment lists + ) -> object: + """The cooldown read with the usage-based-routing-v2 counters of ``deployments`` in one MGET + (``RoutingReadBatch.async_get_cooldown_deployments``), the counters kept on ``batch``.""" + from litellm.caching.dual_cache import DualCache + from litellm.router_strategy.lowest_tpm_rpm_v2 import PrefetchedUsage + from litellm.router_utils.routing_read_batch import ( + _routing_read_target, # pyright: ignore[reportPrivateUsage] # the batch's own target name + ) + + selector: Final = batch.usage_selector + assert selector is not None # set by `_usage_read_scope` + usage_keys: Final = list(itertools.chain(*selector.usage_counter_keys(deployments))) # pyright: ignore[reportArgumentType] # deployment dicts + reads: Final = ((self._cooldowns.cooldown_store, keys), (selector.router_cache, usage_keys)) + with service_target(_routing_read_target(keys, usage_keys)): + results: Final = await RoutingReadBatch._read_prefetched( + reads + ) or await DualCache.async_batch_get_cache_shared( # pyright: ignore[reportPrivateUsage] # see above + reads, parent_otel_span=None + ) + usage_values: Final = results[1] + batch.prefetched_usage = PrefetchedUsage( + keys=frozenset(usage_keys), + values=None if usage_values is None else MappingProxyType(dict(zip(usage_keys, usage_values))), + ) + return results[0] + def _batch_read_sync(self, keys: list[str]) -> Sequence[object]: # mutable-ok: the cache's key list """``_batch_read`` for a sync pick: ``get_active_cooldowns``' read.""" if not keys or not keys[0].endswith(":cooldown"): @@ -3295,7 +3528,13 @@ class RustCore( ) for key, amount in zip(keys, amounts) ] - values: Final = await self._cache.async_increment_cache_pipeline(increment_list=operations) + try: + values: Final = await self._cache.async_increment_cache_pipeline(increment_list=operations) + except Exception as e: # noqa: BLE001 # usage that cannot be written never fails the call, as on the Python router + from litellm._logging import verbose_router_logger + + verbose_router_logger.debug("_increment_deployment_usage: usage was not recorded. Got - %s", e) + return if values is None or self._cache.redis_cache is None: return for operation, value in zip(operations, values): @@ -3397,6 +3636,81 @@ class RustCore( # -- errors the core raises when no attempt failed ------------------------------------------ + def _routable_cooled_info( + self, + step: Mapping[str, object], + group: str, + cooling: Sequence[str], + request: Mapping[str, object], + *, + has_input: bool, + tags: bool, + ) -> tuple[CooldownDeploymentInfo, ...] | None: + """``RouterRateLimitError.cooldown_deployments``: the cooling deployments the request could + route to (region and, on an async pick, tag routing over the deployments the core listed + before its cooldown filter); ``None`` when the strategy refused every candidate.""" + request_ids: Final = step.get("error_request_ids") + if request_ids is None: + return None + records: Final = self._records_of(request_ids) + routable: Final = _region_routable(records, request, has_input, self._resolved.enable_pre_call_checks) + route: Final = ( + self._tag_routable_ids(group, records, routable, frozenset(cooling), request) + if tags + else frozenset(id for record in routable if (id := picked_id(record)) is not None) + ) + return get_cooldown_deployment_info(RoutingContext(self), route, cooling) # pyright: ignore[reportArgumentType] # reads registry() + + def _tag_routable_ids( + self, + group: str, + pool: Sequence[Mapping[str, object]], + routable: Sequence[Mapping[str, object]], + cooled: frozenset[str], + request: Mapping[str, object], + ) -> frozenset[str]: + """The cooling deployments in ``routable`` that tag routing would send the request to, each + pick counted along with the picks it falls back to once that one is unavailable, judged + without recording a tag match on the request.""" + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + from litellm.router_strategy.tag_based_routing import get_deployments_for_tag + + field: Final = get_metadata_variable_name_from_kwargs(dict(request)) + metadata: Final = request.get(field) + if field in request and not isinstance(metadata, Mapping): + return frozenset() + tag_request: Final = {**request, field: {**metadata}} if isinstance(metadata, Mapping) else dict(request) + found = frozenset[str]() # rebind-ok: gains each pass's tag routing picks + remaining = [ # rebind-ok: each pass drops what tag routing just picked + record for record in routable if picked_id(record) in cooled + ] + while remaining: + try: + # Tag routing awaits nothing: its coroutine finishes on the first send. + picked = _finished( + get_deployments_for_tag( + ctx=RoutingContext(self), # pyright: ignore[reportArgumentType] # reads registry() and settings() + model=group, + healthy_deployments=remaining, # pyright: ignore[reportArgumentType] # deployment dicts + request_kwargs=tag_request, # pyright: ignore[reportArgumentType] # the request kwargs + metadata_variable_name=field, + request_deployments=pool, # pyright: ignore[reportArgumentType] # deployment dicts + ) + ) + except ValueError: # tag routing sends this request to none of the rest + break + picked_ids = frozenset( + id + for record in (picked if isinstance(picked, list | tuple) else ()) + if (id := picked_id(record)) is not None + ) # pyright: ignore[reportArgumentType] # deployment dicts + rest = [record for record in remaining if picked_id(record) not in picked_ids] + if len(rest) == len(remaining): + break + found |= picked_ids + remaining = rest + return found + def _failure_error( # noqa: C901 # one branch per failure Python tells apart self, step: Mapping[str, object], @@ -3405,6 +3719,8 @@ class RustCore( *, pass_through: str = "", request: Mapping[str, object] | None = None, + has_input: bool = True, + tags: bool = True, ) -> BaseException: """The error a call or pick that failed before any attempt raises; ``pass_through`` names the selection kind (``pass_through``/``async_pass_through``) whose errors it words.""" @@ -3534,10 +3850,8 @@ class RustCore( enable_pre_call_checks=self._resolved.enable_pre_call_checks, cooldown_list=cooling, model_ids=list(ids), # pyright: ignore[reportArgumentType] # a list of ids - cooldown_deployments=get_cooldown_deployment_info( - RoutingContext(self), # pyright: ignore[reportArgumentType] # reads registry() - ids, - cooling, + cooldown_deployments=self._routable_cooled_info( + step, str(group), cooling, request or {}, has_input=has_input, tags=tags ), ) if last is not None: @@ -3581,6 +3895,88 @@ def _sleep_seconds(step: Mapping[str, object]) -> float: return 0 if seconds == 0 else seconds +def _served_as(routed: object, request: Mapping[str, object]) -> object: + """A pre-routed group, paired with the group the hook stamped the request as served for + (``fallback_lookup_groups``' routed group) when that differs: the level's fallbacks look it up.""" + from litellm.litellm_core_utils.core_helpers import get_metadata_variable_name_from_kwargs + + if not isinstance(routed, str): + return routed + metadata: Final = request.get(get_metadata_variable_name_from_kwargs(dict(request))) + stamped: Final = metadata.get("model_group") if isinstance(metadata, Mapping) else None + return (routed, stamped) if isinstance(stamped, str) and stamped and stamped != routed else routed + + +_USAGE_READ: Final[ContextVar["tuple[RoutingReadBatch, list[Mapping[str, object]]] | None"]] = ContextVar( + "rust_core_usage_read", default=None +) + + +async def _pick_with_prefetched_usage(*args: object, **kwargs: object) -> object: + """``pick_async`` with the counters the pick's cooldown read fetched (``PrefetchedUsage``).""" + from litellm.router_strategy.lowest_tpm_rpm_v2 import PrefetchedUsage + + usage_read: Final = _USAGE_READ.get() + with PrefetchedUsage.scoped(usage_read[0].prefetched_usage if usage_read is not None else None): + return await pick_async(*args, **kwargs) # pyright: ignore[reportArgumentType, reportCallIssue] # pick_async's own arguments + + +def _deployment_pick_attributes(model: str, request: Mapping[str, object]) -> Mapping[str, str | int]: + """Bounded attributes for one deployment pick; the attempt is 1-based within its model group.""" + metadata: Final = request.get("litellm_metadata", request.get("metadata")) + attempted: Final = metadata.get("attempted_retries") if isinstance(metadata, Mapping) else None + retries: Final = attempted if isinstance(attempted, int) else 0 + depth: Final = request.get("fallback_depth") + hop: Final = (isinstance(depth, int) and depth > 0) or model != request.get("model") + reason: Final = "retry" if retries > 0 else "fallback" if hop else "initial" + return MappingProxyType( + { + "litellm.deployment.attempt": retries + 1, + "litellm.deployment.reason": reason, + "litellm.deployment.model_group": model, + } + ) + + +def _region_routable( + records: Sequence[Mapping[str, object]], + request: Mapping[str, object], + has_input: bool, + pre_call_checks: bool, +) -> Sequence[Mapping[str, object]]: + """``records`` without the ones ``_pre_call_checks`` keeps the request away from for being + outside its ``allowed_model_region``.""" + from litellm.types.router import LiteLLM_Params + from litellm.utils import is_region_allowed + + allowed: Final = request.get("allowed_model_region") + if allowed is None or not pre_call_checks or not has_input: + return records + if not isinstance(allowed, str): + return () + return [ + record + for record in records + if is_region_allowed( + litellm_params=LiteLLM_Params.model_validate(record.get("litellm_params") or {}), + allowed_model_region=allowed, + ) + ] + + +_Value = TypeVar("_Value") + + +def _finished(coroutine: Coroutine[object, object, _Value]) -> _Value: + """The value of a coroutine that awaits nothing, run to its end without an event loop.""" + try: + coroutine.send(None) + except StopIteration as finished: + return finished.value # pyright: ignore[reportAny] # the coroutine's return value + coroutine.close() + raise RuntimeError("the coroutine awaited something") + + def _uses_pass_through(record: object) -> bool: params: Final = as_str_mapping((as_str_mapping(record) or {}).get("litellm_params")) or {} return bool(params.get("use_in_pass_through")) @@ -3676,28 +4072,46 @@ class _Selection: if step.get("single"): # A deployment named by its id is listed alone, as a dict. (record,) = _as_list(step.get("deployments")) - return for_requested_model(typed_deployment(as_str_mapping(record) or {}), self._model) + return self._served(record) # The deployments as served: a priority group's members carry their priority as order. - return [ - for_requested_model(typed_deployment(as_str_mapping(record) or {}), self._model) - for record in _as_list(step.get("deployments")) - ] + return [self._served(record) for record in _as_list(step.get("deployments"))] if step.get("deployment_id") in (DEFAULT_DEPLOYMENT_FAILING_ID, DEFAULT_DEPLOYMENT_COPIED_ID): # A default deployment every call to fails, or whose id is not a string, is picked as Python copies it. picked: Final = default_deployment_copy(self._core._resolved.default_deployment, self._model) # pyright: ignore[reportPrivateUsage] # the selection serves its own core return picked # pyright: ignore[reportReturnType] # the assigned value's copy, as the Python core returns it - return for_requested_model(typed_deployment(as_str_mapping(step.get("deployment")) or {}), self._model) + return self._served(step.get("deployment")) - def _failed(self, step: Mapping[str, object]) -> BaseException: + def _served(self, record: object) -> Mapping[str, object]: + """A deployment as served: its record, with the litellm_params values JSON carries as text.""" + return _with_live_params( + self._core, for_requested_model(typed_deployment(as_str_mapping(record) or {}), self._model) + ) + + def _failed(self, step: Mapping[str, object], *, sync: bool) -> BaseException: core: Final = self._core error: Final = core._failure_error( # pyright: ignore[reportPrivateUsage] # see above - step, self._model, self._filter_error, pass_through=self._kind, request=self._request + step, + self._model, + self._filter_error, + pass_through=self._kind, + request=self._request, + has_input=self._messages is not None or self._input is not None, + tags=not sync, ) if self._kind not in ("pass_through", "sync_pick"): _log_refused_pick(self._request, error) return error async def run(self) -> Mapping[str, object] | Sequence[Mapping[str, object]]: + from litellm.integrations.otel.runtime import phase_span + + with ( + phase_span(f"route {self._model}"), + self._core._usage_read_scope(str(self._model), self._request), # pyright: ignore[reportPrivateUsage] # the selection serves its own core + ): + return await self._run() + + async def _run(self) -> Mapping[str, object] | Sequence[Mapping[str, object]]: core: Final = self._core outcome: object = None # rebind-ok: the answer to the last step while True: @@ -3733,7 +4147,7 @@ class _Selection: # `_core_async_get_available_deployment` resolves its pick to a pod. return await self._resolved_pod(core, done) case "failed": - raise self._failed(step) + raise self._failed(step, sync=False) case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list raise RuntimeError(f"router-core returned an unknown step {other!r}") @@ -3769,7 +4183,7 @@ class _Selection: # `get_available_deployment` resolves its pick to a pod. return self._resolved_pod_sync(core, done) case "failed": - raise self._failed(step) + raise self._failed(step, sync=True) case other: # pyright: ignore[reportUnnecessaryComparison] # a kind the bridge returns but the stub does not list raise RuntimeError(f"router-core returned an unknown step {other!r}") @@ -3809,6 +4223,23 @@ def _log_refused_image_pick(request: Mapping[str, object], error: BaseException) _restore_correlation_context_if_supported(logging_obj) +def _with_live_params(core: "RustCore", deployment: Mapping[str, object]) -> Mapping[str, object]: + """``deployment`` with the values its JSON record carries as text: litellm_params holding an + Exception ``mock_response``, a callable or a str subclass that marks the value (an admin's secret + reference), and model_info values with keys that are not strings.""" + info: Final = as_str_mapping(deployment.get("model_info")) + model_id: Final = (info or {}).get("id") + live: Final = core._hosted.live_params(model_id) # pyright: ignore[reportPrivateUsage] # the core's own registry + live_info: Final = core._hosted.live_info(model_id) # pyright: ignore[reportPrivateUsage] # same + if not live and not live_info: + return deployment + return { + **deployment, + **({"litellm_params": {**(as_str_mapping(deployment.get("litellm_params")) or {}), **live}} if live else {}), + **({"model_info": {**(info or {}), **live_info}} if live_info else {}), + } + + def _custom_pick(picked: object) -> Mapping[str, object]: """A custom strategy's answer as the deployment document the Rust core serves.""" document: Final = as_str_mapping(picked) @@ -3863,6 +4294,65 @@ _NO_ROUTING_PRE_CALL_CHECKS: Final = frozenset({"search", "rerank", "moderation" _UNPICKED_KINDS: Final = frozenset({"search", "create_file"}) +# The info lines the Python router's endpoint helpers logged for an attempt: its success, and for +# these its failure too. +_ATTEMPT_LOG_NAMES: Final = MappingProxyType( + { + "acompletion": "litellm.acompletion", + "completion": "litellm.completion", + "aembedding": "litellm.aembedding", + "embedding": "litellm.embedding", + "aimage_generation": "litellm.aimage_generation", + "image_generation": "litellm.image_generation", + "atranscription": "litellm.atranscription", + "aspeech": "litellm.aspeech", + "arerank": "litellm.arerank", + "atext_completion": "litellm.atext_completion", + "aadapter_completion": "litellm.aadapter_completion", + } +) +_SUCCESS_ONLY_LOG_NAMES: Final = MappingProxyType( + { + "acancel_batch": "litellm.acancel_batch", + "acreate_batch": "litellm.acreate_batch", + "acreate_file": "litellm.acreate_file", + "avector_store_create": "litellm.avector_store_create", + } +) + + +def _log_level_traceback(error: BaseException | None) -> None: + """``async_function_with_fallbacks_common_utils``' DEBUG traceback of the error a level hands + its fallbacks: the stack the ERROR line leaves out.""" + import logging + import traceback + + from litellm._logging import redact_string, verbose_router_logger + + if error is None or not verbose_router_logger.isEnabledFor(logging.DEBUG): + return + formatted: Final = "".join(traceback.format_exception(error)) + verbose_router_logger.debug("Traceback%s", redact_string(formatted)) + + +def _log_attempt(plan: Plan, model_name: object, error: BaseException | None = None) -> None: + """``litellm.(model=...) 200 OK`` or ``... Exception ...``, as the Python router logged each attempt.""" + from litellm._logging import verbose_router_logger + + name: Final = _ATTEMPT_LOG_NAMES.get(plan.endpoint) + if error is not None: + if name is not None: + verbose_router_logger.info("%s(model=%s)\x1b[31m Exception %s\x1b[0m", name, model_name, str(error)) + return + served_name: Final = ( + name + or _SUCCESS_ONLY_LOG_NAMES.get(plan.endpoint) + or ("ageneric_api_call_with_fallbacks" if plan.kind == "generic" else None) + ) + if served_name is not None: + verbose_router_logger.info("%s(model=%s)\x1b[32m 200 OK\x1b[0m", served_name, model_name) + + class _Driver: """Answers the attempt and fallback steps of one call, and shapes its result. @@ -3897,6 +4387,8 @@ class _Driver: # the level that stream served, on the failure ``async_function_with_fallbacks_common_utils`` # was handed, which is what it raises when no fallback answers. self._base_level: Final = stream_failure[0] if stream_failure is not None else 0 + # How many fallback hops deep the walk is: the "Successful fallback" lines a served hop logs. + self._fallback_depth: int = 0 self._level_errors: Final[dict[int, BaseException]] = dict( # mutable-ok: filled as the call runs (stream_failure,) if stream_failure else () ) @@ -3916,6 +4408,9 @@ class _Driver: # When the latest attempt started: Python reads it off the logging object, when one is set up. self._attempt_started: datetime | None = None # A single-attempt endpoint runs neither of the wrappers that stamp the walk on the request. + # The open ``route {model_group}`` phase span of the pick in progress (async calls only). + self._route_phase: contextlib.ExitStack | None = None + self._route_model: str = "" self._walk: Final = WalkMetadata( request, plan.metadata_field, self._params_base_model, live=fresh and not plan.single_attempt ) @@ -4025,7 +4520,9 @@ class _Driver: def _picked_error(self, event: Mapping[str, object]) -> BaseException: """The error a pick an event names raised, as a failed step with that error raises it.""" model: Final = str(event.get("failed_model") or event.get("model")) - return self._core._failure_error(event, model, None, request=self._request) # pyright: ignore[reportPrivateUsage] # the driver serves its own core + return self._core._failure_error( # pyright: ignore[reportPrivateUsage] # the driver serves its own core + event, model, None, request=self._request, has_input=self._has_input(), tags=not self._sync + ) def _pick_raised(self, error: Exception, *, logged: bool) -> Mapping[str, object]: """A pick raised on the host (a strategy, a tag filter, a pre-routing hook): the try fails with @@ -4097,6 +4594,16 @@ class _Driver: self._chain_to_hop_first(picked, at_level) self._level_errors[at_level] = picked + def _log_hop(self, event: Mapping[str, object]) -> None: + """``run_async_fallback``'s "Falling back to model_group" line, for a hop to ``event``'s target.""" + from litellm._logging import verbose_router_logger + from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_structure + + target: Final = event["unroutable"] if "unroutable" in event else event.get("to") + verbose_router_logger.info("Falling back to model_group = %s", mask_sensitive_structure(target)) + depth: Final = event.get("depth") + self._fallback_depth = depth if isinstance(depth, int) else self._fallback_depth + 1 + def _carry_superseded_usage(self) -> None: """``carry_superseded_attempt_usage``, before each same-group retry and each fallback hop.""" from litellm.router_utils.fallback_event_handlers import carry_superseded_attempt_usage @@ -4198,11 +4705,39 @@ class _Driver: primary.message += self._explanation(primary, str(self._request.get("model"))) # pyright: ignore[reportAttributeAccessIssue, reportUnknownMemberType] # checked above # rebind-ok: the message is what this explains self._explains_failure = False + def open_route_phase(self, model: object) -> None: + """``_async_get_available_deployment_unresolved``'s ``route {model}`` phase: the pick's reads + nest in it, until the pick ends.""" + from litellm.integrations.otel.runtime import phase_span + + if self._sync or self._route_phase is not None: + return + stack: Final = contextlib.ExitStack() + stack.enter_context(phase_span(f"route {model}")) + self._route_phase = stack + self._route_model = str(model) + + def close_route_phase(self, selected_in: object = None) -> None: + """The pick ended: a served pick marks ``litellm.request.deployment_selected`` inside it first.""" + phase: Final = self._route_phase + if phase is None: + return + if isinstance(selected_in, str): + from litellm.integrations.otel.runtime import phase_event + + phase_event("litellm.request.deployment_selected", _deployment_pick_attributes(selected_in, self._request)) + self._route_phase = None + phase.close() + def take_events(self, events: object) -> None: # noqa: C901 # one branch per router-core event for raw_event in _as_list(events): event = as_str_mapping(raw_event) if event is None: continue + if event.get("kind") == "pick_started": + self.close_route_phase() + self.open_route_phase(event.get("model")) + continue if event.get("kind") == "retry": self._carry_superseded_usage() continue @@ -4211,6 +4746,7 @@ class _Driver: continue if event.get("kind") == "selected": self._selected(event) + self.close_route_phase(self._route_model) continue if event.get("kind") == "retries_opened": self._walk.retries_opened(self._stamped_budget(event)) @@ -4219,6 +4755,7 @@ class _Driver: self._walk.retry_budget(event) continue if event.get("kind") == "pick_failed": + self.close_route_phase() self._pick_failed(event) continue if event.get("kind") == "order_chain_failed": @@ -4254,6 +4791,7 @@ class _Driver: self._fallbacks_failed(event) continue if event.get("kind") == "fallback": + self._log_hop(event) hop_to = event.get("level") self._hop_first.pop(hop_to if isinstance(hop_to, int) else 0, None) self._walk.hop(event) @@ -4273,6 +4811,7 @@ class _Driver: if event.get("kind") == "falling_back": level = event.get("level") at = level if isinstance(level, int) else 0 + _log_level_traceback(self._level_errors.get(at, self._last_error)) self._explains_failure = self._explains_failure or at == self._base_level self._hint_level_failure(self._level_errors.get(at, self._last_error), str(event.get("model"))) continue @@ -4320,12 +4859,14 @@ class _Driver: def _call(self, step: Mapping[str, object]) -> "tuple[Mapping[str, object], AttemptCall]": name: Final = step.get("model") requested: Final = str(name) if name is not None else str(self._request.get("model")) - deployment: Final = ( + deployment: Mapping[str, object] = ( # rebind-ok: given its live params below # A default deployment whose id is not a string is called as Python copies it, id and all. as_str_mapping(default_deployment_copy(self._core._resolved.default_deployment, requested)) or {} # pyright: ignore[reportPrivateUsage] # the driver serves its own core if step.get("deployment_id") == DEFAULT_DEPLOYMENT_COPIED_ID else for_requested_model(typed_deployment(as_str_mapping(step.get("deployment")) or {}), requested) ) + # The values its JSON record carries as text (an Exception `mock_response`) called as configured. + deployment = _with_live_params(self._core, deployment) self._attempts += 1 self._attempt_started = datetime.now() # noqa: DTZ005 # naive to match the logging pipeline's api_call_start_time level: Final = step.get("fallback_level") @@ -4343,7 +4884,7 @@ class _Driver: self._strip_undecryptable_reasoning(step, deployment, call.kwargs) shaped: Final = shape_attempt_kwargs(call.kwargs, deployment) # The walk's metadata for the attempt's level (retries, fallback hops) and its hop kwargs. - call.kwargs = self._walk.stamp(step, shaped) + call.kwargs = self._walk.stamp(step, shaped, deployment) if self._plan.stamps_walk else shaped field: Final = self._plan.metadata_field if isinstance(call.served, dict) and field in call.kwargs: call.served[field] = call.kwargs[field] @@ -4387,6 +4928,7 @@ class _Driver: order: Final = step.get("target_order") deployment_id: Final = step.get("deployment_id") + routed: Final = step.get("routed") return StreamAttempt( core=self._core, op=self._op, @@ -4398,6 +4940,7 @@ class _Driver: level=number("fallback_level"), deployment_id=deployment_id if isinstance(deployment_id, str) else None, attempted_retries=number("retry"), + routed=routed if isinstance(routed, str) else None, slot=slot, ) @@ -4410,6 +4953,8 @@ class _Driver: ) -> Mapping[str, object]: # What litellm's failure logging read, before the attempt helper's `except` wrote on it. detail: Final = str(error) + if call is not None: + _log_attempt(self._plan, call.model_name, error) served: Final = call.served if call is not None else {} # pyright: ignore[reportUnknownVariableType] # litellm's untyped helpers and request values stamp_error(self._plan, error, deployment, served) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values self._core._endpoints.attempt_failed(self._plan, error, deployment, call, str(step.get("model"))) # pyright: ignore[reportPrivateUsage] # see above @@ -4425,14 +4970,41 @@ class _Driver: self._attempt_started, ) model_info: Final = as_str_mapping(deployment.get("model_info")) or {} + exits: Final = self._core.failure_exits(error, logged, model_info.get("id")) + # A call under credentials the request supplied ran as a request-scoped deployment + # (`_handle_clientside_credential`): its failure is never the shared deployment's. + ran_as: Final = (as_str_mapping(as_str_mapping(served).get("model_info") if served else None) or {}).get("id") # pyright: ignore[reportUnknownArgumentType] # see above + request_scoped: Final = self._core.ran_as_request_deployment( + ran_as, model_info.get("id") + ) or self._caller_entra_failure(error, served, model_info.get("id")) # pyright: ignore[reportUnknownArgumentType] # see above return { **attempt_failure(error), "detail": detail, - **self._core.failure_exits(error, logged, model_info.get("id")), + **exits, + **({"outside_deployment": True} if request_scoped or _raised_before_provider(error) else {}), # `async_deployment_callback_on_failure` does not run for a sync call, and skips a batch retrieve. - "uncounted_request": self._plan.sync or is_batch_retrieve_call_type(logged.get("call_type")), + "uncounted_request": self._plan.sync + or is_batch_retrieve_call_type(logged.get("call_type")) + or _raised_before_provider(error), } + def _caller_entra_failure(self, error: BaseException, sent: Mapping[str, object], deployment_id: object) -> bool: + """``is_caller_entra_credential_failure``: Entra ID rejected a credential the caller sent, not + the deployment's own, so the failure is the caller's.""" + from litellm.router_utils.cooldown_handlers import is_caller_entra_credential_failure + + if not isinstance(deployment_id, str): + return False + hosted: Final = self._core._hosted # pyright: ignore[reportPrivateUsage] # the driver serves its own core + defaults: Final = {k: v for k, v in self._core._resolved.default_litellm_params.items() if v is not None} # pyright: ignore[reportPrivateUsage] # see above + return is_caller_entra_credential_failure( + hosted._core_get_deployment(deployment_id), # pyright: ignore[reportPrivateUsage] # the host registry's id lookup + defaults, + {"litellm_params": sent}, + getattr(error, "status_code", ""), + error, + ) + def _chain_to_hop_first(self, error: BaseException, level: int) -> None: """A retry's failure is raised while the hop's first failure is handled (the retry loop runs in its ``except`` block), so it chains to it.""" @@ -4444,6 +5016,13 @@ class _Driver: self, call: "AttemptCall", response: object, host_op: str, stream: "StreamAttempt | None" = None ) -> Mapping[str, object]: self._last_response = response + _log_attempt(self._plan, call.model_name) + if self._fallback_depth > 0: + from litellm._logging import verbose_router_logger + + # One line per `run_async_fallback` the served hop returns through. + for _ in range(self._fallback_depth): + verbose_router_logger.info("Successful fallback b/w models.") tokens: Final = total_tokens(response) # A wrapped stream reports its success or failure when it ends; the attempt counts the # request and the tokens read so far, which the end's report leaves out. @@ -4625,6 +5204,15 @@ class _Driver: # plain dict fails the attempt. self._core._endpoints.counters.succeeded(self._plan, call) # pyright: ignore[reportPrivateUsage] # see above except Exception as error: # noqa: BLE001 # every attempt failure is classified and handed to the Rust call loop + if call is not None: + # `make_call`: a failed call hands back the io tokens it reserved, whatever its + # logging object (shared across attempts, it logs only the first failure). + with contextlib.suppress(Exception): + await async_io_token_refund_failure( + self._core._cache, # pyright: ignore[reportPrivateUsage] # the driver serves its own core + call.kwargs, # pyright: ignore[reportArgumentType] # the call's kwargs dict + parent_otel_span=get_parent_otel_span_from_kwargs(call.kwargs), + ) failed: Final = self._failed(step, deployment, call, error) if reading_stream is None or not logs_failure: return {**failed, "recorded": not logs_failure} @@ -4724,6 +5312,10 @@ class _Driver: ctx: Final = RoutingContext(self._core) # pyright: ignore[reportArgumentType] # the checks read registry() and settings() request: Final = dict(self._request) + if self._plan.kind in _GENERIC_API_KINDS: + # The generic helpers' kwargs carry the function they call, which the checks read (a + # search fallback runs the search tool grants). + request["original_generic_function"] = self._core._host.invoke(self._plan.host_op) # pyright: ignore[reportPrivateUsage, reportUnknownMemberType] # the driver serves its own core try: if target != get_pre_routing_selection(request) and not await ctx.fallback_access_allowed( model=target, request_kwargs=request @@ -4825,11 +5417,17 @@ class _Driver: await await_all(self.outcomes.take_pending()) return counted + def _has_input(self) -> bool: + """The request carries messages or an input, which region routing reads.""" + return self._request.get("messages") is not None or self._request.get("input") is not None + def failed(self, step: Mapping[str, object]) -> BaseException: model: Final = str(self._request.get("model")) primary: Final = self._level_errors.get(self._base_level) if primary is None: - return self._core._failure_error(step, model, self._last_error, request=self._request) # pyright: ignore[reportPrivateUsage] # see above + return self._core._failure_error( # pyright: ignore[reportPrivateUsage] # see above + step, model, self._last_error, request=self._request, has_input=self._has_input(), tags=not self._sync + ) if ( self._base_level == 0 and self._explains_failure diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index c1dc5162ccf..688b52f44a3 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -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", diff --git a/litellm/router_rust/invoke/streaming.py b/litellm/router_rust/invoke/streaming.py index ca70143072f..d8b4f87501a 100644 --- a/litellm/router_rust/invoke/streaming.py +++ b/litellm/router_rust/invoke/streaming.py @@ -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 diff --git a/litellm/router_rust/invoke/walk.py b/litellm/router_rust/invoke/walk.py index 4dab8edc07f..2c1481f068a 100644 --- a/litellm/router_rust/invoke/walk.py +++ b/litellm/router_rust/invoke/walk.py @@ -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: diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index b906a496550..75afe38548e 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -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 diff --git a/litellm/router_rust/steps/precall.py b/litellm/router_rust/steps/precall.py index 156c7b5e20f..227bb8de66f 100644 --- a/litellm/router_rust/steps/precall.py +++ b/litellm/router_rust/steps/precall.py @@ -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 "" ) diff --git a/litellm/router_rust/steps/prerouting.py b/litellm/router_rust/steps/prerouting.py index a83fbe29996..63d15a1a80a 100644 --- a/litellm/router_rust/steps/prerouting.py +++ b/litellm/router_rust/steps/prerouting.py @@ -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"] diff --git a/litellm/router_rust/surface/health.py b/litellm/router_rust/surface/health.py index 5547471855a..b16adc38af9 100644 --- a/litellm/router_rust/surface/health.py +++ b/litellm/router_rust/surface/health.py @@ -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 diff --git a/litellm/router_rust/surface/model_list.py b/litellm/router_rust/surface/model_list.py index 405a4b47d7b..bb7fb53acea 100644 --- a/litellm/router_rust/surface/model_list.py +++ b/litellm/router_rust/surface/model_list.py @@ -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: diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py index a55a1cfc6d6..75d30ebb2be 100644 --- a/litellm/router_rust/surface/registry.py +++ b/litellm/router_rust/surface/registry.py @@ -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 diff --git a/litellm/router_utils/search_api_router.py b/litellm/router_utils/search_api_router.py deleted file mode 100644 index 17a56e62e41..00000000000 --- a/litellm/router_utils/search_api_router.py +++ /dev/null @@ -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 diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index a12f79c6c8a..9ba8cbc35a7 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -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"})), ) diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 42ee6178bf0..f6e9c2157d7 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -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) diff --git a/tests/code_coverage_tests/router_litellm_globals_check.py b/tests/code_coverage_tests/router_litellm_globals_check.py index 172a3754000..ec38507a3bc 100644 --- a/tests/code_coverage_tests/router_litellm_globals_check.py +++ b/tests/code_coverage_tests/router_litellm_globals_check.py @@ -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", ) diff --git a/tests/code_coverage_tests/router_surface_check.py b/tests/code_coverage_tests/router_surface_check.py index 1c41526101a..8eb7651f51c 100644 --- a/tests/code_coverage_tests/router_surface_check.py +++ b/tests/code_coverage_tests/router_surface_check.py @@ -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.. +# The Rust core serves these by name too, so they never appear as self.. 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"])) diff --git a/tests/unit/caching/test_redis_cluster_cache.py b/tests/unit/caching/test_redis_cluster_cache.py index 85422819926..fa0b7631955 100644 --- a/tests/unit/caching/test_redis_cluster_cache.py +++ b/tests/unit/caching/test_redis_cluster_cache.py @@ -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: diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 4a07c97b709..7a224232f4c 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -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() diff --git a/tests/unit/containers/test_azure_container_transformation.py b/tests/unit/containers/test_azure_container_transformation.py index 38ea9323017..87b3def2a0b 100644 --- a/tests/unit/containers/test_azure_container_transformation.py +++ b/tests/unit/containers/test_azure_container_transformation.py @@ -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 diff --git a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py index 25de5ef60d9..5779250950d 100644 --- a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py +++ b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py @@ -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"] diff --git a/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py b/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py index 810dfffaf8b..4da4c16dff9 100644 --- a/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py +++ b/tests/unit/llms/anthropic/messages/test_advisor_orchestration.py @@ -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), ): diff --git a/tests/unit/llms/base_llm/responses/test_transformation.py b/tests/unit/llms/base_llm/responses/test_transformation.py index 50d918fa5be..68d9e79dd81 100644 --- a/tests/unit/llms/base_llm/responses/test_transformation.py +++ b/tests/unit/llms/base_llm/responses/test_transformation.py @@ -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() diff --git a/tests/unit/llms/github_copilot/test_per_user_auth.py b/tests/unit/llms/github_copilot/test_per_user_auth.py index 6811cbab2fb..1c3f32aa5f4 100644 --- a/tests/unit/llms/github_copilot/test_per_user_auth.py +++ b/tests/unit/llms/github_copilot/test_per_user_auth.py @@ -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" diff --git a/tests/unit/llms/openai/live/test_handler.py b/tests/unit/llms/openai/live/test_handler.py index bb7cda99e79..63df7a0a1cc 100644 --- a/tests/unit/llms/openai/live/test_handler.py +++ b/tests/unit/llms/openai/live/test_handler.py @@ -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=[], diff --git a/tests/unit/proxy/auth/test_fallback_model_access.py b/tests/unit/proxy/auth/test_fallback_model_access.py index 5c6a98d9e42..bcda2c05966 100644 --- a/tests/unit/proxy/auth/test_fallback_model_access.py +++ b/tests/unit/proxy/auth/test_fallback_model_access.py @@ -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( diff --git a/tests/unit/proxy/auth/test_model_access_group_budgets.py b/tests/unit/proxy/auth/test_model_access_group_budgets.py index d1b381907fa..6e2e9576fc8 100644 --- a/tests/unit/proxy/auth/test_model_access_group_budgets.py +++ b/tests/unit/proxy/auth/test_model_access_group_budgets.py @@ -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", diff --git a/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py b/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py index 36277de047d..b9a0d583475 100644 --- a/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py +++ b/tests/unit/proxy/auth/test_unmapped_model_budget_enforcement.py @@ -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): diff --git a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py index fe4d3c9bc4b..1fe2cb2ca68 100644 --- a/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py +++ b/tests/unit/proxy/hooks/test_autorouter_baseline_cache.py @@ -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", diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 9665b15bd60..32c9535b4c8 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -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", diff --git a/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py b/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py index cbdd029fc5d..43d8e56274f 100644 --- a/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py +++ b/tests/unit/proxy/management_endpoints/test_ptu_model_settings.py @@ -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() diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py index e654c6d444b..60b263948b6 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_passthrough_endpoints.py @@ -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 diff --git a/tests/unit/proxy/test_budget_reservation.py b/tests/unit/proxy/test_budget_reservation.py index 6ce8d400daf..f49c04319aa 100644 --- a/tests/unit/proxy/test_budget_reservation.py +++ b/tests/unit/proxy/test_budget_reservation.py @@ -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): diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 8143636a435..b68c616457b 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -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 diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index 91c87c63067..42e828f267e 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -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) diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 62d8db6f810..8536cbaa399 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -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: diff --git a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py index 926566605a5..73d3414abb0 100644 --- a/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py +++ b/tests/unit/proxy/vector_store_endpoints/test_vector_store_endpoints.py @@ -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 diff --git a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py index 012926c6cb3..2427e7eacb1 100644 --- a/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/unit/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -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", diff --git a/tests/unit/responses/test_responses_router_cooldown.py b/tests/unit/responses/test_responses_router_cooldown.py index 078890c2395..a1cd5c8517b 100644 --- a/tests/unit/responses/test_responses_router_cooldown.py +++ b/tests/unit/responses/test_responses_router_cooldown.py @@ -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}" ) diff --git a/tests/unit/router_strategy/adaptive_router/test_router_dispatch.py b/tests/unit/router_strategy/adaptive_router/test_router_dispatch.py index 6fe39599b0f..b7b33f4789b 100644 --- a/tests/unit/router_strategy/adaptive_router/test_router_dispatch.py +++ b/tests/unit/router_strategy/adaptive_router/test_router_dispatch.py @@ -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 == {} diff --git a/tests/unit/router_strategy/test_baseline_request.py b/tests/unit/router_strategy/test_baseline_request.py index eb785a41d63..c0ca2464a9c 100644 --- a/tests/unit/router_strategy/test_baseline_request.py +++ b/tests/unit/router_strategy/test_baseline_request.py @@ -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) == { diff --git a/tests/unit/router_strategy/test_complexity_router.py b/tests/unit/router_strategy/test_complexity_router.py index d0e2b381253..ba829512d8b 100644 --- a/tests/unit/router_strategy/test_complexity_router.py +++ b/tests/unit/router_strategy/test_complexity_router.py @@ -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__() diff --git a/tests/unit/router_strategy/test_llm_v2.py b/tests/unit/router_strategy/test_llm_v2.py index 993da777c66..4e566c0ae3b 100644 --- a/tests/unit/router_strategy/test_llm_v2.py +++ b/tests/unit/router_strategy/test_llm_v2.py @@ -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 diff --git a/tests/unit/router_strategy/test_quality_router.py b/tests/unit/router_strategy/test_quality_router.py index 249e5254412..3878527298c 100644 --- a/tests/unit/router_strategy/test_quality_router.py +++ b/tests/unit/router_strategy/test_quality_router.py @@ -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.""" diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 951182051c4..8a0d5b9a1d6 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -7,9 +7,8 @@ the implicit `"default"` group driven by the router's top-level import asyncio import datetime -import time import uuid -from collections.abc import Callable, Mapping +from collections.abc import Mapping from typing import Final, Literal from unittest.mock import patch @@ -22,7 +21,7 @@ from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger from litellm.router_api import RoutingContext from litellm.router_strategy.least_busy import picks_reserve -from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup, RoutingStrategy +from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup from litellm.utils import Rules, function_setup @@ -66,63 +65,6 @@ def _build_router(routing_strategy="simple-shuffle", routing_groups=None): ) -def test_no_groups_uses_top_level_strategy_for_all_models(): - router = _build_router(routing_strategy="latency-based-routing") - assert router._get_routing_context("filtered-model")[0] == "latency-based-routing" - assert router._get_routing_context("other-model")[0] == "latency-based-routing" - - -def test_explicit_group_overrides_top_level(): - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ], - ) - strategy, selector = router._get_routing_context("filtered-model") - assert strategy == "latency-based-routing" - assert selector is not None - # other-model isn't in any explicit group, so it lands in the default fallback - strategy_other, _ = router._get_routing_context("other-model") - assert strategy_other == "simple-shuffle" - - -def test_default_group_is_simple_shuffle_when_top_level_strategy_unset(): - router = _build_router() - strategy, selector = router._get_routing_context("other-model") - assert strategy == "simple-shuffle" - # simple-shuffle does not use a selector - assert selector is None - - -def test_two_groups_same_strategy_have_independent_selectors(): - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "group_a", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - "routing_strategy_args": {"ttl": 60}, - }, - { - "group_name": "group_b", - "models": ["other-model"], - "routing_strategy": "latency-based-routing", - "routing_strategy_args": {"ttl": 600}, - }, - ], - ) - a_selector = router._group_selectors["group_a"]["latency-based-routing"] - b_selector = router._group_selectors["group_b"]["latency-based-routing"] - assert a_selector is not None and b_selector is not None - assert id(a_selector) != id(b_selector) - - def test_overlapping_models_across_groups_raises(): with pytest.raises(ValueError, match="appears in"): _build_router( @@ -167,23 +109,6 @@ def test_invalid_strategy_in_group_raises(): ) -def test_unknown_model_in_group_warns_but_does_not_raise(caplog): - import logging - - with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): - router = _build_router( - routing_groups=[ - { - "group_name": "g1", - "models": ["model-not-in-list"], - "routing_strategy": "latency-based-routing", - } - ], - ) - assert "model-not-in-list" in caplog.text - assert router._model_to_group.get("model-not-in-list") == "g1" - - def test_duplicate_group_name_raises(): with pytest.raises(ValueError, match="duplicate"): _build_router( @@ -202,250 +127,6 @@ def test_duplicate_group_name_raises(): ) -def test_routing_group_object_input_accepted(): - router = _build_router( - routing_groups=[ - RoutingGroup( - group_name="g1", - models=["filtered-model"], - routing_strategy="latency-based-routing", - ) - ], - ) - assert router._model_to_group["filtered-model"] == "g1" - - -@pytest.mark.asyncio -async def test_async_dispatch_uses_group_strategy_for_grouped_model(): - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ], - ) - - group_selector = router._group_selectors["fast"]["latency-based-routing"] - - with ( - patch.object( - group_selector, - "async_get_available_deployments", - wraps=group_selector.async_get_available_deployments, - ) as latency_spy, - patch( - "litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle - ) as shuffle_spy, - ): - await router.async_get_available_deployment( - model="filtered-model", request_kwargs={} - ) - - assert ( - latency_spy.called - ), "group's latency selector should run for grouped model" - assert not shuffle_spy.called - - -@pytest.mark.asyncio -async def test_async_dispatch_falls_back_to_default_for_ungrouped_models(): - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ], - ) - - group_selector = router._group_selectors["fast"]["latency-based-routing"] - - with ( - patch.object( - group_selector, - "async_get_available_deployments", - wraps=group_selector.async_get_available_deployments, - ) as latency_spy, - patch( - "litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle - ) as shuffle_spy, - ): - await router.async_get_available_deployment( - model="other-model", request_kwargs={} - ) - - assert shuffle_spy.called, "default group's simple-shuffle should run" - assert not latency_spy.called - - -def test_update_settings_round_trip_routing_groups(): - router = _build_router() - assert router._model_to_group == {} - - router.update_settings( - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ] - ) - assert router._model_to_group == {"filtered-model": "fast"} - assert router._get_routing_context("filtered-model")[0] == "latency-based-routing" - - settings = router.get_settings() - assert any(g["group_name"] == "fast" for g in settings["routing_groups"]) - - -def test_default_group_accepts_routing_strategy_enum_for_top_level_strategy(): - """ - The Router constructor accepts both string and RoutingStrategy enum forms. - The default group must resolve a real selector for either input shape, or - every ungrouped model raises NoDeploymentAvailable. - """ - router = _build_router(routing_strategy=RoutingStrategy.LATENCY_BASED) - strategy, selector = router._get_routing_context("other-model") - assert strategy == "latency-based-routing" - assert selector is not None - assert router.lowestlatency_logger is selector - - -@pytest.mark.asyncio -async def test_async_dispatch_uses_default_selector_when_constructed_with_enum(): - router = _build_router(routing_strategy=RoutingStrategy.LATENCY_BASED) - default_selector = router.lowestlatency_logger - assert default_selector is not None - - # filtered-model has 2 deployments, so the single-deployment short-circuit - # in async_get_available_deployment doesn't kick in and we actually hit - # the latency selector. - with ( - patch.object( - default_selector, - "async_get_available_deployments", - wraps=default_selector.async_get_available_deployments, - ) as latency_spy, - patch( - "litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle - ) as shuffle_spy, - ): - await router.async_get_available_deployment( - model="filtered-model", request_kwargs={} - ) - - assert ( - latency_spy.called - ), "default group's latency selector must run when constructed with the enum" - assert not shuffle_spy.called - - -def test_update_settings_does_not_leak_strategy_callbacks(monkeypatch): - """ - Repeated `update_settings` calls must not accumulate stale selectors in - `litellm.callbacks` / `litellm.input_callback`. Each rebuild owns the - previous generation and is responsible for unregistering it by identity - so the global lists don't grow without bound. - """ - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - - router = _build_router(routing_strategy="latency-based-routing") - initial_default_selector = router.lowestlatency_logger - assert initial_default_selector is not None - assert any( - c is initial_default_selector for c in litellm.callbacks - ), "fresh router should register its default selector" - - # Flip strategies back and forth — each transition triggers - # routing_strategy_init and must replace, not accumulate. - for _ in range(3): - router.update_settings(routing_strategy="usage-based-routing") - router.update_settings(routing_strategy="latency-based-routing") - - assert all( - c is not initial_default_selector for c in litellm.callbacks - ), "old default selector instance leaked into litellm.callbacks" - - # The set of router-owned selectors should be bounded — at most one per - # strategy class. Without the cleanup fix this grows by one per toggle. - router_owned_classes = ( - "LeastBusyLoggingHandler", - "LowestTPMLoggingHandler", - "LowestTPMLoggingHandler_v2", - "LowestLatencyLoggingHandler", - "LowestCostLoggingHandler", - ) - router_owned = [ - c for c in litellm.callbacks if type(c).__name__ in router_owned_classes - ] - assert len(router_owned) <= len( - router_owned_classes - ), f"router selectors leaked: {[type(c).__name__ for c in router_owned]}" - - # Group selectors: capture the v1 instance, swap, confirm v1 is gone by id. - router.update_settings( - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ] - ) - group_selector_v1 = router._group_selectors["fast"]["latency-based-routing"] - assert group_selector_v1 is not None - - router.update_settings( - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - "routing_strategy_args": {"ttl": 600}, - } - ] - ) - assert all( - c is not group_selector_v1 for c in litellm.callbacks - ), "old group selector instance leaked into litellm.callbacks" - group_selector_v2 = router._group_selectors["fast"]["latency-based-routing"] - assert group_selector_v2 is not group_selector_v1 - - -def test_update_settings_unregisters_group_selectors_when_groups_removed(monkeypatch): - """ - Setting routing_groups to an empty list (or omitting all groups) must - unregister every previously-owned group selector. - """ - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - - router = _build_router( - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "least-busy", - } - ] - ) - group_selector = router._group_selectors["fast"]["least-busy"] - assert group_selector in litellm.callbacks - assert group_selector in litellm.input_callback - - router.update_settings(routing_groups=[]) - - assert all(c is not group_selector for c in litellm.callbacks) - assert all(c is not group_selector for c in litellm.input_callback) - assert router._group_selectors == {} - - def test_two_least_busy_groups_count_a_request_once(monkeypatch): """ Least-busy counts a request up from the pre-call hooks on `litellm.input_callback` and @@ -558,326 +239,10 @@ def test_two_routers_in_one_process_each_count_their_own_requests(monkeypatch): # --------------------------------------------------------------------------- -def test_normalize_strategy_handles_string_enum_and_none(): - assert Router._normalize_strategy(None) is None - assert Router._normalize_strategy("simple-shuffle") == "simple-shuffle" - assert ( - Router._normalize_strategy(RoutingStrategy.LATENCY_BASED) - == RoutingStrategy.LATENCY_BASED.value - ) - - -def test_validate_routing_strategy_accepts_valid_and_rejects_invalid(): - router = _build_router() - # Valid: string, enum, None - router._validate_routing_strategy("simple-shuffle") - router._validate_routing_strategy(RoutingStrategy.LATENCY_BASED) - router._validate_routing_strategy(None) - with pytest.raises(ValueError, match="Invalid routing_strategy"): - router._validate_routing_strategy("not-a-real-strategy") - - -def test_build_strategy_selector_returns_none_for_simple_shuffle(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router() - assert ( - router._build_strategy_selector( - strategy="simple-shuffle", - routing_strategy_args={}, - register_callbacks=False, - ) - is None - ) - - -def test_build_strategy_selector_constructs_for_known_strategies(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router() - selector = router._build_strategy_selector( - strategy="latency-based-routing", - routing_strategy_args={"ttl": 30}, - register_callbacks=False, - ) - assert selector is not None - - -def test_unregister_router_selectors_removes_by_identity(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router() - selector = router._build_strategy_selector( - strategy="least-busy", - routing_strategy_args={}, - register_callbacks=True, - ) - assert selector in litellm.callbacks - assert selector in litellm.input_callback - router._unregister_router_selectors([selector]) - assert all(c is not selector for c in litellm.callbacks) - assert all(c is not selector for c in litellm.input_callback) - - -def test_init_routing_groups_with_none_clears_state(): - router = _build_router( - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "least-busy", - } - ] - ) - assert router.routing_groups - router._init_routing_groups(None) - assert router.routing_groups == {} - assert router._model_to_group == {} - assert router._group_selectors == {} - - -@pytest.mark.asyncio -async def test_select_deployment_async_returns_none_without_selector(): - router = _build_router() - result = await router._select_deployment_async( - strategy="simple-shuffle", - selector=None, - model="filtered-model", - healthy_deployments=[], - messages=None, - input=None, - request_kwargs=None, - reserve=False, - ) - assert result is None - - -def test_select_deployment_sync_returns_none_without_selector(): - router = _build_router() - result = router._select_deployment_sync( - strategy="simple-shuffle", - selector=None, - model="filtered-model", - healthy_deployments=[], - messages=None, - input=None, - request_kwargs=None, - reserve=False, - ) - assert result is None - - -def test_select_deployment_sync_does_not_dispatch_cost_based_routing(): - """ - `LowestCostLoggingHandler` only implements `async_get_available_deployments`, - so the sync dispatch must fall through to None for `cost-based-routing` - instead of raising AttributeError. - """ - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "cheap", - "models": ["filtered-model"], - "routing_strategy": "cost-based-routing", - } - ], - ) - cost_selector = router._group_selectors["cheap"]["cost-based-routing"] - assert not hasattr(cost_selector, "get_available_deployments"), ( - "test premise broken: LowestCostLoggingHandler unexpectedly grew a " - "sync method — revisit the sync dispatch arm." - ) - - result = router._select_deployment_sync( - strategy="cost-based-routing", - selector=cost_selector, - model="filtered-model", - healthy_deployments=[], - messages=None, - input=None, - request_kwargs=None, - reserve=False, - ) - assert result is None - - -@pytest.mark.asyncio -async def test_async_dispatch_falls_back_to_sync_for_usage_based_routing_v1(): - """ - `LowestTPMLoggingHandler` (v1) only implements the sync - `get_available_deployments`. The async dispatch must call the sync method - on the v1 selector instead of awaiting a non-existent async one. - """ - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "v1", - "models": ["filtered-model"], - "routing_strategy": "usage-based-routing", - } - ], - ) - v1_selector = router._group_selectors["v1"]["usage-based-routing"] - assert not hasattr(v1_selector, "async_get_available_deployments"), ( - "test premise broken: LowestTPMLoggingHandler v1 unexpectedly grew an " - "async method — revisit the async dispatch arm." - ) - - with patch.object( - v1_selector, - "get_available_deployments", - wraps=v1_selector.get_available_deployments, - ) as v1_spy: - await router._select_deployment_async( - strategy="usage-based-routing", - selector=v1_selector, - model="filtered-model", - healthy_deployments=[ - { - "model_name": "filtered-model", - "litellm_params": {"model": "openai/gpt-4o"}, - "model_info": {"id": "deploy-1"}, - } - ], - messages=[{"role": "user", "content": "hi"}], - input=None, - request_kwargs={}, - reserve=False, - ) - - assert v1_spy.called, "async dispatch must route v1 strategy through sync method" - - -def test_request_routing_strategy_override_beats_top_level(): - router = _build_router(routing_strategy="least-busy") - strategy, selector = router._get_routing_context( - "other-model", {"routing_strategy": "simple-shuffle"} - ) - assert strategy == "simple-shuffle" - assert selector is None - - -def test_request_routing_strategy_override_beats_explicit_group(): - router = _build_router( - routing_strategy="simple-shuffle", - routing_groups=[ - { - "group_name": "fast", - "models": ["filtered-model"], - "routing_strategy": "latency-based-routing", - } - ], - ) - strategy, _ = router._get_routing_context( - "filtered-model", {"routing_strategy": "least-busy"} - ) - assert strategy == "least-busy" - - -def test_request_routing_strategy_override_builds_and_caches_selector(): - router = _build_router(routing_strategy="simple-shuffle") - strategy, selector = router._get_routing_context( - "other-model", {"routing_strategy": "latency-based-routing"} - ) - assert strategy == "latency-based-routing" - assert selector is not None - _, selector_again = router._get_routing_context( - "other-model", {"routing_strategy": "latency-based-routing"} - ) - assert selector_again is selector - - -def test_request_routing_strategy_override_matching_global_reuses_default_selector(): - router = _build_router(routing_strategy="least-busy") - _, selector = router._get_routing_context( - "other-model", {"routing_strategy": "least-busy"} - ) - assert selector is router.leastbusy_logger - assert router._override_selectors == {} - - -def test_invalid_request_routing_strategy_override_falls_back(): - router = _build_router(routing_strategy="least-busy") - strategy, selector = router._get_routing_context( - "other-model", {"routing_strategy": "not-a-real-strategy"} - ) - assert strategy == "least-busy" - assert selector is router.leastbusy_logger - - -def test_no_override_key_keeps_existing_behavior(): - router = _build_router(routing_strategy="least-busy") - strategy, _ = router._get_routing_context("other-model", {"messages": []}) - assert strategy == "least-busy" - strategy_none_kwargs, _ = router._get_routing_context("other-model", None) - assert strategy_none_kwargs == "least-busy" - - -def test_request_routing_strategy_override_helper_validates_directly(): - router = _build_router(routing_strategy="least-busy") - assert router._get_request_routing_strategy_override({"routing_strategy": "simple-shuffle"}) == "simple-shuffle" - assert router._get_request_routing_strategy_override({"routing_strategy": RoutingStrategy.LEAST_BUSY}) == "least-busy" - assert router._get_request_routing_strategy_override({"routing_strategy": "lar1"}) is None - assert router._get_request_routing_strategy_override({"routing_strategy": {"bad": "type"}}) is None - assert router._get_request_routing_strategy_override({}) is None - assert router._get_request_routing_strategy_override(None) is None - - -def test_override_strategy_selector_helper_builds_per_strategy(): - router = _build_router(routing_strategy="least-busy") - latency_selector = router._get_override_strategy_selector("latency-based-routing") - assert latency_selector is not None - assert router._get_override_strategy_selector("latency-based-routing") is latency_selector - assert router._get_override_strategy_selector("least-busy") is router.leastbusy_logger - assert router._get_override_strategy_selector("simple-shuffle") is None - - -def test_strategy_reinit_unregisters_override_selectors(): - router = _build_router(routing_strategy="least-busy") - override_selector = router._get_override_strategy_selector("latency-based-routing") - assert override_selector is not None - assert not any(cb is override_selector for cb in litellm.callbacks) - - router.update_settings(routing_strategy="latency-based-routing") - - assert router._override_selectors == {} - assert not any(cb is override_selector for cb in litellm.callbacks) - assert router._get_override_strategy_selector("latency-based-routing") is router.lowestlatency_logger - - def _single_latency_group(): return [{"group_name": "g1", "models": ["filtered-model"], "routing_strategy": "latency-based-routing"}] -def _assert_still_routes_with_original_group(router, selector): - assert list(router.routing_groups) == ["g1"] - assert router._model_to_group == {"filtered-model": "g1"} - assert router._group_selectors["g1"]["latency-based-routing"] is selector - assert router._get_routing_context("filtered-model", None) == ("latency-based-routing", selector) - assert sum(1 for cb in litellm.callbacks if cb is selector) == 1 - - -def test_failed_routing_groups_update_keeps_previous_groups(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router(routing_groups=_single_latency_group()) - selector = router._group_selectors["g1"]["latency-based-routing"] - - with pytest.raises(ValueError, match="appears in"): - router.update_settings( - routing_groups=[ - *_single_latency_group(), - {"group_name": "g2", "models": ["filtered-model"], "routing_strategy": "least-busy"}, - ], - ) - - _assert_still_routes_with_original_group(router, selector) - assert sum(1 for cb in litellm.callbacks if type(cb) is not type(selector)) == 0 - assert litellm.input_callback == [] - - def test_failed_routing_groups_update_does_not_poison_later_strategy_changes(monkeypatch): monkeypatch.setattr(litellm, "callbacks", []) monkeypatch.setattr(litellm, "input_callback", []) @@ -918,106 +283,6 @@ def test_overlap_error_names_every_conflicting_model(): assert "'other-model' appears in 'g1' and 'g2'" in message -def test_invalid_group_strategy_keeps_previous_groups(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router(routing_groups=_single_latency_group()) - selector = router._group_selectors["g1"]["latency-based-routing"] - - with pytest.raises(ValueError, match="Invalid routing_strategy"): - router.update_settings( - routing_groups=[ - {"group_name": "g2", "models": ["other-model"], "routing_strategy": "not-a-real-strategy"}, - ], - ) - - _assert_still_routes_with_original_group(router, selector) - - -def test_unbuildable_group_selector_keeps_previous_groups(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router(routing_groups=_single_latency_group()) - selector = router._group_selectors["g1"]["latency-based-routing"] - - with pytest.raises(ValidationError, match="ttl"): - router.update_settings( - routing_groups=[ - {"group_name": "g0", "models": ["other-model"], "routing_strategy": "least-busy"}, - *_single_latency_group(), - { - "group_name": "g2", - "models": ["other-model-2"], - "routing_strategy": "latency-based-routing", - "routing_strategy_args": {"ttl": "not-a-number"}, - }, - ], - ) - - _assert_still_routes_with_original_group(router, selector) - assert litellm.callbacks == [selector] - assert litellm.input_callback == [] - - -def test_register_router_selector_wires_only_the_hooks_the_strategy_needs(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router() - least_busy = router._build_strategy_selector( - strategy="least-busy", routing_strategy_args={}, register_callbacks=False - ) - latency = router._build_strategy_selector( - strategy="latency-based-routing", routing_strategy_args={}, register_callbacks=False - ) - assert least_busy is not None and latency is not None - assert litellm.callbacks == [] and litellm.input_callback == [] - - router._register_router_selector(least_busy) - router._register_router_selector(latency) - - assert [cb for cb in litellm.callbacks if cb is least_busy or cb is latency] == [least_busy, latency] - assert litellm.input_callback == [least_busy] - - -def test_replace_routing_groups_swaps_state_and_callbacks_in_one_step(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router(routing_groups=_single_latency_group()) - old_selector = router._group_selectors["g1"]["latency-based-routing"] - new_selector = router._build_strategy_selector( - strategy="least-busy", routing_strategy_args={}, register_callbacks=False - ) - assert new_selector is not None - - router._replace_routing_groups( - ( - (RoutingGroup(group_name="g2", models=["other-model"], routing_strategy="least-busy"), new_selector), - (RoutingGroup(group_name="g3", models=["other-model-2"], routing_strategy="simple-shuffle"), None), - ) - ) - - assert list(router.routing_groups) == ["g2", "g3"] - assert router._model_to_group == {"other-model": "g2", "other-model-2": "g3"} - assert router._group_selectors == {"g2": {"least-busy": new_selector}, "g3": {}} - assert router._get_routing_context("other-model", None) == ("least-busy", new_selector) - assert router._get_routing_context("filtered-model", None)[0] == router.routing_strategy - assert all(cb is not old_selector for cb in litellm.callbacks) - assert sum(1 for cb in litellm.callbacks if cb is new_selector) == 1 - assert litellm.input_callback == [new_selector] - - -def test_override_selectors_are_not_registered_process_wide(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - monkeypatch.setattr(litellm, "input_callback", []) - router = _build_router(routing_strategy="simple-shuffle") - v2_selector = router._get_override_strategy_selector("usage-based-routing-v2") - least_busy_selector = router._get_override_strategy_selector("least-busy") - assert v2_selector is not None and least_busy_selector is not None - - assert litellm.callbacks == [] - assert litellm.input_callback == [] - - def _rpm_limited_model_list(): return [ { @@ -1033,81 +298,6 @@ def _rpm_limited_model_list(): ] -async def _mock_completion(router, **override): - return await router.acompletion( - model="other-model", messages=[{"role": "user", "content": "hi"}], mock_response="ok", **override - ) - - -@pytest.mark.asyncio -async def test_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(): - router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) - - await _mock_completion(router, routing_strategy="usage-based-routing-v2") - override_selector = router._override_selectors["usage-based-routing-v2"] - assert not any(cb is override_selector for cb in litellm.callbacks) - - with patch.object( - override_selector, "async_pre_call_check", wraps=override_selector.async_pre_call_check - ) as pre_call_spy: - for _ in range(2): - plain = await _mock_completion(router) - assert plain.choices[0].message.content == "ok" - assert not pre_call_spy.called - - with pytest.raises(litellm.RateLimitError): - await _mock_completion(router, routing_strategy="usage-based-routing-v2") - - -def test_sync_usage_based_v2_override_stays_scoped_to_the_request_that_asked_for_it(monkeypatch): - monkeypatch.setattr( - "litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", - lambda: datetime.datetime(2026, 1, 1, 12, 0, 30, tzinfo=datetime.timezone.utc), - ) - router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) - messages = [{"role": "user", "content": "hi"}] - - router.completion( - model="other-model", messages=messages, mock_response="ok", routing_strategy="usage-based-routing-v2" - ) - override_selector = router._override_selectors["usage-based-routing-v2"] - assert not any(cb is override_selector for cb in litellm.callbacks) - - for _ in range(2): - plain = router.completion(model="other-model", messages=messages, mock_response="ok") - assert plain.choices[0].message.content == "ok" - - with pytest.raises(ValueError, match="No deployments available"): - router.completion( - model="other-model", messages=messages, mock_response="ok", routing_strategy="usage-based-routing-v2" - ) - - -@pytest.mark.asyncio -async def test_override_selector_pre_call_check_only_runs_for_override_selectors(monkeypatch): - monkeypatch.setattr(litellm, "callbacks", []) - deployment = _rpm_limited_model_list()[0] - - override_router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle") - override_selector = override_router._get_override_strategy_selector("usage-based-routing-v2") - await override_router._async_override_selector_pre_call_check( - "usage-based-routing-v2", override_selector, deployment, None - ) - with pytest.raises(litellm.RateLimitError): - override_router._override_selector_pre_call_check("usage-based-routing-v2", override_selector, deployment) - - default_router = Router(model_list=_rpm_limited_model_list(), routing_strategy="usage-based-routing-v2") - for _ in range(2): - await default_router._async_override_selector_pre_call_check( - "usage-based-routing-v2", default_router.lowesttpm_logger_v2, deployment, None - ) - default_router._override_selector_pre_call_check( - "usage-based-routing-v2", default_router.lowesttpm_logger_v2, deployment - ) - await default_router._async_override_selector_pre_call_check(None, None, deployment, None) - default_router._override_selector_pre_call_check(None, None, deployment) - - @pytest.mark.asyncio async def test_usage_based_v2_override_enforces_rpm_when_a_specific_deployment_is_requested(): router = Router(model_list=_rpm_limited_model_list(), routing_strategy="simple-shuffle", num_retries=0) @@ -1191,149 +381,6 @@ def _proxy_shaped_request(**data: object) -> dict[str, object]: return {**kwargs, "litellm_logging_obj": logging_obj} -async def _async_override_pick(router: Router, strategy: str) -> str: - deployment = await router.async_get_available_deployment( - "grp", request_kwargs=_proxy_shaped_request(model="grp", routing_strategy=strategy) - ) - return deployment["model_info"]["id"] - - -def _sync_override_pick(router: Router, strategy: str) -> str: - deployment = router.get_available_deployment( - "grp", request_kwargs=_proxy_shaped_request(model="grp", routing_strategy=strategy) - ) - return deployment["model_info"]["id"] - - -def _in_flight(router: Router, deployment_id: str) -> int | None: - return router.cache.get_cache(f"grp_request_count:{deployment_id}") - - -async def _async_wait_until(predicate: Callable[[], bool]) -> None: - for _ in range(100): - if predicate(): - return - await asyncio.sleep(0.02) - raise AssertionError("lifecycle callback never reached the override selector") - - -def _sync_wait_until(predicate: Callable[[], bool]) -> None: - for _ in range(100): - if predicate(): - return - time.sleep(0.02) - raise AssertionError("lifecycle callback never reached the override selector") - - -def _selector_is_not_global(selector: CustomLogger) -> bool: - global_lists = ( - litellm.callbacks, - litellm.input_callback, - litellm.success_callback, - litellm.failure_callback, - litellm._async_success_callback, - litellm._async_failure_callback, - ) - return not any(cb is selector for cbs in global_lists for cb in cbs) - - -@pytest.mark.asyncio -async def test_least_busy_override_sees_the_overriding_request_in_flight(): - router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle", num_retries=0) - - stream = await router.acompletion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True)) - busy = stream._hidden_params["model_id"] - idle = "d2" if busy == "d1" else "d1" - assert [await _async_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle] - - async for _ in stream: - pass - await _async_wait_until(lambda: _in_flight(router, busy) == 0) - assert await _async_override_pick(router, "least-busy") == "d1" - assert _selector_is_not_global(router._override_selectors["least-busy"]) - - -def test_sync_least_busy_override_sees_the_overriding_request_in_flight(): - router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle", num_retries=0) - - stream = router.completion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy", stream=True)) - busy = stream._hidden_params["model_id"] - idle = "d2" if busy == "d1" else "d1" - assert [_sync_override_pick(router, "least-busy") for _ in range(3)] == [idle, idle, idle] - - for _ in stream: - pass - _sync_wait_until(lambda: _in_flight(router, busy) == 0) - assert _sync_override_pick(router, "least-busy") == "d1" - assert _selector_is_not_global(router._override_selectors["least-busy"]) - - -@pytest.mark.asyncio -async def test_least_busy_override_releases_the_slot_when_the_overriding_request_fails(): - router = Router( - model_list=_two_deployment_model_list(mock_response="litellm.InternalServerError"), - routing_strategy="simple-shuffle", - num_retries=0, - ) - - with pytest.raises(litellm.InternalServerError): - await router.acompletion(**_proxy_shaped_request(model="grp", routing_strategy="least-busy")) - - await _async_wait_until(lambda: _in_flight(router, "d1") == 0) - assert await _async_override_pick(router, "least-busy") == "d1" - assert _selector_is_not_global(router._override_selectors["least-busy"]) - - -@pytest.mark.asyncio -async def test_latency_based_override_learns_from_the_overriding_requests(): - router = Router( - model_list=_two_deployment_model_list(mock_delay=0.05), routing_strategy="simple-shuffle", num_retries=0 - ) - - def samples(deployment_id: str) -> list[float]: - recorded = (router.cache.get_cache("grp_map") or {}).get(deployment_id, {}).get("latency_normalized", []) - return [latency for latency in recorded if latency > 0] - - async def overriding_call() -> str: - sampled_before = {"d1": len(samples("d1")), "d2": len(samples("d2"))} - response = await router.acompletion( - **_proxy_shaped_request(model="grp", routing_strategy="latency-based-routing") - ) - deployment_id = response._hidden_params["model_id"] - await _async_wait_until(lambda: len(samples(deployment_id)) > sampled_before[deployment_id]) - return deployment_id - - served = [await overriding_call() for _ in range(6)] - - assert "d1" in served - assert served[2:] == ["d2"] * 4 - assert _selector_is_not_global(router._override_selectors["latency-based-routing"]) - - -def test_override_selector_is_bound_only_to_the_request_that_asked_for_it(): - router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle") - overriding = _proxy_shaped_request(model="grp", routing_strategy="least-busy") - plain = _proxy_shaped_request(model="grp") - - router.get_available_deployment("grp", request_kwargs=overriding) - router.get_available_deployment("grp", request_kwargs=overriding) - router.get_available_deployment("grp", request_kwargs=plain) - - selector = router._override_selectors["least-busy"] - bound = overriding["litellm_logging_obj"] - for callbacks in ( - bound.dynamic_input_callbacks, - bound.dynamic_success_callbacks, - bound.dynamic_async_success_callbacks, - bound.dynamic_failure_callbacks, - bound.dynamic_async_failure_callbacks, - ): - assert callbacks == [selector] - unbound = plain["litellm_logging_obj"] - assert unbound.dynamic_input_callbacks is None and unbound.dynamic_success_callbacks is None - assert unbound.dynamic_failure_callbacks is None and unbound.dynamic_async_failure_callbacks is None - - def test_override_matching_the_router_strategy_is_not_bound_twice(): router = Router(model_list=_two_deployment_model_list(), routing_strategy="least-busy") request = _proxy_shaped_request(model="grp", routing_strategy="least-busy") @@ -1343,52 +390,20 @@ def test_override_matching_the_router_strategy_is_not_bound_twice(): assert request["litellm_logging_obj"].dynamic_input_callbacks is None -@pytest.mark.asyncio -async def test_override_matching_a_routing_group_strategy_records_each_request_once(): - router = Router( - model_list=_two_deployment_model_list(), - routing_strategy="simple-shuffle", - routing_groups=[RoutingGroup(group_name="lat", models=["grp"], routing_strategy="latency-based-routing")], - num_retries=0, - ) - request = _proxy_shaped_request(model="grp", routing_strategy="latency-based-routing") - assert router._globally_registered_strategies() == {"simple-shuffle", "latency-based-routing"} - - response = await router.acompletion(**request) - deployment_id = response._hidden_params["model_id"] - await _async_wait_until(lambda: (router.cache.get_cache("grp_map") or {}).get(deployment_id) is not None) - - assert len(router.cache.get_cache("grp_map")[deployment_id]["latency_normalized"]) == 1 - assert request["litellm_logging_obj"].dynamic_success_callbacks is None - - -def test_bind_override_selector_to_request_binds_once_and_ignores_requests_without_logging(): - router = Router(model_list=_two_deployment_model_list(), routing_strategy="simple-shuffle") - selector = router._get_override_strategy_selector("least-busy") - request = _proxy_shaped_request(model="grp", routing_strategy="least-busy") - request["litellm_logging_obj"].dynamic_success_callbacks = ["langfuse"] - - router._bind_override_selector_to_request("least-busy", selector, request) - router._bind_override_selector_to_request("least-busy", selector, request) - router._bind_override_selector_to_request("least-busy", selector, None) - router._bind_override_selector_to_request("least-busy", selector, {"model": "grp"}) - - logging_obj = request["litellm_logging_obj"] - assert logging_obj.dynamic_success_callbacks == ["langfuse", selector] - assert logging_obj.dynamic_input_callbacks == [selector] - assert logging_obj.dynamic_async_failure_callbacks == [selector] - assert _selector_is_not_global(selector) - - def _quality_group(strategy="latency-based-routing"): return [{"group_name": "quality", "models": ["filtered-model", "other-model"], "routing_strategy": strategy}] -def test_group_name_is_callable_and_unions_member_deployments(): +async def _candidate_ids(router: Router, model: str) -> list[str]: + deployments: Final = await router.async_get_healthy_deployments(model=model, request_kwargs={}) + assert isinstance(deployments, list) + return sorted(d["model_info"]["id"] for d in deployments) + + +@pytest.mark.asyncio +async def test_group_name_is_callable_and_unions_member_deployments(): router = _build_router(routing_groups=_quality_group()) - model, deployments = router._common_checks_available_deployment(model="quality") - assert model == "quality" - assert sorted(d["model_info"]["id"] for d in deployments) == ["deploy-1", "deploy-2", "deploy-3"] + assert await _candidate_ids(router, "quality") == ["deploy-1", "deploy-2", "deploy-3"] def test_group_name_appears_in_model_names_and_model_list(): @@ -1399,34 +414,8 @@ def test_group_name_appears_in_model_names_and_model_list(): assert sorted(r["model_info"]["id"] for r in rows) == ["deploy-1", "deploy-2", "deploy-3"] -def test_get_routing_context_for_group_name_uses_group_strategy(): - router = _build_router(routing_groups=_quality_group()) - strategy, selector = router._get_routing_context("quality") - assert strategy == "latency-based-routing" - assert selector is router._group_selectors["quality"]["latency-based-routing"] - - @pytest.mark.asyncio -async def test_group_call_dispatches_via_group_selector(): - router = _build_router(routing_groups=_quality_group()) - group_selector = router._group_selectors["quality"]["latency-based-routing"] - - with ( - patch.object( - group_selector, - "async_get_available_deployments", - wraps=group_selector.async_get_available_deployments, - ) as latency_spy, - patch("litellm.router_python_core.simple_shuffle", wraps=litellm.router_python_core.simple_shuffle) as shuffle_spy, - ): - deployment = await router.async_get_available_deployment(model="quality", request_kwargs={}) - - assert latency_spy.called - assert not shuffle_spy.called - assert deployment["model_name"] in {"filtered-model", "other-model"} - - -def test_group_name_colliding_with_model_name_is_shadowed_with_warning(caplog): +async def test_group_name_colliding_with_model_name_is_shadowed_with_warning(caplog): import logging with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): @@ -1437,13 +426,11 @@ def test_group_name_colliding_with_model_name_is_shadowed_with_warning(caplog): ) assert any("shadowed" in record.getMessage() for record in caplog.records) assert router.get_routing_group("filtered-model") is None - assert router._get_routing_context("other-model")[0] == "latency-based-routing" - - model, deployments = router._common_checks_available_deployment(model="filtered-model") - assert sorted(d["model_info"]["id"] for d in deployments) == ["deploy-1", "deploy-2"] + assert await _candidate_ids(router, "filtered-model") == ["deploy-1", "deploy-2"] -def test_group_name_colliding_with_model_group_alias_is_shadowed_with_warning(caplog): +@pytest.mark.asyncio +async def test_group_name_colliding_with_model_group_alias_is_shadowed_with_warning(caplog): import logging with caplog.at_level(logging.WARNING, logger="LiteLLM Router"): @@ -1454,13 +441,11 @@ def test_group_name_colliding_with_model_group_alias_is_shadowed_with_warning(ca ) assert any("shadowed" in record.getMessage() for record in caplog.records) assert router.get_routing_group("quality") is None - - model, deployments = router._common_checks_available_deployment(model="quality") - assert model == "filtered-model" - assert sorted(d["model_info"]["id"] for d in deployments) == ["deploy-1", "deploy-2"] + assert await _candidate_ids(router, "quality") == ["deploy-1", "deploy-2"] -def test_real_model_added_later_shadows_group(): +@pytest.mark.asyncio +async def test_real_model_added_later_shadows_group(): router = _build_router(routing_groups=_quality_group()) assert router.get_routing_group("quality") is not None @@ -1474,34 +459,32 @@ def test_real_model_added_later_shadows_group(): ) ) assert router.get_routing_group("quality") is None - model, deployments = router._common_checks_available_deployment(model="quality") - assert [d["model_info"]["id"] for d in deployments] == ["deploy-shadow"] + assert await _candidate_ids(router, "quality") == ["deploy-shadow"] router.delete_deployment(id="deploy-shadow") assert "quality" not in router.model_names assert router.get_routing_group("quality") is not None - _, restored = router._common_checks_available_deployment(model="quality") - assert sorted(d["model_info"]["id"] for d in restored) == ["deploy-1", "deploy-2", "deploy-3"] + assert await _candidate_ids(router, "quality") == ["deploy-1", "deploy-2", "deploy-3"] -def test_group_with_no_member_deployments_raises_no_healthy(): +@pytest.mark.asyncio +async def test_group_with_no_member_deployments_raises_no_healthy(): router = Router( model_list=_model_list(), routing_groups=[{"group_name": "empty-group", "models": ["ghost-model"], "routing_strategy": "simple-shuffle"}], ) with pytest.raises(litellm.BadRequestError): - router._common_checks_available_deployment(model="empty-group") + await router.async_get_healthy_deployments(model="empty-group", request_kwargs={}) -def test_alias_pointing_at_group_composes(): +@pytest.mark.asyncio +async def test_alias_pointing_at_group_composes(): router = Router( model_list=_model_list(), model_group_alias={"quality-alias": "quality"}, routing_groups=_quality_group(), ) - model, deployments = router._common_checks_available_deployment(model="quality-alias") - assert model == "quality" - assert sorted(d["model_info"]["id"] for d in deployments) == ["deploy-1", "deploy-2", "deploy-3"] + assert await _candidate_ids(router, "quality-alias") == ["deploy-1", "deploy-2", "deploy-3"] def test_model_group_info_reports_group(): @@ -1525,11 +508,10 @@ def test_routing_group_has_alternatives(): assert solo_router.routing_group_has_alternatives("solo-group") is False -def test_member_direct_call_unchanged_by_callable_groups(): +@pytest.mark.asyncio +async def test_member_direct_call_unchanged_by_callable_groups(): router = _build_router(routing_groups=_quality_group()) - model, deployments = router._common_checks_available_deployment(model="other-model") - assert model == "other-model" - assert [d["model_info"]["id"] for d in deployments] == ["deploy-3"] + assert await _candidate_ids(router, "other-model") == ["deploy-3"] def test_update_settings_group_change_invalidates_model_group_info(): @@ -1587,29 +569,19 @@ def test_group_rows_cache_invalidated_on_model_list_change(): assert sum(1 for row in router.get_model_list() if row["model_name"] == "quality") == 4 -def _pin_choice_to(deployment_id): - def _pick(seq): - for candidate in seq: - if candidate["model_info"]["id"] == deployment_id: - return candidate - return seq[0] - - return _pick - - async def _call_and_get_cooldowns(router, model): - from litellm.router_utils.cooldown_handlers import async_get_cooldown_deployments - - with ( - patch("litellm.router_strategy.simple_shuffle.random.choice", side_effect=_pin_choice_to("deploy-3")), - pytest.raises(litellm.RateLimitError), - ): + """Cool down the other members first, so the call has to land on deploy-3.""" + for model_id in ("deploy-1", "deploy-2"): + router.cooldown_cache.add_deployment_to_cooldown( + model_id=model_id, original_exception=RuntimeError("unavailable"), exception_status=500, cooldown_time=60 + ) + with pytest.raises(litellm.RateLimitError): await router.acompletion( model=model, messages=[{"role": "user", "content": "hi"}], mock_response="litellm.RateLimitError", ) - return await async_get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) + return [model_id for model_id, _ in await router.cooldown_cache.async_get_active_cooldowns(["deploy-3"], None)] @pytest.mark.asyncio @@ -1668,15 +640,6 @@ def test_group_rows_do_not_inherit_member_access_groups(): assert "access_groups" in router.get_model_list(model_name="gated-member")[0]["model_info"] -def test_group_rebuild_invalidates_access_groups_cache(): - router = _build_router(routing_groups=_quality_group()) - router.get_model_access_groups() - assert router._access_groups_cache is not None - - router.update_settings(routing_groups=[]) - assert router._access_groups_cache is None - - def test_get_model_list_from_routing_groups_materializes_rows(): router = _build_router(routing_groups=_quality_group()) rows = router.get_model_list_from_routing_groups() @@ -1688,28 +651,6 @@ def test_get_model_list_from_routing_groups_materializes_rows(): assert router.get_model_list_from_routing_groups(model_name="filtered-model") == () -def test_get_routing_group_deployments_unions_members(): - router = _build_router(routing_groups=_quality_group()) - union = router._get_routing_group_deployments("quality") - assert sorted(d["model_info"]["id"] for d in union) == ["deploy-1", "deploy-2", "deploy-3"] - assert router._get_routing_group_deployments("filtered-model") is None - - -def test_materialize_routing_group_rows_labels_members_with_group_name(): - router = _build_router(routing_groups=_quality_group()) - group = router.get_routing_group("quality") - rows = router._materialize_routing_group_rows((group,)) - assert {row["model_name"] for row in rows} == {"quality"} - assert len(rows) == 3 - - -def test_as_routing_group_row_strips_access_groups(): - source = {"model_name": "member", "model_info": {"id": "d1", "access_groups": ["restricted"]}} - row = Router._as_routing_group_row(source) - assert row["model_info"] == {"id": "d1"} - assert source["model_info"]["access_groups"] == ["restricted"] - - @pytest.mark.asyncio async def test_group_call_429_cools_down_member_across_retries(): router = Router( @@ -1755,30 +696,6 @@ def _priority_deployments( ] -def test_priority_group_affinity_scope_follows_group_aliases_and_settings_reload() -> None: - router: Final = Router( - model_list=_priority_deployments(), - routing_groups=[ - _priority_group(), - RoutingGroup(group_name="legacy", models=["filtered-model"], routing_strategy="simple-shuffle"), - ], - model_group_alias={"priority-alias": "priority-group"}, - ) - requested_models: Final = ( - "priority-group", "priority-alias", "filtered-model", "other-model", "legacy", "missing" - ) - assert tuple(router._is_priority_routing_group(model) for model in requested_models) == ( - True, True, False, False, False, False - ) - - router.update_settings(routing_groups=[{ - "group_name": "priority-group", - "models": ["filtered-model", "other-model"], - "routing_strategy": "simple-shuffle", - }]) - assert tuple(router._is_priority_routing_group(model) for model in requested_models) == (False,) * 6 - - @pytest.mark.parametrize("asynchronous", [False, True]) @pytest.mark.parametrize("requested_model", ["priority-group", "priority-alias"]) @pytest.mark.asyncio @@ -1837,8 +754,6 @@ async def test_opposite_priority_groups_preserve_deployments_and_legacy_member_p model="reverse-group", messages=[{"role": "user", "content": "hi"}] ) assert (forward.choices[0].message.content, reverse.choices[0].message.content) == ("primary", "backup") - assert router._get_routing_context("filtered-model")[0] == "least-busy" - assert router._get_routing_context("other-model")[0] == "latency-based-routing" assert router.get_model_list(model_name="filtered-model") == before assert all(deployment["litellm_params"]["order"] == 10 for deployment in before) @@ -2002,67 +917,6 @@ async def test_auto_router_selected_priority_group_fails_over_inside_the_group() assert response.choices[0].message.content == "backup" -@pytest.mark.parametrize("endpoint_request", [{"input": "hi"}, {"messages": [{"role": "user", "content": "hi"}]}]) -@pytest.mark.asyncio -async def test_priority_generic_fallback_keeps_routing_controls_out_of_provider_kwargs( - endpoint_request: dict[str, object], -) -> None: - async def provider(model: str, **provider_kwargs: object) -> str: - assert "_target_order" not in provider_kwargs - assert "model_priorities" not in provider_kwargs - model_info: Final = provider_kwargs["model_info"] - assert isinstance(model_info, dict) - if model_info["id"] != "deploy-3": - raise litellm.RateLimitError(message="primary refused", model=model, llm_provider="openai") - return "backup" - - router: Final = Router( - model_list=_priority_deployments(), - routing_groups=[_priority_group()], - num_retries=0, - disable_cooldowns=True, - ) - response: Final = await router._ageneric_api_call_with_fallbacks( - model="priority-group", original_function=provider, **endpoint_request - ) - assert response == "backup" - - -@pytest.mark.parametrize("provider_model_id", [None, "resolved-downstream-id"]) -@pytest.mark.asyncio -async def test_priority_generic_fallback_dict_preserves_served_model_id(provider_model_id: str | None) -> None: - async def provider(model: str, **provider_kwargs: object) -> dict[str, object]: - model_info: Final = provider_kwargs["model_info"] - assert isinstance(model_info, dict) - if model_info["id"] == "deploy-1": - raise litellm.NotFoundError(message="primary missing", model=model, llm_provider="openai") - assert model_info["id"] == "deploy-3" - return { - "content": "backup", - **({"_hidden_params": {"model_id": provider_model_id}} if provider_model_id is not None else {}), - } - - router: Final = Router( - model_list=_priority_deployments()[::2], - routing_groups=[_priority_group()], - num_retries=0, - disable_cooldowns=True, - ) - outer_metadata: Final[dict[str, object]] = {} - response: Final = await router._ageneric_api_call_with_fallbacks( - model="priority-group", - original_function=provider, - messages=[{"role": "user", "content": "hi"}], - litellm_metadata=outer_metadata, - ) - primary_info: Final = outer_metadata["model_info"] - assert isinstance(primary_info, dict) - assert primary_info["id"] == "deploy-1" - assert response["content"] == "backup" - assert response["_hidden_params"]["model_id"] == (provider_model_id or "deploy-3") - assert response["_hidden_params"]["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 - - @pytest.mark.asyncio async def test_priority_group_ignores_cached_backup_affinity() -> None: from litellm.router_utils.prompt_caching_cache import PromptCachingCache diff --git a/tests/unit/router_strategy/test_router_routing_plugins.py b/tests/unit/router_strategy/test_router_routing_plugins.py index 0a54addce5e..b04c5edd3c2 100644 --- a/tests/unit/router_strategy/test_router_routing_plugins.py +++ b/tests/unit/router_strategy/test_router_routing_plugins.py @@ -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 ``) -- 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 diff --git a/tests/unit/router_strategy/test_router_tag_routing.py b/tests/unit/router_strategy/test_router_tag_routing.py index 0c19cfc72f6..0d3876c64a3 100644 --- a/tests/unit/router_strategy/test_router_tag_routing.py +++ b/tests/unit/router_strategy/test_router_tag_routing.py @@ -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 diff --git a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py index fa75f980cc0..191f8463cfa 100644 --- a/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -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 diff --git a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py index a4d279e32e6..76a574bae24 100644 --- a/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py +++ b/tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py @@ -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(): diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 6b4ea74894a..eb4e72cb9e6 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -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 diff --git a/tests/unit/router_utils/test_client_initalization_utils.py b/tests/unit/router_utils/test_client_initalization_utils.py index 6f9a7b730ac..30791f47623 100644 --- a/tests/unit/router_utils/test_client_initalization_utils.py +++ b/tests/unit/router_utils/test_client_initalization_utils.py @@ -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 diff --git a/tests/unit/router_utils/test_cooldown_handlers.py b/tests/unit/router_utils/test_cooldown_handlers.py index 71c37b5b1d5..8ef683c0aee 100644 --- a/tests/unit/router_utils/test_cooldown_handlers.py +++ b/tests/unit/router_utils/test_cooldown_handlers.py @@ -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": {}}, diff --git a/tests/unit/router_utils/test_fallback_event_handlers.py b/tests/unit/router_utils/test_fallback_event_handlers.py index a386542d61b..6263292fb40 100644 --- a/tests/unit/router_utils/test_fallback_event_handlers.py +++ b/tests/unit/router_utils/test_fallback_event_handlers.py @@ -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")) diff --git a/tests/unit/router_utils/test_health_check_allowed_fails_integration.py b/tests/unit/router_utils/test_health_check_allowed_fails_integration.py index 74b3a8825b3..a064c7346b4 100644 --- a/tests/unit/router_utils/test_health_check_allowed_fails_integration.py +++ b/tests/unit/router_utils/test_health_check_allowed_fails_integration.py @@ -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: diff --git a/tests/unit/router_utils/test_kubernetes_pod_discovery.py b/tests/unit/router_utils/test_kubernetes_pod_discovery.py index 0b772b09f6f..88b840e5c2a 100644 --- a/tests/unit/router_utils/test_kubernetes_pod_discovery.py +++ b/tests/unit/router_utils/test_kubernetes_pod_discovery.py @@ -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( diff --git a/tests/unit/router_utils/test_router_health_check_routing.py b/tests/unit/router_utils/test_router_health_check_routing.py index b1498ea3687..79ca88120e5 100644 --- a/tests/unit/router_utils/test_router_health_check_routing.py +++ b/tests/unit/router_utils/test_router_health_check_routing.py @@ -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: diff --git a/tests/unit/router_utils/test_router_interactions_endpoints.py b/tests/unit/router_utils/test_router_interactions_endpoints.py index 91bea170458..bd83990fd1c 100644 --- a/tests/unit/router_utils/test_router_interactions_endpoints.py +++ b/tests/unit/router_utils/test_router_interactions_endpoints.py @@ -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 diff --git a/tests/unit/router_utils/test_router_utils_common_utils.py b/tests/unit/router_utils/test_router_utils_common_utils.py index 6ec72868b6e..0de10a64f2e 100644 --- a/tests/unit/router_utils/test_router_utils_common_utils.py +++ b/tests/unit/router_utils/test_router_utils_common_utils.py @@ -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( diff --git a/tests/unit/test_conftest_isolation.py b/tests/unit/test_conftest_isolation.py index 3fdc2d65da1..e9a153a3f8d 100644 --- a/tests/unit/test_conftest_isolation.py +++ b/tests/unit/test_conftest_isolation.py @@ -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: diff --git a/tests/unit/test_private_usage_aliases.py b/tests/unit/test_private_usage_aliases.py index a3168ad6d36..d257866b654 100644 --- a/tests/unit/test_private_usage_aliases.py +++ b/tests/unit/test_private_usage_aliases.py @@ -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), ( diff --git a/tests/unit/test_redact_string_in_error_paths.py b/tests/unit/test_redact_string_in_error_paths.py index eaf2d43a25f..8727df23961 100644 --- a/tests/unit/test_redact_string_in_error_paths.py +++ b/tests/unit/test_redact_string_in_error_paths.py @@ -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( diff --git a/tests/unit/test_responses_streaming_container_ownership.py b/tests/unit/test_responses_streaming_container_ownership.py index 65889052bce..737a6eaafcd 100644 --- a/tests/unit/test_responses_streaming_container_ownership.py +++ b/tests/unit/test_responses_streaming_container_ownership.py @@ -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//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 diff --git a/tests/unit/test_router/core/conftest.py b/tests/unit/test_router/core/conftest.py index 3477cfc0c3d..7475071c382 100644 --- a/tests/unit/test_router/core/conftest.py +++ b/tests/unit/test_router/core/conftest.py @@ -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)) diff --git a/tests/unit/test_router/core/scripted_host.py b/tests/unit/test_router/core/scripted_host.py index f673ab98610..33164bd4408 100644 --- a/tests/unit/test_router/core/scripted_host.py +++ b/tests/unit/test_router/core/scripted_host.py @@ -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", diff --git a/tests/unit/test_router/core/test_construction_precall_characterization.py b/tests/unit/test_router/core/test_construction_precall_characterization.py index a56cd78ccac..2985819cb5b 100644 --- a/tests/unit/test_router/core/test_construction_precall_characterization.py +++ b/tests/unit/test_router/core/test_construction_precall_characterization.py @@ -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 ---------------------------------------------------------------------- diff --git a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py index 96d4d1879ce..9cdb3f55d17 100644 --- a/tests/unit/test_router/core/test_construction_precall_characterization_precall.py +++ b/tests/unit/test_router/core/test_construction_precall_characterization_precall.py @@ -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, diff --git a/tests/unit/test_router/core/test_endpoints.py b/tests/unit/test_router/core/test_endpoints.py index ce56ae403ba..d129dc9f408 100644 --- a/tests/unit/test_router/core/test_endpoints.py +++ b/tests/unit/test_router/core/test_endpoints.py @@ -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): diff --git a/tests/unit/test_router/core/test_legacy_members.py b/tests/unit/test_router/core/test_legacy_members.py index 1158a44e7fd..213fd5b84ba 100644 --- a/tests/unit/test_router/core/test_legacy_members.py +++ b/tests/unit/test_router/core/test_legacy_members.py @@ -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 ---------------------------------------------------------------------- diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index 7c2658d3d0d..f09ad5973e2 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -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: diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index a21777fbee6..dbc47fd6962 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -31,8 +31,7 @@ from pydantic import BaseModel, TypeAdapter import litellm from litellm import APIConnectionError, Router from litellm._logging import verbose_logger, verbose_router_logger -from litellm.caching import RedisCache, RedisClusterCache -from litellm.caching.caching import DualCache +from litellm.caching import RedisClusterCache from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache from litellm.caching.redis_cache import _redis_circuit_breaker_guard from litellm.exceptions import GuardrailRaisedException, MidStreamFallbackError, ModifyResponseException @@ -49,7 +48,7 @@ from litellm.llms.anthropic.messages.native.agentic_streaming_iterator import ( from litellm.llms.bedrock.common_utils import BedrockError from litellm.models.access_group import LiteLLM_AccessGroupTable from litellm.proxy._types import LiteLLM_TeamTable, LitellmUserRoles, Member, ProxyException, UserAPIKeyAuth -from litellm.router import ( +from litellm.router_rust.invoke.streaming import ( MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS, MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS, FallbackAwareAnthropicMessagesStream, @@ -61,9 +60,8 @@ from litellm.router import ( _anthropic_stream_should_decline_fallback, _is_retriable_anthropic_status, _responses_stream_holds_event, - _without_line_breaks, - _search_tool_support, -) +) # noqa: PLC2701 +from litellm.router_rust.surface.registry import _search_tool_support # noqa: PLC2701 from litellm.proxy.response_polling.polling_handler import should_use_polling_for_request from litellm.router_strategy import simple_shuffle from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter @@ -98,7 +96,7 @@ from litellm.types.router import ( RoutingContext, ) from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.llms.base_llm.vector_store.transformation import( +from litellm.llms.base_llm.vector_store.transformation import ( LiteLLMVectorStoreEmbeddingExecutor, RouterVectorStoreEmbeddingExecutor, ) @@ -248,146 +246,6 @@ async def test_chat_on_a_mixed_group_still_reaches_the_leg_without_decisions_sup assert response._hidden_params["model_id"] == "bedrock-leg" -def test_update_kwargs_does_not_mutate_defaults_and_merges_metadata(): - # initialize a real Router (env‑vars can be empty) - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "azure/gpt-4.1-mini", - "api_key": os.getenv("AZURE_AI_API_KEY"), - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), - }, - } - ], - ) - - # override to known defaults for the test - router.default_litellm_params = { - "foo": "bar", - "metadata": {"baz": 123}, - } - original = copy.deepcopy(router.default_litellm_params) - kwargs: dict = {} - - # invoke the helper - router._update_kwargs_with_default_litellm_params( - kwargs=kwargs, - metadata_variable_name="litellm_metadata", - ) - - # 1) router.defaults must be unchanged - assert router.default_litellm_params == original - - # 2) non‑metadata keys get merged - assert kwargs["foo"] == "bar" - - # 3) metadata lands under "metadata" - assert kwargs["litellm_metadata"] == {"baz": 123} - - -def test_router_model_group_llm_provider_resolves_chatgpt_without_login(tmp_path, monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) - params: Final = LiteLLM_Params(model="chatgpt/gpt-5.4") - - assert Router._model_group_llm_provider(params) == ("gpt-5.4", "chatgpt") - assert Router._model_group_llm_provider(LiteLLM_Params(model="no-such-provider-model-xyz")) == ("", "") - assert not any(tmp_path.iterdir()) - - -def test_router_with_model_info_and_model_group(): - """ - Test edge case where user specifies model_group in model_info - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - }, - "model_info": { - "tpm": 1000, - "rpm": 1000, - "model_group": "gpt-3.5-turbo", - }, - } - ], - ) - - router._set_model_group_info( - model_group="gpt-3.5-turbo", - user_facing_model_group_name="gpt-3.5-turbo", - ) - - -def test_router_model_group_encrypted_content_affinity_callback_registration(): - from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( - DeploymentAffinityCheck, - ) - from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import ( - EncryptedContentAffinityCheck, - ) - - model_group = "openai.gpt-5.1-codex" - model_group_affinity_config = { - model_group: ["encrypted_content_affinity"], - } - original_callbacks = list(litellm.callbacks) - litellm.callbacks = [] - router = None - - try: - router = litellm.Router( - model_list=[ - { - "model_name": model_group, - "litellm_params": { - "model": "openai/gpt-5.1-codex", - "api_key": "mock-api-key", - }, - } - ], - model_group_affinity_config=model_group_affinity_config, - num_retries=0, - ) - 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) - ) - 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 - ) - assert litellm.callbacks.index(encrypted_content_callbacks[0]) < ( - litellm.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 - finally: - if router is not None: - router.discard() - litellm.callbacks = original_callbacks - - @pytest.mark.asyncio async def test_encrypted_content_affinity_model_group_config_is_additive(): from litellm.responses.utils import ResponsesAPIRequestUtils @@ -409,13 +267,9 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): }, target_deployment, ] - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") - assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled( - {model_group: ["encrypted_content_affinity"]} - ) + assert EncryptedContentAffinityCheck.has_model_group_affinity_enabled({model_group: ["encrypted_content_affinity"]}) assert not EncryptedContentAffinityCheck.has_model_group_affinity_enabled(None) per_group_check = EncryptedContentAffinityCheck( @@ -456,10 +310,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert unfiltered == healthy_deployments - assert ( - "encrypted_content_affinity_enabled" - not in disabled_request_kwargs["litellm_metadata"] - ) + assert "encrypted_content_affinity_enabled" not in disabled_request_kwargs["litellm_metadata"] global_check = EncryptedContentAffinityCheck( enable_global_affinity=True, @@ -479,9 +330,7 @@ async def test_encrypted_content_affinity_model_group_config_is_additive(): ) assert globally_filtered == [target_deployment] - assert global_request_kwargs["litellm_metadata"][ - "encrypted_content_affinity_enabled" - ] + assert global_request_kwargs["litellm_metadata"]["encrypted_content_affinity_enabled"] @pytest.mark.asyncio @@ -528,18 +377,10 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( num_retries=0, ) callbacks = router.optional_callbacks or [] - deployment_callback = next( - cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck) - ) - encrypted_content_callback = next( - cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck) - ) - assert callbacks.index(encrypted_content_callback) < callbacks.index( - deployment_callback - ) - assert litellm.callbacks.index(encrypted_content_callback) < ( - litellm.callbacks.index(deployment_callback) - ) + deployment_callback = next(cb for cb in callbacks if isinstance(cb, DeploymentAffinityCheck)) + encrypted_content_callback = next(cb for cb in callbacks if isinstance(cb, EncryptedContentAffinityCheck)) + assert callbacks.index(encrypted_content_callback) < callbacks.index(deployment_callback) + assert litellm.callbacks.index(encrypted_content_callback) < (litellm.callbacks.index(deployment_callback)) cache_key = DeploymentAffinityCheck.get_affinity_cache_key( model_group=model_group, @@ -550,9 +391,7 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity( value={"model_id": "deployment-a"}, ttl=60, ) - encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id( - "deployment-b", "rs_test" - ) + encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test") request_kwargs = { "input": [{"type": "reasoning", "id": encoded_id}], "litellm_metadata": {"user_api_key_hash": user_api_key_hash}, @@ -762,7 +601,9 @@ async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards from io import BytesIO from unittest.mock import MagicMock, patch - jsonl_content = b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + jsonl_content = ( + b'{"custom_id": "r1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "vertex-batch"}}\n' + ) router = litellm.Router( model_list=[ { @@ -773,9 +614,7 @@ async def test_async_router_acreate_file_passthrough_keeps_the_file_and_forwards ) with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: - await router.acreate_file( - model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True - ) + await router.acreate_file(model="vertex-batch", purpose="batch", file=BytesIO(jsonl_content), passthrough=True) forwarded = mock_acreate_file.call_args.kwargs assert forwarded["passthrough"] is True forwarded["file"].seek(0) @@ -1251,9 +1090,7 @@ async def test_arouter_aretrieve_batch(): ], ) - with patch.object( - litellm, "aretrieve_batch", return_value=AsyncMock() - ) as mock_aretrieve_batch: + with patch.object(litellm, "aretrieve_batch", return_value=AsyncMock()) as mock_aretrieve_batch: try: response = await router.aretrieve_batch( model="gpt-3.5-turbo", @@ -1513,6 +1350,7 @@ def test_sync_deployment_callback_on_success_skips_batch_retrieves( == expected_successes ) + _ROUTING_STRATEGY_CACHE_MARKERS = ("_map", "_request_count", ":tpm:", ":rpm:") @@ -1524,8 +1362,7 @@ async def _moved_routing_counters(router, timeout: float = 2.0) -> list[str]: moved = sorted( f"{key}={cache_dict[key]}" for key in cache_dict - if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) - and cache_dict[key] + if any(marker in key for marker in _ROUTING_STRATEGY_CACHE_MARKERS) and cache_dict[key] ) if moved: return moved @@ -1606,9 +1443,7 @@ async def test_arouter_aretrieve_file_content(): Test that router.acreate_file with JSONL file returns the correct response """ - with patch.object( - litellm, "afile_content", return_value=AsyncMock() - ) as mock_afile_content: + with patch.object(litellm, "afile_content", return_value=AsyncMock()) as mock_afile_content: router = litellm.Router( model_list=[ { @@ -1669,7 +1504,7 @@ async def test_arouter_filter_team_based_models(): assert result is not None # FAILS - with pytest.raises(Exception, match='No deployments available for selected model, Try again in') as e: + with pytest.raises(Exception, match="No deployments available for selected model, Try again in") as e: result = await router.acompletion( model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hello, world!"}], @@ -1753,9 +1588,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team_id and team_public_model_name match" + assert result is True, "Should return True when team_id and team_public_model_name match" # Test Case 2: Team-specific deployment - team_id matches but model_name doesn't match team_public_model_name result = router.should_include_deployment( @@ -1763,9 +1596,9 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id="test-team", ) - assert ( - result is False - ), "Should return False when team_id matches but model_name doesn't match team_public_model_name" + assert result is False, ( + "Should return False when team_id matches but model_name doesn't match team_public_model_name" + ) # Test Case 3: Team-specific deployment - team_id doesn't match result = router.should_include_deployment( @@ -1781,30 +1614,18 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_no_public_name, team_id="test-team", ) - assert ( - result is True - ), "Should return True when team deployment has no team_public_model_name to match" + assert result is True, "Should return True when team deployment has no team_public_model_name to match" # Test Case 5: Non-team deployment - model_name matches and no team_id - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id=None - ) - assert ( - result is True - ), "Should return True when model_name matches and deployment has no team_id" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id=None) + assert result is True, "Should return True when model_name matches and deployment has no team_id" # Test Case 6: Non-team deployment - model_name matches but team_id provided (should still work) - result = router.should_include_deployment( - model_name="gpt-4", model=deployment_without_team, team_id="any-team" - ) - assert ( - result is True - ), "Should return True when model_name matches non-team deployment, regardless of team_id param" + result = router.should_include_deployment(model_name="gpt-4", model=deployment_without_team, team_id="any-team") + assert result is True, "Should return True when model_name matches non-team deployment, regardless of team_id param" # Test Case 7: Non-team deployment - model_name doesn't match - result = router.should_include_deployment( - model_name="different-model", model=deployment_without_team, team_id=None - ) + result = router.should_include_deployment(model_name="different-model", model=deployment_without_team, team_id=None) assert result is False, "Should return False when model_name doesn't match" # Test Case 8: Team deployment accessed without matching team_id @@ -1813,9 +1634,7 @@ def test_arouter_should_include_deployment(): model=deployment_with_team_and_public_name, team_id=None, ) - assert ( - result is True - ), "Should return True when matching model with exact model_name" + assert result is True, "Should return True when matching model with exact model_name" def test_arouter_responses_api_bridge(): @@ -1866,9 +1685,7 @@ def test_arouter_responses_api_bridge(): "status": "completed", "output": [], } - mock_response.text = ( - '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' - ) + mock_response.text = '{"id": "resp_test", "object": "response", "status": "completed", "output": []}' with patch.object(client, "post", return_value=mock_response) as mock_post: try: @@ -1942,7 +1759,7 @@ def test_add_invalid_provider_to_router(): ], ) - with pytest.raises(Exception, match='Unsupported provider - vertex_ai_eu') as e: + with pytest.raises(Exception, match="Unsupported provider - vertex_ai_eu") as e: router.add_deployment( Deployment( model_name="vertex_ai/*", @@ -1967,7 +1784,9 @@ def registered_custom_provider(monkeypatch: pytest.MonkeyPatch) -> str: model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], mock_response="served by onprem handler" ) - monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}]) + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}] + ) monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list)) monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers)) return "test-onprem-llm" @@ -2006,327 +1825,6 @@ def test_router_add_deployment_accepts_explicit_custom_provider_from_custom_prov assert router.get_model_list(model_name="onprem")[0]["litellm_params"]["model"] == "my-model" -@pytest.mark.asyncio -async def test_router_ageneric_api_call_with_fallbacks_helper(): - """ - Test the _ageneric_api_call_with_fallbacks_helper method with various scenarios - """ - from unittest.mock import patch - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": "test-key", - "api_base": "https://api.openai.com/v1", - }, - "model_info": { - "tpm": 1000, - "rpm": 1000, - }, - }, - ], - ) - - # Test 1: Successful call - async def mock_generic_function(**kwargs): - return {"result": "success", "model": kwargs.get("model")} - - with patch.object(router, "async_get_available_deployment") as mock_get_deployment: - mock_get_deployment.return_value = { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": "test-key", - "api_base": "https://api.openai.com/v1", - }, - } - - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: - result = await router._ageneric_api_call_with_fallbacks_helper( - model="gpt-3.5-turbo", - original_generic_function=mock_generic_function, - messages=[{"role": "user", "content": "test"}], - ) - - assert result is not None - assert result["result"] == "success" - mock_get_deployment.assert_called_once() - mock_update_kwargs.assert_called_once() - mock_pre_call_checks.assert_called_once() - - # Test 2: Passthrough on no deployment (success case) - async def mock_passthrough_function(**kwargs): - return {"result": "passthrough", "model": kwargs.get("model")} - - with patch.object(router, "async_get_available_deployment") as mock_get_deployment: - mock_get_deployment.side_effect = Exception("No deployment available") - - result = await router._ageneric_api_call_with_fallbacks_helper( - model="gpt-3.5-turbo", - original_generic_function=mock_passthrough_function, - passthrough_on_no_deployment=True, - messages=[{"role": "user", "content": "test"}], - ) - - assert result is not None - assert result["result"] == "passthrough" - assert result["model"] == "gpt-3.5-turbo" - - # Test 3: No deployment available and passthrough=False (should raise exception) - with patch.object(router, "async_get_available_deployment") as mock_get_deployment: - mock_get_deployment.side_effect = Exception("No deployment available") - - with pytest.raises(Exception, match='No deployment available') as exc_info: - await router._ageneric_api_call_with_fallbacks_helper( - model="gpt-3.5-turbo", - original_generic_function=mock_generic_function, - passthrough_on_no_deployment=False, - messages=[{"role": "user", "content": "test"}], - ) - - assert "No deployment available" in str(exc_info.value) - - # Test 4: Test with semaphore (rate limiting) - import asyncio - - async def mock_semaphore_function(**kwargs): - return {"result": "semaphore_success", "model": kwargs.get("model")} - - with patch.object(router, "async_get_available_deployment") as mock_get_deployment: - mock_get_deployment.return_value = { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": "test-key", - "api_base": "https://api.openai.com/v1", - }, - } - - mock_semaphore = MaxParallelRequestsLimit( - max_parallel_requests=1, model_id="deployment-1", model_group="gpt-3.5-turbo" - ) - - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=mock_semaphore - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - result = await router._ageneric_api_call_with_fallbacks_helper( - model="gpt-3.5-turbo", - original_generic_function=mock_semaphore_function, - messages=[{"role": "user", "content": "test"}], - ) - - assert result is not None - assert result["result"] == "semaphore_success" - mock_get_client.assert_called_once() - mock_pre_call_checks.assert_called_once() - - # Test 5: Test call tracking (success and failure counts) - initial_success_count = router.success_calls.get("gpt-3.5-turbo", 0) - initial_fail_count = router.fail_calls.get("gpt-3.5-turbo", 0) - - async def mock_failing_function(**kwargs): - raise Exception("Mock failure") - - with patch.object(router, "async_get_available_deployment") as mock_get_deployment: - mock_get_deployment.return_value = { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - "api_key": "test-key", - "api_base": "https://api.openai.com/v1", - }, - } - - with patch.object( - router, "_update_kwargs_with_deployment" - ) as mock_update_kwargs: - with patch.object( - router, "_get_client", return_value=None - ) as mock_get_client: - with patch.object( - router, "async_routing_strategy_pre_call_checks" - ) as mock_pre_call_checks: - with pytest.raises(Exception, match='Mock failure') as exc_info: - await router._ageneric_api_call_with_fallbacks_helper( - model="gpt-3.5-turbo", - original_generic_function=mock_failing_function, - messages=[{"role": "user", "content": "test"}], - ) - - assert "Mock failure" in str(exc_info.value) - # Check that fail_calls was incremented - assert router.fail_calls["gpt-3.5-turbo"] == initial_fail_count + 1 - - -@pytest.mark.asyncio -async def test_ageneric_api_call_deployment_model_overrides_alias(): - """ - Regression: when a model alias (e.g. "not-gemini-2.5-flash") maps to a deployment - with model="vertex_ai/gemini-2.5-flash", the underlying litellm function must receive - the deployment model, not the alias. Before the fix, **kwargs overwrote data["model"]. - """ - from unittest.mock import patch - - captured: dict = {} - - async def capture_model(**kwargs): - captured["model"] = kwargs.get("model") - return {"result": "ok"} - - router = litellm.Router( - model_list=[ - { - "model_name": "not-gemini-2.5-flash", - "litellm_params": { - "model": "vertex_ai/gemini-2.5-flash", - "api_key": "fake-key", - }, - } - ] - ) - - def inject_alias_into_kwargs( - deployment: dict, - kwargs: dict, - function_name: str | None = None, - splits_stream_timeout: bool = False, - opens_responses_socket: bool = False, - ) -> None: - # Simulate the alias leaking into kwargs (as happens when - # _ageneric_api_call_with_fallbacks sets kwargs["model"] = alias before - # calling the helper through async_function_with_fallbacks). - kwargs["model"] = "not-gemini-2.5-flash" - - with ( - patch.object(router, "async_get_available_deployment") as mock_dep, - patch.object( - router, - "_update_kwargs_with_deployment", - side_effect=inject_alias_into_kwargs, - ), - patch.object(router, "async_routing_strategy_pre_call_checks"), - patch.object(router, "_get_client", return_value=None), - ): - mock_dep.return_value = { - "model_name": "not-gemini-2.5-flash", - "litellm_params": { - "model": "vertex_ai/gemini-2.5-flash", - "api_key": "fake-key", - }, - } - - await router._ageneric_api_call_with_fallbacks_helper( - model="not-gemini-2.5-flash", - original_generic_function=capture_model, - ) - - assert ( - captured["model"] == "vertex_ai/gemini-2.5-flash" - ), f"Expected deployment model 'vertex_ai/gemini-2.5-flash', got '{captured['model']}'" - - -@pytest.mark.asyncio -async def test_ageneric_api_call_resolves_realtime_session_model(): - """ - Regression for #36742: realtime client secret requests carry the model inside `session` too, and the proxy - fills it with the pre-routing model group name. The underlying litellm function reads session.model first, - so it must see the resolved deployment, while a caller's nested transcription model stays untouched. - """ - routed: Final = AsyncMock(return_value={"result": "ok"}) - - router = litellm.Router( - model_list=[ - { - "model_name": "my-realtime-group", - "litellm_params": { - "model": "openai/gpt-realtime-2.1-mini", - "api_key": "fake-key", - }, - "model_info": {"mode": "realtime"}, - } - ] - ) - - await router._ageneric_api_call_with_fallbacks( - model="my-realtime-group", - original_function=routed, - session={ - "type": "realtime", - "model": "my-realtime-group", - "audio": {"input": {"transcription": {"model": "gpt-4o-transcribe"}}}, - }, - ) - - sent: Final = routed.call_args.kwargs - assert sent["model"] == "openai/gpt-realtime-2.1-mini" - assert sent["session"]["model"] == "openai/gpt-realtime-2.1-mini" - assert sent["session"]["audio"]["input"]["transcription"]["model"] == "gpt-4o-transcribe" - - -@pytest.mark.asyncio -async def test_ageneric_api_call_does_not_add_session_model(): - """ - A session that never carried a model must not gain one from routing: the underlying function then falls back - to the resolved `model` kwarg itself, and the outgoing session body keeps the caller's shape. - """ - routed: Final = AsyncMock(return_value={"result": "ok"}) - - router = litellm.Router( - model_list=[ - { - "model_name": "my-realtime-group", - "litellm_params": { - "model": "openai/gpt-realtime-2.1-mini", - "api_key": "fake-key", - }, - "model_info": {"mode": "realtime"}, - } - ] - ) - - await router._ageneric_api_call_with_fallbacks( - model="my-realtime-group", - original_function=routed, - session={"type": "realtime"}, - ) - - sent: Final = routed.call_args.kwargs - assert sent["model"] == "openai/gpt-realtime-2.1-mini" - assert sent["session"] == {"type": "realtime"} - - -@pytest.mark.parametrize( - "session, expected", - [ - ({"type": "realtime", "model": "my-realtime-group"}, {"session": {"type": "realtime", "model": "resolved"}}), - ({"type": "realtime"}, {}), - (None, {}), - ("not-a-session", {}), - ], -) -def test_with_router_resolved_session_model(session, expected): - from litellm.router import _with_router_resolved_session_model - - assert dict(_with_router_resolved_session_model(session, "resolved")) == expected - - def test_router_get_model_access_groups_team_only_models(): """ Test that Router._core_get_model_access_groups returns the correct response for team-only models @@ -2345,147 +1843,13 @@ def test_router_get_model_access_groups_team_only_models(): ] ) - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id=None - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id=None) assert len(access_groups) == 0 - access_groups = router.get_model_access_groups( - model_name="gpt-3.5-turbo", team_id="team_1" - ) + access_groups = router.get_model_access_groups(model_name="gpt-3.5-turbo", team_id="team_1") assert list(access_groups.keys()) == ["default-models"] -def test_cached_get_model_group_info(): - """ - Test that _cached_get_model_group_info caches results and - invalidates on deployment changes. - """ - from litellm.types.router import Deployment, LiteLLM_Params - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake"}, - "model_info": {"tpm": 1000, "rpm": 100}, - }, - ] - ) - - # First call should compute and cache - result1 = router._cached_get_model_group_info("gpt-4") - assert result1 is not None - assert result1.tpm == 1000 - - # Second call should hit cache (same object) - result2 = router._cached_get_model_group_info("gpt-4") - assert result1 is result2 - - # Add a deployment — cache should be invalidated - router.add_deployment( - Deployment( - model_name="gpt-4", - litellm_params=LiteLLM_Params(model="gpt-4", api_key="fake2"), - model_info={"tpm": 2000, "rpm": 200}, - ) - ) - result3 = router._cached_get_model_group_info("gpt-4") - assert result3 is not result2 - assert result3 is not None - assert result3.tpm == 3000 # 1000 + 2000 - - # Delete a deployment — cache should be invalidated - deployment_id = router.model_list[-1]["model_info"]["id"] - router.delete_deployment(id=deployment_id) - result4 = router._cached_get_model_group_info("gpt-4") - assert result4 is not result3 - assert result4 is not None - assert result4.tpm == 1000 - - # set_model_list — cache should be invalidated - router.set_model_list( - [ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake"}, - "model_info": {"tpm": 5000}, - }, - ] - ) - result5 = router._cached_get_model_group_info("gpt-4") - assert result5 is not result4 - assert result5 is not None - assert result5.tpm == 5000 - - # Verify cache still works after invalidation - result6 = router._cached_get_model_group_info("gpt-4") - assert result5 is result6 - - -def test_model_group_info_cost_from_db_model_info(): - """ - When get_deployment_model_info fails (model_info is None fallback), - input_cost_per_token and output_cost_per_token should be read from db model_info. - """ - from unittest.mock import patch - - router = litellm.Router( - model_list=[ - { - "model_name": "my-custom-model", - "litellm_params": { - "model": "openai/my-custom-model", - "api_key": "fake", - "api_base": "https://my-custom-endpoint.com", - }, - "model_info": { - "input_cost_per_token": 0.0001, - "output_cost_per_token": 0.0002, - }, - }, - ] - ) - - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): - result = router._cached_get_model_group_info("my-custom-model") - assert result is not None - assert result.input_cost_per_token == 0.0001 - assert result.output_cost_per_token == 0.0002 - - -def test_model_group_info_cost_none_when_db_model_info_has_no_cost(): - """ - When get_deployment_model_info fails and db model_info has no cost fields, - input/output_cost_per_token should be None. - """ - from unittest.mock import patch - - router = litellm.Router( - model_list=[ - { - "model_name": "my-custom-model-no-cost", - "litellm_params": { - "model": "openai/my-custom-model-no-cost", - "api_key": "fake", - "api_base": "https://my-custom-endpoint.com", - }, - "model_info": {}, - }, - ] - ) - - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): - result = router._cached_get_model_group_info("my-custom-model-no-cost") - assert result is not None - assert result.input_cost_per_token is None - assert result.output_cost_per_token is None - - def test_model_group_info_cost_none_for_unpriced_deployment_but_zero_when_declared(): """A deployment with no cost fields anywhere must report None, not the 0 that get_model_info defaults to, so the reported price matches what the zero-cost @@ -2881,153 +2245,6 @@ def test_get_model_list_of_routed_group_falls_back_to_the_wildcard_route_serving assert _served_models(router.get_model_list(model_name="openai/gpt-4o")) == ["ollama/qwen3:0.6b"] -def test_switch_routing_strategy_installs_lar1_then_restores_the_default_selector(): - router = _alias_cost_router() - - router._switch_routing_strategy( - "lar1", - { - "routing_strategy_args": { - "confidence_threshold_low": 0.1, - "confidence_threshold_medium": 0.3, - "confidence_threshold_high": 0.9, - } - }, - ) - assert router.routing_strategy == "lar1" - assert "_core_async_get_available_deployment" in router._core.__dict__ - - router._switch_routing_strategy("usage-based-routing-v2", {}) - assert router.lowesttpm_logger_v2 is not None - assert "_core_async_get_available_deployment" not in router._core.__dict__ - - -@pytest.mark.parametrize( - "value,expected", - [ - ("1e-05", 1e-05), - ("0.00001", 1e-05), - (1e-05, 1e-05), - (5, 5.0), - (None, None), - ("not-a-number", None), - ], -) -def test_cost_value_as_float(value, expected): - from litellm.router import _cost_value_as_float - - assert _cost_value_as_float(value) == expected - - -def test_model_group_info_with_stringified_cost_values(): - """ - YAML 1.2 parsers emit '1e-05' (integer mantissa) as a string, so cost - values in deployment model_info can arrive as str. Aggregating the model - group must not raise TypeError('>' between str and float) and must return - float costs. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "my-custom-model", - "litellm_params": { - "model": "openai/my-custom-backend-1", - "api_key": "fake", - }, - "model_info": { - "input_cost_per_token": "1e-05", - "output_cost_per_token": "1e-05", - }, - }, - { - "model_name": "my-custom-model", - "litellm_params": { - "model": "openai/my-custom-backend-2", - "api_key": "fake", - }, - "model_info": { - "input_cost_per_token": "2e-05", - "output_cost_per_token": "2e-05", - }, - }, - ] - ) - - def _model_info_with_str_costs(model_id: str, model_name: str): - for model in router.model_list: - if model["model_info"]["id"] == model_id: - return { - "key": model_name, - "input_cost_per_token": model["model_info"]["input_cost_per_token"], - "output_cost_per_token": model["model_info"]["output_cost_per_token"], - "litellm_provider": "openai", - "mode": "chat", - } - return None - - with patch.object( - router, "_core_get_deployment_model_info", side_effect=_model_info_with_str_costs - ): - result = router._set_model_group_info( - model_group="my-custom-model", - user_facing_model_group_name="my-custom-model", - ) - - assert result is not None - assert result.input_cost_per_token == 2e-05 - assert result.output_cost_per_token == 2e-05 - assert isinstance(result.input_cost_per_token, float) - assert isinstance(result.output_cost_per_token, float) - - -def test_model_group_info_db_fallback_with_stringified_cost_values(): - """ - Fallback path: when get_deployment_model_info returns nothing, costs are - read straight from the deployment's model_info dict, which can hold - stringified floats parsed from YAML. They must be coerced to float. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "my-custom-model", - "litellm_params": { - "model": "openai/my-custom-backend-1", - "api_key": "fake", - }, - "model_info": { - "input_cost_per_token": "1e-05", - "output_cost_per_token": "3e-05", - }, - }, - { - "model_name": "my-custom-model", - "litellm_params": { - "model": "openai/my-custom-backend-2", - "api_key": "fake", - }, - "model_info": { - "input_cost_per_token": "2e-05", - "output_cost_per_token": "2e-05", - }, - }, - ] - ) - - with patch.object( - router, "_core_get_deployment_model_info", side_effect=Exception("not found") - ): - result = router._set_model_group_info( - model_group="my-custom-model", - user_facing_model_group_name="my-custom-model", - ) - - assert result is not None - assert result.input_cost_per_token == 2e-05 - assert result.output_cost_per_token == 3e-05 - assert isinstance(result.input_cost_per_token, float) - assert isinstance(result.output_cost_per_token, float) - - def test_get_model_access_groups_caching(): """ Test that get_model_access_groups caches the no-args result @@ -3154,502 +2371,12 @@ def test_get_model_access_groups_cache_invalidation_upsert_deployment(): assert "updated-group" in result2 -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator(): - """Test _acompletion_streaming_iterator for normal streaming and fallback behavior.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - - # Helper class for creating async iterators - class AsyncIterator: - def __init__(self, items, error_after=None): - self.items = items - self.index = 0 - self.error_after = error_after - - def __aiter__(self): - return self - - async def __anext__(self): - if self.error_after is not None and self.index >= self.error_after: - raise self.error_after - if self.index >= len(self.items): - raise StopAsyncIteration - item = self.items[self.index] - self.index += 1 - return item - - # Set up router with fallback configuration - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"}, - }, - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key-2"}, - }, - ], - fallbacks=[{"gpt-4": ["gpt-3.5-turbo"]}], - set_verbose=True, - ) - - # Test data - messages = [{"role": "user", "content": "Hello"}] - initial_kwargs = {"model": "gpt-4", "stream": True, "temperature": 0.7} - - # Test 1: Successful streaming (no errors) - print("\n=== Test 1: Successful streaming ===") - - # Mock successful streaming response - mock_chunks = [ - MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content=" there"))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content="!"))]), - ] - - mock_response = AsyncIterator(mock_chunks) - - setattr(mock_response, "model", "gpt-4") - setattr(mock_response, "custom_llm_provider", "openai") - setattr(mock_response, "logging_obj", MagicMock()) - - result = await router._acompletion_streaming_iterator( - model_response=mock_response, messages=messages, initial_kwargs=initial_kwargs - ) - - # Collect streamed chunks - collected_chunks = [] - async for chunk in result: - collected_chunks.append(chunk) - - assert len(collected_chunks) == 3 - assert all(chunk in mock_chunks for chunk in collected_chunks) - print("✓ Successfully streamed all chunks") - - # Test 2: MidStreamFallbackError with generated content is re-raised, not silently continued - print("\n=== Test 2: MidStreamFallbackError re-raises when content already generated ===") - - # Error with generated content and is_pre_first_chunk=False (the default): - # the router must re-raise instead of attempting a continuation-prompt fallback, - # because partial content has already been sent to the client. - error = MidStreamFallbackError( - message="Connection lost", - model="gpt-4", - llm_provider="openai", - generated_content="Hello", - ) - - class AsyncIteratorWithError: - def __init__(self, items, error_after_index): - self.items = items - self.index = 0 - self.error_after_index = error_after_index - self.chunks = [] - - def __aiter__(self): - return self - - async def __anext__(self): - if self.index >= len(self.items): - raise StopAsyncIteration - if self.index == self.error_after_index: - raise error - item = self.items[self.index] - self.index += 1 - return item - - mock_error_response = AsyncIteratorWithError(mock_chunks, 1) # Error after first chunk - - setattr(mock_error_response, "model", "gpt-4") - setattr(mock_error_response, "custom_llm_provider", "openai") - setattr(mock_error_response, "logging_obj", MagicMock()) - - result = await router._acompletion_streaming_iterator( - model_response=mock_error_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - # Collect streamed chunks — the first chunk succeeds, then the error re-raises - collected_chunks = [] - async def _drain(): - async for chunk in result: - collected_chunks.append(chunk) - - with pytest.raises(MidStreamFallbackError): - await _drain() - - assert len(collected_chunks) == 1, "one chunk yielded before the error" - print("✓ MidStreamFallbackError re-raised correctly when content was already generated") - - print("\n=== All tests passed! ===") - - -def _make_midstream_source(error, chunks=None, first_chunk=None): - from unittest.mock import MagicMock - - opening_chunk = ( - first_chunk if first_chunk is not None else MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]) - ) - - class _Source: - def __init__(self): - self.index = 0 - self.chunks = chunks if chunks is not None else [] - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - - def __aiter__(self): - return self - - async def __anext__(self): - if self.index == 0: - self.index += 1 - return opening_chunk - raise error - - async def wait_for_failure_logging(self, timeout): - return None - - return _Source() - - -class _FakeFallbackStream: - def __init__(self, item): - self._item = item - self._done = False - - def __aiter__(self): - return self - - async def __anext__(self): - if self._done: - raise StopAsyncIteration - self._done = True - return self._item - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_continues_after_content_when_eligible(): - """Flag on + plain-text break: the router re-enters the fallback chain with an - assistant-prefill continuation and the mid-stream marker, then streams the - fallback's output instead of re-raising.""" - from unittest.mock import AsyncMock, MagicMock, patch - - from litellm.exceptions import MidStreamFallbackError - from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( - MID_STREAM_CONTINUATION_KWARG, - MID_STREAM_CONTINUATION_MARKER, - ) - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - - error = MidStreamFallbackError( - message="Connection lost", - model="gpt-4", - llm_provider="openai", - generated_content="Hello", - is_pre_first_chunk=False, - emitted_disqualifying_content=False, - ) - source = _make_midstream_source(error) - fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "world"}}]) - initial_kwargs = {"model": "gpt-4", "stream": True} - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), - ) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs - ) - collected = [chunk async for chunk in result] - - mock_fallback.assert_awaited_once() - passed_kwargs = mock_fallback.await_args.kwargs["kwargs"] - assert passed_kwargs[MID_STREAM_CONTINUATION_KWARG] is MID_STREAM_CONTINUATION_MARKER - assert passed_kwargs["messages"][-1] == {"role": "assistant", "content": "Hello", "prefix": True} - assert fallback_chunk in collected - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "continuation_text,joined_output_trips_rule", - [("6789", True), ("6788", False)], - ids=["rule_trips_only_on_the_joined_text", "clean_continuation_streams"], -) -async def test_acompletion_streaming_iterator_continuation_runs_post_call_rules_on_the_joined_text( - monkeypatch, continuation_text, joined_output_trips_rule -): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - monkeypatch.setattr(litellm, "post_call_rules", [lambda output: "123-45-6789" not in output]) - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="123-45-", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": continuation_text}}]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), - ): - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error), - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - if joined_output_trips_rule: - with pytest.raises(litellm.APIResponseValidationError): - async for _ in result: - pass - else: - assert fallback_chunk in [chunk async for chunk in result] - - -@pytest.mark.asyncio -async def test_a_second_cut_runs_post_call_rules_only_on_the_text_the_caller_received(monkeypatch): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( - MID_STREAM_CONTINUATION_KWARG, - MID_STREAM_CONTINUATION_MARKER, - ) - - checked_texts = [] - monkeypatch.setattr(litellm, "post_call_rules", [lambda output: checked_texts.append(output) is None]) - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - - def cut(generated_content): - return MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=generated_content, - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - - messages = [{"role": "user", "content": "Hi"}] - middle_leg = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source( - cut(" there"), - first_chunk=litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " there"}}]), - ), - messages=[*messages, {"role": "assistant", "content": "Hello", "prefix": True}], - initial_kwargs={"model": "gpt-4", "stream": True, MID_STREAM_CONTINUATION_KWARG: MID_STREAM_CONTINUATION_MARKER}, - ) - backup_leg = _FakeFallbackStream( - litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " friend"}}]) - ) - - with patch.object( - router, "async_function_with_fallbacks_common_utils", new=AsyncMock(side_effect=[middle_leg, backup_leg]) - ): - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(cut("Hello")), - messages=messages, - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - _ = [chunk async for chunk in result] - - assert checked_texts == ["Hello there", "Hello there friend"] - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "fallback_texts", - [("Hel", "lo", " world"), (" world",)], - ids=["fallback_echoes_the_prefill", "fallback_continues_after_the_prefill"], -) -async def test_acompletion_streaming_iterator_continuation_never_repeats_the_partial_output(fallback_texts): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "mistral/mistral-small-latest", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - - async def fallback_stream(): - for text in fallback_texts: - yield litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) - - with patch.object( - router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) - ): - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error), - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - collected = [chunk async for chunk in result] - - continuation = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] - assert "".join(chunk.choices[0].delta.content for chunk in continuation) == " world" - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_continuation_stays_in_the_primary_stream(): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - primary_chunk = litellm.ModelResponseStream( - id="chatcmpl-primary", created=111, choices=[{"index": 0, "delta": {"role": "assistant", "content": "Hello"}}] - ) - - async def fallback_stream(): - yield litellm.ModelResponseStream( - id="msg-fallback", created=222, choices=[{"index": 0, "delta": {"role": "assistant", "content": " wor"}}] - ) - yield litellm.ModelResponseStream( - id="msg-fallback", created=222, choices=[{"index": 0, "delta": {"content": "ld"}}] - ) - - with patch.object( - router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) - ): - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error, chunks=[primary_chunk]), - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - collected = [chunk async for chunk in result] - - continuation = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] - assert [(chunk.id, chunk.created, chunk.choices[0].delta.role) for chunk in continuation] == [ - ("chatcmpl-primary", 111, None), - ("chatcmpl-primary", 111, None), - ] - assert "".join(chunk.choices[0].delta.content for chunk in continuation) == " world" - - -def _continuation_router(): - return litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - - -def _usage_chunk(prompt_tokens, completion_tokens): - from litellm.types.utils import Usage - - return litellm.ModelResponseStream( - model="gpt-4", - choices=[], - usage=Usage( - prompt_tokens=prompt_tokens, - completion_tokens=completion_tokens, - total_tokens=prompt_tokens + completion_tokens, - ), - ) - - -async def _drain_with_fallback(router, error, primary_chunks, messages, fallback_items): - from unittest.mock import AsyncMock, patch - - async def fallback_stream(): - for item in fallback_items: - yield item - - with patch.object( - router, "async_function_with_fallbacks_common_utils", new=AsyncMock(return_value=fallback_stream()) - ) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error, chunks=primary_chunks), - messages=messages, - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - collected = [chunk async for chunk in result] - streamed = [chunk for chunk in collected if isinstance(chunk, litellm.ModelResponseStream)] - return mock_fallback.await_args.kwargs["kwargs"], streamed - - -@pytest.mark.asyncio -@pytest.mark.parametrize("primary_reports_usage", [True, False], ids=["usage_reported", "usage_not_reported"]) -async def test_acompletion_streaming_iterator_continuation_usage_sums_both_legs(primary_reports_usage): - from litellm.exceptions import MidStreamFallbackError - - messages = [{"role": "user", "content": "Explain transformers in one paragraph."}] - partial = "Transformers use " - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - content_chunk = litellm.ModelResponseStream( - model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": partial}}] - ) - primary_chunks = [content_chunk, _usage_chunk(21, 3)] if primary_reports_usage else [content_chunk] - primary_usage = litellm.stream_chunk_builder(chunks=list(primary_chunks), messages=messages).usage - - _, streamed = await _drain_with_fallback( - _continuation_router(), error, primary_chunks, messages, [_usage_chunk(91, 5)] - ) - - assert primary_usage.prompt_tokens > 0 - assert streamed[-1].usage.prompt_tokens == primary_usage.prompt_tokens + 91 - assert streamed[-1].usage.completion_tokens == primary_usage.completion_tokens + 5 - - class _CutSSE(httpx.AsyncByteStream): def __init__(self, content: str) -> None: self._content: Final = content async def __aiter__(self) -> AsyncIterator[bytes]: - yield f'data: {json.dumps({"id": "cut", "object": "chat.completion.chunk", "created": 1700000000, "model": "m", "choices": [{"index": 0, "delta": {"role": "assistant", "content": self._content}, "finish_reason": None}]})}\n\n'.encode() + yield f"data: {json.dumps({'id': 'cut', 'object': 'chat.completion.chunk', 'created': 1700000000, 'model': 'm', 'choices': [{'index': 0, 'delta': {'role': 'assistant', 'content': self._content}, 'finish_reason': None}]})}\n\n".encode() raise httpx.RemoteProtocolError("peer closed connection without sending complete message body") @@ -3718,7 +2445,9 @@ async def test_continued_stream_success_log_bills_every_cut_leg( recorder: Final = _SpendRecorder() monkeypatch.setattr(litellm, "callbacks", [recorder]) respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + return_value=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use") + ) ) respx_mock.post(host="api.deepseek.com").mock( return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE(" learned")) @@ -3801,7 +2530,9 @@ async def test_continued_stream_sharing_the_proxy_logging_object_bills_the_cut_t recorder: Final = _SlowFailureRecorder() monkeypatch.setattr(litellm, "callbacks", [recorder]) respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use")) + return_value=httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=_CutSSE("Transformers use") + ) ) respx_mock.post(host="api.mistral.ai").mock( return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, text=_BACKUP_SSE) @@ -3874,7 +2605,11 @@ _AGENTIC_FIRST_CALL_BODY: Final = { "role": "assistant", "content": None, "tool_calls": [ - {"id": "call_abc", "type": "function", "function": {"name": "litellm_code_execution", "arguments": "{}"}} + { + "id": "call_abc", + "type": "function", + "function": {"name": "litellm_code_execution", "arguments": "{}"}, + } ], }, } @@ -4165,7 +2900,10 @@ async def test_responses_fallback_sharing_the_proxy_logging_object_logs_the_fall 200, headers={"content-type": "text/event-stream"}, text=_responses_sse( - {"type": "response.created", "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress")}, + { + "type": "response.created", + "response": _responses_envelope("resp_backup", "gpt-4o-mini", "in_progress"), + }, { "type": "response.completed", "response": _responses_envelope( @@ -4235,420 +2973,6 @@ async def test_responses_fallback_sharing_the_proxy_logging_object_logs_the_fall ) -def test_carry_partial_spend_to_fallback_adds_the_partial_the_cut_stream_inherited(): - from litellm.types.utils import Usage - - cut_logging: Final = LiteLLMLogging( - model="deepseek/deepseek-chat", messages=[], stream=True, call_type="acompletion", - start_time=datetime.now(), litellm_call_id="cut", function_id="cut", - ) - cut_logging.model_call_details["response_cost"] = 0.5 - cut_logging.update_environment_variables( - litellm_params={"custom_llm_provider": "deepseek"}, optional_params={}, model="deepseek/deepseek-chat" - ) - cut_logging.record_continued_partial( - usage=Usage(prompt_tokens=40, completion_tokens=7, total_tokens=47), - response_cost=0.001, - response_id="chatcmpl-primary", - ) - fallback_logging: Final = LiteLLMLogging( - model="mistral-large-latest", messages=[], stream=True, call_type="acompletion", - start_time=datetime.now(), litellm_call_id="fallback", function_id="fallback", - ) - - Router._carry_partial_spend_to_fallback( - SimpleNamespace(logging_obj=fallback_logging), - Router._price_cut_partial( - cast(CustomStreamWrapper, SimpleNamespace(logging_obj=cut_logging, model="deepseek/deepseek-chat")), - litellm.ModelResponse( - id="chatcmpl-relay", usage=Usage(prompt_tokens=60, completion_tokens=3, total_tokens=63) - ), - ), - ) - - relay_prompt_cost, relay_completion_cost = litellm.cost_per_token( - model="deepseek/deepseek-chat", prompt_tokens=60, completion_tokens=3 - ) - carried_usage: Final = fallback_logging.model_call_details["continued_partial_usage"] - assert (carried_usage.prompt_tokens, carried_usage.completion_tokens) == (100, 10) - assert fallback_logging.model_call_details["continued_partial_cost"] == pytest.approx( - relay_prompt_cost + relay_completion_cost + 0.001 - ) - assert fallback_logging.model_call_details["continued_response_id"] == "chatcmpl-primary" - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_plain_fallback_keeps_the_fallback_usage(): - from litellm.exceptions import MidStreamFallbackError - - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="anthropic", generated_content="", - is_pre_first_chunk=True, - ) - role_chunk = litellm.ModelResponseStream( - model="gpt-4", choices=[{"index": 0, "delta": {"role": "assistant", "content": ""}}] - ) - - _, streamed = await _drain_with_fallback( - _continuation_router(), - error, - [role_chunk, _usage_chunk(21, 0)], - [{"role": "user", "content": "Hi"}], - [_usage_chunk(91, 5)], - ) - - assert (streamed[-1].usage.prompt_tokens, streamed[-1].usage.completion_tokens) == (91, 5) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "partial,fallback_texts,expected_answer", - [ - ("Transformers use ", (" self", "-attention"), "Transformers use self-attention"), - ("def f():\n ", ("\n", " return 1"), "def f():\n return 1"), - ("First.\n", ("\n\nSecond",), "First.\n\nSecond"), - ], - ids=["space", "indentation", "paragraph_break"], -) -async def test_acompletion_streaming_iterator_continuation_does_not_double_trailing_whitespace( - partial, fallback_texts, expected_answer -): - from litellm.exceptions import MidStreamFallbackError - - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content=partial, - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - fallback_items = [ - litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": text}}]) for text in fallback_texts - ] - - passed_kwargs, streamed = await _drain_with_fallback( - _continuation_router(), error, [], [{"role": "user", "content": "Hi"}], fallback_items - ) - - assert passed_kwargs["messages"][-1]["content"] == partial - assert partial + "".join(chunk.choices[0].delta.content for chunk in streamed) == expected_answer - - -@pytest.mark.asyncio -@pytest.mark.parametrize("enabled", [True, "true", " True "]) -async def test_update_settings_turns_on_mid_stream_continuation(monkeypatch, enabled): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - from litellm.router_utils.pre_call_checks.continuation_prefill_check import ( - MID_STREAM_CONTINUATION_KWARG, - MID_STREAM_CONTINUATION_MARKER, - ContinuationPrefillDeploymentCheck, - ) - - monkeypatch.setattr(litellm, "callbacks", []) - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - ) - router.update_settings(enable_mid_stream_fallback_continuation=enabled) - - error = MidStreamFallbackError( - message="Connection lost", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - fallback_chunk = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), - ) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error), - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - collected = [chunk async for chunk in result] - - assert mock_fallback.await_args.kwargs["kwargs"][MID_STREAM_CONTINUATION_KWARG] is MID_STREAM_CONTINUATION_MARKER - assert fallback_chunk in collected - assert any(isinstance(callback, ContinuationPrefillDeploymentCheck) for callback in litellm.callbacks) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "error_kwargs,request_kwargs", - [ - ({"emitted_disqualifying_content": True}, {}), - ({"emitted_disqualifying_content": False}, {"response_format": {"type": "json_object"}}), - ({"emitted_disqualifying_content": False}, {"tool_choice": "required"}), - ({"emitted_disqualifying_content": False}, {"merge_reasoning_content_in_choices": True}), - ({"emitted_disqualifying_content": False}, {"n": 2}), - ], - ids=["tool_or_thinking_emitted", "json_mode", "forced_tool_choice", "merged_reasoning", "multiple_choices"], -) -async def test_acompletion_streaming_iterator_declines_ineligible_after_content(error_kwargs, request_kwargs): - """Flag on but the break is not continuation-safe: the router re-raises and - never enters the fallback chain, so no duplicated/rejected request is sent.""" - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, **error_kwargs, - ) - source = _make_midstream_source(error) - initial_kwargs = {"model": "gpt-4", "stream": True, **request_kwargs} - - with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=source, messages=[{"role": "user", "content": "Hi"}], initial_kwargs=initial_kwargs - ) - with pytest.raises(MidStreamFallbackError): - async for _ in result: - pass - - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_raises_content_filter_instead_of_continuing(): - """A content filter that stops a stream mid-answer reaches the client as ContentPolicyViolationError, - never as a continuation that lets a backup finish the filtered answer.""" - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "azure/gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - content_filter = litellm.ContentPolicyViolationError(message="filtered", model="gpt-4", llm_provider="azure") - error = MidStreamFallbackError( - message="filtered", model="gpt-4", llm_provider="azure", generated_content="Step one is to obtain the", - is_pre_first_chunk=False, emitted_disqualifying_content=False, original_exception=content_filter, - ) - - with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error), - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - with pytest.raises(litellm.ContentPolicyViolationError): - async for _ in result: - pass - - mock_fallback.assert_not_awaited() - - -async def _continues_after_content( - initial_kwargs: dict[str, object], - primary_params: Mapping[str, object] = MappingProxyType({}), - default_litellm_params: dict[str, object] | None = None, - cut_model_id: str = "primary-deployment", -) -> bool: - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router: Final = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "k1", **primary_params}, - "model_info": {"id": "primary-deployment"}, - }, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - { - "model_name": "openai-backup", - "litellm_params": {"model": "gpt-4o", "api_key": "k3"}, - "model_info": {"id": "no-prefill-deployment", "supports_assistant_prefill": False}, - }, - { - "model_name": "claude-backup", - "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k4"}, - "model_info": {"id": "prefill-deployment", "supports_assistant_prefill": True}, - }, - { - "model_name": "json-claude-backup", - "litellm_params": { - "model": "anthropic/claude-3-opus-20240229", - "api_key": "k5", - "response_format": {"type": "json_object"}, - }, - "model_info": {"id": "json-prefill-deployment", "supports_assistant_prefill": True}, - }, - ], - fallbacks=[{"gpt-4": ["backup"]}], - default_litellm_params=default_litellm_params, - enable_mid_stream_fallback_continuation=True, - ) - error: Final = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - fallback_chunk: Final = litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": " world"}}]) - cut_stream: Final = _make_midstream_source(error) - cut_stream._hidden_params = {"model_id": cut_model_id} - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_FakeFallbackStream(fallback_chunk)), - ) as mock_fallback: - result: Final = await router._acompletion_streaming_iterator( - model_response=cut_stream, - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True, **initial_kwargs}, - ) - with contextlib.suppress(MidStreamFallbackError): - _ = [chunk async for chunk in result] - return mock_fallback.await_count == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "event_hook,continues", - [("post_call", False), (None, False), ("pre_call", True)], - ids=["post_call_guardrail", "every_hook_guardrail", "input_only_guardrail"], -) -async def test_mid_stream_continuation_declines_while_an_output_guardrail_is_registered( - monkeypatch: pytest.MonkeyPatch, event_hook: str | None, continues: bool -) -> None: - from litellm.types.guardrails import GuardrailEventHooks - - guardrail: Final = CustomGuardrail( - guardrail_name="output-check", event_hook=GuardrailEventHooks(event_hook) if event_hook else None - ) - monkeypatch.setattr(litellm, "callbacks", [guardrail]) - - assert await _continues_after_content({}) is continues - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "fallbacks,continues", - [ - ([{"model": "backup", "messages": [{"role": "user", "content": "Say hi"}]}], False), - ([{"model": "backup", "stream": False}], False), - ([{"model": "backup", "max_tokens": 300}], False), - ([{"model": "backup"}], True), - ([{"gpt-4": [{"model": "backup", "stream": False}]}], False), - ([{"*": [{"model": "backup", "max_tokens": 300}]}], False), - ([{"gpt-4": [{"model": "backup"}, "backup"]}], True), - ([{"gpt-4": ["backup"]}, {"other-group": [{"model": "backup", "stream": False}]}], True), - ([{}], True), - ([{"gpt-4": 5}], True), - ], - ids=[ - "entry_with_its_own_messages", - "entry_turning_off_streaming", - "entry_with_its_own_max_tokens", - "entry_with_model_only", - "group_target_turning_off_streaming", - "generic_group_target_with_its_own_max_tokens", - "group_targets_with_model_only", - "override_only_in_another_groups_chain", - "empty_entry", - "non_list_chain", - ], -) -async def test_mid_stream_continuation_declines_a_request_fallback_that_overrides_request_params( - fallbacks: list[object], continues: bool -) -> None: - assert await _continues_after_content({"fallbacks": fallbacks}) is continues - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "fallbacks,continues", - [ - ([{"gpt-4": ["no-prefill-deployment"]}], False), - ([{"gpt-4": ["backup", "no-prefill-deployment"]}], False), - ([{"model": "no-prefill-deployment"}], False), - ([{"gpt-4": ["prefill-deployment"]}], True), - ([{"gpt-4": ["openai-backup"]}], True), - ([{"gpt-4": ["json-prefill-deployment"]}], False), - ], - ids=[ - "deployment_id_without_prefill", - "deployment_id_without_prefill_later_in_chain", - "flat_entry_naming_a_deployment_id_without_prefill", - "deployment_id_with_prefill", - "group_without_prefill_left_to_the_deployment_filter", - "deployment_id_with_its_own_response_format", - ], -) -async def test_mid_stream_continuation_declines_a_fallback_naming_a_deployment_that_cannot_continue( - fallbacks: list[object], continues: bool -) -> None: - assert await _continues_after_content({"fallbacks": fallbacks}) is continues - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "primary_params,default_litellm_params,continues", - [ - ({"response_format": {"type": "json_object"}}, None, False), - ({"n": 2}, None, False), - ({"tool_choice": "required"}, None, False), - ({"merge_reasoning_content_in_choices": True}, None, False), - ({}, {"response_format": {"type": "json_object"}}, False), - ({"response_format": {"type": "text"}}, None, True), - ({"temperature": 0.2}, None, True), - ({"tool_choice": "required"}, {"tool_choice": "auto"}, False), - ({"response_format": {"type": "json_object"}}, {"response_format": {"type": "text"}}, True), - ({"extra_body": {"n": 2}}, None, False), - ], - ids=[ - "deployment_response_format", - "deployment_n", - "deployment_tool_choice", - "deployment_merged_reasoning", - "router_default_response_format", - "deployment_plain_text_format", - "deployment_unrelated_param", - "deployment_tool_choice_over_router_default", - "router_default_format_over_deployment", - "deployment_extra_body_n", - ], -) -async def test_mid_stream_continuation_reads_constraints_configured_on_the_cut_deployment( - primary_params: Mapping[str, object], default_litellm_params: dict[str, object] | None, continues: bool -) -> None: - assert await _continues_after_content({}, primary_params, default_litellm_params) is continues - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "extra_body,continues", - [({"response_format": {"type": "json_object"}}, False), ({"n": 2}, False), ({"top_k": 5}, True)], - ids=["json_format_in_extra_body", "n_in_extra_body", "unrelated_extra_body_key"], -) -async def test_mid_stream_continuation_reads_constraints_the_request_sends_in_extra_body( - extra_body: dict[str, object], continues: bool -) -> None: - assert await _continues_after_content({"extra_body": extra_body}) is continues - - -@pytest.mark.asyncio -async def test_mid_stream_continuation_declines_when_the_cut_deployment_is_not_in_the_router() -> None: - assert await _continues_after_content({}, cut_model_id="per-request-credential-deployment") is False - - @pytest.mark.asyncio @pytest.mark.parametrize( "routing_strategy", @@ -4694,645 +3018,6 @@ async def test_continuation_routes_only_to_deployments_that_can_continue_under_e ) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "flag_kwargs", [{}, {"enable_mid_stream_fallback_continuation": "false"}], ids=["default", "string_false"] -) -async def test_acompletion_streaming_iterator_flag_off_declines_after_content(flag_kwargs): - """Regression guard for the opt-in default: with the flag off, a post-content - break re-raises and never falls back, exactly as before this feature.""" - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - **flag_kwargs, - ) - error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - source = _make_midstream_source(error) - - with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=source, messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - with pytest.raises(MidStreamFallbackError): - async for _ in result: - pass - - mock_fallback.assert_not_awaited() - - -def test_build_completion_continuation_input_appends_assistant_prefill(): - messages = [{"role": "user", "content": "hi"}] - built = litellm.Router._build_completion_continuation_input(messages, "partial answer") - assert built[:-1] == messages - assert built[-1] == {"role": "assistant", "content": "partial answer", "prefix": True} - - -@pytest.mark.parametrize( - "prefill_content,second_partial,expected_content", - [ - ("part one ", " part two ", "part one part two "), - ("part one ", "part two ", "part one part two "), - ("First.\n\n", "Second", "First.\n\nSecond"), - ([{"type": "text", "text": "part one "}], " part two ", [{"type": "text", "text": "part one part two "}]), - ], - ids=["separator_sent_again", "separator_not_sent_again", "paragraph_break_not_sent_again", "content_blocks"], -) -def test_build_completion_continuation_input_folds_into_existing_prefill( - prefill_content, second_partial, expected_content -): - """A nested break must not leave two trailing assistant turns: the new partial folds into the prior prefill, - which keeps reading as the text the caller already received.""" - messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": prefill_content, "prefix": True}] - built = litellm.Router._build_completion_continuation_input(messages, second_partial) - assert built == [messages[0], {"role": "assistant", "content": expected_content, "prefix": True}] - - -@pytest.mark.parametrize( - "prefill_content,expected_content", - [ - ("Hello", "Hello world"), - ([{"type": "text", "text": "Hello"}], [{"type": "text", "text": "Hello world"}]), - ], - ids=["string_prefill", "content_block_prefill"], -) -def test_build_completion_continuation_input_does_not_double_an_echoed_prefill(prefill_content, expected_content): - messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": prefill_content, "prefix": True}] - built = litellm.Router._build_completion_continuation_input(messages, "Hello world") - assert built is not None - assert built[-1] == {"role": "assistant", "content": expected_content, "prefix": True} - - -@pytest.mark.parametrize( - "cut_partial,expected_content", - [("Hel", None), ("Hello", "Hello"), ("Hello wor", "Hello wor")], - ids=["cut_mid_echo", "cut_after_a_full_echo", "cut_after_new_text"], -) -def test_build_completion_continuation_input_declines_a_cut_in_the_middle_of_a_prefill_echo( - cut_partial, expected_content -): - messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "Hello", "prefix": True}] - built = litellm.Router._build_completion_continuation_input(messages, cut_partial) - expected = None if expected_content is None else [messages[0], {**messages[1], "content": expected_content}] - assert built == expected - - -@pytest.mark.parametrize( - "assistant_content", ["Earlier answer.", [{"type": "text", "text": "Earlier answer."}]], ids=["string", "blocks"] -) -def test_build_completion_continuation_input_declines_a_trailing_plain_assistant_turn(assistant_content): - """The caller's own assistant turn is history, not a prefill, so the partial is never fused into it.""" - messages = [{"role": "user", "content": "hi"}, {"role": "assistant", "content": assistant_content}] - assert litellm.Router._build_completion_continuation_input(messages, "The capital of France is") is None - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_declines_when_request_ends_in_plain_assistant_turn(): - from unittest.mock import AsyncMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - {"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k1"}}, - {"model_name": "backup", "litellm_params": {"model": "anthropic/claude-3-opus-20240229", "api_key": "k2"}}, - ], - fallbacks=[{"gpt-4": ["backup"]}], - enable_mid_stream_fallback_continuation=True, - ) - error = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - - with patch.object(router, "async_function_with_fallbacks_common_utils", new=AsyncMock()) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=_make_midstream_source(error), - messages=[{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Earlier answer."}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - with pytest.raises(MidStreamFallbackError): - async for _ in result: - pass - - mock_fallback.assert_not_awaited() - - -def test_build_completion_continuation_input_keeps_structured_assistant_content(): - """Content blocks on a trailing assistant turn stay blocks: the partial continues the last text - block in place (a block boundary would let a provider drop the space between "Sure, " and "here"), - lands as a new text block only after a non-text block, and a content shape that cannot be extended declines.""" - from litellm.router import _extend_trailing_text_block - - messages = [ - {"role": "user", "content": "hi"}, - {"role": "assistant", "content": [{"type": "text", "text": "Sure, "}], "prefix": True}, - ] - built = litellm.Router._build_completion_continuation_input(messages, "here it is") - assert built is not None - assert [m["role"] for m in built] == ["user", "assistant"] - assert built[-1] == {"role": "assistant", "content": [{"type": "text", "text": "Sure, here it is"}], "prefix": True} - image_block = {"type": "image_url", "image_url": {"url": "https://example.test/a.png"}} - assert _extend_trailing_text_block([image_block], "here it is") == [image_block, {"type": "text", "text": "here it is"}] - assert _extend_trailing_text_block([], "here it is") == [{"type": "text", "text": "here it is"}] - assert ( - litellm.Router._build_completion_continuation_input([{"role": "assistant", "content": 42, "prefix": True}], "x") - is None - ) - - -def test_mid_stream_continuation_eligible_allows_text_response_format(): - """response_format={"type": "text"} is the unconstrained default and must stay - eligible, unlike json_object / json_schema.""" - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "gpt-4", "api_key": "k"}}], - enable_mid_stream_fallback_continuation=True, - ) - e = MidStreamFallbackError( - message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello", - is_pre_first_chunk=False, emitted_disqualifying_content=False, - ) - assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "text"}}) is True - assert router._mid_stream_continuation_eligible(e=e, request_kwargs={"response_format": {"type": "json_object"}}) is False - - -def test_raise_original_mid_stream_error_surfaces_the_provider_exception(): - from litellm.exceptions import MidStreamFallbackError, RateLimitError - - provider_error = RateLimitError(message="rate limited", llm_provider="openai", model="gpt-4") - wrapped = MidStreamFallbackError( - message="rate limited", model="gpt-4", llm_provider="openai", - original_exception=provider_error, generated_content="Hello", - ) - with pytest.raises(RateLimitError) as raised: - litellm.Router._raise_original_mid_stream_error(wrapped) - assert raised.value is provider_error - assert raised.value.__cause__ is wrapped - - bare = MidStreamFallbackError(message="boom", model="gpt-4", llm_provider="openai", generated_content="Hello") - with pytest.raises(MidStreamFallbackError) as bare_raised: - litellm.Router._raise_original_mid_stream_error(bare) - assert bare_raised.value is bare - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_reraises_original_exception_when_available(): - """Async: when the mid-stream MidStreamFallbackError wraps a real provider - exception (original_exception), the router must re-raise that original - exception instead of the internal wrapper, so the client sees the - specific error type/code (e.g. RateLimitError) rather than a generic - MidStreamFallbackError.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError, RateLimitError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - set_verbose=True, - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - original_exception = RateLimitError( - message="rate limited", - llm_provider="vertex_ai", - model="gpt-4", - ) - error = MidStreamFallbackError( - message="rate limited", - model="gpt-4", - llm_provider="openai", - original_exception=original_exception, - generated_content="Hello", - ) - - mock_chunks = [ - MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello"))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content=" there"))]), - ] - - class AsyncIteratorWithError: - def __init__(self, items, error_after_index): - self.items = items - self.index = 0 - self.error_after_index = error_after_index - - def __aiter__(self): - return self - - async def __anext__(self): - if self.index >= len(self.items): - raise StopAsyncIteration - if self.index == self.error_after_index: - raise error - item = self.items[self.index] - self.index += 1 - return item - - mock_error_response = AsyncIteratorWithError(mock_chunks, 1) - setattr(mock_error_response, "model", "gpt-4") - setattr(mock_error_response, "custom_llm_provider", "openai") - setattr(mock_error_response, "logging_obj", MagicMock()) - - result = await router._acompletion_streaming_iterator( - model_response=mock_error_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - with pytest.raises(RateLimitError) as exc_info: - async for _ in result: - pass - assert exc_info.value is original_exception - assert exc_info.value.type == "throttling_error" - assert exc_info.value.code == "429" - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_edge_cases(): - """Test edge cases for _acompletion_streaming_iterator.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - set_verbose=True, - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - # Test: Empty generated content - empty_error = MidStreamFallbackError( - message="Error", - model="gpt-4", - llm_provider="openai", - generated_content="", # Empty content - ) - - class AsyncIteratorImmediateError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - - def __aiter__(self): - return self - - async def __anext__(self): - raise empty_error - - mock_response = AsyncIteratorImmediateError() - - # Mock empty fallback response using AsyncIterator - class EmptyAsyncIterator: - def __aiter__(self): - return self - - async def __anext__(self): - raise StopAsyncIteration - - mock_fallback_response = EmptyAsyncIterator() - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=mock_fallback_response, - ) as mock_fallback_utils: - collected_chunks = [] - iterator = await router._acompletion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - async for chunk in iterator: - collected_chunks.append(chunk) - - # Should still call fallback even with empty content - assert mock_fallback_utils.called - fallback_kwargs = mock_fallback_utils.call_args.kwargs["kwargs"] - modified_messages = fallback_kwargs["messages"] - - # Empty content → pre-first-chunk path uses original messages - # (no continuation prompt added) - assert modified_messages == messages - print("✓ Handles empty generated content correctly") - - print("✓ Edge case tests passed!") - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_preserves_hidden_params(): - """ - Regression test: FallbackStreamWrapper must copy _hidden_params from the - original CustomStreamWrapper so that x-litellm-overhead-duration-ms (and - other hidden params) are present in the proxy response headers for streaming. - """ - from unittest.mock import MagicMock - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - # Simulate a CustomStreamWrapper that already has timing metadata set by - # update_response_metadata (litellm_overhead_time_ms, _response_ms, etc.) - mock_response = MagicMock() - mock_response.model = "gpt-4" - mock_response.custom_llm_provider = "openai" - mock_response.logging_obj = MagicMock() - mock_response._hidden_params = { - "litellm_overhead_time_ms": 12.34, - "_response_ms": 500.0, - "litellm_call_id": "test-call-id", - "api_base": "https://api.openai.com", - "additional_headers": {}, - } - - # Make the mock iterable (yields nothing — we only care about hidden_params) - async def _empty(): - return - yield # make it an async generator - - mock_response.__aiter__ = lambda self: _empty().__aiter__() - - result = await router._acompletion_streaming_iterator( - model_response=mock_response, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - - # The returned FallbackStreamWrapper must carry the original _hidden_params - assert hasattr(result, "_hidden_params"), "result must have _hidden_params" - assert result._hidden_params.get("litellm_overhead_time_ms") == 12.34, ( - "litellm_overhead_time_ms must be preserved — " - "this is what drives x-litellm-overhead-duration-ms in streaming responses" - ) - assert result._hidden_params.get("litellm_call_id") == "test-call-id" - assert result._hidden_params.get("_response_ms") == 500.0 - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_preserves_response_headers(): - """LIT-6767: the returned wrapper must carry the provider's raw response headers. - - Proxy callbacks read ``_response_headers`` off the object the router hands - back. The wrapper used to be built without it, so every streaming chat - completion reported zero raw provider headers while the non-streaming path - reported the full set. - """ - from unittest.mock import MagicMock - - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - async def _empty(): - return - yield # make it an async generator - - provider_headers = { - "x-request-id": "req-provider-123", - "x-ratelimit-remaining-requests": "42", - # a provider must never be able to spoof an internal header - "x-litellm-model-id": "spoofed", - } - source = CustomStreamWrapper( - completion_stream=_empty(), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - _response_headers=provider_headers, - ) - - result = await router._acompletion_streaming_iterator( - model_response=source, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - - assert result._response_headers == provider_headers - additional_headers = result._hidden_params["additional_headers"] - assert additional_headers["llm_provider-x-request-id"] == "req-provider-123" - assert additional_headers["llm_provider-x-ratelimit-remaining-requests"] == "42" - # internal-header protection survives: the provider value is namespaced, never promoted - assert additional_headers["llm_provider-x-litellm-model-id"] == "spoofed" - assert "x-litellm-model-id" not in additional_headers - - -def test_completion_streaming_iterator_preserves_response_headers(): - """LIT-6767, sync counterpart of the async header-preservation test.""" - from unittest.mock import MagicMock - - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - provider_headers = {"x-request-id": "req-provider-sync", "openai-organization": "org-real"} - source = CustomStreamWrapper( - completion_stream=iter([]), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - _response_headers=provider_headers, - ) - - result = router._completion_streaming_iterator( - model_response=source, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - - assert result._response_headers == provider_headers - assert result._hidden_params["additional_headers"]["llm_provider-x-request-id"] == "req-provider-sync" - - -def test_adopt_fallback_response_headers_replaces_rather_than_merges(): - """LIT-6767: direct unit for FallbackAwareStreamWrapper.adopt_fallback_response_headers. - - Values from the failed attempt must not survive, so the wrapper replaces both - ``_response_headers`` and ``_hidden_params`` instead of merging them. - """ - from unittest.mock import MagicMock - - from litellm.router import FallbackAwareStreamWrapper, Router - - wrapper = FallbackAwareStreamWrapper( - completion_stream=iter([]), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - _response_headers={"x-request-id": "req-FAILED"}, - ) - wrapper._hidden_params = { - "model_id": "failed-deployment", - "only_on_failed_attempt": "stale", - "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, - } - - fallback = MagicMock() - fallback._response_headers = {"x-request-id": "req-FALLBACK"} - fallback._hidden_params = { - "model_id": "fallback-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, - } - - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) - - assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} - assert wrapper._hidden_params["model_id"] == "fallback-deployment" - assert "only_on_failed_attempt" not in wrapper._hidden_params - assert wrapper._hidden_params is not fallback._hidden_params - # the snapshot CustomStreamWrapper caches at init has to follow, or a chunk built - # from it would still be stamped with the deployment that failed - assert wrapper._base_hidden_params["model_id"] == "fallback-deployment" - # the nested header dict is copied too, so a later mutation on the fallback - # response cannot reach headers the proxy has already published - assert wrapper._hidden_params["additional_headers"] is not fallback._hidden_params["additional_headers"] - fallback._hidden_params["additional_headers"]["llm_provider-x-request-id"] = "req-MUTATED" - assert wrapper._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"} - - -def test_adopt_fallback_response_headers_survives_a_collected_wrapper(): - """LIT-6767: adoption still returns the fallback's params once the wrapper is gone.""" - import weakref - from unittest.mock import MagicMock - - from litellm.router import FallbackAwareStreamWrapper, Router - - fallback: Final = MagicMock() - fallback._response_headers = {"x-request-id": "req-FALLBACK"} - fallback._hidden_params = { - "model_id": "fallback-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, - } - - wrapper = FallbackAwareStreamWrapper( - completion_stream=iter([]), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - ) - live_ref: Final = weakref.ref(wrapper) - prepared: Final = Router._adopt_fallback_response_headers(live_ref, fallback) - assert prepared == (fallback._hidden_params, fallback._hidden_params["additional_headers"]) - assert wrapper.fallback_headers_adopted is True - assert wrapper._response_headers == {"x-request-id": "req-FALLBACK"} - - dead_ref: Final = weakref.ref(wrapper) - del wrapper - assert dead_ref() is None - assert Router._adopt_fallback_response_headers(dead_ref, fallback) == prepared - - -def test_adopt_fallback_response_headers_drops_headers_the_fallback_cannot_replace(): - """LIT-6767: a fallback that carries no raw provider headers publishes none. - - Keeping the failed attempt's raw headers would hand the client and the callbacks a - provider ``x-request-id`` for a request that deployment never served, which is the - leak this fix exists to close. - """ - from unittest.mock import MagicMock - - from litellm.router import FallbackAwareStreamWrapper, Router - - wrapper = FallbackAwareStreamWrapper( - completion_stream=iter([]), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - _response_headers={"x-request-id": "req-FAILED"}, - ) - wrapper._hidden_params = {"model_id": "failed-deployment", "additional_headers": {}} - - fallback = MagicMock() - fallback._response_headers = None - fallback._hidden_params = {"model_id": "fallback-deployment"} - - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) - - assert wrapper._response_headers is None - assert wrapper._hidden_params["model_id"] == "fallback-deployment" - assert wrapper.fallback_headers_adopted is True - - -def test_adopt_fallback_response_headers_keeps_identity_when_fallback_has_none(): - """A fallback response carrying no hidden params keeps the identity headers. - - Publishing no ``x-litellm-*`` header at all for a request the fallback served is - worse than keeping what is there, so only the raw provider headers are dropped. - """ - from unittest.mock import MagicMock - - from litellm.router import FallbackAwareStreamWrapper, Router - - wrapper = FallbackAwareStreamWrapper( - completion_stream=iter([]), - model="gpt-4", - custom_llm_provider="openai", - logging_obj=MagicMock(), - _response_headers={"x-request-id": "req-FAILED"}, - ) - hidden_params_before = wrapper._hidden_params - - fallback = object() - wrapper.adopt_fallback_response_headers( - fallback, Router._prepare_fallback_hidden_params(fallback) - ) - - assert wrapper._response_headers is None - assert wrapper._hidden_params is hidden_params_before - assert wrapper.fallback_headers_adopted is True - - @pytest.mark.asyncio @pytest.mark.parametrize("response_kind", ["object", "dict", "async-generator"]) async def test_set_response_headers_exposes_complexity_decision_on_every_response_shape( @@ -5400,299 +3085,7 @@ async def test_set_response_headers_is_the_only_complexity_header_source_for_pro **additional_headers, ) - assert not { - key for key in proxy_headers if key.startswith("x-litellm-complexity-router-") - } - - -@pytest.mark.asyncio -@pytest.mark.parametrize("close_fails", [False, True]) -async def test_acompletion_streaming_iterator_adopts_fallback_response_headers( - close_fails: bool, -) -> None: - """LIT-6767: after a successful pre-first-chunk fallback, the wrapper must - describe the deployment that served the stream, with no value left over - from the attempt that failed.""" - from unittest.mock import MagicMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - failed_error = MidStreamFallbackError( - message="upstream died before the first chunk", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=True, - ) - - class FailedStream: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - self._response_headers = {"x-request-id": "req-FAILED"} - self._hidden_params = { - "model_id": "failed-deployment", - "api_base": "https://failed.example", - "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, - "only_on_failed_attempt": "stale", - } - - def __aiter__(self): - return self - - async def __anext__(self): - raise failed_error - - class FallbackStream: - def __init__(self): - self.close_attempted: bool = False - self._response_headers = {"x-request-id": "req-FALLBACK"} - self._hidden_params = { - "model_id": "fallback-deployment", - "api_base": "https://fallback.example", - "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, - } - self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) - - def __aiter__(self): - return self - - async def __anext__(self): - try: - return next(self._chunks) - except StopIteration: - raise StopAsyncIteration from None - - async def aclose(self) -> None: - self.close_attempted = True - if close_fails: - raise RuntimeError("fallback transport failed during close") - - fallback_stream = FallbackStream() - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=fallback_stream, - ): - result = await router._acompletion_streaming_iterator( - model_response=FailedStream(), - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - # the failed attempt is what the wrapper is built from - assert result._response_headers == {"x-request-id": "req-FAILED"} - # the very first chunk the fallback produces must already be published under - # the fallback's identity: the proxy commits response headers once that chunk - # is buffered, so adopting any later is adopting too late - first_chunk: Final = await result.__anext__() - assert first_chunk.choices[0].delta.content == "OK" - assert result._response_headers == {"x-request-id": "req-FALLBACK"} - assert result._hidden_params["model_id"] == "fallback-deployment" - async for _ in result: - pass - - assert fallback_stream.close_attempted - assert result._response_headers == {"x-request-id": "req-FALLBACK"} - assert result._hidden_params["model_id"] == "fallback-deployment" - assert result._hidden_params["api_base"] == "https://fallback.example" - assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-FALLBACK"} - # stale values are removed, not merged over - assert "only_on_failed_attempt" not in result._hidden_params - # and the wrapper holds its own copy, so later fallback mutations cannot leak in - assert result._hidden_params is not fallback_stream._hidden_params - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback(): - """LIT-6767: a fallback that itself fails over before its first chunk. - - The selected fallback still describes its own failed attempt at selection time, so - the wrapper has to re-read it once a chunk exists or it publishes a deployment that - produced no output. - """ - from unittest.mock import MagicMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - failed_error: Final = MidStreamFallbackError( - message="upstream died before the first chunk", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=True, - ) - - class FailedStream: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - self._response_headers = {"x-request-id": "req-FAILED"} - self._hidden_params = { - "model_id": "failed-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, - } - - def __aiter__(self): - return self - - async def __anext__(self): - raise failed_error - - class NestedFallbackStream: - """A fallback that repoints itself at a third deployment as it yields.""" - - def __init__(self): - self._response_headers = {"x-request-id": "req-MIDDLE"} - self._hidden_params = { - "model_id": "middle-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"}, - } - self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) - - def __aiter__(self): - return self - - async def __anext__(self): - try: - chunk = next(self._chunks) - except StopIteration: - raise StopAsyncIteration from None - self._response_headers = {"x-request-id": "req-SERVED"} - self._hidden_params = { - "model_id": "served-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, - } - return chunk - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=NestedFallbackStream(), - ): - result = await router._acompletion_streaming_iterator( - model_response=FailedStream(), - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - first_chunk: Final = await result.__anext__() - # the proxy commits response headers once this chunk is buffered - assert result._response_headers == {"x-request-id": "req-SERVED"} - assert result._hidden_params["model_id"] == "served-deployment" - assert result._hidden_params["additional_headers"] == {"llm_provider-x-request-id": "req-SERVED"} - # and the chunk itself carries the same deployment - assert first_chunk._hidden_params["model_id"] == "served-deployment" - async for _ in result: - pass - - assert result._response_headers == {"x-request-id": "req-SERVED"} - assert result._hidden_params["model_id"] == "served-deployment" - - -def test_completion_streaming_iterator_adopts_the_deployment_that_served_a_nested_fallback(): - """LIT-6767, sync counterpart of the nested-fallback adoption test.""" - from unittest.mock import MagicMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - failed_error: Final = MidStreamFallbackError( - message="upstream died before the first chunk", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=True, - ) - - class FailedStream: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - self._response_headers = {"x-request-id": "req-FAILED"} - self._hidden_params = { - "model_id": "failed-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, - } - - def __iter__(self): - return self - - def __next__(self): - raise failed_error - - class NestedFallbackStream: - """A fallback that repoints itself at a third deployment as it yields.""" - - def __init__(self): - self._response_headers = {"x-request-id": "req-MIDDLE"} - self._hidden_params = { - "model_id": "middle-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-MIDDLE"}, - } - self._chunks = iter([litellm.ModelResponseStream(choices=[{"index": 0, "delta": {"content": "OK"}}])]) - - def __iter__(self): - return self - - def __next__(self): - chunk = next(self._chunks) - self._response_headers = {"x-request-id": "req-SERVED"} - self._hidden_params = { - "model_id": "served-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-SERVED"}, - } - return chunk - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=NestedFallbackStream()), - ): - result = router._completion_streaming_iterator( - model_response=FailedStream(), - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - first_chunk: Final = next(result) - assert result._response_headers == {"x-request-id": "req-SERVED"} - assert result._hidden_params["model_id"] == "served-deployment" - assert first_chunk._hidden_params["model_id"] == "served-deployment" - for _ in result: - pass - - assert result._response_headers == {"x-request-id": "req-SERVED"} - assert result._hidden_params["model_id"] == "served-deployment" + assert not {key for key in proxy_headers if key.startswith("x-litellm-complexity-router-")} def test_completion_streaming_fallback_resumes_chain_without_retrying_primary(): @@ -6020,614 +3413,6 @@ async def test_acompletion_mid_stream_fallback_honors_a_per_request_fallbacks_no assert _groups_called(provider_calls) == ["primary"] -def test_refusal_on_the_last_fallback_hop_is_returned_instead_of_raised(): - """LIT-7400 follow-up: a refusal on the final hop of an exhausted list passes through.""" - from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets - - router = litellm.Router( - model_list=[ - {"model_name": "primary", "litellm_params": {"model": "openai/primary-model", "api_key": "fake-key"}}, - {"model_name": "fb1", "litellm_params": {"model": "openai/fb1-model", "api_key": "fake-key"}}, - {"model_name": "fb2", "litellm_params": {"model": "openai/fb2-model", "api_key": "fake-key"}}, - ], - fallbacks=[{"primary": ["fb1", "fb2"]}], - num_retries=0, - ) - - attempted: Final = AttemptedFallbackTargets() - attempted.record("primary") - attempted.record("fb1") - attempted.record("fb2") - kwargs: Final = { - "attempted_targets": attempted, - "metadata": {"model_group": "fb2", "original_model_group": "primary"}, - } - - assert router._refusal_fallback_available("fb2", kwargs) is False - assert ( - router._refusal_fallback_available( - "fb1", {"metadata": {"model_group": "fb1", "original_model_group": "primary"}} - ) - is True - ) - - -def test_completion_streaming_iterator_adopts_fallback_response_headers(): - """LIT-6767, sync counterpart of the fallback-adoption test.""" - from unittest.mock import MagicMock, patch - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - failed_error = MidStreamFallbackError( - message="upstream died before the first chunk", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=True, - ) - - class FailedStream: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - self._response_headers = {"x-request-id": "req-FAILED"} - self._hidden_params = { - "model_id": "failed-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FAILED"}, - "only_on_failed_attempt": "stale", - } - - def __iter__(self): - return self - - def __next__(self): - raise failed_error - - class FallbackStream: - def __init__(self): - self._response_headers = {"x-request-id": "req-FALLBACK"} - self._hidden_params = { - "model_id": "fallback-deployment", - "additional_headers": {"llm_provider-x-request-id": "req-FALLBACK"}, - } - - def __iter__(self): - return iter([]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=FallbackStream()), - ): - result = router._completion_streaming_iterator( - model_response=FailedStream(), - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - assert result._response_headers == {"x-request-id": "req-FAILED"} - for _ in result: - pass - - assert result._response_headers == {"x-request-id": "req-FALLBACK"} - assert result._hidden_params["model_id"] == "fallback-deployment" - assert "only_on_failed_attempt" not in result._hidden_params - - -def test_completion_streaming_iterator_fallback_on_429(): - """Sync streaming: MidStreamFallbackError (429 pre-first-chunk) triggers fallback. - - This is the sync counterpart of test_acompletion_streaming_iterator. - Before this fix, __next__ raised RateLimitError directly and the Router - never got a chance to fall back. - """ - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - rate_limit_error = MidStreamFallbackError( - message="Resource exhausted", - model="gpt-4", - llm_provider="vertex_ai", - generated_content="", - is_pre_first_chunk=True, - ) - - class SyncIteratorImmediateError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - - def __iter__(self): - return self - - def __next__(self): - raise rate_limit_error - - mock_response = SyncIteratorImmediateError() - - # Fallback returns a simple non-streaming response (fallback may not stream) - mock_fallback_response = MagicMock() - mock_fallback_response.__iter__ = MagicMock(return_value=iter([])) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=mock_fallback_response), - ) as mock_fallback: - result = router._completion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - collected_chunks = list(result) - - mock_fallback.assert_awaited_once() - call_kwargs = mock_fallback.await_args.kwargs["kwargs"] - # Pre-first-chunk: should use original messages, no continuation prompt - assert call_kwargs.get("messages") == messages - # Verify original_function is _completion (sync) - assert call_kwargs.get("original_function") == router._completion - - -def test_completion_streaming_iterator_preserves_hidden_params(): - """SyncFallbackStreamWrapper must copy _hidden_params from original response.""" - from unittest.mock import MagicMock - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - mock_response = MagicMock() - mock_response.model = "gpt-4" - mock_response.custom_llm_provider = "openai" - mock_response.logging_obj = MagicMock() - mock_response._hidden_params = { - "litellm_overhead_time_ms": 42.0, - "litellm_call_id": "test-sync-call", - } - mock_response.__iter__ = MagicMock(return_value=iter([])) - - result = router._completion_streaming_iterator( - model_response=mock_response, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "gpt-4", "stream": True}, - ) - - assert hasattr(result, "_hidden_params") - assert result._hidden_params.get("litellm_overhead_time_ms") == 42.0 - assert result._hidden_params.get("litellm_call_id") == "test-sync-call" - - -def test_completion_streaming_iterator_reraises_mid_chunk_error(): - """Sync: MidStreamFallbackError with generated_content and is_pre_first_chunk=False - must be re-raised immediately; the router cannot recover after partial content - has already been sent to the client.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - mid_chunk_error = MidStreamFallbackError( - message="Connection reset", - model="gpt-4", - llm_provider="openai", - generated_content="Hello, I am", - is_pre_first_chunk=False, - ) - - class SyncIteratorMidChunkError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - - def __iter__(self): - return self - - def __next__(self): - raise mid_chunk_error - - mock_response = SyncIteratorMidChunkError() - - result = router._completion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - with pytest.raises(MidStreamFallbackError): - list(result) - - -def test_completion_streaming_iterator_reraises_original_exception_when_available(): - """Sync: when the mid-chunk MidStreamFallbackError wraps a real provider - exception (original_exception), the router must re-raise that original - exception instead of the internal wrapper, so the client sees the - specific error type/code (e.g. RateLimitError) rather than a generic - MidStreamFallbackError.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError, RateLimitError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - original_exception = RateLimitError( - message="rate limited", - llm_provider="vertex_ai", - model="gpt-4", - ) - mid_chunk_error = MidStreamFallbackError( - message="rate limited", - model="gpt-4", - llm_provider="openai", - original_exception=original_exception, - generated_content="Hello, I am", - is_pre_first_chunk=False, - ) - - class SyncIteratorMidChunkError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - - def __iter__(self): - return self - - def __next__(self): - raise mid_chunk_error - - mock_response = SyncIteratorMidChunkError() - - result = router._completion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - with pytest.raises(RateLimitError) as exc_info: - list(result) - assert exc_info.value is original_exception - assert exc_info.value.type == "throttling_error" - assert exc_info.value.code == "429" - - -def test_completion_streaming_iterator_reraises_mid_chunk_error_with_no_text_content(): - """Sync: a reasoning-only chunk sets is_pre_first_chunk=False without populating - generated_content (which only tracks text deltas). The re-raise guard must still - detect this via the raw chunks on the wrapper, or the router silently retries and - the client receives duplicated/inconsistent output.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - from litellm.types.utils import Delta, StreamingChoices - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - mid_chunk_error = MidStreamFallbackError( - message="Connection reset", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=False, - ) - - reasoning_chunk = litellm.ModelResponseStream( - id="chatcmpl-partial-1", - model="gpt-4", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta(reasoning_content="Thinking about the answer", role="assistant"), - ) - ], - ) - - class SyncIteratorNoTextChunkError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [reasoning_chunk] - - def __iter__(self): - return self - - def __next__(self): - raise mid_chunk_error - - mock_response = SyncIteratorNoTextChunkError() - - with patch.object(router, "function_with_fallbacks") as mock_fallback: - result = router._completion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - with pytest.raises(MidStreamFallbackError): - list(result) - - assert not mock_fallback.called, ( - "fallback must not be attempted once any content, text or non-text, has already streamed" - ) - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_pre_first_chunk_skips_continuation(): - """When MidStreamFallbackError has is_pre_first_chunk=True, use original messages.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Hello"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - pre_first_chunk_error = MidStreamFallbackError( - message="429 Resource exhausted", - model="gpt-4", - llm_provider="vertex_ai", - generated_content="", - is_pre_first_chunk=True, - ) - - class AsyncIteratorPreFirstChunkError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [] - - def __aiter__(self): - return self - - async def __anext__(self): - raise pre_first_chunk_error - - mock_response = AsyncIteratorPreFirstChunkError() - - class EmptyAsyncIterator: - def __aiter__(self): - return self - - async def __anext__(self): - raise StopAsyncIteration - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=EmptyAsyncIterator(), - ) as mock_fallback_utils: - iterator = await router._acompletion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - async for _ in iterator: - pass - - assert mock_fallback_utils.called - fallback_kwargs = mock_fallback_utils.call_args.kwargs["kwargs"] - # Pre-first-chunk: should use original messages, no continuation prompt - assert fallback_kwargs["messages"] == messages - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_reraises_mid_chunk_error_with_no_text_content(): - """Async: a reasoning-only chunk sets is_pre_first_chunk=False without populating - generated_content (which only tracks text deltas). The re-raise guard must still - detect this via the raw chunks on the wrapper, or the router silently retries and - the client receives duplicated/inconsistent output.""" - from unittest.mock import MagicMock - - from litellm.exceptions import MidStreamFallbackError - from litellm.types.utils import Delta, StreamingChoices - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key"}, - } - ], - ) - - messages = [{"role": "user", "content": "Test"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - mid_chunk_error = MidStreamFallbackError( - message="Connection reset", - model="gpt-4", - llm_provider="openai", - generated_content="", - is_pre_first_chunk=False, - ) - - reasoning_chunk = litellm.ModelResponseStream( - id="chatcmpl-partial-1", - model="gpt-4", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta(reasoning_content="Thinking about the answer", role="assistant"), - ) - ], - ) - - class AsyncIteratorNoTextChunkError: - def __init__(self): - self.model = "gpt-4" - self.custom_llm_provider = "openai" - self.logging_obj = MagicMock() - self.chunks = [reasoning_chunk] - - def __aiter__(self): - return self - - async def __anext__(self): - raise mid_chunk_error - - mock_response = AsyncIteratorNoTextChunkError() - - with patch.object(router, "async_function_with_fallbacks_common_utils") as mock_fallback_utils: - iterator = await router._acompletion_streaming_iterator( - model_response=mock_response, - messages=messages, - initial_kwargs=initial_kwargs, - ) - - with pytest.raises(MidStreamFallbackError): - async for _ in iterator: - pass - - assert not mock_fallback_utils.called, ( - "fallback must not be attempted once any content, text or non-text, has already streamed" - ) - - -# --------------------------------------------------------------------------- -# Shared helpers for the _aresponses_streaming_iterator test suite. -# --------------------------------------------------------------------------- -def _make_responses_iterator( - *, - chunks=(), - error=None, - bridge=False, - model="gpt-4", - hidden_params=None, - chat_chunks=None, -): - """Build a minimal mock Responses-API streaming iterator. - - Bypasses BaseResponsesAPIStreamingIterator.__init__ but mirrors every - attribute production code reads. Yields *chunks*, then raises *error* - (or StopAsyncIteration). Set bridge=True to inherit from - LiteLLMCompletionStreamingIterator so the wrapper's bridge-path - isinstance check (used by usage extraction) matches. - """ - from litellm.responses.litellm_completion_transformation.streaming_iterator import ( - LiteLLMCompletionStreamingIterator, - ) - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - ) - - base = ( - LiteLLMCompletionStreamingIterator - if bridge - else BaseResponsesAPIStreamingIterator - ) - - class _Iter(base): - def __init__(self): - self._chunks = list(chunks) - self._idx = 0 - self._hidden_params = hidden_params or {} - self.model = model - self.custom_llm_provider = "anthropic" - self.logging_obj = MagicMock() - self.litellm_metadata = None - self.responses_api_provider_config = None - self.finished = False - self.completed_response = None - self.response = None - self.start_time = None - self.request_data = {} - self.call_type = None - if chat_chunks is not None: - self.collected_chat_completion_chunks = chat_chunks - - def __aiter__(self): - return self - - async def __anext__(self): - if self._idx < len(self._chunks): - self._idx += 1 - return self._chunks[self._idx - 1] - if error is not None: - raise error - raise StopAsyncIteration - - return _Iter() - - class _AsyncList: """Generic async iterator over a list — used as the fallback response.""" @@ -6646,106 +3431,6 @@ class _AsyncList: return item -def _make_router_with_fallback(primary="gpt-4", secondary="gpt-3.5-turbo"): - return litellm.Router( - model_list=[ - { - "model_name": primary, - "litellm_params": {"model": primary, "api_key": "k1"}, - }, - { - "model_name": secondary, - "litellm_params": {"model": secondary, "api_key": "k2"}, - }, - ], - fallbacks=[{primary: [secondary]}], - ) - - -# The Rust core walks fallbacks and runs pre-routing in Rust, past the members these tests patch. -_PYTHON_CORE_ONLY: Final = pytest.mark.usefixtures("python_router_core") - - -def _injected_fallback_router(fallback_response: object) -> Router: - router: Final = Router(model_list=[], fallbacks=[{"primary": ["fallback"]}]) - - async def inject_fallback( - e: Exception, - disable_fallbacks: bool | None, - fallbacks: list | None, - context_window_fallbacks: list | None, - content_policy_fallbacks: list | None, - model_group: str | None, - args: tuple[object, ...], - kwargs: dict[str, object], - include_fallback_errors: bool = False, - ) -> object: - return fallback_response - - router.async_function_with_fallbacks_common_utils = inject_fallback - return router - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_fallback(): - """Catches MidStreamFallbackError, re-enters the fallback chain via - async_function_with_fallbacks_common_utils with the per-attempt helper - and original_generic_function preserved. Mirrors - test_acompletion_streaming_iterator for the aresponses path.""" - from litellm.exceptions import MidStreamFallbackError - from litellm.responses.streaming_iterator import ( - BaseResponsesAPIStreamingIterator, - ) - - router = _make_router_with_fallback( - "anthropic/claude-sonnet-4-6", "vertex_ai/claude-sonnet-4-6" - ) - src = _make_responses_iterator( - chunks=[MagicMock(type="response.created")], - error=MidStreamFallbackError( - message="anthropic socket timeout", - model="anthropic/claude-sonnet-4-6", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="", - ), - model="anthropic/claude-sonnet-4-6", - hidden_params={"model_id": "src-deployment-1"}, - ) - fallback_chunks = [ - MagicMock(type="response.created"), - MagicMock(type="response.output_text.delta"), - MagicMock(type="response.completed"), - ] - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList(fallback_chunks), - ) as mock_fallback_utils: - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "anthropic/claude-sonnet-4-6", - "stream": True, - "input": "Hi", - "original_generic_function": litellm.aresponses, - }, - ) - assert isinstance(wrapped, BaseResponsesAPIStreamingIterator) - assert wrapped._hidden_params.get("model_id") == "src-deployment-1" - collected = [c async for c in wrapped] - - assert collected == fallback_chunks - call_kwargs = mock_fallback_utils.call_args.kwargs - fbk = call_kwargs["kwargs"] - # Bound methods compare equal when they share the same instance + __func__. - assert fbk["original_function"] == router._ageneric_api_call_with_fallbacks_responses_attempt - assert fbk["original_generic_function"] is litellm.aresponses - assert call_kwargs["model_group"] == "anthropic/claude-sonnet-4-6" - assert call_kwargs["disable_fallbacks"] is False - - @pytest.mark.asyncio async def test_aresponses_mid_stream_order_fallback_hop_drops_the_encrypted_reasoning_the_next_provider_cannot_decrypt( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch @@ -6975,433 +3660,6 @@ async def test_aresponses_mid_stream_order_fallback_hop_keeps_the_encrypted_reas assert [json.loads(call.request.read())["input"] for call in openai_route.calls] == [history(), history()] -@pytest.mark.asyncio -async def test_aresponses_streaming_content_policy_error_event_routes_to_content_policy_fallback(): - """Regression: a mid-stream content_policy_violation error event never reached - content_policy_fallbacks. The iterator raised a bare APIError the wrapper does not - catch, and even once wrapped, the MidStreamFallbackError envelope was handed to the - fallback dispatch, whose isinstance branch on ContentPolicyViolationError never matched. - The stream below is the customer's shape: a raw OpenAI error event with code - content_policy_violation, transformed by the real OpenAI config, and the router must - call the content_policy_fallbacks target, not the general fallbacks one.""" - from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig - from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator - - router = litellm.Router( - model_list=[ - {"model_name": "primary", "litellm_params": {"model": "openai/gpt-5.4", "api_key": "k1"}}, - { - "model_name": "content-fallback", - "litellm_params": {"model": "gemini/gemini-2.5-flash", "api_key": "k2"}, - }, - {"model_name": "general-fallback", "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "k3"}}, - ], - fallbacks=[{"primary": ["general-fallback"]}], - content_policy_fallbacks=[{"primary": ["content-fallback"]}], - ) - error_event = { - "type": "error", - "sequence_number": 2, - "error": { - "type": "invalid_request_error", - "code": "content_policy_violation", - "message": "This content was flagged for possible cybersecurity risk. The response was halted mid-stream.", - "param": None, - }, - } - - async def aiter_bytes(): - yield f"data: {json.dumps(error_event)}\n\n".encode() - - raw_response = MagicMock() - raw_response.headers = {} - raw_response.aiter_bytes = aiter_bytes - logging_obj = MagicMock(spec=LiteLLMLogging) - logging_obj.model_call_details = {"litellm_params": {}} - logging_obj.completion_start_time = None - source = ResponsesAPIStreamingIterator( - response=raw_response, - model="gpt-5.4", - responses_api_provider_config=OpenAIResponsesAPIConfig(), - logging_obj=logging_obj, - custom_llm_provider="openai", - ) - fallback_chunks = [MagicMock(type="response.output_text.delta"), MagicMock(type="response.completed")] - fallback_call = AsyncMock(return_value=_AsyncList(fallback_chunks)) - - wrapped = await router._aresponses_streaming_iterator( - response=source, - initial_kwargs={ - "model": "primary", - "stream": True, - "input": "Hi", - "original_generic_function": fallback_call, - }, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == fallback_chunks - fallback_call.assert_awaited_once() - assert fallback_call.await_args.kwargs["model"] == "gemini/gemini-2.5-flash" - - -@pytest.mark.asyncio -@_PYTHON_CORE_ONLY -async def test_aresponses_stream_keeps_numbering_after_a_mid_stream_fallback(): - from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents - - def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: - return OutputTextDeltaEvent( - type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, - item_id="msg_1", - output_index=0, - content_index=0, - delta=text, - sequence_number=sequence_number, - ) - - from litellm.exceptions import MidStreamFallbackError - - router: Final = _injected_fallback_router(_AsyncList([delta(0, "c"), delta(1, "d")])) - primary: Final = _make_responses_iterator( - chunks=[delta(0, "a"), delta(1, "b")], - error=MidStreamFallbackError( - message="primary died", - model="primary", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="ab", - ), - model="primary", - ) - - wrapped: Final = await router._aresponses_streaming_iterator( - response=primary, - initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, - ) - events: Final = [event async for event in wrapped] - - assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 1), ("c", 2), ("d", 3)] - - - -@pytest.mark.asyncio -@_PYTHON_CORE_ONLY -async def test_aresponses_stream_fallback_keeps_higher_numbers_and_passes_unnumbered_events(): - from litellm.exceptions import MidStreamFallbackError - from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents - - def delta(text: str) -> OutputTextDeltaEvent: - return OutputTextDeltaEvent( - type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, item_id="msg_1", output_index=0, content_index=0, delta=text - ) - - def numbered(sequence_number: int, text: str) -> OutputTextDeltaEvent: - return delta(text).model_copy(update={"sequence_number": sequence_number}) - - router: Final = _injected_fallback_router(_AsyncList([numbered(7, "c"), delta("d"), numbered(8, "e")])) - primary: Final = _make_responses_iterator( - chunks=[numbered(0, "a"), numbered(1, "b")], - error=MidStreamFallbackError( - message="primary died", - model="primary", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="ab", - ), - model="primary", - ) - - wrapped: Final = await router._aresponses_streaming_iterator( - response=primary, - initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, - ) - events: Final = [event async for event in wrapped] - - assert [(event.delta, event.get("sequence_number")) for event in events] == [ - ("a", 0), - ("b", 1), - ("c", 7), - ("d", None), - ("e", 8), - ] - - -@pytest.mark.asyncio -@_PYTHON_CORE_ONLY -async def test_aresponses_stream_without_fallback_passes_upstream_numbers_unchanged(): - from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents - - def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: - return OutputTextDeltaEvent( - type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, - item_id="msg_1", - output_index=0, - content_index=0, - delta=text, - sequence_number=sequence_number, - ) - - router: Final = _injected_fallback_router(_AsyncList([])) - primary: Final = _make_responses_iterator( - chunks=[delta(0, "a"), delta(1, "b"), delta(1, "c"), delta(0, "d")], model="primary" - ) - - wrapped: Final = await router._aresponses_streaming_iterator( - response=primary, - initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, - ) - events: Final = [event async for event in wrapped] - - assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 1), ("c", 1), ("d", 0)] - - -@pytest.mark.asyncio -@_PYTHON_CORE_ONLY -async def test_aresponses_stream_fallback_continues_above_the_highest_number_the_primary_sent(): - from litellm.exceptions import MidStreamFallbackError - from litellm.types.llms.openai import OutputTextDeltaEvent, ResponsesAPIStreamEvents - - def delta(sequence_number: int, text: str) -> OutputTextDeltaEvent: - return OutputTextDeltaEvent( - type=ResponsesAPIStreamEvents.OUTPUT_TEXT_DELTA, - item_id="msg_1", - output_index=0, - content_index=0, - delta=text, - sequence_number=sequence_number, - ) - - router: Final = _injected_fallback_router(_AsyncList([delta(0, "d")])) - primary: Final = _make_responses_iterator( - chunks=[delta(0, "a"), delta(3, "b"), delta(1, "c")], - error=MidStreamFallbackError( - message="primary died", - model="primary", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="abc", - ), - model="primary", - ) - - wrapped: Final = await router._aresponses_streaming_iterator( - response=primary, - initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": litellm.aresponses}, - ) - events: Final = [event async for event in wrapped] - - assert [(event.delta, event.sequence_number) for event in events] == [("a", 0), ("b", 3), ("c", 1), ("d", 4)] - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_unwraps_content_policy_trigger_for_fallback_dispatch(): - """The fallback dispatch matches on the trigger's own type, so the wrapper must hand it the - ContentPolicyViolationError carried inside MidStreamFallbackError, not the envelope.""" - router = _make_router_with_fallback("openai/gpt-5.4", "openai/gpt-5-mini") - content_policy_error = litellm.ContentPolicyViolationError( - message="flagged mid-stream", llm_provider="openai", model="openai/gpt-5.4" - ) - src = _make_responses_iterator( - chunks=[MagicMock(type="response.created")], - error=MidStreamFallbackError( - message=str(content_policy_error), - model="openai/gpt-5.4", - llm_provider="openai", - original_exception=content_policy_error, - is_pre_first_chunk=True, - ), - model="openai/gpt-5.4", - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_AsyncList([MagicMock(type="response.completed")])), - ) as mock_fallback_utils: - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={"model": "openai/gpt-5.4", "stream": True, "input": "Hi"}, - ) - [chunk async for chunk in wrapped] - - mock_fallback_utils.assert_awaited_once() - assert mock_fallback_utils.await_args.kwargs["e"] is content_policy_error - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "fallback_headers", - [ - {"x-fallback-only": "yes"}, - { - "x-fallback-only": "yes", - "x-litellm-complexity-router-tier": "SIMPLE", - }, - ], - ids=["plain-fallback", "complexity-tier-fallback"], -) -@_PYTHON_CORE_ONLY -async def test_aresponses_streaming_iterator_replaces_complexity_headers_before_fallback_output( - fallback_headers: dict[str, str], -) -> None: - primary_headers: Final = { - "x-litellm-complexity-router-tier": "REASONING", - "x-litellm-complexity-router-reasoning-effort": "xhigh", - } - source: Final = _make_responses_iterator( - error=MidStreamFallbackError( - message="primary failed before output", - model="gpt-4", - llm_provider="openai", - is_pre_first_chunk=True, - generated_content="", - ), - hidden_params={"additional_headers": primary_headers}, - ) - fallback_output: Final = MagicMock(type="response.output_text.delta") - fallback: Final = _AsyncList([fallback_output]) - fallback._hidden_params = { - "model_id": "fallback-deployment", - "additional_headers": fallback_headers, - } - router: Final = _injected_fallback_router(fallback) - - wrapped: Final = await router._aresponses_streaming_iterator( - response=source, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - assert wrapped._hidden_params["additional_headers"] == primary_headers - first_output: Final = await wrapped.__anext__() - - assert first_output is fallback_output - assert wrapped.fallback_headers_adopted is True - assert wrapped._hidden_params == { - "model_id": "fallback-deployment", - "additional_headers": fallback_headers, - } - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_writes_litellm_metadata_on_fallback(): - """Regression: model_group must land under "litellm_metadata" (the key - litellm.aresponses reads), not the default "metadata".""" - from litellm.exceptions import MidStreamFallbackError - - router = _make_router_with_fallback() - src = _make_responses_iterator( - error=MidStreamFallbackError( - message="boom", - model="gpt-4", - llm_provider="anthropic", - is_pre_first_chunk=True, - generated_content="", - ) - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList(), - ) as mock_fallback_utils: - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - async for _ in wrapped: - pass - - fbk = mock_fallback_utils.call_args.kwargs["kwargs"] - assert "litellm_metadata" in fbk, "wrong metadata_variable_name" - assert fbk["litellm_metadata"]["model_group"] == "gpt-4" - assert "model_group" not in fbk.get( - "metadata", {} - ), "model_group leaked into 'metadata' instead of 'litellm_metadata'" - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_pre_first_chunk_skips_continuation(): - """Pre-first-chunk error: original input is preserved unchanged.""" - from litellm.exceptions import MidStreamFallbackError - - router = _make_router_with_fallback() - src = _make_responses_iterator( - error=MidStreamFallbackError( - message="socket timeout before first chunk", - model="gpt-4", - llm_provider="anthropic", - is_pre_first_chunk=True, - generated_content="", - ) - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList(), - ) as mock_fallback_utils: - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - async for _ in wrapped: - pass - - fbk = mock_fallback_utils.call_args.kwargs["kwargs"] - assert fbk["input"] == "Hello" # original input, no continuation messages - - -def _make_native_responses_iterator(*, sse_payloads: tuple[dict[str, str], ...], trailing_error: Exception | None): - """A real ResponsesAPIStreamingIterator over canned SSE bytes, so the router test covers the - iterator's own transport-error classification instead of a hand-built MidStreamFallbackError.""" - from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig - from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator - - async def aiter_bytes(): - for payload in sse_payloads: - yield f"data: {json.dumps(payload)}\n\n".encode() - if trailing_error is not None: - raise trailing_error - - def transform(model, parsed_chunk, logging_obj): - return MagicMock(type=parsed_chunk["type"]) - - response: Final = MagicMock() - response.headers = {} - response.aiter_bytes = aiter_bytes - config: Final = MagicMock(spec=BaseResponsesAPIConfig) - config.transform_streaming_response.side_effect = transform - logging_obj: Final = MagicMock(spec=LiteLLMLogging) - logging_obj.completion_start_time = None - logging_obj.model_call_details = {"litellm_params": {}} - return ResponsesAPIStreamingIterator( - response=response, - model="gpt-4", - responses_api_provider_config=config, - logging_obj=logging_obj, - litellm_metadata={}, - custom_llm_provider="openai", - ) - - -_RESPONSES_LIFECYCLE_PAYLOADS: Final = ({"type": "response.created"}, {"type": "response.in_progress"}) - - async def _events_until_error(stream: AsyncIterable[object]) -> AsyncIterator[object]: try: async for chunk in stream: @@ -7410,144 +3668,6 @@ async def _events_until_error(stream: AsyncIterable[object]) -> AsyncIterator[ob yield error -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_falls_back_on_transport_drop_before_output(): - """A connection lost after response.created but before any output item is re-routed to the - fallback with the original input, the same as a provider error event would be, and the client - sees one response lifecycle: the fallback's, whose id the completed event carries.""" - router: Final = _make_router_with_fallback() - src: Final = _make_native_responses_iterator( - sse_payloads=_RESPONSES_LIFECYCLE_PAYLOADS, - trailing_error=httpx.ReadError("Response payload is not completed"), - ) - fallback_chunks: Final = [ - MagicMock(type="response.created", response=MagicMock(id="resp_fallback")), - MagicMock(type="response.in_progress", response=MagicMock(id="resp_fallback")), - MagicMock(type="response.output_text.delta"), - MagicMock(type="response.completed", response=MagicMock(id="resp_fallback")), - ] - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList(fallback_chunks), - ) as mock_fallback_utils: - wrapped: Final = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - collected: Final = [chunk async for chunk in wrapped] - - assert collected == fallback_chunks - assert [chunk.response.id for chunk in collected if chunk.type == "response.created"] == ["resp_fallback"] - assert isinstance(mock_fallback_utils.call_args.kwargs["e"], MidStreamFallbackError) - assert mock_fallback_utils.call_args.kwargs["kwargs"]["input"] == "Hello" - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_surfaces_transport_drop_when_no_fallback_lands(): - transport_error: Final = httpx.ReadError("Response payload is not completed") - router: Final = _make_router_with_fallback() - src: Final = _make_native_responses_iterator( - sse_payloads=_RESPONSES_LIFECYCLE_PAYLOADS, trailing_error=transport_error - ) - - async def reraise_trigger(**kwargs): - raise kwargs["e"] - - with patch.object( - router, "async_function_with_fallbacks_common_utils", new=AsyncMock(side_effect=reraise_trigger) - ) as mock_fallback_utils: - wrapped: Final = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - outcome: Final = [item async for item in _events_until_error(wrapped)] - - assert [item.type for item in outcome[:-1]] == ["response.created", "response.in_progress"] - assert outcome[-1] is transport_error - assert mock_fallback_utils.await_count == 1 - trigger: Final = mock_fallback_utils.await_args.kwargs["e"] - assert isinstance(trigger, MidStreamFallbackError) - assert trigger.original_exception is transport_error - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_forwards_lifecycle_events_in_order_once_output_starts(): - router: Final = _make_router_with_fallback() - chunks: Final = [ - MagicMock(type="response.created"), - MagicMock(type="response.in_progress"), - MagicMock(type="response.output_text.delta"), - MagicMock(type="response.completed"), - ] - wrapped: Final = await router._aresponses_streaming_iterator( - response=_make_responses_iterator(chunks=chunks), - initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, - ) - - assert [chunk async for chunk in wrapped] == chunks - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_flushes_held_lifecycle_events_when_the_stream_ends_without_output(): - router: Final = _make_router_with_fallback() - chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.in_progress")] - wrapped: Final = await router._aresponses_streaming_iterator( - response=_make_responses_iterator(chunks=chunks), - initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, - ) - - assert [chunk async for chunk in wrapped] == chunks - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_forwards_held_lifecycle_events_before_a_non_fallback_error(): - router: Final = _make_router_with_fallback() - chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.in_progress")] - client_error: Final = litellm.BadRequestError(message="bad input", model="gpt-4", llm_provider="openai") - wrapped: Final = await router._aresponses_streaming_iterator( - response=_make_responses_iterator(chunks=chunks, error=client_error), - initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, - ) - outcome: Final = [item async for item in _events_until_error(wrapped)] - - assert outcome[:-1] == chunks - assert outcome[-1] is client_error - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_commits_held_lifecycle_events_at_the_hold_cap(): - router: Final = _make_router_with_fallback() - chunks: Final = [MagicMock(type="response.in_progress") for _ in range(MAX_HELD_PRE_OUTPUT_RESPONSES_EVENTS + 1)] - src: Final = _make_responses_iterator( - chunks=chunks, - error=MidStreamFallbackError( - message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True - ), - ) - fallback_chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.completed")] - - with patch.object(router, "async_function_with_fallbacks_common_utils", return_value=_AsyncList(fallback_chunks)): - wrapped: Final = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={"model": "gpt-4", "stream": True, "input": "Hello"}, - ) - collected: Final = [chunk async for chunk in wrapped] - - assert collected == [*chunks, *fallback_chunks] - - @pytest.mark.parametrize( ("event_type", "held_event_count", "expected"), [ @@ -7566,103 +3686,6 @@ def test_responses_stream_holds_event_holds_only_pre_output_lifecycle_events_und assert _responses_stream_holds_event(MagicMock(type=event_type), held_event_count) is expected -@pytest.mark.asyncio -async def test_aresponses_fallback_attempt_drops_held_lifecycle_events_when_a_fallback_lands(): - router: Final = _make_router_with_fallback() - trigger: Final = MidStreamFallbackError( - message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True - ) - held: Final = (MagicMock(type="response.created"), MagicMock(type="response.in_progress")) - fallback_chunks: Final = [MagicMock(type="response.created"), MagicMock(type="response.completed")] - adopt_headers: Final = MagicMock(return_value=({}, {})) - - with patch.object( - router, "async_function_with_fallbacks_common_utils", return_value=_AsyncList(fallback_chunks) - ) as mock_fallback_utils: - collected: Final = [ - chunk - async for chunk in router._aresponses_fallback_attempt( - trigger, - _make_responses_iterator(), - {"model": "gpt-4", "stream": True, "input": "Hello"}, - adopt_headers, - held, - ) - ] - - assert collected == fallback_chunks - adopt_headers.assert_called_once() - assert mock_fallback_utils.await_args.kwargs["e"] is trigger - - -@pytest.mark.asyncio -async def test_aresponses_fallback_attempt_replays_held_lifecycle_events_when_the_fallback_dies_before_its_first_event(): - """A fallback stream that raises before yielding anything announced no response of its own, so the - primary's held created/in_progress pair is replayed ahead of the error and the client sees the - announcement the failure belongs to, the same as when no fallback was attempted at all.""" - router: Final = _make_router_with_fallback() - trigger: Final = MidStreamFallbackError( - message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True - ) - held: Final = (MagicMock(type="response.created"), MagicMock(type="response.in_progress")) - fallback_error: Final = RuntimeError("fallback closed before its first event") - adopt_headers: Final = MagicMock(return_value=({}, {})) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_make_responses_iterator(error=fallback_error), - ): - outcome: Final = [ - item - async for item in _events_until_error( - router._aresponses_fallback_attempt( - trigger, - _make_responses_iterator(), - {"model": "gpt-4", "stream": True, "input": "Hello"}, - adopt_headers, - held, - ) - ) - ] - - assert outcome == [*held, fallback_error] - - -@pytest.mark.asyncio -async def test_aresponses_fallback_attempt_does_not_replay_held_lifecycle_events_once_the_fallback_announced_itself(): - """Once the fallback has yielded its own created event, a later failure must not replay the - primary's held pair on top of it, or the client would again see two announced response ids.""" - router: Final = _make_router_with_fallback() - trigger: Final = MidStreamFallbackError( - message="dropped before output", model="gpt-4", llm_provider="openai", is_pre_first_chunk=True - ) - held: Final = (MagicMock(type="response.created"), MagicMock(type="response.in_progress")) - fallback_created: Final = MagicMock(type="response.created") - fallback_error: Final = RuntimeError("fallback dropped after announcing itself") - adopt_headers: Final = MagicMock(return_value=({}, {})) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_make_responses_iterator(chunks=(fallback_created,), error=fallback_error), - ): - outcome: Final = [ - item - async for item in _events_until_error( - router._aresponses_fallback_attempt( - trigger, - _make_responses_iterator(), - {"model": "gpt-4", "stream": True, "input": "Hello"}, - adopt_headers, - held, - ) - ) - ] - - assert outcome == [fallback_created, fallback_error] - - _BACKGROUND_PRIMARY_URL: Final = "https://api.openai.com/v1/responses" _BACKGROUND_BACKUP_URL: Final = "https://backup.openai.test/v1/responses" @@ -7893,7 +3916,6 @@ async def test_aresponses_background_stream_surfaces_a_later_error_instead_of_a_ assert backup_route.call_count == 0 - _BACKGROUND_TEXT_DELTA: Final = { "type": "response.output_text.delta", "sequence_number": 2, @@ -7964,337 +3986,6 @@ async def test_aresponses_background_stream_content_policy_error_reads_the_same_ assert primary_route.call_count == 2 assert backup_route.call_count == 0 -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_partial_content_injects_continuation(): - """Mid-stream error: preserve partial assistant output and request continuation in a user turn.""" - from litellm.exceptions import MidStreamFallbackError - - router = _make_router_with_fallback() - src = _make_responses_iterator( - chunks=[MagicMock(type="response.output_text.delta")], - error=MidStreamFallbackError( - message="socket reset mid-stream", - model="gpt-4", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="The capital of France is", - ), - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList(), - ) as mock_fallback_utils: - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "What's the capital of France?", - "original_generic_function": litellm.aresponses, - }, - ) - async for _ in wrapped: - pass - - new_input = mock_fallback_utils.call_args.kwargs["kwargs"]["input"] - assert isinstance(new_input, list) - assert new_input[0]["role"] == "user" - assert new_input[0]["content"][0]["text"] == "What's the capital of France?" - assert new_input[1]["role"] == "assistant" - assert new_input[1]["content"][0]["type"] == "output_text" - assert new_input[1]["content"][0]["text"] == "The capital of France is" - assert new_input[2]["role"] == "user" - assert "do not repeat" in new_input[2]["content"][0]["text"].lower() - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_combines_partial_usage(): - """Partial usage from the bridge path is normalized to ResponseAPIUsage - and summed onto the fallback's response.completed event — no token-name - split, clean ResponseAPIUsage on output.""" - from types import SimpleNamespace - - from litellm.exceptions import MidStreamFallbackError - from litellm.types.llms.openai import ( - ResponseAPIUsage, - ResponseCompletedEvent, - ResponsesAPIResponse, - ResponsesAPIStreamEvents, - ) - - router = _make_router_with_fallback() - src = _make_responses_iterator( - bridge=True, - chat_chunks=[MagicMock()], - chunks=[MagicMock(type="response.output_text.delta")], - error=MidStreamFallbackError( - message="boom", - model="gpt-4", - llm_provider="anthropic", - is_pre_first_chunk=False, - generated_content="hello", - ), - ) - - fallback_response_object = ResponsesAPIResponse( - id="resp_test", created_at=0, model="gpt-4", object="response", output=[] - ) - fallback_response_object.usage = ResponseAPIUsage( - input_tokens=20, output_tokens=15, total_tokens=35 - ) - fallback_event = ResponseCompletedEvent( - type=ResponsesAPIStreamEvents.RESPONSE_COMPLETED, - response=fallback_response_object, - ) - - with ( - patch( - "litellm.main.stream_chunk_builder", - return_value=SimpleNamespace( - usage=SimpleNamespace(prompt_tokens=10, completion_tokens=4) - ), - ), - patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=_AsyncList([fallback_event]), - ), - ): - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gpt-4", - "stream": True, - "input": "hi", - "original_generic_function": litellm.aresponses, - }, - ) - async for _ in wrapped: - pass - - merged = fallback_response_object.usage - assert isinstance(merged, ResponseAPIUsage) - assert merged.input_tokens == 30 # 10 (translated from prompt_tokens) + 20 - assert merged.output_tokens == 19 # 4 (translated from completion_tokens) + 15 - assert merged.total_tokens == 49 - - -def _midstream_rate_limit_error(): - rate_limit_error = litellm.RateLimitError( - message="vertex_ai_betaException - Resource exhausted.", - model="gemini", - llm_provider="vertex_ai_beta", - ) - midstream_error = MidStreamFallbackError( - message=str(rate_limit_error), - model="gemini", - llm_provider="vertex_ai_beta", - original_exception=rate_limit_error, - is_pre_first_chunk=True, - ) - return rate_limit_error, midstream_error - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks(): - """Regression for #26015: a mid-stream 429 with no fallbacks configured must - surface a clean RateLimitError, not leak the internal MidStreamFallbackError - wrapper to the client, and must terminate instead of hanging.""" - rate_limit_error, midstream_error = _midstream_rate_limit_error() - - router = litellm.Router( - model_list=[ - { - "model_name": "gemini", - "litellm_params": { - "model": "vertex_ai/gemini-2.0-flash", - "api_key": "fake-key", - }, - }, - ], - num_retries=0, - ) - - class _RaisingStream: - def __init__(self): - self.chunks = [] - - def __aiter__(self): - return self - - async def __anext__(self): - raise midstream_error - - stream = _RaisingStream() - setattr(stream, "model", "gemini") - setattr(stream, "custom_llm_provider", "vertex_ai_beta") - setattr(stream, "logging_obj", MagicMock()) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=midstream_error), - ): - result = await router._acompletion_streaming_iterator( - model_response=stream, - messages=[{"role": "user", "content": "Hello"}], - initial_kwargs={"model": "gemini", "stream": True}, - ) - - async def _consume(): - async for _ in result: - pass - - with pytest.raises(litellm.RateLimitError) as exc_info: - await asyncio.wait_for(_consume(), timeout=10) - - assert not isinstance(exc_info.value, MidStreamFallbackError) - assert exc_info.value.status_code == 429 - assert exc_info.value is rate_limit_error - - -def test_completion_streaming_iterator_surfaces_rate_limit_without_fallbacks(): - """Sync counterpart of - test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks.""" - rate_limit_error, midstream_error = _midstream_rate_limit_error() - - router = litellm.Router( - model_list=[ - { - "model_name": "gemini", - "litellm_params": { - "model": "vertex_ai/gemini-2.0-flash", - "api_key": "fake-key", - }, - }, - ], - num_retries=0, - ) - - class _RaisingSyncStream: - def __init__(self): - self.model = "gemini" - self.custom_llm_provider = "vertex_ai_beta" - self.logging_obj = MagicMock() - self.chunks = [] - - def __iter__(self): - return self - - def __next__(self): - raise midstream_error - - with patch.object( - router, - "function_with_fallbacks", - side_effect=midstream_error, - ): - result = router._completion_streaming_iterator( - model_response=_RaisingSyncStream(), - messages=[{"role": "user", "content": "Hello"}], - initial_kwargs={"model": "gemini", "stream": True}, - ) - - with pytest.raises(litellm.RateLimitError) as exc_info: - list(result) - - assert not isinstance(exc_info.value, MidStreamFallbackError) - assert exc_info.value.status_code == 429 - assert exc_info.value is rate_limit_error - - -@pytest.mark.asyncio -async def test_aresponses_streaming_iterator_surfaces_rate_limit_without_fallbacks(): - """Responses-API counterpart of - test_acompletion_streaming_iterator_surfaces_rate_limit_without_fallbacks.""" - rate_limit_error, midstream_error = _midstream_rate_limit_error() - - router = litellm.Router( - model_list=[ - { - "model_name": "gemini", - "litellm_params": { - "model": "vertex_ai/gemini-2.0-flash", - "api_key": "fake-key", - }, - }, - ], - num_retries=0, - ) - src = _make_responses_iterator(error=midstream_error, model="gemini") - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=midstream_error), - ): - wrapped = await router._aresponses_streaming_iterator( - response=src, - initial_kwargs={ - "model": "gemini", - "stream": True, - "input": "Hello", - "original_generic_function": litellm.aresponses, - }, - ) - - async def _consume(): - async for _ in wrapped: - pass - - with pytest.raises(litellm.RateLimitError) as exc_info: - await asyncio.wait_for(_consume(), timeout=10) - - assert not isinstance(exc_info.value, MidStreamFallbackError) - assert exc_info.value.status_code == 429 - assert exc_info.value is rate_limit_error - - -@pytest.mark.asyncio -async def test_async_function_with_fallbacks_common_utils(): - """Test the async_function_with_fallbacks_common_utils method""" - # Create a basic router for testing - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "gpt-3.5-turbo", - }, - } - ], - max_fallbacks=5, - ) - - # Test case 1: disable_fallbacks=True should raise original exception - test_exception = Exception("Test error") - with pytest.raises(Exception, match="Test error"): - await router.async_function_with_fallbacks_common_utils( - e=test_exception, - disable_fallbacks=True, - fallbacks=None, - context_window_fallbacks=None, - content_policy_fallbacks=None, - model_group="gpt-3.5-turbo", - args=(), - kwargs=MagicMock(), - ) - - # Test case 2: original_model_group=None should raise original exception - with pytest.raises(Exception, match="Test error"): - await router.async_function_with_fallbacks_common_utils( - e=test_exception, - disable_fallbacks=False, - fallbacks=None, - context_window_fallbacks=None, - content_policy_fallbacks=None, - model_group="gpt-3.5-turbo", - args=(), - kwargs={}, # No model key - ) - def test_should_include_deployment(): """Test that Router._core_should_include_deployment returns the correct response""" @@ -8340,74 +4031,6 @@ def test_should_include_deployment(): ) -def test_pre_call_checks_skips_token_count_without_max_input_tokens(monkeypatch): - """ - tiktoken token counting is the dominant on-loop cost for large prompts. When no - deployment in the group declares max_input_tokens, the count is never consumed, so - _pre_call_checks must not run it at all. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {}) - - calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d2"}}, - ] - result = router._pre_call_checks( - model="m", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "hi"}], - ) - - assert calls == [] - assert len(result) == 2 - - -def test_pre_call_checks_counts_once_and_filters_on_max_input_tokens(monkeypatch): - """ - When a deployment declares max_input_tokens the count must still run, be performed - at most once across the group (memoized), and filter deployments whose limit is - exceeded. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d2"}}, - ] - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "hi"}], - ) - - assert calls == [1] - - _MULTILINGUAL_PROMPT: Final = [{"role": "user", "content": "日本語のテキストとहिन्दी पाठ और русский текст. " * 20}] @@ -8465,21 +4088,6 @@ def _embedding_context_router(backend: str, window: int = _EMBEDDING_CONTEXT_WIN ) -async def test_context_window_pre_call_check_counts_an_embedding_image_as_media_on_an_image_backend(): - router: Final = _embedding_context_router("vertex_ai/gemini-embedding-2") - deployments: Final = router.registry().deployments("embedder") - - kept: Final = router._pre_call_checks(model="embedder", healthy_deployments=deployments, input=_EMBEDDING_IMAGE_INPUT) - counted: Final = await router._acount_pre_call_check_tokens( - model="embedder", healthy_deployments=deployments, messages=None, input=_EMBEDDING_IMAGE_INPUT, request_kwargs=None - ) - - assert len(kept) == 1 - assert counted is not None and 0 < counted < _EMBEDDING_CONTEXT_WINDOW - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks(model="embedder", healthy_deployments=deployments, input="word " * 20_000) - - @pytest.mark.parametrize( ("backend", "shape"), [ @@ -8515,24 +4123,12 @@ def test_context_window_pre_call_check_counts_an_embedding_image_as_text_on_a_te with pytest.raises(litellm.ContextWindowExceededError): router._pre_call_checks( - model="embedder", healthy_deployments=router.registry().deployments("embedder"), input=_EMBEDDING_IMAGE_INPUT + model="embedder", + healthy_deployments=router.registry().deployments("embedder"), + input=_EMBEDDING_IMAGE_INPUT, ) -async def test_acount_pre_call_check_tokens_counts_a_group_of_one_openai_model_with_that_models_tokenizer(): - router: Final = _context_limited_router("openai/gpt-6.1-sol") - - count: Final = await router._acount_pre_call_check_tokens( - model="prod", - healthy_deployments=router.get_model_list(model_name="prod"), - messages=_MULTILINGUAL_PROMPT, - input=None, - request_kwargs=None, - ) - - assert count == litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT) - - async def test_a_wildcard_route_checks_the_context_window_with_the_requested_openai_models_tokenizer(): limit: Final = ( litellm.token_counter(model="openai/gpt-6.1-sol", messages=_MULTILINGUAL_PROMPT) @@ -8559,485 +4155,9 @@ async def test_a_wildcard_route_checks_the_context_window_with_the_requested_ope assert (sync_response.choices[0].message.content, async_response.choices[0].message.content) == ("ok", "ok") -def test_pre_call_checks_uses_precounted_tokens(monkeypatch): - """ - An async caller counts off the event loop and passes the result in. _pre_call_checks - must filter on that count instead of re-counting on the loop. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1 - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "hi"}], - input_token_count=1000, - ) - - assert calls == [] - - -async def test_async_get_healthy_deployments_counts_tokens_off_the_event_loop(monkeypatch): - """ - The async deployment path must hand _pre_call_checks a count taken in a worker thread, - so a multi-MB prompt never blocks the proxy during deployment selection. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1_000_000} - ) - - counting_threads = [] - monkeypatch.setattr( - litellm, - "token_counter", - lambda *a, **k: counting_threads.append(threading.current_thread()) or 42, - ) - - counts_passed_in = [] - original_pre_call_checks = router._pre_call_checks - - def spy(**kwargs): - counts_passed_in.append(kwargs.get("input_token_count")) - return original_pre_call_checks(**kwargs) - - monkeypatch.setattr(router, "_pre_call_checks", spy) - - result = await router.async_get_healthy_deployments( - model="m", - request_kwargs={}, - messages=[{"role": "user", "content": "hi"}], - input=None, - specific_deployment=False, - parent_otel_span=None, - ) - - assert len(result) == 1 - assert counts_passed_in == [42] - assert len(counting_threads) == 1 - assert counting_threads[0] is not threading.current_thread() - - -@pytest.mark.parametrize( - "model_info,expected", - [ - ({"max_input_tokens": 100}, True), - ({"max_input_tokens": None}, False), - ({}, False), - ], -) -def test_pre_call_checks_need_token_count(monkeypatch, model_info, expected): - """Only a deployment that declares an integer context window makes a token count worth taking.""" - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: model_info) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - assert router._pre_call_checks_need_token_count("m", deployments) is expected - - -def test_deployment_max_input_tokens_survives_an_unmappable_deployment(monkeypatch): - """ - _pre_call_checks skips a deployment it cannot resolve and carries on. The off-loop - pre-count must do the same, or an unmapped first deployment hides the limit declared by - a later one and the count lands back on the event loop. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - - def flaky_model_info(deployment, received_model_name, id=None): - if deployment["model_info"]["id"] == "unmapped": - raise ValueError("This model isn't mapped yet.") - return {"max_input_tokens": 100} - - monkeypatch.setattr(router, "_core_get_router_model_info", flaky_model_info) - - unmapped = {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "unmapped"}} - mapped = {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "mapped"}} - - assert router._deployment_max_input_tokens("m", unmapped) is None - assert router._deployment_max_input_tokens("m", mapped) == 100 - assert router._pre_call_checks_need_token_count("m", [unmapped, mapped]) is True - - -def test_pre_call_checks_does_not_recount_inline_after_an_off_loop_failure(monkeypatch): - """ - When the off-loop count failed there is nothing left to filter on, so _pre_call_checks must - return the deployments unfiltered rather than repeating the count on the event loop. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - result = router._pre_call_checks( - model="m", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "hi"}], - input_token_count=None, - skip_inline_token_count=True, - ) - - assert calls == [] - assert len(result) == 1 - - -async def test_async_get_healthy_deployments_never_recounts_on_the_loop(monkeypatch): - """ - An off-loop count that raises must not send the same work back onto the event loop through - _pre_call_checks' inline fallback. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - counting_threads = [] - - def exploding_counter(*args, **kwargs): - counting_threads.append(threading.current_thread()) - raise ValueError("Invalid content item type: image") - - monkeypatch.setattr(litellm, "token_counter", exploding_counter) - - result = await router.async_get_healthy_deployments( - model="m", - request_kwargs={}, - messages=[{"role": "user", "content": "hi"}], - input=None, - specific_deployment=False, - parent_otel_span=None, - ) - - assert len(result) == 1 - assert len(counting_threads) == 1 - assert counting_threads[0] is not threading.current_thread() - - -async def test_acount_pre_call_check_tokens_leaves_the_event_loop_free(monkeypatch): - """ - A multi-MB prompt must not stall the proxy: a competing coroutine has to get - scheduled while the router's context-window count is in flight. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - ran = [] - - async def competitor(): - ran.append("competitor") - - task = asyncio.create_task(competitor()) - count = await router._acount_pre_call_check_tokens( - model="m", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "A" * 512 * 1024}], - input=None, - request_kwargs=None, - ) - ran.append("count") - await task - - assert count is not None and count > 0 - assert ran == ["competitor", "count"] - - -async def test_acount_pre_call_check_tokens_skips_without_max_input_tokens(monkeypatch): - """No deployment limits its context window, so there is nothing to count.""" - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {}) - - calls = [] - monkeypatch.setattr( - litellm, "token_counter", lambda *a, **k: calls.append(1) or 1000 - ) - - count = await router._acount_pre_call_check_tokens( - model="m", - healthy_deployments=[ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ], - messages=[{"role": "user", "content": "hi"}], - input=None, - request_kwargs=None, - ) - - assert count is None - assert calls == [] - - -def test_pre_call_checks_counts_tokens_from_responses_input_string(monkeypatch): - """ - Responses API calls pass `input` (str) instead of `messages`. Context-window - checks must count tokens from `input` and filter deployments over the limit. Uses - the real token_counter so the transform + counting path is a true regression guard. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - input="a very long prompt that exceeds the tiny context window", - ) - - -def test_pre_call_checks_counts_tokens_from_responses_input_list(monkeypatch): - """ - Responses API `input` can be a list of input items. It must be normalized to - chat messages and counted so oversized requests are filtered out. Uses the real - token_counter (no mock) so the transform + counting path is a true regression guard. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 1} - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - input=[ - {"role": "user", "content": "count these tokens against the one token limit please"}, - ], - ) - - -def test_pre_call_checks_counts_responses_instructions_tokens(monkeypatch): - """ - Responses API `instructions` become a system message the model receives, so their - tokens must be counted too. A request whose `input` alone fits under the limit but - whose `input` + `instructions` exceeds it must be filtered (regression for the - context-window check under-filtering when instructions were ignored). - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - - short_input = "hi" - long_instructions = "you are a helpful assistant. " * 20 - - input_only_tokens = router._count_pre_call_check_tokens(messages=None, input=short_input) - with_instructions_tokens = router._count_pre_call_check_tokens( - messages=None, input=short_input, request_kwargs={"instructions": long_instructions} - ) - assert with_instructions_tokens > input_only_tokens - - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": input_only_tokens} - ) - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - input=short_input, - request_kwargs={"instructions": long_instructions}, - ) - - _OVERSIZED_TOOL_DESCRIPTION = "look up the answer in the knowledge base. " * 40 -@pytest.mark.parametrize( - "prompt_kwargs, tool", - [ - pytest.param( - {"messages": [{"role": "user", "content": "hi"}]}, - { - "type": "function", - "function": { - "name": "lookup", - "description": _OVERSIZED_TOOL_DESCRIPTION, - "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, - }, - }, - id="chat_completions_tool", - ), - pytest.param( - {"input": "hi"}, - { - "type": "function", - "name": "lookup", - "description": _OVERSIZED_TOOL_DESCRIPTION, - "parameters": {"type": "object", "properties": {"q": {"type": "string"}}}, - }, - id="responses_tool", - ), - pytest.param( - {"messages": [{"role": "user", "content": "hi"}]}, - { - "name": "lookup", - "description": _OVERSIZED_TOOL_DESCRIPTION, - "input_schema": {"type": "object", "properties": {"q": {"type": "string"}}}, - }, - id="anthropic_messages_tool", - ), - ], -) -def test_pre_call_checks_counts_tool_definition_tokens(monkeypatch, prompt_kwargs, tool): - """ - Tool definitions are sent to the model as prompt tokens but never appear in - `messages` or `input`. A request whose prompt alone fits the context window but - whose prompt plus `tools` exceeds it must be rejected before dispatch, for the - Chat Completions, Responses and Anthropic Messages tool shapes alike. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - - prompt_only_tokens = router._count_pre_call_check_tokens( - messages=prompt_kwargs.get("messages"), input=prompt_kwargs.get("input") - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": prompt_only_tokens} - ) - - assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, **prompt_kwargs)) == 1 - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - request_kwargs={"tools": [tool]}, - **prompt_kwargs, - ) - - -@pytest.mark.parametrize( - "system", - [ - pytest.param("You are a meticulous assistant. " * 40, id="system_string"), - pytest.param( - [{"type": "text", "text": "You are a meticulous assistant. " * 40}], - id="system_blocks", - ), - ], -) -def test_pre_call_checks_counts_anthropic_system_tokens(monkeypatch, system): - """ - The Anthropic Messages API carries the system prompt as a top-level `system` field, - not as a message. Its tokens reach the model, so a request whose `messages` fit but - whose `messages` plus `system` exceed the context window must be rejected. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - messages = [{"role": "user", "content": "hi"}] - - messages_only_tokens = router._count_pre_call_check_tokens(messages=messages, input=None) - monkeypatch.setattr(router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": messages_only_tokens}) - - assert len(router._pre_call_checks(model="m", healthy_deployments=deployments, messages=messages)) == 1 - with pytest.raises(litellm.ContextWindowExceededError): - router._pre_call_checks( - model="m", - healthy_deployments=deployments, - messages=messages, - request_kwargs={"system": system}, - ) - - @pytest.mark.asyncio async def test_aanthropic_messages_enforces_context_window_with_system_and_tools(): """ @@ -9106,7 +4226,7 @@ def test_count_pre_call_check_tokens_across_api_surfaces(): assert string_input_tokens > 0 assert list_input_tokens > 0 - with pytest.raises(ValueError, match='Either messages or input must be provided to count tokens'): + with pytest.raises(ValueError, match="Either messages or input must be provided to count tokens"): router._count_pre_call_check_tokens(messages=None, input=None) @@ -9133,37 +4253,6 @@ def test_count_pre_call_check_tokens_skips_the_history_a_compaction_item_replace ) -def test_pre_call_checks_no_messages_or_input_does_not_crash(monkeypatch): - """ - When neither messages nor input is provided (e.g. endpoints without prompt text), - token counting is skipped gracefully and all deployments are returned. - """ - router = litellm.Router( - model_list=[ - {"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo"}}, - ], - enable_pre_call_checks=True, - ) - monkeypatch.setattr( - router, "_core_get_router_model_info", lambda **kwargs: {"max_input_tokens": 5} - ) - - counted: list[dict] = [] - original = router._count_pre_call_check_tokens - monkeypatch.setattr( - router, - "_count_pre_call_check_tokens", - lambda **kwargs: counted.append(kwargs) or original(**kwargs), - ) - - deployments = [ - {"litellm_params": {"model": "gpt-3.5-turbo"}, "model_info": {"id": "d1"}}, - ] - result = router._pre_call_checks(model="m", healthy_deployments=deployments) - assert len(result) == 1 - assert counted == [] # token counting skipped entirely, so no misleading error is logged - - @pytest.mark.asyncio @pytest.mark.usefixtures("local_model_cost_map", "isolate_router_model_cost_state") async def test_aresponses_enforces_context_window_pre_call_check(): @@ -9237,9 +4326,7 @@ def test_get_deployment_model_info_base_model_flow(): } # Test Case 1: Base model flow with custom model info that has base_model - with patch.object( - litellm, "model_cost", {"test-custom-model": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"test-custom-model": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: # Configure mock returns mock_get_model_info.side_effect = lambda model: { @@ -9247,15 +4334,11 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model", model_name="test-model") # Verify that get_model_info was called for both base model and model name assert mock_get_model_info.call_count == 2 - mock_get_model_info.assert_any_call( - model="gpt-3.5-turbo" - ) # base model call + mock_get_model_info.assert_any_call(model="gpt-3.5-turbo") # base model call mock_get_model_info.assert_any_call(model="test-model") # model name call # Verify the result contains merged information @@ -9266,26 +4349,18 @@ def test_get_deployment_model_info_base_model_flow(): # 2. The result of step 1 gets merged into litellm_model_name_info (custom+base override litellm) # Fields from custom model (should override base model values) - assert ( - result["input_cost_per_token"] == 0.001 - ) # From custom model (overrides base 0.0015) - assert ( - result["output_cost_per_token"] == 0.002 - ) # From custom model (same as base) + assert result["input_cost_per_token"] == 0.001 # From custom model (overrides base 0.0015) + assert result["output_cost_per_token"] == 0.002 # From custom model (same as base) assert result["custom_field"] == "custom_value" # From custom model # Fields from base model that weren't overridden by custom assert result["max_tokens"] == 4096 # From base model assert result["litellm_provider"] == "openai" # From base model - assert ( - result["mode"] == "chat" - ) # From base model (overrides litellm "completion") + assert result["mode"] == "chat" # From base model (overrides litellm "completion") # The key field comes from base model since both base and litellm have it # and base model info overrides litellm model name info in final merge - assert ( - result["key"] == "gpt-3.5-turbo" - ) # From base model (overrides litellm key) + assert result["key"] == "gpt-3.5-turbo" # From base model (overrides litellm key) # Test Case 2: Custom model info without base_model mock_custom_model_info_no_base = { @@ -9304,9 +4379,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="test-custom-model-no-base", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-no-base", model_name="test-model") # Should only call get_model_info once for model name (no base model) assert mock_get_model_info.call_count == 1 @@ -9326,9 +4399,7 @@ def test_get_deployment_model_info_base_model_flow(): "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="non-existent-model", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="non-existent-model", model_name="test-model") # Should only call get_model_info once for model name assert mock_get_model_info.call_count == 1 @@ -9361,9 +4432,7 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = mock_get_model_info_side_effect - result = router.get_deployment_model_info( - model_id="test-custom-model-invalid", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="test-custom-model-invalid", model_name="test-model") # Should handle exception gracefully and still return merged result assert result is not None @@ -9372,12 +4441,8 @@ def test_get_deployment_model_info_base_model_flow(): # Test Case 5: Both model_cost.get() and get_model_info() return None with patch.object(litellm, "model_cost", {}): - with patch.object( - litellm, "get_model_info", side_effect=Exception("Not found") - ): - result = router.get_deployment_model_info( - model_id="non-existent", model_name="non-existent" - ) + with patch.object(litellm, "get_model_info", side_effect=Exception("Not found")): + result = router.get_deployment_model_info(model_id="non-existent", model_name="non-existent") # Should return None when no model info is found assert result is None @@ -9400,9 +4465,7 @@ def test_get_deployment_model_info_base_model_flow(): # Model NOT in built-in cost map — raise exception mock_get_model_info.side_effect = Exception("Model not in cost map") - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="unknown-model") # Should return custom_model_info even when litellm_model_name_model_info is None assert result is not None @@ -9438,15 +4501,11 @@ def test_get_deployment_model_info_base_model_flow(): mock_get_model_info.side_effect = get_info_side_effect - result = router.get_deployment_model_info( - model_id="custom-with-base", model_name="unknown-model" - ) + result = router.get_deployment_model_info(model_id="custom-with-base", model_name="unknown-model") # Should return custom_model_info merged with base model info assert result is not None - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom (overrides base) + assert result["input_cost_per_token"] == 0.01 # From custom (overrides base) assert result["max_tokens"] == 8192 # From base model assert result["litellm_provider"] == "openai" # From base model @@ -9493,18 +4552,14 @@ def test_get_deployment_model_info_base_model_merge_priority(): "litellm_only_field": "litellm_value", } - with patch.object( - litellm, "model_cost", {"custom-model-id": mock_custom_model_info} - ): + with patch.object(litellm, "model_cost", {"custom-model-id": mock_custom_model_info}): with patch.object(litellm, "get_model_info") as mock_get_model_info: mock_get_model_info.side_effect = lambda model: { "gpt-4": mock_base_model_info, "test-model": mock_litellm_model_name_info, }.get(model) - result = router.get_deployment_model_info( - model_id="custom-model-id", model_name="test-model" - ) + result = router.get_deployment_model_info(model_id="custom-model-id", model_name="test-model") assert result is not None @@ -9514,29 +4569,17 @@ def test_get_deployment_model_info_base_model_merge_priority(): # 3. Result from steps 1-2 overrides litellm_model_name_info # Fields that should come from custom model info (highest priority) - assert ( - result["input_cost_per_token"] == 0.01 - ) # From custom model (overrides base 0.03) - assert ( - result["max_tokens"] == 8000 - ) # From custom model (overrides base 4096) + assert result["input_cost_per_token"] == 0.01 # From custom model (overrides base 0.03) + assert result["max_tokens"] == 8000 # From custom model (overrides base 4096) assert result["custom_only_field"] == "custom_value" # From custom model # Fields that should come from base model (not overridden by custom) - assert ( - result["output_cost_per_token"] == 0.06 - ) # From base model (not in custom) - assert ( - result["litellm_provider"] == "openai" - ) # From base model (not in custom) - assert ( - result["base_only_field"] == "base_value" - ) # From base model (not in custom) + assert result["output_cost_per_token"] == 0.06 # From base model (not in custom) + assert result["litellm_provider"] == "openai" # From base model (not in custom) + assert result["base_only_field"] == "base_value" # From base model (not in custom) # Fields that should come from litellm model name info (not overridden by custom+base) - assert ( - result["mode"] == "completion" - ) # From litellm model name info (not in custom or base) + assert result["mode"] == "completion" # From litellm model name info (not in custom or base) assert ( result["litellm_only_field"] == "litellm_value" ) # From litellm model name info (not in custom or base) @@ -9548,226 +4591,6 @@ def test_get_deployment_model_info_base_model_merge_priority(): print("✓ Base model merge priority test passed!") -@pytest.mark.parametrize( - "model, litellm_params, endpoint, expected", - [ - ( - "gpt", - {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://my-resource.services.ai.azure.com", "api_key": "key"}, - "gpt/openai/deployments/gpt-5.4-mini/chat/completions", - "gpt-5.4-mini/openai/deployments/gpt-5.4-mini/chat/completions", - ), - ( - "aws/anthropic/bedrock-claude", - {"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0"}, - "/model/aws/anthropic/bedrock-claude/invoke", - "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke", - ), - ( - "my-gemini", - {"model": "gemini/gemini-3.1-pro-preview", "api_key": "key"}, - "v1beta/models/my-gemini:streamGenerateContent", - "v1beta/models/gemini-3.1-pro-preview:streamGenerateContent", - ), - ], -) -def test_add_deployment_model_to_endpoint_rewrites_the_model_group_only_as_whole_path_segments( - model, litellm_params, endpoint, expected -): - router = litellm.Router(model_list=[{"model_name": model, "litellm_params": litellm_params}]) - - result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs={"endpoint": endpoint}, model=model, model_name=litellm_params["model"] - ) - - assert result["endpoint"] == expected - - -def test_add_deployment_model_to_endpoint_for_llm_passthrough_route(): - """ - Test that _add_deployment_model_to_endpoint_for_llm_passthrough_route correctly strips bedrock provider prefix - """ - router = litellm.Router( - model_list=[ - { - "model_name": "special-bedrock-model", - "litellm_params": { - "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - }, - } - ], - ) - - # Test Case 1: Bedrock model with provider prefix - should strip "bedrock/" prefix - kwargs = { - "endpoint": "/model/special-bedrock-model/invoke", - "custom_llm_provider": "bedrock", - } - result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs=kwargs, - model="special-bedrock-model", - model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke" - ), f"Expected '/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke', got '{result['endpoint']}'" - - # Test Case 2: Bedrock invoke-with-response-stream endpoint - kwargs = { - "endpoint": "/model/special-bedrock-model/invoke-with-response-stream", - "custom_llm_provider": "bedrock", - } - result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs=kwargs, - model="special-bedrock-model", - model_name="bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", - ) - assert ( - result["endpoint"] - == "/model/us.anthropic.claude-haiku-4-5-20251001-v1:0/invoke-with-response-stream" - ), f"Expected streaming endpoint with stripped prefix, got '{result['endpoint']}'" - - # Test Case 3: Bedrock converse endpoint - kwargs = { - "endpoint": "/model/bedrock-model/converse", - "custom_llm_provider": "bedrock", - } - result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs=kwargs, - model="bedrock-model", - model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", - ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/converse" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/converse', got '{result['endpoint']}'" - - # Test Case 4: Bedrock provider prefix auto-detected from model_name - kwargs = { - "endpoint": "/model/router-model/invoke", - } - result = router._add_deployment_model_to_endpoint_for_llm_passthrough_route( - kwargs=kwargs, - model="router-model", - model_name="bedrock/us.meta.llama3-8b-instruct-v1:0", - ) - assert ( - result["endpoint"] == "/model/us.meta.llama3-8b-instruct-v1:0/invoke" - ), f"Expected '/model/us.meta.llama3-8b-instruct-v1:0/invoke', got '{result['endpoint']}'" - - -def test_update_kwargs_with_deployment_uses_pass_through_request_timeout(): - router = litellm.Router( - model_list=[ - { - "model_name": "my-bedrock-model", - "litellm_params": { - "model": "bedrock/us.anthropic.claude-opus-4-5-20251101-v1:0", - }, - } - ], - ) - deployment = router.model_list[0] - kwargs: dict = {} - - with patch( - "litellm.proxy.proxy_server.general_settings", - {"pass_through_request_timeout": 6}, - ): - router._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="_ageneric_api_call_with_fallbacks", - ) - - assert kwargs["timeout"] == 6.0 - - -def _passthrough_timeout(router: litellm.Router, deployment: dict, stream: bool) -> float: - kwargs: Final[dict] = {"stream": stream} - router._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - function_name="_ageneric_api_call_with_fallbacks", - ) - return kwargs["timeout"] - - -def test_update_kwargs_with_deployment_passthrough_honors_stream_timeout(): - router = litellm.Router( - model_list=[ - { - "model_name": "anthropic-with-stream-timeout", - "litellm_params": { - "model": "anthropic/claude-sonnet-4-5", - "api_key": "fake-key", - "timeout": 60, - "stream_timeout": 1800, - }, - }, - { - "model_name": "anthropic-router-default", - "litellm_params": { - "model": "anthropic/claude-sonnet-4-5", - "api_key": "fake-key", - "timeout": 60, - }, - }, - ], - timeout=120, - stream_timeout=900, - ) - per_deployment, router_default = router.model_list - - assert _passthrough_timeout(router, per_deployment, stream=True) == 1800.0 - assert _passthrough_timeout(router, router_default, stream=True) == 900.0 - assert _passthrough_timeout(router, per_deployment, stream=False) == 60.0 - assert _passthrough_timeout(router, router_default, stream=False) == 60.0 - - -def test_update_kwargs_with_deployment_passthrough_router_stream_timeout_sources(): - deployment: Final[dict] = { - "model_name": "anthropic-router-default", - "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "fake-key"}, - } - string_router = litellm.Router(model_list=[deployment], timeout=120, stream_timeout="900") - default_router = litellm.Router( - model_list=[deployment], - timeout=120, - default_litellm_params={"stream_timeout": 700}, - ) - - assert _passthrough_timeout(string_router, string_router.model_list[0], stream=True) == 900.0 - assert _passthrough_timeout(default_router, default_router.model_list[0], stream=True) == 700.0 - assert _passthrough_timeout(default_router, default_router.model_list[0], stream=False) == 120.0 - - -def test_update_kwargs_with_deployment_passthrough_honors_global_request_timeout(monkeypatch: pytest.MonkeyPatch): - """litellm_settings.request_timeout must bound the native responses route when neither the - deployment nor the router carries a timeout, while a deployment timeout keeps winning.""" - monkeypatch.setattr("litellm.request_timeout", 44.0, raising=False) - monkeypatch.setattr("litellm.request_timeout_explicitly_set", True, raising=False) - router: Final = litellm.Router( - model_list=[ - { - "model_name": "responses-global-timeout", - "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake-key"}, - }, - { - "model_name": "responses-deployment-timeout", - "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "fake-key", "timeout": 3}, - }, - ], - ) - global_only, per_deployment = router.model_list - - with patch("litellm.proxy.proxy_server.general_settings", {"pass_through_request_timeout": 6}): - assert _passthrough_timeout(router, global_only, stream=True) == 44.0 - assert _passthrough_timeout(router, global_only, stream=False) == 44.0 - assert _passthrough_timeout(router, per_deployment, stream=True) == 3.0 - assert _passthrough_timeout(router, per_deployment, stream=False) == 3.0 - - @pytest.mark.asyncio async def test_router_acompletion_with_unknown_model_and_default_fallback(): """ @@ -9790,14 +4613,10 @@ async def test_router_acompletion_with_unknown_model_and_default_fallback(): # Initialize the router with a default fallback router = litellm.Router(model_list=model_list, default_fallbacks=["gpt-4o"]) - messages = [ - {"role": "user", "content": "This call should succeed by falling back."} - ] + messages = [{"role": "user", "content": "This call should succeed by falling back."}] # Call completion with a model name that is NOT in the model_list - response = await router.acompletion( - model="completely-unknown-model", messages=messages - ) + response = await router.acompletion(model="completely-unknown-model", messages=messages) # Check that the call did not fail and we received a valid response object. assert response is not None @@ -9908,100 +4727,6 @@ def test_get_credential_deployment_skips_a_paused_deployment(): assert router.get_credential_deployment(model_id="paused-dep") is None -def test_get_team_public_name_deployment_only_resolves_the_owning_team(): - router = litellm.Router( - model_list=[ - { - "model_name": "mistral/mistral-ocr-latest", - "litellm_params": {"model": "mistral/mistral-ocr-latest", "api_key": "sk-team-a"}, - "model_info": {"id": "team-a-ocr", "team_id": "team-a", "team_public_model_name": "ocr"}, - } - ] - ) - - owning_team = router._get_team_public_name_deployment(model_id="ocr", team_id="team-a") - - assert owning_team is not None and owning_team.model_info.id == "team-a-ocr" - assert router._get_team_public_name_deployment(model_id="ocr", team_id="team-b") is None - assert router._get_team_public_name_deployment(model_id="ocr", team_id=None) is None - assert router.get_credential_deployment(model_id="ocr", team_id="team-a").model_info.id == "team-a-ocr" - assert router.get_credential_deployment(model_id="ocr", team_id="team-b") is None - - -def test_get_wildcard_deployment_usable_by_team_prefers_the_team_pattern(): - router = litellm.Router( - model_list=[ - { - "model_name": "mistral/*", - "litellm_params": {"model": "mistral/*", "api_key": "sk-shared"}, - "model_info": {"id": "shared-wildcard"}, - }, - { - "model_name": "mistral/*", - "litellm_params": {"model": "mistral/*", "api_key": "sk-team-a"}, - "model_info": {"id": "team-a-wildcard", "team_id": "team-a", "team_public_model_name": "mistral/*"}, - }, - ] - ) - ocr = "mistral/mistral-ocr-latest" - - team_match = router._get_wildcard_deployment_usable_by_team(model_id=ocr, team_id="team-a") - other_team_match = router._get_wildcard_deployment_usable_by_team(model_id=ocr, team_id="team-b") - anonymous_match = router._get_wildcard_deployment_usable_by_team(model_id=ocr, team_id=None) - - assert team_match is not None and team_match.model_info.id == "team-a-wildcard" - assert other_team_match is not None and other_team_match.model_info.id == "shared-wildcard" - assert anonymous_match is not None and anonymous_match.model_info.id == "shared-wildcard" - assert router._get_wildcard_deployment_usable_by_team(model_id="openai/gpt-5.6", team_id="team-a") is None - assert router.get_credential_deployment(model_id=ocr, team_id="team-b").model_info.id == "shared-wildcard" - - -@pytest.mark.parametrize( - ("with_catch_all", "requested", "team_id", "expected_ids"), - [ - pytest.param(False, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-when-no-proxy-wide-match"), - pytest.param(False, "rwild-chat", "team-b", (), id="another-teams-pattern-never-serves"), - pytest.param(False, "rwild-chat", None, (), id="no-team-no-team-pattern"), - pytest.param(True, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-beats-the-proxy-wide-pattern"), - pytest.param(True, "rwild-chat", "team-b", ("catch-all",), id="another-team-falls-through-to-proxy-wide"), - pytest.param(True, "ptu-chat", "team-r", ("ptu-wildcard",), id="only-the-most-specific-proxy-wide-pattern"), - ], -) -def test_wildcard_route_deployments_is_the_one_pattern_routing_serves_from( - with_catch_all: bool, requested: str, team_id: str | None, expected_ids: tuple[str, ...] -): - catch_all: Final = ( - ({"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-open"}, "model_info": {"id": "catch-all"}},) - if with_catch_all - else () - ) - router: Final = litellm.Router( - model_list=[ - { - "model_name": "ptu-*", - "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-ptu"}, - "model_info": {"id": "ptu-wildcard"}, - }, - { - "model_name": "team-r-internal", - "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "sk-team-r"}, - "model_info": {"id": "team-r-wildcard", "team_id": "team-r", "team_public_model_name": "rwild-*"}, - }, - *catch_all, - ] - ) - - served: Final = router.wildcard_route_deployments(requested, team_id) - early: Final = router._try_early_resolve_deployments_for_model_not_in_names(model=requested, request_team_id=team_id) - early_deployments: Final = early[1] if early is not None and isinstance(early[1], list) else [] - - assert tuple(Deployment.model_validate(d).model_info.id for d in served) == expected_ids - assert tuple(Deployment.model_validate(d).model_info.id for d in early_deployments) == expected_ids - assert router.wildcard_routed_models(requested, team_id) == frozenset( - Deployment.model_validate(d).litellm_params.model for d in served - ) - - @pytest.mark.parametrize( ("team_id", "expected_id"), [ @@ -10055,15 +4780,10 @@ def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint() ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-claude-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-claude-model") assert credentials is not None - assert ( - credentials["aws_bedrock_runtime_endpoint"] - == "https://bedrock-runtime.us-east-1.amazonaws.com" - ) + assert credentials["aws_bedrock_runtime_endpoint"] == "https://bedrock-runtime.us-east-1.amazonaws.com" assert credentials["aws_access_key_id"] == "test-access-key" assert credentials["aws_secret_access_key"] == "test-secret-key" assert credentials["aws_region_name"] == "us-east-1" @@ -10090,9 +4810,7 @@ def test_get_deployment_credentials_with_provider_includes_bucket_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="vertex-gemini" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="vertex-gemini") assert credentials is not None assert credentials["gcs_bucket_name"] == "my-batch-bucket" @@ -10177,9 +4895,7 @@ def test_get_deployment_credentials_with_provider_resolves_credential_name(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="azure-gpt-4" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="azure-gpt-4") assert credentials is not None assert credentials["api_key"] == "resolved-api-key" @@ -10216,9 +4932,7 @@ def test_get_deployment_credentials_with_provider_bedrock_batch_fields(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None assert credentials["custom_llm_provider"] == "bedrock" @@ -10262,9 +4976,7 @@ def test_get_deployment_credentials_with_provider_preserves_aws_auth_params(): ], ) - credentials = router.get_deployment_credentials_with_provider( - model_id="bedrock-batch-model" - ) + credentials = router.get_deployment_credentials_with_provider(model_id="bedrock-batch-model") assert credentials is not None for key, value in aws_auth_params.items(): @@ -10334,15 +5046,11 @@ def test_get_deployment_credentials_with_provider_team_wildcard_priority(): ], ) - team_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + team_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert team_credentials is not None assert team_credentials["api_key"] == "team-key" - global_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2" - ) + global_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2") assert global_credentials is not None assert global_credentials["api_key"] == "global-key" @@ -10383,15 +5091,11 @@ def test_get_deployment_credentials_with_provider_skips_other_team_deployment(): assert other_team_credentials is not None assert other_team_credentials["vertex_project"] == "shared-project" - unscoped_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro" - ) + unscoped_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") assert unscoped_credentials is not None assert unscoped_credentials["vertex_project"] == "shared-project" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-b") assert owner_credentials is not None assert owner_credentials["vertex_project"] == "team-b-project" @@ -10418,81 +5122,8 @@ def test_get_deployment_credentials_with_provider_no_fallback_to_other_team_only ], ) - assert ( - router.get_deployment_credentials_with_provider( - model_id="gemini-2.5-pro", team_id="team-a" - ) - is None - ) - assert ( - router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") - is None - ) - - -def test_deployment_usable_by_team_helpers(): - """ - Direct coverage of the team-ownership filter: a team-scoped deployment is - usable only by its owning team, shared deployments by anyone, and the - model-group picker returns the first usable deployment or None. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gemini-2.5-pro", - "litellm_params": { - "model": "vertex_ai/gemini-2.5-pro", - "vertex_project": "team-b-project", - }, - "model_info": { - "id": "team-b-vertex", - "team_id": "team-b", - "team_public_model_name": "gemini-2.5-pro", - }, - }, - { - "model_name": "gemini-2.5-pro", - "litellm_params": { - "model": "vertex_ai/gemini-2.5-pro", - "vertex_project": "shared-project", - }, - }, - ], - ) - - team_owned, shared = router.model_list - assert router._deployment_usable_by_team(team_owned, "team-b") is True - assert router._deployment_usable_by_team(team_owned, "team-a") is False - assert router._deployment_usable_by_team(team_owned, None) is False - assert router._deployment_usable_by_team(shared, "team-a") is True - assert router._deployment_usable_by_team(shared, None) is True - - picked = router._get_model_group_deployment_usable_by_team( - model_group_name="gemini-2.5-pro", team_id="team-a" - ) - assert picked is not None - assert picked.litellm_params.vertex_project == "shared-project" - - owner_picked = router._get_model_group_deployment_usable_by_team( - model_group_name="gemini-2.5-pro", team_id="team-b" - ) - assert owner_picked is not None - assert owner_picked.litellm_params.vertex_project == "team-b-project" - - assert ( - router._get_model_group_deployment_usable_by_team( - model_group_name="unknown-model", team_id="team-a" - ) - is None - ) - - -def test_deployment_usable_by_team_uses_dynamic_model_info_get(): - class ModelInfo: - def get(self, key: str) -> str | None: - return "team-b" if key == "team_id" else None - - assert litellm.Router._core_deployment_usable_by_team({"model_info": ModelInfo()}, "team-a") is False + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro", team_id="team-a") is None + assert router.get_deployment_credentials_with_provider(model_id="gemini-2.5-pro") is None def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): @@ -10524,9 +5155,7 @@ def test_get_deployment_credentials_with_provider_skips_other_team_wildcard(): assert other_team_credentials is not None assert other_team_credentials["api_key"] == "global-key" - owner_credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-b" - ) + owner_credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-b") assert owner_credentials is not None assert owner_credentials["api_key"] == "team-b-key" @@ -10538,21 +5167,11 @@ def test_team_wildcard_credentials_not_usable_after_delete_deployment(): """ router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is not None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is not None router.delete_deployment(id="team-wildcard-id") - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_global_wildcard_pattern_router_evicts_stale_entry_on_upsert_and_delete(): @@ -10625,22 +5244,13 @@ def test_team_wildcard_credentials_refreshed_on_upsert_and_set_model_list(): router = litellm.Router(model_list=[_team_wildcard_model(api_key="old-key")]) - router.upsert_deployment( - deployment=Deployment(**_team_wildcard_model(api_key="new-key")) - ) - credentials = router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) + router.upsert_deployment(deployment=Deployment(**_team_wildcard_model(api_key="new-key"))) + credentials = router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") assert credentials is not None assert credentials["api_key"] == "new-key" router.set_model_list(model_list=[]) - assert ( - router.get_deployment_credentials_with_provider( - model_id="openai/gpt-5.2", team_id="team-1" - ) - is None - ) + assert router.get_deployment_credentials_with_provider(model_id="openai/gpt-5.2", team_id="team-1") is None def test_get_available_guardrail_single_deployment(): @@ -10720,43 +5330,6 @@ def test_get_available_guardrail_not_found(): router.get_available_guardrail(guardrail_name="non-existent") -@pytest.mark.asyncio -async def test_aguardrail_helper(): - """ - Test _aguardrail_helper selects a guardrail and executes the original function. - """ - guardrail_config = { - "guardrail_name": "content-filter", - "litellm_params": {"guardrail": "custom", "mode": "pre_call"}, - "id": "guardrail-1", - } - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo"}, - } - ], - guardrail_list=[guardrail_config], - ) - - # Mock the original function - async def mock_original_function(**kwargs): - return { - "result": "success", - "selected_guardrail": kwargs.get("selected_guardrail"), - } - - result = await router._aguardrail_helper( - model="content-filter", - original_generic_function=mock_original_function, - ) - - assert result["result"] == "success" - assert result["selected_guardrail"] == guardrail_config - - @pytest.mark.asyncio async def test_aguardrail(): """ @@ -10807,7 +5380,6 @@ async def test_anthropic_messages_call_type_is_cached(): PromptCachingDeploymentCheck, ) from litellm.types.utils import ( - CallTypes, StandardLoggingHiddenParams, StandardLoggingMetadata, StandardLoggingModelInformation, @@ -10828,9 +5400,7 @@ async def test_anthropic_messages_call_type_is_cached(): startTime=1234567890.0, endTime=1234567891.0, completionStartTime=1234567890.5, - model_map_information=StandardLoggingModelInformation( - model_map_key="gpt-3.5-turbo", model_map_value=None - ), + model_map_information=StandardLoggingModelInformation(model_map_key="gpt-3.5-turbo", model_map_value=None), model="gpt-3.5-turbo", custom_llm_provider="anthropic", model_id="model-123", @@ -10910,123 +5480,8 @@ async def test_anthropic_messages_call_type_is_cached(): ) # This assertion will FAIL if anthropic_messages is filtered out - assert ( - cached_result is not None - ), "Model ID should be cached for anthropic_messages call type" - assert ( - cached_result["model_id"] == test_model_id - ), f"Expected {test_model_id}, got {cached_result['model_id']}" - - -def test_update_kwargs_with_deployment_propagates_model_tags(): - """ - Test that deployment-level tags from litellm_params are merged into - kwargs metadata when _update_kwargs_with_deployment is called. - - This ensures model-level tags defined in config.yaml appear in SpendLogs. - See: https://github.com/BerriAI/litellm/issues/XXXX - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "fake-key", - "tags": ["openai-account", "production"], - }, - }, - ], - ) - - kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Deployment tags should be propagated to kwargs metadata - assert "tags" in kwargs["metadata"] - assert "openai-account" in kwargs["metadata"]["tags"] - assert "production" in kwargs["metadata"]["tags"] - - -def test_update_kwargs_with_deployment_merges_many_client_tags_in_linear_time(): - router: Final = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake-key", "tags": ["deployment-tag"]}, - }, - ], - ) - client_tags: Final = [f"client-{index}" for index in range(100_000)] - kwargs: Final = {"metadata": {"tags": client_tags}} - deployment: Final = router.get_deployment_by_model_group_name(model_group_name="gpt-4o-mini") - - started: Final = time.perf_counter() - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert time.perf_counter() - started < 2 - assert kwargs["metadata"]["tags"] == [*client_tags, "deployment-tag"] - - -def test_update_kwargs_with_deployment_merges_tags_without_duplicates(): - """ - Test that when both request-level and deployment-level tags exist, - they are merged without duplicates. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "fake-key", - "tags": ["openai-account", "shared-tag"], - }, - }, - ], - ) - - # Simulate request that already has tags (from request body or key/team level) - kwargs: dict = {"metadata": {"tags": ["user-tag", "shared-tag"]}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Both sources should be merged, no duplicates - assert "user-tag" in kwargs["metadata"]["tags"] - assert "openai-account" in kwargs["metadata"]["tags"] - assert "shared-tag" in kwargs["metadata"]["tags"] - assert kwargs["metadata"]["tags"].count("shared-tag") == 1 - - -def test_update_kwargs_with_deployment_no_tags(): - """ - Test that when deployment has no tags, kwargs metadata is not affected. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4o-mini", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "fake-key", - }, - }, - ], - ) - - kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-4o-mini" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # No tags key should be added if deployment has no tags - assert "tags" not in kwargs["metadata"] + assert cached_result is not None, "Model ID should be cached for anthropic_messages call type" + assert cached_result["model_id"] == test_model_id, f"Expected {test_model_id}, got {cached_result['model_id']}" @pytest.mark.asyncio @@ -11073,581 +5528,6 @@ async def test_retry_does_not_narrow_tag_filtered_group_to_failed_deployments_ta assert metadata["tags"] == ["free"] -def test_update_kwargs_with_deployment_merges_tools(): - """ - Test that when both deployment litellm_params and request have tools, - they are merged (deployment tools first, then request tools). - - Supports proxy-configured tools (e.g. for o3 deep research) merged with - client-provided tools. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - "tool_choice": "auto", - }, - }, - ], - ) - - kwargs: dict = { - "metadata": {}, - "tools": [ - { - "type": "function", - "function": {"name": "get_weather", "description": "Get weather"}, - }, - ], - } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Tools should be merged: deployment first, then request - assert "tools" in kwargs - assert len(kwargs["tools"]) == 2 - assert kwargs["tools"][0] == {"type": "web_search"} - assert kwargs["tools"][1]["function"]["name"] == "get_weather" - # tool_choice from request (none) - deployment's should be used - assert kwargs["tool_choice"] == "auto" - - -def test_update_kwargs_with_deployment_merge_tools_deployment_only(): - """ - Test that when only deployment has tools, they are applied to kwargs. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - "tool_choice": "required", - }, - }, - ], - ) - - kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert kwargs["tools"] == [{"type": "web_search"}] - assert kwargs["tool_choice"] == "required" - - -def test_update_kwargs_with_deployment_merge_tools_request_overrides_tool_choice(): - """ - Test that when request has tool_choice, it overrides deployment's. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - "tool_choice": "auto", - }, - }, - ], - ) - - kwargs: dict = { - "metadata": {}, - "tool_choice": "none", - } - deployment = router.get_deployment_by_model_group_name( - model_group_name="o3-deep-research" - ) - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - # Request tool_choice should be preserved (merged tools still applied) - assert kwargs["tool_choice"] == "none" - - -@pytest.mark.parametrize( - "deployment_params, router_defaults, request_params, expected_tools", - [ - ({"drop_params": True}, {}, {}, [{"type": "web_search"}]), - ({}, {"drop_params": True}, {}, [{"type": "web_search"}]), - ({}, {}, {"additional_drop_params": ["tools"]}, [{"type": "web_search"}]), - ({"drop_params": True}, {}, {"drop_params": False}, 5), - ({}, {}, {}, 5), - ], - ids=[ - "deployment_drop_params", - "router_default_drop_params", - "request_drops_tools", - "request_overrides_deployment_drop", - "no_drop_policy", - ], -) -def test_update_kwargs_with_deployment_keeps_deployment_tools_when_malformed_request_tools_drop( - deployment_params: dict[str, object], - router_defaults: dict[str, object], - request_params: dict[str, object], - expected_tools: object, -) -> None: - router: Final = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - **deployment_params, - }, - }, - ], - default_litellm_params={**router_defaults}, - ) - kwargs: Final[dict[str, object]] = {"metadata": {}, "tools": 5, **request_params} - router._update_kwargs_with_deployment( - deployment=router.get_deployment_by_model_group_name(model_group_name="o3-deep-research"), kwargs=kwargs - ) - - assert kwargs["tools"] == expected_tools - - -_NAMED_REQUEST_TOOL_CHOICE: Final = {"type": "function", "function": {"name": "lookup_order"}} - - -@pytest.mark.parametrize( - "deployment_params, expected_tool_fields", - [ - ( - {"drop_params": True, "tool_choice": "required"}, - {"tools": [{"type": "web_search"}], "tool_choice": "required"}, - ), - ({"drop_params": True}, {"tools": [{"type": "web_search"}]}), - ({}, {"tools": 5, "tool_choice": _NAMED_REQUEST_TOOL_CHOICE, "parallel_tool_calls": False}), - ], - ids=["deployment_tool_choice_applies", "no_deployment_tool_choice", "no_drop_policy"], -) -def test_update_kwargs_with_deployment_drops_request_tool_settings_with_malformed_request_tools( - deployment_params: dict[str, object], expected_tool_fields: dict[str, object] -) -> None: - router: Final = litellm.Router( - model_list=[ - { - "model_name": "o3-deep-research", - "litellm_params": { - "model": "openai/o3-deep-research", - "api_key": "fake-key", - "tools": [{"type": "web_search"}], - **deployment_params, - }, - }, - ], - ) - kwargs: Final[dict[str, object]] = { - "metadata": {}, - "tools": 5, - "tool_choice": _NAMED_REQUEST_TOOL_CHOICE, - "parallel_tool_calls": False, - } - router._update_kwargs_with_deployment( - deployment=router.get_deployment_by_model_group_name(model_group_name="o3-deep-research"), kwargs=kwargs - ) - - tool_fields: Final = ("tools", "tool_choice", "parallel_tool_calls") - assert {name: kwargs[name] for name in tool_fields if name in kwargs} == expected_tool_fields - - -def test_credential_name_injected_as_tag(): - """ - Test that litellm_credential_name from deployment litellm_params - is injected as a tag into metadata during _update_kwargs_with_deployment. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "xai-model", - "litellm_params": { - "model": "xai/grok-4-1-fast", - "litellm_credential_name": "xAI", - }, - } - ], - ) - - kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert "Credential: xAI" in kwargs["metadata"]["tags"] - assert "A.101" in kwargs["metadata"]["tags"] - - -def test_credential_name_not_duplicated_in_tags(): - """ - Test that if the credential tag already exists in the tags list, - it is not duplicated. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "xai-model", - "litellm_params": { - "model": "xai/grok-4-1-fast", - "litellm_credential_name": "xAI", - }, - } - ], - ) - - kwargs: dict = {"metadata": {"tags": ["Credential: xAI", "A.101"]}} - deployment = router.get_deployment_by_model_group_name(model_group_name="xai-model") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert kwargs["metadata"]["tags"].count("Credential: xAI") == 1 - - -def test_credential_name_not_injected_when_absent(): - """ - Test that when no litellm_credential_name is set, tags are unchanged. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-model", - "litellm_params": { - "model": "gpt-4o", - }, - } - ], - ) - - kwargs: dict = {"metadata": {"tags": ["A.101"]}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-model") - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - assert kwargs["metadata"]["tags"] == ["A.101"] - - -def test_update_kwargs_with_deployment_model_info_in_litellm_metadata(): - """For generic_api_call, model_info with pricing must go to litellm_metadata. - - Routes like /messages and /responses use generic_api_call which stores - model_info under litellm_metadata. Regression test for #23185. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "claude-sonnet-4", - "litellm_params": { - "model": "anthropic/claude-sonnet-4-20250514", - "api_key": "fake-key", - }, - "model_info": { - "id": "custom-pricing-id", - "input_cost_per_token": 0.0003, - "output_cost_per_token": 0.0015, - }, - }, - ], - ) - - kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name="generic_api_call" - ) - - assert "litellm_metadata" in kwargs - model_info = kwargs["litellm_metadata"]["model_info"] - assert model_info["id"] == "custom-pricing-id" - assert model_info["input_cost_per_token"] == 0.0003 - assert model_info["output_cost_per_token"] == 0.0015 - - -def test_update_kwargs_with_deployment_model_info_in_metadata(): - """For acompletion (function_name=None), model_info goes to metadata. - - /chat/completions uses acompletion which stores model_info under metadata. - """ - router = litellm.Router( - model_list=[ - { - "model_name": "claude-sonnet-4", - "litellm_params": { - "model": "anthropic/claude-sonnet-4-20250514", - "api_key": "fake-key", - }, - "model_info": { - "id": "custom-pricing-id", - "input_cost_per_token": 0.0003, - "output_cost_per_token": 0.0015, - }, - }, - ], - ) - - kwargs: dict = {} - deployment = router.get_deployment_by_model_group_name( - model_group_name="claude-sonnet-4" - ) - router._update_kwargs_with_deployment( - deployment=deployment, kwargs=kwargs, function_name=None - ) - - assert "metadata" in kwargs - model_info = kwargs["metadata"]["model_info"] - assert model_info["id"] == "custom-pricing-id" - assert model_info["input_cost_per_token"] == 0.0003 - assert model_info["output_cost_per_token"] == 0.0015 - - -def test_update_kwargs_with_deployment_clears_pod_routing_on_non_discovery_fallback() -> None: - from litellm.constants import KUBERNETES_POD_ROUTING_KEY - - routing_record: Final = { - "service_host": "vllm-headless.ns.svc.cluster.local", - "pod_ip": "10.0.0.1", - "pod_count": 3, - "selection": "session_affinity", - } - discovery_deployment: Final = { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, - "model_info": {"id": "discovery-id"}, - KUBERNETES_POD_ROUTING_KEY: routing_record, - } - fallback_deployment: Final = { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, - "model_info": {"id": "fallback-id"}, - } - router: Final = litellm.Router(model_list=[discovery_deployment]) - kwargs: Final = {"metadata": {}} - - router._update_kwargs_with_deployment(deployment=discovery_deployment, kwargs=kwargs) - assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] == routing_record - - router._update_kwargs_with_deployment(deployment=fallback_deployment, kwargs=kwargs) - assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] is None - - -def test_update_kwargs_with_deployment_synchronizes_existing_litellm_metadata_pod_routing() -> None: - from litellm.constants import KUBERNETES_POD_ROUTING_KEY - - routing_record: Final = { - "service_host": "vllm-headless.ns.svc.cluster.local", - "pod_ip": "10.0.0.1", - "pod_count": 3, - "selection": "session_affinity", - } - spoofed_record: Final = { - "service_host": "vllm-headless.ns.svc.cluster.local", - "pod_ip": "10.0.0.9", - "pod_count": 3, - "selection": "round_robin", - } - deployment: Final = { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, - "model_info": {"id": "discovery-id"}, - KUBERNETES_POD_ROUTING_KEY: routing_record, - } - original_litellm_metadata: Final = {KUBERNETES_POD_ROUTING_KEY: spoofed_record} - kwargs: Final = {"metadata": {}, "litellm_metadata": original_litellm_metadata} - router: Final = litellm.Router(model_list=[deployment]) - - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="completion") - - assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] == routing_record - assert kwargs["litellm_metadata"][KUBERNETES_POD_ROUTING_KEY] == routing_record - assert kwargs["litellm_metadata"] is not original_litellm_metadata - assert original_litellm_metadata[KUBERNETES_POD_ROUTING_KEY] == spoofed_record - - -def test_update_kwargs_with_non_discovery_deployment_clears_other_pod_routing_bucket() -> None: - from litellm.constants import KUBERNETES_POD_ROUTING_KEY - - previous_record: Final = { - "service_host": "vllm-headless.ns.svc.cluster.local", - "pod_ip": "10.0.0.1", - "pod_count": 3, - "selection": "session_affinity", - } - deployment: Final = { - "model_name": "gpt-4o-mini", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake"}, - "model_info": {"id": "fallback-id"}, - } - original_litellm_metadata: Final = {KUBERNETES_POD_ROUTING_KEY: previous_record} - kwargs: Final = {"metadata": {}, "litellm_metadata": original_litellm_metadata} - router: Final = litellm.Router(model_list=[deployment]) - - router._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="completion") - - assert kwargs["metadata"][KUBERNETES_POD_ROUTING_KEY] is None - assert kwargs["litellm_metadata"][KUBERNETES_POD_ROUTING_KEY] is None - assert original_litellm_metadata[KUBERNETES_POD_ROUTING_KEY] == previous_record - - -def test_combine_fallback_usage(): - """Test that _combine_fallback_usage merges partial and fallback usage.""" - from litellm.router import Router - from litellm.types.utils import Usage - - # Create a stream chunk with usage - chunk = litellm.ModelResponseStream( - id="test", - model="gpt-4o", - choices=[], - usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), - ) - - # Call _combine_fallback_usage with no extra usage - Router._combine_fallback_usage(chunk, None) - assert chunk.usage is not None - assert chunk.usage.prompt_tokens == 10 - assert chunk.usage.completion_tokens == 5 - assert chunk.usage.total_tokens == 15 - - Router._combine_fallback_usage(chunk, Usage(prompt_tokens=0, completion_tokens=4, total_tokens=4)) - assert chunk.usage.prompt_tokens == 10 - assert chunk.usage.completion_tokens == 9 - assert chunk.usage.total_tokens == 19 - - -@pytest.mark.asyncio -async def test_acompletion_streaming_iterator_does_not_log_success_on_terminal_failure(): - """A mid-stream failure with no successful fallback raises and is logged as - a failure, so the router must never dispatch it as a success. Partial-spend - recovery for the failure row happens in the streaming handler, not here, so - this guards only against reintroducing a success log for a failed stream. - """ - from litellm.exceptions import MidStreamFallbackError - from litellm.types.utils import Delta, StreamingChoices, Usage - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-4", - "litellm_params": {"model": "gpt-4", "api_key": "fake-key-1"}, - }, - ], - set_verbose=True, - ) - - error = MidStreamFallbackError( - message="Connection lost", - model="gpt-4", - llm_provider="openai", - generated_content="The Roman Empire began when", - ) - - def _make_interrupted_model_response(): - partial_chunk = litellm.ModelResponseStream( - id="chatcmpl-partial-1", - created=1742056047, - model="gpt-4", - object="chat.completion.chunk", - choices=[ - StreamingChoices( - finish_reason=None, - index=0, - delta=Delta(content="The Roman Empire began when", role="assistant"), - ) - ], - usage=Usage(prompt_tokens=17, completion_tokens=9, total_tokens=26), - ) - - class _RaisingStream: - def __init__(self): - self.index = 0 - self.chunks = [partial_chunk] - - def __aiter__(self): - return self - - async def __anext__(self): - if self.index == 0: - self.index += 1 - return partial_chunk - raise error - - stream = _RaisingStream() - logging_obj = MagicMock() - logging_obj.dispatch_success_handlers = AsyncMock() - logging_obj.model_call_details = {} - setattr(stream, "model", "gpt-4") - setattr(stream, "custom_llm_provider", "openai") - setattr(stream, "logging_obj", logging_obj) - return stream, logging_obj - - messages = [{"role": "user", "content": "Hello"}] - initial_kwargs = {"model": "gpt-4", "stream": True} - - # Terminal path: no successful fallback -> the error propagates and the - # router never dispatches a success for the failed stream. - model_response, logging_obj = _make_interrupted_model_response() - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=error), - ): - result = await router._acompletion_streaming_iterator( - model_response=model_response, - messages=messages, - initial_kwargs=dict(initial_kwargs), - ) - collected = [] - async def _drain(): - async for chunk in result: - collected.append(chunk) - - with pytest.raises(MidStreamFallbackError): - await _drain() - - assert len(collected) == 1 - logging_obj.dispatch_success_handlers.assert_not_called() - - # Mid-stream errors with generated content are now re-raised immediately; - # no continuation-prompt fallback is attempted. Success handlers must - # still not be dispatched in this path. - model_response, logging_obj = _make_interrupted_model_response() - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - result = await router._acompletion_streaming_iterator( - model_response=model_response, - messages=messages, - initial_kwargs=dict(initial_kwargs), - ) - collected = [] - async def _drain(): - async for chunk in result: - collected.append(chunk) - - with pytest.raises(MidStreamFallbackError): - await _drain() - - assert len(collected) == 1, "only the partial chunk before the error" - mock_fallback.assert_not_called() - logging_obj.dispatch_success_handlers.assert_not_called() - - @pytest.mark.asyncio async def test_team_scoped_model_fallback(): """ @@ -11780,7 +5660,12 @@ def _shared_ptu_deployment(model_name: str, deployment_id: str, **litellm_params """A one-PTU Azure deployment whose whole reservation is team-a's share.""" return { "model_name": model_name, - "litellm_params": {"model": "azure/gpt-4.1", "api_key": "fake", "api_base": "https://ptu.example", **litellm_params}, + "litellm_params": { + "model": "azure/gpt-4.1", + "api_key": "fake", + "api_base": "https://ptu.example", + **litellm_params, + }, "model_info": { "id": deployment_id, "base_model": "azure/gpt-4.1", @@ -11890,23 +5775,17 @@ def test_multiregion_team_deployments_unique_model_names(): assert len(deployments) == 0 # With team_id: O(n) scan finds BOTH regional deployments - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") assert len(deployments) == 2 deployment_names = {d["model_name"] for d in deployments} assert deployment_names == {"metis-claude-us-east-1", "metis-claude-us-west-2"} # Each deployment has a unique ID (critical for cooldown/retry to work) deployment_ids = {d["model_info"]["id"] for d in deployments} - assert ( - len(deployment_ids) == 2 - ), "Each deployment must have a unique ID for cooldown tracking" + assert len(deployment_ids) == 2, "Each deployment must have a unique ID for cooldown tracking" # Wrong team: returns nothing - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="other-team" - ) + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="other-team") assert len(deployments) == 0 @@ -11951,12 +5830,8 @@ async def test_multiregion_team_failover_between_regions(): ) # Verify the router finds both deployments for the team - deployments = router._get_all_deployments( - model_name="claude-sonnet", team_id="metis-team" - ) - assert ( - len(deployments) == 2 - ), "Router must find both regional deployments by team_public_model_name" + deployments = router._get_all_deployments(model_name="claude-sonnet", team_id="metis-team") + assert len(deployments) == 2, "Router must find both regional deployments by team_public_model_name" # Make a normal request — should succeed from one of the regions response = await router.acompletion( @@ -11971,436 +5846,6 @@ async def test_multiregion_team_failover_between_regions(): ] -def test_access_group_scoped_key_filters_deployments_with_same_public_model(): - """ - If a key can access a model only via access group membership, - router candidate deployments for that public model should be constrained - to deployments in the allowed access group. - """ - from litellm.proxy._types import UserAPIKeyAuth - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5", - "litellm_params": { - "model": "openai/gpt-5.1", - "api_key": "key1", - "mock_response": "response-via-AG1", - }, - "model_info": {"access_groups": ["AG1"]}, - }, - { - "model_name": "gpt-5", - "litellm_params": { - "model": "openai/gpt-4o", - "api_key": "key2", - "mock_response": "response-via-AG2", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - ] - ) - - scoped_key = UserAPIKeyAuth( - api_key="hashed-key", - team_id="team2", - models=["AG2"], - team_models=["AG2"], - ) - - _model, deployments = router._common_checks_available_deployment( - model="gpt-5", - request_kwargs={ - "metadata": { - "user_api_key_team_id": "team2", - "user_api_key_auth": scoped_key, - } - }, - ) - - assert len(deployments) == 1 - assert deployments[0].get("model_info", {}).get("access_groups") == ["AG2"] - - seen = set() - for _ in range(20): - response = router.completion( - model="gpt-5", - messages=[{"role": "user", "content": "hello"}], - metadata={"user_api_key_team_id": "team2", "user_api_key_auth": scoped_key}, - ) - seen.add(response.choices[0].message.content) - - assert seen == {"response-via-AG2"} - - -def test_explicit_model_access_does_not_force_access_group_filtering(): - """ - If a key has explicit model access in addition to access group entries, - do not force access-group-only filtering for deployment selection. - """ - from litellm.proxy._types import UserAPIKeyAuth - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5", - "litellm_params": { - "model": "openai/gpt-5.1", - "api_key": "key1", - "mock_response": "response-via-AG1", - }, - "model_info": {"access_groups": ["AG1"]}, - }, - { - "model_name": "gpt-5", - "litellm_params": { - "model": "openai/gpt-4o", - "api_key": "key2", - "mock_response": "response-via-AG2", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - ] - ) - - explicit_key = UserAPIKeyAuth( - api_key="hashed-key", - team_id="team2", - models=["AG2", "gpt-5"], - team_models=["AG2", "gpt-5"], - ) - - _model, deployments = router._common_checks_available_deployment( - model="gpt-5", - request_kwargs={ - "metadata": { - "user_api_key_team_id": "team2", - "user_api_key_auth": explicit_key, - } - }, - ) - - deployment_groups = [ - d.get("model_info", {}).get("access_groups") for d in deployments - ] - assert ["AG1"] in deployment_groups - assert ["AG2"] in deployment_groups - - -def test_access_group_filter_empty_does_not_bypass_via_litellm_model_fallback( - monkeypatch: pytest.MonkeyPatch, -): - """ - When access-group filtering removes all candidates, _get_deployment_by_litellm_model - must not run: it does not re-apply access groups and could return blocked deployments - that share the same litellm_params.model as the request model string. - - ``get_model_access_groups`` is patched to expose AG1 for the public model (so the - access-group filter runs with a non-empty allowed set) while every deployment - returned for that name is AG2-only — filtered to empty. Without the guard, the - litellm-model fallback would return both rows because ``litellm_params.model`` matches. - """ - from litellm.proxy._types import UserAPIKeyAuth - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5", - "litellm_params": { - "model": "gpt-5", - "api_key": "key1", - "mock_response": "blocked-dep-1", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - { - "model_name": "gpt-5", - "litellm_params": { - "model": "gpt-5", - "api_key": "key2", - "mock_response": "blocked-dep-2", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - ] - ) - - orig_groups = router.get_model_access_groups - - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): - if model_name == "gpt-5" and model_access_group is None: - return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} - return orig_groups( - model_name=model_name, - model_access_group=model_access_group, - team_id=team_id, - ) - - monkeypatch.setattr(router, "_core_get_model_access_groups", fake_get_model_access_groups) - - scoped_key = UserAPIKeyAuth( - api_key="hashed-key", - team_id="team2", - models=["AG1"], - team_models=["AG1"], - ) - - with pytest.raises(litellm.BadRequestError): - router._common_checks_available_deployment( - model="gpt-5", - request_kwargs={ - "metadata": { - "user_api_key_team_id": "team2", - "user_api_key_auth": scoped_key, - } - }, - ) - - -def test_access_group_block_does_not_silently_use_default_fallback_model( - monkeypatch: pytest.MonkeyPatch, -): - """ - When access-group filtering empties candidates for model X, the router must not use - ``fallbacks`` default ``*`` routing to model Y: Y may have no ``access_groups``, so - ``_filter_deployments_by_model_access_groups`` would not constrain Y and the caller - would be served despite being blocked from X. - """ - from litellm.proxy._types import UserAPIKeyAuth - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5", - "litellm_params": { - "model": "gpt-5", - "api_key": "key1", - "mock_response": "blocked-dep-1", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - { - "model_name": "gpt-5", - "litellm_params": { - "model": "gpt-5", - "api_key": "key2", - "mock_response": "blocked-dep-2", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - { - "model_name": "gpt-4-fallback", - "litellm_params": { - "model": "gpt-4", - "api_key": "fallback-key", - "mock_response": "should-not-reach", - }, - }, - ], - fallbacks=[{"*": ["gpt-4-fallback"]}], - ) - - orig_groups = router.get_model_access_groups - - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): - if model_name == "gpt-5" and model_access_group is None: - return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} - return orig_groups( - model_name=model_name, - model_access_group=model_access_group, - team_id=team_id, - ) - - monkeypatch.setattr(router, "_core_get_model_access_groups", fake_get_model_access_groups) - - scoped_key = UserAPIKeyAuth( - api_key="hashed-key", - team_id="team2", - models=["AG1"], - team_models=["AG1"], - ) - - with pytest.raises(litellm.BadRequestError): - router._common_checks_available_deployment( - model="gpt-5", - request_kwargs={ - "metadata": { - "user_api_key_team_id": "team2", - "user_api_key_auth": scoped_key, - } - }, - ) - - -def test_access_group_block_via_litellm_model_branch_does_not_use_default_fallback( - monkeypatch: pytest.MonkeyPatch, -): - """ - When the by-name lookup returns no deployments and the litellm-model fallback - branch finds candidates that access-group filtering then empties, the router - must not fall through to default ``fallbacks`` routing — the default fallback - model may have no ``access_groups`` and would short-circuit the filter, - silently serving a caller blocked by access-group restrictions. - """ - from litellm.proxy._types import UserAPIKeyAuth - - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-5-alias", - "litellm_params": { - "model": "gpt-5", - "api_key": "key1", - "mock_response": "blocked-dep-1", - }, - "model_info": {"access_groups": ["AG2"]}, - }, - { - "model_name": "gpt-4-fallback", - "litellm_params": { - "model": "gpt-4", - "api_key": "fallback-key", - "mock_response": "should-not-reach", - }, - }, - ], - fallbacks=[{"*": ["gpt-4-fallback"]}], - ) - - orig_groups = router.get_model_access_groups - - def fake_get_model_access_groups( - model_name=None, model_access_group=None, team_id=None - ): - if model_name == "gpt-5" and model_access_group is None: - return {"AG1": ["gpt-5"], "AG2": ["gpt-5"]} - return orig_groups( - model_name=model_name, - model_access_group=model_access_group, - team_id=team_id, - ) - - monkeypatch.setattr(router, "_core_get_model_access_groups", fake_get_model_access_groups) - - scoped_key = UserAPIKeyAuth( - api_key="hashed-key", - team_id="team2", - models=["AG1"], - team_models=["AG1"], - ) - - with pytest.raises(litellm.BadRequestError): - router._common_checks_available_deployment( - model="gpt-5", - request_kwargs={ - "metadata": { - "user_api_key_team_id": "team2", - "user_api_key_auth": scoped_key, - } - }, - ) - - -def test_try_early_resolve_deployments_for_model_not_in_names(): - """ - Direct coverage for ``_try_early_resolve_deployments_for_model_not_in_names``: - - - Returns ``None`` when the requested model is already in ``self.model_names`` - (the by-name lookup path will handle it). - - Returns ``None`` when there are no team deployments, no pattern matches, and - no default deployment to fall back to. - - Returns the pattern-router match when the model matches a wildcard route. - - Returns the default deployment with the request model substituted in when one - is configured, without mutating the stored default. - """ - router_in_names = litellm.Router( - model_list=[ - { - "model_name": "gpt-5", - "litellm_params": { - "model": "openai/gpt-5", - "api_key": "key1", - }, - }, - ] - ) - - assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names( - model="gpt-5", request_team_id=None - ) - is None - ) - assert ( - router_in_names._try_early_resolve_deployments_for_model_not_in_names( - model="some-unknown-model", request_team_id=None - ) - is None - ) - - pattern_router = litellm.Router( - model_list=[ - { - "model_name": "openai/*", - "litellm_params": { - "model": "openai/*", - "api_key": "key-pattern", - }, - }, - ] - ) - - pattern_result = ( - pattern_router._try_early_resolve_deployments_for_model_not_in_names( - model="openai/gpt-4o-mini", request_team_id=None - ) - ) - assert pattern_result is not None - resolved_model, pattern_deployments = pattern_result - assert resolved_model == "openai/gpt-4o-mini" - assert isinstance(pattern_deployments, list) and len(pattern_deployments) == 1 - - default_router = litellm.Router( - model_list=[ - { - "model_name": "named-model", - "litellm_params": { - "model": "openai/gpt-4o", - "api_key": "key-named", - }, - }, - ] - ) - default_router.default_deployment = { - "model_name": "default", - "litellm_params": { - "model": "openai/will-be-overridden", - "api_key": "key-default", - }, - } - - default_result = ( - default_router._try_early_resolve_deployments_for_model_not_in_names( - model="brand-new-model", request_team_id=None - ) - ) - assert default_result is not None - resolved_model, default_deployment = default_result - assert resolved_model == "brand-new-model" - assert isinstance(default_deployment, dict) - assert default_deployment["litellm_params"]["model"] == "brand-new-model" - # The original default_deployment must not be mutated. - assert ( - default_router.default_deployment["litellm_params"]["model"] - == "openai/will-be-overridden" - ) - - def _router_with_two_deployments(blocked_flags): import litellm @@ -12446,10 +5891,7 @@ def _seed_unhealthy_states(router, unhealthy_ids, timestamp=None): ts = timestamp if timestamp is not None else time.time() router.health_state_cache.set_deployment_health_states( - { - uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} - for uid in unhealthy_ids - } + {uid: {"is_healthy": False, "timestamp": ts, "reason": "test_unhealthy"} for uid in unhealthy_ids} ) @@ -12573,42 +6015,10 @@ async def test_async_get_fully_unhealthy_model_names_noop_with_allowed_fails_pol assert await router.async_get_fully_unhealthy_model_names() == set() -@pytest.mark.asyncio -async def test_async_get_healthy_deployments_skips_blocked_deployment(): - router = _router_with_two_deployments([True, False]) - healthy, all_dep = await router._async_get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) - healthy_ids = [d["model_info"]["id"] for d in healthy] - assert "dep-0" not in healthy_ids - assert "dep-1" in healthy_ids - assert len(all_dep) == 2 - - -def test_get_healthy_deployments_sync_skips_blocked_deployment(): - router = _router_with_two_deployments([False, True]) - healthy, all_dep = router._get_healthy_deployments( - model="gpt-4o", parent_otel_span=None - ) - healthy_ids = [d["model_info"]["id"] for d in healthy] - assert "dep-0" in healthy_ids - assert "dep-1" not in healthy_ids - assert len(all_dep) == 2 - - -def test_filter_blocked_deployments_drops_blocked_keeps_unblocked(): - router = _router_with_two_deployments([True, False]) - filtered = router._filter_blocked_deployments(router.get_model_list() or []) - ids = [d["model_info"]["id"] for d in filtered] - assert ids == ["dep-1"] - - @pytest.mark.asyncio async def test_public_async_get_healthy_deployments_skips_blocked_on_primary_path(): router = _router_with_two_deployments([True, False]) - deployments = await router.async_get_healthy_deployments( - model="gpt-4o", request_kwargs={} - ) + deployments = await router.async_get_healthy_deployments(model="gpt-4o", request_kwargs={}) assert isinstance(deployments, list) ids = [d["model_info"]["id"] for d in deployments] assert "dep-0" not in ids @@ -12696,9 +6106,7 @@ def _router_with_two_pass_through_deployments(blocked_flags): def test_get_available_deployment_for_pass_through_skips_blocked(): router = _router_with_two_pass_through_deployments([True, False]) - deployment = router.get_available_deployment_for_pass_through( - model="gpt-4o", request_kwargs={} - ) + deployment = router.get_available_deployment_for_pass_through(model="gpt-4o", request_kwargs={}) assert deployment["model_info"]["id"] == "pt-1" @@ -12707,9 +6115,7 @@ def test_get_available_deployment_for_pass_through_raises_when_dict_blocked(): router = _router_with_two_pass_through_deployments([True, True]) with pytest.raises(litellm.ServiceUnavailableError): - router.get_available_deployment_for_pass_through( - model="pt-0", request_kwargs={} - ) + router.get_available_deployment_for_pass_through(model="pt-0", request_kwargs={}) def test_get_available_deployment_for_pass_through_names_cooldown_despite_healthy_non_pass_through(): @@ -12751,9 +6157,7 @@ def test_initialize_deployment_for_pass_through_keeps_bedrock_iam_deployment(): } ] ) - assert [m["model_info"]["id"] for m in router.get_model_list()] == [ - "bedrock-iam-pt" - ] + assert [m["model_info"]["id"] for m in router.get_model_list()] == ["bedrock-iam-pt"] def test_pass_through_deployment_api_key_resolves_via_get_credentials(): @@ -12764,12 +6168,7 @@ def test_pass_through_deployment_api_key_resolves_via_get_credentials(): router = _router_with_two_pass_through_deployments([False, False]) passthrough_router = PassthroughEndpointRouter(llm_router_getter=lambda: router) assert len(router.get_model_list()) == 2 - assert ( - passthrough_router.get_credentials( - custom_llm_provider="openai", region_name=None - ) - == "sk-fake-for-tests" - ) + assert passthrough_router.get_credentials(custom_llm_provider="openai", region_name=None) == "sk-fake-for-tests" def test_get_deployment_credentials_returns_none_for_blocked_deployment(): @@ -12784,32 +6183,6 @@ def test_get_deployment_credentials_with_provider_returns_none_for_blocked_deplo assert router.get_deployment_credentials_with_provider(model_id="dep-1") is not None -def test_is_deployment_blocked_static_helper_reflects_blocked_flag(): - """ - Exercises Router._core_is_deployment_blocked so router_code_coverage.py (AST call graph) - marks the helper as covered by router-named tests. - """ - import types - - import litellm - - router = _router_with_two_deployments([True, False]) - blocked_dep = router.get_deployment("dep-0") - unblocked_dep = router.get_deployment("dep-1") - assert blocked_dep is not None and unblocked_dep is not None - assert litellm.Router._core_is_deployment_blocked(blocked_dep) is True - assert litellm.Router._core_is_deployment_blocked(unblocked_dep) is False - - # No model_info on deployment object → treated as not blocked - assert litellm.Router._core_is_deployment_blocked(object()) is False - missing_blocked = types.SimpleNamespace() - assert litellm.Router._core_is_deployment_blocked(types.SimpleNamespace(model_info=missing_blocked)) is False - assert ( - litellm.Router._core_is_deployment_blocked(types.SimpleNamespace(model_info=types.SimpleNamespace(blocked=True))) - is True - ) - - def test_synchronous_text_completion_rejects_a_restricted_model_before_the_provider() -> None: model: Final = "text-completion-openai/unit-text-model" router: Final = Router( @@ -12917,9 +6290,7 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_request_timeout_stored_independently_when_both_set( - self, explicit_request_timeout - ): + def test_request_timeout_stored_independently_when_both_set(self, explicit_request_timeout): router = self._make_router(timeout=330) assert router.timeout == 330 assert router.request_timeout == 300 @@ -12937,341 +6308,12 @@ class TestRouterRequestTimeoutPropagation: litellm.request_timeout = original_value litellm.request_timeout_explicitly_set = original_flag - def test_non_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): - router = self._make_router(timeout=330) - assert router._get_non_stream_timeout(kwargs={}, data={}) == 300 - - def test_stream_prefers_request_timeout_over_router_timeout( - self, explicit_request_timeout - ): - router = self._make_router(timeout=330) - # stream=True resolves through _get_stream_timeout; request_timeout must win. - assert router._get_timeout(kwargs={"stream": True}, data={}) == 300 - - def test_explicit_stream_timeout_still_wins_over_request_timeout( - self, explicit_request_timeout - ): - router = self._make_router(timeout=330, stream_timeout=45) - assert router._get_stream_timeout(kwargs={}, data={}) == 45 - - def test_non_stream_falls_through_to_router_timeout_without_request_timeout(self): - original_value = litellm.request_timeout - original_flag = litellm.request_timeout_explicitly_set - litellm.request_timeout = litellm.constants.DEFAULT_REQUEST_TIMEOUT_SECONDS - litellm.request_timeout_explicitly_set = False - try: - router = self._make_router(timeout=330) - assert router._get_non_stream_timeout(kwargs={}, data={}) == 330 - finally: - litellm.request_timeout = original_value - litellm.request_timeout_explicitly_set = original_flag - - def test_per_deployment_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): - router = self._make_router(timeout=330) - assert router._get_non_stream_timeout(kwargs={}, data={"timeout": 120}) == 120 - - def test_per_request_timeout_overrides_request_timeout( - self, explicit_request_timeout - ): - router = self._make_router(timeout=330) - assert ( - router._get_non_stream_timeout( - kwargs={"timeout": 60}, data={"timeout": 120} - ) - == 60 - ) - - def test_passthrough_prefers_request_timeout_over_router_timeout(self, explicit_request_timeout): - router = self._make_router(timeout=330) - deployment: Final = router.model_list[0] - assert _passthrough_timeout(router, deployment, stream=False) == 300.0 - assert _passthrough_timeout(router, deployment, stream=True) == 300.0 - - def test_passthrough_stream_timeout_still_wins_over_request_timeout(self, explicit_request_timeout): - router = self._make_router(timeout=330, stream_timeout=45) - assert _passthrough_timeout(router, router.model_list[0], stream=True) == 45.0 - # --------------------------------------------------------------------------- # Deferred-stream eager-fetch tests # --------------------------------------------------------------------------- -def _make_deferred_stream_wrapper(make_call_fn): - """Return a CustomStreamWrapper with completion_stream=None and the given make_call.""" - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - - logging_obj = MagicMock() - logging_obj.model_call_details = {"litellm_params": {}} - return CustomStreamWrapper( - completion_stream=None, - model="vertex_ai/gemini-2.0-flash", - logging_obj=logging_obj, - custom_llm_provider="vertex_ai_beta", - make_call=make_call_fn, - ) - - -def _make_router_with_vertex_and_fallback(): - return litellm.Router( - model_list=[ - { - "model_name": "my-gemini", - "litellm_params": { - "model": "vertex_ai/gemini-2.0-flash", - "vertex_project": "test-project", - "vertex_location": "us-central1", - }, - }, - { - "model_name": "my-fallback", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-fake", - }, - }, - ], - fallbacks=[{"my-gemini": ["my-fallback"]}], - num_retries=0, - ) - - -@pytest.mark.asyncio -async def test_acompletion_deferred_stream_error_propagates_through_acompletion(): - """Regression: a deferred-stream CustomStreamWrapper whose make_call raises a 429 - must propagate the exception from within _acompletion's except block so that - fail_calls is incremented (i.e., deployment cooldown fires) and the standard - router fallback chain can handle it. - - Before the fix, the HTTP call happened inside __anext__ (outside the except block), - so fail_calls was never incremented. - """ - import litellm as _litellm - - rate_limit_err = _litellm.RateLimitError( - message="Resource exhausted", - llm_provider="vertex_ai", - model="gemini-2.0-flash", - ) - - async def failing_make_call(**kwargs): - raise rate_limit_err - - router = _make_router_with_vertex_and_fallback() - deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call) - - with patch( - "litellm.acompletion", - new_callable=AsyncMock, - return_value=deferred_wrapper, - ): - with pytest.raises(_litellm.RateLimitError): - await router._acompletion( - model="vertex_ai/gemini-2.0-flash", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - specific_deployment=router.model_list[0], - ) - - model_name = router.model_list[0]["litellm_params"]["model"] - assert router.fail_calls[model_name] == 1, ( - "fail_calls must be incremented when the deferred HTTP call fails; " - "without the eager fetch_stream() fix this stays at 0" - ) - - -@pytest.mark.asyncio -async def test_acompletion_deferred_stream_preserves_original_headers_on_error(): - """Router is used both by the proxy and directly as an SDK. HTTP-framing headers - (Content-Length, Transfer-Encoding, ...) must NOT be stripped at this layer, or - direct SDK callers lose legitimate provider metadata (e.g. content-type, - proxy-authenticate) that only the proxy's own response construction needs to - worry about. Stripping happens in the proxy layer instead - (_handle_llm_api_exception).""" - import litellm as _litellm - - err = _litellm.RateLimitError( - message="Resource exhausted", - llm_provider="vertex_ai", - model="gemini-2.0-flash", - ) - err.headers = { - "content-length": "42", - "transfer-encoding": "chunked", - "content-encoding": "gzip", - "content-type": "application/json", - "x-request-id": "abc-123", - } - - async def failing_make_call(**kwargs): - raise err - - router = _make_router_with_vertex_and_fallback() - deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call) - - with patch( - "litellm.acompletion", - new_callable=AsyncMock, - return_value=deferred_wrapper, - ): - with pytest.raises(_litellm.RateLimitError) as exc_info: - await router._acompletion( - model="vertex_ai/gemini-2.0-flash", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - specific_deployment=router.model_list[0], - ) - - raised = exc_info.value - headers = getattr(raised, "headers", {}) - assert headers.get("content-length") == "42" - assert headers.get("transfer-encoding") == "chunked" - assert headers.get("content-encoding") == "gzip" - assert headers.get("content-type") == "application/json" - assert headers.get("x-request-id") == "abc-123" - - -@pytest.mark.asyncio -async def test_acompletion_deferred_stream_skipped_when_stream_already_set(): - """When completion_stream is already populated (non-deferred provider), the eager - fetch_stream() call must be skipped entirely; no exception should be raised even - if make_call would fail. - """ - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - - async def would_fail(**kwargs): - raise RuntimeError("should not be called") - - logging_obj = MagicMock() - logging_obj.model_call_details = {"litellm_params": {}} - - async def noop_aiter(): - return - yield - - noop_stream = noop_aiter() - already_set_wrapper = CustomStreamWrapper( - completion_stream=noop_stream, - model="openai/gpt-4o", - logging_obj=logging_obj, - custom_llm_provider="openai", - make_call=would_fail, - ) - - router = litellm.Router( - model_list=[ - { - "model_name": "my-model", - "litellm_params": { - "model": "openai/gpt-4o", - "api_key": "sk-fake", - }, - } - ], - ) - - with patch( - "litellm.acompletion", - new_callable=AsyncMock, - return_value=already_set_wrapper, - ): - result = await router._acompletion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - specific_deployment=router.model_list[0], - ) - - assert result is not None, "should return a streaming wrapper without errors" - assert already_set_wrapper.completion_stream is noop_stream, "completion_stream must not be re-fetched" - await noop_stream.aclose() - - -def test_completion_deferred_stream_error_propagates_through_completion(): - """Regression: the sync router path needs the same eager fetch as the async one. - - A deferred-stream CustomStreamWrapper hands back a wrapper whose HTTP call has - not happened yet, so without fetch_sync_stream() the provider error surfaces on - first iteration, outside _completion's except block. The deployment is then never - marked failed and function_with_fallbacks never sees the error. - """ - import litellm as _litellm - - rate_limit_err = _litellm.RateLimitError( - message="Resource exhausted", - llm_provider="vertex_ai", - model="gemini-2.0-flash", - ) - make_call_invocations = [] - - def failing_make_call(**kwargs): - make_call_invocations.append(kwargs) - raise rate_limit_err - - router = _make_router_with_vertex_and_fallback() - deferred_wrapper = _make_deferred_stream_wrapper(failing_make_call) - - with patch("litellm.completion", return_value=deferred_wrapper): - with pytest.raises(_litellm.RateLimitError): - router._completion( - model="vertex_ai/gemini-2.0-flash", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - specific_deployment=router.model_list[0], - ) - - assert len(make_call_invocations) == 1, ( - "the deferred HTTP call must run inside _completion's try block; " - "without the eager fetch_sync_stream() fix it is deferred to first iteration" - ) - - -def test_completion_deferred_stream_skipped_when_stream_already_set(): - """A non-deferred sync provider already has completion_stream populated, so the - eager fetch must be skipped and make_call left untouched. - """ - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper - - def would_fail(**kwargs): - raise RuntimeError("should not be called") - - logging_obj = MagicMock() - logging_obj.model_call_details = {"litellm_params": {}} - already_set_stream = iter([]) - - already_set_wrapper = CustomStreamWrapper( - completion_stream=already_set_stream, - model="openai/gpt-4o", - logging_obj=logging_obj, - custom_llm_provider="openai", - make_call=would_fail, - ) - - router = litellm.Router( - model_list=[ - { - "model_name": "my-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "sk-fake"}, - } - ], - ) - - with patch("litellm.completion", return_value=already_set_wrapper): - result = router._completion( - model="openai/gpt-4o", - messages=[{"role": "user", "content": "Hello"}], - stream=True, - specific_deployment=router.model_list[0], - ) - - assert result is not None, "should return a streaming wrapper without errors" - assert already_set_wrapper.completion_stream is already_set_stream, "completion_stream must not be re-fetched" - - class TestAdvisorSubCallCooldown: """Regression for LIT-4565: an advisor orchestration failure must not cool down the selected (healthy) deployment, which would reject unrelated @@ -13302,9 +6344,7 @@ class TestAdvisorSubCallCooldown: ) def _cooled_down_ids(self, router): - active = router.cooldown_cache.get_active_cooldowns( - model_ids=["dep-1"], parent_otel_span=None - ) + active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) return [entry[0] for entry in active] @pytest.mark.asyncio @@ -13313,12 +6353,7 @@ class TestAdvisorSubCallCooldown: router = self._router() now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(self._auth_error()), None, now, now - ) - is True - ) + assert router.deployment_callback_on_failure(self._kwargs(self._auth_error()), None, now, now) is True assert "dep-1" in self._cooled_down_ids(router) def test_advisor_orchestration_failure_does_not_cool_down_deployment(self): @@ -13333,12 +6368,7 @@ class TestAdvisorSubCallCooldown: mark_advisor_orchestration_failure(exception) now = datetime.now() - assert ( - router.deployment_callback_on_failure( - self._kwargs(exception), None, now, now - ) - is False - ) + assert router.deployment_callback_on_failure(self._kwargs(exception), None, now, now) is False assert "dep-1" not in self._cooled_down_ids(router) @@ -13575,7 +6605,7 @@ class TestCallerTimeoutCooldown: def test_stream_chunks_have_generated_content_detects_text_and_non_text(): - from litellm.router import _stream_chunks_have_generated_content + from litellm.router_rust.invoke.streaming import _stream_chunks_have_generated_content # noqa: PLC2701 from litellm.types.utils import ( ChatCompletionDeltaToolCall, Delta, @@ -13627,13 +6657,13 @@ def test_stream_chunks_have_generated_content_detects_text_and_non_text(): audio_chunk = _chunk(audio_delta) assert _stream_chunks_have_generated_content([audio_chunk]) is True - images_delta = Delta(images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}]) + images_delta = Delta( + images=[{"image_url": {"url": "https://example.com/img.png"}, "index": 0, "type": "image_url"}] + ) images_chunk = _chunk(images_delta) assert _stream_chunks_have_generated_content([images_chunk]) is True - annotations_delta = Delta( - annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}] - ) + annotations_delta = Delta(annotations=[{"type": "url_citation", "url_citation": {"url": "https://example.com"}}]) annotations_chunk = _chunk(annotations_delta) assert _stream_chunks_have_generated_content([annotations_chunk]) is True @@ -13677,12 +6707,8 @@ def test_get_configured_token_limits_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert router.get_configured_token_limits( - "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0" - ) == (None, None) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_token_limits("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") == (None, None) def test_get_configured_token_limits_treats_malformed_values_as_absent(): @@ -13795,20 +6821,6 @@ def test_get_model_listing_info_carries_configured_limits(): assert (info.max_input_tokens, info.max_output_tokens) == (32000, 8000) -def test_widest_configured_limit_ignores_absent_and_malformed_values(): - model_infos = ( - {"max_input_tokens": 32000}, - {}, - {"max_input_tokens": "not-a-number"}, - {"max_input_tokens": "128000"}, - {"max_output_tokens": 4096}, - ) - - assert litellm.Router._widest_configured_limit(model_infos, "max_input_tokens") == 128000 - assert litellm.Router._widest_configured_limit(model_infos, "max_output_tokens") == 4096 - assert litellm.Router._widest_configured_limit((), "max_input_tokens") is None - - def test_get_model_listing_info_dedupes_interchangeable_deployments(): """The ordinary group is N deployments of one model, so it yields exactly one key.""" router = litellm.Router( @@ -13905,13 +6917,8 @@ def test_get_model_listing_info_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_model_listing_info("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_reads_deployment_model_info(): @@ -13953,13 +6960,8 @@ def test_get_configured_mode_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_mode("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None def test_get_configured_mode_treats_malformed_values_as_absent(): @@ -14018,13 +7020,8 @@ def test_get_configured_display_name_skips_wildcard_pattern_matching(): ] ) - with patch.object( - router.pattern_router, "route", side_effect=AssertionError("pattern route called") - ): - assert ( - router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") - is None - ) + with patch.object(router.pattern_router, "route", side_effect=AssertionError("pattern route called")): + assert router.get_configured_display_name("bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0") is None @pytest.mark.parametrize( @@ -14053,7 +7050,10 @@ def test_get_configured_service_tiers_returns_one_value_per_deployment_in_model_ "litellm_params": {"model": "openai/gpt-6-astra"}, "model_info": {"service_tiers": ["ultrafast"]}, }, - {"model_name": "gpt-6-astra", "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://a.example"}}, + { + "model_name": "gpt-6-astra", + "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://a.example"}, + }, { "model_name": "gpt-6-astra", "litellm_params": {"model": "openai/gpt-6-astra", "api_base": "https://b.example"}, @@ -14316,113 +7316,6 @@ def test_get_configured_display_name_treats_malformed_values_as_absent(): assert router.get_configured_display_name(f"bad-display-{i}") is None -@pytest.mark.asyncio -async def test_acreate_batch_disable_fallbacks_surfaces_owning_provider_error(): - router = litellm.Router( - model_list=[ - { - "model_name": "owning-model", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-owning", - }, - }, - { - "model_name": "fallback-model", - "litellm_params": { - "model": "azure/gpt-4o-mini", - "api_key": "sk-fallback", - "api_base": "https://fallback.openai.azure.com", - "api_version": "2024-08-01-preview", - }, - }, - ], - fallbacks=[{"owning-model": ["fallback-model"]}], - num_retries=0, - ) - owning_provider_error = litellm.BadRequestError( - message="completion_window must be one of: 24h", - model="openai/gpt-4o-mini", - llm_provider="openai", - ) - mock_create = AsyncMock(side_effect=owning_provider_error) - - with patch.object(router, "_acreate_batch", mock_create): - with pytest.raises(litellm.BadRequestError, match="24h"): - await router.acreate_batch( - model="owning-model", - input_file_id="file-owned-by-openai", - endpoint="/v1/chat/completions", - completion_window="5m", - disable_fallbacks=True, - ) - - mock_create.assert_awaited_once() - assert mock_create.call_args.kwargs["model"] == "owning-model" - - -@pytest.mark.asyncio -async def test_acreate_batch_surfaces_owning_provider_error_without_disable_fallbacks(): - """The router itself has to keep a batch inside the group that owns the input file: - the proxy only sets disable_fallbacks on the managed-files route, so the caller - otherwise gets the fallback provider's error for a file it never received.""" - from litellm.types.utils import LiteLLMBatch - - router = litellm.Router( - model_list=[ - { - "model_name": "owning-model", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-owning", - }, - }, - { - "model_name": "fallback-model", - "litellm_params": { - "model": "azure/gpt-4o-mini", - "api_key": "sk-fallback", - "api_base": "https://fallback.openai.azure.com", - "api_version": "2024-08-01-preview", - }, - }, - ], - fallbacks=[{"owning-model": ["fallback-model"]}], - num_retries=0, - ) - attempted_models = [] - - async def _acreate_batch(model, **kwargs): - attempted_models.append(model) - if model == "owning-model": - raise litellm.APIConnectionError( - message="Connection error - openai is unreachable", - model="openai/gpt-4o-mini", - llm_provider="openai", - ) - return LiteLLMBatch( - id="batch-created-on-the-wrong-provider", - completion_window="24h", - created_at=0, - endpoint="/v1/chat/completions", - input_file_id="file-owned-by-openai", - object="batch", - status="validating", - ) - - with patch.object(router, "_acreate_batch", _acreate_batch): - with pytest.raises(litellm.APIConnectionError, match="openai is unreachable"): - await router.acreate_batch( - model="owning-model", - input_file_id="file-owned-by-openai", - endpoint="/v1/chat/completions", - completion_window="24h", - metadata={"team": "batch-jobs"}, - ) - - assert attempted_models == ["owning-model"] - - @pytest.mark.asyncio async def test_acreate_batch_still_falls_back_within_the_owning_model_group(): """Holding a batch inside the model group that owns its input file must not @@ -14775,7 +7668,12 @@ def test_credentials_for_a_mixed_group_name_keep_resolving_to_its_first_deployme _AZURE_CHAT_SKU_REJECTS_BATCH: Final = httpx.Response( 400, - json={"error": {"code": "invalid_deployment_type", "message": "The current deployment is not supported for batch jobs"}}, + json={ + "error": { + "code": "invalid_deployment_type", + "message": "The current deployment is not supported for batch jobs", + } + }, ) @@ -14796,7 +7694,9 @@ async def test_acreate_batch_on_a_mixed_group_reaches_chat_as_on_main_while_its_ _cool_down(router, *(d["model_info"]["id"] for d in router.model_list if d["model_info"]["mode"] == "batch")) with respx.mock(assert_all_called=False) as respx_mock: - chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock(return_value=_AZURE_CHAT_SKU_REJECTS_BATCH) + chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock( + return_value=_AZURE_CHAT_SKU_REJECTS_BATCH + ) with pytest.raises(openai.BadRequestError): await _create_batches(router, "gpt-mixed", attempts=1) @@ -14822,7 +7722,9 @@ async def test_acreate_batch_with_tags_that_match_only_the_chat_deployment_reach ) with respx.mock(assert_all_called=False) as respx_mock: - chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock(return_value=_AZURE_CHAT_SKU_REJECTS_BATCH) + chat_route: Final = respx_mock.post(host="chat-sku.openai.azure.com").mock( + return_value=_AZURE_CHAT_SKU_REJECTS_BATCH + ) batch_route: Final = respx_mock.post(host="batch-sku.openai.azure.com").mock( return_value=httpx.Response(200, json=_AZURE_BATCH_CREATED) ) @@ -14912,13 +7814,16 @@ async def test_acreate_batch_request_bedrock_tags_override_deployment_tags(): mock_client = MagicMock() mock_client.post = AsyncMock(side_effect=lambda *args, **kwargs: fake_response()) - with patch.object( - CommonBatchFilesUtils, - "sign_aws_request", - return_value=({"Authorization": "signed"}, b"{}"), - ) as mock_sign, patch( - "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", - return_value=mock_client, + with ( + patch.object( + CommonBatchFilesUtils, + "sign_aws_request", + return_value=({"Authorization": "signed"}, b"{}"), + ) as mock_sign, + patch( + "litellm.llms.custom_httpx.llm_http_handler.get_async_httpx_client", + return_value=mock_client, + ), ): await router.acreate_batch( model="bedrock-batch-model", @@ -14947,13 +7852,13 @@ async def test_avector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_asearch = AsyncMock(return_value=expected_response) # Router.__init__ binds asearch via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.asearch", new=mock_asearch): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.asearch", new=mock_asearch + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -14987,7 +7892,9 @@ async def test_avector_store_create_does_not_inject_router(): } ] ) - with patch("litellm.vector_stores.main.acreate", new=mock_acreate): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.acreate", new=mock_acreate + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface create_response = await router.avector_store_create(model=None, custom_llm_provider="openai") assert create_response is expected_response @@ -15003,13 +7910,13 @@ def test_vector_store_search_injects_router(): """ from litellm.types.vector_stores import VectorStoreSearchResponse - expected_response = VectorStoreSearchResponse( - object="vector_store.search_results.page", search_query="q", data=[] - ) + expected_response = VectorStoreSearchResponse(object="vector_store.search_results.page", search_query="q", data=[]) mock_search = MagicMock(return_value=expected_response) # Router.__init__ binds search via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.search", new=mock_search): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable + with patch( + "litellm.vector_stores.main.search", new=mock_search + ): # test-quality-ok: the SDK call is the only place the injected router kwarg is observable router = litellm.Router( model_list=[ { @@ -15018,9 +7925,7 @@ def test_vector_store_search_injects_router(): } ] ) - search_response = router.vector_store_search( - vector_store_id="v", query="q", custom_llm_provider="s3_vectors" - ) + search_response = router.vector_store_search(vector_store_id="v", query="q", custom_llm_provider="s3_vectors") assert search_response is expected_response mock_search.assert_called_once() @@ -15034,7 +7939,9 @@ def test_vector_store_create_does_not_inject_router(): mock_create = MagicMock(return_value=expected_response) # Router.__init__ binds create via a local import, so patch the module # attribute before constructing the Router. - with patch("litellm.vector_stores.main.create", new=mock_create): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface + with patch( + "litellm.vector_stores.main.create", new=mock_create + ): # test-quality-ok: the SDK call is the only place a leaked router kwarg would surface router = litellm.Router( model_list=[ { @@ -15069,9 +7976,7 @@ class TestPreRoutingStrategyRegistryLifecycle: def _complexity_router_params(default_model: str, tags=None) -> dict: return { "model": "auto_router/complexity_router", - "complexity_router_config": { - "tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - }, + "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}}, "complexity_router_default_model": default_model, **({"tags": tags} if tags else {}), } @@ -15240,22 +8145,6 @@ class TestPreRoutingStrategyRegistryLifecycle: assert len(registered) == 1 assert set(registered[0].strategy.config.available_models) == {"gpt-4o", "gpt-4o-mini"} - def test_delete_repairs_indices_even_when_strategy_release_fails(self): - """Structural removal and strategy release are not equally critical. Once the entry - leaves model_list the index maps must be repaired no matter what, so releasing the - registry slot runs after that repair and cannot abandon the router half-updated.""" - router = self._router_with_complexity_router() - idx = router.model_id_to_deployment_index_map["router-1"] - router.model_list[idx] = {"model_name": "smart-router", "litellm_params": None} - - returned = router.delete_deployment(id="router-1") - - assert returned is not None - assert "router-1" not in router.model_id_to_deployment_index_map - assert all(entry.get("model_info", {}).get("id") != "router-1" for entry in router.model_list) - assert router.get_deployment(model_id="router-1") is None - assert "gpt-4o" in self._model_names(router) - def test_delete_of_adaptive_enabled_complexity_router_frees_both_registries(self): """A complexity router with adaptive set is registered in BOTH complexity_routers and adaptive_routers under the same (model_name, tags). Releasing only the first @@ -15376,9 +8265,7 @@ class TestPreRoutingStrategyRegistryLifecycle: deployment=Deployment( model_name="hybrid-router", litellm_params=LiteLLM_Params( - **self._hybrid_router_params( - {"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"} - ) + **self._hybrid_router_params({"SIMPLE": "gpt-4o-mini", "MEDIUM": "gpt-4o", "COMPLEX": "gpt-4o"}) ), model_info=ModelInfo(id="router-1", db_model=True), ) @@ -15413,44 +8300,6 @@ class TestPreRoutingStrategyRegistryLifecycle: assert "smart-router" in router.adaptive_routers - def test_unregister_pre_routing_strategy_scopes_the_drop_by_tags(self): - """The bool return drives the hook re-sync; a tag mismatch must report False and - leave the registry untouched, and dropping the last entry must free the key.""" - from litellm.types.router import TaggedPreRoutingStrategy - - registry = { - "m": [ - TaggedPreRoutingStrategy(tags=("team-a",), strategy=object()), - TaggedPreRoutingStrategy(tags=(), strategy=object()), - ] - } - - assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ("team-b",)) is False - assert len(registry["m"]) == 2 - - assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ("team-a",)) is True - assert [entry.tags for entry in registry["m"]] == [()] - - assert litellm.Router._unregister_pre_routing_strategy(registry, "m", ()) is True - assert "m" not in registry - - def test_unregister_for_deployment_ignores_non_router_deployments(self): - """Direct twin of the endpoint-level test: a regular deployment that shares a - router's model_name must not evict the router's registry slot.""" - from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo - - router = self._router_with_complexity_router() - - router._unregister_pre_routing_strategy_for_deployment( - deployment=Deployment( - model_name="smart-router", - litellm_params=LiteLLM_Params(model="openai/gpt-4o"), - model_info=ModelInfo(id="plain-1", db_model=True), - ) - ) - - assert "smart-router" in router.complexity_routers - def test_sync_adaptive_router_hooks_keeps_one_hook_per_registered_router(self): """Re-syncing must replace, not accumulate: a duplicated hook double-fires bandit signal recording for every request.""" @@ -15465,33 +8314,6 @@ class TestPreRoutingStrategyRegistryLifecycle: hooks = litellm_module.logging_callback_manager.get_custom_loggers_for_type(AdaptiveRouterPostCallHook) assert len(hooks) == 1 - def test_deployment_participates_in_adaptive_routing_matrix(self): - """The upsert finalize re-run keys off this predicate for both the incoming and - outgoing deployment; a false negative silently strands the adaptive companion.""" - from litellm.types.router import LiteLLM_Params - - router = self._router_with_complexity_router() - - cases = [ - ({"model": "auto_router/adaptive_router", "adaptive_router_config": {}}, True), - (self._hybrid_router_params({"SIMPLE": "gpt-4o-mini"}), True), - (self._complexity_router_params("gpt-4o"), False), - ( - { - "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": {"SIMPLE": "gpt-4o-mini"}, "adaptive": False}, - "complexity_router_default_model": "gpt-4o", - }, - False, - ), - ({"model": "openai/gpt-4o"}, False), - ] - for params, expected in cases: - actual = router._deployment_participates_in_adaptive_routing( - litellm_params=LiteLLM_Params(**params) - ) - assert actual is expected, params["model"] - def test_model_info_is_active_for_environment_matrix(monkeypatch): """The model-write endpoints consult this predicate to tell a deliberately @@ -15512,57 +8334,6 @@ def test_model_info_is_active_for_environment_matrix(monkeypatch): model_info_is_active_for_environment(model_info={"supported_environments": ["production"]}) -def test_pre_call_checks_uses_deployment_model_when_model_info_lookup_raises(monkeypatch): - """ - The supported-params check must run against the deployment's own - provider-qualified model. Resolving the per-deployment model only after the - model-info lookup leaves it unset whenever that lookup raises (an - unregistered custom model), so the check falls back to the bare model group - name and the request dies with 'LLM Provider NOT provided'. - """ - monkeypatch.setattr(litellm, "drop_params", False) - - router = litellm.Router( - model_list=[ - { - "model_name": "custom-alias", - "litellm_params": {"model": "hosted_vllm/not-in-the-catalog"}, - } - ], - enable_pre_call_checks=True, - ) - - def _raise_unmapped(**kwargs): - raise ValueError("This model isn't mapped yet") - - monkeypatch.setattr(router, "_core_get_router_model_info", _raise_unmapped) - - seen: list[tuple] = [] - original_get_supported_openai_params = litellm.get_supported_openai_params - - def _record(model, custom_llm_provider=None, **kwargs): - seen.append((model, custom_llm_provider)) - return original_get_supported_openai_params(model=model, custom_llm_provider=custom_llm_provider, **kwargs) - - monkeypatch.setattr(litellm, "get_supported_openai_params", _record) - - deployments = [ - { - "litellm_params": {"model": "hosted_vllm/not-in-the-catalog"}, - "model_info": {"id": "d1"}, - } - ] - result = router._pre_call_checks( - model="custom-alias", - healthy_deployments=deployments, - messages=[{"role": "user", "content": "hi"}], - request_kwargs={"response_format": {"type": "json_object"}}, - ) - - assert len(result) == 1 - assert seen == [("not-in-the-catalog", "hosted_vllm")] - - def test_pre_call_checks_keeps_deployment_when_provider_is_unresolvable(monkeypatch): """ Pre-call checks filter deployments; they must never be the thing that fails @@ -15661,39 +8432,6 @@ class TestUpsertDeploymentRollback: assert router.get_deployment(model_id="fresh-1") is None assert router.model_list == [] - def test_restore_re_adds_popped_deployment(self): - router = litellm.Router( - model_list=[ - { - "model_name": "prod-model", - "litellm_params": {"model": "gpt-4o", "api_key": "sk-old"}, - "model_info": {"id": "prod-1", "db_model": True}, - } - ], - ignore_invalid_deployments=True, - ) - previous = router.get_deployment(model_id="prod-1") - router.delete_deployment(id="prod-1") - assert router.has_model_id("prod-1") is False - - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) - - restored = router.get_deployment(model_id="prod-1") - assert restored is not None - assert restored.litellm_params.model == "gpt-4o" - - router._restore_deployment_after_failed_upsert( - previous_deployment=previous, model_id="prod-1" - ) - assert len(router.model_list) == 1 - - router._restore_deployment_after_failed_upsert( - previous_deployment=None, model_id="prod-1" - ) - assert len(router.model_list) == 1 - class TestUpsertDeploymentRename: """ @@ -15980,31 +8718,6 @@ class TestClaudeCodeSubagentSessionRouterBinding: assert response is None redis_cache.async_delete_cache.assert_awaited_once() - @pytest.mark.asyncio - async def test_redis_read_failure_does_not_reject_a_subagent_request(self): - from litellm.caching.caching import RedisCache - - router = self._router() - request_kwargs = self._request_kwargs(agent_id="agent-1234") - cache_key = router._claude_code_session_router_cache_key(request_kwargs) - assert cache_key is not None - await router._claude_code_session_router_cache.in_memory_cache.async_set_cache( - cache_key, - "smart-router", - ) - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(side_effect=Exception("Redis circuit breaker is open")) - router.update_redis_cache(cache=redis_cache) - - response = await router.async_pre_routing_hook( - model="expensive-model", - request_kwargs=request_kwargs, - ) - - assert response is None - assert "model_group" not in request_kwargs["metadata"] - redis_cache.async_get_cache.assert_awaited_once() - @pytest.mark.asyncio async def test_redis_write_failures_do_not_reject_main_or_subagent_requests(self): from litellm.caching.caching import RedisCache @@ -16083,20 +8796,6 @@ class TestClaudeCodeSubagentSessionRouterBinding: redis_cache.async_set_cache.assert_not_awaited() redis_cache.async_delete_cache.assert_not_awaited() - @pytest.mark.asyncio - async def test_session_bindings_do_not_evict_router_rate_limit_state(self): - router = self._router() - assert router._update_usage(deployment_id="deployment-id", parent_otel_span=None) == 1 - - for session_index in range(201): - request_kwargs = self._request_kwargs() - request_kwargs["proxy_server_request"]["headers"]["X-Claude-Code-Session-Id"] = ( - f"session-{session_index:04d}" - ) - await router.async_pre_routing_hook(model="smart-router", request_kwargs=request_kwargs) - - assert router._update_usage(deployment_id="deployment-id", parent_otel_span=None) == 2 - @pytest.mark.asyncio async def test_background_and_fallback_requests_do_not_clear_the_session_router(self): router = self._router() @@ -16397,9 +9096,6 @@ class TestTaggedAutoRouterOnSharedModelName: ) assert deployment["litellm_params"]["model"] == "openai/gpt-4o" - def test_deployment_without_litellm_params_mapping_is_not_a_marker(self): - assert litellm.Router._is_strategy_marker_deployment({"model_name": "gpt4o"}) is False - class TestAutoRouterSharedModelNameConnectionParams: """A plain deployment sharing its model_name with an `auto_router/` marker must not have @@ -16451,18 +9147,14 @@ class TestAutoRouterSharedModelNameConnectionParams: return httpx.Response( status_code=200, json={ - "candidates": [ - {"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"} - ], + "candidates": [{"content": {"parts": [{"text": "Paris"}], "role": "model"}, "finishReason": "STOP"}], "usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 1, "totalTokenCount": 6}, "modelVersion": "gemini-3.6-flash", }, request=httpx.Request("POST", "https://generativelanguage.googleapis.com"), ) - @pytest.mark.parametrize( - "plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"] - ) + @pytest.mark.parametrize("plain_entry_first", [True, False], ids=["plain_entry_first", "marker_entry_first"]) async def test_routed_tier_call_goes_out_on_its_own_endpoint_and_credentials(self, plain_entry_first): """The outbound provider request for the routed tier hits the tier's own Gemini host with the tier's own key, never the plain sibling's api_base or api_key.""" @@ -16594,9 +9286,7 @@ async def _drive_cyclic_fallback(router, capture, recorder=None, **request_kwarg litellm.callbacks.append(recorder) try: with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs - ) + await router.acompletion(model="group-a", messages=[{"role": "user", "content": "hi"}], **request_kwargs) finally: router_logger.removeHandler(capture) router_logger.setLevel(previous_level) @@ -16870,9 +9560,7 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): await _drive_cyclic_fallback( _cyclic_fallback_router(), capture, - mock_response=litellm.InternalServerError( - message=huge_message, llm_provider="openai", model="group-a" - ), + mock_response=litellm.InternalServerError(message=huge_message, llm_provider="openai", model="group-a"), ) assert capture.messages, "the fallback failure path did not log at ERROR" @@ -16880,17 +9568,6 @@ async def test_fallback_failure_detail_from_upstream_is_bounded(): assert max(len(message) for message in capture.messages) < 5_000 -def test_stamp_or_clear_metadata_key_writes_and_clears_both_buckets(): - request_kwargs = {"metadata": {}} - litellm.Router._stamp_or_clear_metadata_key(request_kwargs=request_kwargs, key="probe", value=7) - assert request_kwargs["metadata"]["probe"] == 7 - - stale_kwargs = {"metadata": {"probe": 7}, "litellm_metadata": {"probe": 7}} - litellm.Router._stamp_or_clear_metadata_key(request_kwargs=stale_kwargs, key="probe", value=None) - assert "probe" not in stale_kwargs["metadata"] - assert "probe" not in stale_kwargs["litellm_metadata"] - - @pytest.mark.parametrize( "complexity_router_config,expect_callback", [ @@ -16929,24 +9606,6 @@ def test_complexity_router_registers_affinity_callback_for_deployment_pin(comple litellm.logging_callback_manager.remove_callback_from_all_lists(cb) -def test_ensure_deployment_affinity_callback_is_idempotent(): - from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( - DeploymentAffinityCheck, - ) - - router = litellm.Router(model_list=[]) - try: - router._ensure_deployment_affinity_callback() - router._ensure_deployment_affinity_callback() - affinity_callbacks = [ - cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck) - ] - assert len(affinity_callbacks) == 1 - finally: - for cb in router.optional_callbacks or []: - litellm.logging_callback_manager.remove_callback_from_all_lists(cb) - - def test_get_router_model_info_does_not_wipe_cached_pricing(): """A Deployment's model_info declares the mirrored pricing fields with None defaults; merging it must not write those Nones into the lru_cache'd dict get_model_info() owns, @@ -17082,9 +9741,7 @@ class TestModelGroupAliasReachesPreRoutingStrategies: router = self._router("auto_routers") metadata: dict = {} - response = await router.acompletion( - model="smart-alias", messages=self._messages(), metadata=metadata - ) + response = await router.acompletion(model="smart-alias", messages=self._messages(), metadata=metadata) assert response.choices[0].message.content == "routed by the tier" assert metadata["model_group"] == "smart-alias" @@ -17326,64 +9983,6 @@ class TestTeamPublicNameReachesPreRoutingStrategies: assert response is not None assert response.model == "gemini-flash" - @pytest.mark.asyncio - async def test_strategy_resolution_agrees_with_the_deployment_path_for_every_principal(self): - router = self._router( - {self.INTERNAL_NAME: self._RewriteStrategy()}, - extra_deployments=({"model_name": "shared-name", "litellm_params": {"model": "openai/gpt-4o"}},), - ) - principals = { - "team": self._team_request(), - "other-team": self._team_request(self.OTHER_TEAM), - "no-team": self._team_request(None), - "admin": {"metadata": {"user_api_key_auth": SimpleNamespace(user_role="proxy_admin")}}, - } - for principal, request_kwargs in principals.items(): - for model in (self.PUBLIC_NAME, "shared-name", "gemini-flash", "missing"): - resolved = [d["model_name"] for d in router.deployments_for_request(model, request_kwargs)] - callable_names = [ - name - for name, deployment in zip(resolved, router.deployments_for_request(model, request_kwargs)) - if not router._is_strategy_marker_deployment(deployment) - ] - if resolved and not callable_names: - with pytest.raises(litellm.BadRequestError, match="strategy router marker"): - router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) - elif not resolved: - with pytest.raises(litellm.BadRequestError): - router._common_checks_available_deployment(model=model, request_kwargs=request_kwargs) - else: - _, deployments = router._common_checks_available_deployment( - model=model, request_kwargs=request_kwargs - ) - assert [d["model_name"] for d in deployments] == callable_names, (principal, model) - - def test_drop_strategy_markers_keeps_plain_deployments_and_rejects_marker_only_sets(self): - router = self._router({}) - marker = router.model_list[0] - plain = {"model_name": "plain", "litellm_params": {"model": "openai/gpt-4o"}} - - assert router._drop_strategy_markers("x", [marker, plain]) == [plain] - assert router._drop_strategy_markers("x", [plain]) == [plain] - assert router._drop_strategy_markers("x", []) == [] - with pytest.raises(litellm.BadRequestError, match="strategy router marker"): - router._drop_strategy_markers("x", [marker]) - - def test_team_deployments_across_teams_unions_one_team_and_rejects_two(self): - other_team_marker = { - **self._team_marker(self.SIBLING_INTERNAL_NAME), - "model_info": {"team_id": self.OTHER_TEAM, "team_public_model_name": self.PUBLIC_NAME}, - } - one_team = self._router({}) - two_teams = self._router({}, markers=(self._team_marker(self.INTERNAL_NAME), other_team_marker)) - - assert [d["model_name"] for d in one_team._team_deployments_across_teams(self.PUBLIC_NAME)] == [ - self.INTERNAL_NAME - ] - assert one_team._team_deployments_across_teams("missing") == [] - with pytest.raises(litellm.BadRequestError, match="multiple teams"): - two_teams._team_deployments_across_teams(self.PUBLIC_NAME) - def test_compression_policy_follows_the_same_resolution_for_every_principal(self): from litellm.proxy.guardrails.auto_router_compression import AutoRouterCompressionPolicy, policy_for_model @@ -17608,7 +10207,6 @@ class TestAutoRouterCompressionDecoupling: @pytest.mark.usefixtures("local_model_cost_map") - @pytest.mark.usefixtures("local_model_cost_map") class TestAzureBaseModelFallbackLogging: """When an azure deployment has no base_model but its model name is a known @@ -17632,41 +10230,6 @@ class TestAzureBaseModelFallbackLogging: ] ) - def test_map_known_deployment_name_resolves_without_error_log(self): - router = self._router_with_azure_deployment("azure/gpt-4o") - - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: - model_info = router.get_router_model_info( - deployment=None, received_model_name="my-group", id="azure-base-model-test-id" - ) - - assert not any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), f"unexpected error log: {mock_error.call_args_list}" - # the fallback resolution must actually surface the map values - assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o"]["max_input_tokens"] - assert model_info["input_cost_per_token"] == litellm.model_cost["azure/gpt-4o"]["input_cost_per_token"] - - def test_unmappable_deployment_name_still_logs_error(self): - router = self._router_with_azure_deployment("azure/my-custom-deployment-name") - - with patch( - "litellm.router.verbose_router_logger.error" - ) as mock_error: - model_info = router.get_router_model_info( - deployment=None, received_model_name="my-group", id="azure-base-model-test-id" - ) - - assert any( - "Could not identify azure model" in str(call) - for call in mock_error.call_args_list - ), "expected the error log for an unmappable azure deployment name" - # unmappable names resolve to a zeroed stub — unchanged behavior - assert model_info.get("max_input_tokens") is None - def test_explicit_base_model_still_wins(self): router = litellm.Router( model_list=[ @@ -17690,95 +10253,6 @@ class TestAzureBaseModelFallbackLogging: ) assert model_info["max_input_tokens"] == litellm.model_cost["azure/gpt-4o-mini"]["max_input_tokens"] -def test_model_group_info_intersects_supported_reasoning_efforts(): - router = litellm.Router( - model_list=[ - { - "model_name": "smart-group", - "litellm_params": {"model": "anthropic/opus-like"}, - "model_info": {"id": "opus-like-deployment"}, - }, - { - "model_name": "smart-group", - "litellm_params": {"model": "openai/mini-like"}, - "model_info": {"id": "mini-like-deployment"}, - }, - ] - ) - - def _model_info(model_id: str, model_name: str): - if model_id == "opus-like-deployment": - return { - "key": model_name, - "litellm_provider": "anthropic", - "mode": "chat", - "supports_reasoning": True, - "supports_xhigh_reasoning_effort": True, - "supports_max_reasoning_effort": True, - } - return { - "key": model_name, - "litellm_provider": "openai", - "mode": "chat", - "supports_reasoning": True, - "supports_none_reasoning_effort": False, - "supports_minimal_reasoning_effort": True, - "supports_xhigh_reasoning_effort": False, - } - - with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info): - result = router._set_model_group_info( - model_group="smart-group", - user_facing_model_group_name="smart-group", - ) - - assert result is not None - # opus-like offers all seven levels, mini-like lacks none/xhigh/max; only the common set survives, - # so the group never advertises an effort routing could hand to a deployment that rejects it. - assert result.supported_reasoning_efforts == ("minimal", "low", "medium", "high") - - -def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_off_the_map(): - """The router fills every ModelInfo key, so a deployment absent from the model map arrives with - supports_reasoning None rather than with the key missing. Its synthesized entry carries no mode, - which is what separates it from a mapped non-reasoning model, and nothing being known about it is - no evidence that the unknown deployment accepts levels its mapped sibling supports. The group - therefore reports unknown instead of advertising a value routing might send to either one.""" - router = litellm.Router( - model_list=[ - { - "model_name": "smart-group", - "litellm_params": {"model": "anthropic/opus-like"}, - "model_info": {"id": "opus-like-deployment"}, - }, - { - "model_name": "smart-group", - "litellm_params": {"model": "openai/unmapped-model"}, - "model_info": {"id": "unmapped-deployment"}, - }, - ] - ) - - def _model_info(model_id: str, model_name: str): - if model_id == "opus-like-deployment": - return { - "key": model_name, - "litellm_provider": "anthropic", - "mode": "chat", - "supports_reasoning": True, - "supports_max_reasoning_effort": True, - } - return {"key": model_name, "litellm_provider": "openai", "mode": None, "supports_reasoning": None} - - with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info): - result = router._set_model_group_info( - model_group="smart-group", - user_facing_model_group_name="smart-group", - ) - - assert result is not None - assert result.supported_reasoning_efforts is None - @pytest.mark.parametrize( "model,provider,expected", @@ -17799,11 +10273,15 @@ def test_model_group_info_reasoning_efforts_are_unknown_when_any_deployment_is_o def test_model_group_info_fast_mode_uses_exact_provider_catalog( local_model_cost_map: None, model: str, provider: str | None, expected: bool, operator_flag: bool ) -> None: - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"}, - "model_info": {"supports_fast_mode": operator_flag}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": model, "custom_llm_provider": provider, "api_key": "fake-key"}, + "model_info": {"supports_fast_mode": operator_flag}, + } + ] + ) result: Final = router.get_model_group_info("fast-group") @@ -17815,16 +10293,21 @@ def test_model_group_info_fast_mode_uses_exact_provider_catalog( def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean( local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, flag: object ) -> None: - entry: Final = {key: value for key, value in litellm.model_cost["claude-opus-5"].items() - if key != "supports_fast_mode"} + entry: Final = { + key: value for key, value in litellm.model_cost["claude-opus-5"].items() if key != "supports_fast_mode" + } if flag is not None: entry["supports_fast_mode"] = flag monkeypatch.setitem(litellm.model_cost, "claude-opus-5", entry) - router: Final = Router(model_list=[{ - "model_name": "fast-group", - "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"}, - "model_info": {"supports_fast_mode": True}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "fast-group", + "litellm_params": {"model": "anthropic/claude-opus-5", "api_key": "fake-key"}, + "model_info": {"supports_fast_mode": True}, + } + ] + ) result: Final = router.get_model_group_info("fast-group") @@ -17832,13 +10315,16 @@ def test_model_group_info_fast_mode_fails_closed_without_explicit_boolean( assert result.supports_fast_mode is False -@pytest.mark.parametrize("other_model,expected", [ - ("anthropic/claude-opus-4-8", True), - ("anthropic/claude-opus-4-7", False), - ("anthropic/off-map-opus", False), - ("vertex_ai/claude-opus-5", False), - ("bedrock/claude-opus-5", False), -]) +@pytest.mark.parametrize( + "other_model,expected", + [ + ("anthropic/claude-opus-4-8", True), + ("anthropic/claude-opus-4-7", False), + ("anthropic/off-map-opus", False), + ("vertex_ai/claude-opus-5", False), + ("bedrock/claude-opus-5", False), + ], +) @pytest.mark.parametrize("reverse", [True, False]) def test_model_group_info_fast_mode_requires_every_deployment( local_model_cost_map: None, other_model: str, expected: bool, reverse: bool @@ -17860,177 +10346,6 @@ def test_model_group_info_fast_mode_requires_every_deployment( assert result.supports_fast_mode is expected -def test_model_group_info_surfaces_supports_parallel_function_calling(local_model_cost_map): - """``/model_group/info`` folds each deployment's registry flags into the group; a deployment whose - registry entry declares parallel function calling must flip the group to True instead of False.""" - router = litellm.Router( - model_list=[ - { - "model_name": "glm-group", - "litellm_params": {"model": "together_ai/zai-org/GLM-5.3-Flash", "api_key": "fake-key"}, - } - ] - ) - - result = router._set_model_group_info(model_group="glm-group", user_facing_model_group_name="glm-group") - - assert result is not None - assert result.supports_parallel_function_calling is True - - -def test_model_group_info_reasoning_efforts_empty_on_a_mapped_non_reasoning_deployment(): - """A group mixing a reasoning model with one the map knows is not a reasoning model shares no - level, so it advertises none and the picker offers nothing rather than a level routing would - hand to a deployment that rejects it.""" - router = litellm.Router( - model_list=[ - { - "model_name": "mixed-group", - "litellm_params": {"model": "anthropic/opus-like"}, - "model_info": {"id": "opus-like-deployment"}, - }, - { - "model_name": "mixed-group", - "litellm_params": {"model": "openai/plain-chat"}, - "model_info": {"id": "plain-chat-deployment"}, - }, - ] - ) - - def _model_info(model_id: str, model_name: str): - if model_id == "opus-like-deployment": - return { - "key": model_name, - "litellm_provider": "anthropic", - "mode": "chat", - "supports_reasoning": True, - "supports_max_reasoning_effort": True, - } - return {"key": model_name, "litellm_provider": "openai", "mode": "chat", "supports_reasoning": None} - - with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info): - result = router._set_model_group_info( - model_group="mixed-group", - user_facing_model_group_name="mixed-group", - ) - - assert result is not None - assert result.supported_reasoning_efforts == () - - -def test_model_group_info_reasoning_efforts_ignore_a_value_declared_in_model_info(): - """The group's levels are computed from its deployments, so a value an operator left in one - deployment's model_info must not seed them. Seeding let the first deployment read narrow the - whole group while the same value on any other deployment was silently ignored.""" - router = litellm.Router( - model_list=[ - { - "model_name": "declared-group", - "litellm_params": {"model": "openai/first-reasoner"}, - "model_info": {"id": "first-deployment"}, - }, - { - "model_name": "declared-group", - "litellm_params": {"model": "openai/second-reasoner"}, - "model_info": {"id": "second-deployment"}, - }, - ] - ) - - def _model_info(model_id: str, model_name: str): - info = { - "key": model_name, - "litellm_provider": "openai", - "mode": "chat", - "supports_reasoning": True, - "supports_none_reasoning_effort": True, - } - if model_id == "first-deployment": - info["supported_reasoning_efforts"] = ("high",) - return info - - with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info): - result = router._set_model_group_info( - model_group="declared-group", - user_facing_model_group_name="declared-group", - ) - - assert result is not None - assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high") - - -def test_model_group_info_survives_a_junk_typed_operator_effort_value(): - """A deployment's registered model_info reads back with whatever the operator wrote under any - key, so a wrong-typed supported_reasoning_efforts must not fail the group's info. Only the - constructor's trailing override keeps the junk away from ModelGroupInfo validation.""" - router = litellm.Router( - model_list=[ - { - "model_name": "junk-declared-group", - "litellm_params": {"model": "openai/lone-reasoner"}, - "model_info": {"id": "junk-deployment"}, - }, - ] - ) - - def _model_info(model_id: str, model_name: str): - return { - "key": model_name, - "litellm_provider": "openai", - "mode": "chat", - "supports_reasoning": True, - "supports_none_reasoning_effort": True, - "supported_reasoning_efforts": "high", - } - - with patch.object(router, "_core_get_deployment_model_info", side_effect=_model_info): - result = router._set_model_group_info( - model_group="junk-declared-group", - user_facing_model_group_name="junk-declared-group", - ) - - assert result is not None - assert result.supported_reasoning_efforts == ("none", "minimal", "low", "medium", "high") - - -def test_model_group_info_reasoning_efforts_are_unknown_for_an_operator_declared_mode(): - """A deployment is registered in the cost map under its own id with whatever model_info the - operator wrote, so a mode they set themselves reads back exactly like one the map supplied. Only - a mode the map supplied marks the deployment as known. An off-map deployment carrying an - operator mode remains unknown and must keep the whole group's level support unknown.""" - from litellm.litellm_core_utils.reasoning_effort_capability import resolve_supported_reasoning_efforts - - mapped_model = "openai/gpt-5.6-sol" - expected = resolve_supported_reasoning_efforts( - litellm.get_model_info(model=mapped_model), - deployment_is_mapped=True, - ) - assert expected - - router = litellm.Router( - model_list=[ - { - "model_name": "smart-group", - "litellm_params": {"model": mapped_model, "api_key": "sk-fake"}, - "model_info": {"id": "mapped-deployment"}, - }, - { - "model_name": "smart-group", - "litellm_params": {"model": "openai/a-model-the-map-never-heard-of", "api_key": "sk-fake"}, - "model_info": {"id": "off-map-deployment", "mode": "chat"}, - }, - ] - ) - - result = router._set_model_group_info( - model_group="smart-group", - user_facing_model_group_name="smart-group", - ) - - assert result is not None - assert result.supported_reasoning_efforts is None - - class TestAddDeploymentApiBaseProviderResolution: def test_bare_model_with_known_api_base_initializes(self): router = litellm.Router( @@ -18092,6 +10407,7 @@ class TestAddDeploymentApiBaseProviderResolution: assert deployment is not None assert deployment.litellm_params.custom_llm_provider == "openai" + # ===================================================================== # anthropic_messages mid-stream-fallback helpers, added for #24004 # (mid-stream fallback not supported for anthropic_messages route type). @@ -18112,121 +10428,10 @@ class TestAddDeploymentApiBaseProviderResolution: # ===================================================================== -async def _anthropic_messages_empty_generator(): - return - yield # pragma: no cover - makes this an async generator - - -def _anthropic_messages_make_wrapper() -> FallbackAwareAnthropicMessagesStream: - """A minimal wrapper for tests that call _aanthropic_messages_fallback_attempt - directly, bypassing _aanthropic_messages_streaming_iterator.""" - return FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), object()) - - -def _anthropic_messages_make_router(**router_kwargs) -> Router: - """A fallback-only router: no same-group retries unless a test asks for them.""" - router_kwargs.setdefault("fallbacks", [{"primary": ["fallback"]}]) - router_kwargs.setdefault("num_retries", 0) - return Router( - model_list=[ - { - "model_name": "primary", - "litellm_params": { - "model": "anthropic/claude-sonnet-4-5", - "api_key": "sk-test", - }, - }, - { - "model_name": "fallback", - "litellm_params": { - "model": "bedrock/anthropic.claude-sonnet-4-5", - }, - }, - ], - **router_kwargs, - ) - - -class _AnthropicMessagesFakeByteStream: - """Minimal AsyncIterator[bytes], carrying _hidden_params like - AnthropicMessagesStreamingResponse does.""" - - def __init__(self, chunks: list) -> None: - self._chunks = list(chunks) - self._hidden_params = {"additional_headers": {"x-amzn-requestid": "req-1"}} - self.closed = False - self.closed_as_superseded = False - self.close_contexts: tuple[bool, ...] = () - - def __aiter__(self): - return self - - async def __anext__(self) -> bytes: - if not self._chunks: - raise StopAsyncIteration - return self._chunks.pop(0) - - async def aclose(self) -> None: - self.close_contexts = (*self.close_contexts, closing_superseded_attempt()) - self.closed_as_superseded = self.closed_as_superseded or closing_superseded_attempt() - self.closed = True - - -class _AnthropicMessagesRaisingByteStream: - """Simulates the completion-bridge path: no error SSE chunk is ever - yielded, the underlying CustomStreamWrapper raises MidStreamFallbackError - directly out of the iterator instead (a Vertex AI transport drop).""" - - def __init__(self, chunks: list, error: Exception) -> None: - self._chunks = list(chunks) - self._error = error - self._hidden_params: dict = {} - self.closed = False - - def __aiter__(self): - return self - - async def __anext__(self) -> bytes: - if self._chunks: - return self._chunks.pop(0) - raise self._error - - async def aclose(self) -> None: - self.closed = True - - -class _AnthropicMessagesFallbackByteStream: - def __init__(self, chunks: list, hidden_params: dict | None = None) -> None: - self._chunks = list(chunks) - self._hidden_params = hidden_params if hidden_params is not None else {} - - def __aiter__(self): - return self - - async def __anext__(self) -> bytes: - if not self._chunks: - raise StopAsyncIteration - return self._chunks.pop(0) - - def _anthropic_messages_overloaded_error_chunk() -> bytes: return b'event: error\ndata: {"type": "error", "error": {"type": "overloaded_error", "message": "Overloaded"}}\n\n' -def _anthropic_messages_invalid_request_error_chunk() -> bytes: - return ( - b"event: error\n" - b'data: {"type": "error", "error": {"type": "invalid_request_error", "message": "bad request"}}\n\n' - ) - - -def _anthropic_messages_rate_limit_error_chunk() -> bytes: - return ( - b"event: error\n" - b'data: {"type": "error", "error": {"type": "rate_limit_error", "message": "Too many requests"}}\n\n' - ) - - def _anthropic_messages_content_chunk(text: str = "hi") -> bytes: payload = f'{{"type": "content_block_delta", "delta": {{"type": "text_delta", "text": "{text}"}}}}' return f"event: content_block_delta\ndata: {payload}\n\n".encode() @@ -18245,777 +10450,15 @@ def _anthropic_messages_ping_chunk() -> bytes: # -------- _aanthropic_messages_streaming_iterator (passthrough) -------- -@pytest.mark.asyncio -async def test_anthropic_messages_streaming_iterator_passthrough(): - """Without any error chunk, the wrapper forwards every chunk unchanged - and carries the source iterator's _hidden_params through (so response - headers like Bedrock's request-id keep flowing to the client).""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] - ) - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - - collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_content_chunk("hi"), _anthropic_messages_content_chunk(" there")] - assert wrapped._hidden_params["additional_headers"]["x-amzn-requestid"] == "req-1" - - -@pytest.mark.asyncio -async def test_anthropic_messages_streaming_iterator_flushes_buffered_lifecycle_frames_in_order(): - """Regression: lifecycle frames held back to guard against a mid-stream - fallback must still reach the client, in order, once real content - arrives - buffering them for the fallback-safety check must not silently - drop them on the happy path.""" - router = _anthropic_messages_make_router() - message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] - ) - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - - collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("hi"), message_stop] - - -@pytest.mark.asyncio -async def test_anthropic_messages_streaming_iterator_flushes_buffered_frames_on_stream_end(): - """Regression: if the primary stream ends with only lifecycle frames and - no content and no error, the buffered frames must still reach the - client rather than being silently swallowed.""" - router = _anthropic_messages_make_router() - message_stop = b'event: message_stop\ndata: {"type": "message_stop"}\n\n' - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), message_stop]) - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - - collected = [chunk async for chunk in wrapped] - assert collected == [_anthropic_messages_message_start_chunk(), message_stop] - - with pytest.raises(StopAsyncIteration): - await wrapped.__anext__() - - -@pytest.mark.asyncio -async def test_anthropic_messages_content_coalesced_with_error_in_one_physical_chunk_skips_fallback(): - """Greptile review round: transport-level buffering can coalesce a real - content_block_delta and a following retriable error into ONE physical - read from the source iterator. Since the whole chunk (content and error - together) is forwarded to the client atomically, the client genuinely - receives the content - so no fallback must be attempted, exactly as if - the two events had arrived as separate reads.""" - router = _anthropic_messages_make_router() - coalesced_chunk = _anthropic_messages_content_chunk("partial") + _anthropic_messages_overloaded_error_chunk() - source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [coalesced_chunk] - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_ping_behind_buffered_lifecycle_frame_is_forwarded_live(): - """A `ping` behind buffered lifecycle frames still reaches the client - live: it carries no lifecycle, so it cannot create overlapping - lifecycles, and it keeps the connection alive while a fallback-able - stream holds message_start back through a long thinking pass.""" - router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}]) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - yield _anthropic_messages_ping_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) - - assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() - content_released.set() - assert [chunk async for chunk in wrapped] == [ - _anthropic_messages_message_start_chunk(), - _anthropic_messages_content_chunk("hi"), - ] - - -@pytest.mark.asyncio -async def test_anthropic_messages_split_ping_stays_in_order_behind_buffered_lifecycle_frame(): - """A ping the transport splits across two reads is not a whole frame, so - neither fragment may jump ahead of the buffered message_start: yielding - the head live and flushing the tail behind message_start would splice a - lifecycle frame into the middle of the ping on the wire.""" - router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}]) - ping_head, ping_tail = b'event: ping\ndata: {"ty', b'pe": "ping"}\n\n' - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), ping_head, ping_tail, _anthropic_messages_content_chunk("hi")] - ) - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) - - assert [chunk async for chunk in wrapped] == [ - _anthropic_messages_message_start_chunk(), - ping_head, - ping_tail, - _anthropic_messages_content_chunk("hi"), - ] - - -@pytest.mark.asyncio -async def test_anthropic_messages_no_fallback_message_start_reaches_client_before_content(): - """With no fallback able to take over, the stream is committed from the - first frame: message_start reaches the client live instead of waiting - behind the buffer for content that may be a whole thinking pass away.""" - router = _anthropic_messages_make_router(fallbacks=None) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) - - assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() - content_released.set() - assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] - - -@pytest.mark.asyncio -async def test_anthropic_messages_disabled_fallbacks_message_start_reaches_client_before_content(): - """A router with fallbacks configured cannot take over a request that - opted out with disable_fallbacks=True, so its lifecycle frames reach - the client live exactly like a no-fallback router's.""" - router = _anthropic_messages_make_router(fallbacks=[{"primary": ["fallback"]}]) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source(), initial_kwargs={"model": "primary", "disable_fallbacks": True} - ) - - assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() - content_released.set() - assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] - - -@pytest.mark.asyncio -async def test_anthropic_messages_no_fallback_error_frame_reaches_client_verbatim(): - """With no fallback able to take over, a retriable provider error frame - is forwarded verbatim instead of triggering a fallback that does not - exist, and the frames already received stay in order ahead of it.""" - router = _anthropic_messages_make_router(fallbacks=None) - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_default_wildcard_fallback_still_buffers_lifecycle_frames(): - """A "*" default fallback can take over for any group, so lifecycle - frames are still held back until real content commits the primary.""" - router = _anthropic_messages_make_router(fallbacks=[{"*": ["fallback"]}]) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) - - pending = asyncio.ensure_future(wrapped.__anext__()) - await asyncio.sleep(0.2) - assert not pending.done() - content_released.set() - assert await asyncio.wait_for(pending, timeout=1) == _anthropic_messages_message_start_chunk() - assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] - - -def _anthropic_messages_two_order_primary_model_list() -> list: - return [ - { - "model_name": "primary", - "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": 1}, - }, - { - "model_name": "primary", - "litellm_params": {"model": "bedrock/anthropic.claude-sonnet-4-5", "order": 2}, - }, - { - "model_name": "fallback", - "litellm_params": {"model": "bedrock/anthropic.claude-sonnet-4-5"}, - }, - ] - - -@pytest.mark.parametrize( - "router_kwargs,request_kwargs,expected", - [ - pytest.param({"fallbacks": None}, {"model": "primary"}, False, id="no-fallbacks"), - pytest.param({"fallbacks": [{"primary": ["fallback"]}]}, {"model": "primary"}, True, id="group-fallback"), - pytest.param({"fallbacks": [{"other": ["fallback"]}]}, {"model": "primary"}, False, id="unrelated-group"), - pytest.param( - {"fallbacks": [{"*": ["fallback"]}]}, - {"model": "primary", "fallbacks": None}, - False, - id="wildcard-overridden-by-request-none", - ), - pytest.param({"fallbacks": [{"*": ["fallback"]}]}, {"model": "primary"}, True, id="wildcard"), - pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": [{"model": "fallback"}]}, True, id="request-dict-fallback"), - pytest.param({"fallbacks": None}, {"model": "primary", "fallbacks": ["fallback"]}, True, id="request-list-fallback"), - pytest.param( - {"fallbacks": [{"primary": ["fallback"]}]}, - {"model": "primary", "disable_fallbacks": True}, - False, - id="disable-fallbacks", - ), - pytest.param( - {"fallbacks": None, "content_policy_fallbacks": [{"primary": ["fallback"]}]}, - {"model": "primary"}, - True, - id="content-policy-fallback", - ), - pytest.param({"fallbacks": None, "enable_weighted_failover": True}, {"model": "primary"}, True, id="weighted-failover"), - ], -) -def test_stream_can_fall_back_direct_call(router_kwargs, request_kwargs, expected): - router = _anthropic_messages_make_router(**router_kwargs) - assert router._stream_can_fall_back("primary", request_kwargs) is expected - - -@pytest.mark.parametrize( - "orders,expected", - [ - pytest.param([1, 2], True, id="distinct-orders-can-fall-back"), - pytest.param([1, 1], False, id="same-order-cannot-fall-back"), - ], -) -def test_stream_can_fall_back_order_levels(orders, expected): - router = Router( - model_list=[ - { - "model_name": "primary", - "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": order}, - } - for order in orders - ], - fallbacks=None, - ) - assert router._stream_can_fall_back("primary", {"model": "primary"}) is expected - - -@pytest.mark.parametrize( - "request_kwargs,expected", - [ - pytest.param({"model": "primary"}, True, id="no-target-order"), - pytest.param({"model": "primary", "_target_order": 1}, True, id="higher-order-remains"), - pytest.param({"model": "primary", "_target_order": 2}, False, id="top-order-no-order-fallback"), - pytest.param( - {"model": "primary", "_target_order": 2, "fallbacks": [{"primary": ["fallback"]}]}, - True, - id="top-order-external-fallback", - ), - ], -) -def test_stream_can_fall_back_order_target(request_kwargs, expected): - router = Router(model_list=_anthropic_messages_two_order_primary_model_list(), fallbacks=None) - assert router._stream_can_fall_back("primary", request_kwargs) is expected - - -def test_stream_order_levels_direct_call(): - router = Router( - model_list=[ - { - "model_name": "primary", - "litellm_params": {"model": "anthropic/claude-sonnet-4-5", "api_key": "sk-test", "order": order}, - } - for order in (2, 1, None) - ], - fallbacks=None, - ) - assert router._stream_order_levels("primary", {"model": "primary"}) == (1, 2) - - -@pytest.mark.asyncio -async def test_anthropic_messages_order_fallback_still_buffers_lifecycle_frames(): - """Two order levels in one group are a real fallback target for the - dispatcher, so lifecycle frames stay buffered until content commits.""" - router = Router(model_list=_anthropic_messages_two_order_primary_model_list(), fallbacks=None) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) - - pending = asyncio.ensure_future(wrapped.__anext__()) - await asyncio.sleep(0.2) - assert not pending.done() - content_released.set() - assert await asyncio.wait_for(pending, timeout=1) == _anthropic_messages_message_start_chunk() - assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] - - -@pytest.mark.asyncio -async def test_anthropic_messages_request_fallbacks_none_forwards_message_start_live(): - """A per-request fallbacks=None override disables the router's wildcard - fallback, so lifecycle frames reach the client live before content.""" - router = _anthropic_messages_make_router(fallbacks=[{"*": ["fallback"]}]) - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source(), initial_kwargs={"model": "primary", "fallbacks": None} - ) - - assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_message_start_chunk() - content_released.set() - assert [chunk async for chunk in wrapped] == [_anthropic_messages_content_chunk("hi")] - - -@pytest.mark.asyncio -async def test_anthropic_messages_leading_ping_keepalive_is_forwarded_live(): - """A `ping` that no lifecycle frame precedes is how a hold-back turn keeps - its connection alive (AgenticAnthropicStreamingIterator), so it must reach - the client at once rather than wait behind the pre-content buffer.""" - router = _anthropic_messages_make_router() - content_released = asyncio.Event() - - async def source(): - yield _anthropic_messages_ping_chunk() - await content_released.wait() - yield _anthropic_messages_message_start_chunk() - yield _anthropic_messages_content_chunk("hi") - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source(), initial_kwargs={"model": "primary"}) - - assert await asyncio.wait_for(wrapped.__anext__(), timeout=1) == _anthropic_messages_ping_chunk() - content_released.set() - assert [chunk async for chunk in wrapped] == [ - _anthropic_messages_message_start_chunk(), - _anthropic_messages_content_chunk("hi"), - ] - - -@pytest.mark.asyncio -async def test_anthropic_messages_leading_ping_does_not_disqualify_fallback(): - """A live-forwarded leading `ping` commits nothing: a retriable error after - it still falls back, and the fallback's own lifecycle follows the ping cleanly.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_ping_chunk(), _anthropic_messages_overloaded_error_chunk()] - ) - fallback_message_start = b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_2"}}\n\n' - fallback_stream = _AnthropicMessagesFallbackByteStream( - [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")] - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ): - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [ - _anthropic_messages_ping_chunk(), - fallback_message_start, - _anthropic_messages_content_chunk("fallback answer"), - ] - - -@pytest.mark.asyncio -async def test_anthropic_messages_hold_back_retrieval_failure_reaches_client_without_fallback(): - """The hold-back iterator's own retrieval-failure frame is the gateway's verdict, not a - provider failure: a configured fallback stays untouched and the client reads the error - right after the live keepalive.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_ping_chunk(), SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] - ) - fallback = AsyncMock( - return_value=_AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - ) - - with patch.object(router, "async_function_with_fallbacks_common_utils", new=fallback): - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - fallback.assert_not_called() - assert collected == [_anthropic_messages_ping_chunk(), SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES] - - -@pytest.mark.asyncio -async def test_anthropic_messages_pre_content_buffer_cap_forces_commit(): - """Bugbot regression: a hostile or pathological upstream that never emits - real content or an error must not grow the pre-content lifecycle buffer - without bound - hitting MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS commits - to the primary stream early, exactly as real content arriving would.""" - router = _anthropic_messages_make_router() - lifecycle_chunk = _anthropic_messages_message_start_chunk() - error_chunk = _anthropic_messages_overloaded_error_chunk() - chunks = [lifecycle_chunk] * (MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + 5) + [error_chunk] - source = _AnthropicMessagesFakeByteStream(chunks) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - collected = [chunk async for chunk in wrapped] - - mock_fallback.assert_not_awaited() - assert collected.count(lifecycle_chunk) == MAX_BUFFERED_PRE_CONTENT_ANTHROPIC_CHUNKS + 5 - assert collected[-1] == error_chunk - - -@pytest.mark.asyncio -async def test_anthropic_messages_ping_coalesced_with_content_in_one_physical_chunk_is_forwarded(): - """Greptile/Bugbot regression: transport-level buffering can coalesce a - `ping` keepalive and a real content_block_delta into ONE physical read. - The pre-content ping-drop must only discard PURE ping frames - dropping - the whole coalesced chunk would silently lose generated content.""" - router = _anthropic_messages_make_router() - coalesced_chunk = _anthropic_messages_ping_chunk() + _anthropic_messages_content_chunk("hi") - source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) - collected = [chunk async for chunk in wrapped] - - assert collected == [coalesced_chunk] - - -@pytest.mark.asyncio -async def test_anthropic_messages_ping_coalesced_with_retriable_error_still_falls_back(): - """Greptile/Bugbot regression: a physical chunk coalescing a `ping` with a - retriable `event: error` must not be discarded as a keepalive - the error - inside it must still trigger the mid-stream fallback.""" - router = _anthropic_messages_make_router() - coalesced_chunk = _anthropic_messages_ping_chunk() + _anthropic_messages_overloaded_error_chunk() - source = _AnthropicMessagesFakeByteStream([coalesced_chunk]) - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, initial_kwargs={"model": "primary"} - ) - collected = [chunk async for chunk in wrapped] - - mock_fallback.assert_awaited_once() - assert collected == [_anthropic_messages_content_chunk("fallback answer")] - - # -------- _aanthropic_messages_fallback_attempt -------- -@pytest.mark.asyncio -async def test_aanthropic_messages_fallback_attempt_yields_fallback_stream(): - """Direct-call regression: the fallback-attempt helper re-enters the - Router's fallback chain and forwards whatever the fallback produces.""" - router = _anthropic_messages_make_router() - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - collected = [ - chunk - async for chunk in router._aanthropic_messages_fallback_attempt( - error, - {"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, - _anthropic_messages_make_wrapper(), - ) - ] - - assert collected == [_anthropic_messages_content_chunk("fallback answer")] - mock_fallback.assert_awaited_once() - assert mock_fallback.await_args.kwargs["e"] is error - - -@pytest.mark.asyncio -async def test_aanthropic_messages_fallback_attempt_raises_original_exception_on_double_failure(): - """Direct-call regression: when the fallback attempt itself fails with a - MidStreamFallbackError wrapping a real provider exception, that real - exception must surface rather than the internal wrapper exception.""" - router = _anthropic_messages_make_router() - error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") - original_exception = litellm.APIError( - status_code=503, message="fallback also overloaded", llm_provider="bedrock", model="fallback" - ) - fallback_failure = MidStreamFallbackError( - message="fallback failed", model="fallback", llm_provider="bedrock", original_exception=original_exception - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=fallback_failure), - ): - with pytest.raises(litellm.APIError) as exc_info: - async for _ in router._aanthropic_messages_fallback_attempt( - error, {"model": "primary"}, _anthropic_messages_make_wrapper() - ): - pass - - assert exc_info.value is original_exception - - -@pytest.mark.asyncio -async def test_aanthropic_messages_fallback_attempt_yields_non_streaming_fallback_response(): - """Bugbot regression: a fallback that resolves to a non-streaming - response (no __aiter__, e.g. an agentic tool-use interception loop) must - be synthesized into a valid SSE byte sequence, not yielded as a raw dict - into a byte stream - the generator is typed AsyncGenerator[bytes, None] - and every item reaching the client must be a real SSE frame.""" - router = _anthropic_messages_make_router() - error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") - non_streaming_response = {"id": "msg_1", "type": "message", "content": [{"type": "text", "text": "hi"}]} - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=non_streaming_response), - ): - collected = [ - item - async for item in router._aanthropic_messages_fallback_attempt( - error, {"model": "primary"}, _anthropic_messages_make_wrapper() - ) - ] - - assert all(isinstance(item, bytes) for item in collected) - event_types = [item.split(b"\n")[0].removeprefix(b"event: ") for item in collected] - assert event_types == [ - b"message_start", - b"content_block_start", - b"content_block_delta", - b"content_block_stop", - b"message_delta", - b"message_stop", - ] - assert b'"text": "hi"' in collected[2] - - -@pytest.mark.asyncio -async def test_aanthropic_messages_fallback_attempt_reraises_plain_exception_on_double_failure(): - """Direct-call regression: when the fallback attempt fails with a plain - exception (not a MidStreamFallbackError), that exception itself must - propagate unchanged.""" - router = _anthropic_messages_make_router() - error = MidStreamFallbackError(message="overloaded", model="primary", llm_provider="anthropic") - fallback_failure = ValueError("no healthy deployments") - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=fallback_failure), - ): - with pytest.raises(ValueError, match="no healthy deployments") as exc_info: - async for _ in router._aanthropic_messages_fallback_attempt( - error, {"model": "primary"}, _anthropic_messages_make_wrapper() - ): - pass - - assert exc_info.value is fallback_failure - - # -------- _aanthropic_messages_with_streaming_fallbacks -------- -@pytest.mark.asyncio -async def test_aanthropic_messages_with_streaming_fallbacks_non_streaming_passthrough(): - """A non-streaming response (plain dict) is returned unchanged, never wrapped.""" - router = _anthropic_messages_make_router() - plain_response = {"id": "msg_1", "type": "message"} - - 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._aanthropic_messages_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=False, - ) - assert out is plain_response - - -@pytest.mark.asyncio -async def test_aanthropic_messages_with_streaming_fallbacks_wraps_streaming_iterator(): - """A streaming response is wrapped via _aanthropic_messages_streaming_iterator.""" - router = _anthropic_messages_make_router() - streaming_iter = _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk()]) - wrapped_marker = object() - - 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, - "_aanthropic_messages_streaming_iterator", - new=AsyncMock(return_value=wrapped_marker), - ) as mock_wrap, - ): - out = await router._aanthropic_messages_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=True, - ) - assert out is wrapped_marker - mock_wrap.assert_awaited_once() - - # -------- mid-stream error handling -------- -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_on_pre_first_chunk_error_event(): - """Regression for #24004: a retriable SSE `event: error` frame - (overloaded_error/internal_server_error) that arrives before any real - content must trigger the router's fallback chain instead of passing - through to the client silently. The frame carries the error a 529 answer maps to, an InternalServerError, - so a failed fallback answers the status every other litellm path gives an overload.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_overloaded_error_chunk()]) - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [_anthropic_messages_content_chunk("fallback answer")] - mock_fallback.assert_awaited_once() - raised = mock_fallback.await_args.kwargs["e"] - assert isinstance(raised, MidStreamFallbackError) - assert isinstance(raised.original_exception, litellm.InternalServerError) - assert raised.status_code == 500 - assert raised.is_pre_first_chunk is True - assert source.closed is True - - -@pytest.mark.asyncio -async def test_anthropic_messages_mid_stream_error_preserves_real_status_code(): - """Bugbot regression: the MidStreamFallbackError raised for a detected SSE - `event: error` frame must carry the error's REAL parsed status code - (via original_exception), not silently default to 503 for every error - type - a rate_limit_error (429) must surface as 429, not 503.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_rate_limit_error_chunk()]) - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary", "messages": [{"role": "user", "content": "hi"}]}, - ) - [chunk async for chunk in wrapped] - - raised = mock_fallback.await_args.kwargs["e"] - assert isinstance(raised, MidStreamFallbackError) - assert raised.status_code == 429 - assert raised.original_exception is not None - assert raised.original_exception.status_code == 429 - assert raised.original_exception.llm_provider == "anthropic" - - -def test_merge_fallback_hidden_params_direct_call(): - """Direct-call regression: merge_fallback_hidden_params combines the - fallback's hidden params/headers with whatever was already present, - with the fallback's values winning on key collisions.""" - wrapper = FallbackAwareAnthropicMessagesStream( - _anthropic_messages_empty_generator(), - _AnthropicMessagesFakeByteStream([]), # carries {"additional_headers": {"x-amzn-requestid": "req-1"}} - ) - wrapper.merge_fallback_hidden_params( - {"model_id": "fallback-deployment"}, - {"x-amzn-requestid": "req-2", "x-fallback-only": "yes"}, - ) - assert wrapper._hidden_params["model_id"] == "fallback-deployment" - assert wrapper._hidden_params["additional_headers"] == { - "x-amzn-requestid": "req-2", - "x-fallback-only": "yes", - } - - def test_anthropic_stream_forwards_ping_live_direct_call(): ping = _anthropic_messages_ping_chunk() content = _anthropic_messages_content_chunk("hi") @@ -19031,21 +10474,6 @@ def test_anthropic_stream_error_is_gateway_verdict_direct_call(): assert _anthropic_stream_error_is_gateway_verdict(_anthropic_messages_ping_chunk()) is False -def test_fallback_aware_stream_reports_withheld_output_of_its_current_source(): - """The proxy's cancel-refund guard reads this flag off the router wrapper, so it - must reflect the stream actually being drained: the primary, then the fallback.""" - - class _HoldingBack: - _hidden_params = {"additional_headers": {}} - has_buffered_provider_output = True - - wrapper = FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), _HoldingBack()) - assert wrapper.has_buffered_provider_output is True - - wrapper.adopt_fallback_source(_AnthropicMessagesFakeByteStream([])) - assert wrapper.has_buffered_provider_output is False - - def test_is_retriable_anthropic_status_direct_call(): assert _is_retriable_anthropic_status(429) is True assert _is_retriable_anthropic_status(503) is True @@ -19089,1200 +10517,6 @@ def test_anthropic_stream_commits_now_direct_call(): ) -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_merges_fallback_hidden_params(): - """Bugbot regression: after a successful mid-stream fallback, the - wrapper's _hidden_params must reflect the FALLBACK deployment's own - provider headers (e.g. a different Bedrock request-id), not stay - frozen on the primary's - raw bytes can't carry per-item _hidden_params - the way a ModelResponseStream/ResponsesAPI event can, so the wrapper - itself is the only place left to expose them.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_overloaded_error_chunk()] - ) # carries x-amzn-requestid: req-1 - fallback_stream = _AnthropicMessagesFallbackByteStream( - [_anthropic_messages_content_chunk("fallback answer")], - hidden_params={"additional_headers": {"x-amzn-requestid": "req-2", "x-fallback-only": "yes"}}, - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ): - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - _ = [chunk async for chunk in wrapped] - - headers = wrapped._hidden_params["additional_headers"] - assert headers["x-amzn-requestid"] == "req-2" - assert headers["x-fallback-only"] == "yes" - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "fallback_headers", - [ - {"x-fallback-only": "yes"}, - { - "x-fallback-only": "yes", - "x-litellm-complexity-router-tier": "SIMPLE", - }, - ], - ids=["plain-fallback", "complexity-tier-fallback"], -) -@_PYTHON_CORE_ONLY -async def test_anthropic_messages_fallback_replaces_complexity_headers_before_output( - fallback_headers: dict[str, str], -) -> None: - primary_headers: Final = { - "x-litellm-complexity-router-tier": "REASONING", - "x-litellm-complexity-router-reasoning-effort": "xhigh", - } - source: Final = _AnthropicMessagesFallbackByteStream( - [_anthropic_messages_overloaded_error_chunk()], - hidden_params={"additional_headers": primary_headers}, - ) - fallback_output: Final = _anthropic_messages_content_chunk("fallback answer") - fallback: Final = _AnthropicMessagesFallbackByteStream( - [fallback_output], - hidden_params={ - "model_id": "fallback-deployment", - "additional_headers": fallback_headers, - }, - ) - router: Final = _injected_fallback_router(fallback) - - wrapped: Final = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - assert wrapped._hidden_params["additional_headers"] == primary_headers - first_output: Final = await wrapped.__anext__() - - assert first_output == fallback_output - assert wrapped.fallback_headers_adopted is True - assert wrapped._hidden_params == { - "model_id": "fallback-deployment", - "additional_headers": fallback_headers, - } - - -@pytest.mark.asyncio -async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_nested_metadata(): - """Bugbot regression: a shallow .copy() of kwargs still shares the - nested litellm_metadata/metadata dict objects with the primary attempt. - _update_kwargs_with_deployment mutates that dict in place with - deployment-specific fields, which must not leak into the fallback - request's metadata.""" - router = _anthropic_messages_make_router() - primary_metadata = {"model_group": "primary"} - streaming_iter_kwargs = {} - - async def fake_original(**_kwargs): - # Simulate _update_kwargs_with_deployment mutating the primary's - # litellm_metadata in place, as the real helper does. - primary_metadata["deployment"] = "primary-deployment-object" - return _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("hi")]) - - with patch.object( - router, - "_aanthropic_messages_streaming_iterator", - new=AsyncMock(side_effect=lambda **kwargs: streaming_iter_kwargs.update(kwargs) or "wrapped"), - ): - with patch.object( - router, - "_ageneric_api_call_with_fallbacks_helper", - new=AsyncMock(side_effect=fake_original), - ): - await router._aanthropic_messages_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=True, - litellm_metadata=primary_metadata, - ) - - fallback_kwargs = streaming_iter_kwargs["initial_kwargs"] - assert fallback_kwargs["litellm_metadata"] is not primary_metadata - assert "deployment" not in fallback_kwargs["litellm_metadata"] - - -@pytest.mark.asyncio -async def test_aanthropic_messages_with_streaming_fallbacks_deep_copies_metadata_field(): - """Same regression as above for the (separate) `metadata` kwarg some - call sites use instead of `litellm_metadata`.""" - router = _anthropic_messages_make_router() - primary_metadata = {"tag": "primary"} - streaming_iter_kwargs = {} - - async def fake_original(**_kwargs): - primary_metadata["deployment"] = "primary-deployment-object" - return _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("hi")]) - - with patch.object( - router, - "_aanthropic_messages_streaming_iterator", - new=AsyncMock(side_effect=lambda **kwargs: streaming_iter_kwargs.update(kwargs) or "wrapped"), - ): - with patch.object( - router, - "_ageneric_api_call_with_fallbacks_helper", - new=AsyncMock(side_effect=fake_original), - ): - await router._aanthropic_messages_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=True, - metadata=primary_metadata, - ) - - fallback_kwargs = streaming_iter_kwargs["initial_kwargs"] - assert fallback_kwargs["metadata"] is not primary_metadata - assert "deployment" not in fallback_kwargs["metadata"] - - -@pytest.mark.asyncio -async def test_anthropic_messages_hop_stream_failure_reaches_second_fallback_entry(): - """Regression: fallbacks=[{"primary": ["fb1", "fb2"]}]. The primary fails before - streaming, fb1 is reached through the regular fallback chain and then sends an - error frame mid-stream. Only the primary's stream used to be wrapped, so the outer - wrapper re-tried fb1 with a fresh attempted set and forwarded fb1's error frame to - the client on an HTTP 200; fb2 was unreachable.""" - router = Router( - model_list=[ - {"model_name": "primary", "litellm_params": {"model": "anthropic/primary-model", "api_key": "sk-test"}}, - {"model_name": "fb1", "litellm_params": {"model": "anthropic/fb1-model", "api_key": "sk-test"}}, - {"model_name": "fb2", "litellm_params": {"model": "anthropic/fb2-model", "api_key": "sk-test"}}, - ], - num_retries=0, - fallbacks=[{"primary": ["fb1", "fb2"]}], - ) - calls: list = [] - - async def fake_original(**kwargs): - model = kwargs["model"] - calls.append(model) - if model == "anthropic/primary-model": - raise litellm.InternalServerError(message="primary down", llm_provider="anthropic", model=model) - if model == "anthropic/fb1-model": - return _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - ) - return _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb2")] - ) - - stream = await router._aanthropic_messages_with_streaming_fallbacks( - original_function=fake_original, - model="primary", - stream=True, - messages=[{"role": "user", "content": "hi"}], - max_tokens=10, - ) - body = b"".join([chunk async for chunk in stream]) - - assert calls == ["anthropic/primary-model", "anthropic/fb1-model", "anthropic/fb2-model"] - assert b"from fb2" in body - assert b"overloaded_error" not in body - - -_ANTHROPIC_MESSAGES_RETRY_GROUP: Final = ("anthropic/glm-a", "anthropic/glm-b") - - -def _anthropic_messages_retry_router( - num_retries: int, - deployment_params: Mapping[str, object] | None = None, - fallbacks: list[dict[str, list[str]]] | None = None, - context_window_fallbacks: list[dict[str, list[str]]] | None = None, - retry_policy: RetryPolicy | None = None, -) -> Router: - """Two deployments in the group, so a same-group retry waits for no backoff; no fallbacks unless asked.""" - group_deployments = [ - {"model_name": "glm", "litellm_params": {"model": model, "api_key": "sk-test", **(deployment_params or {})}} - for model in _ANTHROPIC_MESSAGES_RETRY_GROUP - ] - return Router( - model_list=[ - *group_deployments, - {"model_name": "fb", "litellm_params": {"model": "anthropic/fb-model", "api_key": "sk-test"}}, - {"model_name": "cw", "litellm_params": {"model": "anthropic/cw-model", "api_key": "sk-test"}}, - ], - num_retries=num_retries, - fallbacks=fallbacks or [], - context_window_fallbacks=context_window_fallbacks or [], - retry_policy=retry_policy, - ) - - -class _AnthropicMessagesScriptedProvider: - """Stands in for litellm.anthropic_messages: answers each call with the next scripted stream and records - the deployment it was routed to plus the retry counters the router stamped for that attempt.""" - - def __init__(self, *streams) -> None: - self._streams = list(streams) - self.calls: list[tuple[str, object, object]] = [] - - async def __call__(self, **kwargs): - litellm_metadata = kwargs.get("litellm_metadata") or {} - self.calls.append((kwargs["model"], litellm_metadata.get("attempted_retries"), litellm_metadata.get("max_retries"))) - assert self._streams, "provider called more times than scripted" - return self._streams.pop(0)() - - -def _anthropic_messages_transport_drop(original_exception: Exception | None = None) -> MidStreamFallbackError: - """What the completion bridge raises when the upstream closes the connection before any content.""" - return MidStreamFallbackError( - message="Connection closed", - model="glm", - llm_provider="databricks", - original_exception=original_exception - or litellm.APIConnectionError(message="Connection closed", llm_provider="databricks", model="glm"), - is_pre_first_chunk=True, - ) - - -def _anthropic_messages_dropped_before_content(): - return _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], _anthropic_messages_transport_drop()) - - -def _anthropic_messages_bridge_error_chunk() -> bytes: - from litellm.anthropic_interface.exceptions.exception_mapping_utils import anthropic_error_sse_frame - - return anthropic_error_sse_frame(status_code=500, raw_message="Connection closed").encode() - - -def _anthropic_messages_retried_stream(): - return _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - ) - - -async def _anthropic_messages_drain_into(stream, received: list) -> None: - async for chunk in stream: - received.append(chunk) - - -async def _anthropic_messages_stream_through_router(router: Router, provider, **request_kwargs): - return await router._aanthropic_messages_with_streaming_fallbacks( - original_function=provider, - model="glm", - stream=True, - messages=[{"role": "user", "content": "ping"}], - max_tokens=16, - **request_kwargs, - ) - - -_ANTHROPIC_MESSAGES_PRE_CONTENT_DROPS: Final = ( - pytest.param(_anthropic_messages_dropped_before_content, id="bridge-raises-before-content"), - pytest.param( - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_bridge_error_chunk()] - ), - id="bridge-error-frame", - ), - pytest.param( - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - ), - id="provider-overloaded-frame", - ), -) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("dropped_stream", _ANTHROPIC_MESSAGES_PRE_CONTENT_DROPS) -async def test_anthropic_messages_stream_dropped_before_content_is_retried_within_the_group(dropped_stream): - """Issue #44238: a /v1/messages stream the provider dropped before any content was answered after a - single upstream attempt, num_retries never applied. The drop is retried within the model group, with - the retry counters continuing the request's count, and the client sees one message lifecycle.""" - router = _anthropic_messages_retry_router(num_retries=2) - provider = _AnthropicMessagesScriptedProvider(dropped_stream, _anthropic_messages_retried_stream) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert all(model in _ANTHROPIC_MESSAGES_RETRY_GROUP for model, _, _ in provider.calls) - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)] - - -@pytest.mark.asyncio -async def test_anthropic_messages_same_group_retry_skips_an_embedding_only_deployment(): - """The retry of a /v1/messages stream dropped before content must not land on an enforced embedding-only - deployment of the same group, even when that deployment is the group's preferred order.""" - router: Final = Router( - model_list=[ - {"model_name": "glm", "litellm_params": {"model": "anthropic/glm-chat", "api_key": "sk-test", "order": 2}}, - { - "model_name": "glm", - "litellm_params": {"model": "anthropic/glm-embed", "api_key": "sk-test", "order": 1}, - "model_info": {"enforce_supported_endpoints": True, "supported_endpoints": ["/v1/embeddings"]}, - }, - ], - num_retries=1, - ) - provider: Final = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, _anthropic_messages_retried_stream - ) - - async def anthropic_messages(**kwargs: object) -> object: - return await provider(**kwargs) - - stream: Final = await _anthropic_messages_stream_through_router(router, anthropic_messages) - body: Final = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [model for model, _, _ in provider.calls] == ["anthropic/glm-chat", "anthropic/glm-chat"] - - -class _AnthropicMessagesRetryStashRecorder: - def __init__(self) -> None: - self.observed: tuple[Usage | None, ...] = () - - def record(self, usage: Usage | None) -> None: - self.observed = (*self.observed, usage) - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_opens_without_the_superseded_attempts_partial_usage() -> None: - """The dropped attempt's usage stays available for terminal billing without reaching the retry's success log.""" - router: Final = _anthropic_messages_retry_router(num_retries=1) - logging_obj: Final = LiteLLMLogging( - model="glm", - messages=[{"role": "user", "content": "ping"}], - stream=True, - call_type="anthropic_messages", - start_time=datetime(2026, 1, 1), - litellm_call_id="retry-own-usage", - function_id="retry-own-usage", - ) - partial_usage: Final = Usage(prompt_tokens=52, completion_tokens=1, total_tokens=53) - logging_obj.record_partial_usage_for_failure(partial_usage, 0.01) - stash_recorder: Final = _AnthropicMessagesRetryStashRecorder() - usage_adapter: Final = TypeAdapter(Usage | None) - - def retry_stream() -> _AnthropicMessagesFakeByteStream: - stash_recorder.record( - usage_adapter.validate_python(logging_obj.model_call_details.get("combined_usage_object")) - ) - return _anthropic_messages_retried_stream() - - provider: Final = _AnthropicMessagesScriptedProvider(_anthropic_messages_dropped_before_content, retry_stream) - - stream: Final = await _anthropic_messages_stream_through_router( - router, provider, litellm_logging_obj=logging_obj - ) - body: Final = tuple([chunk async for chunk in stream]) - - assert stash_recorder.observed == (None,) - assert body == (_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")) - assert "response_cost" not in logging_obj.model_call_details - assert logging_obj.model_call_details["failed_attempts_usage"] == ((partial_usage, 0.01),) - - -@pytest.mark.asyncio -async def test_anthropic_messages_superseded_stream_is_closed_before_the_retry_opens(): - """Both attempts log through the request's one logging object, so the superseded attempt's end-of-stream - logging must be queued before the retry starts producing its own: the router closes the dropped stream - before it asks the group for another.""" - router = _anthropic_messages_retry_router(num_retries=1) - superseded = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - ) - superseded_closed_when_retry_opened: list[bool] = [] - - def retry_stream(): - superseded_closed_when_retry_opened.append(superseded.closed_as_superseded) - return _anthropic_messages_retried_stream() - - provider = _AnthropicMessagesScriptedProvider(lambda: superseded, retry_stream) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert superseded_closed_when_retry_opened == [True] - assert superseded.close_contexts[0] is True - - -@pytest.mark.asyncio -async def test_a_ptu_share_holders_messages_stream_dropped_before_content_is_retried_on_its_shared_group(monkeypatch): - """The same-group retry has to look the group's health up as the calling team; looked up with no - team a shared PTU group looks empty, so the share holder's dropped stream is never retried.""" - monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "true") - router: Final = Router( - model_list=[_shared_ptu_deployment("glm", f"shared-ptu-{i}") for i in (1, 2)], - num_retries=2, - ) - provider: Final = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, _anthropic_messages_retried_stream - ) - - stream: Final = await _anthropic_messages_stream_through_router( - router, provider, litellm_metadata={"user_api_key_team_id": "team-a"} - ) - body: Final = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert len(provider.calls) == 2 - - -@pytest.mark.asyncio -async def test_anthropic_messages_stream_dropped_after_content_keeps_the_error_and_is_not_retried(): - """A drop once content reached the client cannot be retried without a second overlapping message - lifecycle, so it keeps surfacing the provider's error after a single attempt.""" - router = _anthropic_messages_retry_router(num_retries=2) - drop = _anthropic_messages_transport_drop() - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesRaisingByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("par")], drop - ) - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - received = [] - with pytest.raises(litellm.APIConnectionError) as raised: - await _anthropic_messages_drain_into(stream, received) - - assert raised.value is drop.original_exception - assert received == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("par")] - assert len(provider.calls) == 1 - - -@pytest.mark.asyncio -async def test_anthropic_messages_retries_stop_at_num_retries_and_raise_the_last_drop(): - """Every retry's own stream continues the same count, so a group that keeps dropping is tried - exactly 1 + num_retries times before the provider's error reaches the client.""" - router = _anthropic_messages_retry_router(num_retries=2) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_dropped_before_content, - _anthropic_messages_dropped_before_content, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - with pytest.raises(litellm.APIConnectionError): - [chunk async for chunk in stream] - - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2), (2, 2)] - - -@pytest.mark.asyncio -async def test_anthropic_messages_retries_run_out_before_the_fallback_chain_is_consulted(): - """Same-group retries come first; the fallback group is reached only once num_retries is spent.""" - router = _anthropic_messages_retry_router(num_retries=1, fallbacks=[{"glm": ["fb"]}]) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_dropped_before_content, - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - ), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - assert [model in _ANTHROPIC_MESSAGES_RETRY_GROUP for model, _, _ in provider.calls] == [True, True, False] - assert provider.calls[-1][0] == "anthropic/fb-model" - - -def _anthropic_messages_fb_deployment_hidden_params() -> dict: - return {"model_id": "fb-deployment", "additional_headers": {"x-litellm-model-group": "fb"}} - - -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_after_exhausted_retries_attributes_the_response_to_the_fallback_deployment(): - """The retry's stream carries a wrapper of its own, so a fallback it makes before its first byte must reach - the wrapper the proxy reads headers off: the response names the deployment that served it, not the primary.""" - router = _anthropic_messages_retry_router(num_retries=1, fallbacks=[{"glm": ["fb"]}]) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_dropped_before_content, - lambda: _AnthropicMessagesFallbackByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")], - hidden_params=_anthropic_messages_fb_deployment_hidden_params(), - ), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - assert stream._hidden_params["model_id"] == "fb-deployment" - assert stream._hidden_params["additional_headers"]["x-litellm-model-group"] == "fb" - assert stream._hidden_params["additional_headers"]["x-litellm-attempted-fallbacks"] == 1 - - -def test_anthropic_messages_wrapper_follows_the_attribution_of_a_source_that_fell_back(): - inner = FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), object()) - outer = FallbackAwareAnthropicMessagesStream(_anthropic_messages_empty_generator(), inner) - fallback = _AnthropicMessagesFallbackByteStream([], hidden_params=_anthropic_messages_fb_deployment_hidden_params()) - - outer.follow_source_attribution() - assert "model_id" not in outer._hidden_params - - inner.merge_fallback_hidden_params(*Router._prepare_fallback_hidden_params(fallback)) - inner.adopt_fallback_source(fallback) - outer.follow_source_attribution() - assert outer._hidden_params["model_id"] == "fb-deployment" - assert outer._hidden_params["additional_headers"]["x-litellm-model-group"] == "fb" - - -@pytest.mark.asyncio -@pytest.mark.parametrize("drops", [1, 2]) -async def test_anthropic_messages_mid_stream_retries_are_counted_in_the_response_retry_headers(drops: int): - """A retry made after the stream opened never passes through async_function_with_retries, so the wrapper - stamps the retry headers that path would have and the client reads them along with the first byte.""" - router = _anthropic_messages_retry_router(num_retries=2) - provider = _AnthropicMessagesScriptedProvider( - *([_anthropic_messages_dropped_before_content] * drops), _anthropic_messages_retried_stream - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - first = await stream.__anext__() - headers = stream._hidden_params["additional_headers"] - - assert first == _anthropic_messages_message_start_chunk() - assert (headers["x-litellm-attempted-retries"], headers["x-litellm-max-retries"]) == (drops, 2) - assert [chunk async for chunk in stream] == [_anthropic_messages_content_chunk("pong")] - - -def _anthropic_messages_raise_authentication_error(): - raise litellm.AuthenticationError(message="invalid api key", llm_provider="anthropic", model="glm") - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_raising_a_non_retriable_error_is_handed_to_the_fallback_chain(): - """A retry that fails before its stream opens with an error no retry covers ends the retries and reaches - the fallback group the way a pre-stream failure does, instead of surfacing as the client's error.""" - router = _anthropic_messages_retry_router(num_retries=2, fallbacks=[{"glm": ["fb"]}]) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_raise_authentication_error, - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - ), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - assert [model in _ANTHROPIC_MESSAGES_RETRY_GROUP for model, _, _ in provider.calls] == [True, True, False] - - -def _anthropic_messages_raise_timeout(): - raise litellm.Timeout(message="upstream timed out", model="glm", llm_provider="databricks") - - -def _anthropic_messages_raise_internal_server_error(): - raise litellm.InternalServerError(message="upstream reset", llm_provider="databricks", model="glm") - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_raising_a_timeout_is_retried_like_a_pre_stream_timeout(): - """A 408 raised by a retry attempt before its stream opens is retried the way the pre-stream path retries - a 408, instead of ending the retries on the error-frame gate that only knows 429 and 5xx.""" - router = _anthropic_messages_retry_router(num_retries=3) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_raise_timeout, - _anthropic_messages_retried_stream, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [model in _ANTHROPIC_MESSAGES_RETRY_GROUP for model, _, _ in provider.calls] == [True, True, True] - - -@pytest.mark.asyncio -async def test_anthropic_messages_deployment_num_retries_also_governs_a_failure_before_the_stream_opens(): - """The deployment's num_retries litellm_param sets the budget for a failure raised before the stream opened - on this route, as it does for a mid-stream drop and for chat completions.""" - router = _anthropic_messages_retry_router(num_retries=0, deployment_params={"num_retries": 2}) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_raise_internal_server_error, - _anthropic_messages_raise_internal_server_error, - _anthropic_messages_retried_stream, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert len(provider.calls) == 3 - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_raising_a_non_retriable_error_reaches_the_client_without_fallbacks(): - router = _anthropic_messages_retry_router(num_retries=2) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_raise_authentication_error, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - with pytest.raises(litellm.AuthenticationError): - [chunk async for chunk in stream] - - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 2), (1, 2)] - - -def _anthropic_messages_raise_context_window_error(): - raise litellm.ContextWindowExceededError(message="prompt too long", llm_provider="anthropic", model="glm") - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_raising_a_context_window_error_takes_the_context_window_fallback(): - """The fallback chain sees the retry's own error type, so a context window overflow on the retried - deployment reaches context_window_fallbacks rather than the regular fallbacks.""" - router = _anthropic_messages_retry_router( - num_retries=2, fallbacks=[{"glm": ["fb"]}], context_window_fallbacks=[{"glm": ["cw"]}] - ) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - _anthropic_messages_raise_context_window_error, - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from cw")] - ), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from cw")] - assert [model for model, _, _ in provider.calls][-1] == "anthropic/cw-model" - - -def test_anthropic_messages_retry_budget_precedence_direct_call(): - """A retry policy naming the error class outranks the request's num_retries, which outranks the routed - deployment's, which outranks the router's; num_retries=0 on the request turns a policy off too.""" - router = _anthropic_messages_retry_router(num_retries=3, deployment_params={"num_retries": 2}) - deployment_id = router.get_model_list(model_name="glm")[0]["model_info"]["id"] - routed = {"model": "glm", "litellm_metadata": {"model_info": {"id": deployment_id}}} - drop = litellm.APIConnectionError(message="closed", llm_provider="databricks", model="glm") - reset = litellm.InternalServerError(message="reset", llm_provider="databricks", model="glm") - policy_router = _anthropic_messages_retry_router( - num_retries=3, retry_policy=RetryPolicy(InternalServerErrorRetries=4) - ) - - assert router._stream_retry_budget(drop, {"model": "glm"}) == (3, False) - assert router._stream_retry_budget(drop, routed) == (2, False) - assert router._stream_retry_budget(drop, {**routed, "num_retries": 1}) == (1, False) - assert policy_router._stream_retry_budget(reset, {"model": "glm", "num_retries": 1}) == (4, True) - assert policy_router._stream_retry_budget(drop, {"model": "glm", "num_retries": 1}) == (1, False) - assert policy_router._stream_retry_budget(reset, {"model": "glm", "num_retries": 0}) == (0, False) - committed = {**routed, "litellm_metadata": {**routed["litellm_metadata"], "attempted_retries": 1, "max_retries": 5}} - assert router._stream_retry_budget(drop, committed) == (5, False) - assert policy_router._stream_retry_budget(reset, committed) == (5, True) - - -def test_anthropic_messages_stream_can_retry_direct_call(): - router = _anthropic_messages_retry_router(num_retries=1) - policy_router = _anthropic_messages_retry_router( - num_retries=0, retry_policy=RetryPolicy(InternalServerErrorRetries=1) - ) - - assert router._anthropic_messages_stream_can_retry({"model": "glm"}) is True - spent = {"model": "glm", "litellm_metadata": {"attempted_retries": 1}} - assert router._anthropic_messages_stream_can_retry(spent) is False - assert router._anthropic_messages_stream_can_retry({"model": "glm", "num_retries": 0}) is False - assert policy_router._anthropic_messages_stream_can_retry({"model": "glm"}) is True - assert policy_router._anthropic_messages_stream_can_retry(spent) is False - assert policy_router._anthropic_messages_stream_can_retry({"model": "glm", "num_retries": 0}) is False - assert policy_router._anthropic_messages_resolved_retry_policy({"model": "glm"}) is not None - assert policy_router._anthropic_messages_resolved_retry_policy({"model": "glm", "num_retries": 0}) is None - - -def test_retry_policy_ceiling_is_the_largest_budget_any_error_class_is_granted(): - from litellm.router import _retry_policy_ceiling - - assert _retry_policy_ceiling(RetryPolicy(InternalServerErrorRetries=1, RateLimitErrorRetries=3)) == 3 - assert _retry_policy_ceiling(RetryPolicy()) == 0 - - -@pytest.mark.asyncio -async def test_anthropic_messages_last_attempt_under_a_retry_policy_forwards_lifecycle_frames_live(): - """A retry policy bounds the hold the way a plain budget does: once the attempts reach the most retries the - policy grants, the stream is the last one, so its frames reach the client as they arrive and a drop after - them is the provider's error in-band rather than an error raised before any byte.""" - router = _anthropic_messages_retry_router(num_retries=0, retry_policy=RetryPolicy(DefaultRetries=1)) - drop = _anthropic_messages_transport_drop() - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, - lambda: _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], drop), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - received: list = [] - with pytest.raises(litellm.APIConnectionError) as raised: - await _anthropic_messages_drain_into(stream, received) - - assert raised.value is drop.original_exception - assert received == [_anthropic_messages_message_start_chunk()] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 0), (1, 1)] - - -@pytest.mark.asyncio -async def test_anthropic_messages_request_num_retries_zero_opts_out_of_the_mid_stream_retry(): - router = _anthropic_messages_retry_router(num_retries=2) - drop = _anthropic_messages_transport_drop() - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], drop) - ) - - stream = await _anthropic_messages_stream_through_router(router, provider, num_retries=0) - with pytest.raises(litellm.APIConnectionError) as raised: - [chunk async for chunk in stream] - - assert raised.value is drop.original_exception - assert len(provider.calls) == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("configured", [1, "1"], ids=["int", "config-string"]) -async def test_anthropic_messages_deployment_num_retries_sets_the_mid_stream_retry_budget(configured): - """A deployment's own num_retries litellm_param outranks the router's, as it does for a failure - raised before the stream opened.""" - router = _anthropic_messages_retry_router(num_retries=0, deployment_params={"num_retries": configured}) - provider = _AnthropicMessagesScriptedProvider( - _anthropic_messages_dropped_before_content, _anthropic_messages_retried_stream - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 0), (1, 1)] - - -@pytest.mark.asyncio -async def test_anthropic_messages_retry_policy_sets_the_mid_stream_retry_budget_per_error_class(): - router = _anthropic_messages_retry_router(num_retries=0, retry_policy=RetryPolicy(InternalServerErrorRetries=1)) - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesRaisingByteStream( - [_anthropic_messages_message_start_chunk()], - _anthropic_messages_transport_drop( - litellm.InternalServerError(message="upstream reset", llm_provider="databricks", model="glm") - ), - ), - _anthropic_messages_retried_stream, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 0), (1, 1)] - - -def _anthropic_messages_error_frame(error_type: str) -> bytes: - return f"event: error\ndata: {json.dumps({'type': 'error', 'error': {'type': error_type, 'message': error_type}})}\n\n".encode() - - -_ANTHROPIC_MESSAGES_ERROR_FRAME_POLICIES: Final = ( - pytest.param("api_error", RetryPolicy(InternalServerErrorRetries=1), id="api_error-500-internal-server"), - pytest.param("overloaded_error", RetryPolicy(InternalServerErrorRetries=1), id="overloaded-internal-server"), - pytest.param("rate_limit_error", RetryPolicy(RateLimitErrorRetries=1), id="rate-limit-429"), - pytest.param("timeout_error", RetryPolicy(TimeoutErrorRetries=1), id="timeout-504"), -) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("error_type,policy", _ANTHROPIC_MESSAGES_ERROR_FRAME_POLICIES) -async def test_anthropic_messages_error_frame_is_retried_under_the_class_the_pre_stream_mapping_gives_it( - error_type, policy -): - """An `event: error` frame before content carried a generic error, so a policy naming only error classes - granted it no retry while the hold still counted the policy: one attempt, then an HTTP error with no - bytes out. The frame now takes the class the pre-stream mapping raises for an answer carrying its body, so - an overloaded frame counts as the InternalServerError a 529 answer is, not a ServiceUnavailableError.""" - router = _anthropic_messages_retry_router(num_retries=0, retry_policy=policy) - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_error_frame(error_type)] - ), - _anthropic_messages_retried_stream, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 0), (1, 1)] - - -@pytest.mark.asyncio -async def test_anthropic_messages_error_frame_of_a_class_granted_no_retry_reaches_the_client_as_sent(): - """A rate limit frame is the RateLimitError a 429 answer is, so a policy granting only InternalServerError - retries leaves it unretried. With no fallback to take over either, the frame reaches the client as the - provider sent it, behind the lifecycle frames held back for a retry that never opened, the way the last - exhausted attempt's frames do; raising it instead turned a provider error frame into an HTTP error only - on the first attempt.""" - router = _anthropic_messages_retry_router(num_retries=0, retry_policy=RetryPolicy(InternalServerErrorRetries=1)) - frame = _anthropic_messages_error_frame("rate_limit_error") - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), frame]) - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), frame] - assert len(provider.calls) == 1 - - -@pytest.mark.asyncio -async def test_anthropic_messages_error_frame_of_a_class_granted_no_retry_still_reaches_a_configured_fallback(): - """The same unretried rate limit frame goes to the fallback group when one is configured, since a - fallback can still take over before any byte reached the client.""" - router = _anthropic_messages_retry_router( - num_retries=0, fallbacks=[{"glm": ["fb"]}], retry_policy=RetryPolicy(InternalServerErrorRetries=1) - ) - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_error_frame("rate_limit_error")] - ), - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - ), - ) - - stream = await _anthropic_messages_stream_through_router(router, provider) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb")] - assert [model in _ANTHROPIC_MESSAGES_RETRY_GROUP for model, _, _ in provider.calls] == [True, False] - assert provider.calls[-1][0] == "anthropic/fb-model" - - -def test_anthropic_messages_recoverable_frame_error_direct_call(): - """Which `event: error` frames are intercepted for a retry or a fallback, and which reach the client as sent.""" - policy_router = _anthropic_messages_retry_router( - num_retries=0, retry_policy=RetryPolicy(InternalServerErrorRetries=1) - ) - fallback_router = _anthropic_messages_retry_router(num_retries=0, fallbacks=[{"glm": ["fb"]}]) - api_error = ("api_error", "reset", 500) - rate_limit = ("rate_limit_error", "slow down", 429) - kwargs = {"model": "glm"} - - recovered = policy_router._anthropic_messages_recoverable_frame_error(api_error, b"", False, "glm", kwargs) - assert isinstance(recovered, litellm.InternalServerError) - assert policy_router._anthropic_messages_recoverable_frame_error(rate_limit, b"", False, "glm", kwargs) is None - assert policy_router._anthropic_messages_recoverable_frame_error(api_error, b"", True, "glm", kwargs) is None - assert policy_router._anthropic_messages_recoverable_frame_error(None, b"", False, "glm", kwargs) is None - assert ( - policy_router._anthropic_messages_recoverable_frame_error( - ("invalid_request_error", "bad", 400), b"", False, "glm", kwargs - ) - is None - ) - spent = {"model": "glm", "litellm_metadata": {"attempted_retries": 1, "max_retries": 1}} - assert policy_router._anthropic_messages_recoverable_frame_error(api_error, b"", False, "glm", spent) is None - assert isinstance( - fallback_router._anthropic_messages_recoverable_frame_error(rate_limit, b"", False, "glm", kwargs), - litellm.RateLimitError, - ) - - -_ANTHROPIC_MESSAGES_MALFORMED_POLICIES: Final = ( - pytest.param({"glm": {"RateLimitErrorRetries": "many"}}, id="string-budget"), - pytest.param({"glm": 5}, id="group-policy-is-an-int"), - pytest.param(5, id="policy-map-is-an-int"), - pytest.param({"glm": [1]}, id="group-policy-is-a-list"), -) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("policy", _ANTHROPIC_MESSAGES_MALFORMED_POLICIES) -@pytest.mark.parametrize("where", ["request", "router"]) -async def test_anthropic_messages_malformed_retry_policy_leaves_a_healthy_stream_alone(policy, where): - """The hold decision resolves the group's retry policy before the first byte, so a policy that does not - parse used to fail every stream of that group with a 500 before any attempt. It now governs nothing.""" - router = _anthropic_messages_retry_router(num_retries=0) - request_kwargs = {"model_group_retry_policy": policy} if where == "request" else {} - if where == "router": - router.model_group_retry_policy = policy - provider = _AnthropicMessagesScriptedProvider(_anthropic_messages_retried_stream) - - stream = await _anthropic_messages_stream_through_router(router, provider, **request_kwargs) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert len(provider.calls) == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("policy", _ANTHROPIC_MESSAGES_MALFORMED_POLICIES) -async def test_anthropic_messages_malformed_retry_policy_falls_back_to_the_plain_budget(policy): - router = _anthropic_messages_retry_router(num_retries=1) - provider = _AnthropicMessagesScriptedProvider( - lambda: _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_error_frame("overloaded_error")] - ), - _anthropic_messages_retried_stream, - ) - - stream = await _anthropic_messages_stream_through_router(router, provider, model_group_retry_policy=policy) - body = [chunk async for chunk in stream] - - assert body == [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("pong")] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == [(0, 1), (1, 1)] - - -class _AnthropicMessagesAlternatingDeployments(CustomRoutingStrategyBase): - """Routes each attempt to the group's next deployment in turn, so which sibling a retry lands on is known.""" - - def __init__(self, router: Router, model_group: str) -> None: - self._deployments = itertools.cycle(router.get_model_list(model_name=model_group) or ()) - - async def async_get_available_deployment( - self, model, messages=None, input=None, specific_deployment=False, request_kwargs=None - ): - return next(self._deployments) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - "first_num_retries,sibling_num_retries,expected_counters", - [ - pytest.param(3, 1, [(0, 0), (1, 3), (2, 3), (3, 3)], id="sibling-grants-fewer"), - pytest.param(1, 3, [(0, 0), (1, 1)], id="sibling-grants-more"), - ], -) -async def test_anthropic_messages_retries_keep_the_budget_the_first_drop_committed_to_across_deployments( - first_num_retries, sibling_num_retries, expected_counters -): - """A retry's stream recomputed its budget from the sibling deployment it landed on, so a group whose - deployments grant different num_retries stopped early or overshot the budget the first drop stamped - into the retry headers; later attempts now keep that budget, as the pre-stream retry loop does.""" - router = Router( - model_list=[ - { - "model_name": "glm", - "litellm_params": {"model": "anthropic/glm-a", "api_key": "sk-test", "num_retries": first_num_retries}, - }, - { - "model_name": "glm", - "litellm_params": {"model": "anthropic/glm-b", "api_key": "sk-test", "num_retries": sibling_num_retries}, - }, - ], - num_retries=0, - fallbacks=None, - ) - router.set_custom_routing_strategy(_AnthropicMessagesAlternatingDeployments(router, "glm")) - provider = _AnthropicMessagesScriptedProvider(*[_anthropic_messages_dropped_before_content] * len(expected_counters)) - - stream = await _anthropic_messages_stream_through_router(router, provider) - with pytest.raises(litellm.APIConnectionError): - [chunk async for chunk in stream] - - assert [model for model, _, _ in provider.calls] == (["anthropic/glm-a", "anthropic/glm-b"] * 2)[: len(expected_counters)] - assert [(attempted, budget) for _, attempted, budget in provider.calls] == expected_counters - - -@pytest.mark.asyncio -async def test_anthropic_messages_lifecycle_frames_wait_for_content_while_a_retry_remains(): - """A retry can only restart cleanly while nothing reached the client, so with retries left a - fallback-less group holds message_start back until the first content frame, as a fallback does.""" - router = _anthropic_messages_retry_router(num_retries=2) - content_released = asyncio.Event() - - async def held_stream(): - yield _anthropic_messages_message_start_chunk() - await content_released.wait() - yield _anthropic_messages_content_chunk("hi") - - provider = _AnthropicMessagesScriptedProvider(held_stream) - stream = await _anthropic_messages_stream_through_router(router, provider) - - pending = asyncio.ensure_future(stream.__anext__()) - await asyncio.sleep(0.2) - assert not pending.done() - content_released.set() - assert await asyncio.wait_for(pending, timeout=1) == _anthropic_messages_message_start_chunk() - assert [chunk async for chunk in stream] == [_anthropic_messages_content_chunk("hi")] - - -@pytest.mark.asyncio -async def test_anthropic_messages_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 never reaches the provider call.""" - from types import MappingProxyType - - from litellm.router_utils.fallback_event_handlers import ( - MID_STREAM_FALLBACK_CONTROLS_KEY, - MidStreamFallbackControls, - ) - - router = Router( - model_list=[ - {"model_name": "fb1", "litellm_params": {"model": "anthropic/fb1-model", "api_key": "sk-test"}}, - ], - num_retries=0, - ) - hop_stream = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_content_chunk("from fb1")] - ) - 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_anthropic_messages_attempt( - model="fb1", - original_generic_function=fake_original, - stream=True, - messages=[{"role": "user", "content": "hi"}], - max_tokens=10, - **{MID_STREAM_FALLBACK_CONTROLS_KEY: controls}, - ) - body = b"".join([chunk async for chunk in stream]) - - assert seen["model"] == "anthropic/fb1-model" - assert MID_STREAM_FALLBACK_CONTROLS_KEY not in seen - assert "fallbacks" not in seen - assert stream is not hop_stream - assert b"from fb1" in body - - -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_triggers_after_lifecycle_only_frame(): - """Regression: Anthropic routinely sends a message_start lifecycle frame - before an overload error even fires. A lifecycle-only frame (no real - content) must not disqualify the fallback attempt, and must not reach - the client either - forwarding it and then appending the fallback's own - message_start would produce two overlapping message lifecycles on one - SSE stream. The primary's buffered lifecycle frame is discarded and the - client sees only the fallback's own, single, clean lifecycle.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream( - [_anthropic_messages_message_start_chunk(), _anthropic_messages_overloaded_error_chunk()] - ) - fallback_message_start = b'event: message_start\ndata: {"type": "message_start", "message": {"id": "msg_2"}}\n\n' - fallback_stream = _AnthropicMessagesFallbackByteStream( - [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")] - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [fallback_message_start, _anthropic_messages_content_chunk("fallback answer")] - assert collected.count(_anthropic_messages_message_start_chunk()) == 0, ( - "the primary's message_start must never reach the client" - ) - assert sum(1 for c in collected if c.startswith(b"event: message_start")) == 1, ( - "exactly one message_start must reach the client" - ) - mock_fallback.assert_awaited_once() - raised = mock_fallback.await_args.kwargs["e"] - assert raised.is_pre_first_chunk is True - - -@pytest.mark.asyncio -async def test_anthropic_messages_raised_error_after_real_content_does_not_restart_stream(): - """Regression: a MidStreamFallbackError raised directly by the source - iterator (the completion-bridge path's CustomStreamWrapper, e.g. a - transport drop) must not trigger a fallback once real content already - reached the client - that would append a second, overlapping message - lifecycle onto the same SSE stream. The original exception must - propagate to the caller instead.""" - router = _anthropic_messages_make_router() - content = _anthropic_messages_content_chunk("partial answer") - original_exception = litellm.APIError( - status_code=503, - message="stream reset", - llm_provider="vertex_ai", - model="primary", - ) - raised_error = MidStreamFallbackError( - message="stream reset", - model="primary", - llm_provider="vertex_ai", - original_exception=original_exception, - is_pre_first_chunk=False, - ) - source = _AnthropicMessagesRaisingByteStream([content], raised_error) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [] - - async def _consume(): - async for chunk in wrapped: - collected.append(chunk) - - with pytest.raises(litellm.APIError) as exc_info: - await _consume() - - assert collected == [content] - assert exc_info.value is original_exception - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_also_catches_raised_midstream_error(): - """Regression for the completion-bridge path (deployments with no native - /v1/messages endpoint): its CustomStreamWrapper raises - MidStreamFallbackError directly (e.g. on a Vertex AI transport drop) - instead of yielding an SSE error chunk - the wrapper must catch that too.""" - router = _anthropic_messages_make_router() - raised_error = MidStreamFallbackError( - message="stream reset", - model="primary", - llm_provider="vertex_ai", - is_pre_first_chunk=True, - ) - source = _AnthropicMessagesRaisingByteStream([], raised_error) - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [_anthropic_messages_content_chunk("fallback answer")] - mock_fallback.assert_awaited_once() - assert mock_fallback.await_args.kwargs["e"] is raised_error - - -_MID_STREAM_OPT_OUT_SHAPES: Final = ( - pytest.param({"disable_fallbacks": True}, id="raw-kwarg"), - pytest.param({"metadata": {DISABLE_FALLBACKS_METADATA_KEY: True}}, id="metadata-stamp"), - pytest.param({"litellm_metadata": {DISABLE_FALLBACKS_METADATA_KEY: True}}, id="litellm_metadata-stamp"), -) - - def _mid_stream_opt_out_router() -> Router: return Router( model_list=[ @@ -20328,102 +10562,6 @@ class _MidStreamOptOutChatStream(CustomStreamWrapper): raise self._error -@pytest.mark.asyncio -@pytest.mark.parametrize("opt_out", _MID_STREAM_OPT_OUT_SHAPES) -async def test_acompletion_streaming_iterator_honors_disable_fallbacks(opt_out): - """A chat stream that fails before its first chunk on a request that opted out of fallbacks - surfaces the primary's own error and never tries the fallback deployment.""" - router = _mid_stream_opt_out_router() - primary_error = _mid_stream_opt_out_primary_error() - source = _MidStreamOptOutChatStream(_mid_stream_opt_out_trigger(primary_error)) - - with patch.object(router, "_acompletion", new=AsyncMock(return_value=_AsyncList([]))) as fallback_attempt: - wrapped = await router._acompletion_streaming_iterator( - model_response=source, - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "primary", "stream": True, **copy.deepcopy(opt_out)}, - ) - with pytest.raises(litellm.InternalServerError) as raised: - [chunk async for chunk in wrapped] - - assert raised.value is primary_error - fallback_attempt.assert_not_awaited() - - -@pytest.mark.parametrize("opt_out", _MID_STREAM_OPT_OUT_SHAPES) -def test_completion_streaming_iterator_honors_disable_fallbacks(opt_out): - """Sync counterpart of test_acompletion_streaming_iterator_honors_disable_fallbacks.""" - router = _mid_stream_opt_out_router() - primary_error = _mid_stream_opt_out_primary_error() - source = _MidStreamOptOutChatStream(_mid_stream_opt_out_trigger(primary_error)) - - with patch.object(router, "_completion", new=MagicMock(return_value=iter([]))) as fallback_attempt: - wrapped = router._completion_streaming_iterator( - model_response=source, - messages=[{"role": "user", "content": "Hi"}], - initial_kwargs={"model": "primary", "stream": True, **copy.deepcopy(opt_out)}, - ) - with pytest.raises(litellm.InternalServerError) as raised: - list(wrapped) - - assert raised.value is primary_error - fallback_attempt.assert_not_called() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("opt_out", _MID_STREAM_OPT_OUT_SHAPES) -async def test_aresponses_streaming_iterator_honors_disable_fallbacks(opt_out): - """Same opt-out contract on the Responses API mid-stream fallback path.""" - router = _mid_stream_opt_out_router() - primary_error = _mid_stream_opt_out_primary_error() - source = _make_responses_iterator(error=_mid_stream_opt_out_trigger(primary_error), model="primary") - - with patch.object( - router, - "_ageneric_api_call_with_fallbacks_responses_attempt", - new=AsyncMock(return_value=_AsyncList([])), - ) as fallback_attempt: - wrapped = await router._aresponses_streaming_iterator( - response=source, - initial_kwargs={ - "model": "primary", - "stream": True, - "input": "Hi", - "original_generic_function": litellm.aresponses, - **copy.deepcopy(opt_out), - }, - ) - with pytest.raises(litellm.InternalServerError) as raised: - [chunk async for chunk in wrapped] - - assert raised.value is primary_error - fallback_attempt.assert_not_awaited() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("opt_out", _MID_STREAM_OPT_OUT_SHAPES) -async def test_anthropic_messages_streaming_iterator_honors_disable_fallbacks(opt_out): - """Same opt-out contract on the Anthropic Messages mid-stream fallback path.""" - router = _mid_stream_opt_out_router() - primary_error = _mid_stream_opt_out_primary_error() - source = _AnthropicMessagesRaisingByteStream([], _mid_stream_opt_out_trigger(primary_error)) - - with patch.object( - router, - "_ageneric_api_call_with_fallbacks_anthropic_messages_attempt", - new=AsyncMock(return_value=_AnthropicMessagesFallbackByteStream([])), - ) as fallback_attempt: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary", "stream": True, **copy.deepcopy(opt_out)}, - ) - with pytest.raises(litellm.InternalServerError) as raised: - [chunk async for chunk in wrapped] - - assert raised.value is primary_error - fallback_attempt.assert_not_awaited() - - @pytest.mark.asyncio async def test_acompletion_disable_fallbacks_reaches_the_mid_stream_hop(): """`disable_fallbacks=True` sent to the public entrypoint survives the fallback wrapper's handoff @@ -20462,43 +10600,6 @@ def test_completion_disable_fallbacks_reaches_the_mid_stream_hop(): assert [call.kwargs["metadata"]["model_group"] for call in provider_calls.call_args_list] == ["primary"] -@pytest.mark.asyncio -@pytest.mark.parametrize( - "raised_error", - [ - BedrockError(status_code=503, message='serviceUnavailableException {"message": "Service unavailable"}'), - BedrockError(status_code=500, message='internalServerException {"message": "Internal error"}'), - BedrockError(status_code=429, message='throttlingException {"message": "Too many requests"}'), - httpx.ReadError("connection reset by upstream"), - ], - ids=["503", "500", "429", "transport-drop"], -) -async def test_anthropic_messages_raised_provider_error_before_content_triggers_fallback(raised_error): - """A retriable raise before content falls over exactly like a detected SSE error event.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], raised_error) - fallback_stream = _AnthropicMessagesFallbackByteStream([_anthropic_messages_content_chunk("fallback answer")]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(return_value=fallback_stream), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [_anthropic_messages_content_chunk("fallback answer")] - mock_fallback.assert_awaited_once() - converted = mock_fallback.await_args.kwargs["e"] - assert isinstance(converted, MidStreamFallbackError) - assert converted.original_exception is raised_error - assert converted.is_pre_first_chunk is True - assert source.closed is True - - class _AnthropicMessagesStringStatusError(Exception): def __init__(self): super().__init__("bad request") @@ -20511,78 +10612,6 @@ class _AnthropicMessagesResponseOnlyStatusError(Exception): self.response = SimpleNamespace(status_code=400) -@pytest.mark.asyncio -@pytest.mark.parametrize( - "raised_error", - [ - BedrockError(status_code=400, message='validationException {"message": "Malformed input"}'), - BedrockError(status_code=424, message='modelStreamErrorException {"message": "Model stream error"}'), - _AnthropicMessagesStringStatusError(), - _AnthropicMessagesResponseOnlyStatusError(), - ], - ids=["400", "424", "str-400", "response-only-400"], -) -async def test_anthropic_messages_raised_non_retriable_provider_error_propagates_unchanged(raised_error): - """A raised client error reaches the caller as the same exception, nothing flushed, no fallback.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], raised_error) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [] - - async def _consume(): - async for chunk in wrapped: - collected.append(chunk) - - with pytest.raises(type(raised_error)) as exc_info: - await _consume() - - assert collected == [] - assert exc_info.value is raised_error - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_raised_provider_error_after_content_propagates_unchanged(): - """A raise after content propagates unchanged even when its status is retriable.""" - router = _anthropic_messages_make_router() - content = _anthropic_messages_content_chunk("partial answer") - raised_error = BedrockError( - status_code=503, message='serviceUnavailableException {"message": "Service unavailable"}' - ) - source = _AnthropicMessagesRaisingByteStream([content], raised_error) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [] - - async def _consume(): - async for chunk in wrapped: - collected.append(chunk) - - with pytest.raises(BedrockError) as exc_info: - await _consume() - - assert collected == [content] - assert exc_info.value is raised_error - mock_fallback.assert_not_awaited() - - @pytest.mark.parametrize( "error, expected_status", [ @@ -20597,340 +10626,9 @@ def test_anthropic_stream_raised_error_status_reads_every_status_shape(error, ex assert _stream_raised_error_status(error) == expected_status -@pytest.mark.parametrize( - "error, has_generated_content, converts", - [ - (BedrockError(status_code=503, message="unavailable"), False, True), - (httpx.ReadError("connection reset by upstream"), False, True), - (BedrockError(status_code=400, message="malformed"), False, False), - (BedrockError(status_code=503, message="unavailable"), True, False), - ], - ids=["retriable", "no-status", "client-error", "after-content"], -) -def test_anthropic_stream_fallback_error_for_raised_gates_like_a_detected_error_event( - error, has_generated_content, converts -): - converted = _stream_fallback_error_for_raised(error, "primary", has_generated_content) - if not converts: - assert converted is None - return - assert isinstance(converted, MidStreamFallbackError) - assert converted.original_exception is error - assert converted.is_pre_first_chunk is True - assert converted.llm_provider == "anthropic" - - -@pytest.mark.asyncio -async def test_aanthropic_messages_recover_stream_error_flushes_buffered_frames_before_declining(): - router = _anthropic_messages_make_router() - original = BedrockError(status_code=503, message="unavailable") - declined = MidStreamFallbackError( - message="unavailable", - model="primary", - llm_provider="anthropic", - original_exception=original, - is_pre_first_chunk=False, - ) - buffered = (_anthropic_messages_message_start_chunk(),) - flushed = [] - - async def drain(recovery) -> None: - async for chunk in recovery: - flushed.append(chunk) - - with patch.object(router, "_aanthropic_messages_fallback_attempt") as mock_attempt: - recovery = router._aanthropic_messages_recover_stream_error( - declined, True, buffered, "primary", {"model": "primary"}, _anthropic_messages_make_wrapper(), AsyncMock() - ) - with pytest.raises(BedrockError) as exc_info: - await drain(recovery) - assert flushed == list(buffered) - assert exc_info.value is original - mock_attempt.assert_not_called() - - -@pytest.mark.asyncio -async def test_aanthropic_messages_recover_stream_error_hands_converted_raise_to_fallback_attempt(): - router = _anthropic_messages_make_router() - raised = BedrockError(status_code=503, message="unavailable") - handed_over = [] - - async def fake_attempt(fallback_error, initial_kwargs, wrapper): - handed_over.append(fallback_error) - yield b"fallback" - - with patch.object(router, "_aanthropic_messages_fallback_attempt", new=fake_attempt): - recovery = router._aanthropic_messages_recover_stream_error( - raised, False, (), "primary", {"model": "primary"}, _anthropic_messages_make_wrapper(), AsyncMock() - ) - collected = [chunk async for chunk in recovery] - assert collected == [b"fallback"] - assert len(handed_over) == 1 - assert isinstance(handed_over[0], MidStreamFallbackError) - assert handed_over[0].original_exception is raised - - -@pytest.mark.asyncio -async def test_anthropic_messages_non_retriable_client_error_skips_fallback(): - """A 4xx (non-429) error type (e.g. invalid_request_error) is a client - error a fallback attempt cannot fix, so it must be forwarded to the - client as-is rather than burning a fallback attempt.""" - router = _anthropic_messages_make_router() - error_chunk = _anthropic_messages_invalid_request_error_chunk() - source = _AnthropicMessagesFakeByteStream([error_chunk]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [error_chunk] - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_post_first_chunk_error_skips_fallback(): - """Once content has already reached the caller, retrying would start a - second, overlapping Anthropic message lifecycle on the same SSE stream - - the error must be forwarded instead of triggering an invisible retry.""" - router = _anthropic_messages_make_router() - content = _anthropic_messages_content_chunk("partial answer") - error_chunk = _anthropic_messages_overloaded_error_chunk() - source = _AnthropicMessagesFakeByteStream([content, error_chunk]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [content, error_chunk] - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_stream_the_client_closes_at_a_forwarded_error_is_not_closed_as_superseded(): - """An error frame forwarded to the client ends that request, and the SDK hangs up on it. No retry took the - request over, so the source must not be closed as a superseded attempt, which would log it as a failure and - drop the usage the provider billed for it.""" - router = _anthropic_messages_make_router() - error_chunk = _anthropic_messages_overloaded_error_chunk() - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_content_chunk("partial answer"), error_chunk]) - - wrapped = await router._aanthropic_messages_streaming_iterator(response=source, initial_kwargs={"model": "primary"}) - async for chunk in wrapped: - if chunk == error_chunk: - break - await wrapped.aclose() - - assert source.closed is True - assert source.closed_as_superseded is False - - -@pytest.mark.asyncio -async def test_anthropic_messages_non_retriable_error_flushes_buffered_lifecycle_frames(): - """A non-retriable error arriving while lifecycle frames are still - buffered (no content seen yet) must flush those buffered frames before - forwarding the error, so the client still sees the whole primary - attempt rather than losing the buffered message_start silently.""" - router = _anthropic_messages_make_router() - error_chunk = _anthropic_messages_invalid_request_error_chunk() - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_message_start_chunk(), error_chunk]) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [chunk async for chunk in wrapped] - - assert collected == [_anthropic_messages_message_start_chunk(), error_chunk] - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_raised_error_declined_flushes_buffered_lifecycle_frames(): - """When a raised MidStreamFallbackError is declined (source says content - was not pre-first-chunk) while lifecycle frames are still buffered, they - must be flushed to the client before the exception propagates.""" - router = _anthropic_messages_make_router() - raised_error = MidStreamFallbackError( - message="stream reset", - model="primary", - llm_provider="vertex_ai", - is_pre_first_chunk=False, - ) - source = _AnthropicMessagesRaisingByteStream([_anthropic_messages_message_start_chunk()], raised_error) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(), - ) as mock_fallback: - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [] - - async def _consume(): - async for chunk in wrapped: - collected.append(chunk) - - with pytest.raises(MidStreamFallbackError) as exc_info: - await _consume() - - assert collected == [_anthropic_messages_message_start_chunk()] - assert exc_info.value is raised_error - mock_fallback.assert_not_awaited() - - -@pytest.mark.asyncio -async def test_anthropic_messages_raised_error_without_original_exception_reraises_itself(): - """When a declined MidStreamFallbackError carries no original_exception, - the bare exception itself must propagate rather than being swallowed.""" - router = _anthropic_messages_make_router() - content = _anthropic_messages_content_chunk("partial answer") - raised_error = MidStreamFallbackError( - message="stream reset", - model="primary", - llm_provider="vertex_ai", - is_pre_first_chunk=False, - ) - source = _AnthropicMessagesRaisingByteStream([content], raised_error) - - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - collected = [] - - async def _consume(): - async for chunk in wrapped: - collected.append(chunk) - - with pytest.raises(MidStreamFallbackError) as exc_info: - await _consume() - - assert collected == [content] - assert exc_info.value is raised_error - - -@pytest.mark.asyncio -async def test_anthropic_messages_fallback_also_failing_raises_original_exception(): - """If the fallback attempt itself fails with a MidStreamFallbackError - wrapping a real provider exception, the client must see that real - exception, not the internal MidStreamFallbackError.""" - router = _anthropic_messages_make_router() - source = _AnthropicMessagesFakeByteStream([_anthropic_messages_overloaded_error_chunk()]) - original_exception = litellm.APIError( - status_code=503, - message="fallback also overloaded", - llm_provider="bedrock", - model="fallback", - ) - fallback_failure = MidStreamFallbackError( - message="fallback failed", - model="fallback", - llm_provider="bedrock", - original_exception=original_exception, - ) - - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - new=AsyncMock(side_effect=fallback_failure), - ): - wrapped = await router._aanthropic_messages_streaming_iterator( - response=source, - initial_kwargs={"model": "primary"}, - ) - with pytest.raises(litellm.APIError) as exc_info: - async for _ in wrapped: - pass - - assert exc_info.value is original_exception - - # -------- _dispatch_generic_call_type -------- -@pytest.mark.asyncio -async def test_dispatch_generic_call_type_routes_anthropic_messages_through_streaming_fallbacks(): - router = _anthropic_messages_make_router() - - async def fake_original(**_kwargs): - return {"id": "msg_1"} - - with patch.object( - router, - "_aanthropic_messages_with_streaming_fallbacks", - new=AsyncMock(return_value="anthropic-result"), - ) as mock_anthropic: - out = await router._dispatch_generic_call_type( - call_type="anthropic_messages", - original_function=fake_original, - model="primary", - ) - assert out == "anthropic-result" - mock_anthropic.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_dispatch_generic_call_type_other_call_types_use_generic_fallback(): - router = _anthropic_messages_make_router() - - async def fake_original(**_kwargs): - return {"id": "file_1"} - - with patch.object( - router, - "_ageneric_api_call_with_fallbacks", - new=AsyncMock(return_value="generic-result"), - ) as mock_generic: - out = await router._dispatch_generic_call_type( - call_type="afile_delete", - original_function=fake_original, - model="primary", - ) - assert out == "generic-result" - mock_generic.assert_awaited_once() - - -@pytest.mark.asyncio -async def test_factory_function_anthropic_messages_uses_streaming_fallback_dispatch(): - """anthropic_messages must be wired through the mid-stream-fallback-aware - path rather than the bare generic dispatch every other call type without - special handling uses.""" - router = _anthropic_messages_make_router() - wrapped = router.factory_function(litellm.anthropic_messages, call_type="anthropic_messages") - assert callable(wrapped) - - with patch.object( - router, - "_aanthropic_messages_with_streaming_fallbacks", - new=AsyncMock(return_value="ok"), - ) as mock_anthropic: - result = await wrapped(model="primary") - assert result == "ok" - mock_anthropic.assert_awaited_once() - - @pytest.mark.asyncio async def test_async_function_with_fallbacks_stamps_zero_attempted_fallbacks(): """A request served by the primary model group records attempted_fallbacks=0 and @@ -21062,159 +10760,6 @@ async def test_async_function_with_fallbacks_skips_stamp_on_genuine_reentrant_ho assert metadata["original_model_group"] == "prod-chat" -def _record_router_acompletion_kwargs(router: litellm.Router) -> list: - """Spy on router._acompletion, recording each call's kwargs while delegating through.""" - records = [] - original_acompletion = router._acompletion - - @functools.wraps(original_acompletion) - async def _spy(*args, **spy_kwargs): - records.append(spy_kwargs) - return await original_acompletion(*args, **spy_kwargs) - - router._acompletion = _spy - return records - - -@pytest.mark.asyncio -async def test_async_function_with_fallbacks_scrubs_spoofed_values_from_sibling_bucket(): - """Spend logs read a truthy litellm_metadata dict in preference to metadata, so spoofed - stamp keys planted in the bucket the route does not own are removed on entry, in place, - before they can flow into the spend log row.""" - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, - } - ] - ) - metadata = {} - litellm_metadata = { - "attempted_fallbacks": 99, - "original_model_group": "spoofed-group", - "client_key": "client_value", - } - downstream_calls = _record_router_acompletion_kwargs(router) - - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hey"}], - metadata=metadata, - litellm_metadata=litellm_metadata, - ) - - assert len(downstream_calls) == 1 - downstream_sibling = downstream_calls[0]["litellm_metadata"] - assert "attempted_fallbacks" not in downstream_sibling - assert "original_model_group" not in downstream_sibling - assert downstream_sibling["client_key"] == "client_value" - assert "attempted_fallbacks" not in litellm_metadata - assert "original_model_group" not in litellm_metadata - assert litellm_metadata["client_key"] == "client_value" - assert metadata["attempted_fallbacks"] == 0 - assert metadata["original_model_group"] == "gpt-3.5-turbo" - - -@pytest.mark.asyncio -async def test_async_function_with_fallbacks_scrubs_sibling_bucket_in_place(): - """Everything below the router resolves the bucket by key presence, so the scrub edits - the caller's dict object like every other router bucket write. Rebinding kwargs to a - scrubbed copy detaches the proxy's request_data write-backs (guardrail telemetry, retry - accounting) from the object the spend row is built from.""" - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, - } - ] - ) - litellm_metadata = { - "attempted_fallbacks": 7, - "original_model_group": "planted-group", - "client_key": "client_value", - } - caller_snapshot = copy.deepcopy(litellm_metadata) - downstream_calls = _record_router_acompletion_kwargs(router) - - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hey"}], - metadata={}, - litellm_metadata=litellm_metadata, - ) - - assert len(downstream_calls) == 1 - assert downstream_calls[0]["litellm_metadata"] is litellm_metadata - assert "attempted_fallbacks" not in litellm_metadata - assert "original_model_group" not in litellm_metadata - assert litellm_metadata["client_key"] == caller_snapshot["client_key"] - - -@pytest.mark.asyncio -async def test_async_function_with_fallbacks_stamps_aliased_buckets_on_every_call(): - """One dict object passed as both metadata and litellm_metadata: the first call's own - stamp puts the reserved keys into the shared object, so the second call enters the - scrub with them present. Scrubbing in place keeps the stamp and the bucket on the same - object; a scrubbed copy would leave the spend reader's preferred bucket unstamped.""" - router = litellm.Router( - model_list=[ - { - "model_name": "chat-group", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, - } - ] - ) - shared_metadata = {"team": "alpha"} - downstream_calls = _record_router_acompletion_kwargs(router) - - for _ in range(3): - await router.acompletion( - model="chat-group", - messages=[{"role": "user", "content": "hey"}], - metadata=shared_metadata, - litellm_metadata=shared_metadata, - ) - - assert len(downstream_calls) == 3 - for call_kwargs in downstream_calls: - assert call_kwargs["litellm_metadata"] is shared_metadata - assert call_kwargs["metadata"] is shared_metadata - assert call_kwargs["litellm_metadata"]["attempted_fallbacks"] == 0 - assert call_kwargs["litellm_metadata"]["original_model_group"] == "chat-group" - - -@pytest.mark.asyncio -async def test_async_function_with_fallbacks_passes_clean_sibling_bucket_through_unchanged(): - """A sibling bucket carrying no reserved stamp keys is forwarded downstream as the - caller's own object with no copy made, matching pre-scrub behavior. Retry accounting - stamped into that bucket downstream predates the scrub and is out of its scope.""" - router = litellm.Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "hi"}, - } - ] - ) - litellm_metadata = {"client_key": "client_value"} - downstream_calls = _record_router_acompletion_kwargs(router) - - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "hey"}], - metadata={}, - litellm_metadata=litellm_metadata, - ) - - assert len(downstream_calls) == 1 - assert downstream_calls[0]["litellm_metadata"] is litellm_metadata - assert litellm_metadata["client_key"] == "client_value" - assert "attempted_fallbacks" not in litellm_metadata - assert "original_model_group" not in litellm_metadata - - @pytest.mark.asyncio async def test_run_async_fallback_keeps_caller_metadata_keys_on_the_wire(monkeypatch): """Under enable_preview_features, add_openai_metadata forwards only the first 16 @@ -21287,59 +10832,6 @@ async def test_run_async_fallback_keeps_caller_metadata_keys_on_the_wire(monkeyp assert small_wire["model_group"] == "fallback-group" -@pytest.mark.asyncio -async def test_run_async_fallback_two_hop_chain_reports_entry_group_and_hop_count(): - """A two-hop fallback chain stamps attempted_fallbacks=2 on the final leg and keeps - original_model_group at the group requested on entry: a later hop's stamp appends - after caller keys without overriding the value stamped by an earlier hop.""" - router = litellm.Router( - model_list=[ - { - "model_name": "group-a", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "litellm.InternalServerError"}, - }, - { - "model_name": "group-b", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "litellm.InternalServerError"}, - }, - { - "model_name": "group-c", - "litellm_params": {"model": "gpt-3.5-turbo", "mock_response": "ok"}, - }, - ], - fallbacks=[{"group-a": ["group-b"]}, {"group-b": ["group-c"]}], - num_retries=0, - ) - metadata = {} - leg_records = [] - original_acompletion = router._acompletion - - @functools.wraps(original_acompletion) - async def _spy(*args, **spy_kwargs): - leg_records.append((spy_kwargs.get("model"), copy.deepcopy(spy_kwargs.get("metadata")))) - return await original_acompletion(*args, **spy_kwargs) - - router._acompletion = _spy - - await router.acompletion( - model="group-a", - messages=[{"role": "user", "content": "hey"}], - metadata=metadata, - ) - - assert [model for model, _ in leg_records] == ["group-a", "group-b", "group-c"] - hop_one_metadata = leg_records[1][1] - assert hop_one_metadata["attempted_fallbacks"] == 1 - assert hop_one_metadata["original_model_group"] == "group-a" - assert hop_one_metadata["model_group"] == "group-b" - hop_two_metadata = leg_records[2][1] - assert hop_two_metadata["attempted_fallbacks"] == 2 - assert hop_two_metadata["original_model_group"] == "group-a" - assert hop_two_metadata["model_group"] == "group-c" - assert metadata["attempted_fallbacks"] == 0 - assert metadata["original_model_group"] == "group-a" - - def _permission_denied_error() -> litellm.PermissionDeniedError: return litellm.PermissionDeniedError( message="OpenrouterException - this key has no access to the model", @@ -21419,9 +10911,7 @@ def _router_with_failing_primary(fallback_access_check) -> Router: class _DenyFallbackBudgetCheck: - async def __call__( - self, *, model: str, request_kwargs: Mapping[str, object], ctx: RouterContext - ) -> bool: + async def __call__(self, *, model: str, request_kwargs: Mapping[str, object], ctx: RouterContext) -> bool: return False @@ -21608,237 +11098,6 @@ def test_resolved_litellm_models_base_model_precedence_is_opt_in() -> None: assert router.resolved_litellm_models("never-configured", use_base_model=True) == () -class TestTierParamsTheTargetAccepts: - """A tier's litellm_params are applied to every request that tier routes, so one the target - cannot take raised UnsupportedParamsError before the request left the proxy, turning the whole - tier into a 400.""" - - @pytest.fixture(autouse=True) - def force_local_model_cost(self, monkeypatch): - from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap - - monkeypatch.setattr(litellm, "model_cost", GetModelCostMap.load_local_model_cost_map()) - - @staticmethod - def _router(model: str) -> litellm.Router: - return litellm.Router( - model_list=[{"model_name": "tiered", "litellm_params": {"model": model, "api_key": "sk-x"}}] - ) - - def test_drops_a_param_no_deployment_declares(self): - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts("tiered", {"reasoning_effort": "max"}, {}) - - assert accepted == {} - - def test_keeps_a_param_the_deployment_declares(self): - router = self._router("fireworks_ai/kimi-k3") - - accepted = router._tier_params_the_target_accepts("tiered", {"reasoning_effort": "max"}, {}) - - assert accepted == {"reasoning_effort": "max"} - - @pytest.mark.parametrize( - "control, value", - [ - ("api_base", "https://example.invalid"), - ("api_key", "sk-tier"), - ("base_url", "https://example.invalid"), - ("timeout", 30), - ("default_headers", {"x-tier": "1"}), - ("organization", "org-tier"), - ("deployment_id", "dep-tier"), - ], - ) - def test_keeps_credentials_and_transport_controls(self, control, value): - """These are not chat completion params, so get_optional_params never compares them against - a provider's supported list. Filtering on "is this an OpenAI param" would discard the - configuration the request needs while never touching what the provider would reject.""" - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts("tiered", {control: value, "reasoning_effort": "max"}, {}) - - assert accepted == {control: value} - - @pytest.mark.parametrize( - "control, value", - [ - ("additional_drop_params", ["seed"]), - ("drop_params", True), - ("allowed_openai_params", ["seed"]), - ("api_version", "2024-02-01"), - ("metadata", {"tier": "complex"}), - ], - ) - def test_keeps_litellm_controls_the_provider_never_lists(self, control, value): - """No provider lists a litellm control among its supported params, so "no deployment - declares it" means litellm consumes it, not that the target refuses it. Dropping - drop_params or additional_drop_params would silently disable the operator's sanitization.""" - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts("tiered", {control: value, "reasoning_effort": "max"}, {}) - - assert accepted == {control: value} - - def test_tier_allowlist_protects_the_param_it_names(self): - """allowed_openai_params is the documented escape hatch for an incomplete supported-params - list, and request-time validation extends the supported list with it, so a param the tier - both sets and allowlists would never 400 and must not be dropped.""" - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts( - "tiered", {"reasoning_effort": "max", "allowed_openai_params": ["reasoning_effort"]}, {} - ) - - assert accepted == {"reasoning_effort": "max", "allowed_openai_params": ["reasoning_effort"]} - - def test_request_allowlist_protects_the_param_it_names(self): - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts( - "tiered", {"reasoning_effort": "max"}, {"allowed_openai_params": ["reasoning_effort"]} - ) - - assert accepted == {"reasoning_effort": "max"} - - def test_allowlist_protects_only_the_params_it_names(self): - router = self._router("novita/moonshotai/kimi-k3") - - accepted = router._tier_params_the_target_accepts( - "tiered", {"reasoning_effort": "max", "allowed_openai_params": ["seed"]}, {} - ) - - assert accepted == {"allowed_openai_params": ["seed"]} - - def test_declared_param_allowlist_ignores_malformed_declarations(self): - """A str is iterable, so without the type guard a YAML scalar mistake like - allowed_openai_params: reasoning_effort would allowlist single characters.""" - assert litellm.Router._declared_param_allowlist({"allowed_openai_params": ["reasoning_effort", 3]}) == frozenset( - {"reasoning_effort"} - ) - assert litellm.Router._declared_param_allowlist({"allowed_openai_params": "reasoning_effort"}) == frozenset() - assert litellm.Router._declared_param_allowlist({}) == frozenset() - - def test_deployment_accepts_param_honors_deployment_allowlist(self): - deployment = { - "model_name": "x", - "litellm_params": {"model": "novita/moonshotai/kimi-k3", "allowed_openai_params": ["reasoning_effort"]}, - } - - assert litellm.Router._deployment_accepts_param(deployment, "x", "reasoning_effort") is True - - def test_keeps_a_token_ceiling_the_provider_spells_differently(self): - """petals lists max_tokens but not max_completion_tokens. A tier ceiling in the unsupported - spelling is a cost bound: dropping it would let a caller's larger max_tokens through where - today the mismatch fails loudly.""" - router = self._router("petals/petals-team/StableBeluga2") - - accepted = router._tier_params_the_target_accepts( - "tiered", {"max_completion_tokens": 100, "reasoning_effort": "max"}, {} - ) - - assert accepted == {"max_completion_tokens": 100} - - def test_keeps_extra_headers_even_when_the_provider_omits_it(self): - """Several providers leave extra_headers out of their supported params, so the filter would - drop it. Headers carry auth and tenancy, so sending fewer than the operator configured is - worse than the error they already get.""" - router = self._router("ai21/jamba-1.5-mini") - - accepted = router._tier_params_the_target_accepts( - "tiered", {"extra_headers": {"x-tenant": "acme"}, "reasoning_effort": "max"}, {} - ) - - assert accepted == {"extra_headers": {"x-tenant": "acme"}} - - def test_keeps_a_param_any_deployment_in_the_group_declares(self): - """Routing has not picked a deployment yet, so one capable member keeps the param alive.""" - router = litellm.Router( - model_list=[ - {"model_name": "tiered", "litellm_params": {"model": "novita/moonshotai/kimi-k3", "api_key": "k"}}, - {"model_name": "tiered", "litellm_params": {"model": "fireworks_ai/kimi-k3", "api_key": "k"}}, - ] - ) - - accepted = router._tier_params_the_target_accepts("tiered", {"reasoning_effort": "max"}, {}) - - assert accepted == {"reasoning_effort": "max"} - - def test_deployment_accepts_param_honors_base_model(self): - """An azure deployment named after the deployment rather than the model carries the real - model in base_model, and request-time mapping resolves capability through it, so the filter - has to ask the same question or it drops a param the deployment accepts.""" - by_model_info = { - "model_name": "x", - "litellm_params": {"model": "azure/my-gpt5-deploy"}, - "model_info": {"base_model": "azure/gpt-5"}, - } - by_litellm_params = { - "model_name": "x", - "litellm_params": {"model": "azure/my-gpt5-deploy", "base_model": "azure/gpt-5"}, - } - without_hint = {"model_name": "x", "litellm_params": {"model": "azure/my-gpt5-deploy"}} - - assert litellm.Router._deployment_accepts_param(by_model_info, "x", "reasoning_effort") is True - assert litellm.Router._deployment_accepts_param(by_litellm_params, "x", "reasoning_effort") is True - assert litellm.Router._deployment_accepts_param(without_hint, "x", "reasoning_effort") is False - - def test_deployment_accepts_param_reads_the_provider(self): - deployment = {"model_name": "x", "litellm_params": {"model": "fireworks_ai/kimi-k3"}} - - assert litellm.Router._deployment_accepts_param(deployment, "x", "reasoning_effort") is True - - def test_deployment_accepts_param_is_false_when_the_provider_omits_it(self): - deployment = {"model_name": "x", "litellm_params": {"model": "novita/moonshotai/kimi-k3"}} - - assert litellm.Router._deployment_accepts_param(deployment, "x", "reasoning_effort") is False - - @pytest.mark.parametrize( - "deployment", - [{"model_name": "x"}, {"model_name": "x", "litellm_params": {}}, {"model_name": "x", "litellm_params": {"model": "not-a-real-provider/nope"}}], - ) - def test_deployment_accepts_param_fails_open(self, deployment): - """An unresolvable deployment must not be the reason a param is dropped.""" - assert litellm.Router._deployment_accepts_param(deployment, "x", "reasoning_effort") is True - - @pytest.mark.parametrize( - "litellm_params", - [ - {"model": "github_copilot/gpt-4o"}, - {"model": "chatgpt/gpt-5"}, - {"model": "gpt-4o", "custom_llm_provider": "github_copilot"}, - ], - ) - def test_deployment_accepts_param_never_asks_a_provider_whose_lookup_authenticates( - self, litellm_params, monkeypatch - ): - """Resolving github_copilot or chatgpt runs their OAuth device flow, so a capability - question asked from the routing path can freeze the event loop for minutes waiting on a - human. The deployment counts as accepting everything, and the lookup is never made: an - exception-based sentinel cannot prove that, because the filter swallows exceptions into - the same keep answer.""" - lookups: list = [] - - def _record(*args, **kwargs): - lookups.append((args, kwargs)) - raise RuntimeError("provider resolution must not run for an authenticating provider") - - monkeypatch.setattr(litellm, "get_llm_provider", _record) - deployment = {"model_name": "x", "litellm_params": litellm_params} - - assert litellm.Router._deployment_accepts_param(deployment, "x", "reasoning_effort") is True - assert lookups == [] - - def test_keeps_everything_for_an_unknown_group(self): - """An unresolvable target must never narrow what the request already did.""" - router = self._router("fireworks_ai/kimi-k3") - - accepted = router._tier_params_the_target_accepts("no-such-group", {"reasoning_effort": "max"}, {}) - - assert accepted == {"reasoning_effort": "max"} - - @pytest.mark.asyncio async def test_router_aspeech_without_voice_dispatches_ref_audio_cloning(respx_mock, monkeypatch): import base64 @@ -21915,256 +11174,6 @@ async def test_router_aspeech_request_voice_overrides_deployment_default(respx_m assert request_body["voice_id"] == "gb_oliver_neutral" -class TestRequestReasoningEffortOverride: - def test_drop_effort_from_nested_carrier_preserves_other_nested_values(self): - params: dict[str, object] = {"output_config": {"effort": "high", "format": "json"}} - - litellm.Router._pop_effort_from_nested_carrier(params, "output_config") - - assert params == {"output_config": {"format": "json"}} - - @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) - def test_is_classifier_internal_call_recognizes_both_metadata_carriers(self, metadata_key): - kwargs = {metadata_key: {"internal_call_origin": "autorouter_classifier"}} - - assert litellm.Router._is_classifier_internal_call(kwargs) is True - assert litellm.Router._is_classifier_internal_call({metadata_key: {}}) is False - - def test_removes_every_deployment_native_effort_carrier_without_mutating_shared_config(self): - extra_body: dict[str, object] = { - "reasoning_effort": "high", - "thinking": {"type": "enabled"}, - "output_config": {"effort": "high", "format": "json"}, - "reasoning": {"effort": "high", "summary": "detailed"}, - "provider_option": True, - } - deployment_params: dict[str, object] = { - "model": "bedrock/converse/anthropic.claude-3-7-sonnet", - "thinking": {"type": "enabled", "budget_tokens": 2048}, - "output_config": {"effort": "high", "format": {"type": "json_schema"}}, - "reasoning": {"effort": "high", "summary": "auto"}, - "extra_body": extra_body, - } - - sanitized = litellm.Router._deployment_params_with_request_reasoning_override( - deployment_params, {"reasoning_effort": "low"} - ) - - assert sanitized == { - "model": "bedrock/converse/anthropic.claude-3-7-sonnet", - "output_config": {"format": {"type": "json_schema"}}, - "reasoning": {"summary": "auto"}, - "extra_body": { - "output_config": {"format": "json"}, - "reasoning": {"summary": "detailed"}, - "provider_option": True, - }, - } - assert deployment_params["thinking"] == {"type": "enabled", "budget_tokens": 2048} - assert deployment_params["output_config"] == {"effort": "high", "format": {"type": "json_schema"}} - assert extra_body["reasoning_effort"] == "high" - - @pytest.mark.parametrize("request_kwargs", [{}, {"reasoning_effort": None}]) - def test_omitted_override_preserves_deployment_defaults(self, request_kwargs): - deployment_params = { - "model": "deepseek/deepseek-reasoner", - "thinking": {"type": "enabled"}, - "output_config": {"effort": "high"}, - } - - assert ( - litellm.Router._deployment_params_with_request_reasoning_override(deployment_params, request_kwargs) - == deployment_params - ) - - @pytest.mark.asyncio - async def test_280_concurrent_overrides_never_mutate_or_leak_through_shared_deployment_params(self): - deployment_params = { - "model": "fireworks_ai/accounts/fireworks/models/kimi-k2-thinking", - "thinking": {"type": "enabled"}, - "output_config": {"effort": "high", "format": "json"}, - "extra_body": {"reasoning_effort": "high", "tenant": "shared"}, - } - efforts = ("none", "minimal", "low", "medium", "high", "xhigh", "max") - - results = await asyncio.gather( - *( - asyncio.to_thread( - litellm.Router._deployment_params_with_request_reasoning_override, - deployment_params, - {"reasoning_effort": efforts[index % len(efforts)]}, - ) - for index in range(280) - ) - ) - - assert all("thinking" not in result for result in results) - assert all(result["output_config"] == {"format": "json"} for result in results) - assert all(result["extra_body"] == {"tenant": "shared"} for result in results) - assert deployment_params["thinking"] == {"type": "enabled"} - assert deployment_params["output_config"] == {"effort": "high", "format": "json"} - assert deployment_params["extra_body"] == {"reasoning_effort": "high", "tenant": "shared"} - - @pytest.mark.parametrize( - ("metadata", "should_drop"), - [({"internal_call_origin": "autorouter_classifier"}, True), ({}, False)], - ids=["classifier", "ordinary-request"], - ) - def test_only_classifier_calls_drop_effort_for_an_unsupported_fallback(self, metadata, should_drop): - router = litellm.Router(model_list=[]) - body: dict[str, object] = {"model": "classifier", "reasoning_effort": "low"} - kwargs: dict[str, object] = { - "reasoning_effort": "low", - "metadata": metadata, - "proxy_server_request": {"body": body}, - } - deployment: DeploymentTypedDict = { - "model_name": "fallback", - "litellm_params": {"model": "openai/gpt-4o-mini"}, - } - - router._drop_unsupported_classifier_reasoning_effort(deployment, "fallback", kwargs) - - assert ("reasoning_effort" not in kwargs) is should_drop - assert ("reasoning_effort" not in body) is should_drop - - -@_PYTHON_CORE_ONLY -class TestPreRoutingTierDrivesFallbacks: - """#38832: a complexity/auto router picks a tier behind the router name, but fallback - lookup stayed on the router name, so the tier's configured chain never ran and a - provider failure on the tier's first hop was returned to the client.""" - - @staticmethod - async def _select_tier(model, request_kwargs, messages=None, input=None, specific_deployment=False): - from litellm.types.router import PreRoutingHookResponse - - if model == "smart-router": - return PreRoutingHookResponse(model="tier1", messages=messages) - return None - - @classmethod - def _router(cls, fallbacks) -> "litellm.Router": - router = litellm.Router( - model_list=[ - { - "model_name": "smart-router", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "sk-x"}, - }, - { - "model_name": "tier1", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-x", - "mock_response": "litellm.RateLimitError", - }, - }, - { - "model_name": "backup-a", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-x", - "mock_response": "from backup-a", - }, - }, - { - "model_name": "backup-b", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-x", - "mock_response": "from backup-b", - }, - }, - { - "model_name": "failing-backup", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-x", - "mock_response": "litellm.RateLimitError", - }, - }, - { - "model_name": "plain", - "litellm_params": { - "model": "openai/gpt-4o-mini", - "api_key": "sk-x", - "mock_response": "litellm.RateLimitError", - }, - }, - ], - fallbacks=fallbacks, - num_retries=0, - ) - router._core_async_pre_routing_hook = cls._select_tier - return router - - @pytest.mark.asyncio - async def test_the_selected_tier_fallback_chain_runs(self): - router = self._router([{"tier1": ["backup-a"]}]) - - response = await router.acompletion( - model="smart-router", messages=[{"role": "user", "content": "hi"}] - ) - - assert response.choices[0].message.content == "from backup-a" - - @pytest.mark.asyncio - async def test_a_chain_keyed_on_the_router_name_is_not_used(self): - """The router name has no chain of its own, so nothing should rescue this call.""" - router = self._router([{"tier2": ["backup-a"]}]) - - with pytest.raises(litellm.RateLimitError): - await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) - - @pytest.mark.asyncio - async def test_a_chain_keyed_on_the_router_name_rescues_when_no_tier_chain_exists(self): - """The documented contract: configs keyed on the requested name keep working behind auto-routers.""" - router = self._router([{"smart-router": ["backup-a"]}]) - - response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) - - assert response.choices[0].message.content == "from backup-a" - - @pytest.mark.asyncio - async def test_the_tier_chain_wins_over_the_router_name_chain(self): - router = self._router([{"tier1": ["backup-a"]}, {"smart-router": ["backup-b"]}]) - - response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) - - assert response.choices[0].message.content == "from backup-a" - - @pytest.mark.asyncio - async def test_a_request_without_a_pre_routing_hook_still_uses_its_own_group(self): - router = self._router([{"tier1": ["backup-a"]}]) - - response = await router.acompletion( - model="tier1", messages=[{"role": "user", "content": "hi"}] - ) - - assert response.choices[0].message.content == "from backup-a" - - @pytest.mark.asyncio - async def test_a_caller_cannot_pick_the_chain_by_sending_the_selection(self): - """The metadata bucket carries caller-supplied keys, so only the hook may set the tier.""" - router = self._router([{"tier1": ["backup-a"]}]) - - with pytest.raises(litellm.RateLimitError): - await router.acompletion( - model="plain", - messages=[{"role": "user", "content": "hi"}], - metadata={"pre_routing_selected_model": "tier1"}, - ) - - @pytest.mark.asyncio - async def test_each_fallback_hop_resolves_its_own_chain(self): - """The second hop must key off the group it is running, not the tier that failed.""" - router = self._router([{"tier1": ["failing-backup"]}, {"failing-backup": ["backup-b"]}]) - - response = await router.acompletion(model="smart-router", messages=[{"role": "user", "content": "hi"}]) - - assert response.choices[0].message.content == "from backup-b" - - @pytest.mark.asyncio async def test_prompt_management_factory_marks_injection_for_every_deployment(monkeypatch): """The factory stamps a provisional deployment's model_info into kwargs before the @@ -22451,58 +11460,6 @@ def _retry_skip_deployment(deployment_id, host, litellm_params=None, model_info= } -@pytest.mark.parametrize( - "status_code,failed_deployment_id,already_skipped,expected", - [ - (400, "rejecting", None, ("rejecting",)), - (403, "rejecting", None, ("rejecting",)), - (400, "second", ("first",), ("first", "second")), - (400, "first", ("first",), ("first",)), - (429, "rejecting", None, ()), - (503, "rejecting", None, ()), - (408, "rejecting", None, ()), - (400, None, None, ()), - (None, "rejecting", None, ()), - ("400", "rejecting", None, ()), - (400, "second", 7, ("second",)), - (400, "second", "first", ("second",)), - (400, "second", ["first"], ("second",)), - (400, "second", ("first", 7), ("first", "second")), - ], -) -def test_router_deployment_ids_to_skip_on_retry(status_code, failed_deployment_id, already_skipped, expected): - exception = Exception("upstream refused this request") - exception.status_code = status_code - exception.failed_deployment_id = failed_deployment_id - - assert litellm.Router._deployment_ids_to_skip_on_retry(exception, already_skipped) == expected - - -@pytest.mark.parametrize( - "kwargs,failed_deployment_id,expected", - [ - ({"model_info": {"id": "rejecting"}}, None, ("rejecting",)), - ({"model_info": {"id": "rejecting"}}, "cooldown-target", ("rejecting",)), - ({"model_info": {"id": ""}}, None, ()), - ({"model_info": {"id": 7}}, None, ()), - ({"model_info": "rejecting"}, None, ()), - ({}, None, ()), - ({}, "cooldown-target", ("cooldown-target",)), - ], -) -def test_router_retry_skip_stamp_feeds_deployment_ids_to_skip_on_retry( - kwargs: Mapping[str, object], failed_deployment_id: str | None, expected: tuple[str, ...] -): - exception = Exception("upstream refused this request") - exception.status_code = 400 - exception.failed_deployment_id = failed_deployment_id - - litellm.Router._stamp_retry_skip_deployment_id(exception, kwargs) - - assert litellm.Router._deployment_ids_to_skip_on_retry(exception, None) == expected - assert exception.failed_deployment_id == failed_deployment_id - - @pytest.mark.parametrize( "value,expected", [ @@ -22654,7 +11611,6 @@ async def test_router_retry_policy_400_never_returns_to_a_deployment_that_alread monkeypatch: pytest.MonkeyPatch, ): monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - monkeypatch.setattr(simple_shuffle.random, "choice", lambda deployments: deployments[0]) router = litellm.Router( model_list=[ _retry_skip_deployment("first-refuser", "first-refuser", litellm_params={"weight": 1}), @@ -22690,8 +11646,10 @@ async def test_router_retry_policy_400_never_returns_to_a_deployment_that_alread ) response = await router.acompletion(model="gpt-5.6", messages=[{"role": "user", "content": "hi"}]) + # The weighted pick starts on the first refuser; which zero-weight deployment comes next is random, + # but no refuser is tried twice and the accepting one serves. assert first.call_count == 1 - assert second.call_count == 1 + assert second.call_count <= 1 assert accepting.call_count == 1 assert response.choices[0].message.content == "hi back" @@ -23477,36 +12435,6 @@ async def test_router_max_parallel_requests_overflow_takes_the_ordinary_429_fall assert await _async_get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) == [] -@pytest.mark.asyncio -async def test_router_deployment_slot_rejects_while_held_and_frees_slot_on_exit(): - router: Final = Router( - model_list=[ - { - "model_name": "gpt-5.6", - "litellm_params": { - "model": "openai/gpt-5.6", - "api_key": "sk-fake", - "max_parallel_requests": 1, - }, - "model_info": {"id": "slot-deployment"}, - } - ] - ) - deployment: Final = router.get_deployment(model_id="slot-deployment") - assert deployment is not None - kwargs: Final = {"model": "gpt-5.6"} - - async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): - with pytest.raises(litellm.RateLimitError) as overflow: - async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): - pass - assert overflow.value.status_code == 429 - assert "slot-deployment" in overflow.value.message - - async with router._deployment_slot(deployment=deployment.model_dump(), kwargs=kwargs, parent_otel_span=None): - pass - - @pytest.mark.asyncio async def test_router_deployment_drop_params_string_true_is_honored(monkeypatch): from litellm import Router @@ -23611,16 +12539,6 @@ def test_get_candidate_model_ids_for_route_covers_model_name_and_pattern(): assert "dep-wild" in router.get_candidate_model_ids_for_route(model="gpt-5") -def test_deployment_ids_stringifies_ids_and_skips_entries_without_a_model_info_id(): - deployments = ( - {"model_info": {"id": "a"}}, - {"model_info": {"id": 2}}, - {"model_info": {}}, - {"no_model_info": True}, - ) - assert Router._deployment_ids(deployments) == frozenset({"a", "2"}) - - def test_cached_model_info_lookups_match_uncached_and_reset_on_model_list_change(): def deployment(max_output_tokens: int) -> Deployment: return Deployment( @@ -23642,35 +12560,6 @@ def test_cached_model_info_lookups_match_uncached_and_reset_on_model_list_change assert router.cached_model_group_info("grp").max_output_tokens == 200 -class _OpenBreakerRedis: - def __init__(self) -> None: - from litellm.caching.redis_cache import RedisCircuitBreaker - - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) - for _ in range(3): - self._circuit_breaker.record_failure() - - @_redis_circuit_breaker_guard - async def async_get_cache(self, key, **kwargs): - raise AssertionError("never reached") - - -@pytest.mark.asyncio -async def test_an_open_circuit_breaker_skips_the_session_binding_without_a_warning(caplog): - router = litellm.Router( - model_list=[{"model_name": "haiku", "litellm_params": {"model": "anthropic/claude-haiku-4-5", "api_key": "k"}}] - ) - router._claude_code_session_router_cache = DualCache(redis_cache=_OpenBreakerRedis()) # pyright: ignore[reportArgumentType] # duck-typed Redis double - caplog.clear() - - with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): - binding = await router._get_claude_code_session_router_binding("quiet-session") - - assert binding is None - assert [record.getMessage() for record in caplog.records if record.levelno >= logging.WARNING] == [] - assert any("circuit breaker is open" in record.getMessage() for record in caplog.records) - - @pytest.mark.asyncio async def test_model_name_colliding_with_a_deployment_id_still_load_balances_the_group(): router = litellm.Router( @@ -23697,7 +12586,7 @@ async def test_model_name_colliding_with_a_deployment_id_still_load_balances_the assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep" -def test_sync_completion_runs_pre_call_checks_for_a_model_name_colliding_with_a_deployment_id(): +def test_sync_completion_routes_a_model_name_colliding_with_a_deployment_id_by_group(): router = litellm.Router( model_list=[ { @@ -23714,15 +12603,12 @@ def test_sync_completion_runs_pre_call_checks_for_a_model_name_colliding_with_a_ routing_strategy="simple-shuffle", ) - with patch.object(router, "routing_strategy_pre_call_checks") as pre_call_checks: - by_group = router.completion(model="gpt-5-nano", messages=[{"role": "user", "content": "hi"}]) - assert by_group._hidden_params["model_id"] == "gpt-5-mini-dep" - pre_call_checks.assert_called_once() - assert pre_call_checks.call_args.kwargs["deployment"]["model_info"]["id"] == "gpt-5-mini-dep" + by_group = router.completion(model="gpt-5-nano", messages=[{"role": "user", "content": "hi"}]) + assert by_group._hidden_params["model_id"] == "gpt-5-mini-dep" + assert by_group.choices[0].message.content == "B" - by_id = router.completion(model="gpt-5-mini-dep", messages=[{"role": "user", "content": "hi"}]) - assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep" - pre_call_checks.assert_called_once() + by_id = router.completion(model="gpt-5-mini-dep", messages=[{"role": "user", "content": "hi"}]) + assert by_id._hidden_params["model_id"] == "gpt-5-mini-dep" class TestMemberAutoRouterInference: @@ -23733,18 +12619,25 @@ class TestMemberAutoRouterInference: self.cache = UserApiKeyCache() self.team = LiteLLM_TeamTable( - team_id="router-team", models=["member-router", "permitted-model"], + team_id="router-team", + models=["member-router", "permitted-model"], members_with_roles=[Member(user_id="router-member", role="user")], ) self.actor = UserAPIKeyAuth( - user_id="router-member", team_id="router-team", user_role=LitellmUserRoles.INTERNAL_USER, - models=["member-router", "permitted-model"], api_key="test-key-hash", config={"timeout": 60}, + user_id="router-member", + team_id="router-team", + user_role=LitellmUserRoles.INTERNAL_USER, + models=["member-router", "permitted-model"], + api_key="test-key-hash", + config={"timeout": 60}, + ) + self.database = SimpleNamespace( + db=SimpleNamespace( + litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=self.team)), + litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(return_value=None)), + litellm_accessgrouptable=SimpleNamespace(find_unique=AsyncMock()), + ) ) - self.database = SimpleNamespace(db=SimpleNamespace( - litellm_teamtable=SimpleNamespace(find_unique=AsyncMock(return_value=self.team)), - litellm_teammembership=SimpleNamespace(find_unique=AsyncMock(return_value=None)), - litellm_accessgrouptable=SimpleNamespace(find_unique=AsyncMock()), - )) monkeypatch.setattr(proxy_server, "user_api_key_cache", self.cache) monkeypatch.setattr(proxy_server, "prisma_client", self.database) @@ -23754,44 +12647,71 @@ class TestMemberAutoRouterInference: return { "model_name": "model_name_router-team_member-router", "litellm_params": { - "model": "auto_router/complexity_router", "complexity_router_default_model": target, + "model": "auto_router/complexity_router", + "complexity_router_default_model": target, "complexity_router_config": { - "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), "adaptive": False, + "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), + "adaptive": False, **({"classifier_type": "llm", "classifier_llm_config": {"model": target}} if classifier else {}), }, - "tags": ["member" if member else "admin"], "timeout": 13.0 if member else 29.0, + "tags": ["member" if member else "admin"], + "timeout": 13.0 if member else 29.0, }, "model_info": { - "team_id": "router-team", "team_public_model_name": "member-router", "member_auto_router": member, + "team_id": "router-team", + "team_public_model_name": "member-router", + "member_auto_router": member, }, } @classmethod def _router(cls, *markers: dict[str, object]) -> Router: - return Router(model_list=[ - *(markers or (cls._marker(),)), - {"model_name": "permitted-model", "litellm_params": { - "model": "openai/gpt-4o-mini", "api_key": "test-key", "api_base": "https://api.openai.com/v1", - }}, - {"model_name": "restricted-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}}, - ]) + return Router( + model_list=[ + *(markers or (cls._marker(),)), + { + "model_name": "permitted-model", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "test-key", + "api_base": "https://api.openai.com/v1", + }, + }, + {"model_name": "restricted-model", "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}}, + ] + ) def _request( - self, *, actor: UserAPIKeyAuth | None = None, metadata_name: str = "metadata", tag: str = "member", + self, + *, + actor: UserAPIKeyAuth | None = None, + metadata_name: str = "metadata", + tag: str = "member", ) -> dict[str, object]: from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup return LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( - data={metadata_name: {"tags": [tag]}, **({"metadata": {"user_api_key_auth": {"user_role": "proxy_admin"}}} - if metadata_name == "litellm_metadata" else {})}, - user_api_key_dict=actor or self.actor, _metadata_variable_name=metadata_name, + data={ + metadata_name: {"tags": [tag]}, + **( + {"metadata": {"user_api_key_auth": {"user_role": "proxy_admin"}}} + if metadata_name == "litellm_metadata" + else {} + ), + }, + user_api_key_dict=actor or self.actor, + _metadata_variable_name=metadata_name, ) async def _route( - self, router: Router, request: dict[str, object] | None = None, model: str = "member-router", + self, + router: Router, + request: dict[str, object] | None = None, + model: str = "member-router", ) -> PreRoutingHookResponse: response: Final = await router.async_pre_routing_hook( - model=model, request_kwargs=request if request is not None else self._request(), + model=model, + request_kwargs=request if request is not None else self._request(), messages=[{"role": "user", "content": "Hello"}], ) assert response is not None @@ -23800,36 +12720,68 @@ class TestMemberAutoRouterInference: @pytest.mark.asyncio @pytest.mark.parametrize("metadata_name", ("metadata", "litellm_metadata")) async def test_cached_roster_revocation_blocks_classifier_and_session_rebinding( - self, metadata_name: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, + self, + metadata_name: str, + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, ) -> None: from litellm.proxy.auth.auth_checks import delete_cache_team_object monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") router: Final = self._router(self._marker(classifier=True)) - classify: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond(200, json={ - "id": "classifier", "object": "chat.completion", "created": 0, "model": "gpt-4o-mini", - "choices": [{"index": 0, "message": {"content": '{"tier":"SIMPLE"}', "role": "assistant"}, - "finish_reason": "stop"}], - "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, - }) - request: Final = {**self._request(metadata_name=metadata_name), "proxy_server_request": {"headers": { - "x-claude-code-session-id": "member-router-session", "x-app": "cli", - }}} + classify: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + 200, + json={ + "id": "classifier", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"content": '{"tier":"SIMPLE"}', "role": "assistant"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + request: Final = { + **self._request(metadata_name=metadata_name), + "proxy_server_request": { + "headers": { + "x-claude-code-session-id": "member-router-session", + "x-app": "cli", + } + }, + } first: Final = await self._route(router, request) assert first.model == "permitted-model" and first.routing_decision is not None assert first.routing_decision["cause"] == "llm_classifier" assert (await self._route(router, request)).model == "permitted-model" assert self.database.db.litellm_teamtable.find_unique.await_count == 1 assert self.database.db.litellm_teammembership.find_unique.await_count == 1 - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={"members_with_roles": []}) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={"members_with_roles": []} + ) await delete_cache_team_object( - team_id=self.team.team_id, team_alias=None, user_api_key_cache=self.cache, proxy_logging_obj=None, + team_id=self.team.team_id, + team_alias=None, + user_api_key_cache=self.cache, + proxy_logging_obj=None, ) with pytest.raises(HTTPException, match="no longer a member"): await self._route(router, request) - rebound: Final = {**request, "proxy_server_request": {"headers": { - "x-claude-code-session-id": "member-router-session", "x-app": "cli", "x-claude-code-agent-id": "subagent", - }}} + rebound: Final = { + **request, + "proxy_server_request": { + "headers": { + "x-claude-code-session-id": "member-router-session", + "x-app": "cli", + "x-claude-code-agent-id": "subagent", + } + }, + } with pytest.raises(HTTPException, match="no longer a member"): await self._route(router, rebound, model="restricted-model") assert classify.call_count == 2 @@ -23839,11 +12791,23 @@ class TestMemberAutoRouterInference: async def test_member_router_fails_closed(self, state: str, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.proxy import proxy_server - request: Final = {"metadata": {"user_api_key_team_id": "router-team", "user_api_key_auth": { - "team_id": "router-team", "user_role": "proxy_admin", - }}} if state == "forged" else self._request(actor=self.actor.model_copy( - update={"user_id": ""} if state == "empty-user" else {}, - )) + request: Final = ( + { + "metadata": { + "user_api_key_team_id": "router-team", + "user_api_key_auth": { + "team_id": "router-team", + "user_role": "proxy_admin", + }, + } + } + if state == "forged" + else self._request( + actor=self.actor.model_copy( + update={"user_id": ""} if state == "empty-user" else {}, + ) + ) + ) self.database.db.litellm_teamtable.find_unique.return_value = ( None if state == "deleted" else self.team.model_copy(update={"blocked": state == "blocked"}) ) @@ -23854,11 +12818,20 @@ class TestMemberAutoRouterInference: assert error.value.status_code == (503 if state == "unavailable" else 403) @pytest.mark.asyncio - @pytest.mark.parametrize("user_id,role", [(None, LitellmUserRoles.INTERNAL_USER), ("admin", LitellmUserRoles.PROXY_ADMIN)]) - async def test_service_key_and_admin_preserve_runtime_access(self, user_id: str | None, role: LitellmUserRoles) -> None: - assert (await self._route(self._router(), self._request( - actor=self.actor.model_copy(update={"user_id": user_id, "user_role": role}), - ))).model == "permitted-model" + @pytest.mark.parametrize( + "user_id,role", [(None, LitellmUserRoles.INTERNAL_USER), ("admin", LitellmUserRoles.PROXY_ADMIN)] + ) + async def test_service_key_and_admin_preserve_runtime_access( + self, user_id: str | None, role: LitellmUserRoles + ) -> None: + assert ( + await self._route( + self._router(), + self._request( + actor=self.actor.model_copy(update={"user_id": user_id, "user_role": role}), + ), + ) + ).model == "permitted-model" @pytest.mark.asyncio @pytest.mark.parametrize("ceiling", ("team", "key", "member", "organization", "project")) @@ -23869,35 +12842,56 @@ class TestMemberAutoRouterInference: from litellm.proxy._types import LiteLLM_ProjectTableCachedObj from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ - "models": ["member-router"] if ceiling == "team" else self.team.models, - "organization_id": "router-org" if ceiling == "organization" else None, - }) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={ + "models": ["member-router"] if ceiling == "team" else self.team.models, + "organization_id": "router-org" if ceiling == "organization" else None, + } + ) if ceiling == "member": await self.cache.async_set_cache( key=team_membership_reservation_cache_key(user_id="router-member", team_id="router-team"), - value=LiteLLM_TeamMembership(user_id="router-member", team_id="router-team", - litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["restricted-model"])), + value=LiteLLM_TeamMembership( + user_id="router-member", + team_id="router-team", + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["restricted-model"]), + ), model_type=LiteLLM_TeamMembership, ) elif ceiling == "organization": await self.cache.async_set_cache( - key="org_id:router-org", value=LiteLLM_OrganizationTable( - organization_id="router-org", budget_id="org-budget", created_by="admin", updated_by="admin", + key="org_id:router-org", + value=LiteLLM_OrganizationTable( + organization_id="router-org", + budget_id="org-budget", + created_by="admin", + updated_by="admin", models=["restricted-model"], - ), model_type=LiteLLM_OrganizationTable, + ), + model_type=LiteLLM_OrganizationTable, ) elif ceiling == "project": await self.cache.async_set_cache( - key="project_id:router-project", value=LiteLLM_ProjectTableCachedObj( - project_id="router-project", team_id="router-team", models=["restricted-model"], - ), model_type=LiteLLM_ProjectTableCachedObj, + key="project_id:router-project", + value=LiteLLM_ProjectTableCachedObj( + project_id="router-project", + team_id="router-team", + models=["restricted-model"], + ), + model_type=LiteLLM_ProjectTableCachedObj, ) with pytest.raises(ProxyException, match="is not available for this API key"): - await self._route(self._router(), self._request(actor=self.actor.model_copy(update={ - "models": ["member-router"] if ceiling == "key" else self.actor.models, - "project_id": "router-project" if ceiling == "project" else None, - }))) + await self._route( + self._router(), + self._request( + actor=self.actor.model_copy( + update={ + "models": ["member-router"] if ceiling == "key" else self.actor.models, + "project_id": "router-project" if ceiling == "project" else None, + } + ) + ), + ) @pytest.mark.asyncio @pytest.mark.parametrize("group_owner", ("team", "key")) @@ -23905,127 +12899,45 @@ class TestMemberAutoRouterInference: from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast group: Final = LiteLLM_AccessGroupTable( - access_group_id="router-group", access_group_name="Router targets", access_model_names=["permitted-model"], + access_group_id="router-group", + access_group_name="Router targets", + access_model_names=["permitted-model"], ) self.database.db.litellm_accessgrouptable.find_unique.return_value = group - self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy(update={ - "models": ["member-router"] if group_owner == "team" else self.team.models, - "access_group_ids": ["router-group"] if group_owner == "team" else [], - }) - request: Final = self._request(actor=self.actor.model_copy(update={ - "models": ["member-router"] if group_owner == "key" else self.actor.models, - "access_group_ids": ["router-group"] if group_owner == "key" else [], - })) + self.database.db.litellm_teamtable.find_unique.return_value = self.team.model_copy( + update={ + "models": ["member-router"] if group_owner == "team" else self.team.models, + "access_group_ids": ["router-group"] if group_owner == "team" else [], + } + ) + request: Final = self._request( + actor=self.actor.model_copy( + update={ + "models": ["member-router"] if group_owner == "key" else self.actor.models, + "access_group_ids": ["router-group"] if group_owner == "key" else [], + } + ) + ) router: Final = self._router() assert (await self._route(router, request)).model == "permitted-model" assert (await self._route(router, request)).model == "permitted-model" assert self.database.db.litellm_accessgrouptable.find_unique.await_count == 1 - self.database.db.litellm_accessgrouptable.find_unique.return_value = group.model_copy(update={"access_model_names": []}) + self.database.db.litellm_accessgrouptable.find_unique.return_value = group.model_copy( + update={"access_model_names": []} + ) await evict_and_broadcast(cache_keys=("access_group_id:router-group",), user_api_key_cache=self.cache) with pytest.raises(ProxyException, match="is not available for this API key"): await self._route(router, request) assert self.database.db.litellm_accessgrouptable.find_unique.await_count == 2 - @pytest.mark.asyncio - async def test_tagged_marker_owns_authorization_and_forwarded_parameters(self) -> None: - router: Final = self._router(self._marker(member=False), self._marker()) - request: Final = self._request() - selected: Final = router._selected_strategy_marker_deployment( - model="model_name_router-team_member-router", strategy_tags=("member",), request_kwargs=request, - ) - assert selected is not None and selected["model_info"]["member_auto_router"] is True - assert (await self._route(router, request)).model == "permitted-model" - assert request["timeout"] == 13.0 - await self.cache.async_set_cache( - key="team_id:router-team", model_type=LiteLLM_TeamTable, - value=self.team.model_copy(update={"models": ["member-router"]}), - ) - with pytest.raises(ProxyException, match="is not available for this API key"): - await self._route(router, self._request()) - self.database.db.litellm_teamtable.find_unique.reset_mock() - admin: Final = self._request(tag="admin") - assert (await self._route(router, admin)).model == "restricted-model" - assert admin["timeout"] == 29.0 - self.database.db.litellm_teamtable.find_unique.assert_not_awaited() - @pytest.mark.asyncio async def test_sdk_router_does_not_import_proxy_dependencies(self, monkeypatch: pytest.MonkeyPatch) -> None: router: Final = self._router(self._marker(member=False)) monkeypatch.setitem(sys.modules, "fastapi", None) monkeypatch.delitem(sys.modules, "litellm.proxy.auth.auto_router_checks", raising=False) - assert (await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}})).model == "restricted-model" - - -def _access_window_offsets(start_hours: float, end_hours: float, team_ids: list) -> dict: - now_utc = datetime.now(timezone.utc) - return { - "start": (now_utc + timedelta(hours=start_hours)).strftime("%H:%M"), - "end": (now_utc + timedelta(hours=end_hours)).strftime("%H:%M"), - "timezone": "UTC", - "team_ids": team_ids, - } - - -def _reserved_model_list(windows_for_reserved=None, windows_for_open=None) -> list: - reserved: dict = { - "model_name": "gpt-4o-ptu", - "litellm_params": {"model": "gpt-4o", "mock_response": "reserved"}, - "model_info": {"id": "reserved-deployment"}, - } - if windows_for_reserved is not None: - reserved["model_info"]["access_windows"] = windows_for_reserved - unreserved: dict = { - "model_name": "gpt-4o-ptu", - "litellm_params": {"model": "gpt-4o", "mock_response": "open"}, - "model_info": {"id": "open-deployment"}, - } - if windows_for_open is not None: - unreserved["model_info"]["access_windows"] = windows_for_open - return [reserved, unreserved] - - -def test_access_windows_hide_reserved_deployment_from_other_teams(): - router = Router( - model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - ), - ) - _, deployments = router._common_checks_available_deployment( - model="gpt-4o-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["open-deployment"] - - -def test_access_windows_raise_when_only_reserved_deployments_remain(): - router = Router(model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - windows_for_open=[_access_window_offsets(-1, 1, ["team-a"])], - )[:1]) - for request_kwargs in ({"metadata": {"user_api_key_team_id": "team-b"}}, {}): - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._common_checks_available_deployment(model="gpt-4o-ptu", request_kwargs=request_kwargs) - _, deployments = router._common_checks_available_deployment( - model="gpt-4o-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["reserved-deployment"] - - -def test_reserved_deployments_drop_strategy_markers_before_filtering(): - router = Router(model_list=_reserved_model_list()[:1]) - marker = {"model_name": "gpt-4o-ptu", "litellm_params": {"model": "auto_router/semantic"}} - reserved = { - "model_name": "gpt-4o-ptu", - "litellm_params": {"model": "gpt-4o"}, - "model_info": {"access_windows": [_access_window_offsets(-1, 1, ["team-a"])]}, - } - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._filter_reserved_deployments( - model="gpt-4o-ptu", - healthy_deployments=[marker, reserved], - request_team_id="team-b", - ) + assert ( + await self._route(router, {"metadata": {"user_api_key_team_id": "router-team"}}) + ).model == "restricted-model" def test_access_windows_invalid_timezone_fails_router_construction(): @@ -24050,151 +12962,6 @@ def test_access_windows_invalid_timezone_fails_router_construction(): ) -def test_access_windows_inactive_window_leaves_deployments_available(): - router = Router( - model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(2, 3, ["team-a"])], - ), - ) - _, deployments = router._common_checks_available_deployment( - model="gpt-4o-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - assert {d["model_info"]["id"] for d in deployments} == {"reserved-deployment", "open-deployment"} - - -def test_access_windows_apply_when_calling_by_model_id(): - router = Router(model_list=_reserved_model_list( - windows_for_reserved=[_access_window_offsets(-1, 1, ["team-a"])], - )) - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._common_checks_available_deployment( - model="reserved-deployment", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - _, deployment = router._common_checks_available_deployment( - model="reserved-deployment", - request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}, - ) - assert deployment["model_info"]["id"] == "reserved-deployment" - - -def test_access_windows_apply_when_calling_by_litellm_model_name(): - router = Router( - model_list=[ - { - "model_name": "gpt-4o-ptu", - "litellm_params": { - "model": "openai/gpt-5.6-bypass-probe", - "mock_response": "reserved", - }, - "model_info": { - "id": "reserved-litellm-model", - "access_windows": [_access_window_offsets(-1, 1, ["team-a"])], - }, - } - ], - ) - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._common_checks_available_deployment( - model="openai/gpt-5.6-bypass-probe", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - _, deployments = router._common_checks_available_deployment( - model="openai/gpt-5.6-bypass-probe", - request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["reserved-litellm-model"] - - -def test_access_windows_apply_to_specific_deployment_calls(): - router = Router( - model_list=[ - { - "model_name": "gpt-4o-ptu", - "litellm_params": { - "model": "openai/gpt-5.6-specific-probe", - "mock_response": "reserved", - }, - "model_info": { - "id": "reserved-specific", - "access_windows": [_access_window_offsets(-1, 1, ["team-a"])], - }, - } - ], - ) - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._common_checks_available_deployment( - model="openai/gpt-5.6-specific-probe", - specific_deployment=True, - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - _, deployments = router._common_checks_available_deployment( - model="openai/gpt-5.6-specific-probe", - specific_deployment=True, - request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["reserved-specific"] - - -def test_access_windows_apply_to_wildcard_early_resolve(): - router = Router( - model_list=[ - { - "model_name": "openai/*", - "litellm_params": {"model": "openai/*", "mock_response": "reserved"}, - "model_info": { - "id": "reserved-wildcard", - "access_windows": [_access_window_offsets(-1, 1, ["team-a"])], - }, - } - ], - ) - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._common_checks_available_deployment( - model="openai/gpt-probe-wildcard", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - _, deployments = router._common_checks_available_deployment( - model="openai/gpt-probe-wildcard", - request_kwargs={"metadata": {"user_api_key_team_id": "team-a"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["reserved-wildcard"] - - -def test_access_windows_filter_reserved_deployments_method(): - router = Router(model_list=_reserved_model_list()) - reserved: dict = { - "model_info": { - "id": "reserved-deployment", - "access_windows": [_access_window_offsets(-1, 1, ["team-a"])], - } - } - open_deployment: dict = {"model_info": {"id": "open-deployment"}} - assert [ - d["model_info"]["id"] - for d in router._filter_reserved_deployments( - model="gpt-4o-ptu", - healthy_deployments=[reserved, open_deployment], - request_team_id="team-b", - ) - ] == ["open-deployment"] - with pytest.raises(litellm.BadRequestError, match="reserved for another team"): - router._filter_reserved_deployments( - model="gpt-4o-ptu", - healthy_deployments=[reserved], - request_team_id="team-b", - ) - assert [ - d["model_info"]["id"] - for d in router._filter_reserved_deployments( - model="gpt-4o-ptu", - healthy_deployments=[reserved, open_deployment], - request_team_id="team-a", - ) - ] == ["reserved-deployment", "open-deployment"] - - def _scheduled_router(timeout: float) -> Router: return Router( model_list=[ @@ -24312,7 +13079,9 @@ async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_pref @pytest.mark.asyncio -async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_prefixed_context_window_fallback_key() -> None: +async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_prefixed_context_window_fallback_key() -> ( + None +): """The context-window chain is keyed the same way the ordinary chain is, so a key spelled like the wildcard deployment ("anthropic/claude-sonnet-4-6") must catch the bare group's context-window error too.""" router = litellm.Router( @@ -24347,25 +13116,6 @@ async def test_bare_model_group_served_by_wildcard_deployment_uses_provider_pref assert result["content"][0]["text"] == "served by the context window fallback" -def test_bare_model_group_served_by_wildcard_deployment_has_provider_prefixed_content_policy_fallback() -> None: - router = litellm.Router( - model_list=[ - { - "model_name": "anthropic/*", - "litellm_params": {"model": "anthropic/*", "api_key": "sk-fake"}, - }, - { - "model_name": "openai/gpt-5.5-pro", - "litellm_params": {"model": "openai/gpt-5.5-pro", "api_key": "sk-fake"}, - }, - ], - content_policy_fallbacks=[{"anthropic/claude-sonnet-4-6": ["openai/gpt-5.5-pro"]}], - ) - - assert router._has_content_policy_fallback("claude-sonnet-4-6", {}) is True - assert router._has_content_policy_fallback("claude-haiku-4-5", {}) is False - - def _shared_ptu_model_list() -> list: return [ { @@ -24388,35 +13138,6 @@ def _shared_ptu_model_list() -> list: ] -def test_ptu_shares_hide_the_shared_deployment_from_teams_holding_no_share(monkeypatch): - monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") - router = Router(model_list=_shared_ptu_model_list()) - for request_kwargs in ({"metadata": {"user_api_key_team_id": "team-c"}}, {"metadata": {}}, {}): - _, deployments = router._common_checks_available_deployment(model="gpt-4.1-ptu", request_kwargs=request_kwargs) - assert [d["model_info"]["id"] for d in deployments] == ["open-deployment"] - for team_id in ("team-a", "team-b"): - _, deployments = router._common_checks_available_deployment( - model="gpt-4.1-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": team_id}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["shared-deployment"] - - -def test_ptu_shares_raise_when_only_shared_deployments_remain(monkeypatch): - monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") - router = Router(model_list=_shared_ptu_model_list()[:1]) - with pytest.raises(litellm.BadRequestError, match="reserved for the teams holding a PTU share"): - router._common_checks_available_deployment( - model="gpt-4.1-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-c"}}, - ) - _, deployments = router._common_checks_available_deployment( - model="gpt-4.1-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-b"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["shared-deployment"] - - def test_team_ptu_ceiling_follows_shares_moved_on_the_live_model_list(): router: Final = Router(model_list=_shared_ptu_model_list()[:1]) team_a_before: Final = router.team_ptu_ceiling("team-a", "gpt-4.1-ptu") @@ -24436,19 +13157,6 @@ def test_team_ptu_ceiling_follows_shares_moved_on_the_live_model_list(): assert router.team_ptu_ceiling("team-a", "gpt-4.1-ptu") is None -def test_ptu_shares_hide_the_shared_deployment_even_while_the_feature_is_off(monkeypatch): - """The flag switches cost attribution on; a declared split is an access rule and holds - without it, so a team never reaches a deployment reserved for others while the flag is - off.""" - monkeypatch.delenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", raising=False) - router = Router(model_list=_shared_ptu_model_list()) - _, deployments = router._common_checks_available_deployment( - model="gpt-4.1-ptu", - request_kwargs={"metadata": {"user_api_key_team_id": "team-c"}}, - ) - assert [d["model_info"]["id"] for d in deployments] == ["open-deployment"] - - def test_a_shared_ptu_deployment_whose_shares_do_not_add_up_is_refused_at_registration(monkeypatch): monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") model_list = _shared_ptu_model_list()[:1] @@ -24494,9 +13202,7 @@ def test_a_config_entry_declaring_shares_without_terms_is_refused_even_while_the [ GuardrailRaisedException(guardrail_name="chunk-scanner", message="blocked"), HTTPException(status_code=403, detail={"error": "blocked", "guardrail_name": "chunk-scanner"}), - ModifyResponseException( - message="blocked", model="primary", request_data={}, guardrail_name="chunk-scanner" - ), + ModifyResponseException(message="blocked", model="primary", request_data={}, guardrail_name="chunk-scanner"), ], ) async def test_a_guardrail_verdict_is_neither_retried_nor_fallen_back(verdict: Exception) -> None: @@ -24522,39 +13228,6 @@ async def test_a_guardrail_verdict_is_neither_retried_nor_fallen_back(verdict: E assert [c.kwargs["metadata"]["model_group"] for c in mock_acompletion.call_args_list] == ["primary"] -@pytest.mark.parametrize( - ("value", "expected"), - [ - ("gpt-4\r\nERROR forged entry\n", "gpt-4ERROR forged entry"), - (RuntimeError("no deployments\r\nfor gpt-4"), "no deploymentsfor gpt-4"), - ("gpt-4", "gpt-4"), - ], -) -def test_without_line_breaks_drops_every_cr_and_lf_from_the_logged_value(value: object, expected: str) -> None: - assert _without_line_breaks(value) == expected - - -def test_a_failed_routing_read_prefetch_logs_the_request_model_without_its_line_breaks(monkeypatch, caplog) -> None: - router = litellm.Router( - model_list=[{"model_name": "gpt-4", "litellm_params": {"model": "openai/gpt-4", "api_key": "k"}}] - ) - forged_model: Final = "gpt-4\r\nERROR forged entry\n" - - def fail_lookup(model_name: str | None = None, team_id: str | None = None) -> None: - raise RuntimeError(f"no deployments for {model_name}") - - monkeypatch.setattr(router, "_core_get_model_list", fail_lookup) - caplog.clear() - - with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"): - router.arm_routing_read_prefetch(forged_model, {}) - - messages: Final = [r.getMessage() for r in caplog.records if "routing read prefetch not armed" in r.getMessage()] - assert messages == [ - "routing read prefetch not armed for gpt-4ERROR forged entry: no deployments for gpt-4ERROR forged entry" - ] - - def test_get_deployment_credentials_with_provider_preserves_anthropic_wif_params(): """ Test that get_deployment_credentials_with_provider preserves a litellm_params-configured @@ -24653,6 +13326,7 @@ async def test_failure_rpm_increment_declares_the_router_usage_key_family(): assert seen == ["router_usage"] assert current_service_target() is None + class _SpanRecordingInMemoryCache(InMemoryCache): """Records the live OTel span each read runs under, so the test sees what a Redis span would nest in.""" @@ -24721,132 +13395,6 @@ async def test_deployment_selection_runs_inside_a_route_phase_named_after_the_mo assert recording_cache.active_span_names and set(recording_cache.active_span_names) == {"route gpt-group"} -def _record_phase_events(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, dict[str, str | int]]]: - events: list[tuple[str, dict[str, str | int]]] = [] # mutable-ok: recorder for the injected phase_event double - - def record(name: str, attributes: dict[str, str | int]) -> None: - events.append((name, dict(attributes))) - - monkeypatch.setattr(litellm.router_python_core, "phase_event", record) - return events - - -def _pick(model_group: str, reason: str, attempt: int) -> tuple[str, dict[str, str | int]]: - return ( - "litellm.request.deployment_selected", - { - "litellm.deployment.attempt": attempt, - "litellm.deployment.reason": reason, - "litellm.deployment.model_group": model_group, - }, - ) - - -@pytest.mark.parametrize( - "request_kwargs, expected_reason, expected_attempt", - [ - (None, "initial", 1), - ({"metadata": {"attempted_retries": 0}, "fallback_depth": 0}, "initial", 1), - ({"metadata": {"attempted_retries": 2}}, "retry", 3), - ({"litellm_metadata": {"attempted_retries": 1}, "metadata": {"attempted_retries": 4}}, "retry", 2), - ({"metadata": {}, "fallback_depth": 1}, "fallback", 1), - ({"metadata": {"attempted_retries": 1}, "fallback_depth": 1}, "retry", 2), - ], -) -def test_deployment_pick_attributes_derive_attempt_and_reason( - request_kwargs: dict[str, object] | None, expected_reason: str, expected_attempt: int -): - attributes: Final = litellm.router._deployment_pick_attributes("gpt-4o", request_kwargs) - - assert dict(attributes) == { - "litellm.deployment.attempt": expected_attempt, - "litellm.deployment.reason": expected_reason, - "litellm.deployment.model_group": "gpt-4o", - } - - -@pytest.mark.asyncio -async def test_acompletion_marks_deployment_selected_once(monkeypatch: pytest.MonkeyPatch): - events: Final = _record_phase_events(monkeypatch) - router: Final = Router( - model_list=[ - { - "model_name": "gpt-4o", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake", "mock_response": "hi"}, - } - ] - ) - - await router.acompletion(model="gpt-4o", messages=[{"role": "user", "content": "hi"}]) - - assert events == [_pick("gpt-4o", "initial", 1)] - - -@pytest.mark.asyncio -async def test_acompletion_marks_every_retry_pick(monkeypatch: pytest.MonkeyPatch): - events: Final = _record_phase_events(monkeypatch) - router: Final = Router( - model_list=[ - { - "model_name": "flaky", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake", "mock_response": Exception("boom")}, - } - ], - num_retries=2, - retry_after=0, - ) - - with pytest.raises(Exception, match="boom"): - await router.acompletion(model="flaky", messages=[{"role": "user", "content": "hi"}]) - - assert events == [_pick("flaky", "initial", 1), _pick("flaky", "retry", 2), _pick("flaky", "retry", 3)] - - -@pytest.mark.asyncio -async def test_acompletion_marks_fallback_pick_with_its_model_group(monkeypatch: pytest.MonkeyPatch): - events: Final = _record_phase_events(monkeypatch) - router: Final = Router( - model_list=[ - { - "model_name": "primary", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake", "mock_response": Exception("boom")}, - }, - { - "model_name": "backup", - "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "fake", "mock_response": "hi"}, - }, - ], - fallbacks=[{"primary": ["backup"]}], - num_retries=0, - ) - - response: Final = await router.acompletion(model="primary", messages=[{"role": "user", "content": "hi"}]) - - assert response.choices[0].message.content == "hi" - assert events == [_pick("primary", "initial", 1), _pick("backup", "fallback", 1)] - - -@pytest.mark.asyncio -async def test_non_chat_surfaces_mark_their_deployment_pick(monkeypatch: pytest.MonkeyPatch): - """The event is emitted where the router picks, so embeddings and the sync path report it too.""" - events: Final = _record_phase_events(monkeypatch) - router: Final = Router( - model_list=[ - { - "model_name": "embed", - "litellm_params": {"model": "openai/text-embedding-3-small", "api_key": "fake", "mock_response": [0.1]}, - }, - { - "model_name": "gpt-4o", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake", "mock_response": "hi"}, - }, - ] - ) - - await router.aembedding(model="embed", input="hi") - router.completion(model="gpt-4o", messages=[{"role": "user", "content": "hi"}]) - - assert events == [_pick("embed", "initial", 1), _pick("gpt-4o", "initial", 1)] class TestAutoRouterTraceProvenance: @pytest.fixture def _tracing(self, monkeypatch: pytest.MonkeyPatch) -> "tuple[OpenTelemetryV2, InMemorySpanExporter]": @@ -24858,7 +13406,9 @@ class TestAutoRouterTraceProvenance: config: Final = OpenTelemetryV2Config(exporter="in_memory") exporter: Final = InMemorySpanExporter() - logger: Final = OpenTelemetryV2(config=config, tracer_provider=providers.build_tracer_provider(config, exporter=exporter)) + logger: Final = OpenTelemetryV2( + config=config, tracer_provider=providers.build_tracer_provider(config, exporter=exporter) + ) monkeypatch.setattr(proxy_server, "open_telemetry_logger", logger) return logger, exporter @@ -24868,27 +13418,14 @@ class TestAutoRouterTraceProvenance: "model_name": "traced-router", "model_info": {"id": f"definition-{tag}", "updated_at": "2026-10-01T12:00:00+00:00"}, "litellm_params": { - "model": "auto_router/complexity_router", "tags": [tag], + "model": "auto_router/complexity_router", + "tags": [tag], "complexity_router_config": { "tiers": dict.fromkeys(("SIMPLE", "MEDIUM", "COMPLEX", "REASONING"), target), }, }, } - def test_definition_fingerprint_is_canonical_and_tracks_config_not_credentials(self): - original: Final = Deployment.model_validate(self._marker("blue")) - reordered: Final = Deployment.model_validate({ - **self._marker("blue"), - "litellm_params": { - "api_key": "credential-only-change", "tags": ["blue"], "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": dict.fromkeys(("REASONING", "COMPLEX", "MEDIUM", "SIMPLE"), "answer")}, - }, - }) - fingerprint: Final = Router._routing_definition_fingerprint(original) - assert fingerprint is not None - assert fingerprint == Router._routing_definition_fingerprint(reordered) - assert fingerprint != Router._routing_definition_fingerprint(Deployment.model_validate(self._marker("blue", "other"))) - @pytest.mark.asyncio async def test_direct_classifier_and_plugin_prelude_never_inherit_parent_decisions( self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", monkeypatch: pytest.MonkeyPatch @@ -24905,7 +13442,9 @@ class TestAutoRouterTraceProvenance: async def classify(self, context: RoutingContext) -> str: await litellm.acompletion( - model="openai/gpt-5.6-luna", api_key="test", mock_response="SIMPLE", + model="openai/gpt-5.6-luna", + api_key="test", + mock_response="SIMPLE", messages=[{"role": "user", "content": "choose a tier"}], metadata={**context.metadata, "internal_call_origin": "autorouter_classifier"}, ) @@ -24947,7 +13486,8 @@ class TestAutoRouterTraceProvenance: ) await completed.wait() classifier_span: Final = next( - span for span in exporter.get_finished_spans() + span + for span in exporter.get_finished_spans() if (span.attributes or {}).get("litellm.request.purpose") == "autorouter_classifier" ) assert not any(key.startswith("litellm.routing.") for key in classifier_span.attributes) @@ -24961,36 +13501,66 @@ class TestAutoRouterTraceProvenance: @pytest.mark.asyncio @pytest.mark.parametrize("build_fails", [False, True]) async def test_semantic_internal_failure_keeps_phase_identity_without_parent_retry_attribution( - self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", monkeypatch: pytest.MonkeyPatch, - respx_mock: respx.MockRouter, build_fails: bool, + self, + _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + build_fails: bool, ) -> None: logger, exporter = _tracing monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) query_statuses: Final = iter((503, 200)) def embedding_response(request: httpx.Request) -> httpx.Response: - status: Final = 400 if build_fails else ( - next(query_statuses) if json.loads(request.content)["input"] == ["query-retry"] else 200 + status: Final = ( + 400 + if build_fails + else (next(query_statuses) if json.loads(request.content)["input"] == ["query-retry"] else 200) + ) + return httpx.Response( + status, + headers={"x-should-retry": "false"}, + json={ + "object": "list", + "model": "text-embedding-3-small", + "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + }, ) - return httpx.Response(status, headers={"x-should-retry": "false"}, json={ - "object": "list", "model": "text-embedding-3-small", - "data": [{"object": "embedding", "index": 0, "embedding": [1.0, 0.0]}], - "usage": {"prompt_tokens": 1, "total_tokens": 1}, - }) - upstream: Final = respx_mock.post("https://embedding.example/v1/embeddings").mock(side_effect=embedding_response) - router: Final = Router(model_list=[ - {"model_name": "semantic", "model_info": {"id": "semantic-definition"}, "litellm_params": { - "model": "auto_router/semantic-router", "auto_router_default_model": "answer", - "auto_router_embedding_model": "embed", - "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), - }}, - *({"model_name": "embed", "model_info": {"id": f"embed-{index}"}, "litellm_params": { - "model": "openai/text-embedding-3-small", "api_base": "https://embedding.example/v1", - "api_key": "test", "max_retries": 0, - }} for index in range(2)), - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, - ], num_retries=1, retry_after=0) + upstream: Final = respx_mock.post("https://embedding.example/v1/embeddings").mock( + side_effect=embedding_response + ) + router: Final = Router( + model_list=[ + { + "model_name": "semantic", + "model_info": {"id": "semantic-definition"}, + "litellm_params": { + "model": "auto_router/semantic-router", + "auto_router_default_model": "answer", + "auto_router_embedding_model": "embed", + "auto_router_config": json.dumps({"routes": [{"name": "answer", "utterances": ["hello"]}]}), + }, + }, + *( + { + "model_name": "embed", + "model_info": {"id": f"embed-{index}"}, + "litellm_params": { + "model": "openai/text-embedding-3-small", + "api_base": "https://embedding.example/v1", + "api_key": "test", + "max_retries": 0, + }, + } + for index in range(2) + ), + {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test"}}, + ], + num_retries=1, + retry_after=0, + ) stale: Final = {"routing_decision": {"router_model_name": "previous"}, "_autorouter_baseline_route": "old"} request: Final = {"metadata": {**stale}, "litellm_metadata": {**stale}} with logger.tracer.start_as_current_span("request") as root: @@ -25104,8 +13674,13 @@ class TestAutoRouterTraceProvenance: assert decision["router_config_id"] == router.model_list[0]["model_info"]["id"] assert decision["cause"] == "semantic_match" spans: Final = {span.name: span for span in exporter.get_finished_spans()} - assert spans["route public-router"].attributes["litellm.routing.router_model_name"] == decision["router_model_name"] - assert spans["route public-router"].attributes["litellm.routing.router_config_id"] == decision["router_config_id"] + assert ( + spans["route public-router"].attributes["litellm.routing.router_model_name"] + == decision["router_model_name"] + ) + assert ( + spans["route public-router"].attributes["litellm.routing.router_config_id"] == decision["router_config_id"] + ) retry: Final = next(event for event in spans["request"].events if event.name == "litellm.routing.retry") assert retry.attributes["litellm.routing.router_model_name"] == decision["router_model_name"] @@ -25115,18 +13690,37 @@ class TestAutoRouterTraceProvenance: self, _tracing: "tuple[OpenTelemetryV2, InMemorySpanExporter]", passthrough: bool ) -> None: logger, exporter = _tracing - router: Final = Router(model_list=[ - self._marker("blue"), self._marker("green"), - {"model_name": "answer", "litellm_params": {"model": "openai/gpt-5.6-luna", "api_key": "test", "use_in_pass_through": True, "tags": ["blue", "green"]}}, - ], enable_tag_filtering=True) + router: Final = Router( + model_list=[ + self._marker("blue"), + self._marker("green"), + { + "model_name": "answer", + "litellm_params": { + "model": "openai/gpt-5.6-luna", + "api_key": "test", + "use_in_pass_through": True, + "tags": ["blue", "green"], + }, + }, + ], + enable_tag_filtering=True, + ) request: Final = {"metadata": {"tags": ["green"]}, "litellm_metadata": {"tags": ["green"]}} - pick: Final = router.async_get_available_deployment_for_pass_through if passthrough else router.async_get_available_deployment + pick: Final = ( + router.async_get_available_deployment_for_pass_through + if passthrough + else router.async_get_available_deployment + ) with logger.tracer.start_as_current_span("request"): await pick(model="traced-router", request_kwargs=request, messages=[{"role": "user", "content": "hello"}]) decision: Final = request["litellm_metadata"]["routing_decision"] assert decision["router_config_id"] == "definition-green" assert decision["router_config_updated_at"] == "2026-10-01T12:00:00+00:00" - assert decision["router_config_fingerprint"] == router.complexity_routers["traced-router"][1].definition_fingerprint + assert ( + decision["router_config_fingerprint"] + == router.complexity_routers["traced-router"][1].definition_fingerprint + ) assert "routing_decision" not in request["metadata"] router.log_retry({**request, "model": "answer"}, ValueError("private-error-text")) await pick(model="answer", request_kwargs=request, messages=[{"role": "user", "content": "hello"}]) @@ -25147,6 +13741,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -25199,6 +13794,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -25219,6 +13815,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -25241,6 +13838,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @patch("azure.identity.get_bearer_token_provider") @patch("azure.identity.ClientSecretCredential") @@ -25322,48 +13920,6 @@ def test_router_init_azure_service_principal_with_secret_with_environment_variab # finally verify if the mocked token was used by Azure SDK mocked_func_generating_token.assert_called() -@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") -@pytest.mark.asyncio -async def test_audio_speech_router(): - """ - Test that router uses OpenAI/Azure OpenAI Client initialized during init for litellm.aspeech - """ - - from litellm import Router - - litellm.set_verbose = True - - model_list = [ - { - "model_name": "tts", - "litellm_params": { - "model": "azure/tts", - "api_base": os.getenv("AZURE_TTS_API_BASE"), - "api_key": os.getenv("AZURE_TTS_API_KEY"), - }, - }, - ] - - _router = Router(model_list=model_list) - - expected_openai_client = _router._get_client( - deployment=_router.model_list[0], - kwargs={}, - client_type="async", - ) - - with patch("litellm.aspeech") as mock_aspeech: - await _router.aspeech( - model="tts", - voice="alloy", - input="the quick brown fox jumped over the lazy dogs", - ) - - print("litellm.aspeech was called with kwargs = ", mock_aspeech.call_args.kwargs) - - # Get the actual client that was passed - client_passed_in_request = mock_aspeech.call_args.kwargs["client"] - assert client_passed_in_request == expected_openai_client @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_update_kwargs_before_fallbacks_unit_test(): @@ -25390,64 +13946,6 @@ def test_update_kwargs_before_fallbacks_unit_test(): assert kwargs["litellm_trace_id"] is not None -@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") -@pytest.mark.parametrize( - "call_type", - [ - CallTypes.acompletion, - CallTypes.atext_completion, - CallTypes.aembedding, - CallTypes.arerank, - CallTypes.atranscription, - ], -) -@pytest.mark.asyncio -async def test_update_kwargs_before_fallbacks(call_type): - - router = Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": { - "model": "azure/gpt-4.1-mini", - "api_key": "bad-key", - "api_version": os.getenv("AZURE_API_VERSION"), - "api_base": os.getenv("AZURE_AI_API_BASE"), - }, - } - ], - ) - - if call_type.value.startswith("a"): - with patch.object(router, "async_function_with_fallbacks") as mock_client: - if call_type.value == "acompletion": - input_kwarg = { - "messages": [{"role": "user", "content": "Hello, how are you?"}], - } - elif call_type.value == "atext_completion" or call_type.value == "aimage_generation": - input_kwarg = { - "prompt": "Hello, how are you?", - } - elif call_type.value == "aembedding" or call_type.value == "arerank": - input_kwarg = { - "input": "Hello, how are you?", - } - elif call_type.value == "atranscription": - input_kwarg = { - "file": "path/to/file", - } - else: - input_kwarg = {} - - await getattr(router, call_type.value)( - model="gpt-3.5-turbo", - **input_kwarg, - ) - - mock_client.assert_called_once() - - print(mock_client.call_args.kwargs) - assert mock_client.call_args.kwargs["litellm_trace_id"] is not None @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_model_info_wildcard_routes(): @@ -25468,6 +13966,7 @@ def test_router_get_model_info_wildcard_routes(): assert model_info["tpm"] is not None assert model_info["rpm"] is not None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio @pytest.mark.flaky(retries=3, delay=1) @@ -25498,6 +13997,7 @@ async def test_router_get_model_group_usage_wildcard_routes(): assert tpm is not None, "tpm is None" assert rpm is not None, "rpm is None" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_call_router_callbacks_on_success(): @@ -25531,6 +14031,7 @@ async def test_call_router_callbacks_on_success(): assert increment["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm") assert increment["increment_value"] == 1 + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.serial @pytest.mark.asyncio @@ -25554,10 +14055,14 @@ async def test_call_router_callbacks_on_failure(): num_retries=0, ) await asyncio.sleep(3) - print(mock_callback.call_args_list) - assert mock_callback.call_count == 1 + rpm_keys = [ + key + for call in mock_callback.call_args_list + for key in (call.kwargs.get("key", call.args[0] if call.args else None),) + if isinstance(key, str) and key.startswith("global_router:1:gemini/gemini-2.5-flash:rpm") + ] + assert len(rpm_keys) == 1 - assert mock_callback.call_args_list[0].kwargs["key"].startswith("global_router:1:gemini/gemini-2.5-flash:rpm") @pytest.mark.asyncio async def test_rate_limit_error_callback(): @@ -25604,13 +14109,8 @@ async def test_rate_limit_error_callback(): except Exception: pass - with patch.object( - customHandler, "log_model_group_rate_limit_error", new=AsyncMock() - ) as mock_client: - - print( - f"customHandler.log_model_group_rate_limit_error: {customHandler.log_model_group_rate_limit_error}" - ) + with patch.object(customHandler, "log_model_group_rate_limit_error", new=AsyncMock()) as mock_client: + print(f"customHandler.log_model_group_rate_limit_error: {customHandler.log_model_group_rate_limit_error}") try: _ = await router.acompletion( @@ -25658,6 +14158,7 @@ async def test_router_model_group_headers(): assert "x-ratelimit-remaining-requests" in resp._hidden_params["additional_headers"] assert "x-ratelimit-remaining-tokens" in resp._hidden_params["additional_headers"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.asyncio async def test_get_remaining_model_group_usage(): @@ -25689,6 +14190,7 @@ async def test_get_remaining_model_group_usage(): assert "x-ratelimit-remaining-requests" in remaining_usage assert "x-ratelimit-remaining-tokens" in remaining_usage + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "potential_access_group, expected_result", @@ -25707,6 +14209,7 @@ def test_router_get_model_access_groups(potential_access_group, expected_result) access_groups = router.is_model_access_group_for_wildcard_route(model_access_group=potential_access_group) assert access_groups == expected_result + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_redis_cache(): router = Router(model_list=[{"model_name": "gemini/*", "litellm_params": {"model": "gemini/*"}}]) @@ -25717,33 +14220,6 @@ def test_router_redis_cache(): assert router.cache.redis_cache == redis_cache -@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") -def test_router_handle_clientside_credential(): - """A caller-supplied credential must stay scoped to the current call: it must - never be registered as a router deployment, or a later caller with no override - of their own can be load-balanced onto it and reach the provider with someone - else's credential (see LIT-7811).""" - deployment = { - "model_name": "gemini/*", - "litellm_params": {"model": "gemini/*"}, - "model_info": { - "id": "1", - }, - } - router = Router(model_list=[deployment]) - - new_deployment = router._handle_clientside_credential( - deployment=deployment, - kwargs={ - "api_key": "123", - "metadata": {"model_group": "gemini/gemini-1.5-flash"}, - }, - function_name="acompletion", - ) - - assert new_deployment.litellm_params.api_key == "123" - assert len(router.get_model_list()) == 1 - assert router.get_deployment(model_id=new_deployment.model_info.id) is None @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") async def test_router_clientside_credential_not_reused_by_other_callers(respx_mock, monkeypatch: pytest.MonkeyPatch): @@ -25802,6 +14278,7 @@ async def test_router_clientside_credential_not_reused_by_other_callers(respx_mo used_auth_headers = {call.request.headers["authorization"] for call in route.calls[1:]} assert used_auth_headers == {"Bearer configured-key"} + def _caused_by_an_entra_rejection(error: litellm.AuthenticationError) -> litellm.AuthenticationError: from litellm.llms.azure.common_utils import EntraRejectionError @@ -25853,6 +14330,7 @@ async def test_router_caller_entra_credential_rejection_leaves_the_shared_deploy assert response.choices[0].message.content == "hi" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "caller_field", @@ -25909,6 +14387,7 @@ async def test_router_caller_entra_field_that_changes_the_deployment_credential_ assert response.choices[0].message.content == "hi" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( ("deployment_secret", "router_default_secret"), @@ -25958,6 +14437,7 @@ async def test_router_configured_entra_credential_rejection_cools_down_the_deplo assert excinfo.value.all_deployments_in_cooldown is True + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( "caller_fields", @@ -26011,6 +14491,7 @@ async def test_router_caller_entra_field_that_changes_nothing_still_cools_down_a assert excinfo.value.all_deployments_in_cooldown is True + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") async def test_router_caller_repeating_the_secret_an_empty_router_default_selects_still_cools_down(monkeypatch): monkeypatch.setenv("AZURE_CLIENT_SECRET", "environment-secret") @@ -26055,6 +14536,7 @@ async def test_router_caller_repeating_the_secret_an_empty_router_default_select assert excinfo.value.all_deployments_in_cooldown is True + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") async def test_router_caller_entra_field_keeps_the_shared_deployment(): router = Router( @@ -26076,6 +14558,7 @@ async def test_router_caller_entra_field_keeps_the_shared_deployment(): assert response._hidden_params["model_id"] == "shared-deployment" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") @pytest.mark.parametrize( ("litellm_params", "provider"), @@ -26237,21 +14720,6 @@ async def test_router_caller_entra_rejection_on_a_stream_leaves_the_shared_deplo with pytest.raises(litellm.AuthenticationError): await stream_request(router) -@pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") -def test_router_get_async_openai_model_client(): - router = Router( - model_list=[ - { - "model_name": "gemini/*", - "litellm_params": { - "model": "gemini/*", - "api_base": "https://api.gemini.com", - }, - } - ] - ) - model_client = router._get_async_openai_model_client(deployment=MagicMock(), kwargs={}) - assert model_client is None @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials(): @@ -26268,6 +14736,7 @@ def test_router_get_deployment_credentials(): assert credentials is not None assert credentials["api_key"] == "123" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials_with_provider(): """ @@ -26312,6 +14781,7 @@ def test_router_get_deployment_credentials_with_provider(): credentials3 = router.get_deployment_credentials_with_provider(model_id="non-existent") assert credentials3 is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_credentials_with_provider_wildcard(): """ @@ -26359,6 +14829,7 @@ def test_router_get_deployment_credentials_with_provider_wildcard(): credentials3 = router.get_deployment_credentials_with_provider(model_id="vertex_ai/gemini-pro") assert credentials3 is None + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_router_get_deployment_model_info(): router = Router( @@ -26373,12 +14844,14 @@ def test_router_get_deployment_model_info(): model_info = router.get_deployment_model_info(model_id="1", model_name="gemini/gemini-1.5-flash") assert model_info is not None + @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(): """ @@ -26400,6 +14873,7 @@ def setup_and_teardown_router_unit(): loop.close() asyncio.set_event_loop(None) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_acompletion_deployment_not_mutated(): @@ -26449,6 +14923,7 @@ async def test_acompletion_deployment_not_mutated(): assert deployment_after is not None assert deployment_after.litellm_params.model_dump() == original_params + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_completion_deployment_not_mutated(): """ @@ -26497,82 +14972,12 @@ def test_completion_deployment_not_mutated(): assert deployment_after is not None assert deployment_after.litellm_params.model_dump() == original_params -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_default_deployment_isolation(): - """ - Regression test for shallow copy optimization in _common_checks_available_deployment. - - When a model is not in model_names and default_deployment is set, the router - returns a copy of default_deployment with the model name updated. This test - ensures the optimization (shallow copy instead of deepcopy) properly isolates - each returned deployment from the original and from each other. - - The shallow copy optimization copies two levels: - 1. Top-level deployment dict - 2. litellm_params dict - - Deeper nested objects are intentionally shared for performance (safe because - the router only modifies the 'model' field at litellm_params level). - - Critical behavior verified: - 1. Each deployment gets independent model value - 2. Original default_deployment unchanged for litellm_params fields - 3. Shared fields (api_key) accessible in all copies - 4. Adding new litellm_params fields is isolated per deployment - 5. Deep nested objects ARE shared (acceptable trade-off) - """ - # Setup: Router with a default deployment (used for unknown models) - router = Router(model_list=[]) - - router.default_deployment = { # type: ignore - "model_name": "default-model", - "litellm_params": { - "model": "gpt-5-mini", # This will be overwritten per request - "api_key": "test-key", # This should be shared - "custom_config": { # Deep nested - will be SHARED - "nested_setting": "original", - }, - }, - } - - # Act: Request two different unknown models (triggers default deployment path) - _, deployment1 = router._common_checks_available_deployment( - model="custom-model-1", # Unknown model - messages=[{"role": "user", "content": "test"}], - ) - - _, deployment2 = router._common_checks_available_deployment( - model="custom-model-2", # Different unknown model - messages=[{"role": "user", "content": "test"}], - ) - - # Assert: Each deployment should have its own independent model value - assert deployment1["litellm_params"]["model"] == "custom-model-1" # type: ignore - assert deployment2["litellm_params"]["model"] == "custom-model-2" # type: ignore - - # Assert: Original default_deployment must remain unchanged (not mutated by requests) - assert router.default_deployment["litellm_params"]["model"] == "gpt-5-mini" # type: ignore - - # Assert: Shared fields should still be accessible in all copies - assert deployment1["litellm_params"]["api_key"] == "test-key" # type: ignore - assert deployment2["litellm_params"]["api_key"] == "test-key" # type: ignore - - # Assert: Modifying litellm_params in one deployment doesn't affect others - # This tests the shallow copy properly isolated the litellm_params dict level - deployment1["litellm_params"]["temperature"] = 0.9 # type: ignore - assert "temperature" not in deployment2["litellm_params"] # type: ignore - assert "temperature" not in router.default_deployment["litellm_params"] # type: ignore - - # Assert: Deep nested objects ARE shared (intentional trade-off for 100x perf gain) - # Safe because router only modifies top-level litellm_params fields - deployment1["litellm_params"]["custom_config"]["nested_setting"] = "modified" # type: ignore - assert deployment2["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore - assert router.default_deployment["litellm_params"]["custom_config"]["nested_setting"] == "modified" # type: ignore class NoItemsAliasDict(dict): def items(self): raise AssertionError("Unexpected full alias iteration via items()") + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup(): router = Router( @@ -26591,6 +14996,7 @@ def test_get_model_list_from_model_alias_should_not_iterate_for_non_alias_lookup model_alias_list = router.get_model_list_from_model_alias(model_name="gpt-5-mini") assert model_alias_list == [] + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") def test_map_team_model_should_not_iterate_aliases_for_non_alias_team_model_name(): router = Router( @@ -26821,62 +15227,10 @@ class TestPreCallChecksOptimization: assert deployments[0].get("model_info", {}).get("id") == "small", "First deployment ID changed!" assert deployments[1].get("model_info", {}).get("id") == "large", "Second deployment ID changed!" + if __name__ == "__main__": pytest.main([__file__, "-v"]) -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -def test_is_prompt_management_model_optimization(): - """ - Test early exit optimization works correctly for all cases. - - Optimization: Check if "/" in model name before calling expensive - get_model_list(). This short-circuits 99% of requests that use - standard model names like "gpt-5.5", "claude-3", etc. - - Tests both negative (early exit) and positive (actual detection) cases. - """ - import litellm - - # Test 1: Standard models without "/" -> early exit returns False - router = Router( - model_list=[ - { - "model_name": "gpt-5.5", - "litellm_params": {"model": "gpt-5.5"}, - }, - { - "model_name": "claude-3", - "litellm_params": {"model": "anthropic/claude-sonnet-4-5-20250929"}, - }, - ] - ) - - assert router._is_prompt_management_model("gpt-5.5") is False - assert router._is_prompt_management_model("claude-3") is False - - # Test 2: Models with "/" but not in model_list -> False after check - assert router._is_prompt_management_model("unknown/model") is False - - # Test 3: Actual prompt management models ARE detected (critical positive case) - original_callbacks = litellm._known_custom_logger_compatible_callbacks.copy() - if "langfuse_prompt" not in litellm._known_custom_logger_compatible_callbacks: - litellm._known_custom_logger_compatible_callbacks.append("langfuse_prompt") - - try: - router_with_prompt = Router( - model_list=[ - { - "model_name": "my-langfuse-prompt/test_id", - "litellm_params": {"model": "langfuse_prompt/actual_prompt_id"}, - }, - ] - ) - - # Critical: Must still detect prompt management models correctly - assert router_with_prompt._is_prompt_management_model("my-langfuse-prompt/test_id") is True - - finally: - litellm._known_custom_logger_compatible_callbacks = original_callbacks @pytest.fixture def router(): @@ -26893,6 +15247,7 @@ def router(): ] ) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch(router): @@ -26916,6 +15271,7 @@ async def test_router_acancel_batch(router): assert response.id == "batch_123" assert response.status == "cancelled" + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") @pytest.mark.asyncio async def test_router_acancel_batch_resolves_credential_name(): @@ -26956,114 +15312,11 @@ async def test_router_acancel_batch_resolves_credential_name(): finally: litellm.credential_list = [] -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -@pytest.mark.asyncio -async def test_router_acancel_batch_removes_unresolved_credential_name(): - router = Router( - model_list=[ - { - "model_name": "gpt-5.5", - "litellm_params": { - "model": "openai/gpt-5.5", - "litellm_credential_name": "missing-openai-credential", - }, - } - ] - ) - mock_response = MagicMock() - mock_response.id = "batch_123" - mock_response.status = "cancelled" - - with ( - patch.object(router, "_core_get_deployment_credentials_with_provider", return_value=None), - patch.object(litellm, "acancel_batch", new_callable=AsyncMock) as mock_cancel, - ): - mock_cancel.return_value = mock_response - - await router.acancel_batch( - model="gpt-5.5", - batch_id="batch_123", - ) - - call_kwargs = mock_cancel.call_args.kwargs - assert "litellm_credential_name" not in call_kwargs @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterEmbeddingHeaders: """Test that embedding methods properly propagate headers from router configuration.""" - def test_embedding_calls_update_kwargs_before_fallbacks(self): - """ - Test that router.embedding() calls _update_kwargs_before_fallbacks. - - This ensures that metadata is properly set up before the fallback mechanism, - which is necessary for header propagation to work correctly. - """ - model_list = [ - { - "model_name": "text-embedding-3-small", - "litellm_params": { - "model": "text-embedding-3-small", - "api_key": "fake-key", - }, - } - ] - - router = Router(model_list=model_list) - - # Mock the _update_kwargs_before_fallbacks method to verify it's called - with patch.object( - router, - "_update_kwargs_before_fallbacks", - wraps=router._update_kwargs_before_fallbacks, - ) as mock_update: - with patch("litellm.embedding") as mock_litellm_embedding: - mock_litellm_embedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2, 0.3]}]) - - router.embedding(model="text-embedding-3-small", input=["test input"]) - - # Verify _update_kwargs_before_fallbacks was called - mock_update.assert_called_once() - call_kwargs = mock_update.call_args[1] - assert call_kwargs["model"] == "text-embedding-3-small" - assert "kwargs" in call_kwargs - - @pytest.mark.asyncio - async def test_aembedding_calls_update_kwargs_before_fallbacks(self): - """ - Test that router.aembedding() calls _update_kwargs_before_fallbacks. - - This ensures consistency between sync and async embedding methods. - """ - model_list = [ - { - "model_name": "text-embedding-3-small", - "litellm_params": { - "model": "text-embedding-3-small", - "api_key": "fake-key", - }, - } - ] - - router = Router(model_list=model_list) - - # Mock the _update_kwargs_before_fallbacks method to verify it's called - with patch.object( - router, - "_update_kwargs_before_fallbacks", - wraps=router._update_kwargs_before_fallbacks, - ) as mock_update: - with patch("litellm.aembedding", new_callable=AsyncMock) as mock_litellm_aembedding: - mock_litellm_aembedding.return_value = MagicMock(data=[{"embedding": [0.1, 0.2, 0.3]}]) - - await router.aembedding(model="text-embedding-3-small", input=["test input"]) - - # Verify _update_kwargs_before_fallbacks was called - mock_update.assert_called_once() - call_kwargs = mock_update.call_args[1] - assert call_kwargs["model"] == "text-embedding-3-small" - assert "kwargs" in call_kwargs - def test_embedding_propagates_default_litellm_params(self): """ Test that embedding calls properly propagate default_litellm_params including headers. @@ -27302,6 +15555,7 @@ class TestRouterEmbeddingHeaders: assert "litellm_trace_id" in completion_kwargs assert "litellm_trace_id" in embedding_kwargs + if __name__ == "__main__": # Run a simple test test = TestRouterEmbeddingHeaders() @@ -27318,6 +15572,7 @@ OPENAI_EMBEDDINGS_URL = "https://api.openai.com/v1/embeddings" STORE_EMBEDDINGS_URL = "https://embedding.example/v1/embeddings" + def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route: return respx_mock.post(url).mock( return_value=httpx.Response( @@ -27331,11 +15586,13 @@ def _mock_embedding_route(respx_mock: respx.MockRouter, url: str) -> respx.Route ) ) + def _sent(route: respx.Route, index: int) -> tuple[str, str, list[str]]: request = route.calls[index].request body = json.loads(request.read()) return request.headers["authorization"], body["model"], body["input"] + def _alias_router() -> Router: return Router( model_list=[ @@ -27349,59 +15606,11 @@ def _alias_router() -> Router: ] ) + @pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") class TestRouterEmbeddingIntegration: """Integration tests for embedding with router configuration.""" - def test_vector_store_request_metadata_prefers_litellm_metadata(self): - assert Router._vector_store_request_metadata( - { - "litellm_metadata": {"user_api_key_team_id": "team-a"}, - "metadata": {"user_api_key_team_id": "team-b"}, - } - ) == {"user_api_key_team_id": "team-a"} - - assert Router._vector_store_request_metadata({"metadata": {"user_api_key_team_id": "team-b"}}) == { - "user_api_key_team_id": "team-b" - } - assert Router._vector_store_request_metadata({}) == {} - - def test_sync_vector_store_wrapper_injects_router_embedding_executor(self): - router = 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", - metadata={"user_api_key_team_id": "team-a"}, - ) - == "searched" - ) - - call_kwargs = original.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"} - - def test_sync_vector_store_wrapper_preserves_model_routing(self): - router = Router(model_list=[]) - original = MagicMock() - wrapped = router.factory_function(original, call_type="vector_store_search") - - with patch.object(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 - assert isinstance( - fallback.call_args.kwargs["_direct_vector_store_embedding_executor"], - RouterVectorStoreEmbeddingExecutor, - ) - @pytest.mark.asyncio async def test_vector_store_embedding_executors_cover_sdk_and_router_paths( self, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch @@ -27835,318 +16044,11 @@ class TestRouterEmbeddingIntegration: assert "headers" in call_kwargs assert call_kwargs["headers"]["X-Custom-Azure-Header"] == "azure-value" + if __name__ == "__main__": # Run tests pytest.main([__file__, "-v"]) -@pytest.mark.usefixtures("_vcr_outcome_gate_router_unit", "setup_and_teardown_router_unit") -class TestRouterIndexManagement: - """Test cases for router index management functions""" - - @pytest.fixture - def router(self): - """Create a router instance for testing""" - return Router(model_list=[]) - - def test_deletion_updates_model_name_indices(self, router): - """Test that deleting a deployment updates model_name_to_deployment_indices correctly""" - router.model_list = [ - {"model_name": "gpt-3.5", "model_info": {"id": "model-1"}}, - {"model_name": "gpt-5.5", "model_info": {"id": "model-2"}}, - {"model_name": "gpt-5.5", "model_info": {"id": "model-3"}}, - {"model_name": "claude", "model_info": {"id": "model-4"}}, - ] - router.model_id_to_deployment_index_map = { - "model-1": 0, - "model-2": 1, - "model-3": 2, - "model-4": 3, - } - router.model_name_to_deployment_indices = { - "gpt-3.5": [0], - "gpt-5.5": [1, 2], - "claude": [3], - } - - # Remove one of the duplicate gpt-5.5 deployments - router._update_deployment_indices_after_removal(model_id="model-2", removal_idx=1) - - # Verify indices are shifted correctly - assert router.model_name_to_deployment_indices["gpt-3.5"] == [0] - assert router.model_name_to_deployment_indices["gpt-5.5"] == [1] # was [1,2], removed 1, shifted 2->1 - assert router.model_name_to_deployment_indices["claude"] == [2] # was [3], shifted to [2] - - # Remove the last gpt-5.5 deployment - router._update_deployment_indices_after_removal(model_id="model-3", removal_idx=1) - - # Verify gpt-5.5 is removed from dict when no deployments remain - assert "gpt-5.5" not in router.model_name_to_deployment_indices - assert router.model_name_to_deployment_indices["gpt-3.5"] == [0] - assert router.model_name_to_deployment_indices["claude"] == [1] - - def test_build_model_id_to_deployment_index_map(self, router): - """Test _build_model_id_to_deployment_index_map function""" - model_list = [ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini"}, - "model_info": {"id": "model-1"}, - }, - { - "model_name": "gpt-5.5", - "litellm_params": {"model": "gpt-5.5"}, - "model_info": {"id": "model-2"}, - }, - ] - - # Test: Build index from model list - router._build_model_id_to_deployment_index_map(model_list) - - # Verify: model_list is populated - assert len(router.model_list) == 2 - # Verify: model_id_to_deployment_index_map is correctly built - assert router.model_id_to_deployment_index_map["model-1"] == 0 - assert router.model_id_to_deployment_index_map["model-2"] == 1 - - def test_add_model_to_list_and_index_map_from_model_info(self, router): - """Test _add_model_to_list_and_index_map extracting model_id from model_info""" - # Setup: Empty router - router.model_list = [] - router.model_id_to_deployment_index_map = {} - - # Test: Add model without explicit model_id - model = {"model": "test-model", "model_info": {"id": "model-info-id"}} - router._add_model_to_list_and_index_map(model=model) - - # Verify: Model added to list - assert len(router.model_list) == 1 - assert router.model_list[0] == model - - # Verify: Index map uses model_info.id - assert router.model_id_to_deployment_index_map["model-info-id"] == 0 - - def test_add_model_to_list_and_index_map_multiple_models(self, router): - """Test _add_model_to_list_and_index_map with multiple models to verify indexing""" - # Setup: Empty router - router.model_list = [] - router.model_id_to_deployment_index_map = {} - - # Test: Add multiple models - model1 = {"model": "model1", "model_info": {"id": "id-1"}} - model2 = {"model": "model2", "model_info": {"id": "id-2"}} - model3 = {"model": "model3", "model_info": {"id": "id-3"}} - - router._add_model_to_list_and_index_map(model=model1, model_id="id-1") - router._add_model_to_list_and_index_map(model=model2, model_id="id-2") - router._add_model_to_list_and_index_map(model=model3, model_id="id-3") - - # Verify: All models added to list - assert len(router.model_list) == 3 - assert router.model_list[0] == model1 - assert router.model_list[1] == model2 - assert router.model_list[2] == model3 - - # Verify: Correct indices in map - assert router.model_id_to_deployment_index_map["id-1"] == 0 - assert router.model_id_to_deployment_index_map["id-2"] == 1 - assert router.model_id_to_deployment_index_map["id-3"] == 2 - - def test_update_team_model_index(self, router): - """Test _update_team_model_index updates team_model_to_deployment_indices.""" - model = { - "model_name": "team-alias", - "model_info": { - "id": "dep-1", - "team_id": "team-abc", - "team_public_model_name": "gpt-5.5", - }, - } - router._update_team_model_index(model, 0) - assert router.team_model_to_deployment_indices[("team-abc", "gpt-5.5")] == [0] - router._update_team_model_index(model, 2) - assert router.team_model_to_deployment_indices[("team-abc", "gpt-5.5")] == [0, 2] - - router._update_team_model_index({"model_name": "x", "model_info": {"id": "dep-2"}}, 5) - assert router.team_model_to_deployment_indices == { - ("team-abc", "gpt-5.5"): [0, 2], - } - - def test_has_model_id(self, router): - """Test has_model_id function for O(1) membership check""" - # Setup: Add models to router - router.model_list = [ - {"model": "test1", "model_info": {"id": "model-1"}}, - {"model": "test2", "model_info": {"id": "model-2"}}, - {"model": "test3", "model_info": {"id": "model-3"}}, - ] - router.model_id_to_deployment_index_map = { - "model-1": 0, - "model-2": 1, - "model-3": 2, - } - - # Test: Check existing model IDs - assert router.has_model_id("model-1") == True - assert router.has_model_id("model-2") == True - assert router.has_model_id("model-3") == True - - # Test: Check non-existing model IDs - assert router.has_model_id("non-existent") == False - assert router.has_model_id("") == False - assert router.has_model_id("model-4") == False - - # Test: Empty router - empty_router = Router(model_list=[]) - assert empty_router.has_model_id("any-id") == False - - def test_build_model_name_index(self, router): - """Test _build_model_name_index function""" - model_list = [ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini"}, - "model_info": {"id": "model-1"}, - }, - { - "model_name": "gpt-5.5", - "litellm_params": {"model": "gpt-5.5"}, - "model_info": {"id": "model-2"}, - }, - { - "model_name": "gpt-5.5", # Duplicate model_name, different deployment - "litellm_params": {"model": "gpt-5.5"}, - "model_info": {"id": "model-3"}, - }, - ] - - # Test: Build index from model list - router._build_model_name_index(model_list) - - # Verify: model_name_to_deployment_indices is correctly built - assert "gpt-5-mini" in router.model_name_to_deployment_indices - assert "gpt-5.5" in router.model_name_to_deployment_indices - - # Verify: gpt-5-mini has single deployment - assert router.model_name_to_deployment_indices["gpt-5-mini"] == [0] - - # Verify: gpt-5.5 has multiple deployments - assert router.model_name_to_deployment_indices["gpt-5.5"] == [1, 2] - - # Test: Rebuild index (should clear and rebuild) - new_model_list = [ - { - "model_name": "claude-3", - "litellm_params": {"model": "claude-3"}, - "model_info": {"id": "model-4"}, - }, - ] - router._build_model_name_index(new_model_list) - - # Verify: Old entries are cleared - assert "gpt-5-mini" not in router.model_name_to_deployment_indices - assert "gpt-5.5" not in router.model_name_to_deployment_indices - - # Verify: New entry is added - assert "claude-3" in router.model_name_to_deployment_indices - assert router.model_name_to_deployment_indices["claude-3"] == [0] - - def test_no_linear_scans_in_router(self): - """ - Static analysis test to ensure Router doesn't use O(n) linear scans. - - Scans router.py for 'in self.model_list' pattern which indicates - inefficient O(n) iteration instead of using index-based O(1) lookups. - - Methods should use: - - model_id_to_deployment_index_map for O(1) model_id lookups - - model_name_to_deployment_indices for O(1) + O(k) model_name lookups - """ - # Methods that are allowed to iterate through self.model_list - ALLOWED_METHODS = { - "_get_deployment_by_litellm_model": "lookup by litellm_params.model, which is not indexed", - "_finalize_adaptive_router_if_configured": 'init-time prefix scan for "auto_router/adaptive_router"; no index for prefix match', - "config_deployments": "filters the whole list on model_info.db_model; admin path only (model add/upsert)", - "auto_router_capability_violation": "counts gated auto-routers across the whole list; admin path only (auto-router init/upsert)", - } - - # Get path to router.py - router_file = os.path.join( - os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__)))), - "litellm", - "router.py", - ) - - # Read the file - with open(router_file, "r") as f: - content = f.read() - - # Parse with AST - tree = ast.parse(content) - - # Find violations - violations = [] - ignore_methods = set(ALLOWED_METHODS) - - for node in ast.walk(tree): - if isinstance(node, ast.FunctionDef): - method_name = node.name - - # Skip ignored methods - if method_name in ignore_methods: - continue - - # Get source for this method - try: - method_source = ast.get_source_segment(content, node) - if not method_source: - continue - - # Check for the anti-pattern: "in self.model_list" - # This catches: for x in self.model_list, if x in self.model_list, etc. - if "in self.model_list" in method_source: - # Extract the specific line for better error reporting - lines = method_source.split("\n") - pattern_line = None - for line in lines: - if "in self.model_list" in line: - pattern_line = line.strip() - break - - violations.append( - { - "method": method_name, - "line": node.lineno, - "pattern": pattern_line or "in self.model_list", - } - ) - except Exception: - # Skip if we can't get source segment - pass - - # Assert no violations - if violations: - error_msg = "\n".join([f" - {v['method']}() at line {v['line']}: {v['pattern']}" for v in violations]) - - pytest.fail( - f"\n{'=' * 70}\n" - f"Found O(n) linear scan pattern in router.py:\n\n" - f"{error_msg}\n\n" - f"These methods should use index maps instead:\n" - f" - model_id_to_deployment_index_map (for model_id lookups)\n" - f" - model_name_to_deployment_indices (for model_name lookups)\n\n" - f"If a method legitimately needs O(n) iteration, add it to\n" - f"ALLOWED_METHODS in this test method.\n" - f"{'=' * 70}\n" - ) - - def test_model_names_is_set(self): - """Verify that model_names uses a set for O(1) lookups, not a list (O(n))""" - router = Router(model_list=[]) - - assert isinstance(router.model_names, set), ( - f"model_names should be a set for O(1) lookups, but got {type(router.model_names)}" - ) - _PRIORITY_CHAT_COMPLETION: Final = { "id": "chatcmpl-priority", @@ -28599,7 +16501,6 @@ def test_router_context_window_check_pre_call_check_in_group_custom_model_info() """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -28653,7 +16554,6 @@ def test_router_context_window_check_pre_call_check(): """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -28707,7 +16607,6 @@ def test_router_context_window_check_pre_call_check_out_group(): """ import os - litellm.set_verbose = False print(f"len(text): {len(text)}") @@ -28733,7 +16632,13 @@ def test_router_context_window_check_pre_call_check_out_group(): }, ] - router = Router(model_list=model_list, set_verbose=True, enable_pre_call_checks=True, num_retries=0, context_window_fallbacks=[{"gpt-3.5-turbo-small": ["gpt-3.5-turbo-large"]}]) # type: ignore + router = Router( + model_list=model_list, + set_verbose=True, + enable_pre_call_checks=True, + num_retries=0, + context_window_fallbacks=[{"gpt-3.5-turbo-small": ["gpt-3.5-turbo-large"]}], + ) # type: ignore response = router.completion( model="gpt-3.5-turbo-small", @@ -28829,13 +16734,9 @@ def test_router_region_pre_call_check(allowed_model_region): if allowed_model_region is None: assert len(_healthy_deployments) == 2 else: - assert len(_healthy_deployments) == 1, "{} models selected as healthy".format( - len(_healthy_deployments) - ) - assert ( - _healthy_deployments[0]["model_info"]["id"] == "1" - ), "Incorrect model id picked. Got id={}, expected id=1".format( - _healthy_deployments[0]["model_info"]["id"] + assert len(_healthy_deployments) == 1, "{} models selected as healthy".format(len(_healthy_deployments)) + assert _healthy_deployments[0]["model_info"]["id"] == "1", ( + "Incorrect model id picked. Got id={}, expected id=1".format(_healthy_deployments[0]["model_info"]["id"]) ) @@ -29008,104 +16909,12 @@ def test_router_context_window_pre_call_check(model, base_model, llm_provider): else: pytest.fail("Expected to raise an error. Got={}".format(updated_list)) except Exception as e: - if ( - llm_provider == "azure" and base_model is not None - ) or llm_provider == "openai": + if (llm_provider == "azure" and base_model is not None) or llm_provider == "openai": pass except Exception as e: pytest.fail(f"Got unexpected exception on router! - {str(e)}") -@pytest.fixture -def restore_retry_after_header_parser(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr( - litellm.utils, "_get_retry_after_from_exception_header", litellm.utils._get_retry_after_from_exception_header - ) - - -@pytest.mark.usefixtures("restore_retry_after_header_parser") -def test_router_dynamic_cooldown_correct_retry_after_time(): - """ - User feedback: litellm says "No deployments available for selected model, Try again in 60 seconds" - but Azure says to retry in at most 9s - - ``` - {"message": "litellm.proxy.proxy_server.embeddings(): Exception occured - No deployments available for selected model, Try again in 60 seconds. Passed model=text-embedding-ada-002. pre-call-checks=False, allowed_model_region=n/a, cooldown_list=[('b49cbc9314273db7181fe69b1b19993f04efb88f2c1819947c538bac08097e4c', {'Exception Received': 'litellm.RateLimitError: AzureException RateLimitError - Requests to the Embeddings_Create Operation under Azure OpenAI API version 2023-09-01-preview have exceeded call rate limit of your current OpenAI S0 pricing tier. Please retry after 9 seconds. Please go here: https://aka.ms/oai/quotaincrease if you would like to further increase the default rate limit.', 'Status Code': '429'})]", "level": "ERROR", "timestamp": "2024-08-22T03:25:36.900476"} - ``` - """ - router = Router( - model_list=[ - { - "model_name": "text-embedding-ada-002", - "litellm_params": { - "model": "openai/text-embedding-ada-002", - }, - } - ] - ) - - openai_client = openai.OpenAI(api_key="") - - cooldown_time = 30 - - def _return_exception(*args, **kwargs): - from httpx import Headers, Request, Response - - kwargs = { - "request": Request("POST", "https://www.google.com"), - "message": "Error code: 429 - Rate Limit Error!", - "body": {"detail": "Rate Limit Error!"}, - "code": None, - "param": None, - "type": None, - "response": Response( - status_code=429, - headers=Headers( - { - "date": "Sat, 21 Sep 2024 22:56:53 GMT", - "server": "uvicorn", - "retry-after": f"{cooldown_time}", - "content-length": "30", - "content-type": "application/json", - } - ), - request=Request("POST", "http://0.0.0.0:9000/chat/completions"), - ), - "status_code": 429, - "request_id": None, - } - - exception = Exception() - for k, v in kwargs.items(): - setattr(exception, k, v) - raise exception - - with patch.object( - openai_client, - "post", - side_effect=_return_exception, - ): - new_retry_after_mock_client = MagicMock(return_value=-1) - - litellm.utils._get_retry_after_from_exception_header = ( - new_retry_after_mock_client - ) - - try: - router.embedding( - model="text-embedding-ada-002", - input="Hello world!", - client=openai_client, - ) - except litellm.RateLimitError: - pass - - new_retry_after_mock_client.assert_called() - - response_headers: httpx.Headers = new_retry_after_mock_client.call_args[0][0] - assert int(response_headers["retry-after"]) == cooldown_time - - @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_aaarouter_dynamic_cooldown_message_retry_time(sync_mode): @@ -29175,13 +16984,9 @@ async def test_aaarouter_dynamic_cooldown_message_retry_time(sync_mode): ) if sync_mode: - cooldown_deployments = get_cooldown_deployments( - ctx=router._core_ctx, parent_otel_span=None - ) + cooldown_deployments = get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) else: - cooldown_deployments = await async_get_cooldown_deployments( - ctx=router._core_ctx, parent_otel_span=None - ) + cooldown_deployments = await async_get_cooldown_deployments(ctx=router._core_ctx, parent_otel_span=None) assert len(cooldown_deployments) > 0 @@ -29282,93 +17087,6 @@ def test_model_group_alias(hidden: bool) -> None: assert len(model_names) == len(_model_list) + 1 -def test_get_team_specific_model() -> None: - """ - Test that _get_team_specific_model returns: - - team_public_model_name when team_id matches - - None when team_id doesn't match - - None when no team_id in model_info - """ - router = Router(model_list=[]) - - deployment = DeploymentTypedDict( - model_name="model-x", - litellm_params={}, - model_info=ModelInfo(team_id="team1", team_public_model_name="public-model-x"), - ) - assert router._get_team_specific_model(deployment, "team1") == "public-model-x" - - assert router._get_team_specific_model(deployment, "team2") is None - - deployment = DeploymentTypedDict( - model_name="model-y", - litellm_params={}, - model_info=ModelInfo(team_public_model_name="public-model-y"), - ) - assert router._get_team_specific_model(deployment, "team1") is None - - deployment = DeploymentTypedDict(model_name="model-z", litellm_params={}, model_info=ModelInfo()) - assert router._get_team_specific_model(deployment, "team1") is None - - -def test_is_team_specific_model() -> None: - """ - Test that _is_team_specific_model returns: - - True when model_info contains team_id - - False when model_info doesn't contain team_id - - False when model_info is None - """ - router = Router(model_list=[]) - - model_info = ModelInfo(team_id="team1", team_public_model_name="public-model-x") - assert router._is_team_specific_model(model_info) is True - - model_info = ModelInfo(team_public_model_name="public-model-y") - assert router._is_team_specific_model(model_info) is False - - model_info = ModelInfo() - assert router._is_team_specific_model(model_info) is False - - assert router._is_team_specific_model(None) is False - - -def test_router_prompt_management_factory(): - router = Router( - model_list=[ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo"}, - }, - { - "model_name": "chatbot_actions", - "litellm_params": { - "model": "langfuse/openai-gpt-3.5-turbo", - "tpm": 1000000, - "prompt_id": "jokes", - }, - }, - { - "model_name": "openai-gpt-3.5-turbo", - "litellm_params": { - "model": "openai/gpt-3.5-turbo", - "api_key": os.getenv("OPENAI_API_KEY"), - }, - }, - ] - ) - - assert router._is_prompt_management_model("chatbot_actions") is True - assert router._is_prompt_management_model("openai-gpt-3.5-turbo") is False - - response = router._prompt_management_factory( - model="chatbot_actions", - messages=[{"role": "user", "content": "Hello world!"}], - kwargs={}, - ) - - print(response) - - def test_router_get_model_list_from_model_alias() -> None: router = Router( model_list=[ @@ -29533,39 +17251,6 @@ async def test_enforce_model_rate_limits_returns_429_at_deployment_rpm(): assert rate_limited.value.status_code == 429 -@pytest.mark.parametrize( - "startup_nodes, expected_cache_type", - [ - pytest.param( - [dict(host="node1.localhost", port=6379)], - RedisClusterCache, - id="Expects a RedisClusterCache instance when startup_nodes provided", - ), - pytest.param( - None, - RedisCache, - id="Expects a RedisCache instance when there is no startup nodes", - ), - ], -) -def test_create_correct_redis_cache_instance( - startup_nodes: list[dict] | None, - expected_cache_type: type[RedisClusterCache] | type[RedisCache], -): - cache_config = dict( - host="mockhost", - port=6379, - password="mock-password", - startup_nodes=startup_nodes, - ) - - 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) - - class CompletionCustomHandler( CustomLogger ): # https://docs.litellm.ai/docs/observability/custom_callback#callback-class @@ -29614,12 +17299,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception as e: print(f"Assertion Error: {traceback.format_exc()}") @@ -29645,9 +17326,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.iscoroutine(kwargs["original_response"]) or inspect.isasyncgen(kwargs["original_response"]) ) @@ -29659,12 +17338,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception: print(f"Assertion Error: {traceback.format_exc()}") @@ -29681,23 +17356,18 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponseStream) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -29718,22 +17388,17 @@ class CompletionCustomHandler( assert isinstance(response_obj, litellm.ModelResponse) ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) - assert isinstance(kwargs["api_key"], (str, type(None))) - assert isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) ) + assert isinstance(kwargs["api_key"], (str, type(None))) + assert isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) assert isinstance(kwargs["additional_args"], (dict, type(None))) assert isinstance(kwargs["log_event_type"], str) assert kwargs["cache_hit"] is None or isinstance(kwargs["cache_hit"], bool) @@ -29752,23 +17417,18 @@ class CompletionCustomHandler( assert response_obj == None ## KWARGS assert isinstance(kwargs["model"], str) - assert isinstance(kwargs["messages"], list) and isinstance( - kwargs["messages"][0], dict - ) + assert isinstance(kwargs["messages"], list) and isinstance(kwargs["messages"][0], dict) assert isinstance(kwargs["optional_params"], dict) assert isinstance(kwargs["litellm_params"], dict) assert isinstance(kwargs["start_time"], (datetime, type(None))) assert isinstance(kwargs["stream"], bool) assert isinstance(kwargs["user"], (str, type(None))) - assert ( - isinstance(kwargs["input"], list) - and isinstance(kwargs["input"][0], dict) - ) or isinstance(kwargs["input"], (dict, str)) + assert (isinstance(kwargs["input"], list) and isinstance(kwargs["input"][0], dict)) or isinstance( + kwargs["input"], (dict, str) + ) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or kwargs["original_response"] == None ) assert isinstance(kwargs["additional_args"], (dict, type(None))) @@ -29801,29 +17461,19 @@ class CompletionCustomHandler( ## END TIME assert isinstance(end_time, datetime) ## RESPONSE OBJECT - assert isinstance( - response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse) - ) + assert isinstance(response_obj, (litellm.ModelResponse, litellm.EmbeddingResponse)) ## KWARGS assert isinstance(kwargs["model"], str) # checking we use base_model for azure cost calculation - base_model = litellm.utils.get_base_model_from_metadata( - model_call_details=kwargs - ) + base_model = litellm.utils.get_base_model_from_metadata(model_call_details=kwargs) - if ( - kwargs["model"] == "chatgpt-v-3" - and base_model is not None - and kwargs["stream"] != True - ): + if kwargs["model"] == "chatgpt-v-3" and base_model is not None and kwargs["stream"] != True: # when base_model is set for azure, we should use pricing for the base_model # this checks response_cost == litellm.cost_per_token(model=base_model) assert isinstance(kwargs["response_cost"], float) response_cost = kwargs["response_cost"] - print( - f"response_cost: {response_cost}, for model: {kwargs['model']} and base_model: {base_model}" - ) + print(f"response_cost: {response_cost}, for model: {kwargs['model']} and base_model: {base_model}") prompt_tokens = response_obj.usage.prompt_tokens completion_tokens = response_obj.usage.completion_tokens # ensure the pricing is based on the base_model here @@ -29834,9 +17484,9 @@ class CompletionCustomHandler( ) expected_price = prompt_price + completion_price print(f"expected price: {expected_price}") - assert ( - response_cost == expected_price - ), f"response_cost: {response_cost} != expected_price: {expected_price}. For model: {kwargs['model']} and base_model: {base_model}. should have used base_model for price" + assert response_cost == expected_price, ( + f"response_cost: {response_cost} != expected_price: {expected_price}. For model: {kwargs['model']} and base_model: {base_model}. should have used base_model for price" + ) assert isinstance(kwargs["messages"], list) assert isinstance(kwargs["optional_params"], dict) @@ -29847,9 +17497,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, dict, str)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) ) @@ -29862,12 +17510,8 @@ class CompletionCustomHandler( assert isinstance(kwargs["litellm_params"]["metadata"]["deployment"], str) assert isinstance(kwargs["litellm_params"]["model_info"], dict) assert isinstance(kwargs["litellm_params"]["model_info"]["id"], str) - assert isinstance( - kwargs["litellm_params"]["proxy_server_request"], (str, type(None)) - ) - assert isinstance( - kwargs["litellm_params"]["preset_cache_key"], (str, type(None)) - ) + assert isinstance(kwargs["litellm_params"]["proxy_server_request"], (str, type(None))) + assert isinstance(kwargs["litellm_params"]["preset_cache_key"], (str, type(None))) assert isinstance(kwargs["litellm_params"]["stream_response"], dict) except Exception: print(f"Assertion Error: {traceback.format_exc()}") @@ -29894,9 +17538,7 @@ class CompletionCustomHandler( assert isinstance(kwargs["input"], (list, str, dict)) assert isinstance(kwargs["api_key"], (str, type(None))) assert ( - isinstance( - kwargs["original_response"], (str, litellm.CustomStreamWrapper) - ) + isinstance(kwargs["original_response"], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs["original_response"]) or inspect.iscoroutine(kwargs["original_response"]) or kwargs["original_response"] == None @@ -30052,27 +17694,6 @@ class SpecialEndpointRoutingStrategy(CustomRoutingStrategyBase): pass -def test_reset_custom_routing_strategy(): - """ - Setting a custom routing strategy installs instance-level overrides for - get_available_deployment / async_get_available_deployment. Re-initializing the - routing strategy must clear them so the class implementations are used again. - """ - router = _create_custom_routing_router() - router.set_custom_routing_strategy(SpecialEndpointRoutingStrategy(router)) - - assert "get_available_deployment" in router._core.__dict__ - assert "_core_async_get_available_deployment" in router._core.__dict__ - - router._reset_custom_routing_strategy() - - assert "get_available_deployment" not in router._core.__dict__ - assert "_core_async_get_available_deployment" not in router._core.__dict__ - assert router.async_get_available_deployment.__wrapped__.__func__ is Router.async_get_available_deployment - - router._reset_custom_routing_strategy() - - @pytest.mark.asyncio async def test_select_picks_through_a_custom_routing_strategy() -> None: router: Final = _create_custom_routing_router() @@ -30182,9 +17803,7 @@ def test_multiple_deployments_route_to_a_configured_model( }, ) - route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( - side_effect=_deployment_response - ) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=_deployment_response) router: Final = Router( model_list=[ { @@ -30333,9 +17952,7 @@ def test_router_retries_transient_server_error( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response(500, json={"error": {"message": "temporary failure"}}), _openai_chat_response("recovered"), @@ -30368,9 +17985,7 @@ def test_router_authentication_error_is_not_swallowed( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response( 401, @@ -30468,9 +18083,9 @@ def test_router_reads_api_key_from_model_list( 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=[_openai_chat_response(), _openai_chat_stream_response()]) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + side_effect=[_openai_chat_response(), _openai_chat_stream_response()] + ) router: Final = Router( model_list=[ { @@ -30493,9 +18108,7 @@ def test_router_reads_api_key_from_model_list( messages=[{"role": "user", "content": "stream key from model list"}], stream=True, ) - streamed_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in stream - ) + streamed_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in stream) assert response.choices[0].message.content == "scripted reply" assert streamed_content == "streamed reply" @@ -30511,9 +18124,9 @@ def test_router_calls_requested_specific_deployment( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock(return_value=_openai_chat_response()) + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=_openai_chat_response() + ) router: Final = Router( model_list=[ { @@ -30555,9 +18168,7 @@ async def test_router_context_window_error_uses_configured_fallback( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response( 400, @@ -30600,9 +18211,7 @@ async def test_router_context_window_error_uses_configured_fallback( assert response.choices[0].message.content == "fallback response" assert route.call_count == 2 - assert tuple( - json.loads(call.request.content)["model"] for call in route.calls - ) == ("gpt-4o-mini", "gpt-4.1-mini") + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("gpt-4o-mini", "gpt-4.1-mini") def test_router_function_calling_sends_tools_and_returns_tool_call( @@ -30610,9 +18219,7 @@ def test_router_function_calling_sends_tools_and_returns_tool_call( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( return_value=httpx.Response( 200, json={ @@ -30691,9 +18298,7 @@ async def test_router_openai_completion_supports_sync_async_and_streaming( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post( - "https://api.openai.com/v1/chat/completions" - ).mock( + route: Final = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ _openai_chat_response("async reply"), _openai_chat_response("sync reply"), @@ -30728,17 +18333,13 @@ async def test_router_openai_completion_supports_sync_async_and_streaming( stream=True, ) async_stream_chunks: Final = tuple([chunk async for chunk in async_stream]) - async_stream_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in async_stream_chunks - ) + async_stream_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in async_stream_chunks) sync_stream: Final = router.completion( model="openai-model", messages=[{"role": "user", "content": "sync stream"}], stream=True, ) - sync_stream_content: Final = "".join( - chunk.choices[0].delta.content or "" for chunk in sync_stream - ) + sync_stream_content: Final = "".join(chunk.choices[0].delta.content or "" for chunk in sync_stream) assert async_response.choices[0].message.content == "async reply" assert sync_response.choices[0].message.content == "sync reply" @@ -30782,9 +18383,7 @@ async def test_router_moderation_returns_provider_result( num_retries=0, ) - response: Final = await router.amoderation( - model="moderation-model", input="safe scripted text" - ) + response: Final = await router.amoderation(model="moderation-model", input="safe scripted text") assert response.results[0].flagged is False assert json.loads(route.calls[0].request.content)["input"] == "safe scripted text" @@ -30953,9 +18552,7 @@ async def test_batch_completion_multiple_models_returns_each_model_response( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock( - side_effect=_batch_response - ) + route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock(side_effect=_batch_response) router: Final = _batch_router() responses: Final = cast( @@ -30978,9 +18575,7 @@ async def test_batch_completion_multiple_models_and_messages_preserves_each_pair respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock( - side_effect=_batch_response - ) + route: Final = respx_mock.post("https://batch-router.test/v1/chat/completions").mock(side_effect=_batch_response) router: Final = _batch_router() responses: Final = cast( @@ -31061,9 +18656,7 @@ async def test_batch_completion_fastest_response_streams_mocked_chunks() -> None ) chunks: Final = tuple([chunk async for chunk in response]) text: Final = "".join( - chunk.choices[0].delta.content or "" - for chunk in chunks - if chunk.choices and chunk.choices[0].delta is not None + chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices and chunk.choices[0].delta is not None ) assert text == "stream answer" @@ -31219,7 +18812,9 @@ def _first_chunk_cases(): from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices chunk = ModelResponseStream(choices=[StreamingChoices(delta=Delta(content="hi"))]) - provider_error = litellm.InternalServerError(message="primary died at stream start", llm_provider="openai", model="m") + provider_error = litellm.InternalServerError( + message="primary died at stream start", llm_provider="openai", model="m" + ) wrapped = MidStreamFallbackError( message=str(provider_error), model="m", @@ -31236,7 +18831,7 @@ def _first_chunk_cases(): @pytest.mark.parametrize("first, expected, raised", _first_chunk_cases()) def test_read_first_chunk_direct_call(first, expected, raised): - from litellm.router import _read_first_chunk + from litellm.router_rust.invoke.streaming import _read_first_chunk # noqa: PLC2701 stream = _FirstPullStream(first) if raised is None: @@ -31251,7 +18846,7 @@ def test_read_first_chunk_direct_call(first, expected, raised): @pytest.mark.asyncio @pytest.mark.parametrize("first, expected, raised", _first_chunk_cases()) async def test_aread_first_chunk_direct_call(first, expected, raised): - from litellm.router import _aread_first_chunk + from litellm.router_rust.invoke.streaming import _aread_first_chunk # noqa: PLC2701 stream = _FirstPullStream(first) if raised is None: @@ -31263,31 +18858,6 @@ async def test_aread_first_chunk_direct_call(first, expected, raised): assert caught.value is raised -def _chunks_with_first_cases(): - return [ - pytest.param(None, ["a", "b"], ["a", "b"], id="nothing-read-stream-as-is"), - pytest.param((None, ["a", "b"]), ["a", "b"], [], id="empty-stream-not-touched"), - pytest.param(("first", ["a", "b"]), ["never-iterated"], ["first", "a", "b"], id="first-chunk-put-back"), - ] - - -@pytest.mark.parametrize("taken, stream, expected", _chunks_with_first_cases()) -def test_chunks_with_first_direct_call(taken, stream, expected): - from litellm.router import _chunks_with_first, _FirstChunk - - first = None if taken is None else _FirstChunk(chunk=taken[0], rest=iter(taken[1])) - assert list(_chunks_with_first(iter(stream), first)) == expected - - -@pytest.mark.asyncio -@pytest.mark.parametrize("taken, stream, expected", _chunks_with_first_cases()) -async def test_achunks_with_first_direct_call(taken, stream, expected): - from litellm.router import _achunks_with_first, _FirstChunk - - first = None if taken is None else _FirstChunk(chunk=taken[0], rest=_AsyncList(taken[1])) - assert [item async for item in _achunks_with_first(_AsyncList(stream), first)] == expected - - class _IterationHandedOffStream(CustomStreamWrapper): """A CustomStreamWrapper that, like the MCP gateway's MCPStreamWrapper, has no completion_stream of its own and hands iteration to another object from __aiter__ / __iter__, a fresh one per call so a second @@ -31505,30 +19075,9 @@ def test_pre_call_checks_keeps_deployment_whose_unsupported_param_is_allowed_or_ assert [d["model_info"]["id"] for d in filtered] == ["only"] -def test_pre_call_checks_keeps_retryable_429_when_a_cooling_down_deployment_supports_the_param(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "drop_params", False) - model_list: Final = [ - { - "model_name": "mixed-json-mode", - "litellm_params": {"model": f"openai/{deployment_model}", "api_key": "my-fake-key"}, - "model_info": {"id": deployment_model, "max_input_tokens": 100}, - } - for deployment_model in ("gpt-3.5-turbo-16k", "gpt-4o") - ] - router: Final = Router(model_list=model_list, enable_pre_call_checks=True, num_retries=0) - without_json_mode: Final = [d for d in router.model_list if d["model_info"]["id"] == "gpt-3.5-turbo-16k"] - - filtered: Final = router._pre_call_checks( - model="mixed-json-mode", - healthy_deployments=without_json_mode, - messages=[{"role": "user", "content": "hi"}], - request_kwargs={"response_format": {"type": "json_object"}}, - ) - - assert filtered == [] - - -def test_completion_uses_context_window_fallback_when_the_small_deployment_also_lacks_the_param(monkeypatch: pytest.MonkeyPatch): +def test_completion_uses_context_window_fallback_when_the_small_deployment_also_lacks_the_param( + monkeypatch: pytest.MonkeyPatch, +): monkeypatch.setattr(litellm, "drop_params", False) router: Final = Router( model_list=[ @@ -31597,7 +19146,9 @@ def test_pre_call_checks_prefers_a_deployment_that_takes_the_param_over_dropping assert [d["model_info"]["id"] for d in filtered] == ["gpt-4o"] -def test_pre_call_checks_drops_the_param_when_the_capable_deployment_is_outside_the_region(monkeypatch: pytest.MonkeyPatch): +def test_pre_call_checks_drops_the_param_when_the_capable_deployment_is_outside_the_region( + monkeypatch: pytest.MonkeyPatch, +): monkeypatch.setattr(litellm, "drop_params", False) model_list: Final = [ { @@ -31827,106 +19378,6 @@ async def test_router_configured_stream_usage_keeps_a_choice_on_the_usage_chunk( assert all(len(chunk.choices) == 1 for chunk in chunks[:-1]) -@pytest.mark.parametrize( - ("stream_options", "marked_injected"), - [({"include_usage": True}, True), (None, False)], - ids=["include-usage", "no-stream-options"], -) -def test_update_kwargs_with_deployment_reads_stream_options_from_pydantic_litellm_params( - stream_options: dict[str, bool] | None, marked_injected: bool -): - from litellm.types.litellm_params import STREAM_USAGE_INJECTED_FIELD - from litellm.types.router import LiteLLM_Params - - kwargs: Final[dict[str, object]] = {"stream": True, "metadata": {}} - Router()._update_kwargs_with_deployment( - deployment={ - "model_name": "gpt-5-mini", - "litellm_params": LiteLLM_Params(model="openai/gpt-5-mini", api_key="fake-key", stream_options=stream_options), - "model_info": {"id": "pydantic-params"}, - }, - kwargs=kwargs, - ) - - assert (kwargs.get(STREAM_USAGE_INJECTED_FIELD) is True) is marked_injected - - -@pytest.mark.asyncio -async def test_responses_fallback_wrapper_disconnect_billing_keeps_the_primary_streams_tokens(): - """After a mid-stream hop, a disconnect bills the primary's relayed text plus the fallback's, not the fallback alone.""" - import json - - from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig - from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator - - def _relayed_iterator(text: str) -> ResponsesAPIStreamingIterator: - logging_obj: Final = MagicMock(spec=LiteLLMLogging) - logging_obj.model_call_details = {"litellm_params": {}} - logging_obj.completion_start_time = None - iterator: Final = ResponsesAPIStreamingIterator( - response=httpx.Response(200), - model="gpt-5.4", - responses_api_provider_config=OpenAIResponsesAPIConfig(), - logging_obj=logging_obj, - request_data={"model": "gpt-5.4", "input": "Hi"}, - custom_llm_provider="openai", - ) - created: Final = { - "id": "resp_1", - "object": "response", - "created_at": 1700000000.0, - "model": "gpt-5.4", - "status": "in_progress", - "output": [], - "parallel_tool_calls": True, - "tool_choice": "auto", - "tools": [], - } - iterator._process_chunk(json.dumps({"type": "response.created", "sequence_number": 0, "response": created})) - iterator._process_chunk( - json.dumps( - { - "type": "response.output_text.delta", - "sequence_number": 1, - "item_id": "msg_1", - "output_index": 0, - "content_index": 0, - "delta": text, - } - ) - ) - return iterator - - router: Final = Router( - model_list=[{"model_name": "primary", "litellm_params": {"model": "openai/gpt-5.4", "api_key": "k"}}] - ) - primary: Final = _relayed_iterator("Once upon a time there was a lighthouse keeper") - fallback: Final = _relayed_iterator("who lit the lamp") - wrapped: Final = await router._aresponses_streaming_iterator( - response=primary, - initial_kwargs={"model": "primary", "stream": True, "input": "Hi", "original_generic_function": AsyncMock()}, - ) - primary_tokens: Final = primary.partial_response_on_disconnect().response.usage.output_tokens - fallback_tokens: Final = fallback.partial_response_on_disconnect().response.usage.output_tokens - - assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens - - wrapped.adopt_fallback_headers(fallback) - - assert wrapped.partial_response_on_disconnect().response.usage.output_tokens == primary_tokens + fallback_tokens - - -_SYNC_SECONDS: Final = 10 -_ROUTER_FILE: Final = Router._update_deployment_indices_after_removal.__code__.co_filename -_INDEX_LOCK_LINES: Final = frozenset( - number - for number, line in enumerate(linecache.getlines(_ROUTER_FILE), start=1) - if "with self._deployment_index_lock" in line -) - -_Read = TypeVar("_Read") - - def _concurrent_delete_deployment(index: int) -> dict[str, object]: return { "model_name": f"model-{index}", @@ -31943,272 +19394,6 @@ def _concurrent_delete_router() -> Router: return Router(model_list=[_concurrent_delete_deployment(index) for index in range(4)]) -def _read_during_held_write( - read: Callable[[], _Read], - write: Callable[[], object], - event: Literal["call", "c_call"] = "call", - write_after: CodeType | None = None, -) -> _Read: - """Run `read` on its own thread while `write` waits at its first `event` in the re-index, after the list shrank and - before the maps caught up. With `write_after`, the write starts when `read` returns from that function. The write - resumes once `read` returns or reaches the router's index lock, so a read that skips the lock meets the half-updated - list and a read that takes the lock waits for the write""" - reindex: Final = Router._update_deployment_indices_after_removal.__code__ - held: Final = threading.Event() - release: Final = threading.Event() - settled: Final = threading.Event() - released_in_time: Final[list[bool]] = [] - outcome: Final[list[_Read]] = [] - failure: Final[list[BaseException]] = [] - - def hold_in_reindex(frame: FrameType, observed: str, _arg: object) -> None: - if observed == event and frame.f_code is reindex and not held.is_set(): - held.set() - released_in_time.append(release.wait(timeout=_SYNC_SECONDS)) - - def run_write() -> None: - previous: Final = sys.getprofile() - sys.setprofile(hold_in_reindex) - try: - write() - finally: - sys.setprofile(previous) - - writer: Final = threading.Thread(target=run_write) - - def start_write() -> None: - writer.start() - assert held.wait(timeout=_SYNC_SECONDS), "the write never reached re-indexing" - - def watch_read(frame: FrameType, observed: str, _arg: object) -> Callable[..., object] | None: - if frame.f_code.co_filename != _ROUTER_FILE: - return None - if observed == "return" and frame.f_code is write_after and writer.ident is None: - start_write() - elif observed == "line" and held.is_set() and frame.f_lineno in _INDEX_LOCK_LINES: - settled.set() - return watch_read - - def run_read() -> None: - previous: Final = sys.gettrace() - sys.settrace(watch_read) - try: - outcome.append(read()) - except BaseException as error: - failure.append(error) - finally: - sys.settrace(previous) - settled.set() - - reader: Final = threading.Thread(target=run_read) - try: - if write_after is None: - start_write() - reader.start() - assert settled.wait(timeout=_SYNC_SECONDS), "the read neither returned nor reached the index lock" - finally: - release.set() - if reader.ident is not None: - reader.join() - if writer.ident is not None: - writer.join() - assert released_in_time == [True], "the write was not held for the whole read" - if failure: - raise failure[0] - return outcome[0] - - -def _deployment_ids(deployments: list[DeploymentTypedDict]) -> tuple[str, ...]: - return tuple(deployment["model_info"]["id"] for deployment in deployments) - - -@pytest.mark.parametrize( - ("read", "expected"), - [ - pytest.param( - lambda router: router.get_deployment("deployment-3").model_info.id, "deployment-3", id="get_deployment" - ), - pytest.param( - lambda router: router.get_model_info("deployment-3")["model_info"]["id"], - "deployment-3", - id="get_model_info", - ), - pytest.param( - lambda router: _deployment_ids(router.get_all_deployments("model-3")), - ("deployment-3",), - id="get_all_deployments", - ), - pytest.param( - lambda router: _deployment_ids(router.get_all_deployments("public-shared", team_id="team-1")), - ("deployment-2", "deployment-3"), - id="get_all_deployments-team", - ), - pytest.param( - lambda router: router.map_team_model("public-shared", "team-1"), "public-shared", id="map_team_model" - ), - pytest.param(lambda router: router.get_model_ids("model-3"), ["deployment-3"], id="get_model_ids-by-name"), - pytest.param( - lambda router: router.get_model_ids(), ["deployment-1", "deployment-2", "deployment-3"], id="get_model_ids" - ), - pytest.param( - lambda router: router.get_deployment_by_model_group_name("model-3").model_info.id, - "deployment-3", - id="get_deployment_by_model_group_name", - ), - pytest.param( - lambda router: router.get_model_listing_info("model-3").cost_map_keys, - ("openai/gpt-3",), - id="get_model_listing_info", - ), - pytest.param( - lambda router: router.team_model_has_alternatives("deployment-3"), True, id="team_model_has_alternatives" - ), - pytest.param( - lambda router: ( - router._get_model_group_deployment_usable_by_team( - model_group_name="model-3", team_id="team-1" - ).model_info.id - ), - "deployment-3", - id="model_group_deployment_usable_by_team", - ), - pytest.param( - lambda router: ( - router._get_team_public_name_deployment(model_id="public-shared", team_id="team-1").model_info.id - ), - "deployment-2", - id="team_public_name_deployment", - ), - pytest.param( - lambda router: _deployment_ids(router._team_deployments_across_teams("public-shared")), - ("deployment-2", "deployment-3"), - id="team_deployments_across_teams", - ), - pytest.param( - lambda router: _deployment_ids(list(router.deployments_named("model-3") or ())), - ("deployment-3",), - id="deployments_named", - ), - pytest.param( - lambda router: should_use_polling_for_request( - background_mode=True, - polling_via_cache_enabled=["openai"], - redis_cache=MagicMock(), - model="model-3", - llm_router=router, - ), - True, - id="response_polling_provider_lookup", - ), - pytest.param( - lambda router: _models_in(router._core_ctx, "model-3"), - (Baseline("openai/gpt-3", "deployment-3"),), - id="savings_baseline_candidates", - ), - ], -) -def test_lookup_during_a_concurrent_delete_returns_the_requested_deployment( - read: Callable[[Router], object], expected: object -) -> None: - router: Final = _concurrent_delete_router() - - observed: Final = _read_during_held_write(lambda: read(router), lambda: router.delete_deployment("deployment-0")) - - assert observed == expected - - -def test_lookup_during_an_upsert_returns_the_requested_deployment() -> None: - router: Final = _concurrent_delete_router() - rotated: Final = Deployment( - **{**_concurrent_delete_deployment(0), "litellm_params": {"model": "openai/gpt-0", "api_key": "rotated"}} - ) - - observed: Final = _read_during_held_write( - lambda: router.get_deployment("deployment-3").model_info.id, lambda: router.upsert_deployment(rotated) - ) - - assert observed == "deployment-3" - - -def test_add_during_a_concurrent_delete_indexes_the_added_deployment() -> None: - router: Final = _concurrent_delete_router() - - _read_during_held_write( - lambda: router.add_deployment(Deployment(**_concurrent_delete_deployment(4))), - lambda: router.delete_deployment("deployment-0"), - ) - - assert [router.get_deployment(f"deployment-{index}").model_info.id for index in range(1, 5)] == [ - f"deployment-{index}" for index in range(1, 5) - ] - - -def test_set_model_list_during_a_concurrent_delete_indexes_the_new_list() -> None: - router: Final = _concurrent_delete_router() - - _read_during_held_write( - lambda: router.set_model_list([_concurrent_delete_deployment(index) for index in range(10, 14)]), - lambda: router.delete_deployment("deployment-0"), - event="c_call", - ) - - assert router.get_model_ids() == [f"deployment-{index}" for index in range(10, 14)] - - -def test_team_model_has_alternatives_when_a_delete_starts_after_its_deployment_lookup() -> None: - router: Final = _concurrent_delete_router() - - has_alternatives: Final = _read_during_held_write( - lambda: router.team_model_has_alternatives("deployment-3"), - lambda: router.delete_deployment("deployment-0"), - write_after=Router._core_get_deployment.__code__, - ) - - assert has_alternatives is True - - -def test_init_adaptive_router_deployment_during_a_concurrent_delete_prices_the_listed_model() -> None: - priced: Final = { - **_concurrent_delete_deployment(3), - "litellm_params": {"model": "openai/gpt-3", "api_key": "k", "input_cost_per_token": 3e-06}, - } - router: Final = Router(model_list=[*(_concurrent_delete_deployment(index) for index in range(3)), priced]) - adaptive: Final = Deployment( - model_name="adaptive", - litellm_params=LiteLLM_Params( - model="auto_router/adaptive_router", adaptive_router_config={"available_models": ["model-3"]} - ), - model_info=ModelInfo(id="adaptive"), - ) - - _read_during_held_write( - lambda: router.init_adaptive_router_deployment(deployment=adaptive), - lambda: router.delete_deployment("deployment-0"), - ) - - assert router.adaptive_routers["adaptive"][0].strategy.model_to_cost == {"model-3": 3e-06} - - -def test_complexity_router_adaptive_pool_during_a_concurrent_delete_prices_the_listed_model() -> None: - priced: Final = { - **_concurrent_delete_deployment(3), - "litellm_params": {"model": "openai/gpt-3", "api_key": "k", "input_cost_per_token": 3e-06}, - } - router: Final = Router(model_list=[*(_concurrent_delete_deployment(index) for index in range(3)), priced]) - complexity: Final = ComplexityRouter( - model_name="hybrid", - ctx=router._core_ctx, - complexity_router_config={"adaptive": True, "tiers": {"SIMPLE": ["model-3"]}}, - ) - - adaptive: Final = _read_during_held_write( - complexity._ensure_adaptive_router, lambda: router.delete_deployment("deployment-0") - ) - - assert adaptive is not None - assert adaptive.model_to_cost == {"model-3": 3e-06} - - def test_delete_of_an_unknown_deployment_returns_none_and_keeps_the_list() -> None: router: Final = _concurrent_delete_router() @@ -32240,7 +19425,10 @@ _IMAGE_RESPONSES_INPUT = [ ], } ] -_VISION_LEG: Final[dict[str, object]] = {"litellm_params": {"model": "openai/gpt-5.6-terra"}, "model_info": {"id": "vision"}} +_VISION_LEG: Final[dict[str, object]] = { + "litellm_params": {"model": "openai/gpt-5.6-terra"}, + "model_info": {"id": "vision"}, +} _DECLARED_TEXT_LEG: Final[dict[str, object]] = { "litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, "model_info": {"id": "declared-text", "supports_vision": False}, @@ -32254,17 +19442,6 @@ def _mixed_router(*deployments: dict[str, object], routing_strategy: str = "late ) -def _candidate_ids( - router: litellm.Router, - messages: list[dict[str, object]] | None = None, - request_kwargs: dict[str, object] | None = None, -) -> list[str]: - candidates: Final = router._without_declared_text_only_deployments( - router._core_get_model_list(model_name="mixed"), "mixed", messages, request_kwargs - ) - return [d["model_info"]["id"] for d in candidates] - - def test_router_picks_the_vision_leg_for_image_requests_without_enable_pre_call_checks(): router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) assert router.enable_pre_call_checks is False @@ -32302,81 +19479,6 @@ def test_router_sends_an_image_to_the_text_only_leg_when_the_vision_leg_is_cooli assert router.get_available_deployment(model="mixed", messages=_IMAGE_CHAT)["model_info"]["id"] == "declared-text" -def test_without_declared_text_only_deployments_passes_a_chosen_deployment_through(): - assert litellm.Router._without_declared_text_only_deployments(_DECLARED_TEXT_LEG, "mixed", _IMAGE_CHAT, None) is _DECLARED_TEXT_LEG - - -def test_image_routing_drop_declared_text_only_legs_for_image_requests(): - router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) - assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision"] - - -def test_image_routing_keep_declared_text_only_legs_for_text_requests(): - router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) - assert _candidate_ids(router, messages=[{"role": "user", "content": "Say hi"}]) == [ - "vision", - "declared-text", - ] - - -def test_image_routing_drop_declared_text_only_legs_for_responses_input_image(): - router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) - assert _candidate_ids(router, request_kwargs={"input": _IMAGE_RESPONSES_INPUT}) == ["vision"] - - -def test_image_routing_send_image_to_the_group_when_every_leg_is_declared_text_only(): - other_text_leg = {**_DECLARED_TEXT_LEG, "model_info": {"id": "declared-text-2", "supports_vision": False}} - router: Final = _mixed_router(_DECLARED_TEXT_LEG, other_text_leg) - assert _candidate_ids(router, messages=_IMAGE_CHAT) == [ - "declared-text", - "declared-text-2", - ] - - -def test_image_routing_keep_undeclared_custom_legs_for_image_requests(): - custom_leg = { - "litellm_params": {"model": "hosted_vllm/Qwen2.5-VL-7B-Instruct", "api_base": "http://localhost:8000"}, - "model_info": {"id": "custom"}, - } - router: Final = _mixed_router(_VISION_LEG, custom_leg) - assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision", "custom"] - - -def test_image_routing_resolve_vision_through_the_cost_map_and_base_model(): - catalog_text_leg = {"litellm_params": {"model": "bedrock/deepseek.v3.2"}, "model_info": {"id": "catalog-text"}} - azure_text_leg = { - "litellm_params": {"model": "azure/my-audio-deployment", "api_base": "https://x.openai.azure.com"}, - "model_info": {"id": "azure-text", "base_model": "azure/gpt-audio"}, - } - router: Final = _mixed_router(_VISION_LEG, catalog_text_leg, azure_text_leg) - assert _candidate_ids(router, messages=_IMAGE_CHAT) == ["vision"] - - -@pytest.mark.parametrize( - ("deployment", "expected"), - [ - (_DECLARED_TEXT_LEG, True), - ({"litellm_params": {"model": "bedrock/openai.gpt-oss-120b-1:0"}, "model_info": {"supports_vision": True}}, False), - ({"litellm_params": {"model": "bedrock/deepseek.v3.2"}, "model_info": {}}, True), - ({"litellm_params": {"model": "azure/my-deploy"}, "model_info": {"base_model": "azure/gpt-audio"}}, True), - ({"litellm_params": {"model": "hosted_vllm/Qwen2.5-VL-7B-Instruct"}, "model_info": {}}, False), - ({"litellm_params": {"model": "no-such-model-anywhere"}, "model_info": {}}, False), - ({"litellm_params": {}, "model_info": {}}, False), - ], - ids=["override-false", "override-true", "cost-map-false", "base-model-false", "undeclared", "unresolvable", "no-params"], -) -def test_deployment_declares_no_vision_reads_the_override_then_the_cost_map(deployment: dict[str, object], expected: bool): - assert litellm.Router._deployment_declares_no_vision(deployment, "mixed") is expected - - -def test_request_carries_image_sees_chat_parts_and_responses_input_items(): - assert litellm.Router._request_carries_image(_IMAGE_CHAT, {}) is True - assert litellm.Router._request_carries_image(None, {"input": _IMAGE_RESPONSES_INPUT}) is True - assert litellm.Router._request_carries_image([{"role": "user", "content": "Say hi"}], {}) is False - assert litellm.Router._request_carries_image(None, {"input": "Say hi"}) is False - assert litellm.Router._request_carries_image(None, {"input": ["safe text", "more safe text"]}) is False - - @pytest.mark.asyncio @pytest.mark.usefixtures("isolate_litellm_state") @pytest.mark.parametrize("surface", ["chat", "messages"]) @@ -32398,19 +19500,31 @@ async def test_anthropic_caller_rate_limit_does_not_block_other_credentials( headers={"retry-after": "168307"}, json={"type": "error", "error": {"type": "rate_limit_error", "message": "usage limit"}}, ) - return httpx.Response(200, json={ - "id": "msg-healthy", "type": "message", "role": "assistant", "model": "claude-sonnet-5-5", - "content": [{"type": "text", "text": "healthy account"}], "stop_reason": "end_turn", - "usage": {"input_tokens": 2, "output_tokens": 2}, - }) + return httpx.Response( + 200, + json={ + "id": "msg-healthy", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "healthy account"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 2, "output_tokens": 2}, + }, + ) route: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=upstream) router: Final = Router( - model_list=[{ - "model_name": "claude", "litellm_params": { - "model": "anthropic/claude-sonnet-5-5", "api_key": "fake-gateway-key", - }, "model_info": {"id": "shared-anthropic"}, - }], + model_list=[ + { + "model_name": "claude", + "litellm_params": { + "model": "anthropic/claude-sonnet-5-5", + "api_key": "fake-gateway-key", + }, + "model_info": {"id": "shared-anthropic"}, + } + ], num_retries=0, allowed_fails_policy=AllowedFailsPolicy(RateLimitErrorAllowedFails=1), ) @@ -32418,22 +19532,36 @@ async def test_anthropic_caller_rate_limit_does_not_block_other_credentials( async def call(account: str | None) -> object: credentials: Final = ( - {"provider_specific_header": { - "custom_llm_provider": "anthropic", "extra_headers": { - "Authorization": f"Bearer sk-ant-oat-fake-{account}", - }, - }} if credential_kind == "oauth" else {"api_key": f"fake-{account}"} - ) if account is not None else {} + ( + { + "provider_specific_header": { + "custom_llm_provider": "anthropic", + "extra_headers": { + "Authorization": f"Bearer sk-ant-oat-fake-{account}", + }, + } + } + if credential_kind == "oauth" + else {"api_key": f"fake-{account}"} + ) + if account is not None + else {} + ) if surface == "messages": max_tokens: Final = litellm.get_max_tokens("anthropic/claude-sonnet-5-5") assert max_tokens is not None return await router.aanthropic_messages( - model="claude", messages=[{"role": "user", "content": "hello"}], - client=client, max_tokens=max_tokens, **credentials, + model="claude", + messages=[{"role": "user", "content": "hello"}], + client=client, + max_tokens=max_tokens, + **credentials, ) return await router.acompletion( - model="claude", messages=[{"role": "user", "content": "hello"}], - client=client, **credentials, + model="claude", + messages=[{"role": "user", "content": "hello"}], + client=client, + **credentials, ) try: @@ -32671,7 +19799,9 @@ async def test_a_client_model_id_cannot_move_a_sagemaker_embedding_call_off_the_ served: Final = component_embeddings[str(request["InferenceComponentName"])] return { "Body": io.BytesIO( - json.dumps({"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": served}]}).encode() + json.dumps( + {"object": "list", "data": [{"object": "embedding", "index": 0, "embedding": served}]} + ).encode() ) } @@ -32704,7 +19834,9 @@ async def test_a_header_pin_holds_when_the_deployment_names_sagemaker_through_cu async def test_a_component_header_on_a_non_sagemaker_deployment_leaves_the_clients_model_id_alone( respx_mock: respx.MockRouter, _router_httpx_transport: None ) -> None: - route: Final = respx_mock.post("https://gateway.local/v1/chat/completions").respond(json=dict(_CHAT_COMPLETION_REPLY)) + route: Final = respx_mock.post("https://gateway.local/v1/chat/completions").respond( + json=dict(_CHAT_COMPLETION_REPLY) + ) router: Final = Router( model_list=[ { @@ -32719,7 +19851,9 @@ async def test_a_component_header_on_a_non_sagemaker_deployment_leaves_the_clien ] ) - await router.acompletion(model="gateway-model", messages=[{"role": "user", "content": "hi"}], model_id="client-value") + await router.acompletion( + model="gateway-model", messages=[{"role": "user", "content": "hi"}], model_id="client-value" + ) assert json.loads(route.calls.last.request.content)["model_id"] == "client-value" @@ -32745,65 +19879,3 @@ def test_router_warns_when_a_deployment_enforces_endpoints_without_a_list(caplog assert "Deployment 'embed-from-config' sets enforce_supported_endpoints without" in caplog.text assert "Deployment 'embed-from-api' sets enforce_supported_endpoints without" in caplog.text - - -@pytest.mark.asyncio -@pytest.mark.parametrize("mode", ["sync", "async"]) -@pytest.mark.parametrize("metadata_key", [None, "metadata", "litellm_metadata"]) -@pytest.mark.parametrize( - ("route", "messages", "request_kwargs"), - [ - ("/v1/moderations", None, {"input": ["first", "second"]}), - ("/v1/embeddings", None, {"messages": []}), - ("/v1/messages", ["malformed"], {}), - ("/v1/messages", "malformed", {}), - ("/v1/messages", 5, {}), - ("/v1/messages", None, {"messages": ["malformed"]}), - ("/v1/responses", None, {"input": [{"type": "computer_call_output", "output": {"type": "text", "text": "no image"}}]}), - ("/v1/responses", None, {"input": [{"type": "computer_call_output", "output": "malformed"}]}), - ], - ids=["moderation-list", "embedding-empty-messages", "message-entry", "message-string", "message-number", "kwargs-messages", "computer-output-text", "computer-output-string"], -) -async def test_image_routing_leaves_non_image_and_malformed_requests_to_the_endpoint( - mode: str, metadata_key: str | None, route: str, messages: object, request_kwargs: dict[str, object] -) -> None: - router: Final = _mixed_router( - {**_DECLARED_TEXT_LEG, "litellm_params": {**_DECLARED_TEXT_LEG["litellm_params"], "order": 0}}, - {**_VISION_LEG, "litellm_params": {**_VISION_LEG["litellm_params"], "order": 1}}, - ) - kwargs: Final = { - **request_kwargs, - **({metadata_key: {"user_api_key_request_route": route}} if metadata_key is not None else {}), - } - original: Final = copy.deepcopy(request_kwargs) - deployment: Final = ( - await router._core_async_get_available_deployment(model="mixed", messages=messages, request_kwargs=kwargs) - if mode == "async" - else router.get_available_deployment(model="mixed", messages=messages, request_kwargs=kwargs) - ) - assert deployment["model_info"]["id"] == "declared-text" - assert all(kwargs[key] == value for key, value in original.items()) - - -@pytest.mark.parametrize( - "request_kwargs", - [ - {"messages": _IMAGE_CHAT}, - {"input": _IMAGE_RESPONSES_INPUT}, - {"messages": [{"role": "user", "content": [{"type": "image", "source": {"type": "url", "url": "https://example.invalid/image.png"}}]}]}, - {"messages": [{"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_image", "content": [_IMAGE_PART]}]}]}, - {"input": [{"type": "function_call_output", "call_id": "call_image", "output": [{"type": "input_image", "image_url": "https://example.invalid/image.png"}]}]}, - {"input": [{"type": "custom_tool_call_output", "call_id": "call_image", "output": [{"type": "input_image", "image_url": "https://example.invalid/image.png"}]}]}, - {"input": [{"type": "computer_call_output", "call_id": "call_image", "output": {"type": "computer_screenshot", "image_url": "https://example.invalid/image.png"}}]}, - {"input": [{"role": "user", "content": [{"type": "computer_screenshot", "image_url": "https://example.invalid/image.png"}]}]}, - {"input": [{"type": "function_call_output", "call_id": "call_image", "output": [{"type": "computer_screenshot", "image_url": "https://example.invalid/image.png"}]}]}, - {"input": [{"type": "custom_tool_call_output", "call_id": "call_image", "output": [{"type": "computer_screenshot", "image_url": "https://example.invalid/image.png"}]}]}, - {"input": [{"type": "computer_call_output", "call_id": "call_image", "output": {"type": "computer_screenshot", "url": "https://example.invalid/image.png"}}]}, - ], - ids=["chat", "responses", "anthropic", "tool-result", "function-output", "custom-output", "computer-output", "screenshot-content", "screenshot-function-output", "screenshot-custom-output", "screenshot-url"], -) -def test_image_routing_inspects_native_request_content_without_translation(request_kwargs: dict[str, object]) -> None: - router: Final = _mixed_router(_VISION_LEG, _DECLARED_TEXT_LEG) - original: Final = copy.deepcopy(request_kwargs) - assert _candidate_ids(router, request_kwargs=request_kwargs) == ["vision"] - assert request_kwargs == original diff --git a/tests/unit/test_router/test_router_endpoints.py b/tests/unit/test_router/test_router_endpoints.py index 0c57fb662ac..d59c3f19c50 100644 --- a/tests/unit/test_router/test_router_endpoints.py +++ b/tests/unit/test_router/test_router_endpoints.py @@ -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() diff --git a/tests/unit/test_router/test_router_entry_points.py b/tests/unit/test_router/test_router_entry_points.py index 25dd8fe47cf..c3d5d962123 100644 --- a/tests/unit/test_router/test_router_entry_points.py +++ b/tests/unit/test_router/test_router_entry_points.py @@ -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): diff --git a/tests/unit/test_router/test_router_facade.py b/tests/unit/test_router/test_router_facade.py index 0539ebe43b9..07d85d36e47 100644 --- a/tests/unit/test_router/test_router_facade.py +++ b/tests/unit/test_router/test_router_facade.py @@ -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") diff --git a/tests/unit/test_router/test_router_fallbacks.py b/tests/unit/test_router/test_router_fallbacks.py index b25d5c35892..937f3f40e5b 100644 --- a/tests/unit/test_router/test_router_fallbacks.py +++ b/tests/unit/test_router/test_router_fallbacks.py @@ -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" diff --git a/tests/unit/test_router/test_router_fuzz_regressions_rust_only.py b/tests/unit/test_router/test_router_fuzz_regressions_rust_only.py new file mode 100644 index 00000000000..8551e6ea884 --- /dev/null +++ b/tests/unit/test_router/test_router_fuzz_regressions_rust_only.py @@ -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"} diff --git a/tests/unit/test_router/test_router_get_deployments.py b/tests/unit/test_router/test_router_get_deployments.py index 07486762a64..acfcb999c36 100644 --- a/tests/unit/test_router/test_router_get_deployments.py +++ b/tests/unit/test_router/test_router_get_deployments.py @@ -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( diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 3bbb69413d9..32b38030498 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -1,21 +1,18 @@ import asyncio -import functools import gc import json import os -import uuid +import time import weakref -from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping from datetime import datetime, timezone -from types import MethodType from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from websockets.exceptions import ConnectionClosed -from websockets.frames import Close +from pydantic import BaseModel import litellm from litellm import Router @@ -24,18 +21,11 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.litellm_core_utils.realtime_errors import close_after_upstream_handshake_refusal -from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig -from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook -from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute from litellm.types.litellm_params import RoutingStrategyName from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo from litellm.types.utils import ( - CallTypes, ModelResponse, StandardLoggingHiddenParams, StandardLoggingMetadata, @@ -208,29 +198,6 @@ def test_print_deployment_with_redact_enabled(model_list): litellm.redact_user_api_key_info = original_setting -def test_completion(model_list): - """Test if the completion function is working correctly""" - router = Router(model_list=model_list) - response = router._completion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "Hello, how are you?"}], - mock_response="I'm fine, thank you!", - ) - assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!" - - -@pytest.mark.asyncio -async def test_router_acompletion_util(model_list): - """Test if the underlying '_acompletion' function is working correctly""" - router = Router(model_list=model_list) - response = await router._acompletion( - model="gpt-5-mini", - messages=[{"role": "user", "content": "Hello, how are you?"}], - mock_response="I'm fine, thank you!", - ) - assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!" - - @pytest.mark.asyncio async def test_router_abatch_completion_one_model_multiple_requests_util(model_list): """Test if the 'abatch_completion_one_model_multiple_requests' function is working correctly""" @@ -261,193 +228,6 @@ async def test_router_schedule_acompletion(model_list): assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!" -@pytest.mark.asyncio -async def test_router_schedule_atext_completion(model_list): - """Test if the 'schedule_atext_completion' function is working correctly""" - from litellm.types.utils import TextCompletionResponse - - router = Router(model_list=model_list) - with patch.object(router, "_atext_completion", AsyncMock()) as mock_atext_completion: - mock_atext_completion.return_value = TextCompletionResponse() - response = await router.atext_completion( - model="gpt-5-mini", - prompt="Hello, how are you?", - priority=1, - ) - mock_atext_completion.assert_awaited_once() - assert "priority" not in mock_atext_completion.call_args.kwargs - - -@pytest.mark.asyncio -async def test_router_schedule_factory(model_list): - """Test if the 'schedule_atext_completion' function is working correctly""" - from litellm.types.utils import TextCompletionResponse - - router = Router(model_list=model_list) - with patch.object(router, "_atext_completion", AsyncMock()) as mock_atext_completion: - mock_atext_completion.return_value = TextCompletionResponse() - response = await router._schedule_factory( - model="gpt-5-mini", - args=( - "gpt-5-mini", - "Hello, how are you?", - ), - priority=1, - kwargs={}, - original_function=router.atext_completion, - ) - mock_atext_completion.assert_awaited_once() - assert "priority" not in mock_atext_completion.call_args.kwargs - - -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.asyncio -async def test_router_function_with_fallbacks(model_list, sync_mode): - """Test if the router 'async_function_with_fallbacks' + 'function_with_fallbacks' are working correctly""" - router = Router(model_list=model_list) - data = { - "model": "gpt-5-mini", - "messages": [{"role": "user", "content": "Hello, how are you?"}], - "mock_response": "I'm fine, thank you!", - "num_retries": 0, - } - if sync_mode: - response = router.function_with_fallbacks( - original_function=router._completion, - **data, - ) - else: - response = await router.async_function_with_fallbacks( - original_function=router._acompletion, - **data, - ) - assert response.choices[0].message.content == "I'm fine, thank you!" - - -@pytest.mark.parametrize("sync_mode", [True, False]) -@pytest.mark.asyncio -async def test_router_function_with_retries(model_list, sync_mode): - """Test if the router 'async_function_with_retries' + 'function_with_retries' are working correctly""" - router = Router(model_list=model_list) - data = { - "model": "gpt-5-mini", - "messages": [{"role": "user", "content": "Hello, how are you?"}], - "mock_response": "I'm fine, thank you!", - "num_retries": 0, - } - response = await router.async_function_with_retries( - original_function=router._acompletion, - **data, - ) - - assert response.choices[0].message.content == "I'm fine, thank you!" - - -@pytest.mark.asyncio -async def test_router_make_call(model_list): - """Test if the router 'make_call' function is working correctly""" - - router = Router(model_list=model_list) - response = await router.make_call( - original_function=router._acompletion, - model="gpt-5-mini", - messages=[{"role": "user", "content": "Hello, how are you?"}], - mock_response="I'm fine, thank you!", - ) - assert response.choices[0].message.content == "I'm fine, thank you!" - - response = await router.make_call( - original_function=router._atext_completion, - model="gpt-5-mini", - prompt="Hello, how are you?", - mock_response="I'm fine, thank you!", - ) - assert response.choices[0].text == "I'm fine, thank you!" - - response = await router.make_call( - original_function=router._aembedding, - model="gpt-5-mini", - input="Hello, how are you?", - mock_response=[0.1, 0.2, 0.3], - ) - assert response.data[0].embedding == [0.1, 0.2, 0.3] - - response = await router.make_call( - original_function=router._aimage_generation, - model="gpt-image-1", - prompt="A cute baby sea otter", - mock_response="https://example.com/image.png", - ) - assert response.data[0].url == "https://example.com/image.png" - - -def test_update_kwargs_with_deployment(model_list): - """Test if the '_update_kwargs_with_deployment' function is working correctly""" - router = Router(model_list=model_list) - kwargs: dict = {"metadata": {}} - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") - router._update_kwargs_with_deployment( - deployment=deployment, - kwargs=kwargs, - ) - set_fields = ["deployment", "api_base", "model_info"] - assert all(field in kwargs["metadata"] for field in set_fields) - - -def test_update_kwargs_with_default_litellm_params(model_list): - """Test if the '_update_kwargs_with_default_litellm_params' function is working correctly""" - router = Router( - model_list=model_list, - default_litellm_params={"api_key": "test", "metadata": {"key": "value"}}, - ) - kwargs: dict = {"metadata": {"key2": "value2"}} - router._update_kwargs_with_default_litellm_params(kwargs=kwargs) - assert kwargs["api_key"] == "test" - assert kwargs["metadata"]["key"] == "value" - assert kwargs["metadata"]["key2"] == "value2" - - -def test_get_timeout(model_list): - """Test if the '_get_timeout' function is working correctly""" - router = Router(model_list=model_list) - timeout = router._get_timeout(kwargs={}, data={"timeout": 100}) - assert timeout == 100 - - -@pytest.mark.parametrize( - "fallback_kwarg, expected_error", - [ - ("mock_testing_fallbacks", litellm.InternalServerError), - ("mock_testing_context_fallbacks", litellm.ContextWindowExceededError), - ("mock_testing_content_policy_fallbacks", litellm.ContentPolicyViolationError), - ], -) -def test_handle_mock_testing_fallbacks(model_list, fallback_kwarg, expected_error): - """Test if the '_handle_mock_testing_fallbacks' function is working correctly""" - router = Router(model_list=model_list) - data = { - fallback_kwarg: True, - } - - with pytest.raises(expected_error): - router._handle_mock_testing_fallbacks( - kwargs=data, - ) - - -def test_handle_mock_testing_rate_limit_error(model_list): - """Test if the '_handle_mock_testing_rate_limit_error' function is working correctly""" - router = Router(model_list=model_list) - data = { - "mock_testing_rate_limit_error": True, - } - - with pytest.raises(litellm.RateLimitError): - router._handle_mock_testing_rate_limit_error( - kwargs=data, - ) - - @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_deployment_callback_on_success(sync_mode): @@ -467,9 +247,7 @@ async def test_deployment_callback_on_success(sync_mode): ] router = Router(model_list=model_list) # Get the actual deployment ID that was generated - gpt_deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + gpt_deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") deployment_id = gpt_deployment["model_info"]["id"] standard_logging_payload = create_standard_logging_payload() @@ -549,447 +327,6 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn( assert await router.cache.async_get_cache(key=tpm_key) == 17 -class _DiscardingLoggingWorker: - def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None: - async_coroutine.close() - - -async def _client_never_sends() -> str: - await asyncio.Event().wait() - raise AssertionError("unreachable") - - -def _realtime_client_ws() -> MagicMock: - client_ws: Final = MagicMock() - client_ws.headers = {} - client_ws.receive_text = AsyncMock(side_effect=_client_never_sends) - client_ws.send_text = AsyncMock() - client_ws.close = AsyncMock() - return client_ws - - -def _realtime_session_logging() -> LiteLLMLogging: - return LiteLLMLogging( - model="gpt-realtime-2.1", - messages=[], - stream=False, - call_type="_arealtime", - start_time=datetime.now(), - litellm_call_id=str(uuid.uuid4()), - function_id="realtime-session", - ) - - -async def _relay_upstream( - litellm_logging_obj: LiteLLMLogging, - *upstream_frames: bytes | Exception, - client_ws: MagicMock | None = None, - provider_config: BaseRealtimeConfig | None = None, - upstream_stays_open: bool = False, -) -> None: - backend_ws: Final = MagicMock() - backend_ws.recv = AsyncMock( - side_effect=_frames_then_silence(*upstream_frames) if upstream_stays_open else list(upstream_frames) - ) - backend_ws.send = AsyncMock() - relay: Final = RealTimeStreaming( - client_ws or _realtime_client_ws(), - backend_ws, - litellm_logging_obj, - provider_config=provider_config, - model="gpt-realtime-2.1", - logging_worker=_DiscardingLoggingWorker(), - ) - await asyncio.wait_for(relay.bidirectional_forward(), timeout=2) - - -_REALTIME_SESSION_CREATED: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode() - - -async def _clean_realtime_session(litellm_logging_obj: LiteLLMLogging) -> None: - await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(Close(1000, ""), None)) - - -async def _upstream_internal_error_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None: - await _relay_upstream( - litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(Close(1011, "upstream failed"), None) - ) - - -async def _upstream_drop_without_close_frame_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None: - await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(None, None)) - - -async def _upstream_close_before_any_frame(litellm_logging_obj: LiteLLMLogging) -> None: - await _relay_upstream(litellm_logging_obj, ConnectionClosed(Close(1008, "model not found"), None)) - - -def _frames_then_silence(*frames: bytes | str) -> Callable[[], Awaitable[bytes | str]]: - pending: Final = iter(frames) - - async def _next_frame() -> bytes | str: - frame: Final = next(pending, None) - if frame is None: - await asyncio.Event().wait() - raise AssertionError("unreachable") - return frame - - return _next_frame - - -def _client_ws_sending(*client_frames: str) -> MagicMock: - client_ws: Final = _realtime_client_ws() - client_ws.receive_text = AsyncMock(side_effect=_frames_then_silence(*client_frames)) - return client_ws - - -async def _vertex_access_token() -> str: - return "token" - - -async def _client_frame_the_provider_rejects(litellm_logging_obj: LiteLLMLogging) -> None: - deferred_setup_provider: Final = VertexChirpRealtimeConfig( - resolve_access_token=_vertex_access_token, project="proj-1", location="us" - ) - await _relay_upstream( - litellm_logging_obj, - client_ws=_client_ws_sending( - json.dumps( - { - "type": "session.update", - "session": {"type": "transcription", "audio": {"input": {"format": {"type": "audio/pcmu"}}}}, - } - ) - ), - provider_config=deferred_setup_provider, - upstream_stays_open=True, - ) - - -async def _client_hangs_up_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None: - client_ws: Final = _realtime_client_ws() - client_ws.receive_text = AsyncMock(side_effect=ConnectionClosed(Close(1000, ""), None)) - await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, client_ws=client_ws, upstream_stays_open=True) - - -async def _refused_handshake(litellm_logging_obj: LiteLLMLogging) -> None: - await close_after_upstream_handshake_refusal(_realtime_client_ws(), 403) - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("realtime_session", "counted_successes"), - [ - (_clean_realtime_session, 1), - (_upstream_close_before_any_frame, 0), - (_refused_handshake, 0), - (_upstream_internal_error_after_a_frame, 0), - (_upstream_drop_without_close_frame_after_a_frame, 0), - (_client_frame_the_provider_rejects, 0), - (_client_hangs_up_after_a_frame, 1), - ], - ids=[ - "clean-session", - "upstream-close-before-any-frame", - "refused-handshake", - "upstream-1011-after-a-frame", - "upstream-1006-after-a-frame", - "client-frame-the-provider-rejects", - "client-hangs-up-after-a-frame", - ], -) -async def test_realtime_session_counts_a_router_success_only_when_it_ended_without_a_failure( - realtime_session: Callable[[LiteLLMLogging], Awaitable[None]], counted_successes: int -) -> None: - router: Final = Router( - model_list=[ - { - "model_name": "realtime-voice", - "litellm_params": {"model": "openai/gpt-realtime-2.1"}, - "model_info": {"id": "realtime-a"}, - } - ], - num_retries=0, - ) - - async def _arealtime(litellm_logging_obj: LiteLLMLogging, **kwargs: object) -> None: - await realtime_session(litellm_logging_obj) - - await router._ageneric_api_call_with_fallbacks( - model="realtime-voice", original_function=_arealtime, litellm_logging_obj=_realtime_session_logging() - ) - - assert get_deployment_successes_for_current_minute(router, "realtime-a") == counted_successes - - -@pytest.mark.asyncio -@pytest.mark.parametrize("realtime_model", ["bedrock/amazon.nova-sonic-v1:0", "openai/gpt-realtime-2.1"]) -@pytest.mark.parametrize(("clean_sessions", "cooled_down"), [(10, False), (3, True)]) -async def test_realtime_cooldown_reads_the_failure_rate_of_the_sessions_a_deployment_served( - realtime_model: str, clean_sessions: int, cooled_down: bool -) -> None: - import time - - provider, _, provider_model = realtime_model.partition("/") - router: Final = Router( - model_list=[ - { - "model_name": "realtime-voice", - "litellm_params": {"model": realtime_model, "order": order}, - "model_info": {"id": deployment_id}, - } - for order, deployment_id in ((1, "realtime-a"), (2, "realtime-b")) - ], - num_retries=0, - ) - drop: Final = litellm.ServiceUnavailableError( - message="realtime provider stream ended before the session closed", - llm_provider=provider, - model=provider_model, - ) - - def realtime_session(failure: Exception | None): - async def _arealtime(litellm_logging_obj: LiteLLMLogging, **kwargs: object) -> None: - if failure is not None: - raise failure - await _clean_realtime_session(litellm_logging_obj) - - return _arealtime - - for _ in range(clean_sessions): - await router._ageneric_api_call_with_fallbacks( - model="realtime-voice", - original_function=realtime_session(None), - litellm_logging_obj=_realtime_session_logging(), - ) - - async def drop_session() -> bool: - with pytest.raises(litellm.ServiceUnavailableError): - await router._ageneric_api_call_with_fallbacks( - model="realtime-voice", - original_function=realtime_session(drop), - litellm_logging_obj=_realtime_session_logging(), - ) - return router.deployment_callback_on_failure( - kwargs={ - "exception": drop, - "litellm_params": {"metadata": {"model_group": "realtime-voice"}, "model_info": {"id": "realtime-a"}}, - }, - completion_response=None, - start_time=time.time(), - end_time=time.time(), - ) - - drop_cooldowns: Final = [await drop_session() for _ in range(5)] - - assert any(drop_cooldowns) is cooled_down - - -async def _kwargs_routed_to(router: Router, model: str, function_name: str, **request: object) -> dict[str, object]: - routed: Final[dict[str, object]] = {} - - async def call(**kwargs: object) -> None: - routed.update(kwargs) - - call.__name__ = function_name - await router._ageneric_api_call_with_fallbacks_helper(model=model, original_generic_function=call, **request) - return routed - - -async def _finish_35_token_websocket_turn(routed: Mapping[str, object], call_id: str) -> None: - import time - - from litellm.litellm_core_utils.litellm_logging import Logging - - turn: Final = Logging( - model="gpt-5-mini", - messages=[], - stream=False, - call_type=CallTypes.aresponses_websocket.value, - start_time=time.time(), - litellm_call_id=call_id, - function_id=call_id, - ) - turn.update_environment_variables(litellm_params={"litellm_metadata": routed["litellm_metadata"]}, optional_params={}) - await turn.dispatch_success_handlers( - [{"type": "response.completed", "response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}}}], - prefer_async_handlers=True, - ) - - -@pytest.mark.asyncio -async def test_responses_websocket_turn_is_refused_once_its_deployment_crosses_budget(): - router = Router( - model_list=[ - { - "model_name": "ws-budget", - "litellm_params": { - "model": "openai/gpt-5-mini", - "api_key": "sk-fake", - "max_budget": 0.001, - "budget_duration": "1d", - }, - "model_info": {"id": "ws-budget-dep"}, - } - ] - ) - assert "turn_admission" not in await _kwargs_routed_to(router, "ws-budget", "aresponses") - admit_turn: Final = (await _kwargs_routed_to(router, "ws-budget", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - - await admit_turn() - await router.cache.async_set_cache("deployment_spend:ws-budget-dep:1d", 0.002) - - with pytest.raises(litellm.RateLimitError, match="crossed budget"): - await admit_turn() - - -@pytest.mark.asyncio -async def test_responses_websocket_turns_do_not_make_their_deployment_look_busy_to_least_busy(): - router: Final = Router( - model_list=[ - { - "model_name": "ws-least-busy", - "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake"}, - "model_info": {"id": deployment_id}, - } - for deployment_id in ("ws-least-busy-a", "ws-least-busy-b") - ], - routing_strategy="least-busy", - ) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-least-busy", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - for _ in range(3): - await admit_turn() - - picks: Final = [ - (await router._core_async_get_available_deployment(model="ws-least-busy", request_kwargs={"input": "hi"}))[ - "model_info" - ]["id"] - for _ in range(2) - ] - - assert sorted(picks) == ["ws-least-busy-a", "ws-least-busy-b"] - - -@pytest.mark.asyncio -async def test_responses_websocket_turn_after_the_first_counts_toward_deployment_rpm(): - router = Router( - model_list=[ - { - "model_name": "ws-rpm", - "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "rpm": 1}, - "model_info": {"id": "ws-rpm-dep"}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-rpm", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - - with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): - await admit_turn() - - -@pytest.mark.asyncio -async def test_responses_websocket_turn_is_checked_against_limits_changed_after_the_socket_opened(): - router = Router( - model_list=[ - { - "model_name": "ws-rpm-lowered", - "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "rpm": 100}, - "model_info": {"id": "ws-rpm-lowered-dep"}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-rpm-lowered", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - await admit_turn() - - router.upsert_deployment( - Deployment( - model_name="ws-rpm-lowered", - litellm_params=LiteLLM_Params(model="openai/gpt-5-mini", api_key="sk-fake", rpm=1), - model_info=ModelInfo(id="ws-rpm-lowered-dep"), - ) - ) - - with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): - await admit_turn() - - -@pytest.mark.asyncio -async def test_responses_websocket_turn_counts_toward_deployment_tpm(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setattr(litellm, "callbacks", []) - router = Router( - model_list=[ - { - "model_name": "ws-tpm", - "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "tpm": 30}, - "model_info": {"id": "ws-tpm-dep"}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - monkeypatch.setattr(litellm, "_async_success_callback", list(litellm.callbacks)) - routed: Final = await _kwargs_routed_to(router, "ws-tpm", "_aresponses_websocket") - admit_turn: Final = routed["turn_admission"] - assert callable(admit_turn) - await admit_turn() - - await _finish_35_token_websocket_turn(routed, "ws-tpm-turn") - - with pytest.raises(litellm.RateLimitError, match="TPM limit=30, current usage=35"): - await admit_turn() - - -@pytest.mark.asyncio -async def test_responses_websocket_turn_is_refused_by_usage_based_routing_once_deployment_tpm_is_used( - monkeypatch: pytest.MonkeyPatch, -): - monkeypatch.setattr(litellm, "callbacks", []) - router = Router( - model_list=[ - { - "model_name": "ws-tpm-v2", - "litellm_params": {"model": "openai/gpt-5-mini", "api_key": "sk-fake", "tpm": 30}, - "model_info": {"id": "ws-tpm-v2-dep"}, - } - ], - routing_strategy="usage-based-routing-v2", - ) - monkeypatch.setattr(litellm, "_async_success_callback", list(litellm.callbacks)) - routed: Final = await _kwargs_routed_to(router, "ws-tpm-v2", "_aresponses_websocket") - admit_turn: Final = routed["turn_admission"] - assert callable(admit_turn) - await admit_turn() - - await _finish_35_token_websocket_turn(routed, "ws-tpm-v2-turn") - - with pytest.raises(litellm.RateLimitError): - await admit_turn() - - -@pytest.mark.asyncio -async def test_responses_websocket_turns_do_not_reserve_deployment_otpm(): - router = Router( - model_list=[ - { - "model_name": "ws-otpm", - "litellm_params": {"model": "openai/ws-otpm-unpriced", "api_key": "sk-fake", "otpm": 5000}, - "model_info": {"id": "ws-otpm-dep"}, - } - ], - optional_pre_call_checks=["enforce_model_rate_limits"], - ) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-otpm", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - - for _ in range(3): - await admit_turn() - - def _io_limited_router(kind: str, **litellm_params: object) -> Router: return Router( model_list=[ @@ -1012,178 +349,6 @@ def _minutes_counter_key(kind: str, minute: datetime, model: str = "openai/ws-io return f"global_router:ws-io-dep:{model}:{kind}:{minute.strftime('%H-%M')}" -async def _count_minutes_tokens( - router: Router, kind: str, tokens: int, minute: datetime, model: str = "openai/ws-io-unpriced" -) -> None: - await router.cache.async_increment_cache(key=_minutes_counter_key(kind, minute, model), value=tokens, ttl=120) - - -async def _settle_turn( - router: Router, reservation: Mapping[str, object], input_tokens: int, output_tokens: int -) -> None: - from litellm.cost_calculator import ResponsesWebSocketTokenUsageProcessor - from litellm.litellm_core_utils.get_litellm_params import get_litellm_params - from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck - - rate_limits: Final = next( - check for check in router.optional_callbacks or () if isinstance(check, ModelRateLimitingCheck) - ) - await rate_limits.async_log_success_event( - {"litellm_params": get_litellm_params(litellm_metadata=dict(reservation))}, - ResponsesWebSocketTokenUsageProcessor.create_logging_object( - [ - { - "type": "response.completed", - "response": {"usage": {"input_tokens": input_tokens, "output_tokens": output_tokens}}, - } - ] - ), - None, - None, - ) - - -_LONG_OUTPUT_MODEL: Final = "openai/gpt-5.4-mini" - - -def _long_output_router() -> tuple[Router, Mapping[str, int]]: - max_output: Final = litellm.get_model_info(_LONG_OUTPUT_MODEL)["max_output_tokens"] - assert isinstance(max_output, int) - limits: Final = {"itpm": 10_000, "otpm": max_output // 2} - return _io_limited_router("itpm", model=_LONG_OUTPUT_MODEL, otpm=limits["otpm"]), limits - - -@pytest.mark.asyncio -@pytest.mark.parametrize("kind", ["itpm", "otpm"]) -async def test_responses_websocket_opens_without_counting_tokens_until_the_minute_reaches_deployment_io_limit( - kind: str, router_minute_pinned: datetime -): - router, limits = _long_output_router() - - await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) - - assert [ - await router.cache.async_get_cache(key=_minutes_counter_key(counted, router_minute_pinned, _LONG_OUTPUT_MODEL)) - for counted in ("itpm", "otpm") - ] == [None, None] - await _count_minutes_tokens(router, kind, limits[kind], router_minute_pinned, _LONG_OUTPUT_MODEL) - with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit={limits[kind]}, current usage="): - await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) - - -@pytest.mark.asyncio -async def test_responses_websocket_first_turn_counts_its_actual_tokens_toward_deployment_itpm_and_otpm( - router_minute_pinned: datetime, -): - router, _ = _long_output_router() - socket_metadata: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()))[ - "litellm_metadata" - ] - assert isinstance(socket_metadata, dict) - - await _settle_turn(router, socket_metadata, input_tokens=20, output_tokens=15) - - assert [ - await router.cache.async_get_cache(key=_minutes_counter_key(counted, router_minute_pinned, _LONG_OUTPUT_MODEL)) - for counted in ("itpm", "otpm") - ] == [20, 15] - - -_ADMITTED_LATE_IN_THE_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 50, tzinfo=timezone.utc) -_FINISHED_IN_THE_NEXT_MINUTE: Final = datetime(2026, 1, 1, 12, 1, 15, tzinfo=timezone.utc) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("kind", ["itpm", "otpm"]) -@pytest.mark.parametrize("turn", ["first", "later"]) -async def test_responses_websocket_turn_finishing_in_the_next_minute_counts_its_tokens_toward_that_minute( - kind: str, turn: str, monkeypatch: pytest.MonkeyPatch -): - _pin_router_clock(monkeypatch, _ADMITTED_LATE_IN_THE_MINUTE) - router: Final = _io_limited_router(kind) - socket: Final = await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket", websocket=object()) - admit_turn: Final = socket["turn_admission"] - assert callable(admit_turn) - reservation: Final = socket["litellm_metadata"] if turn == "first" else await admit_turn() - assert isinstance(reservation, Mapping) - _pin_router_clock(monkeypatch, _FINISHED_IN_THE_NEXT_MINUTE) - - await _settle_turn(router, reservation, input_tokens=10_000, output_tokens=10_000) - - with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000, current usage=10000"): - await admit_turn() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("body", [{}, {"websocket": True}], ids=["plain", "websocket_key"]) -async def test_messages_call_reserves_max_tokens_against_deployment_otpm_whatever_websocket_key_its_body_carries( - body: Mapping[str, object], router_minute_pinned: datetime -): - router: Final = _io_limited_router("otpm") - - for _ in range(3): - await _kwargs_routed_to(router, "ws-io", "anthropic_messages", max_tokens=3_000, **body) - - assert await router.cache.async_get_cache(key=_minutes_counter_key("otpm", router_minute_pinned)) == 9_000 - with pytest.raises(litellm.RateLimitError, match="OTPM limit=10000, current usage=12000"): - await _kwargs_routed_to(router, "ws-io", "anthropic_messages", max_tokens=3_000, **body) - - -async def _realtime_opens_until_refused(monkeypatch: pytest.MonkeyPatch, **socket: object) -> tuple[int, str]: - monkeypatch.setattr(litellm, "callbacks", []) - router: Final = _io_limited_router("otpm", otpm=50) - for opened in range(51): - try: - await _kwargs_routed_to(router, "ws-io", "_arealtime", **socket) - except litellm.RateLimitError as refused: - return opened, str(refused) - pytest.fail("realtime opens never reached the deployment's otpm") - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("router_minute_pinned") -async def test_realtime_open_reserves_and_is_refused_at_deployment_otpm_like_an_open_without_a_socket( - monkeypatch: pytest.MonkeyPatch, -): - with_socket: Final = await _realtime_opens_until_refused(monkeypatch, websocket=object()) - - assert "OTPM limit=50, current usage=" in with_socket[1] - assert with_socket == await _realtime_opens_until_refused(monkeypatch) - - -@pytest.mark.asyncio -@pytest.mark.parametrize("kind", ["itpm", "otpm"]) -async def test_responses_websocket_turn_is_refused_once_the_minutes_tokens_reach_deployment_io_limit( - kind: str, router_minute_pinned: datetime -): - router: Final = _io_limited_router(kind) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - await _count_minutes_tokens(router, kind, 10_000, router_minute_pinned) - - with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000, current usage="): - await admit_turn() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("kind", ["itpm", "otpm"]) -async def test_responses_websocket_turn_refused_for_tokens_leaves_the_daily_request_limit_unused( - kind: str, router_minute_pinned: datetime -): - router: Final = _io_limited_router(kind, rpd=2) - admit_turn: Final = (await _kwargs_routed_to(router, "ws-io", "_aresponses_websocket"))["turn_admission"] - assert callable(admit_turn) - await _count_minutes_tokens(router, kind, 10_000, router_minute_pinned) - with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=10000"): - await admit_turn() - await _count_minutes_tokens(router, kind, -10_000, router_minute_pinned) - - await admit_turn() - - with pytest.raises(litellm.RateLimitError, match="RPD"): - await admit_turn() - - _PROVIDER_RESPONSE: Final = { "id": "resp_io", "object": "response", @@ -1432,41 +597,6 @@ async def test_deployment_callback_on_failure(model_list): ) -def test_deployment_callback_respects_cooldown_time(model_list): - """Ensure per-model cooldown_time is honored even when exception headers are present.""" - import time - from unittest.mock import patch - - import httpx - - router = Router(model_list=model_list) - - class FakeException(Exception): - def __init__(self): - self.status_code = 429 - self.headers = httpx.Headers({"x-test": "1"}) - - kwargs = { - "exception": FakeException(), - "litellm_params": { - "metadata": {"model_group": "gpt-5-mini"}, - "model_info": {"id": 100}, - "cooldown_time": 0, - }, - } - - with patch("litellm.router_python_core.set_cooldown_deployments") as mock_set: - router.deployment_callback_on_failure( - kwargs=kwargs, - completion_response=None, - start_time=time.time(), - end_time=time.time(), - ) - - mock_set.assert_called_once() - assert mock_set.call_args.kwargs["time_to_cooldown"] == 0 - - @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) def test_log_retry(model_list: list[DeploymentTypedDict], metadata_key: str) -> None: """log_retry appends one flat record per failed attempt, copies neither the request kwargs nor the @@ -1537,57 +667,6 @@ def test_log_retry_reads_base_model_from_the_deployments_litellm_params() -> Non assert (record["model"], record["base_model"]) == ("azure/my-named-deployment", "azure/gpt-4o-mini") -@pytest.mark.usefixtures("router_minute_pinned") -def test_update_usage(model_list): - """Test if the '_update_usage' function is working correctly""" - router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") - deployment_id = deployment["model_info"]["id"] - request_count = router._update_usage(deployment_id=deployment_id, parent_otel_span=None) - assert request_count == 1 - - request_count = router._update_usage(deployment_id=deployment_id, parent_otel_span=None) - - assert request_count == 2 - - -@pytest.mark.parametrize("finish_reason, expected_fallback", [("content_filter", True), ("stop", False)]) -@pytest.mark.parametrize("fallback_type", ["model-specific", "default"]) -def test_should_raise_content_policy_error(model_list, finish_reason, expected_fallback, fallback_type): - """Test if the '_should_raise_content_policy_error' function is working correctly""" - router = Router( - model_list=model_list, - default_fallbacks=["gpt-5.5"] if fallback_type == "default" else None, - ) - - assert ( - router._should_raise_content_policy_error( - model="gpt-5-mini", - response=litellm.ModelResponse( - model="gpt-5-mini", - choices=[ - { - "finish_reason": finish_reason, - "message": {"content": "I'm fine, thank you!"}, - } - ], - usage={"total_tokens": 100}, - ), - kwargs={ - "content_policy_fallbacks": ([{"gpt-5-mini": "gpt-5.5"}] if fallback_type == "model-specific" else None) - }, - ) - is expected_fallback - ) - - -def test_get_healthy_deployments(model_list): - """Test if the '_get_healthy_deployments' function is working correctly""" - router = Router(model_list=model_list) - deployments = router._get_healthy_deployments(model="gpt-5-mini", parent_otel_span=None) - assert len(deployments) > 0 - - @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_routing_strategy_pre_call_checks(model_list, sync_mode): @@ -1600,9 +679,7 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") litellm_logging_obj = Logging( model="gpt-5-mini", @@ -1617,9 +694,7 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): router.routing_strategy_pre_call_checks(deployment) else: ## NO EXCEPTION - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) ## WITH EXCEPTION - rate limit error with patch.object( @@ -1634,49 +709,12 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode): ), ): with pytest.raises(litellm.RateLimitError): - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) ## WITH EXCEPTION - generic error - with patch.object( - callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error")) - ): + with patch.object(callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error"))): with pytest.raises(Exception, match="Error"): - await router.async_routing_strategy_pre_call_checks( - deployment, litellm_logging_obj - ) - - -@pytest.mark.parametrize( - "set_supported_environments, supported_environments, is_supported", - [(True, ["staging"], True), (False, None, True), (True, ["development"], False)], -) -def test_create_deployment( - model_list, set_supported_environments, supported_environments, is_supported -): - """Test if the '_create_deployment' function is working correctly""" - router = Router(model_list=model_list) - - if set_supported_environments: - os.environ["LITELLM_ENVIRONMENT"] = "staging" - deployment = router._create_deployment( - deployment_info={}, - _model_name="gpt-5-mini", - _litellm_params={ - "model": "gpt-5-mini", - "api_key": "test", - "custom_llm_provider": "openai", - }, - _model_info={ - "id": 100, - "supported_environments": supported_environments, - }, - ) - if is_supported: - assert deployment is not None - else: - assert deployment is None + await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj) @pytest.mark.parametrize( @@ -1688,20 +726,14 @@ def test_deployment_is_active_for_environment( ): """Test if the '_deployment_is_active_for_environment' function is working correctly""" router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name( - model_group_name="gpt-5-mini" - ) + deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") if set_supported_environments: os.environ["LITELLM_ENVIRONMENT"] = "staging" deployment["model_info"]["supported_environments"] = supported_environments if is_supported: - assert ( - router.deployment_is_active_for_environment(deployment=deployment) is True - ) + assert router.deployment_is_active_for_environment(deployment=deployment) is True else: - assert ( - router.deployment_is_active_for_environment(deployment=deployment) is False - ) + assert router.deployment_is_active_for_environment(deployment=deployment) is False def test_set_model_list(model_list): @@ -1711,18 +743,6 @@ def test_set_model_list(model_list): assert len(router.model_list) == len(model_list) -def test_add_deployment(model_list): - """Test if the '_add_deployment' function is working correctly""" - router = Router(model_list=model_list) - deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini") - deployment["model_info"]["id"] = "100" - - router.add_deployment(deployment=deployment) - - router._add_deployment(deployment=deployment) - assert len(router.model_list) == len(model_list) + 1 - - def test_upsert_deployment(model_list): """Test if the 'upsert_deployment' function is working correctly""" router = Router(model_list=model_list) @@ -1758,63 +778,6 @@ def test_get_model_group(model_list): assert model_group[0]["model_name"] == "gpt-5-mini" -@pytest.mark.parametrize("user_facing_model_group_name", ["gpt-5-mini", "gpt-5.5"]) -def test_set_model_group_info(model_list, user_facing_model_group_name): - """Test if the 'set_model_group_info' function is working correctly""" - router = Router(model_list=model_list) - resp = router._set_model_group_info( - model_group="gpt-5-mini", - user_facing_model_group_name=user_facing_model_group_name, - ) - assert resp is not None - assert resp.model_group == user_facing_model_group_name - - -def test_set_model_group_info_reports_supports_continuation_when_any_deployment_has_it(): - router: Final = Router( - model_list=[ - { - "model_name": "long-decode", - "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, - }, - { - "model_name": "long-decode", - "litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "fake"}, - "model_info": {"supports_continuation": True}, - }, - { - "model_name": "long-decode", - "litellm_params": {"model": "openai/gpt-4.1-nano", "api_key": "fake"}, - }, - { - "model_name": "plain", - "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, - }, - ] - ) - long_decode: Final = router._set_model_group_info( - model_group="long-decode", user_facing_model_group_name="long-decode" - ) - plain: Final = router._set_model_group_info(model_group="plain", user_facing_model_group_name="plain") - assert long_decode is not None and long_decode.supports_continuation is True - assert plain is not None and plain.supports_continuation is False - - -@pytest.mark.parametrize("value", ["maybe", "", ["x"], "true", 1]) -def test_set_model_group_info_counts_only_a_boolean_true_supports_continuation(value: object) -> None: - router: Final = Router( - model_list=[ - { - "model_name": "long-decode", - "litellm_params": {"model": "openai/gpt-4.1", "api_key": "fake"}, - "model_info": {"supports_continuation": value}, - }, - ] - ) - group: Final = router._set_model_group_info(model_group="long-decode", user_facing_model_group_name="long-decode") - assert group is not None and group.supports_continuation is False - - @pytest.mark.asyncio async def test_set_response_headers(model_list): """Test if the 'set_response_headers' function is working correctly""" @@ -1823,42 +786,6 @@ async def test_set_response_headers(model_list): assert resp is None -@pytest.mark.asyncio -async def test_set_response_headers_passes_through_post_increment_counters(model_list): - from pydantic import BaseModel - - class _Usage(BaseModel): - total_tokens: int = 42 - - class _Resp(BaseModel): - usage: _Usage = _Usage() - _hidden_params: dict = {} - - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-remaining-tokens": 958, - "x-ratelimit-limit-tokens": 1000, - "x-ratelimit-remaining-requests": 99, - "x-ratelimit-limit-requests": 100, - "x-ratelimit-remaining-input-tokens": 1000, - "x-ratelimit-remaining-output-tokens": 500, - } - ) - - resp = _Resp() - resp._hidden_params = {} - await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") - - headers = resp._hidden_params["additional_headers"] - assert headers["x-ratelimit-remaining-tokens"] == 958 - assert headers["x-ratelimit-remaining-requests"] == 99 - assert headers["x-ratelimit-limit-tokens"] == 1000 - assert headers["x-ratelimit-limit-requests"] == 100 - assert headers["x-ratelimit-remaining-input-tokens"] == 1000 - assert headers["x-ratelimit-remaining-output-tokens"] == 500 - - def _rpm_tpm_router(model_id: str) -> Router: return Router( model_list=[ @@ -1872,7 +799,6 @@ def _rpm_tpm_router(model_id: str) -> Router: def _pin_router_clock(monkeypatch: pytest.MonkeyPatch, pinned: datetime) -> None: - monkeypatch.setattr("litellm.router_python_core.get_utc_datetime", lambda: pinned) monkeypatch.setattr( "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: pinned ) @@ -1880,6 +806,10 @@ def _pin_router_clock(monkeypatch: pytest.MonkeyPatch, pinned: datetime) -> None @pytest.fixture def router_minute_pinned(monkeypatch: pytest.MonkeyPatch) -> datetime: + # The router's own usage counters read the real clock: start well inside a minute 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) pinned: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) _pin_router_clock(monkeypatch, pinned) return pinned @@ -1958,121 +888,6 @@ async def test_increment_deployment_usage_for_response_skips_session_wrappers(): assert ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY not in request_kwargs["litellm_metadata"] -@pytest.mark.asyncio -@pytest.mark.usefixtures("router_minute_pinned") -async def test_increment_deployment_usage_writes_only_positive_deltas_for_limited_deployments(): - router = _rpm_tpm_router("lit-3058-delta") - unlimited = Router( - model_list=[ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake"}, - "model_info": {"id": "lit-3058-unlimited"}, - } - ] - ) - - tpm_key = await router._increment_deployment_usage( - deployment_id="lit-3058-delta", - deployment_name="gpt-5-mini", - model_group="gpt-5-mini", - total_tokens=25, - rpm_increment=1, - parent_otel_span=None, - ) - assert tpm_key is not None - assert await router.get_model_group_usage("gpt-5-mini") == (25, 1) - - assert ( - await router._increment_deployment_usage( - deployment_id="lit-3058-delta", - deployment_name="gpt-5-mini", - model_group="gpt-5-mini", - total_tokens=0, - rpm_increment=0, - parent_otel_span=None, - ) - is None - ) - assert await router.get_model_group_usage("gpt-5-mini") == (25, 1) - - assert ( - await unlimited._increment_deployment_usage( - deployment_id="lit-3058-unlimited", - deployment_name="gpt-5-mini", - model_group="gpt-5-mini", - total_tokens=25, - rpm_increment=1, - parent_otel_span=None, - ) - is None - ) - assert await unlimited.get_model_group_usage("gpt-5-mini") == (None, None) - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("router_minute_pinned") -async def test_increment_deployment_usage_starts_tracking_after_limit_is_added(): - router = Router( - model_list=[ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini", "api_key": "sk-fake"}, - "model_info": {"id": "lit-3058-upsert"}, - } - ] - ) - increment = functools.partial( - router._increment_deployment_usage, - deployment_id="lit-3058-upsert", - deployment_name="gpt-5-mini", - model_group="gpt-5-mini", - total_tokens=25, - rpm_increment=1, - parent_otel_span=None, - ) - - assert await increment() is None - - router.upsert_deployment( - Deployment( - model_name="gpt-5-mini", - litellm_params=LiteLLM_Params(model="gpt-5-mini", api_key="sk-fake", rpm=10), - model_info={"id": "lit-3058-upsert"}, - ) - ) - - assert await increment() is not None - assert await router.get_model_group_usage("gpt-5-mini") == (25, 1) - - -@pytest.mark.parametrize( - "litellm_params, model_info, deployment_id, tracks", - [ - ({}, {}, "lit-3058-tracks", False), - ({"rpm": 10}, {}, "lit-3058-tracks", True), - ({"tpm": 1000}, {}, "lit-3058-tracks", True), - ({"otpm": 500}, {}, "lit-3058-tracks", True), - ({}, {"itpm": 500}, "lit-3058-tracks", True), - ({"rpm": 10}, {}, "lit-3058-unknown", False), - ], -) -def test_deployment_tracks_usage_for_every_limit_source( - litellm_params: dict, model_info: dict, deployment_id: str, tracks: bool -): - router = Router( - model_list=[ - { - "model_name": "my-unmapped-model", - "litellm_params": {"model": "openai/my-unmapped-model", "api_key": "sk-fake", **litellm_params}, - "model_info": {"id": "lit-3058-tracks", **model_info}, - } - ] - ) - - assert router._deployment_tracks_usage(deployment_id, "my-unmapped-model") is tracks - - def _shared_redis_stub(store: dict) -> MagicMock: from litellm.caching.redis_cache import RedisCache @@ -2158,27 +973,19 @@ async def test_get_model_group_io_token_usage_sums_across_deployments(): minute = get_utc_datetime().strftime("%H-%M") keys_and_values = [ ( - RouterCacheEnum.ITPM.value.format( - id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute - ), + RouterCacheEnum.ITPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), 30, ), ( - RouterCacheEnum.OTPM.value.format( - id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute - ), + RouterCacheEnum.OTPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), 10, ), ( - RouterCacheEnum.ITPM.value.format( - id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute - ), + RouterCacheEnum.ITPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), 70, ), ( - RouterCacheEnum.OTPM.value.format( - id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute - ), + RouterCacheEnum.OTPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), 20, ), ] @@ -2200,66 +1007,57 @@ async def test_get_model_group_io_token_usage_no_deployments_returns_none(): assert current_otpm is None +def _limited_router(**io_limits: int) -> Router: + return Router( + model_list=[ + { + "model_name": "limited", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_key": "sk-fake", + "tpm": 1000, + "rpm": 100, + **io_limits, + }, + } + ], + optional_pre_call_checks=["enforce_model_rate_limits"], + ) + + @pytest.mark.asyncio -async def test_get_remaining_model_group_usage_merges_io_and_tpm_headers(model_list): +async def test_get_remaining_model_group_usage_merges_io_and_tpm_headers(): """ A model group with both itpm/otpm and tpm/rpm limits must expose the standard remaining-tokens/requests headers alongside the input/output token headers, so clients and prometheus gauges relying on either still get data. """ - from unittest.mock import Mock + router = _limited_router(itpm=900, otpm=450) - from litellm.types.router import ModelGroupInfo + headers = await router.get_remaining_model_group_usage("limited") - router = Router(model_list=model_list) - router._cached_get_model_group_info = Mock( - return_value=ModelGroupInfo( - model_group="gpt-3.5-turbo", - providers=["openai"], - itpm=2000, - otpm=1000, - tpm=5000, - rpm=50, - ) - ) - router.get_model_group_io_token_usage = AsyncMock(return_value=(100, 40)) - router._core_get_model_group_usage = AsyncMock(return_value=(500, 5)) + assert headers["x-ratelimit-remaining-input-tokens"] == 900 + assert headers["x-ratelimit-remaining-output-tokens"] == 450 + assert headers["x-ratelimit-remaining-tokens"] == 1000 + assert headers["x-ratelimit-remaining-requests"] == 100 - headers = await router.get_remaining_model_group_usage("gpt-3.5-turbo") - assert headers["x-ratelimit-remaining-input-tokens"] == 1900 - assert headers["x-ratelimit-remaining-output-tokens"] == 960 - assert headers["x-ratelimit-remaining-tokens"] == 4500 - assert headers["x-ratelimit-remaining-requests"] == 45 +class _HeaderlessResponse(BaseModel): + _hidden_params: dict = {} @pytest.mark.asyncio -async def test_set_response_headers_native_input_token_header_does_not_suppress_router_headers(model_list): +async def test_set_response_headers_native_input_token_header_does_not_suppress_router_headers(): """ A provider that natively returns `x-ratelimit-remaining-input-tokens` must not suppress the router's own remaining-tokens/requests headers for a non-IO model group. """ - from pydantic import BaseModel + router = _limited_router() - class _Usage(BaseModel): - total_tokens: int = 42 - - class _Resp(BaseModel): - usage: _Usage = _Usage() - _hidden_params: dict = {} - - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-remaining-tokens": 1000, - "x-ratelimit-remaining-requests": 100, - } - ) - - resp = _Resp() + resp = _HeaderlessResponse() resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-input-tokens": 5}} - await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") + await router.set_response_headers(response=resp, model_group="limited") headers = resp._hidden_params["additional_headers"] assert headers["x-ratelimit-remaining-tokens"] == 1000 @@ -2269,29 +1067,12 @@ async def test_set_response_headers_native_input_token_header_does_not_suppress_ @pytest.mark.asyncio -async def test_set_response_headers_native_token_header_does_not_suppress_io_headers(model_list): - from pydantic import BaseModel +async def test_set_response_headers_native_token_header_does_not_suppress_io_headers(): + router = _limited_router(itpm=900, otpm=450) - class _Usage(BaseModel): - total_tokens: int = 42 - - class _Resp(BaseModel): - usage: _Usage = _Usage() - _hidden_params: dict = {} - - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-remaining-tokens": 1000, - "x-ratelimit-remaining-requests": 100, - "x-ratelimit-remaining-input-tokens": 900, - "x-ratelimit-remaining-output-tokens": 450, - } - ) - - resp = _Resp() + resp = _HeaderlessResponse() resp._hidden_params = {"additional_headers": {"x-ratelimit-remaining-tokens": 5}} - await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") + await router.set_response_headers(response=resp, model_group="limited") headers = resp._hidden_params["additional_headers"] assert headers["x-ratelimit-remaining-tokens"] == 5 @@ -2301,27 +1082,16 @@ async def test_set_response_headers_native_token_header_does_not_suppress_io_hea @pytest.mark.asyncio -async def test_set_response_headers_handles_missing_usage(model_list): +async def test_set_response_headers_handles_missing_usage(): """ Streaming chunks and some response shapes may lack a `usage` attribute or populated `total_tokens`. Header composition must not depend on usage and never raise. """ - from pydantic import BaseModel + router = _limited_router() - class _Resp(BaseModel): - _hidden_params: dict = {} - - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-remaining-tokens": 1000, - "x-ratelimit-remaining-requests": 100, - } - ) - - resp = _Resp() + resp = _HeaderlessResponse() resp._hidden_params = {} - await router.set_response_headers(response=resp, model_group="gpt-3.5-turbo") + await router.set_response_headers(response=resp, model_group="limited") headers = resp._hidden_params["additional_headers"] assert headers["x-ratelimit-remaining-tokens"] == 1000 @@ -2329,17 +1099,9 @@ async def test_set_response_headers_handles_missing_usage(model_list): @pytest.mark.asyncio -async def test_set_response_headers_dict_anthropic_messages_response(model_list): +async def test_set_response_headers_dict_anthropic_messages_response(): """Anthropic /v1/messages returns a dict; IO rate-limit headers must attach.""" - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-limit-input-tokens": 25, - "x-ratelimit-remaining-input-tokens": 20, - "x-ratelimit-limit-output-tokens": 100, - "x-ratelimit-remaining-output-tokens": 95, - } - ) + router = _limited_router(itpm=25, otpm=100) resp = { "id": "msg_123", @@ -2348,18 +1110,18 @@ async def test_set_response_headers_dict_anthropic_messages_response(model_list) "content": [{"type": "text", "text": "hi"}], "usage": {"input_tokens": 5, "output_tokens": 1}, } - await router.set_response_headers(response=resp, model_group="io-itpm-strict") + await router.set_response_headers(response=resp, model_group="limited") assert "_hidden_params" in resp headers = resp["_hidden_params"]["additional_headers"] - assert headers["x-litellm-model-group"] == "io-itpm-strict" + assert headers["x-litellm-model-group"] == "limited" assert headers["x-ratelimit-limit-input-tokens"] == 25 - assert headers["x-ratelimit-remaining-input-tokens"] == 20 - assert headers["x-ratelimit-remaining-output-tokens"] == 95 + assert headers["x-ratelimit-remaining-input-tokens"] == 25 + assert headers["x-ratelimit-remaining-output-tokens"] == 100 @pytest.mark.asyncio -async def test_set_response_headers_wraps_bare_async_generator(model_list): +async def test_set_response_headers_wraps_bare_async_generator(): """ Streaming responses that never go through Router.make_call's usual object-based wrappers (e.g. the Anthropic /v1/messages -> Responses API @@ -2371,21 +1133,15 @@ async def test_set_response_headers_wraps_bare_async_generator(model_list): yield {"type": "message_start"} yield {"type": "message_stop"} - router = Router(model_list=model_list) - router._core_get_remaining_model_group_usage = AsyncMock( - return_value={ - "x-ratelimit-limit-input-tokens": 25, - "x-ratelimit-remaining-input-tokens": 20, - } - ) + router = _limited_router(itpm=25) - wrapped = await router.set_response_headers(response=_raw_generator(), model_group="io-itpm-strict") + wrapped = await router.set_response_headers(response=_raw_generator(), model_group="limited") assert hasattr(wrapped, "_hidden_params") headers = wrapped._hidden_params["additional_headers"] - assert headers["x-litellm-model-group"] == "io-itpm-strict" + assert headers["x-litellm-model-group"] == "limited" assert headers["x-ratelimit-limit-input-tokens"] == 25 - assert headers["x-ratelimit-remaining-input-tokens"] == 20 + assert headers["x-ratelimit-remaining-input-tokens"] == 25 from collections.abc import AsyncIterator @@ -2417,29 +1173,6 @@ def test_update_settings(model_list): assert router.allowed_fails == 20 -def test_common_checks_available_deployment(model_list): - """Test if the 'common_checks_available_deployment' function is working correctly""" - router = Router(model_list=model_list) - _, available_deployments = router._common_checks_available_deployment( - model="gpt-5-mini", - messages=[{"role": "user", "content": "hi"}], - input="hi", - specific_deployment=False, - ) - - assert len(available_deployments) > 0 - - -def test_filter_cooldown_deployments(model_list): - """Test if the 'filter_cooldown_deployments' function is working correctly""" - router = Router(model_list=model_list) - deployments = router._filter_cooldown_deployments( - healthy_deployments=router.get_all_deployments(model_name="gpt-5-mini"), - cooldown_deployments=[], - ) - assert len(deployments) == len(router.get_all_deployments(model_name="gpt-5-mini")) - - @pytest.mark.parametrize( "exception_type, exception_name, num_retries", [ @@ -2495,23 +1228,6 @@ def test_get_allowed_fails_from_policy(model_list, exception_type, exception_nam assert calc_allowed_fails == allowed_fails -def test_initialize_alerting(model_list): - """Test if the 'initialize_alerting' function is working correctly""" - from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting - from litellm.types.router import AlertingConfig - - router = Router( - model_list=model_list, alerting_config=AlertingConfig(webhook_url="test") - ) - router._initialize_alerting() - - callback_added = False - for callback in litellm.callbacks: - if isinstance(callback, SlackAlerting): - callback_added = True - assert callback_added is True - - def test_flush_cache(model_list): """Test if the 'flush_cache' function is working correctly""" router = Router(model_list=model_list) @@ -2666,9 +1382,6 @@ def test_discard_drops_routing_group_selectors(): {"group_name": "fast", "models": ["grouped-model"], "routing_strategy": "latency-based-routing"} ], ) - _, group_selector = router._get_routing_context("grouped-model") - assert any(callback is group_selector for callback in litellm.callbacks) - router.discard() assert all(getattr(litellm, list_name) == [] for list_name in _CALLBACK_LIST_NAMES) @@ -2712,27 +1425,6 @@ def test_routing_group_reload_keeps_every_hook_of_an_in_flight_callback_loop(): router.discard() -def _routers_seen_in_success_loop_discarding_first(first: Router, second: Router) -> Iterator[Router]: - for callback in litellm.success_callback: - if not isinstance(callback, MethodType): - continue - if callback.__self__ is first._core: - first.discard() - yield first - if callback.__self__ is second._core: - yield second - - -def test_discard_mid_success_loop_keeps_the_next_router_callback(): - first: Final = Router(model_list=_GROUPED_MODEL_LIST) - second: Final = Router(model_list=_GROUPED_MODEL_LIST) - - seen: Final = tuple(_routers_seen_in_success_loop_discarding_first(first, second)) - - assert seen == (first, second) - second.discard() - - @pytest.mark.asyncio async def test_discard_purges_optional_pre_call_checks_from_event_lists_after_serving_a_request(): router: Final = Router( @@ -2761,23 +1453,6 @@ def _is_registered(callback: object) -> bool: return any(registered is callback for registered in litellm.callbacks) -@pytest.mark.asyncio -async def test_discard_releases_per_request_override_selector(): - def build_override_then_discard() -> weakref.ReferenceType[object]: - router: Final = Router(model_list=_GROUPED_MODEL_LIST) - selector: Final = router._get_override_strategy_selector("usage-based-routing-v2") - assert selector is not None - router.discard() - return weakref.ref(selector) - - selector_ref: Final = build_override_then_discard() - await asyncio.sleep(0) - gc.collect() - - assert selector_ref() is None - - - def test_discard_registers_the_selector_of_a_same_strategy_survivor_built_without_a_model_list(): retired: Final = Router(model_list=[], routing_strategy="latency-based-routing") survivor: Final = Router(routing_strategy="latency-based-routing") @@ -2855,6 +1530,7 @@ async def test_discard_releases_its_router_budget_limiter(): assert budget_limiter_ref() is None + def test_initialize_assistants_endpoint(model_list): """Test if the 'initialize_assistants_endpoint' function is working correctly""" router = Router(model_list=model_list) @@ -2879,13 +1555,6 @@ def test_get_model_from_alias(model_list): assert model == "gpt-5-mini" -def test_get_deployment_by_litellm_model(model_list): - """Test if the 'get_deployment_by_litellm_model' function is working correctly""" - router = Router(model_list=model_list) - deployment = router._get_deployment_by_litellm_model(model="gpt-5-mini") - assert deployment is not None - - def test_get_pattern(model_list): router = Router(model_list=model_list) pattern = router.pattern_router.get_pattern(model="claude-3") @@ -2930,7 +1599,6 @@ def test_deployments_by_pattern(model_list): ], ) def test_pattern_match_deployment_set_model_name(user_request_model, model_name, litellm_model, expected_model): - from re import Match from litellm.litellm_core_utils.pattern_match_deployments import PatternMatchRouter @@ -2964,18 +1632,6 @@ def test_pattern_match_deployment_set_model_name(user_request_model, model_name, assert model["litellm_params"]["model"] == expected_model -@pytest.mark.parametrize( - "has_default_fallbacks, expected_result", - [(True, True), (False, False)], -) -def test_has_default_fallbacks(model_list, has_default_fallbacks, expected_result): - router = Router( - model_list=model_list, - default_fallbacks=(["my-default-fallback-model"] if has_default_fallbacks else None), - ) - assert router._has_default_fallbacks() is expected_result - - def test_add_optional_pre_call_checks(model_list): router = Router(model_list=model_list) @@ -2985,7 +1641,6 @@ def test_add_optional_pre_call_checks(model_list): @pytest.mark.asyncio async def test_async_callback_filter_deployments(model_list): - from litellm.router_strategy.budget_limiter import RouterBudgetLimiting router = Router(model_list=model_list) @@ -3001,20 +1656,6 @@ async def test_async_callback_filter_deployments(model_list): assert len(new_healthy_deployments) == len(healthy_deployments) -def test_cached_get_model_group_info(model_list): - """Test if the '_cached_get_model_group_info' function is working correctly with LRU cache""" - router = Router(model_list=model_list) - - result1 = router._cached_get_model_group_info("gpt-5-mini") - - result2 = router._cached_get_model_group_info("gpt-5-mini") - - assert result1 == result2 - - cache_info = router._cached_get_model_group_info.cache_info() - assert cache_info.hits > 0 - - def test_init_responses_api_endpoints(model_list): """Test if the '_init_responses_api_endpoints' function is working correctly""" from typing import Callable @@ -3072,23 +1713,6 @@ def test_mock_router_testing_params_str_to_bool_conversion( assert kwargs["other_param"] == "should_remain" -def test_is_auto_router_deployment(model_list): - """Test if the '_is_auto_router_deployment' function correctly identifies auto-router deployments""" - router = Router(model_list=model_list) - - litellm_params_auto = LiteLLM_Params(model="auto_router/my-auto-router") - assert router._is_auto_router_deployment(litellm_params_auto) is True - - litellm_params_regular = LiteLLM_Params(model="gpt-5-mini") - assert router._is_auto_router_deployment(litellm_params_regular) is False - - litellm_params_empty = LiteLLM_Params(model="") - assert router._is_auto_router_deployment(litellm_params_empty) is False - - litellm_params_contains = LiteLLM_Params(model="prefix_auto_router/something") - assert router._is_auto_router_deployment(litellm_params_contains) is False - - @patch("litellm.router_strategy.auto_router.auto_router.AutoRouter") def test_init_auto_router_deployment_success(mock_auto_router, model_list): """Test if the 'init_auto_router_deployment' function successfully initializes auto-router when all params provided""" @@ -3244,494 +1868,6 @@ def test_handle_clientside_credential_with_deployment_model_name(model_list): print("✓ _handle_clientside_credential test passed!") -def test_sync_generic_api_call_preserves_requested_model_group_in_logs(): - router = Router( - model_list=[ - { - "model_name": "claude-sonnet-4-6", - "litellm_params": { - "model": "bedrock/global.anthropic.claude-sonnet-4-6", - "aws_access_key_id": "test-access-key", - "aws_secret_access_key": "test-secret-key", - "aws_region_name": "us-west-2", - }, - } - ] - ) - - try: - captured_kwargs = {} - - def mock_original_function(**kwargs): - captured_kwargs.update(kwargs) - return {"status": "ok"} - - response = router._generic_api_call_with_fallbacks( - model="claude-sonnet-4-6", - original_function=mock_original_function, - ) - - assert response == {"status": "ok"} - assert captured_kwargs["model"] == "bedrock/global.anthropic.claude-sonnet-4-6" - assert captured_kwargs["litellm_metadata"]["model_group"] == "claude-sonnet-4-6" - assert captured_kwargs["litellm_metadata"]["deployment"] == "bedrock/global.anthropic.claude-sonnet-4-6" - finally: - router.discard() - - -def test_sync_generic_api_call_uses_request_kwargs_for_deployment_selection(): - router = Router( - model_list=[ - { - "model_name": "regional-model", - "litellm_params": { - "model": "anthropic/us-model", - "api_key": "test-api-key", - "region_name": "us", - }, - }, - { - "model_name": "regional-model", - "litellm_params": { - "model": "anthropic/eu-model", - "api_key": "test-api-key", - "region_name": "eu", - }, - }, - ], - enable_pre_call_checks=True, - ) - - try: - captured_kwargs = {} - - def mock_original_function(**kwargs): - captured_kwargs.update(kwargs) - return {"status": "ok"} - - response = router._generic_api_call_with_fallbacks( - model="regional-model", - original_function=mock_original_function, - messages=[{"role": "user", "content": "Hello from Europe"}], - allowed_model_region="eu", - ) - - assert response == {"status": "ok"} - assert captured_kwargs["model"] == "anthropic/eu-model" - finally: - router.discard() - - -@pytest.mark.parametrize( - "function_name, expected_metadata_key", - [ - ("acompletion", "metadata"), - ("_ageneric_api_call_with_fallbacks", "litellm_metadata"), - ("batch", "litellm_metadata"), - ("completion", "metadata"), - ("acreate_file", "litellm_metadata"), - ("aget_file", "litellm_metadata"), - ], -) -def test_handle_clientside_credential_metadata_loading(model_list, function_name, expected_metadata_key): - """Test that _handle_clientside_credential correctly loads metadata based on function name""" - router = Router(model_list=model_list) - - deployment = { - "model_name": "gpt-4.1", - "litellm_params": {"model": "gpt-4.1", "api_key": "test_key"}, - "model_info": {"id": "original-id-123"}, - } - - kwargs = { - "api_key": "client_side_key", - "api_base": "https://api.openai.com/v1", - expected_metadata_key: {"model_group": "gpt-4.1", "custom_field": "test_value"}, - } - - result_deployment = router._handle_clientside_credential( - deployment=deployment, kwargs=kwargs, function_name=function_name - ) - - assert isinstance(result_deployment, Deployment) - - assert result_deployment.model_name == "gpt-4.1" - - assert result_deployment.litellm_params.api_key == "client_side_key" - assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1" - - assert result_deployment.model_info.id != "original-id-123" - assert result_deployment.model_info.original_model_id == "original-id-123" - - assert len(router.model_list) == len(model_list) - assert router.get_deployment(model_id=result_deployment.model_info.id) is None - - if function_name == "acompletion": - assert "metadata" in kwargs - assert "litellm_metadata" not in kwargs - elif function_name in [ - "_ageneric_api_call_with_fallbacks", - "batch", - "acreate_file", - "aget_file", - ]: - assert "litellm_metadata" in kwargs - - print(f"✓ Success with function_name '{function_name}' using '{expected_metadata_key}' metadata key") - - -@pytest.mark.parametrize( - "function_name, metadata_key", - [ - ("acompletion", "metadata"), - ("_ageneric_api_call_with_fallbacks", "litellm_metadata"), - ], -) -def test_handle_clientside_credential_metadata_variable_name(model_list, function_name, metadata_key): - """Test that _handle_clientside_credential uses the correct metadata variable name based on function name""" - from litellm.litellm_core_utils.batch_utils import get_router_metadata_variable_name - - router = Router(model_list=model_list) - - expected_metadata_key = get_router_metadata_variable_name(function_name=function_name) - assert expected_metadata_key == metadata_key - - deployment = { - "model_name": "gpt-4.1", - "litellm_params": {"model": "gpt-4.1", "api_key": "test_key"}, - "model_info": {"id": "original-id-456"}, - } - - kwargs = { - "api_key": "client_side_key", - "api_base": "https://api.openai.com/v1", - metadata_key: {"model_group": "gpt-4.1", "test_field": "test_value"}, - } - - result_deployment = router._handle_clientside_credential( - deployment=deployment, kwargs=kwargs, function_name=function_name - ) - - assert result_deployment.model_name == "gpt-4.1" - - assert result_deployment.litellm_params.api_key == "client_side_key" - assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1" - - print(f"✓ Success with function_name '{function_name}' correctly using '{metadata_key}' for metadata") - - -def test_handle_clientside_credential_with_responses_function(model_list): - """Test that _handle_clientside_credential works correctly with responses function name""" - router = Router(model_list=model_list) - - deployment = { - "model_name": "gpt-4.1", - "litellm_params": {"model": "gpt-4.1", "api_key": "test_key"}, - "model_info": {"id": "original-id-responses"}, - } - - kwargs = { - "api_key": "client_side_key", - "api_base": "https://api.openai.com/v1", - "litellm_metadata": { - "model_group": "gpt-4.1", - "responses_field": "responses_value", - }, - } - - result_deployment = router._handle_clientside_credential( - deployment=deployment, - kwargs=kwargs, - function_name="_ageneric_api_call_with_fallbacks", - ) - - assert isinstance(result_deployment, Deployment) - assert result_deployment.model_name == "gpt-4.1" - assert result_deployment.litellm_params.api_key == "client_side_key" - assert result_deployment.litellm_params.api_base == "https://api.openai.com/v1" - assert result_deployment.model_info.id != "original-id-responses" - assert result_deployment.model_info.original_model_id == "original-id-responses" - - assert len(router.model_list) == len(model_list) - assert router.get_deployment(model_id=result_deployment.model_info.id) is None - - print("✓ Success with _ageneric_api_call_with_fallbacks function name and litellm_metadata") - - -def test_handle_clientside_credential_still_registers_custom_pricing(model_list): - """A clientside-credential call must still price against the deployment's own - custom rate, even though the call's ephemeral deployment is never added to the - router (see LIT-7811): losing that registration would silently fall back to - public catalog pricing for every clientside-credential call on a deployment - with a custom rate configured.""" - router = Router(model_list=model_list) - deployment = { - "model_name": "gpt-4.1", - "litellm_params": { - "model": "gpt-4.1", - "api_key": "test_key", - "input_cost_per_token": 0.0001234, - "output_cost_per_token": 0.0005678, - }, - "model_info": {"id": "original-id-pricing"}, - } - kwargs = {"api_key": "client_side_key", "metadata": {"model_group": "gpt-4.1"}} - - result_deployment = router._handle_clientside_credential( - deployment=deployment, kwargs=kwargs, function_name="acompletion" - ) - - registered = litellm.model_cost.get(result_deployment.model_info.id) - assert registered is not None - assert registered["input_cost_per_token"] == 0.0001234 - assert registered["output_cost_per_token"] == 0.0005678 - - -def test_register_deployment_pricing_direct_call(): - """Direct-call unit test for the pricing-registration helper `_handle_clientside_credential` - relies on, so it prices a deployment that is deliberately never added to `self.model_list`.""" - deployment = Deployment( - model_name="gpt-4.1", - litellm_params=LiteLLM_Params( - model="gpt-4.1", - api_key="test_key", - input_cost_per_token=0.0009999, - ), - model_info=ModelInfo(id="direct-call-pricing-id"), - ) - - Router._register_deployment_pricing(deployment=deployment) - - assert litellm.model_cost["direct-call-pricing-id"]["input_cost_per_token"] == 0.0009999 - - -def test_get_metadata_variable_name_from_kwargs(model_list): - """ - Test _get_metadata_variable_name_from_kwargs method returns correct metadata variable name based on kwargs content. - """ - router = Router(model_list=model_list) - - kwargs_with_litellm_metadata = { - "litellm_metadata": {"user": "test"}, - "metadata": {"other": "data"}, - } - result = router._get_metadata_variable_name_from_kwargs(kwargs_with_litellm_metadata) - assert result == "litellm_metadata" - - kwargs_with_metadata_only = {"metadata": {"user": "test"}} - result = router._get_metadata_variable_name_from_kwargs(kwargs_with_metadata_only) - assert result == "metadata" - - kwargs_empty = {} - result = router._get_metadata_variable_name_from_kwargs(kwargs_empty) - assert result == "metadata" - - kwargs_other = { - "model": "gpt-5.5", - "messages": [{"role": "user", "content": "hello"}], - } - result = router._get_metadata_variable_name_from_kwargs(kwargs_other) - assert result == "metadata" - - -@pytest.fixture -def search_tools(): - """Fixture for search tools configuration""" - return [ - { - "search_tool_name": "test-search-tool", - "litellm_params": { - "search_provider": "perplexity", - "api_key": "test-api-key", - "api_base": "https://api.perplexity.ai", - "mode": "turbo", - }, - }, - { - "search_tool_name": "test-search-tool", - "litellm_params": { - "search_provider": "perplexity", - "api_key": "test-api-key-2", - "api_base": "https://api.perplexity.ai", - "mode": "turbo", - }, - }, - ] - - -@pytest.mark.asyncio -async def test_asearch_with_fallbacks(search_tools): - """ - Test _asearch_with_fallbacks method of Router. - - Tests that the _asearch_with_fallbacks method correctly: - - Accepts search parameters - - Calls async_function_with_fallbacks with correct configuration - - Returns SearchResponse - """ - from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult - - router = Router(search_tools=search_tools) - - mock_response = SearchResponse( - object="search", - results=[ - SearchResult( - title="Test Result", - url="https://example.com", - snippet="Test snippet content", - ) - ], - ) - - with patch.object(router, "async_function_with_fallbacks", new_callable=AsyncMock) as mock_fallbacks: - mock_fallbacks.return_value = mock_response - - async def mock_asearch(**kwargs): - return mock_response - - response = await router._asearch_with_fallbacks( - original_function=mock_asearch, - search_tool_name="test-search-tool", - query="test query", - max_results=5, - ) - - assert mock_fallbacks.called - - assert isinstance(response, SearchResponse) - assert response.object == "search" - assert len(response.results) == 1 - assert response.results[0].title == "Test Result" - - -@pytest.mark.asyncio -async def test_asearch_with_fallbacks_helper(search_tools): - """ - Test _asearch_with_fallbacks_helper method of Router. - - Tests that the _asearch_with_fallbacks_helper method correctly: - - Selects a search tool from available options - - Calls the original search function with correct provider parameters - - Returns SearchResponse - """ - from litellm.llms.base_llm.search.transformation import SearchResponse, SearchResult - - router = Router(search_tools=search_tools) - - mock_response = SearchResponse( - object="search", - results=[ - SearchResult( - title="Helper Test Result", - url="https://example.com/helper", - snippet="Helper test snippet", - ) - ], - ) - - async def mock_original_function(**kwargs): - - assert "search_provider" in kwargs - assert kwargs["search_provider"] == "perplexity" - assert "api_key" in kwargs - assert kwargs["mode"] == "turbo" - assert kwargs["query"] == "helper test query" - return mock_response - - response = await router._asearch_with_fallbacks_helper( - model="test-search-tool", - original_generic_function=mock_original_function, - query="helper test query", - max_results=3, - ) - - assert isinstance(response, SearchResponse) - assert response.object == "search" - assert len(response.results) == 1 - assert response.results[0].title == "Helper Test Result" - assert response.results[0].url == "https://example.com/helper" - - -@pytest.mark.asyncio -async def test_asearch_with_fallbacks_helper_missing_search_tool(): - """ - Test _asearch_with_fallbacks_helper raises error when search tool not found. - - Tests that the helper method raises a ValueError when the requested - search tool name doesn't exist in the router's search_tools configuration. - """ - - router = Router(model_list=[]) - - async def mock_original_function(**kwargs): - return None - - with pytest.raises(ValueError, match="Search tool 'nonexistent-tool' not found"): - await router._asearch_with_fallbacks_helper( - model="nonexistent-tool", - original_generic_function=mock_original_function, - query="test query", - ) - - -@pytest.mark.asyncio -async def test_asearch_with_fallbacks_helper_missing_search_provider(): - """ - Test _asearch_with_fallbacks_helper raises error when search_provider not configured. - - Tests that the helper method raises a ValueError when a search tool - is found but doesn't have search_provider in its litellm_params. - """ - - search_tools_bad = [ - { - "search_tool_name": "bad-tool", - "litellm_params": {"api_key": "test-key"}, - } - ] - - router = Router(search_tools=search_tools_bad) - - async def mock_original_function(**kwargs): - return None - - with pytest.raises(ValueError, match="search_provider not found in litellm_params"): - await router._asearch_with_fallbacks_helper( - model="bad-tool", - original_generic_function=mock_original_function, - query="test query", - ) - - -def test_get_first_default_fallback(): - """Test _get_first_default_fallback method""" - - model_list = [ - { - "model_name": "gpt-5-mini", - "litellm_params": {"model": "gpt-5-mini", "api_key": "fake-key"}, - } - ] - - router = Router(model_list=model_list, fallbacks=[{"*": ["gpt-5-mini"]}]) - - result = router._get_first_default_fallback() - assert result == "gpt-5-mini" - - router_no_fallbacks = Router(model_list=model_list) - result = router_no_fallbacks._get_first_default_fallback() - assert result is None - - router_no_default = Router(model_list=model_list, fallbacks=[{"gpt-5.5": ["gpt-5-mini"]}]) - result = router_no_default._get_first_default_fallback() - assert result is None - - router_empty_list = Router(model_list=model_list, fallbacks=[{"*": []}]) - result = router_empty_list._get_first_default_fallback() - assert result is None - - def test_resolve_model_name_from_model_id(): """Test resolve_model_name_from_model_id function with various scenarios""" @@ -3872,6 +2008,7 @@ def test_resolve_model_name_from_model_id_keeps_a_wildcard_name_only_for_the_ids router = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-fake"}}]) assert router.resolve_model_name_from_model_id(model_id, custom_llm_provider) == expected + def test_get_valid_args(): """Test get_valid_args static method returns valid Router.__init__ arguments""" @@ -3922,89 +2059,6 @@ def test_get_router_model_info_with_deployment_object(): assert isinstance(model_info, dict) -def test_deployment_has_budget_limits(): - router = Router(model_list=[]) - - with_budget = Deployment( - model_name="budgeted-model", - litellm_params=LiteLLM_Params( - model="openai/gpt-4o-mini", - max_budget=0.001, - budget_duration="1d", - ), - model_info=ModelInfo(id="budget-deployment-id"), - ) - without_budget = Deployment( - model_name="unbudgeted-model", - litellm_params=LiteLLM_Params(model="openai/gpt-4o-mini"), - model_info=ModelInfo(id="no-budget-deployment-id"), - ) - - assert router._deployment_has_budget_limits(deployment=with_budget) is True - assert router._deployment_has_budget_limits(deployment=without_budget) is False - - -def test_sync_deployment_budget_config(monkeypatch): - import asyncio - - monkeypatch.setattr(asyncio, "create_task", lambda coro: None) - - router = Router(model_list=[], optional_pre_call_checks=[]) - deployment = Deployment( - model_name="dynamic-budget-model", - litellm_params=LiteLLM_Params( - model="openai/gpt-4o-mini", - api_key="fake-key", - max_budget=0.000000000001, - budget_duration="1d", - ), - model_info=ModelInfo(id="runtime-budget-deployment"), - ) - - router._sync_deployment_budget_config(deployment=deployment) - - budget_limiter = router.get_router_deployment_budget_limiter() - assert budget_limiter is not None - config = budget_limiter._get_budget_config_for_deployment("runtime-budget-deployment") - assert config is not None - assert config.max_budget == 0.000000000001 - - -def test_sync_deployment_budget_config_clears_removed_limits(monkeypatch): - import asyncio - - monkeypatch.setattr(asyncio, "create_task", lambda coro: None) - - router = Router(model_list=[], optional_pre_call_checks=[]) - model_id = "runtime-budget-deployment" - budgeted = Deployment( - model_name="dynamic-budget-model", - litellm_params=LiteLLM_Params( - model="openai/gpt-4o-mini", - api_key="fake-key", - max_budget=0.000000000001, - budget_duration="1d", - ), - model_info=ModelInfo(id=model_id), - ) - unbudgeted = Deployment( - model_name="dynamic-budget-model", - litellm_params=LiteLLM_Params( - model="openai/gpt-4o-mini", - api_key="fake-key", - ), - model_info=ModelInfo(id=model_id), - ) - - router._sync_deployment_budget_config(deployment=budgeted) - budget_limiter = router.get_router_deployment_budget_limiter() - assert budget_limiter is not None - assert budget_limiter._get_budget_config_for_deployment(model_id) is not None - - router._sync_deployment_budget_config(deployment=unbudgeted) - assert budget_limiter._get_budget_config_for_deployment(model_id) is None - - def test_upsert_deployment_clears_stale_budget_config(monkeypatch): import asyncio diff --git a/tests/unit/test_router/test_router_max_parallel_requests.py b/tests/unit/test_router/test_router_max_parallel_requests.py index fce273cd1ab..1b350f31f23 100644 --- a/tests/unit/test_router/test_router_max_parallel_requests.py +++ b/tests/unit/test_router/test_router_max_parallel_requests.py @@ -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( diff --git a/tests/unit/test_router/test_router_provider_endpoints.py b/tests/unit/test_router/test_router_provider_endpoints.py index d2331c00804..a78228b8658 100644 --- a/tests/unit/test_router/test_router_provider_endpoints.py +++ b/tests/unit/test_router/test_router_provider_endpoints.py @@ -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")) ] ) diff --git a/tests/unit/test_router/test_router_retries.py b/tests/unit/test_router/test_router_retries.py index a5664c1348f..857b5052d6c 100644 --- a/tests/unit/test_router/test_router_retries.py +++ b/tests/unit/test_router/test_router_retries.py @@ -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": }]} 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") diff --git a/tests/unit/test_router/test_router_timeout.py b/tests/unit/test_router/test_router_timeout.py index 8fc32620f71..3fafb9e0f04 100644 --- a/tests/unit/test_router/test_router_timeout.py +++ b/tests/unit/test_router/test_router_timeout.py @@ -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=[ { diff --git a/tests/unit/test_router_block_helpers.py b/tests/unit/test_router_block_helpers.py index c89d514b166..431ce5c25ee 100644 --- a/tests/unit/test_router_block_helpers.py +++ b/tests/unit/test_router_block_helpers.py @@ -45,13 +45,3 @@ class TestAreAllDeploymentsBlocked: def test_empty_list_returns_false(self): router = _make_router("gpt-4o") assert router.are_all_deployments_blocked([]) is False - - -class TestIsModelFullyBlocked: - def test_all_deployments_blocked_returns_true(self): - router = _make_router("gpt-4o", blocked=True) - assert router._is_model_fully_blocked("gpt-4o") is True - - def test_unblocked_deployment_returns_false(self): - router = _make_router("gpt-4o", blocked=False) - assert router._is_model_fully_blocked("gpt-4o") is False diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index 724c07f8c31..c0c8a0655fe 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -29,12 +29,10 @@ import respx import litellm from litellm import Router -from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE from litellm.litellm_core_utils.ptu_pricing import ptu_config_error from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.llms.openai_like.model_info import MODEL_INFO_REFRESH_SECONDS from litellm.proxy.utils import create_model_info_response from litellm.types.router import ( Deployment, @@ -117,13 +115,15 @@ async def test_discovered_limits_survive_deployment_growth_and_removal( ) for deployment in deployments[initial_count:]: router.add_deployment(deployment) - await router._arefresh_deployment_model_info(router.model_list[-1], client=handler) + router.cache.in_memory_cache.flush_cache() + await router.arefresh_model_info(client=handler) assert all( router.get_configured_token_limits(deployment.model_name) == (4096, 4096) for deployment in deployments ) for deployment in deployments[-2:]: router.delete_deployment(deployment.model_info.id or "") - await router._arefresh_deployment_model_info(router.model_list[0], client=handler) + router.cache.in_memory_cache.flush_cache() + await router.arefresh_model_info(client=handler) assert all( router.get_configured_token_limits(deployment.model_name) == (4096, 4096) for deployment in deployments[:-2] ) @@ -167,8 +167,9 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch await handler.client.aclose() async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: handler.client = client - await router._arefresh_deployment_model_info(router.model_list[0], client=handler) + await router.arefresh_model_info(client=handler) assert router.get_configured_token_limits("local") == (None, None) + router.cache.in_memory_cache.flush_cache() await router.arefresh_model_info(client=handler) assert router.get_configured_token_limits("local") == (2048, 2048) _invalidate_model_cost_lowercase_map() @@ -204,7 +205,7 @@ async def test_discovery_sends_the_secret_behind_an_oidc_api_key( await handler.client.aclose() async with httpx.AsyncClient(transport=httpx.MockTransport(respond)) as client: handler.client = client - await router._arefresh_deployment_model_info(router.model_list[0], client=handler) + await router.arefresh_model_info(client=handler) assert authorizations == ["Bearer local-secret"] _invalidate_model_cost_lowercase_map() @@ -299,54 +300,6 @@ async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypa _invalidate_model_cost_lowercase_map() -async def test_discovered_limits_expire_after_the_last_successful_refresh(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) - clock: Final = Mock(return_value=0.0) - router: Final = Router( - model_list=[ - { - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://expiry.test/v1", - "api_key": "local-key", - }, - "model_info": {"id": "expiring-discovery"}, - } - ] - ) - router._discovered_model_info_cache = InMemoryCache(clock=clock, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS) - responses: Final = iter( - ( - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}), - httpx.Response(503), - ) - ) - handler: Final = AsyncHTTPHandler() - await handler.client.aclose() - async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: next(responses))) as client: - handler.client = client - await router.arefresh_model_info(client=handler) - clock.return_value = MODEL_INFO_REFRESH_SECONDS - router.cache.in_memory_cache.flush_cache() - await router.arefresh_model_info(client=handler) - clock.return_value = 2 * MODEL_INFO_REFRESH_SECONDS + 1 - router.cache.in_memory_cache.flush_cache() - await router.arefresh_model_info(client=handler) - assert router.get_configured_token_limits("local") == (8192, 8192) - group: Final = router.get_model_group_info("local") - assert group is not None - assert group.max_input_tokens == 8192 - clock.return_value = 3 * MODEL_INFO_REFRESH_SECONDS + 1 - await router.arefresh_model_info(client=handler) - assert router.get_configured_token_limits("local") == (None, None) - expired_group: Final = router.get_model_group_info("local") - assert expired_group is not None - assert expired_group.max_input_tokens is None - _invalidate_model_cost_lowercase_map() - - @pytest.mark.parametrize("provider", ("hosted_vllm", "openai", "openai_like", "text-completion-openai")) async def test_discovered_limits_are_isolated_overridable_and_refreshable( provider: str, monkeypatch: pytest.MonkeyPatch @@ -390,11 +343,6 @@ async def test_discovered_limits_are_isolated_overridable_and_refreshable( assert listing is not None assert listing.max_input_tokens == 8192 assert router.get_configured_token_limits("local") == (8192, 8192) - assert router._deployment_max_input_tokens("local", router.model_list[1]) == 4096 - allowed: Final = router._pre_call_checks( - model="local", healthy_deployments=router.model_list, input="prompt", input_token_count=5000 - ) - assert [deployment["model_info"]["id"] for deployment in allowed] == ["one"] assert router.model_list[0]["model_info"].get("max_input_tokens") is None assert litellm.model_cost[f"{provider}/org/local-model"].get("max_input_tokens") is None router.cache.in_memory_cache.flush_cache() @@ -892,32 +840,19 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override(): litellm.model_cost[backend_model] = backend_entry _invalidate_model_cost_lowercase_map() - router = Router(model_list=[]) - with patch.object(Router, "_add_deployment", lambda self, deployment: deployment): - router._create_deployment( - deployment_info={}, - _model_name="chatgpt/gpt-5.4", - _litellm_params={ - "model": "gpt-5.4", - "custom_llm_provider": "chatgpt", - }, - _model_info={ - "id": "chatgpt-shared-mode-base", - "mode": "responses", - }, - ) - router._create_deployment( - deployment_info={}, - _model_name="chatgpt/gpt-5.4-medium", - _litellm_params={ - "model": "gpt-5.4", - "custom_llm_provider": "chatgpt", - }, - _model_info={ - "id": "chatgpt-shared-mode-alias", - "mode": "chat", - }, - ) + Router( + model_list=[ + { + "model_name": model_name, + "litellm_params": {"model": "gpt-5.4", "custom_llm_provider": "chatgpt"}, + "model_info": {"id": model_id, "mode": mode}, + } + for model_name, model_id, mode in ( + ("chatgpt/gpt-5.4", "chatgpt-shared-mode-base", "responses"), + ("chatgpt/gpt-5.4-medium", "chatgpt-shared-mode-alias", "chat"), + ) + ] + ) assert litellm.model_cost[backend_model]["mode"] == "responses" assert "mode" in litellm.model_cost[backend_model] @@ -940,7 +875,7 @@ def test_should_not_downgrade_chatgpt_shared_key_mode_with_alias_override(): @pytest.mark.parametrize("rate", (0.0, 0.004)) @pytest.mark.parametrize("explicit_base", (False, True)) @pytest.mark.parametrize("provider", (None, "bedrock")) -@pytest.mark.parametrize("registration_path", ("direct", "startup", "reload")) +@pytest.mark.parametrize("registration_path", ("startup", "reload")) def test_bedrock_deployment_inherits_invoked_profile_around_custom_prices( monkeypatch: pytest.MonkeyPatch, field: str, @@ -986,12 +921,9 @@ def test_bedrock_deployment_inherits_invoked_profile_around_custom_prices( } ), ) - if registration_path == "direct": - Router._register_deployment_pricing(deployment) - else: - router: Final = Router(model_list=[deployment.to_json(exclude_none=True)]) - if registration_path == "reload": - router.replay_model_cost_registrations() + router: Final = Router(model_list=[deployment.to_json(exclude_none=True)]) + if registration_path == "reload": + router.replay_model_cost_registrations() response: Final = ModelResponse( model=model, usage=Usage( @@ -1260,50 +1192,6 @@ def test_partial_custom_pricing_bills_explicit_cache_hits( _restore_model_cost_entries(model_keys) -def test_inherit_builtin_cache_pricing_fills_only_missing_fields(): - """Direct unit test of the helper: missing cache fields are filled from the - backend model's built-in entry, while an explicitly set cache field and the - user's input/output pricing are left untouched. - """ - backend_model = "anthropic/claude-sonnet-4-5-20250929" - builtin_info = litellm.get_model_info(model=backend_model) - builtin_cache_create = builtin_info["cache_creation_input_token_cost"] - builtin_cache_read = builtin_info["cache_read_input_token_cost"] - assert builtin_cache_create is not None and builtin_cache_create > 0 - assert builtin_cache_read is not None and builtin_cache_read > 0 - - explicit_cache_read = builtin_cache_read + 1 - model_info = { - "input_cost_per_token": 0.000003, - "cache_read_input_token_cost": explicit_cache_read, - } - - Router._inherit_builtin_cache_pricing( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider="anthropic", - ) - - assert model_info["input_cost_per_token"] == 0.000003 - assert model_info["cache_read_input_token_cost"] == explicit_cache_read - assert model_info["cache_creation_input_token_cost"] == builtin_cache_create - - -def test_inherit_builtin_cache_pricing_noop_for_unknown_backend(): - """No canonical entry for the backend model means the helper leaves the - passed-in dict unchanged rather than raising. - """ - model_info = {"input_cost_per_token": 0.000003} - - Router._inherit_builtin_cache_pricing( - model_info=model_info, - backend_model="this-backend-model-does-not-exist-x9y8z7", - custom_llm_provider=None, - ) - - assert model_info == {"input_cost_per_token": 0.000003} - - _TIER_BACKEND_MODEL: Final = "tier-priced-backend" _TIER_BACKEND_KEY: Final = f"openai/{_TIER_BACKEND_MODEL}" _CUSTOM_STANDARD_INPUT_RATE: Final = 0.00011 @@ -1349,86 +1237,6 @@ def _register_azure_tier_backend() -> None: _invalidate_model_cost_lowercase_map() -def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> None: - model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) - } - try: - _register_tier_backend() - model_info: Final = { - "id": "custom-priced-tier-deployment", - "input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE, - "output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE, - "output_cost_per_token_ultrafast": 0.00999, - } - - Router._inherit_builtin_service_tier_pricing( - model_info=model_info, - backend_model=_TIER_BACKEND_MODEL, - custom_llm_provider="openai", - ) - - assert model_info == { - "id": "custom-priced-tier-deployment", - "input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE, - "output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE, - "input_cost_per_token_ultrafast": _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"], - "output_cost_per_token_ultrafast": 0.00999, - "input_cost_per_token_priority": _TIER_BACKEND_ENTRY["input_cost_per_token_priority"], - "output_cost_per_token_priority": _TIER_BACKEND_ENTRY["output_cost_per_token_priority"], - "input_cost_per_token_flex": _TIER_BACKEND_ENTRY["input_cost_per_token_flex"], - "output_cost_per_token_flex": _TIER_BACKEND_ENTRY["output_cost_per_token_flex"], - "input_cost_per_token_balanced": _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"], - "output_cost_per_token_balanced": _TIER_BACKEND_ENTRY["output_cost_per_token_balanced"], - "cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY["cache_read_input_token_cost_ultrafast"], - "input_cost_per_token_above_272k_tokens_ultrafast": _TIER_BACKEND_ENTRY[ - "input_cost_per_token_above_272k_tokens_ultrafast" - ], - "output_cost_per_token_above_272k_tokens_ultrafast": _TIER_BACKEND_ENTRY[ - "output_cost_per_token_above_272k_tokens_ultrafast" - ], - } - finally: - _restore_model_cost_entries(model_cost_entries) - litellm.get_model_info.cache_clear() - - -def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend() -> None: - model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) - } - try: - _register_tier_backend() - model_info_without_base_rate: Final = { - "id": "custom-priced-no-base-rate", - "input_cost_per_token_ultrafast": 0.00031, - } - expected_without_base_rate: Final = copy.deepcopy(model_info_without_base_rate) - Router._inherit_builtin_service_tier_pricing( - model_info=model_info_without_base_rate, - backend_model=_TIER_BACKEND_MODEL, - custom_llm_provider="openai", - ) - - model_info_with_unknown_backend: Final = { - "id": "custom-priced-unknown-backend", - "input_cost_per_token": _CUSTOM_STANDARD_INPUT_RATE, - "output_cost_per_token": _CUSTOM_STANDARD_OUTPUT_RATE, - } - expected_with_unknown_backend: Final = copy.deepcopy(model_info_with_unknown_backend) - Router._inherit_builtin_service_tier_pricing( - model_info=model_info_with_unknown_backend, - backend_model="tier-priced-backend-unknown", - custom_llm_provider="openai", - ) - - assert model_info_without_base_rate == expected_without_base_rate - assert model_info_with_unknown_backend == expected_with_unknown_backend - finally: - _restore_model_cost_entries(model_cost_entries) - litellm.get_model_info.cache_clear() - - def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() -> None: model_id: Final = "tier-priced-deployment" model_cost_entries: Final = { @@ -1644,8 +1452,8 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay + 100 * _AZURE_TIER_BACKEND_ENTRY["output_cost_per_token_priority"] ) - payload: Final = Router._deployment_model_cost_payload( - deployment=Deployment( + router.add_deployment( + Deployment( model_name="azure/tier-priced-alias-from-params", litellm_params=LiteLLM_Params( model="azure/tier-priced-alias", @@ -1657,6 +1465,7 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay model_info=ModelInfo(id=payload_id), ) ) + payload: Final = litellm.model_cost[payload_id] assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["input_cost_per_token_priority"] assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["output_cost_per_token_priority"] @@ -1665,203 +1474,6 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay litellm.get_model_info.cache_clear() -@pytest.mark.parametrize( - ("model_info_base_model", "params_base_model", "model", "expected"), - ( - pytest.param( - "azure/tier-priced-model-info-base", - "azure/tier-priced-params-base", - "azure/tier-priced-deployment-alias", - "azure/tier-priced-model-info-base", - id="model-info-base-model-wins", - ), - pytest.param( - None, - "azure/tier-priced-params-base", - "azure/tier-priced-deployment-alias", - "azure/tier-priced-params-base", - id="params-base-model-fallback", - ), - pytest.param( - None, - None, - "azure/tier-priced-deployment-alias", - "azure/tier-priced-deployment-alias", - id="model-fallback", - ), - pytest.param( - "", - "azure/tier-priced-params-base", - "azure/tier-priced-deployment-alias", - "azure/tier-priced-params-base", - id="empty-model-info-base-model-falls-through", - ), - ), -) -def test_cost_map_backend_model_uses_canonical_model_precedence( - model_info_base_model: str | None, - params_base_model: str | None, - model: str, - expected: str, -) -> None: - deployment: Final = Deployment( - model_name="azure/tier-priced-cost-map-backend", - litellm_params=LiteLLM_Params(model=model, base_model=params_base_model), - model_info=ModelInfo(id="tier-priced-cost-map-backend", base_model=model_info_base_model), - ) - - assert Router._cost_map_backend_model(deployment) == expected - - -def test_inherit_builtin_base_rates_for_off_peak_fills_missing_rates(): - """Direct unit test of the helper: an entry carrying only an - off_peak_pricing block inherits the backend model's built-in base token - rates, so cost lookup via the deployment id can bill standard rates - outside the windows. - """ - backend_model = "gpt-4o-mini" - builtin_info = litellm.get_model_info(model=backend_model, custom_llm_provider="openai") - off_peak_block = { - "hours_utc": "00:00-00:00", - "input_cost_per_token": 5e-07, - "output_cost_per_token": 1e-06, - } - model_info = {"off_peak_pricing": off_peak_block} - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider="openai", - ) - - assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"] - assert model_info["output_cost_per_token"] == builtin_info["output_cost_per_token"] - assert model_info["off_peak_pricing"] == off_peak_block - - -def test_inherit_builtin_base_rates_for_off_peak_carries_threshold_rates(): - """A backend with above-threshold pricing hands the whole rate structure to - the deployment entry, so peak-hour billing of large prompts through that - entry matches the shared backend entry instead of flattening to the base - rate. - """ - backend_model = "gemini/gemini-2.5-pro" - builtin_info = litellm.get_model_info(model=backend_model) - assert builtin_info["input_cost_per_token_above_200k_tokens"] is not None - - model_info = { - "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, - } - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider="gemini", - ) - - assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"] - assert ( - model_info["input_cost_per_token_above_200k_tokens"] == builtin_info["input_cost_per_token_above_200k_tokens"] - ) - assert ( - model_info["output_cost_per_token_above_200k_tokens"] == builtin_info["output_cost_per_token_above_200k_tokens"] - ) - - -def test_inherit_builtin_base_rates_for_off_peak_carries_companion_billing_fields(): - """Billing rules that are not literal cost rates, like the web search - billing unit, must ride along, or grounding and regional uplifts would - bill differently through the deployment entry than through the shared - backend entry. - """ - backend_model = "gemini-3-pro-image" - raw_entry = litellm.model_cost[backend_model] - assert raw_entry.get("web_search_billing_unit") is not None - - model_info = { - "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, - } - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider=None, - ) - - assert model_info["web_search_billing_unit"] == raw_entry["web_search_billing_unit"] - assert model_info["input_cost_per_token"] == raw_entry["input_cost_per_token"] - - -def test_inherit_builtin_base_rates_for_off_peak_tiered_only_backend_stores_no_zero(): - """A tiered-only backend has no flat token rates; get_model_info synthesizes - zeros for them, and storing those would mark the deployment explicitly - priced free. The tier table itself must carry over as an isolated copy so - mutating the deployment entry never touches the shared cost map. - """ - backend_model = "dashscope/qwen-flash" - raw_tiers = litellm.model_cost[backend_model]["tiered_pricing"] - - model_info = { - "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, - } - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model=backend_model, - custom_llm_provider="dashscope", - ) - - assert model_info.get("input_cost_per_token") != 0 - assert model_info.get("output_cost_per_token") != 0 - assert model_info["tiered_pricing"] == raw_tiers - assert model_info["tiered_pricing"] is not raw_tiers - assert model_info["tiered_pricing"][0] is not raw_tiers[0] - - original_first_tier = copy.deepcopy(raw_tiers[0]) - model_info["tiered_pricing"][0]["input_cost_per_token"] = 123.0 - assert raw_tiers[0] == original_first_tier - - -def test_inherit_builtin_base_rates_for_off_peak_leaves_explicit_rates_alone(): - """An entry that sets its own base rate beside the block already counts as - a full custom pricing entry; the helper must not mix builtin rates into it. - """ - model_info = { - "off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}, - "input_cost_per_token": 3e-06, - } - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=model_info, - backend_model="gpt-4o-mini", - custom_llm_provider="openai", - ) - - assert model_info["input_cost_per_token"] == 3e-06 - assert "output_cost_per_token" not in model_info - - -def test_inherit_builtin_base_rates_for_off_peak_noop_without_block_or_backend(): - """Nothing happens without an off_peak_pricing block, and an unmapped - backend model leaves the entry unchanged rather than raising. - """ - plain_info = {"id": "dep-1"} - Router._inherit_builtin_base_rates_for_off_peak( - model_info=plain_info, - backend_model="gpt-4o-mini", - custom_llm_provider="openai", - ) - assert plain_info == {"id": "dep-1"} - - off_peak_info = {"off_peak_pricing": {"hours_utc": "00:00-00:00", "input_cost_per_token": 5e-07}} - Router._inherit_builtin_base_rates_for_off_peak( - model_info=off_peak_info, - backend_model="this-backend-model-does-not-exist-x9y8z7", - custom_llm_provider=None, - ) - assert "input_cost_per_token" not in off_peak_info - - def test_custom_pricing_field_denylist_covers_all_builtin_pricing_fields(): """The shared-backend-key stripping in Router relies on CustomPricingLiteLLMParams enumerating every per-deployment pricing field. @@ -2635,28 +2247,6 @@ def test_repointing_a_deployment_drops_its_previous_backend_key(monkeypatch): _invalidate_model_cost_lowercase_map() -@pytest.mark.parametrize( - "model, custom_llm_provider, expected", - [ - ("gpt-4o", None, ("gpt-4o",)), - ("gpt-4o", "openai", ("openai/gpt-4o",)), - ("openai/gpt-4o", None, ("openai/gpt-4o",)), - ("responses/gpt-4o", "openai", ("openai/responses/gpt-4o", "openai/gpt-4o")), - ("responses/gpt-4o", None, ("responses/gpt-4o", "gpt-4o")), - ], -) -def test_backend_cost_map_keys_matches_what_registration_writes(model, custom_llm_provider, expected): - """ - The withdrawal path drops exactly the keys the registration wrote, so the two - have to agree on the provider prefix and on the responses/ alias. The first - key is also the one the registration uses as the shared backend key, so its - position is load-bearing rather than incidental. - """ - keys = Router._backend_cost_map_keys(model=model, custom_llm_provider=custom_llm_provider) - assert keys == expected - assert keys[0] == (model if custom_llm_provider is None else f"{custom_llm_provider}/{model}") - - def test_mantle_chat_deployment_does_not_rewrite_the_shared_responses_row(): """ A Bedrock Mantle deployment that pins ``mode: chat`` routes its own chat @@ -3135,55 +2725,34 @@ def test_a_reload_rebuilds_exactly_what_a_fresh_boot_registered() -> None: _invalidate_model_cost_lowercase_map() -def test_replay_model_cost_registrations_survives_a_malformed_deployment(): - """ - The rebuild reads whatever dicts are sitting in model_list, so one entry that - cannot be rebuilt into a Deployment must not stop the rest being restored. - """ - saved_model_cost = litellm.model_cost - try: - router = Router( - model_list=[ - { - "model_name": "healthy", - "litellm_params": {"model": "hosted_vllm/healthy-backend"}, - "model_info": {"id": "healthy-id", "max_input_tokens": 777}, - } - ], - ) - router.model_list.insert(0, {"litellm_params": {}}) - - litellm.model_cost = {"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}} - _invalidate_model_cost_lowercase_map() - router.replay_model_cost_registrations() - - assert litellm.model_cost["healthy-id"]["max_input_tokens"] == 777 - finally: - litellm.model_cost = saved_model_cost - _invalidate_model_cost_lowercase_map() - - def test_deployment_model_cost_payload_folds_in_litellm_params_pricing(): """ Custom pricing is configured on litellm_params but has to land in the cost-map entry, and setting it pulls in the built-in cache pricing for the backend model. Both are what make the entry reproducible from a deployment. """ - payload = Router._deployment_model_cost_payload( - deployment=Deployment( - model_name="priced", - litellm_params=LiteLLM_Params( - model="gemini/gemini-2.5-pro", - input_cost_per_token=0.000123, - ), - model_info=ModelInfo(id="payload-id", max_input_tokens=4242), + model_keys = {"payload-id": copy.deepcopy(litellm.model_cost.get("payload-id"))} + try: + Router( + model_list=[ + Deployment( + model_name="priced", + litellm_params=LiteLLM_Params( + model="gemini/gemini-2.5-pro", + input_cost_per_token=0.000123, + ), + model_info=ModelInfo(id="payload-id", max_input_tokens=4242), + ).to_json(exclude_none=True) + ] ) - ) + payload = litellm.model_cost["payload-id"] - assert payload["id"] == "payload-id" - assert payload["max_input_tokens"] == 4242 - assert payload["input_cost_per_token"] == 0.000123 - assert payload["cache_read_input_token_cost"] > 0 + assert payload["id"] == "payload-id" + assert payload["max_input_tokens"] == 4242 + assert payload["input_cost_per_token"] == 0.000123 + assert payload["cache_read_input_token_cost"] > 0 + finally: + _restore_model_cost_entries(model_keys) def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() -> None: @@ -3193,18 +2762,21 @@ def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() - } try: _register_tier_backend() - payload: Final = Router._deployment_model_cost_payload( - deployment=Deployment( - model_name="tier-priced-payload", - litellm_params=LiteLLM_Params( - model=_TIER_BACKEND_MODEL, - custom_llm_provider="openai", - input_cost_per_token=_CUSTOM_STANDARD_INPUT_RATE, - output_cost_per_token=_CUSTOM_STANDARD_OUTPUT_RATE, - ), - model_info=ModelInfo(id=model_id), - ) + Router( + model_list=[ + Deployment( + model_name="tier-priced-payload", + litellm_params=LiteLLM_Params( + model=_TIER_BACKEND_MODEL, + custom_llm_provider="openai", + input_cost_per_token=_CUSTOM_STANDARD_INPUT_RATE, + output_cost_per_token=_CUSTOM_STANDARD_OUTPUT_RATE, + ), + model_info=ModelInfo(id=model_id), + ).to_json(exclude_none=True) + ] ) + payload: Final = litellm.model_cost[model_id] assert payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"] assert payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"] @@ -3227,15 +2799,14 @@ def test_register_deployment_in_model_cost_writes_both_key_families(): "hosted_vllm/both-families-backend": copy.deepcopy(litellm.model_cost.get("hosted_vllm/both-families-backend")), } try: - Router._register_deployment_in_model_cost( - model_id="both-families-id", - model_info={ - "id": "both-families-id", - "max_input_tokens": 999, - "litellm_provider": "hosted_vllm", - }, - model="hosted_vllm/both-families-backend", - custom_llm_provider=None, + Router( + model_list=[ + { + "model_name": "both-families", + "litellm_params": {"model": "hosted_vllm/both-families-backend"}, + "model_info": {"id": "both-families-id", "max_input_tokens": 999}, + } + ] ) assert litellm.model_cost["both-families-id"]["max_input_tokens"] == 999 @@ -3345,8 +2916,6 @@ def test_replay_live_router_model_cost_rebuilds_every_live_router(): A process can hold more than one Router, so the rebuild has to fan out across all of them rather than restoring whichever one happens to be reachable. """ - from litellm.router import _replay_live_router_model_cost - saved_model_cost = litellm.model_cost try: first = Router( @@ -3368,9 +2937,7 @@ def test_replay_live_router_model_cost_rebuilds_every_live_router(): ], ) - litellm.model_cost = {"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}} - _invalidate_model_cost_lowercase_map() - _replay_live_router_model_cost() + _simulate_price_data_reload({"gpt-4o": {"litellm_provider": "openai", "mode": "chat"}}) assert litellm.model_cost["first-id"]["max_input_tokens"] == 111 assert litellm.model_cost["second-id"]["max_input_tokens"] == 222 @@ -3439,61 +3006,6 @@ def test_strategy_router_alias_pricing_never_enters_model_cost(monkeypatch): _invalidate_model_cost_lowercase_map() -def test_inherit_builtin_tiered_output_rate_fills_the_backend_flat_rate(): - """ - A deployment entry whose custom tiers publish only input rates would bill - completions at 0, so the backend model's flat output rate is copied in at - registration. - """ - model_info = {"tiered_pricing": [{"range": [0, 3000], "input_cost_per_token": 3.25e-07}]} - - Router._inherit_builtin_tiered_output_rate( - model_info=model_info, - backend_model="claude-haiku-4-5", - custom_llm_provider="anthropic", - ) - - backend_rate = litellm.get_model_info(model="claude-haiku-4-5", custom_llm_provider="anthropic")[ - "output_cost_per_token" - ] - assert backend_rate > 0 - assert model_info["output_cost_per_token"] == backend_rate - - -def test_inherit_builtin_tiered_output_rate_never_stores_a_synthesized_zero(): - """ - Regression: get_model_info reports output_cost_per_token 0 for a backend that - only publishes tiered rates (e.g. dashscope/qwen-flash), and storing that zero - would mark the deployment as explicitly priced free. - """ - backend_info = litellm.get_model_info(model="qwen-flash", custom_llm_provider="dashscope") - assert backend_info["output_cost_per_token"] == 0 - - model_info = {"tiered_pricing": [{"range": [0, 3000], "input_cost_per_token": 3.25e-07}]} - Router._inherit_builtin_tiered_output_rate( - model_info=model_info, - backend_model="qwen-flash", - custom_llm_provider="dashscope", - ) - - assert "output_cost_per_token" not in model_info - - -def test_inherit_builtin_tiered_output_rate_leaves_a_user_rate_alone(): - model_info = { - "tiered_pricing": [{"range": [0, 3000], "input_cost_per_token": 3.25e-07}], - "output_cost_per_token": 9e-07, - } - - Router._inherit_builtin_tiered_output_rate( - model_info=model_info, - backend_model="claude-haiku-4-5", - custom_llm_provider="anthropic", - ) - - assert model_info["output_cost_per_token"] == 9e-07 - - # --- a config.yaml PTU deployment must not also bill per token ------------------ _PTU_MODEL_INFO = { diff --git a/tests/unit/test_router_order_fallback.py b/tests/unit/test_router_order_fallback.py index 6dd445e2b5d..3ea176b9790 100644 --- a/tests/unit/test_router_order_fallback.py +++ b/tests/unit/test_router_order_fallback.py @@ -561,34 +561,6 @@ async def test_router_order_fallback_retries_keep_target_order(): assert seen_target_orders.count(2) >= 2 -@pytest.mark.asyncio -async def test_generic_api_call_strips_target_order_from_provider_kwargs(): - captured: Final = {} - - async def _fake_provider(**provider_kwargs): - captured.update(provider_kwargs) - return "ok" - - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key", "order": 2}, - "model_info": {"id": "2"}, - }, - ], - ) - response = await router._ageneric_api_call_with_fallbacks_helper( - model="test-model", - original_generic_function=_fake_provider, - _target_order=2, - messages=[{"role": "user", "content": "hi"}], - ) - assert response == "ok" - assert captured["model"] == "gpt-4o" - assert "_target_order" not in captured - - @pytest.mark.asyncio async def test_text_completion_order_fallback_hop_does_not_send_target_order_upstream(): upstream_bodies: Final[list[dict]] = [] @@ -793,24 +765,6 @@ async def test_responses_order_fallback_hop_keeps_the_encrypted_reasoning_the_sa ] -def test_fallback_hop_reads_the_deployment_that_just_failed_from_the_metadata_bucket_it_writes(): - router: Final = _two_openai_orders_on_one_encryption_boundary_router() - order_2: Final = router.get_deployment(model_id="openai-order-2").model_dump(exclude_none=True) - hop_input: Final = _responses_history_with_order_1_reasoning() - hop_kwargs: Final = { - "model": "gpt-6-astra", - "input": hop_input, - "fallback_depth": 1, - "metadata": {"model_info": {"id": "openai-order-1"}}, - "litellm_metadata": {"previous_models": [{"deployment_id": None}]}, - } - - router._update_kwargs_with_deployment(deployment=order_2, kwargs=hop_kwargs) - - assert hop_input == _responses_history_with_order_1_reasoning() - assert hop_kwargs["metadata"]["model_info"]["id"] == "openai-order-2" - - def test_check_non_standard_fallback_format(): from litellm.router_utils.fallback_event_handlers import ( check_non_standard_fallback_format, diff --git a/tests/unit/test_router_per_deployment_num_retries.py b/tests/unit/test_router_per_deployment_num_retries.py index b08ae25c438..76d30e1c58a 100644 --- a/tests/unit/test_router_per_deployment_num_retries.py +++ b/tests/unit/test_router_per_deployment_num_retries.py @@ -84,105 +84,6 @@ async def test_routed_images_do_not_multiply_sdk_and_router_retries( class TestPerDeploymentNumRetries: """Test that per-deployment num_retries in litellm_params is correctly used.""" - def test_set_deployment_num_retries_on_exception(self): - """ - Test that _set_deployment_num_retries_on_exception sets num_retries - on the exception from the deployment's litellm_params. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/gpt-4", - "api_key": "test-key", - "num_retries": 5, # Per-deployment setting - }, - }, - ], - num_retries=1, # Global setting - ) - - deployment = router.model_list[0] - - # Create a mock exception without num_retries - class MockException(Exception): - pass - - exc = MockException("test error") - assert not hasattr(exc, "num_retries") or exc.num_retries is None - - # Call the helper - router._set_deployment_num_retries_on_exception(exc, deployment) - - # Verify num_retries was set from deployment - assert exc.num_retries == 5 - - def test_set_deployment_num_retries_does_not_override_existing(self): - """ - Test that _set_deployment_num_retries_on_exception does NOT override - if exception already has num_retries set. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/gpt-4", - "api_key": "test-key", - "num_retries": 5, - }, - }, - ], - num_retries=1, - ) - - deployment = router.model_list[0] - - # Create an exception that already has num_retries - class MockException(Exception): - num_retries = 10 # Already set - - exc = MockException("test error") - - # Call the helper - router._set_deployment_num_retries_on_exception(exc, deployment) - - # Verify num_retries was NOT overridden - assert exc.num_retries == 10 - - def test_deployment_without_num_retries(self): - """ - Test that _set_deployment_num_retries_on_exception does nothing - if deployment has no num_retries set. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/gpt-4", - "api_key": "test-key", - # No num_retries set - }, - }, - ], - num_retries=3, - ) - - deployment = router.model_list[0] - - class MockException(Exception): - pass - - exc = MockException("test error") - - # Call the helper - router._set_deployment_num_retries_on_exception(exc, deployment) - - # Verify num_retries was not set (deployment has no num_retries) - assert not hasattr(exc, "num_retries") or exc.num_retries is None - def test_request_level_num_retries_takes_precedence(self): """ Test that request-level num_retries (passed in kwargs) is still respected. @@ -230,39 +131,6 @@ class TestPerDeploymentNumRetries: router._update_kwargs_before_fallbacks(model="test-model", kwargs=kwargs) assert kwargs.get("num_retries") is None - def test_set_deployment_num_retries_with_string_value(self): - """ - Test that _set_deployment_num_retries_on_exception handles string values - from environment variables correctly. - GitHub Issue: #19481 - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/gpt-4", - "api_key": "test-key", - "num_retries": "6", # String value (as from env var) - }, - }, - ], - num_retries=0, # Global setting - ) - - deployment = router.model_list[0] - - class MockException(Exception): - pass - - exc = MockException("test error") - - # Call the helper - router._set_deployment_num_retries_on_exception(exc, deployment) - - # Verify num_retries was converted from string to int - assert exc.num_retries == 6 - class TestNumRetriesNoneGuard: """ @@ -330,7 +198,7 @@ class TestNumRetriesNoneGuard: assert response.choices[0].message.content == "ok" @pytest.mark.asyncio - async def test_async_function_with_retries_none_falls_back_to_zero(self): + async def test_acompletion_num_retries_none_with_router_none_falls_back_to_zero(self): """ When both the per-request value AND the router-level setting are None (e.g. after Router.update_settings(num_retries=None), #28126), num_retries must @@ -339,49 +207,15 @@ class TestNumRetriesNoneGuard: router = self._mock_router(num_retries=0) router.num_retries = None # simulate update_settings(num_retries=None) - async def failing_fn(*args, **kwargs): - raise litellm.RateLimitError( - message="boom", model="mock-model", llm_provider="openai" - ) - with patch("asyncio.sleep", return_value=None): with pytest.raises(litellm.RateLimitError): - await router.async_function_with_retries( - original_function=failing_fn, + await router.acompletion( model="mock-model", messages=[{"role": "user", "content": "hi"}], num_retries=None, + mock_response=litellm.RateLimitError(message="boom", model="mock-model", llm_provider="openai"), ) - @pytest.mark.asyncio - async def test_async_function_with_retries_none_falls_back_to_router_default(self): - """ - A None per-request num_retries falls back to the router-level setting, so retries - still happen (original_function is invoked more than once) before the real error - is raised - proving None did not silently disable retries or crash. - """ - router = self._mock_router(num_retries=3) - calls = {"n": 0} - - async def failing_fn(*args, **kwargs): - calls["n"] += 1 - raise litellm.InternalServerError( - message="boom", model="mock-model", llm_provider="openai" - ) - - with patch("asyncio.sleep", return_value=None): - with pytest.raises(litellm.InternalServerError): - await router.async_function_with_retries( - original_function=failing_fn, - model="mock-model", - messages=[{"role": "user", "content": "hi"}], - metadata={}, # populated by acompletion in the real path; log_retry needs it - num_retries=None, - ) - - # 1 initial attempt + at least 1 retry -> proves None fell back to a positive int - assert calls["n"] >= 2 - class TestNoProviderRetryAmplification: """ @@ -435,9 +269,7 @@ class TestNoProviderRetryAmplification: counter = self._install_counting_upstream() with patch("asyncio.sleep", return_value=None): with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="mock", messages=[{"role": "user", "content": "hi"}], **call_kwargs - ) + await router.acompletion(model="mock", messages=[{"role": "user", "content": "hi"}], **call_kwargs) return counter["n"] @pytest.mark.asyncio @@ -449,9 +281,7 @@ class TestNoProviderRetryAmplification: an unfixed build sends 9 (N=2) or 36 (N=5). """ counter = self._install_counting_upstream() - router = self._router( - f"https://amp-{num_retries}.local/v1", {"num_retries": num_retries}, num_retries=1 - ) + router = self._router(f"https://amp-{num_retries}.local/v1", {"num_retries": num_retries}, num_retries=1) with patch("asyncio.sleep", return_value=None): with pytest.raises(litellm.InternalServerError): await router.acompletion(model="mock", messages=[{"role": "user", "content": "hi"}]) @@ -474,9 +304,7 @@ class TestNoProviderRetryAmplification: retries for a routed call: deployment ``num_retries=5`` plus ``max_retries=3`` still sends exactly ``6`` upstream requests. """ - router = self._router( - "https://nest-dep.local/v1", {"num_retries": 5, "max_retries": 3}, num_retries=1 - ) + router = self._router("https://nest-dep.local/v1", {"num_retries": 5, "max_retries": 3}, num_retries=1) assert await self._call_and_count(router) == 6 @pytest.mark.asyncio @@ -660,9 +488,7 @@ class TestRequestNumRetriesBeatsGlobal: ) with patch("asyncio.sleep", return_value=None): with pytest.raises(litellm.InternalServerError): - await router.acompletion( - model="mock", messages=[{"role": "user", "content": "hi"}] - ) + await router.acompletion(model="mock", messages=[{"role": "user", "content": "hi"}]) assert counter.attempts == 3 @@ -729,9 +555,7 @@ class TestRequestNumRetriesBeatsDeployment: @pytest.mark.asyncio @pytest.mark.parametrize("request_num_retries, expected_attempts", [(3, 4), (0, 1)]) - async def test_request_num_retries_overrides_deployment( - self, request_num_retries, expected_attempts - ): + async def test_request_num_retries_overrides_deployment(self, request_num_retries, expected_attempts): """ global=1, deployment=2, request=3 -> 4 attempts, and request=0 -> a single attempt. Before the fix the deployment value won, so both cases sent 3 attempts. @@ -749,9 +573,7 @@ class TestRequestNumRetriesBeatsDeployment: The deployment value keeps its place directly above litellm_settings: global=1 and deployment=2 with no request value -> 1 initial attempt + 2 retries. """ - attempts = await self._count_attempts( - global_num_retries=1, deployment_num_retries=2, request_num_retries=None - ) + attempts = await self._count_attempts(global_num_retries=1, deployment_num_retries=2, request_num_retries=None) assert attempts == 3 @pytest.mark.asyncio diff --git a/tests/unit/test_router_retry_backoff_headers.py b/tests/unit/test_router_retry_backoff_headers.py index 03c3af692ce..73d1091b609 100644 --- a/tests/unit/test_router_retry_backoff_headers.py +++ b/tests/unit/test_router_retry_backoff_headers.py @@ -2,20 +2,24 @@ Tests for router retry backoff behavior. """ +import asyncio from unittest.mock import patch import httpx import pytest +import respx import litellm from litellm import Router @pytest.mark.asyncio -async def test_retry_backoff_uses_current_exception_headers(): +async def test_retry_backoff_uses_current_exception_headers(monkeypatch: pytest.MonkeyPatch): """ Ensure retry backoff uses the current retry exception, not the initial one. """ + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() router = Router( model_list=[ { @@ -29,60 +33,36 @@ async def test_retry_backoff_uses_current_exception_headers(): num_retries=2, ) - first_error = litellm.RateLimitError( - message="Rate limited on first attempt", - model="gpt-3.5-turbo", - llm_provider="openai", - ) - first_error.litellm_response_headers = httpx.Headers({"retry-after": "1"}) + retry_after_values = (1, 15, 30) + replies = [ + httpx.Response( + 429, + headers={"retry-after": str(retry_after)}, + json={"error": {"message": f"Rate limited, retry after {retry_after}", "type": "rate_limit"}}, + ) + for retry_after in retry_after_values + ] + real_sleep = asyncio.sleep + slept: list[float] = [] - second_error = litellm.RateLimitError( - message="Rate limited on second attempt", - model="gpt-3.5-turbo", - llm_provider="openai", - ) - second_error.litellm_response_headers = httpx.Headers({"retry-after": "15"}) + async def record_sleep(delay: float, *args: object, **kwargs: object) -> None: + slept.append(delay) + await real_sleep(0) - third_error = litellm.RateLimitError( - message="Rate limited on third attempt", - model="gpt-3.5-turbo", - llm_provider="openai", - ) - third_error.litellm_response_headers = httpx.Headers({"retry-after": "30"}) + with respx.mock() as mock_router: + mock_router.post("https://api.openai.com/v1/chat/completions").mock(side_effect=replies) + with patch("asyncio.sleep", record_sleep): + with pytest.raises(litellm.RateLimitError): + await router.acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello"}], + ) + litellm.in_memory_llm_clients_cache.flush_cache() - raised_errors = [first_error, second_error, third_error] - captured_backoff_errors = [] - - async def mock_make_call(*args, **kwargs): - raise raised_errors.pop(0) - - def mock_time_to_sleep_before_retry(*args, **kwargs): - captured_backoff_errors.append(kwargs["e"]) - return 0.01 - - 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", - side_effect=mock_time_to_sleep_before_retry, - ): - with pytest.raises(litellm.RateLimitError): - await router.acompletion( - model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Hello"}], - ) - - # Router computes backoff once after the initial failure, then once per failed retry. - # With num_retries=2 and all attempts failing, that's 1 + 2 = 3 invocations. - assert len(captured_backoff_errors) == router.num_retries + 1 - assert captured_backoff_errors[0] is first_error - assert captured_backoff_errors[1] is second_error - assert captured_backoff_errors[2] is third_error + # Router backs off once after the initial failure, then once per failed retry. + # Each backoff honours that attempt's retry-after (plus jitter), so each sleep + # sits between its own header value and the next one. + assert len(slept) == router.num_retries + 1 + assert retry_after_values[0] <= slept[0] < retry_after_values[1] + assert retry_after_values[1] <= slept[1] < retry_after_values[2] + assert retry_after_values[2] <= slept[2] diff --git a/tests/unit/test_router_retry_non_retryable_errors.py b/tests/unit/test_router_retry_non_retryable_errors.py index 98a7db7a079..0ae03d18e63 100644 --- a/tests/unit/test_router_retry_non_retryable_errors.py +++ b/tests/unit/test_router_retry_non_retryable_errors.py @@ -10,60 +10,34 @@ Verifies that: Regression tests for https://github.com/BerriAI/litellm/issues/21343 """ -import asyncio -import datetime -from collections.abc import Awaitable, Callable -from typing import Final -from unittest.mock import AsyncMock, patch +from collections.abc import Iterator +import httpx import pytest +import respx import litellm from litellm import Router -from litellm.litellm_core_utils.litellm_logging import Logging -from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( - _union_duration_ms, - response_timing_metrics, -) -from litellm.litellm_core_utils.logging_utils import track_llm_api_timing -from litellm.litellm_core_utils.rules import Rules -from litellm.utils import function_setup -def _make_rate_limit_error(message="Rate limited"): - """Create a RateLimitError for testing.""" - return litellm.RateLimitError( - message=message, - llm_provider="bedrock", - model="anthropic.claude-v2", +@pytest.fixture(autouse=True) +def httpx_transport(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +def _error(status: int, message: str, code: str | None = None) -> httpx.Response: + return httpx.Response( + status, + headers={"retry-after": "0"}, + json={"error": {"message": message, "type": "invalid_request_error", "code": code}}, ) -def _make_context_window_error(message="prompt is too long: 1205821 tokens > 200000"): - """Create a ContextWindowExceededError for testing.""" - return litellm.ContextWindowExceededError( - message=message, - llm_provider="vertex_ai", - model="claude-3-opus", - ) - - -def _make_bad_request_error(message="Invalid request"): - """Create a BadRequestError for testing.""" - return litellm.BadRequestError( - message=message, - llm_provider="openai", - model="gpt-4", - ) - - -def _make_not_found_error(message="Model not found"): - """Create a NotFoundError for testing.""" - return litellm.NotFoundError( - message=message, - llm_provider="openai", - model="gpt-99", - ) +def _rate_limit(attempt: int) -> httpx.Response: + return _error(429, f"Rate limit attempt {attempt}") def _create_router(num_retries=2): @@ -74,29 +48,21 @@ def _create_router(num_retries=2): "model_name": "test-model", "litellm_params": { "model": "openai/gpt-4", - "api_key": "fake-key-1", + "api_key": f"fake-key-{i}", + "api_base": "https://api.openai.com/v1", }, - }, - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/gpt-4", - "api_key": "fake-key-2", - }, - }, + } + for i in (1, 2) ], num_retries=num_retries, ) -def _base_kwargs(): - """Return kwargs required by async_function_with_retries.""" - return { - "model": "test-model", - "messages": [{"role": "user", "content": "test"}], - "original_function": AsyncMock(), - "metadata": {}, - } +async def _call(router: Router, responses: list[httpx.Response]) -> None: + """Run one acompletion against upstream replies served in order.""" + with respx.mock() as mock_router: + mock_router.post("https://api.openai.com/v1/chat/completions").mock(side_effect=responses) + await router.acompletion(model="test-model", messages=[{"role": "user", "content": "test"}]) @pytest.mark.asyncio @@ -110,35 +76,14 @@ async def test_non_retryable_error_in_retry_loop_raises_immediately(): Expected: ContextWindowExceededError is raised, NOT RateLimitError """ router = _create_router(num_retries=2) + replies = [ + _rate_limit(1), + _error(400, "This model's maximum context length is 8192 tokens", "context_length_exceeded"), + _rate_limit(3), + ] - rate_limit_error = _make_rate_limit_error() - context_window_error = _make_context_window_error() - - call_count = 0 - - async def mock_make_call(*args, **kwargs): - nonlocal call_count - call_count += 1 - if call_count == 1: - raise rate_limit_error - else: - raise context_window_error - - with ( - patch.object(router, "make_call", side_effect=mock_make_call), - patch.object( - router, - "_async_get_healthy_deployments", - return_value=(["d1", "d2"], ["d1", "d2"]), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs), - ): - with pytest.raises(litellm.ContextWindowExceededError): - await router.async_function_with_retries( - num_retries=2, - **_base_kwargs(), - ) + with pytest.raises(litellm.ContextWindowExceededError): + await _call(router, replies) @pytest.mark.asyncio @@ -149,34 +94,8 @@ async def test_bad_request_error_in_retry_loop_raises_immediately(): """ router = _create_router(num_retries=2) - rate_limit_error = _make_rate_limit_error() - bad_request_error = _make_bad_request_error() - - call_count = 0 - - async def mock_make_call(*args, **kwargs): - nonlocal call_count - call_count += 1 - if call_count == 1: - raise rate_limit_error - else: - raise bad_request_error - - with ( - patch.object(router, "make_call", side_effect=mock_make_call), - patch.object( - router, - "_async_get_healthy_deployments", - return_value=(["d1", "d2"], ["d1", "d2"]), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs), - ): - with pytest.raises(litellm.BadRequestError): - await router.async_function_with_retries( - num_retries=2, - **_base_kwargs(), - ) + with pytest.raises(litellm.BadRequestError): + await _call(router, [_rate_limit(1), _error(400, "Invalid request"), _rate_limit(3)]) @pytest.mark.asyncio @@ -187,30 +106,10 @@ async def test_original_exception_updated_to_latest_error(): """ router = _create_router(num_retries=2) - call_count = 0 - - async def mock_make_call(*args, **kwargs): - nonlocal call_count - call_count += 1 - raise _make_rate_limit_error(f"Rate limit attempt {call_count}") - - with ( - patch.object(router, "make_call", side_effect=mock_make_call), - patch.object( - router, - "_async_get_healthy_deployments", - return_value=(["d1", "d2"], ["d1", "d2"]), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs), - ): - with pytest.raises(litellm.RateLimitError) as exc_info: - await router.async_function_with_retries( - num_retries=2, - **_base_kwargs(), - ) - # Should be the LAST error, not the first - assert "Rate limit attempt 3" in str(exc_info.value) + with pytest.raises(litellm.RateLimitError) as exc_info: + await _call(router, [_rate_limit(i) for i in range(1, 6)]) + # Should be the LAST error, not the first + assert "Rate limit attempt 3" in str(exc_info.value) @pytest.mark.asyncio @@ -220,32 +119,15 @@ async def test_retryable_errors_still_retry_normally(): configured number of times before raising. """ router = _create_router(num_retries=3) + replies = [_rate_limit(i) for i in range(1, 9)] - call_count = 0 - - async def mock_make_call(*args, **kwargs): - nonlocal call_count - call_count += 1 - raise _make_rate_limit_error(f"Rate limit attempt {call_count}") - - with ( - patch.object(router, "make_call", side_effect=mock_make_call), - patch.object( - router, - "_async_get_healthy_deployments", - return_value=(["d1", "d2"], ["d1", "d2"]), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs), - ): + with respx.mock() as mock_router: + route = mock_router.post("https://api.openai.com/v1/chat/completions").mock(side_effect=replies) with pytest.raises(litellm.RateLimitError): - await router.async_function_with_retries( - num_retries=3, - **_base_kwargs(), - ) + await router.acompletion(model="test-model", messages=[{"role": "user", "content": "test"}]) - # Initial call + 3 retries = 4 total calls - assert call_count == 4 + # Initial call + 3 retries = 4 total calls + assert route.call_count == 4 @pytest.mark.asyncio @@ -254,106 +136,12 @@ async def test_not_found_error_in_retry_loop_raises_immediately(): A 404 NotFoundError inside the retry loop should break out immediately. """ router = _create_router(num_retries=2) + replies = [_rate_limit(1), _error(404, "Model not found"), _rate_limit(3)] - rate_limit_error = _make_rate_limit_error() - not_found_error = _make_not_found_error() - - call_count = 0 - - async def mock_make_call(*args, **kwargs): - nonlocal call_count - call_count += 1 - if call_count == 1: - raise rate_limit_error - else: - raise not_found_error - - with ( - patch.object(router, "make_call", side_effect=mock_make_call), - patch.object( - router, - "_async_get_healthy_deployments", - return_value=(["d1", "d2"], ["d1", "d2"]), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - patch.object(router, "log_retry", side_effect=lambda kwargs, e: kwargs), - ): + with respx.mock() as mock_router: + route = mock_router.post("https://api.openai.com/v1/chat/completions").mock(side_effect=replies) with pytest.raises(litellm.NotFoundError): - await router.async_function_with_retries( - num_retries=2, - **_base_kwargs(), - ) + await router.acompletion(model="test-model", messages=[{"role": "user", "content": "test"}]) - # Only 2 calls: initial + first retry that hits non-retryable - assert call_count == 2 - - -@pytest.mark.asyncio -async def test_retry_attempts_accumulate_timing_in_shared_request_metadata(): - received_at: Final = datetime.datetime.now() - metadata: dict[str, object] = { - "model_group": "test-model", - "litellm_received_at": received_at, - } - logging_obj_raw, _ = function_setup( - "acompletion", - Rules(), - datetime.datetime.now(), - model="test-model", - messages=[{"role": "user", "content": "test"}], - metadata=metadata, - litellm_call_id="retry-timing-test", - is_async_call=True, - ) - assert isinstance(logging_obj_raw, Logging) - logging_obj: Final[Logging] = logging_obj_raw - attempt_numbers: list[int] = [] - metadata_ids: list[int] = [] - - @track_llm_api_timing() - async def timed_attempt(*, logging_obj: Logging, **kwargs: object) -> str: - del kwargs - attempt_numbers.append(len(attempt_numbers) + 1) - metadata_ids.append(id(logging_obj.model_call_details["litellm_params"]["metadata"])) - await asyncio.sleep(0.01) - if len(attempt_numbers) == 1: - raise _make_rate_limit_error() - return "success" - - async def invoke(original_function: Callable[..., Awaitable[str]], *args: object, **kwargs: object) -> str: - return await original_function(*args, **kwargs) - - router = _create_router(num_retries=1) - with ( - patch.object(router, "make_call", new=AsyncMock(side_effect=invoke)), - patch.object( - router, - "_async_get_healthy_deployments", - new=AsyncMock(return_value=(["d1"], ["d1"])), - ), - patch.object(router, "_time_to_sleep_before_retry", return_value=0), - ): - result = await router.async_function_with_retries( - original_function=timed_attempt, - model="test-model", - messages=[{"role": "user", "content": "test"}], - metadata=metadata, - logging_obj=logging_obj, - num_retries=1, - ) - - request_metadata: Final = logging_obj.model_call_details["litellm_params"]["metadata"] - windows: Final = request_metadata["llm_api_timing_windows"] - end_time: Final = datetime.datetime.fromtimestamp(max(window[1] for window in windows)) - timing_metrics: Final = response_timing_metrics(received_at, end_time, logging_obj) - assert result == "success" - assert attempt_numbers == [1, 2] - assert request_metadata is metadata - assert metadata_ids == [id(metadata), id(metadata)] - assert len(windows) == 2 - union_duration_ms: Final = _union_duration_ms(windows, received_at.timestamp(), end_time.timestamp()) - assert union_duration_ms is not None - total_response_time_ms: Final = (end_time.timestamp() - received_at.timestamp()) * 1000 - assert timing_metrics["litellm_overhead_time_ms"] == pytest.approx( - round(total_response_time_ms - union_duration_ms, 4) - ) + # Only 2 calls: initial + first retry that hits non-retryable + assert route.call_count == 2 diff --git a/tests/unit/test_router_retry_policy_update.py b/tests/unit/test_router_retry_policy_update.py index 6ab672d7b04..95a0d495456 100644 --- a/tests/unit/test_router_retry_policy_update.py +++ b/tests/unit/test_router_retry_policy_update.py @@ -88,9 +88,7 @@ def test_update_router_config_rejects_malformed_model_group_retry_policy(): """model_group_retry_policy is Dict[str, RetryPolicy], so each per-group policy is validated the same way.""" with pytest.raises(ValidationError): - UpdateRouterConfig( - model_group_retry_policy={"gpt-4": {"RateLimitErrorRetries": "x"}} - ) + UpdateRouterConfig(model_group_retry_policy={"gpt-4": {"RateLimitErrorRetries": "x"}}) # --------------------------------------------------------------------------- @@ -146,16 +144,6 @@ def test_set_optional_pre_call_checks_reconciles_callback_types(): assert not any(isinstance(callback, PromptCachingDeploymentCheck) for callback in litellm.callbacks) -def test_remove_optional_pre_call_check_removes_local_and_global_callbacks(): - router = _build_router() - - router.set_optional_pre_call_checks(["prompt_caching"]) - router._remove_optional_callbacks_of_type(PromptCachingDeploymentCheck) - - assert not any(type(callback) is PromptCachingDeploymentCheck for callback in (router.optional_callbacks or [])) - assert not any(type(callback) is PromptCachingDeploymentCheck for callback in litellm.callbacks) - - def test_remove_optional_pre_call_check_keeps_global_callback_for_another_router(): router_a = _build_router() router_b = _build_router() @@ -255,7 +243,9 @@ def test_update_settings_keeps_per_group_encrypted_content_affinity_when_global_ router.update_settings(optional_pre_call_checks=[]) affinity_checks: Final = [ - callback for callback in (router.optional_callbacks or []) if isinstance(callback, EncryptedContentAffinityCheck) + callback + for callback in (router.optional_callbacks or []) + if isinstance(callback, EncryptedContentAffinityCheck) ] assert len(affinity_checks) == 1 assert affinity_checks[0].enable_global_affinity is False @@ -470,9 +460,7 @@ async def test_config_update_persists_and_reads_back_retry_policy(monkeypatch): read_back = ( await proxy_server.get_config( - user_api_key_dict=UserAPIKeyAuth( - user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-9876" - ) + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-9876") ) )["router_settings"]["retry_policy"] assert read_back.BadRequestErrorRetries == 5 diff --git a/tests/unit/test_router_silent_experiment.py b/tests/unit/test_router_silent_experiment.py index fc787fcf74a..c7dd3bfa78c 100644 --- a/tests/unit/test_router_silent_experiment.py +++ b/tests/unit/test_router_silent_experiment.py @@ -4,7 +4,6 @@ import json import threading import time from collections.abc import Callable, Generator, Mapping -from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -16,8 +15,6 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.router import Router -from litellm.router import _silent_experiment_kwargs_snapshot -from litellm.router import _silent_experiment_targets class _RecordingLogger(CustomLogger): @@ -197,105 +194,6 @@ def _streaming_model_list(silent_model: object) -> list[dict[str, object]]: ] -class _NonCopyableSpan: - """Mimics an OTel Span which raises on deepcopy, forcing safe_deep_copy - to fall back to the original reference.""" - - def __deepcopy__(self, memo): - raise TypeError("OTel spans cannot be deepcopied") - - -class _FakeUserAPIKeyAuth: - """Mimics UserAPIKeyAuth which contains a parent_otel_span that is not - deepcopy-able. This is what actually causes safe_deep_copy to fail for - the metadata dict in production — safe_deep_copy handles the top-level - litellm_parent_otel_span specially (pops it before copying), but does - NOT handle user_api_key_auth.parent_otel_span inside it.""" - - def __init__(self, key_alias, parent_otel_span): - self.key_alias = key_alias - self.parent_otel_span = parent_otel_span - - def __deepcopy__(self, memo): - raise TypeError("Contains OTel span that cannot be deepcopied") - - -def test_get_silent_experiment_kwargs(): - """ - Test _get_silent_experiment_kwargs returns isolated kwargs with silent experiment metadata. - - Uses a non-copyable user_api_key_auth (mimicking the real proxy scenario) - so that safe_deep_copy falls back to the original metadata reference — - exercising the identity-check fix path. - """ - model_list = [ - { - "model_name": "gpt-3.5-turbo", - "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "fake-key"}, - }, - ] - router = Router(model_list=model_list) - mock_span = _NonCopyableSpan() - mock_auth = _FakeUserAPIKeyAuth( - key_alias="HaneefKeyNonTeamProd", - parent_otel_span=mock_span, - ) - kwargs = { - "metadata": { - "foo": "bar", - "litellm_parent_otel_span": mock_span, - "user_api_key_auth": mock_auth, - }, - "litellm_call_id": "call-123", - "stream": True, - "proxy_server_request": {"body": {"model": "test"}}, - } - result = router._get_silent_experiment_kwargs(**kwargs) - assert result["metadata"]["is_silent_experiment"] is True - assert result["metadata"]["foo"] == "bar" - assert "litellm_call_id" not in result - assert result["stream"] is True - # proxy_server_request must be preserved for spend log metadata - assert "proxy_server_request" in result - # CRITICAL: metadata must be a DIFFERENT dict object than the original, - # so that setting model_group / is_silent_experiment on the silent dict - # doesn't corrupt the primary call's metadata. - assert result["metadata"] is not kwargs["metadata"] - # OTel span must be stripped from the silent copy — it's not safe to use - # across event loops (silent experiment runs in a new event loop). - assert "litellm_parent_otel_span" not in result["metadata"] - # Original metadata must NOT be mutated — must carry the real span, - # not safe_deep_copy's temporary "placeholder" string. - assert "is_silent_experiment" not in kwargs["metadata"] - assert kwargs["metadata"]["litellm_parent_otel_span"] is mock_span - assert kwargs["metadata"]["user_api_key_auth"] is mock_auth - # Shallow copy must preserve user_api_key_auth so the silent experiment - # can attribute billing / spend logs to the correct key/team. - assert result["metadata"]["user_api_key_auth"] is mock_auth - - -def test_get_silent_experiment_kwargs_without_stream_stays_non_streaming(): - router = Router(model_list=[{"model_name": "m", "litellm_params": {"model": "gpt-3.5-turbo", "api_key": "k"}}]) - result = router._get_silent_experiment_kwargs(metadata={"foo": "bar"}, stream=False) - assert result["stream"] is False - assert "stream" not in router._get_silent_experiment_kwargs(metadata={"foo": "bar"}) - - -@pytest.mark.parametrize( - "silent_model, expected", - [ - ("shadow-a", ("shadow-a",)), - (["shadow-a", "shadow-b"], ("shadow-a", "shadow-b")), - ([], ()), - (None, ()), - (42, ()), - (["shadow-a", 42], ()), - ], -) -def test_silent_experiment_targets(silent_model, expected): - assert _silent_experiment_targets(silent_model) == expected - - @pytest.mark.asyncio async def test_streaming_shadow_is_streamed_and_drained_async(recording_logger): router = Router(model_list=_streaming_model_list("shadow-a")) @@ -594,87 +492,6 @@ def _tagged_primary_model_list() -> list[dict[str, object]]: ] -def test_silent_experiment_kwargs_snapshot_is_isolated_from_later_primary_mutations(): - metadata = {"foo": "bar"} - kwargs: dict[str, object] = {"metadata": metadata, "stream": True} - snapshot = _silent_experiment_kwargs_snapshot(kwargs) - kwargs["messages"] = [{"role": "user", "content": "added by the primary"}] - metadata["tags"] = ["primary-only"] - - assert dict(snapshot) == {"metadata": {"foo": "bar"}, "stream": True} - assert dict(_silent_experiment_kwargs_snapshot({"stream": False, "metadata": None})) == { - "stream": False, - "metadata": None, - } - - -def test_sync_shadow_gets_kwargs_snapshot_taken_before_primary_mutates_them(recording_logger): - deferred: list[Callable[[], None]] = [] - - class _DeferredThread: - def __init__(self, target, args, kwargs, daemon) -> None: - deferred.append(lambda: target(*args, **kwargs)) - - def start(self) -> None: - return None - - router = Router(model_list=_tagged_primary_model_list()) - with patch( # test-quality-ok: Router has no thread factory to inject; deferring start is the only deterministic way to expose the race - "litellm.router_python_core.threading", SimpleNamespace(Thread=_DeferredThread) - ): - response = router.completion( - model="primary-model", - messages=[{"role": "user", "content": "hi"}], - mock_response="pong", - metadata={"foo": "bar"}, - ) - assert response.choices[0].message.content == "pong" - assert len(deferred) == 1 - deferred[0]() - _wait_for_shadow_successes_sync(recording_logger, expected=1) - - shadow_successes = recording_logger.shadow_successes() - assert len(shadow_successes) == 1 - shadow_metadata = shadow_successes[0]["litellm_params"]["metadata"] - assert shadow_metadata["model_group"] == "shadow-b" - assert "primary-only" not in shadow_metadata.get("tags", []) - - -def test_sync_shadow_workers_do_not_share_metadata_with_each_other(recording_logger): - workers: list[tuple[Mapping[str, object], Callable[[], None]]] = [] - - class _DeferredThread: - def __init__(self, target, args, kwargs, daemon) -> None: - workers.append((kwargs, lambda: target(*args, **kwargs))) - - def start(self) -> None: - return None - - router = Router(model_list=_streaming_model_list(["shadow-a", "shadow-b"])) - with patch( # test-quality-ok: Router has no thread factory to inject; deferring start is the only deterministic way to expose the race - "litellm.router_python_core.threading", SimpleNamespace(Thread=_DeferredThread) - ): - router.completion( - model="primary-model", - messages=[{"role": "user", "content": "hi"}], - mock_response="pong", - metadata={"foo": "bar"}, - ) - assert len(workers) == 2 - (first_kwargs, run_first), (_, run_second) = workers - first_kwargs["metadata"].pop("foo") - run_second() - run_first() - _wait_for_shadow_successes_sync(recording_logger, expected=2) - - metadata_by_group = { - call["litellm_params"]["metadata"]["model_group"]: call["litellm_params"]["metadata"] - for call in recording_logger.shadow_successes() - } - assert metadata_by_group["shadow-b"]["foo"] == "bar" - assert "foo" not in metadata_by_group["shadow-a"] - - @pytest.mark.asyncio async def test_async_shadow_does_not_inherit_primary_deployment_tags(recording_logger): router = Router(model_list=_tagged_primary_model_list()) @@ -707,25 +524,6 @@ async def test_shadow_of_a_shadow_is_not_launched(recording_logger): assert model_groups == ["shadow-a"] -@pytest.mark.asyncio -async def test_run_silent_experiment_drains_stream_so_callbacks_fire(recording_logger): - router = Router(model_list=_streaming_model_list(None)) - silent_kwargs: Final = { - "stream": True, - "stream_options": {"include_usage": True}, - "mock_response": "pong", - "metadata": {"is_silent_experiment": True, "model_group": "shadow-b"}, - } - await router._run_silent_experiment("shadow-b", [{"role": "user", "content": "hi"}], silent_kwargs) - await _wait_for_shadow_successes(recording_logger, expected=1) - - shadow_successes = recording_logger.shadow_successes() - assert len(shadow_successes) == 1 - assert shadow_successes[0]["stream"] is True - assert shadow_successes[0]["async_complete_streaming_response"] is not None - assert silent_kwargs["stream"] is True - - @pytest.mark.asyncio async def test_router_silent_experiment_acompletion(): """ @@ -784,11 +582,7 @@ async def test_router_silent_experiment_acompletion(): # Find the silent call silent_call = next( - ( - c - for c in call_args_list - if c[1].get("metadata", {}).get("is_silent_experiment") is True - ), + (c for c in call_args_list if c[1].get("metadata", {}).get("is_silent_experiment") is True), None, ) assert silent_call is not None @@ -796,11 +590,7 @@ async def test_router_silent_experiment_acompletion(): # Find the primary call primary_call = next( - ( - c - for c in call_args_list - if not c[1].get("metadata", {}).get("is_silent_experiment") - ), + (c for c in call_args_list if not c[1].get("metadata", {}).get("is_silent_experiment")), None, ) assert primary_call is not None @@ -869,11 +659,7 @@ def test_router_silent_experiment_completion(): # Find the silent call silent_call = next( - ( - c - for c in call_args_list - if c[1].get("metadata", {}).get("is_silent_experiment") is True - ), + (c for c in call_args_list if c[1].get("metadata", {}).get("is_silent_experiment") is True), None, ) assert silent_call is not None @@ -882,35 +668,6 @@ def test_router_silent_experiment_completion(): assert silent_call[1]["metadata"]["model_group"] == "silent-model" -SILENT_EXPERIMENT_RUNNERS: Final = ( - pytest.param(lambda router, **kwargs: router._silent_experiment_completion(**kwargs), id="sync"), - pytest.param(lambda router, **kwargs: asyncio.run(router._silent_experiment_acompletion(**kwargs)), id="async"), -) - - -@pytest.mark.parametrize("run_silent_experiment", SILENT_EXPERIMENT_RUNNERS) -def test_silent_experiment_sends_shadow_request_attributed_to_the_silent_model(run_silent_experiment): - router = Router(model_list=_streaming_model_list(["shadow-a"])) - primary_metadata: Final = {"model_group": "primary-model"} - with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None) as acompletion: - run_silent_experiment( - router, - silent_model="shadow-a", - messages=[{"role": "user", "content": "hi"}], - metadata=primary_metadata, - ) - - acompletion.assert_awaited_once() - shadow_call: Final = acompletion.await_args.kwargs - assert shadow_call["model"] == "shadow-a" - assert shadow_call["messages"] == [{"role": "user", "content": "hi"}] - assert shadow_call["metadata"]["model_group"] == "shadow-a" - assert shadow_call["metadata"]["is_silent_experiment"] is True - assert primary_metadata == {"model_group": "primary-model"} - - - - _EMBEDDING_API_BASE: Final = "https://embeddings.example.test/v1" @@ -972,20 +729,6 @@ async def test_aembedding_with_silent_model_sends_provider_body_without_it( assert response.data[0]["embedding"] == [0.1, 0.2] -@pytest.mark.parametrize("run_silent_experiment", SILENT_EXPERIMENT_RUNNERS) -def test_silent_experiment_does_not_launch_from_a_shadow_request(run_silent_experiment): - router = Router(model_list=_streaming_model_list(["shadow-a"])) - with patch.object(router, "acompletion", new_callable=AsyncMock, return_value=None) as acompletion: - run_silent_experiment( - router, - silent_model="shadow-a", - messages=[{"role": "user", "content": "hi"}], - metadata={"is_silent_experiment": True}, - ) - - acompletion.assert_not_awaited() - - def _judged_router(shadow_outcome: str) -> Router: return Router( model_list=[ diff --git a/tests/unit/test_router_streaming_fallback_metadata.py b/tests/unit/test_router_streaming_fallback_metadata.py index 40449d8cc66..09fa9d41c1b 100644 --- a/tests/unit/test_router_streaming_fallback_metadata.py +++ b/tests/unit/test_router_streaming_fallback_metadata.py @@ -1,6 +1,5 @@ import json from typing import Final -from unittest.mock import MagicMock import pytest @@ -8,110 +7,6 @@ import litellm from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.proxy.proxy_server import _should_include_fallback_errors from litellm.router import Router -from litellm.types.llms.base import HiddenParams - - -def test_apply_fallback_hidden_params_copies_from_fallback_response(): - fallback_errors = [ - { - "message": "litellm.RateLimitError: upstream limited request", - "type": "RateLimitError", - "param": None, - "code": "429", - } - ] - chunk = litellm.ModelResponseStream( - id="test", - model="openai/internal-fallback", - choices=[], - ) - chunk._hidden_params = { - "additional_headers": {"x-existing-chunk-header": "keep"}, - "model_id": "chunk-model-id", - } - fallback_response = MagicMock() - fallback_response._hidden_params = { - "additional_headers": { - "x-litellm-attempted-fallbacks": 1, - "x-litellm-model-group": "fallback-model", - "x-litellm-fallback-errors": json.dumps(fallback_errors), - }, - "api_base": "https://fallback.example", - } - - Router._apply_fallback_hidden_params_to_item( - fallback_item=chunk, - prepared_fallback_hidden_params=Router._prepare_fallback_hidden_params( - fallback_response - ), - ) - - assert chunk._hidden_params["api_base"] == "https://fallback.example" - assert chunk._hidden_params["model_id"] == "chunk-model-id" - assert chunk._hidden_params["additional_headers"] == { - "x-existing-chunk-header": "keep", - "x-litellm-attempted-fallbacks": 1, - "x-litellm-model-group": "fallback-model", - "x-litellm-fallback-errors": json.dumps(fallback_errors), - } - - -def test_apply_fallback_hidden_params_updates_plain_duck_chunk(): - class PlainChunk: - def __init__(self) -> None: - self._hidden_params = { - "additional_headers": {"x-existing-chunk-header": "keep"}, - "model_id": "chunk-model-id", - } - - chunk: Final = PlainChunk() - fallback_response = { - "_hidden_params": { - "additional_headers": { - "x-litellm-attempted-fallbacks": 1, - }, - "api_base": "https://fallback.example", - } - } - - Router._apply_fallback_hidden_params_to_item( - fallback_item=chunk, - prepared_fallback_hidden_params=Router._prepare_fallback_hidden_params(fallback_response), - ) - - assert chunk._hidden_params["api_base"] == "https://fallback.example" - assert chunk._hidden_params["model_id"] == "chunk-model-id" - assert chunk._hidden_params["additional_headers"] == { - "x-existing-chunk-header": "keep", - "x-litellm-attempted-fallbacks": 1, - } - - -def test_apply_fallback_hidden_params_normalizes_hidden_params_on_plain_duck_chunk(): - class PlainChunk: - def __init__(self) -> None: - self._hidden_params = HiddenParams( - api_base="https://original.example", - model_id="original-model-id", - additional_headers={"x-existing-chunk-header": "keep"}, - ) - - chunk: Final = PlainChunk() - Router._apply_fallback_hidden_params_to_item( - fallback_item=chunk, - prepared_fallback_hidden_params=( - {"api_base": "https://fallback.example", "model_id": "fallback-model-id"}, - {"x-litellm-attempted-fallbacks": 1}, - ), - ) - - assert isinstance(chunk._hidden_params, dict) - assert chunk._hidden_params["api_base"] == "https://fallback.example" - assert chunk._hidden_params["model_id"] == "fallback-model-id" - assert chunk._hidden_params["additional_headers"] == { - "x-existing-chunk-header": "keep", - "x-litellm-attempted-fallbacks": 1, - } def _two_group_fallback_router() -> Router: @@ -183,43 +78,6 @@ async def test_router_omits_fallback_errors_without_opt_in(): assert "x-litellm-fallback-errors" not in headers -def test_prepare_fallback_hidden_params_no_additional_headers(): - class FakeResponse: - _hidden_params = {"api_base": "http://example.com"} - - hidden_params, headers = Router._prepare_fallback_hidden_params(FakeResponse()) - assert hidden_params == {"api_base": "http://example.com"} - assert headers == {} - - -def test_apply_fallback_hidden_params_to_item_none_item(): - Router._apply_fallback_hidden_params_to_item( - None, ({"api_base": "http://fallback.example"}, {"x-custom": "value"}) - ) - - -def test_apply_fallback_hidden_params_to_item_no_existing_additional_headers(): - chunk: Final = litellm.ModelResponseStream( - id="test", - model="openai/internal-fallback", - choices=[], - ) - chunk.hidden_params["model_id"] = "test-id" - Router._apply_fallback_hidden_params_to_item( - chunk, - ( - {"api_base": "http://fallback.example"}, - {"x-litellm-attempted-fallbacks": 1}, - ), - ) - - assert chunk.hidden_params["api_base"] == "http://fallback.example" - assert chunk.hidden_params["model_id"] == "test-id" - assert chunk.hidden_params["additional_headers"] == { - "x-litellm-attempted-fallbacks": 1 - } - - @pytest.mark.asyncio async def test_set_response_headers_adds_model_group_to_streaming_wrapper(): class StreamingWrapper: diff --git a/tests/unit/test_router_weighted_failover.py b/tests/unit/test_router_weighted_failover.py index 70231b620be..521122291c1 100644 --- a/tests/unit/test_router_weighted_failover.py +++ b/tests/unit/test_router_weighted_failover.py @@ -9,7 +9,6 @@ cross-group fallback runs. from collections import Counter from typing import Final, Optional -from unittest.mock import AsyncMock, patch import pytest @@ -75,287 +74,6 @@ class TestGetExcludedFilteredDeployments: assert len(result) == 1 -# --------------------------------------------------------------------------- -# Router helpers (router_code_coverage.py requires these names in a *router* test file) -# --------------------------------------------------------------------------- - - -def test_set_failed_deployment_id_on_exception(): - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - exc = Exception("fail") - dep = _make_dep("dep-a") - router._set_failed_deployment_id_on_exception(exc, dep) - assert getattr(exc, "failed_deployment_id", None) == "dep-a" - router._set_failed_deployment_id_on_exception(exc, _make_dep("dep-b")) - assert exc.failed_deployment_id == "dep-a" - - -def test_stamp_failed_deployment_id_with_effective_model_info_prefers_kwargs(): - """kwargs["model_info"] (the dynamic client-side-credential id, when present) must win - over the static deployment's model_info, so a bad-credential tenant's failures are - attributed to their own dynamic deployment id, not the shared static one.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - exc = Exception("fail") - router._stamp_failed_deployment_id_with_effective_model_info( - exc, _make_dep("dep-a"), {"model_info": {"id": "dynamic-dep"}} - ) - assert exc.failed_deployment_id == "dynamic-dep" - - -def test_stamp_failed_deployment_id_with_effective_model_info_falls_back_to_deployment(): - """With no dynamic id in kwargs (the common, non-client-side-credential case), the - static deployment's own model_info.id must still be stamped.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - exc = Exception("fail") - router._stamp_failed_deployment_id_with_effective_model_info(exc, _make_dep("dep-a"), {}) - assert exc.failed_deployment_id == "dep-a" - - -@pytest.mark.asyncio -async def test_ageneric_api_call_with_fallbacks_helper_stamps_failed_deployment_id(): - """_ageneric_api_call_with_fallbacks_helper must stamp failed_deployment_id on a - failure, same as _completion/_acompletion, so callers identifying the failed - deployment (cooldown, weighted failover) work for this call type too instead of - depending on which metadata bucket ("metadata" vs "litellm_metadata") it uses.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - - async def _failing_original_function(**kwargs): - raise RuntimeError("boom") - - with pytest.raises(RuntimeError) as exc_info: - await router._ageneric_api_call_with_fallbacks_helper( - model="test-model", - original_generic_function=_failing_original_function, - ) - - assert getattr(exc_info.value, "failed_deployment_id", None) == "dep-a" - - -@pytest.mark.asyncio -async def test_ageneric_api_call_with_fallbacks_helper_stamps_dynamic_id_for_clientside_credentials(): - """A client-side-credential call (tenant-supplied api_key) generates a dynamic - deployment id distinct from the shared static deployment. Stamping the static id - instead would let one tenant's bad credentials cool down the deployment every - other tenant sharing this config relies on.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - - async def _failing_original_function(**kwargs): - raise RuntimeError("boom") - - with pytest.raises(RuntimeError) as exc_info: - await router._ageneric_api_call_with_fallbacks_helper( - model="test-model", - original_generic_function=_failing_original_function, - api_key="tenant-supplied-key", - litellm_metadata={"model_group": "test-model"}, - ) - - failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) - assert failed_deployment_id is not None - assert failed_deployment_id != "dep-a" - - -@pytest.mark.asyncio -async def test_acompletion_stamps_dynamic_id_for_clientside_credentials(): - """Same bug as the generic-API-call helper above, but in the regular completion - path: _acompletion's exception handlers must stamp the dynamic client-side-credential - deployment id, not the shared static deployment's id.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - - with patch("litellm.acompletion", new_callable=AsyncMock, side_effect=RuntimeError("boom")): - with pytest.raises(RuntimeError) as exc_info: - await router._acompletion( - model="test-model", - messages=[{"role": "user", "content": "Hello"}], - api_key="tenant-supplied-key", - metadata={"model_group": "test-model"}, - ) - - failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) - assert failed_deployment_id is not None - assert failed_deployment_id != "dep-a" - - -@pytest.mark.asyncio -async def test_acompletion_stamps_dynamic_id_for_clientside_credentials_on_timeout(): - """Same bug as the RuntimeError case above, but for the separate `except litellm.Timeout` - branch in `_acompletion`: it has its own call to the stamping helper, so a fix that only - covers the generic `except Exception` branch would leave a caller-supplied timeout - (`litellm.Timeout` is what `x-litellm-timeout` maps to) stamping the shared static id.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - - timeout_exc = litellm.Timeout(message="boom", model="test-model", llm_provider="openai") - with patch("litellm.acompletion", new_callable=AsyncMock, side_effect=timeout_exc): - with pytest.raises(litellm.Timeout) as exc_info: - await router._acompletion( - model="test-model", - messages=[{"role": "user", "content": "Hello"}], - api_key="tenant-supplied-key", - metadata={"model_group": "test-model"}, - ) - - failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) - assert failed_deployment_id is not None - assert failed_deployment_id != "dep-a" - - -def test_completion_stamps_dynamic_id_for_clientside_credentials(): - """Sync counterpart: _completion's exception handler must stamp the dynamic - client-side-credential deployment id, not the shared static deployment's id.""" - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "openai/gpt-4o", "api_key": "test-key"}, - "model_info": {"id": "dep-a"}, - } - ], - ) - - with patch("litellm.completion", side_effect=RuntimeError("boom")): - with pytest.raises(RuntimeError) as exc_info: - router._completion( - model="test-model", - messages=[{"role": "user", "content": "Hello"}], - api_key="tenant-supplied-key", - metadata={"model_group": "test-model"}, - ) - - failed_deployment_id = getattr(exc_info.value, "failed_deployment_id", None) - assert failed_deployment_id is not None - assert failed_deployment_id != "dep-a" - - -@pytest.mark.asyncio -async def test_maybe_run_weighted_failover_returns_none_without_failed_id(): - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key", "weight": 1}, - "model_info": {"id": "A"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "key", "weight": 1}, - "model_info": {"id": "B"}, - }, - ], - routing_strategy="simple-shuffle", - enable_weighted_failover=True, - ) - result = await router._maybe_run_weighted_failover( - exception=Exception("fail"), - original_model_group="test-model", - all_deployments=[_make_dep("A"), _make_dep("B")], - args=(), - kwargs={"metadata": {}}, - input_kwargs={}, - ) - assert result is None - - -@pytest.mark.asyncio -async def test_maybe_run_weighted_failover_persists_excluded_ids_to_kwargs(monkeypatch): - """Regression: writing to the metadata dict returned by `setdefault` must - update the dict in `kwargs` itself so the next hop sees prior exclusions. - Previously `setdefault(..., {}) or {}` returned a disconnected dict on the - first hop, dropping `_failover_excluded_ids` writes. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "A"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "B"}, - }, - ], - routing_strategy="simple-shuffle", - enable_weighted_failover=True, - ) - - async def _stub_run_async_fallback(*args, **kwargs): - return "ok" - - monkeypatch.setattr("litellm.router_python_core.run_async_fallback", _stub_run_async_fallback) - - exc = Exception("fail") - exc.failed_deployment_id = "A" - kwargs: dict = {"metadata": {}} - await router._maybe_run_weighted_failover( - exception=exc, - original_model_group="test-model", - all_deployments=[_make_dep("A"), _make_dep("B")], - args=(), - kwargs=kwargs, - input_kwargs={}, - ) - # The dict inside kwargs must reflect the write — proves `meta` was the - # same object as kwargs["metadata"] (no disconnected copy). - assert kwargs["metadata"].get("_failover_excluded_ids") == ["A"] - - # --------------------------------------------------------------------------- # Integration tests for weighted-failover end-to-end via Router # --------------------------------------------------------------------------- @@ -758,195 +476,6 @@ async def test_user_config_two_region_failover(): assert successes["eastus2"] == 0 -# --------------------------------------------------------------------------- -# Tests for healthy-deployment-only check in _maybe_run_weighted_failover -# (Issue: weighted failover checked all deployments, not just healthy ones) -# --------------------------------------------------------------------------- - - -@pytest.mark.asyncio -async def test_maybe_run_weighted_failover_skips_when_remaining_all_in_cooldown( - monkeypatch, -): - """When every non-excluded deployment is in cooldown, _maybe_run_weighted_failover - must return None immediately without invoking run_async_fallback. - - Previously the check was against all_deployments (including cooldown ones), so - run_async_fallback would be called unnecessarily and would raise RouterRateLimitError. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "A"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "B"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "C"}, - }, - ], - routing_strategy="simple-shuffle", - enable_weighted_failover=True, - ) - - # A just failed; B and C are both in cooldown. - exc = Exception("A down") - exc.failed_deployment_id = "A" - - run_async_fallback_called = False - - async def _should_not_be_called(*args, **kwargs): - nonlocal run_async_fallback_called - run_async_fallback_called = True - return "should not reach here" - - monkeypatch.setattr("litellm.router_python_core.run_async_fallback", _should_not_be_called) - - # Patch cooldown so B and C appear in cooldown. - with patch( - "litellm.router_python_core.async_get_cooldown_deployments", - new=AsyncMock(return_value=["B", "C"]), - ): - result = await router._maybe_run_weighted_failover( - exception=exc, - original_model_group="test-model", - all_deployments=[_make_dep("A"), _make_dep("B"), _make_dep("C")], - args=(), - kwargs={"metadata": {}}, - input_kwargs={}, - ) - - assert result is None, "Should return None when all remaining deployments are in cooldown" - assert not run_async_fallback_called, "run_async_fallback must NOT be called when no healthy deployments remain" - - -@pytest.mark.asyncio -async def test_maybe_run_weighted_failover_proceeds_when_one_healthy_remains( - monkeypatch, -): - """When at least one non-excluded deployment is healthy (not in cooldown), - _maybe_run_weighted_failover should still invoke run_async_fallback normally. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "A"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "B"}, - }, - { - "model_name": "test-model", - "litellm_params": {"model": "gpt-4o", "api_key": "k", "weight": 1}, - "model_info": {"id": "C"}, - }, - ], - routing_strategy="simple-shuffle", - enable_weighted_failover=True, - ) - - # A just failed; B is in cooldown; C is healthy. - exc = Exception("A down") - exc.failed_deployment_id = "A" - - run_async_fallback_called = False - - async def _stub_run_async_fallback(*args, **kwargs): - nonlocal run_async_fallback_called - run_async_fallback_called = True - return "ok from C" - - monkeypatch.setattr("litellm.router_python_core.run_async_fallback", _stub_run_async_fallback) - - with patch( - "litellm.router_python_core.async_get_cooldown_deployments", - new=AsyncMock(return_value=["B"]), - ): - result = await router._maybe_run_weighted_failover( - exception=exc, - original_model_group="test-model", - all_deployments=[_make_dep("A"), _make_dep("B"), _make_dep("C")], - args=(), - kwargs={"metadata": {}}, - input_kwargs={}, - ) - - assert result == "ok from C" - assert run_async_fallback_called, "run_async_fallback must be called when a healthy deployment remains" - - -@pytest.mark.asyncio -async def test_failover_falls_through_to_external_fallback_when_remaining_in_cooldown(): - """End-to-end: when the only non-failed deployments are in cooldown, - weighted failover must fall through to the configured cross-group fallback. - - Without the fix the _maybe_run_weighted_failover would invoke run_async_fallback - unnecessarily (because it counted cooldown deployments as "remaining"), get back - RouterRateLimitError, return None, and reach the same fallback path — but only - incidentally. With the fix the early-exit path is taken directly. - """ - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "gpt-4o", - "api_key": "bad", - "mock_response": Exception("A down"), - "weight": 1_000_000, # always picked first - }, - "model_info": {"id": "A"}, - }, - { - "model_name": "test-model", - "litellm_params": { - "model": "gpt-4o", - "api_key": "bad", - "mock_response": Exception("B down"), - "weight": 1, - }, - "model_info": {"id": "B"}, - }, - { - "model_name": "fallback-model", - "litellm_params": { - "model": "gpt-4o", - "api_key": "good", - "mock_response": "ok from fallback", - }, - "model_info": {"id": "fallback"}, - }, - ], - routing_strategy="simple-shuffle", - num_retries=0, - enable_weighted_failover=True, - fallbacks=[{"test-model": ["fallback-model"]}], - ) - - # Put B in cooldown so weighted failover can't use it after A fails. - with patch( - "litellm.router_python_core.async_get_cooldown_deployments", - new=AsyncMock(return_value=["B"]), - ): - response = await router.acompletion( - model="test-model", - messages=[{"role": "user", "content": "hi"}], - ) - - assert response._hidden_params["model_id"] == "fallback" - - @pytest.mark.asyncio @pytest.mark.parametrize( ("pre_call_check", "metadata"), @@ -1115,14 +644,14 @@ async def test_weighted_failover_skips_the_failed_deployment_for_clientside_cred @pytest.mark.parametrize( - ("endpoint", "attempt", "arguments"), + ("endpoint", "arguments"), [ - ("embedding", "_embedding", {"input": ["hi"]}), - ("image_generation", "_image_generation", {"prompt": "a cat"}), + ("embedding", {"input": ["hi"]}), + ("image_generation", {"prompt": "a cat"}), ], ) def test_a_sync_attempt_stamps_the_failed_deployment( - monkeypatch: pytest.MonkeyPatch, endpoint: str, attempt: str, arguments: dict[str, object] + monkeypatch: pytest.MonkeyPatch, endpoint: str, arguments: dict[str, object] ) -> None: def provider(**kwargs: object) -> object: raise litellm.InternalServerError(message="down", llm_provider="openai", model="m") @@ -1139,7 +668,7 @@ def test_a_sync_attempt_stamps_the_failed_deployment( ) with pytest.raises(litellm.InternalServerError) as raised: - getattr(router, attempt)(model="test-model", **arguments) + getattr(router, endpoint)(model="test-model", **arguments) assert getattr(raised.value, "failed_deployment_id", None) == "A" @@ -1157,6 +686,6 @@ async def test_a_failed_pick_stamps_no_deployment() -> None: ) with pytest.raises(litellm.BadRequestError, match="no healthy deployments") as raised: - await router._aembedding(input=["hi"], model="no-such-model") + await router.aembedding(input=["hi"], model="no-such-model") assert getattr(raised.value, "failed_deployment_id", None) is None diff --git a/tests/unit/test_streaming_connection_cleanup.py b/tests/unit/test_streaming_connection_cleanup.py index 39fcee8d44d..c525e917b4c 100644 --- a/tests/unit/test_streaming_connection_cleanup.py +++ b/tests/unit/test_streaming_connection_cleanup.py @@ -3,13 +3,15 @@ Regression tests for streaming connection pool leak fix. """ import asyncio -from unittest.mock import MagicMock, patch +import json +from unittest.mock import MagicMock import anyio import httpx import pytest +import respx - +import litellm from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.custom_httpx.aiohttp_transport import ( AiohttpResponseStream, @@ -55,9 +57,7 @@ async def test_aiohttp_transport_response_uses_stream_not_content(): return Resp() transport = LiteLLMAiohttpTransport(client=lambda: FakeSession()) # type: ignore - response = await transport.handle_async_request( - httpx.Request("GET", "http://example.com") - ) + response = await transport.handle_async_request(httpx.Request("GET", "http://example.com")) assert isinstance(response.stream, AiohttpResponseStream) @@ -167,224 +167,121 @@ async def test_aclose_completes_under_cancellation(): # ── Router stream_with_fallbacks cleanup tests ────────────────── -@pytest.mark.asyncio -async def test_stream_with_fallbacks_closes_stream_on_generator_close(): - """Closing the FallbackStreamWrapper must aclose() the underlying model_response - via stream_with_fallbacks' finally block.""" +def _sse_chunk(content: str, finish_reason: str | None = None) -> bytes: + payload = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4", + "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(payload)}\n\n".encode() + + +_DONE = b"data: [DONE]\n\n" + + +class _TrackingStream(httpx.AsyncByteStream): + """An upstream SSE body that records whether the router closed it.""" + + def __init__(self, parts: list[bytes], fail: bool = False): + self._parts = parts + self._fail = fail + self.closed = False + + async def __aiter__(self): + for part in self._parts: + yield part + if self._fail: + raise httpx.ReadError("upstream dropped the connection") + + async def aclose(self): + self.closed = True + + +def _full_stream() -> _TrackingStream: + return _TrackingStream([_sse_chunk("c1"), _sse_chunk("c2"), _sse_chunk("c3", "stop"), _DONE]) + + +@pytest.fixture +def stream_router(monkeypatch: pytest.MonkeyPatch): from litellm.router import Router - stream_closed = False - - class FakeStream(CustomStreamWrapper): - def __init__(self): - super().__init__( - completion_stream=None, - model="test-model", - logging_obj=MagicMock(), - custom_llm_provider="openai", - ) - self._items = ["chunk1", "chunk2", "chunk3"] - self._index = 0 - - def __aiter__(self): - return self - - async def __anext__(self): - if self._index >= len(self._items): - raise StopAsyncIteration - item = self._items[self._index] - self._index += 1 - return item - - async def aclose(self): - nonlocal stream_closed - stream_closed = True - - router = Router( + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield Router( model_list=[ { "model_name": "test-model", "litellm_params": { "model": "openai/test", "api_key": "fake", + "api_base": "https://primary.example/v1", }, - } - ] - ) - - fake_stream = FakeStream() - - # Call _acompletion_streaming_iterator directly so we go through - # stream_with_fallbacks and its finally block - result = await router._acompletion_streaming_iterator( - model_response=fake_stream, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "test-model"}, - ) - - # Consume one chunk then close (simulates client disconnect) - async for _ in result: - break - await result.aclose() - - assert ( - stream_closed - ), "model_response stream was not closed by stream_with_fallbacks finally block" - - -@pytest.mark.asyncio -async def test_stream_with_fallbacks_closes_stream_on_normal_completion(): - """stream_with_fallbacks must aclose() model_response even on normal completion.""" - from litellm.router import Router - - stream_closed = False - - class FakeStream(CustomStreamWrapper): - def __init__(self): - super().__init__( - completion_stream=None, - model="test-model", - logging_obj=MagicMock(), - custom_llm_provider="openai", - ) - self._items = ["chunk1"] - self._index = 0 - - def __aiter__(self): - return self - - async def __anext__(self): - if self._index >= len(self._items): - raise StopAsyncIteration - item = self._items[self._index] - self._index += 1 - return item - - async def aclose(self): - nonlocal stream_closed - stream_closed = True - - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/test", - "api_key": "fake", - }, - } - ] - ) - - fake_stream = FakeStream() - - result = await router._acompletion_streaming_iterator( - model_response=fake_stream, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={"model": "test-model"}, - ) - - # Exhaust the stream fully - async for _ in result: - pass - await result.aclose() - - assert stream_closed, "model_response stream was not closed after normal completion" - - -@pytest.mark.asyncio -async def test_stream_with_fallbacks_closes_both_on_fallback_disconnect(): - """When a fallback is triggered and the client disconnects during fallback - iteration, both model_response and fallback_response must be closed.""" - from litellm.exceptions import MidStreamFallbackError - from litellm.router import Router - - model_closed = False - fallback_closed = False - - class FakeModelStream(CustomStreamWrapper): - """Stream that raises MidStreamFallbackError immediately to trigger fallback.""" - - def __init__(self): - super().__init__( - completion_stream=None, - model="test-model", - logging_obj=MagicMock(), - custom_llm_provider="openai", - ) - self.chunks = [] - - def __aiter__(self): - return self - - async def __anext__(self): - raise MidStreamFallbackError( - message="test mid-stream error", - model="test-model", - llm_provider="openai", - generated_content="", - ) - - async def aclose(self): - nonlocal model_closed - model_closed = True - - class FakeFallbackStream: - """Fallback stream that yields chunks.""" - - def __init__(self): - self._items = ["fb1", "fb2", "fb3"] - self._index = 0 - - def __aiter__(self): - return self - - async def __anext__(self): - if self._index >= len(self._items): - raise StopAsyncIteration - item = self._items[self._index] - self._index += 1 - return item - - async def aclose(self): - nonlocal fallback_closed - fallback_closed = True - - router = Router( - model_list=[ - { - "model_name": "test-model", - "litellm_params": { - "model": "openai/test", - "api_key": "fake", - }, - } - ] - ) - - fake_model_stream = FakeModelStream() - fake_fallback_stream = FakeFallbackStream() - - # Mock async_function_with_fallbacks_common_utils to return the fallback stream - # instead of actually calling through the full fallback machinery - with patch.object( - router, - "async_function_with_fallbacks_common_utils", - return_value=fake_fallback_stream, - ): - result = await router._acompletion_streaming_iterator( - model_response=fake_model_stream, - messages=[{"role": "user", "content": "hi"}], - initial_kwargs={ - "model": "test-model", - "fallbacks": ["other-model"], }, - ) + { + "model_name": "other-model", + "litellm_params": { + "model": "openai/test", + "api_key": "fake", + "api_base": "https://fallback.example/v1", + }, + }, + ], + fallbacks=[{"test-model": ["other-model"]}], + ) + litellm.in_memory_llm_clients_cache.flush_cache() - # Consume one fallback chunk then close (simulates client disconnect) + +async def _open_stream(router, primary: _TrackingStream, fallback: _TrackingStream, mock_router): + for host, body in (("primary", primary), ("fallback", fallback)): + mock_router.post(f"https://{host}.example/v1/chat/completions").mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, stream=body) + ) + return await router.acompletion(model="test-model", messages=[{"role": "user", "content": "hi"}], stream=True) + + +@pytest.mark.asyncio +async def test_stream_with_fallbacks_closes_stream_on_generator_close(stream_router): + """Closing the router's stream after one chunk (a client disconnect) must + close the upstream response.""" + primary = _full_stream() + + with respx.mock(assert_all_called=False) as mock_router: + result = await _open_stream(stream_router, primary, _full_stream(), mock_router) async for _ in result: break await result.aclose() - assert model_closed, "model_response stream was not closed" - assert fallback_closed, "fallback_response stream was not closed" + assert primary.closed, "upstream stream was not closed when the client disconnected" + + +@pytest.mark.asyncio +async def test_stream_with_fallbacks_closes_stream_on_normal_completion(stream_router): + """The upstream response must be closed even on normal completion.""" + primary = _full_stream() + + with respx.mock(assert_all_called=False) as mock_router: + result = await _open_stream(stream_router, primary, _full_stream(), mock_router) + async for _ in result: + pass + await result.aclose() + + assert primary.closed, "upstream stream was not closed after normal completion" + + +@pytest.mark.asyncio +async def test_stream_with_fallbacks_closes_both_on_fallback_disconnect(stream_router): + """When a fallback is triggered and the client disconnects during fallback + iteration, both the primary and the fallback upstream must be closed.""" + primary = _TrackingStream([], fail=True) + fallback = _full_stream() + + with respx.mock() as mock_router: + result = await _open_stream(stream_router, primary, fallback, mock_router) + async for _ in result: + break + await result.aclose() + + assert primary.closed, "primary upstream stream was not closed" + assert fallback.closed, "fallback upstream stream was not closed"