From fb44188a14729c445037b0e6d1c0d85bbc5728de Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Tue, 10 Nov 2026 20:40:29 +0000 Subject: [PATCH] feat(messages): serve the easy python-only detectors on rust --- .../crates/inference-messages/src/handler.rs | 130 ++++++++++-- .../crates/inference-messages/src/lib.rs | 5 + .../crates/inference-messages/src/types.rs | 5 + .../crates/llms/src/anthropic/beta_headers.rs | 12 +- .../src/anthropic/messages/transformation.rs | 114 +++++++++- .../src/azure_ai/messages/transformation.rs | 6 + .../crates/llms/src/base_llm/messages/mod.rs | 1 + .../llms/src/base_llm/messages/tools.rs | 194 ++++++++++++++++++ .../src/base_llm/messages/transformation.rs | 10 +- .../messages/transformation.rs | 6 + .../bedrock_mantle/messages/transformation.rs | 13 +- .../src/deepseek/messages/transformation.rs | 6 + .../src/edenai/messages/transformation.rs | 6 + .../src/minimax/messages/transformation.rs | 6 + .../src/tencent/messages/transformation.rs | 6 + .../src/vertex_ai/messages/transformation.rs | 6 + litellm-rust/crates/python-bridge/src/http.rs | 18 ++ litellm-rust/crates/python-bridge/src/lib.rs | 3 + .../python-bridge/src/routes/messages/host.rs | 70 +++++-- litellm/__init__.py | 4 +- litellm/caching/caching_handler.py | 23 ++- litellm/litellm_core_utils/core_helpers.py | 14 +- litellm/messages/dispatch.py | 12 +- litellm/rust_bridge/_native.pyi | 25 +-- litellm/rust_bridge/catalog.py | 12 +- litellm/rust_bridge/detectors.py | 17 +- litellm/rust_bridge/host/routes/messages.py | 164 ++++++++++++++- litellm/rust_bridge/messages/detectors.py | 51 ++--- .../sdk/messages/test_native_parity.py | 149 +++++++++++--- .../messages/test_dispatch.py | 74 +++++-- .../litellm_core_utils/test_core_helpers.py | 27 +++ .../rust_bridge/host/routes/test_messages.py | 89 ++++++++ tests/unit/rust_bridge/test_catalog.py | 21 +- tests/unit/rust_bridge/test_detectors.py | 133 +++++++++--- 34 files changed, 1199 insertions(+), 233 deletions(-) create mode 100644 litellm-rust/crates/llms/src/base_llm/messages/tools.rs diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index c7a7d93bf98..ebf4118015a 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -44,6 +44,7 @@ struct ProviderCall { signer: Option, timeout: Option, stream: bool, + whole_sse_frames: bool, } impl ProviderCall { @@ -153,19 +154,15 @@ async fn provider_call( &env_lookup, )?; + let tool_custom_field = config.tool_custom_field(&trimmed.model); let body = config.request_body( - config.transform_anthropic_messages_request(trimmed, &transform_context)?, + tool_custom_field.apply( + config.transform_anthropic_messages_request(trimmed, &transform_context)?, + shaping.settings.drop_params, + ), &validated.headers, shaping.settings.drop_params, )?; - let environment = ValidatedEnvironment { - headers: config.request_headers( - with_python_default_headers(validated.headers, config.default_headers()), - &body, - ), - auth: validated.auth, - }; - let url = config.get_complete_url( api_base.as_deref(), &body.model, @@ -173,6 +170,16 @@ async fn provider_call( body.params.stream == Some(true), &env_lookup, )?; + let environment = ValidatedEnvironment { + headers: config.headers_for_url( + config.request_headers( + with_python_default_headers(validated.headers, config.default_headers()), + &body, + ), + &url, + ), + auth: validated.auth, + }; let request_context = RequestContext { model: body.model.clone(), @@ -217,6 +224,7 @@ async fn provider_call( signer: authenticated.signer, timeout: timeout.or_else(|| shaping.settings.request_timeout()), stream, + whole_sse_frames: shaping.settings.whole_sse_frames, }) } @@ -307,6 +315,7 @@ async fn call_provider( signer, timeout, stream, + whole_sse_frames, } = request; let provider_name = provider.as_str(); log_request_body(provider_name, stream, &wire.body); @@ -328,12 +337,13 @@ async fn call_provider( } let config = provider.config(); if stream { - return Ok(streaming_response( - response, - config.stream_decoder(), - provider_name, - timeout, - )); + let streamed = + streaming_response(response, config.stream_decoder(), provider_name, timeout); + return Ok(if whole_sse_frames { + framed(streamed) + } else { + streamed + }); } let head = ResponseHead { status: response.status().as_u16(), @@ -469,6 +479,52 @@ fn streaming_response( } } +fn framed(response: MessagesCallResponse) -> MessagesCallResponse { + let MessagesCallResponse::Stream { head, chunks } = response else { + return response; + }; + MessagesCallResponse::Stream { + head, + chunks: whole_sse_frames(chunks), + } +} + +fn whole_sse_frames( + chunks: BoxStream<'static, Result>, +) -> BoxStream<'static, Result> { + futures_util::stream::try_unfold( + (Some(chunks), Vec::new()), + |(chunks, mut pending)| async move { + let Some(mut chunks) = chunks else { + return Ok(None); + }; + while let Some(chunk) = chunks.try_next().await? { + pending.extend_from_slice(&chunk); + let end = complete_frames_end(&pending); + if end > 0 { + let rest = pending.split_off(end); + return Ok(Some((Bytes::from(pending), (Some(chunks), rest)))); + } + } + Ok((!pending.is_empty()).then(|| (Bytes::from(pending), (None, Vec::new())))) + }, + ) + .boxed() +} + +fn complete_frames_end(pending: &[u8]) -> usize { + [b"\r\n\r\n".as_slice(), b"\n\n", b"\r\r"] + .iter() + .filter_map(|delimiter| { + pending + .windows(delimiter.len()) + .rposition(|window| window == *delimiter) + .map(|at| at + delimiter.len()) + }) + .max() + .unwrap_or(0) +} + fn decoded_chunks( bytes: ByteStream, decode: StreamDecoder, @@ -1092,4 +1148,48 @@ mod tests { } assert!(chunks.next().await.is_none()); } + + #[rstest] + #[case::no_delimiter(b"data: a".as_slice(), 0)] + #[case::lf(b"data: a\n\ndata: b", 9)] + #[case::crlf(b"data: a\r\n\r\ndata: b", 11)] + #[case::cr(b"data: a\r\rdata: b", 9)] + #[case::last_of_several(b"data: a\n\ndata: b\n\ndata: c", 18)] + fn complete_frames_end_is_after_the_last_delimiter(#[case] pending: &[u8], #[case] end: usize) { + assert_eq!(complete_frames_end(pending), end); + } + + #[rstest] + #[case::split_frame( + vec![&b"data: a\n"[..], b"\ndata: ", b"b\n\n"], + vec![&b"data: a\n\n"[..], b"data: b\n\n"] + )] + #[case::several_frames_in_one_chunk( + vec![&b"data: a\n\ndata: b\n\nda"[..], b"ta: c\n\n"], + vec![&b"data: a\n\ndata: b\n\n"[..], b"data: c\n\n"] + )] + #[case::unterminated_tail_is_flushed( + vec![&b"data: a\n\nda"[..], b"ta: b"], + vec![&b"data: a\n\n"[..], b"data: b"] + )] + #[tokio::test] + async fn whole_sse_frames_regroups_chunks_at_frame_boundaries( + #[case] chunks: Vec<&'static [u8]>, + #[case] expected: Vec<&'static [u8]>, + ) { + let upstream = futures_util::stream::iter( + chunks + .into_iter() + .map(|chunk| Ok::<_, Error>(Bytes::from_static(chunk))), + ) + .boxed(); + let framed: Vec = whole_sse_frames(upstream).try_collect().await.unwrap(); + assert_eq!( + framed, + expected + .into_iter() + .map(Bytes::from_static) + .collect::>() + ); + } } diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index 25ebeaacb49..23dd8f4a920 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -75,6 +75,11 @@ impl MessagesRoute { call: MessagesCall, context: CallContext<'_, impl Interceptors>, ) -> Result { + if let Some(mock) = &call.shaping.settings.mock_response { + return Ok(MessagesCallResponse::Complete(MessagesBody::Unchecked( + mock.clone(), + ))); + } litellm_inference::diagnostic::call(async { let provider = resolve_llm_provider( &call.body.model, diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index 772ea6f7292..b1ac375f7ab 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -71,6 +71,10 @@ pub struct MessagesSettings { pub additional_drop_params: Vec, #[serde(default)] pub request_timeout: Option, + #[serde(default)] + pub mock_response: Option>, + #[serde(default)] + pub whole_sse_frames: bool, } impl MessagesSettings { @@ -148,6 +152,7 @@ mod tests { reasoning_auto_summary: true, additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()], request_timeout: Some(1800.0), + ..MessagesSettings::default() }, capabilities: MessagesModelCapabilities { supports_reasoning: true, diff --git a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs index 7842b478dce..4a3575c9bbc 100644 --- a/litellm-rust/crates/llms/src/anthropic/beta_headers.rs +++ b/litellm-rust/crates/llms/src/anthropic/beta_headers.rs @@ -1,5 +1,5 @@ use litellm_http::request::{with_header, without_headers}; -use litellm_llms_types::providers::anthropic::{BETA_HEADER, BetaProvider, BetaSet}; +use litellm_llms_types::providers::anthropic::{AnthropicBeta, BETA_HEADER, BetaProvider, BetaSet}; use crate::{anthropic::common_utils::existing_betas, base_llm::auth::Headers}; @@ -33,6 +33,16 @@ impl BetaPolicy { } } +pub fn accepted_or_unlisted(betas: BetaSet, provider: BetaProvider) -> BetaSet { + betas + .into_iter() + .filter_map(|beta| match beta { + AnthropicBeta::Other(_) => Some(beta), + known => known.on(provider), + }) + .collect() +} + #[cfg(test)] mod tests { use crate::anthropic::beta_headers::BetaPolicy; diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 18789e5b2be..fb7b782e599 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -1,25 +1,27 @@ use litellm_auth::CredentialPlacement; +use litellm_http::request::{with_header, without_headers}; use litellm_llms_types::{ formats::messages::{ ContentBlock, ContentBlockType, ContextEdit, ContextManagement, Message, MessagesOptionalParams, MessagesRequest, Speed, ThinkingConfig, ThinkingDisplay, }, - providers::anthropic::{AnthropicBeta, BetaProvider, BetaSet, DEFAULT_HEADERS}, + providers::anthropic::{AnthropicBeta, BETA_HEADER, BetaProvider, BetaSet, DEFAULT_HEADERS}, recognized::Recognized, }; use litellm_router_types::LitellmParams; use serde_json::{Map, Value, json}; use super::{handler::shape_anthropic_messages_request, thinking::translate_thinking}; -use crate::base_llm::messages::context::MessagesTransformContext; +use crate::anthropic::beta_headers::accepted_or_unlisted; +use crate::base_llm::messages::{context::MessagesTransformContext, tools::ToolCustomField}; use crate::{ Error, anthropic::common_utils::{ ANTHROPIC_API_BASE_ENV, ANTHROPIC_API_KEY_ENV, ANTHROPIC_AUTH_TOKEN_ENV, - ANTHROPIC_BASE_URL_ENV, OauthHandling, complete_anthropic_url, get_auth_header, - has_advisor_tool, has_anthropic_credential, is_tool_search_used, merge_beta_headers, - optionally_handle_anthropic_oauth, requires_native_compaction_beta, strip_advisor_blocks, - strip_encrypted_reasoning_blocks, + ANTHROPIC_BASE_URL_ENV, OauthHandling, complete_anthropic_url, existing_betas, + get_auth_header, has_advisor_tool, has_anthropic_credential, is_tool_search_used, + merge_beta_headers, optionally_handle_anthropic_oauth, requires_native_compaction_beta, + strip_advisor_blocks, strip_encrypted_reasoning_blocks, }, base_llm::{ auth::AuthScheme, @@ -59,6 +61,12 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { transform_messages_request(request, &context.keeping_display_updates()) } + fn tool_custom_field(&self, model: &str) -> ToolCustomField { + ToolCustomField::Rewrite { + hoist_defer_loading: is_claude_model(model), + } + } + fn secret_names(&self) -> &'static [&'static str] { &[ ANTHROPIC_API_KEY_ENV, @@ -116,6 +124,47 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { provider_feature_betas(request, BetaProvider::Anthropic), ) } + + fn headers_for_url(&self, headers: Headers, url: &str) -> Headers { + if !is_foundry_url(url) { + return headers; + } + let betas = existing_betas(&headers); + if betas.is_empty() { + return headers; + } + let kept = accepted_or_unlisted(betas, BetaProvider::AzureAi); + if kept.is_empty() { + return without_headers(headers, &[BETA_HEADER]); + } + with_header(headers, BETA_HEADER, kept.to_string()) + } +} + +const FOUNDRY_HOST_SUFFIXES: [&str; 3] = [ + "services.ai.azure.com", + "openai.azure.com", + "cognitiveservices.azure.com", +]; + +fn is_foundry_url(url: &str) -> bool { + let absolute = if url.contains("://") { + url.to_string() + } else { + format!("https://{url}") + }; + url::Url::parse(&absolute) + .ok() + .and_then(|parsed| parsed.host_str().map(str::to_ascii_lowercase)) + .is_some_and(|host| { + FOUNDRY_HOST_SUFFIXES + .iter() + .any(|suffix| host == *suffix || host.ends_with(&format!(".{suffix}"))) + }) +} + +pub(crate) fn is_claude_model(model: &str) -> bool { + model.to_lowercase().contains("claude") } pub(crate) fn transform_messages_request( @@ -437,6 +486,7 @@ mod tests { use super::*; use crate::anthropic::common_utils::ENCRYPTED_REASONING_SIGNATURE_PREFIX; + use litellm_http::request::header_values; type Env = &'static [(&'static str, &'static str)]; @@ -1352,4 +1402,56 @@ mod tests { .collect(); assert_eq!(undeclared, Vec::<&String>::new()); } + + const FOUNDRY: &str = "https://team.services.ai.azure.com/anthropic/v1/messages"; + + #[rstest] + #[case::drops_a_beta_foundry_lacks( + FOUNDRY, + "advisor-tool-2026-03-01,context-1m-2025-08-07", + Some("context-1m-2025-08-07") + )] + #[case::forwards_unknown_betas_sorted( + FOUNDRY, + "zzz,files-api-2025-04-14,aaa", + Some("aaa,files-api-2025-04-14,zzz") + )] + #[case::removes_the_header_when_nothing_is_left(FOUNDRY, "advisor-tool-2026-03-01", None)] + #[case::cognitive_services_host( + "team.cognitiveservices.azure.com/anthropic", + "advisor-tool-2026-03-01", + None + )] + #[case::anthropic_host_is_untouched( + "https://api.anthropic.com/v1/messages", + "advisor-tool-2026-03-01", + Some("advisor-tool-2026-03-01") + )] + #[case::lookalike_host_is_untouched( + "https://notservices.ai.azure.com.example/v1/messages", + "advisor-tool-2026-03-01", + Some("advisor-tool-2026-03-01") + )] + fn foundry_urls_filter_betas_through_the_azure_ai_row( + #[case] url: &str, + #[case] betas: &str, + #[case] expected: Option<&str>, + ) { + let filtered = ANTHROPIC_MESSAGES_CONFIG + .headers_for_url(headers(&[("x-a", "1"), (BETA_HEADER, betas)]), url); + assert_eq!( + header_values(&filtered, BETA_HEADER).collect::>(), + expected.into_iter().collect::>() + ); + assert_eq!(header_values(&filtered, "x-a").collect::>(), ["1"]); + } + + #[rstest] + fn foundry_url_without_betas_keeps_the_headers() { + let sent = headers(&[("x-a", "1")]); + assert_eq!( + ANTHROPIC_MESSAGES_CONFIG.headers_for_url(sent.clone(), FOUNDRY), + sent + ); + } } diff --git a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs index d13cfa81a94..43fd3c897cd 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -36,6 +36,12 @@ pub const AZURE_ANTHROPIC_MESSAGES_CONFIG: AzureAnthropicMessagesConfig = AzureAnthropicMessagesConfig; impl BaseMessagesConfig for AzureAnthropicMessagesConfig { + fn tool_custom_field(&self, model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: crate::anthropic::messages::transformation::is_claude_model(model), + } + } + fn shape_request( &self, request: MessagesRequest, diff --git a/litellm-rust/crates/llms/src/base_llm/messages/mod.rs b/litellm-rust/crates/llms/src/base_llm/messages/mod.rs index e9591daa9da..5c1563c7a23 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/mod.rs @@ -2,4 +2,5 @@ pub mod context; pub mod mid_conversation_system; pub mod normalization; pub mod streaming; +pub mod tools; pub mod transformation; diff --git a/litellm-rust/crates/llms/src/base_llm/messages/tools.rs b/litellm-rust/crates/llms/src/base_llm/messages/tools.rs new file mode 100644 index 00000000000..49113872c56 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/messages/tools.rs @@ -0,0 +1,194 @@ +use litellm_llms_types::{ + formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool}, + recognized::Recognized, +}; +use serde_json::{Map, Value}; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ToolCustomField { + Keep, + Rewrite { hoist_defer_loading: bool }, +} + +impl ToolCustomField { + pub fn apply(self, request: MessagesRequest, drop_params: bool) -> MessagesRequest { + match self { + Self::Keep => request, + Self::Rewrite { + hoist_defer_loading, + } => tools_without_custom_field(request, hoist_defer_loading, drop_params), + } + } +} + +fn tools_without_custom_field( + request: MessagesRequest, + hoist: bool, + drop_custom: bool, +) -> MessagesRequest { + let tools = request.params.tools.map(|tools| { + tools + .into_iter() + .map(|tool| tool_without_custom_field(tool, hoist, drop_custom)) + .collect() + }); + MessagesRequest { + params: MessagesOptionalParams { + tools, + ..request.params + }, + ..request + } +} + +fn tool_without_custom_field( + tool: Recognized, + hoist: bool, + drop_custom: bool, +) -> Recognized { + let Ok(Value::Object(fields)) = serde_json::to_value(&tool) else { + return tool; + }; + let Some(custom) = fields.get("custom").cloned() else { + return tool; + }; + let defer_loading = hoist + .then(|| hoistable_defer_loading(&fields, &custom)) + .flatten(); + let without_custom = fields + .iter() + .filter(|(key, _)| *key != "custom") + .map(|(key, value)| (key.clone(), value.clone())); + let rebuilt: Map = match defer_loading { + None if drop_custom => without_custom.collect(), + None => return tool, + Some(defer_loading) => { + let unhoisted: Map = custom + .as_object() + .into_iter() + .flatten() + .filter(|(key, _)| *key != "defer_loading") + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + let kept_custom = (!unhoisted.is_empty() && !drop_custom) + .then(|| ("custom".to_string(), Value::Object(unhoisted))); + without_custom + .chain([("defer_loading".to_string(), Value::Bool(defer_loading))]) + .chain(kept_custom) + .collect() + } + }; + serde_json::from_value(Value::Object(rebuilt)).unwrap_or(tool) +} + +fn hoistable_defer_loading(tool: &Map, custom: &Value) -> Option { + if tool.contains_key("defer_loading") { + return None; + } + custom.as_object()?.get("defer_loading")?.as_bool() +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + const WEATHER: &str = "get_weather"; + + fn tool(fields: Value) -> Value { + merged( + json!({"name": WEATHER, "input_schema": {"type": "object"}}), + fields, + ) + } + + fn tools_after(policy: ToolCustomField, tools: Value, drop_custom: bool) -> Value { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-sonnet-4-5", + "max_tokens": 16, + "messages": [{"role": "user", "content": "hi"}], + "tools": [tools] + })) + .unwrap(); + serde_json::to_value(policy.apply(request, drop_custom)).unwrap()["tools"][0].clone() + } + + fn merged(base: Value, fields: Value) -> Value { + let mut merged = base.as_object().unwrap().clone(); + merged.extend(fields.as_object().unwrap().clone()); + Value::Object(merged) + } + + const HOIST: ToolCustomField = ToolCustomField::Rewrite { + hoist_defer_loading: true, + }; + const NO_HOIST: ToolCustomField = ToolCustomField::Rewrite { + hoist_defer_loading: false, + }; + + #[rstest] + #[case::hoists_defer_loading( + HOIST, + json!({"custom": {"defer_loading": true}}), + false, + json!({"defer_loading": true}) + )] + #[case::keeps_other_custom_keys( + HOIST, + json!({"custom": {"defer_loading": false, "note": "kept"}}), + false, + json!({"defer_loading": false, "custom": {"note": "kept"}}) + )] + #[case::drops_other_custom_keys_with_drop_params( + HOIST, + json!({"custom": {"defer_loading": true, "note": "dropped"}}), + true, + json!({"defer_loading": true}) + )] + #[case::keeps_an_opaque_custom( + HOIST, + json!({"custom": "opaque"}), + false, + json!({"custom": "opaque"}) + )] + #[case::drops_an_opaque_custom_with_drop_params( + HOIST, + json!({"custom": "opaque"}), + true, + json!({}) + )] + #[case::own_defer_loading_wins( + HOIST, + json!({"defer_loading": false, "custom": {"defer_loading": true}}), + false, + json!({"defer_loading": false, "custom": {"defer_loading": true}}) + )] + #[case::not_hoisted_off_claude_hosts( + NO_HOIST, + json!({"custom": {"defer_loading": true}}), + false, + json!({"custom": {"defer_loading": true}}) + )] + fn tool_custom_field_follows_the_python_handler( + #[case] policy: ToolCustomField, + #[case] fields: Value, + #[case] drop_custom: bool, + #[case] expected: Value, + ) { + assert_eq!( + tools_after(policy, tool(fields), drop_custom), + tool(expected) + ); + } + + #[rstest] + fn keep_leaves_the_custom_field() { + let custom = tool(json!({"custom": {"defer_loading": true}})); + assert_eq!( + tools_after(ToolCustomField::Keep, custom.clone(), true), + custom + ); + } +} diff --git a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs index 744ce89afdb..8d0b377ca85 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs @@ -4,7 +4,7 @@ use litellm_llms_types::formats::messages::{ use litellm_router_types::LitellmParams; use serde_json::Value; -use super::context::MessagesTransformContext; +use super::{context::MessagesTransformContext, tools::ToolCustomField}; pub use crate::base_llm::auth::{Headers, ValidatedEnvironment}; use crate::{Error, base_llm::messages::streaming::StreamDecoder}; @@ -35,6 +35,10 @@ pub trait BaseMessagesConfig: Sync { Ok(request) } + fn tool_custom_field(&self, _model: &str) -> ToolCustomField { + ToolCustomField::Keep + } + fn transform_anthropic_messages_response( &self, _model: &str, @@ -77,6 +81,10 @@ pub trait BaseMessagesConfig: Sync { headers } + fn headers_for_url(&self, headers: Headers, _url: &str) -> Headers { + headers + } + fn request_body( &self, request: MessagesRequest, diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs index 5b89345a1c5..0567178d5ed 100644 --- a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -50,6 +50,12 @@ pub const BEDROCK_CLAUDE_PLATFORM_MESSAGES_CONFIG: BedrockClaudePlatformMessages BedrockClaudePlatformMessagesConfig; impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig { + fn tool_custom_field(&self, model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: crate::anthropic::messages::transformation::is_claude_model(model), + } + } + fn shape_request( &self, request: MessagesRequest, diff --git a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs index f139b93e259..2e0121c9571 100644 --- a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs @@ -17,6 +17,7 @@ use serde_json::Value; use crate::{ Error, anthropic::{ + beta_headers::accepted_or_unlisted, common_utils::{existing_betas, merge_beta_headers}, messages::{ handler::shape_anthropic_messages_request, @@ -191,14 +192,10 @@ impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig { .filter_map(Value::as_str) .filter_map(|beta| beta.parse::().ok()) .collect(); - let caller: BetaSet = existing_betas(&headers) - .union(body_betas) - .into_iter() - .filter_map(|beta| match beta { - AnthropicBeta::Other(_) => Some(beta), - _ => beta.on(BetaProvider::BedrockMantle), - }) - .collect(); + let caller = accepted_or_unlisted( + existing_betas(&headers).union(body_betas), + BetaProvider::BedrockMantle, + ); let safeguards: BetaSet = request .params .extra diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs index 0dbcea01439..dce9178e7b9 100644 --- a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -42,6 +42,12 @@ pub const DEEPSEEK_ANTHROPIC_MESSAGES_CONFIG: DeepSeekAnthropicMessagesConfig = DeepSeekAnthropicMessagesConfig; impl BaseMessagesConfig for DeepSeekAnthropicMessagesConfig { + fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: false, + } + } + fn shape_request( &self, request: MessagesRequest, diff --git a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs index d5944484522..75a8b16d1eb 100644 --- a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs @@ -33,6 +33,12 @@ pub struct EdenAIAnthropicMessagesConfig; pub const EDENAI_MESSAGES_CONFIG: EdenAIAnthropicMessagesConfig = EdenAIAnthropicMessagesConfig; impl BaseMessagesConfig for EdenAIAnthropicMessagesConfig { + fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: false, + } + } + fn validate_environment( &self, headers: Headers, diff --git a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs index dafe5fbd16b..aa467619131 100644 --- a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs @@ -31,6 +31,12 @@ pub struct MinimaxMessagesConfig; pub const MINIMAX_MESSAGES_CONFIG: MinimaxMessagesConfig = MinimaxMessagesConfig; impl BaseMessagesConfig for MinimaxMessagesConfig { + fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: false, + } + } + fn shape_request( &self, request: MessagesRequest, diff --git a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs index 621a09f4fb0..97a359826a7 100644 --- a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs @@ -32,6 +32,12 @@ pub struct TencentAnthropicMessagesConfig; pub const TENCENT_MESSAGES_CONFIG: TencentAnthropicMessagesConfig = TencentAnthropicMessagesConfig; impl BaseMessagesConfig for TencentAnthropicMessagesConfig { + fn tool_custom_field(&self, _model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: false, + } + } + fn shape_request( &self, request: MessagesRequest, diff --git a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs index 038617c2e10..0f3ed086eaf 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -47,6 +47,12 @@ pub const VERTEX_ANTHROPIC_MESSAGES_CONFIG: VertexAiPartnerModelsAnthropicMessag VertexAiPartnerModelsAnthropicMessagesConfig; impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { + fn tool_custom_field(&self, model: &str) -> crate::base_llm::messages::tools::ToolCustomField { + crate::base_llm::messages::tools::ToolCustomField::Rewrite { + hoist_defer_loading: crate::anthropic::messages::transformation::is_claude_model(model), + } + } + fn shape_request( &self, request: MessagesRequest, diff --git a/litellm-rust/crates/python-bridge/src/http.rs b/litellm-rust/crates/python-bridge/src/http.rs index c6a8b98747c..510583ee6e0 100644 --- a/litellm-rust/crates/python-bridge/src/http.rs +++ b/litellm-rust/crates/python-bridge/src/http.rs @@ -129,6 +129,24 @@ pub(crate) fn settings_config( Ok(resolved(py, settings)) } +#[pyfunction] +#[pyo3(signature = (ssl_security_level=None, ssl_ecdh_curve=None))] +pub(crate) fn unsupported_tls_settings( + ssl_security_level: Option, + ssl_ecdh_curve: Option, +) -> Vec { + let settings = HttpSettings::from_layers([HttpSettingsLayer { + ssl_security_level, + ssl_ecdh_curve, + ..HttpSettingsLayer::default() + }]); + Resolution::from(&settings) + .unsupported + .iter() + .map(ToString::to_string) + .collect() +} + /// The client config `settings` resolve to, warning once per process about each setting the client cannot apply fn resolved(py: Python<'_>, settings: HttpSettings) -> HttpClientConfig { let resolution = diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 7edf13035c1..45ea95560e1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -35,6 +35,8 @@ mod _native { #[pymodule_export] use crate::errors::{RustBridgeDeclined, RustUpstreamError}; #[pymodule_export] + use crate::http::unsupported_tls_settings; + #[pymodule_export] use crate::logger::NativeDiagnosticProcessor; #[pymodule_export] use crate::model_cost::{ @@ -191,6 +193,7 @@ mod tests { "gil_stats", "process_state_started", "reserve_process_for_forking", + "unsupported_tls_settings", ]; #[cfg(feature = "huggingface")] expected.push("HuggingFaceEncoding"); diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs index 2a5f79ca393..57e9cfccee7 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -67,11 +67,21 @@ fn project_litellm_params<'py>( Ok(litellm_params(fields.into_iter().chain(folded).collect())) } +#[derive(serde::Deserialize)] +struct HostSettings { + #[serde(flatten)] + settings: MessagesSettings, + model: String, + stream_cost_model: Option, + extra_headers: Option>, +} + /// The Python side of the Messages route: projects the prepared arguments and builds the /// public response, chunks and exceptions. pub(super) struct MessagesPythonHost { view: RequestView, cache: PythonCache, + stream_cost_model: Option, } impl MessagesPythonHost { @@ -79,19 +89,33 @@ impl MessagesPythonHost { Self { view, cache: PythonCache::new(asynchronous), + stream_cost_model: None, } } + fn host_settings( + &self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> PyResult { + from_py( + &py.import(ROUTE_HOST_MODULE)? + .getattr("settings")? + .call1((arguments, self.view.request(py)?))?, + ) + } + fn projection( &self, py: Python<'_>, arguments: &Bound<'_, PyDict>, + settings: HostSettings, ) -> PyResult> { let argument = |name: &str| self.view.argument(py, arguments, name); let string = |name: &str| -> PyResult> { argument(name)?.map(|value| value.extract()).transpose() }; - let model = string("model")?.ok_or(ArgumentError::Required("model"))?; + let model = settings.model; let messages = argument("messages")?.ok_or(ArgumentError::Required("messages"))?; let fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?; let body = [ @@ -117,8 +141,17 @@ impl MessagesPythonHost { .call1((&model, api_base.as_deref()))? .extract()?, }; - let shaping = self.shaping(py, &model, custom_llm_provider.as_deref(), arguments)?; - let extra_headers = self.merged_headers(py, arguments)?; + let shaping = MessagesShaping { + capabilities: self.capabilities(py, &model, custom_llm_provider.as_deref())?, + settings: MessagesSettings { + whole_sse_frames: settings.stream_cost_model.is_some(), + ..settings.settings + }, + }; + let extra_headers = match settings.extra_headers { + Some(headers) => Some(headers), + None => self.merged_headers(py, arguments)?, + }; let provider_specific_header = self.provider_specific_header(py, arguments)?; let litellm_params = project_litellm_params(argument, |name| module_global(py, name))?; Ok(messages_body(body).and_then(|body| { @@ -164,25 +197,17 @@ impl MessagesPythonHost { .transpose() } - fn shaping( + fn capabilities( &self, py: Python<'_>, model: &str, custom_llm_provider: Option<&str>, - arguments: &Bound<'_, PyDict>, - ) -> PyResult { - let module = py.import(ROUTE_HOST_MODULE)?; - let capabilities: MessagesModelCapabilities = from_py( + ) -> PyResult { + from_py( &py.import("litellm.rust_bridge.host.model_capabilities")? .getattr("anthropic_model_capabilities")? .call1((model, custom_llm_provider))?, - )?; - let settings: MessagesSettings = - from_py(&module.getattr("settings")?.call1((arguments,))?)?; - Ok(MessagesShaping { - capabilities, - settings, - }) + ) } fn provider(&self, py: Python<'_>) -> String { @@ -230,7 +255,11 @@ impl PythonBinding for MessagesPythonHost { let selection = crate::cache::configure(&mut self.cache, py, arguments, "anthropic_messages") .map_err(InvokeError::Python)?; - self.projection(py, arguments) + let settings = self + .host_settings(py, arguments) + .map_err(InvokeError::Python)?; + self.stream_cost_model = settings.stream_cost_model.clone(); + self.projection(py, arguments, settings) .map_err(|error| InvokeError::Python(self.map_failure(py, error)))? .map_err(InvokeError::Native) .map(|request| (request, selection)) @@ -260,7 +289,14 @@ impl PythonBinding for MessagesPythonHost { } fn encode_chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult> { - Ok(PyBytes::new(py, &chunk).into_any().unbind()) + let frames = PyBytes::new(py, &chunk).into_any(); + let Some(model) = &self.stream_cost_model else { + return Ok(frames.unbind()); + }; + py.import(ROUTE_HOST_MODULE)? + .getattr("inject_stream_cost")? + .call1((frames, model, self.view.request(py)?)) + .map(Bound::unbind) } fn map_error(&self, py: Python<'_>, error: Error) -> PyResult { diff --git a/litellm/__init__.py b/litellm/__init__.py index 8a73f3447b8..049e6a44dbd 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -27,7 +27,7 @@ from typing import ( from collections.abc import Mapping from litellm.types.integrations.datadog import DatadogInitParams from litellm.types.integrations.newrelic import NewRelicInitParams -from litellm.litellm_core_utils.core_helpers import drop_params_env_flag +from litellm.litellm_core_utils.core_helpers import drop_params_env_flag, env_flag from litellm.types.integrations.pointfive import PointFiveInitParams from litellm.types.integrations.zerobus import ZerobusInitParams from litellm._logging import ( @@ -227,7 +227,7 @@ token: Optional[str] = ( ) max_tokens: int = DEFAULT_MAX_TOKENS # OpenAI Defaults drop_params = drop_params_env_flag(os.environ, verbose_logger) -modify_params = bool(os.getenv("LITELLM_MODIFY_PARAMS", False)) +modify_params = env_flag(os.environ, "LITELLM_MODIFY_PARAMS", verbose_logger) bedrock_neutralize_orphaned_tool_blocks: bool = True use_chat_completions_url_for_anthropic_messages: bool = bool( os.getenv("LITELLM_USE_CHAT_COMPLETIONS_URL_FOR_ANTHROPIC_MESSAGES", False) diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index ab5d2332c21..9fb3f542164 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -114,7 +114,14 @@ def _drop_logging_obj_from_kwargs(request_kwargs: dict[str, object]) -> dict[str return {k: v for k, v in request_kwargs.items() if k != "litellm_logging_obj"} -def _is_response_cache_excluded(model: str | None, kwargs: Mapping[str, object]) -> bool: +def is_call_type_supported_by_cache(call_type: str) -> bool: + if litellm.cache is None or litellm.cache.supported_call_types is None: + return False + covering_call_types: Final = ("aresponses", "responses") if call_type == "aresponses" else (call_type,) + return any(name in litellm.cache.supported_call_types for name in covering_call_types) + + +def is_response_cache_excluded(model: str | None, kwargs: Mapping[str, object]) -> bool: from litellm.llms.github_copilot.per_user_auth import is_github_copilot_per_user_request custom_llm_provider: Final = kwargs.get("custom_llm_provider") @@ -395,7 +402,7 @@ class LLMCachingHandler: Raises: None """ - if _is_response_cache_excluded(model=model, kwargs=kwargs): + if is_response_cache_excluded(model=model, kwargs=kwargs): self.request_kwargs = _drop_logging_obj_from_kwargs(kwargs) return None @@ -524,7 +531,7 @@ class LLMCachingHandler: cached_result: Any | None = None # Check if caching should be performed BEFORE doing expensive kwargs copy - if _is_response_cache_excluded(model=model, kwargs=kwargs): + if is_response_cache_excluded(model=model, kwargs=kwargs): self.request_kwargs = _drop_logging_obj_from_kwargs(kwargs) return CachingHandlerResponse(cached_result=None) if litellm.cache is not None and self._is_call_type_supported_by_cache(original_function=original_function): @@ -965,7 +972,7 @@ class LLMCachingHandler: model_value: Final = new_kwargs.get("model") model: Final = model_value if isinstance(model_value, str) else None self.request_kwargs = _drop_logging_obj_from_kwargs(new_kwargs) - if _is_response_cache_excluded(model=model, kwargs=new_kwargs): + if is_response_cache_excluded(model=model, kwargs=new_kwargs): return None if new_kwargs.get("stream") is True and "cache_key" not in new_kwargs: new_kwargs["cache_key"] = litellm.cache.get_cache_key(**new_kwargs) @@ -1323,7 +1330,7 @@ class LLMCachingHandler: ) no_store: Final = cache_options.get("no-store", False) return ( - not _is_response_cache_excluded(model=model, kwargs=kwargs) + not is_response_cache_excluded(model=model, kwargs=kwargs) and self._is_call_type_supported_by_cache(original_function=original_function) and no_store is not True ) @@ -1364,11 +1371,7 @@ class LLMCachingHandler: Returns: bool: True if the call type is supported by the cache, False otherwise. """ - if litellm.cache is None or litellm.cache.supported_call_types is None: - return False - call_type: Final = str(original_function.__name__) - covering_call_types: Final = ("aresponses", "responses") if call_type == "aresponses" else (call_type,) - return any(name in litellm.cache.supported_call_types for name in covering_call_types) + return is_call_type_supported_by_cache(str(original_function.__name__)) async def add_streaming_response_to_cache(self, processed_chunk: ModelResponse) -> None: """ diff --git a/litellm/litellm_core_utils/core_helpers.py b/litellm/litellm_core_utils/core_helpers.py index 8dec2941e02..8f6b57537fe 100644 --- a/litellm/litellm_core_utils/core_helpers.py +++ b/litellm/litellm_core_utils/core_helpers.py @@ -242,19 +242,21 @@ def drop_params_flag(value: object, source: str, logger: logging.Logger) -> bool DROP_PARAMS_ENV_VAR: Final = "LITELLM_DROP_PARAMS" -def drop_params_env_flag(environ: Mapping[str, str], logger: logging.Logger) -> bool: - configured: Final = environ.get(DROP_PARAMS_ENV_VAR, "").strip() +def env_flag(environ: Mapping[str, str], name: str, logger: logging.Logger) -> bool: + configured: Final = environ.get(name, "").strip() if configured == "": return False - normalized: Final = normalize_drop_params(configured) + normalized: Final = lax_bool(configured) if normalized is None: - logger.warning( - "%s=%r is not a flag value, treating it as on. Set it to true or false", DROP_PARAMS_ENV_VAR, configured - ) + logger.warning("%s=%r is not a flag value, treating it as on. Set it to true or false", name, configured) return True return normalized +def drop_params_env_flag(environ: Mapping[str, str], logger: logging.Logger) -> bool: + return env_flag(environ, DROP_PARAMS_ENV_VAR, logger) + + def safe_divide( numerator: float, denominator: float, diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index faa604a25ce..642ae57c03d 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -133,26 +133,26 @@ def messages_target(model: str, custom_llm_provider: str | None, model_info: obj return _route_target(model, resolved_model, resolved_provider, model_info, api_base, compacting=False) +_NATIVE_COMMON_FEATURES: Final = frozenset({Feature.TIMEOUT_VARIANTS, Feature.LATEST_ALIAS}) + _FEATURES: Final[detectors.Detectors] = MappingProxyType( { - **{feature: detect for feature, detect in detectors.COMMON.items() if feature is not Feature.TIMEOUT_VARIANTS}, - Feature.PROVIDER_AFFINITY: route_detectors.provider_affinity, + **{feature: detect for feature, detect in detectors.COMMON.items() if feature not in _NATIVE_COMMON_FEATURES}, Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, - Feature.RESPONSE_CACHE: detectors.response_cache, + Feature.RESPONSE_CACHE: route_detectors.sync_response_cache, Feature.MODIFY_PARAMS: detectors.modify_params, Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity, Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection, - Feature.STREAM_COST_INJECTION: route_detectors.stream_cost_injection, Feature.THINKING_SIGNATURE_RETRY: route_detectors.thinking_signature_retry, - Feature.TOOL_CUSTOM_FIELD: route_detectors.tool_custom_field, } ) _ASYNC_FEATURES: Final[detectors.Detectors] = MappingProxyType( { **_FEATURES, + Feature.MOCK_RESPONSE: route_detectors.streamed_mock_response, + Feature.RESPONSE_CACHE: route_detectors.response_cache, Feature.PRE_REQUEST_HOOK: route_detectors.pre_request_hook, - Feature.FOUNDRY_BETA_ROW: route_detectors.foundry_beta_row, Feature.WEBSEARCH_INTERCEPTION: route_detectors.websearch_interception, Feature.MESSAGES_INTERCEPTOR: route_detectors.intercepted, Feature.STRICT_TOOLS_RETRY: route_detectors.strict_tools_retry, diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 221968e07ed..0f0060eb37d 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -22,6 +22,7 @@ class NativeRedisStoreError(Exception): ... def trace_encode_error(message: str) -> bytes: ... def bundled_model_cost_map() -> bytes: ... def native_redis_breaker_states() -> list[str]: ... +def unsupported_tls_settings(ssl_security_level: str | None = None, ssl_ecdh_curve: str | None = None) -> list[str]: ... @final class NativeCostMapSource: @@ -230,9 +231,7 @@ class NativeRouterUsage: def attach(self, client: NativeRedisClient) -> bool: ... @property def has_redis(self) -> bool: ... - def increment( - self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> Future[float]: ... + def increment(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> Future[float]: ... def increment_deployment( self, deployment_id: str, @@ -241,12 +240,8 @@ class NativeRouterUsage: tokens: float, requests: int, ) -> Future[None]: ... - def get_many( - self, counters: Sequence[_UsageCounter] - ) -> Future[list[float | None]]: ... - def increment_sync( - self, counter: _UsageCounter, value: float, ttl_seconds: int - ) -> float: ... + def get_many(self, counters: Sequence[_UsageCounter]) -> Future[list[float | None]]: ... + def increment_sync(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> float: ... def increment_deployment_sync( self, deployment_id: str, @@ -255,9 +250,7 @@ class NativeRouterUsage: tokens: float, requests: int, ) -> None: ... - def get_many_sync( - self, counters: Sequence[_UsageCounter] - ) -> list[float | None]: ... + def get_many_sync(self, counters: Sequence[_UsageCounter]) -> list[float | None]: ... def flush_local_counters(self) -> None: ... def snapshot(self) -> dict[str, float]: ... @@ -433,7 +426,9 @@ class NativeAuthCache: def get_sync(self, key: str) -> object: ... def set(self, key: str, value: object, ttl_seconds: float | None = None) -> Future[None]: ... def set_sync(self, key: str, value: object, ttl_seconds: float | None = None) -> None: ... - def declare_set(self, key: str, value: object, ttl_seconds: float | None, send: NativeSend) -> NativePending[None]: ... + def declare_set( + self, key: str, value: object, ttl_seconds: float | None, send: NativeSend + ) -> NativePending[None]: ... def declare_delete(self, key: str, send: NativeSend) -> NativePending[None]: ... def set_many(self, entries: Sequence[tuple[str, object]], ttl_seconds: float | None = None) -> Future[None]: ... def write_back(self, entries: Sequence[tuple[str, object, float | None]], send: NativeSend) -> Future[None]: ... @@ -1189,7 +1184,9 @@ class RouterCoreHandle: callbacks: Mapping[str, bool] | None = None, ) -> RouterCall[_SuccessReportSteps]: ... # Count a success against any id, driven like `report_failure`; `done` carries `key`. - def count_success(self, id: str, callbacks: Mapping[str, bool] | None = None) -> RouterCall[_SuccessReportSteps]: ... + def count_success( + self, id: str, callbacks: Mapping[str, bool] | None = None + ) -> RouterCall[_SuccessReportSteps]: ... # [(id, tpm_key, rpm_key)] for this minute, or None for a group with no deployment. def usage_keys(self, model_group: str) -> list[tuple[str, str, str]] | None: ... # Every deployment as a `model_list` entry, in config order. diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index c493617003d..210df7b0c10 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -66,8 +66,6 @@ class Feature(str, Enum): """Extra body fields that Python transforms and merges into the provider request.""" PROVIDER_AFFINITY = "provider_affinity" """A provider affinity header that Python validates and populates from the session id.""" - FOUNDRY_BETA_ROW = "foundry_beta_row" - """A Foundry api_base, where Python drops the anthropic-beta values Foundry rejects.""" PRE_REQUEST_HOOK = "pre_request_hook" """A logger overriding async_pre_request_hook, which may rewrite the request.""" WEBSEARCH_INTERCEPTION = "websearch_interception" @@ -94,12 +92,8 @@ class Feature(str, Enum): """No static key, so Python mints a workload identity federation token.""" CACHE_CONTROL_INJECTION = "cache_control_injection" """Prompt caching injection that Python applies to messages and system.""" - STREAM_COST_INJECTION = "stream_cost_injection" - """litellm.include_cost_in_streaming_usage or the x-litellm-include-cost-in-streaming-usage header, where Python adds cost to streamed usage.""" THINKING_SIGNATURE_RETRY = "thinking_signature_retry" """Replayed thinking blocks, which Python strips and retries once if Anthropic rejects their signature.""" - TOOL_CUSTOM_FIELD = "tool_custom_field" - """Tools carrying a non-spec custom object, which Python drops or hoists for the provider.""" STRICT_TOOLS_RETRY = "strict_tools_retry" """Strict tools under drop_params, which Python sends again without strict if Claude rejects their limits.""" CLIENT_HEADERS = "client_headers" @@ -172,9 +166,7 @@ CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { Feature.MODEL_ALIAS_MAP, } -MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VARIANTS}) | { - Feature.PROVIDER_AFFINITY, - Feature.FOUNDRY_BETA_ROW, +MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VARIANTS, Feature.LATEST_ALIAS}) | { Feature.SAFEGUARDS_CLASSIFIER, Feature.PRE_REQUEST_HOOK, Feature.WEBSEARCH_INTERCEPTION, @@ -183,9 +175,7 @@ MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VAR Feature.MODIFY_PARAMS, Feature.WORKLOAD_IDENTITY, Feature.CACHE_CONTROL_INJECTION, - Feature.STREAM_COST_INJECTION, Feature.THINKING_SIGNATURE_RETRY, - Feature.TOOL_CUSTOM_FIELD, Feature.STRICT_TOOLS_RETRY, } diff --git a/litellm/rust_bridge/detectors.py b/litellm/rust_bridge/detectors.py index c38c3ea7135..3877e5081ba 100644 --- a/litellm/rust_bridge/detectors.py +++ b/litellm/rust_bridge/detectors.py @@ -2,6 +2,7 @@ from __future__ import annotations +import functools import os from collections.abc import Callable, Iterable, Mapping from types import MappingProxyType @@ -168,15 +169,23 @@ def response_cache(request: NativeCall) -> bool: return litellm.cache is not None and cache.get("no-cache") is not True -_NATIVE_ECDH_CURVES: Final = frozenset({"x25519", "prime256v1", "secp256r1", "p-256", "secp384r1", "p-384"}) +@functools.lru_cache(maxsize=32) +def _tls_unsupported(security_level: str | None, ecdh_curve: str | None) -> bool: + from litellm.rust_bridge import _native + + return bool(_native.unsupported_tls_settings(security_level, ecdh_curve)) def tls_policy(_request: NativeCall) -> bool: """Python applies TLS settings or ECDH curves the native client cannot configure.""" - curve: Final = os.getenv("SSL_ECDH_CURVE", litellm.ssl_ecdh_curve) - return bool(os.getenv("SSL_SECURITY_LEVEL", litellm.ssl_security_level)) or ( - bool(curve) and str(curve).strip().lower() not in _NATIVE_ECDH_CURVES + configured: Final[tuple[object, object]] = ( + os.environ.get("SSL_SECURITY_LEVEL", litellm.ssl_security_level), + os.environ.get("SSL_ECDH_CURVE", litellm.ssl_ecdh_curve), ) + if any(value and not isinstance(value, str) for value in configured): # pyright: ignore[reportUnnecessaryIsInstance] # callers assign the ssl settings untyped + return True + security_level, ecdh_curve = (value if isinstance(value, str) else None for value in configured) + return _tls_unsupported(security_level, ecdh_curve) COMMON: Final[Detectors] = MappingProxyType( diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index c6568264edb..0c9ff7f5b00 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -2,11 +2,12 @@ from __future__ import annotations import datetime import errno +import inspect import ipaddress import os import socket from collections.abc import Callable, Mapping, Sequence -from dataclasses import asdict, dataclass +from dataclasses import asdict, dataclass, replace from typing import Final, TypeAlias, cast # noqa: TID251 # narrows native payloads to public TypedDicts import httpx @@ -15,25 +16,43 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils import exception_mapping_utils -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, + normalize_drop_params, + proxy_request_headers, + stream_cost_injection_requested, +) +from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header +from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.model_latest_aliases import is_model_latest_alias +from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header from litellm.litellm_core_utils.reasoning_effort_capability import deployment_reasoning_effort_levels from litellm.litellm_core_utils.request_timeout_resolver import get_configured_request_timeout -from litellm.llms.anthropic.common_utils import AnthropicError, anthropic_pass_through_api_base +from litellm.llms.anthropic.common_utils import AnthropicError, AnthropicModelInfo, anthropic_pass_through_api_base from litellm.llms.anthropic.messages.native import handler +from litellm.llms.anthropic.messages.native import utils as native_utils from litellm.llms.anthropic.messages.native.handler import ( _native_messages_config as native_messages_config, # pyright: ignore[reportPrivateUsage] # the handler's own config lookup ) +from litellm.llms.anthropic.messages.native.handler import ( + _resolve_requested_provider as _requested_provider, # pyright: ignore[reportPrivateUsage] # the handler's own provider resolution +) from litellm.llms.anthropic.messages.native.handler import ( _with_deployment_output_effort as _deployment_output_effort, # pyright: ignore[reportPrivateUsage] # the handler's own effort rounding ) from litellm.llms.anthropic.messages.utils import is_reasoning_auto_summary_enabled from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import get_ssl_configuration +from litellm.llms.custom_httpx.llm_http_handler import ( + _merge_messages_header_sources as _merged_headers, # pyright: ignore[reportPrivateUsage] # the handler's own header merge +) from litellm.rust_bridge import failures from litellm.rust_bridge.host import callbacks_legacy_python, http_transport from litellm.rust_bridge.public_call import known_params, optional_str from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse +from litellm.types.router import GenericLiteLLMParams +from litellm.types.utils import ProviderSpecificHeader _METADATA: Final = TypeAdapter(dict[str, object]) _BODY_FIELDS: Final = frozenset( @@ -71,6 +90,10 @@ class MessagesSettings: reasoning_auto_summary: bool additional_drop_params: Sequence[str] request_timeout: float | None + model: str | None = None + mock_response: Mapping[str, object] | None = None + stream_cost_model: str | None = None + extra_headers: Mapping[str, object] | None = None def response(value: Mapping[str, object]) -> AnthropicMessagesResponse: @@ -129,7 +152,10 @@ _UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) _BODY_CUT: Final = TypeAdapter(tuple[int, int | None, bool]) _ERROR_CLASS: TypeAlias = Callable[..., Exception] _validate_metadata: Final = cast( # cast-ok: the handler's untyped metadata validator - Callable[[dict[str, object]], object], handler.validate_anthropic_api_metadata + Callable[[object], object], handler.validate_anthropic_api_metadata +) +_mock_response: Final = cast( # cast-ok: the handler's mock builder takes untyped messages + Callable[..., Mapping[str, object]], native_utils.mock_response ) _exception_type: Final = cast( # cast-ok: the legacy public exception mapper is untyped Callable[..., Exception], exception_mapping_utils.exception_type @@ -385,12 +411,128 @@ def _drop_params(kwargs: Mapping[str, object]) -> bool: return bool(litellm.drop_params) or normalize_drop_params(kwargs.get("drop_params")) is True -def settings(kwargs: Mapping[str, object]) -> dict[str, object]: - return asdict( - MessagesSettings( - drop_params=_drop_params(kwargs), - reasoning_auto_summary=is_reasoning_auto_summary_enabled(), - additional_drop_params=handler.additional_drop_paths(kwargs), - request_timeout=get_configured_request_timeout(), +def settings(kwargs: Mapping[str, object], request: Mapping[str, object] | None = None) -> dict[str, object]: + configured: Final = MessagesSettings( + drop_params=_drop_params(kwargs), + reasoning_auto_summary=is_reasoning_auto_summary_enabled(), + additional_drop_params=handler.additional_drop_paths(kwargs), + request_timeout=get_configured_request_timeout(), + ) + return asdict(configured if request is None else _with_call_settings(configured, request)) + + +_HANDLER_PARAMS: Final = frozenset( + name + for name, parameter in inspect.signature( + cast(Callable[..., object], handler.anthropic_messages_handler) # noqa: TID251 # cast-ok: reads the handler's parameter names, never calls it + ).parameters.items() + if parameter.kind is not inspect.Parameter.VAR_KEYWORD +) + + +def _handler_params(request: Mapping[str, object]) -> GenericLiteLLMParams: + return GenericLiteLLMParams( + **{name: value for name, value in request.items() if name not in _HANDLER_PARAMS}, + api_key=optional_str(request.get("api_key")), + api_base=optional_str(request.get("api_base")), + custom_llm_provider=optional_str(request.get("custom_llm_provider")), + ) + + +def _requested_model(request: Mapping[str, object]) -> str: + model, _, _, _ = _requested_provider( + str(request["model"]), optional_str(request.get("custom_llm_provider")), _handler_params(request) + ) + return model + + +def _model_without_latest_alias(request: Mapping[str, object], model: str, custom_llm_provider: str | None) -> str: + if not is_model_latest_alias(model, custom_llm_provider): + return model + resolved, provider, _, _ = _requested_provider(model, custom_llm_provider, _handler_params(request)) + prefix, separator, _ = model.partition("/") + return f"{prefix}/{resolved}" if separator and prefix == provider else resolved + + +def _mock(request: Mapping[str, object]) -> Mapping[str, object] | None: + mock: Final = request.get("mock_response") + if not mock or not isinstance(mock, str): + return None + _validate_metadata(request.get("metadata")) + messages: Final = request.get("messages") + max_tokens: Final = request.get("max_tokens") + return dict( + _mock_response( + model=_requested_model(request), + messages=messages if isinstance(messages, list) else [], + max_tokens=max_tokens if isinstance(max_tokens, int) else 0, + mock_response=mock, ) ) + + +def _stream_cost_model(request: Mapping[str, object]) -> str | None: + if request.get("stream") is not True or not stream_cost_injection_requested(proxy_request_headers(request)): + return None + return _requested_model(request) + + +def _affinity_extra_headers( + request: Mapping[str, object], model: str, explicit_provider: str | None +) -> Mapping[str, object] | None: + api_base: Final = optional_str(request.get("api_base")) + anthropic_header: Final = get_provider_affinity_header(AnthropicModelInfo.get_api_base(api_base)) + other_header: Final = get_provider_affinity_header(api_base) + if anthropic_header is None and other_header is None: + return None + provider: Final = explicit_provider or inferred_provider(model, api_base) + default_header: Final = anthropic_header if provider == "anthropic" else other_header + if default_header is None: + return None + headers: Final = as_str_mapping(request.get("headers")) + extra_headers: Final = as_str_mapping(request.get("extra_headers")) + merged: Final = _merged_headers( + headers, + extra_headers, + ProviderSpecificHeaderUtils.get_provider_specific_headers( + provider_specific_header=cast( # cast-ok: the handler hands the caller's value on unvalidated + ProviderSpecificHeader | None, request.get("provider_specific_header") + ), + custom_llm_provider=provider, + native_anthropic_messages=True, + ), + ) + added: Final = add_provider_affinity_header(merged, _handler_params(request), default_header) + affinity: Final = {name: value for name, value in added.items() if name not in merged} + if not affinity: + return None + return dict(_merged_headers(headers, extra_headers, affinity)) + + +def _with_call_settings(configured: MessagesSettings, request: Mapping[str, object]) -> MessagesSettings: + explicit_provider: Final = optional_str(request.get("custom_llm_provider")) + model: Final = _model_without_latest_alias(request, str(request["model"]), explicit_provider) + mock: Final = _mock(request) + if mock is not None: + return replace(configured, model=model, mock_response=mock) + return replace( + configured, + model=model, + stream_cost_model=_stream_cost_model(request), + extra_headers=_affinity_extra_headers(request, model, explicit_provider), + ) + + +def inject_stream_cost(frames: bytes, model: str, request: Mapping[str, object]) -> object: + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + logger: Final = request.get("litellm_logging_obj") + logging_obj: Final = logger if isinstance(logger, LiteLLMLoggingObj) else None + return ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection( + frames, + model, + logging_obj, + request_headers=proxy_request_headers( + as_str_mapping(logging_obj.model_call_details.get("litellm_params")) if logging_obj else None + ), + ) diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 99555bb4852..f2129fe62d1 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -5,25 +5,40 @@ from __future__ import annotations from typing import Final import litellm -from litellm.anthropic_beta_headers_manager import is_foundry_api_base -from litellm.litellm_core_utils.core_helpers import ( - drops_unsupported_params, - proxy_request_headers, - stream_cost_injection_requested, -) -from litellm.llms.opencode.common_utils import is_opencode_api_base -from litellm.rust_bridge.detectors import api_base, overridden, present, streaming +from litellm.litellm_core_utils.core_helpers import drops_unsupported_params +from litellm.rust_bridge.detectors import overridden, present from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str -def provider_affinity(request: NativeCall) -> bool: - return is_opencode_api_base(api_base(request)) +def _response_cache_consulted(request: NativeCall, call_type: str, *, honours_cache_controls: bool) -> bool: + from litellm.caching.caching_handler import is_call_type_supported_by_cache, is_response_cache_excluded + + cache: Final = litellm.cache + kwargs: Final = request.resolved + if ( + cache is None + or is_response_cache_excluded(optional_str(kwargs.get("model")), kwargs) + or not is_call_type_supported_by_cache(call_type) + ): + return False + if honours_cache_controls and ( + kwargs.get("caching") not in (None, True) + or (optional_mapping(kwargs.get("cache")) or {}).get("no-cache") is True + ): + return False + return cache.should_use_cache(**kwargs) is True # pyright: ignore[reportUnknownMemberType] # Cache.should_use_cache takes untyped kwargs -def foundry_beta_row(request: NativeCall) -> bool: - from litellm.llms.anthropic.common_utils import AnthropicModelInfo +def response_cache(request: NativeCall) -> bool: + return _response_cache_consulted(request, "anthropic_messages", honours_cache_controls=True) - return is_foundry_api_base(AnthropicModelInfo.get_api_base(api_base(request))) + +def sync_response_cache(request: NativeCall) -> bool: + return _response_cache_consulted(request, "anthropic_messages_handler", honours_cache_controls=False) + + +def streamed_mock_response(request: NativeCall) -> bool: + return present(request, "mock_response") and request.resolved.get("stream") is True def anthropic_workload_identity(request: NativeCall) -> bool: @@ -55,11 +70,6 @@ def thinking_signature_retry(request: NativeCall) -> bool: return any(_replays_thinking(message) for message in optional_sequence(request.resolved.get("messages")) or ()) -def tool_custom_field(request: NativeCall) -> bool: - tools: Final = optional_sequence(request.resolved.get("tools")) or () - return any("custom" in (optional_mapping(tool) or {}) for tool in tools) - - def strict_tools_retry(request: NativeCall) -> bool: tools: Final = optional_sequence(request.resolved.get("tools")) or () return drops_unsupported_params(request.resolved.get("drop_params")) and any( @@ -67,11 +77,6 @@ def strict_tools_retry(request: NativeCall) -> bool: ) -def stream_cost_injection(request: NativeCall) -> bool: - """Python adds response cost to streaming usage when the global setting or the request header asks for it.""" - return streaming(request) and stream_cost_injection_requested(proxy_request_headers(request.resolved)) - - def pre_request_hook(_request: NativeCall) -> bool: """The async Messages handler runs callbacks that can rewrite the request.""" return overridden("async_pre_request_hook", litellm.callbacks) diff --git a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py index 41008f0d57f..6239ce7bbc0 100644 --- a/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py +++ b/tests/rust-python-harness/strategies/e2e_parity/sdk/messages/test_native_parity.py @@ -34,8 +34,10 @@ import pytest from pydantic import JsonValue, TypeAdapter import litellm +from litellm.caching.caching import Cache, CacheMode from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict from litellm.messages import dispatch +from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.rust_bridge import catalog from litellm.rust_bridge.catalog import Route, RouteRule from litellm.rust_bridge.configuration import Rollout @@ -342,36 +344,18 @@ S15_COMPLETE_DEADLINE: Final = ( "S15: by design, native bounds a complete response by one overall deadline where Python bounds each read, " "so a body trickling in slower than the timeout succeeds in Python and times out natively" ) -F_MOCK: Final = "F: mock_response is Python-only (detector mock_response)" F_THINKING_RETRY: Final = ( "F: retry without thinking after a rejected signature is Python-only (detector thinking_signature_retry)" ) F_CACHE_INJECTION: Final = "F: prompt cache_control injection is Python-only (detector cache_control_injection)" -F_TOOL_CUSTOM: Final = "F: hoisting a tool's custom field is Python-only (detector tool_custom_field)" -F_STREAM_COST: Final = "F: cost injection into streamed usage is Python-only (detector stream_cost_injection)" +F_STREAMED_MOCK: Final = ( + "F: a mock with stream=True returns the mock dict and logs no success in Python (detector streamed_mock_response)" +) KNOWN_GAPS: Final[Mapping[str, Mapping[str, str | None]]] = MappingProxyType( { "response_private_keys": _gaps(outcome=S6_PRIVATE_KEYS), "missing_max_tokens": _gaps(dispatch=A10_NO_MAX_TOKENS), - "mock_response": _gaps( - logged_response_headers=F_MOCK, - request=F_MOCK, - outcome=F_MOCK, - hidden_params=F_MOCK, - logged=F_MOCK, - logged_metadata=F_MOCK, - logged_api_key=F_MOCK, - logged_call=F_MOCK, - ), - "mock_response_rate_limit": _gaps( - logged_response_headers=F_MOCK, - request=F_MOCK, - outcome=F_MOCK, - logged=F_MOCK, - logged_api_key=F_MOCK, - logged_call=F_MOCK, - ), "thinking_signature_rejected_then_retried": _gaps( request=F_THINKING_RETRY, outcome=F_THINKING_RETRY, @@ -381,9 +365,13 @@ KNOWN_GAPS: Final[Mapping[str, Mapping[str, str | None]]] = MappingProxyType( logged_response_headers=F_THINKING_RETRY, ), "cache_control_injection_points": _gaps(request=F_CACHE_INJECTION), - "tool_custom_field": _gaps(request=F_TOOL_CUSTOM), - "stream_cost_injection": _gaps(outcome=F_STREAM_COST), "enable_prompt_caching": _gaps(request=F_CACHE_INJECTION), + "mock_response_streamed": _gaps( + logged=F_STREAMED_MOCK, + logged_metadata=F_STREAMED_MOCK, + logged_response_headers=F_STREAMED_MOCK, + logged_call=F_STREAMED_MOCK, + ), } ) @@ -729,9 +717,63 @@ SYNC_CASES: Final = ( Case("sync_upstream_429", upstream=(_failure(429, "rate_limit_error", "slow down"),), sync=True), ) -FALLBACK_CASES: Final = ( +LATEST_ALIAS: Final = "anthropic/claude-opus-latest" +CUSTOM_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType({**WEATHER_TOOL, "custom": {"defer_loading": True}}) +COST_HEADER: Final[Mapping[str, JsonValue]] = MappingProxyType( + {"headers": {"x-litellm-include-cost-in-streaming-usage": "true"}} +) + +SERVED_NATIVELY_CASES: Final = ( Case("mock_response", _kw(mock_response="Hi from the mock"), upstream=()), Case("mock_response_rate_limit", _kw(mock_response="litellm.RateLimitError"), upstream=()), + Case("mock_response_context_window", _kw(mock_response="litellm.ContextWindowExceededError"), upstream=()), + Case("mock_response_empty_calls_the_provider", _kw(mock_response="")), + Case("mock_response_bad_metadata", _kw(mock_response="Hi", metadata={"user_id": 123}), upstream=()), + Case("latest_alias", model=LATEST_ALIAS), + Case("latest_alias_bare_with_provider", _kw(custom_llm_provider="anthropic"), model="claude-opus-latest"), + Case("latest_alias_stream", _kw(stream=True), upstream=(_stream(TEXT_STREAM),), model=LATEST_ALIAS), + Case( + "latest_alias_upstream_400", + upstream=(_failure(400, "invalid_request_error", "max_tokens: too large"),), + model=LATEST_ALIAS, + ), + Case("tool_custom_field", _kw(tools=[dict(CUSTOM_TOOL)])), + Case( + "tool_custom_field_keeps_other_custom_keys", + _kw(tools=[{**WEATHER_TOOL, "custom": {"defer_loading": False, "note": "kept"}}]), + ), + Case( + "tool_custom_field_dropped_with_drop_params", + _kw(tools=[{**WEATHER_TOOL, "custom": {"defer_loading": True, "note": "dropped"}}], drop_params=True), + ), + Case("tool_custom_field_not_an_object", _kw(tools=[{**WEATHER_TOOL, "custom": "opaque"}])), + Case( + "tool_custom_field_under_own_defer_loading", + _kw(tools=[{**WEATHER_TOOL, "defer_loading": False, "custom": {"defer_loading": True}}]), + ), + Case( + "stream_cost_injection", + _kw(stream=True), + upstream=(_stream(TEXT_STREAM),), + settings=MappingProxyType({"include_cost_in_streaming_usage": True}), + ), + Case( + "stream_cost_injection_from_proxy_header", + _kw(stream=True, proxy_server_request=dict(COST_HEADER)), + upstream=(_stream(TOOL_STREAM, chunk_delay=0.01),), + ), + Case( + "cache_off_by_default", + settings=MappingProxyType({"cache": Cache(mode=CacheMode.default_off)}), + ), + Case( + "cache_skips_messages_call_type", + settings=MappingProxyType({"cache": Cache(supported_call_types=["acompletion"])}), + ), +) + +FALLBACK_CASES: Final = ( + Case("mock_response_streamed", _kw(stream=True, mock_response="Hi from the mock"), upstream=()), Case( "thinking_replay", _kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), @@ -748,7 +790,6 @@ FALLBACK_CASES: Final = ( "cache_control_injection_points", _kw(system="Answer in one word.", cache_control_injection_points=[{"location": "message", "role": "system"}]), ), - Case("tool_custom_field", _kw(tools=[{**WEATHER_TOOL, "custom": {"defer_loading": True}}])), Case( "strict_tools_rejected_then_retried", _kw(tools=[{**WEATHER_TOOL, "strict": True}], drop_params=True), @@ -758,18 +799,20 @@ FALLBACK_CASES: Final = ( ), ), Case("modify_params", _kw(max_tokens=1), settings=MappingProxyType({"modify_params": True})), - Case( - "stream_cost_injection", - _kw(stream=True), - upstream=(_stream(TEXT_STREAM),), - settings=MappingProxyType({"include_cost_in_streaming_usage": True}), - ), Case("enable_prompt_caching", _kw(system="Answer in one word.", enable_prompt_caching=True)), Case("request_timeout_kwarg", _kw(request_timeout=30)), Case("stream_timeout_kwarg", _kw(stream=True, stream_timeout=30), upstream=(_stream(TEXT_STREAM),)), ) -CASES: Final = (*REQUEST_CASES, *RESPONSE_CASES, *FAILURE_CASES, *STREAM_CASES, *SYNC_CASES, *FALLBACK_CASES) +CASES: Final = ( + *REQUEST_CASES, + *RESPONSE_CASES, + *FAILURE_CASES, + *STREAM_CASES, + *SYNC_CASES, + *SERVED_NATIVELY_CASES, + *FALLBACK_CASES, +) @dataclass(frozen=True, slots=True) @@ -1077,3 +1120,45 @@ async def test_native_messages_match_python( assert native.logged == python.logged case _: assert native.callback_kwargs.get(aspect) == python.callback_kwargs.get(aspect) + + +@pytest.mark.parametrize("case", [pytest.param(case, id=case.name) for case in SERVED_NATIVELY_CASES]) +def test_catalog_serves_the_case_natively(case: Case, parity_server: RecordingServer) -> None: + with isolated_callback_registries(), _environment(case.env), _settings(case.settings): + request: Final = dispatch._public_request(dispatch._AMESSAGES, (), _kwargs(case, parity_server)) # pyright: ignore[reportPrivateUsage] # the public binding the async entry routes by + assert request is not None + context: Final = dispatch._context(request, dispatch._ASYNC_FEATURES) # pyright: ignore[reportPrivateUsage] # the async entry's own route context + assert catalog.rollout(context) is Rollout.RUST_OPT_IN, context.features + + +_COST_CASES: Final = tuple( + pytest.param(case, id=case.name) for case in SERVED_NATIVELY_CASES if case.name.startswith("stream_cost_injection") +) + + +def _proxy_relayed(text: str, request_headers: Mapping[str, object]) -> str: + return "".join( + str( + ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection( + frame + "\n\n", "proxy-model-alias", None, request_headers=request_headers + ) + ) + for frame in text.split("\n\n") + if frame + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("case", _COST_CASES) +async def test_both_lanes_carry_one_cost_per_usage_through_the_proxy_relay( + case: Case, parity_server: RecordingServer, native_messages: FallbackProbe +) -> None: + python, native = await _legs(case, parity_server, native_messages) + opted_in: Final = {"x-litellm-include-cost-in-streaming-usage": "true"} + python_text: Final = str(python.outcome["text"]) + native_text: Final = str(native.outcome["text"]) + assert '"cost"' in python_text + assert native_text == python_text + relayed: Final = _proxy_relayed(native_text, opted_in) + assert relayed == _proxy_relayed(python_text, opted_in) + assert all(frame.count('"cost"') <= 1 for frame in relayed.split("\n\n")) diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index b580eca2e42..309ea01fb8a 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -464,6 +464,10 @@ ADVISOR_TOOL: Final = {"type": ANTHROPIC_ADVISOR_TOOL_TYPE, "name": "advisor", " STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": True} +def _handler_cache() -> Cache: + return Cache(supported_call_types=["anthropic_messages_handler"]) # pyright: ignore[reportArgumentType] # the sync handler name the cache gate keys on + + @pytest.mark.parametrize( ("feature", "model", "kwargs", "settings"), ( @@ -476,7 +480,12 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": ), (Feature.CALLER_CLIENT, "anthropic/claude-sonnet-4-5", {"client": object()}, {}), (Feature.AGENTIC_LOOP_HOOK, "anthropic/claude-sonnet-4-5", {}, {"callbacks": [_Agentic()]}), - (Feature.RESPONSE_CACHE, "anthropic/claude-sonnet-4-5", {}, {"cache": Cache()}), + ( + Feature.RESPONSE_CACHE, + "anthropic/claude-sonnet-4-5", + {}, + {"cache": _handler_cache()}, + ), (Feature.TLS_POLICY, "anthropic/claude-sonnet-4-5", {}, {"ssl_security_level": "DEFAULT@SECLEVEL=1"}), (Feature.MODIFY_PARAMS, "anthropic/claude-sonnet-4-5", {}, {"modify_params": True}), ( @@ -492,18 +501,6 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": {"messages": [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "x"}]}]}, {}, ), - ( - Feature.TOOL_CUSTOM_FIELD, - "anthropic/claude-sonnet-4-5", - {"tools": [{"name": "a", "input_schema": {"type": "object"}, "custom": {"eager_input_streaming": True}}]}, - {}, - ), - ( - Feature.STREAM_COST_INJECTION, - "anthropic/claude-sonnet-4-5", - {"stream": True}, - {"include_cost_in_streaming_usage": True}, - ), ), ) def test_each_detected_feature_reaches_its_rule( @@ -619,12 +616,6 @@ async def test_strict_tools_stay_native_unless_python_could_drop_strict( ASYNC_ONLY_FEATURES: Final = ( pytest.param(Feature.PRE_REQUEST_HOOK, {}, {"callbacks": [_Rewrite()]}, id="pre-request-hook"), - pytest.param( - Feature.FOUNDRY_BETA_ROW, - {"api_base": "https://team.services.ai.azure.com/anthropic"}, - {}, - id="foundry-beta-row", - ), pytest.param( Feature.WEBSEARCH_INTERCEPTION, {}, {"callbacks": [WebSearchInterceptionLogger()]}, id="websearch-interception" ), @@ -686,6 +677,51 @@ async def test_handler_hooks_the_sync_entry_never_runs_keep_only_async_calls_on_ assert served_sync is native_result +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("kwargs", "async_on_python"), + ( + pytest.param({"mock_response": "hi"}, False, id="unary"), + pytest.param({"mock_response": "hi", "stream": True}, True, id="streamed"), + ), +) +async def test_async_mock_response_falls_back_only_when_streamed( + monkeypatch: pytest.MonkeyPatch, kwargs: Mapping[str, object], async_on_python: bool +) -> None: + rules: Final[Rules] = ( + RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=frozenset({Feature.MOCK_RESPONSE})), + RouteRule(Route.MESSAGES, Rollout.RUST_REQUIRED), + ) + python_result: Final = response() + native_result: Final = response() + + async def python(*call_args: object, **python_kwargs: object) -> AnthropicMessagesResponse: # kwargs-ok: stub + return python_result + + async def native(request: NativeCall) -> AnthropicMessagesResponse: + return native_result + + monkeypatch.setattr(litellm, "callbacks", []) + served_async: Final = await _ADISPATCH.arun( + (16, MESSAGES, "anthropic/claude-sonnet-4-5"), + kwargs, + python=python, + binding=amessages_binding(native), + native=lambda hook, request, args, native_kwargs: hook(request), + rules=rules, + ) + served_sync: Final = _DISPATCH.run( + (16, MESSAGES, "anthropic/claude-sonnet-4-5"), + kwargs, + python=lambda *call_args, **python_kwargs: python_result, + binding=messages_binding(lambda request: native_result), + native=lambda hook, request, args, native_kwargs: hook(request), + rules=rules, + ) + assert served_async is (python_result if async_on_python else native_result) + assert served_sync is python_result + + @pytest.mark.asyncio async def test_interceptor_feature_follows_the_resolved_provider() -> None: rules: Final[Rules] = ( diff --git a/tests/unit/litellm_core_utils/test_core_helpers.py b/tests/unit/litellm_core_utils/test_core_helpers.py index 02476545cae..ef2048998b1 100644 --- a/tests/unit/litellm_core_utils/test_core_helpers.py +++ b/tests/unit/litellm_core_utils/test_core_helpers.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.core_helpers import ( bind_budget_reservation_to_callbacks, budget_reservation_from_metadata, drop_params_env_flag, + env_flag, drop_params_flag, get_litellm_metadata_from_kwargs, get_or_create_metadata_bucket, @@ -497,6 +498,32 @@ def test_drop_params_env_flag_keeps_a_non_flag_value_on_with_a_warning(configure ) +@pytest.mark.parametrize( + ("configured", "expected"), + [("", False), ("false", False), ("False", False), ("0", False), ("no", False), ("true", True), (" 1 ", True)], +) +def test_env_flag_parses_modify_params_as_a_bool(configured: str, expected: bool, caplog: pytest.LogCaptureFixture): + with caplog.at_level(logging.WARNING, logger="env-flag-test"): + assert ( + env_flag({"LITELLM_MODIFY_PARAMS": configured}, "LITELLM_MODIFY_PARAMS", logging.getLogger("env-flag-test")) + is expected + ) + assert caplog.text == "" + + +def test_env_flag_treats_a_non_flag_value_as_on_with_a_warning(caplog: pytest.LogCaptureFixture): + with caplog.at_level(logging.WARNING, logger="env-flag-test"): + assert ( + env_flag({"LITELLM_MODIFY_PARAMS": "enabled"}, "LITELLM_MODIFY_PARAMS", logging.getLogger("env-flag-test")) + is True + ) + assert "LITELLM_MODIFY_PARAMS='enabled' is not a flag value, treating it as on" in caplog.text + + +def test_env_flag_is_off_when_unset(): + assert env_flag({}, "LITELLM_MODIFY_PARAMS", logging.getLogger("env-flag-test")) is False + + class TestIsExpectedClientError: def test_status_ranges(self): from litellm.litellm_core_utils.core_helpers import is_expected_client_error diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index aa7dcc74841..075ef7c5d6a 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -5,6 +5,7 @@ import pytest import litellm from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS +from litellm.litellm_core_utils.model_latest_aliases import get_model_latest_aliases from litellm.rust_bridge.host.routes import messages as messages_host from litellm.rust_bridge.host.routes.messages import response @@ -43,6 +44,10 @@ def test_settings_project_caller_configuration_without_resolving_a_model(monkeyp "reasoning_auto_summary": True, "additional_drop_params": ("metadata.user_id",), "request_timeout": 1800.0, + "model": None, + "mock_response": None, + "stream_cost_model": None, + "extra_headers": None, } @@ -445,3 +450,87 @@ def test_a_failure_before_any_response_maps_like_the_python_transport_error( f"litellm.InternalServerError: AnthropicException - {text}. Handle with `litellm.InternalServerError`." ) assert "This looks like a bug in LiteLLM" not in str(mapped) + + +_CALL: Final = MappingProxyType({"messages": [{"role": "user", "content": "hi"}], "max_tokens": 8}) +_OPENCODE: Final = "https://opencode.ai/inference/go/anthropic/v1" +_COST_HEADER: Final = MappingProxyType({"headers": {"x-litellm-include-cost-in-streaming-usage": "true"}}) + + +@pytest.mark.parametrize( + ("model", "custom_llm_provider", "expected"), + [ + ( + "anthropic/claude-opus-latest", + None, + f"anthropic/{get_model_latest_aliases()['anthropic']['claude-opus-latest']}", + ), + ("claude-opus-latest", "anthropic", get_model_latest_aliases()["anthropic"]["claude-opus-latest"]), + ("anthropic/claude-sonnet-4-5", None, "anthropic/claude-sonnet-4-5"), + ], +) +def test_settings_resolve_a_latest_alias_keeping_the_provider_prefix( + model: str, custom_llm_provider: str | None, expected: str +) -> None: + settings: Final = messages_host.settings({}, {**_CALL, "model": model, "custom_llm_provider": custom_llm_provider}) + + assert settings["model"] == expected + + +@pytest.mark.parametrize( + ("extra", "expected"), + [ + ({"api_base": _OPENCODE, "session_id": "s-1"}, {"x-opencode-session": "s-1"}), + ( + {"api_base": _OPENCODE, "session_id": "s-1", "headers": {"x-trace": "t"}}, + {"x-trace": "t", "x-opencode-session": "s-1"}, + ), + ({"api_base": _OPENCODE, "session_id": "s-1", "extra_headers": {"x-opencode-session": "mine"}}, None), + ({"api_base": "https://other.example/v1", "session_id": "s-1"}, None), + ({"api_base": _OPENCODE, "session_id": "s-1", "mock_response": "hi"}, None), + ], +) +def test_settings_add_only_the_affinity_header_the_caller_did_not_send( + extra: dict[str, object], expected: dict[str, str] | None +) -> None: + settings: Final = messages_host.settings({}, {**_CALL, "model": "anthropic/claude-sonnet-4-5", **extra}) + + assert settings["extra_headers"] == expected + + +@pytest.mark.parametrize( + ("extra", "expected"), + [ + ({"stream": True, "proxy_server_request": dict(_COST_HEADER)}, "claude-sonnet-4-5"), + ({"proxy_server_request": dict(_COST_HEADER)}, None), + ({"stream": True}, None), + ({"stream": True, "proxy_server_request": dict(_COST_HEADER), "mock_response": "hi"}, None), + ], +) +def test_settings_name_the_cost_model_only_for_an_opted_in_stream( + extra: dict[str, object], expected: str | None +) -> None: + settings: Final = messages_host.settings({}, {**_CALL, "model": "anthropic/claude-sonnet-4-5", **extra}) + + assert settings["stream_cost_model"] == expected + + +def test_settings_carry_the_mock_the_python_handler_returns() -> None: + settings: Final = messages_host.settings( + {}, {**_CALL, "model": "anthropic/claude-sonnet-4-5", "mock_response": "hi"} + ) + + assert isinstance(settings["mock_response"], dict) + assert settings["mock_response"]["content"] == [{"type": "text", "text": "hi"}] + + +def test_settings_infer_no_provider_off_an_affinity_host(monkeypatch: pytest.MonkeyPatch) -> None: + def unexpected(model: str, api_base: str | None) -> str | None: + pytest.fail("the provider is only inferred for an affinity api_base") + + monkeypatch.setattr(messages_host, "inferred_provider", unexpected) + + settings: Final = messages_host.settings({}, {**_CALL, "model": "claude-sonnet-4-5"}) + + assert settings["model"] == "claude-sonnet-4-5" + assert settings["extra_headers"] is None diff --git a/tests/unit/rust_bridge/test_catalog.py b/tests/unit/rust_bridge/test_catalog.py index 1994dd00709..a15a9c9511b 100644 --- a/tests/unit/rust_bridge/test_catalog.py +++ b/tests/unit/rust_bridge/test_catalog.py @@ -20,16 +20,10 @@ from litellm.rust_bridge.public_call import NativeCall @pytest.mark.parametrize( - ("model", "expected"), - ( - ("anthropic/claude-opus-latest", Rollout.PYTHON_ONLY), - (f"anthropic/{get_model_latest_aliases()['anthropic']['claude-opus-latest']}", Rollout.RUST_OPT_IN), - ), + "model", + ("anthropic/claude-opus-latest", f"anthropic/{get_model_latest_aliases()['anthropic']['claude-opus-latest']}"), ) -def test_latest_alias_messages_falls_back_without_changing_the_pinned_model_rollout( - model: str, - expected: Rollout, -) -> None: +def test_latest_alias_messages_stays_native_like_the_pinned_model(model: str) -> None: request: Final = NativeCall(args=(), kwargs={"model": model}, base={}) context: Final = RouteContext( Route.MESSAGES, @@ -41,7 +35,7 @@ def test_latest_alias_messages_falls_back_without_changing_the_pinned_model_roll RouteRule(Route.MESSAGES, Rollout.PYTHON_ONLY, features=catalog.MESSAGES_PYTHON_FEATURES), RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN, providers=frozenset({"anthropic"})), ) - assert catalog.rollout(context, rules) is expected + assert catalog.rollout(context, rules) is Rollout.RUST_OPT_IN @pytest.fixture(autouse=True) @@ -152,6 +146,7 @@ def test_empty_constraints_match_nothing() -> None: ) assert catalog.decision(RouteContext(Route.OCR), rules) is Decision.PYTHON - assert catalog.decision( - RouteContext(Route.OCR, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), rules - ) is Decision.PYTHON + assert ( + catalog.decision(RouteContext(Route.OCR, features=frozenset({Feature.SAFEGUARDS_CLASSIFIER})), rules) + is Decision.PYTHON + ) diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index fccc01dd8e6..f1869eb8b5c 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -1,4 +1,5 @@ from collections.abc import Callable, Mapping +import ssl from copy import deepcopy from datetime import datetime from functools import reduce @@ -8,7 +9,7 @@ import httpx import pytest import litellm -from litellm.caching.caching import Cache +from litellm.caching.caching import Cache, CacheMode from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.responses.utils import ResponsesAPIRequestUtils @@ -136,8 +137,6 @@ def _wrap_nested_level(value: object, depth: int) -> object: (responses_detectors.provider_affinity, {"provider_affinity_header": None, "session_id": "session-1"}, False), (responses_detectors.provider_affinity, {}, False), (responses_detectors.provider_affinity, {"api_base": "https://opencode.ai/inference/go/openai/v1"}, True), - (messages_detectors.provider_affinity, {"api_base": "https://opencode.ai/inference/go/anthropic/v1"}, True), - (messages_detectors.provider_affinity, {"api_base": "https://other.example/v1"}, False), (detectors.streaming, {"stream": True}, True), (detectors.streaming, {"stream": False}, False), (detectors.timeout_variants, {"request_timeout": 5}, True), @@ -208,9 +207,6 @@ def _wrap_nested_level(value: object, depth: int) -> object: {"messages": [{"role": "user", "content": [{"type": "text"}]}]}, False, ), - (messages_detectors.tool_custom_field, {"tools": [{"name": "a"}, {"name": "b", "custom": {}}]}, True), - (messages_detectors.tool_custom_field, {"tools": [{"name": "a", "defer_loading": True}]}, False), - (messages_detectors.tool_custom_field, {}, False), (messages_detectors.strict_tools_retry, {"drop_params": True, "tools": [{"name": "a"}, STRICT_TOOL]}, True), ( messages_detectors.strict_tools_retry, @@ -323,8 +319,6 @@ def test_input_rewrite_preserves_caller_input(input_value: object, expected: boo (detectors.modify_params, "modify_params", True, {}), (detectors.retries_fallbacks, "num_retries", 3, {}), (detectors.cache_control_injection, "enable_anthropic_prompt_caching", True, {}), - (messages_detectors.stream_cost_injection, "include_cost_in_streaming_usage", True, {"stream": True}), - (messages_detectors.stream_cost_injection, "include_cost_in_streaming_usage", "true", {"stream": True}), (chat_detectors.chat_retries_fallbacks, "model_fallbacks", [{"gpt-5": ["gpt-5-mini"]}], {}), (chat_detectors.chat_retries_fallbacks, "num_retries", 3, {}), (chat_detectors.model_alias_map, "model_alias_map", {"gpt-5": "openai/gpt-5"}, {}), @@ -347,30 +341,6 @@ def test_global_settings_are_detected( assert detector(call(**kwargs)) -def test_stream_cost_injection_needs_a_stream(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) - assert not messages_detectors.stream_cost_injection(call()) - - -def test_stream_cost_injection_reads_the_request_header_the_proxy_reads(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", False) - opted_in: Final = {"proxy_server_request": {"headers": {"x-litellm-include-cost-in-streaming-usage": "true"}}} - assert messages_detectors.stream_cost_injection(call(stream=True, **opted_in)) - assert not messages_detectors.stream_cost_injection(call(**opted_in)) - assert not messages_detectors.stream_cost_injection(call(stream=True, proxy_server_request={"headers": {}})) - - -def test_foundry_beta_row_reads_the_api_base_the_anthropic_handler_reads(monkeypatch: pytest.MonkeyPatch) -> None: - for name in ("ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): - monkeypatch.delenv(name, raising=False) - assert messages_detectors.foundry_beta_row(call(api_base="https://team.services.ai.azure.com/anthropic")) - assert messages_detectors.foundry_beta_row(call(base_url="https://team.cognitiveservices.azure.com/anthropic")) - assert not messages_detectors.foundry_beta_row(call(api_base="https://api.anthropic.com")) - assert not messages_detectors.foundry_beta_row(call()) - monkeypatch.setenv("ANTHROPIC_BASE_URL", "https://team.services.ai.azure.com/anthropic") - assert messages_detectors.foundry_beta_row(call()) - - def test_model_alias_map_applies_only_to_the_requested_model(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_alias_map", {"fast": "openai/gpt-5-mini"}) assert chat_detectors.model_alias_map(call(model="fast")) @@ -549,9 +519,76 @@ def test_response_cache_is_skipped_only_by_no_cache( assert detectors.response_cache(call(**kwargs)) is expected +@pytest.mark.parametrize( + ("cache", "kwargs", "expected"), + ( + pytest.param(None, {}, False, id="no-cache-configured"), + pytest.param(Cache(), {}, True, id="default-cache"), + pytest.param(Cache(), {"cache": {"no-cache": True}}, False, id="no-cache-control"), + pytest.param(Cache(mode=CacheMode.default_off), {}, False, id="default-off"), + pytest.param( + Cache(mode=CacheMode.default_off), {"cache": {"use-cache": True}}, True, id="default-off-opted-in" + ), + pytest.param(Cache(supported_call_types=["acompletion"]), {}, False, id="messages-not-cached"), + pytest.param(Cache(supported_call_types=["anthropic_messages"]), {}, True, id="messages-cached"), + pytest.param(Cache(), {"caching": False}, False, id="caching-off"), + pytest.param(Cache(), {"model": "databricks_agent/agent"}, False, id="excluded-provider"), + ), +) +def test_messages_response_cache_fires_only_when_python_touches_the_cache( + monkeypatch: pytest.MonkeyPatch, cache: Cache | None, kwargs: Mapping[str, object], expected: bool +) -> None: + monkeypatch.setattr(litellm, "cache", cache) + assert messages_detectors.response_cache(call(**kwargs)) is expected + + +def _handler_cache() -> Cache: + return Cache(supported_call_types=["anthropic_messages_handler"]) # pyright: ignore[reportArgumentType] # the sync handler name the cache gate keys on + + +@pytest.mark.parametrize( + ("cache", "kwargs", "expected"), + ( + pytest.param(_handler_cache(), {}, True, id="handler-cached"), + pytest.param(Cache(supported_call_types=["anthropic_messages"]), {}, False, id="only-async-cached"), + pytest.param( + _handler_cache(), + {"cache": {"no-cache": True}}, + True, + id="sync-lookup-ignores-no-cache", + ), + ), +) +def test_sync_messages_response_cache_keys_on_the_handler_name( + monkeypatch: pytest.MonkeyPatch, cache: Cache, kwargs: Mapping[str, object], expected: bool +) -> None: + monkeypatch.setattr(litellm, "cache", cache) + assert messages_detectors.sync_response_cache(call(**kwargs)) is expected + + +@pytest.mark.parametrize( + ("kwargs", "expected"), + ( + pytest.param({"mock_response": "hi", "stream": True}, True, id="streamed"), + pytest.param({"mock_response": "hi"}, False, id="unary"), + pytest.param({"mock_response": "hi", "stream": False}, False, id="stream-false"), + pytest.param({"stream": True}, False, id="no-mock"), + ), +) +def test_streamed_mock_response_needs_a_mock_and_stream_true(kwargs: Mapping[str, object], expected: bool) -> None: + assert messages_detectors.streamed_mock_response(call(**kwargs)) is expected + + @pytest.mark.parametrize( ("environment", "setting", "expected"), - ((None, None, False), ("DEFAULT@SECLEVEL=1", None, True), (None, "DEFAULT@SECLEVEL=1", True)), + ( + (None, None, False), + ("DEFAULT@SECLEVEL=1", None, True), + (None, "DEFAULT@SECLEVEL=1", True), + (None, "ECDHE-RSA-AES128-GCM-SHA256", False), + (None, "ECDHE-RSA-AES128-GCM-SHA256:RC4-SHA", True), + ("HIGH", None, False), + ), ) def test_tls_policy_reads_the_security_level_environment_then_the_setting( monkeypatch: pytest.MonkeyPatch, environment: str | None, setting: str | None, expected: bool @@ -567,7 +604,14 @@ def test_tls_policy_reads_the_security_level_environment_then_the_setting( @pytest.mark.parametrize( ("environment", "setting", "expected"), - ((None, "X25519", False), (None, " P-256 ", False), ("secp521r1", None, True), (None, "secp521r1", True)), + ( + (None, "X25519", False), + (None, " P-256 ", False), + ("secp521r1", None, True), + (None, "secp521r1", True), + ("", "secp521r1", False), + (None, "", False), + ), ) def test_tls_policy_flags_only_curves_native_tls_lacks( monkeypatch: pytest.MonkeyPatch, environment: str | None, setting: str | None, expected: bool @@ -579,3 +623,24 @@ def test_tls_policy_flags_only_curves_native_tls_lacks( monkeypatch.setenv("SSL_ECDH_CURVE", environment) monkeypatch.setattr(litellm, "ssl_ecdh_curve", setting) assert detectors.tls_policy(call()) is expected + + +@pytest.mark.parametrize(("setting", "value"), (("ssl_security_level", 1), ("ssl_ecdh_curve", object()))) +def test_tls_policy_hands_an_undecodable_setting_to_python_without_raising( + monkeypatch: pytest.MonkeyPatch, setting: str, value: object +) -> None: + for name in ("SSL_SECURITY_LEVEL", "SSL_ECDH_CURVE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "ssl_security_level", None) + monkeypatch.setattr(litellm, "ssl_ecdh_curve", None) + monkeypatch.setattr(litellm, setting, value) + assert detectors.tls_policy(call()) is True + + +def test_tls_policy_ignores_a_live_ssl_context(monkeypatch: pytest.MonkeyPatch) -> None: + for name in ("SSL_SECURITY_LEVEL", "SSL_ECDH_CURVE"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "ssl_security_level", None) + monkeypatch.setattr(litellm, "ssl_ecdh_curve", None) + monkeypatch.setattr(litellm, "ssl_verify", ssl.create_default_context()) + assert detectors.tls_policy(call()) is False