mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-26 06:03:19 +00:00
feat(rust-messages): inject cache_control breakpoints natively
This commit is contained in:
parent
c7f19d0d22
commit
d96b6be6df
27 changed files with 437 additions and 9 deletions
1
litellm-rust/Cargo.lock
generated
1
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",
|
||||
|
|
|
|||
|
|
@ -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.clone(),
|
||||
},
|
||||
);
|
||||
let injected_count = injected.added;
|
||||
let trimmed =
|
||||
without_additional_drop_params(sanitized, &shaping.settings.additional_drop_params)?;
|
||||
without_additional_drop_params(injected.request, &shaping.settings.additional_drop_params)?;
|
||||
let transformed = config.transform_anthropic_messages_request(
|
||||
trimmed,
|
||||
&MessagesTransformContext::new(shaping.capabilities, shaping.settings.drop_params),
|
||||
|
|
@ -126,6 +153,7 @@ fn prepare_provider_request(
|
|||
environment,
|
||||
timeout: timeout.or_else(|| shaping.settings.request_timeout()),
|
||||
api_key: api_key.map(SecretValue::new),
|
||||
gateway_injected_cache_blocks: injected_count,
|
||||
})
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -1971,6 +1971,7 @@ mod tests {
|
|||
supports_sampling_params: true,
|
||||
supports_speed: false,
|
||||
supports_mid_conversation_system: false,
|
||||
supports_prompt_caching: false,
|
||||
effort_tiers: tiers(false, false, false, false, false, false),
|
||||
}
|
||||
);
|
||||
|
|
|
|||
|
|
@ -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