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(rust): align cache foundation with Python backends

Browse source
This commit is contained in:
Yujong Lee 2026-09-21 12:50:21 -07:00
parent efcafa7f12
commit 991108cc48
21 changed files with 1136 additions and 533 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1,13 +1,14 @@
use std::{
cmp::Reverse,
collections::{BinaryHeap, HashMap},
collections::{BinaryHeap, HashMap, HashSet},
hash::Hash,
sync::{Arc, Mutex},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache,
Error, IncrementOperation,
BaseCache, BatchCache, CacheConnectionResult, CacheConnectionStatus, ClaimCache, CounterCache,
DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache,
};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
@ -219,7 +220,7 @@ where
key: &str,
candidate: V,
eligible: &[V],
kwargs: CacheKwargs,
context: ExactCacheContext,
) -> Result<V, Error> {
if self.max_size_in_memory == 0 {
return Ok(candidate);
@ -239,14 +240,23 @@ where
return Ok(existing.clone());
}
let winner = existing.unwrap_or(candidate);
Self::set_expiration(&mut state, key, now + self.get_ttl(&kwargs));
Self::set_expiration(
&mut state,
key,
now + self.get_ttl(&context).unwrap_or(self.default_ttl),
);
state.values.insert(key.into(), winner.clone());
Ok(winner)
}
}
impl CounterCache for InMemoryCache<f64> {
fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result<f64, Error> {
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
if self.max_size_in_memory == 0 {
return Ok(amount);
}
@ -255,7 +265,11 @@ impl CounterCache for InMemoryCache<f64> {
Self::evict(&mut state, self.max_size_in_memory, now, key);
let value = state.values.get(key).copied().unwrap_or_default() + amount;
if !state.expirations.contains_key(key) {
Self::set_expiration(&mut state, key, now + self.get_ttl(&kwargs));
Self::set_expiration(
&mut state,
key,
now + self.get_ttl(&context).unwrap_or(self.default_ttl),
);
}
state.values.insert(key.into(), value);
Ok(value)
@ -273,10 +287,7 @@ impl InMemoryCache<f64> {
self.increment_cache(
&operation.key,
operation.amount,
CacheKwargs {
ttl: operation.ttl,
..CacheKwargs::default()
},
ExactCacheContext { ttl: operation.ttl },
)
})
.collect()
@ -285,28 +296,26 @@ impl InMemoryCache<f64> {
impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
type Value = V;
type Context = ExactCacheContext;
fn default_ttl(&self) -> Duration {
self.default_ttl
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(self.default_ttl))
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
let ttl = self.get_ttl(&kwargs);
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &ExactCacheContext,
) -> Result<(), Error> {
let ttl = self.get_ttl(context).unwrap_or(self.default_ttl);
self.set_cache(key, value, Some(ttl)).map(|_| ())
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
self.get_cache(key)
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
self.delete_cache(key)
}
fn flush_cache(&self) -> Result<(), Error> {
self.flush_cache()
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
@ -320,6 +329,60 @@ impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
}
}
impl<V: Clone + Send + Sync + 'static> BatchCache for InMemoryCache<V> {}
impl<V: Clone + Send + Sync + 'static> DeleteCache for InMemoryCache<V> {
fn delete_cache(&self, key: &str) -> Result<(), Error> {
InMemoryCache::delete_cache(self, key)
}
}
impl<V: Clone + Send + Sync + 'static> FlushCache for InMemoryCache<V> {
fn flush_cache(&self) -> Result<(), Error> {
InMemoryCache::flush_cache(self)
}
}
impl<V: Clone + Send + Sync + 'static> TtlCache for InMemoryCache<V> {
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
InMemoryCache::async_get_ttl(self, key).await
}
}
impl<T> SetCache for InMemoryCache<HashSet<T>>
where
T: Clone + Eq + Hash + Send + Sync + 'static,
{
type SetValue = T;
type SetResult = Vec<T>;
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
if self.max_size_in_memory == 0 {
return Ok(values);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, key);
let mut stored = state.values.get(key).cloned().unwrap_or_default();
stored.extend(values.iter().cloned());
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
&& measure(&stored)? > limit
{
return Ok(values);
}
if !state.expirations.contains_key(key) {
Self::set_expiration(&mut state, key, now + ttl.unwrap_or(self.default_ttl));
}
state.values.insert(key.into(), stored);
Ok(values)
}
}
#[cfg(test)]
mod tests {
use super::*;
@ -329,7 +392,7 @@ mod tests {
let cache = InMemoryCache::<f64>::new(Some(4), None);
for _ in 0..100 {
cache
.increment_cache("counter", 1.0, CacheKwargs::default())
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
}
assert_eq!(cache.state.lock().unwrap().expiration_heap.len(), 1);

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

@ -1,4 +1,5 @@
use std::{
collections::HashSet,
sync::{
Arc,
atomic::{AtomicU64, Ordering},
@ -7,8 +8,8 @@ use std::{
};
use litellm_cache::{
BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, Error,
IncrementOperation, get_cache, set_cache,
BaseCache, CacheBackend, CacheConnectionStatus, ClaimCache, CounterCache, DeleteCache, Error,
ExactCacheContext, IncrementOperation, SetCache, get_cache, set_cache,
};
use litellm_cache_memory::{CacheWrite, InMemoryCache};
use rstest::{fixture, rstest};
@ -152,39 +153,38 @@ async fn generic_consumers_share_typed_values_and_honor_expiration() {
let clock = clock();
let cache: CacheBackend<InMemoryCache<String>> = Arc::new(cache(clock.clone(), 4));
let reader = Arc::clone(&cache);
let kwargs = CacheKwargs {
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(5)),
..Default::default()
};
set_cache(cache.as_ref(), "sync", "first".into(), kwargs.clone()).unwrap();
set_cache(cache.as_ref(), "sync", "first".into(), &context).unwrap();
assert_eq!(
get_cache(reader.as_ref(), "sync", &kwargs).unwrap(),
get_cache(reader.as_ref(), "sync", &context).unwrap(),
Some("first".into())
);
cache
.batch_cache_write("async", "second".into(), kwargs.clone())
.batch_cache_write("async", "second".into(), context.clone())
.await
.unwrap();
cache
.async_set_cache_pipeline(vec![("batch".into(), "third".into())], kwargs.clone())
.async_set_cache_pipeline(vec![("batch".into(), "third".into())], context.clone())
.await
.unwrap();
drop(cache);
for (key, value) in [("sync", "first"), ("async", "second"), ("batch", "third")] {
assert_eq!(
reader.async_get_cache(key, &kwargs).await.unwrap(),
reader.async_get_cache(key, &context).await.unwrap(),
Some(value.into())
);
}
reader.async_delete_cache("async").await.unwrap();
assert_eq!(
reader.async_get_cache("async", &kwargs).await.unwrap(),
reader.async_get_cache("async", &context).await.unwrap(),
None
);
clock.store(106, Ordering::SeqCst);
assert_eq!(get_cache(reader.as_ref(), "sync", &kwargs).unwrap(), None);
assert_eq!(get_cache(reader.as_ref(), "sync", &context).unwrap(), None);
assert_eq!(
reader.async_get_cache("batch", &kwargs).await.unwrap(),
reader.async_get_cache("batch", &context).await.unwrap(),
None
);
}
@ -196,20 +196,19 @@ fn claims_are_atomic_and_refresh_eligible_winners() {
let clock = clock.clone();
move || Duration::from_secs(clock.load(Ordering::SeqCst))
});
let kwargs = CacheKwargs {
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(10)),
..Default::default()
};
assert_eq!(
cache
.claim_cache("affinity", "first".to_string(), &[], kwargs.clone())
.claim_cache("affinity", "first".to_string(), &[], context.clone())
.unwrap(),
"first"
);
clock.store(103, Ordering::SeqCst);
assert_eq!(
cache
.claim_cache("affinity", "second".to_string(), &[], kwargs.clone())
.claim_cache("affinity", "second".to_string(), &[], context.clone())
.unwrap(),
"first"
);
@ -224,7 +223,7 @@ fn claims_are_atomic_and_refresh_eligible_winners() {
"affinity",
"second".to_string(),
&["first".to_string(), "second".to_string()],
kwargs,
context,
)
.unwrap(),
"first"
@ -239,11 +238,13 @@ fn claims_are_atomic_and_refresh_eligible_winners() {
fn counters_increment_under_one_lock() {
let cache = InMemoryCache::<f64>::default();
assert_eq!(
CounterCache::increment_cache(&cache, "counter", 1.5, CacheKwargs::default()).unwrap(),
CounterCache::increment_cache(&cache, "counter", 1.5, ExactCacheContext::default())
.unwrap(),
1.5
);
assert_eq!(
CounterCache::increment_cache(&cache, "counter", 2.0, CacheKwargs::default()).unwrap(),
CounterCache::increment_cache(&cache, "counter", 2.0, ExactCacheContext::default())
.unwrap(),
3.5
);
}
@ -263,7 +264,7 @@ fn rewriting_an_existing_key_at_capacity_keeps_other_entries(clock: Arc<AtomicU6
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
cache
.claim_cache("cold", "4".into(), &[], CacheKwargs::default())
.claim_cache("cold", "4".into(), &[], ExactCacheContext::default())
.unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
@ -278,7 +279,7 @@ 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())
.increment_cache(key, 1.0, ExactCacheContext::default())
.unwrap();
}
assert_eq!(cache.get_cache("a").unwrap(), Some(2.0));
@ -290,7 +291,7 @@ 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())
.claim_cache("key", "first".into(), &[], ExactCacheContext::default())
.unwrap(),
"first"
);
@ -299,7 +300,7 @@ fn disabled_cache_does_not_retain_claims_or_counters() {
let counters = InMemoryCache::<f64>::new(Some(0), None);
assert_eq!(
counters
.increment_cache("key", 2.0, CacheKwargs::default())
.increment_cache("key", 2.0, ExactCacheContext::default())
.unwrap(),
2.0
);
@ -348,3 +349,20 @@ async fn increment_pipeline_preserves_operation_order() {
);
assert_eq!(cache.get_cache("a").unwrap(), Some(3.0));
}
#[tokio::test]
async fn set_capability_preserves_python_result_and_deduplicates_storage() {
let cache = InMemoryCache::<HashSet<String>>::new(None, None);
let inserted = vec!["a".into(), "a".into(), "b".into()];
assert_eq!(
cache
.async_set_cache_sadd("members", inserted.clone(), None)
.await
.unwrap(),
inserted
);
assert_eq!(
cache.get_cache("members").unwrap(),
Some(HashSet::from(["a".into(), "b".into()]))
);
}

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

