diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index e62a9463cb5..24d8cfcfe54 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -8,6 +8,13 @@ For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all` +## Python ports + +- Prefer mirroring the outer Python structure. Map `litellm/.py` to `crates//src/.rs` in the crate that mirrors that folder +- Keep class and public function names, with Rust acronym casing and no leading underscore for private helpers. A staticmethod-only class becomes free functions. A class with a Python base becomes a struct implementing the matching trait +- Inside functions, design freely. Prefer typed values over dict inspection, new values over mutation, and tagged unions over string modes. Pin behavior with Rust tests +- List skipped behavior under `Known partial ports` in the nearest AGENTS.md and test deliberate divergences + ## Test placement - Never create a `tests.rs` (or `test.rs`) file under `src/`, and never `#[path = "tests.rs"] mod tests;` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 4ddce5edc78..f6eb48b4bc8 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4035,6 +4035,7 @@ dependencies = [ "litellm-http", "litellm-inference", "litellm-inference-testing", + "litellm-integrations", "litellm-llms", "litellm-llms-types", "litellm-router-types", @@ -4141,6 +4142,16 @@ dependencies = [ "wiremock", ] +[[package]] +name = "litellm-integrations" +version = "0.1.0" +dependencies = [ + "litellm-llms-types", + "rstest", + "serde", + "serde_json", +] + [[package]] name = "litellm-llms" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index f04eb115071..752fac70cc5 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -20,6 +20,7 @@ litellm-inference-transcription = { path = "crates/inference-transcription" } litellm-inference-responses = { path = "crates/inference-responses" } litellm-inference-messages = { path = "crates/inference-messages" } litellm-inference-chat = { path = "crates/inference-chat" } +litellm-integrations = { path = "crates/integrations" } litellm-inference-ocr = { path = "crates/inference-ocr" } litellm-inference-testing = { path = "crates/inference-testing" } litellm-gateway-mcp = { path = "crates/gateway-mcp" } diff --git a/litellm-rust/crates/callbacks-legacy-python/python_contract.json b/litellm-rust/crates/callbacks-legacy-python/python_contract.json index 4cb49da396a..f7b48803f29 100644 --- a/litellm-rust/crates/callbacks-legacy-python/python_contract.json +++ b/litellm-rust/crates/callbacks-legacy-python/python_contract.json @@ -21,6 +21,10 @@ "litellm_params", "custom_llm_provider" ], + "record_gateway_injection": [ + "kwargs", + "added" + ], "pre_call": [ "logger", "input", diff --git a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs index 7cc683b50bc..c6f5ac88c4c 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -374,6 +374,12 @@ impl LegacyLogging { wire: Box, context: &RequestContext, ) -> PyResult>> { + if context.gateway_injected_cache_blocks > 0 { + crate::python::Logging::RecordGatewayInjection.call( + py, + (self.call.kwargs(), context.gateway_injected_cache_blocks), + )?; + } let logger = self.logger()?; logger.update_from_kwargs(py, self.call.kwargs(), &wire, context)?; let body = to_py(py, &wire.body)? @@ -722,7 +728,7 @@ if asynchronous: }); } - #[rstest::rstest] + #[rstest] fn kwargs_returned_by_the_pre_call_hook_are_what_the_call_prepares() { Python::initialize(); Python::attach(|py| { @@ -1089,7 +1095,32 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update body: Value, secret_fields: &[&str], ) -> WireRequest { - before_send_bound(&[], script, optional_params, body, secret_fields) + before_send_bound(&[], script, optional_params, body, secret_fields, 0) + } + + #[rstest] + #[case::none_injected(0, None)] + #[case::gateway_injected(2, Some(2))] + fn gateway_cache_injection_is_recorded_before_pre_call( + #[case] injected: usize, + #[case] expected: Option, + ) { + let expected = serde_json::to_value(expected).unwrap(); + before_send_bound( + &[("expected", &expected)], + c" +metadata = {} +kwargs = {'litellm_metadata': metadata} +observed = [] +on_pre_call = lambda additional_args: observed.append(metadata.get('gateway_injection')) +def check(): + assert observed == [expected], observed +", + json!({}), + json!({}), + &[], + injected, + ); } /// [`before_send_with_secrets`] with `bindings` placed in the namespace before `script` runs. @@ -1099,6 +1130,7 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update optional_params: Value, body: Value, secret_fields: &[&str], + gateway_injected_cache_blocks: usize, ) -> WireRequest { Python::initialize(); Python::attach(|py| { @@ -1117,6 +1149,7 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update optional_params, secret_fields: secret_fields.iter().map(|name| name.to_string()).collect(), api_key: Some(SecretValue::new("route-key")), + gateway_injected_cache_blocks, }; let wire = WireRequest { url: "https://provider.invalid/ocr".into(), @@ -1170,6 +1203,7 @@ assert logger.update['custom_llm_provider'] == 'edenai', logger.update optional_params: json!({}), secret_fields: vec![], api_key: Some(SecretValue::new("route-key")), + gateway_injected_cache_blocks: 0, } } @@ -1725,6 +1759,7 @@ def check(): json!({}), Value::Object(body.clone()), &[], + 0, ); prop_assert_eq!(wire.body, edit.sent(&body)); diff --git a/litellm-rust/crates/callbacks-legacy-python/src/python.rs b/litellm-rust/crates/callbacks-legacy-python/src/python.rs index 578ba39b12e..8134d9e1904 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/python.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/python.rs @@ -40,6 +40,8 @@ pub(crate) enum Logging { CustomPricingFields, #[strum(serialize = "update_logging")] Update, + #[strum(serialize = "record_gateway_injection")] + RecordGatewayInjection, #[strum(serialize = "pre_call")] PreCall, #[strum(serialize = "post_call")] diff --git a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs index b85c78d9022..6b42ffd3502 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/test_support.rs @@ -61,6 +61,7 @@ FAKES = { litellm_params=litellm_params, custom_llm_provider=provider, ), + 'record_gateway_injection': lambda kwargs, added: kwargs['litellm_metadata'].update(gateway_injection=added), 'pre_call': lambda logger, input, api_key, additional_args: logger.pre_call(input, api_key, additional_args), 'post_call': lambda logger, original_response, api_key, additional_args: logger.post_call( original_response, api_key, additional_args diff --git a/litellm-rust/crates/host-http/tests/serve.rs b/litellm-rust/crates/host-http/tests/serve.rs index 9a07a6814c9..21a0046ab53 100644 --- a/litellm-rust/crates/host-http/tests/serve.rs +++ b/litellm-rust/crates/host-http/tests/serve.rs @@ -230,6 +230,7 @@ async fn projection_custom_operations_and_hooks_feed_the_http_response(intercept optional_params: json!({}), secret_fields: Vec::new(), api_key: None, + gateway_injected_cache_blocks: 0, }, ) .await?; @@ -538,6 +539,7 @@ async fn hook_rejection_stops_execution_and_is_reported_once( optional_params: json!({}), secret_fields: Vec::new(), api_key: None, + gateway_injected_cache_blocks: 0, }, ) .await?; @@ -652,6 +654,7 @@ async fn unary_calls_use_into_response_after_hooks_and_before_success(intercepto optional_params: json!({}), secret_fields: Vec::new(), api_key: None, + gateway_injected_cache_blocks: 0, }, ) .await?; diff --git a/litellm-rust/crates/host-native/tests/driver.rs b/litellm-rust/crates/host-native/tests/driver.rs index 40b354a692f..ffe497ebdb7 100644 --- a/litellm-rust/crates/host-native/tests/driver.rs +++ b/litellm-rust/crates/host-native/tests/driver.rs @@ -185,6 +185,7 @@ fn dispatching_call() -> TestMachine { optional_params: json!({}), secret_fields: Vec::new(), api_key: None, + gateway_injected_cache_blocks: 0, }, ) .await?; diff --git a/litellm-rust/crates/host-python/tests/hook_chain.rs b/litellm-rust/crates/host-python/tests/hook_chain.rs index 4b520ca2665..a7b8583aab0 100644 --- a/litellm-rust/crates/host-python/tests/hook_chain.rs +++ b/litellm-rust/crates/host-python/tests/hook_chain.rs @@ -284,6 +284,7 @@ fn transformations_feed_each_other_and_notifications_share_final_values( optional_params: serde_json::Value::Null, secret_fields: vec![], api_key: None, + gateway_injected_cache_blocks: 0, }; let wire = Box::new(WireRequest { url: "url".into(), diff --git a/litellm-rust/crates/host-python/tests/support/mod.rs b/litellm-rust/crates/host-python/tests/support/mod.rs index 3e82bcf394f..6b4e382d52d 100644 --- a/litellm-rust/crates/host-python/tests/support/mod.rs +++ b/litellm-rust/crates/host-python/tests/support/mod.rs @@ -125,6 +125,7 @@ pub fn context() -> RequestContext { optional_params: serde_json::json!({}), secret_fields: Vec::new(), api_key: None, + gateway_injected_cache_blocks: 0, } } diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index dc0577ce110..2dab8b442d3 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -30,6 +30,7 @@ pub struct RequestContext { pub secret_fields: Vec, /// The credential the route resolved for the provider call. pub api_key: Option, + pub gateway_injected_cache_blocks: usize, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -217,6 +218,7 @@ mod tests { optional_params: json!({}), secret_fields: Vec::new(), api_key: None, + gateway_injected_cache_blocks: 0, } } diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index a479f856351..6aaf2081344 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -42,6 +42,7 @@ pub(super) async fn execute( optional_params: Value::Object(optional_params), secret_fields: Vec::new(), api_key, + gateway_injected_cache_blocks: 0, }; let authenticated = resolve_auth(auth, environment, &|key| secrets.get(key)).await?; let identity = litellm_host::interceptors::ProviderIdentity { diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index e310a0be9b1..df1d8927a0b 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-integrations.workspace = true bytes.workspace = true futures-util.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 522dd57a9f1..d34ee5d78ae 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -51,6 +51,7 @@ impl MessagesRoute { environment, timeout, api_key, + gateway_injected_cache_blocks, } = request; let request_context = RequestContext { model: body.model.clone(), @@ -58,6 +59,7 @@ impl MessagesRoute { optional_params: serde_json::to_value(&body.params).map_err(serialize_failure)?, secret_fields: Vec::new(), api_key, + gateway_injected_cache_blocks, }; let authenticated = resolve_auth(&self.auth, environment, &|key| std::env::var(key).ok()).await?; diff --git a/litellm-rust/crates/inference-messages/src/prepare.rs b/litellm-rust/crates/inference-messages/src/prepare.rs index 4b7873b4808..482869daf02 100644 --- a/litellm-rust/crates/inference-messages/src/prepare.rs +++ b/litellm-rust/crates/inference-messages/src/prepare.rs @@ -6,6 +6,7 @@ use litellm_core_utils::{ get_provider_specific_headers::get_provider_specific_headers, settings::Lookup, }; use litellm_http::request::with_default_headers; +use litellm_integrations::{InjectionInputs, maybe_inject_cache_control}; use litellm_llms::base_llm::{ auth::ValidatedEnvironment, messages::context::MessagesTransformContext, }; @@ -32,6 +33,7 @@ pub(super) struct ProviderMessagesRequest { pub(super) timeout: Option, /// The caller's own credential, reported to the host beside the wire request. pub(super) api_key: Option, + pub(super) gateway_injected_cache_blocks: usize, } #[tracing::instrument(name = "litellm.prepare", level = "debug", skip_all)] @@ -84,11 +86,36 @@ fn prepare_provider_request( let env_lookup = |key: &str| secrets.get(key); let sanitized = config.shape_request( - MessagesRequest { model, ..body }, + MessagesRequest { + model: model.clone(), + ..body + }, shaping.settings.reasoning_auto_summary, )?; + let one_shot_subagent = + litellm_llms::anthropic::common_utils::is_claude_code_one_shot_subagent_request( + &sanitized, + shaping.settings.user_agent.as_deref(), + ); + let supports_anthropic_cache_control = + litellm_llms::anthropic::common_utils::supports_anthropic_cache_control( + &model, + provider.as_str(), + shaping.capabilities.supports_prompt_caching, + ); + let injected = maybe_inject_cache_control( + sanitized, + &InjectionInputs { + points: &shaping.settings.cache_control_injection_points, + enable_prompt_caching: shaping.settings.enable_prompt_caching, + supports_anthropic_cache_control, + one_shot_subagent, + ttl: shaping.settings.prompt_caching_ttl, + }, + ); + let injected_count = injected.added; let trimmed = - without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?; + without_additional_drop_params(injected.request, &shaping.settings.additional_drop_params)?; let transformed = config.transform_anthropic_messages_request( trimmed, &MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params), @@ -126,6 +153,7 @@ fn prepare_provider_request( environment, timeout: timeout.or_else(|| shaping.settings.request_timeout()), api_key: api_key.map(SecretValue::new), + gateway_injected_cache_blocks: injected_count, }) } diff --git a/litellm-rust/crates/inference-messages/src/types.rs b/litellm-rust/crates/inference-messages/src/types.rs index 21eac0574ed..d2cce569359 100644 --- a/litellm-rust/crates/inference-messages/src/types.rs +++ b/litellm-rust/crates/inference-messages/src/types.rs @@ -4,8 +4,9 @@ use bytes::Bytes; use litellm_host::call::CallOutput; use litellm_llms::{ErrorDetail, base_llm::messages::context::MessagesModelCapabilities}; use litellm_llms_types::{ - formats::messages::{MessagesRequest, MessagesResponse}, + formats::messages::{CacheControlTtl, MessagesRequest, MessagesResponse}, headers::ProviderSpecificHeaders, + integrations::CacheControlInjectionPoint, }; use litellm_router_types::LitellmParams; use serde::{Deserialize, Serialize}; @@ -62,6 +63,14 @@ pub struct MessagesSettings { pub additional_drop_params: Vec, #[serde(default)] pub request_timeout: Option, + #[serde(default)] + pub cache_control_injection_points: Vec, + #[serde(default)] + pub enable_prompt_caching: bool, + #[serde(default)] + pub prompt_caching_ttl: Option, + #[serde(default)] + pub user_agent: Option, } impl MessagesSettings { @@ -116,6 +125,16 @@ mod tests { ..MessagesShaping::default() }, )] + #[case::prompt_caching_capability( + json!({"capabilities": {"supports_prompt_caching": true}}), + MessagesShaping { + capabilities: MessagesModelCapabilities { + supports_prompt_caching: true, + ..MessagesModelCapabilities::default() + }, + ..MessagesShaping::default() + }, + )] #[case::everything_the_python_host_projects( json!({ "capabilities": { @@ -139,6 +158,7 @@ mod tests { reasoning_auto_summary: true, additional_drop_params: vec!["metadata.user_id".to_string(), "thinking".to_string()], request_timeout: Some(1800.0), + ..MessagesSettings::default() }, capabilities: MessagesModelCapabilities { supports_reasoning: true, @@ -149,6 +169,7 @@ mod tests { supports_sampling_params: false, supports_speed: true, supports_mid_conversation_system: false, + supports_prompt_caching: false, effort_tiers: SupportedEffortTiers { minimal: false, low: true, @@ -184,4 +205,32 @@ mod tests { serde_json::to_value(expected.capabilities).unwrap() ); } + + #[rstest] + #[case::five_minutes("5m", CacheControlTtl::FiveMinutes)] + #[case::one_hour("1h", CacheControlTtl::OneHour)] + fn cache_control_settings_are_projected( + #[case] ttl: &str, + #[case] expected_ttl: CacheControlTtl, + ) { + let point = json!({"location": "message", "index": -1, "control": {"type": "ephemeral"}}); + let projected = json!({ + "cache_control_injection_points": [point], + "enable_prompt_caching": true, + "prompt_caching_ttl": ttl, + "user_agent": "claude-cli/1.0" + }); + + let shaping: MessagesShaping = serde_json::from_value(projected.clone()).unwrap(); + let serialized = serde_json::to_value(&shaping).unwrap(); + + assert!(serialized["enable_prompt_caching"].as_bool().unwrap()); + assert_eq!(serialized["prompt_caching_ttl"], json!(ttl)); + assert_eq!(serialized["user_agent"], json!("claude-cli/1.0")); + assert_eq!( + serialized["cache_control_injection_points"], + projected["cache_control_injection_points"] + ); + assert_eq!(shaping.settings.prompt_caching_ttl, Some(expected_ttl)); + } } diff --git a/litellm-rust/crates/inference-ocr/src/handler.rs b/litellm-rust/crates/inference-ocr/src/handler.rs index 24549456117..1a18de334ee 100644 --- a/litellm-rust/crates/inference-ocr/src/handler.rs +++ b/litellm-rust/crates/inference-ocr/src/handler.rs @@ -59,6 +59,7 @@ impl<'a, H> OcrCallHooks<'a, H> { .cloned() .collect(), api_key: request.connection.api_key.clone(), + gateway_injected_cache_blocks: 0, }, } } diff --git a/litellm-rust/crates/inference-responses/src/prepare.rs b/litellm-rust/crates/inference-responses/src/prepare.rs index 05a9772abd4..25d06c48bae 100644 --- a/litellm-rust/crates/inference-responses/src/prepare.rs +++ b/litellm-rust/crates/inference-responses/src/prepare.rs @@ -44,6 +44,7 @@ pub(super) async fn prepare( } _ => None, }, + gateway_injected_cache_blocks: 0, }; let body = config.transform_responses_api_request(model, call.input, call.optional_params)?; Ok(ProviderResponsesRequest { diff --git a/litellm-rust/crates/integrations/AGENTS.md b/litellm-rust/crates/integrations/AGENTS.md new file mode 100644 index 00000000000..35621470f27 --- /dev/null +++ b/litellm-rust/crates/integrations/AGENTS.md @@ -0,0 +1,25 @@ +# rules + +- Mirrors `litellm/integrations/` and owns deterministic integration behavior, independent of provider I/O and host callbacks +- Shared request contracts belong in `litellm-llms-types`; runtime configuration and provider capability checks are inputs from callers +- Function names follow the Python hook's static methods so a drift report diff maps onto them +- Count breakpoints only where the provider reads them: system blocks, message-level marks, content blocks, tools (top level or under `function`) and the top-level `cache_control`. Explicit nulls are not marks +- Each distinct message target is resolved once against a role index, so cost stays linear in points plus messages +- Use inline named `rstest` cases for behavior + +# Known partial ports + +- `tool_config` points reserve one slot of the four-breakpoint cap and are never placed. No native Messages provider builds Bedrock Converse `toolConfig` +- OpenAI prompt-cache dialects (`prompt_cache_breakpoint`, `_litellm_openai_dialect`) and Chat-only behavior (`tool_calls` marks, server tool replay) are not ported +- `extra_body` is not part of the typed Messages request projection and is not included in the census + +# Deliberate divergences + +- A message point with both `index` and `role: system` targets the index. Python sends it to the system prompt +- The system breakpoint lands on the last block that is not an `x-anthropic-billing-header:` block. Python marks the final block, which provider normalization later strips + +# references + +- `litellm/integrations/anthropic_cache_control_hook.py` +- `litellm-rust/crates/llms-types/src/integrations/anthropic_cache_control_hook.rs` +- `tests/rust-python-harness/strategies/drift/drift.json` mapping `anthropic-cache-control-hook` diff --git a/litellm-rust/crates/integrations/Cargo.toml b/litellm-rust/crates/integrations/Cargo.toml new file mode 100644 index 00000000000..8e7eab86bd7 --- /dev/null +++ b/litellm-rust/crates/integrations/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "litellm-integrations" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-llms-types.workspace = true +serde.workspace = true +serde_json.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/integrations/src/anthropic_cache_control_hook.rs b/litellm-rust/crates/integrations/src/anthropic_cache_control_hook.rs new file mode 100644 index 00000000000..c52120dd9df --- /dev/null +++ b/litellm-rust/crates/integrations/src/anthropic_cache_control_hook.rs @@ -0,0 +1,870 @@ +use std::collections::{HashMap, HashSet}; + +use litellm_llms_types::{ + formats::messages::{ + CacheControl, CacheControlTtl, CacheControlType, ContentBlock, ContentBlockType, Message, + MessageContent, MessagesOptionalParams, MessagesRequest, MessagesTool, SystemPrompt, + }, + integrations::{CacheControlInjectionPoint, CacheControlMessageInjectionPoint, Target}, + providers::anthropic::BILLING_HEADER_PREFIX, + recognized::Recognized, +}; +use serde_json::{Map, Value}; + +const MAX_CACHE_CONTROL_BLOCKS: usize = 4; +const CACHE_CONTROL_KEY: &str = "cache_control"; +const SYSTEM_ROLE: &str = "system"; + +pub struct InjectionInputs<'a> { + pub points: &'a [CacheControlInjectionPoint], + pub enable_prompt_caching: bool, + pub supports_anthropic_cache_control: bool, + pub one_shot_subagent: bool, + pub ttl: Option, +} + +pub struct Injected { + pub request: MessagesRequest, + pub added: usize, +} + +pub fn maybe_inject_cache_control( + request: MessagesRequest, + inputs: &InjectionInputs<'_>, +) -> Injected { + if !inputs.points.is_empty() { + return apply_to_anthropic_messages_request(request, inputs.points); + } + let defaults = get_default_injection_points(&request, inputs); + apply_to_anthropic_messages_request(request, &defaults) +} + +fn get_default_injection_points( + request: &MessagesRequest, + inputs: &InjectionInputs<'_>, +) -> Vec { + if !inputs.enable_prompt_caching + || !inputs.supports_anthropic_cache_control + || inputs.one_shot_subagent + || request_has_cache_control(request) + { + return Vec::new(); + } + let control = default_control(inputs.ttl); + vec![ + CacheControlInjectionPoint::Message(CacheControlMessageInjectionPoint { + target: Target::Role(SYSTEM_ROLE.into()), + control: Some(control.clone()), + }), + CacheControlInjectionPoint::Message(CacheControlMessageInjectionPoint { + target: Target::Index(-1), + control: Some(control), + }), + ] +} + +fn default_control(ttl: Option) -> CacheControl { + CacheControl { + ttl: ttl.map(Recognized::Known), + ..CacheControl::default() + } +} + +fn apply_to_anthropic_messages_request( + request: MessagesRequest, + injection_points: &[CacheControlInjectionPoint], +) -> Injected { + if injection_points.is_empty() { + return Injected { request, added: 0 }; + } + let (system_points, message_points): (Vec<_>, Vec<_>) = injection_points + .iter() + .filter_map(message_point) + .partition(|point| is_system_point(point)); + let max_blocks = MAX_CACHE_CONTROL_BLOCKS.saturating_sub( + blocks_reserved_outside_messages(injection_points) + + count_external_cache_breakpoints(&request), + ); + let used = count_request_cache_breakpoints(&request.messages, request.params.system.as_ref()); + + let system_point = system_points.first().filter(|_| { + used < max_blocks + && request + .params + .system + .as_ref() + .and_then(system_marking_index) + .is_some() + }); + let system_marked = system_point.is_some(); + let system = match system_point { + Some(point) => request + .params + .system + .map(|system| system_block_with_breakpoint(system, point.control.as_ref())), + None => request.params.system, + }; + let messages = request + .messages + .into_iter() + .map(with_block_content) + .collect(); + let (messages, added_to_messages) = apply_message_injections( + &message_points, + messages, + max_blocks.saturating_sub(used + usize::from(system_marked)), + ); + Injected { + request: MessagesRequest { + messages, + params: MessagesOptionalParams { + system, + ..request.params + }, + ..request + }, + added: usize::from(system_marked) + added_to_messages, + } +} + +fn message_point(point: &CacheControlInjectionPoint) -> Option<&CacheControlMessageInjectionPoint> { + match point { + CacheControlInjectionPoint::Message(point) => Some(point), + _ => None, + } +} + +fn is_system_point(point: &CacheControlMessageInjectionPoint) -> bool { + matches!(&point.target, Target::Role(role) if role == SYSTEM_ROLE) +} + +fn blocks_reserved_outside_messages(injection_points: &[CacheControlInjectionPoint]) -> usize { + usize::from( + injection_points + .iter() + .any(|point| matches!(point, CacheControlInjectionPoint::ToolConfig { .. })), + ) +} + +fn with_block_content(message: Message) -> Message { + match message.content { + MessageContent::Text(text) => Message { + content: MessageContent::Blocks(vec![ContentBlock::text(text)]), + ..message + }, + MessageContent::Blocks(_) => message, + } +} + +fn apply_message_injections( + points: &[&CacheControlMessageInjectionPoint], + mut messages: Vec, + max_blocks: usize, +) -> (Vec, usize) { + let roles = role_indices(&messages); + let mut seen = HashSet::new(); + let mut added = 0; + for point in points { + if !seen.insert(&point.target) { + continue; + } + for index in resolve_target_indices(&point.target, messages.len(), &roles) { + if added >= max_blocks { + return (messages, added); + } + if !message_has_cache_control(&messages[index]) + && safe_insert_cache_control_in_message( + &mut messages[index], + point.control.as_ref(), + ) + { + added += 1; + } + } + } + (messages, added) +} + +fn role_indices(messages: &[Message]) -> HashMap> { + messages + .iter() + .enumerate() + .fold(HashMap::new(), |mut roles, (index, message)| { + roles.entry(message.role.clone()).or_default().push(index); + roles + }) +} + +fn resolve_target_indices( + target: &Target, + length: usize, + roles: &HashMap>, +) -> Vec { + match target { + Target::None => Vec::new(), + Target::Index(index) => { + let resolved = if *index < 0 { + i64::try_from(length) + .unwrap_or(i64::MAX) + .saturating_add(*index) + } else { + *index + }; + usize::try_from(resolved) + .ok() + .filter(|resolved| *resolved < length) + .into_iter() + .collect() + } + Target::Role(role) => roles.get(role).cloned().unwrap_or_default(), + } +} + +fn safe_insert_cache_control_in_message( + message: &mut Message, + control: Option<&CacheControl>, +) -> bool { + let MessageContent::Blocks(blocks) = &mut message.content else { + return false; + }; + blocks.last_mut().is_some_and(|block| { + block.cache_control = Some(injected_control(control)); + true + }) +} + +fn system_marking_index(system: &SystemPrompt) -> Option { + match system { + SystemPrompt::Text(text) => (!text.starts_with(BILLING_HEADER_PREFIX)).then_some(0), + SystemPrompt::Blocks(blocks) if count_cache_control_in_blocks(blocks) > 0 => None, + SystemPrompt::Blocks(blocks) => blocks + .iter() + .rposition(|block| !is_billing_header_block(block)), + } +} + +fn system_block_with_breakpoint( + system: SystemPrompt, + control: Option<&CacheControl>, +) -> SystemPrompt { + let index = system_marking_index(&system); + let blocks = match system { + SystemPrompt::Text(text) => vec![ContentBlock::text(text)], + SystemPrompt::Blocks(blocks) => blocks, + }; + SystemPrompt::Blocks( + blocks + .into_iter() + .enumerate() + .map(|(position, block)| { + if Some(position) == index { + ContentBlock { + cache_control: Some(injected_control(control)), + ..block + } + } else { + block + } + }) + .collect(), + ) +} + +fn is_billing_header_block(block: &ContentBlock) -> bool { + block.is_type(ContentBlockType::Text) + && block + .text + .as_deref() + .is_some_and(|text| text.starts_with(BILLING_HEADER_PREFIX)) +} + +fn injected_control(control: Option<&CacheControl>) -> CacheControl { + CacheControl { + cache_type: Some( + control + .and_then(|control| control.cache_type.clone()) + .unwrap_or(Recognized::Known(CacheControlType::Ephemeral)), + ), + ttl: control + .and_then(|control| control.ttl.clone()) + .filter(|ttl| ttl.known().is_some()), + scope: control.and_then(|control| control.scope.clone()), + extra: control + .map(|control| control.extra.clone()) + .unwrap_or_default(), + } +} + +fn request_has_cache_control(request: &MessagesRequest) -> bool { + count_request_cache_breakpoints(&request.messages, request.params.system.as_ref()) + + count_external_cache_breakpoints(request) + > 0 +} + +fn count_request_cache_breakpoints(messages: &[Message], system: Option<&SystemPrompt>) -> usize { + let system_blocks = match system { + Some(SystemPrompt::Blocks(blocks)) => count_cache_control_in_blocks(blocks), + Some(SystemPrompt::Text(_)) | None => 0, + }; + system_blocks + + messages + .iter() + .map(count_cache_control_blocks) + .sum::() +} + +fn count_external_cache_breakpoints(request: &MessagesRequest) -> usize { + let tool_blocks = request + .params + .tools + .iter() + .flatten() + .filter(|tool| tool_carries_cache_breakpoint(tool)) + .count(); + usize::from(carries_cache_breakpoint(&request.params.extra)) + tool_blocks +} + +fn count_cache_control_blocks(message: &Message) -> usize { + let content_blocks = match &message.content { + MessageContent::Blocks(blocks) => count_cache_control_in_blocks(blocks), + MessageContent::Text(_) => 0, + }; + usize::from(carries_cache_breakpoint(&message.extra)) + content_blocks +} + +fn message_has_cache_control(message: &Message) -> bool { + count_cache_control_blocks(message) > 0 +} + +fn count_cache_control_in_blocks(blocks: &[ContentBlock]) -> usize { + blocks + .iter() + .filter(|block| block.cache_control.is_some()) + .count() +} + +fn tool_carries_cache_breakpoint(tool: &Recognized) -> bool { + let fields = match tool { + Recognized::Known( + MessagesTool::Advisor { extra } + | MessagesTool::ToolSearchRegex { extra } + | MessagesTool::ToolSearchBm25 { extra }, + ) => extra, + Recognized::Unrecognized(Value::Object(fields)) => fields, + Recognized::Unrecognized(_) => return false, + }; + carries_cache_breakpoint(fields) + || fields + .get("function") + .and_then(Value::as_object) + .is_some_and(carries_cache_breakpoint) +} + +fn carries_cache_breakpoint(fields: &Map) -> bool { + fields + .get(CACHE_CONTROL_KEY) + .is_some_and(|value| !value.is_null()) +} + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + use serde_json::json; + + use super::*; + + fn request(value: Value) -> MessagesRequest { + serde_json::from_value(value).expect("valid Messages request") + } + + fn configured_points(value: Value) -> Vec { + serde_json::from_value(value).expect("valid cache-control injection points") + } + + fn configured(points: &[CacheControlInjectionPoint]) -> InjectionInputs<'_> { + InjectionInputs { + points, + enable_prompt_caching: false, + supports_anthropic_cache_control: false, + one_shot_subagent: false, + ttl: None, + } + } + + fn defaults(ttl: Option) -> InjectionInputs<'static> { + InjectionInputs { + points: &[], + enable_prompt_caching: true, + supports_anthropic_cache_control: true, + one_shot_subagent: false, + ttl, + } + } + + fn inject(request: MessagesRequest, inputs: &InjectionInputs<'_>) -> (Value, usize) { + let injected = maybe_inject_cache_control(request, inputs); + ( + serde_json::to_value(injected.request).unwrap(), + injected.added, + ) + } + + #[fixture] + fn simple_request() -> MessagesRequest { + request(json!({ + "model": "claude-test", + "max_tokens": 32, + "system": "system prompt", + "messages": [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "middle"}, + {"role": "user", "content": "last"} + ] + })) + } + + #[rstest] + #[case::disabled(false, true, false, 0)] + #[case::unsupported(true, false, false, 0)] + #[case::one_shot(true, true, true, 0)] + #[case::enabled(true, true, false, 2)] + fn defaults_require_all_gates( + simple_request: MessagesRequest, + #[case] enabled: bool, + #[case] supported: bool, + #[case] one_shot: bool, + #[case] expected_added: usize, + ) { + let inputs = InjectionInputs { + points: &[], + enable_prompt_caching: enabled, + supports_anthropic_cache_control: supported, + one_shot_subagent: one_shot, + ttl: None, + }; + + let (value, added) = inject(simple_request, &inputs); + + assert_eq!(added, expected_added); + let expected_mark = (expected_added > 0).then(|| json!({"type": "ephemeral"})); + assert_eq!( + value["system"][0].get("cache_control").cloned(), + expected_mark + ); + assert_eq!( + value["messages"][2]["content"][0] + .get("cache_control") + .cloned(), + expected_mark + ); + } + + #[rstest] + #[case::message_block(json!({ + "model": "claude-test", "max_tokens": 32, "system": "system prompt", + "messages": [{"role": "user", "content": [{"type": "text", "text": "marked", "cache_control": {"type": "ephemeral"}}]}] + }))] + #[case::message_level(json!({ + "model": "claude-test", "max_tokens": 32, "system": "system prompt", + "messages": [{"role": "user", "content": "prompt", "cache_control": {"type": "ephemeral"}}] + }))] + #[case::system_block(json!({ + "model": "claude-test", "max_tokens": 32, + "system": [{"type": "text", "text": "system prompt", "cache_control": {"type": "ephemeral"}}], + "messages": [{"role": "user", "content": "prompt"}] + }))] + #[case::top_level(json!({ + "model": "claude-test", "max_tokens": 32, "system": "system prompt", + "messages": [{"role": "user", "content": "prompt"}], + "cache_control": {"type": "ephemeral"} + }))] + #[case::anthropic_tool(json!({ + "model": "claude-test", "max_tokens": 32, "system": "system prompt", + "messages": [{"role": "user", "content": "prompt"}], + "tools": [{"name": "lookup", "input_schema": {"type": "object"}, "cache_control": {"type": "ephemeral"}}] + }))] + #[case::function_tool(json!({ + "model": "claude-test", "max_tokens": 32, "system": "system prompt", + "messages": [{"role": "user", "content": "prompt"}], + "tools": [{"type": "function", "function": {"name": "lookup", "cache_control": {"type": "ephemeral"}}}] + }))] + fn defaults_stand_down_on_client_marks(#[case] request_value: Value) { + let (value, added) = inject(request(request_value.clone()), &defaults(None)); + + assert_eq!(added, 0); + assert_eq!(value, request_value); + } + + #[rstest] + #[case::tool_schema_property(json!({ + "tools": [{"name": "lookup", "input_schema": { + "type": "object", + "properties": {"cache_control": {"type": "string"}} + }}] + }))] + #[case::tool_use_input(json!({ + "messages": [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": [ + {"type": "tool_use", "id": "toolu_1", "name": "lookup", "input": {"cache_control": "x"}} + ]}, + {"role": "user", "content": "last"} + ] + }))] + #[case::metadata(json!({"metadata": {"cache_control": {"type": "ephemeral"}}}))] + #[case::null_top_level(json!({"cache_control": null}))] + #[case::null_message_level(json!({ + "messages": [{"role": "user", "content": "last", "cache_control": null}] + }))] + fn defaults_ignore_cache_control_outside_breakpoint_positions(#[case] overrides: Value) { + let mut value = json!({ + "model": "claude-test", "max_tokens": 32, "system": "system prompt", + "messages": [{"role": "user", "content": "last"}] + }); + value + .as_object_mut() + .unwrap() + .extend(overrides.as_object().unwrap().clone()); + + let (_, added) = inject(request(value), &defaults(None)); + + assert_eq!(added, 2); + } + + #[rstest] + #[case::five_minutes(CacheControlTtl::FiveMinutes, "5m")] + #[case::one_hour(CacheControlTtl::OneHour, "1h")] + fn defaults_apply_ttl( + simple_request: MessagesRequest, + #[case] ttl: CacheControlTtl, + #[case] expected: &str, + ) { + let (value, _) = inject(simple_request, &defaults(Some(ttl))); + + assert_eq!( + value["system"][0]["cache_control"], + json!({"type": "ephemeral", "ttl": expected}) + ); + assert_eq!( + value["messages"][2]["content"][0]["cache_control"], + json!({"type": "ephemeral", "ttl": expected}) + ); + } + + #[rstest] + fn configured_points_omit_invalid_ttl_and_default_missing_or_null_controls() { + let points = configured_points(json!([ + {"location": "message", "index": 0, "control": {"ttl": "invalid"}}, + {"location": "message", "index": 1, "control": null}, + {"location": "message", "index": 2} + ])); + let request = request(json!({ + "model": "claude-test", "max_tokens": 32, + "messages": [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + {"role": "user", "content": "third"} + ] + })); + + let (value, added) = inject(request, &configured(&points)); + + assert_eq!(added, 3); + assert_eq!( + value["messages"][0]["content"][0]["cache_control"], + json!({"type": "ephemeral"}) + ); + assert_eq!( + value["messages"][1]["content"][0]["cache_control"], + json!({"type": "ephemeral"}) + ); + assert_eq!( + value["messages"][2]["content"][0]["cache_control"], + json!({"type": "ephemeral"}) + ); + } + + #[rstest] + fn configured_points_apply_beside_client_marks_and_index_wins_over_system_role() { + let points = + configured_points(json!([{"location": "message", "index": 0, "role": "system"}])); + let request = request(json!({ + "model": "claude-test", "max_tokens": 32, "system": "system prompt", + "messages": [{"role": "user", "content": "prompt"}], + "cache_control": {"type": "ephemeral"} + })); + + let (value, added) = inject(request, &configured(&points)); + + assert_eq!(added, 1); + assert_eq!(value["system"], "system prompt"); + assert_eq!( + value["messages"][0]["content"][0]["cache_control"], + json!({"type": "ephemeral"}) + ); + } + + #[rstest] + #[case::earlier_block(json!([ + {"type": "text", "text": "client marked", "cache_control": {"type": "ephemeral", "ttl": "1h"}}, + {"type": "text", "text": "unmarked tail"} + ]))] + #[case::final_block(json!([ + {"type": "text", "text": "unmarked head"}, + {"type": "text", "text": "client marked", "cache_control": {"type": "ephemeral", "ttl": "1h"}} + ]))] + fn configured_point_leaves_an_already_marked_message_alone(#[case] content: Value) { + let points = configured_points(json!([{"location": "message", "index": 0}])); + let request_value = json!({ + "model": "claude-test", "max_tokens": 32, + "messages": [{"role": "user", "content": content}] + }); + + let (value, added) = inject(request(request_value.clone()), &configured(&points)); + + assert_eq!(added, 0); + assert_eq!(value, request_value); + } + + #[rstest] + #[case::client_marks_fill_the_cap( + json!("system prompt"), + json!([ + {"role": "user", "content": [{"type": "text", "text": "a", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": [{"type": "text", "text": "b", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": [{"type": "text", "text": "c", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": [{"type": "text", "text": "d", "cache_control": {"type": "ephemeral"}}]} + ]) + )] + #[case::system_already_marked_on_an_earlier_block( + json!([ + {"type": "text", "text": "client marked", "cache_control": {"type": "ephemeral", "ttl": "1h"}}, + {"type": "text", "text": "unmarked tail"} + ]), + json!([{"role": "user", "content": [{"type": "text", "text": "prompt"}]}]) + )] + fn system_point_adds_nothing(#[case] system: Value, #[case] messages: Value) { + let points = configured_points(json!([{"location": "message", "role": "system"}])); + let request_value = json!({ + "model": "claude-test", "max_tokens": 32, "system": system, "messages": messages + }); + + let (value, added) = inject(request(request_value.clone()), &configured(&points)); + + assert_eq!(added, 0); + assert_eq!(value, request_value); + } + + #[rstest] + #[case::one_past_the_end(1)] + #[case::far_past_the_end(99)] + #[case::negative_past_the_start(-2)] + fn out_of_range_index_marks_nothing(#[case] index: i64) { + let points = configured_points(json!([{"location": "message", "index": index}])); + let request_value = json!({ + "model": "claude-test", "max_tokens": 32, + "messages": [{"role": "user", "content": [{"type": "text", "text": "only"}]}] + }); + + let (value, added) = inject(request(request_value.clone()), &configured(&points)); + + assert_eq!(added, 0); + assert_eq!(value, request_value); + } + + #[rstest] + fn system_point_is_applied_before_message_points_and_tool_config_reserves_a_slot() { + let points = configured_points(json!([ + {"location": "message", "role": "user"}, + {"location": "message", "role": "system"}, + {"location": "tool_config"} + ])); + let request = request(json!({ + "model": "claude-test", "max_tokens": 32, "system": "system prompt", + "messages": [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + {"role": "user", "content": "third"} + ] + })); + + let (value, added) = inject(request, &configured(&points)); + + assert_eq!(added, 3); + assert!(value["system"][0].get("cache_control").is_some()); + assert!( + value["messages"][0]["content"][0] + .get("cache_control") + .is_some() + ); + assert!( + value["messages"][1]["content"][0] + .get("cache_control") + .is_some() + ); + assert!( + value["messages"][2]["content"][0] + .get("cache_control") + .is_none() + ); + } + + #[rstest] + fn configured_message_point_marks_only_the_final_content_block() { + let points = configured_points(json!([{"location": "message", "index": 0}])); + let request = request(json!({ + "model": "claude-test", "max_tokens": 32, + "messages": [{"role": "user", "content": [ + {"type": "text", "text": "first"}, + {"type": "text", "text": "last"} + ]}] + })); + + let (value, added) = inject(request, &configured(&points)); + + assert_eq!(added, 1); + assert!( + value["messages"][0]["content"][0] + .get("cache_control") + .is_none() + ); + assert_eq!( + value["messages"][0]["content"][1]["cache_control"], + json!({"type": "ephemeral"}) + ); + } + + #[rstest] + #[case::no_external_marks(json!({}), 4)] + #[case::marked_tool(json!({"tools": [{"name": "t", "input_schema": {}, "cache_control": {"type": "ephemeral"}}]}), 3)] + #[case::automatic_caching(json!({"cache_control": {"type": "ephemeral"}}), 3)] + fn gateway_injection_observes_the_four_block_cap( + #[case] overrides: Value, + #[case] expected: usize, + ) { + let points = configured_points(json!([{"location": "message", "role": "user"}])); + let mut value = json!({ + "model": "claude-test", "max_tokens": 32, + "messages": [ + {"role": "user", "content": "first"}, + {"role": "user", "content": "second"}, + {"role": "user", "content": "third"}, + {"role": "user", "content": "fourth"}, + {"role": "user", "content": "fifth"} + ] + }); + value + .as_object_mut() + .unwrap() + .extend(overrides.as_object().unwrap().clone()); + + let (value, added) = inject(request(value), &configured(&points)); + + assert_eq!(added, expected); + assert!( + value["messages"][expected - 1]["content"][0] + .get("cache_control") + .is_some() + ); + assert!( + value["messages"][expected]["content"][0] + .get("cache_control") + .is_none() + ); + } + + #[rstest] + fn repeated_targets_keep_the_first_points_control() { + let points = configured_points(json!([ + {"location": "message", "role": "user", "control": {"type": "ephemeral", "ttl": "1h"}}, + {"location": "message", "role": "user", "control": {"type": "ephemeral", "ttl": "5m"}}, + {"location": "message", "role": "assistant"} + ])); + let request = request(json!({ + "model": "claude-test", "max_tokens": 32, + "messages": [{"role": "user", "content": "only"}] + })); + + let (value, added) = inject(request, &configured(&points)); + + assert_eq!(added, 1); + assert_eq!( + value["messages"][0]["content"][0]["cache_control"], + json!({"type": "ephemeral", "ttl": "1h"}) + ); + } + + #[rstest] + fn all_message_strings_are_normalized_when_points_apply(simple_request: MessagesRequest) { + let points = configured_points(json!([{"location": "message", "index": 99}])); + + let (value, added) = inject(simple_request, &configured(&points)); + + assert_eq!(added, 0); + assert_eq!( + value["messages"][0]["content"], + json!([{"type": "text", "text": "first"}]) + ); + assert_eq!( + value["messages"][1]["content"], + json!([{"type": "text", "text": "middle"}]) + ); + } + + #[rstest] + fn non_message_points_are_never_placed_on_the_request(simple_request: MessagesRequest) { + let points = configured_points(json!([ + {"location": "future_location", "target": "future", "control": {"x": 1}}, + {"location": "tool_config"} + ])); + + let (value, added) = inject(simple_request, &configured(&points)); + + assert_eq!(added, 0); + assert_eq!(value["system"], "system prompt"); + assert!(value.get("cache_control_injection_points").is_none()); + assert!(value.get("tool_config").is_none()); + } + + #[rstest] + #[case::trailing_billing_block( + json!([ + {"type": "text", "text": "Instructions"}, + {"type": "text", "text": "x-anthropic-billing-header: cc_version=1"} + ]), + Some(0), + 1 + )] + #[case::only_billing_blocks( + json!([{"type": "text", "text": "x-anthropic-billing-header: cc_version=1"}]), + None, + 0 + )] + #[case::billing_text(json!("x-anthropic-billing-header: cc_version=1"), None, 0)] + fn system_breakpoint_skips_billing_header_blocks( + #[case] system: Value, + #[case] marked_block: Option, + #[case] expected_added: usize, + ) { + let points = configured_points(json!([{"location": "message", "role": "system"}])); + let request = request(json!({ + "model": "claude-test", "max_tokens": 32, "system": system, + "messages": [{"role": "user", "content": "prompt"}] + })); + + let (value, added) = inject(request, &configured(&points)); + + assert_eq!(added, expected_added); + let marks: Vec = match &value["system"] { + Value::Array(blocks) => blocks + .iter() + .map(|block| block.get("cache_control").is_some()) + .collect(), + _ => Vec::new(), + }; + let expected: Vec = (0..marks.len()) + .map(|index| Some(index) == marked_block) + .collect(); + assert_eq!(marks, expected); + } +} diff --git a/litellm-rust/crates/integrations/src/lib.rs b/litellm-rust/crates/integrations/src/lib.rs new file mode 100644 index 00000000000..55c8d3c9b9a --- /dev/null +++ b/litellm-rust/crates/integrations/src/lib.rs @@ -0,0 +1,3 @@ +pub mod anthropic_cache_control_hook; + +pub use anthropic_cache_control_hook::{Injected, InjectionInputs, maybe_inject_cache_control}; diff --git a/litellm-rust/crates/llms-types/src/formats/chat_completions/content.rs b/litellm-rust/crates/llms-types/src/formats/chat_completions/content.rs index 7298145533c..c6056816455 100644 --- a/litellm-rust/crates/llms-types/src/formats/chat_completions/content.rs +++ b/litellm-rust/crates/llms-types/src/formats/chat_completions/content.rs @@ -161,6 +161,7 @@ mod tests { messages::ContentSource, }; + use crate::{formats::messages::CacheControlType, recognized::Recognized}; use rstest::rstest; use serde_json::{Value, json}; @@ -208,8 +209,8 @@ mod tests { }; assert_eq!(text, "hello"); assert_eq!( - cache_control.as_ref().unwrap().cache_type.as_deref(), - Some("ephemeral") + cache_control.as_ref().unwrap().cache_type, + Some(Recognized::Known(CacheControlType::Ephemeral)) ); assert_eq!(breakpoint.mode, PromptCacheMode::Explicit); assert!(breakpoint.extra.is_empty()); diff --git a/litellm-rust/crates/llms-types/src/formats/messages/request.rs b/litellm-rust/crates/llms-types/src/formats/messages/request.rs index 9dbfa068025..cfd37fb959e 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/request.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/request.rs @@ -114,15 +114,31 @@ impl ContentBlock { #[derive(Default)] pub struct CacheControl { #[serde(rename = "type", skip_serializing_if = "Option::is_none")] - pub cache_type: Option, + pub cache_type: Option>, #[serde(skip_serializing_if = "Option::is_none")] - pub ttl: Option, + pub ttl: Option>, #[serde(skip_serializing_if = "Option::is_none")] pub scope: Option, #[serde(flatten)] pub extra: Map, } +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, strum::VariantArray)] +pub enum CacheControlType { + #[serde(rename = "ephemeral")] + Ephemeral, +} + +#[macro_rules_attribute::apply(crate::wire_type)] +#[derive(Copy, Eq, strum::VariantArray)] +pub enum CacheControlTtl { + #[serde(rename = "5m")] + FiveMinutes, + #[serde(rename = "1h")] + OneHour, +} + #[macro_rules_attribute::apply(crate::wire_type)] pub struct Message { pub role: String, @@ -399,6 +415,42 @@ mod tests { use serde_json::{Map, Value}; + #[rstest] + #[case::known( + json!({"type": "ephemeral", "ttl": "1h"}), + Some(Recognized::Known(CacheControlType::Ephemeral)), + Some(Recognized::Known(CacheControlTtl::OneHour)) + )] + #[case::five_minutes( + json!({"type": "ephemeral", "ttl": "5m"}), + Some(Recognized::Known(CacheControlType::Ephemeral)), + Some(Recognized::Known(CacheControlTtl::FiveMinutes)) + )] + #[case::unknown_values_pass_through( + json!({"type": "persistent", "ttl": 300}), + Some(Recognized::Unrecognized(json!("persistent"))), + Some(Recognized::Unrecognized(json!(300))) + )] + #[case::variant_name_is_not_a_wire_value( + json!({"ttl": "FiveMinutes"}), + None, + Some(Recognized::Unrecognized(json!("FiveMinutes"))) + )] + #[case::absent(json!({}), None, None)] + fn cache_control_types_known_values_and_keeps_unknown_ones( + #[case] input: Value, + #[case] cache_type: Option>, + #[case] ttl: Option>, + ) { + let control: CacheControl = serde_json::from_value(input.clone()).unwrap(); + + assert_eq!( + (control.cache_type.clone(), control.ttl.clone()), + (cache_type, ttl) + ); + assert_eq!(serde_json::to_value(control).unwrap(), input); + } + fn round_trip(value: &Value) -> Value { let parsed: T = serde_json::from_value(value.clone()).unwrap(); serde_json::to_value(parsed).unwrap() diff --git a/litellm-rust/crates/llms-types/src/formats/messages/tools.rs b/litellm-rust/crates/llms-types/src/formats/messages/tools.rs index 6f5acf9fb3c..9f3924c70ff 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/tools.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/tools.rs @@ -568,6 +568,10 @@ mod tests { use crate::formats::messages::{ AdvisorTool, AdvisorToolName, CacheControl, UrlSourceToolReference, }; + use crate::{ + formats::messages::{CacheControlTtl, CacheControlType}, + recognized::Recognized, + }; use serde_json::Map; use crate::formats::messages::AllowedCaller; @@ -819,7 +823,7 @@ mod tests { #[case::text_editor_20250429( json!({"type":"text_editor_20250429","name":"str_replace_based_edit_tool","cache_control":{"type":"ephemeral","ttl":"1h"}}), BuiltinMessagesTool::TextEditor20250429(ClientTool { - cache_control: Some(cache_control("ephemeral", "1h")), + cache_control: Some(cache_control(CacheControlType::Ephemeral, CacheControlTtl::OneHour)), ..client_tool(StrReplaceBasedEditToolName::StrReplaceBasedEditTool) }) )] @@ -1045,7 +1049,7 @@ mod tests { model: String::from("claude-opus-5-5"), allowed_callers: None, cache_control: None, - caching: Some(cache_control("ephemeral", "5m")), + caching: Some(cache_control(CacheControlType::Ephemeral, CacheControlTtl::FiveMinutes)), defer_loading: None, max_tokens: Some(2048), max_uses: None, @@ -1223,10 +1227,10 @@ mod tests { } } - fn cache_control(cache_type: &str, ttl: &str) -> CacheControl { + fn cache_control(cache_type: CacheControlType, ttl: CacheControlTtl) -> CacheControl { CacheControl { - cache_type: Some(String::from(cache_type)), - ttl: Some(String::from(ttl)), + cache_type: Some(Recognized::Known(cache_type)), + ttl: Some(Recognized::Known(ttl)), ..CacheControl::default() } } diff --git a/litellm-rust/crates/llms-types/src/integrations/anthropic_cache_control_hook.rs b/litellm-rust/crates/llms-types/src/integrations/anthropic_cache_control_hook.rs new file mode 100644 index 00000000000..b488dd3b23c --- /dev/null +++ b/litellm-rust/crates/llms-types/src/integrations/anthropic_cache_control_hook.rs @@ -0,0 +1,202 @@ +use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use serde_json::{Map, Value}; + +use crate::formats::messages::CacheControl; + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum Target { + Index(i64), + Role(String), + None, +} + +#[derive(Clone, Debug, PartialEq)] +pub struct CacheControlMessageInjectionPoint { + pub target: Target, + pub control: Option, +} + +#[derive(Clone, Debug, PartialEq)] +pub enum CacheControlInjectionPoint { + Message(CacheControlMessageInjectionPoint), + ToolConfig { control: Option }, + Unknown(Value), +} + +impl<'de> Deserialize<'de> for CacheControlInjectionPoint { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let value = Value::deserialize(deserializer)?; + let Some(object) = value.as_object() else { + return Ok(Self::Unknown(value)); + }; + let location = object.get("location").and_then(Value::as_str); + match location { + Some("message") => { + let index = object.get("index"); + let target = match index { + Some(Value::Number(number)) => number + .as_i64() + .map(Target::Index) + .unwrap_or_else(|| role_target(object)), + Some(Value::String(index)) => index + .parse::() + .map(Target::Index) + .unwrap_or_else(|_| role_target(object)), + _ => role_target(object), + }; + let control = point_control(object).map_err(serde::de::Error::custom)?; + Ok(Self::Message(CacheControlMessageInjectionPoint { + target, + control, + })) + } + Some("tool_config") => { + let control = point_control(object).map_err(serde::de::Error::custom)?; + Ok(Self::ToolConfig { control }) + } + _ => Ok(Self::Unknown(value)), + } + } +} + +impl Serialize for CacheControlInjectionPoint { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + let value = match self { + Self::Message(point) => { + let mut object = Map::new(); + object.insert("location".into(), Value::String("message".into())); + match &point.target { + Target::Index(index) => { + object.insert("index".into(), Value::Number((*index).into())); + } + Target::Role(role) => { + object.insert("role".into(), Value::String(role.clone())); + } + Target::None => {} + } + if let Some(control) = &point.control { + object.insert( + "control".into(), + serde_json::to_value(control).unwrap_or(Value::Null), + ); + } + Value::Object(object) + } + Self::ToolConfig { control } => { + let mut object = Map::new(); + object.insert("location".into(), Value::String("tool_config".into())); + if let Some(control) = control { + object.insert( + "control".into(), + serde_json::to_value(control).unwrap_or(Value::Null), + ); + } + Value::Object(object) + } + Self::Unknown(value) => value.clone(), + }; + value.serialize(serializer) + } +} + +fn point_control(object: &Map) -> serde_json::Result> { + object + .get("control") + .filter(|control| !control.is_null()) + .cloned() + .map(serde_json::from_value) + .transpose() +} + +fn role_target(object: &Map) -> Target { + object + .get("role") + .and_then(Value::as_str) + .map(|role| Target::Role(role.to_string())) + .unwrap_or(Target::None) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + + use super::{CacheControlInjectionPoint, CacheControlMessageInjectionPoint, Target}; + + #[rstest] + #[case::negative_index_string(json!({"location": "message", "index": "-1"}), Target::Index(-1))] + #[case::out_of_range_index(json!({"location": "message", "index": 99}), Target::Index(99))] + #[case::role_only(json!({"location": "message", "role": "user"}), Target::Role("user".into()))] + #[case::invalid_index_falls_back_to_role( + json!({"location": "message", "index": "latest", "role": "assistant"}), + Target::Role("assistant".into()) + )] + #[case::index_precedes_system_role( + json!({"location": "message", "index": 0, "role": "system"}), + Target::Index(0) + )] + #[case::missing_target(json!({"location": "message"}), Target::None)] + fn message_target_deserialization(#[case] input: Value, #[case] expected: Target) { + let point: CacheControlInjectionPoint = serde_json::from_value(input).unwrap(); + + assert!(matches!( + point, + CacheControlInjectionPoint::Message(message) if message.target == expected + )); + } + + #[rstest] + #[case::missing_control( + json!({"location": "message", "index": 0}), + CacheControlInjectionPoint::Message(CacheControlMessageInjectionPoint { + target: Target::Index(0), + control: None, + }) + )] + #[case::message_null_control( + json!({"location": "message", "index": 0, "control": null}), + CacheControlInjectionPoint::Message(CacheControlMessageInjectionPoint { + target: Target::Index(0), + control: None, + }) + )] + #[case::tool_config_without_control( + json!({"location": "tool_config"}), + CacheControlInjectionPoint::ToolConfig { control: None } + )] + #[case::tool_config_null_control( + json!({"location": "tool_config", "control": null}), + CacheControlInjectionPoint::ToolConfig { control: None } + )] + fn known_points_preserve_optional_control( + #[case] input: Value, + #[case] expected: CacheControlInjectionPoint, + ) { + assert_eq!( + serde_json::from_value::(input).unwrap(), + expected + ); + } + + #[rstest] + #[case::message(json!({"location": "message", "index": 0, "control": "ephemeral"}))] + #[case::tool_config(json!({"location": "tool_config", "control": 5}))] + fn malformed_controls_are_rejected(#[case] input: Value) { + assert!(serde_json::from_value::(input).is_err()); + } + + #[rstest] + fn unknown_location_round_trips_unchanged() { + let input = json!({"location": "future", "index": 7, "custom": {"value": true}}); + let point: CacheControlInjectionPoint = serde_json::from_value(input.clone()).unwrap(); + + assert!(matches!(&point, CacheControlInjectionPoint::Unknown(_))); + assert_eq!(serde_json::to_value(point).unwrap(), input); + } +} diff --git a/litellm-rust/crates/llms-types/src/integrations/mod.rs b/litellm-rust/crates/llms-types/src/integrations/mod.rs new file mode 100644 index 00000000000..0da3b631f2f --- /dev/null +++ b/litellm-rust/crates/llms-types/src/integrations/mod.rs @@ -0,0 +1,3 @@ +mod anthropic_cache_control_hook; + +pub use anthropic_cache_control_hook::*; diff --git a/litellm-rust/crates/llms-types/src/lib.rs b/litellm-rust/crates/llms-types/src/lib.rs index 67d933e39cc..f45784ba88a 100644 --- a/litellm-rust/crates/llms-types/src/lib.rs +++ b/litellm-rust/crates/llms-types/src/lib.rs @@ -6,6 +6,7 @@ macro_rules_attribute::attribute_alias! { pub mod formats; pub mod headers; +pub mod integrations; pub mod json_schema; pub mod providers; pub mod recognized; diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic/constants.rs b/litellm-rust/crates/llms-types/src/providers/anthropic/constants.rs index 941d1e75461..291d0085f54 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic/constants.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic/constants.rs @@ -8,6 +8,7 @@ pub const BETA_HEADER: &str = "anthropic-beta"; pub const VERSION_HEADER: &str = "anthropic-version"; pub const DIRECT_BROWSER_ACCESS_HEADER: &str = "anthropic-dangerous-direct-browser-access"; pub const API_VERSION: &str = "2023-06-01"; +pub const BILLING_HEADER_PREFIX: &str = "x-anthropic-billing-header:"; pub const DEFAULT_HEADERS: &[(&str, &str)] = &[ (VERSION_HEADER, API_VERSION), diff --git a/litellm-rust/crates/llms-types/src/providers/minimax/messages.rs b/litellm-rust/crates/llms-types/src/providers/minimax/messages.rs index dbecb82b8f8..24dd1bfca34 100644 --- a/litellm-rust/crates/llms-types/src/providers/minimax/messages.rs +++ b/litellm-rust/crates/llms-types/src/providers/minimax/messages.rs @@ -64,6 +64,7 @@ mod tests { MinimaxMediaDetail, MinimaxMediaSource, MinimaxMessagesContentBlock, }; + use crate::{formats::messages::CacheControlType, recognized::Recognized}; use rstest::rstest; use serde_json::{Value, json}; @@ -97,8 +98,8 @@ mod tests { assert_eq!(options.max_long_side_pixel, Some(1024)); assert_eq!(options.extra["future"], Value::Null); assert_eq!( - video.cache_control.as_ref().unwrap().cache_type.as_deref(), - Some("ephemeral") + video.cache_control.as_ref().unwrap().cache_type, + Some(Recognized::Known(CacheControlType::Ephemeral)) ); } MinimaxMessagesContentBlock::MidConvSystem { text, extra } => { diff --git a/litellm-rust/crates/llms/AGENTS.md b/litellm-rust/crates/llms/AGENTS.md index 3cf916b1721..7dc3d2f6ff5 100644 --- a/litellm-rust/crates/llms/AGENTS.md +++ b/litellm-rust/crates/llms/AGENTS.md @@ -41,7 +41,7 @@ ## Python pairs -- Python paths identify counterparts but do not dictate Rust module names or class hierarchy. Preserve behavior and concepts, not structure +- Prefer mirroring Python's outer module and class structure, following the workspace's Python porting guideline. Preserve behavior while designing function internals freely - Keep operation and parameter names when responsibilities match. Rust types keep the Python name with Rust acronym casing (`BaseOCRConfig` -> `BaseOcrConfig`). Private Python helpers drop the leading underscore - Give Rust-only helpers distinct responsibility names instead of reusing trait method names - Use trait defaults for unchanged inherited behavior and explicit delegation for shared provider behavior. Config traits represent real provider contracts, so do not recreate inheritance with extra traits diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 7991a57353e..714ca5bdda2 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -6,11 +6,11 @@ use litellm_http::request::{ }; use litellm_llms_types::{ formats::messages::{ - ContentBlock, ContentBlockType, EffortLevel, Message, MessageContent, MessagesTool, - SystemPrompt, + ContentBlock, ContentBlockType, EffortLevel, Message, MessageContent, MessagesRequest, + MessagesTool, SystemPrompt, }, providers::anthropic::{ - API_BASE, API_KEY_HEADER, AnthropicBeta, BETA_HEADER, BetaSet, + API_BASE, API_KEY_HEADER, AnthropicBeta, BETA_HEADER, BILLING_HEADER_PREFIX, BetaSet, DIRECT_BROWSER_ACCESS_HEADER, MESSAGES_PATH, }, recognized::Recognized, @@ -18,6 +18,94 @@ use litellm_llms_types::{ use serde::Deserialize; use serde_json::Value; +pub fn supports_anthropic_cache_control( + model: &str, + provider: &str, + supports_prompt_caching: bool, +) -> bool { + supports_prompt_caching + && model.to_ascii_lowercase().contains("claude") + && matches!(provider, "anthropic" | "bedrock" | "vertex_ai" | "azure_ai") +} + +pub fn is_claude_code_one_shot_subagent_request( + request: &MessagesRequest, + user_agent: Option<&str>, +) -> bool { + let Some(user_agent) = user_agent else { + return false; + }; + if (!user_agent.starts_with("claude-cli/") && !user_agent.starts_with("claude-code/")) + || request.messages.len() != 1 + || request + .params + .tools + .as_ref() + .is_some_and(|tools| !tools.is_empty()) + { + return false; + } + request.messages[0].role == "user" + && request + .params + .system + .as_ref() + .is_some_and(is_claude_code_subagent_billing_system) +} + +fn is_claude_code_subagent_billing_system(system: &SystemPrompt) -> bool { + let texts: Option> = match system { + SystemPrompt::Text(text) => Some(vec![text.as_str()]), + SystemPrompt::Blocks(blocks) => blocks + .iter() + .map(|block| { + block + .is_type(ContentBlockType::Text) + .then_some(block.text.as_deref()) + .flatten() + }) + .collect(), + }; + let meaningful: Vec<&str> = texts + .into_iter() + .flatten() + .filter(|text| !text.trim().is_empty()) + .collect(); + if meaningful.is_empty() { + return false; + } + let fields: Option>> = meaningful + .into_iter() + .map(claude_code_billing_fields) + .collect(); + let subagent_values: Vec<&str> = fields + .into_iter() + .flatten() + .flatten() + .filter_map(|(key, value)| (key == "cc_is_subagent").then_some(value)) + .collect(); + subagent_values == ["true"] +} + +fn claude_code_billing_fields(text: &str) -> Option> { + let stripped = text.trim(); + if stripped.contains(['\n', '\r']) { + return None; + } + let fields = stripped.strip_prefix(BILLING_HEADER_PREFIX)?; + let parsed: Option> = fields + .split(';') + .filter(|field| !field.trim().is_empty()) + .map(|field| { + let (key, value) = field.split_once('=')?; + let key = key.trim(); + let value = value.trim(); + (!key.is_empty() && !value.is_empty()).then_some((key, value)) + }) + .collect(); + parsed.filter(|parsed| !parsed.is_empty()) +} + use crate::{ anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX, base_llm::auth::{AuthScheme, Headers}, @@ -384,8 +472,6 @@ pub fn is_encrypted_reasoning_block(block: &ContentBlock) -> bool { field.is_some_and(|value| value.starts_with(ENCRYPTED_REASONING_SIGNATURE_PREFIX)) } -const BILLING_HEADER_PREFIX: &str = "x-anthropic-billing-header:"; - fn is_billing_header_block(block: &ContentBlock) -> bool { block.block_type == Some(ContentBlockType::Text) && block @@ -1885,6 +1971,7 @@ mod tests { supports_sampling_params: true, supports_speed: false, supports_mid_conversation_system: false, + supports_prompt_caching: false, effort_tiers: tiers(false, false, false, false, false, false), } ); @@ -1920,4 +2007,121 @@ mod tests { expected ); } + + #[rstest] + #[case::anthropic("claude-sonnet", "anthropic", true, true)] + #[case::bedrock("claude-sonnet", "bedrock", true, true)] + #[case::vertex("claude-sonnet", "vertex_ai", true, true)] + #[case::azure("claude-sonnet", "azure_ai", true, true)] + #[case::unsupported_provider("claude-sonnet", "deepseek", true, false)] + #[case::unsupported_model("gpt-test", "anthropic", true, false)] + #[case::unsupported_capability("claude-sonnet", "anthropic", false, false)] + fn cache_control_support_requires_model_provider_and_capability( + #[case] model: &str, + #[case] provider: &str, + #[case] capability: bool, + #[case] expected: bool, + ) { + assert_eq!( + supports_anthropic_cache_control(model, provider, capability), + expected + ); + } + + #[rstest] + #[case::claude_cli( + "claude-cli/1.0", + "x-anthropic-billing-header: cc_is_subagent=true", + None, + true + )] + #[case::claude_code( + "claude-code/1.0", + "x-anthropic-billing-header: cc_is_subagent=true", + None, + true + )] + #[case::user_agent_must_start_with_product( + "client claude-code/1.0", + "x-anthropic-billing-header: cc_is_subagent=true", + None, + false + )] + #[case::other_billing_field( + "claude-code/1.0", + "x-anthropic-billing-header: billing=true", + None, + false + )] + #[case::false_subagent_field( + "claude-code/1.0", + "x-anthropic-billing-header: cc_is_subagent=false", + None, + false + )] + #[case::billing_header_with_instructions( + "claude-code/1.0", + "x-anthropic-billing-header: cc_is_subagent=true\nYou are a helpful agent", + None, + false + )] + #[case::nonempty_tools( + "claude-code/1.0", + "x-anthropic-billing-header: cc_is_subagent=true", + Some(json!([{"name": "lookup", "input_schema": {"type": "object"}}])), + false + )] + fn one_shot_subagent_requires_claude_code_request_shape( + #[case] user_agent: &str, + #[case] system: &str, + #[case] tools: Option, + #[case] expected: bool, + ) { + let mut request_value = json!({ + "model": "claude-test", + "max_tokens": 16, + "system": system, + "messages": [{"role": "user", "content": "prompt"}] + }); + if let Some(tools) = tools { + request_value["tools"] = tools; + } + let request: MessagesRequest = serde_json::from_value(request_value).unwrap(); + + assert_eq!( + is_claude_code_one_shot_subagent_request(&request, Some(user_agent)), + expected + ); + } + + #[rstest] + #[case::empty_text_blocks_are_ignored(json!([ + {"type": "text", "text": ""}, + {"type": "text", "text": "x-anthropic-billing-header: cc_is_subagent=true"} + ]), true)] + #[case::billing_block_beside_instructions(json!([ + {"type": "text", "text": "x-anthropic-billing-header: cc_is_subagent=true"}, + {"type": "text", "text": "Answer the delegated question"} + ]), false)] + #[case::non_text_block(json!([ + {"type": "text", "text": "x-anthropic-billing-header: cc_is_subagent=true"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AA=="}} + ]), false)] + fn one_shot_subagent_requires_every_system_text_to_be_a_billing_header( + #[case] system: Value, + #[case] expected: bool, + ) { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-test", + "max_tokens": 16, + "system": system, + "messages": [{"role": "user", "content": "prompt"}] + })) + .unwrap(); + + assert_eq!( + is_claude_code_one_shot_subagent_request(&request, Some("claude-cli/1.0")), + expected + ); + } } 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 e59a63813b0..d72e252d162 100644 --- a/litellm-rust/crates/llms/src/base_llm/messages/context.rs +++ b/litellm-rust/crates/llms/src/base_llm/messages/context.rs @@ -42,6 +42,8 @@ pub struct MessagesModelCapabilities { #[serde(default)] pub supports_mid_conversation_system: bool, #[serde(default)] + pub supports_prompt_caching: bool, + #[serde(default)] pub effort_tiers: SupportedEffortTiers, } @@ -60,6 +62,7 @@ impl Default for MessagesModelCapabilities { supports_sampling_params: true, supports_speed: false, supports_mid_conversation_system: false, + supports_prompt_caching: false, effort_tiers: SupportedEffortTiers::default(), } } diff --git a/litellm/rust_bridge/host/callbacks_legacy_python.py b/litellm/rust_bridge/host/callbacks_legacy_python.py index 54bf6ae5430..d446ac78f4e 100644 --- a/litellm/rust_bridge/host/callbacks_legacy_python.py +++ b/litellm/rust_bridge/host/callbacks_legacy_python.py @@ -218,6 +218,12 @@ def update_logging( ) +def record_gateway_injection(kwargs: Mapping[str, object], added: int) -> None: + from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook + + AnthropicCacheControlHook.record_gateway_injection(kwargs, added) + + def pre_call(logger: LoggingSurface, input: str, api_key: str | None, additional_args: dict[str, object]) -> None: logger.pre_call(input=input, api_key=api_key, additional_args=additional_args) diff --git a/litellm/rust_bridge/host/model_capabilities.py b/litellm/rust_bridge/host/model_capabilities.py index cbf59dfb97b..03254200e86 100644 --- a/litellm/rust_bridge/host/model_capabilities.py +++ b/litellm/rust_bridge/host/model_capabilities.py @@ -32,5 +32,6 @@ def anthropic_model_capabilities(model: str, custom_llm_provider: str | None) -> "supports_sampling_params": AnthropicModelInfo._supports_sampling_params(resolved_model), # pyright: ignore[reportPrivateUsage] # same gate the handler applies "supports_speed": AnthropicConfig._model_supports_speed_param(resolved_model, provider), # pyright: ignore[reportPrivateUsage] # same gate the handler applies "supports_mid_conversation_system": supports("supports_mid_conversation_system"), + "supports_prompt_caching": litellm.utils.supports_prompt_caching(resolved_model, provider), "effort_tiers": {level: tier(level) for level in ("minimal", "low", "medium", "high", "xhigh", "max")}, } diff --git a/litellm/rust_bridge/host/routes/messages.py b/litellm/rust_bridge/host/routes/messages.py index 1024141d5a2..00ac08c57f5 100644 --- a/litellm/rust_bridge/host/routes/messages.py +++ b/litellm/rust_bridge/host/routes/messages.py @@ -15,7 +15,8 @@ from litellm.rust_bridge import failures from litellm.rust_bridge.public_call import known_params, optional_str from litellm.types.llms.anthropic_messages.anthropic_response import AnthropicMessagesResponse -_DROP_PATHS: Final = TypeAdapter(list[object]) +_OBJECTS: Final = TypeAdapter(tuple[object, ...]) +_STR_MAPPING: Final = TypeAdapter(Mapping[str, object]) _BODY_FIELDS: Final = frozenset( { "max_tokens", @@ -50,6 +51,10 @@ class MessagesSettings: reasoning_auto_summary: bool additional_drop_params: Sequence[str] request_timeout: float | None + cache_control_injection_points: Sequence[Mapping[str, object]] + enable_prompt_caching: bool + prompt_caching_ttl: str | None + user_agent: str | None def response(value: Mapping[str, object]) -> AnthropicMessagesResponse: @@ -87,12 +92,41 @@ def _drop_params(kwargs: Mapping[str, object]) -> bool: return bool(litellm.drop_params) or normalize_drop_params(kwargs.get("drop_params")) is True -def _additional_drop_params(kwargs: Mapping[str, object]) -> tuple[str, ...]: +def _objects(value: object) -> tuple[object, ...]: try: - configured: Final = _DROP_PATHS.validate_python(kwargs.get("additional_drop_params")) + return _OBJECTS.validate_python(value) except ValidationError: return () - return tuple(path for path in configured if isinstance(path, str)) + + +def _additional_drop_params(kwargs: Mapping[str, object]) -> tuple[str, ...]: + return tuple(path for path in _objects(kwargs.get("additional_drop_params")) if isinstance(path, str)) + + +def _str_mapping(value: object) -> Mapping[str, object] | None: + try: + return _STR_MAPPING.validate_python(value) + except ValidationError: + return None + + +def _cache_control_injection_points(kwargs: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + points: Final = (_str_mapping(item) for item in _objects(kwargs.get("cache_control_injection_points"))) + return tuple(point for point in points if point is not None) + + +def _user_agent(kwargs: Mapping[str, object]) -> str | None: + request: Final = _str_mapping(kwargs.get("proxy_server_request")) or {} + headers: Final = _str_mapping(request.get("headers")) or {} + return next( + (value for key, value in headers.items() if key.lower() == "user-agent" and isinstance(value, str)), + None, + ) + + +def _prompt_caching_ttl() -> str | None: + value: Final = getattr(litellm, "anthropic_prompt_caching_ttl", None) + return value if value in ("5m", "1h") else None def settings(kwargs: Mapping[str, object]) -> dict[str, object]: @@ -102,5 +136,10 @@ 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(), + cache_control_injection_points=_cache_control_injection_points(kwargs), + enable_prompt_caching=kwargs.get("enable_prompt_caching") is True + or litellm.enable_anthropic_prompt_caching is True, + prompt_caching_ttl=_prompt_caching_ttl(), + user_agent=_user_agent(kwargs), ) ) diff --git a/tests/rust-python-harness/strategies/drift/drift.json b/tests/rust-python-harness/strategies/drift/drift.json index 58dca7c1363..539b15bee36 100644 --- a/tests/rust-python-harness/strategies/drift/drift.json +++ b/tests/rust-python-harness/strategies/drift/drift.json @@ -18,6 +18,19 @@ "enterprise/pyproject.toml" ], "mappings": [ + { + "id": "anthropic-cache-control-hook", + "python_paths": [ + "litellm/integrations/anthropic_cache_control_hook.py", + "litellm/types/integrations/anthropic_cache_control_hook.py" + ], + "rust_paths": [ + "litellm-rust/crates/integrations/src/anthropic_cache_control_hook.rs", + "litellm-rust/crates/llms-types/src/integrations/anthropic_cache_control_hook.rs" + ], + "checkpoint_sha": "2dfeb80e04d971fa00f72c12e0bd4fb42bc3c027", + "checkpoint_state": "aligned" + }, { "id": "messages", "python_paths": [ diff --git a/tests/test_litellm_rust/messages/test_requests.py b/tests/test_litellm_rust/messages/test_requests.py index 92031e72aee..bc569096b6e 100644 --- a/tests/test_litellm_rust/messages/test_requests.py +++ b/tests/test_litellm_rust/messages/test_requests.py @@ -146,3 +146,37 @@ async def test_optional_positional_body_parameters_reach_the_wire(server: Record assert response["id"] == MESSAGES_RESPONSE["id"] body, _ = sent(server) assert (body["metadata"], body["system"], body["temperature"]) == (args[3], args[6], args[7]) + + +@pytest.mark.asyncio +async def test_client_cache_control_is_preserved_and_suppresses_gateway_defaults( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + model: Final = "claude-test-client-cache-control" + monkeypatch.setitem( + litellm.model_cost, + model, + {"litellm_provider": "anthropic", "mode": "chat", "supports_prompt_caching": True}, + ) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + messages: Final = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "client marked", + "cache_control": {"type": "ephemeral", "ttl": "1h"}, + } + ], + } + ] + + await call(server, model=f"anthropic/{model}", messages=messages, system="stable prefix") + + assert_served_natively(server) + body: Final = server.requests[0].body + assert isinstance(body, dict) + assert body["system"] == "stable prefix" + assert body["messages"] == messages + assert messages[0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} diff --git a/tests/test_litellm_rust/messages/test_settings.py b/tests/test_litellm_rust/messages/test_settings.py new file mode 100644 index 00000000000..0319de53c3e --- /dev/null +++ b/tests/test_litellm_rust/messages/test_settings.py @@ -0,0 +1,132 @@ +from typing import Final + +import pytest + +import litellm +from litellm.types.integrations.anthropic_cache_control_hook import GATEWAY_INJECTED_CACHE_METADATA_KEY +from tests.test_litellm_rust.messages.data import call +from tests.test_litellm_rust.support.recording_server import RecordingServer +from tests.test_litellm_rust.support.routing import assert_served_natively + +pytestmark = pytest.mark.requires_rust_extension + + +@pytest.mark.asyncio +async def test_global_prompt_caching_marks_system_and_latest_message( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + model: Final = "claude-test-cache-control" + monkeypatch.setitem( + litellm.model_cost, + model, + {"litellm_provider": "anthropic", "mode": "chat", "supports_prompt_caching": True}, + ) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + monkeypatch.setattr(litellm, "anthropic_prompt_caching_ttl", "invalid") + + await call( + server, + model=f"anthropic/{model}", + messages=[{"role": "user", "content": "latest turn"}], + system="stable prefix", + ) + + assert_served_natively(server) + body: Final = server.requests[0].body + assert isinstance(body, dict) + assert body["system"] == [{"type": "text", "text": "stable prefix", "cache_control": {"type": "ephemeral"}}] + assert body["messages"] == [ + { + "role": "user", + "content": [{"type": "text", "text": "latest turn", "cache_control": {"type": "ephemeral"}}], + } + ] + + +@pytest.mark.asyncio +async def test_request_prompt_caching_flag_uses_global_ttl_setting( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + model: Final = "claude-test-request-cache-control" + monkeypatch.setitem( + litellm.model_cost, + model, + {"litellm_provider": "anthropic", "mode": "chat", "supports_prompt_caching": True}, + ) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + monkeypatch.setattr(litellm, "anthropic_prompt_caching_ttl", "1h") + + await call( + server, + model=f"anthropic/{model}", + messages=[{"role": "user", "content": "latest turn"}], + system="stable prefix", + enable_prompt_caching=True, + ) + + assert_served_natively(server) + body: Final = server.requests[0].body + assert isinstance(body, dict) + assert body["system"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + assert body["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "1h"} + + +@pytest.mark.asyncio +async def test_configured_cache_control_points_are_applied_on_the_native_route( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + model: Final = "claude-test-configured-cache-control" + monkeypatch.setitem( + litellm.model_cost, + model, + {"litellm_provider": "anthropic", "mode": "chat", "supports_prompt_caching": True}, + ) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", False) + + await call( + server, + model=f"anthropic/{model}", + messages=[{"role": "user", "content": "latest turn"}], + cache_control_injection_points=[ + {"location": "message", "index": "-1", "control": {"type": "ephemeral", "ttl": "5m"}} + ], + ) + + assert_served_natively(server) + body: Final = server.requests[0].body + assert isinstance(body, dict) + assert body["messages"][0]["content"][0]["cache_control"] == {"type": "ephemeral", "ttl": "5m"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("system", "expected"), + [ + ("stable prefix", "deployment-1"), + ([{"type": "text", "text": "marked", "cache_control": {"type": "ephemeral"}}], None), + ], +) +async def test_native_injection_names_the_deployment_in_request_metadata( + server: RecordingServer, monkeypatch: pytest.MonkeyPatch, system: object, expected: str | None +) -> None: + model: Final = "claude-test-gateway-injection-metadata" + monkeypatch.setitem( + litellm.model_cost, + model, + {"litellm_provider": "anthropic", "mode": "chat", "supports_prompt_caching": True}, + ) + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + metadata: Final[dict[str, object]] = {"existing": "kept"} + + await call( + server, + model=f"anthropic/{model}", + messages=[{"role": "user", "content": "latest turn"}], + system=system, + litellm_metadata=metadata, + model_info={"id": "deployment-1"}, + ) + + assert_served_natively(server) + assert metadata.get(GATEWAY_INJECTED_CACHE_METADATA_KEY) == expected + assert metadata["existing"] == "kept" diff --git a/tests/unit/rust_bridge/host/routes/test_messages.py b/tests/unit/rust_bridge/host/routes/test_messages.py index abf0e055e73..10c0e1c7cbc 100644 --- a/tests/unit/rust_bridge/host/routes/test_messages.py +++ b/tests/unit/rust_bridge/host/routes/test_messages.py @@ -35,6 +35,8 @@ 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, "enable_anthropic_prompt_caching", False) + monkeypatch.setattr(litellm, "anthropic_prompt_caching_ttl", None) projected: Final = messages_host.settings({"drop_params": "true", "additional_drop_params": ["metadata.user_id"]}) @@ -43,9 +45,64 @@ 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, + "cache_control_injection_points": (), + "enable_prompt_caching": False, + "prompt_caching_ttl": None, + "user_agent": None, } +@pytest.mark.parametrize( + ("global_flag", "request_flag", "expected"), + [ + (False, None, False), + (False, True, True), + (False, False, False), + (True, None, True), + (True, False, True), + ], +) +def test_prompt_caching_uses_global_or_true_request_flag( + monkeypatch: pytest.MonkeyPatch, + global_flag: bool, + request_flag: bool | None, + expected: bool, +) -> None: + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", global_flag) + kwargs: Final = {} if request_flag is None else {"enable_prompt_caching": request_flag} + + assert messages_host.settings(kwargs)["enable_prompt_caching"] is expected + + +@pytest.mark.parametrize( + ("ttl", "expected"), + [("5m", "5m"), ("1h", "1h"), ("5m ", None), ("forever", None), (None, None)], +) +def test_prompt_caching_ttl_accepts_only_supported_values( + monkeypatch: pytest.MonkeyPatch, ttl: object, expected: str | None +) -> None: + monkeypatch.setattr(litellm, "anthropic_prompt_caching_ttl", ttl) + + assert messages_host.settings({})["prompt_caching_ttl"] == expected + + +def test_settings_project_cache_injection_points_and_case_insensitive_user_agent() -> None: + point: Final = {"location": "message", "index": "-1"} + projected: Final = messages_host.settings( + { + "cache_control_injection_points": [point, "ignored"], + "proxy_server_request": {"headers": {"X-Other": "x", "USER-AGENT": "claude-code/1.0"}}, + } + ) + + assert projected["cache_control_injection_points"] == (point,) + assert projected["user_agent"] == "claude-code/1.0" + assert "cache_control_injection_points" not in messages_host.body_params( + {"cache_control_injection_points": [point], "enable_prompt_caching": True} + ) + assert "enable_prompt_caching" not in messages_host.body_params({"enable_prompt_caching": True}) + + @pytest.mark.parametrize( ("timeout", "explicitly_set", "expected"), [ diff --git a/tests/unit/rust_bridge/host/test_model_capabilities.py b/tests/unit/rust_bridge/host/test_model_capabilities.py index 4d73a29d078..371420cb9de 100644 --- a/tests/unit/rust_bridge/host/test_model_capabilities.py +++ b/tests/unit/rust_bridge/host/test_model_capabilities.py @@ -73,3 +73,11 @@ def test_capability_source_observes_runtime_registration_changes(monkeypatch: py assert before["supports_output_config"] is True assert after["supports_reasoning"] is False assert after["supports_output_config"] is False + + +def test_prompt_caching_capability_comes_from_the_model_map(monkeypatch: pytest.MonkeyPatch) -> None: + _flag_model(monkeypatch, "claude-test-cache", supports_prompt_caching=True) + + capabilities: Final = anthropic_model_capabilities("anthropic/claude-test-cache", None) + + assert capabilities["supports_prompt_caching"] is True