diff --git a/litellm-rust/crates/router-core/src/call.rs b/litellm-rust/crates/router-core/src/call.rs index 4dbe8ff3d9f..2dd95df9705 100644 --- a/litellm-rust/crates/router-core/src/call.rs +++ b/litellm-rust/crates/router-core/src/call.rs @@ -289,16 +289,54 @@ fn serving_the_refused_model<'a>( refused: &'a str, healthy: &'a [Arc], ) -> impl Iterator + 'a { - let model = snapshot - .registry - .deployment(refused) - .map(|deployment| deployment_model(&deployment)); + let model = healthy + .iter() + .find(|d| d.id == refused) + .map(|d| deployment_model(d)) + .or_else(|| { + snapshot + .registry + .deployment(refused) + .map(|deployment| deployment_model(&deployment)) + }); healthy .iter() .filter(move |d| d.id == refused || model.as_ref() == Some(&deployment_model(d))) .map(|d| d.id.clone()) } +fn context_window_refusal(failed: &Failed) -> Option { + match &*failed.error { + Error::Attempt(failure) if failure.is_a("ContextWindowExceededError") => Some(Failed { + error: Box::new(Error::Attempt(failure.clone())), + deployment_id: failed.deployment_id.clone(), + routed: failed.routed.clone(), + served_as: failed.served_as.clone(), + }), + _ => None, + } +} + +/// `fallback_lookup_groups`: the group a pre-routing hook selected, the group being served, then +/// the group the request was for. +fn fallback_lookup<'a>( + routed: Option<&'a str>, + served_as: Option<&'a str>, + target: &'a Target, + request: &'a RouterRequest, +) -> Vec<&'a str> { + 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); + } + } + lookup +} + impl Core { pub(crate) async fn call( &self, @@ -775,17 +813,7 @@ 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 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); - } - } + let lookup = fallback_lookup(routed, served_as, target, request); // `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| { @@ -987,6 +1015,11 @@ impl Core { .flatten(); let policy_applies = policy_retries.is_some(); num_retries = policy_retries.unwrap_or(num_retries); + let refusal = if policy_applies { + None + } else { + context_window_refusal(&failed) + }; drop(first_failure); // `_async_get_healthy_deployments` finds no group for a deployment id, so a call by id // counts no deployments when deciding whether to retry. @@ -1021,32 +1054,41 @@ impl Core { // keeps going, else `should_retry_this_error` decides (it raises when nothing is // healthy). let Some(failure) = retried_failure(&failed.error) else { - return Err(failed); + return Err(refusal.unwrap_or(failed)); }; let refused_by = failure .failed_deployment_id .as_ref() .or(failed.deployment_id.as_ref()); - if let Some(refused) = refused_by - && !policy_applies - && failure.is_a("ContextWindowExceededError") - { - let serving: Vec = - serving_the_refused_model(snapshot, refused, &healthy_deployments) - .filter(|id| !inputs.skipped_ids.contains(id)) - .collect(); - inputs.skipped_ids.extend(serving); - } - let untried = refused_by.map_or(0, |refused| { - healthy_deployments - .iter() - .filter(|d| d.id != *refused && !inputs.skipped_ids.contains(&d.id)) - .count() - }); + let untried = match refused_by { + Some(refused) if !policy_applies && failure.is_a("ContextWindowExceededError") => { + let serving: Vec = + serving_the_refused_model(snapshot, refused, &healthy_deployments) + .filter(|id| !inputs.skipped_ids.contains(id)) + .collect(); + inputs.skipped_ids.extend(serving); + self.untried(snapshot, op, request, &target.group, &inputs) + .await + } + _ => 0, + }; + let lookup = fallback_lookup( + failed.routed.as_deref(), + failed.served_as.as_deref(), + target, + request, + ); if !policy_applies - && !self.should_retry(snapshot, request, &failure, (healthy, untried), all) + && !self.should_retry( + snapshot, + request, + &failure, + &lookup, + (healthy, untried), + all, + ) { - return Err(failed); + return Err(refusal.unwrap_or(failed)); } let uncounted = request.uncounted_retries.filter(|_| !policy_applies); if retry == 0 && uncounted.is_none() { @@ -1061,15 +1103,15 @@ impl Core { model: target.group.clone(), cause: RetryCause::of(&failed.error), }); - return Err(Failed { + return Err(refusal.unwrap_or(Failed { error: Box::new(Error::Attempt(AttemptError::host_raised(&["Exception"]))), deployment_id: None, routed: None, served_as: None, - }); + })); } if num_retries == 0 { - return Err(failed); + return Err(refusal.unwrap_or(failed)); } if retry == 0 { self.log_retry(target, at, 0, &failed.error); @@ -1087,7 +1129,7 @@ impl Core { .iter() .all(|d| inputs.skipped_ids.contains(&d.id)) { - return Err(failed); + return Err(refusal.unwrap_or(failed)); } let delay = self.time_to_sleep( snapshot, @@ -1130,7 +1172,7 @@ impl Core { return Ok(response); } Err(next) if next.error.attempt().is_some_and(|f| f.agentic_followup) => { - return Err(next); + return Err(refusal.unwrap_or(next)); } Err(next) => { self.pick_failed(at, target, &next.error); @@ -1138,7 +1180,25 @@ impl Core { } } } - Err(failed) + Err(refusal.unwrap_or(failed)) + } + + async fn untried( + &self, + snapshot: &RouterStateSnapshot, + op: &Op, + request: &RouterRequest, + group: &str, + inputs: &SelectInputs, + ) -> usize { + self.candidates(snapshot, op, request, group, inputs) + .await + .map_or(0, |c| { + c.healthy + .iter() + .filter(|d| !inputs.skipped_ids.contains(&d.id)) + .count() + }) } /// `error` is a pick's: reported as it happens, whatever the level then does with it. @@ -1203,6 +1263,7 @@ impl Core { snapshot: &RouterStateSnapshot, request: &RouterRequest, failure: &AttemptError, + lookup: &[&str], (healthy, untried): (usize, usize), all: usize, ) -> bool { @@ -1213,7 +1274,9 @@ impl Core { healthy, untried, all, - has_context_window_fallbacks: lists.context_window.is_some(), + has_context_window_fallbacks: lists + .context_window + .is_some_and(|list| list.lookup_groups(lookup, lists.keys).is_some()), has_content_policy_fallbacks: lists.content_policy.is_some(), has_regular_fallbacks: lists.has_regular(), }, diff --git a/litellm-rust/crates/router-core/src/call/stream.rs b/litellm-rust/crates/router-core/src/call/stream.rs index e1dabc4ffc3..d89d152b435 100644 --- a/litellm-rust/crates/router-core/src/call/stream.rs +++ b/litellm-rust/crates/router-core/src/call/stream.rs @@ -10,7 +10,7 @@ use std::sync::Arc; use serde::Deserialize; use super::{ - Failed, FallbackChain, FallbackHop, Lists, Target, Tried, group_members, + Failed, FallbackChain, FallbackHop, Lists, Target, Tried, fallback_lookup, group_members, serving_the_refused_model, }; use crate::{ @@ -247,34 +247,39 @@ impl Router { .or(origin.deployment_id.as_deref()); let (budget, policy_applies) = stream_retry_budget(&snapshot, &request, &origin, failure.kind); - let skipped_ids: Vec = match refused_by { - Some(refused) if !policy_applies && failure.is_a("ContextWindowExceededError") => { - serving_the_refused_model(&snapshot, refused, &healthy_deployments).collect() - } - _ => Vec::new(), - }; - let untried = refused_by.map_or(0, |refused| { - healthy_deployments - .iter() - .filter(|d| d.id != refused && !skipped_ids.contains(&d.id)) - .count() - }); - let inputs = SelectInputs { + let refusal = (!policy_applies && failure.is_a("ContextWindowExceededError")) + .then(|| failure.clone()); + let mut inputs = SelectInputs { target_order: target.order, requested: Some(target.group.clone()), fallback_hop: at.level > 0, - skipped_ids, ..SelectInputs::default() }; + let untried = match refused_by { + Some(refused) if refusal.is_some() => { + inputs.skipped_ids = + serving_the_refused_model(&snapshot, refused, &healthy_deployments).collect(); + self.core + .untried(&snapshot, &op, &request, &target.group, &inputs) + .await + } + _ => 0, + }; + let lookup = fallback_lookup(origin.routed.as_deref(), None, &target, &request); let mut last = failure; for attempt in origin.attempted_retries..budget { if last.agentic_followup { break; } if !policy_applies - && !self - .core - .should_retry(&snapshot, &request, &last, (healthy, untried), all) + && !self.core.should_retry( + &snapshot, + &request, + &last, + &lookup, + (healthy, untried), + all, + ) { break; } @@ -320,11 +325,11 @@ impl Router { break; } } - other => return Err(other), + other => return Err(refusal.map_or(other, Error::Attempt)), }, } } - Err(Error::Attempt(last)) + Err(Error::Attempt(refusal.unwrap_or(last))) } } diff --git a/litellm-rust/crates/router-core/src/retry.rs b/litellm-rust/crates/router-core/src/retry.rs index 0af8528e6f1..a73a8e15652 100644 --- a/litellm-rust/crates/router-core/src/retry.rs +++ b/litellm-rust/crates/router-core/src/retry.rs @@ -65,7 +65,8 @@ pub struct RetryContext { pub healthy: usize, pub untried: usize, pub all: usize, - /// `context_window_fallbacks is not None`. + /// A `context_window_fallbacks` chain matches the failed group + /// (`context_window_fallback_model_group`). pub has_context_window_fallbacks: bool, /// `content_policy_fallbacks is not None`. pub has_content_policy_fallbacks: bool, diff --git a/litellm-rust/crates/router-core/tests/call_loop.rs b/litellm-rust/crates/router-core/tests/call_loop.rs index f0911cf3eb8..1870bb8f956 100644 --- a/litellm-rust/crates/router-core/tests/call_loop.rs +++ b/litellm-rust/crates/router-core/tests/call_loop.rs @@ -973,6 +973,139 @@ async fn a_context_window_error_no_deployment_raised_is_not_retried(config: Rout assert_eq!((host.retries(), host.selected().len()), (0, 0)); } +fn refusing_once() -> Arc { + Arc::new(ScriptedHost { + first_failure: Mutex::new(Some(context_window())), + ..ScriptedHost::default() + }) +} + +#[rstest] +#[tokio::test] +async fn a_context_window_error_takes_the_next_order_tier_instead_of_the_refusing_one() { + let host = refusing_once(); + let router = router(tiers(true), json!({"num_retries": 2}), host.clone()); + let response = router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!( + ( + response.deployment_id.as_str(), + host.selected(), + host.retries() + ), + ("p2", vec!["p1".to_owned(), "p2".to_owned()], 0) + ); +} + +#[rstest] +#[tokio::test] +async fn a_context_window_error_on_a_pinned_deployment_is_not_retried(config: RouterConfig) { + let host = refusing_once(); + let router = router(config, json!({"num_retries": 2}), host.clone()); + let pinned = RouterRequest { + pinned_deployment_id: Some("a1".to_owned()), + ..request() + }; + let refused = router.call(Op::from("acompletion"), pinned).await; + assert!( + matches!(refused, Err(Error::Attempt(failed)) if failed.kind == FailureKind::ContextWindowExceeded) + ); + assert_eq!(host.selected(), ["a1"]); +} + +#[rstest] +#[tokio::test] +async fn a_context_window_error_skips_the_wildcard_deployments_serving_that_model() { + let config: RouterConfig = serde_json::from_value(json!({ + "model_list": [ + {"model_name": "openai/*", "litellm_params": {"model": "openai/*"}, "model_info": {"id": "w1"}}, + {"model_name": "openai/*", "litellm_params": {"model": "openai/*"}, "model_info": {"id": "w2"}}, + ] + })) + .unwrap(); + let host = refusing_once(); + let router = router(config, json!({"num_retries": 2}), host.clone()); + let wildcard = RouterRequest { + model: "openai/gpt-4o".to_owned(), + ..RouterRequest::default() + }; + let refused = router.call(Op::from("acompletion"), wildcard).await; + assert!( + matches!(refused, Err(Error::Attempt(failed)) if failed.kind == FailureKind::ContextWindowExceeded) + ); + assert_eq!(host.selected().len(), 1); +} + +#[rstest] +#[tokio::test] +async fn a_context_window_list_for_another_group_leaves_the_retry_in_place(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing("primary", vec![context_window()])); + let router = router( + config, + json!({"num_retries": 2, "context_window_fallbacks": [{"backup": ["primary"]}]}), + host.clone(), + ); + let response = router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!( + (response.model_group.as_str(), response.attempted_retries), + ("primary", 1) + ); +} + +fn provider_bad_request() -> AttemptError { + let mut failed = AttemptError::new(FailureKind::BadRequest, Some(400)); + failed.names_provider = true; + failed +} + +#[rstest] +#[tokio::test] +async fn a_retry_that_fails_after_a_context_window_error_falls_back_on_that_error( + config: RouterConfig, +) { + let config = RouterConfig { + fallback_on_unsupported_params: Some(false), + ..config + }; + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![context_window(), provider_bad_request()], + )); + let router = router( + config, + json!({"num_retries": 2, "fallbacks": [{"primary": ["backup"]}]}), + host.clone(), + ); + let response = router + .call(Op::from("acompletion"), request()) + .await + .unwrap(); + assert_eq!( + (response.deployment_id.as_str(), host.selected().len()), + ("b1", 3) + ); +} + +#[rstest] +#[tokio::test] +async fn a_retry_that_fails_after_a_context_window_error_raises_that_error(config: RouterConfig) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![context_window(), server_error(), server_error()], + )); + let router = router(config, json!({"num_retries": 2}), host.clone()); + let refused = router.call(Op::from("acompletion"), request()).await; + assert!( + matches!(refused, Err(Error::Attempt(failed)) if failed.kind == FailureKind::ContextWindowExceeded) + ); + assert_eq!(host.selected().len(), 3); +} + #[rstest] #[tokio::test] async fn a_failed_agentic_followup_is_not_retried_but_still_falls_back(config: RouterConfig) { @@ -1113,6 +1246,57 @@ async fn a_stream_context_window_error_retries_on_another_deployment(config: Rou ); } +#[rstest] +#[tokio::test] +async fn a_stream_context_window_error_in_an_order_tier_is_not_retried() { + let host = Arc::new(ScriptedHost::default()); + let router = router(tiers(true), json!({}), host.clone()); + let origin = StreamOrigin { + deployment_id: Some("p1".to_owned()), + ..stream_origin() + }; + let refused = router + .retry_stream( + Op::from("acompletion"), + stream_request(), + origin, + context_window(), + ) + .await; + assert!( + matches!(refused, Err(Error::Attempt(failed)) if failed.kind == FailureKind::ContextWindowExceeded) + ); + assert!(host.selected().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn a_stream_retry_that_fails_after_a_context_window_error_raises_that_error( + config: RouterConfig, +) { + let host = Arc::new(ScriptedHost::failing( + "primary", + vec![server_error(), server_error(), server_error()], + )); + let router = router(config, json!({}), host.clone()); + let origin = StreamOrigin { + deployment_id: Some("a1".to_owned()), + ..stream_origin() + }; + let refused = router + .retry_stream( + Op::from("acompletion"), + stream_request(), + origin, + context_window(), + ) + .await; + assert!( + matches!(refused, Err(Error::Attempt(failed)) if failed.kind == FailureKind::ContextWindowExceeded) + ); + assert_eq!(host.selected().len(), 3); +} + fn fallback_list(value: serde_json::Value) -> Fallbacks { serde_json::from_value(value).unwrap() } diff --git a/tests/unit/test_router/test_router_retries.py b/tests/unit/test_router/test_router_retries.py index 857b5052d6c..5d99ac822c7 100644 --- a/tests/unit/test_router/test_router_retries.py +++ b/tests/unit/test_router/test_router_retries.py @@ -439,7 +439,7 @@ def test_raise_context_window_exceeded_error(): def test_raise_context_window_exceeded_error_no_retry(): """ - Do not Retry Context Window Exceeded Error, when context_window_fallbacks is None + Do not retry a Context Window Exceeded Error when no other healthy deployment is left to try """ context_window_error = litellm.ContextWindowExceededError( message="Context window exceeded",