@ -4,14 +4,16 @@ use std::{
};
use litellm_cache::{
BaseCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs,
ClaimCache, CounterCache, Error,
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
ClaimCache, CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
};
use redis::Commands;
mod operations;
pub use operations::{RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation};
pub use operations::{
RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
};
const DEFAULT_TTL: Duration = Duration::from_secs(600);
const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
@ -179,12 +181,7 @@ where
}
fn namespaced_key(&self, key: &str) -> String {
match &self.namespace {
Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
format!("{namespace}:{key}")
}
_ => key.into(),
}
namespaced_key(self.namespace.as_deref(), key)
}
fn namespaced_pattern(&self) -> Result<String, Error> {
@ -260,20 +257,35 @@ where
}
}
fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
match namespace {
Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
format!("{namespace}:{key}")
}
_ => key.into(),
}
}
impl<S, C> BaseCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = S::Value;
type Context = ExactCacheContext;
fn default_ttl(&self) -> Duration {
self.default_ttl
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(self.default_ttl))
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> {
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &ExactCacheContext,
) -> Result<(), Error> {
let payload = self.codec.encode(&value)?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
let ttl = Self::ttl_seconds(self.get_ttl(context).unwrap_or(self.default_ttl));
let key = self.namespaced_key(key);
self.connections.execute(|connection| {
connection
@ -282,7 +294,7 @@ where
})
}
fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = self.connections.execute(|connection| {
connection
@ -292,48 +304,15 @@ where
self.decode_response(value)
}
fn get_cache_batch(
&self,
keys: &[String],
_: &CacheKwargs,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = self.connections.execute(|connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})?;
values
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
fn delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
self.connections
.execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable))
}
fn flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
self.connections
.execute(|connection| Self::flush_matching(connection, &pattern))
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
kwargs: CacheKwargs,
context: ExactCacheContext,
) -> Result<(), Error> {
let payload = self.codec.encode(&value)?;
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
@ -345,7 +324,7 @@ where
async fn async_get_cache(
&self,
key: &str,
_: &CacheKwargs,
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
@ -357,32 +336,10 @@ where
self.decode_response(value)
}
async fn async_get_cache_batch(
&self,
keys: Vec<String>,
_: CacheKwargs,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
values
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
async fn async_set_cache_pipeline(
&self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
context: ExactCacheContext,
) -> Result<(), Error> {
let entries = cache_list
.into_iter()
@ -392,7 +349,7 @@ where
.map(|payload| (self.namespaced_key(&key), payload))
})
.collect::<Result<Vec<_>, _>>()?;
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut pipeline = redis::pipe();
for (key, payload) in entries {
@ -410,22 +367,6 @@ where
.await
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
.await
}
async fn async_flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Self::flush_matching(connection, &pattern)
})
.await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
@ -457,26 +398,120 @@ where
}
}
impl<S, C> BatchCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn batch_get_cache(
&self,
keys: &[String],
_: &ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = self.connections.execute(|connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})?;
values
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
_: ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
values
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
}
impl<S, C> DeleteCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
self.connections
.execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable))
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
.await
}
}
impl<S, C> FlushCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
self.connections
.execute(|connection| Self::flush_matching(connection, &pattern))
}
async fn async_flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Self::flush_matching(connection, &pattern)
})
.await
}
}
impl<S, C> CounterCache for RedisCache<S, C>
where
S: CacheCodec<Value = f64>,
C: redis::ConnectionLike + Send + 'static,
{
fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result<f64, Error> {
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
self.connections
.execute(|connection| increment(connection, key, amount, ttl))
}
async fn async_increment_cache(
async fn async_increment(
&self,
key: &str,
amount: f64,
kwargs: CacheKwargs,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
increment(connection, key, amount, ttl)
})
@ -570,10 +605,10 @@ where
key: &str,
candidate: S::Value,
eligible: &[S::Value],
kwargs: CacheKwargs,
context: ExactCacheContext,
) -> Result<S::Value, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
self.connections
.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl))
}
@ -583,10 +618,10 @@ where
key: &str,
candidate: S::Value,
eligible: Vec<S::Value>,
kwargs: CacheKwargs,
context: ExactCacheContext,
) -> Result<S::Value, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&kwargs));
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
let codec = self.codec.clone();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
claim(connection, &codec, &key, candidate, &eligible, ttl)
@ -599,7 +634,9 @@ where
mod tests {
use std::time::Duration;
use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, JsonCodec};
use litellm_cache::{
BaseCache, CacheCodec, DeleteCache, ExactCacheContext, FlushCache, JsonCodec,
};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
@ -648,10 +685,12 @@ mod tests {
.with_namespace(Some("litellm-cache".into()));
cache
.set_cache("key", value.clone(), CacheKwargs::default())
.set_cache("key", value.clone(), &ExactCacheContext::default())
.unwrap();
assert_eq!(
cache.get_cache("key", &CacheKwargs::default()).unwrap(),
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(value)
);
cache.delete_cache("key").unwrap();

137
litellm-rust/crates/cache-redis/src/cache/operations.rs vendored
View file

@ -1,9 +1,12 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, Error, IncrementOperation};
use litellm_cache::{
CacheCodec, CacheScript, ClientInfoCache, Error, IncrementOperation, QueueCache, ScanCache,
ScriptCache, SetCache, TtlCache,
};
use redis::Commands;
use super::{ConnectionRef, RedisCache};
use super::{ConnectionRef, Connections, RedisCache, namespaced_key};
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
@ -88,6 +91,46 @@ pub enum RedisLpopResult {
Values(Vec<Vec<u8>>),
}
pub struct RedisScript<C> {
connections: Arc<Connections<C>>,
namespace: Option<String>,
source: String,
}
impl<C> CacheScript for RedisScript<C>
where
C: redis::ConnectionLike + Send + 'static,
{
type Argument = RedisArg;
type Output = redis::Value;
async fn invoke(
&self,
keys: Vec<String>,
arguments: Vec<Self::Argument>,
) -> Result<Self::Output, Error> {
let keys = keys
.into_iter()
.map(|key| namespaced_key(self.namespace.as_deref(), &key))
.collect::<Vec<_>>();
let connections = Arc::clone(&self.connections);
let source = self.source.clone();
tokio::task::spawn_blocking(move || {
connections.execute(|connection| {
redis::cmd("EVAL")
.arg(source)
.arg(keys.len())
.arg(keys)
.arg(arguments)
.query(connection)
.map_err(|_| Error::Unavailable)
})
})
.await
.map_err(|_| Error::Unavailable)?
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
@ -498,3 +541,93 @@ fn increment_with_floor(
.query(connection)
.map_err(|_| Error::Unavailable)
}
impl<S, C> TtlCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
RedisCache::async_get_ttl(self, key)
.await
.map(|ttl| ttl.map(|seconds| Duration::from_secs(seconds as u64)))
}
}
impl<S, C> ScanCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
RedisCache::async_scan_iter(self, pattern, count).await
}
}
impl<S, C> ClientInfoCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type ClientList = String;
type Info = String;
fn client_list(&self) -> Result<Self::ClientList, Error> {
RedisCache::client_list(self)
}
fn info(&self) -> Result<Self::Info, Error> {
RedisCache::info(self)
}
}
impl<S, C> SetCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type SetValue = RedisArg;
type SetResult = usize;
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
RedisCache::async_set_cache_sadd(self, key, values, ttl).await
}
}
impl<S, C> QueueCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type QueueValue = RedisArg;
type PopResult = RedisLpopResult;
async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
RedisCache::async_rpush(self, key, values).await
}
async fn async_lpop(&self, key: &str, count: Option<usize>) -> Result<Self::PopResult, Error> {
RedisCache::async_lpop(self, key, count).await
}
}
impl<S, C> ScriptCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Script = RedisScript<C>;
fn async_register_script(&self, source: String) -> Self::Script {
RedisScript {
connections: Arc::clone(&self.connections),
namespace: self.namespace.clone(),
source,
}
}
}

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

