mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-25 05:54:53 +00:00
feat(rust-messages): inject cache_control breakpoints natively
This commit is contained in:
Devin AI
parent
120fa92e55
commit
1d72fee475
42 changed files with 1852 additions and 26 deletions
|
|
@ -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/<path>.py` to `crates/<crate>/src/<path>.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;`
|
||||
|
|
|
|||
11
litellm-rust/Cargo.lock
generated
11
litellm-rust/Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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" }
|
||||
|
|
|
|||
|
|
@ -21,6 +21,10 @@
|
|||
"litellm_params",
|
||||
"custom_llm_provider"
|
||||
],
|
||||
"record_gateway_injection": [
|
||||
"kwargs",
|
||||
"added"
|
||||
],
|
||||
"pre_call": [
|
||||
"logger",
|
||||
"input",
|
||||
|
|
|
|||
|
|
@ -374,6 +374,12 @@ impl LegacyLogging {
|
|||
wire: Box<WireRequest>,
|
||||
context: &RequestContext,
|
||||
) -> PyResult<HookStep<Self, Box<WireRequest>>> {
|
||||
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<usize>,
|
||||
) {
|
||||
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));
|
||||
|
|
|
|||
|
|
@ -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")]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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?;
|
||||
|
|
|
|||
|
|
@ -185,6 +185,7 @@ fn dispatching_call() -> TestMachine {
|
|||
optional_params: json!({}),
|
||||
secret_fields: Vec::new(),
|
||||
api_key: None,
|
||||
gateway_injected_cache_blocks: 0,
|
||||
},
|
||||
)
|
||||
.await?;
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -30,6 +30,7 @@ pub struct RequestContext {
|
|||
pub secret_fields: Vec<String>,
|
||||
/// The credential the route resolved for the provider call.
|
||||
pub api_key: Option<litellm_auth::SecretValue>,
|
||||
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,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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"] }
|
||||
|
|
|
|||
|
|
@ -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?;
|
||||
|
|
|
|||
|
|
@ -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<Duration>,
|
||||
/// The caller's own credential, reported to the host beside the wire request.
|
||||
pub(super) api_key: Option<SecretValue>,
|
||||
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,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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<String>,
|
||||
#[serde(default)]
|
||||
pub request_timeout: Option<f64>,
|
||||
#[serde(default)]
|
||||
pub cache_control_injection_points: Vec<CacheControlInjectionPoint>,
|
||||
#[serde(default)]
|
||||
pub enable_prompt_caching: bool,
|
||||
#[serde(default)]
|
||||
pub prompt_caching_ttl: Option<CacheControlTtl>,
|
||||
#[serde(default)]
|
||||
pub user_agent: Option<String>,
|
||||
}
|
||||
|
||||
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));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -59,6 +59,7 @@ impl<'a, H> OcrCallHooks<'a, H> {
|
|||
.cloned()
|
||||
.collect(),
|
||||
api_key: request.connection.api_key.clone(),
|
||||
gateway_injected_cache_blocks: 0,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
25
litellm-rust/crates/integrations/AGENTS.md
Normal file
25
litellm-rust/crates/integrations/AGENTS.md
Normal file
|
|
@ -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`
|
||||
14
litellm-rust/crates/integrations/Cargo.toml
Normal file
14
litellm-rust/crates/integrations/Cargo.toml
Normal file
|
|
@ -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
|
||||
|
|
@ -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<CacheControlTtl>,
|
||||
}
|
||||
|
||||
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<CacheControlInjectionPoint> {
|
||||
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<CacheControlTtl>) -> 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<Message>,
|
||||
max_blocks: usize,
|
||||
) -> (Vec<Message>, 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<String, Vec<usize>> {
|
||||
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<String, Vec<usize>>,
|
||||
) -> Vec<usize> {
|
||||
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<usize> {
|
||||
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::<usize>()
|
||||
}
|
||||
|
||||
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<MessagesTool>) -> 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<String, Value>) -> 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<CacheControlInjectionPoint> {
|
||||
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<CacheControlTtl>) -> 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<usize>,
|
||||
#[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<bool> = match &value["system"] {
|
||||
Value::Array(blocks) => blocks
|
||||
.iter()
|
||||
.map(|block| block.get("cache_control").is_some())
|
||||
.collect(),
|
||||
_ => Vec::new(),
|
||||
};
|
||||
let expected: Vec<bool> = (0..marks.len())
|
||||
.map(|index| Some(index) == marked_block)
|
||||
.collect();
|
||||
assert_eq!(marks, expected);
|
||||
}
|
||||
}
|
||||
3
litellm-rust/crates/integrations/src/lib.rs
Normal file
3
litellm-rust/crates/integrations/src/lib.rs
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
pub mod anthropic_cache_control_hook;
|
||||
|
||||
pub use anthropic_cache_control_hook::{Injected, InjectionInputs, maybe_inject_cache_control};
|
||||
|
|
@ -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());
|
||||
|
|
|
|||
|
|
@ -114,15 +114,31 @@ impl ContentBlock {
|
|||
#[derive(Default)]
|
||||
pub struct CacheControl {
|
||||
#[serde(rename = "type", skip_serializing_if = "Option::is_none")]
|
||||
pub cache_type: Option<String>,
|
||||
pub cache_type: Option<Recognized<CacheControlType>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub ttl: Option<String>,
|
||||
pub ttl: Option<Recognized<CacheControlTtl>>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub scope: Option<String>,
|
||||
#[serde(flatten)]
|
||||
pub extra: Map<String, Value>,
|
||||
}
|
||||
|
||||
#[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<Recognized<CacheControlType>>,
|
||||
#[case] ttl: Option<Recognized<CacheControlTtl>>,
|
||||
) {
|
||||
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<T: serde::de::DeserializeOwned + serde::Serialize>(value: &Value) -> Value {
|
||||
let parsed: T = serde_json::from_value(value.clone()).unwrap();
|
||||
serde_json::to_value(parsed).unwrap()
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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<CacheControl>,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug, PartialEq)]
|
||||
pub enum CacheControlInjectionPoint {
|
||||
Message(CacheControlMessageInjectionPoint),
|
||||
ToolConfig { control: Option<CacheControl> },
|
||||
Unknown(Value),
|
||||
}
|
||||
|
||||
impl<'de> Deserialize<'de> for CacheControlInjectionPoint {
|
||||
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
|
||||
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::<i64>()
|
||||
.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<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
|
||||
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<String, Value>) -> serde_json::Result<Option<CacheControl>> {
|
||||
object
|
||||
.get("control")
|
||||
.filter(|control| !control.is_null())
|
||||
.cloned()
|
||||
.map(serde_json::from_value)
|
||||
.transpose()
|
||||
}
|
||||
|
||||
fn role_target(object: &Map<String, Value>) -> 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::<CacheControlInjectionPoint>(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::<CacheControlInjectionPoint>(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);
|
||||
}
|
||||
}
|
||||
3
litellm-rust/crates/llms-types/src/integrations/mod.rs
Normal file
3
litellm-rust/crates/llms-types/src/integrations/mod.rs
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
mod anthropic_cache_control_hook;
|
||||
|
||||
pub use anthropic_cache_control_hook::*;
|
||||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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 } => {
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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<Vec<&str>> = 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<Vec<Vec<(&str, &str)>>> = 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<Vec<(&str, &str)>> {
|
||||
let stripped = text.trim();
|
||||
if stripped.contains(['\n', '\r']) {
|
||||
return None;
|
||||
}
|
||||
let fields = stripped.strip_prefix(BILLING_HEADER_PREFIX)?;
|
||||
let parsed: Option<Vec<(&str, &str)>> = 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<Value>,
|
||||
#[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
|
||||
);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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(),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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")},
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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": [
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
132
tests/test_litellm_rust/messages/test_settings.py
Normal file
132
tests/test_litellm_rust/messages/test_settings.py
Normal file
|
|
@ -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"
|
||||
|
|
@ -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"),
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue