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

fix(cache): harden native foundation parity

Browse source
This commit is contained in:
Yujong Lee 2026-09-21 09:53:41 -07:00
parent da402b8aee
commit ef14fdaf33
23 changed files with 1347 additions and 318 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm-rust/Cargo.lock generated
View file

@ -3712,7 +3712,6 @@ dependencies = [
"itoa",
"num-bigint 0.5.1",
"percent-encoding",
"r2d2",
"ryu",
"sha1_smol",
"socket2 0.6.5",

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

@ -1,7 +1,9 @@
use std::cmp::Reverse;
use std::collections::{BinaryHeap, HashMap};
use std::sync::{Arc, Mutex};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use std::{
cmp::Reverse,
collections::{BinaryHeap, HashMap},
sync::{Arc, Mutex},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache,
@ -95,15 +97,13 @@ impl<V: Clone> InMemoryCache<V> {
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now);
let key = key.into();
state.values.insert(key.clone(), value);
Self::evict(&mut state, self.max_size_in_memory, now, &key);
let expiration = state.expirations.get(&key).copied();
if expiration.is_none_or(|expiration| expiration < now) {
let expiration = now + ttl.unwrap_or(self.default_ttl);
state.expirations.insert(key.clone(), expiration);
state.expiration_heap.push(Reverse((expiration, key)));
Self::set_expiration(&mut state, &key, now + ttl.unwrap_or(self.default_ttl));
}
state.values.insert(key, value);
Ok(CacheWrite::Stored)
}
@ -120,6 +120,10 @@ impl<V: Clone> InMemoryCache<V> {
Ok(state.values.get(key).cloned())
}
pub fn max_size_in_memory(&self) -> usize {
self.max_size_in_memory
}
pub fn expires_at(&self, key: &str) -> Result<Option<Duration>, Error> {
Ok(self
.state
@ -144,7 +148,7 @@ impl<V: Clone> InMemoryCache<V> {
Ok(())
}
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration) {
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration, key: &str) {
while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() {
if state.expirations.get(&key).copied() != Some(expiration) {
state.expiration_heap.pop();
@ -155,6 +159,9 @@ impl<V: Clone> InMemoryCache<V> {
break;
}
}
if state.values.contains_key(key) {
return;
}
while state.values.len() >= capacity {
let Some(Reverse((expiration, key))) = state.expiration_heap.pop() else {
break;
@ -165,6 +172,15 @@ impl<V: Clone> InMemoryCache<V> {
}
}
fn set_expiration(state: &mut CacheState<V>, key: &str, expiration: Duration) {
if state.expirations.get(key).copied() != Some(expiration) {
state.expirations.insert(key.into(), expiration);
state
.expiration_heap
.push(Reverse((expiration, key.into())));
}
}
fn remove(state: &mut CacheState<V>, key: &str) {
state.values.remove(key);
state.expirations.remove(key);
@ -182,40 +198,44 @@ where
eligible: &[V],
kwargs: CacheKwargs,
) -> Result<V, Error> {
if self.max_size_in_memory == 0 {
return Ok(candidate);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now);
let winner = match state.values.get(key) {
Some(existing) if eligible.is_empty() => existing.clone(),
Some(existing) if eligible.contains(existing) => existing.clone(),
_ => candidate,
};
let expiration = now + self.get_ttl(&kwargs);
Self::evict(&mut state, self.max_size_in_memory, now, key);
let existing = state
.values
.get(key)
.filter(|existing| eligible.is_empty() || eligible.contains(existing))
.cloned();
// Matches the Redis claim: an unconditional claim only extends its own winner.
if let Some(existing) = &existing
&& eligible.is_empty()
&& *existing != candidate
{
return Ok(existing.clone());
}
let winner = existing.unwrap_or(candidate);
Self::set_expiration(&mut state, key, now + self.get_ttl(&kwargs));
state.values.insert(key.into(), winner.clone());
state.expirations.insert(key.into(), expiration);
state
.expiration_heap
.push(Reverse((expiration, key.into())));
Ok(winner)
}
}
impl CounterCache for InMemoryCache<f64> {
fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result<f64, Error> {
if self.max_size_in_memory == 0 {
return Ok(amount);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now);
Self::evict(&mut state, self.max_size_in_memory, now, key);
let value = state.values.get(key).copied().unwrap_or_default() + amount;
let expiration = state
.expirations
.get(key)
.copied()
.unwrap_or_else(|| now + self.get_ttl(&kwargs));
if !state.expirations.contains_key(key) {
Self::set_expiration(&mut state, key, now + self.get_ttl(&kwargs));
}
state.values.insert(key.into(), value);
state.expirations.insert(key.into(), expiration);
state
.expiration_heap
.push(Reverse((expiration, key.into())));
Ok(value)
}
}
@ -256,3 +276,19 @@ impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn repeated_increments_keep_one_heap_entry_per_expiration() {
let cache = InMemoryCache::<f64>::new(Some(4), None);
for _ in 0..100 {
cache
.increment_cache("counter", 1.0, CacheKwargs::default())
.unwrap();
}
assert_eq!(cache.state.lock().unwrap().expiration_heap.len(), 1);
}
}

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

@ -1,6 +1,10 @@
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use std::{
sync::{
Arc,
atomic::{AtomicU64, Ordering},
},
time::Duration,
};
use litellm_cache::{
BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, Error,
@ -202,6 +206,17 @@ fn claims_are_atomic_and_refresh_eligible_winners() {
.unwrap(),
"first"
);
clock.store(103, Ordering::SeqCst);
assert_eq!(
cache
.claim_cache("affinity", "second".to_string(), &[], kwargs.clone())
.unwrap(),
"first"
);
assert_eq!(
cache.expires_at("affinity").unwrap(),
Some(Duration::from_secs(110))
);
clock.store(105, Ordering::SeqCst);
assert_eq!(
cache
@ -232,3 +247,61 @@ fn counters_increment_under_one_lock() {
3.5
);
}
#[rstest]
fn rewriting_an_existing_key_at_capacity_keeps_other_entries(clock: Arc<AtomicU64>) {
let cache = cache(clock, 2);
cache
.set_cache("hot", "1".into(), Some(Duration::from_secs(10)))
.unwrap();
cache
.set_cache("cold", "2".into(), Some(Duration::from_secs(20)))
.unwrap();
cache.set_cache("cold", "3".into(), None).unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
cache
.claim_cache("cold", "4".into(), &[], CacheKwargs::default())
.unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
cache.set_cache("new", "5".into(), None).unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), None);
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
assert_eq!(cache.get_cache("new").unwrap(), Some("5".into()));
}
#[test]
fn incrementing_an_existing_counter_at_capacity_keeps_every_counter() {
let cache = InMemoryCache::<f64>::new(Some(2), None);
for key in ["a", "b", "a", "b"] {
cache
.increment_cache(key, 1.0, CacheKwargs::default())
.unwrap();
}
assert_eq!(cache.get_cache("a").unwrap(), Some(2.0));
assert_eq!(cache.get_cache("b").unwrap(), Some(2.0));
}
#[test]
fn disabled_cache_does_not_retain_claims_or_counters() {
let claims = InMemoryCache::<String>::new(Some(0), None);
assert_eq!(
claims
.claim_cache("key", "first".into(), &[], CacheKwargs::default())
.unwrap(),
"first"
);
assert_eq!(claims.get_cache("key").unwrap(), None);
let counters = InMemoryCache::<f64>::new(Some(0), None);
assert_eq!(
counters
.increment_cache("key", 2.0, CacheKwargs::default())
.unwrap(),
2.0
);
assert_eq!(counters.get_cache("key").unwrap(), None);
}

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

@ -7,7 +7,7 @@ repository.workspace = true
[dependencies]
litellm-cache.workspace = true
redis = { version = "1.7.0", features = ["r2d2"] }
redis = "1.7.0"
r2d2 = "0.8.10"
tokio.workspace = true

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

@ -1,5 +1,7 @@
use std::sync::{Arc, Mutex};
use std::time::Duration;
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::{
BaseCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs,
@ -11,8 +13,59 @@ const DEFAULT_TTL: Duration = Duration::from_secs(600);
const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
const REDIS_POOL_SIZE: u32 = 16;
struct PooledConnection {
connection: redis::Connection,
failed: bool,
}
/// Pools connections without a checkout PING, which would double every operation's round trips.
/// A timed-out command leaves its reply on the socket while redis still reports the connection
/// open, so any connection whose operation failed is discarded instead of being reused.
struct ConnectionManager(redis::Client);
impl r2d2::ManageConnection for ConnectionManager {
type Connection = PooledConnection;
type Error = redis::RedisError;
fn connect(&self) -> Result<PooledConnection, redis::RedisError> {
let connection = self.0.get_connection()?;
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
Ok(PooledConnection {
connection,
failed: false,
})
}
fn is_valid(&self, connection: &mut PooledConnection) -> Result<(), redis::RedisError> {
redis::cmd("PING").query::<String>(&mut connection.connection)?;
Ok(())
}
fn has_broken(&self, connection: &mut PooledConnection) -> bool {
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
}
}
const INCREMENT_SCRIPT: &str = concat!(
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
"if redis.call('TTL', KEYS[1]) == -1 then ",
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
);
// Compare-and-set against the exact bytes the claim decision was made on.
// ARGV: [1] expected payload or "" when absent, [2] ttl, [3] new payload, [4] refresh ttl.
const CLAIM_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
"elseif current ~= ARGV[1] then return 0; end; ",
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
);
const CLAIM_ATTEMPTS: usize = 8;
enum Connections<C> {
Pool(r2d2::Pool<redis::Client>),
Pool(r2d2::Pool<ConnectionManager>),
Fixed(Mutex<C>),
}
@ -59,14 +112,10 @@ where
) -> Result<T, Error> {
match self {
Self::Pool(pool) => {
let mut connection = pool.get().map_err(|_| Error::Unavailable)?;
connection
.set_read_timeout(Some(REDIS_TIMEOUT))
.map_err(|_| Error::Unavailable)?;
connection
.set_write_timeout(Some(REDIS_TIMEOUT))
.map_err(|_| Error::Unavailable)?;
operation(&mut ConnectionRef(&mut *connection))
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let result = operation(&mut ConnectionRef(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
result
}
Self::Fixed(connection) => {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
@ -90,7 +139,8 @@ impl<S: CacheCodec> RedisCache<S> {
.max_size(REDIS_POOL_SIZE)
.min_idle(Some(0))
.connection_timeout(REDIS_TIMEOUT)
.build(client)
.test_on_check_out(false)
.build(ConnectionManager(client))
.map_err(|_| Error::Unavailable)?;
Ok(Self {
connections: Arc::new(Connections::Pool(pool)),
@ -122,6 +172,10 @@ where
}
}
pub fn namespace(&self) -> Option<&str> {
self.namespace.as_deref()
}
fn namespaced_key(&self, key: &str) -> String {
match &self.namespace {
Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
@ -407,29 +461,105 @@ where
C: redis::ConnectionLike + Send + 'static,
{
fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result<f64, Error> {
const SCRIPT: &str = concat!(
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
"if redis.call('TTL', KEYS[1]) == -1 then ",
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
);
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
self.connections.execute(|connection| {
redis::cmd("EVAL")
.arg(SCRIPT)
.arg(1)
.arg(key)
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
})
self.connections
.execute(|connection| increment(connection, key, amount, ttl))
}
async fn async_increment_cache(
&self,
key: &str,
amount: f64,
kwargs: CacheKwargs,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
increment(connection, key, amount, ttl)
})
.await
}
}
fn increment(
connection: &mut ConnectionRef<'_>,
key: String,
amount: f64,
ttl: u64,
) -> Result<f64, Error> {
redis::cmd("EVAL")
.arg(INCREMENT_SCRIPT)
.arg(1)
.arg(key)
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
}
fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
match value {
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => Ok(Some(bytes)),
redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
_ => Err(Error::InvalidEntry),
}
}
/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's
/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the
/// bytes that decision was made on, retried when another claimant wins the race.
fn claim<S: CacheCodec>(
connection: &mut ConnectionRef<'_>,
codec: &S,
key: &str,
candidate: S::Value,
eligible: &[S::Value],
ttl: u64,
) -> Result<S::Value, Error>
where
S::Value: PartialEq,
{
let payload = codec.encode(&candidate)?;
if payload.is_empty() {
return Err(Error::InvalidEntry);
}
for _ in 0..CLAIM_ATTEMPTS {
let current = stored_bytes(
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)?,
)?
.filter(|bytes| !bytes.is_empty());
let existing = current
.as_deref()
.and_then(|bytes| codec.decode(bytes).ok())
.filter(|existing| eligible.is_empty() || eligible.contains(existing));
let refresh = existing
.as_ref()
.is_some_and(|existing| !eligible.is_empty() || *existing == candidate);
let write: &[u8] = if existing.is_some() { b"" } else { &payload };
let applied = redis::cmd("EVAL")
.arg(CLAIM_SCRIPT)
.arg(1)
.arg(key)
.arg(current.as_deref().unwrap_or_default())
.arg(ttl)
.arg(write)
.arg(u8::from(refresh))
.query::<bool>(connection)
.map_err(|_| Error::Unavailable)?;
if applied {
return Ok(existing.unwrap_or(candidate));
}
}
Err(Error::Unavailable)
}
impl<S, C> ClaimCache for RedisCache<S, C>
where
S: CacheCodec,
S: CacheCodec + Clone + 'static,
S::Value: PartialEq,
C: redis::ConnectionLike + Send + 'static,
{
@ -440,44 +570,38 @@ where
eligible: &[S::Value],
kwargs: CacheKwargs,
) -> Result<S::Value, Error> {
const SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if current == false then redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ",
"return ARGV[1]; end; if #ARGV > 2 then for index = 3, #ARGV do ",
"if current == ARGV[index] then redis.call('EXPIRE', KEYS[1], ARGV[2]); ",
"return current; end; end; redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ",
"return ARGV[1]; end; if current == ARGV[1] then ",
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return current"
);
let key = self.namespaced_key(key);
let candidate = self.codec.encode(&candidate)?;
let eligible = eligible
.iter()
.map(|value| self.codec.encode(value))
.collect::<Result<Vec<_>, _>>()?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
let value = self.connections.execute(|connection| {
redis::cmd("EVAL")
.arg(SCRIPT)
.arg(1)
.arg(key)
.arg(candidate)
.arg(ttl)
.arg(eligible)
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)
})?;
self.decode_response(value)?.ok_or(Error::Unavailable)
self.connections
.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl))
}
async fn async_claim_cache(
&self,
key: &str,
candidate: S::Value,
eligible: Vec<S::Value>,
kwargs: CacheKwargs,
) -> Result<S::Value, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
let codec = self.codec.clone();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
claim(connection, &codec, &key, candidate, &eligible, ttl)
})
.await
}
}
#[cfg(test)]
mod tests {
use super::RedisCache;
use std::time::Duration;
use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, JsonCodec};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
use std::time::Duration;
use super::RedisCache;
fn entry() -> serde_json::Value {
json!({"deployment": "model-a", "cooldown_seconds": 30})

127
litellm-rust/crates/cache-redis/tests/cache.rs
View file

@ -1,8 +1,8 @@
use std::time::Duration;
use litellm_cache::{
BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, Error, JsonCodec,
get_cache, set_cache,
BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, ClaimCache,
CounterCache, Error, JsonCodec, get_cache, set_cache,
};
use litellm_cache_redis::RedisCache;
use redis_test::{MockCmd, MockRedisConnection};
@ -259,3 +259,126 @@ async fn async_flush_deletes_each_scan_page_separately() {
cache.async_flush_cache().await.unwrap();
}
const CLAIM_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
"elseif current ~= ARGV[1] then return 0; end; ",
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
);
fn claim_eval(expected: &str, write: &str, refresh: bool) -> redis::Cmd {
let mut cmd = redis::cmd("EVAL");
cmd.arg(CLAIM_SCRIPT)
.arg(1)
.arg("pin")
.arg(expected)
.arg(600)
.arg(write)
.arg(u8::from(refresh));
cmd
}
#[tokio::test]
async fn claims_match_eligible_values_written_by_another_encoder() {
let python_payload = r#"{"model_id": "a", "deployment": "east"}"#;
let stored = serde_json::json!({"deployment": "east", "model_id": "a"});
let candidate = serde_json::json!({"model_id": "b"});
let connection = MockRedisConnection::new([
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(python_payload)),
MockCmd::new(claim_eval(python_payload, "", true), Ok(1)),
])
.assert_all_commands_consumed();
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
assert_eq!(
cache
.async_claim_cache(
"pin",
candidate,
vec![stored.clone()],
CacheKwargs::default()
)
.await
.unwrap(),
stored
);
}
#[test]
fn claims_retry_when_the_key_changes_and_replace_ineligible_winners() {
let candidate = serde_json::json!({"model_id": "b"});
let payload = r#"{"model_id":"b"}"#;
let connection = MockRedisConnection::new([
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(redis::Value::Nil)),
MockCmd::new(claim_eval("", payload, false), Ok(0)),
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(r#"{"model_id":"gone"}"#)),
MockCmd::new(claim_eval(r#"{"model_id":"gone"}"#, payload, false), Ok(1)),
])
.assert_all_commands_consumed();
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
assert_eq!(
cache
.claim_cache(
"pin",
candidate.clone(),
&[serde_json::json!({"model_id": "a"})],
CacheKwargs::default()
)
.unwrap(),
candidate
);
}
#[test]
fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl() {
let stored = r#"{"model_id": "a"}"#;
let connection = MockRedisConnection::new([
MockCmd::new(redis::cmd("GET").arg("pin"), Ok(stored)),
MockCmd::new(claim_eval(stored, "", false), Ok(1)),
])
.assert_all_commands_consumed();
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new());
assert_eq!(
cache
.claim_cache(
"pin",
serde_json::json!({"model_id": "b"}),
&[],
CacheKwargs::default()
)
.unwrap(),
serde_json::json!({"model_id": "a"})
);
}
#[tokio::test]
async fn async_increment_runs_the_atomic_script() {
let mut eval = redis::cmd("EVAL");
eval.arg(concat!(
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
"if redis.call('TTL', KEYS[1]) == -1 then ",
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
))
.arg(1)
.arg("counter")
.arg(2.5f64)
.arg(600);
let connection =
MockRedisConnection::new([MockCmd::new(eval, Ok("4.5"))]).assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, JsonCodec::<f64>::new());
assert_eq!(
cache
.async_increment_cache("counter", 2.5, CacheKwargs::default())
.await
.unwrap(),
4.5
);
}

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