@ -1,3 +1,7 @@
mod cache;
mod topology;
pub use cache::{RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation};
pub use cache::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
};
pub use topology::{RedisNode, RedisTopology};

14
litellm-rust/crates/cache-redis/src/topology.rs Normal file
View file

@ -0,0 +1,14 @@
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RedisNode {
pub host: String,
pub port: u16,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum RedisTopology {
#[default]
Standalone,
Cluster {
startup_nodes: Vec<RedisNode>,
},
}

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

@ -1,8 +1,9 @@
use std::time::Duration;
use litellm_cache::{
BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, ClaimCache,
CounterCache, Error, IncrementOperation, JsonCodec, get_cache, set_cache,
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheScript, ClaimCache,
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
ScriptCache, get_cache, set_cache,
};
use litellm_cache_redis::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation,
@ -48,12 +49,11 @@ fn generic_helpers_use_the_injected_codec_and_ttl() {
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
let kwargs = CacheKwargs {
let context = ExactCacheContext {
ttl: Some(Duration::from_millis(1500)),
..Default::default()
};
set_cache(&cache, "counter", 7, kwargs.clone()).unwrap();
assert_eq!(get_cache(&cache, "counter", &kwargs).unwrap(), Some(7));
set_cache(&cache, "counter", 7, &context).unwrap();
assert_eq!(get_cache(&cache, "counter", &context).unwrap(), Some(7));
}
#[tokio::test]
@ -83,28 +83,27 @@ async fn async_operations_preserve_codec_ttl_and_missing_values() {
Some(Duration::from_secs(9)),
TaggedByteCodec(42),
);
let kwargs = CacheKwargs::default();
let context = ExactCacheContext::default();
cache
.batch_cache_write("counter", 7, kwargs.clone())
.batch_cache_write("counter", 7, context.clone())
.await
.unwrap();
assert_eq!(
cache.async_get_cache("counter", &kwargs).await.unwrap(),
cache.async_get_cache("counter", &context).await.unwrap(),
Some(7)
);
cache
.async_set_cache_pipeline(
vec![("batch".into(), 8)],
CacheKwargs {
ExactCacheContext {
ttl: Some(Duration::from_millis(1500)),
..Default::default()
},
)
.await
.unwrap();
cache.async_delete_cache("counter").await.unwrap();
assert_eq!(
cache.async_get_cache("counter", &kwargs).await.unwrap(),
cache.async_get_cache("counter", &context).await.unwrap(),
None
);
}
@ -117,30 +116,30 @@ async fn codec_errors_propagate_without_writing_partial_batches() {
])
.assert_all_commands_consumed();
let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42));
let kwargs = CacheKwargs::default();
let context = ExactCacheContext::default();
assert_eq!(
cache.set_cache("invalid", 255, kwargs.clone()),
cache.set_cache("invalid", 255, &context),
Err(Error::InvalidEntry)
);
assert_eq!(
cache.async_set_cache("invalid", 255, kwargs.clone()).await,
cache.async_set_cache("invalid", 255, context.clone()).await,
Err(Error::InvalidEntry)
);
assert_eq!(
cache
.async_set_cache_pipeline(
vec![("valid".into(), 7), ("invalid".into(), 255)],
kwargs.clone(),
context.clone(),
)
.await,
Err(Error::InvalidEntry)
);
assert_eq!(
cache.get_cache("invalid", &kwargs),
cache.get_cache("invalid", &context),
Err(Error::InvalidEntry)
);
assert_eq!(
cache.async_get_cache("invalid", &kwargs).await,
cache.async_get_cache("invalid", &context).await,
Err(Error::InvalidEntry)
);
}
@ -155,12 +154,14 @@ fn namespaces_are_optional_and_existing_prefixes_are_not_duplicated() {
let cache = RedisCache::with_connection(connection, None, JsonCodec::<String>::new())
.with_namespace(Some("team".into()));
assert_eq!(
cache.get_cache("key", &CacheKwargs::default()).unwrap(),
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
None
);
assert_eq!(
cache
.get_cache("team:key", &CacheKwargs::default())
.get_cache("team:key", &ExactCacheContext::default())
.unwrap(),
None
);
@ -221,9 +222,9 @@ async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries() {
assert_eq!(
cache
.async_get_cache_batch(
.async_batch_get_cache(
vec!["hit".into(), "miss".into(), "invalid".into()],
CacheKwargs::default(),
ExactCacheContext::default(),
)
.await
.unwrap(),
@ -327,6 +328,13 @@ async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() {
.arg("team:key"),
Ok("team:key"),
),
MockCmd::new(
redis::cmd("EVAL")
.arg("return KEYS[1]")
.arg(1usize)
.arg("team:key"),
Ok("team:key"),
),
MockCmd::new(redis::cmd("CLIENT").arg("LIST"), Ok("id=1")),
MockCmd::new(redis::cmd("INFO"), Ok("redis_version:7")),
MockCmd::new(redis::cmd("FLUSHALL"), Ok("OK")),
@ -387,6 +395,14 @@ async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() {
.unwrap(),
redis::Value::BulkString(b"team:key".to_vec())
);
assert_eq!(
cache
.async_register_script("return KEYS[1]".into())
.invoke(vec!["key".into()], Vec::new())
.await
.unwrap(),
redis::Value::BulkString(b"team:key".to_vec())
);
assert_eq!(cache.client_list().unwrap(), "id=1");
assert_eq!(cache.info().unwrap(), "redis_version:7");
cache.flushall().unwrap();
@ -602,7 +618,7 @@ async fn claims_match_eligible_values_written_by_another_encoder() {
"pin",
candidate,
vec![stored.clone()],
CacheKwargs::default()
ExactCacheContext::default()
)
.await
.unwrap(),
@ -630,7 +646,7 @@ fn claims_retry_when_the_key_changes_and_replace_ineligible_winners() {
"pin",
candidate.clone(),
&[serde_json::json!({"model_id": "a"})],
CacheKwargs::default()
ExactCacheContext::default()
)
.unwrap(),
candidate
@ -654,7 +670,7 @@ fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl() {
"pin",
serde_json::json!({"model_id": "b"}),
&[],
CacheKwargs::default()
ExactCacheContext::default()
)
.unwrap(),
serde_json::json!({"model_id": "a"})
@ -679,7 +695,7 @@ async fn async_increment_runs_the_atomic_script() {
assert_eq!(
cache
.async_increment_cache("counter", 2.5, CacheKwargs::default())
.async_increment("counter", 2.5, ExactCacheContext::default())
.await
.unwrap(),
4.5

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

@ -1,6 +1,6 @@
use std::{sync::Mutex, time::Duration};
use litellm_cache::{BaseCache, Error};
use litellm_cache::{BaseCache, Error, ExactCacheContext};
use serde_json::Value;
use crate::{CacheEntry, ResponseCache, ResponseCacheRequest};
@ -18,7 +18,7 @@ impl WriteBuffer {
}
}
pub async fn async_store<B: BaseCache<Value = CacheEntry>>(
pub async fn async_store<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>>(
&self,
cache: &ResponseCache<B>,
request: &ResponseCacheRequest,

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

@ -1,6 +1,8 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, Error};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
};
use serde_json::Value;
use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
@ -9,7 +11,7 @@ use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
pub struct ResponseCacheRequest {
pub key: CacheKeyInput,
pub controls: CacheControls,
pub kwargs: CacheKwargs,
pub context: ExactCacheContext,
pub max_age: Option<Duration>,
}
@ -24,17 +26,17 @@ impl ResponseCacheRequest {
default_on: true,
..Default::default()
},
kwargs: CacheKwargs::default(),
context: ExactCacheContext::default(),
max_age: None,
}
}
}
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>> {
pub struct ResponseCache<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
backend: Arc<B>,
}
impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCache<B> {
pub fn new(backend: Arc<B>) -> Self {
Self { backend }
}
@ -43,11 +45,14 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
&self.backend
}
pub fn default_ttl(&self) -> Duration {
self.backend.default_ttl()
pub fn default_ttl(&self) -> Option<Duration> {
self.backend.get_ttl(&ExactCacheContext::default())
}
pub async fn async_flush(&self) -> Result<(), Error> {
pub async fn async_flush(&self) -> Result<(), Error>
where
B: FlushCache,
{
self.backend.async_flush_cache().await
}
@ -65,7 +70,7 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
}
let entry = match self
.backend
.get_cache(&cache_key(&request.key), &request.kwargs)
.get_cache(&cache_key(&request.key), &request.context)
{
Ok(entry) => entry,
Err(Error::InvalidEntry) => None,
@ -84,7 +89,7 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
}
let entry = match self
.backend
.async_get_cache(&cache_key(&request.key), &request.kwargs)
.async_get_cache(&cache_key(&request.key), &request.context)
.await
{
Ok(entry) => entry,
@ -98,7 +103,10 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
&self,
requests: &[ResponseCacheRequest],
now: Duration,
) -> Result<PartialHits, Error> {
) -> Result<PartialHits, Error>
where
B: BatchCache,
{
let readable = requests
.iter()
.enumerate()
@ -109,7 +117,7 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
.map(|(_, request)| cache_key(&request.key))
.collect::<Vec<_>>();
let entries = if let Some((_, request)) = readable.first() {
self.backend.get_cache_batch(&keys, &request.kwargs)?
self.backend.batch_get_cache(&keys, &request.context)?
} else {
Vec::new()
};
@ -120,7 +128,10 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
&self,
requests: &[ResponseCacheRequest],
now: Duration,
) -> Result<PartialHits, Error> {
) -> Result<PartialHits, Error>
where
B: BatchCache,
{
let readable = requests
.iter()
.enumerate()
@ -132,7 +143,7 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
.collect::<Vec<_>>();
let entries = if let Some((_, request)) = readable.first() {
self.backend
.async_get_cache_batch(keys, request.kwargs.clone())
.async_batch_get_cache(keys, request.context.clone())
.await?
} else {
Vec::new()
@ -155,7 +166,7 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
timestamp: Some(now.as_secs_f64()),
response,
},
request.kwargs.clone(),
&request.context,
)
}
@ -175,7 +186,7 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
timestamp: Some(now.as_secs_f64()),
response,
},
request.kwargs.clone(),
request.context.clone(),
)
.await
}
@ -210,26 +221,29 @@ impl<B: BaseCache<Value = CacheEntry>> ResponseCache<B> {
timestamp: Some(now.as_secs_f64()),
response,
},
request.kwargs,
request.context,
)
})
.collect::<Vec<_>>();
let Some((_, _, first_kwargs)) = writable.first() else {
return Ok(());
};
if writable.iter().all(|(_, _, kwargs)| kwargs == first_kwargs) {
let kwargs = first_kwargs.clone();
if writable
.iter()
.all(|(_, _, context)| context == first_kwargs)
{
let context = first_kwargs.clone();
let cache_list = writable
.into_iter()
.map(|(key, entry, _)| (key, entry))
.collect();
return self
.backend
.async_set_cache_pipeline(cache_list, kwargs)
.async_set_cache_pipeline(cache_list, context)
.await;
}
for (key, entry, kwargs) in writable {
self.backend.async_set_cache(&key, entry, kwargs).await?;
for (key, entry, context) in writable {
self.backend.async_set_cache(&key, entry, context).await?;
}
Ok(())
}

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

@ -43,7 +43,7 @@ async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
));
let cache = ResponseCache::new(backend.clone());
let mut request = request();
request.kwargs.ttl = Some(Duration::from_secs(10));
request.context.ttl = Some(Duration::from_secs(10));
request.max_age = Some(Duration::from_secs(5));
cache
.store(
@ -336,7 +336,7 @@ fn response_codec_preserves_values_without_timestamps() {
assert_eq!(entry.response, raw);
let backend = Arc::new(InMemoryCache::default());
BaseCache::set_cache(backend.as_ref(), "tenant:key", entry, Default::default()).unwrap();
BaseCache::set_cache(backend.as_ref(), "tenant:key", entry, &Default::default()).unwrap();
let cache = ResponseCache::new(backend);
assert_eq!(
cache.lookup(&request(), Duration::from_secs(100)).unwrap(),

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

@ -1,7 +1,6 @@
use std::{future::Future, time::Duration};
use serde::{Deserialize, Serialize};
use serde_json::{Map, Value};
use crate::Error;
@ -12,10 +11,25 @@ pub enum BatchEntry<V> {
Invalid,
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct CacheKwargs {
pub trait CacheContext: Clone + Send + Sync + 'static {
fn ttl(&self) -> Option<Duration>;
fn with_ttl(&self, ttl: Option<Duration>) -> Self;
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct ExactCacheContext {
pub ttl: Option<Duration>,
pub extras: Map<String, Value>,
}
impl CacheContext for ExactCacheContext {
fn ttl(&self) -> Option<Duration> {
self.ttl
}
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
Self { ttl }
}
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)]
@ -35,78 +49,44 @@ pub struct CacheConnectionResult {
pub trait BaseCache: Send + Sync {
type Value: Clone + Send + Sync + 'static;
type Context: CacheContext;
fn default_ttl(&self) -> Duration {
Duration::from_secs(60)
}
fn get_ttl(&self, context: &Self::Context) -> Option<Duration>;
fn get_ttl(&self, kwargs: &CacheKwargs) -> Duration {
kwargs.ttl.unwrap_or_else(|| self.default_ttl())
}
fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error>;
fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result<Option<Self::Value>, Error>;
fn get_cache_batch(
fn set_cache(
&self,
keys: &[String],
kwargs: &CacheKwargs,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
keys.iter()
.map(|key| match self.get_cache(key, kwargs) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
})
.collect()
}
key: &str,
value: Self::Value,
context: &Self::Context,
) -> Result<(), Error>;
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error>;
fn async_set_cache(
&self,
key: &str,
value: Self::Value,
kwargs: CacheKwargs,
context: Self::Context,
) -> impl Future<Output = Result<(), Error>> + Send {
async move { self.set_cache(key, value, kwargs) }
async move { self.set_cache(key, value, &context) }
}
fn async_get_cache(
&self,
key: &str,
kwargs: &CacheKwargs,
context: &Self::Context,
) -> impl Future<Output = Result<Option<Self::Value>, Error>> + Send {
async move { self.get_cache(key, kwargs) }
}
fn async_get_cache_batch(
&self,
keys: Vec<String>,
kwargs: CacheKwargs,
) -> impl Future<Output = Result<Vec<BatchEntry<Self::Value>>, Error>> + Send {
async move {
let mut entries = Vec::with_capacity(keys.len());
for key in keys {
entries.push(match self.async_get_cache(&key, &kwargs).await {
Ok(Some(value)) => BatchEntry::Hit(value),
Ok(None) => BatchEntry::Miss,
Err(Error::InvalidEntry) => BatchEntry::Invalid,
Err(error) => return Err(error),
});
}
Ok(entries)
}
async move { self.get_cache(key, context) }
}
fn async_set_cache_pipeline(
&self,
cache_list: Vec<(String, Self::Value)>,
kwargs: CacheKwargs,
entries: Vec<(String, Self::Value)>,
context: Self::Context,
) -> impl Future<Output = Result<(), Error>> + Send {
async move {
for (key, value) in cache_list {
self.async_set_cache(&key, value, kwargs.clone()).await?;
for (key, value) in entries {
self.async_set_cache(&key, value, context.clone()).await?;
}
Ok(())
}
@ -116,21 +96,9 @@ pub trait BaseCache: Send + Sync {
&self,
key: &str,
value: Self::Value,
kwargs: CacheKwargs,
context: Self::Context,
) -> impl Future<Output = Result<(), Error>> + Send {
self.async_set_cache(key, value, kwargs)
}
fn delete_cache(&self, key: &str) -> Result<(), Error>;
fn async_delete_cache(&self, key: &str) -> impl Future<Output = Result<(), Error>> + Send {
async move { self.delete_cache(key) }
}
fn flush_cache(&self) -> Result<(), Error>;
fn async_flush_cache(&self) -> impl Future<Output = Result<(), Error>> + Send {
async move { self.flush_cache() }
self.async_set_cache(key, value, context)
}
fn disconnect(&self) -> impl Future<Output = Result<(), Error>> + Send;

85
litellm-rust/crates/cache/src/cache_type.rs vendored Normal file
View file

@ -0,0 +1,85 @@
use serde::{Deserialize, Serialize};
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq, Hash)]
pub enum CacheType {
#[serde(rename = "local")]
Local,
#[serde(rename = "redis")]
Redis,
#[serde(rename = "redis-semantic")]
RedisSemantic,
#[serde(rename = "valkey-semantic")]
ValkeySemantic,
#[serde(rename = "s3")]
S3,
#[serde(rename = "disk")]
Disk,
#[serde(rename = "qdrant-semantic")]
QdrantSemantic,
#[serde(rename = "azure-blob")]
AzureBlob,
#[serde(rename = "gcs")]
Gcs,
}
impl CacheType {
pub const ALL: [Self; 9] = [
Self::Local,
Self::Redis,
Self::RedisSemantic,
Self::ValkeySemantic,
Self::S3,
Self::Disk,
Self::QdrantSemantic,
Self::AzureBlob,
Self::Gcs,
];
pub const fn as_python_name(self) -> &'static str {
match self {
Self::Local => "local",
Self::Redis => "redis",
Self::RedisSemantic => "redis-semantic",
Self::ValkeySemantic => "valkey-semantic",
Self::S3 => "s3",
Self::Disk => "disk",
Self::QdrantSemantic => "qdrant-semantic",
Self::AzureBlob => "azure-blob",
Self::Gcs => "gcs",
}
}
pub fn from_python_name(value: &str) -> Option<Self> {
Self::ALL
.into_iter()
.find(|cache_type| cache_type.as_python_name() == value)
}
}
#[cfg(test)]
mod tests {
use super::CacheType;
#[test]
fn every_python_cache_type_has_one_round_trip_identity() {
let names = CacheType::ALL.map(CacheType::as_python_name);
assert_eq!(
names,
[
"local",
"redis",
"redis-semantic",
"valkey-semantic",
"s3",
"disk",
"qdrant-semantic",
"azure-blob",
"gcs",
]
);
assert_eq!(
names.map(CacheType::from_python_name),
CacheType::ALL.map(Some)
);
}
}

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

@ -1,23 +1,23 @@
use std::sync::Arc;
pub use crate::BaseCache as Cache;
use crate::{BaseCache, CacheKwargs, Error};
use crate::{BaseCache, Error};
pub fn get_cache<B: BaseCache>(
cache: &B,
key: &str,
kwargs: &CacheKwargs,
context: &B::Context,
) -> Result<Option<B::Value>, Error> {
cache.get_cache(key, kwargs)
cache.get_cache(key, context)
}
pub fn set_cache<B: BaseCache>(
cache: &B,
key: &str,
value: B::Value,
kwargs: CacheKwargs,
context: &B::Context,
) -> Result<(), Error> {
cache.set_cache(key, value, kwargs)
cache.set_cache(key, value, context)
}
pub type CacheBackend<B> = Arc<B>;

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

@ -1,24 +1,77 @@
use std::future::Future;
use std::{future::Future, time::Duration};
use crate::{BaseCache, CacheKwargs, Error};
use crate::{BaseCache, BatchEntry, Error};
#[derive(Clone, Debug, PartialEq)]
pub struct IncrementOperation {
pub key: String,
pub amount: f64,
pub ttl: Option<std::time::Duration>,
pub ttl: Option<Duration>,
}
pub trait BatchCache: BaseCache {
fn batch_get_cache(
&self,
keys: &[String],
context: &Self::Context,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
keys.iter()
.map(|key| match self.get_cache(key, context) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
})
.collect()
}
fn async_batch_get_cache(
&self,
keys: Vec<String>,
context: Self::Context,
) -> impl Future<Output = Result<Vec<BatchEntry<Self::Value>>, Error>> + Send {
async move {
let mut entries = Vec::with_capacity(keys.len());
for key in keys {
entries.push(match self.async_get_cache(&key, &context).await {
Ok(Some(value)) => BatchEntry::Hit(value),
Ok(None) => BatchEntry::Miss,
Err(Error::InvalidEntry) => BatchEntry::Invalid,
Err(error) => return Err(error),
});
}
Ok(entries)
}
}
}
pub trait DeleteCache: BaseCache {
fn delete_cache(&self, key: &str) -> Result<(), Error>;
fn async_delete_cache(&self, key: &str) -> impl Future<Output = Result<(), Error>> + Send {
async move { self.delete_cache(key) }
}
}
pub trait FlushCache: BaseCache {
fn flush_cache(&self) -> Result<(), Error>;
fn async_flush_cache(&self) -> impl Future<Output = Result<(), Error>> + Send {
async move { self.flush_cache() }
}
}
pub trait CounterCache: BaseCache<Value = f64> {
fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result<f64, Error>;
fn increment_cache(&self, key: &str, amount: f64, context: Self::Context)
-> Result<f64, Error>;
fn async_increment_cache(
fn async_increment(
&self,
key: &str,
amount: f64,
kwargs: CacheKwargs,
context: Self::Context,
) -> impl Future<Output = Result<f64, Error>> + Send {
async move { self.increment_cache(key, amount, kwargs) }
async move { self.increment_cache(key, amount, context) }
}
}
@ -31,7 +84,7 @@ where
key: &str,
candidate: Self::Value,
eligible: &[Self::Value],
kwargs: CacheKwargs,
context: Self::Context,
) -> Result<Self::Value, Error>;
fn async_claim_cache(
@ -39,8 +92,78 @@ where
key: &str,
candidate: Self::Value,
eligible: Vec<Self::Value>,
kwargs: CacheKwargs,
context: Self::Context,
) -> impl Future<Output = Result<Self::Value, Error>> + Send {
async move { self.claim_cache(key, candidate, &eligible, kwargs) }
async move { self.claim_cache(key, candidate, &eligible, context) }
}
}
pub trait TtlCache: BaseCache {
fn async_get_ttl(
&self,
key: &str,
) -> impl Future<Output = Result<Option<Duration>, Error>> + Send;
}
pub trait SetCache: BaseCache {
type SetValue: Clone + Send + Sync + 'static;
type SetResult: Send + Sync + 'static;
fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> impl Future<Output = Result<Self::SetResult, Error>> + Send;
}
pub trait QueueCache: BaseCache {
type QueueValue: Clone + Send + Sync + 'static;
type PopResult: Send + Sync + 'static;
fn async_rpush(
&self,
key: &str,
values: Vec<Self::QueueValue>,
) -> impl Future<Output = Result<usize, Error>> + Send;
fn async_lpop(
&self,
key: &str,
count: Option<usize>,
) -> impl Future<Output = Result<Self::PopResult, Error>> + Send;
}
pub trait ScanCache: BaseCache {
fn async_scan_iter(
&self,
pattern: &str,
count: usize,
) -> impl Future<Output = Result<Vec<String>, Error>> + Send;
}
pub trait ClientInfoCache: BaseCache {
type ClientList: Send + Sync + 'static;
type Info: Send + Sync + 'static;
fn client_list(&self) -> Result<Self::ClientList, Error>;
fn info(&self) -> Result<Self::Info, Error>;
}
pub trait CacheScript: Send + Sync + 'static {
type Argument: Clone + Send + Sync + 'static;
type Output: Send + Sync + 'static;
fn invoke(
&self,
keys: Vec<String>,
arguments: Vec<Self::Argument>,
) -> impl Future<Output = Result<Self::Output, Error>> + Send;
}
pub trait ScriptCache: BaseCache {
type Script: CacheScript;
fn async_register_script(&self, source: String) -> Self::Script;
}

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

