From d8927a8e77100e68e6a68fb354d429ad23a28ed2 Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Tue, 3 Nov 2026 05:36:22 +0000 Subject: [PATCH] feat(messages): serve mock_response, thinking retry, cache injection, tool custom fields and stream cost natively --- .../callbacks-legacy-python/src/adapter.rs | 27 ++ .../crates/gateway-inference/src/messages.rs | 1 + .../crates/inference-messages/src/handler.rs | 275 ++++++++++-- .../crates/inference-messages/src/lib.rs | 4 + .../crates/inference-messages/src/types.rs | 10 +- .../tests/bedrock_native/main.rs | 2 +- .../inference-messages/tests/caching.rs | 4 +- .../inference-messages/tests/live/support.rs | 1 + .../inference-messages/tests/messages/main.rs | 1 + .../tests/messages/request.rs | 1 + .../tests/openrouter/main.rs | 2 +- litellm-rust/crates/inference/src/error.rs | 36 +- .../crates/llms/src/anthropic/common_utils.rs | 167 ++++++- .../src/anthropic/messages/transformation.rs | 20 +- .../src/azure_ai/messages/transformation.rs | 4 + .../src/base_llm/messages/http_error_retry.rs | 180 ++++++++ .../crates/llms/src/base_llm/messages/mod.rs | 1 + .../src/base_llm/messages/transformation.rs | 12 +- .../messages/transformation.rs | 8 +- .../bedrock_mantle/messages/transformation.rs | 3 +- .../src/deepseek/messages/transformation.rs | 4 + .../src/minimax/messages/transformation.rs | 8 +- .../src/tencent/messages/transformation.rs | 8 +- .../src/vertex_ai/messages/transformation.rs | 8 +- .../crates/python-bridge/src/errors.rs | 5 + .../python-bridge/src/routes/messages/host.rs | 55 ++- .../python-bridge/src/routes/messages/mod.rs | 33 +- .../crates/python-bridge/src/routes/mod.rs | 11 +- litellm/messages/dispatch.py | 6 +- litellm/rust_bridge/catalog.py | 12 +- litellm/rust_bridge/host/routes/messages.py | 113 +++++ litellm/rust_bridge/messages/detectors.py | 25 +- .../sdk/messages/test_native_parity.py | 415 ++++++++++++++++-- .../messages/test_dispatch.py | 20 - .../rust_bridge/host/routes/test_messages.py | 212 +++++++++ tests/unit/rust_bridge/test_detectors.py | 24 - 36 files changed, 1534 insertions(+), 184 deletions(-) create mode 100644 litellm-rust/crates/llms/src/base_llm/messages/http_error_retry.rs diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 1bb84184e4c..1d0305d5cd7 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -499,6 +499,14 @@ impl LegacyLogging { ) -> PyResult>> { self.end = Some(datetime(py, timing.end_time)?); self.response = Some(response); + if self.answers_stream_request_completely(py)? { + let logger = self.logger()?.object(py); + logger.setattr("stream", true)?; + logger + .getattr("model_call_details")? + .set_item("stream", true)?; + return self.finalize(py); + } if self.runs_deployment_hooks() { return Ok(HookStep::Await( DeploymentHooks::after_success( @@ -593,11 +601,30 @@ impl LegacyLogging { self.response = Some(response.clone_ref(py)); match &self.stream { Some(stream) => self.stream_success(py, stream)?, + None if self.answers_stream_request_completely(py)? => {} None => self.dispatch_success(py)?, } Ok(HookStep::Ready(())) } + /// Python's `@client` returns a complete Messages result to a `stream=True` call (a + /// `mock_response`) as it would a stream, so neither the success deployment hook nor + /// success logging runs for it (`_is_streaming_request`). + fn answers_stream_request_completely(&self, py: Python<'_>) -> PyResult { + if self.operation != LoggingOperation::Messages || self.stream.is_some() { + return Ok(false); + } + let kwargs = self.call.kwargs().bind(py); + let mocked = kwargs.get_item("mock_response")?.is_some_and(|mock| { + mock.cast::() + .is_ok_and(|text| text.len().is_ok_and(|len| len > 0)) + }); + Ok(mocked + && kwargs + .get_item("stream")? + .is_some_and(|stream| stream.is(pyo3::types::PyBool::new(py, true).to_owned()))) + } + pub(crate) fn failed( &mut self, py: Python<'_>, diff --git a/litellm-rust/crates/gateway-inference/src/messages.rs b/litellm-rust/crates/gateway-inference/src/messages.rs index 58ecd40204d..e3f0c2c0bea 100644 --- a/litellm-rust/crates/gateway-inference/src/messages.rs +++ b/litellm-rust/crates/gateway-inference/src/messages.rs @@ -80,6 +80,7 @@ fn project( }) .collect(); Ok(MessagesCall { + mock_response: None, body: messages_body(body)?, api_key: deployment.api_key.clone(), api_base: deployment.api_base.clone(), diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index e0a7aa5adb8..7291fbf869a 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -44,6 +44,9 @@ struct ProviderCall { signer: Option, timeout: Option, stream: bool, + /// Relay the stream in whole SSE frames, as Python's pass-through chunk processor does + /// while it injects cost into streamed usage. + whole_frames: bool, } impl ProviderCall { @@ -98,6 +101,35 @@ pub(super) async fn execute( Ok(cache.finish(output, &source).await) } +/// The errors Python's `mock_response` raises for the `mock_response` text that names them. +const MOCKED_ERRORS: [&str; 3] = [ + "litellm.InternalServerError", + "litellm.ContextWindowExceededError", + "litellm.RateLimitError", +]; + +/// Python's Messages `mock_response`: the canned message carrying `mock` as its text, or the +/// error it names, answered before any provider work or logging update. +pub(super) fn mock_response(mock: &str) -> Result { + if MOCKED_ERRORS.contains(&mock) { + return Err(Error::Mocked(mock.to_string())); + } + let response = serde_json::from_value(serde_json::json!({ + "content": [{"text": mock, "type": "text"}], + "id": "msg_013Zva2CMHLNnXjNJJKqJ2EF", + "model": "claude-sonnet-4-20250514", + "role": "assistant", + "stop_reason": "end_turn", + "stop_sequence": null, + "type": "message", + "usage": {"input_tokens": 2095, "output_tokens": 503}, + })) + .map_err(|err| { + Error::invalid_response(litellm_llms::ErrorDetail::invalid("mock response", err)) + })?; + Ok(response) +} + async fn provider_call( auth: &AuthServices, provider: MessagesProvider, @@ -123,6 +155,9 @@ async fn provider_call( }; let config = provider.config(); let env_lookup = |key: &str| secrets.get(key); + // Python takes `stream` from the caller, not from the body that `additional_drop_params` + // may have trimmed, for the URL, the response handling and logging. + let caller_stream = body.params.stream == Some(true); let sanitized = config.shape_request( MessagesRequest { model, ..body }, @@ -166,7 +201,7 @@ async fn provider_call( api_base.as_deref(), &body.model, &connection_params, - body.params.stream == Some(true), + caller_stream, &env_lookup, )?; @@ -194,7 +229,8 @@ async fn provider_call( ) .await?; let stream = match wire.body.get("stream") { - None | Some(Value::Null) => false, + None => caller_stream, + Some(Value::Null) => false, Some(Value::Bool(stream)) => *stream, Some(value) => { return Err(Error::invalid_request( @@ -213,6 +249,7 @@ async fn provider_call( signer: authenticated.signer, timeout: timeout.or_else(|| shaping.settings.request_timeout()), stream, + whole_frames: shaping.settings.include_cost_in_streaming_usage, }) } @@ -303,32 +340,55 @@ async fn call_provider( signer, timeout, stream, + whole_frames, } = request; let provider_name = provider.as_str(); - log_request_body(provider_name, stream, &wire.body); - let timeout = timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS)); - let response = send( - http, - Authenticated { - headers: wire.headers, - signer, - }, - &wire.url, - &wire.body, - timeout, - stream, - ) - .await?; - if !response.status().is_success() { - return Err(provider_error(response, timeout).await); - } let config = provider.config(); + let timeout = timeout.unwrap_or(Duration::from_secs(MESSAGES_TIMEOUT_SECS)); + let mut body = wire.body; + let mut attempts = 1; + let response = loop { + log_request_body(provider_name, stream, &body); + let response = send( + http, + Authenticated { + headers: wire.headers.clone(), + signer: signer.clone(), + }, + &wire.url, + &body, + timeout, + stream, + ) + .await?; + if response.status().is_success() { + break response; + } + let rejection = provider_rejection(response, timeout).await?; + // Python's `_async_post_anthropic_messages_with_http_error_retry`: a recoverable + // rejection is sent once more with the body the config rewrites, under the same headers. + let retried = (attempts < litellm_llms::base_llm::messages::http_error_retry::MAX_ATTEMPTS) + .then(|| config.body_after_http_error(rejection.status, &rejection.text, &body)) + .flatten(); + match retried { + Some(rewritten) => { + tracing::debug!( + attempt = attempts + 1, + "anthropic messages: retrying after a recoverable HTTP error" + ); + body = rewritten; + attempts += 1; + } + None => return Err(rejection.into_error()), + } + }; if stream { return Ok(streaming_response( response, config.stream_decoder(), provider_name, timeout, + whole_frames, )); } let head = ResponseHead { @@ -392,18 +452,36 @@ fn timed_chunks( }) } +/// An upstream rejection whose body was read. +struct Rejection { + status: u16, + text: String, + headers: Vec<(String, String)>, +} + +impl Rejection { + fn into_error(self) -> Error { + log_error_body(self.status, &self.text); + litellm_http::Error::upstream_with_headers(self.status, self.text, self.headers).into() + } +} + /// The provider's error, its body read under `timeout`: a stream has no overall deadline once /// its head arrives, so an error body that stalls would otherwise hold the call forever. -async fn provider_error(response: reqwest::Response, timeout: Duration) -> Error { +async fn provider_rejection( + response: reqwest::Response, + timeout: Duration, +) -> Result { let status = response.status().as_u16(); let headers = litellm_http::request::header_pairs(response.headers()); match tokio::time::timeout(timeout, response.text()).await { - Ok(Ok(text)) => { - log_error_body(status, &text); - litellm_http::Error::upstream_with_headers(status, text, headers).into() - } - Ok(Err(error)) => litellm_http::Error::from(error).into(), - Err(elapsed) => litellm_http::Error::read_timeout(elapsed).into(), + Ok(Ok(text)) => Ok(Rejection { + status, + text, + headers, + }), + Ok(Err(error)) => Err(litellm_http::Error::from(error).into()), + Err(elapsed) => Err(litellm_http::Error::read_timeout(elapsed).into()), } } @@ -429,6 +507,7 @@ fn streaming_response( decoder: Option, provider: &'static str, timeout: Duration, + whole_frames: bool, ) -> MessagesCallResponse { let headers = response .headers() @@ -438,6 +517,9 @@ fn streaming_response( let chunks = timed_chunks(response, timeout) .inspect_ok(move |chunk| log_chunk(provider, "provider_response", chunk)); let chunks = match decoder { + // A decoded stream is re-encoded one whole event per chunk, and Python's decoded + // hosts relay it without the pass-through cost injection. + None if whole_frames => in_whole_frames(chunks.map_err(Error::from).boxed()), None => chunks.map_err(Error::from).boxed(), Some(decode) => decoded_chunks( chunks.map_err(std::io::Error::other).boxed(), @@ -451,6 +533,68 @@ fn streaming_response( } } +const SSE_FRAME_DELIMITERS: [&[u8]; 3] = [b"\r\n\r\n", b"\n\n", b"\r\r"]; + +/// Python's `split_complete_sse_frames`: the buffered bytes up to the last frame boundary, and +/// the unterminated tail. `pending[..scanned]` holds no whole delimiter, so only the bytes a new +/// delimiter could end in are searched. +fn split_complete_sse_frames_after(pending: &[u8], scanned: usize) -> (&[u8], &[u8]) { + let longest = SSE_FRAME_DELIMITERS + .iter() + .map(|delimiter| delimiter.len()) + .max() + .unwrap_or(0); + let from = scanned.saturating_sub(longest - 1); + let boundary_end = SSE_FRAME_DELIMITERS + .iter() + .filter_map(|delimiter| { + pending[from..] + .windows(delimiter.len()) + .rposition(|window| window == *delimiter) + .map(|start| from + start + delimiter.len()) + }) + .max() + .unwrap_or(0); + pending.split_at(boundary_end) +} + +/// The stream regrouped into whole SSE frames, as Python's pass-through chunk processor yields +/// it while cost injection is on: a read that completes no frame yields nothing, and an +/// unterminated tail is yielded once the upstream ends. +fn in_whole_frames( + chunks: BoxStream<'static, Result>, +) -> BoxStream<'static, Result> { + futures_util::stream::try_unfold( + (chunks, Vec::new(), false), + |(mut chunks, mut pending, ended)| async move { + if ended { + return Ok(None); + } + loop { + match chunks.try_next().await? { + Some(chunk) => { + let scanned = pending.len(); + pending.extend_from_slice(&chunk); + let (complete, tail) = split_complete_sse_frames_after(&pending, scanned); + if complete.is_empty() { + continue; + } + let frames = Bytes::copy_from_slice(complete); + pending = tail.to_vec(); + return Ok(Some((frames, (chunks, pending, false)))); + } + None if pending.is_empty() => return Ok(None), + None => { + let tail = Bytes::from(std::mem::take(&mut pending)); + return Ok(Some((tail, (chunks, pending, true)))); + } + } + } + }, + ) + .boxed() +} + fn decoded_chunks( bytes: ByteStream, decode: StreamDecoder, @@ -535,6 +679,7 @@ mod tests { async fn wire_body(fields: Value, shaping: MessagesShaping) -> Result { provider_call_for(MessagesCall { + mock_response: None, body: body(fields), api_key: Some("sk-test".into()), api_base: Some("https://anthropic.test".into()), @@ -584,7 +729,7 @@ mod tests { .map(|(_, value)| value.to_string()) }; let call = provider_call_with_secrets( - MessagesCall { + MessagesCall { mock_response: None, body: body( json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), ), @@ -620,6 +765,7 @@ mod tests { drop_params: bool, ) -> MessagesCall { MessagesCall { + mock_response: None, body: body( json!({"model": "claude-sonnet-4-5", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), ), @@ -698,6 +844,7 @@ mod tests { shaping: MessagesShaping, ) { let result = provider_call_for(MessagesCall { + mock_response: None, body: body( json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}]}), ), @@ -735,7 +882,7 @@ mod tests { #[case] request_timeout: Option, #[case] expected: Option, ) { - let call = provider_call_for(MessagesCall { + let call = provider_call_for(MessagesCall { mock_response: None, body: body(json!({"model": "claude-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})), api_key: Some("sk-test".into()), api_base: Some("https://anthropic.test".into()), @@ -848,7 +995,7 @@ mod tests { {"custom_llm_provider": "anthropic", "extra_headers": {"x-scoped": "anthropic", "x-priority": "scoped"}} ])) .unwrap(); - let call = provider_call_for(MessagesCall { + let call = provider_call_for(MessagesCall { mock_response: None, body: body( json!({"model": model, "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16}), ), @@ -996,7 +1143,7 @@ mod tests { #[case] path: &str, #[case] service: &str, ) { - let call = provider_call_with_secrets(MessagesCall { + let call = provider_call_with_secrets(MessagesCall { mock_response: None, body: body(json!({"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})), api_key: None, api_base: Some(base.into()), @@ -1063,6 +1210,7 @@ mod tests { Some(streaming::event_stream), "test", Duration::from_secs(5), + false, ) else { panic!("a streaming response returns chunks"); }; @@ -1074,4 +1222,71 @@ mod tests { } assert!(chunks.next().await.is_none()); } + + #[rstest] + #[case::delimiter_straddles_reads( + b"data: 1\r\n", + b"\r\ndata: 2", + b"data: 1\r\n\r\n", + b"data: 2" + )] + #[case::boundary_in_new_bytes(b"data: 1", b"\n\ndata: 2\n\nx", b"data: 1\n\ndata: 2\n\n", b"x")] + #[case::no_boundary_yet(b"data: 1\r", b"\ndata", b"", b"data: 1\r\ndata")] + fn incremental_split_matches_a_full_rescan( + #[case] earlier: &[u8], + #[case] read: &[u8], + #[case] complete: &[u8], + #[case] tail: &[u8], + ) { + let pending = [earlier, read].concat(); + assert_eq!( + split_complete_sse_frames_after(&pending, earlier.len()), + (complete, tail) + ); + assert_eq!( + split_complete_sse_frames_after(&pending, 0), + (complete, tail) + ); + } + + #[rstest] + #[case::no_boundary(b"event: a\ndata: 1", b"", b"event: a\ndata: 1")] + #[case::one_frame(b"data: 1\n\n", b"data: 1\n\n", b"")] + #[case::frame_and_tail(b"data: 1\n\ndata: 2", b"data: 1\n\n", b"data: 2")] + #[case::crlf(b"data: 1\r\n\r\ndata: 2\n\nx", b"data: 1\r\n\r\ndata: 2\n\n", b"x")] + #[case::cr_only(b"data: 1\r\rdata", b"data: 1\r\r", b"data")] + fn sse_frames_split_at_the_last_boundary( + #[case] pending: &[u8], + #[case] complete: &[u8], + #[case] tail: &[u8], + ) { + assert_eq!( + split_complete_sse_frames_after(pending, 0), + (complete, tail) + ); + } + + #[rstest] + #[case::regrouped( + &["event: a\nda", "ta: 1\n\nevent: b\n", "data: 2\n\n"], + &["event: a\ndata: 1\n\n", "event: b\ndata: 2\n\n"], + )] + #[case::tail_after_the_end(&["data: 1\n\nda", "ta"], &["data: 1\n\n", "data"])] + #[case::nothing(&[], &[])] + #[tokio::test] + async fn whole_frames_regroup_the_reads(#[case] reads: &[&str], #[case] expected: &[&str]) { + let reads: Vec> = reads + .iter() + .map(|read| Ok(Bytes::copy_from_slice(read.as_bytes()))) + .collect(); + let frames: Vec = in_whole_frames(futures_util::stream::iter(reads).boxed()) + .try_collect() + .await + .unwrap(); + let expected: Vec = expected + .iter() + .map(|frame| Bytes::copy_from_slice(frame.as_bytes())) + .collect(); + assert_eq!(frames, expected); + } } diff --git a/litellm-rust/crates/inference-messages/src/lib.rs b/litellm-rust/crates/inference-messages/src/lib.rs index 745442d0314..981bd041340 100644 --- a/litellm-rust/crates/inference-messages/src/lib.rs +++ b/litellm-rust/crates/inference-messages/src/lib.rs @@ -89,6 +89,10 @@ impl MessagesRoute { <&'static str>::from(provider.provider), ); let model = provider.model.to_owned(); + if let Some(mock) = &call.mock_response { + return handler::mock_response(mock) + .map(|response| MessagesCallResponse::Complete(Box::new(response))); + } let secrets = self.secrets.resolve(config.config().secret_names()).await?; handler::execute(self, config, model, call, secrets, &context) .boxed() diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index ae346259559..4e88f00ffb4 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -23,6 +23,8 @@ pub struct MessagesCall { pub provider_specific_header: Option, pub timeout: Option, pub shaping: MessagesShaping, + /// The caller's `mock_response` text, answered without calling the provider. + pub mock_response: Option, } pub fn messages_body(body: Map) -> Result { @@ -62,6 +64,10 @@ pub struct MessagesSettings { pub additional_drop_params: Vec, #[serde(default)] pub request_timeout: Option, + /// `litellm.include_cost_in_streaming_usage`: Python relays a stream in whole SSE frames + /// so the host can price each usage event. + #[serde(default)] + pub include_cost_in_streaming_usage: bool, } impl MessagesSettings { @@ -131,7 +137,8 @@ mod tests { "drop_params": true, "reasoning_auto_summary": true, "additional_drop_params": ["metadata.user_id", "thinking"], - "request_timeout": 1800.0 + "request_timeout": 1800.0, + "include_cost_in_streaming_usage": true }), MessagesShaping { settings: MessagesSettings { @@ -139,6 +146,7 @@ mod tests { reasoning_auto_summary: true, additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()], request_timeout: Some(1800.0), + include_cost_in_streaming_usage: true, }, capabilities: MessagesModelCapabilities { supports_reasoning: true, diff --git a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs index f91d5c79f7c..d66cd9efb50 100644 --- a/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs +++ b/litellm-rust/crates/inference-messages/tests/bedrock_native/main.rs @@ -16,7 +16,7 @@ const MODEL: &str = "claude-test"; #[fixture] fn call() -> MessagesCall { - MessagesCall { + MessagesCall { mock_response: None, body: serde_json::from_value(json!({"model": MODEL, "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})).unwrap(), api_key: None, api_base: None, diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 663c87453a9..6b4716c89cf 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -204,7 +204,7 @@ async fn messages_cache_identity_follows_resolved_configuration_and_request_call .revision .store(usize::from(call >= 2), Ordering::SeqCst); let cache = ScopedCache::new(cache.clone(), CacheScope::Shared); - support::messages_route(secrets.clone()).with_cache(cache).execute(MessagesCall { + support::messages_route(secrets.clone()).with_cache(cache).execute(MessagesCall { mock_response: None, body: serde_json::from_value(json!({"model":"anthropic/cache-test-model","messages":[{"role":"user","content":"hello"}],"max_tokens":32})).unwrap(), api_key:None,api_base:None,custom_llm_provider:None,litellm_params:Default::default(),extra_headers:None,provider_specific_header:None,timeout:None,shaping:Default::default(), }, &hooks, None).await.unwrap(); @@ -262,7 +262,7 @@ async fn provider_reported_cost_is_delivered_once_and_never_charged_for_a_cache_ let route = support::messages_route(litellm_inference_testing::no_secrets()) .with_cache(ScopedCache::new(cache, CacheScope::Shared)); let request = || { - MessagesCall { + MessagesCall { mock_response: None, body: serde_json::from_value(json!({"model": "native-test", "messages": [{"role": "user", "content": "hi"}], "max_tokens": 16})).unwrap(), api_key: Some("eden-key".into()), api_base: Some(upstream.uri()), custom_llm_provider: Some("edenai".into()), litellm_params: Default::default(), extra_headers: None, provider_specific_header: None, diff --git a/litellm-rust/crates/inference-messages/tests/live/support.rs b/litellm-rust/crates/inference-messages/tests/live/support.rs index 81e295e34ad..a92da9401eb 100644 --- a/litellm-rust/crates/inference-messages/tests/live/support.rs +++ b/litellm-rust/crates/inference-messages/tests/live/support.rs @@ -97,6 +97,7 @@ pub(super) fn call(provider: &str, payload: Value) -> MessagesCall { )) .expect("live payload must be a valid Messages request"); MessagesCall { + mock_response: None, body, api_key: None, api_base: None, diff --git a/litellm-rust/crates/inference-messages/tests/messages/main.rs b/litellm-rust/crates/inference-messages/tests/messages/main.rs index a01e571ec75..27facde903e 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/main.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/main.rs @@ -74,6 +74,7 @@ fn message_response() -> ResponseTemplate { #[fixture] fn call() -> MessagesCall { MessagesCall { + mock_response: None, body: body(json!({ "model": MODEL, "max_tokens": 16, diff --git a/litellm-rust/crates/inference-messages/tests/messages/request.rs b/litellm-rust/crates/inference-messages/tests/messages/request.rs index 216aa96fbe6..d8ded3740b9 100644 --- a/litellm-rust/crates/inference-messages/tests/messages/request.rs +++ b/litellm-rust/crates/inference-messages/tests/messages/request.rs @@ -494,6 +494,7 @@ async fn unsupported_params_are_dropped_under_drop_params_and_rejected_without_i let shaped = |drop_params: bool| { with_fields( MessagesCall { + mock_response: None, api_key: Some("sk".into()), api_base: Some(upstream.uri()), shaping: MessagesShaping { diff --git a/litellm-rust/crates/inference-messages/tests/openrouter/main.rs b/litellm-rust/crates/inference-messages/tests/openrouter/main.rs index 80243718768..5c4ca5c0a92 100644 --- a/litellm-rust/crates/inference-messages/tests/openrouter/main.rs +++ b/litellm-rust/crates/inference-messages/tests/openrouter/main.rs @@ -13,7 +13,7 @@ use support::*; #[fixture] fn call() -> MessagesCall { - MessagesCall { + MessagesCall { mock_response: None, body: serde_json::from_value(json!({"model": "openrouter/anthropic/test-model", "max_tokens": 16, "messages": [{"role": "user", "content": "hi"}]})).unwrap(), api_key: None, api_base: None, diff --git a/litellm-rust/crates/inference/src/error.rs b/litellm-rust/crates/inference/src/error.rs index 841e0c53799..7772a29db0c 100644 --- a/litellm-rust/crates/inference/src/error.rs +++ b/litellm-rust/crates/inference/src/error.rs @@ -26,6 +26,10 @@ pub enum Error { Host(#[source] MachineFault), #[error("post-call hook failed")] PostCallHook(#[source] Arc), + /// The caller's `mock_response` names an error to raise instead of calling the provider, + /// and the host raises the error its own mock names. + #[error("mock_response raised {0}")] + Mocked(String), } impl From for Error { @@ -55,7 +59,11 @@ impl Error { Self::Llm(error) => error.is_request(), Self::Auth(error) => error.is_request(), Self::Http(error) => error.is_request(), - Self::Declined(_) | Self::Secret(_) | Self::Host(_) | Self::PostCallHook(_) => false, + Self::Declined(_) + | Self::Secret(_) + | Self::Host(_) + | Self::PostCallHook(_) + | Self::Mocked(_) => false, } } @@ -70,7 +78,8 @@ impl Error { | Self::Headers(_) | Self::Secret(_) | Self::Host(_) - | Self::PostCallHook(_) => false, + | Self::PostCallHook(_) + | Self::Mocked(_) => false, } } @@ -85,7 +94,8 @@ impl Error { | Self::Http(_) | Self::Secret(_) | Self::Host(_) - | Self::PostCallHook(_) => false, + | Self::PostCallHook(_) + | Self::Mocked(_) => false, } } @@ -100,7 +110,8 @@ impl Error { | Self::Http(_) | Self::Secret(_) | Self::Host(_) - | Self::PostCallHook(_) => None, + | Self::PostCallHook(_) + | Self::Mocked(_) => None, } } @@ -127,7 +138,16 @@ impl Error { | Self::Declined(_) | Self::Auth(_) | Self::Headers(_) - | Self::Secret(_) => false, + | Self::Secret(_) + | Self::Mocked(_) => false, + } + } + + /// The `mock_response` text that asked for this error. + pub fn mocked(&self) -> Option<&str> { + match self { + Self::Mocked(mock) => Some(mock), + _ => None, } } @@ -142,7 +162,8 @@ impl Error { | Self::Http(_) | Self::Secret(_) | Self::Host(_) - | Self::PostCallHook(_) => None, + | Self::PostCallHook(_) + | Self::Mocked(_) => None, } } @@ -185,7 +206,8 @@ impl Error { | Self::Headers(_) | Self::Secret(_) | Self::Host(_) - | Self::PostCallHook(_) => None, + | Self::PostCallHook(_) + | Self::Mocked(_) => None, } } } diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 13653f438f7..65ce767d14e 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -18,7 +18,7 @@ use litellm_llms_types::{ recognized::Recognized, }; use serde::Deserialize; -use serde_json::Value; +use serde_json::{Map, Value}; use crate::{ anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX, @@ -278,6 +278,112 @@ pub fn has_advisor_tool(tools: Option<&[Recognized]>) -> bool { .any(|tool| matches!(tool, Recognized::Known(MessagesTool::Advisor { .. }))) } +/// What a Claude-compatible host does with a tool's non-spec `custom` object (Python's +/// `tools_for_upstream`). +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum CustomToolField { + /// The tools are sent as the caller wrote them (Bedrock Invoke and Mantle). + Forward, + /// `custom` is hoisted and, under `drop_params`, dropped (Python's + /// `tools_without_custom_field`). + Rewrite { + hoist_defer_loading: bool, + drop_custom: bool, + }, +} + +impl CustomToolField { + /// The base config's policy for `provider`, as Python's `AnthropicMessagesConfig` applies it. + pub fn rewrite(provider: &str, model: &str, drop_params: bool) -> Self { + Self::Rewrite { + hoist_defer_loading: is_claude_model_target(provider, model), + drop_custom: drop_params, + } + } +} + +/// Python's `is_claude_model_target`: a Claude model on a host whose tools take Claude's own +/// `defer_loading` field. +pub fn is_claude_model_target(provider: &str, model: &str) -> bool { + let model = model.to_lowercase(); + model.contains("claude") + && matches!(provider, "anthropic" | "azure_ai" | "vertex_ai" | "bedrock") +} + +fn hoistable_defer_loading(tool: &Map, custom: &Value) -> Option { + let custom = custom.as_object()?; + if tool.contains_key("defer_loading") { + return None; + } + custom.get("defer_loading").and_then(Value::as_bool) +} + +/// Python's `_tool_without_custom_field` for one tool. +fn tool_without_custom_field(tool: Value, hoist_defer_loading: bool, drop_custom: bool) -> Value { + let Value::Object(tool) = tool else { + return tool; + }; + let Some(custom) = tool.get("custom").cloned() else { + return Value::Object(tool); + }; + let without_custom = || { + tool.iter() + .filter(|(key, _)| *key != "custom") + .map(|(key, value)| (key.clone(), value.clone())) + .collect::>() + }; + let defer_loading = hoist_defer_loading + .then(|| hoistable_defer_loading(&tool, &custom)) + .flatten(); + let Some(defer_loading) = defer_loading else { + return Value::Object(if drop_custom { without_custom() } else { tool }); + }; + let unhoisted: Map = custom + .as_object() + .into_iter() + .flatten() + .filter(|(key, _)| *key != "defer_loading") + .map(|(key, value)| (key.clone(), value.clone())) + .collect(); + let mut rewritten = without_custom(); + rewritten.insert("defer_loading".into(), Value::Bool(defer_loading)); + if !unhoisted.is_empty() && !drop_custom { + rewritten.insert("custom".into(), Value::Object(unhoisted)); + } + Value::Object(rewritten) +} + +/// Python's `tools_for_upstream`: each tool's `custom` object handled as `policy` says. +pub fn tools_for_upstream( + tools: Vec>, + policy: CustomToolField, +) -> Result>, serde_json::Error> { + let CustomToolField::Rewrite { + hoist_defer_loading, + drop_custom, + } = policy + else { + return Ok(tools); + }; + tools + .into_iter() + .map(|tool| { + let value = serde_json::to_value(&tool)?; + if !value + .as_object() + .is_some_and(|tool| tool.contains_key("custom")) + { + return Ok(tool); + } + serde_json::from_value(tool_without_custom_field( + value, + hoist_defer_loading, + drop_custom, + )) + }) + .collect() +} + pub fn requires_native_compaction_beta(compaction: Option<&Value>, messages: &[Message]) -> bool { compaction.is_some() || messages.iter().flat_map(Message::blocks).any(|block| { @@ -676,6 +782,65 @@ mod tests { value.map(|tools| serde_json::from_value(tools).unwrap()) } + #[rstest] + #[case::forward_keeps_custom( + CustomToolField::Forward, + json!({"name": "f", "custom": {"defer_loading": true}}), + json!({"name": "f", "custom": {"defer_loading": true}}), + )] + #[case::hoists_defer_loading( + CustomToolField::Rewrite { hoist_defer_loading: true, drop_custom: false }, + json!({"name": "f", "custom": {"defer_loading": true}}), + json!({"name": "f", "defer_loading": true}), + )] + #[case::hoists_and_keeps_the_rest( + CustomToolField::Rewrite { hoist_defer_loading: true, drop_custom: false }, + json!({"name": "f", "custom": {"defer_loading": false, "x": 1}}), + json!({"name": "f", "defer_loading": false, "custom": {"x": 1}}), + )] + #[case::hoists_and_drops_the_rest( + CustomToolField::Rewrite { hoist_defer_loading: true, drop_custom: true }, + json!({"name": "f", "custom": {"defer_loading": true, "x": 1}}), + json!({"name": "f", "defer_loading": true}), + )] + #[case::top_level_defer_loading_wins( + CustomToolField::Rewrite { hoist_defer_loading: true, drop_custom: false }, + json!({"name": "f", "defer_loading": false, "custom": {"defer_loading": true}}), + json!({"name": "f", "defer_loading": false, "custom": {"defer_loading": true}}), + )] + #[case::non_bool_defer_loading_is_not_hoisted( + CustomToolField::Rewrite { hoist_defer_loading: true, drop_custom: false }, + json!({"name": "f", "custom": {"defer_loading": "yes"}}), + json!({"name": "f", "custom": {"defer_loading": "yes"}}), + )] + #[case::not_hoisted_off_claude( + CustomToolField::Rewrite { hoist_defer_loading: false, drop_custom: false }, + json!({"name": "f", "custom": {"defer_loading": true}}), + json!({"name": "f", "custom": {"defer_loading": true}}), + )] + #[case::dropped_under_drop_params( + CustomToolField::Rewrite { hoist_defer_loading: false, drop_custom: true }, + json!({"name": "f", "custom": "anything"}), + json!({"name": "f"}), + )] + fn custom_tool_fields_follow_the_host_policy( + #[case] policy: CustomToolField, + #[case] tool: Value, + #[case] expected: Value, + ) { + let rewritten = tools_for_upstream(tools(Some(json!([tool]))).unwrap(), policy).unwrap(); + assert_eq!(serde_json::to_value(rewritten).unwrap(), json!([expected])); + } + + #[rstest] + #[case::anthropic_claude("anthropic", "claude-sonnet-4-5", true)] + #[case::bedrock_claude("bedrock", "us.anthropic.Claude-opus", true)] + #[case::anthropic_other_model("anthropic", "other-model", false)] + #[case::deepseek_claude_name("deepseek", "claude-sonnet-4-5", false)] + fn claude_model_targets(#[case] provider: &str, #[case] model: &str, #[case] expected: bool) { + assert_eq!(is_claude_model_target(provider, model), expected); + } + fn headers(pairs: &[(&str, &str)]) -> Headers { pairs .iter() diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index 18789e5b2be..43f742d8ce9 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -16,10 +16,10 @@ 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, CustomToolField, 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, tools_for_upstream, }, base_llm::{ auth::AuthScheme, @@ -56,7 +56,9 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { - transform_messages_request(request, &context.keeping_display_updates()) + let custom_tools = + CustomToolField::rewrite("anthropic", &request.model, context.drop_params); + transform_messages_request(request, &context.keeping_display_updates(), custom_tools) } fn secret_names(&self) -> &'static [&'static str] { @@ -121,6 +123,7 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { pub(crate) fn transform_messages_request( request: MessagesRequest, context: &MessagesTransformContext, + custom_tools: CustomToolField, ) -> Result { if request.params.max_tokens.is_none() { return Err(Error::MissingField("max_tokens")); @@ -137,10 +140,17 @@ pub(crate) fn transform_messages_request( } else { strip_advisor_blocks(request.messages) }; + let tools = request + .params + .tools + .map(|tools| tools_for_upstream(tools, custom_tools)) + .transpose() + .map_err(|err| Error::InvalidRequest(crate::ErrorDetail::invalid("Messages tools", err)))?; Ok(MessagesRequest { messages: strip_encrypted_reasoning_blocks(messages), params: MessagesOptionalParams { context_management, + tools, extra: without_null_safeguards(request.params.extra), ..request.params }, 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..1af8486e93e 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -6,6 +6,7 @@ use litellm_llms_types::{ use litellm_router_types::LitellmParams; use litellm_secret_value::SecretValue; +use crate::anthropic::common_utils::CustomToolField; use crate::{ Error, anthropic::{ @@ -60,6 +61,8 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { + let custom_tools = + CustomToolField::rewrite("azure_ai", &request.model, context.drop_params); transform_messages_request( strip_cache_control_scope(normalize_system_role_messages( request, @@ -69,6 +72,7 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig { .supports_mid_conversation_system, )), &context.keeping_display_updates(), + custom_tools, ) } diff --git a/litellm-rust/crates/llms/src/base_llm/messages/http_error_retry.rs b/litellm-rust/crates/llms/src/base_llm/messages/http_error_retry.rs new file mode 100644 index 00000000000..2f23c7c1323 --- /dev/null +++ b/litellm-rust/crates/llms/src/base_llm/messages/http_error_retry.rs @@ -0,0 +1,180 @@ +//! Python's `BaseAnthropicMessagesConfig` recovery from a rejected request: a 400 naming an +//! invalid replayed thinking block is answered by sending the body again without thinking +//! (`should_retry_anthropic_messages_on_http_error` and +//! `transform_anthropic_messages_request_on_http_error`). + +use serde_json::Value; + +/// Python's `max_retry_on_anthropic_messages_http_error`: the HTTP attempts a call may make. +pub const MAX_ATTEMPTS: usize = 2; + +const THINKING_BLOCK_TYPES: [&str; 2] = ["thinking", "redacted_thinking"]; + +/// Python's `is_anthropic_invalid_thinking_block_error`. +pub fn is_invalid_thinking_block_error(error_text: &str) -> bool { + if error_text.is_empty() { + return false; + } + let lower = error_text.to_lowercase(); + if !lower.contains("thinking") { + return false; + } + if lower.contains("signature") && (lower.contains("invalid") || lower.contains("valid string")) + { + return true; + } + lower.contains("must contain thinking") +} + +/// Python's `is_anthropic_thinking_binding_mismatch_error`. +fn is_thinking_binding_mismatch_error(error_text: &str) -> bool { + error_text + .to_lowercase() + .contains("bound to a different conversation") +} + +/// Python's `anthropic_request_errors_on_thinking_binding_mismatch`: the caller pinned the +/// blocks with `prefix_mismatch_behavior: error`, so the retry would undo what it asked for. +fn errors_on_thinking_binding_mismatch(body: &Value) -> bool { + body.get("thinking") + .and_then(Value::as_object) + .and_then(|thinking| thinking.get("block_binding")) + .and_then(Value::as_object) + .and_then(|binding| binding.get("prefix_mismatch_behavior")) + .is_some_and(|behavior| behavior == "error") +} + +/// Python's `should_retry_anthropic_messages_on_http_error` for the thinking case. +pub fn retries_without_thinking(status: u16, error_text: &str, body: &Value) -> bool { + if status != 400 { + return false; + } + let caller_asked_for_this_error = + errors_on_thinking_binding_mismatch(body) && is_thinking_binding_mismatch_error(error_text); + !caller_asked_for_this_error && is_invalid_thinking_block_error(error_text) +} + +fn is_thinking_block(block: &Value) -> bool { + block + .get("type") + .and_then(Value::as_str) + .is_some_and(|kind| THINKING_BLOCK_TYPES.contains(&kind)) +} + +/// Python's `strip_thinking_blocks_from_anthropic_messages`: a message left with an empty +/// content list is dropped, since Anthropic rejects one. +fn without_thinking_blocks(messages: Vec) -> Vec { + messages + .into_iter() + .filter_map(|mut message| { + if let Some(Value::Array(content)) = message.get_mut("content") { + content.retain(|block| !is_thinking_block(block)); + if content.is_empty() { + return None; + } + } + Some(message) + }) + .collect() +} + +/// Python's `strip_thinking_blocks_from_anthropic_messages_request_dict`: the body without its +/// replayed thinking blocks and without the `thinking` param. +pub fn without_thinking(body: Value) -> Value { + let Value::Object(body) = body else { + return body; + }; + Value::Object( + body.into_iter() + .filter(|(key, _)| key != "thinking") + .map(|(key, value)| match (key.as_str(), value) { + ("messages", Value::Array(messages)) => { + (key, Value::Array(without_thinking_blocks(messages))) + } + (_, value) => (key, value), + }) + .collect(), + ) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::json; + + use super::*; + + #[rstest] + #[case::invalid_signature( + "messages.1.content.0: Invalid `signature` in `thinking` block", + true + )] + #[case::not_a_string( + "messages.2.content.0.thinking.signature.str: Input should be a valid string", + true + )] + #[case::empty_thinking( + "messages.1.content.0.thinking: each thinking block must contain thinking", + true + )] + #[case::no_thinking("Invalid `signature` in block", false)] + #[case::unrelated("max_tokens: too large", false)] + #[case::empty("", false)] + fn invalid_thinking_errors_are_recognized(#[case] text: &str, #[case] expected: bool) { + assert_eq!(is_invalid_thinking_block_error(text), expected); + } + + #[rstest] + #[case::rejected_signature(400, "Invalid `signature` in `thinking` block", json!({}), true)] + #[case::other_status(500, "Invalid `signature` in `thinking` block", json!({}), false)] + #[case::pinned_binding( + 400, + "messages.1.content.0: Invalid `signature` in `thinking` block. The block is bound to a different conversation.", + json!({"thinking": {"block_binding": {"prefix_mismatch_behavior": "error"}}}), + false, + )] + #[case::pinned_binding_other_error( + 400, + "Invalid `signature` in `thinking` block", + json!({"thinking": {"block_binding": {"prefix_mismatch_behavior": "error"}}}), + true, + )] + fn retry_follows_the_rejection( + #[case] status: u16, + #[case] text: &str, + #[case] body: Value, + #[case] expected: bool, + ) { + assert_eq!(retries_without_thinking(status, text, &body), expected); + } + + #[test] + fn thinking_is_stripped_from_the_body() { + let body = json!({ + "model": "claude", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "redacted_thinking", "data": "x"}]}, + {"role": "assistant", "content": [ + {"type": "thinking", "thinking": "t", "signature": "s"}, + {"type": "text", "text": "Paris"} + ]}, + "not a message" + ], + "thinking": {"type": "enabled", "budget_tokens": 1024}, + "max_tokens": 64 + }); + assert_eq!( + without_thinking(body), + json!({ + "model": "claude", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": "Paris"}]}, + "not a message" + ], + "max_tokens": 64 + }) + ); + } +} 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..e3c0a8f5095 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/mod.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/mod.rs @@ -1,4 +1,5 @@ pub mod context; +pub mod http_error_retry; pub mod mid_conversation_system; pub mod normalization; pub mod streaming; 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..cbff926ec39 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, http_error_retry}; pub use crate::base_llm::auth::{Headers, ValidatedEnvironment}; use crate::{Error, base_llm::messages::streaming::StreamDecoder}; @@ -91,6 +91,16 @@ pub trait BaseMessagesConfig: Sync { fn wire_body(&self, body: Value) -> Value { body } + + /// The body to send again after the upstream rejected `body` with `status` and + /// `error_text`, or `None` when the rejection stands: Python's + /// `should_retry_anthropic_messages_on_http_error` with + /// `transform_anthropic_messages_request_on_http_error`, which every Messages config + /// inherits. + fn body_after_http_error(&self, status: u16, error_text: &str, body: &Value) -> Option { + http_error_retry::retries_without_thinking(status, error_text, body) + .then(|| http_error_retry::without_thinking(body.clone())) + } } #[cfg(test)] 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..695bf8a7a22 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 @@ -13,6 +13,7 @@ use litellm_llms_types::{ use litellm_router_types::LitellmParams; use serde_json::Value; +use crate::anthropic::common_utils::CustomToolField; use crate::{ Error, ErrorDetail, anthropic::messages::{ @@ -141,9 +142,14 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { + let request_model = request.model.clone(); let body = serde_json::to_value(request).map_err(invalid_request)?; let filtered = serde_json::from_value(filter_body(body, false)).map_err(invalid_request)?; - transform_messages_request(filtered, context) + transform_messages_request( + filtered, + context, + CustomToolField::rewrite("bedrock", &request_model, context.drop_params), + ) } fn secret_names(&self) -> &'static [&'static str] { 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..114d326d060 100644 --- a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs @@ -14,6 +14,7 @@ use litellm_llms_types::{ use litellm_router_types::LitellmParams; use serde_json::Value; +use crate::anthropic::common_utils::CustomToolField; use crate::{ Error, anthropic::{ @@ -123,7 +124,7 @@ impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { - transform_messages_request(request, context) + transform_messages_request(request, context, CustomToolField::Forward) } fn shape_request( diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs index 0dbcea01439..f2eaa5267c0 100644 --- a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -10,6 +10,7 @@ use litellm_llms_types::{ use litellm_router_types::LitellmParams; use serde_json::Value; +use crate::anthropic::common_utils::CustomToolField; use crate::{ Error, anthropic::{ @@ -66,6 +67,8 @@ impl BaseMessagesConfig for DeepSeekAnthropicMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { + let custom_tools = + CustomToolField::rewrite("deepseek", &request.model, context.drop_params); let request = transform_messages_request( MessagesRequest { params: MessagesOptionalParams { @@ -78,6 +81,7 @@ impl BaseMessagesConfig for DeepSeekAnthropicMessagesConfig { ..request }, context, + custom_tools, )?; Ok(MessagesRequest { params: MessagesOptionalParams { diff --git a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs index dafe5fbd16b..0170c1b0982 100644 --- a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs @@ -8,6 +8,7 @@ use litellm_llms_types::{ }; use litellm_router_types::LitellmParams; +use crate::anthropic::common_utils::CustomToolField; use crate::{ Error, anthropic::{ @@ -91,7 +92,12 @@ impl BaseMessagesConfig for MinimaxMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { - let request = transform_messages_request(request, context)?; + let request_model = request.model.clone(); + let request = transform_messages_request( + request, + context, + CustomToolField::rewrite("minimax", &request_model, context.drop_params), + )?; Ok(MessagesRequest { params: MessagesOptionalParams { system: request diff --git a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs index 621a09f4fb0..46e471186bc 100644 --- a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs @@ -8,6 +8,7 @@ use litellm_llms_types::{ }; use litellm_router_types::LitellmParams; +use crate::anthropic::common_utils::CustomToolField; use crate::{ Error, anthropic::{ @@ -96,7 +97,12 @@ impl BaseMessagesConfig for TencentAnthropicMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { - let request = transform_messages_request(request, context)?; + let request_model = request.model.clone(); + let request = transform_messages_request( + request, + context, + CustomToolField::rewrite("tencent", &request_model, context.drop_params), + )?; Ok(MessagesRequest { params: MessagesOptionalParams { system: request 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..fecc5155b88 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -12,6 +12,7 @@ use litellm_llms_types::{ use litellm_router_types::LitellmParams; use serde_json::{Value, json}; +use crate::anthropic::common_utils::CustomToolField; use crate::{ Error, anthropic::{ @@ -77,7 +78,12 @@ impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { request: MessagesRequest, context: &MessagesTransformContext, ) -> Result { - let request = transform_messages_request(request, &context.keeping_display_updates())?; + let request_model = request.model.clone(); + let request = transform_messages_request( + request, + &context.keeping_display_updates(), + CustomToolField::rewrite("vertex_ai", &request_model, context.drop_params), + )?; let request = strip_cache_control_scope(normalize_system_role_messages( request, context diff --git a/litellm-rust/crates/python-bridge/src/errors.rs b/litellm-rust/crates/python-bridge/src/errors.rs index 426cd861729..5408f46ce12 100644 --- a/litellm-rust/crates/python-bridge/src/errors.rs +++ b/litellm-rust/crates/python-bridge/src/errors.rs @@ -34,6 +34,8 @@ const MESSAGES_MISSING_CREDENTIAL: &str = "messages_missing_credential"; /// Python refuses the request only once its logger knows the resolved model and provider (the /// provider config's request transform), so the failure logs them. const MESSAGES_LATE_REJECTION: &str = "messages_late_rejection"; +/// The `mock_response` text whose error the Messages host raises in place of this one. +const MESSAGES_MOCK_ERROR: &str = "messages_mock_error"; pub(crate) fn upstream_error( py: Python<'_>, @@ -137,6 +139,9 @@ pub(crate) fn route_failure( if error.is_missing_credential() { value.setattr(MESSAGES_MISSING_CREDENTIAL, true)?; } + if let Some(mock) = error.mocked() { + value.setattr(MESSAGES_MOCK_ERROR, mock)?; + } native } (RequestMarker::Unmarked, _) => by_fault(error.is_request(), &error), 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 af614f17360..265b1729437 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -72,6 +72,8 @@ fn project_litellm_params<'py>( pub(super) struct MessagesPythonHost { view: RequestView, cache: PythonCache, + /// The model Python prices streamed usage with while it injects cost into a stream. + stream_cost_model: Option>, } impl MessagesPythonHost { @@ -79,6 +81,7 @@ impl MessagesPythonHost { Self { view, cache: PythonCache::new(asynchronous), + stream_cost_model: None, } } @@ -92,8 +95,27 @@ impl MessagesPythonHost { argument(name)?.map(|value| value.extract()).transpose() }; let model = string("model")?.ok_or(ArgumentError::Required("model"))?; - let messages = argument("messages")?.ok_or(ArgumentError::Required("messages"))?; - let fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?; + let mut messages = argument("messages")?.ok_or(ArgumentError::Required("messages"))?; + let mut fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?; + // Python's handler injects prompt-cache breakpoints with its own hook; the host runs + // that same hook so both callers share one injection. + let injected = py + .import(ROUTE_HOST_MODULE)? + .getattr("injected_prompt")? + .call1((self.view.effective_request(py, arguments)?,))?; + if !injected.is_none() { + let injected = injected.cast_into::()?; + if let Some(value) = injected.get_item("messages")? { + messages = value; + } + for name in ["system", "tools"] { + if let Some(value) = injected.get_item(name)? + && !value.is_none() + { + fields.insert(name.to_string(), from_py(&value)?); + } + } + } let body = [ ("model".to_string(), Value::String(model.clone())), ("messages".to_string(), from_py(&messages)?), @@ -120,9 +142,15 @@ impl MessagesPythonHost { let shaping = self.shaping(py, &model, custom_llm_provider.as_deref(), arguments)?; let extra_headers = self.merged_headers(py, arguments)?; let provider_specific_header = self.provider_specific_header(py, arguments)?; + // Python's handler answers a non-empty string `mock_response` itself and sends any + // other value on to the provider. + let mock_response = argument("mock_response")? + .and_then(|value| value.extract::().ok()) + .filter(|mock| !mock.is_empty()); let litellm_params = project_litellm_params(argument, |name| module_global(py, name))?; Ok(messages_body(body).and_then(|body| { Ok(MessagesCall { + mock_response, body, api_key, api_base, @@ -249,14 +277,29 @@ impl PythonBinding for MessagesPythonHost { py: Python<'_>, head: MessagesStreamHead, ) -> PyResult> { - py.import(ROUTE_HOST_MODULE)? + let module = py.import(ROUTE_HOST_MODULE)?; + let model = module + .getattr("stream_cost_model")? + .call1((self.view.request(py)?,))?; + self.stream_cost_model = (!model.is_none()).then(|| model.unbind()); + module .getattr("stream_hidden_params")? .call1((to_py(py, &head.headers)?,)) .map(Bound::unbind) } + /// The chunk as the caller receives it, with the call's cost added to its streamed usage + /// while Python's `include_cost_in_streaming_usage` is on: the route then relays whole SSE + /// frames, as Python's pass-through stream does. fn encode_chunk(&mut self, py: Python<'_>, chunk: Bytes) -> PyResult> { - Ok(PyBytes::new(py, &chunk).into_any().unbind()) + let bytes = PyBytes::new(py, &chunk).into_any(); + let Some(model) = &self.stream_cost_model else { + return Ok(bytes.unbind()); + }; + py.import(ROUTE_HOST_MODULE)? + .getattr("stream_chunk_with_cost")? + .call1((bytes, model.bind(py), self.view.request(py)?)) + .map(Bound::unbind) } fn map_error(&self, py: Python<'_>, error: Error) -> PyResult { @@ -303,9 +346,13 @@ impl PythonOwned for MessagesPythonHost { fn close(&mut self, _: Python<'_>) { self.view.close(); self.cache.close(); + self.stream_cost_model = None; } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { self.view.traverse(visit)?; + if let Some(model) = &self.stream_cost_model { + visit.call(model)?; + } self.cache.traverse(visit) } } diff --git a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs index a9ae01bd7ee..d6bb7e5b334 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/mod.rs @@ -1,14 +1,25 @@ mod host; use host::MessagesPythonHost; -use litellm_callbacks_legacy_python::LoggingOperation; -use pyo3::prelude::*; +use litellm_callbacks_legacy_python::{LoggingOperation, PublicCall}; +use litellm_host_python::{HookChain, PythonCallHooks}; +use pyo3::{prelude::*, types::PyDict}; use super::NativeCall; fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyResult> { let (arguments, hooks) = crate::routes::call_hooks(py, LoggingOperation::Messages, &call, asynchronous)?; + run_messages_with(py, &call, arguments, hooks, asynchronous) +} + +fn run_messages_with( + py: Python<'_>, + call: &NativeCall<'_>, + arguments: Py, + hooks: impl PythonCallHooks + 'static, + asynchronous: bool, +) -> PyResult> { crate::routes::run_public_call( py, arguments, @@ -55,9 +66,25 @@ fn run_messages(py: Python<'_>, call: NativeCall<'_>, asynchronous: bool) -> PyR /// that until both sides serve sync calls. const SYNC_NOT_IMPLEMENTED: &str = "anthropic_messages_handler is not implemented for sync calls"; +/// Python's handler answers a non-empty string `mock_response` itself. +fn mocked(call: &NativeCall<'_>) -> PyResult { + Ok(call + .resolved()? + .get_item("mock_response")? + .and_then(|value| value.extract::().ok()) + .is_some_and(|mock| !mock.is_empty())) +} + +/// Python's sync handler answers a mocked call before it reaches the provider handler that +/// refuses sync calls, and its sync entry point logs nothing, so the mocked call runs without +/// the legacy logging hooks. #[pyfunction] #[pyo3(signature = (call))] -pub(crate) fn messages(call: NativeCall<'_>) -> PyResult> { +pub(crate) fn messages(py: Python<'_>, call: NativeCall<'_>) -> PyResult> { + if mocked(&call)? { + let arguments = PublicCall::capture(&call.base, &call.args, &call.kwargs)?.arguments(py); + return run_messages_with(py, &call, arguments, HookChain::new(), false); + } drop(call); Err(pyo3::exceptions::PyValueError::new_err( SYNC_NOT_IMPLEMENTED, diff --git a/litellm-rust/crates/python-bridge/src/routes/mod.rs b/litellm-rust/crates/python-bridge/src/routes/mod.rs index f51c755c33e..2e1adb13bcf 100644 --- a/litellm-rust/crates/python-bridge/src/routes/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/mod.rs @@ -94,6 +94,15 @@ impl RequestView { present(arguments, self.request(py)?, name) } + /// The keywords the call would send, the hook-rewritten ones over the base. + pub(crate) fn effective_request<'py>( + &self, + py: Python<'py>, + arguments: &Bound<'py, PyDict>, + ) -> PyResult> { + effective_py_args(self.request(py)?, arguments) + } + /// The body params `module.body_params` keeps from the keywords the call would send, so /// each route's Python host decides which params reach the provider. pub(crate) fn body_params( @@ -102,7 +111,7 @@ impl RequestView { arguments: &Bound<'_, PyDict>, module: &str, ) -> PyResult> { - let request = effective_py_args(self.request(py)?, arguments)?; + let request = self.effective_request(py, arguments)?; from_py( &py.import(module)? .getattr("body_params")? diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 4c94f26c6c1..9122c9767c4 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -122,16 +122,12 @@ def messages_target(model: str, custom_llm_provider: str | None, model_info: obj _FEATURES: Final[detectors.Detectors] = MappingProxyType( { - **detectors.COMMON, + **{feature: detector for feature, detector in detectors.COMMON.items() if feature is not Feature.MOCK_RESPONSE}, Feature.PROVIDER_AFFINITY: route_detectors.provider_affinity, Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, Feature.RESPONSE_CACHE: detectors.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, } ) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index 9adfcb60d74..53aed4e7d18 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -94,12 +94,6 @@ 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, 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,7 +166,7 @@ CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { Feature.MODEL_ALIAS_MAP, } -MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { +MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.MOCK_RESPONSE}) | { Feature.PROVIDER_AFFINITY, Feature.FOUNDRY_BETA_ROW, Feature.SAFEGUARDS_CLASSIFIER, @@ -182,10 +176,6 @@ MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { Feature.RESPONSE_CACHE, 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/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index c8b53097f46..7f94aed3dff 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -22,6 +22,7 @@ from litellm.llms.anthropic.messages.native.handler import ( 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.native.utils import prepare_native_messages from litellm.llms.anthropic.messages.utils import is_reasoning_auto_summary_enabled from litellm.rust_bridge import failures from litellm.rust_bridge.host import callbacks_legacy_python @@ -65,6 +66,7 @@ class MessagesSettings: reasoning_auto_summary: bool additional_drop_params: Sequence[str] request_timeout: float | None + include_cost_in_streaming_usage: bool def response(value: Mapping[str, object]) -> AnthropicMessagesResponse: @@ -110,6 +112,70 @@ def _with_deployment_effort(params: Mapping[str, object], request: Mapping[str, ) +_PREPARE: TypeAlias = Callable[..., tuple[object, object]] +# `anthropic_messages`'s named params, which its `**kwargs` never carry. +_NAMED_PARAMS: Final = frozenset( + { + "max_tokens", + "messages", + "model", + "metadata", + "stop_sequences", + "stream", + "system", + "temperature", + "thinking", + "tool_choice", + "tools", + "top_k", + "top_p", + "api_key", + "api_base", + "client", + "custom_llm_provider", + } +) + + +def injected_prompt(request: Mapping[str, object]) -> Mapping[str, object] | None: + """The messages, system and tools Python's handler sends once it injects the configured or + default prompt-cache breakpoints, or None when it injects none. Stamps the call's metadata + bucket with the injection as the handler does.""" + from litellm.integrations.anthropic_cache_control_hook import ( + AnthropicCacheControlHook, + configured_injection_points, + ) + from litellm.llms.anthropic.common_utils import tools_marked_for_a_tool_config_point + + if not ( + litellm.enable_anthropic_prompt_caching is True + or request.get("enable_prompt_caching") is True + or configured_injection_points(request.get("cache_control_injection_points")) + ): + return None + # A shallow copy, so the metadata bucket the stamp writes is the call's own. + kwargs: Final = {name: value for name, value in request.items() if name not in _NAMED_PARAMS} + tools: Final = request.get("tools") + # The handler hands the caller's own objects to the injection, which edits them in place. + messages, system = cast(_PREPARE, prepare_native_messages)( # cast-ok: the caller's unvalidated prompt + request.get("messages"), + request.get("system"), + kwargs, + model=str(request["model"]), + custom_llm_provider=optional_str(request.get("custom_llm_provider")), + tools=tools, + api_base=optional_str(request.get("api_base")), + ) + marked: Final = tools_marked_for_a_tool_config_point(tools, kwargs.pop("cache_control_injection_points", None)) + if marked is not None: + AnthropicCacheControlHook.record_gateway_injection(kwargs, 1) + return { + "messages": messages, + "system": system, + "tools": list(marked) if marked is not None else tools, + } + + def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, object]: from litellm.llms.anthropic.messages.native.streaming_iterator import ( anthropic_messages_stream_hidden_params, @@ -118,6 +184,29 @@ def stream_hidden_params(headers: Sequence[tuple[str, str]]) -> Mapping[str, obj return anthropic_messages_stream_hidden_params(httpx.Headers(list(headers))) +def stream_cost_model(request: Mapping[str, object]) -> str | None: + """The model Python's pass-through stream prices each streamed usage event with, or None when + it adds no cost (`PassThroughStreamingHandler.chunk_processor`).""" + if not litellm.include_cost_in_streaming_usage: + return None + logging_obj: Final = request.get("litellm_logging_obj") + if not isinstance(logging_obj, LiteLLMLoggingObj): + return None + model: Final = logging_obj.model_call_details.get("model") + return model if isinstance(model, str) and model else None + + +def stream_chunk_with_cost(chunk: bytes, model: str, request: Mapping[str, object]) -> object: + """The whole SSE frames of one streamed chunk with the call's cost added to the first usage + event, as Python's pass-through stream relays them.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + + logging_obj: Final = request.get("litellm_logging_obj") + return ProxyBaseLLMRequestProcessing.process_chunk_with_cost_injection( + chunk, model, logging_obj if isinstance(logging_obj, LiteLLMLoggingObj) else None + ) + + _UPSTREAM_ARGS: Final = TypeAdapter(tuple[int, str]) _UPSTREAM_HEADERS: Final = TypeAdapter(list[tuple[str, str]]) _ERROR_CLASS: TypeAlias = Callable[..., Exception] @@ -226,8 +315,31 @@ def _log_resolved_call(request: Mapping[str, object], request_provider: str) -> ) +def _mocked_error(mock: str, request: Mapping[str, object], request_provider: str) -> Exception: + """The error Python's `mock_response` raises for the `mock_response` text that names it.""" + from litellm.llms.anthropic.messages.native import utils as native_utils + + raise_mocked: Final = cast( # cast-ok: the mock's untyped kwargs + Callable[..., object], + native_utils.mock_response, # pyright: ignore[reportUnknownMemberType] # untyped **kwargs + ) + try: + raise_mocked( + model=str(request["model"]).removeprefix(f"{request_provider}/"), + messages=[], + max_tokens=0, + mock_response=mock, + ) + except Exception as raised: # noqa: BLE001 # the mock raises the error it names + return raised + return RuntimeError(f"mock_response={mock!r} names no error") + + def map_failure(error: Exception, request: Mapping[str, object], request_provider: str) -> Exception: """Raises what Python's Messages handler raises for the same failure.""" + mocked: Final = getattr(error, "messages_mock_error", None) + if isinstance(mocked, str): + return _mocked_error(mocked, request, request_provider) if getattr(error, "messages_timeout", None) == "Body": # Python reads a stream's body outside its HTTP client's error mapping, so the transport's # own timeout reaches the caller unmapped. @@ -278,5 +390,6 @@ def settings(kwargs: Mapping[str, object]) -> dict[str, object]: reasoning_auto_summary=is_reasoning_auto_summary_enabled(), additional_drop_params=_additional_drop_params(kwargs), request_timeout=get_configured_request_timeout(), + include_cost_in_streaming_usage=bool(litellm.include_cost_in_streaming_usage), ) ) diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 6a0c867979b..63513c1b2b5 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -8,7 +8,7 @@ import litellm from litellm.anthropic_beta_headers_manager import is_foundry_api_base from litellm.litellm_core_utils.core_helpers import drops_unsupported_params from litellm.llms.opencode.common_utils import is_opencode_api_base -from litellm.rust_bridge.detectors import api_base, overridden, present, streaming +from litellm.rust_bridge.detectors import api_base, overridden, present from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str @@ -38,24 +38,6 @@ def anthropic_workload_identity(request: NativeCall) -> bool: ) -_THINKING_BLOCKS: Final = frozenset({"thinking", "redacted_thinking"}) - - -def _replays_thinking(message: object) -> bool: - content: Final = optional_sequence((optional_mapping(message) or {}).get("content")) or () - return any((optional_mapping(block) or {}).get("type") in _THINKING_BLOCKS for block in content) - - -def thinking_signature_retry(request: NativeCall) -> bool: - """Python may strip replayed thinking blocks and retry after a rejected signature.""" - 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( @@ -63,11 +45,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 is truthy.""" - return streaming(request) and bool(litellm.include_cost_in_streaming_usage) - - 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..9f38c61ae83 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 @@ -270,6 +270,25 @@ THINKING_STREAM: Final = _events( *_message_end("end_turn", 12), ) +CACHED_USAGE_STREAM: Final = _events( + _message_start("msg_01ParityCachedStream"), + *_text_block(0, "Paris"), + ( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": { + "input_tokens": 14, + "output_tokens": 4, + "cache_creation_input_tokens": 2048, + "cache_read_input_tokens": 1024, + }, + }, + ), + ("message_stop", {"type": "message_stop"}), +) + ERROR_MID_STREAM: Final = _events( _message_start("msg_01ParityOverloaded"), *_text_block(0, "Pa"), @@ -283,6 +302,14 @@ SLOW_STREAM: Final = _events( ) +CACHE_TURNS: Final[tuple[JsonValue, ...]] = ( + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": [{"type": "text", "text": "Paris"}]}, + {"role": "user", "content": [{"type": "text", "text": "And of Spain?"}]}, +) +COST_IN_STREAM: Final[Mapping[str, object]] = MappingProxyType({"include_cost_in_streaming_usage": True}) + + def _upstream_error(kind: str, message: str) -> Mapping[str, JsonValue]: return MappingProxyType({"type": "error", "error": {"type": kind, "message": message}}) @@ -342,48 +369,11 @@ 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)" 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, - hidden_params=F_THINKING_RETRY, - logged=F_THINKING_RETRY, - logged_metadata=F_THINKING_RETRY, - 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), } ) @@ -491,6 +481,7 @@ REQUEST_CASES: Final = ( _kw(temperature=0.5, additional_drop_params=["temperature"]), model=OPUS_ADAPTIVE, ), + Case("additional_drop_params_drops_stream", _kw(stream=True, additional_drop_params=["stream"])), Case( "additional_drop_params_nested_path", _kw( @@ -732,6 +723,42 @@ SYNC_CASES: Final = ( FALLBACK_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_internal_server_error", _kw(mock_response="litellm.InternalServerError"), upstream=()), + Case( + "mock_response_context_window_exceeded", + _kw(mock_response="litellm.ContextWindowExceededError"), + upstream=(), + ), + Case("mock_response_unknown_error_name_is_text", _kw(mock_response="litellm.BadRequestError"), upstream=()), + Case("mock_response_stream", _kw(mock_response="Hi from the mock", stream=True), upstream=()), + Case("mock_response_stream_rate_limit", _kw(mock_response="litellm.RateLimitError", stream=True), upstream=()), + Case("mock_response_sync", _kw(mock_response="Hi from the mock"), upstream=(), sync=True), + Case("mock_response_sync_rate_limit", _kw(mock_response="litellm.RateLimitError"), upstream=(), sync=True), + Case("mock_response_sync_stream", _kw(mock_response="Hi from the mock", stream=True), upstream=(), sync=True), + Case("mock_response_empty_string_calls_the_provider", _kw(mock_response="")), + Case("mock_response_not_a_string_calls_the_provider", _kw(mock_response={"content": "hi"})), + Case("mock_response_bare_model", _kw(mock_response="Hi from the mock"), model="claude-sonnet-4-5", upstream=()), + Case( + "mock_response_explicit_provider_rate_limit", + _kw(mock_response="litellm.RateLimitError", custom_llm_provider="anthropic"), + model="claude-sonnet-4-5", + upstream=(), + ), + Case( + "mock_response_with_metadata", + _kw( + mock_response="Hi from the mock", + metadata={"user_id": "user-123"}, + litellm_metadata={"user_api_key": "hashed-key"}, + model_info={"id": "deployment-1"}, + ), + upstream=(), + ), + Case( + "mock_response_without_api_key", + _kw(mock_response="Hi from the mock", api_key=None), + upstream=(), + ), Case( "thinking_replay", _kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), @@ -744,11 +771,159 @@ FALLBACK_CASES: Final = ( _ok(), ), ), + Case( + "thinking_rejected_twice_raises", + _kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), + upstream=( + _failure(400, "invalid_request_error", "messages.1.content.0: Invalid `signature` in `thinking` block"), + _failure(400, "invalid_request_error", "messages.1.content.0: Invalid `signature` in `thinking` block"), + ), + ), + Case( + "thinking_empty_block_rejected_then_retried", + _kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), + upstream=( + _failure( + 400, "invalid_request_error", "messages.1.content.0.thinking: each thinking block must contain thinking" + ), + _ok(), + ), + ), + Case( + "thinking_signature_not_a_string_retried", + _kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), + upstream=( + _failure( + 400, + "invalid_request_error", + "messages.1.content.0.thinking.signature.str: Input should be a valid string", + ), + _ok(), + ), + ), + Case( + "thinking_only_turn_dropped_on_retry", + _kw( + messages=[ + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": [{"type": "redacted_thinking", "data": "EmwKAhgBEgy"}]}, + {"role": "user", "content": "And Germany?"}, + ], + thinking={"type": "enabled", "budget_tokens": 2048}, + max_tokens=4096, + ), + upstream=( + _failure(400, "invalid_request_error", "messages.1.content.0: Invalid `signature` in `thinking` block"), + _ok(), + ), + ), + Case( + "thinking_rejected_then_retried_stream", + _kw( + stream=True, + messages=list(THINKING_REPLAY), + thinking={"type": "enabled", "budget_tokens": 2048}, + max_tokens=4096, + ), + upstream=( + _failure(400, "invalid_request_error", "messages.1.content.0: Invalid `signature` in `thinking` block"), + _stream(TEXT_STREAM), + ), + ), + Case( + "thinking_rejected_then_retry_fails_differently", + _kw(messages=list(THINKING_REPLAY), thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), + upstream=( + _failure(400, "invalid_request_error", "messages.1.content.0: Invalid `signature` in `thinking` block"), + _failure(429, "rate_limit_error", "slow down"), + ), + ), + Case( + "thinking_binding_mismatch_pinned_is_not_retried", + _kw( + messages=list(THINKING_REPLAY), + thinking={"type": "adaptive", "block_binding": {"prefix_mismatch_behavior": "error"}}, + max_tokens=4096, + ), + model=OPUS_ADAPTIVE, + upstream=( + _failure( + 400, + "invalid_request_error", + "messages.1.content.0: Invalid `signature` in `thinking` block. " + "The block is bound to a different conversation.", + ), + _ok(), + ), + ), + Case( + "thinking_binding_mismatch_unpinned_is_retried", + _kw(messages=list(THINKING_REPLAY), thinking={"type": "adaptive"}, max_tokens=4096), + model=OPUS_ADAPTIVE, + upstream=( + _failure( + 400, + "invalid_request_error", + "messages.1.content.0: Invalid `signature` in `thinking` block. " + "The block is bound to a different conversation.", + ), + _ok(), + ), + ), + Case( + "thinking_rejection_without_replayed_thinking", + _kw(thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), + upstream=( + _failure(400, "invalid_request_error", "messages.1.content.0: Invalid `signature` in `thinking` block"), + _ok(), + ), + ), Case( "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( + "tool_custom_field_keeps_unhoisted_keys", + _kw(tools=[{**WEATHER_TOOL, "custom": {"defer_loading": False, "eager_input_streaming": True}}]), + ), + Case( + "tool_custom_field_dropped_with_drop_params", + _kw(tools=[{**WEATHER_TOOL, "custom": {"defer_loading": True, "note": "x"}}], drop_params=True), + ), + Case( + "tool_custom_field_without_defer_loading_dropped", + _kw(tools=[{**WEATHER_TOOL, "custom": {"note": "x"}}], drop_params=True), + ), + Case( + "tool_custom_field_under_global_drop_params", + _kw(tools=[{**WEATHER_TOOL, "custom": {"defer_loading": True, "note": "x"}}]), + settings=MappingProxyType({"drop_params": True}), + ), + Case( + "tool_custom_field_top_level_defer_loading_wins", + _kw(tools=[{**WEATHER_TOOL, "defer_loading": False, "custom": {"defer_loading": True}}]), + ), + Case("tool_custom_field_not_an_object", _kw(tools=[{**WEATHER_TOOL, "custom": "opaque"}])), + Case( + "tool_custom_field_non_bool_defer_loading", + _kw(tools=[{**WEATHER_TOOL, "custom": {"defer_loading": "yes"}}]), + ), + Case( + "tool_custom_field_on_several_tools", + _kw( + tools=[ + {**WEATHER_TOOL, "custom": {"defer_loading": True}}, + {**WEATHER_TOOL, "name": "get_time", "custom": {"defer_loading": False, "x": 1}}, + {"type": "web_search_20250305", "name": "web_search", "custom": {"defer_loading": True}}, + ] + ), + ), + Case( + "tool_custom_field_stream", + _kw(stream=True, tools=[{**WEATHER_TOOL, "custom": {"defer_loading": True}}]), + upstream=(_stream(TEXT_STREAM),), + ), Case( "strict_tools_rejected_then_retried", _kw(tools=[{**WEATHER_TOOL, "strict": True}], drop_params=True), @@ -762,9 +937,173 @@ FALLBACK_CASES: Final = ( "stream_cost_injection", _kw(stream=True), upstream=(_stream(TEXT_STREAM),), - settings=MappingProxyType({"include_cost_in_streaming_usage": True}), + settings=COST_IN_STREAM, + ), + Case( + "stream_cost_injection_tool_use", + _kw(stream=True, tools=[dict(WEATHER_TOOL)]), + upstream=(_stream(TOOL_STREAM),), + settings=COST_IN_STREAM, + ), + Case( + "stream_cost_injection_thinking", + _kw(stream=True, thinking={"type": "enabled", "budget_tokens": 2048}, max_tokens=4096), + upstream=(_stream(THINKING_STREAM),), + settings=COST_IN_STREAM, + ), + Case( + "stream_cost_injection_cache_usage", + _kw(stream=True), + upstream=(_stream(CACHED_USAGE_STREAM),), + settings=COST_IN_STREAM, + ), + Case( + "stream_cost_injection_custom_pricing", + _kw(stream=True, input_cost_per_token=0.001, output_cost_per_token=0.002), + upstream=(_stream(TEXT_STREAM),), + settings=COST_IN_STREAM, + ), + Case( + "stream_cost_injection_error_mid_stream", + _kw(stream=True), + upstream=(_stream(ERROR_MID_STREAM),), + settings=COST_IN_STREAM, + ), + Case( + "stream_cost_injection_upstream_429", + _kw(stream=True), + upstream=(_failure(429, "rate_limit_error", "slow down", **{"retry-after": "3"}),), + settings=COST_IN_STREAM, + ), + Case( + "stream_cost_injection_truthy_string_setting", + _kw(stream=True), + upstream=(_stream(TEXT_STREAM),), + settings=MappingProxyType({"include_cost_in_streaming_usage": "true"}), + ), + Case("stream_cost_injection_not_streamed", settings=COST_IN_STREAM), + Case( + "stream_cost_injection_sync_stream", + _kw(stream=True), + upstream=(_stream(TEXT_STREAM),), + settings=COST_IN_STREAM, + sync=True, ), Case("enable_prompt_caching", _kw(system="Answer in one word.", enable_prompt_caching=True)), + Case( + "cache_injection_several_points", + _kw( + system="Answer in one word.", + messages=list(CACHE_TURNS), + cache_control_injection_points=[ + {"location": "message", "role": "system"}, + {"location": "message", "role": "user"}, + {"location": "message", "index": 1}, + ], + ), + ), + Case( + "cache_injection_negative_index", + _kw(messages=list(CACHE_TURNS), cache_control_injection_points=[{"location": "message", "index": -1}]), + ), + Case( + "cache_injection_out_of_range_index", + _kw(messages=list(CACHE_TURNS), cache_control_injection_points=[{"location": "message", "index": -9}]), + ), + Case( + "cache_injection_with_ttl", + _kw( + system=[{"type": "text", "text": "Answer in one word."}], + cache_control_injection_points=[ + {"location": "message", "role": "system", "control": {"type": "ephemeral", "ttl": "1h"}} + ], + ), + ), + Case( + "cache_injection_beside_client_marks", + _kw( + system=[{"type": "text", "text": "Answer in one word.", "cache_control": {"type": "ephemeral"}}], + messages=list(CACHE_TURNS), + cache_control_injection_points=[{"location": "message", "index": -1}], + ), + ), + Case( + "cache_injection_tool_config", + _kw(tools=[dict(WEATHER_TOOL)], cache_control_injection_points=[{"location": "tool_config"}]), + ), + Case( + "cache_injection_tool_config_and_message", + _kw( + tools=[dict(WEATHER_TOOL)], + cache_control_injection_points=[{"location": "tool_config"}, {"location": "message", "role": "user"}], + ), + ), + Case( + "cache_injection_tool_config_without_tools", + _kw(cache_control_injection_points=[{"location": "tool_config"}]), + ), + Case( + "cache_injection_stamps_metadata", + _kw( + system="Answer in one word.", + cache_control_injection_points=[{"location": "message", "role": "system"}], + litellm_metadata={"user_api_key": "hashed-key"}, + model_info={"id": "deployment-1"}, + ), + ), + Case( + "cache_injection_invalid_points", + _kw(cache_control_injection_points=[{"location": "nowhere"}]), + ), + Case( + "cache_injection_stream", + _kw( + stream=True, + system="Answer in one word.", + cache_control_injection_points=[{"location": "message", "role": "system"}], + ), + upstream=(_stream(TEXT_STREAM),), + ), + Case( + "cache_injection_sync_with_mock", + _kw( + mock_response="Hi", + system="Answer in one word.", + cache_control_injection_points=[{"location": "message", "role": "system"}], + litellm_metadata={"user_api_key": "hashed-key"}, + model_info={"id": "deployment-1"}, + ), + sync=True, + ), + Case( + "enable_prompt_caching_multi_turn", + _kw(system="Answer in one word.", messages=list(CACHE_TURNS), enable_prompt_caching=True), + ), + Case("enable_prompt_caching_false", _kw(system="Answer in one word.", enable_prompt_caching=False)), + Case( + "enable_prompt_caching_stands_down_on_client_marks", + _kw( + system=[{"type": "text", "text": "Answer in one word.", "cache_control": {"type": "ephemeral"}}], + enable_prompt_caching=True, + ), + ), + Case( + "global_prompt_caching", + _kw(system="Answer in one word.", messages=list(CACHE_TURNS)), + settings=MappingProxyType({"enable_anthropic_prompt_caching": True}), + ), + Case( + "global_prompt_caching_stream", + _kw(stream=True, system="Answer in one word."), + upstream=(_stream(TEXT_STREAM),), + settings=MappingProxyType({"enable_anthropic_prompt_caching": True}), + ), + Case( + "global_prompt_caching_sync", + _kw(system="Answer in one word."), + settings=MappingProxyType({"enable_anthropic_prompt_caching": True}), + sync=True, + ), Case("request_timeout_kwarg", _kw(request_timeout=30)), Case("stream_timeout_kwarg", _kw(stream=True, stream_timeout=30), upstream=(_stream(TEXT_STREAM),)), ) diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index 6d65f289524..b4b6e76f44a 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -448,7 +448,6 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": @pytest.mark.parametrize( ("feature", "model", "kwargs", "settings"), ( - (Feature.MOCK_RESPONSE, "anthropic/claude-sonnet-4-5", {"mock_response": "hi"}, {}), ( Feature.PROXY_MCP_TOOLS, "anthropic/claude-sonnet-4-5", @@ -467,25 +466,6 @@ STRICT_TOOL: Final = {"name": "a", "input_schema": {"type": "object"}, "strict": {"anthropic_federation_rule_id": "rule-1", "anthropic_organization_id": "org-1"}, {}, ), - (Feature.CACHE_CONTROL_INJECTION, "anthropic/claude-sonnet-4-5", {"enable_prompt_caching": True}, {}), - ( - Feature.THINKING_SIGNATURE_RETRY, - "anthropic/claude-sonnet-4-5", - {"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( diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index 9d8b09dd1c4..82019c56368 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -1,7 +1,10 @@ +import json +import math from types import MappingProxyType from typing import Final import pytest +from pydantic import TypeAdapter import litellm from litellm.constants import DEFAULT_REQUEST_TIMEOUT_SECONDS @@ -35,6 +38,7 @@ def test_settings_project_caller_configuration_without_resolving_a_model(monkeyp monkeypatch.setattr(litellm, "reasoning_auto_summary", True) monkeypatch.setattr(litellm, "request_timeout", 1800) monkeypatch.setattr(litellm, "request_timeout_explicitly_set", True) + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", "true") projected: Final = messages_host.settings({"drop_params": "true", "additional_drop_params": ["metadata.user_id"]}) @@ -43,6 +47,7 @@ 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, + "include_cost_in_streaming_usage": True, } @@ -245,3 +250,210 @@ def test_a_late_rejection_tells_the_logger_the_resolved_call_first() -> None: assert isinstance(mapped, litellm.BadRequestError) assert logger.model_call_details["custom_llm_provider"] == "anthropic" assert logger.model_call_details["optional_params"] == {"max_tokens": 8} + + +def _logger(model: str) -> object: + import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + logger: Final = Logging( + model=model, + messages=[], + stream=True, + call_type="anthropic_messages", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="stream", + function_id="stream", + ) + logger.model_call_details["model"] = model + return logger + + +@pytest.mark.parametrize( + ("mock", "error_class"), + [ + ("litellm.RateLimitError", litellm.RateLimitError), + ("litellm.InternalServerError", litellm.InternalServerError), + ("litellm.ContextWindowExceededError", litellm.ContextWindowExceededError), + ], +) +def test_a_mocked_error_maps_to_the_error_python_mock_response_raises(mock: str, error_class: type[Exception]) -> None: + native: Final = ValueError("mocked") + native.messages_mock_error = mock # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + + mapped: Final = messages_host.map_failure(native, {"model": "anthropic/claude-sonnet-4-5"}, "anthropic") + + assert type(mapped) is error_class + assert getattr(mapped, "model") == "claude-sonnet-4-5" + assert getattr(mapped, "llm_provider") == "anthropic" + + +def test_a_mock_naming_no_error_maps_to_a_runtime_error_naming_the_mock() -> None: + native: Final = ValueError("mocked") + native.messages_mock_error = "hello" # pyright: ignore[reportAttributeAccessIssue] # marker the native host sets + + mapped: Final = messages_host.map_failure(native, {"model": "anthropic/claude-sonnet-4-5"}, "anthropic") + + assert isinstance(mapped, RuntimeError) + assert "'hello'" in str(mapped) + + +@pytest.mark.parametrize( + ("enabled", "model", "expected"), + [ + (False, "claude-sonnet-4-5", None), + (True, "claude-sonnet-4-5", "claude-sonnet-4-5"), + (True, "", None), + ], +) +def test_stream_cost_model_is_the_logged_model_only_when_cost_is_on( + monkeypatch: pytest.MonkeyPatch, enabled: bool, model: str, expected: str | None +) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", enabled) + + assert messages_host.stream_cost_model({"litellm_logging_obj": _logger(model)}) == expected + + +def test_stream_cost_model_needs_a_logging_object(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + + assert messages_host.stream_cost_model({"litellm_logging_obj": "not a logger"}) is None + assert messages_host.stream_cost_model({}) is None + + +_EVENT: Final = TypeAdapter(dict[str, object]) +_USAGE: Final = TypeAdapter(dict[str, int | float]) +_USAGE_EVENT: Final = { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"input_tokens": 1000, "output_tokens": 500}, +} + + +@pytest.mark.parametrize("with_logger", [True, False]) +def test_stream_chunk_with_cost_adds_the_cost_to_the_usage_event( + monkeypatch: pytest.MonkeyPatch, with_logger: bool +) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + model: Final = "claude-sonnet-4-5" + chunk: Final = f"event: message_delta\ndata: {json.dumps(_USAGE_EVENT)}\n\n".encode() + request: Final[dict[str, object]] = {"litellm_logging_obj": _logger(model)} if with_logger else {} + + priced: Final = messages_host.stream_chunk_with_cost(chunk, model, request) + + assert isinstance(priced, bytes) + data_line: Final = next(line for line in priced.decode().splitlines() if line.startswith("data: ")) + event: Final = _EVENT.validate_json(data_line.removeprefix("data: ")) + usage: Final = dict(_USAGE.validate_python(event["usage"])) + cost: Final = usage.pop("cost") + assert usage == _USAGE_EVENT["usage"] + assert {key: value for key, value in event.items() if key != "usage"} == { + key: value for key, value in _USAGE_EVENT.items() if key != "usage" + } + expected: Final = litellm.completion_cost( + model=model, + completion_response=litellm.ModelResponse( + usage=litellm.Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + ), + ) + assert expected > 0 + assert math.isclose(float(cost), expected) + + +def test_stream_chunk_with_cost_leaves_a_chunk_without_usage_alone(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "include_cost_in_streaming_usage", True) + chunk: Final = b'event: ping\ndata: {"type": "ping"}\n\n' + + assert messages_host.stream_chunk_with_cost(chunk, "claude-sonnet-4-5", {}) == chunk + + +_TOOLS: Final = ( + {"name": "first", "input_schema": {"type": "object"}}, + {"name": "last", "input_schema": {"type": "object"}}, +) + + +def _prompt_request(**extra: object) -> dict[str, object]: + return { + "model": "anthropic/claude-sonnet-4-5", + "custom_llm_provider": "anthropic", + "max_tokens": 8, + "messages": [{"role": "user", "content": [{"type": "text", "text": "hello"}]}], + "system": [{"type": "text", "text": "be brief"}], + **extra, + } + + +def test_injected_prompt_is_none_without_prompt_caching(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + + assert messages_host.injected_prompt(_prompt_request(cache_control_injection_points=[])) is None + + +def test_injected_prompt_marks_the_configured_message_point(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + request: Final = _prompt_request(cache_control_injection_points=[{"location": "message", "role": "user"}]) + + injected: Final = messages_host.injected_prompt(request) + + assert injected is not None + messages: Final = injected["messages"] + assert isinstance(messages, list) + assert messages[0]["content"][-1]["cache_control"] == {"type": "ephemeral"} + assert injected["system"] == [{"type": "text", "text": "be brief"}] + assert injected["tools"] is None + + +@pytest.mark.parametrize("per_request", [True, False]) +def test_injected_prompt_places_default_breakpoints_when_caching_is_enabled( + monkeypatch: pytest.MonkeyPatch, per_request: bool +) -> None: + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", not per_request) + request: Final = _prompt_request(**({"enable_prompt_caching": True} if per_request else {})) + + injected: Final = messages_host.injected_prompt(request) + + assert injected is not None + assert "cache_control" in json.dumps(injected) + assert "cache_control" not in json.dumps(request["messages"]) + json.dumps(request["system"]) + + +def test_injected_prompt_marks_the_last_tool_and_credits_the_deployment(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY + + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + metadata: Final[dict[str, object]] = {} + request: Final = _prompt_request( + tools=list(_TOOLS), + cache_control_injection_points=[{"location": "tool_config"}], + litellm_metadata=metadata, + model_info={"id": "deployment-1"}, + ) + + injected: Final = messages_host.injected_prompt(request) + + assert injected is not None + assert injected["tools"] == [_TOOLS[0], {**_TOOLS[1], "cache_control": {"type": "ephemeral"}}] + assert metadata[GATEWAY_INJECTED_CACHE_METADATA_KEY] == "deployment-1" + assert request["tools"] == list(_TOOLS) + + +def test_injected_prompt_keeps_the_client_tool_breakpoint(monkeypatch: pytest.MonkeyPatch) -> None: + from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY + + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + tools: Final = [{**_TOOLS[0], "cache_control": {"type": "ephemeral"}}, _TOOLS[1]] + metadata: Final[dict[str, object]] = {} + request: Final = _prompt_request( + tools=tools, + cache_control_injection_points=[{"location": "tool_config"}], + litellm_metadata=metadata, + model_info={"id": "deployment-1"}, + ) + + injected: Final = messages_host.injected_prompt(request) + + assert injected is not None + assert injected["tools"] is tools + assert GATEWAY_INJECTED_CACHE_METADATA_KEY not in metadata diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index 3001d23e57e..22d16c4ae47 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -31,8 +31,6 @@ MAPPED_FILE: Final = { } MAPPED_FILE_INPUT: Final = [{"role": "user", "content": [{"type": "input_file", "file_id": "litellm-file"}]}] CACHED_TEXT: Final = {"type": "input_text", "text": "hi", "cache_control": {"type": "ephemeral"}} -THINKING_TURN: Final = {"role": "assistant", "content": [{"type": "thinking", "thinking": "x", "signature": "s"}]} -REDACTED: Final = {"type": "redacted_thinking", "data": "x"} STRICT_TOOL: Final = {"name": "b", "input_schema": {"type": "object"}, "strict": True} TOOL_SEARCH_TYPE: Final = sorted(ANTHROPIC_TOOL_SEARCH_TOOL_TYPES)[0] @@ -194,21 +192,6 @@ def _wrap_nested_level(value: object, depth: int) -> object: (detectors.cache_control_injection, {"enable_prompt_caching": True}, True), (detectors.cache_control_injection, {"cache_control_injection_points": [{"location": "message"}]}, True), (detectors.cache_control_injection, {"cache_control_injection_points": []}, False), - (messages_detectors.thinking_signature_retry, {"messages": [THINKING_TURN]}, True), - ( - messages_detectors.thinking_signature_retry, - {"messages": [{"role": "assistant", "content": [REDACTED]}]}, - True, - ), - (messages_detectors.thinking_signature_retry, {"messages": [{"role": "assistant", "content": "done"}]}, False), - ( - messages_detectors.thinking_signature_retry, - {"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, @@ -321,8 +304,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"}, {}), @@ -345,11 +326,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_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)