From 78b7b677899bc3cd2d353340f8bb4f89061e2693 Mon Sep 17 00:00:00 2001 From: narmstrong-berry Date: Tue, 10 Nov 2026 20:05:34 +0000 Subject: [PATCH] feat(messages): apply prompt caching injection on rust --- .../crates/inference-messages/src/handler.rs | 2 +- .../crates/inference-messages/src/types.rs | 14 +- .../llms/src/anthropic/messages/handler.rs | 42 ++++- .../src/anthropic/messages/transformation.rs | 9 +- .../src/azure_ai/messages/transformation.rs | 6 +- .../llms/src/base_llm/messages/context.rs | 6 + .../src/base_llm/messages/transformation.rs | 22 ++- .../messages/transformation.rs | 9 +- .../anthropic_claude_transformation.rs | 6 +- .../bedrock_mantle/messages/transformation.rs | 6 +- .../src/deepseek/messages/transformation.rs | 9 +- .../src/minimax/messages/transformation.rs | 9 +- .../src/tencent/messages/transformation.rs | 9 +- .../src/vertex_ai/messages/transformation.rs | 6 +- .../python-bridge/src/routes/messages/host.rs | 28 +++- .../llms/anthropic/messages/native/handler.py | 53 +++++-- litellm/messages/dispatch.py | 1 - litellm/rust_bridge/catalog.py | 1 - litellm/rust_bridge/host/routes/messages.py | 110 +++++++++++-- .../sdk/messages/test_native_parity.py | 146 +++++++++++++++++- .../messages/test_dispatch.py | 1 - .../rust_bridge/host/routes/test_messages.py | 97 +++++++++++- 22 files changed, 505 insertions(+), 87 deletions(-) diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 456dac28c80..46b16f1f79b 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -188,7 +188,7 @@ async fn provider_call( let sanitized = config.shape_request( MessagesRequest { model, ..body }, - shaping.settings.reasoning_auto_summary, + shaping.settings.shape_options(), )?; let transform_context = MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params); diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index b1ac375f7ab..fc7bf0f95cb 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -2,7 +2,10 @@ use std::time::Duration; use bytes::Bytes; use litellm_host::call::CallOutput; -use litellm_llms::{ErrorDetail, base_llm::messages::context::MessagesModelCapabilities}; +use litellm_llms::{ + ErrorDetail, + base_llm::messages::context::{MessagesModelCapabilities, MessagesShapeOptions}, +}; use litellm_llms_types::{ formats::messages::{MessagesRequest, MessagesResponse}, headers::ProviderSpecificHeaders, @@ -75,9 +78,18 @@ pub struct MessagesSettings { pub mock_response: Option>, #[serde(default)] pub whole_sse_frames: bool, + #[serde(default)] + pub presanitized: bool, } impl MessagesSettings { + pub fn shape_options(&self) -> MessagesShapeOptions { + MessagesShapeOptions { + reasoning_auto_summary: self.reasoning_auto_summary, + presanitized: self.presanitized, + } + } + pub fn request_timeout(&self) -> Option { self.request_timeout .and_then(|secs| Duration::try_from_secs_f64(secs).ok()) diff --git a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs index 54f924ddec6..92d0eb59ae6 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/handler.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/handler.rs @@ -13,14 +13,19 @@ use crate::{ flatten_unencrypted_web_search_results, sanitize_tool_use_ids, strip_empty_content_blocks, strip_provider_specific_fields, }, + base_llm::messages::context::MessagesShapeOptions, }; pub fn shape_anthropic_messages_request( request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { Ok(MessagesRequest { - messages: sanitize_anthropic_messages(request.messages), + messages: if options.presanitized { + strip_provider_specific_fields(request.messages) + } else { + sanitize_anthropic_messages(request.messages) + }, params: MessagesOptionalParams { metadata: request .params @@ -28,7 +33,10 @@ pub fn shape_anthropic_messages_request( .as_ref() .map(validate_anthropic_api_metadata) .transpose()?, - thinking: with_reasoning_auto_summary(request.params.thinking, reasoning_auto_summary), + thinking: with_reasoning_auto_summary( + request.params.thinking, + options.reasoning_auto_summary, + ), ..request.params }, ..request @@ -299,7 +307,10 @@ mod tests { "thinking": {"type": "enabled", "budget_tokens": 1024}, "safeguards": [{"type": "dangerous_tool_use"}] })), - true, + MessagesShapeOptions { + reasoning_auto_summary: true, + presanitized: false, + }, ) .unwrap(); assert_eq!( @@ -315,4 +326,27 @@ mod tests { }) ); } + + #[test] + fn presanitized_shaping_keeps_blocks_and_strips_only_provider_fields() { + let messages = json!([{"role": "user", "content": [ + {"type": "text", "text": " ", "cache_control": {"type": "ephemeral"}}, + {"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok", "provider_specific_fields": {"x": 1}} + ]}]); + let sanitized = shape_anthropic_messages_request( + request(json!({"model": "m", "messages": messages})), + MessagesShapeOptions { + reasoning_auto_summary: false, + presanitized: true, + }, + ) + .unwrap(); + assert_eq!( + serde_json::to_value(sanitized).unwrap(), + json!({"model": "m", "messages": [{"role": "user", "content": [ + {"type": "text", "text": " ", "cache_control": {"type": "ephemeral"}}, + {"type": "tool_result", "tool_use_id": "functions.Bash:0", "content": "ok"} + ]}]}) + ); + } } diff --git a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs index bb1f63bc029..2d39f444901 100644 --- a/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/messages/transformation.rs @@ -14,7 +14,10 @@ use serde_json::{Map, Value, json}; use super::{handler::shape_anthropic_messages_request, thinking::translate_thinking}; use crate::anthropic::beta_headers::accepted_or_unlisted; -use crate::base_llm::messages::{context::MessagesTransformContext, tools::ToolCustomField}; +use crate::base_llm::messages::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + tools::ToolCustomField, +}; use crate::{ Error, anthropic::common_utils::{ @@ -41,9 +44,9 @@ impl BaseMessagesConfig for AnthropicMessagesConfig { fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn get_complete_url( 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 037d1cdb722..fb50d229bc7 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -24,7 +24,7 @@ use crate::{ base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, messages::{ - context::MessagesTransformContext, + context::{MessagesShapeOptions, MessagesTransformContext}, normalization::{normalize_system_role_messages, strip_cache_control_scope}, transformation::BaseMessagesConfig, }, @@ -58,9 +58,9 @@ impl BaseMessagesConfig for AzureAnthropicMessagesConfig { fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/base_llm/messages/context.rs b/litellm-rust/crates/llms/src/base_llm/messages/context.rs index 3d8b0f41d8c..229475d7178 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/context.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/context.rs @@ -23,6 +23,12 @@ impl SupportedEffortTiers { } } +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct MessagesShapeOptions { + pub reasoning_auto_summary: bool, + pub presanitized: bool, +} + #[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)] pub struct MessagesModelCapabilities { #[serde(default)] 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 cdb9c0fb8d5..72475e9c316 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/transformation.rs @@ -5,7 +5,10 @@ use litellm_llms_types::formats::messages::{ use litellm_router_types::LitellmParams; use serde_json::Value; -use super::{context::MessagesTransformContext, tools::ToolCustomField}; +use super::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + tools::ToolCustomField, +}; pub use crate::base_llm::auth::{Headers, ValidatedEnvironment}; use crate::{Error, base_llm::messages::streaming::StreamDecoder}; @@ -14,7 +17,7 @@ pub trait BaseMessagesConfig: Sync { fn shape_request( &self, request: MessagesRequest, - _reasoning_auto_summary: bool, + _options: MessagesShapeOptions, ) -> Result { Ok(request) } @@ -168,9 +171,10 @@ mod tests { } #[rstest] - #[case::disabled(false)] - #[case::enabled(true)] - fn default_shaping_preserves_provider_policy_inputs(#[case] reasoning_auto_summary: bool) { + fn default_shaping_preserves_provider_policy_inputs( + #[values(false, true)] reasoning_auto_summary: bool, + #[values(false, true)] presanitized: bool, + ) { let request: MessagesRequest = serde_json::from_value(serde_json::json!({ "model": "test-model", "metadata": {"user_id": 7, "extra": "keep"}, @@ -179,7 +183,13 @@ mod tests { })) .unwrap(); assert_eq!( - DefaultsConfig.shape_request(request.clone(), reasoning_auto_summary), + DefaultsConfig.shape_request( + request.clone(), + MessagesShapeOptions { + reasoning_auto_summary, + presanitized + } + ), Ok(request) ); assert_eq!( 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 aedf72e5615..4bf30b56f00 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 @@ -23,7 +23,10 @@ use crate::{ }, base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, - messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig}, + messages::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + transformation::BaseMessagesConfig, + }, }, bedrock::common_utils::validated_region, }; @@ -71,9 +74,9 @@ impl BaseMessagesConfig for BedrockClaudePlatformMessagesConfig { fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn validate_environment( diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs index 77014d068d1..fa037974342 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs @@ -5,7 +5,7 @@ use std::convert::Infallible; use crate::{ anthropic::messages::handler::shape_anthropic_messages_request, - base_llm::messages::context::MessagesTransformContext, + base_llm::messages::context::{MessagesShapeOptions, MessagesTransformContext}, }; use futures_util::StreamExt; use litellm_auth::AwsParams; @@ -118,9 +118,9 @@ impl BaseMessagesConfig for AmazonAnthropicClaudeMessagesConfig { fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn get_complete_url( 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 58d76662977..0e67ea1a139 100644 --- a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs @@ -29,7 +29,7 @@ use crate::{ base_llm::{ auth::AuthScheme, messages::{ - context::MessagesTransformContext, + context::{MessagesShapeOptions, MessagesTransformContext}, transformation::{BaseMessagesConfig, Headers, ValidatedEnvironment}, }, }, @@ -142,9 +142,9 @@ impl BaseMessagesConfig for BedrockMantleAnthropicMessagesConfig { fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn wire_body(&self, body: Value) -> Value { diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs index 2c4945e41bf..ee2aec9efb9 100644 --- a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -23,7 +23,10 @@ use crate::{ }, base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, - messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig}, + messages::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + transformation::BaseMessagesConfig, + }, }, }; @@ -63,9 +66,9 @@ impl BaseMessagesConfig for DeepSeekAnthropicMessagesConfig { fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn get_complete_url( diff --git a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs index b82442d9294..2eaf182e657 100644 --- a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs @@ -22,7 +22,10 @@ use crate::{ }, base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, - messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig}, + messages::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + transformation::BaseMessagesConfig, + }, }, }; @@ -53,9 +56,9 @@ impl BaseMessagesConfig for MinimaxMessagesConfig { fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn validate_environment( diff --git a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs index 9aef16728e0..b3ee28b7474 100644 --- a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs @@ -22,7 +22,10 @@ use crate::{ }, base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, - messages::{context::MessagesTransformContext, transformation::BaseMessagesConfig}, + messages::{ + context::{MessagesShapeOptions, MessagesTransformContext}, + transformation::BaseMessagesConfig, + }, }, }; @@ -54,9 +57,9 @@ impl BaseMessagesConfig for TencentAnthropicMessagesConfig { fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn validate_environment( 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 18e66c5a77e..d824c5a6a22 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/messages/transformation.rs @@ -26,7 +26,7 @@ use crate::{ base_llm::{ auth::{AuthScheme, Headers, ValidatedEnvironment}, messages::{ - context::MessagesTransformContext, + context::{MessagesShapeOptions, MessagesTransformContext}, normalization::{normalize_system_role_messages, strip_cache_control_scope}, transformation::BaseMessagesConfig, }, @@ -68,9 +68,9 @@ impl BaseMessagesConfig for VertexAiPartnerModelsAnthropicMessagesConfig { fn shape_request( &self, request: MessagesRequest, - reasoning_auto_summary: bool, + options: MessagesShapeOptions, ) -> Result { - shape_anthropic_messages_request(request, reasoning_auto_summary) + shape_anthropic_messages_request(request, options) } fn get_complete_url( 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 0e3a94cd4ef..e32884d2b84 100644 --- a/litellm-rust/crates/python-bridge/src/routes/messages/host.rs +++ b/litellm-rust/crates/python-bridge/src/routes/messages/host.rs @@ -2,7 +2,7 @@ use crate::cache::{CacheCall, Cached, PythonCache, Selection}; use litellm_host_python::{PythonHostCalls, PythonOwned}; use bytes::Bytes; -use litellm_host_python::{InvokeError, PythonBinding, from_py, to_py}; +use litellm_host_python::{InvokeError, PythonBinding, effective_py_args, from_py, to_py}; use litellm_inference_messages::{ Error, MessagesCall, MessagesSettings, MessagesShaping, litellm_params, messages_body, route::{Messages, MessagesStreamHead}, @@ -82,6 +82,7 @@ pub(super) struct MessagesPythonHost { view: RequestView, cache: PythonCache, stream_cost_model: Option, + asynchronous: bool, } impl MessagesPythonHost { @@ -90,6 +91,7 @@ impl MessagesPythonHost { view, cache: PythonCache::new(asynchronous), stream_cost_model: None, + asynchronous, } } @@ -105,6 +107,19 @@ impl MessagesPythonHost { ) } + fn body_params( + &self, + py: Python<'_>, + arguments: &Bound<'_, PyDict>, + ) -> PyResult> { + let request = effective_py_args(self.view.request(py)?, arguments)?; + from_py( + &py.import(ROUTE_HOST_MODULE)? + .getattr("body_params")? + .call1((request, self.asynchronous))?, + ) + } + fn projection( &self, py: Python<'_>, @@ -117,10 +132,16 @@ impl MessagesPythonHost { }; let model = settings.model; let messages = argument("messages")?.ok_or(ArgumentError::Required("messages"))?; - let fields = self.view.body_params(py, arguments, ROUTE_HOST_MODULE)?; + let mut fields = self.body_params(py, arguments)?; + let prepared = fields.remove("messages"); + let presanitized = prepared.is_some(); + let messages = match prepared { + Some(messages) => messages, + None => from_py(&messages)?, + }; let body = [ ("model".to_string(), Value::String(model.clone())), - ("messages".to_string(), from_py(&messages)?), + ("messages".to_string(), messages), ] .into_iter() .chain(fields) @@ -145,6 +166,7 @@ impl MessagesPythonHost { capabilities: self.capabilities(py, &model, custom_llm_provider.as_deref())?, settings: MessagesSettings { whole_sse_frames: settings.stream_cost_model.is_some(), + presanitized, ..settings.settings }, }; diff --git a/litellm/llms/anthropic/messages/native/handler.py b/litellm/llms/anthropic/messages/native/handler.py index 59ffb6fef04..6a169137c24 100644 --- a/litellm/llms/anthropic/messages/native/handler.py +++ b/litellm/llms/anthropic/messages/native/handler.py @@ -458,6 +458,40 @@ def _short_circuit_safeguards_evaluator( return request_safeguards_evaluator(kwargs, messages) +def provider_after_pre_request_hooks( + model: str, custom_llm_provider: str | None, request_kwargs: Mapping[str, object] +) -> str | None: + if custom_llm_provider: + return custom_llm_provider + hook_params: Final = request_kwargs.get("litellm_params") + hooked: Final = ( + cast(Mapping[str, object], hook_params).get("custom_llm_provider") # cast-ok: the hooks' own params dict + if isinstance(hook_params, Mapping) + else None + ) + derived: Final = hooked if isinstance(hooked, str) else None + if derived: + return derived + try: + _, resolved, _, _ = litellm.get_llm_provider(model=model) + except Exception: # noqa: BLE001 # an unknown model keeps the provider the hooks left + return derived + return resolved + + +def tools_marked_for_cache_points( + tools: object, + request_kwargs: dict[str, object], # mutable-ok: pops the points and records the mark in the handler's kwargs +) -> Sequence[object] | None: + marked: Final = tools_marked_for_a_tool_config_point( + tools, request_kwargs.pop("cache_control_injection_points", None) + ) + if marked is None: + return None + AnthropicCacheControlHook.record_gateway_injection(request_kwargs, 1) + return list(marked) + + @client async def anthropic_messages( max_tokens: int, @@ -518,16 +552,7 @@ async def anthropic_messages( tool_choice = request_kwargs.pop("tool_choice", tool_choice) top_k = request_kwargs.pop("top_k", top_k) top_p = request_kwargs.pop("top_p", top_p) - # Propagate the provider derived inside pre-request hooks, if not already set. - # The litellm_params dict may have been overwritten by **kwargs in - # _execute_pre_request_hooks, so fall back to get_llm_provider() if needed. - if not custom_llm_provider: - custom_llm_provider = request_kwargs.get("litellm_params", {}).get("custom_llm_provider") - if not custom_llm_provider: - try: - _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=model) - except Exception: - pass + custom_llm_provider = provider_after_pre_request_hooks(model, custom_llm_provider, request_kwargs) # Remove litellm_params from kwargs (only needed for hooks) request_kwargs.pop("litellm_params", None) # Merge back any other modifications @@ -834,12 +859,8 @@ def anthropic_messages_handler( ) request_kwargs: Final = cast(dict[str, object], kwargs) # cast-ok: **kwargs of the handler is untyped - marked_tools: Final = tools_marked_for_a_tool_config_point( - local_vars.get("tools"), request_kwargs.pop("cache_control_injection_points", None) - ) - if marked_tools is not None: - AnthropicCacheControlHook.record_gateway_injection(request_kwargs, 1) - local_vars.update({**request_kwargs, **({"tools": list(marked_tools)} if marked_tools is not None else {})}) + marked_tools: Final = tools_marked_for_cache_points(local_vars.get("tools"), request_kwargs) + local_vars.update({**request_kwargs, **({"tools": marked_tools} if marked_tools is not None else {})}) anthropic_messages_optional_request_params: Final = ( AnthropicMessagesRequestUtils.get_requested_anthropic_messages_optional_param( params=local_vars, diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index 9fac2b70078..9ac52affddf 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -141,7 +141,6 @@ _FEATURES: Final[detectors.Detectors] = MappingProxyType( Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, Feature.MODIFY_PARAMS: detectors.modify_params, Feature.WORKLOAD_IDENTITY: route_detectors.anthropic_workload_identity, - Feature.CACHE_CONTROL_INJECTION: detectors.cache_control_injection, } ) diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index b7cdfaadc9e..c8d3f00166b 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -171,7 +171,6 @@ MESSAGES_PYTHON_FEATURES: Final = (COMMON_PYTHON_FEATURES - {Feature.TIMEOUT_VAR Feature.MESSAGES_INTERCEPTOR, Feature.MODIFY_PARAMS, Feature.WORKLOAD_IDENTITY, - Feature.CACHE_CONTROL_INJECTION, Feature.STRICT_TOOLS_RETRY, } diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index 5385f8f301e..2d4b7ae8eab 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -12,7 +12,7 @@ from typing import Final, TypeAlias, cast # noqa: TID251 # narrows native payl import httpx from aiohttp.http_exceptions import ContentLengthError, TransferEncodingError -from pydantic import TypeAdapter, ValidationError +from pydantic import JsonValue, TypeAdapter, ValidationError import litellm from litellm.litellm_core_utils import exception_mapping_utils @@ -41,15 +41,16 @@ 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.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import get_ssl_configuration from litellm.llms.custom_httpx.llm_http_handler import ( _merge_messages_header_sources as _merged_headers, # pyright: ignore[reportPrivateUsage] # the handler's own header merge ) -from litellm.rust_bridge import failures +from litellm.rust_bridge import detectors, failures from litellm.rust_bridge.host import callbacks_legacy_python, http_transport -from litellm.rust_bridge.public_call import known_params, optional_str +from litellm.rust_bridge.public_call import NativeCall, known_params, optional_str from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ProviderSpecificHeader @@ -113,13 +114,26 @@ def inferred_provider(model: str, api_base: str | None) -> str | None: return provider -def body_params(request: Mapping[str, object]) -> Mapping[str, object]: +def _body_fields(request: Mapping[str, object]) -> Mapping[str, object]: params: Final = known_params(_BODY_FIELDS.difference(handler.additional_drop_paths(request)), request) if deployment_reasoning_effort_levels(request.get("model_info")) is None: return params return _with_deployment_effort(params, request) +def body_params(request: Mapping[str, object], asynchronous: bool) -> Mapping[str, object]: + fields: Final = _body_fields(request) + if not detectors.cache_control_injection(NativeCall(args=(), kwargs=request, base={})): + return fields + messages, system, tools = _prompt_cached(request, asynchronous) + return { + **fields, + "messages": messages, + **({"system": system} if "system" in fields else {}), + **({"tools": tools} if tools is not None and "tools" in fields else {}), + } + + def _with_deployment_effort(params: Mapping[str, object], request: Mapping[str, object]) -> Mapping[str, object]: """Python's handler rounds the requested effort to the levels the deployment declares in `model_info.reasoning_effort_levels`, the one place a deployment, not the model, decides.""" @@ -146,6 +160,81 @@ def cached_function() -> Callable[..., object]: ) +_ENTRY: Final = cast( # cast-ok: only the untyped handler's signature is read + Callable[..., object], + handler.anthropic_messages, # noqa: TID251 # reads the handler's signature, never calls it +) +_HANDLER: Final = cast( # cast-ok: only the untyped handler's signature is read + Callable[..., object], + handler.anthropic_messages_handler, # noqa: TID251 # reads the handler's signature, never calls it +) +_ENTRY_PARAMS: Final = frozenset(inspect.signature(_ENTRY).parameters) - {"kwargs"} +_HANDLER_PARAMS: Final = frozenset(inspect.signature(_HANDLER).parameters) - {"kwargs"} +_JsonBlocks: TypeAlias = list[dict[str, JsonValue]] +_Prepared: TypeAlias = tuple[_JsonBlocks, str | _JsonBlocks | None, Sequence[object] | None] + + +def _prompt_cached(request: Mapping[str, object], asynchronous: bool) -> _Prepared: + model: Final = str(request["model"]) + provider: Final = optional_str(request.get("custom_llm_provider")) + api_base: Final = optional_str(request.get("api_base")) + tools: Final = cast(_JsonBlocks | None, request.get("tools")) # cast-ok: handed on unchecked, as the handler does + messages: Final = cast(_JsonBlocks, request["messages"]) # cast-ok: handed on unchecked, as the handler does + system: Final = cast( # cast-ok: handed on unchecked, as the handler does + str | _JsonBlocks | None, request.get("system") + ) + if not asynchronous: + return _handler_prepared( + request, + messages, + system, + model=model, + provider=provider, + tools=tools, + api_base=api_base, + presanitized=False, + ) + entry_kwargs: Final = {name: value for name, value in request.items() if name not in _ENTRY_PARAMS} + entry_messages, entry_system = prepare_native_messages( + messages, system, entry_kwargs, model=model, custom_llm_provider=provider, tools=tools, api_base=api_base + ) + return _handler_prepared( + entry_kwargs, + entry_messages, + entry_system, + model=model, + provider=handler.provider_after_pre_request_hooks(model, provider, entry_kwargs), + tools=tools, + api_base=api_base, + presanitized=True, + ) + + +def _handler_prepared( + call: Mapping[str, object], + messages: _JsonBlocks, + system: str | _JsonBlocks | None, + *, + model: str, + provider: str | None, + tools: _JsonBlocks | None, + api_base: str | None, + presanitized: bool, +) -> _Prepared: + handler_kwargs: Final = {name: value for name, value in call.items() if name not in _HANDLER_PARAMS} + prepared_messages, prepared_system = prepare_native_messages( + messages, + system, + handler_kwargs, + model=model, + custom_llm_provider=provider, + tools=tools, + api_base=api_base, + presanitized=presanitized, + ) + return prepared_messages, prepared_system, handler.tools_marked_for_cache_points(tools, handler_kwargs) + + 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, @@ -364,8 +453,8 @@ def _log_resolved_call(request: Mapping[str, object], request_provider: str) -> logger, kwargs={name: value for name, value in request.items() if name != "metadata"}, model=str(request["model"]).removeprefix(f"{request_provider}/"), - optional_params=dict(body_params(request)), - litellm_params={"model_info": request.get("model_info"), **body_params(request)}, + optional_params=dict(_body_fields(request)), + litellm_params={"model_info": request.get("model_info"), **_body_fields(request)}, custom_llm_provider=request_provider, ) @@ -428,15 +517,6 @@ def settings(kwargs: Mapping[str, object], request: Mapping[str, object] | None return asdict(configured if request is None else _with_call_settings(configured, request)) -_HANDLER_PARAMS: Final = frozenset( - name - for name, parameter in inspect.signature( - cast(Callable[..., object], handler.anthropic_messages_handler) # noqa: TID251 # cast-ok: reads the handler's parameter names, never calls it - ).parameters.items() - if parameter.kind is not inspect.Parameter.VAR_KEYWORD -) - - def _handler_params(request: Mapping[str, object]) -> GenericLiteLLMParams: return GenericLiteLLMParams( **{name: value for name, value in request.items() if name not in _HANDLER_PARAMS}, 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 af53e2069af..02b29e81c38 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 @@ -349,7 +349,6 @@ 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_CACHE_INJECTION: Final = "F: prompt cache_control injection is Python-only (detector cache_control_injection)" F_STREAMED_MOCK: Final = ( "F: a mock with stream=True returns the mock dict and logs no success in Python (detector streamed_mock_response)" ) @@ -358,8 +357,6 @@ 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), - "cache_control_injection_points": _gaps(request=F_CACHE_INJECTION, logged_request=F_CACHE_INJECTION), - "enable_prompt_caching": _gaps(request=F_CACHE_INJECTION, logged_request=F_CACHE_INJECTION), "mock_response_streamed": _gaps( logged=F_STREAMED_MOCK, logged_metadata=F_STREAMED_MOCK, @@ -713,6 +710,142 @@ SYNC_CASES: Final = ( Case("sync_upstream_429", upstream=(_failure(429, "rate_limit_error", "slow down"),), sync=True), ) +SYSTEM_POINT: Final[Mapping[str, JsonValue]] = MappingProxyType({"location": "message", "role": "system"}) +LAST_MESSAGE_POINT: Final[Mapping[str, JsonValue]] = MappingProxyType({"location": "message", "index": -1}) +SYSTEM_BLOCKS: Final[tuple[JsonValue, ...]] = ( + {"type": "text", "text": "Answer in one word."}, + {"type": "text", "text": "Be terse."}, +) +MARKED_TURNS: Final[tuple[JsonValue, ...]] = ( + {"role": "user", "content": [{"type": "text", "text": "One.", "cache_control": {"type": "ephemeral"}}]}, + {"role": "assistant", "content": [{"type": "text", "text": "Two.", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": [{"type": "text", "text": "Three.", "cache_control": {"type": "ephemeral"}}]}, +) + +PROMPT_CACHING_CASES: Final = ( + Case( + "cache_control_injection_points", + _kw(system="Answer in one word.", cache_control_injection_points=[dict(SYSTEM_POINT)]), + ), + Case("enable_prompt_caching", _kw(system="Answer in one word.", enable_prompt_caching=True)), + Case( + "enable_prompt_caching_without_system", + _kw(messages=list(TOOL_TURNS), tools=[dict(WEATHER_TOOL)], enable_prompt_caching=True), + ), + Case( + "global_prompt_caching_with_ttl", + _kw(system=list(SYSTEM_BLOCKS)), + settings=MappingProxyType({"enable_anthropic_prompt_caching": True, "anthropic_prompt_caching_ttl": "1h"}), + ), + Case( + "global_prompt_caching_stands_down_on_client_marks", + _kw(system=[{"type": "text", "text": "Answer in one word.", "cache_control": {"type": "ephemeral"}}]), + settings=MappingProxyType({"enable_anthropic_prompt_caching": True}), + ), + Case( + "global_prompt_caching_stands_down_on_marked_tools", + _kw(tools=[{**WEATHER_TOOL, "cache_control": {"type": "ephemeral"}}]), + settings=MappingProxyType({"enable_anthropic_prompt_caching": True}), + ), + Case( + "enable_prompt_caching_false_overrides_nothing", + _kw(system="Answer in one word.", enable_prompt_caching=False), + ), + Case( + "injection_point_negative_index", + _kw( + messages=list(TOOL_TURNS), + cache_control_injection_points=[ + {"location": "message", "index": -2, "control": {"type": "ephemeral", "ttl": "1h"}} + ], + ), + ), + Case( + "injection_point_string_index", + _kw(messages=list(TOOL_TURNS), cache_control_injection_points=[{"location": "message", "index": "1"}]), + ), + Case( + "injection_point_out_of_range_index", + _kw(messages=list(TOOL_TURNS), cache_control_injection_points=[{"location": "message", "index": 7}]), + ), + Case( + "injection_point_by_role", + _kw(messages=list(TOOL_TURNS), cache_control_injection_points=[{"location": "message", "role": "user"}]), + ), + Case( + "injection_points_capped_by_client_marks", + _kw( + messages=list(MARKED_TURNS), + system=list(SYSTEM_BLOCKS), + tools=[{**WEATHER_TOOL, "cache_control": {"type": "ephemeral"}}], + cache_control_injection_points=[dict(SYSTEM_POINT), dict(LAST_MESSAGE_POINT)], + ), + ), + Case( + "injection_points_beside_client_marks", + _kw( + messages=[*MARKED_TURNS[:2], {"role": "user", "content": "Three."}], + cache_control_injection_points=[dict(SYSTEM_POINT), dict(LAST_MESSAGE_POINT)], + system="Answer in one word.", + ), + ), + Case( + "tool_config_point_marks_the_last_tool", + _kw( + tools=[dict(WEATHER_TOOL), {**WEATHER_TOOL, "name": "get_time"}], + cache_control_injection_points=[{"location": "tool_config"}, dict(LAST_MESSAGE_POINT)], + ), + ), + Case( + "tool_config_point_without_tools", + _kw(cache_control_injection_points=[{"location": "tool_config"}]), + ), + Case( + "injection_marks_the_gateway_deployment", + _kw( + system="Answer in one word.", + enable_prompt_caching=True, + litellm_metadata={"user_api_key": "hashed"}, + model_info={"id": "deployment-1"}, + ), + ), + Case( + "injection_into_whitespace_only_turn", + _kw( + messages=[dict(USER), {"role": "assistant", "content": "Paris"}, {"role": "user", "content": " "}], + cache_control_injection_points=[dict(LAST_MESSAGE_POINT)], + ), + upstream=(_failure(400, "invalid_request_error", "messages: text content blocks must be non-empty"),), + ), + Case( + "injection_after_empty_blocks_are_stripped", + _kw( + messages=[ + dict(USER), + {"role": "assistant", "content": [{"type": "text", "text": ""}]}, + {"role": "user", "content": "And Germany?"}, + ], + cache_control_injection_points=[{"location": "message", "index": 1}], + ), + ), + Case( + "stream_with_prompt_caching", + _kw(stream=True, system="Answer in one word.", enable_prompt_caching=True), + upstream=(_stream(TEXT_STREAM),), + ), + Case( + "global_prompt_caching_beside_an_unmatched_point", + _kw( + system="Answer in one word.", cache_control_injection_points=[{"location": "message", "role": "assistant"}] + ), + settings=MappingProxyType({"enable_anthropic_prompt_caching": True}), + ), +) + +SYNC_PROMPT_CACHING_CASES: Final = tuple( + replace(case, name=f"sync_{case.name}", sync=True) for case in PROMPT_CACHING_CASES +) + LATEST_ALIAS: Final = "anthropic/claude-opus-latest" CUSTOM_TOOL: Final[Mapping[str, JsonValue]] = MappingProxyType({**WEATHER_TOOL, "custom": {"defer_loading": True}}) COST_HEADER: Final[Mapping[str, JsonValue]] = MappingProxyType( @@ -770,10 +903,6 @@ SERVED_NATIVELY_CASES: Final = ( FALLBACK_CASES: Final = ( Case("mock_response_streamed", _kw(stream=True, mock_response="Hi from the mock"), upstream=()), - Case( - "cache_control_injection_points", - _kw(system="Answer in one word.", cache_control_injection_points=[{"location": "message", "role": "system"}]), - ), Case( "strict_tools_rejected_then_retried", _kw(tools=[{**WEATHER_TOOL, "strict": True}], drop_params=True), @@ -783,7 +912,6 @@ FALLBACK_CASES: Final = ( ), ), Case("modify_params", _kw(max_tokens=1), settings=MappingProxyType({"modify_params": True})), - Case("enable_prompt_caching", _kw(system="Answer in one word.", enable_prompt_caching=True)), Case("request_timeout_kwarg", _kw(request_timeout=30)), Case("stream_timeout_kwarg", _kw(stream=True, stream_timeout=30), upstream=(_stream(TEXT_STREAM),)), ) @@ -898,6 +1026,8 @@ CACHE_CASES: Final = ( CASES: Final = ( *REQUEST_CASES, + *PROMPT_CACHING_CASES, + *SYNC_PROMPT_CACHING_CASES, *RESPONSE_CASES, *FAILURE_CASES, *STREAM_CASES, diff --git a/tests/test_litellm_rust/messages/test_dispatch.py b/tests/test_litellm_rust/messages/test_dispatch.py index ac803294ab7..40e141673db 100644 --- a/tests/test_litellm_rust/messages/test_dispatch.py +++ b/tests/test_litellm_rust/messages/test_dispatch.py @@ -484,7 +484,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}, {}), ), ) 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 075ef7c5d6a..a99601a4bae 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -175,7 +175,7 @@ def test_body_params_round_effort_to_the_levels_the_deployment_declares( "model_info": {"id": "deployment-1", "reasoning_effort_levels": levels}, } - assert messages_host.body_params(request)["output_config"] == {"effort": expected} + assert messages_host.body_params(request, asynchronous=True)["output_config"] == {"effort": expected} def test_body_params_keep_the_requested_effort_without_deployment_levels() -> None: @@ -186,7 +186,7 @@ def test_body_params_keep_the_requested_effort_without_deployment_levels() -> No "output_config": {"effort": "max"}, } - assert messages_host.body_params(request)["output_config"] == {"effort": "max"} + assert messages_host.body_params(request, asynchronous=True)["output_config"] == {"effort": "max"} @pytest.mark.parametrize( @@ -208,7 +208,7 @@ def test_body_params_derive_no_effort_from_a_param_the_deployment_lists( "model_info": {"id": "deployment-1", "reasoning_effort_levels": ["low", "medium", "high"]}, } - assert set(messages_host.body_params(request)).difference({"messages", "model"}) == expected_keys + assert set(messages_host.body_params(request, asynchronous=True)).difference({"messages", "model"}) == expected_keys def _marked(phase: str | None = None, *args: object) -> ValueError: @@ -534,3 +534,94 @@ def test_settings_infer_no_provider_off_an_affinity_host(monkeypatch: pytest.Mon assert settings["model"] == "claude-sonnet-4-5" assert settings["extra_headers"] is None + + +def test_body_params_leave_the_messages_to_the_route_for_a_call_without_prompt_caching( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + + fields: Final = messages_host.body_params( + { + "model": "anthropic/claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "system": "Be terse.", + "enable_prompt_caching": False, + "cache_control_injection_points": [], + }, + asynchronous=True, + ) + + assert fields == {"system": "Be terse."} + + +@pytest.mark.parametrize("asynchronous", [True, False]) +def test_body_params_mark_the_turns_tools_and_deployment_the_handler_marks( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool +) -> None: + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + gateway: Final[dict[str, object]] = {} + caller_metadata: Final = {"user_id": "u"} + tool: Final = {"name": "get_weather", "input_schema": {"type": "object"}} + + fields: Final = messages_host.body_params( + { + "model": "anthropic/claude-sonnet-4-5", + "messages": [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": [{"type": "text", "text": ""}]}, + {"role": "user", "content": "again"}, + ], + "system": "Be terse.", + "tools": [tool], + "metadata": caller_metadata, + "litellm_metadata": gateway, + "model_info": {"id": "deployment-1"}, + "cache_control_injection_points": [ + {"location": "message", "role": "system"}, + {"location": "message", "index": -1, "control": {"type": "ephemeral", "ttl": "1h"}}, + {"location": "tool_config"}, + ], + }, + asynchronous=asynchronous, + ) + + assert fields == { + "messages": [ + {"role": "user", "content": [{"type": "text", "text": "hi"}]}, + { + "role": "user", + "content": [{"type": "text", "text": "again", "cache_control": {"type": "ephemeral", "ttl": "1h"}}], + }, + ], + "system": [{"type": "text", "text": "Be terse.", "cache_control": {"type": "ephemeral"}}], + "tools": [{**tool, "cache_control": {"type": "ephemeral"}}], + "metadata": caller_metadata, + } + assert gateway == {"litellm_gateway_injected_cache": "deployment-1"} + assert caller_metadata == {"user_id": "u"} + + +@pytest.mark.parametrize( + ("asynchronous", "expected_system"), + [ + (True, [{"type": "text", "text": "Be terse.", "cache_control": {"type": "ephemeral"}}]), + (False, "Be terse."), + ], +) +def test_body_params_prepare_prompt_caching_as_the_entry_that_ran( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool, expected_system: object +) -> None: + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + + fields: Final = messages_host.body_params( + { + "model": "anthropic/claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hi"}], + "system": "Be terse.", + "cache_control_injection_points": [{"location": "message", "role": "assistant"}], + }, + asynchronous=asynchronous, + ) + + assert fields["system"] == expected_system