From 0c3a0a208948e97d0da05ec6c7e28672205c2f00 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 21:11:48 -0700 Subject: [PATCH] refactor(cache): separate response policy and host selection --- litellm-rust/Cargo.lock | 4 +- litellm-rust/crates/cache-memory/Cargo.toml | 2 +- litellm-rust/crates/cache-memory/src/cache.rs | 47 +----- .../crates/cache-memory/tests/cache.rs | 80 ++++------ litellm-rust/crates/cache-redis/src/cache.rs | 34 +++-- litellm-rust/crates/cache-response/Cargo.toml | 7 +- litellm-rust/crates/cache-response/README.md | 55 +++++++ .../crates/cache-response/src/caching.rs | 143 ++++++++++++++++++ .../crates/cache-response/src/codec.rs | 4 +- litellm-rust/crates/cache-response/src/lib.rs | 7 +- .../crates/cache-response/src/response.rs | 6 +- .../crates/cache-response/tests/caching.rs | 83 ++++++++++ .../crates/cache-response/tests/response.rs | 40 +++-- litellm-rust/crates/cache/Cargo.toml | 1 - litellm-rust/crates/cache/src/caching.rs | 143 ------------------ litellm-rust/crates/cache/src/lib.rs | 5 +- litellm-rust/crates/cache/tests/caching.rs | 94 +----------- litellm-rust/crates/cache/tests/codec.rs | 14 +- litellm-rust/crates/python-bridge/Cargo.toml | 2 + .../crates/python-bridge/src/cache/facade.rs | 3 +- .../crates/python-bridge/src/cache/mod.rs | 6 +- .../src => python-bridge/src/cache}/native.rs | 33 ++-- tests/test_litellm_rust/test_cache.py | 20 ++- 23 files changed, 426 insertions(+), 407 deletions(-) create mode 100644 litellm-rust/crates/cache-response/README.md create mode 100644 litellm-rust/crates/cache-response/src/caching.rs create mode 100644 litellm-rust/crates/cache-response/tests/caching.rs rename litellm-rust/crates/{cache-response/src => python-bridge/src/cache}/native.rs (75%) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index aa62e4f3770..8b299f5455b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2460,7 +2460,6 @@ dependencies = [ "rstest", "serde", "serde_json", - "sha2 0.10.9", "thiserror 2.0.19", "tokio", ] @@ -2498,6 +2497,7 @@ dependencies = [ "redis-test", "serde", "serde_json", + "sha2 0.10.9", "tokio", ] @@ -2665,6 +2665,8 @@ dependencies = [ "litellm-auth", "litellm-auth-gcp", "litellm-cache", + "litellm-cache-memory", + "litellm-cache-redis", "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core", diff --git a/litellm-rust/crates/cache-memory/Cargo.toml b/litellm-rust/crates/cache-memory/Cargo.toml index d4487573a9a..86ab01564c8 100644 --- a/litellm-rust/crates/cache-memory/Cargo.toml +++ b/litellm-rust/crates/cache-memory/Cargo.toml @@ -7,8 +7,8 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -serde_json.workspace = true [dev-dependencies] +serde_json.workspace = true rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 974cdbe9760..43186faf3f7 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -3,15 +3,12 @@ use std::collections::{BinaryHeap, HashMap}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, -}; +use litellm_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error}; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; const DEFAULT_TTL: Duration = Duration::from_secs(600); type ValueMeasure = Arc Result + Send + Sync>; -type ValueValidator = Arc Result<(), Error> + Send + Sync>; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum CacheWrite { @@ -32,7 +29,6 @@ pub struct InMemoryCache { default_ttl: Duration, max_entry_bytes: Option, measure_value: Option>, - validate_value: Option>, now: Arc Duration + Send + Sync>, } @@ -76,7 +72,6 @@ impl InMemoryCache { default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), max_entry_bytes, measure_value, - validate_value: None, now: Arc::new(now), } } @@ -90,9 +85,6 @@ impl InMemoryCache { if self.max_size_in_memory == 0 { return Ok(CacheWrite::Disabled); } - if let Some(validate) = &self.validate_value { - validate(&value)?; - } if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value) && measure(&value)? > limit { @@ -176,43 +168,6 @@ impl InMemoryCache { } } -impl InMemoryCache { - pub fn response_cache(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { - Self::response_cache_with_clock(capacity, ttl, max_entry_bytes, || { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - }) - } - - pub fn response_cache_with_clock( - capacity: usize, - ttl: Duration, - max_entry_bytes: usize, - now: impl Fn() -> Duration + Send + Sync + 'static, - ) -> Self { - let mut cache = Self::with_clock_and_size_measurement( - Some(capacity), - Some(ttl), - Some(max_entry_bytes), - Some(Arc::new(|entry: &CacheEntry| { - serde_json::to_vec(entry) - .map(|bytes| bytes.len()) - .map_err(|_| Error::InvalidEntry) - })), - now, - ); - cache.validate_value = Some(Arc::new(|entry: &CacheEntry| { - entry - .timestamp - .is_finite() - .then_some(()) - .ok_or(Error::InvalidEntry) - })); - cache - } -} - impl BaseCache for InMemoryCache { type Value = V; diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index ffac9d8ae64..370145bebef 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -3,8 +3,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use litellm_cache::{ - BaseCache, CacheBackend, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, get_cache, - set_cache, + BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, Error, get_cache, set_cache, }; use litellm_cache_memory::{CacheWrite, InMemoryCache}; use rstest::{fixture, rstest}; @@ -87,66 +86,49 @@ fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc } #[test] -fn disabled_size_limited_and_synchronized_response_writes_are_observable() { - let disabled = InMemoryCache::::response_cache(0, Duration::from_secs(60), 80); +fn disabled_size_limited_and_validated_writes_are_observable() { + let cache = |capacity| { + InMemoryCache::with_clock_and_size_measurement( + Some(capacity), + Some(Duration::from_secs(60)), + Some(4), + Some(Arc::new(|value: &String| { + if value.is_empty() { + return Err(Error::InvalidEntry); + } + Ok(value.len()) + })), + || Duration::from_secs(100), + ) + }; + let disabled = cache(0); assert_eq!( - disabled - .set_cache( - "a", - CacheEntry { - timestamp: 1.0, - response: serde_json::json!("x") - }, - None - ) - .unwrap(), + disabled.set_cache("a", "x".into(), None).unwrap(), CacheWrite::Disabled ); - let cache = InMemoryCache::::response_cache(2, Duration::from_secs(60), 80); + let cache = cache(2); assert_eq!( - cache - .set_cache( - "large", - CacheEntry { - timestamp: 1.0, - response: serde_json::json!("x".repeat(100)) - }, - None - ) - .unwrap(), + cache.set_cache("large", "oversized".into(), None).unwrap(), CacheWrite::TooLarge ); - cache - .set_cache( - "small", - CacheEntry { - timestamp: 1.0, - response: serde_json::json!("ok"), - }, - None, - ) - .unwrap(); - assert!(cache.get_cache("small").unwrap().is_some()); + assert_eq!(cache.get_cache("large").unwrap(), None); assert_eq!( - cache - .set_cache( - "invalid", - CacheEntry { - timestamp: f64::NAN, - response: serde_json::json!("bad"), - }, - None, - ) - .unwrap_err(), - Error::InvalidEntry + cache.set_cache("small", "ok".into(), None).unwrap(), + CacheWrite::Stored ); + assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into())); + assert_eq!( + cache.set_cache("invalid", String::new(), None), + Err(Error::InvalidEntry) + ); + assert_eq!(cache.get_cache("invalid").unwrap(), None); cache.delete_cache("small").unwrap(); - cache.flush_cache().unwrap(); + assert_eq!(cache.get_cache("small").unwrap(), None); } #[tokio::test] async fn connection_test_matches_python_result_contract() { - let cache = InMemoryCache::::default(); + let cache = InMemoryCache::::default(); let result = BaseCache::test_connection(&cache).await.unwrap(); assert_eq!(result.status, CacheConnectionStatus::Success); assert_eq!(result.message, "In-memory cache connection test successful"); diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 0faca6cdaaf..d4ca0cf0522 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -238,30 +238,27 @@ where #[cfg(test)] mod tests { use super::RedisCache; - use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKwargs, JsonCodec}; + use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, JsonCodec}; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; use std::time::Duration; - fn entry() -> CacheEntry { - CacheEntry { - timestamp: 123.0, - response: json!({"choices": [{"text": "cached"}]}), - } + fn entry() -> serde_json::Value { + json!({"deployment": "model-a", "cooldown_seconds": 30}) } #[test] fn ttl_seconds_rounds_up_and_keeps_expiration_positive() { assert_eq!( - RedisCache::>::ttl_seconds(Duration::ZERO), + RedisCache::>::ttl_seconds(Duration::ZERO), 1 ); assert_eq!( - RedisCache::>::ttl_seconds(Duration::from_millis(1500)), + RedisCache::>::ttl_seconds(Duration::from_millis(1500)), 2 ); assert_eq!( - RedisCache::>::ttl_seconds(Duration::from_secs(15)), + RedisCache::>::ttl_seconds(Duration::from_secs(15)), 15 ); } @@ -269,7 +266,9 @@ mod tests { #[test] fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() { let value = entry(); - let payload = JsonCodec::::new().encode(&value).unwrap(); + let payload = JsonCodec::::new() + .encode(&value) + .unwrap(); let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("SETEX") @@ -282,8 +281,9 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("litellm-cache".into())); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); cache .set_cache("key", value.clone(), CacheKwargs::default()) @@ -308,8 +308,9 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("litellm-cache".into())); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); cache.flush_cache().unwrap(); } @@ -318,8 +319,9 @@ mod tests { async fn test_connection_runs_ping_off_executor() { let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("litellm-cache".into())); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); assert_eq!( cache.test_connection().await.unwrap().status, diff --git a/litellm-rust/crates/cache-response/Cargo.toml b/litellm-rust/crates/cache-response/Cargo.toml index a0c4a1f74ef..04affb9872d 100644 --- a/litellm-rust/crates/cache-response/Cargo.toml +++ b/litellm-rust/crates/cache-response/Cargo.toml @@ -7,13 +7,14 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -litellm-cache-memory.workspace = true -litellm-cache-redis.workspace = true py_literal = "0.4.0" -redis = "1.7.0" serde.workspace = true serde_json.workspace = true +sha2.workspace = true [dev-dependencies] +litellm-cache-memory.workspace = true +litellm-cache-redis.workspace = true +redis = "1.7.0" redis-test = "1.0.4" tokio.workspace = true diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md new file mode 100644 index 00000000000..309296f7773 --- /dev/null +++ b/litellm-rust/crates/cache-response/README.md @@ -0,0 +1,55 @@ +# Response cache foundation + +`ResponseCache` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache` + +## Ownership + +`litellm-cache` defines typed storage and codec traits. Memory and Redis implement those traits without depending on response policy. Other consumers can store their own value types using the same backend implementations + +`litellm-cache-response` owns response keys, controls, entries, and the Python-compatible response codec. It has no runtime dependency on a specific cache backend or Python + +The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host + +## Native Rust use + +```rust +use std::{sync::Arc, time::Duration}; +use litellm_cache_memory::InMemoryCache; +use litellm_cache_response::{CacheKeyInput, ResponseCache, ResponseCacheRequest}; +use serde_json::json; + +let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); +let request = ResponseCacheRequest::new(CacheKeyInput { + preset: Some("example:key".into()), + ..Default::default() +}); +let now = Duration::from_secs(100); +cache.store(&request, json!({"answer": 7}), now)?; +assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7}))); +``` + +For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Each Redis constructor currently opens its own connection; shared connection pools remain follow-up work + +Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it + +## Python integration boundary + +The extension exposes `NativeCacheHandle`, `CacheResolver`, and captured `CacheBinding` objects for host integration. Memory and Redis handles support synchronous and asynchronous response lookup and storage + +The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution + +Explicit facade registration checks object identity, method overrides, and configuration changes before selecting native execution. Registration does not migrate entries or replace Python methods. Until activation configures one shared service, a registered facade and its native handle can hold separate data. Existing public cache constructors remain on Python + +Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy + +## Adding another backend + +Implement `BaseCache` for the backend with its associated value type, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache` then works without another response implementation. Add a concrete bridge enum variant and constructor only when exposing that backend to Python + +Verify typed values, TTL precedence, missing entries, serialization failures, namespaces, batch ordering, and sync/async behavior. Run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before enabling a public facade + +## Follow-up scope + +Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial batches, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths + +Redis cluster, disk, cloud stores, dual caching, and semantic caching remain follow-ups. Atomic counters, affinity claims, reservations, queues, and pubsub need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees diff --git a/litellm-rust/crates/cache-response/src/caching.rs b/litellm-rust/crates/cache-response/src/caching.rs new file mode 100644 index 00000000000..53b34025ccd --- /dev/null +++ b/litellm-rust/crates/cache-response/src/caching.rs @@ -0,0 +1,143 @@ +use std::time::Duration; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +pub enum CacheMode { + #[default] + #[serde(rename = "default_on")] + DefaultOn, + #[serde(rename = "default_off")] + DefaultOff, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct CacheKeyField { + pub name: String, + pub value: Option, + pub api_parameter: bool, + pub internal_parameter: bool, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +#[serde(default)] +pub struct CacheKeyInput { + pub fields: Vec, + pub preset: Option, + pub namespace: Option, + pub include_provider_parameters: bool, +} + +#[derive(Default)] +pub struct CacheKeyContext { + pub model_group: Option, + pub caching_groups: Vec<(Vec, String)>, + pub file_checksum: Option, + pub file_object_name: Option, + pub metadata_file_name: Option, + pub parameters_file_name: Option, +} + +impl CacheKeyContext { + pub fn apply(self, input: &mut CacheKeyInput) { + let group = self.model_group.as_ref().and_then(|model| { + self.caching_groups + .iter() + .find(|(models, _)| models.contains(model)) + }); + for field in &mut input.fields { + match field.name.as_str() { + "model" => { + field.value = group + .map(|(_, formatted)| formatted.clone()) + .or_else(|| self.model_group.clone()) + .or_else(|| field.value.take()) + } + "file" => { + field.value = self + .file_checksum + .clone() + .or_else(|| self.file_object_name.clone()) + .or_else(|| self.metadata_file_name.clone()) + .or_else(|| self.parameters_file_name.clone()) + } + _ => {} + } + } + } +} + +pub fn get_cache_key(input: &CacheKeyInput) -> String { + cache_key(input) +} + +pub fn cache_key(input: &CacheKeyInput) -> String { + if let Some(preset) = &input.preset { + return preset.clone(); + } + let mut digest = Sha256::new(); + for field in &input.fields { + if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter)) + && let Some(value) = &field.value + { + digest.update(field.name.as_bytes()); + digest.update(b": "); + digest.update(value.as_bytes()); + } + } + let hash = format!("{:x}", digest.finalize()); + input + .namespace + .as_deref() + .filter(|namespace| !namespace.is_empty()) + .map_or(hash.clone(), |namespace| format!("{namespace}:{hash}")) +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)] +pub struct CacheControls { + pub supported_call_type: bool, + pub configured: bool, + pub native_backend: bool, + pub default_on: bool, + pub caching: Option, + pub no_cache: bool, + pub no_store: bool, + #[serde(default)] + pub use_cache: bool, +} + +impl CacheControls { + pub fn reads(self) -> bool { + self.supported_call_type + && self.configured + && self.caching.unwrap_or(true) + && !self.no_cache + && (self.default_on || self.use_cache) + } + + pub fn writes(self) -> bool { + self.supported_call_type + && self.configured + && !self.no_store + && (self.default_on || self.use_cache) + } +} + +pub fn should_use_cache(controls: CacheControls) -> bool { + controls.reads() || controls.writes() +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct CacheEntry { + pub timestamp: f64, + pub response: Value, +} + +impl CacheEntry { + pub fn fresh(&self, now: Duration, max_age: Option) -> bool { + self.timestamp.is_finite() + && max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64()) + } +} diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs index 137cf61267a..f1d55ddefe8 100644 --- a/litellm-rust/crates/cache-response/src/codec.rs +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -1,4 +1,6 @@ -use litellm_cache::{CacheCodec, CacheEntry, Error}; +use litellm_cache::{CacheCodec, Error}; + +use crate::CacheEntry; use serde_json::Value; pub struct ResponseCacheCodec; diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index 454a82e76de..efa0b04b9f7 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -1,7 +1,10 @@ +mod caching; mod codec; -mod native; mod response; +pub use caching::{ + CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode, cache_key, + get_cache_key, should_use_cache, +}; pub use codec::ResponseCacheCodec; -pub use native::NativeResponseCache; pub use response::{ResponseCache, ResponseCacheRequest}; diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index 3e9807b6d1b..987a47f0554 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,8 +1,8 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{ - BaseCache, CacheControls, CacheEntry, CacheKeyInput, CacheKwargs, Error, cache_key, -}; +use litellm_cache::{BaseCache, CacheKwargs, Error}; + +use crate::{CacheControls, CacheEntry, CacheKeyInput, cache_key}; use serde_json::Value; #[derive(Clone)] diff --git a/litellm-rust/crates/cache-response/tests/caching.rs b/litellm-rust/crates/cache-response/tests/caching.rs new file mode 100644 index 00000000000..d403791e421 --- /dev/null +++ b/litellm-rust/crates/cache-response/tests/caching.rs @@ -0,0 +1,83 @@ +use litellm_cache_response::{ + CacheControls, CacheKeyContext, CacheKeyField, CacheKeyInput, cache_key, get_cache_key, +}; +use sha2::{Digest, Sha256}; + +#[test] +fn keys_match_python_order_groups_files_presets_and_namespaces() { + let mut input = CacheKeyInput { + fields: vec![ + CacheKeyField { + name: "model".into(), + value: Some("deployment".into()), + api_parameter: true, + internal_parameter: false, + }, + CacheKeyField { + name: "file".into(), + value: None, + api_parameter: true, + internal_parameter: false, + }, + ], + namespace: Some("team".into()), + ..Default::default() + }; + CacheKeyContext { + model_group: Some("group".into()), + caching_groups: vec![(vec!["group".into()], "['group']".into())], + file_checksum: Some("checksum".into()), + ..Default::default() + } + .apply(&mut input); + assert_eq!( + cache_key(&input), + format!( + "team:{:x}", + Sha256::digest(b"model: ['group']file: checksum") + ) + ); + input.preset = Some("preset".into()); + assert_eq!(get_cache_key(&input), "preset"); +} + +#[test] +fn cache_controls_honor_default_modes_and_directives() { + let enabled = CacheControls { + supported_call_type: true, + configured: true, + default_on: true, + ..Default::default() + }; + assert!(enabled.reads()); + assert!(enabled.writes()); + assert!( + !CacheControls { + default_on: false, + ..enabled + } + .reads() + ); + assert!( + CacheControls { + default_on: false, + use_cache: true, + ..enabled + } + .reads() + ); + assert!( + !CacheControls { + no_cache: true, + ..enabled + } + .reads() + ); + assert!( + !CacheControls { + no_store: true, + ..enabled + } + .writes() + ); +} diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index a4beda9fa4c..88b53fdfe86 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -6,15 +6,23 @@ use std::{ time::Duration, }; -use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKeyField, CacheKeyInput, Error}; +use litellm_cache::{BaseCache, CacheCodec, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ - NativeResponseCache, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, + CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec, + ResponseCacheRequest, }; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; +fn memory() -> Arc>> { + Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new( + Some(8), + Some(Duration::from_secs(600)), + )))) +} + fn request() -> ResponseCacheRequest { ResponseCacheRequest::new(CacheKeyInput { preset: Some("tenant:key".into()), @@ -87,7 +95,7 @@ async fn sync_and_async_consumers_share_keys_ttls_and_freshness() { #[tokio::test] async fn directives_skip_io_and_keep_reads_and_writes_independent() { - let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let cache = memory(); let mut request = request(); let now = Duration::from_secs(100); request.controls.no_store = true; @@ -112,7 +120,7 @@ async fn directives_skip_io_and_keep_reads_and_writes_independent() { } #[tokio::test] -async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_json() { +async fn redis_consumer_reads_python_sync_and_async_envelopes_and_writes_compatible_json() { let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("GET").arg("tenant:key"), @@ -133,7 +141,7 @@ async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_ .assert_all_commands_consumed(); let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec) .with_namespace(Some("tenant".into())); - let cache = NativeResponseCache::Redis(Arc::new(ResponseCache::new(Arc::new(backend)))); + let cache = ResponseCache::new(Arc::new(backend)); let request = request(); let expected = json!({"ok": true, "text": "cached"}); assert_eq!( @@ -154,10 +162,10 @@ async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_ } #[tokio::test] -async fn captured_enum_keeps_the_selected_backend_for_background_writes() { - let original = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); +async fn captured_service_keeps_the_selected_backend_for_background_writes() { + let original = memory(); let captured = original.clone(); - let replacement = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let replacement = memory(); let request = request(); let writer = tokio::spawn({ let request = request.clone(); @@ -186,7 +194,7 @@ async fn captured_enum_keeps_the_selected_backend_for_background_writes() { #[test] fn generated_keys_preserve_namespace_and_explicit_keys() { - let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let cache = memory(); let key = CacheKeyInput { fields: vec![CacheKeyField { name: "model".into(), @@ -199,7 +207,7 @@ fn generated_keys_preserve_namespace_and_explicit_keys() { }; let generated = ResponseCacheRequest::new(key.clone()); let explicit = ResponseCacheRequest::new(CacheKeyInput { - preset: Some(litellm_cache::cache_key(&key)), + preset: Some(litellm_cache_response::cache_key(&key)), ..Default::default() }); cache @@ -288,3 +296,15 @@ fn malformed_memory_entries_are_rejected_by_the_response_consumer() { Error::InvalidEntry ); } + +#[test] +fn response_entries_preserve_the_existing_json_representation() { + let codec = ResponseCacheCodec; + let entry = CacheEntry { + timestamp: 123.0, + response: json!({"choices": [{"text": "cached"}]}), + }; + let bytes = codec.encode(&entry).unwrap(); + assert_eq!(bytes, serde_json::to_vec(&entry).unwrap()); + assert_eq!(codec.decode(&bytes).unwrap(), entry); +} diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml index 350db4b1adb..0c504ab727a 100644 --- a/litellm-rust/crates/cache/Cargo.toml +++ b/litellm-rust/crates/cache/Cargo.toml @@ -8,7 +8,6 @@ repository.workspace = true [dependencies] serde.workspace = true serde_json.workspace = true -sha2.workspace = true thiserror.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs index 1d1df966ba2..39479694f3b 100644 --- a/litellm-rust/crates/cache/src/caching.rs +++ b/litellm-rust/crates/cache/src/caching.rs @@ -1,152 +1,9 @@ use std::sync::Arc; -use std::time::Duration; - -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use sha2::{Digest, Sha256}; use crate::{BaseCache, CacheKwargs, Error}; pub use crate::BaseCache as Cache; -#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] -pub enum CacheMode { - #[default] - #[serde(rename = "default_on")] - DefaultOn, - #[serde(rename = "default_off")] - DefaultOff, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct CacheKeyField { - pub name: String, - pub value: Option, - pub api_parameter: bool, - pub internal_parameter: bool, -} - -#[derive(Clone, Debug, Default, Deserialize, Serialize)] -#[serde(default)] -pub struct CacheKeyInput { - pub fields: Vec, - pub preset: Option, - pub namespace: Option, - pub include_provider_parameters: bool, -} - -#[derive(Default)] -pub struct CacheKeyContext { - pub model_group: Option, - pub caching_groups: Vec<(Vec, String)>, - pub file_checksum: Option, - pub file_object_name: Option, - pub metadata_file_name: Option, - pub parameters_file_name: Option, -} - -impl CacheKeyContext { - pub fn apply(self, input: &mut CacheKeyInput) { - let group = self.model_group.as_ref().and_then(|model| { - self.caching_groups - .iter() - .find(|(models, _)| models.contains(model)) - }); - for field in &mut input.fields { - match field.name.as_str() { - "model" => { - field.value = group - .map(|(_, formatted)| formatted.clone()) - .or_else(|| self.model_group.clone()) - .or_else(|| field.value.take()) - } - "file" => { - field.value = self - .file_checksum - .clone() - .or_else(|| self.file_object_name.clone()) - .or_else(|| self.metadata_file_name.clone()) - .or_else(|| self.parameters_file_name.clone()) - } - _ => {} - } - } - } -} - -pub fn get_cache_key(input: &CacheKeyInput) -> String { - cache_key(input) -} - -pub fn cache_key(input: &CacheKeyInput) -> String { - if let Some(preset) = &input.preset { - return preset.clone(); - } - let mut digest = Sha256::new(); - for field in &input.fields { - if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter)) - && let Some(value) = &field.value - { - digest.update(field.name.as_bytes()); - digest.update(b": "); - digest.update(value.as_bytes()); - } - } - let hash = format!("{:x}", digest.finalize()); - input - .namespace - .as_deref() - .filter(|namespace| !namespace.is_empty()) - .map_or(hash.clone(), |namespace| format!("{namespace}:{hash}")) -} - -#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)] -pub struct CacheControls { - pub supported_call_type: bool, - pub configured: bool, - pub native_backend: bool, - pub default_on: bool, - pub caching: Option, - pub no_cache: bool, - pub no_store: bool, - #[serde(default)] - pub use_cache: bool, -} - -impl CacheControls { - pub fn reads(self) -> bool { - self.supported_call_type - && self.configured - && self.caching.unwrap_or(true) - && !self.no_cache - && (self.default_on || self.use_cache) - } - - pub fn writes(self) -> bool { - self.supported_call_type - && self.configured - && !self.no_store - && (self.default_on || self.use_cache) - } -} - -pub fn should_use_cache(controls: CacheControls) -> bool { - controls.reads() || controls.writes() -} - -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct CacheEntry { - pub timestamp: f64, - pub response: Value, -} - -impl CacheEntry { - pub fn fresh(&self, now: Duration, max_age: Option) -> bool { - self.timestamp.is_finite() - && max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64()) - } -} - pub fn get_cache( cache: &B, key: &str, diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index a1d9d1402bb..4ff02319bdc 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -4,9 +4,6 @@ mod codec; mod error; pub use base_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs}; -pub use caching::{ - Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, - CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache, -}; +pub use caching::{Cache, CacheBackend, get_cache, set_cache}; pub use codec::{CacheCodec, JsonCodec}; pub use error::Error; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 5c250c6b3c9..824de00bdf4 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,17 +1,13 @@ -use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, - CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key, -}; -use sha2::{Digest, Sha256}; +use litellm_cache::{BaseCache, CacheConnectionResult, CacheKwargs, Error}; use std::{sync::Mutex, time::Duration}; struct TestCache { default_ttl: Duration, - writes: Mutex>, + writes: Mutex>, } impl BaseCache for TestCache { - type Value = CacheEntry; + type Value = String; fn default_ttl(&self) -> Duration { self.default_ttl @@ -83,10 +79,7 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() { default_ttl: Duration::from_secs(60), writes: Mutex::default(), }; - let entry = CacheEntry { - timestamp: 123.0, - response: serde_json::json!("cached"), - }; + let entry = String::from("cached"); let kwargs = CacheKwargs { ttl: Some(Duration::from_secs(5)), ..Default::default() @@ -116,82 +109,3 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() { ] ); } - -#[test] -fn keys_match_python_order_groups_files_presets_and_namespaces() { - let mut input = CacheKeyInput { - fields: vec![ - CacheKeyField { - name: "model".into(), - value: Some("deployment".into()), - api_parameter: true, - internal_parameter: false, - }, - CacheKeyField { - name: "file".into(), - value: None, - api_parameter: true, - internal_parameter: false, - }, - ], - namespace: Some("team".into()), - ..Default::default() - }; - CacheKeyContext { - model_group: Some("group".into()), - caching_groups: vec![(vec!["group".into()], "['group']".into())], - file_checksum: Some("checksum".into()), - ..Default::default() - } - .apply(&mut input); - assert_eq!( - cache_key(&input), - format!( - "team:{:x}", - Sha256::digest(b"model: ['group']file: checksum") - ) - ); - input.preset = Some("preset".into()); - assert_eq!(get_cache_key(&input), "preset"); -} - -#[test] -fn cache_controls_honor_default_modes_and_directives() { - let enabled = CacheControls { - supported_call_type: true, - configured: true, - default_on: true, - ..Default::default() - }; - assert!(enabled.reads()); - assert!(enabled.writes()); - assert!( - !CacheControls { - default_on: false, - ..enabled - } - .reads() - ); - assert!( - CacheControls { - default_on: false, - use_cache: true, - ..enabled - } - .reads() - ); - assert!( - !CacheControls { - no_cache: true, - ..enabled - } - .reads() - ); - assert!( - !CacheControls { - no_store: true, - ..enabled - } - .writes() - ); -} diff --git a/litellm-rust/crates/cache/tests/codec.rs b/litellm-rust/crates/cache/tests/codec.rs index dad5398a879..e24545caad6 100644 --- a/litellm-rust/crates/cache/tests/codec.rs +++ b/litellm-rust/crates/cache/tests/codec.rs @@ -1,6 +1,6 @@ use std::collections::BTreeMap; -use litellm_cache::{CacheCodec, CacheEntry, Error, JsonCodec}; +use litellm_cache::{CacheCodec, Error, JsonCodec}; use serde::{Deserialize, Serialize}; use serde_json::json; @@ -25,18 +25,6 @@ fn json_codec_round_trips_typed_domain_values() { ); } -#[test] -fn response_entries_preserve_the_existing_json_representation() { - let codec = JsonCodec::::new(); - let entry = CacheEntry { - timestamp: 123.0, - response: json!({"choices": [{"text": "cached"}]}), - }; - let bytes = codec.encode(&entry).unwrap(); - assert_eq!(bytes, serde_json::to_vec(&entry).unwrap()); - assert_eq!(codec.decode(&bytes).unwrap(), entry); -} - #[test] fn json_codec_rejects_malformed_and_wrongly_typed_entries() { let codec = JsonCodec::::new(); diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 308dfd2dd7a..1eb2ec28036 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -21,6 +21,8 @@ tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true litellm-cache.workspace = true +litellm-cache-memory.workspace = true +litellm-cache-redis.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 32360e72469..eb07118e964 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -1,4 +1,3 @@ -use litellm_cache_response::NativeResponseCache; use litellm_host_python::from_py; use pyo3::{ PyTraverseError, PyVisit, @@ -8,7 +7,7 @@ use pyo3::{ }; use serde_json::Value; -use super::NativeCacheHandle; +use super::{NativeCacheHandle, native::NativeResponseCache}; struct ClassGuard { class: Py, diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index f00cceeb86e..5918967009a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,9 +1,10 @@ mod facade; +mod native; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::{CacheControls, CacheKeyInput, Error}; -use litellm_cache_response::{NativeResponseCache, ResponseCacheRequest}; +use litellm_cache::Error; +use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; use pyo3::{ PyTraverseError, PyVisit, @@ -15,6 +16,7 @@ use serde::Deserialize; use serde_json::Value; use facade::FacadeGuard; +use native::NativeResponseCache; #[derive(Deserialize)] #[serde(deny_unknown_fields)] diff --git a/litellm-rust/crates/cache-response/src/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs similarity index 75% rename from litellm-rust/crates/cache-response/src/native.rs rename to litellm-rust/crates/python-bridge/src/cache/native.rs index c25c61cee82..6af04bfe2b2 100644 --- a/litellm-rust/crates/cache-response/src/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,33 +1,30 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheEntry, Error}; +use litellm_cache::{CacheCodec, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use serde_json::Value; -use crate::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; +use litellm_cache_response::{CacheEntry, ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; -pub enum NativeResponseCache -where - C: redis::ConnectionLike + Send + 'static, -{ +#[derive(Clone)] +pub(super) enum NativeResponseCache { Memory(Arc>>), - Redis(Arc>>), -} - -impl Clone for NativeResponseCache { - fn clone(&self) -> Self { - match self { - Self::Memory(cache) => Self::Memory(Arc::clone(cache)), - Self::Redis(cache) => Self::Redis(Arc::clone(cache)), - } - } + Redis(Arc>>), } impl NativeResponseCache { pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { Self::Memory(Arc::new(ResponseCache::new(Arc::new( - InMemoryCache::response_cache(capacity, ttl, max_entry_bytes), + InMemoryCache::with_clock_and_size_measurement( + Some(capacity), + Some(ttl), + Some(max_entry_bytes), + Some(Arc::new(|entry| { + ResponseCacheCodec.encode(entry).map(|bytes| bytes.len()) + })), + super::now, + ), )))) } @@ -41,7 +38,7 @@ impl NativeResponseCache { } } -impl NativeResponseCache { +impl NativeResponseCache { pub fn kind(&self) -> &'static str { match self { Self::Memory(_) => "memory", diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index d1cea860cb0..493baac228a 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -224,8 +224,24 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd def test_invalid_duration_and_request_shape_fail_before_storage() -> None: binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() for seconds in (-1.0, float("nan"), float("inf")): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) assert binding.lookup(request()) is None - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): _native.NativeCacheHandle.memory(ttl_seconds=-1) + + +async def test_memory_size_policy_is_applied_by_the_native_host() -> None: + handle: Final = _native.NativeCacheHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = _native.CacheResolver(SimpleNamespace(cache=handle)).resolve() + small: Final = {"answer": "ok"} + binding.store(request("small"), small) + assert await binding.async_lookup(request("small")) == small + await binding.async_store(request("large"), {"answer": "x" * 256}) + assert binding.lookup(request("large")) is None + assert binding.lookup(request("small")) == small + disabled: Final = _native.CacheResolver( + SimpleNamespace(cache=_native.NativeCacheHandle.memory(capacity=0)) + ).resolve() + await disabled.async_store(request(), small) + assert await disabled.async_lookup(request()) is None