@ -1,7 +1,8 @@
use std::{sync::Arc, time::Duration};
use crate::{
BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error,
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, ClaimCache,
CounterCache, DeleteCache, Error, FlushCache,
};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
@ -94,19 +95,17 @@ impl<L1, L2> DualCache<L1, L2> {
}
}
fn promotion_kwargs(&self, kwargs: &CacheKwargs) -> CacheKwargs {
CacheKwargs {
ttl: self.promotion_ttl.or(kwargs.ttl),
extras: kwargs.extras.clone(),
}
fn promotion_context<C: CacheContext>(&self, context: &C) -> C {
context.with_ttl(self.promotion_ttl.or(context.ttl()))
}
}
impl<V, L1, L2> DualCache<L1, L2>
impl<V, C, L1, L2> DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
L1: BaseCache<Value = V>,
L2: BaseCache<Value = V>,
C: CacheContext,
L1: BaseCache<Value = V, Context = C>,
L2: BaseCache<Value = V, Context = C>,
{
fn missing(entries: &[BatchEntry<V>]) -> Vec<usize> {
entries
@ -119,7 +118,7 @@ where
fn merge_batch(
&self,
keys: &[String],
kwargs: &CacheKwargs,
context: &C,
mut entries: Vec<BatchEntry<V>>,
missing: Vec<usize>,
remote: Vec<BatchEntry<V>>,
@ -129,8 +128,9 @@ where
}
for (index, entry) in missing.into_iter().zip(remote) {
if let BatchEntry::Hit(value) = &entry {
let promotion_context = self.promotion_context(context);
self.l1
.set_cache(&keys[index], value.clone(), self.promotion_kwargs(kwargs))?;
.set_cache(&keys[index], value.clone(), &promotion_context)?;
}
entries[index] = entry;
}
@ -138,46 +138,105 @@ where
}
}
impl<V, L1, L2> BaseCache for DualCache<L1, L2>
impl<V, C, L1, L2> BaseCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
L1: BaseCache<Value = V>,
L2: BaseCache<Value = V>,
C: CacheContext,
L1: BaseCache<Value = V, Context = C>,
L2: BaseCache<Value = V, Context = C>,
{
type Value = V;
type Context = C;
fn default_ttl(&self) -> Duration {
self.l2.default_ttl()
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
self.l2.get_ttl(context)
}
fn set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> {
fn set_cache(&self, key: &str, value: V, context: &C) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.set_cache(key, value.clone(), kwargs.clone()))?;
self.remote(self.l2.set_cache(key, value.clone(), context))?;
}
self.l1.set_cache(key, value, kwargs)
self.l1.set_cache(key, value, context)
}
fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result<Option<V>, Error> {
if let Some(value) = self.l1.get_cache(key, kwargs)? {
fn get_cache(&self, key: &str, context: &C) -> Result<Option<V>, Error> {
if let Some(value) = self.l1.get_cache(key, context)? {
return Ok(Some(value));
}
if !self.reads_remote() {
return Ok(None);
}
let value = self.remote(self.l2.get_cache(key, kwargs))?.flatten();
let value = self.remote(self.l2.get_cache(key, context))?.flatten();
if let Some(value) = &value {
self.l1
.set_cache(key, value.clone(), self.promotion_kwargs(kwargs))?;
let promotion_context = self.promotion_context(context);
self.l1.set_cache(key, value.clone(), &promotion_context)?;
}
Ok(value)
}
fn get_cache_batch(
async fn async_set_cache(&self, key: &str, value: V, context: C) -> Result<(), Error> {
if self.writes_remote() {
self.remote(
self.l2
.async_set_cache(key, value.clone(), context.clone())
.await,
)?;
}
self.l1.async_set_cache(key, value, context).await
}
async fn async_get_cache(&self, key: &str, context: &C) -> Result<Option<V>, Error> {
if let Some(value) = self.l1.async_get_cache(key, context).await? {
return Ok(Some(value));
}
if !self.reads_remote() {
return Ok(None);
}
let value = self
.remote(self.l2.async_get_cache(key, context).await)?
.flatten();
if let Some(value) = &value {
self.l1
.async_set_cache(key, value.clone(), self.promotion_context(context))
.await?;
}
Ok(value)
}
async fn async_set_cache_pipeline(
&self,
keys: &[String],
kwargs: &CacheKwargs,
) -> Result<Vec<BatchEntry<V>>, Error> {
let entries = self.l1.get_cache_batch(keys, kwargs)?;
entries: Vec<(String, V)>,
context: C,
) -> Result<(), Error> {
if self.writes_remote() {
self.remote(
self.l2
.async_set_cache_pipeline(entries.clone(), context.clone())
.await,
)?;
}
self.l1.async_set_cache_pipeline(entries, context).await
}
async fn disconnect(&self) -> Result<(), Error> {
self.l2.disconnect().await?;
self.l1.disconnect().await
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
self.l2.test_connection().await
}
}
impl<V, C, L1, L2> BatchCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: BatchCache<Value = V, Context = C>,
L2: BatchCache<Value = V, Context = C>,
{
fn batch_get_cache(&self, keys: &[String], context: &C) -> Result<Vec<BatchEntry<V>>, Error> {
let entries = self.l1.batch_get_cache(keys, context)?;
let missing = Self::missing(&entries);
if missing.is_empty() || !self.reads_remote() {
return Ok(entries);
@ -186,49 +245,20 @@ where
.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),
match self.remote(self.l2.batch_get_cache(&remote_keys, context))? {
Some(remote) => self.merge_batch(keys, context, 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(
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
kwargs: CacheKwargs,
context: C,
) -> Result<Vec<BatchEntry<V>>, Error> {
let entries = self
.l1
.async_get_cache_batch(keys.clone(), kwargs.clone())
.async_batch_get_cache(keys.clone(), context.clone())
.await?;
let missing = Self::missing(&entries);
if missing.is_empty() || !self.reads_remote() {
@ -237,29 +267,22 @@ where
let remote_keys = missing.iter().map(|index| keys[*index].clone()).collect();
match self.remote(
self.l2
.async_get_cache_batch(remote_keys, kwargs.clone())
.async_batch_get_cache(remote_keys, context.clone())
.await,
)? {
Some(remote) => self.merge_batch(&keys, &kwargs, entries, missing, remote),
Some(remote) => self.merge_batch(&keys, &context, 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
}
impl<V, C, L1, L2> DeleteCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: DeleteCache<Value = V, Context = C>,
L2: DeleteCache<Value = V, Context = C>,
{
fn delete_cache(&self, key: &str) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.delete_cache(key))?;
@ -273,7 +296,15 @@ where
}
self.l1.async_delete_cache(key).await
}
}
impl<V, C, L1, L2> FlushCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: FlushCache<Value = V, Context = C>,
L2: FlushCache<Value = V, Context = C>,
{
fn flush_cache(&self) -> Result<(), Error> {
if self.writes_remote() {
self.remote(self.l2.flush_cache())?;
@ -287,65 +318,47 @@ where
}
self.l1.async_flush_cache().await
}
async fn disconnect(&self) -> Result<(), Error> {
self.l2.disconnect().await?;
self.l1.disconnect().await
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
self.l2.test_connection().await
}
}
impl<L1, L2> CounterCache for DualCache<L1, L2>
impl<C, L1, L2> CounterCache for DualCache<L1, L2>
where
L1: BaseCache<Value = f64>,
L2: CounterCache,
C: CacheContext,
L1: BaseCache<Value = f64, Context = C>,
L2: CounterCache<Context = C>,
{
fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result<f64, Error> {
let value = self.l2.increment_cache(key, amount, kwargs.clone())?;
self.l1.set_cache(key, value, kwargs)?;
fn increment_cache(&self, key: &str, amount: f64, context: C) -> Result<f64, Error> {
let value = self.l2.increment_cache(key, amount, context.clone())?;
self.l1.set_cache(key, value, &context)?;
Ok(value)
}
async fn async_increment_cache(
&self,
key: &str,
amount: f64,
kwargs: CacheKwargs,
) -> Result<f64, Error> {
async fn async_increment(&self, key: &str, amount: f64, context: C) -> Result<f64, Error> {
let value = self
.l2
.async_increment_cache(key, amount, kwargs.clone())
.async_increment(key, amount, context.clone())
.await?;
self.l1.async_set_cache(key, value, kwargs).await?;
self.l1.async_set_cache(key, value, context).await?;
Ok(value)
}
}
impl<V, L1, L2> ClaimCache for DualCache<L1, L2>
impl<V, C, L1, L2> ClaimCache for DualCache<L1, L2>
where
V: Clone + PartialEq + Send + Sync + 'static,
L1: ClaimCache<Value = V>,
L2: ClaimCache<Value = V>,
C: CacheContext,
L1: ClaimCache<Value = V, Context = C>,
L2: ClaimCache<Value = V, Context = C>,
{
fn claim_cache(
&self,
key: &str,
candidate: V,
eligible: &[V],
kwargs: CacheKwargs,
) -> Result<V, Error> {
fn claim_cache(&self, key: &str, candidate: V, eligible: &[V], context: C) -> Result<V, Error> {
match self.remote(
self.l2
.claim_cache(key, candidate.clone(), eligible, kwargs.clone()),
.claim_cache(key, candidate.clone(), eligible, context.clone()),
)? {
Some(winner) => {
self.l1.set_cache(key, winner.clone(), kwargs)?;
self.l1.set_cache(key, winner.clone(), &context)?;
Ok(winner)
}
None => self.l1.claim_cache(key, candidate, eligible, kwargs),
None => self.l1.claim_cache(key, candidate, eligible, context),
}
}
@ -354,20 +367,22 @@ where
key: &str,
candidate: V,
eligible: Vec<V>,
kwargs: CacheKwargs,
context: C,
) -> Result<V, Error> {
match self.remote(
self.l2
.async_claim_cache(key, candidate.clone(), eligible.clone(), kwargs.clone())
.async_claim_cache(key, candidate.clone(), eligible.clone(), context.clone())
.await,
)? {
Some(winner) => {
self.l1.async_set_cache(key, winner.clone(), kwargs).await?;
self.l1
.async_set_cache(key, winner.clone(), context)
.await?;
Ok(winner)
}
None => {
self.l1
.async_claim_cache(key, candidate, eligible, kwargs)
.async_claim_cache(key, candidate, eligible, context)
.await
}
}

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

@ -1,14 +1,21 @@
mod base_cache;
mod cache_type;
mod caching;
mod capabilities;
mod codec;
pub mod dual;
mod dual;
mod error;
pub use base_cache::{
BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheKwargs,
BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext,
ExactCacheContext,
};
pub use cache_type::CacheType;
pub use caching::{Cache, CacheBackend, get_cache, set_cache};
pub use capabilities::{ClaimCache, CounterCache, IncrementOperation};
pub use capabilities::{
BatchCache, CacheScript, ClaimCache, ClientInfoCache, CounterCache, DeleteCache, FlushCache,
IncrementOperation, QueueCache, ScanCache, ScriptCache, SetCache, TtlCache,
};
pub use codec::{CacheCodec, JsonCodec};
pub use dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy};
pub use error::Error;

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

@ -1,20 +1,69 @@
use std::{sync::Mutex, time::Duration};
use litellm_cache::{BaseCache, CacheConnectionResult, CacheKwargs, Error};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache,
};
struct TestCache {
default_ttl: Duration,
writes: Mutex<Vec<(String, String, CacheKwargs)>>,
writes: Mutex<Vec<(String, String, ExactCacheContext)>>,
}
#[derive(Clone)]
struct SemanticContext {
ttl: Option<Duration>,
query: String,
}
impl CacheContext for SemanticContext {
fn ttl(&self) -> Option<Duration> {
self.ttl
}
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
Self {
ttl,
query: self.query.clone(),
}
}
}
struct SemanticCache;
impl BaseCache for SemanticCache {
type Value = String;
type Context = SemanticContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(&self, _: &str, _: Self::Value, _: &Self::Context) -> Result<(), Error> {
Ok(())
}
fn get_cache(&self, _: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
Ok((context.query == "matching prompt").then(|| "semantic hit".into()))
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
unreachable!()
}
}
impl BaseCache for TestCache {
type Value = String;
type Context = ExactCacheContext;
fn default_ttl(&self) -> Duration {
self.default_ttl
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(self.default_ttl))
}
fn set_cache(&self, _: &str, _: Self::Value, _: CacheKwargs) -> Result<(), Error> {
fn set_cache(&self, _: &str, _: Self::Value, _: &ExactCacheContext) -> Result<(), Error> {
Err(Error::Unavailable)
}
@ -22,7 +71,7 @@ impl BaseCache for TestCache {
&self,
key: &str,
value: Self::Value,
kwargs: CacheKwargs,
context: ExactCacheContext,
) -> Result<(), Error> {
if key == "unavailable" {
return Err(Error::Unavailable);
@ -30,22 +79,14 @@ impl BaseCache for TestCache {
self.writes
.lock()
.unwrap()
.push((key.into(), value, kwargs));
.push((key.into(), value, context));
Ok(())
}
fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result<Option<Self::Value>, Error> {
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
Ok(None)
}
fn delete_cache(&self, _: &str) -> Result<(), Error> {
Ok(())
}
fn flush_cache(&self) -> Result<(), Error> {
Ok(())
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
@ -62,15 +103,26 @@ fn ttl_uses_default_and_allows_per_call_override() {
writes: Mutex::default(),
};
assert_eq!(
cache.get_ttl(&CacheKwargs::default()),
Duration::from_secs(60)
cache.get_ttl(&ExactCacheContext::default()),
Some(Duration::from_secs(60))
);
assert_eq!(
cache.get_ttl(&CacheKwargs {
cache.get_ttl(&ExactCacheContext {
ttl: Some(Duration::from_secs(5)),
..Default::default()
}),
Duration::from_secs(5)
Some(Duration::from_secs(5))
);
}
#[test]
fn associated_context_preserves_backend_specific_lookup_inputs() {
let context = SemanticContext {
ttl: None,
query: "matching prompt".into(),
};
assert_eq!(
get_cache(&SemanticCache, "shared-key", &context).unwrap(),
Some("semantic hit".into())
);
}
@ -81,12 +133,11 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() {
writes: Mutex::default(),
};
let entry = String::from("cached");
let kwargs = CacheKwargs {
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(5)),
..Default::default()
};
cache
.batch_cache_write("single", entry.clone(), kwargs.clone())
.batch_cache_write("single", entry.clone(), context.clone())
.await
.unwrap();
assert_eq!(
@ -97,7 +148,7 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() {
("unavailable".into(), entry.clone()),
("skipped".into(), entry.clone()),
],
kwargs.clone(),
context.clone(),
)
.await,
Err(Error::Unavailable)
@ -105,8 +156,8 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() {
assert_eq!(
*cache.writes.lock().unwrap(),
vec![
("single".into(), entry.clone(), kwargs.clone()),
("first".into(), entry, kwargs),
("single".into(), entry.clone(), context.clone()),
("first".into(), entry, context),
]
);
}

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

@ -4,8 +4,8 @@ use std::{
};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error,
dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy},
BaseCache, BatchCache, CacheConnectionResult, ClaimCache, CounterCache, DeleteCache, DualCache,
Error, ExactCacheContext, FlushCache, ReadPolicy, RemoteFailurePolicy, WritePolicy,
};
struct TestCache<V> {
@ -27,26 +27,21 @@ where
V: Clone + Send + Sync + 'static,
{
type Value = V;
type Context = ExactCacheContext;
fn set_cache(&self, _: &str, value: V, _: CacheKwargs) -> Result<(), Error> {
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(Duration::from_secs(60)))
}
fn set_cache(&self, _: &str, value: V, _: &ExactCacheContext) -> Result<(), Error> {
*self.value.lock().unwrap() = Some(value);
Ok(())
}
fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result<Option<V>, Error> {
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<V>, Error> {
Ok(self.value.lock().unwrap().clone())
}
fn delete_cache(&self, _: &str) -> Result<(), Error> {
*self.value.lock().unwrap() = None;
Ok(())
}
fn flush_cache(&self) -> Result<(), Error> {
*self.value.lock().unwrap() = None;
Ok(())
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
@ -56,8 +51,30 @@ where
}
}
impl<V> BatchCache for TestCache<V> where V: Clone + Send + Sync + 'static {}
impl<V> DeleteCache for TestCache<V>
where
V: Clone + Send + Sync + 'static,
{
fn delete_cache(&self, _: &str) -> Result<(), Error> {
*self.value.lock().unwrap() = None;
Ok(())
}
}
impl<V> FlushCache for TestCache<V>
where
V: Clone + Send + Sync + 'static,
{
fn flush_cache(&self) -> Result<(), Error> {
*self.value.lock().unwrap() = None;
Ok(())
}
}
impl CounterCache for TestCache<f64> {
fn increment_cache(&self, _: &str, amount: f64, _: CacheKwargs) -> Result<f64, Error> {
fn increment_cache(&self, _: &str, amount: f64, _: ExactCacheContext) -> Result<f64, Error> {
if self.fail {
return Err(Error::Unavailable);
}
@ -77,7 +94,7 @@ where
_: &str,
candidate: V,
eligible: &[V],
_: CacheKwargs,
_: ExactCacheContext,
) -> Result<V, Error> {
if self.fail {
return Err(Error::Unavailable);
@ -100,11 +117,12 @@ fn failed_l2_increment_leaves_l1_unchanged() {
let cache = DualCache::new(l1.clone(), Arc::new(TestCache::new(Some(20.0), true)));
assert_eq!(
cache.increment_cache("counter", 2.0, CacheKwargs::default()),
cache.increment_cache("counter", 2.0, ExactCacheContext::default()),
Err(Error::Unavailable)
);
assert_eq!(
l1.get_cache("counter", &CacheKwargs::default()).unwrap(),
l1.get_cache("counter", &ExactCacheContext::default())
.unwrap(),
Some(10.0)
);
}
@ -121,9 +139,8 @@ fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() {
"affinity",
"second".into(),
&["first".into(), "second".into()],
CacheKwargs {
ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
..Default::default()
},
)
.unwrap(),
@ -135,12 +152,17 @@ struct SyncPanics(TestCache<String>);
impl BaseCache for SyncPanics {
type Value = String;
type Context = ExactCacheContext;
fn set_cache(&self, _: &str, _: String, _: CacheKwargs) -> Result<(), Error> {
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
self.0.get_ttl(context)
}
fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> {
panic!("sync L2 write on an async path")
}
fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result<Option<String>, Error> {
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<String>, Error> {
panic!("sync L2 read on an async path")
}
@ -148,54 +170,30 @@ impl BaseCache for SyncPanics {
&self,
key: &str,
value: String,
kwargs: CacheKwargs,
context: ExactCacheContext,
) -> Result<(), Error> {
self.0.set_cache(key, value, kwargs)
self.0.set_cache(key, value, &context)
}
async fn async_get_cache(
&self,
key: &str,
kwargs: &CacheKwargs,
context: &ExactCacheContext,
) -> 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,
}])
self.0.get_cache(key, context)
}
async fn async_set_cache_pipeline(
&self,
cache_list: Vec<(String, String)>,
kwargs: CacheKwargs,
context: ExactCacheContext,
) -> Result<(), Error> {
for (key, value) in cache_list {
self.0.set_cache(&key, value, kwargs.clone())?;
self.0.set_cache(&key, value, &context)?;
}
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(())
}
@ -205,6 +203,36 @@ impl BaseCache for SyncPanics {
}
}
impl BatchCache for SyncPanics {
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
context: ExactCacheContext,
) -> Result<Vec<litellm_cache::BatchEntry<String>>, Error> {
assert_eq!(keys, ["missing"]);
Ok(vec![match self.0.get_cache("missing", &context)? {
Some(value) => litellm_cache::BatchEntry::Hit(value),
None => litellm_cache::BatchEntry::Miss,
}])
}
}
impl DeleteCache for SyncPanics {
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)
}
}
impl FlushCache for SyncPanics {
fn flush_cache(&self) -> Result<(), Error> {
panic!("sync L2 flush on an async path")
}
}
#[tokio::test]
async fn async_operations_use_the_async_l2_methods() {
let l1 = Arc::new(TestCache::new(None, false));
@ -215,36 +243,36 @@ async fn async_operations_use_the_async_l2_methods() {
false,
))),
);
let kwargs = CacheKwargs::default();
let context = ExactCacheContext::default();
assert_eq!(
cache.async_get_cache("missing", &kwargs).await.unwrap(),
cache.async_get_cache("missing", &context).await.unwrap(),
Some("remote".into())
);
assert_eq!(
l1.get_cache("missing", &kwargs).unwrap(),
l1.get_cache("missing", &context).unwrap(),
Some("remote".into())
);
l1.delete_cache("missing").unwrap();
assert_eq!(
cache
.async_get_cache_batch(vec!["missing".into()], kwargs.clone())
.async_batch_get_cache(vec!["missing".into()], context.clone())
.await
.unwrap(),
[litellm_cache::BatchEntry::Hit("remote".to_string())]
);
cache
.async_set_cache("missing", "written".into(), kwargs.clone())
.async_set_cache("missing", "written".into(), context.clone())
.await
.unwrap();
cache
.async_set_cache_pipeline(vec![("missing".into(), "piped".into())], kwargs.clone())
.async_set_cache_pipeline(vec![("missing".into(), "piped".into())], context.clone())
.await
.unwrap();
cache.async_delete_cache("missing").await.unwrap();
assert_eq!(
cache.async_get_cache("missing", &kwargs).await.unwrap(),
cache.async_get_cache("missing", &context).await.unwrap(),
None
);
}
@ -253,20 +281,17 @@ struct Unavailable;
impl BaseCache for Unavailable {
type Value = String;
type Context = ExactCacheContext;
fn set_cache(&self, _: &str, _: String, _: CacheKwargs) -> Result<(), Error> {
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> 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> {
fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result<Option<String>, Error> {
Err(Error::Unavailable)
}
@ -279,13 +304,27 @@ impl BaseCache for Unavailable {
}
}
impl BatchCache for Unavailable {}
impl DeleteCache for Unavailable {
fn delete_cache(&self, _: &str) -> Result<(), Error> {
Err(Error::Unavailable)
}
}
impl FlushCache for Unavailable {
fn flush_cache(&self) -> Result<(), Error> {
Err(Error::Unavailable)
}
}
impl ClaimCache for Unavailable {
fn claim_cache(
&self,
_: &str,
_: String,
_: &[String],
_: CacheKwargs,
_: ExactCacheContext,
) -> Result<String, Error> {
Err(Error::InvalidEntry)
}
@ -293,24 +332,25 @@ impl ClaimCache for Unavailable {
#[test]
fn remote_failure_policy_selects_propagation_or_the_local_tier() {
let kwargs = CacheKwargs::default();
let context = ExactCacheContext::default();
let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable));
assert_eq!(
strict.set_cache("key", "value".into(), kwargs.clone()),
strict.set_cache("key", "value".into(), &context),
Err(Error::Unavailable)
);
assert_eq!(strict.get_cache("key", &kwargs), Err(Error::Unavailable));
assert_eq!(strict.get_cache("key", &context), 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())));
assert_eq!(degraded.get_cache("key", &context), Ok(None));
degraded.set_cache("key", "value".into(), &context).unwrap();
assert_eq!(
degraded.get_cache("key", &context),
Ok(Some("value".into()))
);
degraded.delete_cache("key").unwrap();
assert_eq!(l1.get_cache("key", &kwargs), Ok(None));
assert_eq!(l1.get_cache("key", &context), Ok(None));
}
#[test]
@ -321,7 +361,12 @@ fn claim_fallback_does_not_hide_non_availability_errors() {
)
.with_remote_failure_policy(RemoteFailurePolicy::UseLocal);
assert_eq!(
cache.claim_cache("affinity", "second".into(), &[], CacheKwargs::default()),
cache.claim_cache(
"affinity",
"second".into(),
&[],
ExactCacheContext::default()
),
Err(Error::InvalidEntry)
);
}
@ -332,11 +377,9 @@ fn local_only_policies_never_touch_l2() {
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();
let context = ExactCacheContext::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())));
assert_eq!(cache.get_cache("key", &context), Ok(None));
cache.set_cache("key", "local".into(), &context).unwrap();
assert_eq!(l2.get_cache("key", &context), Ok(Some("remote".into())));
}

22
litellm-rust/crates/python-bridge/src/cache/config.rs vendored
View file

@ -1,5 +1,6 @@
use std::time::Duration;
use litellm_cache::CacheType;
use pyo3::{
exceptions::{PyTypeError, PyValueError},
prelude::*,
@ -127,21 +128,30 @@ impl NativeCacheConfig {
.extract::<String>()?,
};
let backend = facade.getattr("cache")?;
match backend_name.as_str() {
"local" => project_memory(&backend).map(|backend| {
match CacheType::from_python_name(&backend_name) {
Some(CacheType::Local) => project_memory(&backend).map(|backend| {
CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::Memory(backend),
}))
}),
"redis" => match project_redis(&backend)? {
Some(CacheType::Redis) => match project_redis(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::Redis(Box::new(backend)),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
_ => Ok(CacheConfigProjection::Unsupported(
Some(
CacheType::RedisSemantic
| CacheType::ValkeySemantic
| CacheType::S3
| CacheType::Disk
| CacheType::QdrantSemantic
| CacheType::AzureBlob
| CacheType::Gcs,
)
| None => Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
)),
}
@ -149,10 +159,10 @@ impl NativeCacheConfig {
pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> {
if service.default_ttl()
!= match &self.backend {
!= Some(match &self.backend {
CacheBackendConfig::Memory(config) => config.default_ttl,
CacheBackendConfig::Redis(config) => config.default_ttl,
}
})
{
return Some("facade and native backend default TTLs must match");
}

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

@ -53,7 +53,7 @@ impl NativeResponseCache {
}
}
pub fn default_ttl(&self) -> Duration {
pub fn default_ttl(&self) -> Option<Duration> {
match self {
Self::Memory(cache) => cache.default_ttl(),
Self::Redis { cache, .. } => cache.default_ttl(),

2
litellm-rust/crates/python-bridge/src/cache/request.rs vendored
View file

@ -24,7 +24,7 @@ fn request_input(input: RequestInput) -> PyResult<ResponseCacheRequest> {
if let Some(controls) = input.controls {
request.controls = controls;
}
request.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?;
request.context.ttl = input.ttl_seconds.map(duration).transpose()?;
request.max_age = input.max_age_seconds.map(duration).transpose()?;
Ok(request)
}