diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7ae8e983361..caf03f04f7a 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", 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..abffaadff49 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.clone(), + }, + ); + 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/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index c90308a4d9b..714ca5bdda2 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1971,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), } ); 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