diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 48b6730b0c6..6a0bd13196b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3834,6 +3834,7 @@ dependencies = [ "litellm-http", "litellm-inference", "litellm-llms", + "litellm-model-catalog", "litellm-redis-client", "litellm-redis-connect", "litellm-redis-resolve", @@ -3855,6 +3856,7 @@ dependencies = [ "tracing", "uuid", "veil", + "wiremock", ] [[package]] @@ -3926,6 +3928,7 @@ dependencies = [ "litellm-inference-transcription", "litellm-llms", "litellm-llms-types", + "litellm-model-catalog", "litellm-router", "litellm-secret-value", "litellm-secrets", @@ -4570,8 +4573,10 @@ dependencies = [ "litellm-core-utils", "litellm-inference", "litellm-inference-messages", + "litellm-router-core", "litellm-router-types", "rstest", + "serde_json", ] [[package]] diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 69ca561e4b5..7278b4d0f33 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -24,6 +24,7 @@ litellm-host-http.workspace = true litellm-host.workspace = true litellm-http.workspace = true litellm-llms.workspace = true +litellm-model-catalog.workspace = true litellm-router.workspace = true litellm-secrets.workspace = true litellm-llms-types.workspace = true diff --git a/litellm-rust/crates/gateway-inference/src/chat_completions.rs b/litellm-rust/crates/gateway-inference/src/chat_completions.rs index fb8a6f4609d..973ae0a0b56 100644 --- a/litellm-rust/crates/gateway-inference/src/chat_completions.rs +++ b/litellm-rust/crates/gateway-inference/src/chat_completions.rs @@ -6,11 +6,13 @@ use axum::{ extract::State, response::{IntoResponse, Response}, }; +use litellm_inference::provider::resolve_llm_provider; use litellm_inference_chat::types::ChatCompletionsCall; +use litellm_model_catalog::Catalog; use serde_json::{Map, Value}; use crate::{ - Error, Gateway, JsonObject, + Deployment, Error, Gateway, JsonObject, request::{self, Path}, }; @@ -55,6 +57,9 @@ async fn handle( None => route, }; + let max_output_tokens = deployment + .max_output_tokens + .or_else(|| catalog_max_output_tokens(&(gateway.catalog)(), deployment)); let messages = body.get("messages").cloned().unwrap_or_default(); let headers = crate::caching::CacheHeaders::default(); let response = litellm_host_http::serve_unary( @@ -71,6 +76,7 @@ async fn handle( custom_llm_provider: deployment.custom_llm_provider.clone(), extra_headers: None, timeout: deployment.timeout, + max_output_tokens, }, cache_options.policy, ), @@ -82,3 +88,16 @@ async fn handle( .await?; Ok(headers.apply(response)) } + +fn catalog_max_output_tokens(catalog: &Catalog, deployment: &Deployment) -> Option { + let provider = resolve_llm_provider( + &deployment.model, + deployment.custom_llm_provider.as_deref(), + "chat completions", + ) + .ok()?; + let found = catalog + .lookup(provider.model) + .filter(|found| found.matched_key == provider.model)?; + i64::try_from(found.entry.model()?.attributes.max_output_tokens?).ok() +} diff --git a/litellm-rust/crates/gateway-inference/src/lib.rs b/litellm-rust/crates/gateway-inference/src/lib.rs index 43966f2bee7..4a7e68e4c19 100644 --- a/litellm-rust/crates/gateway-inference/src/lib.rs +++ b/litellm-rust/crates/gateway-inference/src/lib.rs @@ -23,14 +23,18 @@ use litellm_inference_ocr::OcrRoute; use litellm_inference_responses::ResponsesRoute; use litellm_inference_transcription::AudioTranscriptionRoute; use litellm_llms::base_llm::ocr::{handler::OcrClient, settings::OcrSettings}; +use litellm_model_catalog::Catalog; use litellm_secrets::source::SecretSource; pub use error::{BodyError, Error}; pub use litellm_router::{Deployment, Router as ModelRouter}; pub use request::{JsonObject, RequestId}; +pub type CatalogSource = Arc Arc + Send + Sync>; + pub struct Gateway { cache: Option>, + catalog: CatalogSource, pub audio_transcription: AudioTranscriptionRoute, pub chat_completions: ChatCompletionsRoute, pub messages: MessagesRoute, @@ -55,11 +59,13 @@ impl Gateway { http: HttpClientConfig, secrets: Arc, models: ModelRouter, + catalog: CatalogSource, ) -> Result { let provider = resources.pool.client(&http, ClientVariant::Provider)?; let auth = resources.auth.clone(); Ok(Self { cache: None, + catalog, audio_transcription: AudioTranscriptionRoute::new( provider.clone(), auth.clone(), diff --git a/litellm-rust/crates/gateway-inference/tests/routes.rs b/litellm-rust/crates/gateway-inference/tests/routes.rs index 24b8c5cc694..ee4a6d06021 100644 --- a/litellm-rust/crates/gateway-inference/tests/routes.rs +++ b/litellm-rust/crates/gateway-inference/tests/routes.rs @@ -6,19 +6,23 @@ use axum::{ body::{Body, Bytes}, http::Request, }; +use litellm_gateway_inference::Deployment; use litellm_http::{ ClientVariant, HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver, }; use litellm_inference::resources::CoreResources; use litellm_inference_chat::{ChatCompletionsRoute, types::ChatCompletionsRequest}; +use litellm_llms::anthropic::chat::transformation::DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS; use rstest::rstest; use serde_json::{Value, json}; use tower::ServiceExt; use wiremock::{ - Mock, MockServer, ResponseTemplate, - matchers::{body_partial_json, method}, + Mock, MockServer, Request as UpstreamRequest, ResponseTemplate, + matchers::{any, body_partial_json, method}, }; +const OUTPUT_CAP_CATALOG: &str = r#"{"test-model": {"litellm_provider": "anthropic", "mode": "chat", "max_output_tokens": 4321}}"#; + #[rstest] #[case::chat("/chat/completions", Some("public/model"))] #[case::versioned_chat("/v1/chat/completions", Some("public/model"))] @@ -67,6 +71,70 @@ async fn chat_aliases_call_core_and_use_the_body_model_before_the_path( ); } +#[rstest] +#[case::catalog_cap("anthropic/test-model", None, json!({}), json!(4321))] +#[case::deployment_cap_over_catalog("anthropic/test-model", Some(1234), json!({}), json!(1234))] +#[case::unmapped_model( + "anthropic/unmapped-model", + None, + json!({}), + json!(DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS) +)] +#[case::differently_cased_model( + "anthropic/TEST-MODEL", + None, + json!({}), + json!(DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS) +)] +#[case::client_null_max_tokens("anthropic/test-model", None, json!({"max_tokens": null}), json!(4321))] +#[case::client_max_tokens("anthropic/test-model", Some(1234), json!({"max_tokens": 16}), json!(16))] +#[tokio::test] +async fn chat_sends_anthropic_max_tokens_from_the_client_then_the_model_output_cap( + #[case] model: &str, + #[case] deployment_cap: Option, + #[case] fields: Value, + #[case] expected: Value, +) { + let upstream = MockServer::start().await; + Mock::given(any()) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "msg_test", "model": "test-model", "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", "usage": {"input_tokens": 1, "output_tokens": 1} + }))) + .mount(&upstream) + .await; + let app = support::app_with_catalog( + Deployment { + max_output_tokens: deployment_cap, + ..support::deployment(model, &upstream.uri()) + }, + OUTPUT_CAP_CATALOG, + ); + let body = fields + .as_object() + .unwrap() + .clone() + .into_iter() + .chain([ + ("model".into(), json!("public/model")), + ( + "messages".into(), + json!([{"role": "user", "content": "hi"}]), + ), + ]) + .collect(); + + let response = support::post(app, "/v1/chat/completions", Value::Object(body)).await; + + assert_eq!(response.status(), 200); + let [sent] = <[UpstreamRequest; 1]>::try_from(upstream.received_requests().await.unwrap()) + .unwrap_or_else(|requests| panic!("one upstream request, saw {}", requests.len())); + assert_eq!( + serde_json::from_slice::(&sent.body).unwrap()["max_tokens"], + expected + ); +} + #[rstest] #[case::streaming(json!({"messages": [{"role": "user", "content": "hi"}], "stream": true}), 501)] #[case::missing_messages(json!({}), 400)] diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 185cbc8594d..bdd2c76d05f 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -13,6 +13,7 @@ use futures_util::future::BoxFuture; use litellm_gateway_inference::{Deployment, Gateway, router}; use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; use litellm_inference::resources::CoreResources; +use litellm_model_catalog::Catalog; use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; @@ -72,33 +73,62 @@ pub fn app_with_cache_for_principal( ) } +pub fn 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() + } +} + +const UNRELATED_CATALOG: &str = r#"{"unrelated-model": {"litellm_provider": "openai"}}"#; + +pub fn app_with_catalog(deployment: Deployment, catalog: &str) -> Router { + gateway_app( + deployment, + Catalog::parse(catalog.as_bytes()).unwrap().value, + litellm_gateway_auth::Permissions::All, + None, + None, + ) +} + fn configured_app( model: &str, api_base: &str, permissions: litellm_gateway_auth::Permissions, cache: Option>, principal: Option, +) -> Router { + gateway_app( + deployment(model, api_base), + Catalog::parse(UNRELATED_CATALOG.as_bytes()).unwrap().value, + permissions, + cache, + principal, + ) +} + +fn gateway_app( + deployment: Deployment, + catalog: Catalog, + 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; let secrets = Arc::new(NoSecrets); let resources = CoreResources::new(pool); + let catalog = Arc::new(catalog); let gateway = Gateway::new( 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(), + Arc::new(move || Arc::clone(&catalog)), ) .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..48e898ad81f 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("the bundled model cost map is invalid")] + BundledModelCatalog(#[source] litellm_model_catalog::Error), } diff --git a/litellm-rust/crates/gateway/src/lib.rs b/litellm-rust/crates/gateway/src/lib.rs index 0a0191b0499..0a4ebc400ee 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; @@ -65,7 +66,19 @@ pub fn build_inference(config: &Config) -> Result, Error> { Arc::new(EnvironmentSecrets::python_compatible(client)), )); let resources = CoreResources::new(pool); - Ok(Arc::new(Gateway::new(resources, http, secrets, router)?)) + let catalog = Arc::new(bundled_catalog()?); + Ok(Arc::new(Gateway::new( + resources, + http, + secrets, + router, + Arc::new(move || Arc::clone(&catalog)), + )?)) +} + +fn bundled_catalog() -> Result { + let loaded = Catalog::parse(BUNDLED_REGISTRY).map_err(Error::BundledModelCatalog)?; + Ok(loaded.value.without_rejected(&loaded.errors).0) } pub fn install_crypto_provider() { diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index 7de7a97c870..d5b3b062e99 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -8,11 +8,13 @@ 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, Mode}; 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::any}; struct LogSink(mpsc::Sender); @@ -320,3 +322,60 @@ async fn readiness_answers_without_credentials_like_the_proxy_with_no_database( json!({"status": "healthy", "db": "Not connected"}) ); } + +fn bundled_anthropic_chat_output_cap() -> (String, u64) { + let catalog = Catalog::parse(BUNDLED_REGISTRY).unwrap().value; + catalog + .model_names() + .filter(|name| !name.contains('/')) + .find_map(|name| { + let attributes = &catalog.lookup(name)?.entry.model()?.attributes; + let anthropic_chat = attributes.litellm_provider.as_deref() == Some("anthropic") + && attributes.mode == Some(Mode::Chat); + anthropic_chat.then_some((name.to_owned(), attributes.max_output_tokens?)) + }) + .unwrap() +} + +#[tokio::test] +async fn chat_sends_the_bundled_output_cap_as_anthropic_max_tokens() { + let (model, output_cap) = bundled_anthropic_chat_output_cap(); + let upstream = MockServer::start().await; + Mock::given(any()) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "id": "msg_test", "model": model, "content": [{"type": "text", "text": "hello"}], + "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 inference = litellm_gateway::build_inference(&config).unwrap(); + let request = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("authorization", "Bearer gateway-key") + .header("content-type", "application/json") + .body(Body::from( + json!({"model": "public-model", "messages": [{"role": "user", "content": "hi"}]}) + .to_string(), + )) + .unwrap(); + + let response = litellm_gateway::router(inference, &config, None, None) + .oneshot(request) + .await + .unwrap(); + + assert_eq!(response.status().as_u16(), 200); + let [sent] = <[wiremock::Request; 1]>::try_from(upstream.received_requests().await.unwrap()) + .unwrap_or_else(|requests| panic!("one upstream request, saw {}", requests.len())); + assert_eq!( + serde_json::from_slice::(&sent.body).unwrap()["max_tokens"], + json!(output_cap) + ); +} diff --git a/litellm-rust/crates/inference-chat/src/common_utils.rs b/litellm-rust/crates/inference-chat/src/common_utils.rs index 1c177d58552..73cc22bcd64 100644 --- a/litellm-rust/crates/inference-chat/src/common_utils.rs +++ b/litellm-rust/crates/inference-chat/src/common_utils.rs @@ -72,6 +72,7 @@ pub(super) fn resolve_request( api_base: call.api_base, extra_headers: call.extra_headers, timeout: call.timeout, + max_output_tokens: call.max_output_tokens, }) } diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index cdcc3e911b9..dc6bf2b36c8 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -11,7 +11,7 @@ use litellm_llms::base_llm::{ }; use litellm_llms_types::formats::chat_completions::ChatCompletionsResponse; use litellm_secrets::source::Secrets; -use serde_json::Value; +use serde_json::{Map, Value}; use super::Error; use crate::{ @@ -45,8 +45,12 @@ pub(super) async fn execute( &call.optional_params, &|key| secrets.get(key), )?; + let optional_params = under_client_params( + config.get_config(call.max_output_tokens), + call.optional_params, + ); let body = config - .transform_request(provider.model, call.messages, call.optional_params.clone())? + .transform_request(provider.model, call.messages, optional_params.clone())? .body; let environment = ValidatedEnvironment { headers: with_default_headers(validated.headers, config.default_headers()), @@ -56,7 +60,7 @@ pub(super) async fn execute( let context = RequestContext { model: provider.model.to_owned(), custom_llm_provider: provider_name.to_owned(), - optional_params: Value::Object(call.optional_params), + optional_params: Value::Object(optional_params), secret_fields: Vec::new(), api_key: call.api_key.map(litellm_secret_value::SecretValue::new), }; @@ -134,6 +138,17 @@ pub(super) async fn execute( .await } +fn under_client_params( + defaults: Map, + client: Map, +) -> Map { + let unset: Vec<(String, Value)> = defaults + .into_iter() + .filter(|(name, _)| client.get(name).is_none_or(Value::is_null)) + .collect(); + client.into_iter().chain(unset).collect() +} + #[cfg(test)] mod tests { use std::sync::Mutex; @@ -148,6 +163,7 @@ mod tests { types::{ChatCompletionsCall, ResolvedChatCompletionsRequest}, }; use litellm_inference::provider::resolve_llm_provider; + use litellm_llms::anthropic::chat::transformation::DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS; const ANTHROPIC_MESSAGE: &str = r#"{"id":"msg_1","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[{"type":"text","text":"hello"}],"stop_reason":"end_turn","stop_sequence":null,"usage":{"input_tokens":11,"output_tokens":4}}"#; @@ -186,6 +202,8 @@ mod tests { fn prepared( api_base: &str, + optional_params: Value, + max_output_tokens: Option, ) -> ( litellm_inference::provider::ResolvedProvider<'static>, &'static dyn BaseConfig, @@ -200,12 +218,13 @@ mod tests { ChatCompletionsCall { model: "anthropic/claude-sonnet-4-5".into(), messages: json!([{"role": "user", "content": "hi"}]), - optional_params: json!({"max_tokens": 16}).as_object().unwrap().clone(), + optional_params: optional_params.as_object().unwrap().clone(), api_key: Some("sk-test".into()), api_base: Some(api_base.into()), custom_llm_provider: None, extra_headers: None, timeout: None, + max_output_tokens, }, config, ) @@ -225,7 +244,8 @@ mod tests { .await; let interceptors = RecordingHooks::default(); - let (provider, config, request) = prepared(&upstream.uri()); + let (provider, config, request) = + prepared(&upstream.uri(), json!({"max_tokens": 16}), None); execute( &Client::for_test(), &AuthServices::default(), @@ -260,6 +280,55 @@ mod tests { assert_eq!(interceptors.raw.into_inner().unwrap(), [ANTHROPIC_MESSAGE]); } + #[rstest] + #[case::output_cap(json!({}), Some(4321), json!(4321))] + #[case::no_output_cap(json!({}), None, json!(DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS))] + #[case::client_value(json!({"max_tokens": 16}), Some(4321), json!(16))] + #[case::client_null(json!({"max_tokens": null}), Some(4321), json!(4321))] + #[tokio::test] + async fn an_unset_max_tokens_takes_the_provider_default_before_the_hooks_run( + #[case] optional_params: Value, + #[case] max_output_tokens: Option, + #[case] expected: Value, + ) { + let upstream = MockServer::start().await; + Mock::given(any()) + .respond_with( + ResponseTemplate::new(200).set_body_raw(ANTHROPIC_MESSAGE, "application/json"), + ) + .mount(&upstream) + .await; + let interceptors = RecordingHooks::default(); + + let (provider, config, request) = + prepared(&upstream.uri(), optional_params, max_output_tokens); + execute( + &Client::for_test(), + &AuthServices::default(), + config, + provider, + request, + std::sync::Arc::new(|_: &str| None), + None, + None, + &interceptors, + None, + ) + .await + .expect("chat completions call succeeds"); + + let [request] = <[Request; 1]>::try_from(upstream.received_requests().await.unwrap()) + .unwrap_or_else(|requests| panic!("one request, saw {}", requests.len())); + let sent: Value = serde_json::from_slice(&request.body).unwrap(); + assert_eq!(sent["max_tokens"], expected); + let [context] = + <[RequestContext; 1]>::try_from(interceptors.contexts.into_inner().unwrap()) + .unwrap_or_else(|seen| { + panic!("before_provider_request runs once, saw {}", seen.len()) + }); + assert_eq!(context.optional_params, json!({"max_tokens": expected})); + } + #[rstest] #[tokio::test] async fn an_upstream_failure_is_not_reported_as_a_received_response() { @@ -270,7 +339,8 @@ mod tests { .await; let interceptors = RecordingHooks::default(); - let (provider, config, request) = prepared(&upstream.uri()); + let (provider, config, request) = + prepared(&upstream.uri(), json!({"max_tokens": 16}), None); let error = execute( &Client::for_test(), &AuthServices::default(), diff --git a/litellm-rust/crates/inference-chat/src/types.rs b/litellm-rust/crates/inference-chat/src/types.rs index 23b1d26de54..7df97f8109d 100644 --- a/litellm-rust/crates/inference-chat/src/types.rs +++ b/litellm-rust/crates/inference-chat/src/types.rs @@ -29,6 +29,7 @@ pub struct ChatCompletionsCall { pub custom_llm_provider: Option, pub extra_headers: Option>, pub timeout: Option, + pub max_output_tokens: Option, } impl From> for ChatCompletionsCall { @@ -42,6 +43,7 @@ impl From> for ChatCompletionsCall { custom_llm_provider: request.custom_llm_provider.map(str::to_owned), extra_headers: request.extra_headers, timeout: request.timeout, + max_output_tokens: None, } } } @@ -53,4 +55,5 @@ pub struct ResolvedChatCompletionsRequest { pub api_base: Option, pub extra_headers: Option>, pub timeout: Option, + pub max_output_tokens: Option, } diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index 164dbce1264..f2dd102eefb 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -48,6 +48,8 @@ const SUPPORTED_PARAMS: &[(&str, &str)] = &[ ("stop", "stop_sequences"), ]; +pub const DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS: i64 = 128_000; + #[derive(Deserialize)] struct TextResponseProjection { model: String, @@ -96,6 +98,13 @@ impl BaseConfig for AnthropicConfig { SUPPORTED_PARAMS } + fn get_config(&self, max_output_tokens: Option) -> Map { + Map::from_iter([( + "max_tokens".to_owned(), + json!(max_output_tokens.unwrap_or(DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS)), + )]) + } + fn get_complete_url( &self, api_base: Option<&str>, @@ -257,7 +266,9 @@ fn anthropic_body( mod tests { use crate::{ Error, - anthropic::chat::transformation::ANTHROPIC_CHAT_COMPLETIONS_CONFIG, + anthropic::chat::transformation::{ + ANTHROPIC_CHAT_COMPLETIONS_CONFIG, DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS, + }, base_llm::{ auth::AuthScheme, chat::transformation::{BaseConfig, ProviderChatResponseData}, @@ -319,6 +330,17 @@ mod tests { ); } + #[rstest] + #[case::output_cap(Some(4321), 4321)] + #[case::negative_output_cap(Some(-5), -5)] + #[case::no_output_cap(None, DEFAULT_ANTHROPIC_CHAT_MAX_TOKENS)] + fn defaults_max_tokens_to_the_output_cap(#[case] cap: Option, #[case] expected: i64) { + assert_eq!( + ANTHROPIC_CHAT_COMPLETIONS_CONFIG.get_config(cap), + params(json!({"max_tokens": expected})) + ); + } + #[test] fn omits_system_when_no_system_message_is_present() { let body = transform( diff --git a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs index 8ba4030540c..981963ebb74 100644 --- a/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/chat/transformation.rs @@ -36,6 +36,10 @@ pub trait BaseConfig: Sync { /// Supported OpenAI parameter names paired with their provider names. fn supported_openai_param_mappings(&self) -> &'static [(&'static str, &'static str)]; + fn get_config(&self, _max_output_tokens: Option) -> Map { + Map::new() + } + fn get_complete_url( &self, api_base: Option<&str>, diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs index 4db9b37cbb8..a5d2abe002c 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs @@ -27,6 +27,7 @@ pub(super) fn project( custom_llm_provider: call.options.custom_llm_provider, extra_headers: call.options.extra_headers, timeout: call.options.timeout, + max_output_tokens: None, }) } diff --git a/litellm-rust/crates/router/Cargo.toml b/litellm-rust/crates/router/Cargo.toml index 988d1b68805..d80358c58b0 100644 --- a/litellm-rust/crates/router/Cargo.toml +++ b/litellm-rust/crates/router/Cargo.toml @@ -10,7 +10,9 @@ litellm-config.workspace = true litellm-core-utils.workspace = true litellm-inference.workspace = true litellm-inference-messages.workspace = true +litellm-router-core.workspace = true litellm-router-types.workspace = true +serde_json.workspace = true [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/router/src/deployment.rs b/litellm-rust/crates/router/src/deployment.rs index d57284bae5a..3775e63c974 100644 --- a/litellm-rust/crates/router/src/deployment.rs +++ b/litellm-rust/crates/router/src/deployment.rs @@ -12,4 +12,5 @@ pub struct Deployment { pub litellm_params: LitellmParams, pub timeout: Option, pub shaping: MessagesShaping, + pub max_output_tokens: Option, } diff --git a/litellm-rust/crates/router/src/lib.rs b/litellm-rust/crates/router/src/lib.rs index b75c1a14446..d33d0ec67e5 100644 --- a/litellm-rust/crates/router/src/lib.rs +++ b/litellm-rust/crates/router/src/lib.rs @@ -2,9 +2,10 @@ mod deployment; use std::collections::HashMap; -use litellm_config::{Config, Spelled}; +use litellm_config::{Config, Object, Spelled}; use litellm_core_utils::{serde_compat::parse_pydantic_bool, settings::Lookup}; use litellm_inference_messages::{MessagesSettings, MessagesShaping}; +use litellm_router_core::group_info::coerce_token_limit; pub use deployment::Deployment; @@ -24,6 +25,15 @@ fn is_on(flag: Option<&Spelled>, env: &impl Lookup) -> bool { } } +fn configured_max_output_tokens(model_info: &Object) -> Option { + coerce_token_limit( + model_info + .get("max_output_tokens") + .and_then(|value| serde_json::to_value(value).ok()) + .as_ref(), + ) +} + impl Router { pub fn from_config(config: &Config, env: &impl Lookup) -> Self { let drops_params_globally = is_on(config.litellm_settings.drop_params.as_ref(), env); @@ -51,6 +61,7 @@ impl Router { }, ..MessagesShaping::default() }, + max_output_tokens: configured_max_output_tokens(&model.model_info), ..Deployment::default() }, ) @@ -124,4 +135,32 @@ mod tests { expected ); } + + #[rstest] + #[case::integer("\n model_info:\n max_output_tokens: 4321", Some(4321))] + #[case::numeric_string("\n model_info:\n max_output_tokens: \" 4_321 \"", Some(4321))] + #[case::float("\n model_info:\n max_output_tokens: 4321.9", Some(4321))] + #[case::zero("\n model_info:\n max_output_tokens: 0", Some(0))] + #[case::negative("\n model_info:\n max_output_tokens: -5", Some(-5))] + #[case::absent("", None)] + #[case::null("\n model_info:\n max_output_tokens: null", None)] + #[case::thousands_separator("\n model_info:\n max_output_tokens: \"4,321\"", None)] + #[case::boolean("\n model_info:\n max_output_tokens: true", None)] + fn configured_max_output_tokens_reach_the_deployment( + #[case] model_info: &str, + #[case] expected: Option, + ) { + let config = Config::from_yaml(&format!( + "model_list:\n - model_name: public-model\n litellm_params:\n model: provider/model{model_info}" + )) + .unwrap(); + + assert_eq!( + Router::from_config(&config, &|_: &str| None::) + .get("public-model") + .unwrap() + .max_output_tokens, + expected + ); + } } diff --git a/litellm-rust/crates/router/tests/router.rs b/litellm-rust/crates/router/tests/router.rs index 03e78ab9d11..1f4351713db 100644 --- a/litellm-rust/crates/router/tests/router.rs +++ b/litellm-rust/crates/router/tests/router.rs @@ -100,6 +100,7 @@ fn programmatic_deployments_preserve_overrides_and_last_entry_wins() { custom_llm_provider: Some("test-provider".into()), litellm_params: Default::default(), timeout: Some(Duration::from_secs(7)), + max_output_tokens: Some(4321), shaping: MessagesShaping { settings: MessagesSettings { drop_params: true, @@ -120,5 +121,6 @@ fn programmatic_deployments_preserve_overrides_and_last_entry_wins() { assert_eq!(selected.api_base, deployment.api_base); assert_eq!(selected.custom_llm_provider, deployment.custom_llm_provider); assert_eq!(selected.timeout, deployment.timeout); + assert_eq!(selected.max_output_tokens, deployment.max_output_tokens); assert_eq!(selected.shaping, deployment.shaping); }