@ -28,17 +28,21 @@ 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. Sync operations check out independent connections from a bounded pool, while async callers move that blocking work off the executor
For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers, including counters and claims, move that blocking work off the executor. The pool skips the checkout PING and instead discards any connection whose command failed
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 single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring
The extension keeps a private test harness for memory and Redis single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring. No bridge-only cache type is part of the public API
Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec
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, effective TTL, and configuration changes before selecting native execution. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. 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
Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend
The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and configuration changes before selecting native execution. It does not compare Redis connection settings. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service 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
@ -52,4 +56,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na
Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths
Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache now provides L2-first counters and atomic affinity claims with local fallback, but public Router integration remains follow-up work. Reservations, queues, and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees
Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations, queues, and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees

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

@ -1,8 +1,9 @@
use litellm_cache::{CacheCodec, Error};
use crate::CacheEntry;
use serde_json::Value;
use crate::CacheEntry;
#[derive(Clone, Copy, Debug, Default)]
pub struct ResponseCacheCodec;
impl CacheCodec for ResponseCacheCodec {
@ -15,7 +16,18 @@ impl CacheCodec for ResponseCacheCodec {
{
return Err(Error::InvalidEntry);
}
serde_json::to_vec(value).map_err(|_| Error::InvalidEntry)
// Python reads a `response` that is either a dict or a serialized string, so every
// other shape is written serialized. A string on the wire is therefore always a
// serialized response, which keeps string-valued responses unambiguous.
if value.timestamp.is_none() || value.response.is_object() {
return serde_json::to_vec(value).map_err(|_| Error::InvalidEntry);
}
let response = serde_json::to_string(&value.response).map_err(|_| Error::InvalidEntry)?;
serde_json::to_vec(&CacheEntry {
timestamp: value.timestamp,
response: Value::String(response),
})
.map_err(|_| Error::InvalidEntry)
}
fn decode(&self, bytes: &[u8]) -> Result<CacheEntry, Error> {
@ -30,7 +42,10 @@ impl CacheCodec for ResponseCacheCodec {
let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else {
return Err(Error::InvalidEntry);
};
let response = value.get("response").cloned().ok_or(Error::InvalidEntry)?;
let response = match value.get("response").ok_or(Error::InvalidEntry)? {
Value::String(text) => decode_value(text)?,
response => response.clone(),
};
Ok(CacheEntry {
timestamp: Some(timestamp),
response,
@ -38,7 +53,7 @@ impl CacheCodec for ResponseCacheCodec {
}
}
pub(crate) fn decode_value(text: &str) -> Result<Value, Error> {
fn decode_value(text: &str) -> Result<Value, Error> {
if let Ok(value) = serde_json::from_str(text) {
return Ok(value);
}

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

@ -1,9 +1,9 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, Error};
use serde_json::Value;
use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
use serde_json::Value;
#[derive(Clone)]
pub struct ResponseCacheRequest {
@ -39,6 +39,10 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
Self { backend }
}
pub fn backend(&self) -> &B {
&self.backend
}
pub fn default_ttl(&self) -> Duration {
self.backend.default_ttl()
}
@ -67,7 +71,7 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
Err(Error::InvalidEntry) => None,
Err(error) => return Err(error),
};
Self::fresh_or_miss(entry, now, request.max_age)
Ok(Self::fresh_or_miss(entry, now, request.max_age))
}
pub async fn async_lookup(
@ -87,7 +91,7 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
Err(Error::InvalidEntry) => None,
Err(error) => return Err(error),
};
Self::fresh_or_miss(entry, now, request.max_age)
Ok(Self::fresh_or_miss(entry, now, request.max_age))
}
pub fn lookup_batch(
@ -180,11 +184,26 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
&self,
entries: Vec<(ResponseCacheRequest, Value)>,
now: Duration,
) -> Result<(), Error> {
self.async_store_entries(
entries
.into_iter()
.map(|(request, response)| (request, response, now))
.collect(),
)
.await
}
/// Stores entries that each carry the time they were produced, so a deferred write keeps
/// the freshness of its original response.
pub async fn async_store_entries(
&self,
entries: Vec<(ResponseCacheRequest, Value, Duration)>,
) -> Result<(), Error> {
let writable = entries
.into_iter()
.filter(|(request, _)| request.controls.writes())
.map(|(request, response)| {
.filter(|(request, _, _)| request.controls.writes())
.map(|(request, response, now)| {
(
cache_key(&request.key),
CacheEntry {
@ -227,7 +246,7 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
let mut values = vec![None; requests.len()];
for ((index, request), entry) in readable.into_iter().zip(entries) {
let response = match entry {
BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age)?,
BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age),
BatchEntry::Miss | BatchEntry::Invalid => None,
};
values[index] = response;
@ -239,24 +258,9 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
entry: Option<CacheEntry>,
now: Duration,
max_age: Option<Duration>,
) -> Result<Option<Value>, Error> {
match Self::fresh_response(entry, now, max_age) {
Err(Error::InvalidEntry) => Ok(None),
result => result,
}
}
fn fresh_response(
entry: Option<CacheEntry>,
now: Duration,
max_age: Option<Duration>,
) -> Result<Option<Value>, Error> {
) -> Option<Value> {
entry
.filter(|entry| entry.fresh(now, max_age))
.map(|entry| match (entry.timestamp, entry.response) {
(Some(_), Value::String(text)) => crate::codec::decode_value(&text),
(_, value) => Ok(value),
})
.transpose()
.map(|entry| entry.response)
}
}

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

@ -273,22 +273,45 @@ async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis() {
}
#[test]
fn malformed_memory_entries_are_treated_as_misses() {
let backend = Arc::new(InMemoryCache::default());
BaseCache::set_cache(
backend.as_ref(),
"tenant:key",
CacheEntry {
fn string_responses_round_trip_through_typed_and_wire_backends() {
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
let now = Duration::from_secs(100);
for response in [json!("hello world"), json!("123"), json!("null")] {
cache.store(&request(), response.clone(), now).unwrap();
assert_eq!(
cache.lookup(&request(), now).unwrap(),
Some(response.clone())
);
let wire = ResponseCacheCodec
.encode(&CacheEntry {
timestamp: Some(100.0),
response: response.clone(),
})
.unwrap();
assert_eq!(ResponseCacheCodec.decode(&wire).unwrap().response, response);
}
}
#[test]
fn non_object_responses_are_written_as_python_readable_serialized_strings() {
let wire = ResponseCacheCodec
.encode(&CacheEntry {
timestamp: Some(100.0),
response: json!("not a serialized response"),
},
Default::default(),
)
.unwrap();
let cache = ResponseCache::new(backend);
response: json!([1, 2]),
})
.unwrap();
assert_eq!(
cache.lookup(&request(), Duration::from_secs(100)).unwrap(),
None
serde_json::from_slice::<serde_json::Value>(&wire).unwrap(),
json!({"timestamp": 100.0, "response": "[1,2]"})
);
assert_eq!(
ResponseCacheCodec.decode(&wire).unwrap().response,
json!([1, 2])
);
assert_eq!(
ResponseCacheCodec.decode(br#"{"timestamp": 100.0, "response": "not serialized"}"#),
Err(Error::InvalidEntry)
);
}
@ -367,3 +390,27 @@ async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses() {
None
);
}
#[tokio::test]
async fn deferred_entries_keep_the_time_they_were_produced() {
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
let mut request = request();
request.max_age = Some(Duration::from_secs(10));
cache
.async_store_entries(vec![(
request.clone(),
json!({"answer": 7}),
Duration::from_secs(100),
)])
.await
.unwrap();
assert_eq!(
cache.lookup(&request, Duration::from_secs(110)).unwrap(),
Some(json!({"answer": 7}))
);
assert_eq!(
cache.lookup(&request, Duration::from_secs(111)).unwrap(),
None
);
}

3
litellm-rust/crates/cache/src/base_cache.rs vendored
View file

@ -1,5 +1,4 @@
use std::future::Future;
use std::time::Duration;
use std::{future::Future, time::Duration};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};

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

@ -1,8 +1,7 @@
use std::sync::Arc;
use crate::{BaseCache, CacheKwargs, Error};
pub use crate::BaseCache as Cache;
use crate::{BaseCache, CacheKwargs, Error};
pub fn get_cache<B: BaseCache>(
cache: &B,

8
litellm-rust/crates/cache/src/codec.rs vendored
View file

@ -14,6 +14,14 @@ pub trait CacheCodec: Send + Sync {
pub struct JsonCodec<V>(PhantomData<fn() -> V>);
impl<V> Clone for JsonCodec<V> {
fn clone(&self) -> Self {
*self
}
}
impl<V> Copy for JsonCodec<V> {}
impl<V> Default for JsonCodec<V> {
fn default() -> Self {
Self::new()

302
litellm-rust/crates/cache/src/dual.rs vendored
View file

@ -1,15 +1,140 @@
use std::sync::Arc;
use std::{sync::Arc, time::Duration};
use crate::{BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error};
use crate::{
BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error,
};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum ReadPolicy {
#[default]
LocalThenRemote,
LocalOnly,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum WritePolicy {
#[default]
Both,
LocalOnly,
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub enum RemoteFailurePolicy {
#[default]
Propagate,
UseLocal,
}
pub struct DualCache<L1, L2> {
l1: Arc<L1>,
l2: Arc<L2>,
read_policy: ReadPolicy,
write_policy: WritePolicy,
remote_failure_policy: RemoteFailurePolicy,
promotion_ttl: Option<Duration>,
}
impl<L1, L2> DualCache<L1, L2> {
pub fn new(l1: Arc<L1>, l2: Arc<L2>) -> Self {
Self { l1, l2 }
Self {
l1,
l2,
read_policy: ReadPolicy::default(),
write_policy: WritePolicy::default(),
remote_failure_policy: RemoteFailurePolicy::default(),
promotion_ttl: None,
}
}
pub fn with_read_policy(self, read_policy: ReadPolicy) -> Self {
Self {
read_policy,
..self
}
}
pub fn with_write_policy(self, write_policy: WritePolicy) -> Self {
Self {
write_policy,
..self
}
}
pub fn with_remote_failure_policy(self, remote_failure_policy: RemoteFailurePolicy) -> Self {
Self {
remote_failure_policy,
..self
}
}
pub fn with_promotion_ttl(self, promotion_ttl: Duration) -> Self {
Self {
promotion_ttl: Some(promotion_ttl),
..self
}
}
fn reads_remote(&self) -> bool {
self.read_policy == ReadPolicy::LocalThenRemote
}
fn writes_remote(&self) -> bool {
self.write_policy == WritePolicy::Both
}
fn remote<T>(&self, result: Result<T, Error>) -> Result<Option<T>, Error> {
match result {
Ok(value) => Ok(Some(value)),
Err(Error::Unavailable)
if self.remote_failure_policy == RemoteFailurePolicy::UseLocal =>
{
Ok(None)
}
Err(error) => Err(error),
}
}
fn promotion_kwargs(&self, kwargs: &CacheKwargs) -> CacheKwargs {
CacheKwargs {
ttl: self.promotion_ttl.or(kwargs.ttl),
extras: kwargs.extras.clone(),
}
}
}
impl<V, L1, L2> DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
L1: BaseCache<Value = V>,
L2: BaseCache<Value = V>,
{
fn missing(entries: &[BatchEntry<V>]) -> Vec<usize> {
entries
.iter()
.enumerate()
.filter_map(|(index, entry)| (!matches!(entry, BatchEntry::Hit(_))).then_some(index))
.collect()
}
fn merge_batch(
&self,
keys: &[String],
kwargs: &CacheKwargs,
mut entries: Vec<BatchEntry<V>>,
missing: Vec<usize>,
remote: Vec<BatchEntry<V>>,
) -> Result<Vec<BatchEntry<V>>, Error> {
if missing.len() != remote.len() {
return Err(Error::Unavailable);
}
for (index, entry) in missing.into_iter().zip(remote) {
if let BatchEntry::Hit(value) = &entry {
self.l1
.set_cache(&keys[index], value.clone(), self.promotion_kwargs(kwargs))?;
}
entries[index] = entry;
}
Ok(entries)
}
}
@ -21,12 +146,14 @@ where
{
type Value = V;
fn default_ttl(&self) -> std::time::Duration {
fn default_ttl(&self) -> Duration {
self.l2.default_ttl()
}
fn set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> {
self.l2.set_cache(key, value.clone(), kwargs.clone())?;
if self.writes_remote() {
self.remote(self.l2.set_cache(key, value.clone(), kwargs.clone()))?;
}
self.l1.set_cache(key, value, kwargs)
}
@ -34,25 +161,130 @@ where
if let Some(value) = self.l1.get_cache(key, kwargs)? {
return Ok(Some(value));
}
let value = self.l2.get_cache(key, kwargs)?;
if !self.reads_remote() {
return Ok(None);
}
let value = self.remote(self.l2.get_cache(key, kwargs))?.flatten();
if let Some(value) = &value {
self.l1.set_cache(key, value.clone(), kwargs.clone())?;
self.l1
.set_cache(key, value.clone(), self.promotion_kwargs(kwargs))?;
}
Ok(value)
}
fn get_cache_batch(
&self,
keys: &[String],
kwargs: &CacheKwargs,
) -> Result<Vec<BatchEntry<V>>, Error> {
let entries = self.l1.get_cache_batch(keys, kwargs)?;
let missing = Self::missing(&entries);
if missing.is_empty() || !self.reads_remote() {
return Ok(entries);
}
let remote_keys = missing
.iter()
.map(|index| keys[*index].clone())
.collect::<Vec<_>>();
match self.remote(self.l2.get_cache_batch(&remote_keys, kwargs))? {
Some(remote) => self.merge_batch(keys, kwargs, entries, missing, remote),
None => Ok(entries),
}
}
async fn async_set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> {
if self.writes_remote() {
self.remote(
self.l2
.async_set_cache(key, value.clone(), kwargs.clone())
.await,
)?;
}
self.l1.async_set_cache(key, value, kwargs).await
}
async fn async_get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result<Option<V>, Error> {
if let Some(value) = self.l1.async_get_cache(key, kwargs).await? {
return Ok(Some(value));
}
if !self.reads_remote() {
return Ok(None);
}
let value = self
.remote(self.l2.async_get_cache(key, kwargs).await)?
.flatten();
if let Some(value) = &value {
self.l1
.async_set_cache(key, value.clone(), self.promotion_kwargs(kwargs))
.await?;
}
Ok(value)
}
async fn async_get_cache_batch(
&self,
keys: Vec<String>,
kwargs: CacheKwargs,
) -> Result<Vec<BatchEntry<V>>, Error> {
let entries = self
.l1
.async_get_cache_batch(keys.clone(), kwargs.clone())
.await?;
let missing = Self::missing(&entries);
if missing.is_empty() || !self.reads_remote() {
return Ok(entries);
}
let remote_keys = missing.iter().map(|index| keys[*index].clone()).collect();
match self.remote(
self.l2
.async_get_cache_batch(remote_keys, kwargs.clone())
.await,
)? {
Some(remote) => self.merge_batch(&keys, &kwargs, entries, missing, remote),
None => Ok(entries),
}
}
async fn async_set_cache_pipeline(
&self,
cache_list: Vec<(String, V)>,
kwargs: CacheKwargs,
) -> Result<(), Error> {
if self.writes_remote() {
self.remote(
self.l2
.async_set_cache_pipeline(cache_list.clone(), kwargs.clone())
.await,
)?;
}
self.l1.async_set_cache_pipeline(cache_list, kwargs).await
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.l2.delete_cache(key)?;
if self.writes_remote() {
self.remote(self.l2.delete_cache(key))?;
}
self.l1.delete_cache(key)
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.async_delete_cache(key).await)?;
}
self.l1.async_delete_cache(key).await
}
fn flush_cache(&self) -> Result<(), Error> {
self.l2.flush_cache()?;
if self.writes_remote() {
self.remote(self.l2.flush_cache())?;
}
self.l1.flush_cache()
}
async fn async_flush_cache(&self) -> Result<(), Error> {
self.l2.async_flush_cache().await?;
if self.writes_remote() {
self.remote(self.l2.async_flush_cache().await)?;
}
self.l1.async_flush_cache().await
}
@ -76,6 +308,20 @@ where
self.l1.set_cache(key, value, kwargs)?;
Ok(value)
}
async fn async_increment_cache(
&self,
key: &str,
amount: f64,
kwargs: CacheKwargs,
) -> Result<f64, Error> {
let value = self
.l2
.async_increment_cache(key, amount, kwargs.clone())
.await?;
self.l1.async_set_cache(key, value, kwargs).await?;
Ok(value)
}
}
impl<V, L1, L2> ClaimCache for DualCache<L1, L2>
@ -91,15 +337,39 @@ where
eligible: &[V],
kwargs: CacheKwargs,
) -> Result<V, Error> {
match self
.l2
.claim_cache(key, candidate.clone(), eligible, kwargs.clone())
{
Ok(winner) => {
match self.remote(
self.l2
.claim_cache(key, candidate.clone(), eligible, kwargs.clone()),
)? {
Some(winner) => {
self.l1.set_cache(key, winner.clone(), kwargs)?;
Ok(winner)
}
Err(_) => self.l1.claim_cache(key, candidate, eligible, kwargs),
None => self.l1.claim_cache(key, candidate, eligible, kwargs),
}
}
async fn async_claim_cache(
&self,
key: &str,
candidate: V,
eligible: Vec<V>,
kwargs: CacheKwargs,
) -> Result<V, Error> {
match self.remote(
self.l2
.async_claim_cache(key, candidate.clone(), eligible.clone(), kwargs.clone())
.await,
)? {
Some(winner) => {
self.l1.async_set_cache(key, winner.clone(), kwargs).await?;
Ok(winner)
}
None => {
self.l1
.async_claim_cache(key, candidate, eligible, kwargs)
.await
}
}
}
}

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

@ -1,6 +1,7 @@
use litellm_cache::{BaseCache, CacheConnectionResult, CacheKwargs, Error};
use std::{sync::Mutex, time::Duration};
use litellm_cache::{BaseCache, CacheConnectionResult, CacheKwargs, Error};
struct TestCache {
default_ttl: Duration,
writes: Mutex<Vec<(String, String, CacheKwargs)>>,

216
litellm-rust/crates/cache/tests/dual.rs vendored
View file

@ -4,7 +4,8 @@ use std::{
};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, dual::DualCache,
BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error,
dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy},
};
struct TestCache<V> {
@ -111,7 +112,8 @@ fn failed_l2_increment_leaves_l1_unchanged() {
#[test]
fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() {
let l1 = Arc::new(TestCache::new(Some("first".to_string()), false));
let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true)));
let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true)))
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
assert_eq!(
cache
@ -128,3 +130,213 @@ fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() {
"first"
);
}
struct SyncPanics(TestCache<String>);
impl BaseCache for SyncPanics {
type Value = String;
fn set_cache(&self, _: &str, _: String, _: CacheKwargs) -> Result<(), Error> {
panic!("sync L2 write on an async path")
}
fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result<Option<String>, Error> {
panic!("sync L2 read on an async path")
}
async fn async_set_cache(
&self,
key: &str,
value: String,
kwargs: CacheKwargs,
) -> Result<(), Error> {
self.0.set_cache(key, value, kwargs)
}
async fn async_get_cache(
&self,
key: &str,
kwargs: &CacheKwargs,
) -> Result<Option<String>, Error> {
self.0.get_cache(key, kwargs)
}
async fn async_get_cache_batch(
&self,
keys: Vec<String>,
kwargs: CacheKwargs,
) -> Result<Vec<litellm_cache::BatchEntry<String>>, Error> {
assert_eq!(keys, ["missing"]);
Ok(vec![match self.0.get_cache("missing", &kwargs)? {
Some(value) => litellm_cache::BatchEntry::Hit(value),
None => litellm_cache::BatchEntry::Miss,
}])
}
async fn async_set_cache_pipeline(
&self,
cache_list: Vec<(String, String)>,
kwargs: CacheKwargs,
) -> Result<(), Error> {
for (key, value) in cache_list {
self.0.set_cache(&key, value, kwargs.clone())?;
}
Ok(())
}
fn delete_cache(&self, _: &str) -> Result<(), Error> {
panic!("sync L2 delete on an async path")
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
self.0.delete_cache(key)
}
fn flush_cache(&self) -> Result<(), Error> {
panic!("sync L2 flush on an async path")
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
unreachable!()
}
}
#[tokio::test]
async fn async_operations_use_the_async_l2_methods() {
let l1 = Arc::new(TestCache::new(None, false));
let cache = DualCache::new(
l1.clone(),
Arc::new(SyncPanics(TestCache::new(
Some("remote".to_string()),
false,
))),
);
let kwargs = CacheKwargs::default();
assert_eq!(
cache.async_get_cache("missing", &kwargs).await.unwrap(),
Some("remote".into())
);
assert_eq!(
l1.get_cache("missing", &kwargs).unwrap(),
Some("remote".into())
);
l1.delete_cache("missing").unwrap();
assert_eq!(
cache
.async_get_cache_batch(vec!["missing".into()], kwargs.clone())
.await
.unwrap(),
[litellm_cache::BatchEntry::Hit("remote".to_string())]
);
cache
.async_set_cache("missing", "written".into(), kwargs.clone())
.await
.unwrap();
cache
.async_set_cache_pipeline(vec![("missing".into(), "piped".into())], kwargs.clone())
.await
.unwrap();
cache.async_delete_cache("missing").await.unwrap();
assert_eq!(
cache.async_get_cache("missing", &kwargs).await.unwrap(),
None
);
}
struct Unavailable;
impl BaseCache for Unavailable {
type Value = String;
fn set_cache(&self, _: &str, _: String, _: CacheKwargs) -> Result<(), Error> {
Err(Error::Unavailable)
}
fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result<Option<String>, Error> {
Err(Error::Unavailable)
}
fn delete_cache(&self, _: &str) -> Result<(), Error> {
Err(Error::Unavailable)
}
fn flush_cache(&self) -> Result<(), Error> {
Err(Error::Unavailable)
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
unreachable!()
}
}
impl ClaimCache for Unavailable {
fn claim_cache(
&self,
_: &str,
_: String,
_: &[String],
_: CacheKwargs,
) -> Result<String, Error> {
Err(Error::InvalidEntry)
}
}
#[test]
fn remote_failure_policy_selects_propagation_or_the_local_tier() {
let kwargs = CacheKwargs::default();
let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable));
assert_eq!(
strict.set_cache("key", "value".into(), kwargs.clone()),
Err(Error::Unavailable)
);
assert_eq!(strict.get_cache("key", &kwargs), Err(Error::Unavailable));
let l1 = Arc::new(TestCache::new(None, false));
let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable))
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
assert_eq!(degraded.get_cache("key", &kwargs), Ok(None));
degraded
.set_cache("key", "value".into(), kwargs.clone())
.unwrap();
assert_eq!(degraded.get_cache("key", &kwargs), Ok(Some("value".into())));
degraded.delete_cache("key").unwrap();
assert_eq!(l1.get_cache("key", &kwargs), Ok(None));
}
#[test]
fn claim_fallback_does_not_hide_non_availability_errors() {
let cache = DualCache::new(
Arc::new(TestCache::new(Some("first".to_string()), false)),
Arc::new(Unavailable),
)
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
assert_eq!(
cache.claim_cache("affinity", "second".into(), &[], CacheKwargs::default()),
Err(Error::InvalidEntry)
);
}
#[test]
fn local_only_policies_never_touch_l2() {
let l2 = Arc::new(TestCache::new(Some("remote".to_string()), false));
let cache = DualCache::new(Arc::new(TestCache::new(None, false)), l2.clone())
.with_read_policy(ReadPolicy::LocalOnly)
.with_write_policy(WritePolicy::LocalOnly);
let kwargs = CacheKwargs::default();
assert_eq!(cache.get_cache("key", &kwargs), Ok(None));
cache
.set_cache("key", "local".into(), kwargs.clone())
.unwrap();
assert_eq!(l2.get_cache("key", &kwargs), Ok(Some("remote".into())));
}

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

@ -1,3 +1,5 @@
use std::time::Duration;
use litellm_host_python::from_py;
use pyo3::{
PyTraverseError, PyVisit,
@ -6,9 +8,8 @@ use pyo3::{
types::{PyDict, PyTuple, PyType},
};
use serde_json::Value;
use std::time::Duration;
use super::{NativeCacheHandle, native::NativeResponseCache};
use super::{CacheTestHandle, native::NativeResponseCache};
struct ClassGuard {
class: Py<PyType>,
@ -130,9 +131,10 @@ impl FacadeGuard {
pub(super) fn capture(
py: Python<'_>,
facade: &Bound<'_, PyAny>,
kind: &str,
native_default_ttl: Duration,
service: &NativeResponseCache,
) -> PyResult<Self> {
let kind = service.kind();
let native_default_ttl: Duration = service.default_ttl();
let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?;
if !facade.get_type().is(&cache_type) {
return Err(PyTypeError::new_err(
@ -158,6 +160,23 @@ impl FacadeGuard {
"facade and native backend default TTLs must match",
));
}
let namespace = match backend.getattr_opt("namespace")? {
Some(namespace) => namespace.extract::<Option<String>>()?,
None => None,
}
.filter(|namespace| !namespace.is_empty());
if kind == "redis" && namespace.as_deref() != service.namespace() {
return Err(PyTypeError::new_err(
"facade and native backend namespaces must match",
));
}
if let Some(capacity) = service.capacity()
&& backend.getattr("max_size_in_memory")?.extract::<usize>()? != capacity
{
return Err(PyTypeError::new_err(
"facade and native backend capacities must match",
));
}
Ok(Self {
outer: ObjectGuard::capture(
py,
@ -169,6 +188,7 @@ impl FacadeGuard {
"namespace",
"supported_call_types",
"redis_flush_size",
"semantic_cache_scope",
],
)?,
backend: ObjectGuard::capture(
@ -179,6 +199,8 @@ impl FacadeGuard {
"default_ttl",
"max_size_in_memory",
"max_size_per_item",
"redis_kwargs",
"redis_flush_size",
],
)?,
})
@ -208,7 +230,7 @@ pub(super) fn resolve(
let Some(handle) = dict.get_item("_native_cache_handle")? else {
return Ok(None);
};
let Ok(handle) = handle.extract::<PyRef<'_, NativeCacheHandle>>() else {
let Ok(handle) = handle.extract::<PyRef<'_, CacheTestHandle>>() else {
return Ok(None);
};
let Some(guard) = &handle.guard else {

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

@ -3,21 +3,21 @@ mod native;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use facade::FacadeGuard;
use litellm_cache::Error;
use litellm_cache_response::{CacheControls, CacheKeyInput, PartialHits, ResponseCacheRequest};
use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py};
use native::NativeResponseCache;
use pyo3::{
PyTraverseError, PyVisit,
exceptions::{PyRuntimeError, PyTypeError, PyValueError},
prelude::*,
types::{PyDict, PyList},
types::{PyDict, PyList, PyTuple},
};
use serde::Deserialize;
use serde_json::Value;
use crate::python_settings::PythonSettings;
use facade::FacadeGuard;
use native::NativeResponseCache;
const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60);
@ -84,14 +84,14 @@ fn cache_error(error: Error) -> PyErr {
}
}
#[pyclass(frozen)]
pub(crate) struct NativeCacheHandle {
#[pyclass(frozen, name = "_CacheTestHandle")]
pub(crate) struct CacheTestHandle {
service: NativeResponseCache,
guard: Option<FacadeGuard>,
pid: u32,
}
impl NativeCacheHandle {
impl CacheTestHandle {
fn service(&self) -> PyResult<NativeResponseCache> {
if self.pid != std::process::id() {
return Err(PyRuntimeError::new_err(
@ -103,7 +103,7 @@ impl NativeCacheHandle {
}
#[pymethods]
impl NativeCacheHandle {
impl CacheTestHandle {
#[staticmethod]
#[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))]
fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult<Self> {
@ -140,9 +140,9 @@ impl NativeCacheHandle {
self.service.kind()
}
fn bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> {
fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> {
let service = self.service()?;
let guard = FacadeGuard::capture(py, facade, self.backend(), service.default_ttl())?;
let guard = FacadeGuard::capture(py, facade, &service)?;
let service = service.with_redis_flush_size(
facade
.getattr("redis_flush_size")?
@ -173,7 +173,7 @@ enum CacheBinding {
PythonCallback(Py<PyAny>),
}
#[pyclass(frozen, name = "CacheBinding")]
#[pyclass(frozen, name = "_CacheTestBinding")]
pub(crate) struct ResolvedCache {
binding: CacheBinding,
pid: u32,
@ -285,12 +285,15 @@ impl ResolvedCache {
}
}
/// Native bindings return `{values, missing_indices}`. The built-in `Cache` API has no batch
/// read, so a Python callback receives one `get_cache(**kwargs)` call per request, in order,
/// and the results come back as a list.
#[pyo3(signature = (requests, *, callback_kwargs=None))]
fn lookup_batch(
&self,
py: Python<'_>,
requests: &Bound<'_, PyAny>,
callback_kwargs: Option<&Bound<'_, PyDict>>,
callback_kwargs: Option<&Bound<'_, PyAny>>,
) -> PyResult<Py<PyAny>> {
self.check_process()?;
match &self.binding {
@ -305,14 +308,17 @@ impl ResolvedCache {
.map_err(cache_error)?;
to_py(py, &response)
}
CacheBinding::PythonCallback(object) => object
.bind(py)
.call_method(
"batch_get_cache",
(callback_keys(py, requests)?,),
Some(self::callback_kwargs(callback_kwargs)?),
)
.map(Bound::unbind),
CacheBinding::PythonCallback(object) => {
let results = PyList::empty(py);
for kwargs in batch_callback_kwargs(requests, callback_kwargs)? {
results.append(object.bind(py).call_method(
"get_cache",
(),
Some(&kwargs),
)?)?;
}
Ok(results.into_any().unbind())
}
}
}
@ -364,7 +370,7 @@ impl ResolvedCache {
&self,
py: Python<'py>,
requests: &Bound<'py, PyAny>,
callback_kwargs: Option<&Bound<'py, PyDict>>,
callback_kwargs: Option<&Bound<'py, PyAny>>,
) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
match &self.binding {
@ -381,20 +387,30 @@ impl ResolvedCache {
cache_error,
)
}
CacheBinding::PythonCallback(object) => object.bind(py).call_method(
"async_batch_get_cache",
(callback_keys(py, requests)?,),
Some(self::callback_kwargs(callback_kwargs)?),
),
CacheBinding::PythonCallback(object) => {
let awaitables = batch_callback_kwargs(requests, callback_kwargs)?
.iter()
.map(|kwargs| {
object
.bind(py)
.call_method("async_get_cache", (), Some(kwargs))
})
.collect::<PyResult<Vec<_>>>()?;
py.import("asyncio")?
.call_method1("gather", PyTuple::new(py, awaitables)?)
}
}
}
#[pyo3(signature = (requests, responses, *, callback_kwargs=None))]
/// A Python callback receives the caller's original result through `callback_result`, because
/// the built-in `Cache.async_add_cache_pipeline` splits the batch itself.
#[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))]
fn async_store_batch<'py>(
&self,
py: Python<'py>,
requests: &Bound<'py, PyAny>,
responses: &Bound<'py, PyAny>,
callback_result: Option<&Bound<'py, PyAny>>,
callback_kwargs: Option<&Bound<'py, PyDict>>,
) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
@ -417,20 +433,14 @@ impl ResolvedCache {
)
}
CacheBinding::PythonCallback(object) => {
let keys = callback_keys(py, requests)?;
let responses = responses.try_iter()?.collect::<PyResult<Vec<_>>>()?;
if keys.len() != responses.len() {
return Err(PyValueError::new_err(
"batch cache requests and responses must have equal lengths",
));
}
let cache_list = PyList::empty(py);
for (key, response) in keys.iter().zip(responses) {
cache_list.append((key, response))?;
}
let result = callback_result.ok_or_else(|| {
PyTypeError::new_err(
"Python cache callbacks require their original callback_result",
)
})?;
object.bind(py).call_method(
"async_set_cache_pipeline",
(cache_list,),
"async_add_cache_pipeline",
(result,),
Some(self::callback_kwargs(callback_kwargs)?),
)
}
@ -445,8 +455,17 @@ impl ResolvedCache {
let service = service.clone();
run_async(py, async move { service.async_flush().await }, cache_error)
}
// The built-in `Cache` facade has no flush of its own; its backend does.
CacheBinding::PythonCallback(object) => {
object.bind(py).call_method0("flush_cache")?;
let object = object.bind(py);
let backend = match object.getattr_opt("cache")? {
Some(backend) if !backend.is_none() => backend,
_ => object.clone(),
};
if backend.hasattr("async_flush_cache")? {
return backend.call_method0("async_flush_cache");
}
backend.call_method0("flush_cache")?;
ready_none(py)
}
}
@ -464,7 +483,7 @@ impl ResolvedCache {
cache_error,
)
}
CacheBinding::PythonCallback(object) => object.bind(py).call_method0("test_connection"),
CacheBinding::PythonCallback(object) => object.bind(py).call_method0("ping"),
}
}
@ -484,16 +503,25 @@ fn callback_kwargs<'a, 'py>(
})
}
fn callback_keys<'py>(
py: Python<'py>,
fn batch_callback_kwargs<'py>(
requests: &Bound<'py, PyAny>,
) -> PyResult<Bound<'py, PyList>> {
PyList::new(
py,
self::requests(requests)?
.into_iter()
.map(|request| litellm_cache_response::cache_key(&request.key)),
)
kwargs: Option<&Bound<'py, PyAny>>,
) -> PyResult<Vec<Bound<'py, PyDict>>> {
let kwargs = kwargs
.ok_or_else(|| {
PyTypeError::new_err(
"Python cache callbacks require one original callback_kwargs mapping per request",
)
})?
.try_iter()?
.map(|item| Ok(item?.cast_into::<PyDict>()?))
.collect::<PyResult<Vec<_>>>()?;
if kwargs.len() != requests.len()? {
return Err(PyValueError::new_err(
"batch cache requests and callback_kwargs must have equal lengths",
));
}
Ok(kwargs)
}
fn ready_none(py: Python<'_>) -> PyResult<Bound<'_, PyAny>> {
@ -512,13 +540,13 @@ fn ready_value<'py, T: serde::Serialize>(
Ok(future)
}
#[pyclass(frozen)]
pub(crate) struct CacheResolver {
#[pyclass(frozen, name = "_CacheTestResolver")]
pub(crate) struct CacheTestResolver {
namespace: Py<PyAny>,
}
#[pymethods]
impl CacheResolver {
impl CacheTestResolver {
#[new]
fn new(namespace: Py<PyAny>) -> Self {
Self { namespace }
@ -528,7 +556,7 @@ impl CacheResolver {
let object = self.namespace.bind(py).getattr("cache")?;
let binding = if object.is_none() {
CacheBinding::Disabled
} else if let Ok(handle) = object.extract::<PyRef<'_, NativeCacheHandle>>() {
} else if let Ok(handle) = object.extract::<PyRef<'_, CacheTestHandle>>() {
CacheBinding::Native(handle.service()?)
} else if let Some(service) = facade::resolve(py, &object)? {
CacheBinding::Native(service)

37
litellm-rust/crates/python-bridge/src/cache/native.rs vendored
View file

@ -3,12 +3,11 @@ use std::{sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::RedisCache;
use serde_json::Value;
use tokio::sync::Mutex;
use litellm_cache_response::{
CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
};
use serde_json::Value;
use tokio::sync::Mutex;
#[derive(Clone)]
pub(super) enum NativeResponseCache {
@ -21,7 +20,7 @@ pub(super) enum NativeResponseCache {
pub(super) struct RedisWriteBuffer {
flush_size: usize,
entries: Mutex<Vec<(ResponseCacheRequest, Value)>>,
entries: Mutex<Vec<(ResponseCacheRequest, Value, Duration)>>,
}
impl NativeResponseCache {
@ -67,6 +66,20 @@ impl NativeResponseCache {
}
}
pub fn namespace(&self) -> Option<&str> {
match self {
Self::Memory(_) => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
}
}
pub fn capacity(&self) -> Option<usize> {
match self {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. } => None,
}
}
pub fn with_redis_flush_size(self, flush_size: Option<usize>) -> Self {
match self {
Self::Redis { cache, .. } => Self::Redis {
@ -145,19 +158,15 @@ impl NativeResponseCache {
} => {
let pending = {
let mut entries = buffer.entries.lock().await;
entries.push((request.clone(), response));
entries.push((request.clone(), response, now));
(entries.len() >= buffer.flush_size).then(|| std::mem::take(&mut *entries))
};
let Some(pending) = pending else {
return Ok(());
};
if let Err(error) = cache.async_store_batch(pending.clone(), now).await {
let mut entries = buffer.entries.lock().await;
let current = std::mem::take(&mut *entries);
*entries = pending.into_iter().chain(current).collect();
return Err(error);
// A failed flush drops its batch, as Python does. Requeueing would grow the
// buffer and re-send an ever larger pipeline on every write during an outage.
match pending {
Some(pending) => cache.async_store_entries(pending).await,
None => Ok(()),
}
Ok(())
}
}
}

16
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -10,8 +10,7 @@ mod token_counter;
#[pymodule(gil_used = true)]
mod _native {
#[pymodule_export]
use crate::cache::{CacheResolver, NativeCacheHandle, ResolvedCache};
use crate::cache::{CacheTestHandle, CacheTestResolver, ResolvedCache};
#[cfg(feature = "panic-test")]
#[pymodule_export]
use crate::diagnostics::_panic_for_test;
@ -35,6 +34,16 @@ mod _native {
use crate::token_counter::TokenCounter;
#[pymodule_export]
use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking};
use pyo3::{prelude::*, types::PyModule};
#[pymodule_init]
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
let py = module.py();
let dict = module.dict();
dict.set_item("_CacheTestHandle", py.get_type::<CacheTestHandle>())?;
dict.set_item("_CacheTestResolver", py.get_type::<CacheTestResolver>())?;
dict.set_item("_CacheTestBinding", py.get_type::<ResolvedCache>())
}
}
use pyo3::prelude::*;
@ -68,9 +77,6 @@ mod tests {
"achat_completions",
"ResponsesWebSocketConnection",
"TokenCounter",
"CacheResolver",
"NativeCacheHandle",
"CacheBinding",
"gil_stats",
"process_state_started",
"reserve_process_for_forking",

26
litellm/rust_bridge/_native.pyi
View file

@ -94,25 +94,25 @@ class ResponsesWebSocketConnection:
def close(self) -> Future[None]: ...
@final
class NativeCacheHandle:
class _CacheTestHandle:
def __new__(cls, _uninstantiable: Never, /) -> Never: ...
@staticmethod
def memory(
*, capacity: int = 200, ttl_seconds: float = 600.0, max_entry_bytes: int = 1048576
) -> NativeCacheHandle: ...
) -> _CacheTestHandle: ...
@staticmethod
def redis(url: str, *, ttl_seconds: float | None = None, namespace: str | None = None) -> NativeCacheHandle: ...
def redis(url: str, *, ttl_seconds: float | None = None, namespace: str | None = None) -> _CacheTestHandle: ...
@property
def backend(self) -> str: ...
def bind_facade(self, facade: object) -> None: ...
def _bind_facade(self, facade: object) -> None: ...
@final
class CacheResolver:
def __new__(cls, namespace: object) -> CacheResolver: ...
def resolve(self) -> CacheBinding: ...
class _CacheTestResolver:
def __new__(cls, namespace: object) -> _CacheTestResolver: ...
def resolve(self) -> _CacheTestBinding: ...
@final
class CacheBinding:
class _CacheTestBinding:
def __new__(cls, _uninstantiable: Never, /) -> Never: ...
@property
def kind(self) -> str: ...
@ -130,7 +130,7 @@ class CacheBinding:
self,
requests: Sequence[Mapping[str, object]],
*,
callback_kwargs: dict[str, object] | None = None,
callback_kwargs: Sequence[dict[str, object]] | None = None,
) -> object: ...
def async_lookup(
self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None
@ -146,17 +146,18 @@ class CacheBinding:
self,
requests: Sequence[Mapping[str, object]],
*,
callback_kwargs: dict[str, object] | None = None,
callback_kwargs: Sequence[dict[str, object]] | None = None,
) -> Awaitable[object]: ...
def async_store_batch(
self,
requests: Sequence[Mapping[str, object]],
responses: Sequence[object],
*,
callback_result: object = None,
callback_kwargs: dict[str, object] | None = None,
) -> Awaitable[object]: ...
def async_flush(self) -> Awaitable[None]: ...
def ping(self) -> Awaitable[dict[str, object] | None]: ...
def ping(self) -> Awaitable[object]: ...
@final
class TokenCounter:
@ -174,10 +175,7 @@ def process_state_started() -> bool: ...
def reserve_process_for_forking() -> None: ...
__all__ = [
"CacheBinding",
"CacheResolver",
"ForkedAfterNativeRuntimeStarted",
"NativeCacheHandle",
"ProcessReservedForForking",
"ResponsesWebSocketConnection",
"RustBridgeDeclined",

2
tests/test_litellm/caching/test_dual_cache.py
View file

@ -6,10 +6,10 @@ from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache, _redis_circuit_breaker_guard, _redis_circuit_breaker_guard_sync
from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE
from litellm.types.caching import RedisPipelineIncrementOperation

176
tests/test_litellm_rust/test_cache.py
View file

@ -15,7 +15,7 @@ import pytest
import redis
import litellm
from litellm.caching.caching import Cache
from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.rust_bridge import _native
from litellm.types.caching import LiteLLMCacheType
@ -50,20 +50,43 @@ def test_existing_constructor_and_global_are_unchanged() -> None:
assert type(facade.cache) is InMemoryCache
assert "_native_cache_handle" not in vars(facade)
with rebound(litellm, "cache", facade):
resolver: Final = _native.CacheResolver(litellm)
resolver: Final = _native._CacheTestResolver(litellm)
assert resolver.resolve().kind == "python_callback"
resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"})
assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7}
def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None:
resolver: Final = _native._CacheTestResolver(litellm)
enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30)
enabled: Final = litellm.cache
assert isinstance(enabled, Cache)
assert enabled.ttl == 30
assert resolver.resolve().kind == "python_callback"
enable_cache(type=LiteLLMCacheType.LOCAL, ttl=60)
assert litellm.cache is enabled
update_cache(type=LiteLLMCacheType.LOCAL, ttl=60)
updated: Final = litellm.cache
assert isinstance(updated, Cache)
assert updated is not enabled
assert updated.ttl == 60
disable_cache()
assert litellm.cache is None
assert resolver.resolve().kind == "disabled"
async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None:
namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.memory())
resolver: Final = _native.CacheResolver(namespace)
namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.memory())
resolver: Final = _native._CacheTestResolver(namespace)
selected: Final = resolver.resolve()
assert selected.kind == "native"
selected.store(request(), {"answer": 1})
assert await selected.async_lookup(request()) == {"answer": 1}
with rebound(namespace, "cache", _native.NativeCacheHandle.memory()):
with rebound(namespace, "cache", _native._CacheTestHandle.memory()):
replacement: Final = resolver.resolve()
await selected.async_store(request(), {"answer": 2})
assert replacement.lookup(request()) is None
@ -96,7 +119,7 @@ async def test_python_callback_preserves_identity_caller_task_context_and_errors
raise failure
namespace: Final = SimpleNamespace(cache=CustomCache())
binding: Final = _native.CacheResolver(namespace).resolve()
binding: Final = _native._CacheTestResolver(namespace).resolve()
assert binding.kind == "python_callback"
assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel
assert context.get() == "callback"
@ -117,7 +140,7 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None:
finally:
finished.set()
binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve()
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve()
async def lookup() -> object:
return await binding.async_lookup(None, callback_kwargs={})
@ -132,9 +155,9 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None:
def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None:
facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
handle: Final = _native.NativeCacheHandle.memory()
handle.bind_facade(facade)
resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade))
handle: Final = _native._CacheTestHandle.memory()
handle._bind_facade(facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
native: Final = resolver.resolve()
assert native.kind == "native"
native.store(request(), {"source": "native"})
@ -165,16 +188,18 @@ def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not
class CustomCache(Cache):
pass
handle: Final = _native.NativeCacheHandle.memory()
handle: Final = _native._CacheTestHandle.memory()
with pytest.raises(TypeError):
handle.bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL))
handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL))
facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
handle.bind_facade(facade)
resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade))
handle._bind_facade(facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
with rebound(facade, "cache", InMemoryCache()):
assert resolver.resolve().kind == "python_callback"
with rebound(facade, "ttl", 12):
assert resolver.resolve().kind == "python_callback"
with rebound(facade, "semantic_cache_scope", "end_user"):
assert resolver.resolve().kind == "python_callback"
def custom_key(**_kwargs: object) -> str:
return "custom"
@ -193,7 +218,7 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None:
def cyclic_reference() -> weakref.ReferenceType[CustomCache]:
callback: Final = CustomCache()
namespace: Final = SimpleNamespace(cache=callback)
binding: Final = _native.CacheResolver(namespace).resolve()
binding: Final = _native._CacheTestResolver(namespace).resolve()
setattr(callback, "binding", binding)
return weakref.ref(callback)
@ -204,8 +229,8 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None:
async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None:
client: Final = redis.Redis.from_url(redis_url)
namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url, namespace="team"))
binding: Final = _native.CacheResolver(namespace).resolve()
namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url, namespace="team"))
binding: Final = _native._CacheTestResolver(namespace).resolve()
response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None}
envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)}
client.set("team:sync", str(envelope))
@ -227,33 +252,33 @@ 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()
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve()
for seconds in (-1.0, float("nan"), float("inf")):
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, match="cache durations must be finite and nonnegative"):
_native.NativeCacheHandle.memory(ttl_seconds=-1)
_native._CacheTestHandle.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()
handle: Final = _native._CacheTestHandle.memory(capacity=2, max_entry_bytes=128)
binding: Final = _native._CacheTestResolver(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))
disabled: Final = _native._CacheTestResolver(
SimpleNamespace(cache=_native._CacheTestHandle.memory(capacity=0))
).resolve()
await disabled.async_store(request(), small)
assert await disabled.async_lookup(request()) is None
async def test_native_batch_lookup_and_store_report_partial_hits() -> None:
binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve()
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve()
requests: Final = [request("hit"), request("miss"), request("disabled")]
requests[2]["controls"] = {
"supported_call_type": True,
@ -275,50 +300,72 @@ async def test_native_batch_lookup_and_store_report_partial_hits() -> None:
}
async def test_python_batch_callbacks_receive_keys_and_key_value_pairs() -> None:
first: Final = object()
second: Final = object()
class CustomCache:
def batch_get_cache(self, keys: list[str], *, marker: object) -> tuple[list[str], object]:
return keys, marker
async def async_batch_get_cache(self, keys: list[str], *, marker: object) -> tuple[list[str], object]:
return keys, marker
async def async_set_cache_pipeline(
self, cache_list: list[tuple[str, object]], *, marker: object
) -> tuple[list[tuple[str, object]], object]:
return cache_list, marker
async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None:
result: Final = object()
marker: Final = object()
binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve()
requests: Final = [request("first"), request("second")]
assert binding.lookup_batch(requests, callback_kwargs={"marker": marker}) == (["first", "second"], marker)
assert await binding.async_lookup_batch(requests, callback_kwargs={"marker": marker}) == (
["first", "second"],
marker,
)
class CustomCache(Cache):
def get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object:
return ("sync", kwargs)
async def async_get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object:
return ("async", kwargs)
async def async_add_cache_pipeline(
self, result: object, dynamic_cache_object: object = None, **kwargs: object
) -> object:
return result, kwargs
binding: Final = _native._CacheTestResolver(
SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL))
).resolve()
assert binding.kind == "python_callback"
requests: Final = [request("first"), request("second")]
kwargs: Final = [{"cache_key": "first"}, {"cache_key": "second"}]
assert binding.lookup_batch(requests, callback_kwargs=kwargs) == [("sync", kwargs[0]), ("sync", kwargs[1])]
assert await binding.async_lookup_batch(requests, callback_kwargs=kwargs) == [
("async", kwargs[0]),
("async", kwargs[1]),
]
with pytest.raises(ValueError, match="equal lengths"):
binding.lookup_batch(requests, callback_kwargs=kwargs[:1])
with pytest.raises(TypeError, match="callback_result"):
await binding.async_store_batch(requests, [1, 2], callback_kwargs={"marker": marker})
stored: Final = cast(
tuple[list[tuple[str, object]], object],
await binding.async_store_batch(
requests,
[first, second],
callback_kwargs={"marker": marker},
),
tuple[object, dict[str, object]],
await binding.async_store_batch(requests, [1, 2], callback_result=result, callback_kwargs={"marker": marker}),
)
assert [key for key, _ in stored[0]] == ["first", "second"]
assert stored[1] is marker
assert stored[0][0][1] is first
assert stored[0][1][1] is second
assert stored[0] is result
assert stored[1] == {"marker": marker}
async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None:
async def ping() -> str:
return "pong"
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
cache.cache.set_cache("key", "value")
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=cache)).resolve()
assert binding.kind == "python_callback"
setattr(cache.cache, "ping", ping)
assert await binding.ping() == "pong"
await binding.async_flush()
assert cache.cache.get_cache("key") is None
def test_facade_registration_rejects_mismatched_capacity() -> None:
facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
with pytest.raises(TypeError, match="capacities must match"):
_native._CacheTestHandle.memory(capacity=7)._bind_facade(facade)
async def test_redis_handle_reads_the_python_default_ttl(redis_url: str) -> None:
client: Final = redis.Redis.from_url(redis_url)
with rebound(litellm, "default_redis_ttl", 7):
binding: Final = _native.CacheResolver(
SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url))
binding: Final = _native._CacheTestResolver(
SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url))
).resolve()
await binding.async_store(request("native-default"), {"value": 1})
@ -336,11 +383,16 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None:
redis_flush_size=2,
)
with pytest.raises(TypeError, match="default TTLs must match"):
_native.NativeCacheHandle.redis(redis_url, ttl_seconds=61).bind_facade(facade)
_native.NativeCacheHandle.redis(redis_url).bind_facade(facade)
binding: Final = _native.CacheResolver(SimpleNamespace(cache=facade)).resolve()
_native._CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade)
with pytest.raises(TypeError, match="namespaces must match"):
_native._CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade)
_native._CacheTestHandle.redis(redis_url)._bind_facade(facade)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
client: Final = redis.Redis.from_url(redis_url)
with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}):
assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback"
await binding.async_store(request("first"), {"value": 1})
assert client.get("first") is None
await binding.async_store(request("second"), {"value": 2})