From 95cf7066d16186e94a8f27828ac38db55ef45cf7 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 20:38:13 -0700 Subject: [PATCH 01/17] refactor(cache): use static dispatch and typed backend codecs --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/cache-memory/src/cache.rs | 23 +-- .../crates/cache-memory/tests/cache.rs | 47 ++++- litellm-rust/crates/cache-redis/Cargo.toml | 2 +- litellm-rust/crates/cache-redis/src/cache.rs | 171 ++++++++---------- .../crates/cache-redis/tests/cache.rs | 152 +++++++++++++++- litellm-rust/crates/cache/Cargo.toml | 1 + litellm-rust/crates/cache/src/base_cache.rs | 53 +++--- litellm-rust/crates/cache/src/caching.rs | 16 +- litellm-rust/crates/cache/src/codec.rs | 42 +++++ litellm-rust/crates/cache/src/lib.rs | 6 +- litellm-rust/crates/cache/tests/caching.rs | 70 ++++++- litellm-rust/crates/cache/tests/codec.rs | 53 ++++++ 13 files changed, 480 insertions(+), 157 deletions(-) create mode 100644 litellm-rust/crates/cache/src/codec.rs create mode 100644 litellm-rust/crates/cache/tests/codec.rs 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); +} From 081c93908f1e2750c37beb0aa4e87660aa983d4e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 20:55:45 -0700 Subject: [PATCH 02/17] feat(cache): add native response cache and Python binding foundations --- litellm-rust/Cargo.lock | 100 ++++- litellm-rust/Cargo.toml | 2 + litellm-rust/crates/cache-redis/src/cache.rs | 83 +++-- .../crates/cache-redis/tests/cache.rs | 78 ++-- litellm-rust/crates/cache-response/Cargo.toml | 19 + .../crates/cache-response/src/codec.rs | 100 +++++ litellm-rust/crates/cache-response/src/lib.rs | 7 + .../crates/cache-response/src/native.rs | 97 +++++ .../crates/cache-response/src/response.rs | 124 +++++++ .../crates/cache-response/tests/response.rs | 290 +++++++++++++++ litellm-rust/crates/cache/src/caching.rs | 1 + litellm-rust/crates/cache/src/error.rs | 2 + litellm-rust/crates/python-bridge/Cargo.toml | 3 + .../crates/python-bridge/src/cache/facade.rs | 210 +++++++++++ .../crates/python-bridge/src/cache/mod.rs | 350 ++++++++++++++++++ litellm-rust/crates/python-bridge/src/lib.rs | 6 + litellm/rust_bridge/_native.pyi | 49 ++- tests/test_litellm_rust/test_cache.py | 231 ++++++++++++ 18 files changed, 1702 insertions(+), 50 deletions(-) create mode 100644 litellm-rust/crates/cache-response/Cargo.toml create mode 100644 litellm-rust/crates/cache-response/src/codec.rs create mode 100644 litellm-rust/crates/cache-response/src/lib.rs create mode 100644 litellm-rust/crates/cache-response/src/native.rs create mode 100644 litellm-rust/crates/cache-response/src/response.rs create mode 100644 litellm-rust/crates/cache-response/tests/response.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/facade.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/mod.rs create mode 100644 tests/test_litellm_rust/test_cache.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ccfaee44f50..aa62e4f3770 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2486,6 +2486,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-response" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "litellm-cache-memory", + "litellm-cache-redis", + "py_literal", + "redis", + "redis-test", + "serde", + "serde_json", + "tokio", +] + [[package]] name = "litellm-callbacks-legacy-python" version = "0.1.0" @@ -2649,6 +2664,8 @@ dependencies = [ "futures-util", "litellm-auth", "litellm-auth-gcp", + "litellm-cache", + "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core", "litellm-core-utils", @@ -2660,6 +2677,7 @@ dependencies = [ "pyo3", "pyo3-async-runtimes", "rstest", + "serde", "serde_json", "tokio", "tokio-tungstenite", @@ -2948,6 +2966,16 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "num-bigint" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c89e69e7e0f03bea5ef08013795c25018e101932225a656383bd384495ecc367" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-bigint" version = "0.5.1" @@ -2958,6 +2986,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + [[package]] name = "num-conv" version = "0.2.2" @@ -3131,6 +3168,48 @@ version = "2.3.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" +[[package]] +name = "pest" +version = "2.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d45aeb61b4bf818e12d4205f2466f8c4748f85f4fce0146d1c03d69d753f0ad" +dependencies = [ + "memchr", + "ucd-trie", +] + +[[package]] +name = "pest_derive" +version = "2.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89cc5a242e25ed4e7704d0be240f2cfbe20a8c27e7e252d94835be93d92dc39f" +dependencies = [ + "pest", + "pest_generator", +] + +[[package]] +name = "pest_generator" +version = "2.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7abf21475cc3820fe4b2ca2dc2142902f67a02189f3b5b3a229f4febc01a43e5" +dependencies = [ + "pest", + "pest_meta", + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "pest_meta" +version = "2.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adba4db388f687393c18c51348d44a41d870ca9df71a2c98172ea3035dc6936e" +dependencies = [ + "pest", +] + [[package]] name = "pin-project" version = "1.1.13" @@ -3305,6 +3384,19 @@ dependencies = [ "prost", ] +[[package]] +name = "py_literal" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "102df7a3d46db9d3891f178dcc826dc270a6746277a9ae6436f8d29fd490a8e1" +dependencies = [ + "num-bigint 0.4.8", + "num-complex", + "num-traits", + "pest", + "pest_derive", +] + [[package]] name = "pyo3" version = "0.29.2" @@ -3604,7 +3696,7 @@ dependencies = [ "arcstr", "combine", "itoa", - "num-bigint", + "num-bigint 0.5.1", "percent-encoding", "ryu", "sha1_smol", @@ -4955,6 +5047,12 @@ dependencies = [ "syn 2.0.119", ] +[[package]] +name = "ucd-trie" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" + [[package]] name = "unarray" version = "0.1.4" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 2f6f5feb4ad..570d0dd3568 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -28,6 +28,8 @@ litellm-types = { path = "crates/types" } litellm-core-utils = { path = "crates/core-utils" } litellm-cache = { path = "crates/cache" } litellm-cache-memory = { path = "crates/cache-memory" } +litellm-cache-redis = { path = "crates/cache-redis" } +litellm-cache-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" } diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 8d92fdc8f75..0faca6cdaaf 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -7,12 +7,12 @@ use litellm_cache::{ use redis::Commands; const DEFAULT_TTL: Duration = Duration::from_secs(600); -const KEY_PREFIX: &str = "litellm-cache:"; pub struct RedisCache { connection: Arc>, default_ttl: Duration, codec: S, + namespace: Option, } impl RedisCache { @@ -33,6 +33,7 @@ where connection: Arc::new(Mutex::new(connection)), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, + namespace: None, } } @@ -40,13 +41,44 @@ where self.connection.lock().map_err(|_| Error::Unavailable) } - fn namespaced_key(key: &str) -> String { - format!("{KEY_PREFIX}{key}") + pub fn with_namespace(self, namespace: Option) -> Self { + Self { + namespace: namespace.filter(|value| !value.is_empty()), + ..self + } } - fn namespaced_pattern() -> &'static str { - const PATTERN: &str = "litellm-cache:*"; - PATTERN + fn namespaced_key(&self, key: &str) -> String { + match &self.namespace { + Some(namespace) if !key.starts_with(&format!("{namespace}:")) => { + format!("{namespace}:{key}") + } + _ => key.into(), + } + } + + fn namespaced_pattern(&self) -> Result { + let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?; + let escaped: String = namespace + .chars() + .flat_map(|ch| { + if matches!(ch, '*' | '?' | '[' | ']' | '\\') { + vec!['\\', ch] + } else { + vec![ch] + } + }) + .collect(); + Ok(format!("{escaped}:*")) + } + + fn decode_response(&self, value: redis::Value) -> Result, Error> { + match value { + redis::Value::Nil => Ok(None), + redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some), + redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some), + _ => Err(Error::InvalidEntry), + } } fn ttl_seconds(ttl: Duration) -> u64 { @@ -84,28 +116,29 @@ where 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) + .set_ex::<_, _, ()>(self.namespaced_key(key), payload, ttl) .map_err(|_| Error::Unavailable) } fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { - let bytes = self + let value = self .connection()? - .get::<_, Option>>(Self::namespaced_key(key)) + .get::<_, redis::Value>(self.namespaced_key(key)) .map_err(|_| Error::Unavailable)?; - bytes.map(|bytes| self.codec.decode(&bytes)).transpose() + self.decode_response(value) } fn delete_cache(&self, key: &str) -> Result<(), Error> { self.connection()? - .del::<_, ()>(Self::namespaced_key(key)) + .del::<_, ()>(self.namespaced_key(key)) .map_err(|_| Error::Unavailable) } fn flush_cache(&self) -> Result<(), Error> { + let pattern = self.namespaced_pattern()?; let mut connection = self.connection()?; let keys = connection - .scan_match(Self::namespaced_pattern()) + .scan_match(pattern) .map_err(|_| Error::Unavailable)? .collect::>>() .map_err(|_| Error::Unavailable)?; @@ -125,7 +158,7 @@ where kwargs: CacheKwargs, ) -> Result<(), Error> { let payload = self.codec.encode(&value)?; - let key = Self::namespaced_key(key); + 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 @@ -140,15 +173,14 @@ where key: &str, _: &CacheKwargs, ) -> Result, Error> { - let key = Self::namespaced_key(key); - Self::run_blocking(Arc::clone(&self.connection), move |connection| { + let key = self.namespaced_key(key); + let value = Self::run_blocking(Arc::clone(&self.connection), move |connection| { connection - .get::<_, Option>>(key) + .get::<_, redis::Value>(key) .map_err(|_| Error::Unavailable) }) - .await? - .map(|bytes| self.codec.decode(&bytes)) - .transpose() + .await?; + self.decode_response(value) } async fn async_set_cache_pipeline( @@ -161,7 +193,7 @@ where .map(|(key, value)| { self.codec .encode(&value) - .map(|payload| (Self::namespaced_key(&key), payload)) + .map(|payload| (self.namespaced_key(&key), payload)) }) .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); @@ -177,7 +209,7 @@ where } async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { - let key = Self::namespaced_key(key); + let key = self.namespaced_key(key); Self::run_blocking(Arc::clone(&self.connection), move |connection| { connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) }) @@ -250,7 +282,8 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); cache .set_cache("key", value.clone(), CacheKwargs::default()) @@ -275,7 +308,8 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); cache.flush_cache().unwrap(); } @@ -284,7 +318,8 @@ mod tests { async fn test_connection_runs_ping_off_executor() { let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); 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 fe15fcd975b..d5bba19a8bd 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -34,15 +34,12 @@ fn generic_helpers_use_the_injected_codec_and_ttl() { let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("SETEX") - .arg("litellm-cache:counter") + .arg("counter") .arg(2) .arg([42u8, 7].as_slice()), Ok("OK"), ), - MockCmd::new( - redis::cmd("GET").arg("litellm-cache:counter"), - Ok(vec![42u8, 7]), - ), + MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])), ]) .assert_all_commands_consumed(); let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); @@ -59,27 +56,21 @@ async fn async_operations_preserve_codec_ttl_and_missing_values() { let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("SETEX") - .arg("litellm-cache:counter") + .arg("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("GET").arg("counter"), Ok(vec![42u8, 7])), MockCmd::new( redis::cmd("SETEX") - .arg("litellm-cache:batch") + .arg("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), - ), + MockCmd::new(redis::cmd("DEL").arg("counter"), Ok(1u32)), + MockCmd::new(redis::cmd("GET").arg("counter"), Ok(redis::Value::Nil)), ]) .assert_all_commands_consumed(); let cache = RedisCache::with_connection( @@ -116,14 +107,8 @@ async fn async_operations_preserve_codec_ttl_and_missing_values() { #[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]), - ), + MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])), + MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])), ]) .assert_all_commands_consumed(); let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); @@ -154,3 +139,48 @@ async fn codec_errors_propagate_without_writing_partial_batches() { Err(Error::InvalidEntry) ); } + +#[test] +fn namespaces_are_optional_and_existing_prefixes_are_not_duplicated() { + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)), + MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + assert_eq!( + cache.get_cache("key", &CacheKwargs::default()).unwrap(), + None + ); + assert_eq!( + cache + .get_cache("team:key", &CacheKwargs::default()) + .unwrap(), + None + ); +} + +#[test] +fn flush_requires_a_namespace_and_escapes_glob_metacharacters() { + let unscoped = RedisCache::with_connection( + MockRedisConnection::new([]).assert_all_commands_consumed(), + None, + JsonCodec::::new(), + ); + assert_eq!(unscoped.flush_cache(), Err(Error::UnscopedFlush)); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(0) + .arg("MATCH") + .arg("team\\*:*"), + Ok(redis_test::redis_value!(["0", ["team*:key"]])), + ), + MockCmd::new(redis::cmd("DEL").arg("team*:key"), Ok(1u32)), + ]) + .assert_all_commands_consumed(); + let scoped = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team*".into())); + scoped.flush_cache().unwrap(); +} diff --git a/litellm-rust/crates/cache-response/Cargo.toml b/litellm-rust/crates/cache-response/Cargo.toml new file mode 100644 index 00000000000..a0c4a1f74ef --- /dev/null +++ b/litellm-rust/crates/cache-response/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "litellm-cache-response" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-cache-memory.workspace = true +litellm-cache-redis.workspace = true +py_literal = "0.4.0" +redis = "1.7.0" +serde.workspace = true +serde_json.workspace = true + +[dev-dependencies] +redis-test = "1.0.4" +tokio.workspace = true diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs new file mode 100644 index 00000000000..137cf61267a --- /dev/null +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -0,0 +1,100 @@ +use litellm_cache::{CacheCodec, CacheEntry, Error}; +use serde_json::Value; + +pub struct ResponseCacheCodec; + +impl CacheCodec for ResponseCacheCodec { + type Value = CacheEntry; + + fn encode(&self, value: &CacheEntry) -> Result, Error> { + if !value.timestamp.is_finite() { + return Err(Error::InvalidEntry); + } + serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) + } + + fn decode(&self, bytes: &[u8]) -> Result { + let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?; + let entry: CacheEntry = + serde_json::from_value(decode_value(text)?).map_err(|_| Error::InvalidEntry)?; + if !entry.timestamp.is_finite() { + return Err(Error::InvalidEntry); + } + Ok(entry) + } +} + +pub(crate) fn decode_value(text: &str) -> Result { + if let Ok(value) = serde_json::from_str(text) { + return Ok(value); + } + check_literal_depth(text)?; + let literal: py_literal::Value = text.parse().map_err(|_| Error::InvalidEntry)?; + literal_value(literal, 0) +} + +fn literal_value(value: py_literal::Value, depth: usize) -> Result { + use py_literal::Value as Literal; + if depth > 128 { + return Err(Error::InvalidEntry); + } + match value { + Literal::String(text) => Ok(Value::String(text)), + Literal::Boolean(value) => Ok(Value::Bool(value)), + Literal::None => Ok(Value::Null), + Literal::Integer(value) => { + serde_json::from_str(&value.to_string()).map_err(|_| Error::InvalidEntry) + } + Literal::Float(value) => serde_json::Number::from_f64(value) + .map(Value::Number) + .ok_or(Error::InvalidEntry), + Literal::List(values) | Literal::Tuple(values) => values + .into_iter() + .map(|value| literal_value(value, depth + 1)) + .collect::, _>>() + .map(Value::Array), + Literal::Dict(entries) => entries + .into_iter() + .map(|(key, value)| { + let Literal::String(key) = key else { + return Err(Error::InvalidEntry); + }; + Ok((key, literal_value(value, depth + 1)?)) + }) + .collect::, _>>() + .map(Value::Object), + _ => Err(Error::InvalidEntry), + } +} + +fn check_literal_depth(text: &str) -> Result<(), Error> { + let mut quote = None; + let mut escaped = false; + let mut depth = 0usize; + for ch in text.chars() { + if escaped { + escaped = false; + continue; + } + if let Some(delimiter) = quote { + if ch == '\\' { + escaped = true; + } else if ch == delimiter { + quote = None; + } + continue; + } + match ch { + '\'' | '"' => quote = Some(ch), + '[' | '{' | '(' => { + depth += 1; + if depth > 128 { + return Err(Error::InvalidEntry); + } + } + ']' | '}' | ')' => depth = depth.saturating_sub(1), + _ => {} + } + } + Ok(()) +} diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs new file mode 100644 index 00000000000..454a82e76de --- /dev/null +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -0,0 +1,7 @@ +mod codec; +mod native; +mod response; + +pub use codec::ResponseCacheCodec; +pub use native::NativeResponseCache; +pub use response::{ResponseCache, ResponseCacheRequest}; diff --git a/litellm-rust/crates/cache-response/src/native.rs b/litellm-rust/crates/cache-response/src/native.rs new file mode 100644 index 00000000000..c25c61cee82 --- /dev/null +++ b/litellm-rust/crates/cache-response/src/native.rs @@ -0,0 +1,97 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::{CacheEntry, Error}; +use litellm_cache_memory::InMemoryCache; +use litellm_cache_redis::RedisCache; +use serde_json::Value; + +use crate::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; + +pub enum NativeResponseCache +where + C: redis::ConnectionLike + Send + 'static, +{ + Memory(Arc>>), + Redis(Arc>>), +} + +impl Clone for NativeResponseCache { + fn clone(&self) -> Self { + match self { + Self::Memory(cache) => Self::Memory(Arc::clone(cache)), + Self::Redis(cache) => Self::Redis(Arc::clone(cache)), + } + } +} + +impl NativeResponseCache { + pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { + Self::Memory(Arc::new(ResponseCache::new(Arc::new( + InMemoryCache::response_cache(capacity, ttl, max_entry_bytes), + )))) + } + + pub fn redis( + url: &str, + ttl: Option, + namespace: Option, + ) -> Result { + let backend = RedisCache::new(url, ttl, ResponseCacheCodec)?.with_namespace(namespace); + Ok(Self::Redis(Arc::new(ResponseCache::new(Arc::new(backend))))) + } +} + +impl NativeResponseCache { + pub fn kind(&self) -> &'static str { + match self { + Self::Memory(_) => "memory", + Self::Redis(_) => "redis", + } + } + + pub fn lookup( + &self, + request: &ResponseCacheRequest, + now: Duration, + ) -> Result, Error> { + match self { + Self::Memory(cache) => cache.lookup(request, now), + Self::Redis(cache) => cache.lookup(request, now), + } + } + + pub fn store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + match self { + Self::Memory(cache) => cache.store(request, response, now), + Self::Redis(cache) => cache.store(request, response, now), + } + } + + pub async fn async_lookup( + &self, + request: &ResponseCacheRequest, + now: Duration, + ) -> Result, Error> { + match self { + Self::Memory(cache) => cache.async_lookup(request, now).await, + Self::Redis(cache) => cache.async_lookup(request, now).await, + } + } + + pub async fn async_store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + match self { + Self::Memory(cache) => cache.async_store(request, response, now).await, + Self::Redis(cache) => cache.async_store(request, response, now).await, + } + } +} diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs new file mode 100644 index 00000000000..3e9807b6d1b --- /dev/null +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -0,0 +1,124 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::{ + BaseCache, CacheControls, CacheEntry, CacheKeyInput, CacheKwargs, Error, cache_key, +}; +use serde_json::Value; + +#[derive(Clone)] +pub struct ResponseCacheRequest { + pub key: CacheKeyInput, + pub controls: CacheControls, + pub kwargs: CacheKwargs, + pub max_age: Option, +} + +impl ResponseCacheRequest { + pub fn new(key: CacheKeyInput) -> Self { + Self { + key, + controls: CacheControls { + configured: true, + supported_call_type: true, + native_backend: true, + default_on: true, + ..Default::default() + }, + kwargs: CacheKwargs::default(), + max_age: None, + } + } +} + +pub struct ResponseCache> { + backend: Arc, +} + +impl> ResponseCache { + pub fn new(backend: Arc) -> Self { + Self { backend } + } + + pub fn lookup( + &self, + request: &ResponseCacheRequest, + now: Duration, + ) -> Result, Error> { + if !request.controls.reads() { + return Ok(None); + } + let entry = self + .backend + .get_cache(&cache_key(&request.key), &request.kwargs)?; + Self::fresh_response(entry, now, request.max_age) + } + + pub async fn async_lookup( + &self, + request: &ResponseCacheRequest, + now: Duration, + ) -> Result, Error> { + if !request.controls.reads() { + return Ok(None); + } + let entry = self + .backend + .async_get_cache(&cache_key(&request.key), &request.kwargs) + .await?; + Self::fresh_response(entry, now, request.max_age) + } + + pub fn store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + if !request.controls.writes() { + return Ok(()); + } + self.backend.set_cache( + &cache_key(&request.key), + CacheEntry { + timestamp: now.as_secs_f64(), + response, + }, + request.kwargs.clone(), + ) + } + + pub async fn async_store( + &self, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + if !request.controls.writes() { + return Ok(()); + } + self.backend + .async_set_cache( + &cache_key(&request.key), + CacheEntry { + timestamp: now.as_secs_f64(), + response, + }, + request.kwargs.clone(), + ) + .await + } + + fn fresh_response( + entry: Option, + now: Duration, + max_age: Option, + ) -> Result, Error> { + entry + .filter(|entry| entry.fresh(now, max_age)) + .map(|entry| match entry.response { + Value::String(text) => crate::codec::decode_value(&text), + value => Ok(value), + }) + .transpose() + } +} diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs new file mode 100644 index 00000000000..a4beda9fa4c --- /dev/null +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -0,0 +1,290 @@ +use std::{ + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; + +use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKeyField, CacheKeyInput, Error}; +use litellm_cache_memory::InMemoryCache; +use litellm_cache_redis::RedisCache; +use litellm_cache_response::{ + NativeResponseCache, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, +}; +use redis_test::{MockCmd, MockRedisConnection}; +use serde_json::json; + +fn request() -> ResponseCacheRequest { + ResponseCacheRequest::new(CacheKeyInput { + preset: Some("tenant:key".into()), + ..Default::default() + }) +} + +#[tokio::test] +async fn sync_and_async_consumers_share_keys_ttls_and_freshness() { + let clock = Arc::new(AtomicU64::new(100)); + let backend = Arc::new(InMemoryCache::with_clock( + Some(8), + Some(Duration::from_secs(600)), + { + let clock = clock.clone(); + move || Duration::from_secs(clock.load(Ordering::SeqCst)) + }, + )); + let cache = ResponseCache::new(backend.clone()); + let mut request = request(); + request.kwargs.ttl = Some(Duration::from_secs(10)); + request.max_age = Some(Duration::from_secs(5)); + cache + .store( + &request, + json!({"choices": [1], "usage": {"total_tokens": 7}}), + Duration::from_secs(100), + ) + .unwrap(); + assert_eq!( + backend.expires_at("tenant:key").unwrap(), + Some(Duration::from_secs(110)) + ); + assert!( + cache + .async_lookup(&request, Duration::from_secs(105)) + .await + .unwrap() + .is_some() + ); + assert_eq!( + cache.lookup(&request, Duration::from_secs(106)).unwrap(), + None + ); + request.max_age = None; + assert_eq!( + cache + .lookup(&request, Duration::from_secs(106)) + .unwrap() + .unwrap()["usage"]["total_tokens"], + 7 + ); + clock.store(111, Ordering::SeqCst); + assert_eq!( + cache + .async_lookup(&request, Duration::from_secs(111)) + .await + .unwrap(), + None + ); + cache + .async_store(&request, json!({"choices": [2]}), Duration::from_secs(111)) + .await + .unwrap(); + assert_eq!( + cache.lookup(&request, Duration::from_secs(111)).unwrap(), + Some(json!({"choices": [2]})) + ); +} + +#[tokio::test] +async fn directives_skip_io_and_keep_reads_and_writes_independent() { + let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let mut request = request(); + let now = Duration::from_secs(100); + request.controls.no_store = true; + cache + .async_store(&request, json!({"v": 1}), now) + .await + .unwrap(); + assert_eq!(cache.lookup(&request, now).unwrap(), None); + request.controls.no_store = false; + request.controls.no_cache = true; + cache.store(&request, json!({"v": 2}), now).unwrap(); + assert_eq!(cache.async_lookup(&request, now).await.unwrap(), None); + request.controls.no_cache = false; + assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2}))); + request.controls.default_on = false; + cache.store(&request, json!({"v": 3}), now).unwrap(); + assert_eq!(cache.lookup(&request, now).unwrap(), None); + request.controls.use_cache = true; + assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2}))); + request.controls.supported_call_type = false; + assert_eq!(cache.lookup(&request, now).unwrap(), None); +} + +#[tokio::test] +async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_json() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("GET").arg("tenant:key"), + Ok(br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.to_vec()), + ), + MockCmd::new( + redis::cmd("GET").arg("tenant:key"), + Ok(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.to_vec()), + ), + MockCmd::new( + redis::cmd("SETEX") + .arg("tenant:key") + .arg(600) + .arg(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()), + Ok("OK"), + ), + ]) + .assert_all_commands_consumed(); + let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec) + .with_namespace(Some("tenant".into())); + let cache = NativeResponseCache::Redis(Arc::new(ResponseCache::new(Arc::new(backend)))); + let request = request(); + let expected = json!({"ok": true, "text": "cached"}); + assert_eq!( + cache.lookup(&request, Duration::from_secs(101)).unwrap(), + Some(expected.clone()) + ); + assert_eq!( + cache + .async_lookup(&request, Duration::from_secs(101)) + .await + .unwrap(), + Some(expected.clone()) + ); + cache + .async_store(&request, expected, Duration::from_secs(100)) + .await + .unwrap(); +} + +#[tokio::test] +async fn captured_enum_keeps_the_selected_backend_for_background_writes() { + let original = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let captured = original.clone(); + let replacement = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let request = request(); + let writer = tokio::spawn({ + let request = request.clone(); + async move { + captured + .async_store( + &request, + json!({"selected": "original"}), + Duration::from_secs(100), + ) + .await + } + }); + writer.await.unwrap().unwrap(); + assert_eq!( + original.lookup(&request, Duration::from_secs(100)).unwrap(), + Some(json!({"selected":"original"})) + ); + assert_eq!( + replacement + .lookup(&request, Duration::from_secs(100)) + .unwrap(), + None + ); +} + +#[test] +fn generated_keys_preserve_namespace_and_explicit_keys() { + let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let key = CacheKeyInput { + fields: vec![CacheKeyField { + name: "model".into(), + value: Some("a".into()), + api_parameter: true, + internal_parameter: false, + }], + namespace: Some("tenant".into()), + ..Default::default() + }; + let generated = ResponseCacheRequest::new(key.clone()); + let explicit = ResponseCacheRequest::new(CacheKeyInput { + preset: Some(litellm_cache::cache_key(&key)), + ..Default::default() + }); + cache + .store(&generated, json!({"value": 7}), Duration::from_secs(100)) + .unwrap(); + assert_eq!( + cache.lookup(&explicit, Duration::from_secs(100)).unwrap(), + Some(json!({"value":7})) + ); +} + +#[test] +fn response_codec_accepts_python_literals_without_executing_code() { + let bytes = br#"{'timestamp': 100.0, 'response': {'text': 'hello \\ world', 'flag': True, 'empty': None, 'list': [1, 2.5]}}"#; + let entry = ResponseCacheCodec.decode(bytes).unwrap(); + assert_eq!( + entry.response, + json!({"text": "hello \\ world", "flag": true, "empty": null, "list": [1, 2.5]}) + ); + for bytes in [ + b"__import__('os').system('false')".as_slice(), + b"{'timestamp': 'invalid', 'response': {}}", + b"{'timestamp': 1e9999, 'response': {}}", + ] { + assert_eq!( + ResponseCacheCodec.decode(bytes).unwrap_err(), + Error::InvalidEntry + ); + } + let deep = format!("{}None{}", "[".repeat(1000), "]".repeat(1000)); + assert_eq!( + ResponseCacheCodec.decode(deep.as_bytes()).unwrap_err(), + Error::InvalidEntry + ); + assert_eq!( + ResponseCacheCodec + .encode(&CacheEntry { + timestamp: f64::NAN, + response: json!({}) + }) + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[tokio::test] +async fn backend_failures_remain_observable_and_disabled_reads_do_not_touch_redis() { + let connection = MockRedisConnection::new([MockCmd::new( + redis::cmd("GET").arg("tenant:key"), + Ok(b"invalid".to_vec()), + )]) + .assert_all_commands_consumed(); + let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec); + let cache = ResponseCache::new(Arc::new(backend)); + let mut request = request(); + request.controls.no_cache = true; + assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None); + request.controls.no_cache = false; + assert_eq!( + cache + .async_lookup(&request, Duration::ZERO) + .await + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[test] +fn malformed_memory_entries_are_rejected_by_the_response_consumer() { + let backend = Arc::new(InMemoryCache::default()); + BaseCache::set_cache( + backend.as_ref(), + "tenant:key", + CacheEntry { + timestamp: 100.0, + response: json!("not a serialized response"), + }, + Default::default(), + ) + .unwrap(); + let cache = ResponseCache::new(backend); + assert_eq!( + cache + .lookup(&request(), Duration::from_secs(100)) + .unwrap_err(), + Error::InvalidEntry + ); +} diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs index 21ebcce29bb..1d1df966ba2 100644 --- a/litellm-rust/crates/cache/src/caching.rs +++ b/litellm-rust/crates/cache/src/caching.rs @@ -27,6 +27,7 @@ pub struct CacheKeyField { } #[derive(Clone, Debug, Default, Deserialize, Serialize)] +#[serde(default)] pub struct CacheKeyInput { pub fields: Vec, pub preset: Option, diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index d447c80f62d..ff3ff6572d4 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -4,4 +4,6 @@ pub enum Error { Unavailable, #[error("invalid cache entry")] InvalidEntry, + #[error("flushing Redis requires an explicit namespace")] + UnscopedFlush, } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index a76b069935f..308dfd2dd7a 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -20,6 +20,9 @@ tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true +litellm-cache.workspace = true +litellm-cache-response.workspace = true +serde.workspace = true litellm-auth.workspace = true litellm-callbacks-legacy-python.workspace = true litellm-core.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs new file mode 100644 index 00000000000..32360e72469 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -0,0 +1,210 @@ +use litellm_cache_response::NativeResponseCache; +use litellm_host_python::from_py; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::PyTypeError, + prelude::*, + types::{PyDict, PyTuple, PyType}, +}; +use serde_json::Value; + +use super::NativeCacheHandle; + +struct ClassGuard { + class: Py, + attributes: Vec<(String, Py)>, +} + +struct ObjectGuard { + reference: Py, + classes: Vec, + config_names: &'static [&'static str], + config: Vec, +} + +pub(super) struct FacadeGuard { + outer: ObjectGuard, + backend: ObjectGuard, +} + +impl ObjectGuard { + fn capture( + py: Python<'_>, + object: &Bound<'_, PyAny>, + config_names: &'static [&'static str], + ) -> PyResult { + let classes = object + .get_type() + .getattr("__mro__")? + .cast_into::()? + .iter() + .map(|class| { + let class = class.cast_into::()?; + let attributes = class + .getattr("__dict__")? + .call_method0("items")? + .try_iter()? + .map(|item| item?.extract::<(String, Py)>()) + .collect::>>()?; + Ok(ClassGuard { + class: class.unbind(), + attributes, + }) + }) + .collect::>>()?; + let guard = Self { + reference: py + .import("weakref")? + .getattr("ref")? + .call1((object,))? + .unbind(), + classes, + config_names, + config: Self::config(object, config_names)?, + }; + if !guard.matches(py, object)? { + return Err(PyTypeError::new_err( + "native facade registration requires unmodified built-in methods", + )); + } + Ok(guard) + } + + fn config(object: &Bound<'_, PyAny>, names: &[&str]) -> PyResult> { + names + .iter() + .map(|name| match object.getattr(*name) { + Ok(value) => from_py(&value), + Err(error) + if error.is_instance_of::(object.py()) => + { + Ok(Value::Null) + } + Err(error) => Err(error), + }) + .collect() + } + + fn matches(&self, py: Python<'_>, object: &Bound<'_, PyAny>) -> PyResult { + if !self.reference.bind(py).call0()?.is(object) { + return Ok(false); + } + let mro = object + .get_type() + .getattr("__mro__")? + .cast_into::()?; + if mro.len() != self.classes.len() { + return Ok(false); + } + let instance = object.getattr("__dict__")?.cast_into::()?; + for (class, expected) in mro.iter().zip(&self.classes) { + if !class.is(expected.class.bind(py)) { + return Ok(false); + } + let attributes = class.getattr("__dict__")?; + if attributes.len()? != expected.attributes.len() { + return Ok(false); + } + for (name, value) in &expected.attributes { + if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) { + return Ok(false); + } + } + } + Ok(Self::config(object, self.config_names)? == self.config) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reference)?; + for class in &self.classes { + visit.call(&class.class)?; + for (_, value) in &class.attributes { + visit.call(value)?; + } + } + Ok(()) + } +} + +impl FacadeGuard { + pub(super) fn capture(py: Python<'_>, facade: &Bound<'_, PyAny>, kind: &str) -> PyResult { + let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?; + if !facade.get_type().is(&cache_type) { + return Err(PyTypeError::new_err( + "only exact built-in Cache facades can be registered", + )); + } + let (module, name, cache_kind) = match kind { + "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), + "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"), + _ => unreachable!(), + }; + let backend = facade.getattr("cache")?; + if facade.getattr("type")?.extract::()? != cache_kind + || !backend.get_type().is(&py.import(module)?.getattr(name)?) + { + return Err(PyTypeError::new_err( + "facade and native backend types must match", + )); + } + Ok(Self { + outer: ObjectGuard::capture( + py, + facade, + &[ + "type", + "mode", + "ttl", + "namespace", + "supported_call_types", + "redis_flush_size", + ], + )?, + backend: ObjectGuard::capture( + py, + &backend, + &[ + "namespace", + "default_ttl", + "max_size_in_memory", + "max_size_per_item", + ], + )?, + }) + } + + fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult { + Ok(self.outer.matches(py, facade)? + && self.backend.matches(py, &facade.getattr("cache")?)?) + } + + pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + self.outer.traverse(&visit)?; + self.backend.traverse(&visit) + } +} + +pub(super) fn resolve( + py: Python<'_>, + facade: &Bound<'_, PyAny>, +) -> PyResult> { + let Ok(dict) = facade + .getattr("__dict__") + .and_then(|dict| dict.cast_into::().map_err(Into::into)) + else { + return Ok(None); + }; + let Some(handle) = dict.get_item("_native_cache_handle")? else { + return Ok(None); + }; + let Ok(handle) = handle.extract::>() else { + return Ok(None); + }; + let Some(guard) = &handle.guard else { + return Ok(None); + }; + if !guard.matches(py, facade).unwrap_or(false) { + return Ok(None); + } + handle.service().map(Some) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs new file mode 100644 index 00000000000..f00cceeb86e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -0,0 +1,350 @@ +mod facade; + +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use litellm_cache::{CacheControls, CacheKeyInput, Error}; +use litellm_cache_response::{NativeResponseCache, ResponseCacheRequest}; +use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyRuntimeError, PyTypeError, PyValueError}, + prelude::*, + types::PyDict, +}; +use serde::Deserialize; +use serde_json::Value; + +use facade::FacadeGuard; + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RequestInput { + key: CacheKeyInput, + controls: Option, + ttl_seconds: Option, + max_age_seconds: Option, +} + +fn request(value: &Bound<'_, PyAny>) -> PyResult { + let input: RequestInput = from_py(value)?; + let mut request = ResponseCacheRequest::new(input.key); + if let Some(controls) = input.controls { + request.controls = controls; + } + request.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?; + request.max_age = input.max_age_seconds.map(duration).transpose()?; + Ok(request) +} + +fn duration(seconds: f64) -> PyResult { + Duration::try_from_secs_f64(seconds) + .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) +} + +fn now() -> Duration { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() +} + +fn cache_error(error: Error) -> PyErr { + match error { + Error::InvalidEntry => PyValueError::new_err(error.to_string()), + _ => PyRuntimeError::new_err(error.to_string()), + } +} + +#[pyclass(frozen)] +pub(crate) struct NativeCacheHandle { + service: NativeResponseCache, + guard: Option, + pid: u32, +} + +impl NativeCacheHandle { + fn service(&self) -> PyResult { + if self.pid != std::process::id() { + return Err(PyRuntimeError::new_err( + "native cache handles must be recreated after fork", + )); + } + Ok(self.service.clone()) + } +} + +#[pymethods] +impl NativeCacheHandle { + #[staticmethod] + #[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))] + fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult { + Ok(Self { + service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes), + guard: None, + pid: std::process::id(), + }) + } + + #[staticmethod] + #[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))] + fn redis( + py: Python<'_>, + url: String, + ttl_seconds: Option, + namespace: Option, + ) -> PyResult { + let ttl = ttl_seconds.map(duration).transpose()?; + let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + + #[getter] + fn backend(&self) -> &'static str { + self.service.kind() + } + + fn bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { + let service = self.service()?; + let guard = FacadeGuard::capture(py, facade, self.backend())?; + let handle = Py::new( + py, + Self { + service, + guard: Some(guard), + pid: self.pid, + }, + )?; + facade.setattr("_native_cache_handle", handle) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Some(guard) = &self.guard { + guard.traverse(visit)?; + } + Ok(()) + } +} + +enum CacheBinding { + Disabled, + Native(NativeResponseCache), + PythonCallback(Py), +} + +#[pyclass(frozen, name = "CacheBinding")] +pub(crate) struct ResolvedCache { + binding: CacheBinding, + pid: u32, +} + +impl ResolvedCache { + fn check_process(&self) -> PyResult<()> { + if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() { + return Err(PyRuntimeError::new_err( + "native cache bindings must be resolved again after fork", + )); + } + Ok(()) + } + + pub(crate) fn lookup_step( + &self, + py: Python<'_>, + input: &Bound<'_, PyAny>, + kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + self.check_process()?; + let awaitable = match &self.binding { + CacheBinding::Disabled => ready_none(py)?, + CacheBinding::Native(service) => { + let request = request(input)?; + let service = service.clone(); + run_async( + py, + async move { service.async_lookup(&request, now()).await }, + cache_error, + )? + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method( + "async_get_cache", + (), + Some(callback_kwargs(kwargs)?), + )?, + }; + Ok(ExecutionStep::Await(awaitable.unbind())) + } +} + +#[pymethods] +impl ResolvedCache { + #[getter] + fn kind(&self) -> &'static str { + match self.binding { + CacheBinding::Disabled => "disabled", + CacheBinding::Native(_) => "native", + CacheBinding::PythonCallback(_) => "python_callback", + } + } + + #[pyo3(signature = (request, *, callback_kwargs=None))] + fn lookup( + &self, + py: Python<'_>, + request: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => Ok(py.None()), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let service = service.clone(); + let response = release_gil(py, move || service.lookup(&request, now())) + .map_err(cache_error)?; + to_py(py, &response) + } + CacheBinding::PythonCallback(object) => object + .bind(py) + .call_method( + "get_cache", + (), + Some(self::callback_kwargs(callback_kwargs)?), + ) + .map(Bound::unbind), + } + } + + #[pyo3(signature = (request, response, *, callback_kwargs=None))] + fn store( + &self, + py: Python<'_>, + request: &Bound<'_, PyAny>, + response: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult<()> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => Ok(()), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let response: Value = from_py(response)?; + let service = service.clone(); + release_gil(py, move || service.store(&request, response, now())) + .map_err(cache_error) + } + CacheBinding::PythonCallback(object) => object + .bind(py) + .call_method( + "add_cache", + (response,), + Some(self::callback_kwargs(callback_kwargs)?), + ) + .map(|_| ()), + } + } + + #[pyo3(signature = (request, *, callback_kwargs=None))] + fn async_lookup<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)? + else { + unreachable!() + }; + Ok(awaitable.into_bound(py)) + } + + #[pyo3(signature = (request, response, *, callback_kwargs=None))] + fn async_store<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + response: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let response: Value = from_py(response)?; + let service = service.clone(); + run_async( + py, + async move { service.async_store(&request, response, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method( + "async_add_cache", + (response,), + Some(self::callback_kwargs(callback_kwargs)?), + ), + } + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + if let CacheBinding::PythonCallback(object) = &self.binding { + visit.call(object)?; + } + Ok(()) + } +} + +fn callback_kwargs<'a, 'py>( + kwargs: Option<&'a Bound<'py, PyDict>>, +) -> PyResult<&'a Bound<'py, PyDict>> { + kwargs.ok_or_else(|| { + PyTypeError::new_err("Python cache callbacks require their original callback_kwargs") + }) +} + +fn ready_none(py: Python<'_>) -> PyResult> { + let future = py + .import("asyncio")? + .call_method0("get_running_loop")? + .call_method0("create_future")?; + future.call_method1("set_result", (py.None(),))?; + Ok(future) +} + +#[pyclass(frozen)] +pub(crate) struct CacheResolver { + namespace: Py, +} + +#[pymethods] +impl CacheResolver { + #[new] + fn new(namespace: Py) -> Self { + Self { namespace } + } + + pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult { + let object = self.namespace.bind(py).getattr("cache")?; + let binding = if object.is_none() { + CacheBinding::Disabled + } else if let Ok(handle) = object.extract::>() { + CacheBinding::Native(handle.service()?) + } else if let Some(service) = facade::resolve(py, &object)? { + CacheBinding::Native(service) + } else { + CacheBinding::PythonCallback(object.unbind()) + }; + Ok(ResolvedCache { + binding, + pid: std::process::id(), + }) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.namespace) + } +} diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 46f98736aa1..621c111a35b 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -1,3 +1,4 @@ +mod cache; mod credentials; mod diagnostics; mod errors; @@ -9,6 +10,8 @@ mod token_counter; #[pymodule(gil_used = true)] mod _native { + #[pymodule_export] + use crate::cache::{CacheResolver, NativeCacheHandle, ResolvedCache}; #[cfg(feature = "panic-test")] #[pymodule_export] use crate::diagnostics::_panic_for_test; @@ -65,6 +68,9 @@ mod tests { "achat_completions", "ResponsesWebSocketConnection", "TokenCounter", + "CacheResolver", + "NativeCacheHandle", + "CacheBinding", "gil_stats", "process_state_started", "reserve_process_for_forking", diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 05a6df6d5af..4fd2f0829a3 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,5 +1,5 @@ from asyncio import Future -from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence from typing import Never, final from litellm.llms.base_llm.ocr.transformation import OCRResponse @@ -93,6 +93,50 @@ class ResponsesWebSocketConnection: def recv_text(self) -> Future[str | None]: ... def close(self) -> Future[None]: ... +@final +class NativeCacheHandle: + def __new__(cls, _uninstantiable: Never, /) -> Never: ... + @staticmethod + def memory( + *, capacity: int = 200, ttl_seconds: float = 600.0, max_entry_bytes: int = 1048576 + ) -> NativeCacheHandle: ... + @staticmethod + def redis(url: str, *, ttl_seconds: float | None = None, namespace: str | None = None) -> NativeCacheHandle: ... + @property + def backend(self) -> str: ... + def bind_facade(self, facade: object) -> None: ... + +@final +class CacheResolver: + def __new__(cls, namespace: object) -> CacheResolver: ... + def resolve(self) -> CacheBinding: ... + +@final +class CacheBinding: + def __new__(cls, _uninstantiable: Never, /) -> Never: ... + @property + def kind(self) -> str: ... + def lookup( + self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None + ) -> object: ... + def store( + self, + request: Mapping[str, object] | None, + response: object, + *, + callback_kwargs: dict[str, object] | None = None, + ) -> None: ... + def async_lookup( + self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None + ) -> Awaitable[object]: ... + def async_store( + self, + request: Mapping[str, object] | None, + response: object, + *, + callback_kwargs: dict[str, object] | None = None, + ) -> Awaitable[object]: ... + @final class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... @@ -109,7 +153,10 @@ def process_state_started() -> bool: ... def reserve_process_for_forking() -> None: ... __all__ = [ + "CacheBinding", + "CacheResolver", "ForkedAfterNativeRuntimeStarted", + "NativeCacheHandle", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RustBridgeDeclined", diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py new file mode 100644 index 00000000000..d1cea860cb0 --- /dev/null +++ b/tests/test_litellm_rust/test_cache.py @@ -0,0 +1,231 @@ +import asyncio +import contextvars +import gc +import json +import threading +import time +import weakref +from collections.abc import Generator +from types import SimpleNamespace +from typing import Final, Protocol, cast + +import fakeredis +import pytest +import redis + +import litellm +from litellm.caching.caching import Cache +from litellm.caching.in_memory_cache import InMemoryCache +from litellm.rust_bridge import _native +from litellm.types.caching import LiteLLMCacheType +from tests.test_litellm_rust.support.isolation import rebound + +pytestmark: Final = pytest.mark.requires_rust_extension + + +class CacheLookup(Protocol): + def get_cache(self, **kwargs: object) -> object: ... + + +def request(key: str = "key") -> dict[str, object]: + return {"key": {"preset": key}} + + +@pytest.fixture +def redis_url() -> Generator[str]: + server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis") + worker: Final = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + try: + yield f"redis://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + server.server_close() + worker.join(timeout=5) + + +def test_existing_constructor_and_global_are_unchanged() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + assert type(facade.cache) is InMemoryCache + assert "_native_cache_handle" not in vars(facade) + with rebound(litellm, "cache", facade): + resolver: Final = _native.CacheResolver(litellm) + assert resolver.resolve().kind == "python_callback" + resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) + assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} + + +async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: + namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.memory()) + resolver: Final = _native.CacheResolver(namespace) + selected: Final = resolver.resolve() + assert selected.kind == "native" + selected.store(request(), {"answer": 1}) + assert await selected.async_lookup(request()) == {"answer": 1} + with rebound(namespace, "cache", _native.NativeCacheHandle.memory()): + replacement: Final = resolver.resolve() + await selected.async_store(request(), {"answer": 2}) + assert replacement.lookup(request()) is None + assert selected.lookup(request()) == {"answer": 2} + with rebound(namespace, "cache", None): + disabled: Final = resolver.resolve() + assert disabled.kind == "disabled" + assert disabled.lookup(None) is None + await disabled.async_store(None, object()) + assert await disabled.async_lookup(None) is None + assert selected.lookup(request()) == {"answer": 2} + + +async def test_python_callback_preserves_identity_caller_task_context_and_errors() -> None: + context: Final = contextvars.ContextVar("cache_context", default="caller") + caller: Final = asyncio.current_task() + sentinel: Final = object() + failure: Final = RuntimeError("callback failed") + + class CustomCache: + async def async_get_cache(self, *, marker: object) -> object: + assert marker is sentinel + assert asyncio.current_task() is caller + context.set("callback") + return marker + + async def async_add_cache(self, response: object, *, marker: object) -> None: + assert response is sentinel + assert marker is sentinel + raise failure + + namespace: Final = SimpleNamespace(cache=CustomCache()) + binding: Final = _native.CacheResolver(namespace).resolve() + assert binding.kind == "python_callback" + assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel + assert context.get() == "callback" + with pytest.raises(RuntimeError) as caught: + await binding.async_store(None, sentinel, callback_kwargs={"marker": sentinel}) + assert caught.value is failure + + +async def test_callback_cancellation_stays_in_the_callers_task() -> None: + entered: Final = asyncio.Event() + finished: Final = asyncio.Event() + + class CustomCache: + async def async_get_cache(self) -> None: + entered.set() + try: + await asyncio.Future() + finally: + finished.set() + + binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve() + + async def lookup() -> object: + return await binding.async_lookup(None, callback_kwargs={}) + + task: Final = asyncio.create_task(lookup()) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert finished.is_set() + + +def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + handle: Final = _native.NativeCacheHandle.memory() + handle.bind_facade(facade) + resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade)) + native: Final = resolver.resolve() + assert native.kind == "native" + native.store(request(), {"source": "native"}) + assert native.lookup(request()) == {"source": "native"} + assert cast(CacheLookup, facade).get_cache(cache_key="key") is None + sentinel: Final = object() + + def outer_override(**_kwargs: object) -> object: + return sentinel + + def backend_override(*_args: object, **_kwargs: object) -> dict[str, str]: + return {"source": "override"} + + with rebound(facade, "get_cache", outer_override): + fallback: Final = resolver.resolve() + assert fallback.kind == "python_callback" + assert fallback.lookup(None, callback_kwargs={"cache_key": "key"}) is sentinel + assert resolver.resolve().kind == "python_callback" + delattr(facade, "get_cache") + assert resolver.resolve().kind == "native" + with rebound(facade.cache, "get_cache", backend_override): + backend_fallback: Final = resolver.resolve() + assert backend_fallback.kind == "python_callback" + assert backend_fallback.lookup(None, callback_kwargs={"cache_key": "key"}) == {"source": "override"} + + +def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not_bypassed() -> None: + class CustomCache(Cache): + pass + + handle: Final = _native.NativeCacheHandle.memory() + with pytest.raises(TypeError): + handle.bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + handle.bind_facade(facade) + resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade)) + with rebound(facade, "cache", InMemoryCache()): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "ttl", 12): + assert resolver.resolve().kind == "python_callback" + + def custom_key(**_kwargs: object) -> str: + return "custom" + + with rebound(facade, "get_cache_key", custom_key): + assert resolver.resolve().kind == "python_callback" + assert resolver.resolve().kind == "python_callback" + delattr(facade, "get_cache_key") + assert resolver.resolve().kind == "native" + + +def test_resolver_and_callback_cycles_can_be_collected() -> None: + class CustomCache: + pass + + def cyclic_reference() -> weakref.ReferenceType[CustomCache]: + callback: Final = CustomCache() + namespace: Final = SimpleNamespace(cache=callback) + binding: Final = _native.CacheResolver(namespace).resolve() + setattr(callback, "binding", binding) + return weakref.ref(callback) + + reference: Final = cyclic_reference() + gc.collect() + assert reference() is None + + +async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: + client: Final = redis.Redis.from_url(redis_url) + namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url, namespace="team")) + binding: Final = _native.CacheResolver(namespace).resolve() + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} + client.set("team:sync", str(envelope)) + client.set("team:async", json.dumps({"timestamp": time.time(), "response": response})) + assert binding.lookup(request("sync")) == response + assert await binding.async_lookup(request("team:async")) == response + await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) + stored: Final = client.get("team:native") + assert isinstance(stored, bytes) + assert json.loads(stored)["response"] == response + assert 0 < client.ttl("team:native") <= 12 + assert client.get("litellm-cache:team:native") is None + assert client.get("team:team:async") is None + client.close() + + +def test_invalid_duration_and_request_shape_fail_before_storage() -> None: + binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() + for seconds in (-1.0, float("nan"), float("inf")): + with pytest.raises(ValueError): + binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) + assert binding.lookup(request()) is None + with pytest.raises(ValueError): + _native.NativeCacheHandle.memory(ttl_seconds=-1) From 0c3a0a208948e97d0da05ec6c7e28672205c2f00 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Sun, 20 Sep 2026 21:11:48 -0700 Subject: [PATCH 03/17] refactor(cache): separate response policy and host selection --- litellm-rust/Cargo.lock | 4 +- litellm-rust/crates/cache-memory/Cargo.toml | 2 +- litellm-rust/crates/cache-memory/src/cache.rs | 47 +----- .../crates/cache-memory/tests/cache.rs | 80 ++++------ litellm-rust/crates/cache-redis/src/cache.rs | 34 +++-- litellm-rust/crates/cache-response/Cargo.toml | 7 +- litellm-rust/crates/cache-response/README.md | 55 +++++++ .../crates/cache-response/src/caching.rs | 143 ++++++++++++++++++ .../crates/cache-response/src/codec.rs | 4 +- litellm-rust/crates/cache-response/src/lib.rs | 7 +- .../crates/cache-response/src/response.rs | 6 +- .../crates/cache-response/tests/caching.rs | 83 ++++++++++ .../crates/cache-response/tests/response.rs | 40 +++-- litellm-rust/crates/cache/Cargo.toml | 1 - litellm-rust/crates/cache/src/caching.rs | 143 ------------------ litellm-rust/crates/cache/src/lib.rs | 5 +- litellm-rust/crates/cache/tests/caching.rs | 94 +----------- litellm-rust/crates/cache/tests/codec.rs | 14 +- litellm-rust/crates/python-bridge/Cargo.toml | 2 + .../crates/python-bridge/src/cache/facade.rs | 3 +- .../crates/python-bridge/src/cache/mod.rs | 6 +- .../src => python-bridge/src/cache}/native.rs | 33 ++-- tests/test_litellm_rust/test_cache.py | 20 ++- 23 files changed, 426 insertions(+), 407 deletions(-) create mode 100644 litellm-rust/crates/cache-response/README.md create mode 100644 litellm-rust/crates/cache-response/src/caching.rs create mode 100644 litellm-rust/crates/cache-response/tests/caching.rs rename litellm-rust/crates/{cache-response/src => python-bridge/src/cache}/native.rs (75%) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index aa62e4f3770..8b299f5455b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2460,7 +2460,6 @@ dependencies = [ "rstest", "serde", "serde_json", - "sha2 0.10.9", "thiserror 2.0.19", "tokio", ] @@ -2498,6 +2497,7 @@ dependencies = [ "redis-test", "serde", "serde_json", + "sha2 0.10.9", "tokio", ] @@ -2665,6 +2665,8 @@ dependencies = [ "litellm-auth", "litellm-auth-gcp", "litellm-cache", + "litellm-cache-memory", + "litellm-cache-redis", "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core", diff --git a/litellm-rust/crates/cache-memory/Cargo.toml b/litellm-rust/crates/cache-memory/Cargo.toml index d4487573a9a..86ab01564c8 100644 --- a/litellm-rust/crates/cache-memory/Cargo.toml +++ b/litellm-rust/crates/cache-memory/Cargo.toml @@ -7,8 +7,8 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -serde_json.workspace = true [dev-dependencies] +serde_json.workspace = true rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 974cdbe9760..43186faf3f7 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -3,15 +3,12 @@ use std::collections::{BinaryHeap, HashMap}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, -}; +use litellm_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error}; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; const DEFAULT_TTL: Duration = Duration::from_secs(600); type ValueMeasure = Arc Result + Send + Sync>; -type ValueValidator = Arc Result<(), Error> + Send + Sync>; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub enum CacheWrite { @@ -32,7 +29,6 @@ pub struct InMemoryCache { default_ttl: Duration, max_entry_bytes: Option, measure_value: Option>, - validate_value: Option>, now: Arc Duration + Send + Sync>, } @@ -76,7 +72,6 @@ impl InMemoryCache { default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), max_entry_bytes, measure_value, - validate_value: None, now: Arc::new(now), } } @@ -90,9 +85,6 @@ impl InMemoryCache { if self.max_size_in_memory == 0 { return Ok(CacheWrite::Disabled); } - if let Some(validate) = &self.validate_value { - validate(&value)?; - } if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value) && measure(&value)? > limit { @@ -176,43 +168,6 @@ impl InMemoryCache { } } -impl InMemoryCache { - pub fn response_cache(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { - Self::response_cache_with_clock(capacity, ttl, max_entry_bytes, || { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - }) - } - - pub fn response_cache_with_clock( - capacity: usize, - ttl: Duration, - max_entry_bytes: usize, - now: impl Fn() -> Duration + Send + Sync + 'static, - ) -> Self { - let mut cache = Self::with_clock_and_size_measurement( - Some(capacity), - Some(ttl), - Some(max_entry_bytes), - Some(Arc::new(|entry: &CacheEntry| { - serde_json::to_vec(entry) - .map(|bytes| bytes.len()) - .map_err(|_| Error::InvalidEntry) - })), - now, - ); - cache.validate_value = Some(Arc::new(|entry: &CacheEntry| { - entry - .timestamp - .is_finite() - .then_some(()) - .ok_or(Error::InvalidEntry) - })); - cache - } -} - impl BaseCache for InMemoryCache { type Value = V; diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index ffac9d8ae64..370145bebef 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -3,8 +3,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use litellm_cache::{ - BaseCache, CacheBackend, CacheConnectionStatus, CacheEntry, CacheKwargs, Error, get_cache, - set_cache, + BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, Error, get_cache, set_cache, }; use litellm_cache_memory::{CacheWrite, InMemoryCache}; use rstest::{fixture, rstest}; @@ -87,66 +86,49 @@ fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc } #[test] -fn disabled_size_limited_and_synchronized_response_writes_are_observable() { - let disabled = InMemoryCache::::response_cache(0, Duration::from_secs(60), 80); +fn disabled_size_limited_and_validated_writes_are_observable() { + let cache = |capacity| { + InMemoryCache::with_clock_and_size_measurement( + Some(capacity), + Some(Duration::from_secs(60)), + Some(4), + Some(Arc::new(|value: &String| { + if value.is_empty() { + return Err(Error::InvalidEntry); + } + Ok(value.len()) + })), + || Duration::from_secs(100), + ) + }; + let disabled = cache(0); assert_eq!( - disabled - .set_cache( - "a", - CacheEntry { - timestamp: 1.0, - response: serde_json::json!("x") - }, - None - ) - .unwrap(), + disabled.set_cache("a", "x".into(), None).unwrap(), CacheWrite::Disabled ); - let cache = InMemoryCache::::response_cache(2, Duration::from_secs(60), 80); + let cache = cache(2); assert_eq!( - cache - .set_cache( - "large", - CacheEntry { - timestamp: 1.0, - response: serde_json::json!("x".repeat(100)) - }, - None - ) - .unwrap(), + cache.set_cache("large", "oversized".into(), None).unwrap(), CacheWrite::TooLarge ); - cache - .set_cache( - "small", - CacheEntry { - timestamp: 1.0, - response: serde_json::json!("ok"), - }, - None, - ) - .unwrap(); - assert!(cache.get_cache("small").unwrap().is_some()); + assert_eq!(cache.get_cache("large").unwrap(), None); assert_eq!( - cache - .set_cache( - "invalid", - CacheEntry { - timestamp: f64::NAN, - response: serde_json::json!("bad"), - }, - None, - ) - .unwrap_err(), - Error::InvalidEntry + cache.set_cache("small", "ok".into(), None).unwrap(), + CacheWrite::Stored ); + assert_eq!(cache.get_cache("small").unwrap(), Some("ok".into())); + assert_eq!( + cache.set_cache("invalid", String::new(), None), + Err(Error::InvalidEntry) + ); + assert_eq!(cache.get_cache("invalid").unwrap(), None); cache.delete_cache("small").unwrap(); - cache.flush_cache().unwrap(); + assert_eq!(cache.get_cache("small").unwrap(), None); } #[tokio::test] async fn connection_test_matches_python_result_contract() { - let cache = InMemoryCache::::default(); + let cache = InMemoryCache::::default(); let result = BaseCache::test_connection(&cache).await.unwrap(); assert_eq!(result.status, CacheConnectionStatus::Success); assert_eq!(result.message, "In-memory cache connection test successful"); diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 0faca6cdaaf..d4ca0cf0522 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -238,30 +238,27 @@ where #[cfg(test)] mod tests { use super::RedisCache; - use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKwargs, JsonCodec}; + use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, JsonCodec}; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; use std::time::Duration; - fn entry() -> CacheEntry { - CacheEntry { - timestamp: 123.0, - response: json!({"choices": [{"text": "cached"}]}), - } + fn entry() -> serde_json::Value { + json!({"deployment": "model-a", "cooldown_seconds": 30}) } #[test] fn ttl_seconds_rounds_up_and_keeps_expiration_positive() { assert_eq!( - RedisCache::>::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 ); } @@ -269,7 +266,9 @@ mod tests { #[test] fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() { let value = entry(); - let payload = JsonCodec::::new().encode(&value).unwrap(); + let payload = JsonCodec::::new() + .encode(&value) + .unwrap(); let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("SETEX") @@ -282,8 +281,9 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("litellm-cache".into())); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); cache .set_cache("key", value.clone(), CacheKwargs::default()) @@ -308,8 +308,9 @@ mod tests { MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), ]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("litellm-cache".into())); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); cache.flush_cache().unwrap(); } @@ -318,8 +319,9 @@ mod tests { async fn test_connection_runs_ping_off_executor() { let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))]) .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("litellm-cache".into())); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("litellm-cache".into())); assert_eq!( cache.test_connection().await.unwrap().status, diff --git a/litellm-rust/crates/cache-response/Cargo.toml b/litellm-rust/crates/cache-response/Cargo.toml index a0c4a1f74ef..04affb9872d 100644 --- a/litellm-rust/crates/cache-response/Cargo.toml +++ b/litellm-rust/crates/cache-response/Cargo.toml @@ -7,13 +7,14 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -litellm-cache-memory.workspace = true -litellm-cache-redis.workspace = true py_literal = "0.4.0" -redis = "1.7.0" serde.workspace = true serde_json.workspace = true +sha2.workspace = true [dev-dependencies] +litellm-cache-memory.workspace = true +litellm-cache-redis.workspace = true +redis = "1.7.0" redis-test = "1.0.4" tokio.workspace = true diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md new file mode 100644 index 00000000000..309296f7773 --- /dev/null +++ b/litellm-rust/crates/cache-response/README.md @@ -0,0 +1,55 @@ +# Response cache foundation + +`ResponseCache` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache` + +## Ownership + +`litellm-cache` defines typed storage and codec traits. Memory and Redis implement those traits without depending on response policy. Other consumers can store their own value types using the same backend implementations + +`litellm-cache-response` owns response keys, controls, entries, and the Python-compatible response codec. It has no runtime dependency on a specific cache backend or Python + +The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host + +## Native Rust use + +```rust +use std::{sync::Arc, time::Duration}; +use litellm_cache_memory::InMemoryCache; +use litellm_cache_response::{CacheKeyInput, ResponseCache, ResponseCacheRequest}; +use serde_json::json; + +let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); +let request = ResponseCacheRequest::new(CacheKeyInput { + preset: Some("example:key".into()), + ..Default::default() +}); +let now = Duration::from_secs(100); +cache.store(&request, json!({"answer": 7}), now)?; +assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7}))); +``` + +For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Each Redis constructor currently opens its own connection; shared connection pools remain follow-up work + +Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it + +## Python integration boundary + +The extension exposes `NativeCacheHandle`, `CacheResolver`, and captured `CacheBinding` objects for host integration. Memory and Redis handles support synchronous and asynchronous response lookup and storage + +The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution + +Explicit facade registration checks object identity, method overrides, and configuration changes before selecting native execution. Registration does not migrate entries or replace Python methods. Until activation configures one shared service, a registered facade and its native handle can hold separate data. Existing public cache constructors remain on Python + +Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy + +## Adding another backend + +Implement `BaseCache` for the backend with its associated value type, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache` then works without another response implementation. Add a concrete bridge enum variant and constructor only when exposing that backend to Python + +Verify typed values, TTL precedence, missing entries, serialization failures, namespaces, batch ordering, and sync/async behavior. Run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before enabling a public facade + +## Follow-up scope + +Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial batches, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths + +Redis cluster, disk, cloud stores, dual caching, and semantic caching remain follow-ups. Atomic counters, affinity claims, reservations, queues, and pubsub need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees diff --git a/litellm-rust/crates/cache-response/src/caching.rs b/litellm-rust/crates/cache-response/src/caching.rs new file mode 100644 index 00000000000..53b34025ccd --- /dev/null +++ b/litellm-rust/crates/cache-response/src/caching.rs @@ -0,0 +1,143 @@ +use std::time::Duration; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +pub enum CacheMode { + #[default] + #[serde(rename = "default_on")] + DefaultOn, + #[serde(rename = "default_off")] + DefaultOff, +} + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct CacheKeyField { + pub name: String, + pub value: Option, + pub api_parameter: bool, + pub internal_parameter: bool, +} + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +#[serde(default)] +pub struct CacheKeyInput { + pub fields: Vec, + pub preset: Option, + pub namespace: Option, + pub include_provider_parameters: bool, +} + +#[derive(Default)] +pub struct CacheKeyContext { + pub model_group: Option, + pub caching_groups: Vec<(Vec, String)>, + pub file_checksum: Option, + pub file_object_name: Option, + pub metadata_file_name: Option, + pub parameters_file_name: Option, +} + +impl CacheKeyContext { + pub fn apply(self, input: &mut CacheKeyInput) { + let group = self.model_group.as_ref().and_then(|model| { + self.caching_groups + .iter() + .find(|(models, _)| models.contains(model)) + }); + for field in &mut input.fields { + match field.name.as_str() { + "model" => { + field.value = group + .map(|(_, formatted)| formatted.clone()) + .or_else(|| self.model_group.clone()) + .or_else(|| field.value.take()) + } + "file" => { + field.value = self + .file_checksum + .clone() + .or_else(|| self.file_object_name.clone()) + .or_else(|| self.metadata_file_name.clone()) + .or_else(|| self.parameters_file_name.clone()) + } + _ => {} + } + } + } +} + +pub fn get_cache_key(input: &CacheKeyInput) -> String { + cache_key(input) +} + +pub fn cache_key(input: &CacheKeyInput) -> String { + if let Some(preset) = &input.preset { + return preset.clone(); + } + let mut digest = Sha256::new(); + for field in &input.fields { + if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter)) + && let Some(value) = &field.value + { + digest.update(field.name.as_bytes()); + digest.update(b": "); + digest.update(value.as_bytes()); + } + } + let hash = format!("{:x}", digest.finalize()); + input + .namespace + .as_deref() + .filter(|namespace| !namespace.is_empty()) + .map_or(hash.clone(), |namespace| format!("{namespace}:{hash}")) +} + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)] +pub struct CacheControls { + pub supported_call_type: bool, + pub configured: bool, + pub native_backend: bool, + pub default_on: bool, + pub caching: Option, + pub no_cache: bool, + pub no_store: bool, + #[serde(default)] + pub use_cache: bool, +} + +impl CacheControls { + pub fn reads(self) -> bool { + self.supported_call_type + && self.configured + && self.caching.unwrap_or(true) + && !self.no_cache + && (self.default_on || self.use_cache) + } + + pub fn writes(self) -> bool { + self.supported_call_type + && self.configured + && !self.no_store + && (self.default_on || self.use_cache) + } +} + +pub fn should_use_cache(controls: CacheControls) -> bool { + controls.reads() || controls.writes() +} + +#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] +pub struct CacheEntry { + pub timestamp: f64, + pub response: Value, +} + +impl CacheEntry { + pub fn fresh(&self, now: Duration, max_age: Option) -> bool { + self.timestamp.is_finite() + && max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64()) + } +} diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs index 137cf61267a..f1d55ddefe8 100644 --- a/litellm-rust/crates/cache-response/src/codec.rs +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -1,4 +1,6 @@ -use litellm_cache::{CacheCodec, CacheEntry, Error}; +use litellm_cache::{CacheCodec, Error}; + +use crate::CacheEntry; use serde_json::Value; pub struct ResponseCacheCodec; diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index 454a82e76de..efa0b04b9f7 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -1,7 +1,10 @@ +mod caching; mod codec; -mod native; mod response; +pub use caching::{ + CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode, cache_key, + get_cache_key, should_use_cache, +}; pub use codec::ResponseCacheCodec; -pub use native::NativeResponseCache; pub use response::{ResponseCache, ResponseCacheRequest}; diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index 3e9807b6d1b..987a47f0554 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,8 +1,8 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{ - BaseCache, CacheControls, CacheEntry, CacheKeyInput, CacheKwargs, Error, cache_key, -}; +use litellm_cache::{BaseCache, CacheKwargs, Error}; + +use crate::{CacheControls, CacheEntry, CacheKeyInput, cache_key}; use serde_json::Value; #[derive(Clone)] diff --git a/litellm-rust/crates/cache-response/tests/caching.rs b/litellm-rust/crates/cache-response/tests/caching.rs new file mode 100644 index 00000000000..d403791e421 --- /dev/null +++ b/litellm-rust/crates/cache-response/tests/caching.rs @@ -0,0 +1,83 @@ +use litellm_cache_response::{ + CacheControls, CacheKeyContext, CacheKeyField, CacheKeyInput, cache_key, get_cache_key, +}; +use sha2::{Digest, Sha256}; + +#[test] +fn keys_match_python_order_groups_files_presets_and_namespaces() { + let mut input = CacheKeyInput { + fields: vec![ + CacheKeyField { + name: "model".into(), + value: Some("deployment".into()), + api_parameter: true, + internal_parameter: false, + }, + CacheKeyField { + name: "file".into(), + value: None, + api_parameter: true, + internal_parameter: false, + }, + ], + namespace: Some("team".into()), + ..Default::default() + }; + CacheKeyContext { + model_group: Some("group".into()), + caching_groups: vec![(vec!["group".into()], "['group']".into())], + file_checksum: Some("checksum".into()), + ..Default::default() + } + .apply(&mut input); + assert_eq!( + cache_key(&input), + format!( + "team:{:x}", + Sha256::digest(b"model: ['group']file: checksum") + ) + ); + input.preset = Some("preset".into()); + assert_eq!(get_cache_key(&input), "preset"); +} + +#[test] +fn cache_controls_honor_default_modes_and_directives() { + let enabled = CacheControls { + supported_call_type: true, + configured: true, + default_on: true, + ..Default::default() + }; + assert!(enabled.reads()); + assert!(enabled.writes()); + assert!( + !CacheControls { + default_on: false, + ..enabled + } + .reads() + ); + assert!( + CacheControls { + default_on: false, + use_cache: true, + ..enabled + } + .reads() + ); + assert!( + !CacheControls { + no_cache: true, + ..enabled + } + .reads() + ); + assert!( + !CacheControls { + no_store: true, + ..enabled + } + .writes() + ); +} diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index a4beda9fa4c..88b53fdfe86 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -6,15 +6,23 @@ use std::{ time::Duration, }; -use litellm_cache::{BaseCache, CacheCodec, CacheEntry, CacheKeyField, CacheKeyInput, Error}; +use litellm_cache::{BaseCache, CacheCodec, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ - NativeResponseCache, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, + CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec, + ResponseCacheRequest, }; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; +fn memory() -> Arc>> { + Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new( + Some(8), + Some(Duration::from_secs(600)), + )))) +} + fn request() -> ResponseCacheRequest { ResponseCacheRequest::new(CacheKeyInput { preset: Some("tenant:key".into()), @@ -87,7 +95,7 @@ async fn sync_and_async_consumers_share_keys_ttls_and_freshness() { #[tokio::test] async fn directives_skip_io_and_keep_reads_and_writes_independent() { - let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let cache = memory(); let mut request = request(); let now = Duration::from_secs(100); request.controls.no_store = true; @@ -112,7 +120,7 @@ async fn directives_skip_io_and_keep_reads_and_writes_independent() { } #[tokio::test] -async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_json() { +async fn redis_consumer_reads_python_sync_and_async_envelopes_and_writes_compatible_json() { let connection = MockRedisConnection::new([ MockCmd::new( redis::cmd("GET").arg("tenant:key"), @@ -133,7 +141,7 @@ async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_ .assert_all_commands_consumed(); let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec) .with_namespace(Some("tenant".into())); - let cache = NativeResponseCache::Redis(Arc::new(ResponseCache::new(Arc::new(backend)))); + let cache = ResponseCache::new(Arc::new(backend)); let request = request(); let expected = json!({"ok": true, "text": "cached"}); assert_eq!( @@ -154,10 +162,10 @@ async fn redis_enum_reads_python_sync_and_async_envelopes_and_writes_compatible_ } #[tokio::test] -async fn captured_enum_keeps_the_selected_backend_for_background_writes() { - let original = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); +async fn captured_service_keeps_the_selected_backend_for_background_writes() { + let original = memory(); let captured = original.clone(); - let replacement = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let replacement = memory(); let request = request(); let writer = tokio::spawn({ let request = request.clone(); @@ -186,7 +194,7 @@ async fn captured_enum_keeps_the_selected_backend_for_background_writes() { #[test] fn generated_keys_preserve_namespace_and_explicit_keys() { - let cache = NativeResponseCache::memory(8, Duration::from_secs(600), 1024); + let cache = memory(); let key = CacheKeyInput { fields: vec![CacheKeyField { name: "model".into(), @@ -199,7 +207,7 @@ fn generated_keys_preserve_namespace_and_explicit_keys() { }; let generated = ResponseCacheRequest::new(key.clone()); let explicit = ResponseCacheRequest::new(CacheKeyInput { - preset: Some(litellm_cache::cache_key(&key)), + preset: Some(litellm_cache_response::cache_key(&key)), ..Default::default() }); cache @@ -288,3 +296,15 @@ fn malformed_memory_entries_are_rejected_by_the_response_consumer() { Error::InvalidEntry ); } + +#[test] +fn response_entries_preserve_the_existing_json_representation() { + let codec = ResponseCacheCodec; + let entry = CacheEntry { + timestamp: 123.0, + response: json!({"choices": [{"text": "cached"}]}), + }; + let bytes = codec.encode(&entry).unwrap(); + assert_eq!(bytes, serde_json::to_vec(&entry).unwrap()); + assert_eq!(codec.decode(&bytes).unwrap(), entry); +} diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml index 350db4b1adb..0c504ab727a 100644 --- a/litellm-rust/crates/cache/Cargo.toml +++ b/litellm-rust/crates/cache/Cargo.toml @@ -8,7 +8,6 @@ repository.workspace = true [dependencies] serde.workspace = true serde_json.workspace = true -sha2.workspace = true thiserror.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs index 1d1df966ba2..39479694f3b 100644 --- a/litellm-rust/crates/cache/src/caching.rs +++ b/litellm-rust/crates/cache/src/caching.rs @@ -1,152 +1,9 @@ use std::sync::Arc; -use std::time::Duration; - -use serde::{Deserialize, Serialize}; -use serde_json::Value; -use sha2::{Digest, Sha256}; use crate::{BaseCache, CacheKwargs, Error}; pub use crate::BaseCache as Cache; -#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] -pub enum CacheMode { - #[default] - #[serde(rename = "default_on")] - DefaultOn, - #[serde(rename = "default_off")] - DefaultOff, -} - -#[derive(Clone, Debug, Deserialize, Serialize)] -pub struct CacheKeyField { - pub name: String, - pub value: Option, - pub api_parameter: bool, - pub internal_parameter: bool, -} - -#[derive(Clone, Debug, Default, Deserialize, Serialize)] -#[serde(default)] -pub struct CacheKeyInput { - pub fields: Vec, - pub preset: Option, - pub namespace: Option, - pub include_provider_parameters: bool, -} - -#[derive(Default)] -pub struct CacheKeyContext { - pub model_group: Option, - pub caching_groups: Vec<(Vec, String)>, - pub file_checksum: Option, - pub file_object_name: Option, - pub metadata_file_name: Option, - pub parameters_file_name: Option, -} - -impl CacheKeyContext { - pub fn apply(self, input: &mut CacheKeyInput) { - let group = self.model_group.as_ref().and_then(|model| { - self.caching_groups - .iter() - .find(|(models, _)| models.contains(model)) - }); - for field in &mut input.fields { - match field.name.as_str() { - "model" => { - field.value = group - .map(|(_, formatted)| formatted.clone()) - .or_else(|| self.model_group.clone()) - .or_else(|| field.value.take()) - } - "file" => { - field.value = self - .file_checksum - .clone() - .or_else(|| self.file_object_name.clone()) - .or_else(|| self.metadata_file_name.clone()) - .or_else(|| self.parameters_file_name.clone()) - } - _ => {} - } - } - } -} - -pub fn get_cache_key(input: &CacheKeyInput) -> String { - cache_key(input) -} - -pub fn cache_key(input: &CacheKeyInput) -> String { - if let Some(preset) = &input.preset { - return preset.clone(); - } - let mut digest = Sha256::new(); - for field in &input.fields { - if (field.api_parameter || (input.include_provider_parameters && !field.internal_parameter)) - && let Some(value) = &field.value - { - digest.update(field.name.as_bytes()); - digest.update(b": "); - digest.update(value.as_bytes()); - } - } - let hash = format!("{:x}", digest.finalize()); - input - .namespace - .as_deref() - .filter(|namespace| !namespace.is_empty()) - .map_or(hash.clone(), |namespace| format!("{namespace}:{hash}")) -} - -#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize)] -pub struct CacheControls { - pub supported_call_type: bool, - pub configured: bool, - pub native_backend: bool, - pub default_on: bool, - pub caching: Option, - pub no_cache: bool, - pub no_store: bool, - #[serde(default)] - pub use_cache: bool, -} - -impl CacheControls { - pub fn reads(self) -> bool { - self.supported_call_type - && self.configured - && self.caching.unwrap_or(true) - && !self.no_cache - && (self.default_on || self.use_cache) - } - - pub fn writes(self) -> bool { - self.supported_call_type - && self.configured - && !self.no_store - && (self.default_on || self.use_cache) - } -} - -pub fn should_use_cache(controls: CacheControls) -> bool { - controls.reads() || controls.writes() -} - -#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] -pub struct CacheEntry { - pub timestamp: f64, - pub response: Value, -} - -impl CacheEntry { - pub fn fresh(&self, now: Duration, max_age: Option) -> bool { - self.timestamp.is_finite() - && max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64()) - } -} - pub fn get_cache( cache: &B, key: &str, diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index a1d9d1402bb..4ff02319bdc 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -4,9 +4,6 @@ mod codec; mod error; pub use base_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs}; -pub use caching::{ - Cache, CacheBackend, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, - CacheMode, cache_key, get_cache, get_cache_key, set_cache, should_use_cache, -}; +pub use caching::{Cache, CacheBackend, get_cache, set_cache}; pub use codec::{CacheCodec, JsonCodec}; pub use error::Error; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 5c250c6b3c9..824de00bdf4 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,17 +1,13 @@ -use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, - CacheKeyInput, CacheKwargs, Error, cache_key, get_cache_key, -}; -use sha2::{Digest, Sha256}; +use litellm_cache::{BaseCache, CacheConnectionResult, CacheKwargs, Error}; use std::{sync::Mutex, time::Duration}; struct TestCache { default_ttl: Duration, - writes: Mutex>, + writes: Mutex>, } impl BaseCache for TestCache { - type Value = CacheEntry; + type Value = String; fn default_ttl(&self) -> Duration { self.default_ttl @@ -83,10 +79,7 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() { default_ttl: Duration::from_secs(60), writes: Mutex::default(), }; - let entry = CacheEntry { - timestamp: 123.0, - response: serde_json::json!("cached"), - }; + let entry = String::from("cached"); let kwargs = CacheKwargs { ttl: Some(Duration::from_secs(5)), ..Default::default() @@ -116,82 +109,3 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() { ] ); } - -#[test] -fn keys_match_python_order_groups_files_presets_and_namespaces() { - let mut input = CacheKeyInput { - fields: vec![ - CacheKeyField { - name: "model".into(), - value: Some("deployment".into()), - api_parameter: true, - internal_parameter: false, - }, - CacheKeyField { - name: "file".into(), - value: None, - api_parameter: true, - internal_parameter: false, - }, - ], - namespace: Some("team".into()), - ..Default::default() - }; - CacheKeyContext { - model_group: Some("group".into()), - caching_groups: vec![(vec!["group".into()], "['group']".into())], - file_checksum: Some("checksum".into()), - ..Default::default() - } - .apply(&mut input); - assert_eq!( - cache_key(&input), - format!( - "team:{:x}", - Sha256::digest(b"model: ['group']file: checksum") - ) - ); - input.preset = Some("preset".into()); - assert_eq!(get_cache_key(&input), "preset"); -} - -#[test] -fn cache_controls_honor_default_modes_and_directives() { - let enabled = CacheControls { - supported_call_type: true, - configured: true, - default_on: true, - ..Default::default() - }; - assert!(enabled.reads()); - assert!(enabled.writes()); - assert!( - !CacheControls { - default_on: false, - ..enabled - } - .reads() - ); - assert!( - CacheControls { - default_on: false, - use_cache: true, - ..enabled - } - .reads() - ); - assert!( - !CacheControls { - no_cache: true, - ..enabled - } - .reads() - ); - assert!( - !CacheControls { - no_store: true, - ..enabled - } - .writes() - ); -} diff --git a/litellm-rust/crates/cache/tests/codec.rs b/litellm-rust/crates/cache/tests/codec.rs index dad5398a879..e24545caad6 100644 --- a/litellm-rust/crates/cache/tests/codec.rs +++ b/litellm-rust/crates/cache/tests/codec.rs @@ -1,6 +1,6 @@ use std::collections::BTreeMap; -use litellm_cache::{CacheCodec, CacheEntry, Error, JsonCodec}; +use litellm_cache::{CacheCodec, Error, JsonCodec}; use serde::{Deserialize, Serialize}; use serde_json::json; @@ -25,18 +25,6 @@ fn json_codec_round_trips_typed_domain_values() { ); } -#[test] -fn response_entries_preserve_the_existing_json_representation() { - let codec = JsonCodec::::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(); diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 308dfd2dd7a..1eb2ec28036 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -21,6 +21,8 @@ tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] bytes.workspace = true litellm-cache.workspace = true +litellm-cache-memory.workspace = true +litellm-cache-redis.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 32360e72469..eb07118e964 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -1,4 +1,3 @@ -use litellm_cache_response::NativeResponseCache; use litellm_host_python::from_py; use pyo3::{ PyTraverseError, PyVisit, @@ -8,7 +7,7 @@ use pyo3::{ }; use serde_json::Value; -use super::NativeCacheHandle; +use super::{NativeCacheHandle, native::NativeResponseCache}; struct ClassGuard { class: Py, diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index f00cceeb86e..5918967009a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,9 +1,10 @@ mod facade; +mod native; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::{CacheControls, CacheKeyInput, Error}; -use litellm_cache_response::{NativeResponseCache, ResponseCacheRequest}; +use litellm_cache::Error; +use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; use pyo3::{ PyTraverseError, PyVisit, @@ -15,6 +16,7 @@ use serde::Deserialize; use serde_json::Value; use facade::FacadeGuard; +use native::NativeResponseCache; #[derive(Deserialize)] #[serde(deny_unknown_fields)] diff --git a/litellm-rust/crates/cache-response/src/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs similarity index 75% rename from litellm-rust/crates/cache-response/src/native.rs rename to litellm-rust/crates/python-bridge/src/cache/native.rs index c25c61cee82..6af04bfe2b2 100644 --- a/litellm-rust/crates/cache-response/src/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,33 +1,30 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheEntry, Error}; +use litellm_cache::{CacheCodec, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use serde_json::Value; -use crate::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; +use litellm_cache_response::{CacheEntry, ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; -pub enum NativeResponseCache -where - C: redis::ConnectionLike + Send + 'static, -{ +#[derive(Clone)] +pub(super) enum NativeResponseCache { Memory(Arc>>), - Redis(Arc>>), -} - -impl Clone for NativeResponseCache { - fn clone(&self) -> Self { - match self { - Self::Memory(cache) => Self::Memory(Arc::clone(cache)), - Self::Redis(cache) => Self::Redis(Arc::clone(cache)), - } - } + Redis(Arc>>), } impl NativeResponseCache { pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { Self::Memory(Arc::new(ResponseCache::new(Arc::new( - InMemoryCache::response_cache(capacity, ttl, max_entry_bytes), + InMemoryCache::with_clock_and_size_measurement( + Some(capacity), + Some(ttl), + Some(max_entry_bytes), + Some(Arc::new(|entry| { + ResponseCacheCodec.encode(entry).map(|bytes| bytes.len()) + })), + super::now, + ), )))) } @@ -41,7 +38,7 @@ impl NativeResponseCache { } } -impl NativeResponseCache { +impl NativeResponseCache { pub fn kind(&self) -> &'static str { match self { Self::Memory(_) => "memory", diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index d1cea860cb0..493baac228a 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -224,8 +224,24 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd def test_invalid_duration_and_request_shape_fail_before_storage() -> None: binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() for seconds in (-1.0, float("nan"), float("inf")): - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) assert binding.lookup(request()) is None - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): _native.NativeCacheHandle.memory(ttl_seconds=-1) + + +async def test_memory_size_policy_is_applied_by_the_native_host() -> None: + handle: Final = _native.NativeCacheHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = _native.CacheResolver(SimpleNamespace(cache=handle)).resolve() + small: Final = {"answer": "ok"} + binding.store(request("small"), small) + assert await binding.async_lookup(request("small")) == small + await binding.async_store(request("large"), {"answer": "x" * 256}) + assert binding.lookup(request("large")) is None + assert binding.lookup(request("small")) == small + disabled: Final = _native.CacheResolver( + SimpleNamespace(cache=_native.NativeCacheHandle.memory(capacity=0)) + ).resolve() + await disabled.async_store(request(), small) + assert await disabled.async_lookup(request()) is None From cc8cadebb46d22286166bb9d4152b58bac11d91b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 07:56:20 -0700 Subject: [PATCH 04/17] fix(cache): close native parity gaps --- litellm-rust/Cargo.lock | 22 ++ litellm-rust/crates/cache-memory/src/cache.rs | 54 ++- .../crates/cache-memory/tests/cache.rs | 51 ++- litellm-rust/crates/cache-redis/Cargo.toml | 3 +- litellm-rust/crates/cache-redis/src/cache.rs | 356 +++++++++++++++--- .../crates/cache-redis/tests/cache.rs | 79 +++- litellm-rust/crates/cache-response/README.md | 8 +- .../crates/cache-response/src/caching.rs | 10 +- .../crates/cache-response/src/codec.rs | 24 +- .../crates/cache-response/src/embedding.rs | 22 ++ litellm-rust/crates/cache-response/src/lib.rs | 2 + .../crates/cache-response/src/response.rs | 164 +++++++- .../crates/cache-response/tests/caching.rs | 7 + .../crates/cache-response/tests/response.rs | 87 ++++- litellm-rust/crates/cache/src/base_cache.rs | 45 +++ litellm-rust/crates/cache/src/capabilities.rs | 39 ++ litellm-rust/crates/cache/src/dual.rs | 105 ++++++ litellm-rust/crates/cache/src/lib.rs | 7 +- litellm-rust/crates/cache/tests/dual.rs | 130 +++++++ .../crates/python-bridge/python_settings.json | 3 + .../crates/python-bridge/src/cache/facade.rs | 14 +- .../crates/python-bridge/src/cache/mod.rs | 178 ++++++++- .../crates/python-bridge/src/cache/native.rs | 129 ++++++- .../python-bridge/src/python_settings.rs | 5 +- litellm/caching/dual_cache.py | 79 ++-- litellm/rust_bridge/_native.pyi | 21 ++ litellm/rust_bridge/settings.py | 11 + tests/test_litellm/caching/test_dual_cache.py | 79 ++-- tests/test_litellm_rust/test_cache.py | 65 ++++ 29 files changed, 1615 insertions(+), 184 deletions(-) create mode 100644 litellm-rust/crates/cache-response/src/embedding.rs create mode 100644 litellm-rust/crates/cache/src/capabilities.rs create mode 100644 litellm-rust/crates/cache/src/dual.rs create mode 100644 litellm-rust/crates/cache/tests/dual.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 8b299f5455b..725d2cfef41 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2479,6 +2479,7 @@ name = "litellm-cache-redis" version = "0.1.0" dependencies = [ "litellm-cache", + "r2d2", "redis", "redis-test", "serde_json", @@ -3564,6 +3565,17 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "r2d2" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" +dependencies = [ + "log", + "parking_lot", + "scheduled-thread-pool", +] + [[package]] name = "rand" version = "0.8.7" @@ -3700,6 +3712,7 @@ dependencies = [ "itoa", "num-bigint 0.5.1", "percent-encoding", + "r2d2", "ryu", "sha1_smol", "socket2 0.6.5", @@ -4096,6 +4109,15 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "scheduled-thread-pool" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" +dependencies = [ + "parking_lot", +] + [[package]] name = "schemars" version = "0.9.0" diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 43186faf3f7..074c8dcb0fd 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -3,7 +3,10 @@ use std::collections::{BinaryHeap, HashMap}; use std::sync::{Arc, Mutex}; use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error}; +use litellm_cache::{ + BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, + Error, +}; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; const DEFAULT_TTL: Duration = Duration::from_secs(600); @@ -168,6 +171,55 @@ impl InMemoryCache { } } +impl ClaimCache for InMemoryCache +where + V: Clone + PartialEq + Send + Sync + 'static, +{ + fn claim_cache( + &self, + key: &str, + candidate: V, + eligible: &[V], + kwargs: CacheKwargs, + ) -> Result { + let now = (self.now)(); + let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + Self::evict(&mut state, self.max_size_in_memory, now); + let winner = match state.values.get(key) { + Some(existing) if eligible.is_empty() => existing.clone(), + Some(existing) if eligible.contains(existing) => existing.clone(), + _ => candidate, + }; + let expiration = now + self.get_ttl(&kwargs); + state.values.insert(key.into(), winner.clone()); + state.expirations.insert(key.into(), expiration); + state + .expiration_heap + .push(Reverse((expiration, key.into()))); + Ok(winner) + } +} + +impl CounterCache for InMemoryCache { + fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + let now = (self.now)(); + let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; + Self::evict(&mut state, self.max_size_in_memory, now); + let value = state.values.get(key).copied().unwrap_or_default() + amount; + let expiration = state + .expirations + .get(key) + .copied() + .unwrap_or_else(|| now + self.get_ttl(&kwargs)); + state.values.insert(key.into(), value); + state.expirations.insert(key.into(), expiration); + state + .expiration_heap + .push(Reverse((expiration, key.into()))); + Ok(value) + } +} + impl BaseCache for InMemoryCache { type Value = V; diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index 370145bebef..e5831dfd6d5 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -3,7 +3,8 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use litellm_cache::{ - BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, Error, get_cache, set_cache, + BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, Error, + get_cache, set_cache, }; use litellm_cache_memory::{CacheWrite, InMemoryCache}; use rstest::{fixture, rstest}; @@ -183,3 +184,51 @@ async fn generic_consumers_share_typed_values_and_honor_expiration() { None ); } + +#[test] +fn claims_are_atomic_and_refresh_eligible_winners() { + let clock = clock(); + let cache = InMemoryCache::with_clock(Some(4), Some(Duration::from_secs(60)), { + let clock = clock.clone(); + move || Duration::from_secs(clock.load(Ordering::SeqCst)) + }); + let kwargs = CacheKwargs { + ttl: Some(Duration::from_secs(10)), + ..Default::default() + }; + assert_eq!( + cache + .claim_cache("affinity", "first".to_string(), &[], kwargs.clone()) + .unwrap(), + "first" + ); + clock.store(105, Ordering::SeqCst); + assert_eq!( + cache + .claim_cache( + "affinity", + "second".to_string(), + &["first".to_string(), "second".to_string()], + kwargs, + ) + .unwrap(), + "first" + ); + assert_eq!( + cache.expires_at("affinity").unwrap(), + Some(Duration::from_secs(115)) + ); +} + +#[test] +fn counters_increment_under_one_lock() { + let cache = InMemoryCache::::default(); + assert_eq!( + CounterCache::increment_cache(&cache, "counter", 1.5, CacheKwargs::default()).unwrap(), + 1.5 + ); + assert_eq!( + CounterCache::increment_cache(&cache, "counter", 2.0, CacheKwargs::default()).unwrap(), + 3.5 + ); +} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index a60813b6260..1a1bb505a7c 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,7 +7,8 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis = "1.7.0" +redis = { version = "1.7.0", features = ["r2d2"] } +r2d2 = "0.8.10" tokio.workspace = true [dev-dependencies] diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index d4ca0cf0522..281ffe17534 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -1,15 +1,83 @@ -use std::sync::{Arc, Mutex, MutexGuard}; +use std::sync::{Arc, Mutex}; use std::time::Duration; use litellm_cache::{ - BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, Error, + BaseCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, + ClaimCache, CounterCache, Error, }; use redis::Commands; const DEFAULT_TTL: Duration = Duration::from_secs(600); +const REDIS_TIMEOUT: Duration = Duration::from_secs(5); +const REDIS_POOL_SIZE: u32 = 16; + +enum Connections { + Pool(r2d2::Pool), + Fixed(Mutex), +} + +struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); + +impl redis::ConnectionLike for ConnectionRef<'_> { + fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { + self.0.req_packed_command(cmd) + } + + fn req_packed_commands( + &mut self, + cmd: &[u8], + offset: usize, + count: usize, + ) -> redis::RedisResult> { + self.0.req_packed_commands(cmd, offset, count) + } + + fn get_db(&self) -> i64 { + self.0.get_db() + } + + fn supports_pipelining(&self) -> bool { + self.0.supports_pipelining() + } + + fn check_connection(&mut self) -> bool { + self.0.check_connection() + } + + fn is_open(&self) -> bool { + self.0.is_open() + } +} + +impl Connections +where + C: redis::ConnectionLike + Send + 'static, +{ + fn execute( + &self, + operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, + ) -> Result { + match self { + Self::Pool(pool) => { + let mut connection = pool.get().map_err(|_| Error::Unavailable)?; + connection + .set_read_timeout(Some(REDIS_TIMEOUT)) + .map_err(|_| Error::Unavailable)?; + connection + .set_write_timeout(Some(REDIS_TIMEOUT)) + .map_err(|_| Error::Unavailable)?; + operation(&mut ConnectionRef(&mut *connection)) + } + Self::Fixed(connection) => { + let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; + operation(&mut ConnectionRef(&mut *connection)) + } + } + } +} pub struct RedisCache { - connection: Arc>, + connections: Arc>, default_ttl: Duration, codec: S, namespace: Option, @@ -18,8 +86,18 @@ pub struct RedisCache { 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, codec)) + let pool = r2d2::Pool::builder() + .max_size(REDIS_POOL_SIZE) + .min_idle(Some(0)) + .connection_timeout(REDIS_TIMEOUT) + .build(client) + .map_err(|_| Error::Unavailable)?; + Ok(Self { + connections: Arc::new(Connections::Pool(pool)), + default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), + codec, + namespace: None, + }) } } @@ -30,17 +108,13 @@ where { pub fn with_connection(connection: C, default_ttl: Option, codec: S) -> Self { Self { - connection: Arc::new(Mutex::new(connection)), + connections: Arc::new(Connections::Fixed(Mutex::new(connection))), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, namespace: None, } } - fn connection(&self) -> Result, Error> { - self.connection.lock().map_err(|_| Error::Unavailable) - } - pub fn with_namespace(self, namespace: Option) -> Self { Self { namespace: namespace.filter(|value| !value.is_empty()), @@ -72,6 +146,29 @@ where Ok(format!("{escaped}:*")) } + fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> { + let mut cursor = 0u64; + loop { + let (next_cursor, keys): (u64, Vec) = redis::cmd("SCAN") + .cursor_arg(cursor) + .arg("MATCH") + .arg(pattern) + .arg("COUNT") + .arg(1000) + .query(connection) + .map_err(|_| Error::Unavailable)?; + if !keys.is_empty() { + connection + .del::<_, usize>(keys) + .map_err(|_| Error::Unavailable)?; + } + if next_cursor == 0 { + return Ok(()); + } + cursor = next_cursor; + } + } + fn decode_response(&self, value: redis::Value) -> Result, Error> { match value { redis::Value::Nil => Ok(None), @@ -81,23 +178,29 @@ where } } + fn decode_batch_response(&self, value: redis::Value) -> Result, Error> { + match self.decode_response(value) { + Ok(Some(value)) => Ok(BatchEntry::Hit(value)), + Ok(None) => Ok(BatchEntry::Miss), + Err(Error::InvalidEntry) => Ok(BatchEntry::Invalid), + Err(error) => Err(error), + } + } + fn ttl_seconds(ttl: Duration) -> u64 { ttl.as_secs() .saturating_add(u64::from(ttl.subsec_nanos() > 0)) .max(1) } - async fn run_blocking(connection: Arc>, operation: F) -> Result + async fn run_blocking(connections: Arc>, operation: F) -> Result where T: Send + 'static, - F: FnOnce(&mut C) -> Result + Send + 'static, + F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, { - 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 || connections.execute(operation)) + .await + .map_err(|_| Error::Unavailable)? } } @@ -115,40 +218,55 @@ where fn set_cache(&self, key: &str, value: Self::Value, kwargs: CacheKwargs) -> Result<(), Error> { 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) - .map_err(|_| Error::Unavailable) + let key = self.namespaced_key(key); + self.connections.execute(|connection| { + connection + .set_ex::<_, _, ()>(key, payload, ttl) + .map_err(|_| Error::Unavailable) + }) } fn get_cache(&self, key: &str, _: &CacheKwargs) -> Result, Error> { - let value = self - .connection()? - .get::<_, redis::Value>(self.namespaced_key(key)) - .map_err(|_| Error::Unavailable)?; + let key = self.namespaced_key(key); + let value = self.connections.execute(|connection| { + connection + .get::<_, redis::Value>(key) + .map_err(|_| Error::Unavailable) + })?; 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> { - self.connection()? - .del::<_, ()>(self.namespaced_key(key)) - .map_err(|_| Error::Unavailable) + 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()?; - let mut connection = self.connection()?; - let keys = connection - .scan_match(pattern) - .map_err(|_| Error::Unavailable)? - .collect::>>() - .map_err(|_| Error::Unavailable)?; - if keys.is_empty() { - return Ok(()); - } - connection - .del::<_, usize>(keys) - .map(|_| ()) - .map_err(|_| Error::Unavailable) + self.connections + .execute(|connection| Self::flush_matching(connection, &pattern)) } async fn async_set_cache( @@ -160,7 +278,7 @@ where 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| { + Self::run_blocking(Arc::clone(&self.connections), move |connection| { connection .set_ex::<_, _, ()>(key, payload, ttl) .map_err(|_| Error::Unavailable) @@ -174,7 +292,7 @@ where _: &CacheKwargs, ) -> Result, Error> { let key = self.namespaced_key(key); - let value = Self::run_blocking(Arc::clone(&self.connection), move |connection| { + let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { connection .get::<_, redis::Value>(key) .map_err(|_| Error::Unavailable) @@ -183,6 +301,28 @@ 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)>, @@ -197,41 +337,137 @@ where }) .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); - Self::run_blocking(Arc::clone(&self.connection), move |connection| { + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut pipeline = redis::pipe(); for (key, payload) in entries { - connection - .set_ex::<_, _, ()>(key, payload, ttl) - .map_err(|_| Error::Unavailable)?; + pipeline + .cmd("SETEX") + .arg(key) + .arg(ttl) + .arg(payload) + .ignore(); } - Ok(()) + pipeline + .query::<()>(connection) + .map_err(|_| Error::Unavailable) }) .await } 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| { + 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(()) } async fn test_connection(&self) -> Result { - Self::run_blocking(Arc::clone(&self.connection), |connection| { - redis::cmd("PING") - .query::(connection) + match Self::run_blocking(Arc::clone(&self.connections), |connection| { + Ok(match redis::cmd("PING").query::(connection) { + Ok(_) => CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, + }, + Err(error) => CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }, + }) + }) + .await + { + Ok(result) => Ok(result), + Err(error) => Ok(CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }), + } + } +} + +impl CounterCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + const SCRIPT: &str = concat!( + "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", + "if redis.call('TTL', KEYS[1]) == -1 then ", + "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" + ); + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + self.connections.execute(|connection| { + redis::cmd("EVAL") + .arg(SCRIPT) + .arg(1) + .arg(key) + .arg(amount) + .arg(ttl) + .query(connection) .map_err(|_| Error::Unavailable) }) - .await?; - Ok(CacheConnectionResult { - status: CacheConnectionStatus::Success, - message: "Redis cache connection test successful".into(), - error: None, - }) + } +} + +impl ClaimCache for RedisCache +where + S: CacheCodec, + S::Value: PartialEq, + C: redis::ConnectionLike + Send + 'static, +{ + fn claim_cache( + &self, + key: &str, + candidate: S::Value, + eligible: &[S::Value], + kwargs: CacheKwargs, + ) -> Result { + const SCRIPT: &str = concat!( + "local current = redis.call('GET', KEYS[1]); ", + "if current == false then redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ", + "return ARGV[1]; end; if #ARGV > 2 then for index = 3, #ARGV do ", + "if current == ARGV[index] then redis.call('EXPIRE', KEYS[1], ARGV[2]); ", + "return current; end; end; redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ", + "return ARGV[1]; end; if current == ARGV[1] then ", + "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return current" + ); + let key = self.namespaced_key(key); + let candidate = self.codec.encode(&candidate)?; + let eligible = eligible + .iter() + .map(|value| self.codec.encode(value)) + .collect::, _>>()?; + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let value = self.connections.execute(|connection| { + redis::cmd("EVAL") + .arg(SCRIPT) + .arg(1) + .arg(key) + .arg(candidate) + .arg(ttl) + .arg(eligible) + .query::(connection) + .map_err(|_| Error::Unavailable) + })?; + self.decode_response(value)?.ok_or(Error::Unavailable) } } @@ -302,7 +538,9 @@ mod tests { redis::cmd("SCAN") .cursor_arg(0) .arg("MATCH") - .arg("litellm-cache:*"), + .arg("litellm-cache:*") + .arg("COUNT") + .arg(1000), Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])), ), MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index d5bba19a8bd..6aef9bb36bf 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -1,6 +1,9 @@ use std::time::Duration; -use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, Error, JsonCodec, get_cache, set_cache}; +use litellm_cache::{ + BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, Error, JsonCodec, + get_cache, set_cache, +}; use litellm_cache_redis::RedisCache; use redis_test::{MockCmd, MockRedisConnection}; @@ -174,7 +177,9 @@ fn flush_requires_a_namespace_and_escapes_glob_metacharacters() { redis::cmd("SCAN") .cursor_arg(0) .arg("MATCH") - .arg("team\\*:*"), + .arg("team\\*:*") + .arg("COUNT") + .arg(1000), Ok(redis_test::redis_value!(["0", ["team*:key"]])), ), MockCmd::new(redis::cmd("DEL").arg("team*:key"), Ok(1u32)), @@ -184,3 +189,73 @@ fn flush_requires_a_namespace_and_escapes_glob_metacharacters() { .with_namespace(Some("team*".into())); scoped.flush_cache().unwrap(); } + +#[tokio::test] +async fn connection_failures_use_the_python_result_contract() { + let error = redis::RedisError::from((redis::ErrorKind::Io, "connection refused")); + let connection = + MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Err::(error))]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); + + let result = cache.test_connection().await.unwrap(); + assert_eq!(result.status, CacheConnectionStatus::Failed); + assert!(result.message.starts_with("Redis connection failed:")); + assert!(result.error.is_some()); +} + +#[tokio::test] +async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries() { + let connection = MockRedisConnection::new([MockCmd::new( + redis::cmd("MGET").arg("hit").arg("miss").arg("invalid"), + Ok(vec![ + redis::Value::BulkString(vec![42, 7]), + redis::Value::Nil, + redis::Value::BulkString(vec![99, 7]), + ]), + )]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); + + assert_eq!( + cache + .async_get_cache_batch( + vec!["hit".into(), "miss".into(), "invalid".into()], + CacheKwargs::default(), + ) + .await + .unwrap(), + vec![BatchEntry::Hit(7), BatchEntry::Miss, BatchEntry::Invalid] + ); +} + +#[tokio::test] +async fn async_flush_deletes_each_scan_page_separately() { + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(0) + .arg("MATCH") + .arg("team:*") + .arg("COUNT") + .arg(1000), + Ok(redis_test::redis_value!(["7", ["team:a", "team:b"]])), + ), + MockCmd::new(redis::cmd("DEL").arg("team:a").arg("team:b"), Ok(2u32)), + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(7) + .arg("MATCH") + .arg("team:*") + .arg("COUNT") + .arg(1000), + Ok(redis_test::redis_value!(["0", ["team:c"]])), + ), + MockCmd::new(redis::cmd("DEL").arg("team:c"), Ok(1u32)), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + + cache.async_flush_cache().await.unwrap(); +} diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 309296f7773..4e6694c191f 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -28,17 +28,17 @@ cache.store(&request, json!({"answer": 7}), now)?; assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7}))); ``` -For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Each Redis constructor currently opens its own connection; shared connection pools remain follow-up work +For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers move that blocking work off the executor Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it ## Python integration boundary -The extension exposes `NativeCacheHandle`, `CacheResolver`, and captured `CacheBinding` objects for host integration. Memory and Redis handles support synchronous and asynchronous response lookup and storage +The extension exposes `NativeCacheHandle`, `CacheResolver`, and captured `CacheBinding` objects for host integration. Memory and Redis handles support single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution -Explicit facade registration checks object identity, method overrides, and configuration changes before selecting native execution. Registration does not migrate entries or replace Python methods. Until activation configures one shared service, a registered facade and its native handle can hold separate data. Existing public cache constructors remain on Python +Explicit facade registration checks object identity, method overrides, effective TTL, and configuration changes before selecting native execution. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. Registration does not migrate entries or replace Python methods. Until activation configures one shared service, a registered facade and its native handle can hold separate data. Existing public cache constructors remain on Python Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy @@ -50,6 +50,6 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na ## Follow-up scope -Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial batches, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths +Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths Redis cluster, disk, cloud stores, dual caching, and semantic caching remain follow-ups. Atomic counters, affinity claims, reservations, queues, and pubsub need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees diff --git a/litellm-rust/crates/cache-response/src/caching.rs b/litellm-rust/crates/cache-response/src/caching.rs index 53b34025ccd..afae4dfe4a5 100644 --- a/litellm-rust/crates/cache-response/src/caching.rs +++ b/litellm-rust/crates/cache-response/src/caching.rs @@ -120,6 +120,7 @@ impl CacheControls { pub fn writes(self) -> bool { self.supported_call_type && self.configured + && self.caching.unwrap_or(true) && !self.no_store && (self.default_on || self.use_cache) } @@ -131,13 +132,16 @@ pub fn should_use_cache(controls: CacheControls) -> bool { #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct CacheEntry { - pub timestamp: f64, + #[serde(skip_serializing_if = "Option::is_none")] + pub timestamp: Option, pub response: Value, } impl CacheEntry { pub fn fresh(&self, now: Duration, max_age: Option) -> bool { - self.timestamp.is_finite() - && max_age.is_none_or(|age| now.as_secs_f64() - self.timestamp <= age.as_secs_f64()) + self.timestamp.is_none_or(|timestamp| { + timestamp.is_finite() + && max_age.is_none_or(|age| now.as_secs_f64() - timestamp <= age.as_secs_f64()) + }) } } diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs index f1d55ddefe8..eaba3d0c349 100644 --- a/litellm-rust/crates/cache-response/src/codec.rs +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -9,7 +9,10 @@ impl CacheCodec for ResponseCacheCodec { type Value = CacheEntry; fn encode(&self, value: &CacheEntry) -> Result, Error> { - if !value.timestamp.is_finite() { + if value + .timestamp + .is_some_and(|timestamp| !timestamp.is_finite()) + { return Err(Error::InvalidEntry); } serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) @@ -17,12 +20,21 @@ impl CacheCodec for ResponseCacheCodec { fn decode(&self, bytes: &[u8]) -> Result { let text = std::str::from_utf8(bytes).map_err(|_| Error::InvalidEntry)?; - let entry: CacheEntry = - serde_json::from_value(decode_value(text)?).map_err(|_| Error::InvalidEntry)?; - if !entry.timestamp.is_finite() { + let value = decode_value(text)?; + let Some(timestamp) = value.get("timestamp") else { + return Ok(CacheEntry { + timestamp: None, + response: value, + }); + }; + let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else { return Err(Error::InvalidEntry); - } - Ok(entry) + }; + let response = value.get("response").cloned().ok_or(Error::InvalidEntry)?; + Ok(CacheEntry { + timestamp: Some(timestamp), + response, + }) } } diff --git a/litellm-rust/crates/cache-response/src/embedding.rs b/litellm-rust/crates/cache-response/src/embedding.rs new file mode 100644 index 00000000000..d1f8a2bc0a6 --- /dev/null +++ b/litellm-rust/crates/cache-response/src/embedding.rs @@ -0,0 +1,22 @@ +use serde::Serialize; +use serde_json::Value; + +#[derive(Clone, Debug, PartialEq, Serialize)] +pub struct PartialHits { + pub values: Vec>, + pub missing_indices: Vec, +} + +impl PartialHits { + pub fn new(values: Vec>) -> Self { + let missing_indices = values + .iter() + .enumerate() + .filter_map(|(index, value)| value.is_none().then_some(index)) + .collect(); + Self { + values, + missing_indices, + } + } +} diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index efa0b04b9f7..72a507f8ee9 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -1,5 +1,6 @@ mod caching; mod codec; +mod embedding; mod response; pub use caching::{ @@ -7,4 +8,5 @@ pub use caching::{ get_cache_key, should_use_cache, }; pub use codec::ResponseCacheCodec; +pub use embedding::PartialHits; pub use response::{ResponseCache, ResponseCacheRequest}; diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index 987a47f0554..8f0fe953df6 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,8 +1,8 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{BaseCache, CacheKwargs, Error}; +use litellm_cache::{BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, Error}; -use crate::{CacheControls, CacheEntry, CacheKeyInput, cache_key}; +use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; use serde_json::Value; #[derive(Clone)] @@ -39,6 +39,18 @@ impl> ResponseCache { Self { backend } } + pub fn default_ttl(&self) -> Duration { + self.backend.default_ttl() + } + + pub async fn async_flush(&self) -> Result<(), Error> { + self.backend.async_flush_cache().await + } + + pub async fn test_connection(&self) -> Result { + self.backend.test_connection().await + } + pub fn lookup( &self, request: &ResponseCacheRequest, @@ -47,10 +59,15 @@ impl> ResponseCache { if !request.controls.reads() { return Ok(None); } - let entry = self + let entry = match self .backend - .get_cache(&cache_key(&request.key), &request.kwargs)?; - Self::fresh_response(entry, now, request.max_age) + .get_cache(&cache_key(&request.key), &request.kwargs) + { + Ok(entry) => entry, + Err(Error::InvalidEntry) => None, + Err(error) => return Err(error), + }; + Self::fresh_or_miss(entry, now, request.max_age) } pub async fn async_lookup( @@ -61,11 +78,62 @@ impl> ResponseCache { if !request.controls.reads() { return Ok(None); } - let entry = self + let entry = match self .backend .async_get_cache(&cache_key(&request.key), &request.kwargs) - .await?; - Self::fresh_response(entry, now, request.max_age) + .await + { + Ok(entry) => entry, + Err(Error::InvalidEntry) => None, + Err(error) => return Err(error), + }; + Self::fresh_or_miss(entry, now, request.max_age) + } + + pub fn lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + let readable = requests + .iter() + .enumerate() + .filter(|(_, request)| request.controls.reads()) + .collect::>(); + let keys = readable + .iter() + .map(|(_, request)| cache_key(&request.key)) + .collect::>(); + let entries = if let Some((_, request)) = readable.first() { + self.backend.get_cache_batch(&keys, &request.kwargs)? + } else { + Vec::new() + }; + Self::partial_hits(requests, readable, entries, now) + } + + pub async fn async_lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + let readable = requests + .iter() + .enumerate() + .filter(|(_, request)| request.controls.reads()) + .collect::>(); + let keys = readable + .iter() + .map(|(_, request)| cache_key(&request.key)) + .collect::>(); + let entries = if let Some((_, request)) = readable.first() { + self.backend + .async_get_cache_batch(keys, request.kwargs.clone()) + .await? + } else { + Vec::new() + }; + Self::partial_hits(requests, readable, entries, now) } pub fn store( @@ -80,7 +148,7 @@ impl> ResponseCache { self.backend.set_cache( &cache_key(&request.key), CacheEntry { - timestamp: now.as_secs_f64(), + timestamp: Some(now.as_secs_f64()), response, }, request.kwargs.clone(), @@ -100,7 +168,7 @@ impl> ResponseCache { .async_set_cache( &cache_key(&request.key), CacheEntry { - timestamp: now.as_secs_f64(), + timestamp: Some(now.as_secs_f64()), response, }, request.kwargs.clone(), @@ -108,6 +176,76 @@ impl> ResponseCache { .await } + pub async fn async_store_batch( + &self, + entries: Vec<(ResponseCacheRequest, Value)>, + now: Duration, + ) -> Result<(), Error> { + let writable = entries + .into_iter() + .filter(|(request, _)| request.controls.writes()) + .map(|(request, response)| { + ( + cache_key(&request.key), + CacheEntry { + timestamp: Some(now.as_secs_f64()), + response, + }, + request.kwargs, + ) + }) + .collect::>(); + let Some((_, _, first_kwargs)) = writable.first() else { + return Ok(()); + }; + if writable.iter().all(|(_, _, kwargs)| kwargs == first_kwargs) { + let kwargs = 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) + .await; + } + for (key, entry, kwargs) in writable { + self.backend.async_set_cache(&key, entry, kwargs).await?; + } + Ok(()) + } + + fn partial_hits( + requests: &[ResponseCacheRequest], + readable: Vec<(usize, &ResponseCacheRequest)>, + entries: Vec>, + now: Duration, + ) -> Result { + if readable.len() != entries.len() { + return Err(Error::Unavailable); + } + let mut values = vec![None; requests.len()]; + for ((index, request), entry) in readable.into_iter().zip(entries) { + let response = match entry { + BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age)?, + BatchEntry::Miss | BatchEntry::Invalid => None, + }; + values[index] = response; + } + Ok(PartialHits::new(values)) + } + + fn fresh_or_miss( + entry: Option, + now: Duration, + max_age: Option, + ) -> Result, Error> { + match Self::fresh_response(entry, now, max_age) { + Err(Error::InvalidEntry) => Ok(None), + result => result, + } + } + fn fresh_response( entry: Option, now: Duration, @@ -115,9 +253,9 @@ impl> ResponseCache { ) -> Result, Error> { entry .filter(|entry| entry.fresh(now, max_age)) - .map(|entry| match entry.response { - Value::String(text) => crate::codec::decode_value(&text), - value => Ok(value), + .map(|entry| match (entry.timestamp, entry.response) { + (Some(_), Value::String(text)) => crate::codec::decode_value(&text), + (_, value) => Ok(value), }) .transpose() } diff --git a/litellm-rust/crates/cache-response/tests/caching.rs b/litellm-rust/crates/cache-response/tests/caching.rs index d403791e421..0e8ce9b3b1d 100644 --- a/litellm-rust/crates/cache-response/tests/caching.rs +++ b/litellm-rust/crates/cache-response/tests/caching.rs @@ -80,4 +80,11 @@ fn cache_controls_honor_default_modes_and_directives() { } .writes() ); + assert!( + !CacheControls { + caching: Some(false), + ..enabled + } + .writes() + ); } diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 88b53fdfe86..e4a04b3dec5 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -245,7 +245,7 @@ fn response_codec_accepts_python_literals_without_executing_code() { assert_eq!( ResponseCacheCodec .encode(&CacheEntry { - timestamp: f64::NAN, + timestamp: Some(f64::NAN), response: json!({}) }) .unwrap_err(), @@ -254,7 +254,7 @@ fn response_codec_accepts_python_literals_without_executing_code() { } #[tokio::test] -async fn backend_failures_remain_observable_and_disabled_reads_do_not_touch_redis() { +async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis() { let connection = MockRedisConnection::new([MockCmd::new( redis::cmd("GET").arg("tenant:key"), Ok(b"invalid".to_vec()), @@ -267,22 +267,19 @@ async fn backend_failures_remain_observable_and_disabled_reads_do_not_touch_redi assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None); request.controls.no_cache = false; assert_eq!( - cache - .async_lookup(&request, Duration::ZERO) - .await - .unwrap_err(), - Error::InvalidEntry + cache.async_lookup(&request, Duration::ZERO).await.unwrap(), + None ); } #[test] -fn malformed_memory_entries_are_rejected_by_the_response_consumer() { +fn malformed_memory_entries_are_treated_as_misses() { let backend = Arc::new(InMemoryCache::default()); BaseCache::set_cache( backend.as_ref(), "tenant:key", CacheEntry { - timestamp: 100.0, + timestamp: Some(100.0), response: json!("not a serialized response"), }, Default::default(), @@ -290,10 +287,8 @@ fn malformed_memory_entries_are_rejected_by_the_response_consumer() { .unwrap(); let cache = ResponseCache::new(backend); assert_eq!( - cache - .lookup(&request(), Duration::from_secs(100)) - .unwrap_err(), - Error::InvalidEntry + cache.lookup(&request(), Duration::from_secs(100)).unwrap(), + None ); } @@ -301,10 +296,74 @@ fn malformed_memory_entries_are_rejected_by_the_response_consumer() { fn response_entries_preserve_the_existing_json_representation() { let codec = ResponseCacheCodec; let entry = CacheEntry { - timestamp: 123.0, + timestamp: Some(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 response_codec_preserves_values_without_timestamps() { + let codec = ResponseCacheCodec; + let raw = json!({"choices": [{"text": "legacy"}]}); + let entry = codec.decode(&serde_json::to_vec(&raw).unwrap()).unwrap(); + assert_eq!(entry.timestamp, None); + assert_eq!(entry.response, raw); + + let backend = Arc::new(InMemoryCache::default()); + 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(), + Some(json!({"choices": [{"text": "legacy"}]})) + ); +} + +#[tokio::test] +async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses() { + let cache = memory(); + let requests = ["hit", "miss", "disabled"].map(|key| { + ResponseCacheRequest::new(CacheKeyInput { + preset: Some(key.into()), + ..Default::default() + }) + }); + cache + .store(&requests[0], json!({"value": 1}), Duration::from_secs(100)) + .unwrap(); + let mut requests = requests.to_vec(); + requests[2].controls.caching = Some(false); + + let partial = cache + .async_lookup_batch(&requests, Duration::from_secs(100)) + .await + .unwrap(); + assert_eq!(partial.values, vec![Some(json!({"value": 1})), None, None]); + assert_eq!(partial.missing_indices, vec![1, 2]); + + cache + .async_store_batch( + vec![ + (requests[1].clone(), json!({"value": 2})), + (requests[2].clone(), json!({"value": 3})), + ], + Duration::from_secs(100), + ) + .await + .unwrap(); + assert_eq!( + cache + .lookup(&requests[1], Duration::from_secs(100)) + .unwrap(), + Some(json!({"value": 2})) + ); + requests[2].controls.caching = None; + assert_eq!( + cache + .lookup(&requests[2], Duration::from_secs(100)) + .unwrap(), + None + ); +} diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 1891e417cb0..5bc1ebc2945 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -6,6 +6,13 @@ use serde_json::{Map, Value}; use crate::Error; +#[derive(Clone, Debug, PartialEq)] +pub enum BatchEntry { + Hit(V), + Miss, + Invalid, +} + #[derive(Clone, Debug, Default, PartialEq)] pub struct CacheKwargs { pub ttl: Option, @@ -42,6 +49,21 @@ pub trait BaseCache: Send + Sync { fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error>; + fn get_cache_batch( + &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() + } + fn async_set_cache( &self, key: &str, @@ -59,6 +81,25 @@ pub trait BaseCache: Send + Sync { 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) + } + } + fn async_set_cache_pipeline( &self, cache_list: Vec<(String, Self::Value)>, @@ -89,6 +130,10 @@ pub trait BaseCache: Send + Sync { fn flush_cache(&self) -> Result<(), Error>; + fn async_flush_cache(&self) -> impl Future> + Send { + async move { self.flush_cache() } + } + fn disconnect(&self) -> impl Future> + Send; fn test_connection(&self) -> impl Future> + Send; diff --git a/litellm-rust/crates/cache/src/capabilities.rs b/litellm-rust/crates/cache/src/capabilities.rs new file mode 100644 index 00000000000..0b9deab1f5a --- /dev/null +++ b/litellm-rust/crates/cache/src/capabilities.rs @@ -0,0 +1,39 @@ +use std::future::Future; + +use crate::{BaseCache, CacheKwargs, Error}; + +pub trait CounterCache: BaseCache { + fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result; + + fn async_increment_cache( + &self, + key: &str, + amount: f64, + kwargs: CacheKwargs, + ) -> impl Future> + Send { + async move { self.increment_cache(key, amount, kwargs) } + } +} + +pub trait ClaimCache: BaseCache +where + Self::Value: PartialEq, +{ + fn claim_cache( + &self, + key: &str, + candidate: Self::Value, + eligible: &[Self::Value], + kwargs: CacheKwargs, + ) -> Result; + + fn async_claim_cache( + &self, + key: &str, + candidate: Self::Value, + eligible: Vec, + kwargs: CacheKwargs, + ) -> impl Future> + Send { + async move { self.claim_cache(key, candidate, &eligible, kwargs) } + } +} diff --git a/litellm-rust/crates/cache/src/dual.rs b/litellm-rust/crates/cache/src/dual.rs new file mode 100644 index 00000000000..17a2c430ddd --- /dev/null +++ b/litellm-rust/crates/cache/src/dual.rs @@ -0,0 +1,105 @@ +use std::sync::Arc; + +use crate::{BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error}; + +pub struct DualCache { + l1: Arc, + l2: Arc, +} + +impl DualCache { + pub fn new(l1: Arc, l2: Arc) -> Self { + Self { l1, l2 } + } +} + +impl BaseCache for DualCache +where + V: Clone + Send + Sync + 'static, + L1: BaseCache, + L2: BaseCache, +{ + type Value = V; + + fn default_ttl(&self) -> std::time::Duration { + self.l2.default_ttl() + } + + fn set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> { + self.l2.set_cache(key, value.clone(), kwargs.clone())?; + self.l1.set_cache(key, value, kwargs) + } + + fn get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error> { + if let Some(value) = self.l1.get_cache(key, kwargs)? { + return Ok(Some(value)); + } + let value = self.l2.get_cache(key, kwargs)?; + if let Some(value) = &value { + self.l1.set_cache(key, value.clone(), kwargs.clone())?; + } + Ok(value) + } + + fn delete_cache(&self, key: &str) -> Result<(), Error> { + self.l2.delete_cache(key)?; + self.l1.delete_cache(key) + } + + fn flush_cache(&self) -> Result<(), Error> { + self.l2.flush_cache()?; + self.l1.flush_cache() + } + + async fn async_flush_cache(&self) -> Result<(), Error> { + self.l2.async_flush_cache().await?; + 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 +where + 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)?; + Ok(value) + } +} + +impl ClaimCache for DualCache +where + V: Clone + PartialEq + Send + Sync + 'static, + L1: ClaimCache, + L2: ClaimCache, +{ + fn claim_cache( + &self, + key: &str, + candidate: V, + eligible: &[V], + kwargs: CacheKwargs, + ) -> Result { + match self + .l2 + .claim_cache(key, candidate.clone(), eligible, kwargs.clone()) + { + Ok(winner) => { + self.l1.set_cache(key, winner.clone(), kwargs)?; + Ok(winner) + } + Err(_) => self.l1.claim_cache(key, candidate, eligible, kwargs), + } + } +} diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index 4ff02319bdc..ed67eb2fe15 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -1,9 +1,14 @@ mod base_cache; mod caching; +mod capabilities; mod codec; +pub mod dual; mod error; -pub use base_cache::{BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs}; +pub use base_cache::{ + BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, +}; pub use caching::{Cache, CacheBackend, get_cache, set_cache}; +pub use capabilities::{ClaimCache, CounterCache}; pub use codec::{CacheCodec, JsonCodec}; pub use error::Error; diff --git a/litellm-rust/crates/cache/tests/dual.rs b/litellm-rust/crates/cache/tests/dual.rs new file mode 100644 index 00000000000..1be1556734c --- /dev/null +++ b/litellm-rust/crates/cache/tests/dual.rs @@ -0,0 +1,130 @@ +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{ + BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, dual::DualCache, +}; + +struct TestCache { + value: Mutex>, + fail: bool, +} + +impl TestCache { + fn new(value: Option, fail: bool) -> Self { + Self { + value: Mutex::new(value), + fail, + } + } +} + +impl BaseCache for TestCache +where + V: Clone + Send + Sync + 'static, +{ + type Value = V; + + fn set_cache(&self, _: &str, value: V, _: CacheKwargs) -> Result<(), Error> { + *self.value.lock().unwrap() = Some(value); + Ok(()) + } + + fn get_cache(&self, _: &str, _: &CacheKwargs) -> 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(()) + } + + async fn test_connection(&self) -> Result { + unreachable!() + } +} + +impl CounterCache for TestCache { + fn increment_cache(&self, _: &str, amount: f64, _: CacheKwargs) -> Result { + if self.fail { + return Err(Error::Unavailable); + } + let mut value = self.value.lock().unwrap(); + let incremented = value.unwrap_or_default() + amount; + *value = Some(incremented); + Ok(incremented) + } +} + +impl ClaimCache for TestCache +where + V: Clone + PartialEq + Send + Sync + 'static, +{ + fn claim_cache( + &self, + _: &str, + candidate: V, + eligible: &[V], + _: CacheKwargs, + ) -> Result { + if self.fail { + return Err(Error::Unavailable); + } + let mut value = self.value.lock().unwrap(); + let winner = match value.as_ref() { + Some(existing) if eligible.is_empty() || eligible.contains(existing) => { + existing.clone() + } + _ => candidate, + }; + *value = Some(winner.clone()); + Ok(winner) + } +} + +#[test] +fn failed_l2_increment_leaves_l1_unchanged() { + let l1 = Arc::new(TestCache::new(Some(10.0), false)); + let cache = DualCache::new(l1.clone(), Arc::new(TestCache::new(Some(20.0), true))); + + assert_eq!( + cache.increment_cache("counter", 2.0, CacheKwargs::default()), + Err(Error::Unavailable) + ); + assert_eq!( + l1.get_cache("counter", &CacheKwargs::default()).unwrap(), + Some(10.0) + ); +} + +#[test] +fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() { + let l1 = Arc::new(TestCache::new(Some("first".to_string()), false)); + let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true))); + + assert_eq!( + cache + .claim_cache( + "affinity", + "second".into(), + &["first".into(), "second".into()], + CacheKwargs { + ttl: Some(Duration::from_secs(60)), + ..Default::default() + }, + ) + .unwrap(), + "first" + ); +} diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index 0af55083bef..15d0d603ac7 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -22,5 +22,8 @@ ], "secret_manager": [ "readable" + ], + "cache_settings": [ + "default_redis_ttl" ] } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index eb07118e964..58550c2987d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -6,6 +6,7 @@ use pyo3::{ types::{PyDict, PyTuple, PyType}, }; use serde_json::Value; +use std::time::Duration; use super::{NativeCacheHandle, native::NativeResponseCache}; @@ -126,7 +127,12 @@ impl ObjectGuard { } impl FacadeGuard { - pub(super) fn capture(py: Python<'_>, facade: &Bound<'_, PyAny>, kind: &str) -> PyResult { + pub(super) fn capture( + py: Python<'_>, + facade: &Bound<'_, PyAny>, + kind: &str, + native_default_ttl: Duration, + ) -> PyResult { let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?; if !facade.get_type().is(&cache_type) { return Err(PyTypeError::new_err( @@ -146,6 +152,12 @@ impl FacadeGuard { "facade and native backend types must match", )); } + let python_default_ttl = backend.getattr("default_ttl")?.extract::()?; + if python_default_ttl != native_default_ttl.as_secs_f64() { + return Err(PyTypeError::new_err( + "facade and native backend default TTLs must match", + )); + } Ok(Self { outer: ObjectGuard::capture( py, diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 5918967009a..4bab03fe243 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -4,7 +4,7 @@ mod native; use std::time::{Duration, SystemTime, UNIX_EPOCH}; use litellm_cache::Error; -use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; +use litellm_cache_response::{CacheControls, CacheKeyInput, PartialHits, ResponseCacheRequest}; use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; use pyo3::{ PyTraverseError, PyVisit, @@ -15,9 +15,26 @@ use pyo3::{ use serde::Deserialize; use serde_json::Value; +use crate::python_settings::PythonSettings; use facade::FacadeGuard; use native::NativeResponseCache; +const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); + +#[derive(FromPyObject)] +struct PythonCacheSettings { + default_redis_ttl: Option, +} + +fn redis_default_ttl(py: Python<'_>) -> PyResult { + let settings: PythonCacheSettings = PythonSettings::Cache.read(py)?.extract()?; + settings + .default_redis_ttl + .map(duration) + .transpose() + .map(|ttl| ttl.unwrap_or(PYTHON_REDIS_DEFAULT_TTL)) +} + #[derive(Deserialize)] #[serde(deny_unknown_fields)] struct RequestInput { @@ -29,6 +46,10 @@ struct RequestInput { fn request(value: &Bound<'_, PyAny>) -> PyResult { let input: RequestInput = from_py(value)?; + request_input(input) +} + +fn request_input(input: RequestInput) -> PyResult { let mut request = ResponseCacheRequest::new(input.key); if let Some(controls) = input.controls { request.controls = controls; @@ -38,6 +59,13 @@ fn request(value: &Bound<'_, PyAny>) -> PyResult { Ok(request) } +fn requests(value: &Bound<'_, PyAny>) -> PyResult> { + from_py::>(value)? + .into_iter() + .map(request_input) + .collect() +} + fn duration(seconds: f64) -> PyResult { Duration::try_from_secs_f64(seconds) .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) @@ -94,7 +122,10 @@ impl NativeCacheHandle { ttl_seconds: Option, namespace: Option, ) -> PyResult { - let ttl = ttl_seconds.map(duration).transpose()?; + let ttl = Some(match ttl_seconds { + Some(seconds) => duration(seconds)?, + None => redis_default_ttl(py)?, + }); let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) .map_err(cache_error)?; Ok(Self { @@ -111,7 +142,12 @@ impl NativeCacheHandle { fn bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { let service = self.service()?; - let guard = FacadeGuard::capture(py, facade, self.backend())?; + let guard = FacadeGuard::capture(py, facade, self.backend(), service.default_ttl())?; + let service = service.with_redis_flush_size( + facade + .getattr("redis_flush_size")? + .extract::>()?, + ); let handle = Py::new( py, Self { @@ -249,6 +285,37 @@ impl ResolvedCache { } } + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn lookup_batch( + &self, + py: Python<'_>, + requests: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + to_py(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + let response = release_gil(py, move || service.lookup_batch(&requests, now())) + .map_err(cache_error)?; + to_py(py, &response) + } + CacheBinding::PythonCallback(object) => object + .bind(py) + .call_method( + "batch_get_cache", + (), + Some(self::callback_kwargs(callback_kwargs)?), + ) + .map(Bound::unbind), + } + } + #[pyo3(signature = (request, *, callback_kwargs=None))] fn async_lookup<'py>( &self, @@ -292,6 +359,102 @@ impl ResolvedCache { } } + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn async_lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + ready_value(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + run_async( + py, + async move { service.async_lookup_batch(&requests, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method( + "async_batch_get_cache", + (), + Some(self::callback_kwargs(callback_kwargs)?), + ), + } + } + + #[pyo3(signature = (requests, responses, *, callback_kwargs=None))] + fn async_store_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + responses: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let responses: Vec = from_py(responses)?; + if requests.len() != responses.len() { + return Err(PyValueError::new_err( + "batch cache requests and responses must have equal lengths", + )); + } + let entries = requests.into_iter().zip(responses).collect(); + let service = service.clone(); + run_async( + py, + async move { service.async_store_batch(entries, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method( + "async_set_cache_pipeline", + (responses,), + Some(self::callback_kwargs(callback_kwargs)?), + ), + } + } + + fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async(py, async move { service.async_flush().await }, cache_error) + } + CacheBinding::PythonCallback(object) => { + object.bind(py).call_method0("flush_cache")?; + ready_none(py) + } + } + } + + fn ping<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async( + py, + async move { service.test_connection().await }, + cache_error, + ) + } + CacheBinding::PythonCallback(object) => object.bind(py).call_method0("test_connection"), + } + } + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { if let CacheBinding::PythonCallback(object) = &self.binding { visit.call(object)?; @@ -309,11 +472,18 @@ fn callback_kwargs<'a, 'py>( } fn ready_none(py: Python<'_>) -> PyResult> { + ready_value(py, &()) +} + +fn ready_value<'py, T: serde::Serialize>( + py: Python<'py>, + value: &T, +) -> PyResult> { let future = py .import("asyncio")? .call_method0("get_running_loop")? .call_method0("create_future")?; - future.call_method1("set_result", (py.None(),))?; + future.call_method1("set_result", (to_py(py, value)?,))?; Ok(future) } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 6af04bfe2b2..20891719550 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,16 +1,27 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, Error}; +use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use serde_json::Value; +use tokio::sync::Mutex; -use litellm_cache_response::{CacheEntry, ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; +use litellm_cache_response::{ + CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, +}; #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), - Redis(Arc>>), + Redis { + cache: Arc>>, + buffer: Option>, + }, +} + +pub(super) struct RedisWriteBuffer { + flush_size: usize, + entries: Mutex>, } impl NativeResponseCache { @@ -34,7 +45,10 @@ impl NativeResponseCache { namespace: Option, ) -> Result { let backend = RedisCache::new(url, ttl, ResponseCacheCodec)?.with_namespace(namespace); - Ok(Self::Redis(Arc::new(ResponseCache::new(Arc::new(backend))))) + Ok(Self::Redis { + cache: Arc::new(ResponseCache::new(Arc::new(backend))), + buffer: None, + }) } } @@ -42,7 +56,29 @@ impl NativeResponseCache { pub fn kind(&self) -> &'static str { match self { Self::Memory(_) => "memory", - Self::Redis(_) => "redis", + Self::Redis { .. } => "redis", + } + } + + pub fn default_ttl(&self) -> Duration { + match self { + Self::Memory(cache) => cache.default_ttl(), + Self::Redis { cache, .. } => cache.default_ttl(), + } + } + + pub fn with_redis_flush_size(self, flush_size: Option) -> Self { + match self { + Self::Redis { cache, .. } => Self::Redis { + cache, + buffer: flush_size.map(|flush_size| { + Arc::new(RedisWriteBuffer { + flush_size: flush_size.max(1), + entries: Mutex::new(Vec::new()), + }) + }), + }, + memory => memory, } } @@ -53,7 +89,7 @@ impl NativeResponseCache { ) -> Result, Error> { match self { Self::Memory(cache) => cache.lookup(request, now), - Self::Redis(cache) => cache.lookup(request, now), + Self::Redis { cache, .. } => cache.lookup(request, now), } } @@ -65,7 +101,18 @@ impl NativeResponseCache { ) -> Result<(), Error> { match self { Self::Memory(cache) => cache.store(request, response, now), - Self::Redis(cache) => cache.store(request, response, now), + Self::Redis { cache, .. } => cache.store(request, response, now), + } + } + + pub fn lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + match self { + Self::Memory(cache) => cache.lookup_batch(requests, now), + Self::Redis { cache, .. } => cache.lookup_batch(requests, now), } } @@ -76,7 +123,7 @@ impl NativeResponseCache { ) -> Result, Error> { match self { Self::Memory(cache) => cache.async_lookup(request, now).await, - Self::Redis(cache) => cache.async_lookup(request, now).await, + Self::Redis { cache, .. } => cache.async_lookup(request, now).await, } } @@ -88,7 +135,71 @@ impl NativeResponseCache { ) -> Result<(), Error> { match self { Self::Memory(cache) => cache.async_store(request, response, now).await, - Self::Redis(cache) => cache.async_store(request, response, now).await, + Self::Redis { + cache, + buffer: None, + } => cache.async_store(request, response, now).await, + Self::Redis { + cache, + buffer: Some(buffer), + } => { + let pending = { + let mut entries = buffer.entries.lock().await; + entries.push((request.clone(), response)); + (entries.len() >= buffer.flush_size).then(|| std::mem::take(&mut *entries)) + }; + let Some(pending) = pending else { + return Ok(()); + }; + if let Err(error) = cache.async_store_batch(pending.clone(), now).await { + let mut entries = buffer.entries.lock().await; + let current = std::mem::take(&mut *entries); + *entries = pending.into_iter().chain(current).collect(); + return Err(error); + } + Ok(()) + } + } + } + + pub async fn async_lookup_batch( + &self, + requests: &[ResponseCacheRequest], + now: Duration, + ) -> Result { + match self { + Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, + Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, + } + } + + pub async fn async_store_batch( + &self, + entries: Vec<(ResponseCacheRequest, Value)>, + now: Duration, + ) -> Result<(), Error> { + match self { + Self::Memory(cache) => cache.async_store_batch(entries, now).await, + Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, + } + } + + pub async fn async_flush(&self) -> Result<(), Error> { + match self { + Self::Memory(cache) => cache.async_flush().await, + Self::Redis { cache, buffer } => { + if let Some(buffer) = buffer { + buffer.entries.lock().await.clear(); + } + cache.async_flush().await + } + } + } + + pub async fn test_connection(&self) -> Result { + match self { + Self::Memory(cache) => cache.test_connection().await, + Self::Redis { cache, .. } => cache.test_connection().await, } } } diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 7ac23a05542..90819c5b3fc 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -8,15 +8,17 @@ pub(crate) enum PythonSettings { UrlPolicy, ProviderDefaults, SecretManager, + Cache, } impl PythonSettings { #[cfg(test)] - pub(crate) const ALL: [Self; 4] = [ + pub(crate) const ALL: [Self; 5] = [ Self::Http, Self::UrlPolicy, Self::ProviderDefaults, Self::SecretManager, + Self::Cache, ]; pub(crate) fn name(self) -> &'static str { @@ -25,6 +27,7 @@ impl PythonSettings { Self::UrlPolicy => "url_policy", Self::ProviderDefaults => "provider_defaults", Self::SecretManager => "secret_manager", + Self::Cache => "cache_settings", } } diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 66be77dbb40..64a5af2c618 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -12,7 +12,7 @@ import logging import time from collections.abc import Sequence from threading import Lock -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, TypeVar if TYPE_CHECKING: from litellm.types.caching import RedisPipelineIncrementOperation @@ -34,14 +34,20 @@ else: from collections import OrderedDict +_KeyT = TypeVar("_KeyT") +_ValueT = TypeVar("_ValueT") -class LimitedSizeOrderedDict(OrderedDict): - def __init__(self, *args, max_size=100, **kwargs): - super().__init__(*args, **kwargs) + +class LimitedSizeOrderedDict(OrderedDict[_KeyT, _ValueT]): + def __init__(self, *, max_size: int = 100) -> None: + super().__init__() self.max_size = max_size - def __setitem__(self, key, value): - # If inserting a new key exceeds max size, remove the oldest item + def __setitem__(self, key: _KeyT, value: _ValueT) -> None: + if key in self: + super().__setitem__(key, value) + self.move_to_end(key) + return if len(self) >= self.max_size: self.popitem(last=False) super().__setitem__(key, value) @@ -68,7 +74,9 @@ class DualCache(BaseCache): self.in_memory_cache = in_memory_cache or InMemoryCache() # If redis_cache is not provided, use the default RedisCache self.redis_cache = redis_cache - self.last_redis_batch_access_time = LimitedSizeOrderedDict(max_size=default_max_redis_batch_cache_size) + self.last_redis_batch_access_time: LimitedSizeOrderedDict[str, float] = LimitedSizeOrderedDict( + max_size=default_max_redis_batch_cache_size + ) self._last_redis_batch_access_time_lock = Lock() self.redis_batch_cache_expiry = ( default_redis_batch_cache_expiry or litellm.default_redis_batch_cache_expiry or 10 @@ -131,7 +139,7 @@ class DualCache(BaseCache): except Exception as e: print_verbose(e) - def increment_cache(self, key, value: int, local_only: bool = False, **kwargs) -> int: + def increment_cache(self, key, value: int, local_only: bool = False, **kwargs) -> float: """ Key - the key in cache @@ -140,14 +148,15 @@ class DualCache(BaseCache): Returns - int - the incremented value """ try: - result: int = value - if self.in_memory_cache is not None: - result = self.in_memory_cache.increment_cache(key, value, **kwargs) - if self.redis_cache is not None and local_only is False: - result = self.redis_cache.increment_cache(key, value, **kwargs) + result: Final = self.redis_cache.increment_cache(key, value, **kwargs) + if self.in_memory_cache is not None: + self.in_memory_cache.set_cache(key, result, **kwargs) + return result - return result + if self.in_memory_cache is not None: + return self.in_memory_cache.increment_cache(key, value, **kwargs) + return value except Exception as e: verbose_logger.error("LiteLLM Cache: Excepton async add_cache: %s", e) raise e @@ -421,29 +430,30 @@ class DualCache(BaseCache): Returns - the incremented value, or None if no cache backend is available (in_memory_cache is None and Redis failed/is absent). """ - result: float | None = None try: - if self.in_memory_cache is not None: - result = await self.in_memory_cache.async_increment(key, value, **kwargs) - if self.redis_cache is not None and local_only is False: - result = await self.redis_cache.async_increment( + result: Final = await self.redis_cache.async_increment( key, value, parent_otel_span=parent_otel_span, ttl=kwargs.get("ttl", None), refresh_ttl=refresh_ttl, ) + if self.in_memory_cache is not None: + await self.in_memory_cache.async_set_cache(key, result, **kwargs) + return result - return result + if self.in_memory_cache is not None: + return await self.in_memory_cache.async_increment(key, value, **kwargs) + return None except Exception as e: log_redis_failure( verbose_logger, logging.WARNING, - "Redis async_increment_cache failed, falling back to in-memory result", + "Redis async_increment_cache failed; local counter unchanged", e, ) - return result + return None async def async_increment_cache_pipeline( self, @@ -452,29 +462,32 @@ class DualCache(BaseCache): parent_otel_span: Span | None = None, **kwargs, ) -> list[float] | None: - result: list[float] | None = None try: - if self.in_memory_cache is not None: - result = await self.in_memory_cache.async_increment_pipeline( - increment_list=increment_list, - parent_otel_span=parent_otel_span, - ) - if self.redis_cache is not None and local_only is False: - result = await self.redis_cache.async_increment_pipeline( + result: Final = await self.redis_cache.async_increment_pipeline( increment_list=increment_list, parent_otel_span=parent_otel_span, ) + if result is not None and self.in_memory_cache is not None: + await self.in_memory_cache.async_set_cache_pipeline( + cache_list=tuple((increment["key"], value) for increment, value in zip(increment_list, result)) + ) + return result - return result + if self.in_memory_cache is not None: + return await self.in_memory_cache.async_increment_pipeline( + increment_list=increment_list, + parent_otel_span=parent_otel_span, + ) + return None except Exception as e: log_redis_failure( verbose_logger, logging.WARNING, - "Redis async_increment_cache_pipeline failed, falling back to in-memory result", + "Redis async_increment_cache_pipeline failed; local counters unchanged", e, ) - return result + return None async def async_set_cache_sadd(self, key, value: list, local_only: bool = False, **kwargs) -> None: """ diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 4fd2f0829a3..68b1742dc41 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -126,6 +126,12 @@ class CacheBinding: *, callback_kwargs: dict[str, object] | None = None, ) -> None: ... + def lookup_batch( + self, + requests: Sequence[Mapping[str, object]], + *, + callback_kwargs: dict[str, object] | None = None, + ) -> object: ... def async_lookup( self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None ) -> Awaitable[object]: ... @@ -136,6 +142,21 @@ class CacheBinding: *, callback_kwargs: dict[str, object] | None = None, ) -> Awaitable[object]: ... + def async_lookup_batch( + self, + requests: Sequence[Mapping[str, object]], + *, + callback_kwargs: dict[str, object] | None = None, + ) -> Awaitable[object]: ... + def async_store_batch( + self, + requests: Sequence[Mapping[str, object]], + responses: Sequence[object], + *, + callback_kwargs: dict[str, object] | None = None, + ) -> Awaitable[object]: ... + def async_flush(self) -> Awaitable[None]: ... + def ping(self) -> Awaitable[dict[str, object] | None]: ... @final class TokenCounter: diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 3aa2d742862..862a116496d 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -36,6 +36,11 @@ class SecretManager: readable: bool +@dataclass(frozen=True, slots=True) +class CacheSettings: + default_redis_ttl: float | None + + def warn(message: str) -> None: from litellm._logging import verbose_logger @@ -50,6 +55,12 @@ def secret_manager() -> SecretManager: return SecretManager(readable=_should_read_secret_from_secret_manager()) +def cache_settings() -> CacheSettings: + import litellm + + return CacheSettings(default_redis_ttl=litellm.default_redis_ttl) + + def provider_defaults() -> ProviderDefaults: import litellm diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 5f59de9cca5..149c9b34bd5 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -6,18 +6,16 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest -from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE -from litellm.caching.dual_cache import DualCache +from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, _redis_circuit_breaker_guard, _redis_circuit_breaker_guard_sync +from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from litellm.types.caching import RedisPipelineIncrementOperation @pytest.mark.asyncio async def test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads(): - dual_cache = DualCache( - redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10) keys = ["shared_a", "shared_b"] start_gate = asyncio.Event() @@ -44,9 +42,7 @@ async def test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads @pytest.mark.asyncio async def test_dual_cache_async_batch_get_cache_rolls_back_redis_reservation_on_error(): - dual_cache = DualCache( - redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10) keys = ["shared_a", "shared_b"] with patch.object( @@ -116,9 +112,7 @@ def test_dual_cache_batch_get_cache_only_reads_missing_keys_from_redis(): def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): mock_redis = _redis_mock_for_sync_batch({"absent_key": None}) - dual_cache = DualCache( - in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) first = dual_cache.batch_get_cache(keys=["absent_key"]) second = dual_cache.batch_get_cache(keys=["absent_key"]) @@ -131,9 +125,7 @@ def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): mock_redis = MagicMock(spec=RedisCache) mock_redis.batch_get_cache.side_effect = RuntimeError("redis unavailable") - dual_cache = DualCache( - in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) first_result = dual_cache.batch_get_cache(keys=["shared_a"]) second_result = dual_cache.batch_get_cache(keys=["shared_a"]) @@ -146,9 +138,7 @@ def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): def test_dual_cache_batch_get_cache_returns_memory_only_when_redis_read_is_throttled(): mock_redis = _redis_mock_for_sync_batch({"throttled_key": "redis_value"}) - dual_cache = DualCache( - in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 - ) + dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) dual_cache.last_redis_batch_access_time["throttled_key"] = time.time() result = dual_cache.batch_get_cache(keys=["throttled_key"]) @@ -257,9 +247,7 @@ async def test_dual_cache_batch_redis_backfill_injects_default_in_memory_ttl(): default_in_memory_ttl, same as the single-key path.""" in_memory_cache = InMemoryCache(default_ttl=600) mock_redis = MagicMock(spec=RedisCache) - mock_redis.async_batch_get_cache = AsyncMock( - return_value={"batch_backfill_key": "redis_value"} - ) + mock_redis.async_batch_get_cache = AsyncMock(return_value={"batch_backfill_key": "redis_value"}) dual_cache = DualCache( in_memory_cache=in_memory_cache, redis_cache=mock_redis, @@ -371,9 +359,7 @@ async def test_circuit_breaker_open_skips_redis(): class FakeRedis: def __init__(self): - self._circuit_breaker = RedisCircuitBreaker( - failure_threshold=3, recovery_timeout=60 - ) + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) self._circuit_breaker._state = "open" self._circuit_breaker._opened_at = time.time() self.call_count = 0 @@ -426,9 +412,7 @@ def test_circuit_breaker_half_open_concurrent_calls_are_fast_failed(): # All subsequent concurrent callers: HALF_OPEN → fast-fail (return True) for _ in range(10): - assert ( - cb.is_open() is True - ), "concurrent callers should be fast-failed in HALF_OPEN" + assert cb.is_open() is True, "concurrent callers should be fast-failed in HALF_OPEN" def test_circuit_breaker_disabled_never_opens(): @@ -472,9 +456,7 @@ async def test_circuit_breaker_disabled_guard_always_calls_method(): class FakeRedis: def __init__(self): - self._circuit_breaker = RedisCircuitBreaker( - failure_threshold=1, recovery_timeout=60, enabled=False - ) + self._circuit_breaker = RedisCircuitBreaker(failure_threshold=1, recovery_timeout=60, enabled=False) self.call_count = 0 @_redis_circuit_breaker_guard @@ -512,6 +494,30 @@ async def test_async_increment_cache_returns_none_when_no_in_memory_cache_and_re ) +@pytest.mark.asyncio +async def test_failed_redis_increment_does_not_change_the_local_counter(): + memory = InMemoryCache() + memory.set_cache("counter", 10) + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment = AsyncMock(side_effect=RuntimeError("redis down")) + cache = DualCache(in_memory_cache=memory, redis_cache=redis_cache) + + assert await cache.async_increment_cache("counter", 2) is None + assert memory.get_cache("counter") == 10 + + +@pytest.mark.asyncio +async def test_successful_redis_increment_replaces_the_local_counter_with_the_authoritative_value(): + memory = InMemoryCache() + memory.set_cache("counter", 10) + redis_cache = MagicMock(spec=RedisCache) + redis_cache.async_increment = AsyncMock(return_value=42.0) + cache = DualCache(in_memory_cache=memory, redis_cache=redis_cache) + + assert await cache.async_increment_cache("counter", 2) == 42.0 + assert memory.get_cache("counter") == 42.0 + + def test_dual_cache_late_attach_redis_wires_writes_and_ttl_sync(): """ Typical lazy startup (sync): DualCache runs with in-memory only, then Redis @@ -742,7 +748,7 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo assert [(r.levelno, r.getMessage()) for r in visible] == [ ( logging.WARNING, - "Redis async_increment_cache_pipeline failed, falling back to in-memory result:" + "Redis async_increment_cache_pipeline failed; local counters unchanged:" " Timeout reading from 127.0.0.1:6379", ) ] @@ -756,7 +762,7 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo assert [(r.levelno, r.getMessage()) for r in caplog.records] == [ ( logging.WARNING, - "Redis async_increment_cache failed, falling back to in-memory result: Timeout reading from 127.0.0.1:6379" + "Redis async_increment_cache failed; local counter unchanged: Timeout reading from 127.0.0.1:6379" " (199 more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", ) ] @@ -791,3 +797,14 @@ async def test_async_delete_cache_keys_on_empty_list_touches_no_backend(): await dual_cache.async_delete_cache_keys([]) redis_cache.delete_cache_keys.assert_not_awaited() + + +def test_limited_ordered_dict_refreshes_recency_without_evicting_another_key(): + tracker = LimitedSizeOrderedDict(max_size=2) + tracker["hot"] = 1 + tracker["cold"] = 2 + + tracker["hot"] = 3 + tracker["new"] = 4 + + assert list(tracker.items()) == [("hot", 3), ("new", 4)] diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 493baac228a..7456a1499f7 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -8,6 +8,7 @@ import weakref from collections.abc import Generator from types import SimpleNamespace from typing import Final, Protocol, cast +from urllib.parse import urlparse import fakeredis import pytest @@ -209,8 +210,12 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} client.set("team:sync", str(envelope)) client.set("team:async", json.dumps({"timestamp": time.time(), "response": response})) + client.set("team:raw", json.dumps(response)) + client.set("team:invalid", "not a cache entry") assert binding.lookup(request("sync")) == response assert await binding.async_lookup(request("team:async")) == response + assert binding.lookup(request("raw")) == response + assert await binding.async_lookup(request("invalid")) is None await binding.async_store({**request("native"), "ttl_seconds": 12.0}, response) stored: Final = client.get("team:native") assert isinstance(stored, bytes) @@ -245,3 +250,63 @@ async def test_memory_size_policy_is_applied_by_the_native_host() -> None: ).resolve() await disabled.async_store(request(), small) assert await disabled.async_lookup(request()) is None + + +async def test_native_batch_lookup_and_store_report_partial_hits() -> None: + binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() + requests: Final = [request("hit"), request("miss"), request("disabled")] + requests[2]["controls"] = { + "supported_call_type": True, + "configured": True, + "native_backend": True, + "default_on": True, + "caching": False, + "no_cache": False, + "no_store": False, + "use_cache": False, + } + await binding.async_store_batch(requests, [{"value": 1}, {"value": 2}, {"value": 3}]) + + partial: Final = await binding.async_lookup_batch(requests) + + assert partial == { + "values": [{"value": 1}, {"value": 2}, None], + "missing_indices": [2], + } + + +async def test_redis_handle_reads_the_python_default_ttl(redis_url: str) -> None: + client: Final = redis.Redis.from_url(redis_url) + with rebound(litellm, "default_redis_ttl", 7): + binding: Final = _native.CacheResolver( + SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url)) + ).resolve() + await binding.async_store(request("native-default"), {"value": 1}) + + assert 0 < client.ttl("native-default") <= 7 + client.close() + + +async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: + parsed: Final = urlparse(redis_url) + with rebound(litellm, "default_redis_ttl", 60): + facade: Final = Cache( + type=LiteLLMCacheType.REDIS, + host=parsed.hostname, + port=str(parsed.port), + redis_flush_size=2, + ) + with pytest.raises(TypeError, match="default TTLs must match"): + _native.NativeCacheHandle.redis(redis_url, ttl_seconds=61).bind_facade(facade) + _native.NativeCacheHandle.redis(redis_url).bind_facade(facade) + binding: Final = _native.CacheResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(redis_url) + + await binding.async_store(request("first"), {"value": 1}) + assert client.get("first") is None + await binding.async_store(request("second"), {"value": 2}) + + assert client.get("first") is not None + assert client.get("second") is not None + await facade.cache.disconnect() + client.close() From 9783b7a377009982abd9541b305d466590770252 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 07:56:58 -0700 Subject: [PATCH 05/17] docs(cache): align native cache follow-up scope --- litellm-rust/crates/cache-response/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 4e6694c191f..d8ffd6d6a50 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -52,4 +52,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths -Redis cluster, disk, cloud stores, dual caching, and semantic caching remain follow-ups. Atomic counters, affinity claims, reservations, queues, and pubsub need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees +Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache now provides L2-first counters and atomic affinity claims with local fallback, but public Router integration remains follow-up work. Reservations, queues, and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees From da402b8aeee24203fad87a646ce3b5fec7b9ff51 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 08:27:17 -0700 Subject: [PATCH 06/17] fix(cache): preserve batch callback contracts --- .../crates/python-bridge/src/cache/mod.rs | 41 ++++++++++--- litellm/caching/dual_cache.py | 61 +++++++++---------- tests/test_litellm/caching/test_dual_cache.py | 28 +-------- tests/test_litellm_rust/test_cache.py | 39 ++++++++++++ 4 files changed, 102 insertions(+), 67 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 4bab03fe243..1e2e42600ff 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -10,7 +10,7 @@ use pyo3::{ PyTraverseError, PyVisit, exceptions::{PyRuntimeError, PyTypeError, PyValueError}, prelude::*, - types::PyDict, + types::{PyDict, PyList}, }; use serde::Deserialize; use serde_json::Value; @@ -309,7 +309,7 @@ impl ResolvedCache { .bind(py) .call_method( "batch_get_cache", - (), + (callback_keys(py, requests)?,), Some(self::callback_kwargs(callback_kwargs)?), ) .map(Bound::unbind), @@ -383,7 +383,7 @@ impl ResolvedCache { } CacheBinding::PythonCallback(object) => object.bind(py).call_method( "async_batch_get_cache", - (), + (callback_keys(py, requests)?,), Some(self::callback_kwargs(callback_kwargs)?), ), } @@ -416,11 +416,24 @@ impl ResolvedCache { cache_error, ) } - CacheBinding::PythonCallback(object) => object.bind(py).call_method( - "async_set_cache_pipeline", - (responses,), - Some(self::callback_kwargs(callback_kwargs)?), - ), + CacheBinding::PythonCallback(object) => { + let keys = callback_keys(py, requests)?; + let responses = responses.try_iter()?.collect::>>()?; + if keys.len() != responses.len() { + return Err(PyValueError::new_err( + "batch cache requests and responses must have equal lengths", + )); + } + let cache_list = PyList::empty(py); + for (key, response) in keys.iter().zip(responses) { + cache_list.append((key, response))?; + } + object.bind(py).call_method( + "async_set_cache_pipeline", + (cache_list,), + Some(self::callback_kwargs(callback_kwargs)?), + ) + } } } @@ -471,6 +484,18 @@ fn callback_kwargs<'a, 'py>( }) } +fn callback_keys<'py>( + py: Python<'py>, + requests: &Bound<'py, PyAny>, +) -> PyResult> { + PyList::new( + py, + self::requests(requests)? + .into_iter() + .map(|request| litellm_cache_response::cache_key(&request.key)), + ) +} + fn ready_none(py: Python<'_>) -> PyResult> { ready_value(py, &()) } diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 64a5af2c618..04c82232784 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -139,7 +139,7 @@ class DualCache(BaseCache): except Exception as e: print_verbose(e) - def increment_cache(self, key, value: int, local_only: bool = False, **kwargs) -> float: + def increment_cache(self, key, value: int, local_only: bool = False, **kwargs) -> int: """ Key - the key in cache @@ -148,15 +148,14 @@ class DualCache(BaseCache): Returns - int - the incremented value """ try: - if self.redis_cache is not None and local_only is False: - result: Final = self.redis_cache.increment_cache(key, value, **kwargs) - if self.in_memory_cache is not None: - self.in_memory_cache.set_cache(key, result, **kwargs) - return result - + result: int = value if self.in_memory_cache is not None: - return self.in_memory_cache.increment_cache(key, value, **kwargs) - return value + result = self.in_memory_cache.increment_cache(key, value, **kwargs) + + if self.redis_cache is not None and local_only is False: + result = self.redis_cache.increment_cache(key, value, **kwargs) + + return result except Exception as e: verbose_logger.error("LiteLLM Cache: Excepton async add_cache: %s", e) raise e @@ -430,30 +429,29 @@ class DualCache(BaseCache): Returns - the incremented value, or None if no cache backend is available (in_memory_cache is None and Redis failed/is absent). """ + result: float | None = None try: + if self.in_memory_cache is not None: + result = await self.in_memory_cache.async_increment(key, value, **kwargs) + if self.redis_cache is not None and local_only is False: - result: Final = await self.redis_cache.async_increment( + result = await self.redis_cache.async_increment( key, value, parent_otel_span=parent_otel_span, ttl=kwargs.get("ttl", None), refresh_ttl=refresh_ttl, ) - if self.in_memory_cache is not None: - await self.in_memory_cache.async_set_cache(key, result, **kwargs) - return result - if self.in_memory_cache is not None: - return await self.in_memory_cache.async_increment(key, value, **kwargs) - return None + return result except Exception as e: log_redis_failure( verbose_logger, logging.WARNING, - "Redis async_increment_cache failed; local counter unchanged", + "Redis async_increment_cache failed, falling back to in-memory result", e, ) - return None + return result async def async_increment_cache_pipeline( self, @@ -462,32 +460,29 @@ class DualCache(BaseCache): parent_otel_span: Span | None = None, **kwargs, ) -> list[float] | None: + result: list[float] | None = None try: - if self.redis_cache is not None and local_only is False: - result: Final = await self.redis_cache.async_increment_pipeline( - increment_list=increment_list, - parent_otel_span=parent_otel_span, - ) - if result is not None and self.in_memory_cache is not None: - await self.in_memory_cache.async_set_cache_pipeline( - cache_list=tuple((increment["key"], value) for increment, value in zip(increment_list, result)) - ) - return result - if self.in_memory_cache is not None: - return await self.in_memory_cache.async_increment_pipeline( + result = await self.in_memory_cache.async_increment_pipeline( increment_list=increment_list, parent_otel_span=parent_otel_span, ) - return None + + if self.redis_cache is not None and local_only is False: + result = await self.redis_cache.async_increment_pipeline( + increment_list=increment_list, + parent_otel_span=parent_otel_span, + ) + + return result except Exception as e: log_redis_failure( verbose_logger, logging.WARNING, - "Redis async_increment_cache_pipeline failed; local counters unchanged", + "Redis async_increment_cache_pipeline failed, falling back to in-memory result", e, ) - return None + return result async def async_set_cache_sadd(self, key, value: list, local_only: bool = False, **kwargs) -> None: """ diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 149c9b34bd5..05a4ef68e0b 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -494,30 +494,6 @@ async def test_async_increment_cache_returns_none_when_no_in_memory_cache_and_re ) -@pytest.mark.asyncio -async def test_failed_redis_increment_does_not_change_the_local_counter(): - memory = InMemoryCache() - memory.set_cache("counter", 10) - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_increment = AsyncMock(side_effect=RuntimeError("redis down")) - cache = DualCache(in_memory_cache=memory, redis_cache=redis_cache) - - assert await cache.async_increment_cache("counter", 2) is None - assert memory.get_cache("counter") == 10 - - -@pytest.mark.asyncio -async def test_successful_redis_increment_replaces_the_local_counter_with_the_authoritative_value(): - memory = InMemoryCache() - memory.set_cache("counter", 10) - redis_cache = MagicMock(spec=RedisCache) - redis_cache.async_increment = AsyncMock(return_value=42.0) - cache = DualCache(in_memory_cache=memory, redis_cache=redis_cache) - - assert await cache.async_increment_cache("counter", 2) == 42.0 - assert memory.get_cache("counter") == 42.0 - - def test_dual_cache_late_attach_redis_wires_writes_and_ttl_sync(): """ Typical lazy startup (sync): DualCache runs with in-memory only, then Redis @@ -748,7 +724,7 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo assert [(r.levelno, r.getMessage()) for r in visible] == [ ( logging.WARNING, - "Redis async_increment_cache_pipeline failed; local counters unchanged:" + "Redis async_increment_cache_pipeline failed, falling back to in-memory result:" " Timeout reading from 127.0.0.1:6379", ) ] @@ -762,7 +738,7 @@ async def test_redis_timeouts_falling_back_to_memory_log_once_per_interval(caplo assert [(r.levelno, r.getMessage()) for r in caplog.records] == [ ( logging.WARNING, - "Redis async_increment_cache failed; local counter unchanged: Timeout reading from 127.0.0.1:6379" + "Redis async_increment_cache failed, falling back to in-memory result: Timeout reading from 127.0.0.1:6379" " (199 more Redis timeouts since the previous Redis timeout line were logged at DEBUG)", ) ] diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 7456a1499f7..d38583dad0a 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -275,6 +275,45 @@ async def test_native_batch_lookup_and_store_report_partial_hits() -> None: } +async def test_python_batch_callbacks_receive_keys_and_key_value_pairs() -> None: + first: Final = object() + second: Final = object() + + class CustomCache: + def batch_get_cache(self, keys: list[str], *, marker: object) -> tuple[list[str], object]: + return keys, marker + + async def async_batch_get_cache(self, keys: list[str], *, marker: object) -> tuple[list[str], object]: + return keys, marker + + async def async_set_cache_pipeline( + self, cache_list: list[tuple[str, object]], *, marker: object + ) -> tuple[list[tuple[str, object]], object]: + return cache_list, marker + + marker: Final = object() + binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve() + requests: Final = [request("first"), request("second")] + + assert binding.lookup_batch(requests, callback_kwargs={"marker": marker}) == (["first", "second"], marker) + assert await binding.async_lookup_batch(requests, callback_kwargs={"marker": marker}) == ( + ["first", "second"], + marker, + ) + stored: Final = cast( + tuple[list[tuple[str, object]], object], + await binding.async_store_batch( + requests, + [first, second], + callback_kwargs={"marker": marker}, + ), + ) + assert [key for key, _ in stored[0]] == ["first", "second"] + assert stored[1] is marker + assert stored[0][0][1] is first + assert stored[0][1][1] is second + + async def test_redis_handle_reads_the_python_default_ttl(redis_url: str) -> None: client: Final = redis.Redis.from_url(redis_url) with rebound(litellm, "default_redis_ttl", 7): From ef14fdaf339006a1f2bcb443ecc14225554d1568 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 09:53:41 -0700 Subject: [PATCH 07/17] fix(cache): harden native foundation parity --- litellm-rust/Cargo.lock | 1 - litellm-rust/crates/cache-memory/src/cache.rs | 98 ++++-- .../crates/cache-memory/tests/cache.rs | 79 ++++- litellm-rust/crates/cache-redis/Cargo.toml | 2 +- litellm-rust/crates/cache-redis/src/cache.rs | 236 ++++++++++---- .../crates/cache-redis/tests/cache.rs | 127 +++++++- litellm-rust/crates/cache-response/README.md | 12 +- .../crates/cache-response/src/codec.rs | 25 +- .../crates/cache-response/src/response.rs | 50 +-- .../crates/cache-response/tests/response.rs | 75 ++++- litellm-rust/crates/cache/src/base_cache.rs | 3 +- litellm-rust/crates/cache/src/caching.rs | 3 +- litellm-rust/crates/cache/src/codec.rs | 8 + litellm-rust/crates/cache/src/dual.rs | 302 +++++++++++++++++- litellm-rust/crates/cache/tests/caching.rs | 3 +- litellm-rust/crates/cache/tests/dual.rs | 216 ++++++++++++- .../crates/python-bridge/src/cache/facade.rs | 32 +- .../crates/python-bridge/src/cache/mod.rs | 136 ++++---- .../crates/python-bridge/src/cache/native.rs | 37 ++- litellm-rust/crates/python-bridge/src/lib.rs | 16 +- litellm/rust_bridge/_native.pyi | 26 +- tests/test_litellm/caching/test_dual_cache.py | 2 +- tests/test_litellm_rust/test_cache.py | 176 ++++++---- 23 files changed, 1347 insertions(+), 318 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 725d2cfef41..2d6fb6c3082 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3712,7 +3712,6 @@ dependencies = [ "itoa", "num-bigint 0.5.1", "percent-encoding", - "r2d2", "ryu", "sha1_smol", "socket2 0.6.5", diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 074c8dcb0fd..a9814ff6fd6 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -1,7 +1,9 @@ -use std::cmp::Reverse; -use std::collections::{BinaryHeap, HashMap}; -use std::sync::{Arc, Mutex}; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use std::{ + cmp::Reverse, + collections::{BinaryHeap, HashMap}, + sync::{Arc, Mutex}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; use litellm_cache::{ BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, @@ -95,15 +97,13 @@ impl InMemoryCache { } let now = (self.now)(); let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; - Self::evict(&mut state, self.max_size_in_memory, now); let key = key.into(); - state.values.insert(key.clone(), value); + Self::evict(&mut state, self.max_size_in_memory, now, &key); let expiration = state.expirations.get(&key).copied(); if expiration.is_none_or(|expiration| expiration < now) { - let expiration = now + ttl.unwrap_or(self.default_ttl); - state.expirations.insert(key.clone(), expiration); - state.expiration_heap.push(Reverse((expiration, key))); + Self::set_expiration(&mut state, &key, now + ttl.unwrap_or(self.default_ttl)); } + state.values.insert(key, value); Ok(CacheWrite::Stored) } @@ -120,6 +120,10 @@ impl InMemoryCache { Ok(state.values.get(key).cloned()) } + pub fn max_size_in_memory(&self) -> usize { + self.max_size_in_memory + } + pub fn expires_at(&self, key: &str) -> Result, Error> { Ok(self .state @@ -144,7 +148,7 @@ impl InMemoryCache { Ok(()) } - fn evict(state: &mut CacheState, capacity: usize, now: Duration) { + fn evict(state: &mut CacheState, capacity: usize, now: Duration, key: &str) { while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() { if state.expirations.get(&key).copied() != Some(expiration) { state.expiration_heap.pop(); @@ -155,6 +159,9 @@ impl InMemoryCache { break; } } + if state.values.contains_key(key) { + return; + } while state.values.len() >= capacity { let Some(Reverse((expiration, key))) = state.expiration_heap.pop() else { break; @@ -165,6 +172,15 @@ impl InMemoryCache { } } + fn set_expiration(state: &mut CacheState, key: &str, expiration: Duration) { + if state.expirations.get(key).copied() != Some(expiration) { + state.expirations.insert(key.into(), expiration); + state + .expiration_heap + .push(Reverse((expiration, key.into()))); + } + } + fn remove(state: &mut CacheState, key: &str) { state.values.remove(key); state.expirations.remove(key); @@ -182,40 +198,44 @@ where eligible: &[V], kwargs: CacheKwargs, ) -> Result { + if self.max_size_in_memory == 0 { + return Ok(candidate); + } let now = (self.now)(); let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; - Self::evict(&mut state, self.max_size_in_memory, now); - let winner = match state.values.get(key) { - Some(existing) if eligible.is_empty() => existing.clone(), - Some(existing) if eligible.contains(existing) => existing.clone(), - _ => candidate, - }; - let expiration = now + self.get_ttl(&kwargs); + Self::evict(&mut state, self.max_size_in_memory, now, key); + let existing = state + .values + .get(key) + .filter(|existing| eligible.is_empty() || eligible.contains(existing)) + .cloned(); + // Matches the Redis claim: an unconditional claim only extends its own winner. + if let Some(existing) = &existing + && eligible.is_empty() + && *existing != candidate + { + return Ok(existing.clone()); + } + let winner = existing.unwrap_or(candidate); + Self::set_expiration(&mut state, key, now + self.get_ttl(&kwargs)); state.values.insert(key.into(), winner.clone()); - state.expirations.insert(key.into(), expiration); - state - .expiration_heap - .push(Reverse((expiration, key.into()))); Ok(winner) } } impl CounterCache for InMemoryCache { fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { + if self.max_size_in_memory == 0 { + return Ok(amount); + } let now = (self.now)(); let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; - Self::evict(&mut state, self.max_size_in_memory, now); + Self::evict(&mut state, self.max_size_in_memory, now, key); let value = state.values.get(key).copied().unwrap_or_default() + amount; - let expiration = state - .expirations - .get(key) - .copied() - .unwrap_or_else(|| now + self.get_ttl(&kwargs)); + if !state.expirations.contains_key(key) { + Self::set_expiration(&mut state, key, now + self.get_ttl(&kwargs)); + } state.values.insert(key.into(), value); - state.expirations.insert(key.into(), expiration); - state - .expiration_heap - .push(Reverse((expiration, key.into()))); Ok(value) } } @@ -256,3 +276,19 @@ impl BaseCache for InMemoryCache { }) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn repeated_increments_keep_one_heap_entry_per_expiration() { + let cache = InMemoryCache::::new(Some(4), None); + for _ in 0..100 { + cache + .increment_cache("counter", 1.0, CacheKwargs::default()) + .unwrap(); + } + assert_eq!(cache.state.lock().unwrap().expiration_heap.len(), 1); + } +} diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index e5831dfd6d5..22c5595da52 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -1,6 +1,10 @@ -use std::sync::Arc; -use std::sync::atomic::{AtomicU64, Ordering}; -use std::time::Duration; +use std::{ + sync::{ + Arc, + atomic::{AtomicU64, Ordering}, + }, + time::Duration, +}; use litellm_cache::{ BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, Error, @@ -202,6 +206,17 @@ fn claims_are_atomic_and_refresh_eligible_winners() { .unwrap(), "first" ); + clock.store(103, Ordering::SeqCst); + assert_eq!( + cache + .claim_cache("affinity", "second".to_string(), &[], kwargs.clone()) + .unwrap(), + "first" + ); + assert_eq!( + cache.expires_at("affinity").unwrap(), + Some(Duration::from_secs(110)) + ); clock.store(105, Ordering::SeqCst); assert_eq!( cache @@ -232,3 +247,61 @@ fn counters_increment_under_one_lock() { 3.5 ); } + +#[rstest] +fn rewriting_an_existing_key_at_capacity_keeps_other_entries(clock: Arc) { + let cache = cache(clock, 2); + cache + .set_cache("hot", "1".into(), Some(Duration::from_secs(10))) + .unwrap(); + cache + .set_cache("cold", "2".into(), Some(Duration::from_secs(20))) + .unwrap(); + + cache.set_cache("cold", "3".into(), None).unwrap(); + assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into())); + assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into())); + + cache + .claim_cache("cold", "4".into(), &[], CacheKwargs::default()) + .unwrap(); + assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into())); + + cache.set_cache("new", "5".into(), None).unwrap(); + assert_eq!(cache.get_cache("hot").unwrap(), None); + assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into())); + assert_eq!(cache.get_cache("new").unwrap(), Some("5".into())); +} + +#[test] +fn incrementing_an_existing_counter_at_capacity_keeps_every_counter() { + let cache = InMemoryCache::::new(Some(2), None); + for key in ["a", "b", "a", "b"] { + cache + .increment_cache(key, 1.0, CacheKwargs::default()) + .unwrap(); + } + assert_eq!(cache.get_cache("a").unwrap(), Some(2.0)); + assert_eq!(cache.get_cache("b").unwrap(), Some(2.0)); +} + +#[test] +fn disabled_cache_does_not_retain_claims_or_counters() { + let claims = InMemoryCache::::new(Some(0), None); + assert_eq!( + claims + .claim_cache("key", "first".into(), &[], CacheKwargs::default()) + .unwrap(), + "first" + ); + assert_eq!(claims.get_cache("key").unwrap(), None); + + let counters = InMemoryCache::::new(Some(0), None); + assert_eq!( + counters + .increment_cache("key", 2.0, CacheKwargs::default()) + .unwrap(), + 2.0 + ); + assert_eq!(counters.get_cache("key").unwrap(), None); +} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index 1a1bb505a7c..f123e774158 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis = { version = "1.7.0", features = ["r2d2"] } +redis = "1.7.0" r2d2 = "0.8.10" tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 281ffe17534..b8f0d3857e2 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -1,5 +1,7 @@ -use std::sync::{Arc, Mutex}; -use std::time::Duration; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; use litellm_cache::{ BaseCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, @@ -11,8 +13,59 @@ const DEFAULT_TTL: Duration = Duration::from_secs(600); const REDIS_TIMEOUT: Duration = Duration::from_secs(5); const REDIS_POOL_SIZE: u32 = 16; +struct PooledConnection { + connection: redis::Connection, + failed: bool, +} + +/// Pools connections without a checkout PING, which would double every operation's round trips. +/// A timed-out command leaves its reply on the socket while redis still reports the connection +/// open, so any connection whose operation failed is discarded instead of being reused. +struct ConnectionManager(redis::Client); + +impl r2d2::ManageConnection for ConnectionManager { + type Connection = PooledConnection; + type Error = redis::RedisError; + + fn connect(&self) -> Result { + let connection = self.0.get_connection()?; + connection.set_read_timeout(Some(REDIS_TIMEOUT))?; + connection.set_write_timeout(Some(REDIS_TIMEOUT))?; + Ok(PooledConnection { + connection, + failed: false, + }) + } + + fn is_valid(&self, connection: &mut PooledConnection) -> Result<(), redis::RedisError> { + redis::cmd("PING").query::(&mut connection.connection)?; + Ok(()) + } + + fn has_broken(&self, connection: &mut PooledConnection) -> bool { + connection.failed || !redis::ConnectionLike::is_open(&connection.connection) + } +} + +const INCREMENT_SCRIPT: &str = concat!( + "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", + "if redis.call('TTL', KEYS[1]) == -1 then ", + "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" +); + +// Compare-and-set against the exact bytes the claim decision was made on. +// ARGV: [1] expected payload or "" when absent, [2] ttl, [3] new payload, [4] refresh ttl. +const CLAIM_SCRIPT: &str = concat!( + "local current = redis.call('GET', KEYS[1]); ", + "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", + "elseif current ~= ARGV[1] then return 0; end; ", + "if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ", + "elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1" +); +const CLAIM_ATTEMPTS: usize = 8; + enum Connections { - Pool(r2d2::Pool), + Pool(r2d2::Pool), Fixed(Mutex), } @@ -59,14 +112,10 @@ where ) -> Result { match self { Self::Pool(pool) => { - let mut connection = pool.get().map_err(|_| Error::Unavailable)?; - connection - .set_read_timeout(Some(REDIS_TIMEOUT)) - .map_err(|_| Error::Unavailable)?; - connection - .set_write_timeout(Some(REDIS_TIMEOUT)) - .map_err(|_| Error::Unavailable)?; - operation(&mut ConnectionRef(&mut *connection)) + let mut pooled = pool.get().map_err(|_| Error::Unavailable)?; + let result = operation(&mut ConnectionRef(&mut pooled.connection)); + pooled.failed = matches!(result, Err(Error::Unavailable)); + result } Self::Fixed(connection) => { let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; @@ -90,7 +139,8 @@ impl RedisCache { .max_size(REDIS_POOL_SIZE) .min_idle(Some(0)) .connection_timeout(REDIS_TIMEOUT) - .build(client) + .test_on_check_out(false) + .build(ConnectionManager(client)) .map_err(|_| Error::Unavailable)?; Ok(Self { connections: Arc::new(Connections::Pool(pool)), @@ -122,6 +172,10 @@ where } } + pub fn namespace(&self) -> Option<&str> { + self.namespace.as_deref() + } + fn namespaced_key(&self, key: &str) -> String { match &self.namespace { Some(namespace) if !key.starts_with(&format!("{namespace}:")) => { @@ -407,29 +461,105 @@ where C: redis::ConnectionLike + Send + 'static, { fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result { - const SCRIPT: &str = concat!( - "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", - "if redis.call('TTL', KEYS[1]) == -1 then ", - "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" - ); let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); - self.connections.execute(|connection| { - redis::cmd("EVAL") - .arg(SCRIPT) - .arg(1) - .arg(key) - .arg(amount) - .arg(ttl) - .query(connection) - .map_err(|_| Error::Unavailable) - }) + self.connections + .execute(|connection| increment(connection, key, amount, ttl)) } + + async fn async_increment_cache( + &self, + key: &str, + amount: f64, + kwargs: CacheKwargs, + ) -> Result { + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + increment(connection, key, amount, ttl) + }) + .await + } +} + +fn increment( + connection: &mut ConnectionRef<'_>, + key: String, + amount: f64, + ttl: u64, +) -> Result { + redis::cmd("EVAL") + .arg(INCREMENT_SCRIPT) + .arg(1) + .arg(key) + .arg(amount) + .arg(ttl) + .query(connection) + .map_err(|_| Error::Unavailable) +} + +fn stored_bytes(value: redis::Value) -> Result>, Error> { + match value { + redis::Value::Nil => Ok(None), + redis::Value::BulkString(bytes) => Ok(Some(bytes)), + redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())), + _ => Err(Error::InvalidEntry), + } +} + +/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's +/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the +/// bytes that decision was made on, retried when another claimant wins the race. +fn claim( + connection: &mut ConnectionRef<'_>, + codec: &S, + key: &str, + candidate: S::Value, + eligible: &[S::Value], + ttl: u64, +) -> Result +where + S::Value: PartialEq, +{ + let payload = codec.encode(&candidate)?; + if payload.is_empty() { + return Err(Error::InvalidEntry); + } + for _ in 0..CLAIM_ATTEMPTS { + let current = stored_bytes( + connection + .get::<_, redis::Value>(key) + .map_err(|_| Error::Unavailable)?, + )? + .filter(|bytes| !bytes.is_empty()); + let existing = current + .as_deref() + .and_then(|bytes| codec.decode(bytes).ok()) + .filter(|existing| eligible.is_empty() || eligible.contains(existing)); + let refresh = existing + .as_ref() + .is_some_and(|existing| !eligible.is_empty() || *existing == candidate); + let write: &[u8] = if existing.is_some() { b"" } else { &payload }; + let applied = redis::cmd("EVAL") + .arg(CLAIM_SCRIPT) + .arg(1) + .arg(key) + .arg(current.as_deref().unwrap_or_default()) + .arg(ttl) + .arg(write) + .arg(u8::from(refresh)) + .query::(connection) + .map_err(|_| Error::Unavailable)?; + if applied { + return Ok(existing.unwrap_or(candidate)); + } + } + Err(Error::Unavailable) } impl ClaimCache for RedisCache where - S: CacheCodec, + S: CacheCodec + Clone + 'static, S::Value: PartialEq, C: redis::ConnectionLike + Send + 'static, { @@ -440,44 +570,38 @@ where eligible: &[S::Value], kwargs: CacheKwargs, ) -> Result { - const SCRIPT: &str = concat!( - "local current = redis.call('GET', KEYS[1]); ", - "if current == false then redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ", - "return ARGV[1]; end; if #ARGV > 2 then for index = 3, #ARGV do ", - "if current == ARGV[index] then redis.call('EXPIRE', KEYS[1], ARGV[2]); ", - "return current; end; end; redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]); ", - "return ARGV[1]; end; if current == ARGV[1] then ", - "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return current" - ); let key = self.namespaced_key(key); - let candidate = self.codec.encode(&candidate)?; - let eligible = eligible - .iter() - .map(|value| self.codec.encode(value)) - .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); - let value = self.connections.execute(|connection| { - redis::cmd("EVAL") - .arg(SCRIPT) - .arg(1) - .arg(key) - .arg(candidate) - .arg(ttl) - .arg(eligible) - .query::(connection) - .map_err(|_| Error::Unavailable) - })?; - self.decode_response(value)?.ok_or(Error::Unavailable) + self.connections + .execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl)) + } + + async fn async_claim_cache( + &self, + key: &str, + candidate: S::Value, + eligible: Vec, + kwargs: CacheKwargs, + ) -> Result { + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(self.get_ttl(&kwargs)); + let codec = self.codec.clone(); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + claim(connection, &codec, &key, candidate, &eligible, ttl) + }) + .await } } #[cfg(test)] mod tests { - use super::RedisCache; + use std::time::Duration; + use litellm_cache::{BaseCache, CacheCodec, CacheKwargs, JsonCodec}; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; - use std::time::Duration; + + use super::RedisCache; fn entry() -> serde_json::Value { json!({"deployment": "model-a", "cooldown_seconds": 30}) diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index 6aef9bb36bf..6a61b80b84c 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -1,8 +1,8 @@ use std::time::Duration; use litellm_cache::{ - BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, Error, JsonCodec, - get_cache, set_cache, + BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, ClaimCache, + CounterCache, Error, JsonCodec, get_cache, set_cache, }; use litellm_cache_redis::RedisCache; use redis_test::{MockCmd, MockRedisConnection}; @@ -259,3 +259,126 @@ async fn async_flush_deletes_each_scan_page_separately() { cache.async_flush_cache().await.unwrap(); } + +const CLAIM_SCRIPT: &str = concat!( + "local current = redis.call('GET', KEYS[1]); ", + "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", + "elseif current ~= ARGV[1] then return 0; end; ", + "if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ", + "elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1" +); + +fn claim_eval(expected: &str, write: &str, refresh: bool) -> redis::Cmd { + let mut cmd = redis::cmd("EVAL"); + cmd.arg(CLAIM_SCRIPT) + .arg(1) + .arg("pin") + .arg(expected) + .arg(600) + .arg(write) + .arg(u8::from(refresh)); + cmd +} + +#[tokio::test] +async fn claims_match_eligible_values_written_by_another_encoder() { + let python_payload = r#"{"model_id": "a", "deployment": "east"}"#; + let stored = serde_json::json!({"deployment": "east", "model_id": "a"}); + let candidate = serde_json::json!({"model_id": "b"}); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("GET").arg("pin"), Ok(python_payload)), + MockCmd::new(claim_eval(python_payload, "", true), Ok(1)), + ]) + .assert_all_commands_consumed(); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()); + + assert_eq!( + cache + .async_claim_cache( + "pin", + candidate, + vec![stored.clone()], + CacheKwargs::default() + ) + .await + .unwrap(), + stored + ); +} + +#[test] +fn claims_retry_when_the_key_changes_and_replace_ineligible_winners() { + let candidate = serde_json::json!({"model_id": "b"}); + let payload = r#"{"model_id":"b"}"#; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("GET").arg("pin"), Ok(redis::Value::Nil)), + MockCmd::new(claim_eval("", payload, false), Ok(0)), + MockCmd::new(redis::cmd("GET").arg("pin"), Ok(r#"{"model_id":"gone"}"#)), + MockCmd::new(claim_eval(r#"{"model_id":"gone"}"#, payload, false), Ok(1)), + ]) + .assert_all_commands_consumed(); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()); + + assert_eq!( + cache + .claim_cache( + "pin", + candidate.clone(), + &[serde_json::json!({"model_id": "a"})], + CacheKwargs::default() + ) + .unwrap(), + candidate + ); +} + +#[test] +fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl() { + let stored = r#"{"model_id": "a"}"#; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("GET").arg("pin"), Ok(stored)), + MockCmd::new(claim_eval(stored, "", false), Ok(1)), + ]) + .assert_all_commands_consumed(); + let cache = + RedisCache::with_connection(connection, None, JsonCodec::::new()); + + assert_eq!( + cache + .claim_cache( + "pin", + serde_json::json!({"model_id": "b"}), + &[], + CacheKwargs::default() + ) + .unwrap(), + serde_json::json!({"model_id": "a"}) + ); +} + +#[tokio::test] +async fn async_increment_runs_the_atomic_script() { + let mut eval = redis::cmd("EVAL"); + eval.arg(concat!( + "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", + "if redis.call('TTL', KEYS[1]) == -1 then ", + "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" + )) + .arg(1) + .arg("counter") + .arg(2.5f64) + .arg(600); + let connection = + MockRedisConnection::new([MockCmd::new(eval, Ok("4.5"))]).assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); + + assert_eq!( + cache + .async_increment_cache("counter", 2.5, CacheKwargs::default()) + .await + .unwrap(), + 4.5 + ); +} diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index d8ffd6d6a50..013cf4b7baf 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -28,17 +28,21 @@ cache.store(&request, json!({"answer": 7}), now)?; assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7}))); ``` -For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers move that blocking work off the executor +For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers, including counters and claims, move that blocking work off the executor. The pool skips the checkout PING and instead discards any connection whose command failed Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it ## Python integration boundary -The extension exposes `NativeCacheHandle`, `CacheResolver`, and captured `CacheBinding` objects for host integration. Memory and Redis handles support single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring +The extension keeps a private test harness for memory and Redis single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring. No bridge-only cache type is part of the public API + +Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution -Explicit facade registration checks object identity, method overrides, effective TTL, and configuration changes before selecting native execution. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. Registration does not migrate entries or replace Python methods. Until activation configures one shared service, a registered facade and its native handle can hold separate data. Existing public cache constructors remain on Python +Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend + +The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and configuration changes before selecting native execution. It does not compare Redis connection settings. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy @@ -52,4 +56,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths -Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache now provides L2-first counters and atomic affinity claims with local fallback, but public Router integration remains follow-up work. Reservations, queues, and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees +Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations, queues, and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees diff --git a/litellm-rust/crates/cache-response/src/codec.rs b/litellm-rust/crates/cache-response/src/codec.rs index eaba3d0c349..6b0f29e0a58 100644 --- a/litellm-rust/crates/cache-response/src/codec.rs +++ b/litellm-rust/crates/cache-response/src/codec.rs @@ -1,8 +1,9 @@ use litellm_cache::{CacheCodec, Error}; - -use crate::CacheEntry; use serde_json::Value; +use crate::CacheEntry; + +#[derive(Clone, Copy, Debug, Default)] pub struct ResponseCacheCodec; impl CacheCodec for ResponseCacheCodec { @@ -15,7 +16,18 @@ impl CacheCodec for ResponseCacheCodec { { return Err(Error::InvalidEntry); } - serde_json::to_vec(value).map_err(|_| Error::InvalidEntry) + // Python reads a `response` that is either a dict or a serialized string, so every + // other shape is written serialized. A string on the wire is therefore always a + // serialized response, which keeps string-valued responses unambiguous. + if value.timestamp.is_none() || value.response.is_object() { + return serde_json::to_vec(value).map_err(|_| Error::InvalidEntry); + } + let response = serde_json::to_string(&value.response).map_err(|_| Error::InvalidEntry)?; + serde_json::to_vec(&CacheEntry { + timestamp: value.timestamp, + response: Value::String(response), + }) + .map_err(|_| Error::InvalidEntry) } fn decode(&self, bytes: &[u8]) -> Result { @@ -30,7 +42,10 @@ impl CacheCodec for ResponseCacheCodec { let Some(timestamp) = timestamp.as_f64().filter(|timestamp| timestamp.is_finite()) else { return Err(Error::InvalidEntry); }; - let response = value.get("response").cloned().ok_or(Error::InvalidEntry)?; + let response = match value.get("response").ok_or(Error::InvalidEntry)? { + Value::String(text) => decode_value(text)?, + response => response.clone(), + }; Ok(CacheEntry { timestamp: Some(timestamp), response, @@ -38,7 +53,7 @@ impl CacheCodec for ResponseCacheCodec { } } -pub(crate) fn decode_value(text: &str) -> Result { +fn decode_value(text: &str) -> Result { if let Ok(value) = serde_json::from_str(text) { return Ok(value); } diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index 8f0fe953df6..f18a48863c5 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,9 +1,9 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, Error}; +use serde_json::Value; use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; -use serde_json::Value; #[derive(Clone)] pub struct ResponseCacheRequest { @@ -39,6 +39,10 @@ impl> ResponseCache { Self { backend } } + pub fn backend(&self) -> &B { + &self.backend + } + pub fn default_ttl(&self) -> Duration { self.backend.default_ttl() } @@ -67,7 +71,7 @@ impl> ResponseCache { Err(Error::InvalidEntry) => None, Err(error) => return Err(error), }; - Self::fresh_or_miss(entry, now, request.max_age) + Ok(Self::fresh_or_miss(entry, now, request.max_age)) } pub async fn async_lookup( @@ -87,7 +91,7 @@ impl> ResponseCache { Err(Error::InvalidEntry) => None, Err(error) => return Err(error), }; - Self::fresh_or_miss(entry, now, request.max_age) + Ok(Self::fresh_or_miss(entry, now, request.max_age)) } pub fn lookup_batch( @@ -180,11 +184,26 @@ impl> ResponseCache { &self, entries: Vec<(ResponseCacheRequest, Value)>, now: Duration, + ) -> Result<(), Error> { + self.async_store_entries( + entries + .into_iter() + .map(|(request, response)| (request, response, now)) + .collect(), + ) + .await + } + + /// Stores entries that each carry the time they were produced, so a deferred write keeps + /// the freshness of its original response. + pub async fn async_store_entries( + &self, + entries: Vec<(ResponseCacheRequest, Value, Duration)>, ) -> Result<(), Error> { let writable = entries .into_iter() - .filter(|(request, _)| request.controls.writes()) - .map(|(request, response)| { + .filter(|(request, _, _)| request.controls.writes()) + .map(|(request, response, now)| { ( cache_key(&request.key), CacheEntry { @@ -227,7 +246,7 @@ impl> ResponseCache { let mut values = vec![None; requests.len()]; for ((index, request), entry) in readable.into_iter().zip(entries) { let response = match entry { - BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age)?, + BatchEntry::Hit(entry) => Self::fresh_or_miss(Some(entry), now, request.max_age), BatchEntry::Miss | BatchEntry::Invalid => None, }; values[index] = response; @@ -239,24 +258,9 @@ impl> ResponseCache { entry: Option, now: Duration, max_age: Option, - ) -> Result, Error> { - match Self::fresh_response(entry, now, max_age) { - Err(Error::InvalidEntry) => Ok(None), - result => result, - } - } - - fn fresh_response( - entry: Option, - now: Duration, - max_age: Option, - ) -> Result, Error> { + ) -> Option { entry .filter(|entry| entry.fresh(now, max_age)) - .map(|entry| match (entry.timestamp, entry.response) { - (Some(_), Value::String(text)) => crate::codec::decode_value(&text), - (_, value) => Ok(value), - }) - .transpose() + .map(|entry| entry.response) } } diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index e4a04b3dec5..7c69e5a1d55 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -273,22 +273,45 @@ async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis() { } #[test] -fn malformed_memory_entries_are_treated_as_misses() { - let backend = Arc::new(InMemoryCache::default()); - BaseCache::set_cache( - backend.as_ref(), - "tenant:key", - CacheEntry { +fn string_responses_round_trip_through_typed_and_wire_backends() { + let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); + let now = Duration::from_secs(100); + for response in [json!("hello world"), json!("123"), json!("null")] { + cache.store(&request(), response.clone(), now).unwrap(); + assert_eq!( + cache.lookup(&request(), now).unwrap(), + Some(response.clone()) + ); + + let wire = ResponseCacheCodec + .encode(&CacheEntry { + timestamp: Some(100.0), + response: response.clone(), + }) + .unwrap(); + assert_eq!(ResponseCacheCodec.decode(&wire).unwrap().response, response); + } +} + +#[test] +fn non_object_responses_are_written_as_python_readable_serialized_strings() { + let wire = ResponseCacheCodec + .encode(&CacheEntry { timestamp: Some(100.0), - response: json!("not a serialized response"), - }, - Default::default(), - ) - .unwrap(); - let cache = ResponseCache::new(backend); + response: json!([1, 2]), + }) + .unwrap(); assert_eq!( - cache.lookup(&request(), Duration::from_secs(100)).unwrap(), - None + serde_json::from_slice::(&wire).unwrap(), + json!({"timestamp": 100.0, "response": "[1,2]"}) + ); + assert_eq!( + ResponseCacheCodec.decode(&wire).unwrap().response, + json!([1, 2]) + ); + assert_eq!( + ResponseCacheCodec.decode(br#"{"timestamp": 100.0, "response": "not serialized"}"#), + Err(Error::InvalidEntry) ); } @@ -367,3 +390,27 @@ async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses() { None ); } + +#[tokio::test] +async fn deferred_entries_keep_the_time_they_were_produced() { + let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); + let mut request = request(); + request.max_age = Some(Duration::from_secs(10)); + cache + .async_store_entries(vec![( + request.clone(), + json!({"answer": 7}), + Duration::from_secs(100), + )]) + .await + .unwrap(); + + assert_eq!( + cache.lookup(&request, Duration::from_secs(110)).unwrap(), + Some(json!({"answer": 7})) + ); + assert_eq!( + cache.lookup(&request, Duration::from_secs(111)).unwrap(), + None + ); +} diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 5bc1ebc2945..d6ef8052c4f 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::time::Duration; +use std::{future::Future, time::Duration}; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/cache/src/caching.rs b/litellm-rust/crates/cache/src/caching.rs index 39479694f3b..4ec94244ac5 100644 --- a/litellm-rust/crates/cache/src/caching.rs +++ b/litellm-rust/crates/cache/src/caching.rs @@ -1,8 +1,7 @@ use std::sync::Arc; -use crate::{BaseCache, CacheKwargs, Error}; - pub use crate::BaseCache as Cache; +use crate::{BaseCache, CacheKwargs, Error}; pub fn get_cache( cache: &B, diff --git a/litellm-rust/crates/cache/src/codec.rs b/litellm-rust/crates/cache/src/codec.rs index 09bee6032f6..6d47c682406 100644 --- a/litellm-rust/crates/cache/src/codec.rs +++ b/litellm-rust/crates/cache/src/codec.rs @@ -14,6 +14,14 @@ pub trait CacheCodec: Send + Sync { pub struct JsonCodec(PhantomData V>); +impl Clone for JsonCodec { + fn clone(&self) -> Self { + *self + } +} + +impl Copy for JsonCodec {} + impl Default for JsonCodec { fn default() -> Self { Self::new() diff --git a/litellm-rust/crates/cache/src/dual.rs b/litellm-rust/crates/cache/src/dual.rs index 17a2c430ddd..9858c5e2748 100644 --- a/litellm-rust/crates/cache/src/dual.rs +++ b/litellm-rust/crates/cache/src/dual.rs @@ -1,15 +1,140 @@ -use std::sync::Arc; +use std::{sync::Arc, time::Duration}; -use crate::{BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error}; +use crate::{ + BaseCache, BatchEntry, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, +}; + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum ReadPolicy { + #[default] + LocalThenRemote, + LocalOnly, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum WritePolicy { + #[default] + Both, + LocalOnly, +} + +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub enum RemoteFailurePolicy { + #[default] + Propagate, + UseLocal, +} pub struct DualCache { l1: Arc, l2: Arc, + read_policy: ReadPolicy, + write_policy: WritePolicy, + remote_failure_policy: RemoteFailurePolicy, + promotion_ttl: Option, } impl DualCache { pub fn new(l1: Arc, l2: Arc) -> Self { - Self { l1, l2 } + Self { + l1, + l2, + read_policy: ReadPolicy::default(), + write_policy: WritePolicy::default(), + remote_failure_policy: RemoteFailurePolicy::default(), + promotion_ttl: None, + } + } + + pub fn with_read_policy(self, read_policy: ReadPolicy) -> Self { + Self { + read_policy, + ..self + } + } + + pub fn with_write_policy(self, write_policy: WritePolicy) -> Self { + Self { + write_policy, + ..self + } + } + + pub fn with_remote_failure_policy(self, remote_failure_policy: RemoteFailurePolicy) -> Self { + Self { + remote_failure_policy, + ..self + } + } + + pub fn with_promotion_ttl(self, promotion_ttl: Duration) -> Self { + Self { + promotion_ttl: Some(promotion_ttl), + ..self + } + } + + fn reads_remote(&self) -> bool { + self.read_policy == ReadPolicy::LocalThenRemote + } + + fn writes_remote(&self) -> bool { + self.write_policy == WritePolicy::Both + } + + fn remote(&self, result: Result) -> Result, Error> { + match result { + Ok(value) => Ok(Some(value)), + Err(Error::Unavailable) + if self.remote_failure_policy == RemoteFailurePolicy::UseLocal => + { + Ok(None) + } + Err(error) => Err(error), + } + } + + fn promotion_kwargs(&self, kwargs: &CacheKwargs) -> CacheKwargs { + CacheKwargs { + ttl: self.promotion_ttl.or(kwargs.ttl), + extras: kwargs.extras.clone(), + } + } +} + +impl DualCache +where + V: Clone + Send + Sync + 'static, + L1: BaseCache, + L2: BaseCache, +{ + fn missing(entries: &[BatchEntry]) -> Vec { + entries + .iter() + .enumerate() + .filter_map(|(index, entry)| (!matches!(entry, BatchEntry::Hit(_))).then_some(index)) + .collect() + } + + fn merge_batch( + &self, + keys: &[String], + kwargs: &CacheKwargs, + mut entries: Vec>, + missing: Vec, + remote: Vec>, + ) -> Result>, Error> { + if missing.len() != remote.len() { + return Err(Error::Unavailable); + } + for (index, entry) in missing.into_iter().zip(remote) { + if let BatchEntry::Hit(value) = &entry { + self.l1 + .set_cache(&keys[index], value.clone(), self.promotion_kwargs(kwargs))?; + } + entries[index] = entry; + } + Ok(entries) } } @@ -21,12 +146,14 @@ where { type Value = V; - fn default_ttl(&self) -> std::time::Duration { + fn default_ttl(&self) -> Duration { self.l2.default_ttl() } fn set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> { - self.l2.set_cache(key, value.clone(), kwargs.clone())?; + if self.writes_remote() { + self.remote(self.l2.set_cache(key, value.clone(), kwargs.clone()))?; + } self.l1.set_cache(key, value, kwargs) } @@ -34,25 +161,130 @@ where if let Some(value) = self.l1.get_cache(key, kwargs)? { return Ok(Some(value)); } - let value = self.l2.get_cache(key, kwargs)?; + if !self.reads_remote() { + return Ok(None); + } + let value = self.remote(self.l2.get_cache(key, kwargs))?.flatten(); if let Some(value) = &value { - self.l1.set_cache(key, value.clone(), kwargs.clone())?; + self.l1 + .set_cache(key, value.clone(), self.promotion_kwargs(kwargs))?; } Ok(value) } + fn get_cache_batch( + &self, + keys: &[String], + kwargs: &CacheKwargs, + ) -> Result>, Error> { + let entries = self.l1.get_cache_batch(keys, kwargs)?; + let missing = Self::missing(&entries); + if missing.is_empty() || !self.reads_remote() { + return Ok(entries); + } + let remote_keys = missing + .iter() + .map(|index| keys[*index].clone()) + .collect::>(); + match self.remote(self.l2.get_cache_batch(&remote_keys, kwargs))? { + Some(remote) => self.merge_batch(keys, kwargs, entries, missing, remote), + None => Ok(entries), + } + } + + async fn async_set_cache(&self, key: &str, value: V, kwargs: CacheKwargs) -> Result<(), Error> { + if self.writes_remote() { + self.remote( + self.l2 + .async_set_cache(key, value.clone(), kwargs.clone()) + .await, + )?; + } + self.l1.async_set_cache(key, value, kwargs).await + } + + async fn async_get_cache(&self, key: &str, kwargs: &CacheKwargs) -> Result, Error> { + if let Some(value) = self.l1.async_get_cache(key, kwargs).await? { + return Ok(Some(value)); + } + if !self.reads_remote() { + return Ok(None); + } + let value = self + .remote(self.l2.async_get_cache(key, kwargs).await)? + .flatten(); + if let Some(value) = &value { + self.l1 + .async_set_cache(key, value.clone(), self.promotion_kwargs(kwargs)) + .await?; + } + Ok(value) + } + + async fn async_get_cache_batch( + &self, + keys: Vec, + kwargs: CacheKwargs, + ) -> Result>, Error> { + let entries = self + .l1 + .async_get_cache_batch(keys.clone(), kwargs.clone()) + .await?; + let missing = Self::missing(&entries); + if missing.is_empty() || !self.reads_remote() { + return Ok(entries); + } + let remote_keys = missing.iter().map(|index| keys[*index].clone()).collect(); + match self.remote( + self.l2 + .async_get_cache_batch(remote_keys, kwargs.clone()) + .await, + )? { + Some(remote) => self.merge_batch(&keys, &kwargs, entries, missing, remote), + None => Ok(entries), + } + } + + async fn async_set_cache_pipeline( + &self, + cache_list: Vec<(String, V)>, + kwargs: CacheKwargs, + ) -> Result<(), Error> { + if self.writes_remote() { + self.remote( + self.l2 + .async_set_cache_pipeline(cache_list.clone(), kwargs.clone()) + .await, + )?; + } + self.l1.async_set_cache_pipeline(cache_list, kwargs).await + } + fn delete_cache(&self, key: &str) -> Result<(), Error> { - self.l2.delete_cache(key)?; + if self.writes_remote() { + self.remote(self.l2.delete_cache(key))?; + } self.l1.delete_cache(key) } + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { + if self.writes_remote() { + self.remote(self.l2.async_delete_cache(key).await)?; + } + self.l1.async_delete_cache(key).await + } + fn flush_cache(&self) -> Result<(), Error> { - self.l2.flush_cache()?; + if self.writes_remote() { + self.remote(self.l2.flush_cache())?; + } self.l1.flush_cache() } async fn async_flush_cache(&self) -> Result<(), Error> { - self.l2.async_flush_cache().await?; + if self.writes_remote() { + self.remote(self.l2.async_flush_cache().await)?; + } self.l1.async_flush_cache().await } @@ -76,6 +308,20 @@ where self.l1.set_cache(key, value, kwargs)?; Ok(value) } + + async fn async_increment_cache( + &self, + key: &str, + amount: f64, + kwargs: CacheKwargs, + ) -> Result { + let value = self + .l2 + .async_increment_cache(key, amount, kwargs.clone()) + .await?; + self.l1.async_set_cache(key, value, kwargs).await?; + Ok(value) + } } impl ClaimCache for DualCache @@ -91,15 +337,39 @@ where eligible: &[V], kwargs: CacheKwargs, ) -> Result { - match self - .l2 - .claim_cache(key, candidate.clone(), eligible, kwargs.clone()) - { - Ok(winner) => { + match self.remote( + self.l2 + .claim_cache(key, candidate.clone(), eligible, kwargs.clone()), + )? { + Some(winner) => { self.l1.set_cache(key, winner.clone(), kwargs)?; Ok(winner) } - Err(_) => self.l1.claim_cache(key, candidate, eligible, kwargs), + None => self.l1.claim_cache(key, candidate, eligible, kwargs), + } + } + + async fn async_claim_cache( + &self, + key: &str, + candidate: V, + eligible: Vec, + kwargs: CacheKwargs, + ) -> Result { + match self.remote( + self.l2 + .async_claim_cache(key, candidate.clone(), eligible.clone(), kwargs.clone()) + .await, + )? { + Some(winner) => { + self.l1.async_set_cache(key, winner.clone(), kwargs).await?; + Ok(winner) + } + None => { + self.l1 + .async_claim_cache(key, candidate, eligible, kwargs) + .await + } } } } diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 824de00bdf4..5c27aa9ff4e 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,6 +1,7 @@ -use litellm_cache::{BaseCache, CacheConnectionResult, CacheKwargs, Error}; use std::{sync::Mutex, time::Duration}; +use litellm_cache::{BaseCache, CacheConnectionResult, CacheKwargs, Error}; + struct TestCache { default_ttl: Duration, writes: Mutex>, diff --git a/litellm-rust/crates/cache/tests/dual.rs b/litellm-rust/crates/cache/tests/dual.rs index 1be1556734c..2e1e72c7119 100644 --- a/litellm-rust/crates/cache/tests/dual.rs +++ b/litellm-rust/crates/cache/tests/dual.rs @@ -4,7 +4,8 @@ use std::{ }; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, dual::DualCache, + BaseCache, CacheConnectionResult, CacheKwargs, ClaimCache, CounterCache, Error, + dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy}, }; struct TestCache { @@ -111,7 +112,8 @@ fn failed_l2_increment_leaves_l1_unchanged() { #[test] fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() { let l1 = Arc::new(TestCache::new(Some("first".to_string()), false)); - let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true))); + let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true))) + .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); assert_eq!( cache @@ -128,3 +130,213 @@ fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() { "first" ); } + +struct SyncPanics(TestCache); + +impl BaseCache for SyncPanics { + type Value = String; + + fn set_cache(&self, _: &str, _: String, _: CacheKwargs) -> Result<(), Error> { + panic!("sync L2 write on an async path") + } + + fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result, Error> { + panic!("sync L2 read on an async path") + } + + async fn async_set_cache( + &self, + key: &str, + value: String, + kwargs: CacheKwargs, + ) -> Result<(), Error> { + self.0.set_cache(key, value, kwargs) + } + + async fn async_get_cache( + &self, + key: &str, + kwargs: &CacheKwargs, + ) -> Result, 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, + }]) + } + + async fn async_set_cache_pipeline( + &self, + cache_list: Vec<(String, String)>, + kwargs: CacheKwargs, + ) -> Result<(), Error> { + for (key, value) in cache_list { + self.0.set_cache(&key, value, kwargs.clone())?; + } + Ok(()) + } + + fn delete_cache(&self, _: &str) -> Result<(), Error> { + panic!("sync L2 delete on an async path") + } + + async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { + self.0.delete_cache(key) + } + + fn flush_cache(&self) -> Result<(), Error> { + panic!("sync L2 flush on an async path") + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + unreachable!() + } +} + +#[tokio::test] +async fn async_operations_use_the_async_l2_methods() { + let l1 = Arc::new(TestCache::new(None, false)); + let cache = DualCache::new( + l1.clone(), + Arc::new(SyncPanics(TestCache::new( + Some("remote".to_string()), + false, + ))), + ); + let kwargs = CacheKwargs::default(); + + assert_eq!( + cache.async_get_cache("missing", &kwargs).await.unwrap(), + Some("remote".into()) + ); + assert_eq!( + l1.get_cache("missing", &kwargs).unwrap(), + Some("remote".into()) + ); + + l1.delete_cache("missing").unwrap(); + assert_eq!( + cache + .async_get_cache_batch(vec!["missing".into()], kwargs.clone()) + .await + .unwrap(), + [litellm_cache::BatchEntry::Hit("remote".to_string())] + ); + cache + .async_set_cache("missing", "written".into(), kwargs.clone()) + .await + .unwrap(); + cache + .async_set_cache_pipeline(vec![("missing".into(), "piped".into())], kwargs.clone()) + .await + .unwrap(); + cache.async_delete_cache("missing").await.unwrap(); + assert_eq!( + cache.async_get_cache("missing", &kwargs).await.unwrap(), + None + ); +} + +struct Unavailable; + +impl BaseCache for Unavailable { + type Value = String; + + fn set_cache(&self, _: &str, _: String, _: CacheKwargs) -> Result<(), Error> { + Err(Error::Unavailable) + } + + fn get_cache(&self, _: &str, _: &CacheKwargs) -> Result, Error> { + Err(Error::Unavailable) + } + + fn delete_cache(&self, _: &str) -> Result<(), Error> { + Err(Error::Unavailable) + } + + fn flush_cache(&self) -> Result<(), Error> { + Err(Error::Unavailable) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + unreachable!() + } +} + +impl ClaimCache for Unavailable { + fn claim_cache( + &self, + _: &str, + _: String, + _: &[String], + _: CacheKwargs, + ) -> Result { + Err(Error::InvalidEntry) + } +} + +#[test] +fn remote_failure_policy_selects_propagation_or_the_local_tier() { + let kwargs = CacheKwargs::default(); + let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable)); + assert_eq!( + strict.set_cache("key", "value".into(), kwargs.clone()), + Err(Error::Unavailable) + ); + assert_eq!(strict.get_cache("key", &kwargs), Err(Error::Unavailable)); + + let l1 = Arc::new(TestCache::new(None, false)); + let degraded = DualCache::new(l1.clone(), Arc::new(Unavailable)) + .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); + assert_eq!(degraded.get_cache("key", &kwargs), Ok(None)); + degraded + .set_cache("key", "value".into(), kwargs.clone()) + .unwrap(); + assert_eq!(degraded.get_cache("key", &kwargs), Ok(Some("value".into()))); + degraded.delete_cache("key").unwrap(); + assert_eq!(l1.get_cache("key", &kwargs), Ok(None)); +} + +#[test] +fn claim_fallback_does_not_hide_non_availability_errors() { + let cache = DualCache::new( + Arc::new(TestCache::new(Some("first".to_string()), false)), + Arc::new(Unavailable), + ) + .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); + assert_eq!( + cache.claim_cache("affinity", "second".into(), &[], CacheKwargs::default()), + Err(Error::InvalidEntry) + ); +} + +#[test] +fn local_only_policies_never_touch_l2() { + let l2 = Arc::new(TestCache::new(Some("remote".to_string()), false)); + let cache = DualCache::new(Arc::new(TestCache::new(None, false)), l2.clone()) + .with_read_policy(ReadPolicy::LocalOnly) + .with_write_policy(WritePolicy::LocalOnly); + let kwargs = CacheKwargs::default(); + + assert_eq!(cache.get_cache("key", &kwargs), Ok(None)); + cache + .set_cache("key", "local".into(), kwargs.clone()) + .unwrap(); + assert_eq!(l2.get_cache("key", &kwargs), Ok(Some("remote".into()))); +} diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 58550c2987d..2ad13ce200f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -1,3 +1,5 @@ +use std::time::Duration; + use litellm_host_python::from_py; use pyo3::{ PyTraverseError, PyVisit, @@ -6,9 +8,8 @@ use pyo3::{ types::{PyDict, PyTuple, PyType}, }; use serde_json::Value; -use std::time::Duration; -use super::{NativeCacheHandle, native::NativeResponseCache}; +use super::{CacheTestHandle, native::NativeResponseCache}; struct ClassGuard { class: Py, @@ -130,9 +131,10 @@ impl FacadeGuard { pub(super) fn capture( py: Python<'_>, facade: &Bound<'_, PyAny>, - kind: &str, - native_default_ttl: Duration, + service: &NativeResponseCache, ) -> PyResult { + let kind = service.kind(); + let native_default_ttl: Duration = service.default_ttl(); let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?; if !facade.get_type().is(&cache_type) { return Err(PyTypeError::new_err( @@ -158,6 +160,23 @@ impl FacadeGuard { "facade and native backend default TTLs must match", )); } + let namespace = match backend.getattr_opt("namespace")? { + Some(namespace) => namespace.extract::>()?, + None => None, + } + .filter(|namespace| !namespace.is_empty()); + if kind == "redis" && namespace.as_deref() != service.namespace() { + return Err(PyTypeError::new_err( + "facade and native backend namespaces must match", + )); + } + if let Some(capacity) = service.capacity() + && backend.getattr("max_size_in_memory")?.extract::()? != capacity + { + return Err(PyTypeError::new_err( + "facade and native backend capacities must match", + )); + } Ok(Self { outer: ObjectGuard::capture( py, @@ -169,6 +188,7 @@ impl FacadeGuard { "namespace", "supported_call_types", "redis_flush_size", + "semantic_cache_scope", ], )?, backend: ObjectGuard::capture( @@ -179,6 +199,8 @@ impl FacadeGuard { "default_ttl", "max_size_in_memory", "max_size_per_item", + "redis_kwargs", + "redis_flush_size", ], )?, }) @@ -208,7 +230,7 @@ pub(super) fn resolve( let Some(handle) = dict.get_item("_native_cache_handle")? else { return Ok(None); }; - let Ok(handle) = handle.extract::>() else { + let Ok(handle) = handle.extract::>() else { return Ok(None); }; let Some(guard) = &handle.guard else { diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 1e2e42600ff..f83788f6036 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -3,21 +3,21 @@ mod native; use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use facade::FacadeGuard; use litellm_cache::Error; use litellm_cache_response::{CacheControls, CacheKeyInput, PartialHits, ResponseCacheRequest}; use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; +use native::NativeResponseCache; use pyo3::{ PyTraverseError, PyVisit, exceptions::{PyRuntimeError, PyTypeError, PyValueError}, prelude::*, - types::{PyDict, PyList}, + types::{PyDict, PyList, PyTuple}, }; use serde::Deserialize; use serde_json::Value; use crate::python_settings::PythonSettings; -use facade::FacadeGuard; -use native::NativeResponseCache; const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); @@ -84,14 +84,14 @@ fn cache_error(error: Error) -> PyErr { } } -#[pyclass(frozen)] -pub(crate) struct NativeCacheHandle { +#[pyclass(frozen, name = "_CacheTestHandle")] +pub(crate) struct CacheTestHandle { service: NativeResponseCache, guard: Option, pid: u32, } -impl NativeCacheHandle { +impl CacheTestHandle { fn service(&self) -> PyResult { if self.pid != std::process::id() { return Err(PyRuntimeError::new_err( @@ -103,7 +103,7 @@ impl NativeCacheHandle { } #[pymethods] -impl NativeCacheHandle { +impl CacheTestHandle { #[staticmethod] #[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))] fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult { @@ -140,9 +140,9 @@ impl NativeCacheHandle { self.service.kind() } - fn bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { + fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { let service = self.service()?; - let guard = FacadeGuard::capture(py, facade, self.backend(), service.default_ttl())?; + let guard = FacadeGuard::capture(py, facade, &service)?; let service = service.with_redis_flush_size( facade .getattr("redis_flush_size")? @@ -173,7 +173,7 @@ enum CacheBinding { PythonCallback(Py), } -#[pyclass(frozen, name = "CacheBinding")] +#[pyclass(frozen, name = "_CacheTestBinding")] pub(crate) struct ResolvedCache { binding: CacheBinding, pid: u32, @@ -285,12 +285,15 @@ impl ResolvedCache { } } + /// Native bindings return `{values, missing_indices}`. The built-in `Cache` API has no batch + /// read, so a Python callback receives one `get_cache(**kwargs)` call per request, in order, + /// and the results come back as a list. #[pyo3(signature = (requests, *, callback_kwargs=None))] fn lookup_batch( &self, py: Python<'_>, requests: &Bound<'_, PyAny>, - callback_kwargs: Option<&Bound<'_, PyDict>>, + callback_kwargs: Option<&Bound<'_, PyAny>>, ) -> PyResult> { self.check_process()?; match &self.binding { @@ -305,14 +308,17 @@ impl ResolvedCache { .map_err(cache_error)?; to_py(py, &response) } - CacheBinding::PythonCallback(object) => object - .bind(py) - .call_method( - "batch_get_cache", - (callback_keys(py, requests)?,), - Some(self::callback_kwargs(callback_kwargs)?), - ) - .map(Bound::unbind), + CacheBinding::PythonCallback(object) => { + let results = PyList::empty(py); + for kwargs in batch_callback_kwargs(requests, callback_kwargs)? { + results.append(object.bind(py).call_method( + "get_cache", + (), + Some(&kwargs), + )?)?; + } + Ok(results.into_any().unbind()) + } } } @@ -364,7 +370,7 @@ impl ResolvedCache { &self, py: Python<'py>, requests: &Bound<'py, PyAny>, - callback_kwargs: Option<&Bound<'py, PyDict>>, + callback_kwargs: Option<&Bound<'py, PyAny>>, ) -> PyResult> { self.check_process()?; match &self.binding { @@ -381,20 +387,30 @@ impl ResolvedCache { cache_error, ) } - CacheBinding::PythonCallback(object) => object.bind(py).call_method( - "async_batch_get_cache", - (callback_keys(py, requests)?,), - Some(self::callback_kwargs(callback_kwargs)?), - ), + CacheBinding::PythonCallback(object) => { + let awaitables = batch_callback_kwargs(requests, callback_kwargs)? + .iter() + .map(|kwargs| { + object + .bind(py) + .call_method("async_get_cache", (), Some(kwargs)) + }) + .collect::>>()?; + py.import("asyncio")? + .call_method1("gather", PyTuple::new(py, awaitables)?) + } } } - #[pyo3(signature = (requests, responses, *, callback_kwargs=None))] + /// A Python callback receives the caller's original result through `callback_result`, because + /// the built-in `Cache.async_add_cache_pipeline` splits the batch itself. + #[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))] fn async_store_batch<'py>( &self, py: Python<'py>, requests: &Bound<'py, PyAny>, responses: &Bound<'py, PyAny>, + callback_result: Option<&Bound<'py, PyAny>>, callback_kwargs: Option<&Bound<'py, PyDict>>, ) -> PyResult> { self.check_process()?; @@ -417,20 +433,14 @@ impl ResolvedCache { ) } CacheBinding::PythonCallback(object) => { - let keys = callback_keys(py, requests)?; - let responses = responses.try_iter()?.collect::>>()?; - if keys.len() != responses.len() { - return Err(PyValueError::new_err( - "batch cache requests and responses must have equal lengths", - )); - } - let cache_list = PyList::empty(py); - for (key, response) in keys.iter().zip(responses) { - cache_list.append((key, response))?; - } + let result = callback_result.ok_or_else(|| { + PyTypeError::new_err( + "Python cache callbacks require their original callback_result", + ) + })?; object.bind(py).call_method( - "async_set_cache_pipeline", - (cache_list,), + "async_add_cache_pipeline", + (result,), Some(self::callback_kwargs(callback_kwargs)?), ) } @@ -445,8 +455,17 @@ impl ResolvedCache { let service = service.clone(); run_async(py, async move { service.async_flush().await }, cache_error) } + // The built-in `Cache` facade has no flush of its own; its backend does. CacheBinding::PythonCallback(object) => { - object.bind(py).call_method0("flush_cache")?; + let object = object.bind(py); + let backend = match object.getattr_opt("cache")? { + Some(backend) if !backend.is_none() => backend, + _ => object.clone(), + }; + if backend.hasattr("async_flush_cache")? { + return backend.call_method0("async_flush_cache"); + } + backend.call_method0("flush_cache")?; ready_none(py) } } @@ -464,7 +483,7 @@ impl ResolvedCache { cache_error, ) } - CacheBinding::PythonCallback(object) => object.bind(py).call_method0("test_connection"), + CacheBinding::PythonCallback(object) => object.bind(py).call_method0("ping"), } } @@ -484,16 +503,25 @@ fn callback_kwargs<'a, 'py>( }) } -fn callback_keys<'py>( - py: Python<'py>, +fn batch_callback_kwargs<'py>( requests: &Bound<'py, PyAny>, -) -> PyResult> { - PyList::new( - py, - self::requests(requests)? - .into_iter() - .map(|request| litellm_cache_response::cache_key(&request.key)), - ) + kwargs: Option<&Bound<'py, PyAny>>, +) -> PyResult>> { + let kwargs = kwargs + .ok_or_else(|| { + PyTypeError::new_err( + "Python cache callbacks require one original callback_kwargs mapping per request", + ) + })? + .try_iter()? + .map(|item| Ok(item?.cast_into::()?)) + .collect::>>()?; + if kwargs.len() != requests.len()? { + return Err(PyValueError::new_err( + "batch cache requests and callback_kwargs must have equal lengths", + )); + } + Ok(kwargs) } fn ready_none(py: Python<'_>) -> PyResult> { @@ -512,13 +540,13 @@ fn ready_value<'py, T: serde::Serialize>( Ok(future) } -#[pyclass(frozen)] -pub(crate) struct CacheResolver { +#[pyclass(frozen, name = "_CacheTestResolver")] +pub(crate) struct CacheTestResolver { namespace: Py, } #[pymethods] -impl CacheResolver { +impl CacheTestResolver { #[new] fn new(namespace: Py) -> Self { Self { namespace } @@ -528,7 +556,7 @@ impl CacheResolver { let object = self.namespace.bind(py).getattr("cache")?; let binding = if object.is_none() { CacheBinding::Disabled - } else if let Ok(handle) = object.extract::>() { + } else if let Ok(handle) = object.extract::>() { CacheBinding::Native(handle.service()?) } else if let Some(service) = facade::resolve(py, &object)? { CacheBinding::Native(service) diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 20891719550..3fc8f61dff6 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -3,12 +3,11 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; -use serde_json::Value; -use tokio::sync::Mutex; - use litellm_cache_response::{ CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, }; +use serde_json::Value; +use tokio::sync::Mutex; #[derive(Clone)] pub(super) enum NativeResponseCache { @@ -21,7 +20,7 @@ pub(super) enum NativeResponseCache { pub(super) struct RedisWriteBuffer { flush_size: usize, - entries: Mutex>, + entries: Mutex>, } impl NativeResponseCache { @@ -67,6 +66,20 @@ impl NativeResponseCache { } } + pub fn namespace(&self) -> Option<&str> { + match self { + Self::Memory(_) => None, + Self::Redis { cache, .. } => cache.backend().namespace(), + } + } + + pub fn capacity(&self) -> Option { + match self { + Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), + Self::Redis { .. } => None, + } + } + pub fn with_redis_flush_size(self, flush_size: Option) -> Self { match self { Self::Redis { cache, .. } => Self::Redis { @@ -145,19 +158,15 @@ impl NativeResponseCache { } => { let pending = { let mut entries = buffer.entries.lock().await; - entries.push((request.clone(), response)); + entries.push((request.clone(), response, now)); (entries.len() >= buffer.flush_size).then(|| std::mem::take(&mut *entries)) }; - let Some(pending) = pending else { - return Ok(()); - }; - if let Err(error) = cache.async_store_batch(pending.clone(), now).await { - let mut entries = buffer.entries.lock().await; - let current = std::mem::take(&mut *entries); - *entries = pending.into_iter().chain(current).collect(); - return Err(error); + // A failed flush drops its batch, as Python does. Requeueing would grow the + // buffer and re-send an ever larger pipeline on every write during an outage. + match pending { + Some(pending) => cache.async_store_entries(pending).await, + None => Ok(()), } - Ok(()) } } } diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index 621c111a35b..bd62c5aadf1 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -10,8 +10,7 @@ mod token_counter; #[pymodule(gil_used = true)] mod _native { - #[pymodule_export] - use crate::cache::{CacheResolver, NativeCacheHandle, ResolvedCache}; + use crate::cache::{CacheTestHandle, CacheTestResolver, ResolvedCache}; #[cfg(feature = "panic-test")] #[pymodule_export] use crate::diagnostics::_panic_for_test; @@ -35,6 +34,16 @@ mod _native { use crate::token_counter::TokenCounter; #[pymodule_export] use litellm_host_python::{ForkedAfterNativeRuntimeStarted, ProcessReservedForForking}; + use pyo3::{prelude::*, types::PyModule}; + + #[pymodule_init] + fn init(module: &Bound<'_, PyModule>) -> PyResult<()> { + let py = module.py(); + let dict = module.dict(); + dict.set_item("_CacheTestHandle", py.get_type::())?; + dict.set_item("_CacheTestResolver", py.get_type::())?; + dict.set_item("_CacheTestBinding", py.get_type::()) + } } use pyo3::prelude::*; @@ -68,9 +77,6 @@ mod tests { "achat_completions", "ResponsesWebSocketConnection", "TokenCounter", - "CacheResolver", - "NativeCacheHandle", - "CacheBinding", "gil_stats", "process_state_started", "reserve_process_for_forking", diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 68b1742dc41..ab4639bc876 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -94,25 +94,25 @@ class ResponsesWebSocketConnection: def close(self) -> Future[None]: ... @final -class NativeCacheHandle: +class _CacheTestHandle: def __new__(cls, _uninstantiable: Never, /) -> Never: ... @staticmethod def memory( *, capacity: int = 200, ttl_seconds: float = 600.0, max_entry_bytes: int = 1048576 - ) -> NativeCacheHandle: ... + ) -> _CacheTestHandle: ... @staticmethod - def redis(url: str, *, ttl_seconds: float | None = None, namespace: str | None = None) -> NativeCacheHandle: ... + def redis(url: str, *, ttl_seconds: float | None = None, namespace: str | None = None) -> _CacheTestHandle: ... @property def backend(self) -> str: ... - def bind_facade(self, facade: object) -> None: ... + def _bind_facade(self, facade: object) -> None: ... @final -class CacheResolver: - def __new__(cls, namespace: object) -> CacheResolver: ... - def resolve(self) -> CacheBinding: ... +class _CacheTestResolver: + def __new__(cls, namespace: object) -> _CacheTestResolver: ... + def resolve(self) -> _CacheTestBinding: ... @final -class CacheBinding: +class _CacheTestBinding: def __new__(cls, _uninstantiable: Never, /) -> Never: ... @property def kind(self) -> str: ... @@ -130,7 +130,7 @@ class CacheBinding: self, requests: Sequence[Mapping[str, object]], *, - callback_kwargs: dict[str, object] | None = None, + callback_kwargs: Sequence[dict[str, object]] | None = None, ) -> object: ... def async_lookup( self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None @@ -146,17 +146,18 @@ class CacheBinding: self, requests: Sequence[Mapping[str, object]], *, - callback_kwargs: dict[str, object] | None = None, + callback_kwargs: Sequence[dict[str, object]] | None = None, ) -> Awaitable[object]: ... def async_store_batch( self, requests: Sequence[Mapping[str, object]], responses: Sequence[object], *, + callback_result: object = None, callback_kwargs: dict[str, object] | None = None, ) -> Awaitable[object]: ... def async_flush(self) -> Awaitable[None]: ... - def ping(self) -> Awaitable[dict[str, object] | None]: ... + def ping(self) -> Awaitable[object]: ... @final class TokenCounter: @@ -174,10 +175,7 @@ def process_state_started() -> bool: ... def reserve_process_for_forking() -> None: ... __all__ = [ - "CacheBinding", - "CacheResolver", "ForkedAfterNativeRuntimeStarted", - "NativeCacheHandle", "ProcessReservedForForking", "ResponsesWebSocketConnection", "RustBridgeDeclined", diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index 05a4ef68e0b..d34d23ca1d7 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -6,10 +6,10 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest +from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, _redis_circuit_breaker_guard, _redis_circuit_breaker_guard_sync -from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE from litellm.types.caching import RedisPipelineIncrementOperation diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index d38583dad0a..cf46a0566f7 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -15,7 +15,7 @@ import pytest import redis import litellm -from litellm.caching.caching import Cache +from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache from litellm.caching.in_memory_cache import InMemoryCache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType @@ -50,20 +50,43 @@ def test_existing_constructor_and_global_are_unchanged() -> None: assert type(facade.cache) is InMemoryCache assert "_native_cache_handle" not in vars(facade) with rebound(litellm, "cache", facade): - resolver: Final = _native.CacheResolver(litellm) + resolver: Final = _native._CacheTestResolver(litellm) assert resolver.resolve().kind == "python_callback" resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} +def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None: + resolver: Final = _native._CacheTestResolver(litellm) + + enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30) + enabled: Final = litellm.cache + assert isinstance(enabled, Cache) + assert enabled.ttl == 30 + assert resolver.resolve().kind == "python_callback" + + enable_cache(type=LiteLLMCacheType.LOCAL, ttl=60) + assert litellm.cache is enabled + + update_cache(type=LiteLLMCacheType.LOCAL, ttl=60) + updated: Final = litellm.cache + assert isinstance(updated, Cache) + assert updated is not enabled + assert updated.ttl == 60 + + disable_cache() + assert litellm.cache is None + assert resolver.resolve().kind == "disabled" + + async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: - namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.memory()) - resolver: Final = _native.CacheResolver(namespace) + namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.memory()) + resolver: Final = _native._CacheTestResolver(namespace) selected: Final = resolver.resolve() assert selected.kind == "native" selected.store(request(), {"answer": 1}) assert await selected.async_lookup(request()) == {"answer": 1} - with rebound(namespace, "cache", _native.NativeCacheHandle.memory()): + with rebound(namespace, "cache", _native._CacheTestHandle.memory()): replacement: Final = resolver.resolve() await selected.async_store(request(), {"answer": 2}) assert replacement.lookup(request()) is None @@ -96,7 +119,7 @@ async def test_python_callback_preserves_identity_caller_task_context_and_errors raise failure namespace: Final = SimpleNamespace(cache=CustomCache()) - binding: Final = _native.CacheResolver(namespace).resolve() + binding: Final = _native._CacheTestResolver(namespace).resolve() assert binding.kind == "python_callback" assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel assert context.get() == "callback" @@ -117,7 +140,7 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: finally: finished.set() - binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve() + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() async def lookup() -> object: return await binding.async_lookup(None, callback_kwargs={}) @@ -132,9 +155,9 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle: Final = _native.NativeCacheHandle.memory() - handle.bind_facade(facade) - resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade)) + handle: Final = _native._CacheTestHandle.memory() + handle._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) native: Final = resolver.resolve() assert native.kind == "native" native.store(request(), {"source": "native"}) @@ -165,16 +188,18 @@ def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not class CustomCache(Cache): pass - handle: Final = _native.NativeCacheHandle.memory() + handle: Final = _native._CacheTestHandle.memory() with pytest.raises(TypeError): - handle.bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) + handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle.bind_facade(facade) - resolver: Final = _native.CacheResolver(SimpleNamespace(cache=facade)) + handle._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) with rebound(facade, "cache", InMemoryCache()): assert resolver.resolve().kind == "python_callback" with rebound(facade, "ttl", 12): assert resolver.resolve().kind == "python_callback" + with rebound(facade, "semantic_cache_scope", "end_user"): + assert resolver.resolve().kind == "python_callback" def custom_key(**_kwargs: object) -> str: return "custom" @@ -193,7 +218,7 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: def cyclic_reference() -> weakref.ReferenceType[CustomCache]: callback: Final = CustomCache() namespace: Final = SimpleNamespace(cache=callback) - binding: Final = _native.CacheResolver(namespace).resolve() + binding: Final = _native._CacheTestResolver(namespace).resolve() setattr(callback, "binding", binding) return weakref.ref(callback) @@ -204,8 +229,8 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: client: Final = redis.Redis.from_url(redis_url) - namespace: Final = SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url, namespace="team")) - binding: Final = _native.CacheResolver(namespace).resolve() + namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url, namespace="team")) + binding: Final = _native._CacheTestResolver(namespace).resolve() response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} client.set("team:sync", str(envelope)) @@ -227,33 +252,33 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd def test_invalid_duration_and_request_shape_fail_before_storage() -> None: - binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() for seconds in (-1.0, float("nan"), float("inf")): with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) assert binding.lookup(request()) is None with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): - _native.NativeCacheHandle.memory(ttl_seconds=-1) + _native._CacheTestHandle.memory(ttl_seconds=-1) async def test_memory_size_policy_is_applied_by_the_native_host() -> None: - handle: Final = _native.NativeCacheHandle.memory(capacity=2, max_entry_bytes=128) - binding: Final = _native.CacheResolver(SimpleNamespace(cache=handle)).resolve() + handle: Final = _native._CacheTestHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() small: Final = {"answer": "ok"} binding.store(request("small"), small) assert await binding.async_lookup(request("small")) == small await binding.async_store(request("large"), {"answer": "x" * 256}) assert binding.lookup(request("large")) is None assert binding.lookup(request("small")) == small - disabled: Final = _native.CacheResolver( - SimpleNamespace(cache=_native.NativeCacheHandle.memory(capacity=0)) + disabled: Final = _native._CacheTestResolver( + SimpleNamespace(cache=_native._CacheTestHandle.memory(capacity=0)) ).resolve() await disabled.async_store(request(), small) assert await disabled.async_lookup(request()) is None async def test_native_batch_lookup_and_store_report_partial_hits() -> None: - binding: Final = _native.CacheResolver(SimpleNamespace(cache=_native.NativeCacheHandle.memory())).resolve() + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() requests: Final = [request("hit"), request("miss"), request("disabled")] requests[2]["controls"] = { "supported_call_type": True, @@ -275,50 +300,72 @@ async def test_native_batch_lookup_and_store_report_partial_hits() -> None: } -async def test_python_batch_callbacks_receive_keys_and_key_value_pairs() -> None: - first: Final = object() - second: Final = object() - - class CustomCache: - def batch_get_cache(self, keys: list[str], *, marker: object) -> tuple[list[str], object]: - return keys, marker - - async def async_batch_get_cache(self, keys: list[str], *, marker: object) -> tuple[list[str], object]: - return keys, marker - - async def async_set_cache_pipeline( - self, cache_list: list[tuple[str, object]], *, marker: object - ) -> tuple[list[tuple[str, object]], object]: - return cache_list, marker - +async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: + result: Final = object() marker: Final = object() - binding: Final = _native.CacheResolver(SimpleNamespace(cache=CustomCache())).resolve() - requests: Final = [request("first"), request("second")] - assert binding.lookup_batch(requests, callback_kwargs={"marker": marker}) == (["first", "second"], marker) - assert await binding.async_lookup_batch(requests, callback_kwargs={"marker": marker}) == ( - ["first", "second"], - marker, - ) + class CustomCache(Cache): + def get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: + return ("sync", kwargs) + + async def async_get_cache(self, dynamic_cache_object: object = None, **kwargs: object) -> object: + return ("async", kwargs) + + async def async_add_cache_pipeline( + self, result: object, dynamic_cache_object: object = None, **kwargs: object + ) -> object: + return result, kwargs + + binding: Final = _native._CacheTestResolver( + SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL)) + ).resolve() + assert binding.kind == "python_callback" + requests: Final = [request("first"), request("second")] + kwargs: Final = [{"cache_key": "first"}, {"cache_key": "second"}] + + assert binding.lookup_batch(requests, callback_kwargs=kwargs) == [("sync", kwargs[0]), ("sync", kwargs[1])] + assert await binding.async_lookup_batch(requests, callback_kwargs=kwargs) == [ + ("async", kwargs[0]), + ("async", kwargs[1]), + ] + with pytest.raises(ValueError, match="equal lengths"): + binding.lookup_batch(requests, callback_kwargs=kwargs[:1]) + with pytest.raises(TypeError, match="callback_result"): + await binding.async_store_batch(requests, [1, 2], callback_kwargs={"marker": marker}) stored: Final = cast( - tuple[list[tuple[str, object]], object], - await binding.async_store_batch( - requests, - [first, second], - callback_kwargs={"marker": marker}, - ), + tuple[object, dict[str, object]], + await binding.async_store_batch(requests, [1, 2], callback_result=result, callback_kwargs={"marker": marker}), ) - assert [key for key, _ in stored[0]] == ["first", "second"] - assert stored[1] is marker - assert stored[0][0][1] is first - assert stored[0][1][1] is second + assert stored[0] is result + assert stored[1] == {"marker": marker} + + +async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: + async def ping() -> str: + return "pong" + + cache: Final = Cache(type=LiteLLMCacheType.LOCAL) + cache.cache.set_cache("key", "value") + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=cache)).resolve() + assert binding.kind == "python_callback" + + setattr(cache.cache, "ping", ping) + assert await binding.ping() == "pong" + await binding.async_flush() + assert cache.cache.get_cache("key") is None + + +def test_facade_registration_rejects_mismatched_capacity() -> None: + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + with pytest.raises(TypeError, match="capacities must match"): + _native._CacheTestHandle.memory(capacity=7)._bind_facade(facade) async def test_redis_handle_reads_the_python_default_ttl(redis_url: str) -> None: client: Final = redis.Redis.from_url(redis_url) with rebound(litellm, "default_redis_ttl", 7): - binding: Final = _native.CacheResolver( - SimpleNamespace(cache=_native.NativeCacheHandle.redis(redis_url)) + binding: Final = _native._CacheTestResolver( + SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url)) ).resolve() await binding.async_store(request("native-default"), {"value": 1}) @@ -336,11 +383,16 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: redis_flush_size=2, ) with pytest.raises(TypeError, match="default TTLs must match"): - _native.NativeCacheHandle.redis(redis_url, ttl_seconds=61).bind_facade(facade) - _native.NativeCacheHandle.redis(redis_url).bind_facade(facade) - binding: Final = _native.CacheResolver(SimpleNamespace(cache=facade)).resolve() + _native._CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) + with pytest.raises(TypeError, match="namespaces must match"): + _native._CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) + _native._CacheTestHandle.redis(redis_url)._bind_facade(facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() client: Final = redis.Redis.from_url(redis_url) + with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): + assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + await binding.async_store(request("first"), {"value": 1}) assert client.get("first") is None await binding.async_store(request("second"), {"value": 2}) From dc5f0c58a4decfdd6227fbf3af661bf6458ab04e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 10:21:54 -0700 Subject: [PATCH 08/17] feat(cache): add native backend operation primitives --- litellm-rust/Cargo.lock | 2 + litellm-rust/crates/cache-memory/src/cache.rs | 42 +- .../crates/cache-memory/tests/cache.rs | 45 +- litellm-rust/crates/cache-redis/Cargo.toml | 2 +- litellm-rust/crates/cache-redis/src/cache.rs | 4 + .../cache-redis/src/cache/operations.rs | 500 ++++++++++++++++++ litellm-rust/crates/cache-redis/src/lib.rs | 2 +- .../crates/cache-redis/tests/cache.rs | 307 ++++++++++- litellm-rust/crates/cache-response/README.md | 6 +- litellm-rust/crates/cache/src/capabilities.rs | 7 + litellm-rust/crates/cache/src/lib.rs | 2 +- 11 files changed, 910 insertions(+), 9 deletions(-) create mode 100644 litellm-rust/crates/cache-redis/src/cache/operations.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 2d6fb6c3082..ed4ae4e3353 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3712,6 +3712,8 @@ dependencies = [ "itoa", "num-bigint 0.5.1", "percent-encoding", + "rustls 0.23.42", + "rustls-native-certs", "ryu", "sha1_smol", "socket2 0.6.5", diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index a9814ff6fd6..77635893640 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -7,7 +7,7 @@ use std::{ use litellm_cache::{ BaseCache, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, - Error, + Error, IncrementOperation, }; const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200; @@ -134,6 +134,25 @@ impl InMemoryCache { .copied()) } + pub async fn async_get_ttl(&self, key: &str) -> Result, Error> { + self.expires_at(key) + } + + pub async fn async_get_oldest_n_keys(&self, count: usize) -> Result, Error> { + let state = self.state.lock().map_err(|_| Error::Unavailable)?; + let mut expirations = state + .expirations + .iter() + .map(|(key, expiration)| (key.clone(), *expiration)) + .collect::>(); + expirations.sort_unstable_by_key(|(_, expiration)| *expiration); + Ok(expirations + .into_iter() + .take(count) + .map(|(key, _)| key) + .collect()) + } + pub fn delete_cache(&self, key: &str) -> Result<(), Error> { let mut state = self.state.lock().map_err(|_| Error::Unavailable)?; Self::remove(&mut state, key); @@ -240,6 +259,27 @@ impl CounterCache for InMemoryCache { } } +impl InMemoryCache { + pub async fn async_increment_pipeline( + &self, + operations: Vec, + ) -> Result, Error> { + operations + .into_iter() + .map(|operation| { + self.increment_cache( + &operation.key, + operation.amount, + CacheKwargs { + ttl: operation.ttl, + ..CacheKwargs::default() + }, + ) + }) + .collect() + } +} + impl BaseCache for InMemoryCache { type Value = V; diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index 22c5595da52..c44590b63ca 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -8,7 +8,7 @@ use std::{ use litellm_cache::{ BaseCache, CacheBackend, CacheConnectionStatus, CacheKwargs, ClaimCache, CounterCache, Error, - get_cache, set_cache, + IncrementOperation, get_cache, set_cache, }; use litellm_cache_memory::{CacheWrite, InMemoryCache}; use rstest::{fixture, rstest}; @@ -305,3 +305,46 @@ fn disabled_cache_does_not_retain_claims_or_counters() { ); assert_eq!(counters.get_cache("key").unwrap(), None); } + +#[tokio::test] +async fn ttl_and_oldest_key_operations_use_the_stored_expirations() { + let clock = Arc::new(AtomicU64::new(100)); + let cache = cache(clock, 3); + cache + .set_cache("later", "2".into(), Some(Duration::from_secs(20))) + .unwrap(); + cache + .set_cache("first", "1".into(), Some(Duration::from_secs(10))) + .unwrap(); + + assert_eq!( + cache.async_get_ttl("first").await.unwrap(), + Some(Duration::from_secs(110)) + ); + assert_eq!(cache.async_get_oldest_n_keys(1).await.unwrap(), ["first"]); + assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None); +} + +#[tokio::test] +async fn increment_pipeline_preserves_operation_order() { + let cache = InMemoryCache::::new(Some(3), None); + assert_eq!( + cache + .async_increment_pipeline(vec![ + IncrementOperation { + key: "a".into(), + amount: 1.0, + ttl: Some(Duration::from_secs(10)), + }, + IncrementOperation { + key: "a".into(), + amount: 2.0, + ttl: Some(Duration::from_secs(20)), + }, + ]) + .await + .unwrap(), + [1.0, 3.0] + ); + assert_eq!(cache.get_cache("a").unwrap(), Some(3.0)); +} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index f123e774158..5818f75ff3d 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -7,7 +7,7 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true -redis = "1.7.0" +redis = { version = "1.7.0", features = ["tls-rustls"] } r2d2 = "0.8.10" tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index b8f0d3857e2..23b1fabbab4 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -9,6 +9,10 @@ use litellm_cache::{ }; use redis::Commands; +mod operations; + +pub use operations::{RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation}; + const DEFAULT_TTL: Duration = Duration::from_secs(600); const REDIS_TIMEOUT: Duration = Duration::from_secs(5); const REDIS_POOL_SIZE: u32 = 16; diff --git a/litellm-rust/crates/cache-redis/src/cache/operations.rs b/litellm-rust/crates/cache-redis/src/cache/operations.rs new file mode 100644 index 00000000000..f8c2bf4078c --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/cache/operations.rs @@ -0,0 +1,500 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::{CacheCodec, Error, IncrementOperation}; +use redis::Commands; + +use super::{ConnectionRef, RedisCache}; + +const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!( + "local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ", + "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ", + "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", + "return count" +); +const SET_MAX_SCRIPT: &str = concat!( + "local current = redis.call('GET', KEYS[1]); ", + "if current == false or tonumber(current) < tonumber(ARGV[1]) then ", + "redis.call('SET', KEYS[1], ARGV[1]); ", + "if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", + "return ARGV[1]; end; return current" +); + +#[derive(Clone, Debug, PartialEq)] +pub enum RedisArg { + Bytes(Vec), + Integer(i64), + Float(f64), +} + +impl From<&str> for RedisArg { + fn from(value: &str) -> Self { + Self::Bytes(value.as_bytes().to_vec()) + } +} + +impl From for RedisArg { + fn from(value: String) -> Self { + Self::Bytes(value.into_bytes()) + } +} + +impl From> for RedisArg { + fn from(value: Vec) -> Self { + Self::Bytes(value) + } +} + +impl From for RedisArg { + fn from(value: i64) -> Self { + Self::Integer(value) + } +} + +impl From for RedisArg { + fn from(value: f64) -> Self { + Self::Float(value) + } +} + +impl redis::ToRedisArgs for RedisArg { + fn write_redis_args(&self, out: &mut W) + where + W: ?Sized + redis::RedisWrite, + { + match self { + Self::Bytes(value) => value.write_redis_args(out), + Self::Integer(value) => value.write_redis_args(out), + Self::Float(value) => value.write_redis_args(out), + } + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct RedisRpushOperation { + pub key: String, + pub values: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct RedisLpopOperation { + pub key: String, + pub count: Option, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum RedisLpopResult { + Missing, + Value(Vec), + Values(Vec>), +} + +impl RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + pub async fn delete_cache_keys(&self, keys: Vec) -> Result { + if keys.is_empty() { + return Ok(0); + } + let keys = keys + .into_iter() + .map(|key| self.namespaced_key(&key)) + .collect::>(); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + connection.del(keys).map_err(|_| Error::Unavailable) + }) + .await + } + + pub fn batch_get_counts(&self, keys: &[String]) -> 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(count).collect() + } + + pub async fn async_batch_get_counts( + &self, + keys: Vec, + ) -> 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(count).collect() + } + + pub fn sync_ping(&self) -> Result { + self.connections.execute(|connection| { + redis::cmd("PING") + .query::(connection) + .map(|response| response == "PONG") + .map_err(|_| Error::Unavailable) + }) + } + + pub async fn ping(&self) -> Result { + Self::run_blocking(Arc::clone(&self.connections), |connection| { + redis::cmd("PING") + .query::(connection) + .map(|response| response == "PONG") + .map_err(|_| Error::Unavailable) + }) + .await + } + + pub async fn async_get_ttl(&self, key: &str) -> Result, Error> { + let key = self.namespaced_key(key); + let ttl = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("TTL") + .arg(key) + .query::(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + Ok((ttl >= 0).then_some(ttl)) + } + + pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, Error> { + let pattern = format!("{}*", self.namespaced_key(pattern)); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut cursor = 0u64; + let mut matches = Vec::new(); + loop { + let (next_cursor, keys): (u64, Vec) = redis::cmd("SCAN") + .cursor_arg(cursor) + .arg("MATCH") + .arg(&pattern) + .arg("COUNT") + .arg(count) + .query(connection) + .map_err(|_| Error::Unavailable)?; + matches.extend(keys); + if matches.len() >= count || next_cursor == 0 { + matches.truncate(count); + return Ok(matches); + } + cursor = next_cursor; + } + }) + .await + } + + pub async fn async_set_cache_sadd( + &self, + key: &str, + values: Vec, + ttl: Option, + ) -> Result { + if values.is_empty() { + return Err(Error::InvalidEntry); + } + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl)); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut pipeline = redis::pipe(); + pipeline.cmd("SADD").arg(&key).arg(values); + pipeline.cmd("EXPIRE").arg(&key).arg(ttl).ignore(); + pipeline + .query::<(usize,)>(connection) + .map(|(added,)| added) + .map_err(|_| Error::Unavailable) + }) + .await + } + + pub async fn async_rpush(&self, key: &str, values: Vec) -> Result { + if values.is_empty() { + return Err(Error::InvalidEntry); + } + let key = self.namespaced_key(key); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("RPUSH") + .arg(key) + .arg(values) + .query(connection) + .map_err(|_| Error::Unavailable) + }) + .await + } + + pub async fn async_rpush_pipeline( + &self, + operations: Vec, + ) -> Result, Error> { + let operations = operations + .into_iter() + .map(|operation| { + if operation.values.is_empty() { + return Err(Error::InvalidEntry); + } + Ok((self.namespaced_key(&operation.key), operation.values)) + }) + .collect::, _>>()?; + if operations.is_empty() { + return Ok(Vec::new()); + } + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut pipeline = redis::pipe(); + for (key, values) in operations { + pipeline.cmd("RPUSH").arg(key).arg(values); + } + pipeline.query(connection).map_err(|_| Error::Unavailable) + }) + .await + } + + pub async fn async_lpop( + &self, + key: &str, + count: Option, + ) -> Result { + let key = self.namespaced_key(key); + let multiple = count.is_some(); + let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut command = redis::cmd("LPOP"); + command.arg(key); + if let Some(count) = count { + command.arg(count); + } + command + .query::(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + lpop_result(value, multiple) + } + + pub async fn async_lpop_pipeline( + &self, + operations: Vec, + ) -> Result, Error> { + let operations = operations + .into_iter() + .map(|operation| (self.namespaced_key(&operation.key), operation.count)) + .collect::>(); + if operations.is_empty() { + return Ok(Vec::new()); + } + let multiple = operations + .iter() + .map(|(_, count)| count.is_some()) + .collect::>(); + let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut pipeline = redis::pipe(); + for (key, count) in operations { + let command = pipeline.cmd("LPOP").arg(key); + if let Some(count) = count { + command.arg(count); + } + } + pipeline + .query::>(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + values + .into_iter() + .zip(multiple) + .map(|(value, multiple)| lpop_result(value, multiple)) + .collect() + } + + pub async fn async_eval( + &self, + script: String, + keys: Vec, + arguments: Vec, + ) -> Result { + let keys = keys + .into_iter() + .map(|key| self.namespaced_key(&key)) + .collect::>(); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("EVAL") + .arg(script) + .arg(keys.len()) + .arg(keys) + .arg(arguments) + .query(connection) + .map_err(|_| Error::Unavailable) + }) + .await + } + + pub fn client_list(&self) -> Result { + self.connections.execute(|connection| { + redis::cmd("CLIENT") + .arg("LIST") + .query(connection) + .map_err(|_| Error::Unavailable) + }) + } + + pub fn info(&self) -> Result { + self.connections.execute(|connection| { + redis::cmd("INFO") + .query(connection) + .map_err(|_| Error::Unavailable) + }) + } + + pub fn flushall(&self) -> Result<(), Error> { + self.connections.execute(|connection| { + redis::cmd("FLUSHALL") + .query(connection) + .map_err(|_| Error::Unavailable) + }) + } +} + +impl RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + pub fn increment_with_floor( + &self, + key: &str, + amount: i64, + ttl: Duration, + ) -> Result { + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(ttl); + self.connections + .execute(|connection| increment_with_floor(connection, key, amount, ttl)) + } + + pub async fn async_increment_pipeline( + &self, + operations: Vec, + ) -> Result, Error> { + let operations = operations + .into_iter() + .map(|operation| { + ( + self.namespaced_key(&operation.key), + operation.amount, + operation.ttl.map(Self::ttl_seconds), + ) + }) + .collect::>(); + if operations.is_empty() { + return Ok(Vec::new()); + } + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let mut pipeline = redis::pipe(); + for (key, amount, ttl) in operations { + pipeline.cmd("INCRBYFLOAT").arg(&key).arg(amount); + if let Some(ttl) = ttl { + pipeline.cmd("EXPIRE").arg(key).arg(ttl).ignore(); + } + } + pipeline.query(connection).map_err(|_| Error::Unavailable) + }) + .await + } + + pub async fn async_increment_with_floor( + &self, + key: &str, + amount: i64, + ttl: Duration, + ) -> Result { + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(ttl); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + increment_with_floor(connection, key, amount, ttl) + }) + .await + } + + pub async fn async_set_max( + &self, + key: &str, + value: f64, + ttl: Option, + ) -> Result { + let key = self.namespaced_key(key); + let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl)); + Self::run_blocking(Arc::clone(&self.connections), move |connection| { + redis::cmd("EVAL") + .arg(SET_MAX_SCRIPT) + .arg(1) + .arg(key) + .arg(value) + .arg(ttl) + .query(connection) + .map_err(|_| Error::Unavailable) + }) + .await + } +} + +fn redis_bytes(value: redis::Value) -> Result, Error> { + match value { + redis::Value::BulkString(bytes) => Ok(bytes), + redis::Value::SimpleString(text) => Ok(text.into_bytes()), + _ => Err(Error::InvalidEntry), + } +} + +fn lpop_result(value: redis::Value, multiple: bool) -> Result { + match value { + redis::Value::Nil => Ok(RedisLpopResult::Missing), + redis::Value::Array(values) if multiple => values + .into_iter() + .map(redis_bytes) + .collect::, _>>() + .map(RedisLpopResult::Values), + value if !multiple => redis_bytes(value).map(RedisLpopResult::Value), + _ => Err(Error::InvalidEntry), + } +} + +fn count(value: redis::Value) -> Result, Error> { + match value { + redis::Value::Nil => Ok(None), + redis::Value::Int(value) => Ok(Some(value)), + redis::Value::BulkString(value) => std::str::from_utf8(&value) + .ok() + .and_then(|value| value.parse().ok()) + .map(Some) + .ok_or(Error::InvalidEntry), + redis::Value::SimpleString(value) => { + value.parse().map(Some).map_err(|_| Error::InvalidEntry) + } + _ => Err(Error::InvalidEntry), + } +} + +fn increment_with_floor( + connection: &mut ConnectionRef<'_>, + key: String, + amount: i64, + ttl: u64, +) -> Result { + redis::cmd("EVAL") + .arg(INCREMENT_WITH_FLOOR_SCRIPT) + .arg(1) + .arg(key) + .arg(amount) + .arg(ttl) + .query(connection) + .map_err(|_| Error::Unavailable) +} diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs index 37b35c5ea4a..2548c7ac3c6 100644 --- a/litellm-rust/crates/cache-redis/src/lib.rs +++ b/litellm-rust/crates/cache-redis/src/lib.rs @@ -1,3 +1,3 @@ mod cache; -pub use cache::RedisCache; +pub use cache::{RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation}; diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index 6a61b80b84c..3d755d0f44c 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -2,9 +2,11 @@ use std::time::Duration; use litellm_cache::{ BaseCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheKwargs, ClaimCache, - CounterCache, Error, JsonCodec, get_cache, set_cache, + CounterCache, Error, IncrementOperation, JsonCodec, get_cache, set_cache, +}; +use litellm_cache_redis::{ + RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, }; -use litellm_cache_redis::RedisCache; use redis_test::{MockCmd, MockRedisConnection}; struct TaggedByteCodec(u8); @@ -260,6 +262,307 @@ async fn async_flush_deletes_each_scan_page_separately() { cache.async_flush_cache().await.unwrap(); } +#[tokio::test] +async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() { + let mut sadd_pipeline = redis::pipe(); + sadd_pipeline + .cmd("SADD") + .arg("team:members") + .arg("a") + .arg("b") + .cmd("EXPIRE") + .arg("team:members") + .arg(600u64) + .ignore(); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("MGET").arg("team:count").arg("team:missing"), + Ok(redis_test::redis_value!(["7", nil])), + ), + MockCmd::new( + redis::cmd("MGET").arg("team:count").arg("team:missing"), + Ok(redis_test::redis_value!(["7", nil])), + ), + MockCmd::new(redis::cmd("PING"), Ok("PONG")), + MockCmd::new(redis::cmd("PING"), Ok("PONG")), + MockCmd::new(redis::cmd("TTL").arg("team:missing"), Ok(-2i64)), + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(0) + .arg("MATCH") + .arg("team:job-*") + .arg("COUNT") + .arg(25), + Ok(redis_test::redis_value!(["4", ["team:job-a"]])), + ), + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(4) + .arg("MATCH") + .arg("team:job-*") + .arg("COUNT") + .arg(25), + Ok(redis_test::redis_value!(["0", ["team:job-b"]])), + ), + MockCmd::new( + redis::cmd("DEL").arg("team:job-a").arg("team:job-b"), + Ok(2u32), + ), + MockCmd::with_values( + sadd_pipeline, + Ok(vec![redis::Value::Int(2), redis::Value::Int(1)]), + ), + MockCmd::new( + redis::cmd("RPUSH").arg("team:queue").arg("a").arg("b"), + Ok(2u32), + ), + MockCmd::new( + redis::cmd("LPOP").arg("team:queue").arg(2usize), + Ok(redis_test::redis_value!(["a", "b"])), + ), + 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")), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + + assert_eq!( + cache + .batch_get_counts(&["count".into(), "missing".into()]) + .unwrap(), + [Some(7), None] + ); + assert_eq!( + cache + .async_batch_get_counts(vec!["count".into(), "missing".into()]) + .await + .unwrap(), + [Some(7), None] + ); + assert!(cache.sync_ping().unwrap()); + assert!(cache.ping().await.unwrap()); + assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None); + assert_eq!( + cache.async_scan_iter("job-", 25).await.unwrap(), + ["team:job-a", "team:job-b"] + ); + assert_eq!( + cache + .delete_cache_keys(vec!["job-a".into(), "job-b".into()]) + .await + .unwrap(), + 2 + ); + assert_eq!( + cache + .async_set_cache_sadd("members", vec!["a".into(), "b".into()], None) + .await + .unwrap(), + 2 + ); + assert_eq!( + cache + .async_rpush("queue", vec!["a".into(), "b".into()]) + .await + .unwrap(), + 2 + ); + assert_eq!( + cache.async_lpop("queue", Some(2)).await.unwrap(), + RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()]) + ); + assert_eq!( + cache + .async_eval("return KEYS[1]".into(), 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(); +} + +#[tokio::test] +async fn direct_redis_pipelines_preserve_operation_order() { + let mut rpush_pipeline = redis::pipe(); + rpush_pipeline + .cmd("RPUSH") + .arg("team:a") + .arg("one") + .cmd("RPUSH") + .arg("team:b") + .arg("two"); + let mut lpop_pipeline = redis::pipe(); + lpop_pipeline + .cmd("LPOP") + .arg("team:a") + .arg(2usize) + .cmd("LPOP") + .arg("team:b"); + let connection = MockRedisConnection::new([ + MockCmd::with_values( + rpush_pipeline, + Ok(vec![redis::Value::Int(1), redis::Value::Int(2)]), + ), + MockCmd::with_values( + lpop_pipeline, + Ok(vec![redis_test::redis_value!(["one"]), redis::Value::Nil]), + ), + ]) + .assert_all_commands_consumed(); + let queue = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + + assert_eq!( + queue + .async_rpush_pipeline(vec![ + RedisRpushOperation { + key: "a".into(), + values: vec![RedisArg::from("one")], + }, + RedisRpushOperation { + key: "b".into(), + values: vec![RedisArg::from("two")], + }, + ]) + .await + .unwrap(), + [1, 2] + ); + assert_eq!( + queue + .async_lpop_pipeline(vec![ + RedisLpopOperation { + key: "a".into(), + count: Some(2), + }, + RedisLpopOperation { + key: "b".into(), + count: None, + }, + ]) + .await + .unwrap(), + [ + RedisLpopResult::Values(vec![b"one".to_vec()]), + RedisLpopResult::Missing, + ] + ); + + let mut increment_pipeline = redis::pipe(); + increment_pipeline + .cmd("INCRBYFLOAT") + .arg("team:counter") + .arg(1.5f64) + .cmd("EXPIRE") + .arg("team:counter") + .arg(10u64) + .ignore() + .cmd("INCRBYFLOAT") + .arg("team:counter") + .arg(2.0f64); + let connection = MockRedisConnection::new([MockCmd::with_values( + increment_pipeline, + Ok(vec![ + redis::Value::BulkString(b"1.5".to_vec()), + redis::Value::Int(1), + redis::Value::BulkString(b"3.5".to_vec()), + ]), + )]) + .assert_all_commands_consumed(); + let counters = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + assert_eq!( + counters + .async_increment_pipeline(vec![ + IncrementOperation { + key: "counter".into(), + amount: 1.5, + ttl: Some(Duration::from_secs(10)), + }, + IncrementOperation { + key: "counter".into(), + amount: 2.0, + ttl: None, + }, + ]) + .await + .unwrap(), + [1.5, 3.5] + ); +} + +const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!( + "local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ", + "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ", + "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", + "return count" +); +const SET_MAX_SCRIPT: &str = concat!( + "local current = redis.call('GET', KEYS[1]); ", + "if current == false or tonumber(current) < tonumber(ARGV[1]) then ", + "redis.call('SET', KEYS[1], ARGV[1]); ", + "if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", + "return ARGV[1]; end; return current" +); + +#[tokio::test] +async fn counter_repairs_are_atomic_and_use_default_ttl() { + let floor = || { + redis::cmd("EVAL") + .arg(INCREMENT_WITH_FLOOR_SCRIPT) + .arg(1) + .arg("team:counter") + .arg(-2i64) + .arg(30u64) + .clone() + }; + let connection = MockRedisConnection::new([ + MockCmd::new(floor(), Ok(0i64)), + MockCmd::new(floor(), Ok(0i64)), + MockCmd::new( + redis::cmd("EVAL") + .arg(SET_MAX_SCRIPT) + .arg(1) + .arg("team:counter") + .arg(4.5f64) + .arg(600u64), + Ok("4.5"), + ), + ]) + .assert_all_commands_consumed(); + let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) + .with_namespace(Some("team".into())); + + assert_eq!( + cache + .increment_with_floor("counter", -2, Duration::from_secs(30)) + .unwrap(), + 0 + ); + assert_eq!( + cache + .async_increment_with_floor("counter", -2, Duration::from_secs(30)) + .await + .unwrap(), + 0 + ); + assert_eq!( + cache.async_set_max("counter", 4.5, None).await.unwrap(), + 4.5 + ); +} + const CLAIM_SCRIPT: &str = concat!( "local current = redis.call('GET', KEYS[1]); ", "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 013cf4b7baf..f9e7a148706 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -42,10 +42,12 @@ The resolver reads the namespace's `cache` attribute each time it resolves. A ca Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend -The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and configuration changes before selecting native execution. It does not compare Redis connection settings. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python +The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy +The Redis backend also provides the primitives needed to preserve its direct Python surface later: TLS URLs, ping, bulk delete, counter batches, TTL, scan, set membership, raw queue push and pop, queue and counter pipelines, counter floor and maximum operations, script evaluation, client information, namespaced flush, and full flush. These are backend operations only and are not exported to Python by this PR. Memory provides TTL, oldest-key, and counter-pipeline operations + ## Adding another backend Implement `BaseCache` for the backend with its associated value type, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache` then works without another response implementation. Add a concrete bridge enum variant and constructor only when exposing that backend to Python @@ -56,4 +58,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths -Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations, queues, and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees +Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees diff --git a/litellm-rust/crates/cache/src/capabilities.rs b/litellm-rust/crates/cache/src/capabilities.rs index 0b9deab1f5a..1613a5786f4 100644 --- a/litellm-rust/crates/cache/src/capabilities.rs +++ b/litellm-rust/crates/cache/src/capabilities.rs @@ -2,6 +2,13 @@ use std::future::Future; use crate::{BaseCache, CacheKwargs, Error}; +#[derive(Clone, Debug, PartialEq)] +pub struct IncrementOperation { + pub key: String, + pub amount: f64, + pub ttl: Option, +} + pub trait CounterCache: BaseCache { fn increment_cache(&self, key: &str, amount: f64, kwargs: CacheKwargs) -> Result; diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index ed67eb2fe15..cdd5589aa39 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -9,6 +9,6 @@ pub use base_cache::{ BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheKwargs, }; pub use caching::{Cache, CacheBackend, get_cache, set_cache}; -pub use capabilities::{ClaimCache, CounterCache}; +pub use capabilities::{ClaimCache, CounterCache, IncrementOperation}; pub use codec::{CacheCodec, JsonCodec}; pub use error::Error; From 595711829aa2b0a00b2573070d2fc97a850bbc5a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 10:30:42 -0700 Subject: [PATCH 09/17] wip --- litellm-rust/crates/cache-response/README.md | 4 +- .../crates/cache-response/src/buffer.rs | 46 ++ litellm-rust/crates/cache-response/src/lib.rs | 2 + .../crates/cache-response/tests/response.rs | 70 ++- .../crates/python-bridge/src/cache/binding.rs | 294 +++++++++ .../python-bridge/src/cache/callback.rs | 169 ++++++ .../crates/python-bridge/src/cache/facade.rs | 2 +- .../crates/python-bridge/src/cache/future.rs | 18 + .../crates/python-bridge/src/cache/handle.rs | 106 ++++ .../crates/python-bridge/src/cache/mod.rs | 570 +----------------- .../crates/python-bridge/src/cache/native.rs | 35 +- .../crates/python-bridge/src/cache/request.rs | 48 ++ .../python-bridge/src/cache/resolver.rs | 39 ++ 13 files changed, 810 insertions(+), 593 deletions(-) create mode 100644 litellm-rust/crates/cache-response/src/buffer.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/binding.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/callback.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/future.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/handle.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/request.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/resolver.rs diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index f9e7a148706..9863c46783f 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -6,9 +6,9 @@ `litellm-cache` defines typed storage and codec traits. Memory and Redis implement those traits without depending on response policy. Other consumers can store their own value types using the same backend implementations -`litellm-cache-response` owns response keys, controls, entries, and the Python-compatible response codec. It has no runtime dependency on a specific cache backend or Python +`litellm-cache-response` owns response keys, controls, entries, the Python-compatible response codec, and `WriteBuffer`, the backend-neutral deferred-write policy. It has no runtime dependency on a specific cache backend or Python -The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host +The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum, which only dispatches. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host ## Native Rust use diff --git a/litellm-rust/crates/cache-response/src/buffer.rs b/litellm-rust/crates/cache-response/src/buffer.rs new file mode 100644 index 00000000000..68af5278c8c --- /dev/null +++ b/litellm-rust/crates/cache-response/src/buffer.rs @@ -0,0 +1,46 @@ +use std::{sync::Mutex, time::Duration}; + +use litellm_cache::{BaseCache, Error}; +use serde_json::Value; + +use crate::{CacheEntry, ResponseCache, ResponseCacheRequest}; + +/// Defers async writes until `flush_size` entries are pending, then stores them as one batch. +pub struct WriteBuffer { + flush_size: usize, + entries: Mutex>, +} + +impl WriteBuffer { + pub fn new(flush_size: usize) -> Self { + Self { + flush_size: flush_size.max(1), + entries: Mutex::new(Vec::new()), + } + } + + pub async fn async_store>( + &self, + cache: &ResponseCache, + request: &ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Result<(), Error> { + let pending = { + let mut entries = self.entries.lock().map_err(|_| Error::Unavailable)?; + entries.push((request.clone(), response, now)); + (entries.len() >= self.flush_size).then(|| std::mem::take(&mut *entries)) + }; + // A failed flush drops its batch, as Python does. Requeueing would grow the + // buffer and re-send an ever larger pipeline on every write during an outage. + match pending { + Some(pending) => cache.async_store_entries(pending).await, + None => Ok(()), + } + } + + pub fn clear(&self) -> Result<(), Error> { + self.entries.lock().map_err(|_| Error::Unavailable)?.clear(); + Ok(()) + } +} diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index 72a507f8ee9..91b36ebe24b 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -1,8 +1,10 @@ +mod buffer; mod caching; mod codec; mod embedding; mod response; +pub use buffer::WriteBuffer; pub use caching::{ CacheControls, CacheEntry, CacheKeyContext, CacheKeyField, CacheKeyInput, CacheMode, cache_key, get_cache_key, should_use_cache, diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 7c69e5a1d55..94b25626f86 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -11,7 +11,7 @@ use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec, - ResponseCacheRequest, + ResponseCacheRequest, WriteBuffer, }; use redis_test::{MockCmd, MockRedisConnection}; use serde_json::json; @@ -414,3 +414,71 @@ async fn deferred_entries_keep_the_time_they_were_produced() { None ); } + +#[tokio::test] +async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() { + let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); + let buffer = WriteBuffer::new(2); + let mut first = request(); + first.max_age = Some(Duration::from_secs(10)); + let mut second = request(); + second.key.preset = Some("tenant:other".into()); + + buffer + .async_store( + &cache, + &first, + json!({"answer": 7}), + Duration::from_secs(100), + ) + .await + .unwrap(); + assert_eq!( + cache.lookup(&first, Duration::from_secs(100)).unwrap(), + None + ); + + buffer + .async_store( + &cache, + &second, + json!({"answer": 8}), + Duration::from_secs(200), + ) + .await + .unwrap(); + assert_eq!( + cache.lookup(&first, Duration::from_secs(110)).unwrap(), + Some(json!({"answer": 7})) + ); + assert_eq!( + cache.lookup(&first, Duration::from_secs(111)).unwrap(), + None + ); + assert_eq!( + cache.lookup(&second, Duration::from_secs(200)).unwrap(), + Some(json!({"answer": 8})) + ); +} + +#[tokio::test] +async fn write_buffer_clear_drops_pending_entries() { + let cache = ResponseCache::new(Arc::new(InMemoryCache::default())); + let buffer = WriteBuffer::new(2); + let mut other = request(); + other.key.preset = Some("tenant:other".into()); + let now = Duration::from_secs(100); + + buffer + .async_store(&cache, &request(), json!({"answer": 7}), now) + .await + .unwrap(); + buffer.clear().unwrap(); + buffer + .async_store(&cache, &other, json!({"answer": 8}), now) + .await + .unwrap(); + + assert_eq!(cache.lookup(&request(), now).unwrap(), None); + assert_eq!(cache.lookup(&other, now).unwrap(), None); +} diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs new file mode 100644 index 00000000000..44c133d4611 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -0,0 +1,294 @@ +use litellm_cache_response::PartialHits; +use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyRuntimeError, PyValueError}, + prelude::*, + types::PyDict, +}; +use serde_json::Value; + +use super::{ + cache_error, + callback::PythonCallback, + future::{ready_none, ready_value}, + native::NativeResponseCache, + request::{now, request, requests}, +}; + +pub(super) enum CacheBinding { + Disabled, + Native(NativeResponseCache), + PythonCallback(PythonCallback), +} + +#[pyclass(frozen, name = "_CacheTestBinding")] +pub(crate) struct ResolvedCache { + binding: CacheBinding, + pid: u32, +} + +impl ResolvedCache { + pub(super) fn new(binding: CacheBinding) -> Self { + Self { + binding, + pid: std::process::id(), + } + } + + fn check_process(&self) -> PyResult<()> { + if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() { + return Err(PyRuntimeError::new_err( + "native cache bindings must be resolved again after fork", + )); + } + Ok(()) + } + + pub(crate) fn lookup_step( + &self, + py: Python<'_>, + input: &Bound<'_, PyAny>, + kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult { + self.check_process()?; + let awaitable = match &self.binding { + CacheBinding::Disabled => ready_none(py)?, + CacheBinding::Native(service) => { + let request = request(input)?; + let service = service.clone(); + run_async( + py, + async move { service.async_lookup(&request, now()).await }, + cache_error, + )? + } + CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?, + }; + Ok(ExecutionStep::Await(awaitable.unbind())) + } +} + +#[pymethods] +impl ResolvedCache { + #[getter] + fn kind(&self) -> &'static str { + match self.binding { + CacheBinding::Disabled => "disabled", + CacheBinding::Native(_) => "native", + CacheBinding::PythonCallback(_) => "python_callback", + } + } + + #[pyo3(signature = (request, *, callback_kwargs=None))] + fn lookup( + &self, + py: Python<'_>, + request: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => Ok(py.None()), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let service = service.clone(); + let response = release_gil(py, move || service.lookup(&request, now())) + .map_err(cache_error)?; + to_py(py, &response) + } + CacheBinding::PythonCallback(callback) => { + callback.lookup(py, callback_kwargs).map(Bound::unbind) + } + } + } + + #[pyo3(signature = (request, response, *, callback_kwargs=None))] + fn store( + &self, + py: Python<'_>, + request: &Bound<'_, PyAny>, + response: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult<()> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => Ok(()), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let response: Value = from_py(response)?; + let service = service.clone(); + release_gil(py, move || service.store(&request, response, now())) + .map_err(cache_error) + } + CacheBinding::PythonCallback(callback) => callback.store(py, response, callback_kwargs), + } + } + + /// Native bindings return `{values, missing_indices}`, while a Python callback returns the + /// list of its per-request results. + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn lookup_batch( + &self, + py: Python<'_>, + requests: &Bound<'_, PyAny>, + callback_kwargs: Option<&Bound<'_, PyAny>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + to_py(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + let response = release_gil(py, move || service.lookup_batch(&requests, now())) + .map_err(cache_error)?; + to_py(py, &response) + } + CacheBinding::PythonCallback(callback) => callback + .lookup_batch(py, requests, callback_kwargs) + .map(Bound::unbind), + } + } + + #[pyo3(signature = (request, *, callback_kwargs=None))] + fn async_lookup<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)? + else { + unreachable!() + }; + Ok(awaitable.into_bound(py)) + } + + #[pyo3(signature = (request, response, *, callback_kwargs=None))] + fn async_store<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + response: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let request = self::request(request)?; + let response: Value = from_py(response)?; + let service = service.clone(); + run_async( + py, + async move { service.async_store(&request, response, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => { + callback.async_store(py, response, callback_kwargs) + } + } + } + + #[pyo3(signature = (requests, *, callback_kwargs=None))] + fn async_lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + callback_kwargs: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => { + let requests = self::requests(requests)?; + ready_value(py, &PartialHits::new(vec![None; requests.len()])) + } + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let service = service.clone(); + run_async( + py, + async move { service.async_lookup_batch(&requests, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => { + callback.async_lookup_batch(py, requests, callback_kwargs) + } + } + } + + /// A Python callback receives the caller's original result through `callback_result`. + #[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))] + fn async_store_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + responses: &Bound<'py, PyAny>, + callback_result: Option<&Bound<'py, PyAny>>, + callback_kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let requests = self::requests(requests)?; + let responses: Vec = from_py(responses)?; + if requests.len() != responses.len() { + return Err(PyValueError::new_err( + "batch cache requests and responses must have equal lengths", + )); + } + let entries = requests.into_iter().zip(responses).collect(); + let service = service.clone(); + run_async( + py, + async move { service.async_store_batch(entries, now()).await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => { + callback.async_store_batch(py, callback_result, callback_kwargs) + } + } + } + + fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async(py, async move { service.async_flush().await }, cache_error) + } + CacheBinding::PythonCallback(callback) => callback.async_flush(py), + } + } + + fn ping<'py>(&self, py: Python<'py>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Disabled => ready_none(py), + CacheBinding::Native(service) => { + let service = service.clone(); + run_async( + py, + async move { service.test_connection().await }, + cache_error, + ) + } + CacheBinding::PythonCallback(callback) => callback.ping(py), + } + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + if let CacheBinding::PythonCallback(callback) = &self.binding { + callback.traverse(&visit)?; + } + Ok(()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/callback.rs b/litellm-rust/crates/python-bridge/src/cache/callback.rs new file mode 100644 index 00000000000..318f9d02080 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/callback.rs @@ -0,0 +1,169 @@ +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyTypeError, PyValueError}, + prelude::*, + types::{PyDict, PyList, PyTuple}, +}; + +use super::future::ready_none; + +/// A custom Python cache object, driven through the built-in `Cache` API so a `Cache` subclass +/// works unchanged. +pub(super) struct PythonCallback(Py); + +impl PythonCallback { + pub(super) fn new(object: Py) -> Self { + Self(object) + } + + pub(super) fn lookup<'py>( + &self, + py: Python<'py>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.0 + .bind(py) + .call_method("get_cache", (), Some(callback_kwargs(kwargs)?)) + } + + pub(super) fn async_lookup<'py>( + &self, + py: Python<'py>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.0 + .bind(py) + .call_method("async_get_cache", (), Some(callback_kwargs(kwargs)?)) + } + + pub(super) fn store( + &self, + py: Python<'_>, + response: &Bound<'_, PyAny>, + kwargs: Option<&Bound<'_, PyDict>>, + ) -> PyResult<()> { + self.0 + .bind(py) + .call_method("add_cache", (response,), Some(callback_kwargs(kwargs)?)) + .map(|_| ()) + } + + pub(super) fn async_store<'py>( + &self, + py: Python<'py>, + response: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + self.0.bind(py).call_method( + "async_add_cache", + (response,), + Some(callback_kwargs(kwargs)?), + ) + } + + /// The built-in `Cache` API has no batch read, so the callback receives one + /// `get_cache(**kwargs)` call per request, in order, and the results come back as a list. + pub(super) fn lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let results = PyList::empty(py); + for kwargs in batch_callback_kwargs(requests, kwargs)? { + results.append( + self.0 + .bind(py) + .call_method("get_cache", (), Some(&kwargs))?, + )?; + } + Ok(results.into_any()) + } + + pub(super) fn async_lookup_batch<'py>( + &self, + py: Python<'py>, + requests: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyAny>>, + ) -> PyResult> { + let awaitables = batch_callback_kwargs(requests, kwargs)? + .iter() + .map(|kwargs| { + self.0 + .bind(py) + .call_method("async_get_cache", (), Some(kwargs)) + }) + .collect::>>()?; + py.import("asyncio")? + .call_method1("gather", PyTuple::new(py, awaitables)?) + } + + /// Receives the caller's original result, because the built-in + /// `Cache.async_add_cache_pipeline` splits the batch itself. + pub(super) fn async_store_batch<'py>( + &self, + py: Python<'py>, + result: Option<&Bound<'py, PyAny>>, + kwargs: Option<&Bound<'py, PyDict>>, + ) -> PyResult> { + let result = result.ok_or_else(|| { + PyTypeError::new_err("Python cache callbacks require their original callback_result") + })?; + self.0.bind(py).call_method( + "async_add_cache_pipeline", + (result,), + Some(callback_kwargs(kwargs)?), + ) + } + + /// The built-in `Cache` facade has no flush of its own; its backend does. + pub(super) fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { + let object = self.0.bind(py); + let backend = match object.getattr_opt("cache")? { + Some(backend) if !backend.is_none() => backend, + _ => object.clone(), + }; + if backend.hasattr("async_flush_cache")? { + return backend.call_method0("async_flush_cache"); + } + backend.call_method0("flush_cache")?; + ready_none(py) + } + + pub(super) fn ping<'py>(&self, py: Python<'py>) -> PyResult> { + self.0.bind(py).call_method0("ping") + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } +} + +fn callback_kwargs<'a, 'py>( + kwargs: Option<&'a Bound<'py, PyDict>>, +) -> PyResult<&'a Bound<'py, PyDict>> { + kwargs.ok_or_else(|| { + PyTypeError::new_err("Python cache callbacks require their original callback_kwargs") + }) +} + +fn batch_callback_kwargs<'py>( + requests: &Bound<'py, PyAny>, + kwargs: Option<&Bound<'py, PyAny>>, +) -> PyResult>> { + let kwargs = kwargs + .ok_or_else(|| { + PyTypeError::new_err( + "Python cache callbacks require one original callback_kwargs mapping per request", + ) + })? + .try_iter()? + .map(|item| Ok(item?.cast_into::()?)) + .collect::>>()?; + if kwargs.len() != requests.len()? { + return Err(PyValueError::new_err( + "batch cache requests and callback_kwargs must have equal lengths", + )); + } + Ok(kwargs) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 2ad13ce200f..19220bf868f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -9,7 +9,7 @@ use pyo3::{ }; use serde_json::Value; -use super::{CacheTestHandle, native::NativeResponseCache}; +use super::{handle::CacheTestHandle, native::NativeResponseCache}; struct ClassGuard { class: Py, diff --git a/litellm-rust/crates/python-bridge/src/cache/future.rs b/litellm-rust/crates/python-bridge/src/cache/future.rs new file mode 100644 index 00000000000..42593eee1f4 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/future.rs @@ -0,0 +1,18 @@ +use litellm_host_python::to_py; +use pyo3::prelude::*; + +pub(super) fn ready_none(py: Python<'_>) -> PyResult> { + ready_value(py, &()) +} + +pub(super) fn ready_value<'py, T: serde::Serialize>( + py: Python<'py>, + value: &T, +) -> PyResult> { + let future = py + .import("asyncio")? + .call_method0("get_running_loop")? + .call_method0("create_future")?; + future.call_method1("set_result", (to_py(py, value)?,))?; + Ok(future) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs new file mode 100644 index 00000000000..42d7f2c2f3d --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -0,0 +1,106 @@ +use std::time::Duration; + +use litellm_host_python::release_gil; +use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; + +use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; +use crate::python_settings::PythonSettings; + +const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); + +#[derive(FromPyObject)] +struct PythonCacheSettings { + default_redis_ttl: Option, +} + +fn redis_default_ttl(py: Python<'_>) -> PyResult { + let settings: PythonCacheSettings = PythonSettings::Cache.read(py)?.extract()?; + settings + .default_redis_ttl + .map(duration) + .transpose() + .map(|ttl| ttl.unwrap_or(PYTHON_REDIS_DEFAULT_TTL)) +} + +#[pyclass(frozen, name = "_CacheTestHandle")] +pub(crate) struct CacheTestHandle { + service: NativeResponseCache, + pub(super) guard: Option, + pid: u32, +} + +impl CacheTestHandle { + pub(super) fn service(&self) -> PyResult { + if self.pid != std::process::id() { + return Err(PyRuntimeError::new_err( + "native cache handles must be recreated after fork", + )); + } + Ok(self.service.clone()) + } +} + +#[pymethods] +impl CacheTestHandle { + #[staticmethod] + #[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))] + fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult { + Ok(Self { + service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes), + guard: None, + pid: std::process::id(), + }) + } + + #[staticmethod] + #[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))] + fn redis( + py: Python<'_>, + url: String, + ttl_seconds: Option, + namespace: Option, + ) -> PyResult { + let ttl = Some(match ttl_seconds { + Some(seconds) => duration(seconds)?, + None => redis_default_ttl(py)?, + }); + let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + + #[getter] + fn backend(&self) -> &'static str { + self.service.kind() + } + + fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { + let service = self.service()?; + let guard = FacadeGuard::capture(py, facade, &service)?; + let service = service.with_redis_flush_size( + facade + .getattr("redis_flush_size")? + .extract::>()?, + ); + let handle = Py::new( + py, + Self { + service, + guard: Some(guard), + pid: self.pid, + }, + )?; + facade.setattr("_native_cache_handle", handle) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Some(guard) = &self.guard { + guard.traverse(visit)?; + } + Ok(()) + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index f83788f6036..7955ed934b2 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,81 +1,21 @@ +mod binding; +mod callback; mod facade; +mod future; +mod handle; mod native; +mod request; +mod resolver; -use std::time::{Duration, SystemTime, UNIX_EPOCH}; - -use facade::FacadeGuard; use litellm_cache::Error; -use litellm_cache_response::{CacheControls, CacheKeyInput, PartialHits, ResponseCacheRequest}; -use litellm_host_python::{ExecutionStep, from_py, release_gil, run_async, to_py}; -use native::NativeResponseCache; use pyo3::{ - PyTraverseError, PyVisit, - exceptions::{PyRuntimeError, PyTypeError, PyValueError}, + exceptions::{PyRuntimeError, PyValueError}, prelude::*, - types::{PyDict, PyList, PyTuple}, }; -use serde::Deserialize; -use serde_json::Value; -use crate::python_settings::PythonSettings; - -const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); - -#[derive(FromPyObject)] -struct PythonCacheSettings { - default_redis_ttl: Option, -} - -fn redis_default_ttl(py: Python<'_>) -> PyResult { - let settings: PythonCacheSettings = PythonSettings::Cache.read(py)?.extract()?; - settings - .default_redis_ttl - .map(duration) - .transpose() - .map(|ttl| ttl.unwrap_or(PYTHON_REDIS_DEFAULT_TTL)) -} - -#[derive(Deserialize)] -#[serde(deny_unknown_fields)] -struct RequestInput { - key: CacheKeyInput, - controls: Option, - ttl_seconds: Option, - max_age_seconds: Option, -} - -fn request(value: &Bound<'_, PyAny>) -> PyResult { - let input: RequestInput = from_py(value)?; - request_input(input) -} - -fn request_input(input: RequestInput) -> PyResult { - let mut request = ResponseCacheRequest::new(input.key); - if let Some(controls) = input.controls { - request.controls = controls; - } - request.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?; - request.max_age = input.max_age_seconds.map(duration).transpose()?; - Ok(request) -} - -fn requests(value: &Bound<'_, PyAny>) -> PyResult> { - from_py::>(value)? - .into_iter() - .map(request_input) - .collect() -} - -fn duration(seconds: f64) -> PyResult { - Duration::try_from_secs_f64(seconds) - .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) -} - -fn now() -> Duration { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() -} +pub(crate) use self::{ + binding::ResolvedCache, handle::CacheTestHandle, resolver::CacheTestResolver, +}; fn cache_error(error: Error) -> PyErr { match error { @@ -83,493 +23,3 @@ fn cache_error(error: Error) -> PyErr { _ => PyRuntimeError::new_err(error.to_string()), } } - -#[pyclass(frozen, name = "_CacheTestHandle")] -pub(crate) struct CacheTestHandle { - service: NativeResponseCache, - guard: Option, - pid: u32, -} - -impl CacheTestHandle { - fn service(&self) -> PyResult { - if self.pid != std::process::id() { - return Err(PyRuntimeError::new_err( - "native cache handles must be recreated after fork", - )); - } - Ok(self.service.clone()) - } -} - -#[pymethods] -impl CacheTestHandle { - #[staticmethod] - #[pyo3(signature = (*, capacity=200, ttl_seconds=600.0, max_entry_bytes=1048576))] - fn memory(capacity: usize, ttl_seconds: f64, max_entry_bytes: usize) -> PyResult { - Ok(Self { - service: NativeResponseCache::memory(capacity, duration(ttl_seconds)?, max_entry_bytes), - guard: None, - pid: std::process::id(), - }) - } - - #[staticmethod] - #[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))] - fn redis( - py: Python<'_>, - url: String, - ttl_seconds: Option, - namespace: Option, - ) -> PyResult { - let ttl = Some(match ttl_seconds { - Some(seconds) => duration(seconds)?, - None => redis_default_ttl(py)?, - }); - let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) - .map_err(cache_error)?; - Ok(Self { - service, - guard: None, - pid: std::process::id(), - }) - } - - #[getter] - fn backend(&self) -> &'static str { - self.service.kind() - } - - fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { - let service = self.service()?; - let guard = FacadeGuard::capture(py, facade, &service)?; - let service = service.with_redis_flush_size( - facade - .getattr("redis_flush_size")? - .extract::>()?, - ); - let handle = Py::new( - py, - Self { - service, - guard: Some(guard), - pid: self.pid, - }, - )?; - facade.setattr("_native_cache_handle", handle) - } - - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - if let Some(guard) = &self.guard { - guard.traverse(visit)?; - } - Ok(()) - } -} - -enum CacheBinding { - Disabled, - Native(NativeResponseCache), - PythonCallback(Py), -} - -#[pyclass(frozen, name = "_CacheTestBinding")] -pub(crate) struct ResolvedCache { - binding: CacheBinding, - pid: u32, -} - -impl ResolvedCache { - fn check_process(&self) -> PyResult<()> { - if matches!(self.binding, CacheBinding::Native(_)) && self.pid != std::process::id() { - return Err(PyRuntimeError::new_err( - "native cache bindings must be resolved again after fork", - )); - } - Ok(()) - } - - pub(crate) fn lookup_step( - &self, - py: Python<'_>, - input: &Bound<'_, PyAny>, - kwargs: Option<&Bound<'_, PyDict>>, - ) -> PyResult { - self.check_process()?; - let awaitable = match &self.binding { - CacheBinding::Disabled => ready_none(py)?, - CacheBinding::Native(service) => { - let request = request(input)?; - let service = service.clone(); - run_async( - py, - async move { service.async_lookup(&request, now()).await }, - cache_error, - )? - } - CacheBinding::PythonCallback(object) => object.bind(py).call_method( - "async_get_cache", - (), - Some(callback_kwargs(kwargs)?), - )?, - }; - Ok(ExecutionStep::Await(awaitable.unbind())) - } -} - -#[pymethods] -impl ResolvedCache { - #[getter] - fn kind(&self) -> &'static str { - match self.binding { - CacheBinding::Disabled => "disabled", - CacheBinding::Native(_) => "native", - CacheBinding::PythonCallback(_) => "python_callback", - } - } - - #[pyo3(signature = (request, *, callback_kwargs=None))] - fn lookup( - &self, - py: Python<'_>, - request: &Bound<'_, PyAny>, - callback_kwargs: Option<&Bound<'_, PyDict>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => Ok(py.None()), - CacheBinding::Native(service) => { - let request = self::request(request)?; - let service = service.clone(); - let response = release_gil(py, move || service.lookup(&request, now())) - .map_err(cache_error)?; - to_py(py, &response) - } - CacheBinding::PythonCallback(object) => object - .bind(py) - .call_method( - "get_cache", - (), - Some(self::callback_kwargs(callback_kwargs)?), - ) - .map(Bound::unbind), - } - } - - #[pyo3(signature = (request, response, *, callback_kwargs=None))] - fn store( - &self, - py: Python<'_>, - request: &Bound<'_, PyAny>, - response: &Bound<'_, PyAny>, - callback_kwargs: Option<&Bound<'_, PyDict>>, - ) -> PyResult<()> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => Ok(()), - CacheBinding::Native(service) => { - let request = self::request(request)?; - let response: Value = from_py(response)?; - let service = service.clone(); - release_gil(py, move || service.store(&request, response, now())) - .map_err(cache_error) - } - CacheBinding::PythonCallback(object) => object - .bind(py) - .call_method( - "add_cache", - (response,), - Some(self::callback_kwargs(callback_kwargs)?), - ) - .map(|_| ()), - } - } - - /// Native bindings return `{values, missing_indices}`. The built-in `Cache` API has no batch - /// read, so a Python callback receives one `get_cache(**kwargs)` call per request, in order, - /// and the results come back as a list. - #[pyo3(signature = (requests, *, callback_kwargs=None))] - fn lookup_batch( - &self, - py: Python<'_>, - requests: &Bound<'_, PyAny>, - callback_kwargs: Option<&Bound<'_, PyAny>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => { - let requests = self::requests(requests)?; - to_py(py, &PartialHits::new(vec![None; requests.len()])) - } - CacheBinding::Native(service) => { - let requests = self::requests(requests)?; - let service = service.clone(); - let response = release_gil(py, move || service.lookup_batch(&requests, now())) - .map_err(cache_error)?; - to_py(py, &response) - } - CacheBinding::PythonCallback(object) => { - let results = PyList::empty(py); - for kwargs in batch_callback_kwargs(requests, callback_kwargs)? { - results.append(object.bind(py).call_method( - "get_cache", - (), - Some(&kwargs), - )?)?; - } - Ok(results.into_any().unbind()) - } - } - } - - #[pyo3(signature = (request, *, callback_kwargs=None))] - fn async_lookup<'py>( - &self, - py: Python<'py>, - request: &Bound<'py, PyAny>, - callback_kwargs: Option<&Bound<'py, PyDict>>, - ) -> PyResult> { - let ExecutionStep::Await(awaitable) = self.lookup_step(py, request, callback_kwargs)? - else { - unreachable!() - }; - Ok(awaitable.into_bound(py)) - } - - #[pyo3(signature = (request, response, *, callback_kwargs=None))] - fn async_store<'py>( - &self, - py: Python<'py>, - request: &Bound<'py, PyAny>, - response: &Bound<'py, PyAny>, - callback_kwargs: Option<&Bound<'py, PyDict>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let request = self::request(request)?; - let response: Value = from_py(response)?; - let service = service.clone(); - run_async( - py, - async move { service.async_store(&request, response, now()).await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => object.bind(py).call_method( - "async_add_cache", - (response,), - Some(self::callback_kwargs(callback_kwargs)?), - ), - } - } - - #[pyo3(signature = (requests, *, callback_kwargs=None))] - fn async_lookup_batch<'py>( - &self, - py: Python<'py>, - requests: &Bound<'py, PyAny>, - callback_kwargs: Option<&Bound<'py, PyAny>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => { - let requests = self::requests(requests)?; - ready_value(py, &PartialHits::new(vec![None; requests.len()])) - } - CacheBinding::Native(service) => { - let requests = self::requests(requests)?; - let service = service.clone(); - run_async( - py, - async move { service.async_lookup_batch(&requests, now()).await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => { - let awaitables = batch_callback_kwargs(requests, callback_kwargs)? - .iter() - .map(|kwargs| { - object - .bind(py) - .call_method("async_get_cache", (), Some(kwargs)) - }) - .collect::>>()?; - py.import("asyncio")? - .call_method1("gather", PyTuple::new(py, awaitables)?) - } - } - } - - /// A Python callback receives the caller's original result through `callback_result`, because - /// the built-in `Cache.async_add_cache_pipeline` splits the batch itself. - #[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))] - fn async_store_batch<'py>( - &self, - py: Python<'py>, - requests: &Bound<'py, PyAny>, - responses: &Bound<'py, PyAny>, - callback_result: Option<&Bound<'py, PyAny>>, - callback_kwargs: Option<&Bound<'py, PyDict>>, - ) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let requests = self::requests(requests)?; - let responses: Vec = from_py(responses)?; - if requests.len() != responses.len() { - return Err(PyValueError::new_err( - "batch cache requests and responses must have equal lengths", - )); - } - let entries = requests.into_iter().zip(responses).collect(); - let service = service.clone(); - run_async( - py, - async move { service.async_store_batch(entries, now()).await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => { - let result = callback_result.ok_or_else(|| { - PyTypeError::new_err( - "Python cache callbacks require their original callback_result", - ) - })?; - object.bind(py).call_method( - "async_add_cache_pipeline", - (result,), - Some(self::callback_kwargs(callback_kwargs)?), - ) - } - } - } - - fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let service = service.clone(); - run_async(py, async move { service.async_flush().await }, cache_error) - } - // The built-in `Cache` facade has no flush of its own; its backend does. - CacheBinding::PythonCallback(object) => { - let object = object.bind(py); - let backend = match object.getattr_opt("cache")? { - Some(backend) if !backend.is_none() => backend, - _ => object.clone(), - }; - if backend.hasattr("async_flush_cache")? { - return backend.call_method0("async_flush_cache"); - } - backend.call_method0("flush_cache")?; - ready_none(py) - } - } - } - - fn ping<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - match &self.binding { - CacheBinding::Disabled => ready_none(py), - CacheBinding::Native(service) => { - let service = service.clone(); - run_async( - py, - async move { service.test_connection().await }, - cache_error, - ) - } - CacheBinding::PythonCallback(object) => object.bind(py).call_method0("ping"), - } - } - - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - if let CacheBinding::PythonCallback(object) = &self.binding { - visit.call(object)?; - } - Ok(()) - } -} - -fn callback_kwargs<'a, 'py>( - kwargs: Option<&'a Bound<'py, PyDict>>, -) -> PyResult<&'a Bound<'py, PyDict>> { - kwargs.ok_or_else(|| { - PyTypeError::new_err("Python cache callbacks require their original callback_kwargs") - }) -} - -fn batch_callback_kwargs<'py>( - requests: &Bound<'py, PyAny>, - kwargs: Option<&Bound<'py, PyAny>>, -) -> PyResult>> { - let kwargs = kwargs - .ok_or_else(|| { - PyTypeError::new_err( - "Python cache callbacks require one original callback_kwargs mapping per request", - ) - })? - .try_iter()? - .map(|item| Ok(item?.cast_into::()?)) - .collect::>>()?; - if kwargs.len() != requests.len()? { - return Err(PyValueError::new_err( - "batch cache requests and callback_kwargs must have equal lengths", - )); - } - Ok(kwargs) -} - -fn ready_none(py: Python<'_>) -> PyResult> { - ready_value(py, &()) -} - -fn ready_value<'py, T: serde::Serialize>( - py: Python<'py>, - value: &T, -) -> PyResult> { - let future = py - .import("asyncio")? - .call_method0("get_running_loop")? - .call_method0("create_future")?; - future.call_method1("set_result", (to_py(py, value)?,))?; - Ok(future) -} - -#[pyclass(frozen, name = "_CacheTestResolver")] -pub(crate) struct CacheTestResolver { - namespace: Py, -} - -#[pymethods] -impl CacheTestResolver { - #[new] - fn new(namespace: Py) -> Self { - Self { namespace } - } - - pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult { - let object = self.namespace.bind(py).getattr("cache")?; - let binding = if object.is_none() { - CacheBinding::Disabled - } else if let Ok(handle) = object.extract::>() { - CacheBinding::Native(handle.service()?) - } else if let Some(service) = facade::resolve(py, &object)? { - CacheBinding::Native(service) - } else { - CacheBinding::PythonCallback(object.unbind()) - }; - Ok(ResolvedCache { - binding, - pid: std::process::id(), - }) - } - - fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&self.namespace) - } -} diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 3fc8f61dff6..a718d07b286 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -4,25 +4,19 @@ use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ - CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, + CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, }; use serde_json::Value; -use tokio::sync::Mutex; #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), Redis { cache: Arc>>, - buffer: Option>, + buffer: Option>, }, } -pub(super) struct RedisWriteBuffer { - flush_size: usize, - entries: Mutex>, -} - impl NativeResponseCache { pub fn memory(capacity: usize, ttl: Duration, max_entry_bytes: usize) -> Self { Self::Memory(Arc::new(ResponseCache::new(Arc::new( @@ -33,7 +27,7 @@ impl NativeResponseCache { Some(Arc::new(|entry| { ResponseCacheCodec.encode(entry).map(|bytes| bytes.len()) })), - super::now, + super::request::now, ), )))) } @@ -84,12 +78,7 @@ impl NativeResponseCache { match self { Self::Redis { cache, .. } => Self::Redis { cache, - buffer: flush_size.map(|flush_size| { - Arc::new(RedisWriteBuffer { - flush_size: flush_size.max(1), - entries: Mutex::new(Vec::new()), - }) - }), + buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))), }, memory => memory, } @@ -155,19 +144,7 @@ impl NativeResponseCache { Self::Redis { cache, buffer: Some(buffer), - } => { - let pending = { - let mut entries = buffer.entries.lock().await; - entries.push((request.clone(), response, now)); - (entries.len() >= buffer.flush_size).then(|| std::mem::take(&mut *entries)) - }; - // A failed flush drops its batch, as Python does. Requeueing would grow the - // buffer and re-send an ever larger pipeline on every write during an outage. - match pending { - Some(pending) => cache.async_store_entries(pending).await, - None => Ok(()), - } - } + } => buffer.async_store(cache, request, response, now).await, } } @@ -198,7 +175,7 @@ impl NativeResponseCache { Self::Memory(cache) => cache.async_flush().await, Self::Redis { cache, buffer } => { if let Some(buffer) = buffer { - buffer.entries.lock().await.clear(); + buffer.clear()?; } cache.async_flush().await } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs new file mode 100644 index 00000000000..d8793abd115 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -0,0 +1,48 @@ +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; +use litellm_host_python::from_py; +use pyo3::{exceptions::PyValueError, prelude::*}; +use serde::Deserialize; + +#[derive(Deserialize)] +#[serde(deny_unknown_fields)] +struct RequestInput { + key: CacheKeyInput, + controls: Option, + ttl_seconds: Option, + max_age_seconds: Option, +} + +pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { + let input: RequestInput = from_py(value)?; + request_input(input) +} + +fn request_input(input: RequestInput) -> PyResult { + let mut request = ResponseCacheRequest::new(input.key); + if let Some(controls) = input.controls { + request.controls = controls; + } + request.kwargs.ttl = input.ttl_seconds.map(duration).transpose()?; + request.max_age = input.max_age_seconds.map(duration).transpose()?; + Ok(request) +} + +pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { + from_py::>(value)? + .into_iter() + .map(request_input) + .collect() +} + +pub(super) fn duration(seconds: f64) -> PyResult { + Duration::try_from_secs_f64(seconds) + .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) +} + +pub(super) fn now() -> Duration { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() +} diff --git a/litellm-rust/crates/python-bridge/src/cache/resolver.rs b/litellm-rust/crates/python-bridge/src/cache/resolver.rs new file mode 100644 index 00000000000..ef6f142e0a1 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/resolver.rs @@ -0,0 +1,39 @@ +use pyo3::{PyTraverseError, PyVisit, prelude::*}; + +use super::{ + binding::{CacheBinding, ResolvedCache}, + callback::PythonCallback, + facade, + handle::CacheTestHandle, +}; + +#[pyclass(frozen, name = "_CacheTestResolver")] +pub(crate) struct CacheTestResolver { + namespace: Py, +} + +#[pymethods] +impl CacheTestResolver { + #[new] + fn new(namespace: Py) -> Self { + Self { namespace } + } + + pub(crate) fn resolve(&self, py: Python<'_>) -> PyResult { + let object = self.namespace.bind(py).getattr("cache")?; + let binding = if object.is_none() { + CacheBinding::Disabled + } else if let Ok(handle) = object.extract::>() { + CacheBinding::Native(handle.service()?) + } else if let Some(service) = facade::resolve(py, &object)? { + CacheBinding::Native(service) + } else { + CacheBinding::PythonCallback(PythonCallback::new(object.unbind())) + }; + Ok(ResolvedCache::new(binding)) + } + + fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.namespace) + } +} From 1b6b704ddd639898895b1add55cdb9579a278c61 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 11:15:44 -0700 Subject: [PATCH 10/17] refactor(cache): keep native foundation isolated --- litellm-rust/crates/cache-response/README.md | 2 +- .../crates/python-bridge/python_settings.json | 3 - .../crates/python-bridge/src/cache/handle.rs | 28 +------- .../python-bridge/src/python_settings.rs | 5 +- litellm/caching/dual_cache.py | 22 ++---- litellm/rust_bridge/_native.pyi | 68 +------------------ litellm/rust_bridge/settings.py | 11 --- tests/test_litellm/caching/test_dual_cache.py | 49 +++++++------ tests/test_litellm_rust/test_cache.py | 14 +--- 9 files changed, 42 insertions(+), 160 deletions(-) diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 9863c46783f..56c1646d343 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -42,7 +42,7 @@ The resolver reads the namespace's `cache` attribute each time it resolves. A ca Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend -The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Redis defaults come from the Python settings snapshot, including `litellm.default_redis_ttl`, and buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python +The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings, including `litellm.default_redis_ttl` and SSL options. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy diff --git a/litellm-rust/crates/python-bridge/python_settings.json b/litellm-rust/crates/python-bridge/python_settings.json index 15d0d603ac7..0af55083bef 100644 --- a/litellm-rust/crates/python-bridge/python_settings.json +++ b/litellm-rust/crates/python-bridge/python_settings.json @@ -22,8 +22,5 @@ ], "secret_manager": [ "readable" - ], - "cache_settings": [ - "default_redis_ttl" ] } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 42d7f2c2f3d..8251b3df06c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,26 +1,7 @@ -use std::time::Duration; - use litellm_host_python::release_gil; use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; -use crate::python_settings::PythonSettings; - -const PYTHON_REDIS_DEFAULT_TTL: Duration = Duration::from_secs(60); - -#[derive(FromPyObject)] -struct PythonCacheSettings { - default_redis_ttl: Option, -} - -fn redis_default_ttl(py: Python<'_>) -> PyResult { - let settings: PythonCacheSettings = PythonSettings::Cache.read(py)?.extract()?; - settings - .default_redis_ttl - .map(duration) - .transpose() - .map(|ttl| ttl.unwrap_or(PYTHON_REDIS_DEFAULT_TTL)) -} #[pyclass(frozen, name = "_CacheTestHandle")] pub(crate) struct CacheTestHandle { @@ -53,17 +34,14 @@ impl CacheTestHandle { } #[staticmethod] - #[pyo3(signature = (url, *, ttl_seconds=None, namespace=None))] + #[pyo3(signature = (url, *, ttl_seconds=60.0, namespace=None))] fn redis( py: Python<'_>, url: String, - ttl_seconds: Option, + ttl_seconds: f64, namespace: Option, ) -> PyResult { - let ttl = Some(match ttl_seconds { - Some(seconds) => duration(seconds)?, - None => redis_default_ttl(py)?, - }); + let ttl = Some(duration(ttl_seconds)?); let service = release_gil(py, move || NativeResponseCache::redis(&url, ttl, namespace)) .map_err(cache_error)?; Ok(Self { diff --git a/litellm-rust/crates/python-bridge/src/python_settings.rs b/litellm-rust/crates/python-bridge/src/python_settings.rs index 90819c5b3fc..7ac23a05542 100644 --- a/litellm-rust/crates/python-bridge/src/python_settings.rs +++ b/litellm-rust/crates/python-bridge/src/python_settings.rs @@ -8,17 +8,15 @@ pub(crate) enum PythonSettings { UrlPolicy, ProviderDefaults, SecretManager, - Cache, } impl PythonSettings { #[cfg(test)] - pub(crate) const ALL: [Self; 5] = [ + pub(crate) const ALL: [Self; 4] = [ Self::Http, Self::UrlPolicy, Self::ProviderDefaults, Self::SecretManager, - Self::Cache, ]; pub(crate) fn name(self) -> &'static str { @@ -27,7 +25,6 @@ impl PythonSettings { Self::UrlPolicy => "url_policy", Self::ProviderDefaults => "provider_defaults", Self::SecretManager => "secret_manager", - Self::Cache => "cache_settings", } } diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index 04c82232784..66be77dbb40 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -12,7 +12,7 @@ import logging import time from collections.abc import Sequence from threading import Lock -from typing import TYPE_CHECKING, Any, Final, TypeVar +from typing import TYPE_CHECKING, Any, Final if TYPE_CHECKING: from litellm.types.caching import RedisPipelineIncrementOperation @@ -34,20 +34,14 @@ else: from collections import OrderedDict -_KeyT = TypeVar("_KeyT") -_ValueT = TypeVar("_ValueT") - -class LimitedSizeOrderedDict(OrderedDict[_KeyT, _ValueT]): - def __init__(self, *, max_size: int = 100) -> None: - super().__init__() +class LimitedSizeOrderedDict(OrderedDict): + def __init__(self, *args, max_size=100, **kwargs): + super().__init__(*args, **kwargs) self.max_size = max_size - def __setitem__(self, key: _KeyT, value: _ValueT) -> None: - if key in self: - super().__setitem__(key, value) - self.move_to_end(key) - return + def __setitem__(self, key, value): + # If inserting a new key exceeds max size, remove the oldest item if len(self) >= self.max_size: self.popitem(last=False) super().__setitem__(key, value) @@ -74,9 +68,7 @@ class DualCache(BaseCache): self.in_memory_cache = in_memory_cache or InMemoryCache() # If redis_cache is not provided, use the default RedisCache self.redis_cache = redis_cache - self.last_redis_batch_access_time: LimitedSizeOrderedDict[str, float] = LimitedSizeOrderedDict( - max_size=default_max_redis_batch_cache_size - ) + self.last_redis_batch_access_time = LimitedSizeOrderedDict(max_size=default_max_redis_batch_cache_size) self._last_redis_batch_access_time_lock = Lock() self.redis_batch_cache_expiry = ( default_redis_batch_cache_expiry or litellm.default_redis_batch_cache_expiry or 10 diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index ab4639bc876..05a6df6d5af 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,5 +1,5 @@ from asyncio import Future -from collections.abc import AsyncIterator, Awaitable, Coroutine, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence from typing import Never, final from litellm.llms.base_llm.ocr.transformation import OCRResponse @@ -93,72 +93,6 @@ class ResponsesWebSocketConnection: def recv_text(self) -> Future[str | None]: ... def close(self) -> Future[None]: ... -@final -class _CacheTestHandle: - def __new__(cls, _uninstantiable: Never, /) -> Never: ... - @staticmethod - def memory( - *, capacity: int = 200, ttl_seconds: float = 600.0, max_entry_bytes: int = 1048576 - ) -> _CacheTestHandle: ... - @staticmethod - def redis(url: str, *, ttl_seconds: float | None = None, namespace: str | None = None) -> _CacheTestHandle: ... - @property - def backend(self) -> str: ... - def _bind_facade(self, facade: object) -> None: ... - -@final -class _CacheTestResolver: - def __new__(cls, namespace: object) -> _CacheTestResolver: ... - def resolve(self) -> _CacheTestBinding: ... - -@final -class _CacheTestBinding: - def __new__(cls, _uninstantiable: Never, /) -> Never: ... - @property - def kind(self) -> str: ... - def lookup( - self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None - ) -> object: ... - def store( - self, - request: Mapping[str, object] | None, - response: object, - *, - callback_kwargs: dict[str, object] | None = None, - ) -> None: ... - def lookup_batch( - self, - requests: Sequence[Mapping[str, object]], - *, - callback_kwargs: Sequence[dict[str, object]] | None = None, - ) -> object: ... - def async_lookup( - self, request: Mapping[str, object] | None, *, callback_kwargs: dict[str, object] | None = None - ) -> Awaitable[object]: ... - def async_store( - self, - request: Mapping[str, object] | None, - response: object, - *, - callback_kwargs: dict[str, object] | None = None, - ) -> Awaitable[object]: ... - def async_lookup_batch( - self, - requests: Sequence[Mapping[str, object]], - *, - callback_kwargs: Sequence[dict[str, object]] | None = None, - ) -> Awaitable[object]: ... - def async_store_batch( - self, - requests: Sequence[Mapping[str, object]], - responses: Sequence[object], - *, - callback_result: object = None, - callback_kwargs: dict[str, object] | None = None, - ) -> Awaitable[object]: ... - def async_flush(self) -> Awaitable[None]: ... - def ping(self) -> Awaitable[object]: ... - @final class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... diff --git a/litellm/rust_bridge/settings.py b/litellm/rust_bridge/settings.py index 862a116496d..3aa2d742862 100644 --- a/litellm/rust_bridge/settings.py +++ b/litellm/rust_bridge/settings.py @@ -36,11 +36,6 @@ class SecretManager: readable: bool -@dataclass(frozen=True, slots=True) -class CacheSettings: - default_redis_ttl: float | None - - def warn(message: str) -> None: from litellm._logging import verbose_logger @@ -55,12 +50,6 @@ def secret_manager() -> SecretManager: return SecretManager(readable=_should_read_secret_from_secret_manager()) -def cache_settings() -> CacheSettings: - import litellm - - return CacheSettings(default_redis_ttl=litellm.default_redis_ttl) - - def provider_defaults() -> ProviderDefaults: import litellm diff --git a/tests/test_litellm/caching/test_dual_cache.py b/tests/test_litellm/caching/test_dual_cache.py index d34d23ca1d7..5f59de9cca5 100644 --- a/tests/test_litellm/caching/test_dual_cache.py +++ b/tests/test_litellm/caching/test_dual_cache.py @@ -7,7 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.constants import DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE -from litellm.caching.dual_cache import DualCache, LimitedSizeOrderedDict +from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache, _redis_circuit_breaker_guard, _redis_circuit_breaker_guard_sync from litellm.types.caching import RedisPipelineIncrementOperation @@ -15,7 +15,9 @@ from litellm.types.caching import RedisPipelineIncrementOperation @pytest.mark.asyncio async def test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads(): - dual_cache = DualCache(redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10) + dual_cache = DualCache( + redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 + ) keys = ["shared_a", "shared_b"] start_gate = asyncio.Event() @@ -42,7 +44,9 @@ async def test_dual_cache_async_batch_get_cache_coalesces_concurrent_redis_reads @pytest.mark.asyncio async def test_dual_cache_async_batch_get_cache_rolls_back_redis_reservation_on_error(): - dual_cache = DualCache(redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10) + dual_cache = DualCache( + redis_cache=MagicMock(spec=RedisCache), default_redis_batch_cache_expiry=10 + ) keys = ["shared_a", "shared_b"] with patch.object( @@ -112,7 +116,9 @@ def test_dual_cache_batch_get_cache_only_reads_missing_keys_from_redis(): def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): mock_redis = _redis_mock_for_sync_batch({"absent_key": None}) - dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) first = dual_cache.batch_get_cache(keys=["absent_key"]) second = dual_cache.batch_get_cache(keys=["absent_key"]) @@ -125,7 +131,9 @@ def test_dual_cache_batch_get_cache_throttles_repeat_redis_reads(): def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): mock_redis = MagicMock(spec=RedisCache) mock_redis.batch_get_cache.side_effect = RuntimeError("redis unavailable") - dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) first_result = dual_cache.batch_get_cache(keys=["shared_a"]) second_result = dual_cache.batch_get_cache(keys=["shared_a"]) @@ -138,7 +146,9 @@ def test_dual_cache_batch_get_cache_rolls_back_redis_reservation_on_error(): def test_dual_cache_batch_get_cache_returns_memory_only_when_redis_read_is_throttled(): mock_redis = _redis_mock_for_sync_batch({"throttled_key": "redis_value"}) - dual_cache = DualCache(in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10) + dual_cache = DualCache( + in_memory_cache=InMemoryCache(), redis_cache=mock_redis, default_redis_batch_cache_expiry=10 + ) dual_cache.last_redis_batch_access_time["throttled_key"] = time.time() result = dual_cache.batch_get_cache(keys=["throttled_key"]) @@ -247,7 +257,9 @@ async def test_dual_cache_batch_redis_backfill_injects_default_in_memory_ttl(): default_in_memory_ttl, same as the single-key path.""" in_memory_cache = InMemoryCache(default_ttl=600) mock_redis = MagicMock(spec=RedisCache) - mock_redis.async_batch_get_cache = AsyncMock(return_value={"batch_backfill_key": "redis_value"}) + mock_redis.async_batch_get_cache = AsyncMock( + return_value={"batch_backfill_key": "redis_value"} + ) dual_cache = DualCache( in_memory_cache=in_memory_cache, redis_cache=mock_redis, @@ -359,7 +371,9 @@ async def test_circuit_breaker_open_skips_redis(): class FakeRedis: def __init__(self): - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=3, recovery_timeout=60) + self._circuit_breaker = RedisCircuitBreaker( + failure_threshold=3, recovery_timeout=60 + ) self._circuit_breaker._state = "open" self._circuit_breaker._opened_at = time.time() self.call_count = 0 @@ -412,7 +426,9 @@ def test_circuit_breaker_half_open_concurrent_calls_are_fast_failed(): # All subsequent concurrent callers: HALF_OPEN → fast-fail (return True) for _ in range(10): - assert cb.is_open() is True, "concurrent callers should be fast-failed in HALF_OPEN" + assert ( + cb.is_open() is True + ), "concurrent callers should be fast-failed in HALF_OPEN" def test_circuit_breaker_disabled_never_opens(): @@ -456,7 +472,9 @@ async def test_circuit_breaker_disabled_guard_always_calls_method(): class FakeRedis: def __init__(self): - self._circuit_breaker = RedisCircuitBreaker(failure_threshold=1, recovery_timeout=60, enabled=False) + self._circuit_breaker = RedisCircuitBreaker( + failure_threshold=1, recovery_timeout=60, enabled=False + ) self.call_count = 0 @_redis_circuit_breaker_guard @@ -773,14 +791,3 @@ async def test_async_delete_cache_keys_on_empty_list_touches_no_backend(): await dual_cache.async_delete_cache_keys([]) redis_cache.delete_cache_keys.assert_not_awaited() - - -def test_limited_ordered_dict_refreshes_recency_without_evicting_another_key(): - tracker = LimitedSizeOrderedDict(max_size=2) - tracker["hot"] = 1 - tracker["cold"] = 2 - - tracker["hot"] = 3 - tracker["new"] = 4 - - assert list(tracker.items()) == [("hot", 3), ("new", 4)] diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index cf46a0566f7..796a0ec36ac 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -361,18 +361,6 @@ def test_facade_registration_rejects_mismatched_capacity() -> None: _native._CacheTestHandle.memory(capacity=7)._bind_facade(facade) -async def test_redis_handle_reads_the_python_default_ttl(redis_url: str) -> None: - client: Final = redis.Redis.from_url(redis_url) - with rebound(litellm, "default_redis_ttl", 7): - binding: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url)) - ).resolve() - await binding.async_store(request("native-default"), {"value": 1}) - - assert 0 < client.ttl("native-default") <= 7 - client.close() - - async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: parsed: Final = urlparse(redis_url) with rebound(litellm, "default_redis_ttl", 60): @@ -386,7 +374,7 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: _native._CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) with pytest.raises(TypeError, match="namespaces must match"): _native._CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) - _native._CacheTestHandle.redis(redis_url)._bind_facade(facade) + _native._CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() client: Final = redis.Redis.from_url(redis_url) From 22995d1575ec7b6e9ee2712efaf24ef45a59e07a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 11:23:47 -0700 Subject: [PATCH 11/17] fix(cache): remove redundant source comments --- litellm-rust/crates/cache-memory/src/cache.rs | 1 - litellm-rust/crates/cache-redis/src/cache.rs | 2 -- litellm-rust/crates/cache-response/src/buffer.rs | 1 - litellm-rust/crates/python-bridge/src/cache/binding.rs | 3 --- litellm-rust/crates/python-bridge/src/cache/callback.rs | 7 ------- 5 files changed, 14 deletions(-) diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 77635893640..45c638f5178 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -228,7 +228,6 @@ where .get(key) .filter(|existing| eligible.is_empty() || eligible.contains(existing)) .cloned(); - // Matches the Redis claim: an unconditional claim only extends its own winner. if let Some(existing) = &existing && eligible.is_empty() && *existing != candidate diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 23b1fabbab4..2249966dd79 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -57,8 +57,6 @@ const INCREMENT_SCRIPT: &str = concat!( "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" ); -// Compare-and-set against the exact bytes the claim decision was made on. -// ARGV: [1] expected payload or "" when absent, [2] ttl, [3] new payload, [4] refresh ttl. const CLAIM_SCRIPT: &str = concat!( "local current = redis.call('GET', KEYS[1]); ", "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", diff --git a/litellm-rust/crates/cache-response/src/buffer.rs b/litellm-rust/crates/cache-response/src/buffer.rs index 68af5278c8c..1fd2bb809de 100644 --- a/litellm-rust/crates/cache-response/src/buffer.rs +++ b/litellm-rust/crates/cache-response/src/buffer.rs @@ -5,7 +5,6 @@ use serde_json::Value; use crate::{CacheEntry, ResponseCache, ResponseCacheRequest}; -/// Defers async writes until `flush_size` entries are pending, then stores them as one batch. pub struct WriteBuffer { flush_size: usize, entries: Mutex>, diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs index 44c133d4611..ad64b24d3c1 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -125,8 +125,6 @@ impl ResolvedCache { } } - /// Native bindings return `{values, missing_indices}`, while a Python callback returns the - /// list of its per-request results. #[pyo3(signature = (requests, *, callback_kwargs=None))] fn lookup_batch( &self, @@ -222,7 +220,6 @@ impl ResolvedCache { } } - /// A Python callback receives the caller's original result through `callback_result`. #[pyo3(signature = (requests, responses, *, callback_result=None, callback_kwargs=None))] fn async_store_batch<'py>( &self, diff --git a/litellm-rust/crates/python-bridge/src/cache/callback.rs b/litellm-rust/crates/python-bridge/src/cache/callback.rs index 318f9d02080..492e0329672 100644 --- a/litellm-rust/crates/python-bridge/src/cache/callback.rs +++ b/litellm-rust/crates/python-bridge/src/cache/callback.rs @@ -7,8 +7,6 @@ use pyo3::{ use super::future::ready_none; -/// A custom Python cache object, driven through the built-in `Cache` API so a `Cache` subclass -/// works unchanged. pub(super) struct PythonCallback(Py); impl PythonCallback { @@ -61,8 +59,6 @@ impl PythonCallback { ) } - /// The built-in `Cache` API has no batch read, so the callback receives one - /// `get_cache(**kwargs)` call per request, in order, and the results come back as a list. pub(super) fn lookup_batch<'py>( &self, py: Python<'py>, @@ -98,8 +94,6 @@ impl PythonCallback { .call_method1("gather", PyTuple::new(py, awaitables)?) } - /// Receives the caller's original result, because the built-in - /// `Cache.async_add_cache_pipeline` splits the batch itself. pub(super) fn async_store_batch<'py>( &self, py: Python<'py>, @@ -116,7 +110,6 @@ impl PythonCallback { ) } - /// The built-in `Cache` facade has no flush of its own; its backend does. pub(super) fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { let object = self.0.bind(py); let backend = match object.getattr_opt("cache")? { From a7bc8e373e0d5b582003b2aafe62292139c61e12 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 11:56:04 -0700 Subject: [PATCH 12/17] feat(cache): project Python backend configuration --- litellm-rust/crates/cache-memory/src/cache.rs | 4 + .../crates/python-bridge/src/cache/config.rs | 586 ++++++++++++++++++ .../crates/python-bridge/src/cache/facade.rs | 45 +- .../crates/python-bridge/src/cache/mod.rs | 1 + .../crates/python-bridge/src/cache/native.rs | 7 + 5 files changed, 616 insertions(+), 27 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/cache/config.rs diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 45c638f5178..54d831378e4 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -124,6 +124,10 @@ impl InMemoryCache { self.max_size_in_memory } + pub fn max_entry_bytes(&self) -> Option { + self.max_entry_bytes + } + pub fn expires_at(&self, key: &str) -> Result, Error> { Ok(self .state diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs new file mode 100644 index 00000000000..637bdab4055 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -0,0 +1,586 @@ +use std::time::Duration; + +use pyo3::{ + exceptions::{PyOverflowError, PyTypeError, PyValueError}, + prelude::*, + types::{PyAny, PyDict}, +}; + +use super::{native::NativeResponseCache, request::duration}; + +#[derive(PartialEq)] +pub(super) struct CachePolicy { + pub(super) mode: String, + pub(super) ttl: Option, + pub(super) namespace: Option, + pub(super) supported_call_types: Option>, + pub(super) redis_flush_size: Option, + pub(super) semantic_cache_scope: String, +} + +#[derive(PartialEq)] +pub(super) struct MemoryCacheConfig { + pub(super) default_ttl: Duration, + pub(super) capacity: usize, + pub(super) max_entry_bytes: usize, +} + +#[derive(Debug, PartialEq)] +pub(super) enum RedisProtocol { + Resp2, + Resp3, +} + +#[derive(Debug, PartialEq)] +pub(super) enum CertificateRequirement { + None, + Optional, + Required, +} + +#[derive(PartialEq)] +pub(super) struct RedisTlsConfig { + pub(super) certificate_requirement: CertificateRequirement, + pub(super) check_hostname: bool, + pub(super) ca_certificate: Option, + pub(super) ca_data: Option>, + pub(super) client_certificate: Option, + pub(super) client_key: Option, +} + +#[derive(PartialEq)] +pub(super) struct RedisConnectionConfig { + pub(super) host: String, + pub(super) port: u16, + pub(super) database: i64, + pub(super) username: Option, + pub(super) password: Option, + pub(super) protocol: RedisProtocol, + pub(super) pool_size: usize, + pub(super) read_timeout: Option, + pub(super) connect_timeout: Option, + pub(super) socket_keepalive: Option, + pub(super) health_check_interval: Duration, + pub(super) client_name: Option, + pub(super) tls: Option, +} + +#[derive(PartialEq)] +pub(super) struct RedisCacheConfig { + pub(super) default_ttl: Duration, + pub(super) namespace: Option, + pub(super) flush_size: usize, + pub(super) connection: RedisConnectionConfig, +} + +#[derive(PartialEq)] +pub(super) enum CacheBackendConfig { + Memory(MemoryCacheConfig), + Redis(Box), +} + +#[derive(PartialEq)] +pub(super) struct NativeCacheConfig { + pub(super) policy: CachePolicy, + pub(super) backend: CacheBackendConfig, +} + +pub(super) enum UnsupportedCacheConfig { + Backend(String), + RedisMode(&'static str), + RedisOption(String), +} + +impl UnsupportedCacheConfig { + pub(super) fn message(&self) -> String { + match self { + Self::Backend(backend) => { + format!("native cache backend {backend:?} is not implemented") + } + Self::RedisMode(mode) => format!("native Redis {mode} mode is not implemented"), + Self::RedisOption(option) => { + format!("native Redis option {option:?} is not implemented") + } + } + } +} + +pub(super) enum CacheConfigProjection { + Native(Box), + Unsupported(UnsupportedCacheConfig), +} + +impl NativeCacheConfig { + pub(super) fn project(facade: &Bound<'_, PyAny>) -> PyResult { + let backend_name = facade.getattr("type")?.extract::()?; + let policy = CachePolicy { + mode: facade.getattr("mode")?.extract::()?, + ttl: optional_duration(facade.getattr("ttl")?)?, + namespace: optional_string(facade.getattr("namespace")?)?, + supported_call_types: facade + .getattr("supported_call_types")? + .extract::>>()?, + redis_flush_size: facade + .getattr("redis_flush_size")? + .extract::>()?, + semantic_cache_scope: facade + .getattr("semantic_cache_scope")? + .extract::()?, + }; + let backend = facade.getattr("cache")?; + match backend_name.as_str() { + "local" => project_memory(&backend).map(|backend| { + CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::Memory(backend), + })) + }), + "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( + UnsupportedCacheConfig::Backend(backend_name), + )), + } + } + + pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { + if service.default_ttl() + != 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"); + } + match &self.backend { + CacheBackendConfig::Memory(config) if service.kind() != "memory" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::Memory(config) if service.capacity() != Some(config.capacity) => { + Some("facade and native backend capacities must match") + } + CacheBackendConfig::Memory(config) + if service.max_entry_bytes() != Some(config.max_entry_bytes) => + { + Some("facade and native backend item limits must match") + } + CacheBackendConfig::Memory(_) => None, + CacheBackendConfig::Redis(_) if service.kind() != "redis" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::Redis(config) => (service.namespace() + != config.namespace.as_deref()) + .then_some("facade and native backend namespaces must match"), + } + } +} + +fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { + let max_size_kib = backend.getattr("max_size_per_item")?.extract::()?; + Ok(MemoryCacheConfig { + default_ttl: duration(backend.getattr("default_ttl")?.extract::()?)?, + capacity: backend.getattr("max_size_in_memory")?.extract::()?, + max_entry_bytes: max_size_kib + .checked_mul(1024) + .ok_or_else(|| PyOverflowError::new_err("memory cache item limit is too large"))?, + }) +} + +fn project_redis( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let source = backend.getattr("redis_kwargs")?.cast_into::()?; + if has_value(&source, "startup_nodes")? { + return Ok(Err(UnsupportedCacheConfig::RedisMode("cluster"))); + } + if has_value(&source, "sentinel_nodes")? { + return Ok(Err(UnsupportedCacheConfig::RedisMode("sentinel"))); + } + for key in [ + "credential_provider", + "redis_connect_func", + "connection_pool", + ] { + if has_value(&source, key)? { + return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + } + } + for key in [ + "retry", + "retry_on_error", + "socket_keepalive_options", + "unix_socket_path", + "cache", + "cache_config", + "event_dispatcher", + "ssl_ca_path", + "ssl_password", + "ssl_min_version", + "ssl_ciphers", + "ssl_validate_ocsp", + "ssl_validate_ocsp_stapled", + "ssl_ocsp_context", + "ssl_ocsp_expected_cert", + ] { + if has_value(&source, key)? { + return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + } + } + for key in ["retry_on_timeout", "single_connection_client"] { + if optional_coerced_bool(&source, key)?.unwrap_or(false) { + return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + } + } + + let client = backend.getattr("redis_client")?; + let pool = client.getattr("connection_pool")?; + let pool_class = class_identity(&pool)?; + if pool_class != ("redis.connection".to_owned(), "ConnectionPool".to_owned()) { + return Ok(Err(UnsupportedCacheConfig::RedisMode("custom pool"))); + } + let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; + for key in ["credential_provider", "redis_connect_func"] { + if has_value(&resolved, key)? { + return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + } + } + let connection_class = resolved + .get_item("connection_class")? + .unwrap_or(pool.getattr("connection_class")?); + let connection_class = ( + connection_class + .getattr("__module__")? + .extract::()?, + connection_class + .getattr("__qualname__")? + .extract::()?, + ); + let tls = match connection_class { + (module, name) if module == "redis.connection" && name == "Connection" => None, + (module, name) if module == "redis.connection" && name == "SSLConnection" => { + Some(project_tls(&resolved)?) + } + _ => return Ok(Err(UnsupportedCacheConfig::RedisMode("custom connection"))), + }; + + let protocol = match optional_u8(&resolved, "protocol")?.unwrap_or(2) { + 2 => RedisProtocol::Resp2, + 3 => RedisProtocol::Resp3, + value => { + return Err(PyValueError::new_err(format!( + "unsupported Redis protocol version {value}" + ))); + } + }; + let health_check_interval = + duration(optional_f64(&resolved, "health_check_interval")?.unwrap_or(0.0))?; + Ok(Ok(RedisCacheConfig { + default_ttl: duration(backend.getattr("default_ttl")?.extract::()?)?, + namespace: optional_attribute_string(backend, "namespace")?, + flush_size: backend.getattr("redis_flush_size")?.extract::()?, + connection: RedisConnectionConfig { + host: required_string(&resolved, "host")?, + port: required_u16(&resolved, "port")?, + database: optional_i64(&resolved, "db")?.unwrap_or(0), + username: optional_dict_string(&resolved, "username")?, + password: optional_dict_string(&resolved, "password")?, + protocol, + pool_size: pool.getattr("max_connections")?.extract::()?, + read_timeout: optional_dict_duration(&resolved, "socket_timeout")?, + connect_timeout: optional_dict_duration(&resolved, "socket_connect_timeout")?, + socket_keepalive: optional_bool(&resolved, "socket_keepalive")?, + health_check_interval, + client_name: optional_dict_string(&resolved, "client_name")?, + tls, + }, + })) +} + +fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { + Ok(RedisTlsConfig { + certificate_requirement: certificate_requirement(values)?, + check_hostname: optional_bool(values, "ssl_check_hostname")?.unwrap_or(false), + ca_certificate: optional_dict_string(values, "ssl_ca_certs")?, + ca_data: optional_bytes(values, "ssl_ca_data")?, + client_certificate: optional_dict_string(values, "ssl_certfile")?, + client_key: optional_dict_string(values, "ssl_keyfile")?, + }) +} + +fn certificate_requirement(values: &Bound<'_, PyDict>) -> PyResult { + let Some(value) = values.get_item("ssl_cert_reqs")? else { + return Ok(CertificateRequirement::Required); + }; + if value.is_none() { + return Ok(CertificateRequirement::Required); + } + if let Ok(number) = value.extract::() { + return match number { + 0 => Ok(CertificateRequirement::None), + 1 => Ok(CertificateRequirement::Optional), + 2 => Ok(CertificateRequirement::Required), + _ => Err(PyValueError::new_err( + "invalid Redis TLS certificate requirement", + )), + }; + } + match value.str()?.to_str()?.to_ascii_lowercase().as_str() { + "none" | "cert_none" => Ok(CertificateRequirement::None), + "optional" | "cert_optional" => Ok(CertificateRequirement::Optional), + "required" | "cert_required" => Ok(CertificateRequirement::Required), + _ => Err(PyValueError::new_err( + "invalid Redis TLS certificate requirement", + )), + } +} + +fn class_identity(value: &Bound<'_, PyAny>) -> PyResult<(String, String)> { + let class = value.get_type(); + Ok(( + class.getattr("__module__")?.extract::()?, + class.getattr("__qualname__")?.extract::()?, + )) +} + +fn optional_duration(value: Bound<'_, PyAny>) -> PyResult> { + value.extract::>()?.map(duration).transpose() +} + +fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { + match value.getattr(name) { + Ok(value) => optional_string(value), + Err(error) if error.is_instance_of::(value.py()) => { + Ok(None) + } + Err(error) => Err(error), + } +} + +fn optional_string(value: Bound<'_, PyAny>) -> PyResult> { + Ok(value + .extract::>()? + .filter(|value| !value.is_empty())) +} + +fn has_value(values: &Bound<'_, PyDict>, key: &str) -> PyResult { + Ok(values.get_item(key)?.is_some_and(|value| !value.is_none())) +} + +fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult { + values + .get_item(key)? + .ok_or_else(|| PyTypeError::new_err(format!("Redis connection is missing {key:?}")))? + .extract::() +} + +fn required_u16(values: &Bound<'_, PyDict>, key: &str) -> PyResult { + values + .get_item(key)? + .ok_or_else(|| PyTypeError::new_err(format!("Redis connection is missing {key:?}")))? + .extract::() +} + +fn optional_dict_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + match values.get_item(key)? { + Some(value) if !value.is_none() => optional_string(value), + _ => Ok(None), + } +} + +fn optional_bytes(values: &Bound<'_, PyDict>, key: &str) -> PyResult>> { + let Some(value) = values.get_item(key)? else { + return Ok(None); + }; + if value.is_none() { + return Ok(None); + } + if let Ok(bytes) = value.extract::>() { + return Ok(Some(bytes)); + } + Ok(Some(value.extract::()?.into_bytes())) +} + +fn optional_f64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + match values.get_item(key)? { + Some(value) => value.extract::>(), + None => Ok(None), + } +} + +fn optional_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + match values.get_item(key)? { + Some(value) => value.extract::>(), + None => Ok(None), + } +} + +fn optional_u8(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + match values.get_item(key)? { + Some(value) => value.extract::>(), + None => Ok(None), + } +} + +fn optional_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + match values.get_item(key)? { + Some(value) => value.extract::>(), + None => Ok(None), + } +} + +fn optional_coerced_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + let Some(value) = values.get_item(key)? else { + return Ok(None); + }; + if value.is_none() { + return Ok(None); + } + if let Ok(text) = value.extract::() { + return Ok(Some(matches!( + text.to_ascii_lowercase().as_str(), + "true" | "1" | "yes" + ))); + } + value.extract::().map(Some) +} + +fn optional_dict_duration(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { + optional_f64(values, key)?.map(duration).transpose() +} + +#[cfg(test)] +mod tests { + use std::ffi::CString; + + use pyo3::{prelude::*, types::PyDict}; + + use super::{ + CacheBackendConfig, CacheConfigProjection, CertificateRequirement, NativeCacheConfig, + RedisProtocol, + }; + use crate::cache::native::NativeResponseCache; + + fn facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> { + let locals = PyDict::new(py); + py.run( + &CString::new(format!( + "from types import SimpleNamespace\n\ + ConnectionPool = type('ConnectionPool', (), {{'__module__': 'redis.connection'}})\n\ + Connection = type('Connection', (), {{'__module__': 'redis.connection'}})\n\ + SSLConnection = type('SSLConnection', (), {{'__module__': 'redis.connection'}})\n\ + {body}" + )) + .unwrap(), + None, + Some(&locals), + ) + .unwrap(); + locals.get_item("facade").unwrap().unwrap() + } + + #[test] + fn projects_effective_memory_configuration() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(default_ttl=913, max_size_in_memory=37, max_size_per_item=8)\n\ + facade = SimpleNamespace(type='local', mode='default-on', ttl=11.5, namespace=None, supported_call_types=['completion'], redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("memory cache should be supported"); + }; + assert_eq!( + config.policy.ttl.unwrap(), + std::time::Duration::from_secs_f64(11.5) + ); + let CacheBackendConfig::Memory(memory) = config.backend else { + panic!("expected memory configuration"); + }; + assert_eq!(memory.default_ttl, std::time::Duration::from_secs(913)); + assert_eq!(memory.capacity, 37); + assert_eq!(memory.max_entry_bytes, 8192); + let matching = + NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8192); + let mismatched = + NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8191); + let matching_config = NativeCacheConfig { + policy: config.policy, + backend: CacheBackendConfig::Memory(memory), + }; + assert_eq!(matching_config.service_mismatch(&matching), None); + assert_eq!( + matching_config.service_mismatch(&mismatched), + Some("facade and native backend item limits must match") + ); + }); + } + + #[test] + fn projects_resolved_redis_tls_configuration() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "pool = ConnectionPool()\n\ + pool.connection_class = SSLConnection\n\ + pool.max_connections = 29\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6380, 'db': 4, 'username': 'user', 'password': 'secret', 'protocol': 3, 'socket_timeout': 7.5, 'socket_connect_timeout': 2, 'socket_keepalive': True, 'health_check_interval': 15, 'client_name': 'litellm', 'ssl_cert_reqs': 'optional', 'ssl_check_hostname': True, 'ssl_ca_certs': '/ca.pem', 'ssl_ca_data': 'CA DATA', 'ssl_certfile': '/client.pem', 'ssl_keyfile': '/client.key'}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\ + facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("Redis cache should be supported"); + }; + let CacheBackendConfig::Redis(redis) = config.backend else { + panic!("expected Redis configuration"); + }; + assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777)); + assert_eq!(redis.namespace.as_deref(), Some("team")); + assert_eq!(redis.flush_size, 31); + assert_eq!(redis.connection.host, "cache.internal"); + assert_eq!(redis.connection.port, 6380); + assert_eq!(redis.connection.database, 4); + assert_eq!(redis.connection.protocol, RedisProtocol::Resp3); + assert_eq!(redis.connection.pool_size, 29); + let tls = redis.connection.tls.unwrap(); + assert_eq!( + tls.certificate_requirement, + CertificateRequirement::Optional + ); + assert!(tls.check_hostname); + assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem")); + assert_eq!(tls.ca_data.as_deref(), Some(b"CA DATA".as_slice())); + assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem")); + assert_eq!(tls.client_key.as_deref(), Some("/client.key")); + }); + } + + #[test] + fn dynamic_redis_auth_stays_on_python() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\ + facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("dynamic authentication must stay on Python"); + }; + assert!(reason.message().contains("credential_provider")); + }); + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 19220bf868f..6f54d6a121a 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -1,5 +1,3 @@ -use std::time::Duration; - use litellm_host_python::from_py; use pyo3::{ PyTraverseError, PyVisit, @@ -9,7 +7,11 @@ use pyo3::{ }; use serde_json::Value; -use super::{handle::CacheTestHandle, native::NativeResponseCache}; +use super::{ + config::{CacheConfigProjection, NativeCacheConfig}, + handle::CacheTestHandle, + native::NativeResponseCache, +}; struct ClassGuard { class: Py, @@ -26,6 +28,7 @@ struct ObjectGuard { pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, + config: NativeCacheConfig, } impl ObjectGuard { @@ -134,7 +137,6 @@ impl FacadeGuard { service: &NativeResponseCache, ) -> PyResult { let kind = service.kind(); - let native_default_ttl: Duration = service.default_ttl(); let cache_type = py.import("litellm.caching.caching")?.getattr("Cache")?; if !facade.get_type().is(&cache_type) { return Err(PyTypeError::new_err( @@ -154,28 +156,14 @@ impl FacadeGuard { "facade and native backend types must match", )); } - let python_default_ttl = backend.getattr("default_ttl")?.extract::()?; - if python_default_ttl != native_default_ttl.as_secs_f64() { - return Err(PyTypeError::new_err( - "facade and native backend default TTLs must match", - )); - } - let namespace = match backend.getattr_opt("namespace")? { - Some(namespace) => namespace.extract::>()?, - None => None, - } - .filter(|namespace| !namespace.is_empty()); - if kind == "redis" && namespace.as_deref() != service.namespace() { - return Err(PyTypeError::new_err( - "facade and native backend namespaces must match", - )); - } - if let Some(capacity) = service.capacity() - && backend.getattr("max_size_in_memory")?.extract::()? != capacity - { - return Err(PyTypeError::new_err( - "facade and native backend capacities must match", - )); + let config = match NativeCacheConfig::project(facade)? { + CacheConfigProjection::Native(config) => *config, + CacheConfigProjection::Unsupported(reason) => { + return Err(PyTypeError::new_err(reason.message())); + } + }; + if let Some(message) = config.service_mismatch(service) { + return Err(PyTypeError::new_err(message)); } Ok(Self { outer: ObjectGuard::capture( @@ -203,12 +191,15 @@ impl FacadeGuard { "redis_flush_size", ], )?, + config, }) } fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult { + let projected = NativeCacheConfig::project(facade)?; Ok(self.outer.matches(py, facade)? - && self.backend.matches(py, &facade.getattr("cache")?)?) + && self.backend.matches(py, &facade.getattr("cache")?)? + && matches!(projected, CacheConfigProjection::Native(config) if *config == self.config)) } pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 7955ed934b2..aec08610f6e 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,5 +1,6 @@ mod binding; mod callback; +mod config; mod facade; mod future; mod handle; diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index a718d07b286..6a3835ac84d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -74,6 +74,13 @@ impl NativeResponseCache { } } + pub fn max_entry_bytes(&self) -> Option { + match self { + Self::Memory(cache) => cache.backend().max_entry_bytes(), + Self::Redis { .. } => None, + } + } + pub fn with_redis_flush_size(self, flush_size: Option) -> Self { match self { Self::Redis { cache, .. } => Self::Redis { From cda890438297ac133df56a93d69e0e3eac65340b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:03:13 -0700 Subject: [PATCH 13/17] fix(cache): keep native wheel within size budget --- .../crates/python-bridge/src/cache/config.rs | 84 ++++++++----------- 1 file changed, 33 insertions(+), 51 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 637bdab4055..37eee2048e0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -1,7 +1,7 @@ use std::time::Duration; use pyo3::{ - exceptions::{PyOverflowError, PyTypeError, PyValueError}, + exceptions::{PyTypeError, PyValueError}, prelude::*, types::{PyAny, PyDict}, }; @@ -43,7 +43,7 @@ pub(super) struct RedisTlsConfig { pub(super) certificate_requirement: CertificateRequirement, pub(super) check_hostname: bool, pub(super) ca_certificate: Option, - pub(super) ca_data: Option>, + pub(super) ca_data: Option, pub(super) client_certificate: Option, pub(super) client_key: Option, } @@ -86,21 +86,21 @@ pub(super) struct NativeCacheConfig { } pub(super) enum UnsupportedCacheConfig { - Backend(String), - RedisMode(&'static str), - RedisOption(String), + Backend, + RedisTopology, + RedisCredentials, + RedisConnection, + RedisOption, } impl UnsupportedCacheConfig { - pub(super) fn message(&self) -> String { + pub(super) fn message(&self) -> &'static str { match self { - Self::Backend(backend) => { - format!("native cache backend {backend:?} is not implemented") - } - Self::RedisMode(mode) => format!("native Redis {mode} mode is not implemented"), - Self::RedisOption(option) => { - format!("native Redis option {option:?} is not implemented") - } + Self::Backend => "native cache backend is not implemented", + Self::RedisTopology => "native Redis topology is not implemented", + Self::RedisCredentials => "native Redis credentials require Python", + Self::RedisConnection => "native Redis connection type is not implemented", + Self::RedisOption => "native Redis configuration requires Python", } } } @@ -143,7 +143,7 @@ impl NativeCacheConfig { Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, _ => Ok(CacheConfigProjection::Unsupported( - UnsupportedCacheConfig::Backend(backend_name), + UnsupportedCacheConfig::Backend, )), } } @@ -187,7 +187,7 @@ fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { capacity: backend.getattr("max_size_in_memory")?.extract::()?, max_entry_bytes: max_size_kib .checked_mul(1024) - .ok_or_else(|| PyOverflowError::new_err("memory cache item limit is too large"))?, + .ok_or_else(|| PyValueError::new_err("memory cache item limit is too large"))?, }) } @@ -196,20 +196,19 @@ fn project_redis( ) -> PyResult> { let source = backend.getattr("redis_kwargs")?.cast_into::()?; if has_value(&source, "startup_nodes")? { - return Ok(Err(UnsupportedCacheConfig::RedisMode("cluster"))); + return Ok(Err(UnsupportedCacheConfig::RedisTopology)); } if has_value(&source, "sentinel_nodes")? { - return Ok(Err(UnsupportedCacheConfig::RedisMode("sentinel"))); + return Ok(Err(UnsupportedCacheConfig::RedisTopology)); } - for key in [ - "credential_provider", - "redis_connect_func", - "connection_pool", - ] { + for key in ["credential_provider", "redis_connect_func"] { if has_value(&source, key)? { - return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); } } + if has_value(&source, "connection_pool")? { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } for key in [ "retry", "retry_on_error", @@ -228,12 +227,12 @@ fn project_redis( "ssl_ocsp_expected_cert", ] { if has_value(&source, key)? { - return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + return Ok(Err(UnsupportedCacheConfig::RedisOption)); } } for key in ["retry_on_timeout", "single_connection_client"] { if optional_coerced_bool(&source, key)?.unwrap_or(false) { - return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + return Ok(Err(UnsupportedCacheConfig::RedisOption)); } } @@ -241,12 +240,12 @@ fn project_redis( let pool = client.getattr("connection_pool")?; let pool_class = class_identity(&pool)?; if pool_class != ("redis.connection".to_owned(), "ConnectionPool".to_owned()) { - return Ok(Err(UnsupportedCacheConfig::RedisMode("custom pool"))); + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); } let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; for key in ["credential_provider", "redis_connect_func"] { if has_value(&resolved, key)? { - return Ok(Err(UnsupportedCacheConfig::RedisOption(key.to_owned()))); + return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); } } let connection_class = resolved @@ -265,17 +264,13 @@ fn project_redis( (module, name) if module == "redis.connection" && name == "SSLConnection" => { Some(project_tls(&resolved)?) } - _ => return Ok(Err(UnsupportedCacheConfig::RedisMode("custom connection"))), + _ => return Ok(Err(UnsupportedCacheConfig::RedisConnection)), }; let protocol = match optional_u8(&resolved, "protocol")?.unwrap_or(2) { 2 => RedisProtocol::Resp2, 3 => RedisProtocol::Resp3, - value => { - return Err(PyValueError::new_err(format!( - "unsupported Redis protocol version {value}" - ))); - } + _ => return Err(PyValueError::new_err("unsupported Redis protocol version")), }; let health_check_interval = duration(optional_f64(&resolved, "health_check_interval")?.unwrap_or(0.0))?; @@ -306,7 +301,7 @@ fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { certificate_requirement: certificate_requirement(values)?, check_hostname: optional_bool(values, "ssl_check_hostname")?.unwrap_or(false), ca_certificate: optional_dict_string(values, "ssl_ca_certs")?, - ca_data: optional_bytes(values, "ssl_ca_data")?, + ca_data: optional_dict_string(values, "ssl_ca_data")?, client_certificate: optional_dict_string(values, "ssl_certfile")?, client_key: optional_dict_string(values, "ssl_keyfile")?, }) @@ -374,14 +369,14 @@ fn has_value(values: &Bound<'_, PyDict>, key: &str) -> PyResult { fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult { values .get_item(key)? - .ok_or_else(|| PyTypeError::new_err(format!("Redis connection is missing {key:?}")))? + .ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))? .extract::() } fn required_u16(values: &Bound<'_, PyDict>, key: &str) -> PyResult { values .get_item(key)? - .ok_or_else(|| PyTypeError::new_err(format!("Redis connection is missing {key:?}")))? + .ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))? .extract::() } @@ -392,19 +387,6 @@ fn optional_dict_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult, key: &str) -> PyResult>> { - let Some(value) = values.get_item(key)? else { - return Ok(None); - }; - if value.is_none() { - return Ok(None); - } - if let Ok(bytes) = value.extract::>() { - return Ok(Some(bytes)); - } - Ok(Some(value.extract::()?.into_bytes())) -} - fn optional_f64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) => value.extract::>(), @@ -560,7 +542,7 @@ mod tests { ); assert!(tls.check_hostname); assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem")); - assert_eq!(tls.ca_data.as_deref(), Some(b"CA DATA".as_slice())); + assert_eq!(tls.ca_data.as_deref(), Some("CA DATA")); assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem")); assert_eq!(tls.client_key.as_deref(), Some("/client.key")); }); @@ -580,7 +562,7 @@ mod tests { else { panic!("dynamic authentication must stay on Python"); }; - assert!(reason.message().contains("credential_provider")); + assert_eq!(reason.message(), "native Redis credentials require Python"); }); } } From 59dbbe5ce71f382c64790ac5823c3f93a14fc8b5 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:10:06 -0700 Subject: [PATCH 14/17] fix(cache): outline Python configuration extraction --- .../crates/python-bridge/src/cache/config.rs | 66 ++++++++++++------- 1 file changed, 43 insertions(+), 23 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 37eee2048e0..bcacf7b4e36 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -3,7 +3,7 @@ use std::time::Duration; use pyo3::{ exceptions::{PyTypeError, PyValueError}, prelude::*, - types::{PyAny, PyDict}, + types::{PyAny, PyDict, PyString}, }; use super::{native::NativeResponseCache, request::duration}; @@ -111,6 +111,7 @@ pub(super) enum CacheConfigProjection { } impl NativeCacheConfig { + #[inline(never)] pub(super) fn project(facade: &Bound<'_, PyAny>) -> PyResult { let backend_name = facade.getattr("type")?.extract::()?; let policy = CachePolicy { @@ -180,6 +181,7 @@ impl NativeCacheConfig { } } +#[inline(never)] fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { let max_size_kib = backend.getattr("max_size_per_item")?.extract::()?; Ok(MemoryCacheConfig { @@ -191,6 +193,7 @@ fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { }) } +#[inline(never)] fn project_redis( backend: &Bound<'_, PyAny>, ) -> PyResult> { @@ -238,8 +241,7 @@ fn project_redis( let client = backend.getattr("redis_client")?; let pool = client.getattr("connection_pool")?; - let pool_class = class_identity(&pool)?; - if pool_class != ("redis.connection".to_owned(), "ConnectionPool".to_owned()) { + if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? { return Ok(Err(UnsupportedCacheConfig::RedisConnection)); } let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; @@ -251,20 +253,12 @@ fn project_redis( let connection_class = resolved .get_item("connection_class")? .unwrap_or(pool.getattr("connection_class")?); - let connection_class = ( - connection_class - .getattr("__module__")? - .extract::()?, - connection_class - .getattr("__qualname__")? - .extract::()?, - ); - let tls = match connection_class { - (module, name) if module == "redis.connection" && name == "Connection" => None, - (module, name) if module == "redis.connection" && name == "SSLConnection" => { - Some(project_tls(&resolved)?) - } - _ => return Ok(Err(UnsupportedCacheConfig::RedisConnection)), + let tls = if class_is(&connection_class, "redis.connection", "Connection")? { + None + } else if class_is(&connection_class, "redis.connection", "SSLConnection")? { + Some(project_tls(&resolved)?) + } else { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); }; let protocol = match optional_u8(&resolved, "protocol")?.unwrap_or(2) { @@ -296,6 +290,7 @@ fn project_redis( })) } +#[inline(never)] fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { Ok(RedisTlsConfig { certificate_requirement: certificate_requirement(values)?, @@ -307,6 +302,7 @@ fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { }) } +#[inline(never)] fn certificate_requirement(values: &Bound<'_, PyDict>) -> PyResult { let Some(value) = values.get_item("ssl_cert_reqs")? else { return Ok(CertificateRequirement::Required); @@ -334,18 +330,31 @@ fn certificate_requirement(values: &Bound<'_, PyDict>) -> PyResult) -> PyResult<(String, String)> { - let class = value.get_type(); - Ok(( - class.getattr("__module__")?.extract::()?, - class.getattr("__qualname__")?.extract::()?, - )) +#[inline(never)] +fn instance_class_is(value: &Bound<'_, PyAny>, module: &str, name: &str) -> PyResult { + class_is(value.get_type().as_any(), module, name) } +#[inline(never)] +fn class_is(value: &Bound<'_, PyAny>, module: &str, name: &str) -> PyResult { + Ok(value + .getattr("__module__")? + .cast_into::()? + .to_str()? + == module + && value + .getattr("__qualname__")? + .cast_into::()? + .to_str()? + == name) +} + +#[inline(never)] fn optional_duration(value: Bound<'_, PyAny>) -> PyResult> { value.extract::>()?.map(duration).transpose() } +#[inline(never)] fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { match value.getattr(name) { Ok(value) => optional_string(value), @@ -356,16 +365,19 @@ fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult) -> PyResult> { Ok(value .extract::>()? .filter(|value| !value.is_empty())) } +#[inline(never)] fn has_value(values: &Bound<'_, PyDict>, key: &str) -> PyResult { Ok(values.get_item(key)?.is_some_and(|value| !value.is_none())) } +#[inline(never)] fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult { values .get_item(key)? @@ -373,6 +385,7 @@ fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult { .extract::() } +#[inline(never)] fn required_u16(values: &Bound<'_, PyDict>, key: &str) -> PyResult { values .get_item(key)? @@ -380,6 +393,7 @@ fn required_u16(values: &Bound<'_, PyDict>, key: &str) -> PyResult { .extract::() } +#[inline(never)] fn optional_dict_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) if !value.is_none() => optional_string(value), @@ -387,6 +401,7 @@ fn optional_dict_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) => value.extract::>(), @@ -394,6 +409,7 @@ fn optional_f64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> } } +#[inline(never)] fn optional_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) => value.extract::>(), @@ -401,6 +417,7 @@ fn optional_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> } } +#[inline(never)] fn optional_u8(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) => value.extract::>(), @@ -408,6 +425,7 @@ fn optional_u8(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { } } +#[inline(never)] fn optional_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { Some(value) => value.extract::>(), @@ -415,6 +433,7 @@ fn optional_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult } } +#[inline(never)] fn optional_coerced_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { let Some(value) = values.get_item(key)? else { return Ok(None); @@ -431,6 +450,7 @@ fn optional_coerced_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult().map(Some) } +#[inline(never)] fn optional_dict_duration(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { optional_f64(values, key)?.map(duration).transpose() } From 14febcc8788d0c83094a096d612827d56925fb5f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:22:18 -0700 Subject: [PATCH 15/17] fix(rust): shrink cache configuration bridge --- .../crates/python-bridge/src/cache/config.rs | 56 +++++++++---------- .../crates/python-bridge/src/cache/facade.rs | 6 +- 2 files changed, 27 insertions(+), 35 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index bcacf7b4e36..6e694c20705 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -8,7 +8,7 @@ use pyo3::{ use super::{native::NativeResponseCache, request::duration}; -#[derive(PartialEq)] +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct CachePolicy { pub(super) mode: String, pub(super) ttl: Option, @@ -18,7 +18,6 @@ pub(super) struct CachePolicy { pub(super) semantic_cache_scope: String, } -#[derive(PartialEq)] pub(super) struct MemoryCacheConfig { pub(super) default_ttl: Duration, pub(super) capacity: usize, @@ -38,7 +37,7 @@ pub(super) enum CertificateRequirement { Required, } -#[derive(PartialEq)] +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct RedisTlsConfig { pub(super) certificate_requirement: CertificateRequirement, pub(super) check_hostname: bool, @@ -48,7 +47,7 @@ pub(super) struct RedisTlsConfig { pub(super) client_key: Option, } -#[derive(PartialEq)] +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct RedisConnectionConfig { pub(super) host: String, pub(super) port: u16, @@ -65,7 +64,7 @@ pub(super) struct RedisConnectionConfig { pub(super) tls: Option, } -#[derive(PartialEq)] +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct RedisCacheConfig { pub(super) default_ttl: Duration, pub(super) namespace: Option, @@ -73,13 +72,12 @@ pub(super) struct RedisCacheConfig { pub(super) connection: RedisConnectionConfig, } -#[derive(PartialEq)] pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), } -#[derive(PartialEq)] +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct NativeCacheConfig { pub(super) policy: CachePolicy, pub(super) backend: CacheBackendConfig, @@ -261,7 +259,7 @@ fn project_redis( return Ok(Err(UnsupportedCacheConfig::RedisConnection)); }; - let protocol = match optional_u8(&resolved, "protocol")?.unwrap_or(2) { + let protocol = match optional_i64(&resolved, "protocol")?.unwrap_or(2) { 2 => RedisProtocol::Resp2, 3 => RedisProtocol::Resp3, _ => return Err(PyValueError::new_err("unsupported Redis protocol version")), @@ -274,7 +272,8 @@ fn project_redis( flush_size: backend.getattr("redis_flush_size")?.extract::()?, connection: RedisConnectionConfig { host: required_string(&resolved, "host")?, - port: required_u16(&resolved, "port")?, + port: u16::try_from(required_i64(&resolved, "port")?) + .map_err(|_| PyValueError::new_err("invalid Redis port"))?, database: optional_i64(&resolved, "db")?.unwrap_or(0), username: optional_dict_string(&resolved, "username")?, password: optional_dict_string(&resolved, "password")?, @@ -320,14 +319,20 @@ fn certificate_requirement(values: &Bound<'_, PyDict>) -> PyResult Ok(CertificateRequirement::None), - "optional" | "cert_optional" => Ok(CertificateRequirement::Optional), - "required" | "cert_required" => Ok(CertificateRequirement::Required), - _ => Err(PyValueError::new_err( - "invalid Redis TLS certificate requirement", - )), + let text = value.str()?; + let text = text.to_str()?; + if text.eq_ignore_ascii_case("none") || text.eq_ignore_ascii_case("cert_none") { + return Ok(CertificateRequirement::None); } + if text.eq_ignore_ascii_case("optional") || text.eq_ignore_ascii_case("cert_optional") { + return Ok(CertificateRequirement::Optional); + } + if text.eq_ignore_ascii_case("required") || text.eq_ignore_ascii_case("cert_required") { + return Ok(CertificateRequirement::Required); + } + Err(PyValueError::new_err( + "invalid Redis TLS certificate requirement", + )) } #[inline(never)] @@ -386,11 +391,11 @@ fn required_string(values: &Bound<'_, PyDict>, key: &str) -> PyResult { } #[inline(never)] -fn required_u16(values: &Bound<'_, PyDict>, key: &str) -> PyResult { +fn required_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult { values .get_item(key)? .ok_or_else(|| PyTypeError::new_err("Redis connection is incomplete"))? - .extract::() + .extract::() } #[inline(never)] @@ -417,14 +422,6 @@ fn optional_i64(values: &Bound<'_, PyDict>, key: &str) -> PyResult> } } -#[inline(never)] -fn optional_u8(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { - match values.get_item(key)? { - Some(value) => value.extract::>(), - None => Ok(None), - } -} - #[inline(never)] fn optional_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult> { match values.get_item(key)? { @@ -442,10 +439,9 @@ fn optional_coerced_bool(values: &Bound<'_, PyDict>, key: &str) -> PyResult() { - return Ok(Some(matches!( - text.to_ascii_lowercase().as_str(), - "true" | "1" | "yes" - ))); + return Ok(Some( + text == "1" || text.eq_ignore_ascii_case("true") || text.eq_ignore_ascii_case("yes"), + )); } value.extract::().map(Some) } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 6f54d6a121a..83508356263 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -28,7 +28,6 @@ struct ObjectGuard { pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, - config: NativeCacheConfig, } impl ObjectGuard { @@ -191,15 +190,12 @@ impl FacadeGuard { "redis_flush_size", ], )?, - config, }) } fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult { - let projected = NativeCacheConfig::project(facade)?; Ok(self.outer.matches(py, facade)? - && self.backend.matches(py, &facade.getattr("cache")?)? - && matches!(projected, CacheConfigProjection::Native(config) if *config == self.config)) + && self.backend.matches(py, &facade.getattr("cache")?)?) } pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { From efcafa7f120341187991bd2655cdec3257459b3b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:30:17 -0700 Subject: [PATCH 16/17] fix(rust): invalidate changed Redis pool settings --- .../crates/python-bridge/src/cache/facade.rs | 69 ++++++++++++++++++- tests/test_litellm_rust/test_cache.py | 4 ++ 2 files changed, 70 insertions(+), 3 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 83508356263..f2f86c14b37 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -25,9 +25,17 @@ struct ObjectGuard { config: Vec, } +struct RedisPoolGuard { + reference: Py, + connection_class: Py, + connection_kwargs: Py, + max_connections: usize, +} + pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, + redis_pool: Option, } impl ObjectGuard { @@ -129,6 +137,45 @@ impl ObjectGuard { } } +impl RedisPoolGuard { + fn capture(backend: &Bound<'_, PyAny>) -> PyResult { + let pool = backend + .getattr("redis_client")? + .getattr("connection_pool")?; + Ok(Self { + reference: pool.clone().unbind(), + connection_class: pool.getattr("connection_class")?.unbind(), + connection_kwargs: pool + .getattr("connection_kwargs")? + .call_method0("copy")? + .unbind(), + max_connections: pool.getattr("max_connections")?.extract::()?, + }) + } + + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { + let pool = backend + .getattr("redis_client")? + .getattr("connection_pool")?; + Ok(self.reference.bind(py).is(&pool) + && self + .connection_class + .bind(py) + .is(&pool.getattr("connection_class")?) + && self.max_connections == pool.getattr("max_connections")?.extract::()? + && self + .connection_kwargs + .bind(py) + .eq(pool.getattr("connection_kwargs")?)?) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reference)?; + visit.call(&self.connection_class)?; + visit.call(&self.connection_kwargs) + } +} + impl FacadeGuard { pub(super) fn capture( py: Python<'_>, @@ -190,17 +237,33 @@ impl FacadeGuard { "redis_flush_size", ], )?, + redis_pool: (kind == "redis") + .then(|| RedisPoolGuard::capture(&backend)) + .transpose()?, }) } fn matches(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult { - Ok(self.outer.matches(py, facade)? - && self.backend.matches(py, &facade.getattr("cache")?)?) + if !self.outer.matches(py, facade)? { + return Ok(false); + } + let backend = facade.getattr("cache")?; + if !self.backend.matches(py, &backend)? { + return Ok(false); + } + match &self.redis_pool { + Some(guard) => guard.matches(py, &backend), + None => Ok(true), + } } pub(super) fn traverse(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { self.outer.traverse(&visit)?; - self.backend.traverse(&visit) + self.backend.traverse(&visit)?; + if let Some(guard) = &self.redis_pool { + guard.traverse(&visit)?; + } + Ok(()) } } diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 796a0ec36ac..c35cb1a20fb 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -381,6 +381,10 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + pool: Final = facade.cache.redis_client.connection_pool + with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}): + assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + await binding.async_store(request("first"), {"value": 1}) assert client.get("first") is None await binding.async_store(request("second"), {"value": 2}) From 991108cc4865efa649d0871966c275dc2593125e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 12:50:21 -0700 Subject: [PATCH 17/17] refactor(rust): align cache foundation with Python backends --- litellm-rust/crates/cache-memory/src/cache.rs | 113 ++++++-- .../crates/cache-memory/tests/cache.rs | 64 +++-- litellm-rust/crates/cache-redis/src/cache.rs | 243 ++++++++++------- .../cache-redis/src/cache/operations.rs | 137 +++++++++- litellm-rust/crates/cache-redis/src/lib.rs | 6 +- .../crates/cache-redis/src/topology.rs | 14 + .../crates/cache-redis/tests/cache.rs | 68 +++-- .../crates/cache-response/src/buffer.rs | 4 +- .../crates/cache-response/src/response.rs | 58 ++-- .../crates/cache-response/tests/response.rs | 4 +- litellm-rust/crates/cache/src/base_cache.rs | 106 +++----- litellm-rust/crates/cache/src/cache_type.rs | 85 ++++++ litellm-rust/crates/cache/src/caching.rs | 10 +- litellm-rust/crates/cache/src/capabilities.rs | 143 +++++++++- litellm-rust/crates/cache/src/dual.rs | 255 +++++++++--------- litellm-rust/crates/cache/src/lib.rs | 13 +- litellm-rust/crates/cache/tests/caching.rs | 105 ++++++-- litellm-rust/crates/cache/tests/dual.rs | 215 +++++++++------ .../crates/python-bridge/src/cache/config.rs | 22 +- .../crates/python-bridge/src/cache/native.rs | 2 +- .../crates/python-bridge/src/cache/request.rs | 2 +- 21 files changed, 1136 insertions(+), 533 deletions(-) create mode 100644 litellm-rust/crates/cache-redis/src/topology.rs create mode 100644 litellm-rust/crates/cache/src/cache_type.rs 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) }