From 8140fdf54f31a417d8cb186f2074df62a3cf490c Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Sat, 14 Nov 2026 08:51:52 -0800 Subject: [PATCH] fix(gateway)!: apply a deployment's Messages params and fill model capabilities from the cost map (#51834) * fix(gateway): apply a deployment's sampling params on /v1/messages The gateway ignored temperature, top_p, top_k, max_tokens, and stop_sequences set in a deployment's litellm_params, while the Python proxy sends them upstream. Merge them under the client body, so a value the client sets, null included, still wins. * fix(gateway): apply a deployment's non-sampling litellm_params on /v1/messages The gateway dropped thinking, tools, tool_choice, system, and the other Anthropic Messages params a deployment sets in litellm_params, while the Python proxy sends them upstream. Every key Python forwards now applies as a default under the client body, so a client value, null included, still wins A deployment's tools go first and the client's tools follow, as the Python router merges them. Client tools that are not an array stay as sent, and extra_body or any key outside that set still never reaches the provider * test(gateway): cover a deployment tools value that is not a list on /v1/messages * fix(gateway): fill model capabilities from the cost map so deployment params are shaped per model Ports the capability fill from #50780 onto the moved litellm/rust crates, so a deployment's thinking, reasoning_effort, speed, and temperature are shaped for the model instead of reaching Claude 5.5 models untranslated and failing with a 400. Also forwards a deployment's thread param, which Python's AnthropicMessagesRequestOptionalParams lists * fix(gateway): keep MCP server tokens out of the debug request log A deployment's mcp_servers now reach the provider body, so the debug provider-request log would print each server's authorization_token. The logged copy replaces it with REDACTED while the provider still receives the real token * fix(gateway): resolve os.environ/ references and honor additional_drop_params in deployment params A deployment's untyped litellm_params now resolve os.environ/NAME at config load the way the Python proxy's config loader does, including nested objects and true/false words, so a system prompt or MCP server token kept in the environment reaches the provider resolved. A deployment's additional_drop_params now reach the Messages shaping, as on the Python proxy --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/rust/Cargo.lock | 7 + litellm/rust/crates/auth-aws/src/aws.rs | 2 +- .../crates/gateway-inference/src/messages.rs | 65 ++- .../gateway-inference/tests/messages.rs | 313 +++++++++- .../gateway-inference/tests/support/mod.rs | 64 +- litellm/rust/crates/gateway/Cargo.toml | 2 + litellm/rust/crates/gateway/src/error.rs | 2 + litellm/rust/crates/gateway/src/lib.rs | 4 +- litellm/rust/crates/gateway/tests/server.rs | 148 +++++ .../crates/inference-messages/src/handler.rs | 28 +- .../tests/messages/stream.rs | 33 ++ litellm/rust/crates/llms/Cargo.toml | 1 + .../llms/src/anthropic/messages/thinking.rs | 122 +++- .../src/anthropic/messages/transformation.rs | 76 ++- litellm/rust/crates/llms/src/anthropic/mod.rs | 1 + .../llms/src/anthropic/model_capabilities.rs | 546 ++++++++++++++++++ .../crates/llms/src/bedrock/common_utils.rs | 119 ++++ litellm/rust/crates/llms/src/bedrock/mod.rs | 2 +- litellm/rust/crates/model-catalog/Cargo.toml | 1 + .../rust/crates/model-catalog/src/catalog.rs | 12 +- .../rust/crates/model-catalog/src/fallback.rs | 99 ++++ .../crates/model-catalog/src/resolve/mod.rs | 56 ++ .../crates/router-types/src/litellm_params.rs | 6 + litellm/rust/crates/router/Cargo.toml | 3 + litellm/rust/crates/router/src/lib.rs | 167 +++++- litellm/rust/crates/router/tests/router.rs | 26 +- 26 files changed, 1854 insertions(+), 51 deletions(-) create mode 100644 litellm/rust/crates/llms/src/anthropic/model_capabilities.rs diff --git a/litellm/rust/Cargo.lock b/litellm/rust/Cargo.lock index 798e85effa2..edb3a07a8ba 100644 --- a/litellm/rust/Cargo.lock +++ b/litellm/rust/Cargo.lock @@ -3864,6 +3864,7 @@ dependencies = [ "litellm-http", "litellm-inference", "litellm-llms", + "litellm-model-catalog", "litellm-redis-client", "litellm-redis-connect", "litellm-redis-resolve", @@ -3885,6 +3886,7 @@ dependencies = [ "tracing", "uuid", "veil", + "wiremock", ] [[package]] @@ -4343,6 +4345,7 @@ dependencies = [ "litellm-host", "litellm-http", "litellm-llms-types", + "litellm-model-catalog", "litellm-python-compat", "litellm-router-types", "litellm-secret-value", @@ -4380,6 +4383,7 @@ name = "litellm-model-catalog" version = "0.1.0" dependencies = [ "encoding_rs", + "fancy-regex 0.19.2", "indexmap 2.14.0", "jiff", "jsonschema", @@ -4624,8 +4628,11 @@ dependencies = [ "litellm-core-utils", "litellm-inference", "litellm-inference-messages", + "litellm-llms", + "litellm-model-catalog", "litellm-router-types", "rstest", + "serde_json", ] [[package]] diff --git a/litellm/rust/crates/auth-aws/src/aws.rs b/litellm/rust/crates/auth-aws/src/aws.rs index e6e32754b4f..8508ac0c47c 100644 --- a/litellm/rust/crates/auth-aws/src/aws.rs +++ b/litellm/rust/crates/auth-aws/src/aws.rs @@ -694,7 +694,7 @@ pub fn bedrock_model_id_and_region(model: &str) -> (String, Option) { (stripped.to_string(), region) } -fn is_bedrock_region(value: &str) -> bool { +pub fn is_bedrock_region(value: &str) -> bool { value.len() > 3 && value.contains('-') && value diff --git a/litellm/rust/crates/gateway-inference/src/messages.rs b/litellm/rust/crates/gateway-inference/src/messages.rs index 58ecd40204d..56e283d63b9 100644 --- a/litellm/rust/crates/gateway-inference/src/messages.rs +++ b/litellm/rust/crates/gateway-inference/src/messages.rs @@ -25,6 +25,31 @@ const OPENROUTER_API_HEADERS: [&str; 2] = [ ]; const ANTHROPIC_API_HEADER_PROVIDERS: &str = "anthropic,bedrock,bedrock_mantle,vertex_ai,openrouter,deepseek"; +const DEPLOYMENT_DEFAULT_PARAMS: [&str; 23] = [ + "max_tokens", + "metadata", + "stop_sequences", + "stream", + "system", + "temperature", + "thinking", + "tool_choice", + "top_k", + "inference_geo", + "top_p", + "mcp_servers", + "context_management", + "compaction", + "container", + "output_format", + "speed", + "output_config", + "cache_control", + "reasoning_effort", + "safeguards", + "service_tier", + "thread", +]; pub async fn create( State(gateway): State>, @@ -72,12 +97,15 @@ fn project( body: Map, headers: &HeaderMap, ) -> Result { - let body = body - .into_iter() - .map(|(name, value)| match name.as_str() { - "model" => (name, Value::from(deployment.model.as_str())), - _ => (name, value), - }) + let tools = + merged_tools(deployment, body.get("tools")).map(|tools| ("tools".to_owned(), tools)); + let client_params = body.into_iter().map(|(name, value)| match name.as_str() { + "model" => (name, Value::from(deployment.model.as_str())), + _ => (name, value), + }); + let body = deployment_default_params(deployment) + .chain(client_params) + .chain(tools) .collect(); Ok(MessagesCall { body: messages_body(body)?, @@ -92,6 +120,31 @@ fn project( }) } +fn deployment_default_params( + deployment: &Deployment, +) -> impl Iterator + '_ { + DEPLOYMENT_DEFAULT_PARAMS.iter().filter_map(|name| { + let value = deployment.litellm_params.extra.get(*name)?; + Some(((*name).to_owned(), value.clone())) + }) +} + +fn merged_tools(deployment: &Deployment, client_tools: Option<&Value>) -> Option { + let deployment_tools = deployment.litellm_params.extra.get("tools")?.as_array()?; + let client_tools: &[Value] = match client_tools.unwrap_or(&Value::Null) { + Value::Null => &[], + Value::Array(tools) => tools, + _ => return None, + }; + Some( + deployment_tools + .iter() + .chain(client_tools) + .cloned() + .collect(), + ) +} + fn anthropic_api_headers(headers: &HeaderMap) -> Option { let scopes = [ ( diff --git a/litellm/rust/crates/gateway-inference/tests/messages.rs b/litellm/rust/crates/gateway-inference/tests/messages.rs index 7d6d1f160d1..4406f8a9467 100644 --- a/litellm/rust/crates/gateway-inference/tests/messages.rs +++ b/litellm/rust/crates/gateway-inference/tests/messages.rs @@ -5,13 +5,15 @@ use axum::{ http::Request, }; use rstest::rstest; -use serde_json::json; +use serde_json::{Value, json}; use tower::ServiceExt; use wiremock::{ Mock, MockServer, ResponseTemplate, matchers::{body_json, header, method, path}, }; +const MESSAGE_SSE: &str = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + #[rstest] #[case::anthropic("anthropic/test-model", "/v1/messages", true)] #[case::azure("azure_ai/test-model", "/anthropic/v1/messages", false)] @@ -27,7 +29,7 @@ async fn messages_reaches_the_provider_and_preserves_json_or_sse( let message = json!({"id": "msg_test", "type": "message", "role": "assistant", "model": "test-model", "content": [{"type": "text", "text": "hello"}], "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1}}); - let sse = "event: message_start\ndata: {\"type\":\"message_start\"}\n\nevent: message_stop\ndata: {\"type\":\"message_stop\"}\n\n"; + let sse = MESSAGE_SSE; let template = if streaming { ResponseTemplate::new(200).set_body_raw(sse, "text/event-stream") } else { @@ -75,6 +77,313 @@ async fn messages_reaches_the_provider_and_preserves_json_or_sse( assert!(!requests[0].headers.contains_key("x-request-id")); } +#[rstest] +#[case::the_defaults_fill_what_the_client_omits( + json!({}), + json!({"temperature": 0.2, "top_p": 0.9, "top_k": 5, "max_tokens": 7, "stop_sequences": ["default"]}), +)] +#[case::client_values_win( + json!({"temperature": 0.7, "top_p": 0.8, "top_k": 6, "max_tokens": 16, "stop_sequences": ["client"]}), + json!({"temperature": 0.7, "top_p": 0.8, "top_k": 6, "max_tokens": 16, "stop_sequences": ["client"]}), +)] +#[case::a_client_null_clears_each_default( + json!({"temperature": null, "top_p": null, "top_k": null, "stop_sequences": null, "max_tokens": 16}), + json!({"max_tokens": 16}), +)] +#[tokio::test] +async fn deployment_sampling_params_sit_under_the_client_body( + #[case] client_params: Value, + #[case] expected_params: Value, +) { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .and(body_json(messages_request("test-model", expected_params))) + .respond_with(message_response()) + .expect(1) + .mount(&upstream) + .await; + let app = support::app_with_litellm_params( + "anthropic/test-model", + &upstream.uri(), + json!({ + "model": "anthropic/test-model", + "temperature": 0.2, + "top_p": 0.9, + "top_k": 5, + "max_tokens": 7, + "stop_sequences": ["default"], + "a_litellm_only_setting": "stays-out-of-the-provider-body", + }), + ); + + let response = support::post( + app, + "/v1/messages", + messages_request("public/model", client_params), + ) + .await; + + assert_eq!(response.status(), 200); +} + +#[rstest] +#[case::max_tokens("max_tokens", json!(7))] +#[case::metadata("metadata", json!({"user_id": "deployment-user"}))] +#[case::stop_sequences("stop_sequences", json!(["default"]))] +#[case::stream("stream", json!(false))] +#[case::system("system", json!("Answer in one word."))] +#[case::temperature("temperature", json!(0.2))] +#[case::thinking("thinking", json!({"type": "enabled", "budget_tokens": 1024}))] +#[case::tool_choice("tool_choice", json!({"type": "auto"}))] +#[case::top_k("top_k", json!(5))] +#[case::inference_geo("inference_geo", json!("us"))] +#[case::top_p("top_p", json!(0.9))] +#[case::mcp_servers( + "mcp_servers", + json!([{"type": "url", "url": "https://mcp.example.com/sse", "name": "example"}]), +)] +#[case::context_management( + "context_management", + json!({"edits": [{"type": "clear_tool_uses_20250919"}]}), +)] +#[case::compaction("compaction", json!({"type": "summarize"}))] +#[case::container( + "container", + json!({"skills": [{"type": "anthropic", "skill_id": "pptx", "version": "latest"}]}), +)] +#[case::output_format( + "output_format", + json!({"type": "json_schema", "schema": {"type": "object", "properties": {}}}), +)] +#[case::speed("speed", json!("fast"))] +#[case::output_config("output_config", json!({"effort": "low"}))] +#[case::cache_control("cache_control", json!({"type": "ephemeral"}))] +#[case::reasoning_effort("reasoning_effort", json!("low"))] +#[case::safeguards( + "safeguards", + json!([{"type": "dangerous_tool_use", "classifier_context": {"v": 1}}]), +)] +#[case::service_tier("service_tier", json!("auto"))] +#[case::thread("thread", json!({"type": "create"}))] +#[tokio::test] +async fn a_deployment_default_reaches_the_provider_as_the_same_client_value( + #[case] name: &str, + #[case] value: Value, +) { + let required = required_params_besides(name); + let from_deployment = upstream_body(json!({name: value.clone()}), required.clone()).await; + let from_client = upstream_body(json!({}), with_param(required, name, value)).await; + let baseline = upstream_body(json!({}), json!({"max_tokens": 16})).await; + + assert_eq!(from_deployment, from_client); + assert_ne!(from_deployment, baseline); +} + +#[rstest] +#[case::a_client_thinking_wins( + json!({"thinking": {"type": "enabled", "budget_tokens": 1024}}), + json!({"thinking": {"type": "disabled"}}), + json!({"thinking": {"type": "disabled"}}), +)] +#[case::a_client_null_clears_the_thinking_default( + json!({"thinking": {"type": "enabled", "budget_tokens": 1024}}), + json!({"thinking": null}), + json!({}), +)] +#[case::a_client_system_wins( + json!({"system": "From the deployment."}), + json!({"system": "From the client."}), + json!({"system": "From the client."}), +)] +#[case::a_client_null_clears_the_system_default( + json!({"system": "From the deployment."}), + json!({"system": null}), + json!({}), +)] +#[case::deployment_tools_come_before_client_tools( + json!({"tools": [tool("deployment_lookup")]}), + json!({"tools": [tool("client_lookup")]}), + json!({"tools": [tool("deployment_lookup"), tool("client_lookup")]}), +)] +#[case::deployment_tools_fill_omitted_tools( + json!({"tools": [tool("deployment_lookup")]}), + json!({}), + json!({"tools": [tool("deployment_lookup")]}), +)] +#[case::deployment_tools_replace_null_tools( + json!({"tools": [tool("deployment_lookup")]}), + json!({"tools": null}), + json!({"tools": [tool("deployment_lookup")]}), +)] +#[case::a_client_tool_choice_wins( + json!({"tools": [tool("deployment_lookup")], "tool_choice": {"type": "auto"}}), + json!({"tool_choice": {"type": "any"}}), + json!({"tools": [tool("deployment_lookup")], "tool_choice": {"type": "any"}}), +)] +#[case::extra_body_stays_out( + json!({"extra_body": {"deployment_only": true}}), + json!({}), + json!({}), +)] +#[case::an_unknown_key_stays_out( + json!({"a_litellm_only_setting": "stays-out-of-the-provider-body"}), + json!({}), + json!({}), +)] +#[case::deployment_tools_that_are_not_a_list_stay_out( + json!({"tools": tool("deployment_lookup")}), + json!({}), + json!({}), +)] +#[tokio::test] +async fn a_deployment_param_reaches_the_provider_as_its_client_equivalent( + #[case] deployment_params: Value, + #[case] client_params: Value, + #[case] client_equivalent: Value, +) { + assert_eq!( + upstream_body( + deployment_params, + with_param(client_params, "max_tokens", json!(16)) + ) + .await, + upstream_body( + json!({}), + with_param(client_equivalent, "max_tokens", json!(16)) + ) + .await, + ); +} + +#[rstest] +#[tokio::test] +async fn deployment_tools_leave_client_tools_that_are_not_an_array_as_sent() { + let upstream = MockServer::start().await; + let app = support::app_with_litellm_params( + "anthropic/test-model", + &upstream.uri(), + deployment_params(json!({"tools": [tool("deployment_lookup")]})), + ); + + let response = support::post( + app, + "/v1/messages", + messages_request( + "public/model", + json!({"max_tokens": 16, "tools": tool("client_lookup")}), + ), + ) + .await; + + assert_eq!(response.status(), 400); + assert!(upstream.received_requests().await.unwrap().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn a_deployment_stream_default_streams_a_request_that_omits_stream() { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .and(body_json(messages_request( + "test-model", + json!({"max_tokens": 16, "stream": true}), + ))) + .respond_with(ResponseTemplate::new(200).set_body_raw(MESSAGE_SSE, "text/event-stream")) + .expect(1) + .mount(&upstream) + .await; + let app = support::app_with_litellm_params( + "anthropic/test-model", + &upstream.uri(), + deployment_params(json!({"stream": true})), + ); + + let response = support::post( + app, + "/v1/messages", + messages_request("public/model", json!({"max_tokens": 16})), + ) + .await; + + assert_eq!(response.status(), 200); + assert_eq!(response.headers()["content-type"], "text/event-stream"); + assert_eq!( + to_bytes(response.into_body(), 4096).await.unwrap(), + MESSAGE_SSE + ); +} + +async fn upstream_body(deployment: Value, client: Value) -> Value { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(message_response()) + .mount(&upstream) + .await; + let app = support::app_with_shaping( + "anthropic/test-model", + &upstream.uri(), + deployment_params(deployment), + json!({"capabilities": { + "supports_reasoning": true, + "supports_adaptive_thinking": true, + "supports_output_config": true, + "supports_speed": true, + "effort_tiers": {"low": true, "medium": true, "high": true}, + }}), + ); + let response = support::post( + app, + "/v1/messages", + messages_request("public/model", client), + ) + .await; + assert_eq!(response.status(), 200, "{}", support::json(response).await); + upstream.received_requests().await.unwrap()[0] + .body_json() + .unwrap() +} + +fn message_response() -> ResponseTemplate { + ResponseTemplate::new(200).set_body_json(json!({ + "id": "msg_test", "type": "message", "role": "assistant", "model": "test-model", + "content": [{"type": "text", "text": "hello"}], "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1}, + })) +} + +fn tool(name: &str) -> Value { + json!({"name": name, "description": "Looks something up.", "input_schema": {"type": "object"}}) +} + +fn required_params_besides(name: &str) -> Value { + match name { + "max_tokens" => json!({}), + _ => json!({"max_tokens": 16}), + } +} + +fn with_param(params: Value, name: &str, value: Value) -> Value { + let mut with_param = params; + with_param[name] = value; + with_param +} + +fn deployment_params(params: Value) -> Value { + let mut litellm_params = params; + litellm_params["model"] = json!("anthropic/test-model"); + litellm_params +} + +fn messages_request(model: &str, params: Value) -> Value { + let mut request = params; + request["model"] = json!(model); + request["messages"] = json!([{"role": "user", "content": "hi"}]); + request +} + #[tokio::test] async fn invalid_messages_stays_an_anthropic_error() { let response = support::post( diff --git a/litellm/rust/crates/gateway-inference/tests/support/mod.rs b/litellm/rust/crates/gateway-inference/tests/support/mod.rs index 185cbc8594d..c3be972f021 100644 --- a/litellm/rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm/rust/crates/gateway-inference/tests/support/mod.rs @@ -72,12 +72,63 @@ pub fn app_with_cache_for_principal( ) } +pub fn app_with_litellm_params(model: &str, api_base: &str, litellm_params: Value) -> Router { + app_with_shaping( + model, + api_base, + litellm_params, + Value::Object(Default::default()), + ) +} + +pub fn app_with_shaping( + model: &str, + api_base: &str, + litellm_params: Value, + shaping: Value, +) -> Router { + deployment_app( + Deployment { + litellm_params: serde_json::from_value(litellm_params).unwrap(), + shaping: serde_json::from_value(shaping).unwrap(), + ..test_deployment(model, api_base) + }, + litellm_gateway_auth::Permissions::All, + None, + None, + ) +} + +fn test_deployment(model: &str, api_base: &str) -> Deployment { + Deployment { + model: model.into(), + api_base: Some(api_base.into()), + api_key: Some("test-key".into()), + timeout: Some(Duration::from_secs(5)), + ..Default::default() + } +} + fn configured_app( model: &str, api_base: &str, permissions: litellm_gateway_auth::Permissions, cache: Option>, principal: Option, +) -> Router { + deployment_app( + test_deployment(model, api_base), + permissions, + cache, + principal, + ) +} + +fn deployment_app( + deployment: Deployment, + permissions: litellm_gateway_auth::Permissions, + cache: Option>, + principal: Option, ) -> Router { let pool = Arc::new(HttpClientPool::new(Arc::new(PublicDnsResolver))); let http = Resolution::from(&HttpSettings::default()).config; @@ -87,18 +138,7 @@ fn configured_app( resources, http, secrets, - [( - "public/model".into(), - Deployment { - model: model.into(), - api_base: Some(api_base.into()), - api_key: Some("test-key".into()), - timeout: Some(Duration::from_secs(5)), - ..Default::default() - }, - )] - .into_iter() - .collect(), + [("public/model".into(), deployment)].into_iter().collect(), ) .unwrap(); let gateway = match cache { diff --git a/litellm/rust/crates/gateway/Cargo.toml b/litellm/rust/crates/gateway/Cargo.toml index a14df9b9c53..bea1ea1a70c 100644 --- a/litellm/rust/crates/gateway/Cargo.toml +++ b/litellm/rust/crates/gateway/Cargo.toml @@ -30,6 +30,7 @@ litellm-redis-types.workspace = true rustls.workspace = true litellm-http.workspace = true litellm-llms.workspace = true +litellm-model-catalog.workspace = true litellm-secrets.workspace = true litellm-tracing.workspace = true serde_json.workspace = true @@ -48,3 +49,4 @@ rstest.workspace = true tempfile.workspace = true tokio = { workspace = true, features = ["sync"] } tower = { version = "0.5", features = ["util"] } +wiremock.workspace = true diff --git a/litellm/rust/crates/gateway/src/error.rs b/litellm/rust/crates/gateway/src/error.rs index c267e79353d..779656fb76c 100644 --- a/litellm/rust/crates/gateway/src/error.rs +++ b/litellm/rust/crates/gateway/src/error.rs @@ -20,4 +20,6 @@ pub enum Error { RedisConnect(#[from] litellm_redis_connect::Error), #[error("MCP requires working gateway authentication")] Auth(#[from] litellm_gateway_auth::Error), + #[error(transparent)] + ModelCatalog(#[from] litellm_model_catalog::Error), } diff --git a/litellm/rust/crates/gateway/src/lib.rs b/litellm/rust/crates/gateway/src/lib.rs index 0a0191b0499..43bba5dbf56 100644 --- a/litellm/rust/crates/gateway/src/lib.rs +++ b/litellm/rust/crates/gateway/src/lib.rs @@ -24,6 +24,7 @@ use litellm_http::{ media::PublicDnsResolver, }; use litellm_inference::resources::CoreResources; +use litellm_model_catalog::{BUNDLED_REGISTRY, Catalog}; use litellm_secrets::source::EnvironmentSecrets; use litellm_tracing::ByteChunk; use uuid::Uuid; @@ -37,7 +38,8 @@ pub fn build_inference(config: &Config) -> Result, Error> { .map(|value| value.expose().to_owned()) .or_else(|| std::env::var(name).ok()) }; - let router = ModelRouter::from_config(config, &lookup); + let catalog = Catalog::parse(BUNDLED_REGISTRY)?.value; + let router = ModelRouter::from_config(config, &lookup, &catalog); let settings = &config.litellm_settings; let http_settings = HttpSettings::from_layers([ HttpSettingsLayer::from_environment(&lookup), diff --git a/litellm/rust/crates/gateway/tests/server.rs b/litellm/rust/crates/gateway/tests/server.rs index 7de7a97c870..19ef7e24eb6 100644 --- a/litellm/rust/crates/gateway/tests/server.rs +++ b/litellm/rust/crates/gateway/tests/server.rs @@ -8,11 +8,16 @@ use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; use litellm_config::Config; use litellm_gateway_inference::{Error, Gateway}; use litellm_http::ClientVariant; +use litellm_model_catalog::{BUNDLED_REGISTRY, Catalog}; use litellm_tracing::{Logger, Metadata, Record, Sink}; use rstest::{fixture, rstest}; use serde_json::{Value, json}; use tokio::{net::TcpListener, sync::oneshot, time::timeout}; use tower::ServiceExt; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{method, path}, +}; struct LogSink(mpsc::Sender); @@ -320,3 +325,146 @@ async fn readiness_answers_without_credentials_like_the_proxy_with_no_database( json!({"status": "healthy", "db": "Not connected"}) ); } + +async fn forwarded_body(deployment_params: &str, client_params: Value) -> Value { + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "msg_test", "type": "message", "role": "assistant", "model": "claude-unlisted-9", + "content": [{"type": "text", "text": "hi"}], "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1} + }))) + .mount(&upstream) + .await; + let config = Config::from_yaml(&format!( + "environment_variables:\n SYSTEM_PROMPT: Answer in French.\n MCP_TOKEN: mcp-token\nmodel_list:\n - model_name: public-model\n litellm_params:\n model: anthropic/claude-unlisted-9\n api_base: {}\n api_key: test-key\n{deployment_params}general_settings:\n master_key: gateway-key\n", + upstream.uri() + )) + .unwrap(); + let mut body = json!({ + "model": "public-model", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], + }); + body.as_object_mut() + .unwrap() + .extend(client_params.as_object().cloned().unwrap_or_default()); + let request = Request::post("/v1/messages") + .header("authorization", "Bearer gateway-key") + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(); + + let response = litellm_gateway::router( + litellm_gateway::build_inference(&config).unwrap(), + &config, + None, + None, + ) + .oneshot(request) + .await + .unwrap(); + + assert_eq!(response.status().as_u16(), 200); + serde_json::from_slice(&upstream.received_requests().await.unwrap()[0].body).unwrap() +} + +#[rstest] +#[case::system( + " system: os.environ/SYSTEM_PROMPT\n", + "/system", + json!("Answer in French.") +)] +#[case::mcp_server_token( + " mcp_servers:\n - type: url\n url: https://mcp.example.com/sse\n name: docs\n authorization_token: os.environ/MCP_TOKEN\n", + "/mcp_servers/0/authorization_token", + json!("mcp-token") +)] +#[tokio::test] +async fn a_deployment_param_from_the_environment_reaches_the_provider_resolved( + #[case] deployment_params: &str, + #[case] pointer: &str, + #[case] expected: Value, +) { + let forwarded = forwarded_body(deployment_params, json!({})).await; + + assert_eq!(forwarded.pointer(pointer), Some(&expected)); +} + +#[rstest] +#[case::listed(" additional_drop_params: [temperature]\n", None)] +#[case::unlisted("", Some(json!(0.5)))] +#[tokio::test] +async fn a_deployment_additional_drop_param_never_reaches_the_provider( + #[case] deployment_params: &str, + #[case] expected: Option, +) { + let forwarded = forwarded_body(deployment_params, json!({"temperature": 0.5})).await; + + assert_eq!(forwarded.get("temperature").cloned(), expected); +} + +#[rstest] +#[case::fast_mode("supports_speed", "speed", json!("fast"))] +#[case::adaptive_thinking("supports_adaptive_thinking", "thinking", json!({"type": "adaptive"}))] +#[tokio::test] +async fn forwards_a_param_the_catalog_says_the_model_supports( + #[case] capability: &str, + #[case] param: &str, + #[case] value: Value, +) { + let catalog = Catalog::parse(BUNDLED_REGISTRY).unwrap().value; + let (model, _) = catalog + .served() + .find(|(name, fields)| { + !name.contains('/') + && fields.get("litellm_provider").and_then(Value::as_str) == Some("anthropic") + && fields.get(capability) == Some(&Value::Bool(true)) + }) + .unwrap(); + let upstream = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/messages")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "msg_test", "type": "message", "role": "assistant", "model": model, + "content": [{"type": "text", "text": "hi"}], "stop_reason": "end_turn", + "usage": {"input_tokens": 1, "output_tokens": 1} + }))) + .mount(&upstream) + .await; + let config = Config::from_yaml(&format!( + "model_list:\n - model_name: public-model\n litellm_params:\n model: anthropic/{model}\n api_base: {}\n api_key: test-key\ngeneral_settings:\n master_key: gateway-key\n", + upstream.uri() + )) + .unwrap(); + let body: serde_json::Map = [ + ("model", json!("public-model")), + ("max_tokens", json!(8)), + ("messages", json!([{"role": "user", "content": "hi"}])), + (param, value.clone()), + ] + .into_iter() + .map(|(key, field)| (key.to_owned(), field)) + .collect(); + let request = Request::post("/v1/messages") + .header("authorization", "Bearer gateway-key") + .header("content-type", "application/json") + .body(Body::from(Value::Object(body).to_string())) + .unwrap(); + + let response = litellm_gateway::router( + litellm_gateway::build_inference(&config).unwrap(), + &config, + None, + None, + ) + .oneshot(request) + .await + .unwrap(); + + assert_eq!(response.status().as_u16(), 200); + let forwarded: Value = + serde_json::from_slice(&upstream.received_requests().await.unwrap()[0].body).unwrap(); + assert_eq!(forwarded[param], value); +} diff --git a/litellm/rust/crates/inference-messages/src/handler.rs b/litellm/rust/crates/inference-messages/src/handler.rs index d54d4f5f8d1..ffa1a1ea569 100644 --- a/litellm/rust/crates/inference-messages/src/handler.rs +++ b/litellm/rust/crates/inference-messages/src/handler.rs @@ -25,7 +25,7 @@ use litellm_llms_types::formats::messages::{ MessagesOptionalParams, MessagesRequest, MessagesResponse, }; use litellm_secrets::source::Secrets; -use litellm_tracing::ByteChunk; +use litellm_tracing::{ByteChunk, REDACTED}; use serde_json::{Map, Value}; use super::{ @@ -1028,7 +1028,31 @@ fn decoded_chunks( } fn log_request_body(provider: &str, stream: bool, body: &serde_json::Value) { - tracing::debug!(provider, stream, body = %body, "provider request"); + tracing::debug!(provider, stream, body = %without_mcp_tokens(body), "provider request"); +} + +fn without_mcp_tokens(body: &Value) -> std::borrow::Cow<'_, Value> { + let Some(servers) = body.get("mcp_servers").and_then(Value::as_array) else { + return std::borrow::Cow::Borrowed(body); + }; + let mut redacted = body.clone(); + redacted["mcp_servers"] = servers.iter().map(without_mcp_token).collect(); + std::borrow::Cow::Owned(redacted) +} + +fn without_mcp_token(server: &Value) -> Value { + match server { + Value::Object(fields) if fields.contains_key("authorization_token") => Value::Object( + fields + .iter() + .map(|(key, value)| match key.as_str() { + "authorization_token" => (key.clone(), Value::from(REDACTED)), + _ => (key.clone(), value.clone()), + }) + .collect(), + ), + other => other.clone(), + } } fn log_response_body(body: &str) { diff --git a/litellm/rust/crates/inference-messages/tests/messages/stream.rs b/litellm/rust/crates/inference-messages/tests/messages/stream.rs index c9738fa2a0d..0a20ae8f718 100644 --- a/litellm/rust/crates/inference-messages/tests/messages/stream.rs +++ b/litellm/rust/crates/inference-messages/tests/messages/stream.rs @@ -224,6 +224,39 @@ async fn debug_trace_keeps_provider_input_and_every_stream_chunk(call: MessagesC assert!(!format!("{records:?}").contains("sk-ant")); } +#[rstest] +#[tokio::test] +async fn debug_trace_hides_mcp_server_tokens_that_still_reach_the_provider(call: MessagesCall) { + let upstream = upstream([sse_response()]).await; + let servers = json!([{ + "type": "url", + "url": "https://mcp.example.com/sse", + "name": "docs", + "authorization_token": "mcp-server-token" + }]); + let call = with_fields(call, json!({"mcp_servers": servers})); + let host = RecordingStreamHost::new(streaming(call, upstream.uri()), usize::MAX); + let (sender, receiver) = mpsc::channel(); + + Logger::new(TraceSink(sender)) + .instrument(stream_through(&host)) + .await + .unwrap(); + + let records: Vec<(String, Value)> = receiver.try_iter().collect(); + let request = records + .iter() + .find(|(message, _)| message == "provider request") + .unwrap(); + let logged: Value = serde_json::from_str(request.1["body"].as_str().unwrap()).unwrap(); + assert_eq!(logged["mcp_servers"][0]["name"], "docs"); + assert!(!format!("{records:?}").contains("mcp-server-token")); + let sent: Value = upstream.received_requests().await.unwrap()[0] + .body_json() + .unwrap(); + assert_eq!(sent["mcp_servers"], servers); +} + #[rstest] #[case::at_open(1)] #[case::after_the_first_chunk(2)] diff --git a/litellm/rust/crates/llms/Cargo.toml b/litellm/rust/crates/llms/Cargo.toml index d59cd1be1d2..f779d4edf66 100644 --- a/litellm/rust/crates/llms/Cargo.toml +++ b/litellm/rust/crates/llms/Cargo.toml @@ -23,6 +23,7 @@ litellm-http.workspace = true litellm-secrets.workspace = true litellm-python-compat.workspace = true litellm-router-types.workspace = true +litellm-model-catalog.workspace = true litellm-tracing.workspace = true base64.workspace = true bytes.workspace = true diff --git a/litellm/rust/crates/llms/src/anthropic/messages/thinking.rs b/litellm/rust/crates/llms/src/anthropic/messages/thinking.rs index c6f7c1679ef..ac7d04e8659 100644 --- a/litellm/rust/crates/llms/src/anthropic/messages/thinking.rs +++ b/litellm/rust/crates/llms/src/anthropic/messages/thinking.rs @@ -539,21 +539,57 @@ fn translate_adaptive_effort_for_non_adaptive_model( }) } +pub(crate) fn drop_unsupported_effort( + request: MessagesRequest, + capabilities: &MessagesModelCapabilities, +) -> MessagesRequest { + let rejected = |level: EffortLevel| !accepts_effort(capabilities, level); + let drops_reasoning_effort = capabilities.supports_adaptive_thinking + && matches!( + &request.params.reasoning_effort, + Some(Recognized::Known(effort)) if output_effort(*effort).is_some_and(rejected) + ); + let reshaped_to_budget = !capabilities.supports_adaptive_thinking + && matches!(known_thinking(&request), Some(ThinkingConfig::Adaptive(_))); + let drops_output_effort = supports_effort_param(capabilities) + && !reshaped_to_budget + && matches!(known_effort(&request), Some(Recognized::Known(level)) if rejected(*level)); + MessagesRequest { + params: MessagesOptionalParams { + reasoning_effort: request + .params + .reasoning_effort + .filter(|_| !drops_reasoning_effort), + output_config: if drops_output_effort { + without_effort(request.params.output_config) + } else { + request.params.output_config + }, + ..request.params + }, + ..request + } +} + fn drop_incompatible_temperature_for_thinking( request: MessagesRequest, context: &ThinkingContext, ) -> MessagesRequest { - if context.capabilities.supports_adaptive_thinking { - return request; - } let pinned = request .params .temperature .as_ref() .is_some_and(|temperature| !is_one(temperature)); - let thinking_enabled = matches!(known_thinking(&request), Some(ThinkingConfig::Enabled(_))); - let effort_enabled = known_effort(&request).is_some(); - if !pinned || !(thinking_enabled || effort_enabled) { + let thinking_on = matches!( + known_thinking(&request), + Some(ThinkingConfig::Enabled(_) | ThinkingConfig::Adaptive(_)) + ); + let droppable = if context.capabilities.supports_adaptive_thinking { + context.drop_params && thinking_on + } else { + thinking_on || known_effort(&request).is_some() + }; + if !pinned || !droppable { return request; } MessagesRequest { @@ -608,6 +644,17 @@ mod tests { translate_thinking(request(fields), &context(capabilities)) } + fn translate_dropping_params( + capabilities: MessagesModelCapabilities, + fields: Value, + ) -> Result { + let context = ThinkingContext { + drop_params: true, + ..context(capabilities) + }; + translate_thinking(request(fields), &context) + } + fn overridden_budgets(overrides: &[(&str, &str)]) -> ThinkingBudgets { let env = |name: &str| { overrides @@ -1354,6 +1401,47 @@ mod tests { ); } + #[rstest] + #[case::claude_code_payload( + opus_4_7(), + claude_code_payload("medium", 8192), + 0.0, + claude_code_payload("medium", 8192) + )] + #[case::legacy_thinking_kept_on_legacy_capable_model( + sonnet_4_6(), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}), + 0.5, + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) + )] + #[case::bare_adaptive_thinking( + sonnet_4_6(), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + 0.5, + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + )] + #[case::reasoning_effort_synthesized_adaptive_thinking( + sonnet_4_6(), + serde_json::json!({"max_tokens": 8192, "reasoning_effort": "high"}), + 0.0, + serde_json::json!({ + "max_tokens": 8192, + "thinking": {"type": "adaptive", "display": "summarized"}, + "output_config": {"effort": "high"} + }) + )] + fn pinned_temperature_is_dropped_under_drop_params_when_thinking_survives_on_adaptive_model( + #[case] capabilities: MessagesModelCapabilities, + #[case] input: Value, + #[case] temperature: f64, + #[case] expected: Value, + ) { + assert_eq!( + translate_dropping_params(capabilities, with_temperature(input, temperature)), + Ok(request(expected)) + ); + } + #[rstest] #[case::temperature_one_with_enabled_thinking( haiku_4_5(), @@ -1386,18 +1474,36 @@ mod tests { 0.0, serde_json::json!({"max_tokens": 8192, "output_config": {"format": {"type": "json_schema"}}}) )] - #[case::adaptive_model( + #[case::claude_code_payload_on_adaptive_model_without_drop_params( opus_4_7(), claude_code_payload("medium", 8192), 0.0, claude_code_payload("medium", 8192) )] - #[case::legacy_thinking_on_adaptive_model( + #[case::adaptive_thinking_without_drop_params( + sonnet_4_6(), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}), + 0.5, + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "adaptive"}}) + )] + #[case::legacy_thinking_on_adaptive_model_without_drop_params( sonnet_4_6(), serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}), 0.0, serde_json::json!({"max_tokens": 8192, "thinking": {"type": "enabled", "budget_tokens": 4096}}) )] + #[case::effort_without_thinking_on_adaptive_model( + sonnet_4_6(), + serde_json::json!({"max_tokens": 8192, "output_config": {"effort": "high"}}), + 0.5, + serde_json::json!({"max_tokens": 8192, "output_config": {"effort": "high"}}) + )] + #[case::disabled_thinking_on_adaptive_model( + sonnet_4_6(), + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "disabled"}}), + 0.5, + serde_json::json!({"max_tokens": 8192, "thinking": {"type": "disabled"}}) + )] fn temperature_is_kept( #[case] capabilities: MessagesModelCapabilities, #[case] input: Value, diff --git a/litellm/rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm/rust/crates/llms/src/anthropic/messages/transformation.rs index bbb9b4e4c17..224bffa2e0b 100644 --- a/litellm/rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm/rust/crates/llms/src/anthropic/messages/transformation.rs @@ -15,7 +15,7 @@ use serde_json::{Map, Value, json}; use super::{ handler::shape_anthropic_messages_request, scalar::{is_one, python_str, python_value}, - thinking::translate_thinking, + thinking::{drop_unsupported_effort, translate_thinking}, }; use crate::anthropic::beta_headers::accepted_or_unlisted; use crate::base_llm::messages::{ @@ -599,6 +599,11 @@ pub fn drop_unsupported_params( context: &MessagesTransformContext, ) -> Result { let capabilities = &context.thinking.capabilities; + let request = if context.drop_params { + drop_unsupported_effort(request, capabilities) + } else { + request + }; let model = request.model.clone(); let reject = |param: &str, value: String, hint: &str| -> Result<(), Error> { if context.drop_params { @@ -711,6 +716,7 @@ mod tests { use super::*; use crate::anthropic::common_utils::ENCRYPTED_REASONING_SIGNATURE_PREFIX; + use crate::base_llm::messages::context::SupportedEffortTiers; use litellm_http::request::header_values; type Env = &'static [(&'static str, &'static str)]; @@ -813,6 +819,30 @@ mod tests { } } + #[fixture] + fn adaptive_without_xhigh() -> MessagesModelCapabilities { + MessagesModelCapabilities { + supports_reasoning: true, + supports_adaptive_thinking: true, + supports_legacy_thinking: true, + supports_output_config: true, + effort_tiers: SupportedEffortTiers { + max: true, + ..Default::default() + }, + ..Default::default() + } + } + + #[fixture] + fn effort_without_adaptive() -> MessagesModelCapabilities { + MessagesModelCapabilities { + supports_reasoning: true, + supports_output_config: true, + ..Default::default() + } + } + #[rstest] #[case::alone(json!({"max_tokens": null}))] #[case::ahead_of_the_param_gate(json!({"max_tokens": null, "speed": "fast"}))] @@ -846,6 +876,16 @@ mod tests { )] #[case::speed_on_a_fast_mode_model(fast_mode(), false, json!({"speed": "fast"}))] #[case::speed_on_a_fast_mode_model_under_drop_params(fast_mode(), true, json!({"speed": "fast"}))] + #[case::effort_tier_the_model_lacks_is_left_to_the_provider( + adaptive_without_xhigh(), + false, + json!({"output_config": {"effort": "xhigh"}}) + )] + #[case::effort_tier_the_model_has_under_drop_params( + adaptive_without_xhigh(), + true, + json!({"output_config": {"effort": "max"}}) + )] #[case::native_context_management_edits(unmapped(), false, json!({"context_management": {"edits": [{ "type": "clear_tool_uses_20250919", "trigger": {"type": "input_tokens", "value": 30000}, @@ -900,6 +940,21 @@ mod tests { json!({"speed": "fast", "temperature": 1.0}), json!({"temperature": 1.0}) )] + #[case::reasoning_effort_tier_the_adaptive_model_lacks( + adaptive_without_xhigh(), + json!({"reasoning_effort": "xhigh"}), + json!({}) + )] + #[case::output_config_effort_tier_the_adaptive_model_lacks( + adaptive_without_xhigh(), + json!({"output_config": {"effort": "xhigh", "format": {"type": "json_schema"}}}), + json!({"output_config": {"format": {"type": "json_schema"}}}) + )] + #[case::output_config_effort_tier_the_effort_model_lacks( + effort_without_adaptive(), + json!({"output_config": {"effort": "max"}}), + json!({}) + )] fn removed_params_are_dropped_under_drop_params( #[case] capabilities: MessagesModelCapabilities, #[case] fields: Value, @@ -908,6 +963,25 @@ mod tests { assert_eq!(transform(fields, capabilities, true), Ok(body(expected))); } + #[rstest] + fn drop_params_keeps_the_effort_sized_budget_on_a_model_without_adaptive_thinking( + effort_without_adaptive: MessagesModelCapabilities, + ) { + let adaptive = json!({"max_tokens": 64000, "thinking": {"type": "adaptive"}}); + let with_xhigh = merged( + adaptive.clone(), + json!({"output_config": {"effort": "xhigh"}}), + ); + + let dropped = transform(with_xhigh.clone(), effort_without_adaptive, true); + + assert_eq!( + dropped, + transform(with_xhigh, effort_without_adaptive, false) + ); + assert_ne!(dropped, transform(adaptive, effort_without_adaptive, true)); + } + #[rstest] #[case::temperature( sampling_removed(), diff --git a/litellm/rust/crates/llms/src/anthropic/mod.rs b/litellm/rust/crates/llms/src/anthropic/mod.rs index 72c390cfe21..78b894be0d7 100644 --- a/litellm/rust/crates/llms/src/anthropic/mod.rs +++ b/litellm/rust/crates/llms/src/anthropic/mod.rs @@ -5,5 +5,6 @@ pub mod batches; pub mod chat; pub mod count_tokens; pub mod messages; +pub mod model_capabilities; pub const ANTHROPIC_OAUTH_TOKEN_PREFIX: &str = "sk-ant-oat"; diff --git a/litellm/rust/crates/llms/src/anthropic/model_capabilities.rs b/litellm/rust/crates/llms/src/anthropic/model_capabilities.rs new file mode 100644 index 00000000000..1f31473db45 --- /dev/null +++ b/litellm/rust/crates/llms/src/anthropic/model_capabilities.rs @@ -0,0 +1,546 @@ +use std::sync::LazyLock; + +use crate::base_llm::messages::context::{MessagesModelCapabilities, SupportedEffortTiers}; +use crate::bedrock::common_utils::{ + bedrock_base_model, is_bedrock_application_inference_profile_arn, without_routing_prefixes, +}; +use fancy_regex::Regex; +use litellm_auth_aws::is_bedrock_region; +use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; +use litellm_model_catalog::{Catalog, wire::Fields}; + +const SAMPLING_RESTRICTED_NAMES: [&str; 9] = [ + "fable", "opus-4-7", "opus_4_7", "opus-4.7", "opus_4.7", "opus-4-8", "opus_4_8", "opus-4.8", + "opus_4.8", +]; +const SAMPLING_TOLERANT_PROVIDERS: [&str; 1] = ["openrouter"]; +const BACKSTOP_PREFIXES: [&str; 4] = [ + "bedrock/converse/", + "bedrock/invoke/", + "bedrock/", + "vertex_ai/", +]; + +static VERSION_SUFFIX: LazyLock = + LazyLock::new(|| Regex::new(r"-v\d+(?::\d+)?$").expect("valid regex")); +static PROFILE_MINOR: LazyLock = + LazyLock::new(|| Regex::new(r":\d+$").expect("valid regex")); +static DATED_RELEASE: LazyLock = + LazyLock::new(|| Regex::new(r"-\d{8}$").expect("valid regex")); +static DOTTED_VERSION: LazyLock = + LazyLock::new(|| Regex::new(r"(\d)\.(\d)").expect("valid regex")); +static TRAILING_NUMBER: LazyLock = + LazyLock::new(|| Regex::new(r"-\d+$").expect("valid regex")); +static TRAILING_DATE: LazyLock = + LazyLock::new(|| Regex::new(r"-\d{4}-\d{2}-\d{2}$").expect("valid regex")); + +pub fn anthropic_model_capabilities( + catalog: &Catalog, + model: &str, + custom_llm_provider: Option<&str>, + base_model: Option<&str>, + declared: &dyn Fn(&str) -> Option, +) -> MessagesModelCapabilities { + let (resolved, provider) = get_custom_llm_provider(model, custom_llm_provider) + .map_or((model, "anthropic"), |resolved| { + (resolved.model, resolved.custom_llm_provider) + }); + let (model, name) = match base_model.map(without_routing_prefixes) { + Some(base) + if is_bedrock(provider) && is_bedrock_application_inference_profile_arn(resolved) => + { + (base, base) + } + _ => (model, resolved), + }; + let probe = Probe { + catalog, + declared, + fallbacks: backstop_candidates(model), + info: catalog.model_info( + &candidate_names(name, provider) + .each_ref() + .map(String::as_str), + provider, + ), + bare: catalog.lookup(name).map(|found| found.entry.fields()), + }; + let thinking_always_on = probe.supports("thinking_always_on"); + MessagesModelCapabilities { + supports_reasoning: probe.supports("supports_reasoning"), + supports_adaptive_thinking: probe.supports("supports_adaptive_thinking"), + thinking_always_on, + thinking_on_by_default: thinking_always_on || probe.supports("thinking_on_by_default"), + supports_legacy_thinking: probe.supports("supports_legacy_thinking"), + supports_output_config: probe.supports("supports_output_config"), + supports_sampling_params: declared("supports_sampling_params").unwrap_or_else(|| { + SAMPLING_TOLERANT_PROVIDERS.contains(&provider) + || 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"), + effort_tiers: SupportedEffortTiers { + minimal: probe.supports("supports_minimal_reasoning_effort"), + low: probe.supports("supports_low_reasoning_effort"), + medium: probe.supports("supports_medium_reasoning_effort"), + high: probe.supports("supports_high_reasoning_effort"), + xhigh: probe.supports("supports_xhigh_reasoning_effort"), + max: probe.supports("supports_max_reasoning_effort"), + }, + } +} + +struct Probe<'a> { + catalog: &'a Catalog, + declared: &'a dyn Fn(&str) -> Option, + fallbacks: Vec, + info: Option, + bare: Option<&'a Fields>, +} + +impl Probe<'_> { + fn supports(&self, capability: &str) -> bool { + if let Some(supported) = (self.declared)(capability) { + return supported; + } + if let Some(supported) = flag(self.info.as_ref(), capability) { + return supported; + } + if self.info.is_some() && flag(self.bare, capability) == Some(true) { + return true; + } + backstop(self.catalog, &self.fallbacks, capability) == Some(true) + } +} + +fn flag(fields: Option<&Fields>, capability: &str) -> Option { + fields?.get(capability)?.as_bool() +} + +fn exact<'a>(catalog: &'a Catalog, name: &str) -> Option<&'a Fields> { + catalog + .lookup(name) + .filter(|found| found.matched_key == name) + .map(|found| found.entry.fields()) +} + +fn is_bedrock(provider: &str) -> bool { + matches!(provider, "bedrock" | "bedrock_converse") +} + +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)), + ) + }; + [ + combined, + model.to_owned(), + 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 = 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(), + "vertex_ai" | "gemini" | "databricks" => TRAILING_NUMBER.replace(model, "").into_owned(), + _ => TRAILING_DATE.replace(model, "").into_owned(), + } +} + +fn backstop(catalog: &Catalog, candidates: &[String], capability: &str) -> Option { + candidates + .iter() + .find_map(|name| flag(exact(catalog, name), capability)) +} + +fn backstop_candidates(model: &str) -> Vec { + let base = bedrock_base_model(model); + let primary: Vec = std::iter::once(model) + .chain( + BACKSTOP_PREFIXES + .iter() + .filter_map(|prefix| model.strip_prefix(prefix)), + ) + .map(str::to_owned) + .chain([base.to_owned(), format!("bedrock/{base}")]) + .collect(); + let normalized: Vec = primary + .iter() + .flat_map(|name| normalized_names(name)) + .collect(); + primary.into_iter().chain(normalized).collect() +} + +fn normalized_names(name: &str) -> [String; 6] { + let undotted = DOTTED_VERSION.replace_all(name, "$1-$2"); + let unversioned = VERSION_SUFFIX.replace(&undotted, "").into_owned(); + [ + VERSION_SUFFIX.replace(name, "").into_owned(), + PROFILE_MINOR.replace(name, "").into_owned(), + DATED_RELEASE.replace(name, "").into_owned(), + undotted.into_owned(), + DATED_RELEASE.replace(&unversioned, "").into_owned(), + name.rfind('@') + .filter(|at| *at > 0) + .map_or(name, |at| &name[..at]) + .to_owned(), + ] +} + +#[cfg(test)] +mod tests { + use litellm_model_catalog::{BUNDLED_REGISTRY, Catalog}; + use rstest::{fixture, rstest}; + use serde_json::Value; + + use super::anthropic_model_capabilities; + + const CATALOG: &str = r#"{ + "claude-test-5": { + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": true, + "supports_speed": true, + "supports_xhigh_reasoning_effort": true, + "aliases": ["claude-test-5-latest"] + }, + "claude-slow-5": { + "litellm_provider": "anthropic", + "mode": "chat", + "supports_reasoning": false, + "supports_adaptive_thinking": false + }, + "claude-embed-5": {"litellm_provider": "anthropic", "mode": "embedding"}, + "claude-nosample-5": {"litellm_provider": "anthropic", "mode": "chat", "supports_sampling_params": false}, + "anthropic/claude-nosample-5": {"litellm_provider": "anthropic", "mode": "chat", "supports_sampling_params": false}, + "claude-default-think-5": {"litellm_provider": "anthropic", "mode": "chat", "thinking_on_by_default": true}, + "claude-always-think-5": {"litellm_provider": "anthropic", "mode": "chat", "thinking_always_on": true}, + "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", + "supports_reasoning": true, + "supports_speed": true + }, + "fallback_generalizations": { + "rules": [ + {"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", + "pattern": "claude", + "model_info": {"supports_adaptive_thinking": true}, + "fill_missing_for_providers": ["anthropic"] + } + ] + } + }"#; + + #[fixture] + fn catalog() -> Catalog { + Catalog::parse(CATALOG.as_bytes()).unwrap().value + } + + #[rstest] + #[case::row("anthropic/claude-test-5", None, (true, true, true, true))] + #[case::alias("anthropic/claude-test-5-latest", None, (true, true, true, true))] + #[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::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))] + fn reads_capabilities_from_the_catalog( + catalog: Catalog, + #[case] model: &str, + #[case] custom_llm_provider: Option<&str>, + #[case] expected: (bool, bool, bool, bool), + ) { + let capabilities = + anthropic_model_capabilities(&catalog, model, custom_llm_provider, None, &|_| None); + + assert_eq!( + ( + capabilities.supports_reasoning, + capabilities.supports_adaptive_thinking, + capabilities.supports_speed, + capabilities.supports_sampling_params, + ), + expected + ); + } + + #[rstest] + #[case::default_on("anthropic/claude-default-think-5", true)] + #[case::always_on("anthropic/claude-always-think-5", true)] + #[case::off("anthropic/claude-test-5", false)] + fn thinking_on_by_default_follows_the_row( + catalog: Catalog, + #[case] model: &str, + #[case] expected: bool, + ) { + assert_eq!( + anthropic_model_capabilities(&catalog, model, None, None, &|_| None) + .thinking_on_by_default, + expected + ); + } + + #[rstest] + fn effort_tiers_follow_the_row(catalog: Catalog) { + let tiers = + anthropic_model_capabilities(&catalog, "anthropic/claude-test-5", None, None, &|_| { + None + }) + .effort_tiers; + + assert_eq!((tiers.xhigh, tiers.low), (true, false)); + } + + #[rstest] + #[case::declares_a_model_the_map_lacks( + "anthropic/claude-other-7", + &[("supports_reasoning", true)], + (true, true, false, true) + )] + #[case::turns_off_a_filled_flag( + "anthropic/claude-test-5", + &[("supports_adaptive_thinking", false)], + (true, false, true, true) + )] + #[case::turns_off_sampling("anthropic/claude-test-5", &[("supports_sampling_params", false)], (true, true, true, false))] + #[case::speed_still_needs_its_row("anthropic/claude-slow-5", &[("supports_speed", true)], (false, false, false, true))] + #[case::undeclared_keeps_the_row("anthropic/claude-test-5", &[], (true, true, true, true))] + fn declared_flags_win_over_the_catalog( + catalog: Catalog, + #[case] model: &str, + #[case] declared: &[(&str, bool)], + #[case] expected: (bool, bool, bool, bool), + ) { + let lookup = |capability: &str| { + declared + .iter() + .find(|(name, _)| *name == capability) + .map(|(_, supported)| *supported) + }; + let capabilities = anthropic_model_capabilities(&catalog, model, None, None, &lookup); + + assert_eq!( + ( + capabilities.supports_reasoning, + capabilities.supports_adaptive_thinking, + capabilities.supports_speed, + capabilities.supports_sampling_params, + ), + expected + ); + } + + #[rstest] + #[case::anthropic_row("anthropic/claude-nosample-5", &[], false)] + #[case::anthropic_name("anthropic/claude-fable-test", &[], false)] + #[case::openrouter_row("openrouter/anthropic/claude-nosample-5", &[], true)] + #[case::openrouter_name("openrouter/anthropic/claude-fable-test", &[], true)] + #[case::openrouter_declared_off( + "openrouter/anthropic/claude-fable-test", + &[("supports_sampling_params", false)], + false + )] + fn sampling_params_are_gated_only_where_the_provider_rejects_them( + catalog: Catalog, + #[case] model: &str, + #[case] declared: &[(&str, bool)], + #[case] expected: bool, + ) { + let lookup = |capability: &str| { + declared + .iter() + .find(|(name, _)| *name == capability) + .map(|(_, supported)| *supported) + }; + + assert_eq!( + anthropic_model_capabilities(&catalog, model, None, None, &lookup) + .supports_sampling_params, + expected + ); + } + + const PROFILE_ARN: &str = + "bedrock/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abc123"; + + #[rstest] + #[case::base_model_names_the_profile(PROFILE_ARN, Some("anthropic.claude-test-5-v1:0"), true)] + #[case::prefixed_base_model(PROFILE_ARN, Some("bedrock/anthropic.claude-test-5-v1:0"), true)] + #[case::profile_without_base_model(PROFILE_ARN, None, false)] + #[case::named_model_keeps_its_own_row( + "bedrock/anthropic.claude-test-5-v1:0", + Some("anthropic.claude-slow-5"), + true + )] + fn application_inference_profile_reads_its_base_model( + catalog: Catalog, + #[case] model: &str, + #[case] base_model: Option<&str>, + #[case] expected: bool, + ) { + assert_eq!( + anthropic_model_capabilities(&catalog, model, None, base_model, &|_| None) + .supports_reasoning, + expected + ); + } + + #[rstest] + fn every_bundled_anthropic_row_shapes_with_its_own_flags() { + let catalog = Catalog::parse(BUNDLED_REGISTRY).unwrap().value; + let mismatches: Vec = catalog + .served() + .filter(|(name, fields)| { + !name.contains('/') + && fields.get("litellm_provider").and_then(Value::as_str) == Some("anthropic") + }) + .flat_map(|(name, fields)| { + let capabilities = anthropic_model_capabilities( + &catalog, + &format!("anthropic/{name}"), + None, + None, + &|_| None, + ); + [ + ("supports_reasoning", capabilities.supports_reasoning), + ( + "supports_adaptive_thinking", + capabilities.supports_adaptive_thinking, + ), + ( + "supports_output_config", + capabilities.supports_output_config, + ), + ( + "supports_sampling_params", + capabilities.supports_sampling_params, + ), + ("supports_speed", capabilities.supports_speed), + ( + "supports_legacy_thinking", + capabilities.supports_legacy_thinking, + ), + ("thinking_always_on", capabilities.thinking_always_on), + ( + "thinking_on_by_default", + capabilities.thinking_on_by_default, + ), + ( + "supports_minimal_reasoning_effort", + capabilities.effort_tiers.minimal, + ), + ( + "supports_low_reasoning_effort", + capabilities.effort_tiers.low, + ), + ( + "supports_medium_reasoning_effort", + capabilities.effort_tiers.medium, + ), + ( + "supports_high_reasoning_effort", + capabilities.effort_tiers.high, + ), + ( + "supports_mid_conversation_system", + capabilities.supports_mid_conversation_system, + ), + ( + "supports_xhigh_reasoning_effort", + capabilities.effort_tiers.xhigh, + ), + ( + "supports_max_reasoning_effort", + capabilities.effort_tiers.max, + ), + ] + .into_iter() + .filter(|(flag, shaped)| { + fields + .get(*flag) + .and_then(Value::as_bool) + .is_some_and(|row| row != *shaped) + }) + .map(move |(flag, _)| format!("{name} {flag}")) + }) + .collect(); + + assert_eq!(mismatches, Vec::::new()); + } +} diff --git a/litellm/rust/crates/llms/src/bedrock/common_utils.rs b/litellm/rust/crates/llms/src/bedrock/common_utils.rs index 6c6fb8562f0..05f6871529a 100644 --- a/litellm/rust/crates/llms/src/bedrock/common_utils.rs +++ b/litellm/rust/crates/llms/src/bedrock/common_utils.rs @@ -1,5 +1,88 @@ +use litellm_auth_aws::is_bedrock_region; + use crate::{Error, ErrorDetail}; +const ROUTE_PREFIXES: [&str; 3] = ["bedrock/converse/", "bedrock/", "converse/"]; +const ROUTING_PREFIXES: [&str; 8] = [ + "bedrock/", + "chat_completions/", + "converse/", + "invoke/", + "openai/", + "mantle/", + "nova-2/", + "nova/", +]; +const CROSS_REGION_INFERENCE_PREFIXES: [&str; 7] = + ["global", "us", "eu", "apac", "jp", "au", "us-gov"]; + +pub fn bedrock_base_model(model: &str) -> &str { + let unrouted = ROUTE_PREFIXES + .into_iter() + .find_map(|prefix| model.strip_prefix(prefix)) + .unwrap_or(model); + if unrouted.starts_with("nova-2/") { + return "amazon.nova-2-custom"; + } + if unrouted.starts_with("nova/") { + return "amazon.nova-custom"; + } + let named = without_throughput_suffix(model_name_from_arn(without_routing_prefixes(model))); + if let Some((region, rest)) = named.split_once('.') + && CROSS_REGION_INFERENCE_PREFIXES.contains(®ion) + { + return rest; + } + match named.split_once('/') { + Some((region, rest)) if is_bedrock_region(region) => rest, + _ => named, + } +} + +pub fn is_bedrock_application_inference_profile_arn(model: &str) -> bool { + model.contains(":application-inference-profile/") +} + +pub fn without_routing_prefixes(model: &str) -> &str { + ROUTING_PREFIXES.into_iter().fold(model, |model, prefix| { + model.strip_prefix(prefix).unwrap_or(model) + }) +} + +fn model_name_from_arn(model: &str) -> &str { + if !model.to_lowercase().contains("arn") { + return model; + } + model.rsplit('/').next().unwrap_or(model) +} + +fn without_throughput_suffix(model: &str) -> &str { + let unthrottled = model + .strip_suffix('k') + .and_then(|rest| rest.rsplit_once(':')) + .filter(|(head, units)| { + is_digits(units) + && head + .rsplit_once(':') + .is_some_and(|(_, minor)| is_digits(minor)) + }) + .map_or(model, |(head, _)| head); + unthrottled + .strip_suffix(']') + .and_then(|rest| rest.rsplit_once('[')) + .filter(|(_, tag)| { + !tag.is_empty() + && tag + .chars() + .all(|char| char.is_alphanumeric() || char == '_') + }) + .map_or(unthrottled, |(head, _)| head) +} + +fn is_digits(value: &str) -> bool { + !value.is_empty() && value.bytes().all(|byte| byte.is_ascii_digit()) +} + pub(crate) fn validated_region(region: String) -> Result { if !region.is_empty() && region @@ -14,3 +97,39 @@ pub(crate) fn validated_region(region: String) -> Result { actual: region.into(), })) } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::bedrock_base_model; + + #[rstest] + #[case::cross_region("us.anthropic.claude-x-v1:0", "anthropic.claude-x-v1:0")] + #[case::converse_route("bedrock/converse/anthropic.claude-x", "anthropic.claude-x")] + #[case::mantle_route( + "bedrock/mantle/anthropic.claude-sonnet-5", + "anthropic.claude-sonnet-5" + )] + #[case::invoke_route_cross_region("invoke/global.anthropic.claude-x", "anthropic.claude-x")] + #[case::throughput_suffix("anthropic.claude-x-v2:0:51k", "anthropic.claude-x-v2:0")] + #[case::version_pair_without_throughput("anthropic.claude-x:1:2", "anthropic.claude-x:1:2")] + #[case::context_tag("anthropic.claude-x[1m]", "anthropic.claude-x")] + #[case::region_path("us-east-1/anthropic.claude-x", "anthropic.claude-x")] + #[case::region_path_keeps_cross_region( + "bedrock/us-east-1/us.anthropic.claude-x", + "us.anthropic.claude-x" + )] + #[case::profile_arn( + "arn:aws:bedrock:us-east-1:123:application-inference-profile/us.anthropic.claude-x", + "anthropic.claude-x" + )] + #[case::nova_2_custom( + "bedrock/nova-2/arn:aws:bedrock:us-east-1:1:custom-model/x", + "amazon.nova-2-custom" + )] + #[case::nova_custom("bedrock/nova/arn:aws:x", "amazon.nova-custom")] + fn base_model_matches_the_python_resolution(#[case] model: &str, #[case] expected: &str) { + assert_eq!(bedrock_base_model(model), expected); + } +} diff --git a/litellm/rust/crates/llms/src/bedrock/mod.rs b/litellm/rust/crates/llms/src/bedrock/mod.rs index e3160e547b4..b2dc8937a8a 100644 --- a/litellm/rust/crates/llms/src/bedrock/mod.rs +++ b/litellm/rust/crates/llms/src/bedrock/mod.rs @@ -1,5 +1,5 @@ pub mod audio_transcription; pub mod chat; pub mod claude_platform; -pub(crate) mod common_utils; +pub mod common_utils; pub mod messages; diff --git a/litellm/rust/crates/model-catalog/Cargo.toml b/litellm/rust/crates/model-catalog/Cargo.toml index d85bb799924..1e807cf3788 100644 --- a/litellm/rust/crates/model-catalog/Cargo.toml +++ b/litellm/rust/crates/model-catalog/Cargo.toml @@ -14,6 +14,7 @@ litellm-error.workspace = true litellm-tracing.workspace = true encoding_rs = "0.8.35" +fancy-regex.workspace = true indexmap = { version = "2.14.0", features = ["serde"] } jiff = { version = "0.2.37", default-features = false, features = ["std", "tzdb-bundle-always"] } macro_rules_attribute.workspace = true diff --git a/litellm/rust/crates/model-catalog/src/catalog.rs b/litellm/rust/crates/model-catalog/src/catalog.rs index 05725d06863..70ec96cc557 100644 --- a/litellm/rust/crates/model-catalog/src/catalog.rs +++ b/litellm/rust/crates/model-catalog/src/catalog.rs @@ -7,7 +7,7 @@ use serde::Deserialize; use serde_json::Value; use crate::error::{Error, Loaded, Rejects, Repr, pointer_child, pointer_for_key}; -use crate::fallback::{FallbackGeneralizations, FallbackRule}; +use crate::fallback::{CapabilityRules, FallbackGeneralizations, FallbackRule}; use crate::model::Model; use crate::store::{Rows, Store}; use crate::wire::{EntryWire, Fields}; @@ -63,6 +63,7 @@ pub struct Catalog { store: Store, sections: Vec
, fallback_generalizations: Option, + capability_rules: CapabilityRules, } /// Models and aliases, an alias served as its owner's row @@ -246,6 +247,10 @@ impl Catalog { }) .ok() }); + let capability_rules = fallback_generalizations + .as_ref() + .map(|generalizations| CapabilityRules::new(&generalizations.rules)) + .unwrap_or_default(); Ok(Loaded { value: Self { store: Store::new(CatalogRows { @@ -255,6 +260,7 @@ impl Catalog { }), sections, fallback_generalizations, + capability_rules, }, errors: problems, }) @@ -481,6 +487,10 @@ impl Catalog { pub fn fallback_rules(&self) -> Option<&[FallbackRule]> { Some(self.fallback_generalizations.as_ref()?.rules.as_slice()) } + + pub fn capability_rules(&self) -> &CapabilityRules { + &self.capability_rules + } } struct RowRead { diff --git a/litellm/rust/crates/model-catalog/src/fallback.rs b/litellm/rust/crates/model-catalog/src/fallback.rs index f5a0cadb6b9..8f0b3b7e8b6 100644 --- a/litellm/rust/crates/model-catalog/src/fallback.rs +++ b/litellm/rust/crates/model-catalog/src/fallback.rs @@ -1,6 +1,10 @@ +use fancy_regex::{Regex, RegexBuilder}; +use serde::Deserialize; use serde_json::Value; use std::collections::BTreeMap; +use crate::wire::Fields; + /// One regex rule generalizing unknown model ids to known families. #[macro_rules_attribute::apply(crate::wire_type)] #[derive(Eq, Hash)] @@ -20,3 +24,98 @@ pub struct FallbackRule { pub struct FallbackGeneralizations { pub rules: Vec, } + +#[derive(Clone, Debug, Default)] +pub struct CapabilityRules(Vec); + +#[derive(Clone, Debug)] +struct CapabilityRule { + pattern: Regex, + model_info: Fields, + fill_missing_for_providers: Vec, +} + +impl CapabilityRules { + pub fn new(rules: &[FallbackRule]) -> Self { + Self(rules.iter().filter_map(CapabilityRule::compile).collect()) + } + + pub fn matching(&self, model: &str) -> Option { + let matched: Vec<&CapabilityRule> = + self.0.iter().filter(|rule| rule.matches(model)).collect(); + (!matched.is_empty()).then(|| merged(matched)) + } + + pub fn filling(&self, model: &str, provider: &str) -> Fields { + merged( + self.0 + .iter() + .filter(|rule| rule.fills(provider) && rule.matches(model)), + ) + } +} + +impl CapabilityRule { + fn compile(rule: &FallbackRule) -> Option { + let Some(Value::Object(model_info)) = rule.extra.get("model_info") else { + return None; + }; + if rule.extra.contains_key("extends") || model_info.contains_key("litellm_provider") { + return None; + } + let fill_missing_for_providers = match rule.extra.get("fill_missing_for_providers") { + Some(providers) => Vec::::deserialize(providers).ok()?, + None => Vec::new(), + }; + Some(Self { + pattern: pattern(rule).ok()?, + model_info: model_info.clone(), + fill_missing_for_providers, + }) + } + + fn matches(&self, model: &str) -> bool { + self.pattern.is_match(model).unwrap_or(false) + } + + fn fills(&self, provider: &str) -> bool { + self.fill_missing_for_providers + .iter() + .any(|listed| listed == provider) + } +} + +fn pattern(rule: &FallbackRule) -> Result { + RegexBuilder::new(&rule.pattern) + .case_insensitive(true) + .build() +} + +fn merged<'a>(rules: impl IntoIterator) -> Fields { + rules + .into_iter() + .flat_map(|rule| rule.model_info.clone()) + .collect() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::pattern; + use crate::{BUNDLED_REGISTRY, Catalog}; + + #[rstest] + fn every_bundled_fallback_pattern_compiles() { + let catalog = Catalog::parse(BUNDLED_REGISTRY).unwrap().value; + let broken: Vec<&str> = catalog + .fallback_rules() + .unwrap_or_default() + .iter() + .filter(|rule| pattern(rule).is_err()) + .map(|rule| rule.name.as_str()) + .collect(); + + assert_eq!(broken, Vec::<&str>::new()); + } +} diff --git a/litellm/rust/crates/model-catalog/src/resolve/mod.rs b/litellm/rust/crates/model-catalog/src/resolve/mod.rs index fbaead622aa..438ea17d4ab 100644 --- a/litellm/rust/crates/model-catalog/src/resolve/mod.rs +++ b/litellm/rust/crates/model-catalog/src/resolve/mod.rs @@ -2,6 +2,12 @@ mod image; pub use image::*; +use serde_json::Value; + +use crate::{Catalog, wire::Fields}; + +const BACKFILL_MODES: [&str; 2] = ["chat", "responses"]; + /// The seven names main's `get_model_info` tries, built by Python, walked in this order by Rust #[derive(Clone, Debug, Default, PartialEq, Eq)] pub struct Candidates { @@ -44,3 +50,53 @@ pub enum Resolution { tried: Vec, }, } + +impl Catalog { + pub fn model_info(&self, names: &[&str], provider: &str) -> Option { + let registered = names.iter().find_map(|name| { + let found = self.lookup(name)?; + let entry_provider = found + .entry + .field("litellm_provider") + .and_then(Value::as_str); + provider_matches(entry_provider, provider).then(|| { + self.with_missing_filled( + found.matched_key, + found.entry.fields(), + entry_provider.unwrap_or_default(), + ) + }) + }); + if registered.is_some() || names.iter().any(|name| self.lookup(name).is_some()) { + return registered; + } + names + .iter() + .find_map(|name| self.capability_rules().matching(name)) + } + + fn with_missing_filled(&self, key: &str, fields: &Fields, entry_provider: &str) -> Fields { + let backfills = fields.get("mode").is_none_or(|mode| { + mode.as_str() + .is_some_and(|mode| BACKFILL_MODES.contains(&mode)) + }); + let missing = if backfills { + self.capability_rules().filling(key, entry_provider) + } else { + Fields::new() + }; + missing.into_iter().chain(fields.clone()).collect() + } +} + +fn provider_matches(entry_provider: Option<&str>, provider: &str) -> bool { + let Some(entry_provider) = entry_provider else { + return true; + }; + provider.is_empty() + || entry_provider == provider + || matches!(provider, "vertex_ai" | "fireworks_ai") && entry_provider.starts_with(provider) + || provider.starts_with("bedrock") && entry_provider.starts_with("bedrock") + || matches!(provider, "litellm_proxy" | "github" | "nadir") + || provider == "azure_ai" && matches!(entry_provider, "azure" | "openai") +} diff --git a/litellm/rust/crates/router-types/src/litellm_params.rs b/litellm/rust/crates/router-types/src/litellm_params.rs index aa2c47f6e3c..021216dd63b 100644 --- a/litellm/rust/crates/router-types/src/litellm_params.rs +++ b/litellm/rust/crates/router-types/src/litellm_params.rs @@ -87,6 +87,12 @@ impl Deref for ExtraParams { } } +impl FromIterator<(String, Value)> for ExtraParams { + fn from_iter>(entries: I) -> Self { + Self(entries.into_iter().collect()) + } +} + impl fmt::Debug for ExtraParams { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter.debug_list().entries(self.0.keys()).finish() diff --git a/litellm/rust/crates/router/Cargo.toml b/litellm/rust/crates/router/Cargo.toml index 988d1b68805..a3bb8ee4ea8 100644 --- a/litellm/rust/crates/router/Cargo.toml +++ b/litellm/rust/crates/router/Cargo.toml @@ -10,7 +10,10 @@ litellm-config.workspace = true litellm-core-utils.workspace = true litellm-inference.workspace = true litellm-inference-messages.workspace = true +litellm-llms.workspace = true +litellm-model-catalog.workspace = true litellm-router-types.workspace = true +serde_json.workspace = true [dev-dependencies] rstest.workspace = true diff --git a/litellm/rust/crates/router/src/lib.rs b/litellm/rust/crates/router/src/lib.rs index 6955b8a927a..1f018967bd9 100644 --- a/litellm/rust/crates/router/src/lib.rs +++ b/litellm/rust/crates/router/src/lib.rs @@ -3,8 +3,15 @@ mod deployment; use std::collections::HashMap; use litellm_config::{Config, Spelled}; -use litellm_core_utils::{serde_compat::parse_pydantic_bool, settings::Lookup}; +use litellm_core_utils::{ + serde_compat::{parse_pydantic_bool, parse_str_bool}, + settings::Lookup, +}; use litellm_inference_messages::{MessagesSettings, MessagesShaping}; +use litellm_llms::anthropic::model_capabilities::anthropic_model_capabilities; +use litellm_model_catalog::Catalog; +use litellm_router_types::{ExtraParams, LitellmParams}; +use serde_json::{Map, Value}; pub use deployment::Deployment; @@ -24,8 +31,54 @@ fn is_on(flag: Option<&Spelled>, env: &impl Lookup) -> bool { } } +fn environment_value(name: &str, env: &impl Lookup) -> Value { + env.get(name).map_or(Value::Null, |value| { + parse_str_bool(&value).map_or(Value::String(value), Value::Bool) + }) +} + +fn with_environment(value: &Value, env: &impl Lookup) -> Value { + match value { + Value::String(text) => text + .strip_prefix("os.environ/") + .map_or_else(|| value.clone(), |name| environment_value(name, env)), + Value::Object(fields) => Value::Object(fields_with_environment(fields, env)), + Value::Array(items) => Value::Array( + items + .iter() + .map(|item| match item { + Value::Object(fields) => Value::Object(fields_with_environment(fields, env)), + _ => item.clone(), + }) + .collect(), + ), + _ => value.clone(), + } +} + +fn fields_with_environment(fields: &Map, env: &impl Lookup) -> Map { + fields + .iter() + .map(|(name, value)| (name.clone(), with_environment(value, env))) + .collect() +} + +fn additional_drop_params(extra: &ExtraParams) -> Vec { + extra + .get("additional_drop_params") + .and_then(Value::as_array) + .map(|names| { + names + .iter() + .filter_map(Value::as_str) + .map(str::to_owned) + .collect() + }) + .unwrap_or_default() +} + impl Router { - pub fn from_config(config: &Config, env: &impl Lookup) -> Self { + pub fn from_config(config: &Config, env: &impl Lookup, catalog: &Catalog) -> Self { let drops_params_globally = is_on(config.litellm_settings.drop_params.as_ref(), env); config .model_list @@ -33,6 +86,15 @@ impl Router { .map(|model| { let drop_params = drops_params_globally || is_on(model.litellm_params.drop_params.as_ref(), env); + let litellm_params = LitellmParams { + extra: model + .litellm_params + .extra + .iter() + .map(|(name, value)| (name.clone(), with_environment(value, env))) + .collect(), + ..model.litellm_params.clone() + }; ( model.model_name.clone(), Deployment { @@ -44,15 +106,32 @@ impl Router { .map(|value| value.expose().to_string()), api_base: model.litellm_params.api_base.clone(), custom_llm_provider: model.litellm_params.custom_llm_provider.clone(), - litellm_params: model.litellm_params.clone(), shaping: MessagesShaping { settings: MessagesSettings { drop_params, drops_unsupported_params: drop_params, + additional_drop_params: additional_drop_params( + &litellm_params.extra, + ), ..MessagesSettings::default() }, - ..MessagesShaping::default() + capabilities: anthropic_model_capabilities( + catalog, + &model.litellm_params.model, + model.litellm_params.custom_llm_provider.as_deref(), + litellm_params + .extra + .get("base_model") + .and_then(|value| value.as_str()), + &|capability| { + model + .model_info + .get(capability) + .and_then(|value| value.as_bool()) + }, + ), }, + litellm_params, ..Deployment::default() }, ) @@ -74,7 +153,9 @@ impl FromIterator<(String, Deployment)> for Router { #[cfg(test)] mod tests { use litellm_config::Config; + use litellm_model_catalog::Catalog; use rstest::rstest; + use serde_json::{Value, json}; use super::Router; @@ -115,9 +196,13 @@ mod tests { )) .unwrap(); - let settings = Router::from_config(&config, &|name: &str| { - (name == "DROP_PARAMS_FLAG").then(|| " Yes ".to_owned()) - }) + let settings = Router::from_config( + &config, + &|name: &str| (name == "DROP_PARAMS_FLAG").then(|| " Yes ".to_owned()), + &Catalog::parse(br#"{"provider-model": {"litellm_provider": "provider"}}"#) + .unwrap() + .value, + ) .get("public-model") .unwrap() .shaping @@ -129,4 +214,72 @@ mod tests { (expected, expected) ); } + + #[rstest] + #[case::text(json!("os.environ/PROMPT"), json!("Answer in French."))] + #[case::true_word(json!("os.environ/ON_WORD"), json!(true))] + #[case::false_word(json!("os.environ/OFF_WORD"), json!(false))] + #[case::unset(json!("os.environ/UNSET"), Value::Null)] + #[case::plain(json!("Answer in French."), json!("Answer in French."))] + #[case::nested_object(json!({"inner": "os.environ/PROMPT"}), json!({"inner": "Answer in French."}))] + #[case::object_in_a_list( + json!([{"authorization_token": "os.environ/PROMPT"}]), + json!([{"authorization_token": "Answer in French."}]) + )] + #[case::text_in_a_list(json!(["os.environ/PROMPT"]), json!(["os.environ/PROMPT"]))] + fn deployment_params_resolve_os_environ_references_as_python_does( + #[case] configured: Value, + #[case] expected: Value, + ) { + let config = Config::from_yaml(&format!( + "model_list:\n - model_name: public-model\n litellm_params:\n model: provider/model\n system: {configured}" + )) + .unwrap(); + let environment = |name: &str| match name { + "PROMPT" => Some("Answer in French.".to_owned()), + "ON_WORD" => Some(" True ".to_owned()), + "OFF_WORD" => Some("false".to_owned()), + _ => None, + }; + + let router = Router::from_config( + &config, + &environment, + &Catalog::parse(br#"{"provider-model": {"litellm_provider": "provider"}}"#) + .unwrap() + .value, + ); + + assert_eq!( + router.get("public-model").unwrap().litellm_params.extra["system"], + expected + ); + } + + #[rstest] + #[case::declared("\n model_info:\n supports_reasoning: true", true)] + #[case::declared_off("\n model_info:\n supports_reasoning: false", false)] + #[case::undeclared("", false)] + fn model_info_flags_reach_the_deployment(#[case] model_info: &str, #[case] expected: bool) { + let config = Config::from_yaml(&format!( + "model_list:\n - model_name: public-model\n litellm_params:\n model: anthropic/claude-unlisted-9{model_info}" + )) + .unwrap(); + + assert_eq!( + Router::from_config( + &config, + &|_: &str| None, + &Catalog::parse(br#"{"provider-model": {"litellm_provider": "provider"}}"#) + .unwrap() + .value + ) + .get("public-model") + .unwrap() + .shaping + .capabilities + .supports_reasoning, + expected + ); + } } diff --git a/litellm/rust/crates/router/tests/router.rs b/litellm/rust/crates/router/tests/router.rs index 03e78ab9d11..1ab0727e25f 100644 --- a/litellm/rust/crates/router/tests/router.rs +++ b/litellm/rust/crates/router/tests/router.rs @@ -2,8 +2,16 @@ use std::time::Duration; 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() + .value +} #[rstest] #[case::minimal("")] @@ -11,12 +19,12 @@ use rstest::rstest; "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); + 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; @@ -33,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 @@ -44,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); + let router = Router::from_config(&config, &|_: &str| None, &catalog); let params = &router.get("public-model").unwrap().litellm_params; assert_eq!( @@ -63,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 @@ -74,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); + 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) + Router::from_config(&config, &|_: &str| None, &catalog) .get("") .is_none() );