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)