scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(gateway): default max_tokens for Anthropic chat completions

Browse source 
When neither the client nor the deployment set max_tokens, the Rust gateway sent Anthropic chat requests without it and Anthropic answered 400 "max_tokens: Field required". The gateway now fills max_tokens the same way the Python SDK does, taking the deployment's model_info.max_output_tokens, then the bundled cost map's max_output_tokens for the exact model name, then 128000

A client max_tokens always wins, and a client null counts as unset. The bundled cost map is parsed once when the gateway starts and is injected into the inference gateway
This commit is contained in:
mateo-berri 2026-11-07 16:39:39 -08:00
parent 7617e66acf
commit 0a7f656ef1
20 changed files with 374 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
litellm-rust/Cargo.lock generated
View file

@ -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]]

1
litellm-rust/crates/gateway-inference/Cargo.toml
View file

@ -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

21
litellm-rust/crates/gateway-inference/src/chat_completions.rs
View file

@ -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<i64> {
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()
}

6
litellm-rust/crates/gateway-inference/src/lib.rs
View file

@ -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<dyn Fn() -> Arc<Catalog> + Send + Sync>;
pub struct Gateway {
cache: Option<Arc<dyn litellm_cache_response::ResponseCacheService>>,
catalog: CatalogSource,
pub audio_transcription: AudioTranscriptionRoute,
pub chat_completions: ChatCompletionsRoute,
pub messages: MessagesRoute,
@ -55,11 +59,13 @@ impl Gateway {
http: HttpClientConfig,
secrets: Arc<dyn SecretSource>,
models: ModelRouter,
catalog: CatalogSource,
) -> Result<Self, litellm_http::Error> {
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(),

72
litellm-rust/crates/gateway-inference/tests/routes.rs
View file

@ -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<i64>,
#[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::<Value>(&sent.body).unwrap()["max_tokens"],
expected
);
}
#[rstest]
#[case::streaming(json!({"messages": [{"role": "user", "content": "hi"}], "stream": true}), 501)]
#[case::missing_messages(json!({}), 400)]

54
litellm-rust/crates/gateway-inference/tests/support/mod.rs
View file

@ -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<Arc<dyn litellm_cache_response::ResponseCacheService>>,
principal: Option<litellm_gateway_auth::Principal>,
) -> 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<Arc<dyn litellm_cache_response::ResponseCacheService>>,
principal: Option<litellm_gateway_auth::Principal>,
) -> 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 {

2
litellm-rust/crates/gateway/Cargo.toml
View file

@ -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

2
litellm-rust/crates/gateway/src/error.rs
View file

@ -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),
}

15
litellm-rust/crates/gateway/src/lib.rs
View file

@ -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<Arc<Gateway>, 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<Catalog, Error> {
let loaded = Catalog::parse(BUNDLED_REGISTRY).map_err(Error::BundledModelCatalog)?;
Ok(loaded.value.without_rejected(&loaded.errors).0)
}
pub fn install_crypto_provider() {

59
litellm-rust/crates/gateway/tests/server.rs
View file

@ -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<Value>);
@ -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::<Value>(&sent.body).unwrap()["max_tokens"],
json!(output_cap)
);
}

1
litellm-rust/crates/inference-chat/src/common_utils.rs
View file

@ -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,
})
}

82
litellm-rust/crates/inference-chat/src/handler.rs
View file

@ -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<String, Value>,
client: Map<String, Value>,
) -> Map<String, Value> {
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<i64>,
) -> (
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<i64>,
#[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(),

3
litellm-rust/crates/inference-chat/src/types.rs
View file

@ -29,6 +29,7 @@ pub struct ChatCompletionsCall {
pub custom_llm_provider: Option<String>,
pub extra_headers: Option<Map<String, Value>>,
pub timeout: Option<Duration>,
pub max_output_tokens: Option<i64>,
}
impl From<ChatCompletionsRequest<'_>> for ChatCompletionsCall {
@ -42,6 +43,7 @@ impl From<ChatCompletionsRequest<'_>> 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<String>,
pub extra_headers: Option<Map<String, Value>>,
pub timeout: Option<Duration>,
pub max_output_tokens: Option<i64>,
}

24
litellm-rust/crates/llms/src/anthropic/chat/transformation.rs
View file

@ -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<i64>) -> Map<String, Value> {
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<i64>, #[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(

4
litellm-rust/crates/llms/src/base_llm/chat/transformation.rs
View file

@ -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<i64>) -> Map<String, Value> {
Map::new()
}
fn get_complete_url(
&self,
api_base: Option<&str>,

1
litellm-rust/crates/python-bridge/src/routes/chat_completions/host.rs
View file

@ -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,
})
}

2
litellm-rust/crates/router/Cargo.toml
View file

@ -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

1
litellm-rust/crates/router/src/deployment.rs
View file

@ -12,4 +12,5 @@ pub struct Deployment {
pub litellm_params: LitellmParams,
pub timeout: Option<Duration>,
pub shaping: MessagesShaping,
pub max_output_tokens: Option<i64>,
}

41
litellm-rust/crates/router/src/lib.rs
View file

@ -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<bool>>, env: &impl Lookup) -> bool {
}
}
fn configured_max_output_tokens(model_info: &Object) -> Option<i64> {
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<i64>,
) {
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::<String>)
.get("public-model")
.unwrap()
.max_output_tokens,
expected
);
}
}

2
litellm-rust/crates/router/tests/router.rs
View file

@ -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);
}