diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 8c35a0be0b4..ccfaee44f50 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2462,6 +2462,7 @@ dependencies = [ "serde_json", "sha2 0.10.9", "thiserror 2.0.19", + "tokio", ] [[package]] diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 1908ff44a81..974cdbe9760 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -4,8 +4,7 @@ use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs, - Error, + BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, }; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; @@ -214,8 +213,8 @@ impl InMemoryCache { } } -impl BaseCache for InMemoryCache { - type Value = CacheEntry; +impl BaseCache for InMemoryCache { + type Value = V; fn default_ttl(&self) -> Duration { self.default_ttl @@ -238,17 +237,15 @@ impl BaseCache for InMemoryCache { self.flush_cache() } - fn disconnect(&self) -> CacheFuture<'_, ()> { - Box::pin(async { Ok(()) }) + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) } - fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> { - Box::pin(async { - Ok(CacheConnectionResult { - status: CacheConnectionStatus::Success, - message: "In-memory cache connection test successful".into(), - error: None, - }) + async fn test_connection(&self) -> Result { + Ok(CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "In-memory cache connection test successful".into(), + error: None, }) } } diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index aaf82641db7..ffac9d8ae64 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -2,7 +2,10 @@ use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; -use litellm_cache::{BaseCache, CacheConnectionStatus, CacheEntry, Error}; +use litellm_cache::{ + BaseCache, CacheBackend, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, get_cache, + set_cache, +}; use litellm_cache_memory::{CacheWrite, InMemoryCache}; use rstest::{fixture, rstest}; @@ -156,3 +159,45 @@ async fn connection_test_matches_python_result_contract() { }) ); } + +#[tokio::test] +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 { + ttl: Some(Duration::from_secs(5)), + ..Default::default() + }; + set_cache(cache.as_ref(), "sync", "first".into(), kwargs.clone()).unwrap(); + assert_eq!( + get_cache(reader.as_ref(), "sync", &kwargs).unwrap(), + Some("first".into()) + ); + cache + .batch_cache_write("async", "second".into(), kwargs.clone()) + .await + .unwrap(); + cache + .async_set_cache_pipeline(vec![("batch".into(), "third".into())], kwargs.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(), + Some(value.into()) + ); + } + reader.async_delete_cache("async").await.unwrap(); + assert_eq!( + reader.async_get_cache("async", &kwargs).await.unwrap(), + None + ); + clock.store(106, Ordering::SeqCst); + assert_eq!(get_cache(reader.as_ref(), "sync", &kwargs).unwrap(), None); + assert_eq!( + reader.async_get_cache("batch", &kwargs).await.unwrap(), + None + ); +} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index 933b0feaae4..a60813b6260 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -8,8 +8,8 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true redis = "1.7.0" -serde_json.workspace = true tokio.workspace = true [dev-dependencies] redis-test = "1.0.4" +serde_json.workspace = true diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 69dee6c6363..8d92fdc8f75 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -2,35 +2,37 @@ use std::sync::{Arc, Mutex, MutexGuard}; use std::time::Duration; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheFuture, CacheKwargs, - Error, + BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error, }; use redis::Commands; const DEFAULT_TTL: Duration = Duration::from_secs(600); const KEY_PREFIX: &str = "litellm-cache:"; -pub struct RedisCache { +pub struct RedisCache { connection: Arc>, default_ttl: Duration, + codec: S, } -impl RedisCache { - pub fn new(url: &str, default_ttl: Option) -> Result { +impl RedisCache { + pub fn new(url: &str, default_ttl: Option, codec: S) -> Result { let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; let connection = client.get_connection().map_err(|_| Error::Unavailable)?; - Ok(Self::with_connection(connection, default_ttl)) + Ok(Self::with_connection(connection, default_ttl, codec)) } } -impl RedisCache +impl RedisCache where + S: CacheCodec, C: redis::ConnectionLike + Send + 'static, { - fn with_connection(connection: C, default_ttl: Option) -> Self { + pub fn with_connection(connection: C, default_ttl: Option, codec: S) -> Self { Self { connection: Arc::new(Mutex::new(connection)), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), + codec, } } @@ -47,48 +49,39 @@ where PATTERN } - fn encode(value: &CacheEntry) -> Result, Error> { - serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) - } - - fn decode(value: Vec) -> Result { - serde_json::from_slice(&value).map_err(|_| Error::InvalidEntry) - } - fn ttl_seconds(ttl: Duration) -> u64 { ttl.as_secs() .saturating_add(u64::from(ttl.subsec_nanos() > 0)) .max(1) } - fn run_blocking(connection: Arc>, operation: F) -> CacheFuture<'static, T> + async fn run_blocking(connection: Arc>, operation: F) -> Result where T: Send + 'static, F: FnOnce(&mut C) -> Result + Send + 'static, { - Box::pin(async move { - tokio::task::spawn_blocking(move || { - let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; - operation(&mut connection) - }) - .await - .map_err(|_| Error::Unavailable)? + tokio::task::spawn_blocking(move || { + let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; + operation(&mut connection) }) + .await + .map_err(|_| Error::Unavailable)? } } -impl BaseCache for RedisCache +impl BaseCache for RedisCache where + S: CacheCodec, C: redis::ConnectionLike + Send + 'static, { - type Value = CacheEntry; + type Value = S::Value; fn default_ttl(&self) -> Duration { self.default_ttl } fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> { - let payload = Self::encode(&value)?; + let payload = self.codec.encode(&value)?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); self.connection()? .set_ex::<_, _, ()>(Self::namespaced_key(key), payload, ttl) @@ -96,11 +89,11 @@ where } fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { - self.connection()? + let bytes = self + .connection()? .get::<_, Option>>(Self::namespaced_key(key)) - .map_err(|_| Error::Unavailable)? - .map(Self::decode) - .transpose() + .map_err(|_| Error::Unavailable)?; + bytes.map(|bytes| self.codec.decode(&bytes)).transpose() } fn delete_cache(&self, key: &str) -> Result<(), Error> { @@ -125,86 +118,87 @@ where .map_err(|_| Error::Unavailable) } - fn async_set_cache<'a>( - &'a self, - key: &'a str, + async fn async_set_cache( + &self, + key: &str, value: Self::Value, kwargs: CacheKwargs, - ) -> CacheFuture<'a, ()> { - let payload = Self::encode(&value); + ) -> Result<(), Error> { + let payload = self.codec.encode(&value)?; let key = Self::namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); Self::run_blocking(Arc::clone(&self.connection), move |connection| { connection - .set_ex::<_, _, ()>(key, payload?, ttl) + .set_ex::<_, _, ()>(key, payload, ttl) .map_err(|_| Error::Unavailable) }) + .await } - fn async_get_cache<'a>( - &'a self, - key: &'a str, - _: &'a CacheKwargs, - ) -> CacheFuture<'a, Option> { + async fn async_get_cache( + &self, + key: &str, + _: &CacheKwargs, + ) -> Result, Error> { let key = Self::namespaced_key(key); - Box::pin(async move { - Self::run_blocking(Arc::clone(&self.connection), move |connection| { - connection - .get::<_, Option>>(key) - .map_err(|_| Error::Unavailable) - }) - .await? - .map(Self::decode) - .transpose() + Self::run_blocking(Arc::clone(&self.connection), move |connection| { + connection + .get::<_, Option>>(key) + .map_err(|_| Error::Unavailable) }) + .await? + .map(|bytes| self.codec.decode(&bytes)) + .transpose() } - fn async_set_cache_pipeline<'a>( - &'a self, + async fn async_set_cache_pipeline( + &self, cache_list: Vec<(String, Self::Value)>, kwargs: CacheKwargs, - ) -> CacheFuture<'a, ()> { + ) -> Result<(), Error> { let entries = cache_list .into_iter() .map(|(key, value)| { - Self::encode(&value).map(|payload| (Self::namespaced_key(&key), payload)) + self.codec + .encode(&value) + .map(|payload| (Self::namespaced_key(&key), payload)) }) - .collect::, _>>(); + .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); Self::run_blocking(Arc::clone(&self.connection), move |connection| { - for (key, payload) in entries? { + for (key, payload) in entries { connection .set_ex::<_, _, ()>(key, payload, ttl) .map_err(|_| Error::Unavailable)?; } Ok(()) }) + .await } - fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> { + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { let key = Self::namespaced_key(key); Self::run_blocking(Arc::clone(&self.connection), move |connection| { connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) }) + .await } - fn disconnect(&self) -> CacheFuture<'_, ()> { - Box::pin(async { Ok(()) }) + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) } - fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> { - Box::pin(async move { - Self::run_blocking(Arc::clone(&self.connection), |connection| { - redis::cmd("PING") - .query::(connection) - .map_err(|_| Error::Unavailable) - }) - .await?; - Ok(CacheConnectionResult { - status: CacheConnectionStatus::Success, - message: "Redis cache connection test successful".into(), - error: None, - }) + async fn test_connection(&self) -> Result { + Self::run_blocking(Arc::clone(&self.connection), |connection| { + redis::cmd("PING") + .query::(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + Ok(CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, }) } } @@ -212,7 +206,7 @@ where #[cfg(test)] mod tests { use super::RedisCache; - use litellm_cache::{BaseCache, CacheEntry, CacheKwargs}; + use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKwargs, JsonCodec}; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; use std::time::Duration; @@ -224,33 +218,18 @@ mod tests { } } - #[test] - fn cache_entries_round_trip_through_json() { - let entry = entry(); - let encoded = RedisCache::::encode(&entry).unwrap(); - assert_eq!( - RedisCache::::decode(encoded).unwrap(), - entry - ); - } - - #[test] - fn invalid_json_is_rejected() { - assert!(RedisCache::::decode(b"not json".to_vec()).is_err()); - } - #[test] fn ttl_seconds_rounds_up_and_keeps_expiration_positive() { assert_eq!( - RedisCache::::ttl_seconds(Duration::ZERO), + RedisCache::>::ttl_seconds(Duration::ZERO), 1 ); assert_eq!( - RedisCache::::ttl_seconds(Duration::from_millis(1500)), + RedisCache::>::ttl_seconds(Duration::from_millis(1500)), 2 ); assert_eq!( - RedisCache::::ttl_seconds(Duration::from_secs(15)), + RedisCache::>::ttl_seconds(Duration::from_secs(15)), 15 ); } @@ -258,7 +237,7 @@ mod tests { #[test] fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() { let value = entry(); - let payload = RedisCache::::encode(&value).unwrap(); + let payload = JsonCodec::::new().encode(&value).unwrap(); let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("SETEX") @@ -271,7 +250,7 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); cache .set_cache("key", value.clone(), CacheKwargs::default()) @@ -296,7 +275,7 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); cache.flush_cache().unwrap(); } @@ -305,7 +284,7 @@ mod tests { async fn test_connection_runs_ping_off_executor() { let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); assert_eq!( cache.test_connection().await.unwrap().status, diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index 76f73145da8..fe15fcd975b 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -1,6 +1,156 @@ +use std::time::Duration; + +use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, Error, JsonCodec, get_cache, set_cache}; use litellm_cache_redis::RedisCache; +use redis_test::{MockCmd, MockRedisConnection}; + +struct TaggedByteCodec(u8); + +impl CacheCodec for TaggedByteCodec { + type Value = u8; + + fn encode(&self, value: &u8) -> Result, Error> { + if *value > 127 { + return Err(Error::InvalidEntry); + } + Ok(vec![self.0, *value]) + } + + fn decode(&self, bytes: &[u8]) -> Result { + match bytes { + [tag, value] if *tag == self.0 => Ok(*value), + _ => Err(Error::InvalidEntry), + } + } +} #[test] fn constructor_rejects_invalid_urls() { - assert!(RedisCache::new("not a redis url", None).is_err()); + assert!(RedisCache::new("not a redis url", None, JsonCodec::::new()).is_err()); +} + +#[test] +fn generic_helpers_use_the_injected_codec_and_ttl() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SETEX") + .arg("litellm-cache:counter") + .arg(2) + .arg([42u8, 7].as_slice()), + Ok("OK"), + ), + MockCmd::new( + redis::cmd("GET").arg("litellm-cache:counter"), + Ok(vec![42u8, 7]), + ), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); + let kwargs = CacheKwargs { + 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)); +} + +#[tokio::test] +async fn async_operations_preserve_codec_ttl_and_missing_values() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SETEX") + .arg("litellm-cache:counter") + .arg(9) + .arg([42u8, 7].as_slice()), + Ok("OK"), + ), + MockCmd::new( + redis::cmd("GET").arg("litellm-cache:counter"), + Ok(vec![42u8, 7]), + ), + MockCmd::new( + redis::cmd("SETEX") + .arg("litellm-cache:batch") + .arg(2) + .arg([42u8, 8].as_slice()), + Ok("OK"), + ), + MockCmd::new(redis::cmd("DEL").arg("litellm-cache:counter"), Ok(1u32)), + MockCmd::new( + redis::cmd("GET").arg("litellm-cache:counter"), + Ok(redis::Value::Nil), + ), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection( + connection, + Some(Duration::from_secs(9)), + TaggedByteCodec(42), + ); + let kwargs = CacheKwargs::default(); + cache + .batch_cache_write("counter", 7, kwargs.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache("counter", &kwargs).await.unwrap(), + Some(7) + ); + cache + .async_set_cache_pipeline( + vec![("batch".into(), 8)], + CacheKwargs { + 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(), + None + ); +} + +#[tokio::test] +async fn codec_errors_propagate_without_writing_partial_batches() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("GET").arg("litellm-cache:invalid"), + Ok(vec![99u8, 7]), + ), + MockCmd::new( + redis::cmd("GET").arg("litellm-cache:invalid"), + Ok(vec![99u8, 7]), + ), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); + let kwargs = CacheKwargs::default(); + assert_eq!( + cache.set_cache("invalid", 255, kwargs.clone()), + Err(Error::InvalidEntry) + ); + assert_eq!( + cache.async_set_cache("invalid", 255, kwargs.clone()).await, + Err(Error::InvalidEntry) + ); + assert_eq!( + cache + .async_set_cache_pipeline( + vec![("valid".into(), 7), ("invalid".into(), 255)], + kwargs.clone(), + ) + .await, + Err(Error::InvalidEntry) + ); + assert_eq!( + cache.get_cache("invalid", &kwargs), + Err(Error::InvalidEntry) + ); + assert_eq!( + cache.async_get_cache("invalid", &kwargs).await, + Err(Error::InvalidEntry) + ); } diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml index a14c4294aa0..350db4b1adb 100644 --- a/litellm-rust/crates/cache/Cargo.toml +++ b/litellm-rust/crates/cache/Cargo.toml @@ -13,3 +13,4 @@ thiserror.workspace = true [dev-dependencies] rstest.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 2ba8ff92ebd..1891e417cb0 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -1,5 +1,4 @@ use std::future::Future; -use std::pin::Pin; use std::time::Duration; use serde::{Deserialize, Serialize}; @@ -7,8 +6,6 @@ use serde_json::{Map, Value}; use crate::Error; -pub type CacheFuture<'a, T> = Pin> + Send + 'a>>; - #[derive(Clone, Debug, Default, PartialEq)] pub struct CacheKwargs { pub ttl: Option, @@ -45,54 +42,54 @@ pub trait BaseCache: Send + Sync { fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error>; - fn async_set_cache<'a>( - &'a self, - key: &'a str, + fn async_set_cache( + &self, + key: &str, value: Self::Value, kwargs: CacheKwargs, - ) -> CacheFuture<'a, ()> { - Box::pin(async move { self.set_cache(key, value, kwargs) }) + ) -> impl Future> + Send { + async move { self.set_cache(key, value, kwargs) } } - fn async_get_cache<'a>( - &'a self, - key: &'a str, - kwargs: &'a CacheKwargs, - ) -> CacheFuture<'a, Option> { - Box::pin(async move { self.get_cache(key, kwargs) }) + fn async_get_cache( + &self, + key: &str, + kwargs: &CacheKwargs, + ) -> impl Future, Error>> + Send { + async move { self.get_cache(key, kwargs) } } - fn async_set_cache_pipeline<'a>( - &'a self, + fn async_set_cache_pipeline( + &self, cache_list: Vec<(String, Self::Value)>, kwargs: CacheKwargs, - ) -> CacheFuture<'a, ()> { - Box::pin(async move { + ) -> impl Future> + Send { + async move { for (key, value) in cache_list { - self.set_cache(&key, value, kwargs.clone())?; + self.async_set_cache(&key, value, kwargs.clone()).await?; } Ok(()) - }) + } } - fn batch_cache_write<'a>( - &'a self, - key: &'a str, + fn batch_cache_write( + &self, + key: &str, value: Self::Value, kwargs: CacheKwargs, - ) -> CacheFuture<'a, ()> { + ) -> impl Future> + Send { self.async_set_cache(key, value, kwargs) } fn delete_cache(&self, key: &str) -> Result<(), Error>; - fn async_delete_cache<'a>(&'a self, key: &'a str) -> CacheFuture<'a, ()> { - Box::pin(async move { self.delete_cache(key) }) + fn async_delete_cache(&self, key: &str) -> impl Future> + Send { + async move { self.delete_cache(key) } } fn flush_cache(&self) -> Result<(), Error>; - fn disconnect(&self) -> CacheFuture<'_, ()>; + fn disconnect(&self) -> impl Future> + Send; - fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult>; + fn test_connection(&self) -> impl Future> + Send; } diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs index 1aab6ee8e91..21ebcce29bb 100644 --- a/litellm-rust/crates/cache/src/caching.rs +++ b/litellm-rust/crates/cache/src/caching.rs @@ -146,21 +146,21 @@ impl CacheEntry { } } -pub fn get_cache( - cache: &dyn BaseCache, +pub fn get_cache( + cache: &B, key: &str, kwargs: &CacheKwargs, -) -> Result, Error> { +) -> Result, Error> { cache.get_cache(key, kwargs) } -pub fn set_cache( - cache: &dyn BaseCache, +pub fn set_cache( + cache: &B, key: &str, - entry: CacheEntry, + value: B::Value, kwargs: CacheKwargs, ) -> Result<(), Error> { - cache.set_cache(key, entry, kwargs) + cache.set_cache(key, value, kwargs) } -pub type CacheBackend = Arc>; +pub type CacheBackend = Arc; diff --git a/litellm-rust/crates/cache/src/codec.rs b/litellm-rust/crates/cache/src/codec.rs new file mode 100644 index 00000000000..09bee6032f6 --- /dev/null +++ b/litellm-rust/crates/cache/src/codec.rs @@ -0,0 +1,42 @@ +use std::marker::PhantomData; + +use serde::{Serialize, de::DeserializeOwned}; + +use crate::Error; + +pub trait CacheCodec: Send + Sync { + type Value: Clone + Send + Sync + 'static; + + fn encode(&self, value: &Self::Value) -> Result, Error>; + + fn decode(&self, bytes: &[u8]) -> Result; +} + +pub struct JsonCodec(PhantomData V>); + +impl Default for JsonCodec { + fn default() -> Self { + Self::new() + } +} + +impl JsonCodec { + pub const fn new() -> Self { + Self(PhantomData) + } +} + +impl CacheCodec for JsonCodec +where + V: Clone + Send + Sync + Serialize + DeserializeOwned + 'static, +{ + type Value = V; + + fn encode(&self, value: &Self::Value) -> Result, Error> { + serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) + } + + fn decode(&self, bytes: &[u8]) -> Result { + serde_json::from_slice(bytes).map_err(|_| Error::InvalidEntry) + } +} diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index d0fe3de15cd..a1d9d1402bb 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -1,12 +1,12 @@ mod base_cache; mod caching; +mod codec; mod error; -pub use base_cache::{ - BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheFuture, CacheKwargs, -}; +pub use base_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs}; pub use caching::{ Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache, }; +pub use codec::{CacheCodec, JsonCodec}; pub use error::Error; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 1192fc9a2b0..5c250c6b3c9 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,12 +1,13 @@ use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheFuture, CacheKeyContext, - CacheKeyField, CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key, + BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, + CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key, }; use sha2::{Digest, Sha256}; -use std::time::Duration; +use std::{sync::Mutex, time::Duration}; struct TestCache { default_ttl: Duration, + writes: Mutex>, } impl BaseCache for TestCache { @@ -17,6 +18,22 @@ impl BaseCache for TestCache { } fn set_cache(&self, _: &str, _: Self::Value, _: CacheKwargs) -> Result<(), Error> { + Err(Error::Unavailable) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + kwargs: CacheKwargs, + ) -> Result<(), Error> { + if key == "unavailable" { + return Err(Error::Unavailable); + } + self.writes + .lock() + .unwrap() + .push((key.into(), value, kwargs)); Ok(()) } @@ -32,11 +49,11 @@ impl BaseCache for TestCache { Ok(()) } - fn disconnect(&self) -> CacheFuture<'_, ()> { - Box::pin(async { Ok(()) }) + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) } - fn test_connection(&self) -> CacheFuture<'_, CacheConnectionResult> { + async fn test_connection(&self) -> Result { unreachable!() } } @@ -45,6 +62,7 @@ impl BaseCache for TestCache { fn ttl_uses_default_and_allows_per_call_override() { let cache = TestCache { default_ttl: Duration::from_secs(60), + writes: Mutex::default(), }; assert_eq!( cache.get_ttl(&CacheKwargs::default()), @@ -59,6 +77,46 @@ fn ttl_uses_default_and_allows_per_call_override() { ); } +#[tokio::test] +async fn default_batch_operations_use_async_writes_and_stop_on_failure() { + let cache = TestCache { + default_ttl: Duration::from_secs(60), + writes: Mutex::default(), + }; + let entry = CacheEntry { + timestamp: 123.0, + response: serde_json::json!("cached"), + }; + let kwargs = CacheKwargs { + ttl: Some(Duration::from_secs(5)), + ..Default::default() + }; + cache + .batch_cache_write("single", entry.clone(), kwargs.clone()) + .await + .unwrap(); + assert_eq!( + cache + .async_set_cache_pipeline( + vec![ + ("first".into(), entry.clone()), + ("unavailable".into(), entry.clone()), + ("skipped".into(), entry.clone()), + ], + kwargs.clone(), + ) + .await, + Err(Error::Unavailable) + ); + assert_eq!( + *cache.writes.lock().unwrap(), + vec![ + ("single".into(), entry.clone(), kwargs.clone()), + ("first".into(), entry, kwargs), + ] + ); +} + #[test] fn keys_match_python_order_groups_files_presets_and_namespaces() { let mut input = CacheKeyInput { diff --git a/litellm-rust/crates/cache/tests/codec.rs b/litellm-rust/crates/cache/tests/codec.rs new file mode 100644 index 00000000000..dad5398a879 --- /dev/null +++ b/litellm-rust/crates/cache/tests/codec.rs @@ -0,0 +1,53 @@ +use std::collections::BTreeMap; + +use litellm_cache::{CacheCodec, CacheEntry, Error, JsonCodec}; +use serde::{Deserialize, Serialize}; +use serde_json::json; + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +struct RoutingState { + deployment: String, + cooldown_seconds: u64, +} + +#[test] +fn json_codec_round_trips_typed_domain_values() { + let codec = JsonCodec::::new(); + let value = RoutingState { + deployment: "deployment-a".into(), + cooldown_seconds: 30, + }; + let bytes = codec.encode(&value).unwrap(); + assert_eq!(codec.decode(&bytes).unwrap(), value); + assert_eq!( + serde_json::from_slice::(&bytes).unwrap(), + json!({"deployment": "deployment-a", "cooldown_seconds": 30}) + ); +} + +#[test] +fn response_entries_preserve_the_existing_json_representation() { + let codec = JsonCodec::::new(); + let entry = CacheEntry { + timestamp: 123.0, + response: json!({"choices": [{"text": "cached"}]}), + }; + let bytes = codec.encode(&entry).unwrap(); + assert_eq!(bytes, serde_json::to_vec(&entry).unwrap()); + assert_eq!(codec.decode(&bytes).unwrap(), entry); +} + +#[test] +fn json_codec_rejects_malformed_and_wrongly_typed_entries() { + let codec = JsonCodec::::new(); + for bytes in [b"not json".as_slice(), br#"{"deployment":12}"#.as_slice()] { + assert_eq!(codec.decode(bytes).unwrap_err(), Error::InvalidEntry); + } +} + +#[test] +fn json_codec_propagates_encoding_errors() { + let codec = JsonCodec::>::new(); + let value = BTreeMap::from([((1, 2), "invalid JSON object key".into())]); + assert_eq!(codec.encode(&value).unwrap_err(), Error::InvalidEntry); +}