diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index e62a9463cb5..24d8cfcfe54 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -8,6 +8,13 @@ For fieldless enums, derive `strum::VariantArray` and use `VARIANTS` instead of Use the derived conversions directly (`<&'static str>::from(x)` / `.into()`, `str::parse`) with no `as_str`/`parse` wrapper that only delegates, and spell each variant with explicit `#[strum(serialize = "...")]` instead of `serialize_all` +## Python ports + +- Prefer mirroring the outer Python structure. Map `litellm/.py` to `crates//src/.rs` in the crate that mirrors that folder +- Keep class and public function names, with Rust acronym casing and no leading underscore for private helpers. A staticmethod-only class becomes free functions. A class with a Python base becomes a struct implementing the matching trait +- Inside functions, design freely. Prefer typed values over dict inspection, new values over mutation, and tagged unions over string modes. Pin behavior with Rust tests +- List skipped behavior under `Known partial ports` in the nearest AGENTS.md and test deliberate divergences + ## Test placement - Never create a `tests.rs` (or `test.rs`) file under `src/`, and never `#[path = "tests.rs"] mod tests;` diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 4ddce5edc78..7ae8e983361 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -4141,6 +4141,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" @@ -4188,6 +4198,7 @@ dependencies = [ "serde_json", "serde_with", "strum", + "unicode-properties", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index f04eb115071..1608f6d45a8 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -49,6 +49,7 @@ litellm-secrets-cyberark = { path = "crates/secrets-cyberark" } litellm-http = { path = "crates/http" } litellm-llms = { path = "crates/llms" } litellm-llms-types = { path = "crates/llms-types" } +litellm-integrations = { path = "crates/integrations" } litellm-core-utils = { path = "crates/core-utils" } litellm-db = { path = "crates/db" } litellm-db-testing = { path = "crates/db-testing" } diff --git a/litellm-rust/crates/integrations/AGENTS.md b/litellm-rust/crates/integrations/AGENTS.md new file mode 100644 index 00000000000..24b7c8a38a8 --- /dev/null +++ b/litellm-rust/crates/integrations/AGENTS.md @@ -0,0 +1,20 @@ +# 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 returned for downstream handling while a slot remains. 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 +- Python integer strings use the default digit limit. Host overrides of that limit are not yet projected + +# references + +- `litellm/integrations/anthropic_cache_control_hook.py` +- `litellm-rust/crates/llms-types/src/integrations/anthropic_cache_control_hook.rs` diff --git a/litellm-rust/crates/integrations/Cargo.toml b/litellm-rust/crates/integrations/Cargo.toml new file mode 100644 index 00000000000..8e7eab86bd7 --- /dev/null +++ b/litellm-rust/crates/integrations/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "litellm-integrations" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-llms-types.workspace = true +serde.workspace = true +serde_json.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/integrations/src/anthropic_cache_control_hook.rs b/litellm-rust/crates/integrations/src/anthropic_cache_control_hook.rs new file mode 100644 index 00000000000..5c7f5408d6b --- /dev/null +++ b/litellm-rust/crates/integrations/src/anthropic_cache_control_hook.rs @@ -0,0 +1,1064 @@ +use std::collections::{HashMap, HashSet}; + +use litellm_llms_types::{ + formats::messages::{ + CacheControl, CacheControlTtl, CacheControlType, ContentBlock, Message, MessageContent, + MessagesOptionalParams, MessagesRequest, MessagesTool, SystemPrompt, + }, + integrations::{CacheControlInjectionPoint, CacheControlMessageInjectionPoint}, + message_role::MessageRole, + recognized::Recognized, + serde_compat::ParsedInteger, +}; +use serde_json::{Map, Value}; + +const MAX_CACHE_CONTROL_BLOCKS: usize = 4; +const CACHE_CONTROL_KEY: &str = "cache_control"; + +pub struct InjectionInputs<'a> { + pub points: &'a [CacheControlInjectionPoint], + pub enable_prompt_caching: bool, + pub supports_anthropic_cache_control: bool, + pub one_shot_subagent: bool, + pub ttl: Option, +} + +pub struct Injected { + pub request: MessagesRequest, + pub added: usize, + pub remaining_points: Vec, +} + +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +enum Target<'a> { + Index(i64), + Role(&'a str), + None, +} + +pub fn maybe_inject_cache_control( + request: MessagesRequest, + inputs: &InjectionInputs<'_>, +) -> Injected { + if inputs.points.iter().any(|point| match point { + CacheControlInjectionPoint::Unknown(value) => value.is_object(), + _ => true, + }) { + return apply_to_anthropic_messages_request(request, inputs.points); + } + let defaults = get_default_injection_points(&request, inputs); + apply_to_anthropic_messages_request(request, &defaults) +} + +fn get_default_injection_points( + request: &MessagesRequest, + inputs: &InjectionInputs<'_>, +) -> Vec { + if !inputs.enable_prompt_caching + || !inputs.supports_anthropic_cache_control + || inputs.one_shot_subagent + || request_has_cache_control(request) + { + return Vec::new(); + } + let control = default_control(inputs.ttl.clone()); + vec![ + CacheControlInjectionPoint::Message(CacheControlMessageInjectionPoint::new( + None, + Some(MessageRole::System), + Some(control.clone()), + )), + CacheControlInjectionPoint::Message(CacheControlMessageInjectionPoint::new( + Some(-1), + None, + Some(control), + )), + ] +} + +fn default_control(ttl: Option) -> CacheControl { + CacheControl { + cache_type: Some(Some(CacheControlType::Ephemeral)), + ttl: ttl + .filter(|ttl| matches!(ttl, CacheControlTtl::FiveMinutes | CacheControlTtl::OneHour)) + .map(Some), + ..CacheControl::default() + } +} + +fn apply_to_anthropic_messages_request( + request: MessagesRequest, + injection_points: &[CacheControlInjectionPoint], +) -> Injected { + if injection_points.is_empty() { + return Injected { + request, + added: 0, + remaining_points: Vec::new(), + }; + } + let (system_points, message_points): (Vec<_>, Vec<_>) = injection_points + .iter() + .filter_map(message_point) + .partition(|point| is_system_point(point)); + let external_breakpoints = count_external_cache_breakpoints(&request); + let max_blocks = MAX_CACHE_CONTROL_BLOCKS + .saturating_sub(blocks_reserved_outside_messages(injection_points) + external_breakpoints); + 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())), + 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)), + ); + let added = usize::from(system_marked) + added_to_messages; + Injected { + request: MessagesRequest { + messages, + params: MessagesOptionalParams { + system, + ..request.params + }, + ..request + }, + added, + remaining_points: points_with_a_slot_left( + injection_points, + used + added + external_breakpoints, + ), + } +} + +fn message_point(point: &CacheControlInjectionPoint) -> Option<&CacheControlMessageInjectionPoint> { + match point { + CacheControlInjectionPoint::Message(point) => Some(point), + _ => None, + } +} + +fn is_system_point(point: &CacheControlMessageInjectionPoint) -> bool { + point.role().map(AsRef::as_ref) == Some(MessageRole::System.as_ref()) +} + +fn points_with_a_slot_left( + points: &[CacheControlInjectionPoint], + breakpoints_on_wire: usize, +) -> Vec { + points + .iter() + .filter(|point| match point { + CacheControlInjectionPoint::Message(_) => false, + CacheControlInjectionPoint::ToolConfig(_) => { + breakpoints_on_wire < MAX_CACHE_CONTROL_BLOCKS + } + CacheControlInjectionPoint::Unknown(value) => value.is_object(), + }) + .cloned() + .collect() +} + +fn blocks_reserved_outside_messages(injection_points: &[CacheControlInjectionPoint]) -> usize { + usize::from( + injection_points + .iter() + .any(|point| matches!(point, CacheControlInjectionPoint::ToolConfig(_))), + ) +} + +fn with_block_content(message: Message) -> Message { + match message.content { + MessageContent::Text(text) => Message { + content: MessageContent::Blocks(vec![ContentBlock::text(text)]), + ..message + }, + MessageContent::Blocks(_) => message, + } +} + +fn apply_message_injections( + points: &[&CacheControlMessageInjectionPoint], + mut messages: Vec, + max_blocks: usize, +) -> (Vec, usize) { + let roles = role_indices(&messages); + let mut seen = HashSet::new(); + let mut added = 0; + for point in points { + let target = message_target(point); + if !seen.insert(target) { + continue; + } + for index in resolve_target_indices(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()) + { + added += 1; + } + } + } + (messages, added) +} + +fn role_indices(messages: &[Message]) -> HashMap> { + messages + .iter() + .enumerate() + .fold(HashMap::new(), |mut roles, (index, message)| { + roles.entry(message.role.clone()).or_default().push(index); + roles + }) +} + +fn message_target(point: &CacheControlMessageInjectionPoint) -> Target<'_> { + match point.parsed_index() { + ParsedInteger::Value(index) => Target::Index(index), + ParsedInteger::OutOfRange => Target::None, + ParsedInteger::Invalid => point + .role() + .map(|role| Target::Role(role.as_ref())) + .unwrap_or(Target::None), + } +} + +fn resolve_target_indices<'a>( + target: Target<'_>, + length: usize, + roles: &'a HashMap>, +) -> impl Iterator + 'a { + let (index, matching_roles): (Option, &[usize]) = match target { + Target::None => (None, &[]), + 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), + &[], + ) + } + Target::Role(role) => (None, roles.get(role).map(Vec::as_slice).unwrap_or_default()), + }; + index.into_iter().chain(matching_roles.iter().copied()) +} + +fn safe_insert_cache_control_in_message( + message: &mut Message, + control: Option<&CacheControl>, +) -> bool { + let MessageContent::Blocks(blocks) = &mut message.content else { + return false; + }; + blocks.last_mut().is_some_and(|block| { + block.cache_control = Some(injected_control(control)); + true + }) +} + +fn system_marking_index(system: &SystemPrompt) -> Option { + match system { + SystemPrompt::Text(_) => Some(0), + SystemPrompt::Blocks(blocks) if count_cache_control_in_blocks(blocks) > 0 => None, + SystemPrompt::Blocks(blocks) => blocks.len().checked_sub(1), + } +} + +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 injected_control(control: Option<&CacheControl>) -> CacheControl { + control + .filter(|control| { + control.cache_type.is_some() + || control.ttl.is_some() + || control.scope.is_some() + || !control.extra.is_empty() + }) + .cloned() + .unwrap_or_else(|| default_control(None)) +} + +fn request_has_cache_control(request: &MessagesRequest) -> bool { + count_request_cache_breakpoints(&request.messages, request.params.system.as_ref()) + + count_external_cache_breakpoints(request) + > 0 +} + +fn count_request_cache_breakpoints(messages: &[Message], system: Option<&SystemPrompt>) -> usize { + let system_blocks = match system { + Some(SystemPrompt::Blocks(blocks)) => count_cache_control_in_blocks(blocks), + Some(SystemPrompt::Text(_)) | None => 0, + }; + system_blocks + + messages + .iter() + .map(count_cache_control_blocks) + .sum::() +} + +fn count_external_cache_breakpoints(request: &MessagesRequest) -> usize { + let tool_blocks = request + .params + .tools + .iter() + .flatten() + .filter(|tool| tool_carries_cache_breakpoint(tool)) + .count(); + usize::from(carries_cache_breakpoint(&request.params.extra)) + tool_blocks +} + +fn count_cache_control_blocks(message: &Message) -> usize { + let content_blocks = match &message.content { + MessageContent::Blocks(blocks) => count_cache_control_in_blocks(blocks), + MessageContent::Text(_) => 0, + }; + usize::from(carries_cache_breakpoint(&message.extra)) + content_blocks +} + +fn message_has_cache_control(message: &Message) -> bool { + count_cache_control_blocks(message) > 0 +} + +fn count_cache_control_in_blocks(blocks: &[ContentBlock]) -> usize { + blocks + .iter() + .filter(|block| block.cache_control.is_some()) + .count() +} + +fn tool_carries_cache_breakpoint(tool: &Recognized) -> bool { + let fields = match tool { + Recognized::Known( + MessagesTool::Advisor { extra } + | MessagesTool::ToolSearchRegex { extra } + | MessagesTool::ToolSearchBm25 { extra }, + ) => extra, + Recognized::Unrecognized(Value::Object(fields)) => fields, + Recognized::Unrecognized(_) => return false, + }; + carries_cache_breakpoint(fields) + || fields + .get("function") + .and_then(Value::as_object) + .is_some_and(carries_cache_breakpoint) +} + +fn carries_cache_breakpoint(fields: &Map) -> bool { + fields + .get(CACHE_CONTROL_KEY) + .is_some_and(|value| !value.is_null()) +} + +#[cfg(test)] +mod tests { + use rstest::{fixture, rstest}; + use serde_json::json; + + use super::*; + use litellm_llms_types::serde_compat::PYTHON_DEFAULT_INT_MAX_STR_DIGITS; + + fn request(value: Value) -> MessagesRequest { + serde_json::from_value(value).expect("valid Messages request") + } + + fn configured_points(value: Value) -> Vec { + serde_json::from_value(value).expect("valid cache-control injection points") + } + + fn configured(points: &[CacheControlInjectionPoint]) -> InjectionInputs<'_> { + InjectionInputs { + points, + enable_prompt_caching: false, + supports_anthropic_cache_control: false, + one_shot_subagent: false, + ttl: None, + } + } + + fn defaults(ttl: Option) -> InjectionInputs<'static> { + InjectionInputs { + points: &[], + enable_prompt_caching: true, + supports_anthropic_cache_control: true, + one_shot_subagent: false, + ttl, + } + } + + fn inject(request: MessagesRequest, inputs: &InjectionInputs<'_>) -> (Value, usize) { + let injected = maybe_inject_cache_control(request, inputs); + ( + serde_json::to_value(injected.request).unwrap(), + injected.added, + ) + } + + #[fixture] + fn simple_request() -> MessagesRequest { + request(json!({ + "model": "claude-test", + "max_tokens": 32, + "system": "system prompt", + "messages": [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "middle"}, + {"role": "user", "content": "last"} + ] + })) + } + + #[rstest] + #[case::primitives_only(json!([null, 7, true, "point", []]), 2)] + #[case::unknown_object(json!([null, {"location": "future"}]), 0)] + #[case::message_among_primitives(json!([null, {"location": "message", "index": " 0 "}, 7]), 1)] + fn configured_points_filter_non_objects_before_selecting_defaults( + simple_request: MessagesRequest, + #[case] points: Value, + #[case] expected_added: usize, + ) { + let points = configured_points(points); + let inputs = InjectionInputs { + points: &points, + ..defaults(None) + }; + + let (value, added) = inject(simple_request, &inputs); + + assert_eq!(added, expected_added); + let marked = match expected_added { + 2 => &value["messages"][2]["content"][0], + 1 => &value["messages"][0]["content"][0], + _ => return, + }; + assert_eq!(marked["cache_control"], json!({"type": "ephemeral"})); + } + + #[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::non_object_points_fall_back(json!([null, 5, "x"]))] + fn non_object_configured_points_fall_back_to_defaults( + simple_request: MessagesRequest, + #[case] point_values: Value, + ) { + let points = configured_points(point_values); + let inputs = InjectionInputs { + points: &points, + enable_prompt_caching: true, + supports_anthropic_cache_control: true, + one_shot_subagent: false, + ttl: None, + }; + let actual = inject(simple_request.clone(), &inputs); + let expected = inject(simple_request, &defaults(None)); + + assert_eq!(actual, expected); + assert_eq!(actual.1, 2); + } + + #[rstest] + fn object_unknown_configured_points_prevent_defaults(simple_request: MessagesRequest) { + let points = configured_points(json!([{"location": "future"}])); + let inputs = InjectionInputs { + points: &points, + enable_prompt_caching: true, + supports_anthropic_cache_control: true, + one_shot_subagent: false, + ttl: None, + }; + let (actual, added) = inject(simple_request, &inputs); + + assert_eq!(added, 0); + assert!( + actual["messages"][2]["content"][0] + .get("cache_control") + .is_none() + ); + } + + #[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] + #[case::unknown_ttl(json!({"type": "ephemeral", "ttl": "future"}), json!({"type": "ephemeral", "ttl": "future"}))] + #[case::missing_type(json!({"ttl": "1h"}), json!({"ttl": "1h"}))] + #[case::null_type(json!({"type": null, "ttl": "1h"}), json!({"type": null, "ttl": "1h"}))] + #[case::null_ttl_and_scope(json!({"type": "ephemeral", "ttl": null, "scope": null}), json!({"type": "ephemeral", "ttl": null, "scope": null}))] + #[case::extension(json!({"custom": {"enabled": true}}), json!({"custom": {"enabled": true}}))] + #[case::empty_control(json!({}), json!({"type": "ephemeral"}))] + #[case::null_control(Value::Null, json!({"type": "ephemeral"}))] + fn configured_controls_preserve_fields(#[case] control: Value, #[case] expected: Value) { + let points = + configured_points(json!([{"location": "message", "index": 0, "control": control}])); + let request = request(json!({ + "model": "claude-test", "max_tokens": 32, + "messages": [{"role": "user", "content": "prompt"}] + })); + + let (value, added) = inject(request, &configured(&points)); + + assert_eq!(added, 1); + assert_eq!( + value["messages"][0]["content"][0]["cache_control"], + expected + ); + } + + #[rstest] + #[case::unknown_ttl(CacheControlTtl::Other("future".into()))] + #[case::empty_ttl(CacheControlTtl::Other(String::new()))] + fn defaults_omit_unknown_ttl(simple_request: MessagesRequest, #[case] ttl: CacheControlTtl) { + let actual = inject(simple_request.clone(), &defaults(Some(ttl))); + let expected = inject(simple_request, &defaults(None)); + + assert_eq!(actual, expected); + } + + #[rstest] + #[case::oversized_string(json!("9223372036854775808"), None)] + #[case::oversized_number(json!(u64::MAX), None)] + #[case::negative_overflow(json!("-9223372036854775809"), None)] + #[case::underscores(json!("0_0"), Some(0))] + #[case::python_digit_limit(json!("0".repeat(PYTHON_DEFAULT_INT_MAX_STR_DIGITS + 1)), Some(1))] + #[case::python_unicode_digit_limit(json!("٠".repeat(PYTHON_DEFAULT_INT_MAX_STR_DIGITS + 1)), Some(1))] + #[case::unicode_digits(json!("٠"), Some(0))] + #[case::invalid_index_uses_role(json!("0__0"), Some(1))] + #[case::missing_index_uses_role(Value::Null, Some(1))] + fn configured_indices_follow_python_resolution( + #[case] index: Value, + #[case] expected_target: Option, + ) { + let points = configured_points( + json!([{"location": "message", "index": index, "role": "assistant"}]), + ); + let request = request(json!({ + "model": "claude-test", "max_tokens": 32, + "messages": [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "second"} + ] + })); + + let (value, added) = inject(request, &configured(&points)); + + assert_eq!(added, usize::from(expected_target.is_some())); + assert_eq!( + value["messages"][0]["content"][0] + .get("cache_control") + .is_some(), + expected_target == Some(0) + ); + assert_eq!( + value["messages"][1]["content"][0] + .get("cache_control") + .is_some(), + expected_target == Some(1) + ); + } + + #[rstest] + #[case::with_system(Some("system prompt"))] + #[case::without_system(None)] + fn configured_system_role_precedes_index_beside_client_marks(#[case] system: Option<&str>) { + let points = + configured_points(json!([{"location": "message", "index": 0, "role": "system"}])); + let request = request(json!({ + "model": "claude-test", "max_tokens": 32, "system": system, + "messages": [{"role": "user", "content": "prompt"}], + "cache_control": {"type": "ephemeral"} + })); + + let (value, added) = inject(request, &configured(&points)); + + assert_eq!(added, usize::from(system.is_some())); + assert_eq!( + value["system"][0].get("cache_control").cloned(), + system.map(|_| json!({"type": "ephemeral"})) + ); + assert!( + value["messages"][0]["content"][0] + .get("cache_control") + .is_none() + ); + } + + #[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 injected = maybe_inject_cache_control(simple_request, &configured(&points)); + let value = serde_json::to_value(injected.request).unwrap(); + + assert_eq!(injected.added, 0); + assert_eq!(injected.remaining_points, points); + assert_eq!(value["system"], "system prompt"); + assert!(value.get("cache_control_injection_points").is_none()); + assert!(value.get("tool_config").is_none()); + } + + #[rstest] + #[case::slot_available(3, 0, json!([ + {"location": "future", "custom": {"keep": true}}, + {"location": "tool_config", "control": null, "custom": "keep"} + ]))] + #[case::cap_full(4, 0, json!([{"location": "future", "custom": {"keep": true}}]))] + #[case::cap_exceeded(5, 0, json!([{"location": "future", "custom": {"keep": true}}]))] + #[case::cap_full_from_tools(3, 1, json!([{"location": "future", "custom": {"keep": true}}]))] + fn pending_points_preserve_extensions_and_only_forward_tools_with_a_slot( + #[case] marks: usize, + #[case] tool_marks: usize, + #[case] expected: Value, + ) { + let points = configured_points(json!([ + null, + {"location": "message", "index": 0}, + {"location": "future", "custom": {"keep": true}}, + {"location": "tool_config", "control": null, "custom": "keep"}, + 7 + ])); + let request = request(json!({ + "model": "claude-test", "max_tokens": 32, + "tools": vec![json!({ + "name": "lookup", "input_schema": {}, "cache_control": {"type": "ephemeral"} + }); tool_marks], + "messages": [{"role": "user", "content": vec![json!({ + "type": "text", "text": "marked", "cache_control": {"type": "ephemeral"} + }); marks]}] + })); + + let injected = maybe_inject_cache_control(request, &configured(&points)); + + assert_eq!(injected.added, 0); + assert_eq!( + serde_json::to_value(injected.remaining_points).unwrap(), + expected + ); + } + + #[rstest] + #[case::trailing_billing_block( + json!([ + {"type": "text", "text": "Instructions"}, + {"type": "text", "text": "x-anthropic-billing-header: cc_version=1"} + ]), + 1 + )] + #[case::only_billing_blocks( + json!([{"type": "text", "text": "x-anthropic-billing-header: cc_version=1"}]), + 0 + )] + #[case::billing_text(json!("x-anthropic-billing-header: cc_version=1"), 0)] + fn system_breakpoint_matches_python_final_block_placement( + #[case] system: Value, + #[case] marked_block: 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, 1); + assert_eq!( + value["system"][marked_block]["cache_control"], + json!({"type": "ephemeral"}) + ); + let blocks = value["system"].as_array().unwrap(); + assert_eq!( + blocks + .iter() + .filter(|block| block.get("cache_control").is_some()) + .count(), + 1 + ); + } +} diff --git a/litellm-rust/crates/integrations/src/lib.rs b/litellm-rust/crates/integrations/src/lib.rs new file mode 100644 index 00000000000..55c8d3c9b9a --- /dev/null +++ b/litellm-rust/crates/integrations/src/lib.rs @@ -0,0 +1,3 @@ +pub mod anthropic_cache_control_hook; + +pub use anthropic_cache_control_hook::{Injected, InjectionInputs, maybe_inject_cache_control}; diff --git a/litellm-rust/crates/llms-types/Cargo.toml b/litellm-rust/crates/llms-types/Cargo.toml index 4f190f34026..2119489d34b 100644 --- a/litellm-rust/crates/llms-types/Cargo.toml +++ b/litellm-rust/crates/llms-types/Cargo.toml @@ -16,6 +16,7 @@ serde.workspace = true serde_json.workspace = true serde_with.workspace = true strum.workspace = true +unicode-properties = { version = "0.1.4", default-features = false, features = ["general-category"] } [dev-dependencies] rstest.workspace = true diff --git a/litellm-rust/crates/llms-types/src/formats/chat_completions/content.rs b/litellm-rust/crates/llms-types/src/formats/chat_completions/content.rs index 7298145533c..ba080515368 100644 --- a/litellm-rust/crates/llms-types/src/formats/chat_completions/content.rs +++ b/litellm-rust/crates/llms-types/src/formats/chat_completions/content.rs @@ -161,6 +161,7 @@ mod tests { messages::ContentSource, }; + use crate::formats::messages::CacheControlType; 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(Some(CacheControlType::Ephemeral)) ); assert_eq!(breakpoint.mode, PromptCacheMode::Explicit); assert!(breakpoint.extra.is_empty()); diff --git a/litellm-rust/crates/llms-types/src/formats/messages/metadata.rs b/litellm-rust/crates/llms-types/src/formats/messages/metadata.rs index 8e92c4ce3d3..df64bf6f659 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/metadata.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/metadata.rs @@ -152,48 +152,6 @@ pub struct Safeguard { pub extra: Map, } -#[derive( - Clone, - Debug, - PartialEq, - Eq, - strum::Display, - strum::EnumString, - strum::AsRefStr, - serde_with::DeserializeFromStr, - serde_with::SerializeDisplay, -)] -#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] -#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))] -pub enum MessageRole { - #[strum(serialize = "user")] - User, - #[strum(serialize = "assistant")] - Assistant, - #[strum(serialize = "system")] - System, - #[strum(default, transparent)] - Other(String), -} - -impl MessageRole { - pub fn as_str(&self) -> &str { - self.as_ref() - } -} - -impl From for MessageRole { - fn from(value: String) -> Self { - value.parse().unwrap_or_else(|never| match never {}) - } -} - -impl From for String { - fn from(value: MessageRole) -> Self { - value.to_string() - } -} - #[derive( Clone, Debug, @@ -350,9 +308,7 @@ mod tests { use crate::formats::messages::McpServer; - use crate::formats::messages::{ - MessageRole, MessageType, MessagesCompaction, MessagesContainer, - }; + use crate::formats::messages::{MessageType, MessagesCompaction, MessagesContainer}; use crate::formats::messages::{ MessagesDiagnostics, MessagesDiagnosticsParam, MessagesMetadata, @@ -513,15 +469,6 @@ mod tests { assert!(serde_json::from_value::(json!({"type":"json_schema"})).is_err()); } - #[rstest] - #[case::user(json!("user"))] - #[case::assistant(json!("assistant"))] - #[case::system(json!("system"))] - #[case::future(json!("future_role"))] - fn message_roles_round_trip(#[case] wire: Value) { - round_trip::(wire); - } - #[rstest] #[case::message(json!("message"))] #[case::future(json!("future_type"))] diff --git a/litellm-rust/crates/llms-types/src/formats/messages/request.rs b/litellm-rust/crates/llms-types/src/formats/messages/request.rs index 9dbfa068025..529ead3b282 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/request.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/request.rs @@ -113,16 +113,86 @@ impl ContentBlock { #[macro_rules_attribute::apply(crate::wire_type)] #[derive(Default)] pub struct CacheControl { - #[serde(rename = "type", skip_serializing_if = "Option::is_none")] - pub cache_type: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub ttl: Option, - #[serde(skip_serializing_if = "Option::is_none")] - pub scope: Option, + #[serde( + rename = "type", + default, + skip_serializing_if = "Option::is_none", + with = "serde_with::rust::double_option" + )] + #[cfg_attr(feature = "schema", schemars(with = "Option"))] + pub cache_type: Option>, + #[serde( + default, + skip_serializing_if = "Option::is_none", + with = "serde_with::rust::double_option" + )] + #[cfg_attr(feature = "schema", schemars(with = "Option"))] + pub ttl: Option>, + #[serde( + default, + skip_serializing_if = "Option::is_none", + with = "serde_with::rust::double_option" + )] + #[cfg_attr(feature = "schema", schemars(with = "Option"))] + pub scope: Option>, #[serde(flatten)] pub extra: Map, } +#[derive( + Clone, + Debug, + PartialEq, + Eq, + Hash, + strum::Display, + strum::EnumString, + serde_with::DeserializeFromStr, + serde_with::SerializeDisplay, +)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))] +pub enum CacheControlType { + #[strum(serialize = "ephemeral")] + Ephemeral, + #[strum(default, transparent)] + Other(String), +} + +impl From for CacheControlType { + fn from(value: String) -> Self { + value.parse().unwrap_or_else(|never| match never {}) + } +} + +#[derive( + Clone, + Debug, + PartialEq, + Eq, + Hash, + strum::Display, + strum::EnumString, + serde_with::DeserializeFromStr, + serde_with::SerializeDisplay, +)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))] +pub enum CacheControlTtl { + #[strum(serialize = "5m")] + FiveMinutes, + #[strum(serialize = "1h")] + OneHour, + #[strum(default, transparent)] + Other(String), +} + +impl From for CacheControlTtl { + fn from(value: String) -> Self { + value.parse().unwrap_or_else(|never| match never {}) + } +} + #[macro_rules_attribute::apply(crate::wire_type)] pub struct Message { pub role: String, @@ -399,11 +469,94 @@ mod tests { use serde_json::{Map, Value}; + #[rstest] + #[case::known( + json!({"type": "ephemeral", "ttl": "1h"}), + Some(CacheControlType::Ephemeral), + Some(CacheControlTtl::OneHour) + )] + #[case::five_minutes( + json!({"type": "ephemeral", "ttl": "5m"}), + Some(CacheControlType::Ephemeral), + Some(CacheControlTtl::FiveMinutes) + )] + #[case::unknown_values_pass_through( + json!({"type": "persistent", "ttl": "future"}), + Some(CacheControlType::Other("persistent".into())), + Some(CacheControlTtl::Other("future".into())) + )] + #[case::variant_name_is_not_a_wire_value( + json!({"ttl": "FiveMinutes"}), + None, + Some(CacheControlTtl::Other("FiveMinutes".into())) + )] + #[case::empty_strings(json!({"type": "", "ttl": ""}), Some(CacheControlType::Other("".into())), Some(CacheControlTtl::Other("".into())))] + #[case::absent(json!({}), None, None)] + fn cache_control_types_known_values_and_keeps_unknown_ones( + #[case] input: Value, + #[case] cache_type: Option, + #[case] ttl: Option, + ) { + let control: CacheControl = serde_json::from_value(input.clone()).unwrap(); + + assert_eq!( + (control.cache_type.clone(), control.ttl.clone()), + (cache_type.map(Some), ttl.map(Some)) + ); + assert_eq!(serde_json::to_value(control).unwrap(), input); + } + + #[rstest] + #[case::null_type(json!({"type": null}))] + #[case::null_ttl(json!({"type": "ephemeral", "ttl": null}))] + #[case::null_scope(json!({"scope": null}))] + #[case::all_null(json!({"type": null, "ttl": null, "scope": null}))] + fn cache_control_preserves_explicit_null_fields(#[case] input: Value) { + let control: CacheControl = serde_json::from_value(input.clone()).unwrap(); + + assert_eq!(serde_json::to_value(control).unwrap(), input); + } + + #[rstest] + #[case::number_type(json!({"type": 42}))] + #[case::object_type(json!({"type": {"nested": true}}))] + #[case::boolean_type(json!({"type": true}))] + #[case::number_ttl(json!({"ttl": 300}))] + #[case::object_ttl(json!({"ttl": {"nested": true}}))] + #[case::boolean_ttl(json!({"ttl": false}))] + fn cache_control_rejects_non_string_type_and_ttl(#[case] input: Value) { + assert!(serde_json::from_value::(input).is_err()); + } + fn round_trip(value: &Value) -> Value { let parsed: T = serde_json::from_value(value.clone()).unwrap(); serde_json::to_value(parsed).unwrap() } + #[rstest] + fn malformed_cache_control_strings_are_rejected( + #[values("type", "ttl")] field: &str, + #[values(json!(42), json!(true), json!([]), json!({"seconds": 300}))] value: Value, + ) { + let request = json!({ + "model": "claude-test", "max_tokens": 16, + "messages": [{"role": "user", "content": [{ + "type": "text", "text": "prompt", "cache_control": {field: value} + }]}] + }); + + assert!(serde_json::from_value::(request).is_err()); + } + + #[cfg(feature = "schema")] + #[rstest] + #[case::cache_type(schemars::schema_for!(CacheControlType).to_value())] + #[case::ttl(schemars::schema_for!(CacheControlTtl).to_value())] + fn cache_control_enum_schema_accepts_unknown_strings(#[case] schema: Value) { + assert_eq!(schema.get("type"), Some(&json!("string"))); + assert!(schema.get("enum").is_none()); + } + #[rstest] #[case::text(json!({"type": "text", "text": "hi"}))] #[case::text_with_citations_and_cache_control(json!({ diff --git a/litellm-rust/crates/llms-types/src/formats/messages/tools.rs b/litellm-rust/crates/llms-types/src/formats/messages/tools.rs index 6f5acf9fb3c..eedbdc684ad 100644 --- a/litellm-rust/crates/llms-types/src/formats/messages/tools.rs +++ b/litellm-rust/crates/llms-types/src/formats/messages/tools.rs @@ -568,6 +568,7 @@ mod tests { use crate::formats::messages::{ AdvisorTool, AdvisorToolName, CacheControl, UrlSourceToolReference, }; + use crate::formats::messages::{CacheControlTtl, CacheControlType}; use serde_json::Map; use crate::formats::messages::AllowedCaller; @@ -819,7 +820,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 +1046,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 +1224,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(Some(cache_type)), + ttl: Some(Some(ttl)), ..CacheControl::default() } } diff --git a/litellm-rust/crates/llms-types/src/integrations/anthropic_cache_control_hook.rs b/litellm-rust/crates/llms-types/src/integrations/anthropic_cache_control_hook.rs new file mode 100644 index 00000000000..245c5013afa --- /dev/null +++ b/litellm-rust/crates/llms-types/src/integrations/anthropic_cache_control_hook.rs @@ -0,0 +1,220 @@ +use serde::{Deserialize, Deserializer, Serialize}; +use serde_json::{Map, Value}; + +use crate::{ + formats::messages::CacheControl, + message_role::MessageRole, + recognized::Recognized, + serde_compat::{PYTHON_DEFAULT_INT_MAX_STR_DIGITS, ParsedInteger, parse_python_integer}, +}; + +#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)] +#[serde(tag = "location", rename = "message")] +pub struct CacheControlMessageInjectionPoint { + #[serde( + default, + skip_serializing_if = "Option::is_none", + with = "serde_with::rust::double_option" + )] + control: Option>, + #[serde( + default, + skip_serializing_if = "Option::is_none", + with = "serde_with::rust::double_option" + )] + index: Option>, + #[serde( + default, + skip_serializing_if = "Option::is_none", + with = "serde_with::rust::double_option" + )] + role: Option>>, + #[serde(flatten)] + extra: Map, +} + +impl CacheControlMessageInjectionPoint { + pub fn new( + index: Option, + role: Option, + control: Option, + ) -> Self { + Self { + control: control.map(Some), + index: index.map(|index| Some(CacheControlIndex::Number(index.into()))), + role: role.map(|role| Some(Recognized::Known(role))), + extra: Map::new(), + } + } + + pub fn parsed_index(&self) -> ParsedInteger { + match self.index.as_ref().and_then(Option::as_ref) { + Some(CacheControlIndex::Number(number)) => number + .as_i64() + .map(ParsedInteger::Value) + .unwrap_or(ParsedInteger::OutOfRange), + Some(CacheControlIndex::String(index)) => { + parse_python_integer(index, PYTHON_DEFAULT_INT_MAX_STR_DIGITS) + } + Some(CacheControlIndex::Boolean(value)) => ParsedInteger::Value(i64::from(*value)), + Some(CacheControlIndex::Other(_)) | None => ParsedInteger::Invalid, + } + } + + pub fn role(&self) -> Option<&MessageRole> { + self.role + .as_ref() + .and_then(Option::as_ref) + .and_then(Recognized::known) + } + + pub fn control(&self) -> Option<&CacheControl> { + self.control.as_ref().and_then(Option::as_ref) + } +} + +#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)] +#[serde(untagged)] +enum CacheControlIndex { + Number(serde_json::Number), + String(String), + Boolean(bool), + Other(Value), +} + +#[derive(Clone, Debug, PartialEq, Deserialize, Serialize)] +#[serde(tag = "location", rename = "tool_config")] +pub struct CacheControlToolConfigInjectionPoint { + #[serde( + default, + skip_serializing_if = "Option::is_none", + with = "serde_with::rust::double_option" + )] + control: Option>, + #[serde(flatten)] + extra: Map, +} + +#[derive(Clone, Debug, PartialEq, Serialize)] +#[serde(untagged)] +pub enum CacheControlInjectionPoint { + Message(CacheControlMessageInjectionPoint), + ToolConfig(CacheControlToolConfigInjectionPoint), + Unknown(Value), +} + +impl<'de> Deserialize<'de> for CacheControlInjectionPoint { + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + let value = Value::deserialize(deserializer)?; + match value.get("location").and_then(Value::as_str) { + Some("message") => serde_json::from_value(value).map(Self::Message), + Some("tool_config") => serde_json::from_value(value).map(Self::ToolConfig), + _ => return Ok(Self::Unknown(value)), + } + .map_err(serde::de::Error::custom) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use serde_json::{Value, json}; + + use super::{CacheControlInjectionPoint, CacheControlMessageInjectionPoint}; + use crate::message_role::MessageRole; + use crate::serde_compat::ParsedInteger; + + #[rstest] + #[case::negative_index_string(json!({"location": "message", "index": "-1"}), ParsedInteger::Value(-1))] + #[case::padded_index_string(json!({"location": "message", "index": " \t0\n"}), ParsedInteger::Value(0))] + #[case::out_of_range_index(json!({"location": "message", "index": 99}), ParsedInteger::Value(99))] + #[case::oversized_string(json!({"location": "message", "index": "9223372036854775808", "role": "user"}), ParsedInteger::OutOfRange)] + #[case::oversized_number(json!({"location": "message", "index": u64::MAX, "role": "user"}), ParsedInteger::OutOfRange)] + #[case::underscores(json!({"location": "message", "index": "0_0"}), ParsedInteger::Value(0))] + #[case::unicode_digits(json!({"location": "message", "index": "٠"}), ParsedInteger::Value(0))] + #[case::signed_and_padded(json!({"location": "message", "index": " +0_1 "}), ParsedInteger::Value(1))] + #[case::negative_limit(json!({"location": "message", "index": "-9223372036854775808"}), ParsedInteger::Value(i64::MIN))] + #[case::negative_overflow(json!({"location": "message", "index": "-9223372036854775809", "role": "user"}), ParsedInteger::OutOfRange)] + #[case::boolean_index(json!({"location": "message", "index": true}), ParsedInteger::Value(1))] + #[case::mixed_digit_scripts(json!({"location": "message", "index": "٠_１"}), ParsedInteger::Value(1))] + #[case::mathematical_digit(json!({"location": "message", "index": "𝟙"}), ParsedInteger::Value(1))] + #[case::double_underscore(json!({"location": "message", "index": "0__1", "role": "user"}), ParsedInteger::Invalid)] + #[case::leading_underscore(json!({"location": "message", "index": "_1", "role": "user"}), ParsedInteger::Invalid)] + #[case::trailing_underscore(json!({"location": "message", "index": "1_", "role": "user"}), ParsedInteger::Invalid)] + #[case::sign_without_digits(json!({"location": "message", "index": "+", "role": "user"}), ParsedInteger::Invalid)] + #[case::embedded_sign(json!({"location": "message", "index": "1+0", "role": "user"}), ParsedInteger::Invalid)] + #[case::non_decimal_digit(json!({"location": "message", "index": "²", "role": "user"}), ParsedInteger::Invalid)] + #[case::role_only(json!({"location": "message", "role": "user"}), ParsedInteger::Invalid)] + #[case::invalid_index_falls_back_to_role( + json!({"location": "message", "index": "latest", "role": "assistant"}), + ParsedInteger::Invalid + )] + #[case::index_precedes_system_role( + json!({"location": "message", "index": 0, "role": "system"}), + ParsedInteger::Value(0) + )] + #[case::missing_target(json!({"location": "message"}), ParsedInteger::Invalid)] + fn message_index_deserialization(#[case] input: Value, #[case] expected: ParsedInteger) { + let point: CacheControlInjectionPoint = serde_json::from_value(input).unwrap(); + + assert!(matches!( + point, + CacheControlInjectionPoint::Message(message) if message.parsed_index() == expected + )); + } + + #[rstest] + #[case::missing_control(json!({"location": "message", "index": 0}))] + #[case::message_null_control(json!({"location": "message", "index": 0, "control": null}))] + #[case::tool_config_without_control(json!({"location": "tool_config"}))] + #[case::tool_config_null_control(json!({"location": "tool_config", "control": null}))] + fn known_points_preserve_optional_control(#[case] input: Value) { + let point: CacheControlInjectionPoint = serde_json::from_value(input.clone()).unwrap(); + + assert_eq!(serde_json::to_value(point).unwrap(), input); + } + + #[rstest] + #[case::message(json!({"location": "message", "index": 0, "control": "ephemeral"}))] + #[case::tool_config(json!({"location": "tool_config", "control": 5}))] + fn malformed_controls_are_rejected(#[case] input: Value) { + assert!(serde_json::from_value::(input).is_err()); + } + + #[rstest] + fn unknown_location_round_trips_unchanged() { + let input = json!({"location": "future", "index": 7, "custom": {"value": true}}); + let point: CacheControlInjectionPoint = serde_json::from_value(input.clone()).unwrap(); + + assert!(matches!(&point, CacheControlInjectionPoint::Unknown(_))); + assert_eq!(serde_json::to_value(point).unwrap(), input); + } + + #[rstest] + #[case::role_preserved_when_index_wins(json!({"location": "message", "index": 0, "role": "assistant"}))] + #[case::unparsed_index_preserved_when_role_wins(json!({"location": "message", "index": "latest", "role": "assistant"}))] + #[case::message(json!({"location": "message", "index": " 0 ", "role": "user", "control": null, "custom": "keep"}))] + #[case::tool_config(json!({"location": "tool_config", "control": null, "custom": "keep"}))] + #[case::null_control_fields(json!({"location": "message", "index": null, "role": null, "control": {"type": null, "ttl": null, "scope": null}}))] + #[case::unknown_target_types(json!({"location": "message", "index": [1], "role": {"custom": true}}))] + #[case::unicode_index(json!({"location": "message", "index": " +٠_１ ", "role": "assistant"}))] + fn known_points_preserve_wire_fields(#[case] input: Value) { + let point: CacheControlInjectionPoint = serde_json::from_value(input.clone()).unwrap(); + + assert_eq!(serde_json::to_value(point).unwrap(), input); + } + + #[rstest] + fn constructed_points_preserve_both_index_and_role() { + let point = + CacheControlMessageInjectionPoint::new(Some(0), Some(MessageRole::System), None); + + assert_eq!( + serde_json::to_value(point).unwrap(), + json!({"location": "message", "index": 0, "role": "system"}) + ); + } +} diff --git a/litellm-rust/crates/llms-types/src/integrations/mod.rs b/litellm-rust/crates/llms-types/src/integrations/mod.rs new file mode 100644 index 00000000000..0da3b631f2f --- /dev/null +++ b/litellm-rust/crates/llms-types/src/integrations/mod.rs @@ -0,0 +1,3 @@ +mod anthropic_cache_control_hook; + +pub use anthropic_cache_control_hook::*; diff --git a/litellm-rust/crates/llms-types/src/lib.rs b/litellm-rust/crates/llms-types/src/lib.rs index 67d933e39cc..975f7b4fc74 100644 --- a/litellm-rust/crates/llms-types/src/lib.rs +++ b/litellm-rust/crates/llms-types/src/lib.rs @@ -6,7 +6,9 @@ macro_rules_attribute::attribute_alias! { pub mod formats; pub mod headers; +pub mod integrations; pub mod json_schema; +pub mod message_role; pub mod providers; pub mod recognized; pub mod serde_compat; diff --git a/litellm-rust/crates/llms-types/src/message_role.rs b/litellm-rust/crates/llms-types/src/message_role.rs new file mode 100644 index 00000000000..cfc653cf966 --- /dev/null +++ b/litellm-rust/crates/llms-types/src/message_role.rs @@ -0,0 +1,74 @@ +#[derive( + Clone, + Debug, + PartialEq, + Eq, + strum::Display, + strum::EnumString, + strum::AsRefStr, + serde_with::DeserializeFromStr, + serde_with::SerializeDisplay, +)] +#[cfg_attr(feature = "schema", derive(schemars::JsonSchema))] +#[cfg_attr(feature = "schema", schemars(from = "String", into = "String"))] +pub enum MessageRole { + #[strum(serialize = "user")] + User, + #[strum(serialize = "assistant")] + Assistant, + #[strum(serialize = "system")] + System, + #[strum(default, transparent)] + Other(String), +} + +impl From for MessageRole { + fn from(value: String) -> Self { + value.parse().unwrap_or_else(|never| match never {}) + } +} + +impl From for String { + fn from(value: MessageRole) -> Self { + value.to_string() + } +} + +#[cfg(test)] +mod tests { + use super::MessageRole; + use rstest::rstest; + use serde_json::{Value, json}; + + #[rstest] + #[case::user("user", MessageRole::User)] + #[case::assistant("assistant", MessageRole::Assistant)] + #[case::system("system", MessageRole::System)] + #[case::unknown("future_role", MessageRole::Other("future_role".into()))] + #[case::case_sensitive("User", MessageRole::Other("User".into()))] + #[case::empty("", MessageRole::Other(String::new()))] + fn roles_preserve_wire_values(#[case] wire: &str, #[case] expected: MessageRole) { + let role: MessageRole = serde_json::from_value(json!(wire)).unwrap(); + + assert_eq!(role, expected); + assert_eq!(wire.parse::().unwrap(), role); + assert_eq!(role.as_ref(), wire); + assert_eq!(serde_json::to_value(role).unwrap(), json!(wire)); + } + + #[rstest] + fn malformed_roles_are_rejected( + #[values(json!(7), json!(true), json!([]), json!({}), Value::Null)] wire: Value, + ) { + assert!(serde_json::from_value::(wire).is_err()); + } + + #[cfg(feature = "schema")] + #[rstest] + fn role_schema_accepts_unknown_strings() { + let schema = schemars::schema_for!(MessageRole).to_value(); + + assert_eq!(schema.get("type"), Some(&json!("string"))); + assert!(schema.get("enum").is_none()); + } +} diff --git a/litellm-rust/crates/llms-types/src/providers/anthropic/constants.rs b/litellm-rust/crates/llms-types/src/providers/anthropic/constants.rs index 941d1e75461..291d0085f54 100644 --- a/litellm-rust/crates/llms-types/src/providers/anthropic/constants.rs +++ b/litellm-rust/crates/llms-types/src/providers/anthropic/constants.rs @@ -8,6 +8,7 @@ pub const BETA_HEADER: &str = "anthropic-beta"; pub const VERSION_HEADER: &str = "anthropic-version"; pub const DIRECT_BROWSER_ACCESS_HEADER: &str = "anthropic-dangerous-direct-browser-access"; pub const API_VERSION: &str = "2023-06-01"; +pub const BILLING_HEADER_PREFIX: &str = "x-anthropic-billing-header:"; pub const DEFAULT_HEADERS: &[(&str, &str)] = &[ (VERSION_HEADER, API_VERSION), diff --git a/litellm-rust/crates/llms-types/src/providers/minimax/messages.rs b/litellm-rust/crates/llms-types/src/providers/minimax/messages.rs index dbecb82b8f8..b9a89b9e092 100644 --- a/litellm-rust/crates/llms-types/src/providers/minimax/messages.rs +++ b/litellm-rust/crates/llms-types/src/providers/minimax/messages.rs @@ -64,6 +64,7 @@ mod tests { MinimaxMediaDetail, MinimaxMediaSource, MinimaxMessagesContentBlock, }; + use crate::formats::messages::CacheControlType; 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(Some(CacheControlType::Ephemeral)) ); } MinimaxMessagesContentBlock::MidConvSystem { text, extra } => { diff --git a/litellm-rust/crates/llms-types/src/serde_compat.rs b/litellm-rust/crates/llms-types/src/serde_compat.rs index d32ca80f862..a3c43ef12fb 100644 --- a/litellm-rust/crates/llms-types/src/serde_compat.rs +++ b/litellm-rust/crates/llms-types/src/serde_compat.rs @@ -3,10 +3,66 @@ use serde::{ de::{Error, Visitor}, }; use serde_with::DeserializeAs; +use unicode_properties::{GeneralCategory, UnicodeGeneralCategory}; pub struct LaxI64; pub struct FiniteF64; +pub const PYTHON_DEFAULT_INT_MAX_STR_DIGITS: usize = 4_300; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum ParsedInteger { + Value(i64), + OutOfRange, + Invalid, +} + +pub fn parse_python_integer(value: &str, max_digits: usize) -> ParsedInteger { + let Some(normalized) = value + .trim() + .chars() + .map(|character| match character { + '+' | '-' | '_' => Some(character), + _ => decimal_digit(character), + }) + .collect::>() + else { + return ParsedInteger::Invalid; + }; + let digits = normalized.strip_prefix(['+', '-']).unwrap_or(&normalized); + if !digits + .split('_') + .all(|part| !part.is_empty() && part.bytes().all(|byte| byte.is_ascii_digit())) + || (max_digits > 0 && digits.bytes().filter(u8::is_ascii_digit).count() > max_digits) + { + return ParsedInteger::Invalid; + } + normalized + .replace('_', "") + .parse::() + .map(ParsedInteger::Value) + .unwrap_or(ParsedInteger::OutOfRange) +} + +fn decimal_digit(character: char) -> Option { + if character.is_ascii_digit() { + return Some(character); + } + if character.general_category() != GeneralCategory::DecimalNumber { + return None; + } + // Unicode decimal digits form consecutive groups of ten + let preceding_digits = (0..u32::from(character)) + .rev() + .take_while(|code| { + char::from_u32(*code).is_some_and(|character| { + character.general_category() == GeneralCategory::DecimalNumber + }) + }) + .count(); + char::from_digit(u32::try_from(preceding_digits % 10).ok()?, 10) +} + impl<'de> DeserializeAs<'de, i64> for LaxI64 { fn deserialize_as>(deserializer: D) -> Result { deserializer.deserialize_any(Self) @@ -89,19 +145,13 @@ fn integer_string(value: &str) -> Option { } None => value, }; - if integer.starts_with('_') || integer.ends_with('_') || integer.contains("__") { + if !integer.is_ascii() || integer.trim() != integer { return None; } - let digits = integer.strip_prefix(['+', '-']).unwrap_or(integer); - if digits.is_empty() - || digits.starts_with('_') - || !digits - .bytes() - .all(|byte| byte.is_ascii_digit() || byte == b'_') - { - return None; + match parse_python_integer(integer, 0) { + ParsedInteger::Value(value) => Some(value), + ParsedInteger::OutOfRange | ParsedInteger::Invalid => None, } - integer.replace('_', "").parse().ok() } fn integral_float(value: f64) -> Option { @@ -114,7 +164,9 @@ fn integral_float(value: f64) -> Option { #[cfg(test)] mod tests { - use crate::serde_compat::{FiniteF64, LaxI64}; + use super::{ + FiniteF64, LaxI64, PYTHON_DEFAULT_INT_MAX_STR_DIGITS, ParsedInteger, parse_python_integer, + }; use rstest::rstest; use serde::{Deserialize, Serialize}; @@ -174,6 +226,8 @@ mod tests { #[case::exponent(json!("1e3"))] #[case::missing_fraction(json!("2."))] #[case::missing_integer(json!(".0"))] + #[case::space_before_fraction(json!("2 .0"))] + #[case::unicode_digit(json!("١"))] #[case::leading_separator(json!("_2"))] #[case::repeated_separator(json!("2__0"))] #[case::fraction(json!(2.5))] @@ -201,4 +255,41 @@ mod tests { let numbers: Numbers = serde_json::from_value(json!({"float": input})).unwrap(); assert_eq!(numbers.float, Some(expected)); } + + #[rstest] + #[case::padded(" \t+0_1\n", ParsedInteger::Value(1))] + #[case::negative("-1", ParsedInteger::Value(-1))] + #[case::mixed_scripts("٠_１", ParsedInteger::Value(1))] + #[case::mathematical_digit("𝟙", ParsedInteger::Value(1))] + #[case::minimum("-9223372036854775808", ParsedInteger::Value(i64::MIN))] + #[case::positive_overflow("9223372036854775808", ParsedInteger::OutOfRange)] + #[case::negative_overflow("-9223372036854775809", ParsedInteger::OutOfRange)] + #[case::empty("", ParsedInteger::Invalid)] + #[case::sign_only("+", ParsedInteger::Invalid)] + #[case::leading_separator("_1", ParsedInteger::Invalid)] + #[case::trailing_separator("1_", ParsedInteger::Invalid)] + #[case::repeated_separator("1__0", ParsedInteger::Invalid)] + #[case::embedded_sign("1+0", ParsedInteger::Invalid)] + #[case::superscript("²", ParsedInteger::Invalid)] + #[case::fraction("1.0", ParsedInteger::Invalid)] + fn python_integer_strings(#[case] input: &str, #[case] expected: ParsedInteger) { + assert_eq!( + parse_python_integer(input, PYTHON_DEFAULT_INT_MAX_STR_DIGITS), + expected + ); + } + + #[rstest] + #[case::at_limit("0".repeat(PYTHON_DEFAULT_INT_MAX_STR_DIGITS), PYTHON_DEFAULT_INT_MAX_STR_DIGITS, ParsedInteger::Value(0))] + #[case::over_limit("0".repeat(PYTHON_DEFAULT_INT_MAX_STR_DIGITS + 1), PYTHON_DEFAULT_INT_MAX_STR_DIGITS, ParsedInteger::Invalid)] + #[case::custom_limit("0".repeat(4), 3, ParsedInteger::Invalid)] + #[case::unlimited("0".repeat(PYTHON_DEFAULT_INT_MAX_STR_DIGITS + 1), 0, ParsedInteger::Value(0))] + #[case::sign_and_separators(" +٠_０ ".into(), 2, ParsedInteger::Value(0))] + fn python_integer_digit_limits( + #[case] input: String, + #[case] max_digits: usize, + #[case] expected: ParsedInteger, + ) { + assert_eq!(parse_python_integer(&input, max_digits), expected); + } } diff --git a/litellm-rust/crates/llms/AGENTS.md b/litellm-rust/crates/llms/AGENTS.md index 3cf916b1721..7dc3d2f6ff5 100644 --- a/litellm-rust/crates/llms/AGENTS.md +++ b/litellm-rust/crates/llms/AGENTS.md @@ -41,7 +41,7 @@ ## Python pairs -- Python paths identify counterparts but do not dictate Rust module names or class hierarchy. Preserve behavior and concepts, not structure +- Prefer mirroring Python's outer module and class structure, following the workspace's Python porting guideline. Preserve behavior while designing function internals freely - Keep operation and parameter names when responsibilities match. Rust types keep the Python name with Rust acronym casing (`BaseOCRConfig` -> `BaseOcrConfig`). Private Python helpers drop the leading underscore - Give Rust-only helpers distinct responsibility names instead of reusing trait method names - Use trait defaults for unchanged inherited behavior and explicit delegation for shared provider behavior. Config traits represent real provider contracts, so do not recreate inheritance with extra traits diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 7991a57353e..c90308a4d9b 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -6,11 +6,11 @@ use litellm_http::request::{ }; use litellm_llms_types::{ formats::messages::{ - ContentBlock, ContentBlockType, EffortLevel, Message, MessageContent, MessagesTool, - SystemPrompt, + ContentBlock, ContentBlockType, EffortLevel, Message, MessageContent, MessagesRequest, + MessagesTool, SystemPrompt, }, providers::anthropic::{ - API_BASE, API_KEY_HEADER, AnthropicBeta, BETA_HEADER, BetaSet, + API_BASE, API_KEY_HEADER, AnthropicBeta, BETA_HEADER, BILLING_HEADER_PREFIX, BetaSet, DIRECT_BROWSER_ACCESS_HEADER, MESSAGES_PATH, }, recognized::Recognized, @@ -18,6 +18,94 @@ use litellm_llms_types::{ use serde::Deserialize; use serde_json::Value; +pub fn supports_anthropic_cache_control( + model: &str, + provider: &str, + supports_prompt_caching: bool, +) -> bool { + supports_prompt_caching + && model.to_ascii_lowercase().contains("claude") + && matches!(provider, "anthropic" | "bedrock" | "vertex_ai" | "azure_ai") +} + +pub fn is_claude_code_one_shot_subagent_request( + request: &MessagesRequest, + user_agent: Option<&str>, +) -> bool { + let Some(user_agent) = user_agent else { + return false; + }; + if (!user_agent.starts_with("claude-cli/") && !user_agent.starts_with("claude-code/")) + || request.messages.len() != 1 + || request + .params + .tools + .as_ref() + .is_some_and(|tools| !tools.is_empty()) + { + return false; + } + request.messages[0].role == "user" + && request + .params + .system + .as_ref() + .is_some_and(is_claude_code_subagent_billing_system) +} + +fn is_claude_code_subagent_billing_system(system: &SystemPrompt) -> bool { + let texts: Option> = match system { + SystemPrompt::Text(text) => Some(vec![text.as_str()]), + SystemPrompt::Blocks(blocks) => blocks + .iter() + .map(|block| { + block + .is_type(ContentBlockType::Text) + .then_some(block.text.as_deref()) + .flatten() + }) + .collect(), + }; + let meaningful: Vec<&str> = texts + .into_iter() + .flatten() + .filter(|text| !text.trim().is_empty()) + .collect(); + if meaningful.is_empty() { + return false; + } + let fields: Option>> = meaningful + .into_iter() + .map(claude_code_billing_fields) + .collect(); + let subagent_values: Vec<&str> = fields + .into_iter() + .flatten() + .flatten() + .filter_map(|(key, value)| (key == "cc_is_subagent").then_some(value)) + .collect(); + subagent_values == ["true"] +} + +fn claude_code_billing_fields(text: &str) -> Option> { + let stripped = text.trim(); + if stripped.contains(['\n', '\r']) { + return None; + } + let fields = stripped.strip_prefix(BILLING_HEADER_PREFIX)?; + let parsed: Option> = fields + .split(';') + .filter(|field| !field.trim().is_empty()) + .map(|field| { + let (key, value) = field.split_once('=')?; + let key = key.trim(); + let value = value.trim(); + (!key.is_empty() && !value.is_empty()).then_some((key, value)) + }) + .collect(); + parsed.filter(|parsed| !parsed.is_empty()) +} + use crate::{ anthropic::ANTHROPIC_OAUTH_TOKEN_PREFIX, base_llm::auth::{AuthScheme, Headers}, @@ -384,8 +472,6 @@ pub fn is_encrypted_reasoning_block(block: &ContentBlock) -> bool { field.is_some_and(|value| value.starts_with(ENCRYPTED_REASONING_SIGNATURE_PREFIX)) } -const BILLING_HEADER_PREFIX: &str = "x-anthropic-billing-header:"; - fn is_billing_header_block(block: &ContentBlock) -> bool { block.block_type == Some(ContentBlockType::Text) && block @@ -1920,4 +2006,121 @@ mod tests { expected ); } + + #[rstest] + #[case::anthropic("claude-sonnet", "anthropic", true, true)] + #[case::bedrock("claude-sonnet", "bedrock", true, true)] + #[case::vertex("claude-sonnet", "vertex_ai", true, true)] + #[case::azure("claude-sonnet", "azure_ai", true, true)] + #[case::unsupported_provider("claude-sonnet", "deepseek", true, false)] + #[case::unsupported_model("gpt-test", "anthropic", true, false)] + #[case::unsupported_capability("claude-sonnet", "anthropic", false, false)] + fn cache_control_support_requires_model_provider_and_capability( + #[case] model: &str, + #[case] provider: &str, + #[case] capability: bool, + #[case] expected: bool, + ) { + assert_eq!( + supports_anthropic_cache_control(model, provider, capability), + expected + ); + } + + #[rstest] + #[case::claude_cli( + "claude-cli/1.0", + "x-anthropic-billing-header: cc_is_subagent=true", + None, + true + )] + #[case::claude_code( + "claude-code/1.0", + "x-anthropic-billing-header: cc_is_subagent=true", + None, + true + )] + #[case::user_agent_must_start_with_product( + "client claude-code/1.0", + "x-anthropic-billing-header: cc_is_subagent=true", + None, + false + )] + #[case::other_billing_field( + "claude-code/1.0", + "x-anthropic-billing-header: billing=true", + None, + false + )] + #[case::false_subagent_field( + "claude-code/1.0", + "x-anthropic-billing-header: cc_is_subagent=false", + None, + false + )] + #[case::billing_header_with_instructions( + "claude-code/1.0", + "x-anthropic-billing-header: cc_is_subagent=true\nYou are a helpful agent", + None, + false + )] + #[case::nonempty_tools( + "claude-code/1.0", + "x-anthropic-billing-header: cc_is_subagent=true", + Some(json!([{"name": "lookup", "input_schema": {"type": "object"}}])), + false + )] + fn one_shot_subagent_requires_claude_code_request_shape( + #[case] user_agent: &str, + #[case] system: &str, + #[case] tools: Option, + #[case] expected: bool, + ) { + let mut request_value = json!({ + "model": "claude-test", + "max_tokens": 16, + "system": system, + "messages": [{"role": "user", "content": "prompt"}] + }); + if let Some(tools) = tools { + request_value["tools"] = tools; + } + let request: MessagesRequest = serde_json::from_value(request_value).unwrap(); + + assert_eq!( + is_claude_code_one_shot_subagent_request(&request, Some(user_agent)), + expected + ); + } + + #[rstest] + #[case::empty_text_blocks_are_ignored(json!([ + {"type": "text", "text": ""}, + {"type": "text", "text": "x-anthropic-billing-header: cc_is_subagent=true"} + ]), true)] + #[case::billing_block_beside_instructions(json!([ + {"type": "text", "text": "x-anthropic-billing-header: cc_is_subagent=true"}, + {"type": "text", "text": "Answer the delegated question"} + ]), false)] + #[case::non_text_block(json!([ + {"type": "text", "text": "x-anthropic-billing-header: cc_is_subagent=true"}, + {"type": "image", "source": {"type": "base64", "media_type": "image/png", "data": "AA=="}} + ]), false)] + fn one_shot_subagent_requires_every_system_text_to_be_a_billing_header( + #[case] system: Value, + #[case] expected: bool, + ) { + let request: MessagesRequest = serde_json::from_value(json!({ + "model": "claude-test", + "max_tokens": 16, + "system": system, + "messages": [{"role": "user", "content": "prompt"}] + })) + .unwrap(); + + assert_eq!( + is_claude_code_one_shot_subagent_request(&request, Some("claude-cli/1.0")), + expected + ); + } }