diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 54d831378e4..85850c1d925 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -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 { 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 { - fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + fn increment_cache( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + ) -> Result { if self.max_size_in_memory == 0 { return Ok(amount); } @@ -255,7 +265,11 @@ impl CounterCache for InMemoryCache { 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 { self.increment_cache( &operation.key, operation.amount, - CacheKwargs { - ttl: operation.ttl, - ..CacheKwargs::default() - }, + ExactCacheContext { ttl: operation.ttl }, ) }) .collect() @@ -285,28 +296,26 @@ impl InMemoryCache { impl BaseCache for InMemoryCache { type Value = V; + type Context = ExactCacheContext; - fn default_ttl(&self) -> Duration { - self.default_ttl + fn get_ttl(&self, context: &Self::Context) -> Option { + 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, Error> { + fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, 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 BaseCache for InMemoryCache { } } +impl BatchCache for InMemoryCache {} + +impl DeleteCache for InMemoryCache { + fn delete_cache(&self, key: &str) -> Result<(), Error> { + InMemoryCache::delete_cache(self, key) + } +} + +impl FlushCache for InMemoryCache { + fn flush_cache(&self) -> Result<(), Error> { + InMemoryCache::flush_cache(self) + } +} + +impl TtlCache for InMemoryCache { + async fn async_get_ttl(&self, key: &str) -> Result, Error> { + InMemoryCache::async_get_ttl(self, key).await + } +} + +impl SetCache for InMemoryCache> +where + T: Clone + Eq + Hash + Send + Sync + 'static, +{ + type SetValue = T; + type SetResult = Vec; + + async fn async_set_cache_sadd( + &self, + key: &str, + values: Vec, + ttl: Option, + ) -> Result { + 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::::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); diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index c44590b63ca..0df0319b990 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -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> = 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::::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::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::::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::::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::>::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()])) + ); +} diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 2249966dd79..a960c383bf4 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -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 { @@ -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 BaseCache for RedisCache 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 { + 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, Error> { + fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, 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>, Error> { - let keys = keys - .iter() - .map(|key| self.namespaced_key(key)) - .collect::>(); - let values = self.connections.execute(|connection| { - redis::cmd("MGET") - .arg(keys) - .query::>(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, 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, - _: CacheKwargs, - ) -> Result>, Error> { - let keys = keys - .iter() - .map(|key| self.namespaced_key(key)) - .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { - redis::cmd("MGET") - .arg(keys) - .query::>(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::, _>>()?; - 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 BatchCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn batch_get_cache( + &self, + keys: &[String], + _: &ExactCacheContext, + ) -> Result>, Error> { + let keys = keys + .iter() + .map(|key| self.namespaced_key(key)) + .collect::>(); + let values = self.connections.execute(|connection| { + redis::cmd("MGET") + .arg(keys) + .query::>(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, + _: ExactCacheContext, + ) -> Result>, Error> { + let keys = keys + .iter() + .map(|key| self.namespaced_key(key)) + .collect::>(); + let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("MGET") + .arg(keys) + .query::>(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + values + .into_iter() + .map(|value| self.decode_batch_response(value)) + .collect() + } +} + +impl DeleteCache for RedisCache +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 FlushCache for RedisCache +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 CounterCache for RedisCache where S: CacheCodec, C: redis::ConnectionLike + Send + 'static, { - fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + fn increment_cache( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + ) -> Result { 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 { 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 { 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, - kwargs: CacheKwargs, + context: ExactCacheContext, ) -> Result { 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(); diff --git a/litellm-rust/crates/cache-redis/src/cache/operations.rs b/litellm-rust/crates/cache-redis/src/cache/operations.rs index f8c2bf4078c..d8d9ae24c4c 100644 --- a/litellm-rust/crates/cache-redis/src/cache/operations.rs +++ b/litellm-rust/crates/cache-redis/src/cache/operations.rs @@ -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>), } +pub struct RedisScript { + connections: Arc>, + namespace: Option, + source: String, +} + +impl CacheScript for RedisScript +where + C: redis::ConnectionLike + Send + 'static, +{ + type Argument = RedisArg; + type Output = redis::Value; + + async fn invoke( + &self, + keys: Vec, + arguments: Vec, + ) -> Result { + let keys = keys + .into_iter() + .map(|key| namespaced_key(self.namespace.as_deref(), &key)) + .collect::>(); + 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 RedisCache where S: CacheCodec, @@ -498,3 +541,93 @@ fn increment_with_floor( .query(connection) .map_err(|_| Error::Unavailable) } + +impl TtlCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + async fn async_get_ttl(&self, key: &str) -> Result, Error> { + RedisCache::async_get_ttl(self, key) + .await + .map(|ttl| ttl.map(|seconds| Duration::from_secs(seconds as u64))) + } +} + +impl ScanCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, Error> { + RedisCache::async_scan_iter(self, pattern, count).await + } +} + +impl ClientInfoCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type ClientList = String; + type Info = String; + + fn client_list(&self) -> Result { + RedisCache::client_list(self) + } + + fn info(&self) -> Result { + RedisCache::info(self) + } +} + +impl SetCache for RedisCache +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, + ttl: Option, + ) -> Result { + RedisCache::async_set_cache_sadd(self, key, values, ttl).await + } +} + +impl QueueCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type QueueValue = RedisArg; + type PopResult = RedisLpopResult; + + async fn async_rpush(&self, key: &str, values: Vec) -> Result { + RedisCache::async_rpush(self, key, values).await + } + + async fn async_lpop(&self, key: &str, count: Option) -> Result { + RedisCache::async_lpop(self, key, count).await + } +} + +impl ScriptCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type Script = RedisScript; + + fn async_register_script(&self, source: String) -> Self::Script { + RedisScript { + connections: Arc::clone(&self.connections), + namespace: self.namespace.clone(), + source, + } + } +} diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs index 2548c7ac3c6..98f6bfd8ce5 100644 --- a/litellm-rust/crates/cache-redis/src/lib.rs +++ b/litellm-rust/crates/cache-redis/src/lib.rs @@ -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}; diff --git a/litellm-rust/crates/cache-redis/src/topology.rs b/litellm-rust/crates/cache-redis/src/topology.rs new file mode 100644 index 00000000000..7f4ee48b222 --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/topology.rs @@ -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, + }, +} diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index 3d755d0f44c..337f27984f8 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -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::::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 diff --git a/litellm-rust/crates/cache-response/src/buffer.rs b/litellm-rust/crates/cache-response/src/buffer.rs index 1fd2bb809de..606c21410c7 100644 --- a/litellm-rust/crates/cache-response/src/buffer.rs +++ b/litellm-rust/crates/cache-response/src/buffer.rs @@ -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>( + pub async fn async_store>( &self, cache: &ResponseCache, request: &ResponseCacheRequest, diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index f18a48863c5..e50e68cdabb 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -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, } @@ -24,17 +26,17 @@ impl ResponseCacheRequest { default_on: true, ..Default::default() }, - kwargs: CacheKwargs::default(), + context: ExactCacheContext::default(), max_age: None, } } } -pub struct ResponseCache> { +pub struct ResponseCache> { backend: Arc, } -impl> ResponseCache { +impl> ResponseCache { pub fn new(backend: Arc) -> Self { Self { backend } } @@ -43,11 +45,14 @@ impl> ResponseCache { &self.backend } - pub fn default_ttl(&self) -> Duration { - self.backend.default_ttl() + pub fn default_ttl(&self) -> Option { + 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> ResponseCache { } 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> ResponseCache { } 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> ResponseCache { &self, requests: &[ResponseCacheRequest], now: Duration, - ) -> Result { + ) -> Result + where + B: BatchCache, + { let readable = requests .iter() .enumerate() @@ -109,7 +117,7 @@ impl> ResponseCache { .map(|(_, request)| cache_key(&request.key)) .collect::>(); 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> ResponseCache { &self, requests: &[ResponseCacheRequest], now: Duration, - ) -> Result { + ) -> Result + where + B: BatchCache, + { let readable = requests .iter() .enumerate() @@ -132,7 +143,7 @@ impl> ResponseCache { .collect::>(); 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> ResponseCache { timestamp: Some(now.as_secs_f64()), response, }, - request.kwargs.clone(), + &request.context, ) } @@ -175,7 +186,7 @@ impl> ResponseCache { timestamp: Some(now.as_secs_f64()), response, }, - request.kwargs.clone(), + request.context.clone(), ) .await } @@ -210,26 +221,29 @@ impl> ResponseCache { timestamp: Some(now.as_secs_f64()), response, }, - request.kwargs, + request.context, ) }) .collect::>(); 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(()) } diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 94b25626f86..e4f78dae8b2 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -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(), diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index d6ef8052c4f..8bd69ba5ad6 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -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 { Invalid, } -#[derive(Clone, Debug, Default, PartialEq)] -pub struct CacheKwargs { +pub trait CacheContext: Clone + Send + Sync + 'static { + fn ttl(&self) -> Option; + + fn with_ttl(&self, ttl: Option) -> Self; +} + +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct ExactCacheContext { pub ttl: Option, - pub extras: Map, +} + +impl CacheContext for ExactCacheContext { + fn ttl(&self) -> Option { + self.ttl + } + + fn with_ttl(&self, ttl: Option) -> 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; - 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, Error>; - - fn get_cache_batch( + fn set_cache( &self, - keys: &[String], - kwargs: &CacheKwargs, - ) -> Result>, 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, Error>; fn async_set_cache( &self, key: &str, value: Self::Value, - kwargs: CacheKwargs, + context: Self::Context, ) -> impl Future> + 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, Error>> + Send { - async move { self.get_cache(key, kwargs) } - } - - fn async_get_cache_batch( - &self, - keys: Vec, - kwargs: CacheKwargs, - ) -> impl Future>, 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> + 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> + 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> + Send { - async move { self.delete_cache(key) } - } - - fn flush_cache(&self) -> Result<(), Error>; - - fn async_flush_cache(&self) -> impl Future> + Send { - async move { self.flush_cache() } + self.async_set_cache(key, value, context) } fn disconnect(&self) -> impl Future> + Send; diff --git a/litellm-rust/crates/cache/src/cache_type.rs b/litellm-rust/crates/cache/src/cache_type.rs new file mode 100644 index 00000000000..f0a97c04fd5 --- /dev/null +++ b/litellm-rust/crates/cache/src/cache_type.rs @@ -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::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) + ); + } +} diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs index 4ec94244ac5..fc7f46d943e 100644 --- a/litellm-rust/crates/cache/src/caching.rs +++ b/litellm-rust/crates/cache/src/caching.rs @@ -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( cache: &B, key: &str, - kwargs: &CacheKwargs, + context: &B::Context, ) -> Result, Error> { - cache.get_cache(key, kwargs) + cache.get_cache(key, context) } pub fn set_cache( 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 = Arc; diff --git a/litellm-rust/crates/cache/src/capabilities.rs b/litellm-rust/crates/cache/src/capabilities.rs index 1613a5786f4..f7307e5c7bd 100644 --- a/litellm-rust/crates/cache/src/capabilities.rs +++ b/litellm-rust/crates/cache/src/capabilities.rs @@ -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, + pub ttl: Option, +} + +pub trait BatchCache: BaseCache { + fn batch_get_cache( + &self, + keys: &[String], + context: &Self::Context, + ) -> Result>, 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, + context: Self::Context, + ) -> impl Future>, 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> + Send { + async move { self.delete_cache(key) } + } +} + +pub trait FlushCache: BaseCache { + fn flush_cache(&self) -> Result<(), Error>; + + fn async_flush_cache(&self) -> impl Future> + Send { + async move { self.flush_cache() } + } } pub trait CounterCache: BaseCache { - fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result; + fn increment_cache(&self, key: &str, amount: f64, context: Self::Context) + -> Result; - fn async_increment_cache( + fn async_increment( &self, key: &str, amount: f64, - kwargs: CacheKwargs, + context: Self::Context, ) -> impl Future> + 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; fn async_claim_cache( @@ -39,8 +92,78 @@ where key: &str, candidate: Self::Value, eligible: Vec, - kwargs: CacheKwargs, + context: Self::Context, ) -> impl Future> + 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, 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, + ttl: Option, + ) -> impl Future> + Send; +} + +pub trait QueueCache: BaseCache { + type QueueValue: Clone + Send + Sync + 'static; + type PopResult: Send + Sync + 'static; + + fn async_rpush( + &self, + key: &str, + values: Vec, + ) -> impl Future> + Send; + + fn async_lpop( + &self, + key: &str, + count: Option, + ) -> impl Future> + Send; +} + +pub trait ScanCache: BaseCache { + fn async_scan_iter( + &self, + pattern: &str, + count: usize, + ) -> impl Future, Error>> + Send; +} + +pub trait ClientInfoCache: BaseCache { + type ClientList: Send + Sync + 'static; + type Info: Send + Sync + 'static; + + fn client_list(&self) -> Result; + + fn info(&self) -> Result; +} + +pub trait CacheScript: Send + Sync + 'static { + type Argument: Clone + Send + Sync + 'static; + type Output: Send + Sync + 'static; + + fn invoke( + &self, + keys: Vec, + arguments: Vec, + ) -> impl Future> + Send; +} + +pub trait ScriptCache: BaseCache { + type Script: CacheScript; + + fn async_register_script(&self, source: String) -> Self::Script; +} diff --git a/litellm-rust/crates/cache/src/dual.rs b/litellm-rust/crates/cache/src/dual.rs index 9858c5e2748..d68d4b2b69f 100644 --- a/litellm-rust/crates/cache/src/dual.rs +++ b/litellm-rust/crates/cache/src/dual.rs @@ -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 DualCache { } } - fn promotion_kwargs(&self, kwargs: &CacheKwargs) -> CacheKwargs { - CacheKwargs { - ttl: self.promotion_ttl.or(kwargs.ttl), - extras: kwargs.extras.clone(), - } + fn promotion_context(&self, context: &C) -> C { + context.with_ttl(self.promotion_ttl.or(context.ttl())) } } -impl DualCache +impl DualCache where V: Clone + Send + Sync + 'static, - L1: BaseCache, - L2: BaseCache, + C: CacheContext, + L1: BaseCache, + L2: BaseCache, { fn missing(entries: &[BatchEntry]) -> Vec { entries @@ -119,7 +118,7 @@ where fn merge_batch( &self, keys: &[String], - kwargs: &CacheKwargs, + context: &C, mut entries: Vec>, missing: Vec, remote: Vec>, @@ -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 BaseCache for DualCache +impl BaseCache for DualCache where V: Clone + Send + Sync + 'static, - L1: BaseCache, - L2: BaseCache, + C: CacheContext, + L1: BaseCache, + L2: BaseCache, { type Value = V; + type Context = C; - fn default_ttl(&self) -> Duration { - self.l2.default_ttl() + fn get_ttl(&self, context: &Self::Context) -> Option { + 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, Error> { - if let Some(value) = self.l1.get_cache(key, kwargs)? { + fn get_cache(&self, key: &str, context: &C) -> Result, 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, 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>, 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 { + self.l2.test_connection().await + } +} + +impl BatchCache for DualCache +where + V: Clone + Send + Sync + 'static, + C: CacheContext, + L1: BatchCache, + L2: BatchCache, +{ + fn batch_get_cache(&self, keys: &[String], context: &C) -> Result>, 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::>(); - 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, 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, - kwargs: CacheKwargs, + context: C, ) -> Result>, 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 DeleteCache for DualCache +where + V: Clone + Send + Sync + 'static, + C: CacheContext, + L1: DeleteCache, + L2: DeleteCache, +{ 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 FlushCache for DualCache +where + V: Clone + Send + Sync + 'static, + C: CacheContext, + L1: FlushCache, + L2: FlushCache, +{ 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 { - self.l2.test_connection().await - } } -impl CounterCache for DualCache +impl CounterCache for DualCache where - L1: BaseCache, - L2: CounterCache, + C: CacheContext, + L1: BaseCache, + L2: CounterCache, { - fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { - 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 { + 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 { + async fn async_increment(&self, key: &str, amount: f64, context: C) -> Result { 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 ClaimCache for DualCache +impl ClaimCache for DualCache where V: Clone + PartialEq + Send + Sync + 'static, - L1: ClaimCache, - L2: ClaimCache, + C: CacheContext, + L1: ClaimCache, + L2: ClaimCache, { - fn claim_cache( - &self, - key: &str, - candidate: V, - eligible: &[V], - kwargs: CacheKwargs, - ) -> Result { + fn claim_cache(&self, key: &str, candidate: V, eligible: &[V], context: C) -> Result { 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, - kwargs: CacheKwargs, + context: C, ) -> Result { 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 } } diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index cdd5589aa39..ce9f93b6dc4 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -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; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 5c27aa9ff4e..9180ee9d0dc 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -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>, + writes: Mutex>, +} + +#[derive(Clone)] +struct SemanticContext { + ttl: Option, + query: String, +} + +impl CacheContext for SemanticContext { + fn ttl(&self) -> Option { + self.ttl + } + + fn with_ttl(&self, ttl: Option) -> 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 { + context.ttl + } + + fn set_cache(&self, _: &str, _: Self::Value, _: &Self::Context) -> Result<(), Error> { + Ok(()) + } + + fn get_cache(&self, _: &str, context: &Self::Context) -> Result, Error> { + Ok((context.query == "matching prompt").then(|| "semantic hit".into())) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + 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 { + 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, Error> { + fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, 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), ] ); } diff --git a/litellm-rust/crates/cache/tests/dual.rs b/litellm-rust/crates/cache/tests/dual.rs index 2e1e72c7119..e7e8927f8d0 100644 --- a/litellm-rust/crates/cache/tests/dual.rs +++ b/litellm-rust/crates/cache/tests/dual.rs @@ -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 { @@ -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 { + 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, Error> { + fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, 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 BatchCache for TestCache where V: Clone + Send + Sync + 'static {} + +impl DeleteCache for TestCache +where + V: Clone + Send + Sync + 'static, +{ + fn delete_cache(&self, _: &str) -> Result<(), Error> { + *self.value.lock().unwrap() = None; + Ok(()) + } +} + +impl FlushCache for TestCache +where + V: Clone + Send + Sync + 'static, +{ + fn flush_cache(&self) -> Result<(), Error> { + *self.value.lock().unwrap() = None; + Ok(()) + } +} + impl CounterCache for TestCache { - fn increment_cache(&self, _: &str, amount: f64, _: CacheKwargs) -> Result { + fn increment_cache(&self, _: &str, amount: f64, _: ExactCacheContext) -> Result { if self.fail { return Err(Error::Unavailable); } @@ -77,7 +94,7 @@ where _: &str, candidate: V, eligible: &[V], - _: CacheKwargs, + _: ExactCacheContext, ) -> Result { 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); 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 { + 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, Error> { + fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, 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, Error> { - self.0.get_cache(key, kwargs) - } - - async fn async_get_cache_batch( - &self, - keys: Vec, - kwargs: CacheKwargs, - ) -> Result>, 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, + context: ExactCacheContext, + ) -> Result>, 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 { + context.ttl + } + + fn set_cache(&self, _: &str, _: String, _: &ExactCacheContext) -> Result<(), Error> { Err(Error::Unavailable) } - fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result, 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, 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 { 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()))); } diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 6e694c20705..0e7d6aee11d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -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::()?, }; 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"); } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 6a3835ac84d..a9475429e45 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -53,7 +53,7 @@ impl NativeResponseCache { } } - pub fn default_ttl(&self) -> Duration { + pub fn default_ttl(&self) -> Option { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index d8793abd115..0c5343a63d0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -24,7 +24,7 @@ fn request_input(input: RequestInput) -> PyResult { 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) }