mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-15 09:12:18 +00:00
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>
This commit is contained in:
norkelbyte
GitHub
parent
1f6e401a68
commit
8140fdf54f
26 changed files with 1854 additions and 51 deletions
7
litellm/rust/Cargo.lock
generated
7
litellm/rust/Cargo.lock
generated
|
|
@ -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]]
|
||||
|
|
|
|||
|
|
@ -694,7 +694,7 @@ pub fn bedrock_model_id_and_region(model: &str) -> (String, Option<String>) {
|
|||
(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
|
||||
|
|
|
|||
|
|
@ -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<Arc<Gateway>>,
|
||||
|
|
@ -72,12 +97,15 @@ fn project(
|
|||
body: Map<String, Value>,
|
||||
headers: &HeaderMap,
|
||||
) -> Result<MessagesCall, Error> {
|
||||
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<Item = (String, Value)> + '_ {
|
||||
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<Value> {
|
||||
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<ProviderSpecificHeaders> {
|
||||
let scopes = [
|
||||
(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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<Arc<dyn litellm_cache_response::ResponseCacheService>>,
|
||||
principal: Option<litellm_gateway_auth::Principal>,
|
||||
) -> Router {
|
||||
deployment_app(
|
||||
test_deployment(model, api_base),
|
||||
permissions,
|
||||
cache,
|
||||
principal,
|
||||
)
|
||||
}
|
||||
|
||||
fn deployment_app(
|
||||
deployment: Deployment,
|
||||
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;
|
||||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<Arc<Gateway>, 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),
|
||||
|
|
|
|||
|
|
@ -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<Value>);
|
||||
|
||||
|
|
@ -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<Value>,
|
||||
) {
|
||||
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<String, Value> = [
|
||||
("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);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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)]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<MessagesRequest, Error> {
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -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<MessagesRequest, Error> {
|
||||
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(),
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
|
|
|
|||
546
litellm/rust/crates/llms/src/anthropic/model_capabilities.rs
Normal file
546
litellm/rust/crates/llms/src/anthropic/model_capabilities.rs
Normal file
|
|
@ -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<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"-v\d+(?::\d+)?$").expect("valid regex"));
|
||||
static PROFILE_MINOR: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r":\d+$").expect("valid regex"));
|
||||
static DATED_RELEASE: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"-\d{8}$").expect("valid regex"));
|
||||
static DOTTED_VERSION: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"(\d)\.(\d)").expect("valid regex"));
|
||||
static TRAILING_NUMBER: LazyLock<Regex> =
|
||||
LazyLock::new(|| Regex::new(r"-\d+$").expect("valid regex"));
|
||||
static TRAILING_DATE: LazyLock<Regex> =
|
||||
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<bool>,
|
||||
) -> 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<bool>,
|
||||
fallbacks: Vec<String>,
|
||||
info: Option<Fields>,
|
||||
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<bool> {
|
||||
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<bool> {
|
||||
candidates
|
||||
.iter()
|
||||
.find_map(|name| flag(exact(catalog, name), capability))
|
||||
}
|
||||
|
||||
fn backstop_candidates(model: &str) -> Vec<String> {
|
||||
let base = bedrock_base_model(model);
|
||||
let primary: Vec<String> = 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<String> = 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<String> = 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::<String>::new());
|
||||
}
|
||||
}
|
||||
|
|
@ -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<String, Error> {
|
||||
if !region.is_empty()
|
||||
&& region
|
||||
|
|
@ -14,3 +97,39 @@ pub(crate) fn validated_region(region: String) -> Result<String, Error> {
|
|||
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);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<CatalogRows>,
|
||||
sections: Vec<Section>,
|
||||
fallback_generalizations: Option<FallbackGeneralizations>,
|
||||
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 {
|
||||
|
|
|
|||
|
|
@ -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<FallbackRule>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, Default)]
|
||||
pub struct CapabilityRules(Vec<CapabilityRule>);
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct CapabilityRule {
|
||||
pattern: Regex,
|
||||
model_info: Fields,
|
||||
fill_missing_for_providers: Vec<String>,
|
||||
}
|
||||
|
||||
impl CapabilityRules {
|
||||
pub fn new(rules: &[FallbackRule]) -> Self {
|
||||
Self(rules.iter().filter_map(CapabilityRule::compile).collect())
|
||||
}
|
||||
|
||||
pub fn matching(&self, model: &str) -> Option<Fields> {
|
||||
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<Self> {
|
||||
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::<String>::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<Regex, fancy_regex::Error> {
|
||||
RegexBuilder::new(&rule.pattern)
|
||||
.case_insensitive(true)
|
||||
.build()
|
||||
}
|
||||
|
||||
fn merged<'a>(rules: impl IntoIterator<Item = &'a CapabilityRule>) -> 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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
},
|
||||
}
|
||||
|
||||
impl Catalog {
|
||||
pub fn model_info(&self, names: &[&str], provider: &str) -> Option<Fields> {
|
||||
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")
|
||||
}
|
||||
|
|
|
|||
|
|
@ -87,6 +87,12 @@ impl Deref for ExtraParams {
|
|||
}
|
||||
}
|
||||
|
||||
impl FromIterator<(String, Value)> for ExtraParams {
|
||||
fn from_iter<I: IntoIterator<Item = (String, Value)>>(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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<bool>>, 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<String, Value>, env: &impl Lookup) -> Map<String, Value> {
|
||||
fields
|
||||
.iter()
|
||||
.map(|(name, value)| (name.clone(), with_environment(value, env)))
|
||||
.collect()
|
||||
}
|
||||
|
||||
fn additional_drop_params(extra: &ExtraParams) -> Vec<String> {
|
||||
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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
);
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue