diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 1431027a35a..b9e510f6713 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -326,6 +326,7 @@ async fn readiness_answers_without_credentials_like_the_proxy_with_no_database( ); } +#[rstest] #[tokio::test] async fn forwards_speed_to_a_model_the_catalog_marks_fast() { let catalog = Catalog::parse(BUNDLED_REGISTRY).unwrap().value; diff --git a/litellm-rust/crates/inference-messages/src/model_capabilities.rs b/litellm-rust/crates/inference-messages/src/model_capabilities.rs index 0fe47335375..c4e87b2c6d6 100644 --- a/litellm-rust/crates/inference-messages/src/model_capabilities.rs +++ b/litellm-rust/crates/inference-messages/src/model_capabilities.rs @@ -42,7 +42,7 @@ pub fn anthropic_model_capabilities( }); let probe = Probe { catalog, - model, + fallbacks: backstop_candidates(model), info: catalog.model_info( &candidate_names(name, provider) .each_ref() @@ -57,13 +57,17 @@ pub fn anthropic_model_capabilities( thinking_always_on: probe.supports("thinking_always_on"), supports_legacy_thinking: probe.supports("supports_legacy_thinking"), supports_output_config: probe.supports("supports_output_config"), - supports_sampling_params: backstop(catalog, name, "supports_sampling_params") - .unwrap_or_else(|| { - let lowered = name.to_lowercase(); - !SAMPLING_RESTRICTED_NAMES - .iter() - .any(|restricted| lowered.contains(restricted)) - }), + supports_sampling_params: backstop( + catalog, + &backstop_candidates(name), + "supports_sampling_params", + ) + .unwrap_or_else(|| { + let lowered = name.to_lowercase(); + !SAMPLING_RESTRICTED_NAMES + .iter() + .any(|restricted| lowered.contains(restricted)) + }), supports_speed: provider == "anthropic" && flag(exact(catalog, name), "supports_speed") == Some(true), supports_mid_conversation_system: probe.supports("supports_mid_conversation_system"), @@ -80,7 +84,7 @@ pub fn anthropic_model_capabilities( struct Probe<'a> { catalog: &'a Catalog, - model: &'a str, + fallbacks: Vec, info: Option, bare: Option<&'a Fields>, } @@ -93,7 +97,7 @@ impl Probe<'_> { if self.info.is_some() && flag(self.bare, capability) == Some(true) { return true; } - backstop(self.catalog, self.model, capability) == Some(true) + backstop(self.catalog, &self.fallbacks, capability) == Some(true) } } @@ -112,32 +116,52 @@ fn is_bedrock(provider: &str) -> bool { matches!(provider, "bedrock" | "bedrock_converse") } -fn candidate_names(model: &str, provider: &str) -> [String; 5] { - let stripped = stripped_model(model, provider); +fn candidate_names(model: &str, provider: &str) -> [String; 6] { + let split = split_model(model, provider); + let combined = format!("{provider}/{model}"); + let (region_free_combined, combined_stripped) = if provider == "bedrock_mantle" { + ( + format!("{provider}/{split}"), + format!("{provider}/{}", stripped_model(split, provider)), + ) + } else { + ( + combined.clone(), + format!("{provider}/{}", stripped_model(model, provider)), + ) + }; [ - format!("{provider}/{model}"), + combined, model.to_owned(), - split_model(model, provider).to_owned(), - format!("{provider}/{stripped}"), - stripped, + region_free_combined, + split.to_owned(), + combined_stripped, + stripped_model(model, provider), ] } fn split_model<'a>(model: &'a str, provider: &str) -> &'a str { + if provider == "bedrock_mantle" { + return without_region(model); + } if !is_bedrock(provider) { return model; } let routed = without_routing_prefixes(model); - let region_free = match routed.split_once('/') { - Some((region, rest)) if is_bedrock_region(region) => rest, - _ => routed, - }; + let region_free = without_region(routed); if !region_free.starts_with("arn:") { return routed; } region_free.rsplit('/').next().unwrap_or(region_free) } +fn without_region(model: &str) -> &str { + match model.split_once('/') { + Some((region, rest)) if is_bedrock_region(region) => rest, + _ => model, + } +} + fn stripped_model(model: &str, provider: &str) -> String { match provider { provider if is_bedrock(provider) => bedrock_base_model(model).to_owned(), @@ -146,8 +170,8 @@ fn stripped_model(model: &str, provider: &str) -> String { } } -fn backstop(catalog: &Catalog, model: &str, capability: &str) -> Option { - backstop_candidates(model) +fn backstop(catalog: &Catalog, candidates: &[String], capability: &str) -> Option { + candidates .iter() .find_map(|name| flag(exact(catalog, name), capability)) } @@ -203,13 +227,27 @@ mod tests { "supports_xhigh_reasoning_effort": true, "aliases": ["claude-test-5-latest"] }, - "claude-slow-5": {"litellm_provider": "anthropic", "mode": "chat", "supports_reasoning": false}, + "claude-slow-5": { + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": false, + "supports_adaptive_thinking": false + }, + "claude-embed-5": {"litellm_provider": "anthropic", "mode": "embedding"}, + "anthropic/claude-split-5": {"litellm_provider": "anthropic", "mode": "chat"}, + "claude-split-5": {"litellm_provider": "anthropic", "mode": "chat", "supports_reasoning": true}, "anthropic.claude-test-5-v1:0": { "litellm_provider": "bedrock_converse", "mode": "chat", "supports_reasoning": true, "supports_speed": true }, + "bedrock_mantle/anthropic.claude-mantle-5": { + "litellm_provider": "bedrock_mantle", + "mode": "chat", + "supports_reasoning": true, + "supports_adaptive_thinking": true + }, "vertex_ai/claude-test-5": { "litellm_provider": "vertex_ai-anthropic_models", "mode": "chat", @@ -218,7 +256,8 @@ mod tests { }, "fallback_generalizations": { "rules": [ - {"name": "routing", "pattern": "claude", "model_info": {"litellm_provider": "anthropic", "supports_speed": true}}, + {"name": "routing", "pattern": "claude", "model_info": {"litellm_provider": "anthropic", "supports_reasoning": true}}, + {"name": "derived", "pattern": "claude", "extends": "fill", "model_info": {"supports_reasoning": true}}, {"name": "new-family", "pattern": "claude-(?=newfamily)", "model_info": {"supports_reasoning": true}}, { "name": "fill", @@ -241,10 +280,15 @@ mod tests { #[case::bare_name("claude-test-5", None, (true, true, true, true))] #[case::configured_provider("claude-test-5", Some("anthropic"), (true, true, true, true))] #[case::other_casing("anthropic/Claude-Test-5", None, (true, true, false, true))] - #[case::filled_row("anthropic/claude-slow-5", None, (false, true, false, true))] + #[case::row_over_fill("anthropic/claude-slow-5", None, (false, false, false, true))] + #[case::fill_skips_other_modes("anthropic/claude-embed-5", None, (false, false, false, true))] + #[case::bare_row_under_prefixed_row("anthropic/claude-split-5", None, (true, true, false, true))] + #[case::other_provider_row("azure_ai/claude-test-5", None, (false, false, false, true))] #[case::unknown_family("anthropic/claude-newfamily-6", None, (true, true, false, true))] + #[case::unknown_model_skips_routing_and_extends_rules("anthropic/claude-other-7", None, (false, true, false, true))] #[case::cross_region_bedrock("bedrock/us.anthropic.claude-test-5-v1:0", None, (true, false, false, true))] #[case::region_path_bedrock("bedrock/us-east-1/anthropic.claude-test-5-v1:0", None, (true, false, false, true))] + #[case::region_path_mantle("bedrock_mantle/us-east-1/anthropic.claude-mantle-5", None, (true, true, false, true))] #[case::converse_bedrock("bedrock/converse/anthropic.claude-test-5-v1:0", None, (true, false, false, true))] #[case::vertex_version("vertex_ai/claude-test-5@20261001", None, (true, true, false, true))] #[case::sampling_restricted_name("anthropic/claude-fable-test", None, (false, true, false, false))] diff --git a/litellm-rust/crates/router/tests/router.rs b/litellm-rust/crates/router/tests/router.rs index 6e79e0dd9ef..1ab0727e25f 100644 --- a/litellm-rust/crates/router/tests/router.rs +++ b/litellm-rust/crates/router/tests/router.rs @@ -4,8 +4,9 @@ use litellm_config::Config; use litellm_inference_messages::{MessagesSettings, MessagesShaping}; use litellm_model_catalog::Catalog; use litellm_router::{Deployment, Router}; -use rstest::rstest; +use rstest::{fixture, rstest}; +#[fixture] fn catalog() -> Catalog { Catalog::parse(br#"{"provider-model": {"litellm_provider": "provider"}}"#) .unwrap() @@ -18,12 +19,12 @@ fn catalog() -> Catalog { "api_key: test-key\n api_base: https://provider.example/v1\n custom_llm_provider: test-provider" )] #[case::secret_reference("api_key: os.environ/ROUTER_TEST_API_KEY")] -fn configuration_preserves_deployment_parameters(#[case] parameters: &str) { +fn configuration_preserves_deployment_parameters(catalog: Catalog, #[case] parameters: &str) { let config = Config::from_yaml(&format!( "model_list:\n - model_name: public-model\n litellm_params:\n model: provider/model\n {parameters}" )) .unwrap(); - let router = Router::from_config(&config, &|_: &str| None, &catalog()); + let router = Router::from_config(&config, &|_: &str| None, &catalog); let deployment = router.get(&config.model_list[0].model_name).unwrap(); let params = &config.model_list[0].litellm_params; @@ -40,7 +41,7 @@ fn configuration_preserves_deployment_parameters(#[case] parameters: &str) { } #[rstest] -fn configuration_types_the_provider_params_of_a_deployment() { +fn configuration_types_the_provider_params_of_a_deployment(catalog: Catalog) { let config = Config::from_yaml( "model_list: - model_name: public-model @@ -51,7 +52,7 @@ fn configuration_types_the_provider_params_of_a_deployment() { azure_ad_token: ignored-by-every-group", ) .unwrap(); - let router = Router::from_config(&config, &|_: &str| None, &catalog()); + let router = Router::from_config(&config, &|_: &str| None, &catalog); let params = &router.get("public-model").unwrap().litellm_params; assert_eq!( @@ -70,7 +71,7 @@ fn configuration_types_the_provider_params_of_a_deployment() { #[case::unknown("missing", None)] #[case::provider_name_is_not_an_alias("provider/a", None)] #[case::case_sensitive("PUBLIC-A", None)] -fn lookup_uses_public_names(#[case] name: &str, #[case] expected: Option<&str>) { +fn lookup_uses_public_names(catalog: Catalog, #[case] name: &str, #[case] expected: Option<&str>) { let config = Config::from_yaml( "model_list: - model_name: public-a @@ -81,17 +82,17 @@ fn lookup_uses_public_names(#[case] name: &str, #[case] expected: Option<&str>) model: provider/b", ) .unwrap(); - let router = Router::from_config(&config, &|_: &str| None, &catalog()); + let router = Router::from_config(&config, &|_: &str| None, &catalog); assert_eq!(router.get(name).map(|entry| entry.model.as_str()), expected); } #[rstest] -fn empty_configuration_has_no_deployment() { +fn empty_configuration_has_no_deployment(catalog: Catalog) { let config = Config::from_yaml("model_list: []").unwrap(); assert!( - Router::from_config(&config, &|_: &str| None, &catalog()) + Router::from_config(&config, &|_: &str| None, &catalog) .get("") .is_none() );