scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

refactor(cache): separate response policy and host selection

Browse source
This commit is contained in:
Yujong Lee 2026-09-20 21:11:48 -07:00
parent 081c93908f
commit 0c3a0a2089
23 changed files with 426 additions and 407 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm-rust/Cargo.lock generated
View file

@ -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",

2
litellm-rust/crates/cache-memory/Cargo.toml
View file

@ -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

47
litellm-rust/crates/cache-memory/src/cache.rs
View file

@ -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<V> = Arc<dyn Fn(&V) -> Result<usize, Error> + Send + Sync>;
type ValueValidator<V> = Arc<dyn Fn(&V) -> Result<(), Error> + Send + Sync>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum CacheWrite {
@ -32,7 +29,6 @@ pub struct InMemoryCache<V: Clone> {
default_ttl: Duration,
max_entry_bytes: Option<usize>,
measure_value: Option<ValueMeasure<V>>,
validate_value: Option<ValueValidator<V>>,
now: Arc<dyn Fn() -> Duration + Send + Sync>,
}
@ -76,7 +72,6 @@ impl<V: Clone> InMemoryCache<V> {
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<V: Clone> InMemoryCache<V> {
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<V: Clone> InMemoryCache<V> {
}
}
impl InMemoryCache<CacheEntry> {
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<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
type Value = V;

80
litellm-rust/crates/cache-memory/tests/cache.rs
View file

@ -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<AtomicU64>
}
#[test]
fn disabled_size_limited_and_synchronized_response_writes_are_observable() {
let disabled = InMemoryCache::<CacheEntry>::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::<CacheEntry>::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::<CacheEntry>::default();
let cache = InMemoryCache::<String>::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");

34
litellm-rust/crates/cache-redis/src/cache.rs
View file

@ -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::<JsonCodec<CacheEntry>>::ttl_seconds(Duration::ZERO),
RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::ZERO),
1
);
assert_eq!(
RedisCache::<JsonCodec<CacheEntry>>::ttl_seconds(Duration::from_millis(1500)),
RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::from_millis(1500)),
2
);
assert_eq!(
RedisCache::<JsonCodec<CacheEntry>>::ttl_seconds(Duration::from_secs(15)),
RedisCache::<JsonCodec<serde_json::Value>>::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::<CacheEntry>::new().encode(&value).unwrap();
let payload = JsonCodec::<serde_json::Value>::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::<CacheEntry>::new())
.with_namespace(Some("litellm-cache".into()));
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::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::<CacheEntry>::new())
.with_namespace(Some("litellm-cache".into()));
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::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::<CacheEntry>::new())
.with_namespace(Some("litellm-cache".into()));
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
.with_namespace(Some("litellm-cache".into()));
assert_eq!(
cache.test_connection().await.unwrap().status,

7
litellm-rust/crates/cache-response/Cargo.toml
View file

@ -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

55
litellm-rust/crates/cache-response/README.md Normal file
View file

@ -0,0 +1,55 @@
# Response cache foundation
`ResponseCache<B>` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache<Value = CacheEntry>`
## 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<B>` 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

143
litellm-rust/crates/cache-response/src/caching.rs Normal file
View file

@ -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<String>,
pub api_parameter: bool,
pub internal_parameter: bool,
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
#[serde(default)]
pub struct CacheKeyInput {
pub fields: Vec<CacheKeyField>,
pub preset: Option<String>,
pub namespace: Option<String>,
pub include_provider_parameters: bool,
}
#[derive(Default)]
pub struct CacheKeyContext {
pub model_group: Option<String>,
pub caching_groups: Vec<(Vec<String>, String)>,
pub file_checksum: Option<String>,
pub file_object_name: Option<String>,
pub metadata_file_name: Option<String>,
pub parameters_file_name: Option<String>,
}
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<bool>,
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<Duration>) -> bool {
self.timestamp.is_finite()
&& max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64())
}
}

4
litellm-rust/crates/cache-response/src/codec.rs
View file

@ -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;

7
litellm-rust/crates/cache-response/src/lib.rs
View file

@ -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};

6
litellm-rust/crates/cache-response/src/response.rs
View file

@ -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)]

83
litellm-rust/crates/cache-response/tests/caching.rs Normal file
View file

@ -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()
);
}

40
litellm-rust/crates/cache-response/tests/response.rs
View file

@ -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<ResponseCache<InMemoryCache<CacheEntry>>> {
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);
}

1
litellm-rust/crates/cache/Cargo.toml vendored
View file

@ -8,7 +8,6 @@ repository.workspace = true
[dependencies]
serde.workspace = true
serde_json.workspace = true
sha2.workspace = true
thiserror.workspace = true
[dev-dependencies]

143
litellm-rust/crates/cache/src/caching.rs vendored
View file

@ -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<String>,
pub api_parameter: bool,
pub internal_parameter: bool,
}
#[derive(Clone, Debug, Default, Deserialize, Serialize)]
#[serde(default)]
pub struct CacheKeyInput {
pub fields: Vec<CacheKeyField>,
pub preset: Option<String>,
pub namespace: Option<String>,
pub include_provider_parameters: bool,
}
#[derive(Default)]
pub struct CacheKeyContext {
pub model_group: Option<String>,
pub caching_groups: Vec<(Vec<String>, String)>,
pub file_checksum: Option<String>,
pub file_object_name: Option<String>,
pub metadata_file_name: Option<String>,
pub parameters_file_name: Option<String>,
}
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<bool>,
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<Duration>) -> 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<B: BaseCache>(
cache: &B,
key: &str,

5
litellm-rust/crates/cache/src/lib.rs vendored
View file

@ -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;

94
litellm-rust/crates/cache/tests/caching.rs vendored
View file

@ -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<Vec<(String, CacheEntry, CacheKwargs)>>,
writes: Mutex<Vec<(String, String, CacheKwargs)>>,
}
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()
);
}

14
litellm-rust/crates/cache/tests/codec.rs vendored
View file

@ -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::<CacheEntry>::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::<RoutingState>::new();

2
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -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

3
litellm-rust/crates/python-bridge/src/cache/facade.rs vendored
View file

@ -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<PyType>,

6
litellm-rust/crates/python-bridge/src/cache/mod.rs vendored
View file

@ -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)]

33
litellm-rust/crates/cache-response/src/native.rs → litellm-rust/crates/python-bridge/src/cache/native.rs vendored
View file

@ -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<C = redis::Connection>
where
C: redis::ConnectionLike + Send + 'static,
{
#[derive(Clone)]
pub(super) enum NativeResponseCache {
Memory(Arc<ResponseCache<InMemoryCache<CacheEntry>>>),
Redis(Arc<ResponseCache<RedisCache<ResponseCacheCodec, C>>>),
}
impl<C: redis::ConnectionLike + Send + 'static> Clone for NativeResponseCache<C> {
fn clone(&self) -> Self {
match self {
Self::Memory(cache) => Self::Memory(Arc::clone(cache)),
Self::Redis(cache) => Self::Redis(Arc::clone(cache)),
}
}
Redis(Arc<ResponseCache<RedisCache<ResponseCacheCodec>>>),
}
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<C: redis::ConnectionLike + Send + 'static> NativeResponseCache<C> {
impl NativeResponseCache {
pub fn kind(&self) -> &'static str {
match self {
Self::Memory(_) => "memory",

20
tests/test_litellm_rust/test_cache.py
View file

@ -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