From 8dc960c928614c2276c8f5e7cc8d1658009ba796 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:22:13 +0000 Subject: [PATCH 01/22] feat(cache): add SemanticCacheContext Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache/src/base_cache.rs | 25 +++++++++++++++++++++ litellm-rust/crates/cache/src/lib.rs | 2 +- litellm-rust/crates/cache/tests/caching.rs | 24 +++++++++++++++++++- 3 files changed, 49 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 8bd69ba5ad6..6f961798795 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -32,6 +32,31 @@ impl CacheContext for ExactCacheContext { } } +#[derive(Clone, Debug, Default, PartialEq)] +pub struct SemanticCacheContext { + pub input: Option, + pub messages: Vec, + pub metadata: serde_json::Map, + pub scope: Option, + pub ttl: Option, +} + +impl CacheContext for SemanticCacheContext { + fn ttl(&self) -> Option { + self.ttl + } + + fn with_ttl(&self, ttl: Option) -> Self { + Self { + input: self.input.clone(), + messages: self.messages.clone(), + metadata: self.metadata.clone(), + scope: self.scope.clone(), + ttl, + } + } +} + #[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum CacheConnectionStatus { diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index ce9f93b6dc4..8364c635e3a 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -8,7 +8,7 @@ mod error; pub use base_cache::{ BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext, - ExactCacheContext, + ExactCacheContext, SemanticCacheContext, }; pub use cache_type::CacheType; pub use caching::{Cache, CacheBackend, get_cache, set_cache}; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 9180ee9d0dc..2e65b4eeae5 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,7 +1,8 @@ use std::{sync::Mutex, time::Duration}; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache, + BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, + SemanticCacheContext, get_cache, }; struct TestCache { @@ -126,6 +127,27 @@ fn associated_context_preserves_backend_specific_lookup_inputs() { ); } +#[test] +fn semantic_context_with_ttl_preserves_lookup_inputs() { + let context = SemanticCacheContext { + input: Some(serde_json::json!("text")), + messages: vec![serde_json::json!({"role": "user", "content": "hi"})], + metadata: serde_json::Map::from_iter([( + "key".into(), + serde_json::json!("value"), + )]), + scope: Some("scope".into()), + ttl: None, + }; + let updated = context.with_ttl(Some(Duration::from_secs(30))); + assert_eq!(updated.ttl(), Some(Duration::from_secs(30))); + assert_eq!(updated.input, context.input); + assert_eq!(updated.messages, context.messages); + assert_eq!(updated.metadata, context.metadata); + assert_eq!(updated.scope, context.scope); + assert_eq!(context.with_ttl(None).ttl(), None); +} + #[tokio::test] async fn default_batch_operations_use_async_writes_and_stop_on_failure() { let cache = TestCache { From 1320eeeb41fbcd2a5842889e39f768d446fa6a05 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:23:29 +0000 Subject: [PATCH 02/22] refactor(cache-response): generalize ResponseCache over the backend context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-response/src/response.rs | 50 +++++++++++-------- 1 file changed, 30 insertions(+), 20 deletions(-) diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index e50e68cdabb..a27cc1967d5 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,21 +1,22 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache, + BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, + ExactCacheContext, FlushCache, }; use serde_json::Value; use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; #[derive(Clone)] -pub struct ResponseCacheRequest { +pub struct ResponseCacheRequest { pub key: CacheKeyInput, pub controls: CacheControls, - pub context: ExactCacheContext, + pub context: C, pub max_age: Option, } -impl ResponseCacheRequest { +impl ResponseCacheRequest { pub fn new(key: CacheKeyInput) -> Self { Self { key, @@ -32,11 +33,11 @@ impl ResponseCacheRequest { } } -pub struct ResponseCache> { +pub struct ResponseCache> { backend: Arc, } -impl> ResponseCache { +impl> ResponseCache { pub fn new(backend: Arc) -> Self { Self { backend } } @@ -45,8 +46,11 @@ impl> ResponseCach &self.backend } - pub fn default_ttl(&self) -> Option { - self.backend.get_ttl(&ExactCacheContext::default()) + pub fn default_ttl(&self) -> Option + where + B::Context: Default, + { + self.backend.get_ttl(&B::Context::default()) } pub async fn async_flush(&self) -> Result<(), Error> @@ -62,7 +66,7 @@ impl> ResponseCach pub fn lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -81,7 +85,7 @@ impl> ResponseCach pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -101,7 +105,7 @@ impl> ResponseCach pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -126,7 +130,7 @@ impl> ResponseCach pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -153,7 +157,7 @@ impl> ResponseCach pub fn store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -172,7 +176,7 @@ impl> ResponseCach pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -193,9 +197,12 @@ impl> ResponseCach pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(ResponseCacheRequest, Value)>, now: Duration, - ) -> Result<(), Error> { + ) -> Result<(), Error> + where + B::Context: PartialEq, + { self.async_store_entries( entries .into_iter() @@ -209,8 +216,11 @@ impl> ResponseCach /// the freshness of its original response. pub async fn async_store_entries( &self, - entries: Vec<(ResponseCacheRequest, Value, Duration)>, - ) -> Result<(), Error> { + entries: Vec<(ResponseCacheRequest, Value, Duration)>, + ) -> Result<(), Error> + where + B::Context: PartialEq, + { let writable = entries .into_iter() .filter(|(request, _, _)| request.controls.writes()) @@ -249,8 +259,8 @@ impl> ResponseCach } fn partial_hits( - requests: &[ResponseCacheRequest], - readable: Vec<(usize, &ResponseCacheRequest)>, + requests: &[ResponseCacheRequest], + readable: Vec<(usize, &ResponseCacheRequest)>, entries: Vec>, now: Duration, ) -> Result { From 8d9ab9eeaafe88efdf19fe67809e5fd21b2adb03 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:24:23 +0000 Subject: [PATCH 03/22] feat(cache-redis): expose the pooled connection handling for reuse Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache-redis/src/cache.rs | 93 +++++++++++-------- .../cache-redis/src/cache/operations.rs | 28 +++--- litellm-rust/crates/cache-redis/src/lib.rs | 4 + 3 files changed, 72 insertions(+), 53 deletions(-) diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index a960c383bf4..6388448accc 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -19,7 +19,7 @@ const DEFAULT_TTL: Duration = Duration::from_secs(600); const REDIS_TIMEOUT: Duration = Duration::from_secs(5); const REDIS_POOL_SIZE: u32 = 16; -struct PooledConnection { +pub struct PooledConnection { connection: redis::Connection, failed: bool, } @@ -27,16 +27,19 @@ struct PooledConnection { /// 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); +pub struct ConnectionManager { + client: redis::Client, + timeout: Duration, +} 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))?; + let connection = self.client.get_connection()?; + connection.set_read_timeout(Some(self.timeout))?; + connection.set_write_timeout(Some(self.timeout))?; Ok(PooledConnection { connection, failed: false, @@ -68,12 +71,12 @@ const CLAIM_SCRIPT: &str = concat!( ); const CLAIM_ATTEMPTS: usize = 8; -enum Connections { +pub enum Connections { Pool(r2d2::Pool), Fixed(Mutex), } -struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); +pub struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); impl redis::ConnectionLike for ConnectionRef<'_> { fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { @@ -110,7 +113,23 @@ impl Connections where C: redis::ConnectionLike + Send + 'static, { - fn execute( + pub fn pooled(url: &str, timeout: Duration, pool_size: u32) -> Result { + let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; + let pool = r2d2::Pool::builder() + .max_size(pool_size) + .min_idle(Some(0)) + .connection_timeout(timeout) + .test_on_check_out(false) + .build(ConnectionManager { client, timeout }) + .map_err(|_| Error::Unavailable)?; + Ok(Self::Pool(pool)) + } + + pub fn fixed(connection: C) -> Self { + Self::Fixed(Mutex::new(connection)) + } + + pub fn execute( &self, operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, ) -> Result { @@ -127,6 +146,16 @@ where } } } + + pub async fn run_blocking(connections: Arc, operation: F) -> Result + where + T: Send + 'static, + F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, + { + tokio::task::spawn_blocking(move || connections.execute(operation)) + .await + .map_err(|_| Error::Unavailable)? + } } pub struct RedisCache { @@ -138,16 +167,8 @@ 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 pool = r2d2::Pool::builder() - .max_size(REDIS_POOL_SIZE) - .min_idle(Some(0)) - .connection_timeout(REDIS_TIMEOUT) - .test_on_check_out(false) - .build(ConnectionManager(client)) - .map_err(|_| Error::Unavailable)?; Ok(Self { - connections: Arc::new(Connections::Pool(pool)), + connections: Arc::new(Connections::pooled(url, REDIS_TIMEOUT, REDIS_POOL_SIZE)?), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, namespace: None, @@ -162,7 +183,7 @@ where { pub fn with_connection(connection: C, default_ttl: Option, codec: S) -> Self { Self { - connections: Arc::new(Connections::Fixed(Mutex::new(connection))), + connections: Arc::new(Connections::fixed(connection)), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, namespace: None, @@ -241,20 +262,14 @@ where } fn ttl_seconds(ttl: Duration) -> u64 { - ttl.as_secs() - .saturating_add(u64::from(ttl.subsec_nanos() > 0)) - .max(1) + ttl_seconds(ttl) } +} - async fn run_blocking(connections: Arc>, operation: F) -> Result - where - T: Send + 'static, - F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, - { - tokio::task::spawn_blocking(move || connections.execute(operation)) - .await - .map_err(|_| Error::Unavailable)? - } +pub fn ttl_seconds(ttl: Duration) -> u64 { + ttl.as_secs() + .saturating_add(u64::from(ttl.subsec_nanos() > 0)) + .max(1) } fn namespaced_key(namespace: Option<&str>, key: &str) -> String { @@ -313,7 +328,7 @@ where let payload = self.codec.encode(&value)?; let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection .set_ex::<_, _, ()>(key, payload, ttl) .map_err(|_| Error::Unavailable) @@ -327,7 +342,7 @@ where _: &ExactCacheContext, ) -> Result, Error> { let key = self.namespaced_key(key); - let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection .get::<_, redis::Value>(key) .map_err(|_| Error::Unavailable) @@ -350,7 +365,7 @@ where }) .collect::, _>>()?; let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut pipeline = redis::pipe(); for (key, payload) in entries { pipeline @@ -372,7 +387,7 @@ where } async fn test_connection(&self) -> Result { - match Self::run_blocking(Arc::clone(&self.connections), |connection| { + match Connections::run_blocking(Arc::clone(&self.connections), |connection| { Ok(match redis::cmd("PING").query::(connection) { Ok(_) => CacheConnectionResult { status: CacheConnectionStatus::Success, @@ -433,7 +448,7 @@ where .iter() .map(|key| self.namespaced_key(key)) .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("MGET") .arg(keys) .query::>(connection) @@ -460,7 +475,7 @@ where 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| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) }) .await @@ -480,7 +495,7 @@ where async fn async_flush_cache(&self) -> Result<(), Error> { let pattern = self.namespaced_pattern()?; - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { Self::flush_matching(connection, &pattern) }) .await @@ -512,7 +527,7 @@ where ) -> Result { let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { increment(connection, key, amount, ttl) }) .await @@ -623,7 +638,7 @@ where let key = self.namespaced_key(key); 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| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { claim(connection, &codec, &key, candidate, &eligible, ttl) }) .await diff --git a/litellm-rust/crates/cache-redis/src/cache/operations.rs b/litellm-rust/crates/cache-redis/src/cache/operations.rs index d8d9ae24c4c..f27a7802bab 100644 --- a/litellm-rust/crates/cache-redis/src/cache/operations.rs +++ b/litellm-rust/crates/cache-redis/src/cache/operations.rs @@ -144,7 +144,7 @@ where .into_iter() .map(|key| self.namespaced_key(&key)) .collect::>(); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection.del(keys).map_err(|_| Error::Unavailable) }) .await @@ -172,7 +172,7 @@ where .iter() .map(|key| self.namespaced_key(key)) .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("MGET") .arg(keys) .query::>(connection) @@ -192,7 +192,7 @@ where } pub async fn ping(&self) -> Result { - Self::run_blocking(Arc::clone(&self.connections), |connection| { + Connections::run_blocking(Arc::clone(&self.connections), |connection| { redis::cmd("PING") .query::(connection) .map(|response| response == "PONG") @@ -203,7 +203,7 @@ where 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| { + let ttl = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("TTL") .arg(key) .query::(connection) @@ -215,7 +215,7 @@ where 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| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut cursor = 0u64; let mut matches = Vec::new(); loop { @@ -249,7 +249,7 @@ where } 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| { + Connections::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(); @@ -266,7 +266,7 @@ where return Err(Error::InvalidEntry); } let key = self.namespaced_key(key); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("RPUSH") .arg(key) .arg(values) @@ -292,7 +292,7 @@ where if operations.is_empty() { return Ok(Vec::new()); } - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::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); @@ -309,7 +309,7 @@ where ) -> Result { let key = self.namespaced_key(key); let multiple = count.is_some(); - let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut command = redis::cmd("LPOP"); command.arg(key); if let Some(count) = count { @@ -338,7 +338,7 @@ where .iter() .map(|(_, count)| count.is_some()) .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let values = Connections::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); @@ -368,7 +368,7 @@ where .into_iter() .map(|key| self.namespaced_key(&key)) .collect::>(); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("EVAL") .arg(script) .arg(keys.len()) @@ -440,7 +440,7 @@ where if operations.is_empty() { return Ok(Vec::new()); } - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::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); @@ -461,7 +461,7 @@ where ) -> Result { let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(ttl); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { increment_with_floor(connection, key, amount, ttl) }) .await @@ -475,7 +475,7 @@ where ) -> 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| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("EVAL") .arg(SET_MAX_SCRIPT) .arg(1) diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs index 98f6bfd8ce5..ea75906e9c9 100644 --- a/litellm-rust/crates/cache-redis/src/lib.rs +++ b/litellm-rust/crates/cache-redis/src/lib.rs @@ -1,6 +1,10 @@ mod cache; mod topology; +pub mod connection { + pub use crate::cache::{ConnectionRef, Connections, ttl_seconds}; +} + pub use cache::{ RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript, }; From 6ab121a3e7364178544bcbee529f6c57b3115a0a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:29:49 +0000 Subject: [PATCH 04/22] feat(cache-redis-semantic): add native Redis Semantic cache backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 15 + litellm-rust/Cargo.toml | 1 + .../crates/cache-redis-semantic/Cargo.toml | 21 + .../crates/cache-redis-semantic/src/cache.rs | 599 ++++++++++++++ .../crates/cache-redis-semantic/src/lib.rs | 4 + .../crates/cache-redis-semantic/src/prompt.rs | 95 +++ .../cache-redis-semantic/tests/cache.rs | 751 ++++++++++++++++++ litellm-rust/crates/cache/tests/caching.rs | 9 +- 8 files changed, 1489 insertions(+), 6 deletions(-) create mode 100644 litellm-rust/crates/cache-redis-semantic/Cargo.toml create mode 100644 litellm-rust/crates/cache-redis-semantic/src/cache.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/src/lib.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/src/prompt.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/tests/cache.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..e911d0d9c45 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2486,6 +2486,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "litellm-cache-redis", + "litellm-cache-response", + "r2d2", + "redis", + "redis-test", + "serde_json", + "sha2 0.10.9", + "tokio", +] + [[package]] name = "litellm-cache-response" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 570d0dd3568..05eea6bc299 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -29,6 +29,7 @@ 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-redis-semantic = { path = "crates/cache-redis-semantic" } litellm-cache-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } diff --git a/litellm-rust/crates/cache-redis-semantic/Cargo.toml b/litellm-rust/crates/cache-redis-semantic/Cargo.toml new file mode 100644 index 00000000000..9a8755a189e --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-cache-redis.workspace = true +litellm-cache-response.workspace = true +redis = { version = "1.7.0", features = ["tls-rustls"] } +r2d2 = "0.8.10" +serde_json.workspace = true +sha2.workspace = true +tokio.workspace = true + +[dev-dependencies] +redis-test = "1.0.4" +serde_json.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs new file mode 100644 index 00000000000..26d19d34670 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -0,0 +1,599 @@ +use std::{ + future::Future, + sync::{Arc, OnceLock}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_cache::{ + BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error, + SemanticCacheContext, +}; +use litellm_cache_redis::connection::{ConnectionRef, Connections, ttl_seconds}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +use crate::prompt::prompt_from_context; + +const REDIS_TIMEOUT: Duration = Duration::from_secs(5); +const REDIS_POOL_SIZE: u32 = 16; +const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; +const CACHE_KEY_FIELD: &str = "litellm_cache_key"; +const VECTOR_FIELD: &str = "prompt_vector"; + +pub trait Embedder: Send + Sync + 'static { + fn embed( + &self, + prompt: &str, + metadata: &serde_json::Map, + ) -> Result, Error>; + + fn async_embed( + &self, + prompt: &str, + metadata: &serde_json::Map, + ) -> impl Future, Error>> + Send; +} + +#[derive(Clone, Debug)] +pub struct RedisSemanticConfig { + pub index_name: String, + pub similarity_threshold: f32, +} + +impl Default for RedisSemanticConfig { + fn default() -> Self { + Self { + index_name: DEFAULT_INDEX_NAME.into(), + similarity_threshold: 0.9, + } + } +} + +struct Inner { + index_name: String, + distance_threshold: f64, + resolved_index: OnceLock, + codec: ResponseCacheCodec, + clock: fn() -> f64, +} + +impl Inner { + fn new(config: RedisSemanticConfig) -> Self { + Self { + index_name: config.index_name, + distance_threshold: 1.0 - f64::from(config.similarity_threshold), + resolved_index: OnceLock::new(), + codec: ResponseCacheCodec, + clock: timestamp, + } + } + + fn ensure_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + if let Some(name) = self.resolved_index.get() { + return Ok(name.clone()); + } + let name = match index_compatible(connection, &self.index_name, dims)? { + Some(true) => self.index_name.clone(), + Some(false) => self.isolated_index(connection, dims)?, + None => { + create_index(connection, &self.index_name, dims)?; + self.index_name.clone() + } + }; + let _ = self.resolved_index.set(name.clone()); + Ok(name) + } + + fn isolated_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + let name = format!("{}_isolated", self.index_name); + match index_compatible(connection, &name, dims)? { + Some(true) => Ok(name), + Some(false) => { + redis::cmd("FT.DROPINDEX") + .arg(&name) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + create_index(connection, &name, dims)?; + Ok(name) + } + None => { + create_index(connection, &name, dims)?; + Ok(name) + } + } + } + + fn store( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + value: &CacheEntry, + prompt: &str, + vector: &[f32], + ttl: Option, + ) -> Result<(), Error> { + let index = self.ensure_index(connection, vector.len())?; + let entry_id = entry_id(prompt, tag); + let hash_key = format!("{index}:{entry_id}"); + let response = self.codec.encode(value)?; + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(&entry_id) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg(VECTOR_FIELD) + .arg(vector_buffer(vector)) + .arg("inserted_at") + .arg(format!("{}", (self.clock)())) + .arg("updated_at") + .arg(format!("{}", (self.clock)())) + .arg(CACHE_KEY_FIELD) + .arg(tag) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + if let Some(ttl) = ttl { + redis::cmd("EXPIRE") + .arg(&hash_key) + .arg(ttl_seconds(ttl)) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + } + Ok(()) + } + + fn lookup( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + vector: &[f32], + ) -> Result, Error> { + let index = self.ensure_index(connection, vector.len())?; + let query = format!( + "(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]", + escape_tag(tag) + ); + let result = redis::cmd("FT.SEARCH") + .arg(&index) + .arg(query) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg(CACHE_KEY_FIELD) + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_buffer(vector)) + .query::(connection) + .map_err(|_| Error::Unavailable)?; + let Some(fields) = first_document(&result) else { + return Ok(None); + }; + if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) { + return Ok(None); + } + if number_field(fields, "vector_distance") + .is_none_or(|distance| distance > self.distance_threshold) + { + return Ok(None); + } + let Some(response) = bytes_field(fields, "response") else { + return Ok(None); + }; + self.codec.decode(&response).map(Some) + } +} + +pub struct RedisSemanticCache { + connections: Arc>, + embedder: E, + inner: Arc, +} + +impl RedisSemanticCache { + pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result { + Ok(Self { + connections: Arc::new(Connections::pooled(url, REDIS_TIMEOUT, REDIS_POOL_SIZE)?), + embedder, + inner: Arc::new(Inner::new(config)), + }) + } +} + +impl RedisSemanticCache { + pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self { + Self { + connections: Arc::new(Connections::fixed(connection)), + embedder, + inner: Arc::new(Inner::new(config)), + } + } + + pub fn with_clock(self, clock: fn() -> f64) -> Self { + Self { + inner: Arc::new(Inner { + index_name: self.inner.index_name.clone(), + distance_threshold: self.inner.distance_threshold, + resolved_index: OnceLock::new(), + codec: self.inner.codec, + clock, + }), + ..self + } + } + + fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str { + context.scope.as_deref().unwrap_or(key) + } +} + +impl BaseCache + for RedisSemanticCache +{ + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(()); + }; + let vector = self.embedder.embed(&prompt, &context.metadata)?; + let tag = Self::tag(key, context).to_string(); + self.connections.execute(|connection| { + self.inner + .store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self.embedder.embed(&prompt, &context.metadata)?; + let tag = Self::tag(key, context).to_string(); + self.connections + .execute(|connection| self.inner.lookup(connection, &tag, &vector)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(&context) else { + return Ok(()); + }; + let vector = self + .embedder + .async_embed(&prompt, &context.metadata) + .await?; + let tag = Self::tag(key, &context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + .await + } + + async fn async_get_cache( + &self, + key: &str, + context: &Self::Context, + ) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self + .embedder + .async_embed(&prompt, &context.metadata) + .await?; + let tag = Self::tag(key, context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.lookup(connection, &tag, &vector) + }) + .await + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + match Connections::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()), + }), + } + } +} + +fn timestamp() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or_default() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(CACHE_KEY_FIELD.as_bytes()); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn vector_buffer(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn escape_tag(value: &str) -> String { + value + .chars() + .flat_map(|ch| { + if matches!( + ch, + ',' | '.' + | '<' + | '>' + | '{' + | '}' + | '[' + | ']' + | '\\' + | '"' + | '\'' + | ':' + | ';' + | '!' + | '@' + | '#' + | '$' + | '%' + | '^' + | '&' + | '*' + | '(' + | ')' + | '-' + | '+' + | '=' + | '~' + | '|' + | '/' + | ' ' + | '?' + ) { + vec!['\\', ch] + } else { + vec![ch] + } + }) + .collect() +} + +fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> { + redis::cmd("FT.CREATE") + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg(VECTOR_FIELD) + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg(CACHE_KEY_FIELD) + .arg("TAG") + .arg("SEPARATOR") + .arg(",") + .query::<()>(connection) + .map_err(|_| Error::Unavailable) +} + +fn index_compatible( + connection: &mut ConnectionRef<'_>, + name: &str, + dims: usize, +) -> Result, Error> { + let info = match redis::cmd("FT.INFO") + .arg(name) + .query::(connection) + { + Ok(info) => info, + Err(error) if unknown_index(&error) => return Ok(None), + Err(_) => return Err(Error::Unavailable), + }; + Ok(Some(schema_compatible(&info, dims))) +} + +fn unknown_index(error: &redis::RedisError) -> bool { + let message = error.to_string().to_lowercase(); + message.contains("unknown") && message.contains("index") +} + +fn schema_compatible(info: &redis::Value, dims: usize) -> bool { + let redis::Value::Array(entries) = info else { + return false; + }; + let attributes = entries + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some("attributes")) + .map(|pair| &pair[1]); + let Some(redis::Value::Array(attributes)) = attributes else { + return false; + }; + let fields = attributes + .iter() + .map(|attribute| { + let redis::Value::Array(attribute) = attribute else { + return (None, None, None); + }; + let mut name = None; + let mut field_type = None; + let mut dim = None; + for pair in attribute.as_chunks::<2>().0 { + match string_value(&pair[0]).as_deref() { + Some("identifier") => name = string_value(&pair[1]), + Some("type") => field_type = string_value(&pair[1]), + Some("dim") => dim = number_value(&pair[1]), + _ => {} + } + } + (name, field_type, dim) + }) + .collect::>(); + let has_field = |name: &str, field_type: &str| { + fields + .iter() + .any(|(n, t, _)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type)) + }; + has_field("prompt", "TEXT") + && has_field("response", "TEXT") + && has_field("inserted_at", "NUMERIC") + && has_field("updated_at", "NUMERIC") + && has_field(CACHE_KEY_FIELD, "TAG") + && fields.iter().any(|(n, t, d)| { + n.as_deref() == Some(VECTOR_FIELD) + && t.as_deref() == Some("VECTOR") + && *d == Some(dims as f64) + }) +} + +fn string_value(value: &redis::Value) -> Option { + match value { + redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), + redis::Value::SimpleString(text) => Some(text.clone()), + redis::Value::VerbatimString { text, .. } => Some(text.clone()), + _ => None, + } +} + +fn number_value(value: &redis::Value) -> Option { + match value { + redis::Value::Int(number) => Some(*number as f64), + redis::Value::Double(number) => Some(*number), + _ => string_value(value).and_then(|text| text.parse().ok()), + } +} + +fn first_document(result: &redis::Value) -> Option<&[redis::Value]> { + let redis::Value::Array(items) = result else { + return None; + }; + let [count, _document_id, fields, ..] = items.as_slice() else { + return None; + }; + if !matches!(count, redis::Value::Int(count) if *count > 0) { + return None; + } + match fields { + redis::Value::Array(fields) => Some(fields.as_slice()), + _ => None, + } +} + +fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> { + fields + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some(name)) + .map(|pair| &pair[1]) +} + +fn string_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(string_value) +} + +fn number_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(number_value) +} + +fn bytes_field(fields: &[redis::Value], name: &str) -> Option> { + match field_value(fields, name)? { + redis::Value::BulkString(bytes) => Some(bytes.clone()), + redis::Value::SimpleString(text) => Some(text.clone().into_bytes()), + _ => None, + } +} diff --git a/litellm-rust/crates/cache-redis-semantic/src/lib.rs b/litellm-rust/crates/cache-redis-semantic/src/lib.rs new file mode 100644 index 00000000000..a34603cd18f --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/lib.rs @@ -0,0 +1,4 @@ +mod cache; +mod prompt; + +pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig}; diff --git a/litellm-rust/crates/cache-redis-semantic/src/prompt.rs b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs new file mode 100644 index 00000000000..fc99898d847 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs @@ -0,0 +1,95 @@ +use litellm_cache::SemanticCacheContext; +use serde_json::Value; + +pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { + if !context.messages.is_empty() { + return Some(messages_text(&context.messages)); + } + let input = context.input.as_ref()?; + let mut parts = Vec::new(); + collect_input_text(input, &mut parts); + let prompt = parts.join("\n").trim().to_string(); + (!prompt.is_empty()).then_some(prompt) +} + +fn messages_text(messages: &[Value]) -> String { + let mut text = String::new(); + for message in messages { + let Some(message) = message.as_object() else { + continue; + }; + match message.get("content") { + Some(Value::String(content)) => text.push_str(content), + Some(Value::Array(parts)) => { + for part in parts { + if let Some(text_content) = part.get("text").and_then(Value::as_str) { + text.push_str(text_content); + } + } + } + _ => {} + } + text.push_str(&search_results_text(message.get("search_results"))); + } + text +} + +fn search_results_text(search_results: Option<&Value>) -> String { + let Some(Value::Array(results)) = search_results else { + return String::new(); + }; + let mut text = String::new(); + for result in results { + let Some(result) = result.as_object() else { + continue; + }; + for key in ["source", "title"] { + if let Some(value) = result.get(key).and_then(Value::as_str) { + text.push_str(value); + } + } + if let Some(Value::Array(content)) = result.get("content") { + for block in content { + if let Some(value) = block.get("text").and_then(Value::as_str) { + text.push_str(value); + } + } + } + if let Some(citations) = result.get("citations") { + text.push_str(&citations.to_string()); + } + } + text +} + +fn collect_input_text(value: &Value, parts: &mut Vec) { + match value { + Value::String(text) => { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + } + } + Value::Array(items) => { + for item in items { + collect_input_text(item, parts); + } + } + Value::Object(map) => { + if let Some(content) = map.get("content").filter(|content| !content.is_null()) { + collect_input_text(content, parts); + return; + } + for key in ["text", "output", "input_text", "output_text"] { + if let Some(Value::String(text)) = map.get(key) { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + return; + } + } + } + } + _ => {} + } +} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs new file mode 100644 index 00000000000..77b057ae3b9 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -0,0 +1,751 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{BaseCache, CacheCodec, Error, SemanticCacheContext}; +use litellm_cache_redis_semantic::{Embedder, RedisSemanticCache, RedisSemanticConfig}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use redis_test::{MockCmd, MockRedisConnection}; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; + +const INDEX: &str = "litellm_semantic_cache_index"; + +struct FakeEmbedder { + vectors: HashMap>, + calls: Arc>>, +} + +impl FakeEmbedder { + fn new(vectors: &[(&str, &[f32])]) -> (Self, Arc>>) { + let calls = Arc::new(Mutex::new(Vec::new())); + ( + Self { + vectors: vectors + .iter() + .map(|(prompt, vector)| (prompt.to_string(), vector.to_vec())) + .collect(), + calls: Arc::clone(&calls), + }, + calls, + ) + } +} + +impl Embedder for FakeEmbedder { + fn embed(&self, prompt: &str, _: &serde_json::Map) -> Result, Error> { + self.calls.lock().unwrap().push(prompt.to_string()); + + Ok(self + .vectors + .get(prompt) + .cloned() + .unwrap_or_else(|| vec![0.1, 0.2, 0.3])) + } + + async fn async_embed( + &self, + prompt: &str, + metadata: &serde_json::Map, + ) -> Result, Error> { + self.embed(prompt, metadata) + } +} + +fn config() -> RedisSemanticConfig { + RedisSemanticConfig { + index_name: INDEX.into(), + similarity_threshold: 0.9, + } +} + +fn messages_context(messages: Vec) -> SemanticCacheContext { + SemanticCacheContext { + messages, + ..Default::default() + } +} + +fn entry() -> CacheEntry { + CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "yes"}), + } +} + +fn encoded(entry: &CacheEntry) -> Vec { + ResponseCacheCodec.encode(entry).unwrap() +} + +fn vector_bytes(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(b"litellm_cache_key"); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn s(value: &str) -> redis::Value { + redis::Value::BulkString(value.as_bytes().to_vec()) +} + +fn unknown_index_error() -> redis::RedisError { + redis::RedisError::from((redis::ErrorKind::Extension, "Unknown index name")) +} + +fn attribute(name: &str, field_type: &str, extra: Vec) -> redis::Value { + let mut parts = vec![ + s("identifier"), + s(name), + s("attribute"), + s(name), + s("type"), + s(field_type), + ]; + parts.extend(extra); + redis::Value::Array(parts) +} + +fn index_info(attributes: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("index_name"), + s(INDEX), + s("attributes"), + redis::Value::Array(attributes), + ]) +} + +fn vector_attribute(dims: i64) -> redis::Value { + attribute( + "prompt_vector", + "VECTOR", + vec![ + s("algorithm"), + s("FLAT"), + s("data_type"), + s("FLOAT32"), + s("dim"), + redis::Value::Int(dims), + s("distance_metric"), + s("COSINE"), + ], + ) +} + +fn compatible_info(dims: i64) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector_attribute(dims), + attribute("litellm_cache_key", "TAG", vec![]), + ]) +} + +fn unscoped_info(dims: i64) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector_attribute(dims), + ]) +} + +fn create_index_command(name: &str, dims: usize) -> redis::Cmd { + let mut command = redis::cmd("FT.CREATE"); + command + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg("prompt_vector") + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg("litellm_cache_key") + .arg("TAG") + .arg("SEPARATOR") + .arg(","); + command +} + +fn search_command(index: &str, tag: &str, vector: &[f32]) -> redis::Cmd { + let mut command = redis::cmd("FT.SEARCH"); + command + .arg(index) + .arg(format!( + "(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]" + )) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg("litellm_cache_key") + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_bytes(vector)); + command +} + +fn hit_fields(tag: &str, distance: &str, response: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("entry_id"), + s("stored-id"), + s("prompt"), + s("hello prompt"), + s("response"), + redis::Value::BulkString(response), + s("inserted_at"), + s("1700000000.5"), + s("updated_at"), + s("1700000000.5"), + s("litellm_cache_key"), + s(tag), + s("vector_distance"), + s(distance), + ]) +} + +fn search_result(fields: redis::Value) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::Int(1), + s("litellm_semantic_cache_index:stored-id"), + fields, + ]) +} + +fn empty_result() -> redis::Value { + redis::Value::Array(vec![redis::Value::Int(0)]) +} + +#[test] +fn store_creates_index_and_writes_hash_with_expire() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(INDEX, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(5), Ok(1)), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(5)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + cache.set_cache(tag, value, &context).unwrap(); +} + +#[test] +fn store_without_ttl_skips_expire() { + let prompt = "hello prompt"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "key1"))) + .arg("entry_id") + .arg(entry_id(prompt, "key1")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("key1"), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + "key1", + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn lookup_returns_hit_below_distance_threshold() { + let vector = vec![0.1f32, 0.2, 0.3]; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let hit = cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]), + ) + .unwrap(); + assert_eq!(hit, Some(value)); +} + +#[test] +fn lookup_misses_above_distance_threshold_and_on_tag_mismatch() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.5", encoded(&entry())))), + ), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "other", + "0.05", + encoded(&entry()), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + let context = messages_context(vec![json!({"role": "user", "content": "hello prompt"})]); + + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); +} + +#[test] +fn lookup_returns_invalid_entry_on_malformed_response() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "key1", + "0.05", + b"not json!".to_vec(), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[test] +fn missing_prompt_is_noop_and_never_embeds() { + let connection = MockRedisConnection::new(Vec::::new()).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let context = SemanticCacheContext::default(); + cache.set_cache("key1", entry(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert!(calls.lock().unwrap().is_empty()); +} + +#[test] +fn scope_overrides_key_as_filter_tag() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "scope-a"))) + .arg("entry_id") + .arg(entry_id(prompt, "scope-a")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("scope-a"), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, "scope\\-a", &vector), + Ok(search_result(hit_fields( + "scope-a", + "0.05", + encoded(&value), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = SemanticCacheContext { + scope: Some("scope-a".into()), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + + cache.set_cache("key1", value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(value)); +} + +#[test] +fn incompatible_schema_falls_back_to_isolated_index() { + let prompt = "hello prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(unscoped_info(3))), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{isolated}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn tag_special_characters_are_escaped_in_search_filter() { + let vector = vec![0.1f32, 0.2, 0.3]; + let tag = "a:b, c|d"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "a\\:b\\,\\ c\\|d", &vector), + Ok(empty_result()), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + tag, + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap(), + None + ); +} + +#[test] +fn prompt_extraction_matches_python_message_and_input_shapes() { + let vector = vec![0.1f32, 0.2, 0.3]; + let lookups = 5; + let mut commands = vec![MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Ok(compatible_info(3)), + )]; + for _ in 0..lookups { + commands.push(MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(empty_result()), + )); + } + let connection = MockRedisConnection::new(commands).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + cache + .get_cache( + "key1", + &messages_context(vec![ + json!({"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]}), + json!({"role": "assistant", "content": "reply"}), + ]), + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!(" plain input ")), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some( + json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"]), + ), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!({"output_text": " result text "})), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &messages_context(vec![json!({ + "role": "user", + "content": "question", + "search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}], + })]), + ) + .unwrap(); + + assert_eq!( + *calls.lock().unwrap(), + vec![ + "firstsecondreply", + "plain input", + "nested\ntail", + "result text", + "questionsrctfound{\"a\":1}", + ] + ); +} + +#[test] +fn ttl_passes_through_context_only() { + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection( + MockRedisConnection::new(Vec::::new()), + embedder, + config(), + ); + assert_eq!(cache.get_ttl(&SemanticCacheContext::default()), None); + assert_eq!( + cache.get_ttl(&SemanticCacheContext { + ttl: Some(Duration::from_secs(9)), + ..Default::default() + }), + Some(Duration::from_secs(9)) + ); +} + +#[tokio::test] +async fn async_paths_embed_then_run_blocking_redis_work() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, tag, &vector), + Ok(search_result(hit_fields(tag, "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = messages_context(vec![json!({"role": "user", "content": prompt})]); + + cache + .async_set_cache(tag, value.clone(), context.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache(tag, &context).await.unwrap(), + Some(value) + ); +} + +#[test] +fn live_store_lookup_and_ttl_against_redis_stack() { + let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { + return; + }; + let vector = vec![0.1f32, 0.2, 0.3, 0.4]; + let prompt = "rust semantic cache live prompt"; + let tag = "live-key"; + let index_name = format!("rust_semantic_test_{}", std::process::id()); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::new( + &url, + embedder, + RedisSemanticConfig { + index_name: index_name.clone(), + similarity_threshold: 0.9, + }, + ) + .unwrap(); + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(120)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + let value = entry(); + + cache.set_cache(tag, value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache(tag, &context).unwrap(), Some(value)); + assert_eq!(cache.get_cache("other-key", &context).unwrap(), None); + + let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); + let ttl: i64 = redis::Commands::ttl( + &mut connection, + format!("{index_name}:{}", entry_id(prompt, tag)), + ) + .unwrap(); + assert!( + ttl > 0, + "expected stored hash to carry an expiry, got {ttl}" + ); +} diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 2e65b4eeae5..33171f36f46 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,8 +1,8 @@ use std::{sync::Mutex, time::Duration}; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, - SemanticCacheContext, get_cache, + BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, SemanticCacheContext, + get_cache, }; struct TestCache { @@ -132,10 +132,7 @@ fn semantic_context_with_ttl_preserves_lookup_inputs() { let context = SemanticCacheContext { input: Some(serde_json::json!("text")), messages: vec![serde_json::json!({"role": "user", "content": "hi"})], - metadata: serde_json::Map::from_iter([( - "key".into(), - serde_json::json!("value"), - )]), + metadata: serde_json::Map::from_iter([("key".into(), serde_json::json!("value"))]), scope: Some("scope".into()), ttl: None, }; From a69ebb7ac4584e3a91497e44fd65b4bf86d52815 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:34:39 +0000 Subject: [PATCH 05/22] feat(cache): add the unsupported operation error Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache/src/error.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index ff3ff6572d4..2418ab978dc 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -6,4 +6,6 @@ pub enum Error { InvalidEntry, #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, + #[error("cache backend does not support this operation")] + UnsupportedOperation, } From c311073a178d295a5133d2e68300c9f5f83664bc Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:34:51 +0000 Subject: [PATCH 06/22] feat(cache-redis-semantic): expose backend accessors for bridge binding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-redis-semantic/src/cache.rs | 21 +++++++++++-------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs index 26d19d34670..5aa484356d7 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -41,15 +41,6 @@ pub struct RedisSemanticConfig { pub similarity_threshold: f32, } -impl Default for RedisSemanticConfig { - fn default() -> Self { - Self { - index_name: DEFAULT_INDEX_NAME.into(), - similarity_threshold: 0.9, - } - } -} - struct Inner { index_name: String, distance_threshold: f64, @@ -247,6 +238,18 @@ impl RedisSemanticCache< } } + pub fn embedder(&self) -> &E { + &self.embedder + } + + pub fn index_name(&self) -> &str { + &self.inner.index_name + } + + pub fn similarity_threshold(&self) -> f32 { + (1.0 - self.inner.distance_threshold) as f32 + } + fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str { context.scope.as_deref().unwrap_or(key) } From ac1c4a399ff2dcdf82ce82b93ef95ac748d1d01c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:34:57 +0000 Subject: [PATCH 07/22] refactor(cache-redis-semantic): drop the unused default index name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache-redis-semantic/src/cache.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs index 5aa484356d7..b1440e80de5 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -17,7 +17,6 @@ use crate::prompt::prompt_from_context; const REDIS_TIMEOUT: Duration = Duration::from_secs(5); const REDIS_POOL_SIZE: u32 = 16; -const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; const CACHE_KEY_FIELD: &str = "litellm_cache_key"; const VECTOR_FIELD: &str = "prompt_vector"; From 10b977fe29caccc1a2730568d33c21aa751cddab Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:38:33 +0000 Subject: [PATCH 08/22] feat(python-bridge): serve redis-semantic caches natively Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/cache/config.rs | 67 +++++++- .../python-bridge/src/cache/embedder.rs | 85 +++++++++ .../crates/python-bridge/src/cache/facade.rs | 21 +++ .../crates/python-bridge/src/cache/handle.rs | 43 ++++- .../crates/python-bridge/src/cache/mod.rs | 4 +- .../crates/python-bridge/src/cache/native.rs | 161 ++++++++++++++---- .../crates/python-bridge/src/cache/request.rs | 66 +++++-- 9 files changed, 399 insertions(+), 50 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/cache/embedder.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e911d0d9c45..0030018df34 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2683,6 +2683,7 @@ dependencies = [ "litellm-cache", "litellm-cache-memory", "litellm-cache-redis", + "litellm-cache-redis-semantic", "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1eb2ec28036..635c0942ceb 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -23,6 +23,7 @@ bytes.workspace = true litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true +litellm-cache-redis-semantic.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 0e7d6aee11d..85e400d89a3 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -73,9 +73,23 @@ pub(super) struct RedisCacheConfig { pub(super) connection: RedisConnectionConfig, } +#[allow( + dead_code, + reason = "embedding settings are projected so drift falls back to Python" +)] +pub(super) struct RedisSemanticCacheConfig { + pub(super) redis_url: String, + pub(super) index_name: String, + pub(super) similarity_threshold: f64, + pub(super) embedding_model: String, + pub(super) embedding_max_input_tokens: Option, + pub(super) embedding_timeout: Option, +} + pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + RedisSemantic(Box), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -142,9 +156,14 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::RedisSemantic) => project_redis_semantic(&backend).map(|backend| { + CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::RedisSemantic(Box::new(backend)), + })) + }), Some( - CacheType::RedisSemantic - | CacheType::ValkeySemantic + CacheType::ValkeySemantic | CacheType::S3 | CacheType::Disk | CacheType::QdrantSemantic @@ -159,10 +178,11 @@ impl NativeCacheConfig { pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { if service.default_ttl() - != Some(match &self.backend { - CacheBackendConfig::Memory(config) => config.default_ttl, - CacheBackendConfig::Redis(config) => config.default_ttl, - }) + != match &self.backend { + CacheBackendConfig::Memory(config) => Some(config.default_ttl), + CacheBackendConfig::Redis(config) => Some(config.default_ttl), + CacheBackendConfig::RedisSemantic(_) => None, + } { return Some("facade and native backend default TTLs must match"); } @@ -185,10 +205,45 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), + CacheBackendConfig::RedisSemantic(_) if service.kind() != "redis_semantic" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::RedisSemantic(config) + if service.index_name() != Some(config.index_name.as_str()) => + { + Some("facade and native backend index names must match") + } + CacheBackendConfig::RedisSemantic(config) + if service.similarity_threshold() != Some(config.similarity_threshold as f32) => + { + Some("facade and native backend similarity thresholds must match") + } + CacheBackendConfig::RedisSemantic(_) => None, } } } +#[inline(never)] +pub(super) fn project_redis_semantic( + backend: &Bound<'_, PyAny>, +) -> PyResult { + Ok(RedisSemanticCacheConfig { + redis_url: backend.getattr("_redis_url")?.extract::()?, + index_name: backend + .getattr("_index_name")? + .extract::>()? + .unwrap_or_else(|| "litellm_semantic_cache_index".into()), + similarity_threshold: backend.getattr("similarity_threshold")?.extract::()?, + embedding_model: backend.getattr("embedding_model")?.extract::()?, + embedding_max_input_tokens: backend + .getattr("embedding_max_input_tokens")? + .extract::>()?, + embedding_timeout: backend + .getattr("embedding_timeout")? + .extract::>()?, + }) +} + #[inline(never)] fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { let max_size_kib = backend.getattr("max_size_per_item")?.extract::()?; diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs new file mode 100644 index 00000000000..63e078cd815 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -0,0 +1,85 @@ +use std::future::Future; + +use litellm_cache::Error; +use litellm_cache_redis_semantic::Embedder; +use litellm_host_python::to_py; +use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict}; +use serde_json::{Map, Value}; + +pub(super) struct PythonEmbedder(Py); + +impl PythonEmbedder { + pub(super) fn new(object: Py) -> Self { + Self(object) + } + + pub(super) fn object(&self) -> &Py { + &self.0 + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } + + fn metadata_kwargs<'py>( + py: Python<'py>, + metadata: &Map, + ) -> PyResult> { + let kwargs = PyDict::new(py); + if metadata.is_empty() { + kwargs.set_item("metadata", py.None())?; + } else { + kwargs.set_item("metadata", to_py(py, metadata)?)?; + } + Ok(kwargs) + } + + fn extract(vector: Bound<'_, PyAny>) -> PyResult> { + Ok(vector + .extract::>()? + .into_iter() + .map(|value| value as f32) + .collect()) + } +} + +impl Embedder for PythonEmbedder { + fn embed(&self, prompt: &str, metadata: &Map) -> Result, Error> { + Python::attach(|py| { + let kwargs = Self::metadata_kwargs(py, metadata)?; + Self::extract(self.0.bind(py).call_method( + "_get_embedding", + (prompt,), + Some(&kwargs), + )?) + }) + .map_err(|_| Error::Unavailable) + } + + fn async_embed( + &self, + prompt: &str, + metadata: &Map, + ) -> impl Future, Error>> + Send { + let coroutine = Python::attach(|py| { + let kwargs = Self::metadata_kwargs(py, metadata)?; + self.0 + .bind(py) + .call_method("_get_async_embedding", (prompt,), Some(&kwargs)) + .map(Bound::unbind) + }) + .map_err(|_| Error::Unavailable); + async move { + let coroutine = coroutine?; + let awaited = Python::attach(|py| { + pyo3_async_runtimes::tokio::into_future(coroutine.into_bound(py)) + }) + .map_err(|_| Error::Unavailable)? + .await + .map_err(|_| Error::Unavailable)?; + let vector = Python::attach(|py| awaited.extract::>(py)) + .map_err(|_| Error::Unavailable)?; + Ok(vector.into_iter().map(|value| value as f32).collect()) + } + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index f2f86c14b37..58730857d60 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -192,6 +192,11 @@ impl FacadeGuard { let (module, name, cache_kind) = match kind { "memory" => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), "redis" => ("litellm.caching.redis_cache", "RedisCache", "redis"), + "redis_semantic" => ( + "litellm.caching.redis_semantic_cache", + "RedisSemanticCache", + "redis-semantic", + ), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -211,6 +216,15 @@ impl FacadeGuard { if let Some(message) = config.service_mismatch(service) { return Err(PyTypeError::new_err(message)); } + if kind == "redis_semantic" + && service + .embedder_object() + .is_none_or(|embedder| !backend.is(embedder.bind(py))) + { + return Err(PyTypeError::new_err( + "facade backend must be the native embedder", + )); + } Ok(Self { outer: ObjectGuard::capture( py, @@ -235,6 +249,13 @@ impl FacadeGuard { "max_size_per_item", "redis_kwargs", "redis_flush_size", + "similarity_threshold", + "distance_threshold", + "embedding_model", + "embedding_max_input_tokens", + "embedding_timeout", + "_index_name", + "_redis_url", ], )?, redis_pool: (kind == "redis") diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 8251b3df06c..b61ae59bb58 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,7 +1,15 @@ +use litellm_cache_redis_semantic::RedisSemanticConfig; use litellm_host_python::release_gil; -use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyRuntimeError, PyTypeError}, + prelude::*, +}; -use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; +use super::{ + cache_error, config::project_redis_semantic, embedder::PythonEmbedder, facade::FacadeGuard, + native::NativeResponseCache, request::duration, +}; #[pyclass(frozen, name = "_CacheTestHandle")] pub(crate) struct CacheTestHandle { @@ -51,6 +59,36 @@ impl CacheTestHandle { }) } + #[staticmethod] + fn redis_semantic(py: Python<'_>, backend: Bound<'_, PyAny>) -> PyResult { + let class = py + .import("litellm.caching.redis_semantic_cache")? + .getattr("RedisSemanticCache")?; + if !backend.get_type().is(&class) { + return Err(PyTypeError::new_err( + "native redis-semantic handles require the built-in RedisSemanticCache", + )); + } + let config = project_redis_semantic(&backend)?; + let embedder = PythonEmbedder::new(backend.unbind()); + let service = release_gil(py, move || { + NativeResponseCache::redis_semantic( + &config.redis_url, + embedder, + RedisSemanticConfig { + index_name: config.index_name, + similarity_threshold: config.similarity_threshold as f32, + }, + ) + }) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() @@ -76,6 +114,7 @@ impl CacheTestHandle { } fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + self.service.traverse(&visit)?; if let Some(guard) = &self.guard { guard.traverse(visit)?; } diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index aec08610f6e..4cc87367d91 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,6 +1,7 @@ mod binding; mod callback; mod config; +mod embedder; mod facade; mod future; mod handle; @@ -10,7 +11,7 @@ mod resolver; use litellm_cache::Error; use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, + exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, prelude::*, }; @@ -21,6 +22,7 @@ pub(crate) use self::{ fn cache_error(error: Error) -> PyErr { match error { Error::InvalidEntry => PyValueError::new_err(error.to_string()), + Error::UnsupportedOperation => PyNotImplementedError::new_err(error.to_string()), _ => PyRuntimeError::new_err(error.to_string()), } } diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index a9475429e45..8cd77fa8eb0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,13 +1,17 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; +use litellm_cache::{CacheCodec, CacheConnectionResult, Error, ExactCacheContext}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; +use litellm_cache_redis_semantic::{RedisSemanticCache, RedisSemanticConfig}; use litellm_cache_response::{ CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, }; +use pyo3::{Py, PyAny, PyTraverseError, PyVisit}; use serde_json::Value; +use super::{embedder::PythonEmbedder, request::CacheRequest}; + #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), @@ -15,6 +19,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + RedisSemantic(Arc>>), } impl NativeResponseCache { @@ -43,6 +48,17 @@ impl NativeResponseCache { buffer: None, }) } + + pub fn redis_semantic( + url: &str, + embedder: PythonEmbedder, + config: RedisSemanticConfig, + ) -> Result { + let backend = RedisSemanticCache::new(url, embedder, config)?; + Ok(Self::RedisSemantic(Arc::new(ResponseCache::new(Arc::new( + backend, + ))))) + } } impl NativeResponseCache { @@ -50,6 +66,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", + Self::RedisSemantic(_) => "redis_semantic", } } @@ -57,12 +74,13 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), + Self::RedisSemantic(cache) => cache.default_ttl(), } } pub fn namespace(&self) -> Option<&str> { match self { - Self::Memory(_) => None, + Self::Memory(_) | Self::RedisSemantic(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), } } @@ -70,110 +88,189 @@ impl NativeResponseCache { pub fn capacity(&self) -> Option { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } => None, + Self::Redis { .. } | Self::RedisSemantic(_) => None, } } pub fn max_entry_bytes(&self) -> Option { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), - Self::Redis { .. } => None, + Self::Redis { .. } | Self::RedisSemantic(_) => None, } } + pub fn index_name(&self) -> Option<&str> { + match self { + Self::RedisSemantic(cache) => Some(cache.backend().index_name()), + Self::Memory(_) | Self::Redis { .. } => None, + } + } + + pub fn similarity_threshold(&self) -> Option { + match self { + Self::RedisSemantic(cache) => Some(cache.backend().similarity_threshold()), + Self::Memory(_) | Self::Redis { .. } => None, + } + } + + pub fn embedder_object(&self) -> Option<&Py> { + match self { + Self::RedisSemantic(cache) => Some(cache.backend().embedder().object()), + Self::Memory(_) | Self::Redis { .. } => None, + } + } + + pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Self::RedisSemantic(cache) = self { + cache.backend().embedder().traverse(visit)?; + } + Ok(()) + } + 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(WriteBuffer::new(flush_size))), }, - memory => memory, + other => other, } } - pub fn lookup( - &self, - request: &ResponseCacheRequest, - now: Duration, - ) -> Result, Error> { + fn exact_requests(requests: &[CacheRequest]) -> Vec> { + requests.iter().map(CacheRequest::exact).collect() + } + + pub fn lookup(&self, request: &CacheRequest, now: Duration) -> Result, Error> { match self { - Self::Memory(cache) => cache.lookup(request, now), - Self::Redis { cache, .. } => cache.lookup(request, now), + Self::Memory(cache) => cache.lookup(&request.exact(), now), + Self::Redis { cache, .. } => cache.lookup(&request.exact(), now), + Self::RedisSemantic(cache) => cache.lookup(&request.semantic(), now), } } pub fn store( &self, - request: &ResponseCacheRequest, + request: &CacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.store(request, response, now), - Self::Redis { cache, .. } => cache.store(request, response, now), + Self::Memory(cache) => cache.store(&request.exact(), response, now), + Self::Redis { cache, .. } => cache.store(&request.exact(), response, now), + Self::RedisSemantic(cache) => cache.store(&request.semantic(), response, now), } } pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[CacheRequest], now: Duration, ) -> Result { match self { - Self::Memory(cache) => cache.lookup_batch(requests, now), - Self::Redis { cache, .. } => cache.lookup_batch(requests, now), + Self::Memory(cache) => cache.lookup_batch(&Self::exact_requests(requests), now), + Self::Redis { cache, .. } => cache.lookup_batch(&Self::exact_requests(requests), now), + Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), } } pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &CacheRequest, now: Duration, ) -> Result, Error> { match self { - Self::Memory(cache) => cache.async_lookup(request, now).await, - Self::Redis { cache, .. } => cache.async_lookup(request, now).await, + Self::Memory(cache) => cache.async_lookup(&request.exact(), now).await, + Self::Redis { cache, .. } => cache.async_lookup(&request.exact(), now).await, + Self::RedisSemantic(cache) => cache.async_lookup(&request.semantic(), now).await, } } pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &CacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_store(request, response, now).await, + Self::Memory(cache) => cache.async_store(&request.exact(), response, now).await, Self::Redis { cache, buffer: None, - } => cache.async_store(request, response, now).await, + } => cache.async_store(&request.exact(), response, now).await, Self::Redis { cache, buffer: Some(buffer), - } => buffer.async_store(cache, request, response, now).await, + } => { + buffer + .async_store(cache, &request.exact(), response, now) + .await + } + Self::RedisSemantic(cache) => { + cache.async_store(&request.semantic(), response, now).await + } } } pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[CacheRequest], 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, + Self::Memory(cache) => { + cache + .async_lookup_batch(&Self::exact_requests(requests), now) + .await + } + Self::Redis { cache, .. } => { + cache + .async_lookup_batch(&Self::exact_requests(requests), now) + .await + } + Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), } } pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(CacheRequest, 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, + Self::Memory(cache) => { + cache + .async_store_batch( + entries + .into_iter() + .map(|(request, value)| (request.exact(), value)) + .collect(), + now, + ) + .await + } + Self::Redis { cache, .. } => { + cache + .async_store_batch( + entries + .into_iter() + .map(|(request, value)| (request.exact(), value)) + .collect(), + now, + ) + .await + } + Self::RedisSemantic(cache) => { + cache + .async_store_batch( + entries + .into_iter() + .map(|(request, value)| (request.semantic(), value)) + .collect(), + now, + ) + .await + } } } @@ -186,6 +283,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), } } @@ -193,6 +291,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), } } } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 0c5343a63d0..26e0fe4e62c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -1,9 +1,11 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use litellm_cache::{ExactCacheContext, SemanticCacheContext}; use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; use serde::Deserialize; +use serde_json::{Map, Value}; #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -12,24 +14,68 @@ struct RequestInput { controls: Option, ttl_seconds: Option, max_age_seconds: Option, + input: Option, + messages: Option>, + metadata: Option>, + scope: Option, } -pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { +pub(super) struct CacheRequest { + key: CacheKeyInput, + controls: CacheControls, + ttl: Option, + max_age: Option, + input: Option, + messages: Vec, + metadata: Map, + scope: Option, +} + +impl CacheRequest { + pub(super) fn exact(&self) -> ResponseCacheRequest { + let mut request = ResponseCacheRequest::new(self.key.clone()); + request.controls = self.controls; + request.context.ttl = self.ttl; + request.max_age = self.max_age; + request + } + + pub(super) fn semantic(&self) -> ResponseCacheRequest { + ResponseCacheRequest { + key: self.key.clone(), + controls: self.controls, + context: SemanticCacheContext { + input: self.input.clone(), + messages: self.messages.clone(), + metadata: self.metadata.clone(), + scope: self.scope.clone(), + ttl: self.ttl, + }, + max_age: self.max_age, + } + } +} + +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.context.ttl = input.ttl_seconds.map(duration).transpose()?; - request.max_age = input.max_age_seconds.map(duration).transpose()?; - Ok(request) +fn request_input(input: RequestInput) -> PyResult { + let defaults = ResponseCacheRequest::::new(input.key.clone()); + Ok(CacheRequest { + key: input.key, + controls: input.controls.unwrap_or(defaults.controls), + ttl: input.ttl_seconds.map(duration).transpose()?, + max_age: input.max_age_seconds.map(duration).transpose()?, + input: input.input, + messages: input.messages.unwrap_or_default(), + metadata: input.metadata.unwrap_or_default(), + scope: input.scope, + }) } -pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { +pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { from_py::>(value)? .into_iter() .map(request_input) From 25af094e27ba4b40ceabaccbae944528a807c3cf Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:52:28 +0000 Subject: [PATCH 09/22] fix(python-bridge): allow instance attributes to shadow class defaults Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/python-bridge/src/cache/facade.rs | 5 +- litellm/rust_bridge/_native.pyi | 64 ++++++++++++++++++- 2 files changed, 67 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 58730857d60..d7ec2052dd0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -117,7 +117,10 @@ impl ObjectGuard { return Ok(false); } for (name, value) in &expected.attributes { - if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) { + if !attributes.get_item(name)?.is(value.bind(py)) { + return Ok(false); + } + if instance.contains(name)? && value.bind(py).is_callable() { return Ok(false); } } diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 05a6df6d5af..7eb266d5a09 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,6 +1,6 @@ from asyncio import Future from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence -from typing import Never, final +from typing import Literal, Never, final from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest @@ -93,6 +93,68 @@ 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 = 60.0, + namespace: str | None = None, + ) -> _CacheTestHandle: ... + @staticmethod + def redis_semantic(backend: object) -> _CacheTestHandle: ... + @property + def backend(self) -> Literal["memory", "redis", "redis_semantic"]: ... + def _bind_facade(self, facade: object) -> None: ... + +@final +class _CacheTestBinding: + def __new__(cls, _uninstantiable: Never, /) -> Never: ... + @property + def kind(self) -> Literal["disabled", "native", "python_callback"]: ... + def lookup( + self, request: object, *, callback_kwargs: object = None + ) -> object: ... + def store( + self, request: object, response: object, *, callback_kwargs: object = None + ) -> None: ... + def lookup_batch( + self, requests: object, *, callback_kwargs: object = None + ) -> object: ... + def async_lookup( + self, request: object, *, callback_kwargs: object = None + ) -> Future[object]: ... + def async_store( + self, request: object, response: object, *, callback_kwargs: object = None + ) -> Future[object]: ... + def async_lookup_batch( + self, requests: object, *, callback_kwargs: object = None + ) -> Future[object]: ... + def async_store_batch( + self, + requests: object, + responses: object, + *, + callback_result: object = None, + callback_kwargs: object = None, + ) -> Future[object]: ... + def async_flush(self) -> Future[object]: ... + def ping(self) -> Future[object]: ... + +@final +class _CacheTestResolver: + def __new__(cls, namespace: object) -> _CacheTestResolver: ... + def resolve(self) -> _CacheTestBinding: ... + @final class TokenCounter: def __new__(cls, tokenizer_json: str) -> TokenCounter: ... From 4a0151f77fc3d8b68606a59171412fbe55b4015e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:52:28 +0000 Subject: [PATCH 10/22] test(rust): add redis-semantic native parity fixtures Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm_rust/test_cache.py | 549 ++++++++++++++++++++++++-- 1 file changed, 518 insertions(+), 31 deletions(-) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index c35cb1a20fb..e9f4b99a3b7 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -1,14 +1,19 @@ import asyncio import contextvars import gc +import hashlib import json +import math +import os import threading import time import weakref -from collections.abc import Generator +from collections.abc import Callable, Generator +from contextlib import ExitStack from types import SimpleNamespace from typing import Final, Protocol, cast from urllib.parse import urlparse +from uuid import uuid4 import fakeredis import pytest @@ -17,10 +22,16 @@ import redis import litellm from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_semantic_cache import RedisSemanticCache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType +from litellm.types.llms.custom_llm import CustomLLMItem +from litellm.types.utils import EmbeddingResponse from tests.test_litellm_rust.support.isolation import rebound +_CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name +_CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name + pytestmark: Final = pytest.mark.requires_rust_extension @@ -50,14 +61,14 @@ 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._CacheTestResolver(litellm) + resolver: Final = _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) + resolver: Final = _CacheTestResolver(litellm) enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30) enabled: Final = litellm.cache @@ -80,13 +91,13 @@ def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> Non async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: - namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.memory()) - resolver: Final = _native._CacheTestResolver(namespace) + namespace: Final = SimpleNamespace(cache=_CacheTestHandle.memory()) + resolver: Final = _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._CacheTestHandle.memory()): + with rebound(namespace, "cache", _CacheTestHandle.memory()): replacement: Final = resolver.resolve() await selected.async_store(request(), {"answer": 2}) assert replacement.lookup(request()) is None @@ -119,7 +130,7 @@ async def test_python_callback_preserves_identity_caller_task_context_and_errors raise failure namespace: Final = SimpleNamespace(cache=CustomCache()) - binding: Final = _native._CacheTestResolver(namespace).resolve() + binding: Final = _CacheTestResolver(namespace).resolve() assert binding.kind == "python_callback" assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel assert context.get() == "callback" @@ -140,7 +151,7 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: finally: finished.set() - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() async def lookup() -> object: return await binding.async_lookup(None, callback_kwargs={}) @@ -155,9 +166,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._CacheTestHandle.memory() + handle: Final = _CacheTestHandle.memory() handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) native: Final = resolver.resolve() assert native.kind == "native" native.store(request(), {"source": "native"}) @@ -188,12 +199,12 @@ def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not class CustomCache(Cache): pass - handle: Final = _native._CacheTestHandle.memory() + handle: Final = _CacheTestHandle.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._CacheTestResolver(SimpleNamespace(cache=facade)) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) with rebound(facade, "cache", InMemoryCache()): assert resolver.resolve().kind == "python_callback" with rebound(facade, "ttl", 12): @@ -218,7 +229,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._CacheTestResolver(namespace).resolve() + binding: Final = _CacheTestResolver(namespace).resolve() setattr(callback, "binding", binding) return weakref.ref(callback) @@ -229,8 +240,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._CacheTestHandle.redis(redis_url, namespace="team")) - binding: Final = _native._CacheTestResolver(namespace).resolve() + namespace: Final = SimpleNamespace(cache=_CacheTestHandle.redis(redis_url, namespace="team")) + binding: Final = _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)) @@ -252,33 +263,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._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=_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._CacheTestHandle.memory(ttl_seconds=-1) + _CacheTestHandle.memory(ttl_seconds=-1) async def test_memory_size_policy_is_applied_by_the_native_host() -> None: - handle: Final = _native._CacheTestHandle.memory(capacity=2, max_entry_bytes=128) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + handle: Final = _CacheTestHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = _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._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.memory(capacity=0)) + disabled: Final = _CacheTestResolver( + SimpleNamespace(cache=_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._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() requests: Final = [request("hit"), request("miss"), request("disabled")] requests[2]["controls"] = { "supported_call_type": True, @@ -316,7 +327,7 @@ async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: ) -> object: return result, kwargs - binding: Final = _native._CacheTestResolver( + binding: Final = _CacheTestResolver( SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL)) ).resolve() assert binding.kind == "python_callback" @@ -346,7 +357,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) cache.cache.set_cache("key", "value") - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=cache)).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=cache)).resolve() assert binding.kind == "python_callback" setattr(cache.cache, "ping", ping) @@ -358,7 +369,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> 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) + _CacheTestHandle.memory(capacity=7)._bind_facade(facade) async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: @@ -371,19 +382,19 @@ 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._CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) + _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, ttl_seconds=60)._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + _CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) + _CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) + binding: Final = _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" + assert _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" + assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" await binding.async_store(request("first"), {"value": 1}) assert client.get("first") is None @@ -393,3 +404,479 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: assert client.get("second") is not None await facade.cache.disconnect() client.close() + + +PARAPHRASE_MARKER: Final = " (paraphrase)" +SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic" +SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_" + + +def _normalized(vector: list[float]) -> list[float]: + norm: Final = math.sqrt(sum(component * component for component in vector)) + return [component / norm for component in vector] + + +def _base_embedding(prompt: str) -> list[float]: + digest: Final = hashlib.sha256(prompt.encode("utf-8")).digest() + return _normalized([float(digest[index] + 1) for index in range(8)]) + + +def _semantic_embedding(prompt: str) -> list[float]: + if PARAPHRASE_MARKER not in prompt: + return _base_embedding(prompt) + base: Final = _base_embedding(prompt.replace(PARAPHRASE_MARKER, "").strip()) + pivot: Final = min(range(8), key=lambda index: abs(base[index])) + direction: Final = _normalized( + [ + (1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot] + for index in range(8) + ] + ) + # Rotating an orthogonal unit direction by 0.329 produces ~0.05 cosine distance + return _normalized([base[index] + 0.329 * direction[index] for index in range(8)]) + + +class DeterministicEmbedding(litellm.CustomLLM): + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + + def _respond( + self, + model: str, + input: object, + model_response: EmbeddingResponse, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.calls.append({"model": model, "input": texts}) + model_response.model = model + model_response.data = [ + {"object": "embedding", "index": index, "embedding": _semantic_embedding(str(text))} + for index, text in enumerate(texts) + ] + return model_response + + def embedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + return self._respond(model, input, model_response) + + async def aembedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + return self._respond(model, input, model_response) + + +@pytest.fixture +def semantic_embedding() -> Generator[DeterministicEmbedding]: + handler: Final = DeterministicEmbedding() + with ExitStack() as stack: + stack.enter_context( + rebound( + litellm, + "custom_provider_map", + [ + *litellm.custom_provider_map, + cast( + CustomLLMItem, + {"provider": "semantic-test", "custom_handler": handler}, + ), + ], + ) + ) + stack.enter_context( + rebound( + litellm, + "_custom_providers", # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + [*litellm._custom_providers, "semantic-test"], # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + ) + ) + stack.enter_context( + rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"]) + ) + yield handler + + +@pytest.fixture +def redis_stack() -> Generator[tuple[str, str]]: + url: Final = os.environ.get("LITELLM_REDIS_STACK_URL") + if url is None: + pytest.skip("LITELLM_REDIS_STACK_URL is not set") + index: Final = f"{SEMANTIC_INDEX_PREFIX}{uuid4().hex}" + yield url, index + client: Final = redis.Redis.from_url(url) + try: + client.execute_command("FT.DROPINDEX", index, "DD") # pyright: ignore[reportUnknownMemberType] # redis-py leaves execute_command partially unknown + except redis.RedisError: + pass + client.close() + + +def semantic_request(key: str, prompt: str, **extra: object) -> dict[str, object]: + return { + "key": {"preset": key}, + "messages": [{"role": "user", "content": prompt}], + **extra, + } + + +def semantic_messages(prompt: str) -> list[dict[str, object]]: + return [{"role": "user", "content": prompt}] + + +def semantic_entry_id(prompt: str, tag: str) -> str: + return hashlib.sha256(f"{prompt}litellm_cache_key{tag}".encode()).hexdigest() + + +def semantic_facade(url: str, index: str, *, similarity_threshold: float = 0.8) -> Cache: + facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=similarity_threshold, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(facade) + return facade + + +def test_redis_semantic_constructor_identity_and_provenance( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + backend: Final = cast(RedisSemanticCache, facade.cache) + assert backend.__class__.__module__ == "litellm.caching.redis_semantic_cache" + assert type(backend) is RedisSemanticCache + assert backend._redis_url == url # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend._index_name == index # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend.similarity_threshold == 0.8 + assert backend.embedding_model == SEMANTIC_EMBEDDING_MODEL + handle: Final = cast(object, getattr(facade, "_native_cache_handle")) + assert isinstance(handle, _CacheTestHandle) + assert handle.backend == "redis_semantic" + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + + +def test_redis_semantic_native_and_python_sync_entries_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + response: Final = {"choices": [{"text": "paris"}], "usage": {"total_tokens": 2}} + + binding.store(semantic_request("geo", "what is the capital of france"), response) + + native_hash_key: Final = f"{index}:{semantic_entry_id('what is the capital of france', 'geo')}" + stored: Final = client.hgetall(native_hash_key) + assert set(stored) == { + b"entry_id", + b"prompt", + b"response", + b"prompt_vector", + b"inserted_at", + b"updated_at", + b"litellm_cache_key", + }, stored + assert stored[b"entry_id"].decode() == native_hash_key.split(":", 1)[1] + assert stored[b"prompt"] == b"what is the capital of france" + assert stored[b"litellm_cache_key"] == b"geo" + assert len(stored[b"prompt_vector"]) == 32 + decoded: Final = cast(dict[str, object], json.loads(stored[b"response"])) + assert decoded["response"] == response + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "geo", messages=semantic_messages("what is the capital of france") + ) + == decoded + ) + assert semantic_embedding.calls == [ + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["dimension test"]}, + ] + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "math", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": 42}}), + messages=semantic_messages("what is 6 times 7"), + ) + python_hash_key: Final = f"{index}:{semantic_entry_id('what is 6 times 7', 'math')}" + assert json.loads(cast(bytes, client.hget(python_hash_key, "response"))) == { + "timestamp": 1700000000.0, + "response": {"answer": 42}, + } + assert binding.lookup(semantic_request("math", "what is 6 times 7")) == {"answer": 42} + client.close() + + +async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + await binding.async_store( + semantic_request("async", "name a primary color"), {"answer": "blue"} + ) + hash_key: Final = f"{index}:{semantic_entry_id('name a primary color', 'async')}" + decoded: Final = cast(dict[str, object], json.loads(cast(bytes, client.hget(hash_key, "response")))) + python_read: Final = await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async", messages=semantic_messages("name a primary color") + ) + assert python_read == decoded + + await binding.async_store_batch( + [ + semantic_request("batch-one", "first batch prompt"), + semantic_request("batch-two", "second batch prompt"), + ], + [{"answer": 1}, {"answer": 2}], + ) + expected: Final = { + key: json.loads( + cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response")) + ) + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ) + } + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ): + assert cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + key, messages=semantic_messages(prompt) + ) == expected[key], key + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async-python", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": "python"}}), + messages=semantic_messages("python written prompt"), + ) + assert await binding.async_lookup( + semantic_request("async-python", "python written prompt") + ) == {"answer": "python"} + client.close() + + +def test_redis_semantic_similarity_tag_and_threshold_boundaries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + binding.store(semantic_request("sim", "tell me a joke"), {"answer": "haha"}) + paraphrase: Final = f"tell me a joke{PARAPHRASE_MARKER}" + assert binding.lookup(semantic_request("sim", paraphrase)) == {"answer": "haha"} + assert binding.lookup(semantic_request("sim", "an unrelated question about spreadsheets")) is None + assert binding.lookup(semantic_request("other-key", "tell me a joke")) is None + + strict: Final = semantic_facade(url, index, similarity_threshold=0.99) + strict_binding: Final = _CacheTestResolver(SimpleNamespace(cache=strict)).resolve() + assert strict_binding.lookup(semantic_request("sim", paraphrase)) is None + assert strict_binding.lookup(semantic_request("sim", "tell me a joke")) == {"answer": "haha"} + + +def test_redis_semantic_ttl_is_written_only_when_requested( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store( + {**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1} + ) + expiring: Final = f"{index}:{semantic_entry_id('ttl prompt', 'ttl')}" + assert 0 < client.ttl(expiring) <= 12 + + binding.store(semantic_request("ttl-none", "untimed prompt"), {"answer": 2}) + persistent: Final = f"{index}:{semantic_entry_id('untimed prompt', 'ttl-none')}" + assert client.ttl(persistent) == -1 + + binding.store( + {**semantic_request("ttl-fraction", "fractional prompt"), "ttl_seconds": 1.5}, + {"answer": 3}, + ) + fractional: Final = f"{index}:{semantic_entry_id('fractional prompt', 'ttl-fraction')}" + assert client.ttl(fractional) == 2 + client.close() + + +def test_redis_semantic_malformed_response_is_a_miss_for_both_readers( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(semantic_request("bad", "corrupt me"), {"answer": 1}) + hash_key: Final = f"{index}:{semantic_entry_id('corrupt me', 'bad')}" + client.hset(hash_key, "response", b"{not json") + assert binding.lookup(semantic_request("bad", "corrupt me")) is None + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "bad", messages=semantic_messages("corrupt me") + ) + is None + ) + client.close() + + +async def test_redis_semantic_unsupported_operations_raise_not_implemented( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + with pytest.raises(NotImplementedError): + binding.lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_flush() + with pytest.raises(NotImplementedError): + await binding.ping() + + +def test_redis_semantic_requests_without_prompt_are_noops( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(request("plain"), {"answer": 1}) + assert binding.lookup(request("plain")) is None + assert semantic_embedding.calls == [] + assert client.keys(f"{index}:*") == [] + client.close() + + +def test_redis_semantic_scope_overrides_the_tag_and_isolates_entries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + scoped: Final = {**semantic_request("scoped", "scoped prompt"), "scope": "team-a"} + binding.store(scoped, {"answer": "kept"}) + hash_key: Final = f"{index}:{semantic_entry_id('scoped prompt', 'team-a')}" + assert client.hget(hash_key, "litellm_cache_key") == b"team-a" + assert binding.lookup(scoped) == {"answer": "kept"} + assert binding.lookup(semantic_request("scoped", "scoped prompt")) is None + assert binding.lookup({**scoped, "scope": "team-b"}) is None + client.close() + + +def test_redis_semantic_configuration_drift_falls_back_to_python( + redis_stack: tuple[str, str], + semantic_embedding: DeterministicEmbedding, + monkeypatch: pytest.MonkeyPatch, +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + + with rebound(facade.cache, "similarity_threshold", 0.5): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "semantic_cache_scope", "end_user"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "embedding_model", "other-model"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "_index_name", "other-index"): + assert resolver.resolve().kind == "python_callback" + + def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]: + return _semantic_embedding(prompt) + + monkeypatch.setattr(RedisSemanticCache, "_get_embedding", patched_embedding) + assert resolver.resolve().kind == "python_callback" + + +def test_redis_semantic_handle_rejects_wrong_backends( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + + class CustomSemanticCache(RedisSemanticCache): + pass + + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + _CacheTestHandle.redis_semantic(object()) + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + _CacheTestHandle.redis_semantic( + CustomSemanticCache( + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=f"{index}_subclass", + ) + ) + + facade: Final = semantic_facade(url, index) + with pytest.raises(TypeError, match="backend types must match"): + _CacheTestHandle.redis(url)._bind_facade(facade) + + subclassed_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + subclassed_facade.cache = CustomSemanticCache( # pyright: ignore[reportAttributeAccessIssue] # facade backend slot is not declared + + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=index, + ) + with pytest.raises(TypeError): + _CacheTestHandle.redis_semantic( + subclassed_facade.cache + )._bind_facade(subclassed_facade) + + replacement_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + with pytest.raises(TypeError, match="must be the native embedder"): + _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade) From d3f2ddba050bad5af5ec662ac83bac263919d907 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:58:02 +0000 Subject: [PATCH 11/22] fix(python-bridge): allow instance shadowing only for validated config attributes Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/src/cache/facade.rs | 2 +- tests/test_litellm_rust/test_cache.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index d7ec2052dd0..a6395d086cb 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -120,7 +120,7 @@ impl ObjectGuard { if !attributes.get_item(name)?.is(value.bind(py)) { return Ok(false); } - if instance.contains(name)? && value.bind(py).is_callable() { + if instance.contains(name)? && !self.config_names.contains(&name.as_str()) { return Ok(false); } } diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index e9f4b99a3b7..73a6321011b 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -820,6 +820,8 @@ def test_redis_semantic_configuration_drift_falls_back_to_python( assert resolver.resolve().kind == "python_callback" with rebound(facade.cache, "_index_name", "other-index"): assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "CACHE_KEY_FIELD_NAME", "other-field"): + assert resolver.resolve().kind == "python_callback" def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]: return _semantic_embedding(prompt) From a9ff1de42bfdf85c6ae327113a968bacb914e99d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:09:36 +0000 Subject: [PATCH 12/22] build(rust): switch release LTO to fat for wheel size headroom Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 05eea6bc299..7fa8de05f60 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -75,7 +75,7 @@ veil = "0.3.0" [profile.release] opt-level = 3 -lto = "thin" +lto = "fat" codegen-units = 1 panic = "unwind" debug = false From d2f8e8c83504d0662bf462fa54d3828d4dcc426f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:21:57 +0000 Subject: [PATCH 13/22] fix(cache-redis-semantic): harden index initialization Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-redis-semantic/src/cache.rs | 24 +++- .../cache-redis-semantic/tests/cache.rs | 118 +++++++++++++++++- 2 files changed, 132 insertions(+), 10 deletions(-) diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs index b1440e80de5..cd79f067296 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -71,7 +71,11 @@ impl Inner { Some(true) => self.index_name.clone(), Some(false) => self.isolated_index(connection, dims)?, None => { - create_index(connection, &self.index_name, dims)?; + if create_index(connection, &self.index_name, dims).is_err() + && index_compatible(connection, &self.index_name, dims)? != Some(true) + { + return Err(Error::Unavailable); + } self.index_name.clone() } }; @@ -509,36 +513,46 @@ fn schema_compatible(info: &redis::Value, dims: usize) -> bool { .iter() .map(|attribute| { let redis::Value::Array(attribute) = attribute else { - return (None, None, None); + return (None, None, None, None, None); }; let mut name = None; let mut field_type = None; let mut dim = None; + let mut data_type = None; + let mut distance_metric = None; for pair in attribute.as_chunks::<2>().0 { match string_value(&pair[0]).as_deref() { Some("identifier") => name = string_value(&pair[1]), Some("type") => field_type = string_value(&pair[1]), Some("dim") => dim = number_value(&pair[1]), + Some("data_type") => data_type = string_value(&pair[1]), + Some("distance_metric") => distance_metric = string_value(&pair[1]), _ => {} } } - (name, field_type, dim) + (name, field_type, dim, data_type, distance_metric) }) .collect::>(); let has_field = |name: &str, field_type: &str| { fields .iter() - .any(|(n, t, _)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type)) + .any(|(n, t, ..)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type)) }; has_field("prompt", "TEXT") && has_field("response", "TEXT") && has_field("inserted_at", "NUMERIC") && has_field("updated_at", "NUMERIC") && has_field(CACHE_KEY_FIELD, "TAG") - && fields.iter().any(|(n, t, d)| { + && fields.iter().any(|(n, t, d, data, metric)| { n.as_deref() == Some(VECTOR_FIELD) && t.as_deref() == Some("VECTOR") && *d == Some(dims as f64) + && data + .as_deref() + .is_some_and(|data| data.eq_ignore_ascii_case("float32")) + && metric + .as_deref() + .is_some_and(|metric| metric.eq_ignore_ascii_case("cosine")) }) } diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs index 77b057ae3b9..85a35a033b2 100644 --- a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -124,7 +124,7 @@ fn index_info(attributes: Vec) -> redis::Value { ]) } -fn vector_attribute(dims: i64) -> redis::Value { +fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> redis::Value { attribute( "prompt_vector", "VECTOR", @@ -132,26 +132,34 @@ fn vector_attribute(dims: i64) -> redis::Value { s("algorithm"), s("FLAT"), s("data_type"), - s("FLOAT32"), + s(data_type), s("dim"), redis::Value::Int(dims), s("distance_metric"), - s("COSINE"), + s(distance_metric), ], ) } -fn compatible_info(dims: i64) -> redis::Value { +fn vector_attribute(dims: i64) -> redis::Value { + vector_attribute_with(dims, "FLOAT32", "COSINE") +} + +fn info_with_vector(vector: redis::Value) -> redis::Value { index_info(vec![ attribute("prompt", "TEXT", vec![]), attribute("response", "TEXT", vec![]), attribute("inserted_at", "NUMERIC", vec![]), attribute("updated_at", "NUMERIC", vec![]), - vector_attribute(dims), + vector, attribute("litellm_cache_key", "TAG", vec![]), ]) } +fn compatible_info(dims: i64) -> redis::Value { + info_with_vector(vector_attribute(dims)) +} + fn unscoped_info(dims: i64) -> redis::Value { index_info(vec![ attribute("prompt", "TEXT", vec![]), @@ -537,6 +545,106 @@ fn incompatible_schema_falls_back_to_isolated_index() { .unwrap(); } +#[test] +fn create_index_race_rechecks_schema_and_stores() { + let prompt = "hello prompt"; + let tag = "key1"; + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new( + create_index_command(INDEX, 3), + Err::<&str, _>(redis::RedisError::from(( + redis::ErrorKind::Extension, + "Index already exists", + ))), + ), + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn wrong_distance_metric_falls_back_to_isolated_index() { + let prompt = "hello prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Ok(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2"))), + ), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{isolated}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + #[test] fn tag_special_characters_are_escaped_in_search_filter() { let vector = vec![0.1f32, 0.2, 0.3]; From 974d9f97ff13b4deb7afa19b2ab43387baf597a0 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:22:09 +0000 Subject: [PATCH 14/22] revert(rust): restore thin LTO in the release profile Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 7fa8de05f60..05eea6bc299 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -75,7 +75,7 @@ veil = "0.3.0" [profile.release] opt-level = 3 -lto = "fat" +lto = "thin" codegen-units = 1 panic = "unwind" debug = false From 6237dd51cb8711df528b360100e080d5fc73d6c2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:22:59 +0000 Subject: [PATCH 15/22] fix(cache-redis-semantic): isolate on an incompatible index after a lost create race Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-redis-semantic/src/cache.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs index cd79f067296..4181ce719e3 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -70,14 +70,14 @@ impl Inner { let name = match index_compatible(connection, &self.index_name, dims)? { Some(true) => self.index_name.clone(), Some(false) => self.isolated_index(connection, dims)?, - None => { - if create_index(connection, &self.index_name, dims).is_err() - && index_compatible(connection, &self.index_name, dims)? != Some(true) - { - return Err(Error::Unavailable); - } - self.index_name.clone() - } + None => match create_index(connection, &self.index_name, dims) { + Ok(()) => self.index_name.clone(), + Err(_) => match index_compatible(connection, &self.index_name, dims)? { + Some(true) => self.index_name.clone(), + Some(false) => self.isolated_index(connection, dims)?, + None => return Err(Error::Unavailable), + }, + }, }; let _ = self.resolved_index.set(name.clone()); Ok(name) From 220b981ab4c390fdfce6d8baaad557a2bebb7812 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:53:18 +0000 Subject: [PATCH 16/22] test(rust): deduplicate the merged os import Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/test_litellm_rust/test_cache.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 87e9d76ccb4..a60feb9973a 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -3,7 +3,6 @@ import contextvars import gc import hashlib import json -import os import math import os import threading From 8c8250596473aaef9ba6b1e685eeee3ead42b8a1 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:02:52 +0000 Subject: [PATCH 17/22] fix(python-bridge): await semantic embeddings inline in the caller's task Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-redis-semantic/src/lib.rs | 1 + litellm-rust/crates/python-bridge/Cargo.toml | 2 +- .../crates/python-bridge/src/cache/binding.rs | 22 ++- .../python-bridge/src/cache/embedder.rs | 78 ++++++--- .../crates/python-bridge/src/cache/mod.rs | 1 + .../crates/python-bridge/src/cache/native.rs | 7 + .../crates/python-bridge/src/cache/request.rs | 1 + .../python-bridge/src/cache/semantic.rs | 165 ++++++++++++++++++ tests/test_litellm_rust/test_cache.py | 56 ++++++ 9 files changed, 308 insertions(+), 25 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/cache/semantic.rs diff --git a/litellm-rust/crates/cache-redis-semantic/src/lib.rs b/litellm-rust/crates/cache-redis-semantic/src/lib.rs index a34603cd18f..51d0b4ba5f3 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/lib.rs @@ -2,3 +2,4 @@ mod cache; mod prompt; pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig}; +pub use prompt::prompt_from_context; diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index b28ddc50181..93ce5828489 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -39,7 +39,7 @@ litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true serde_json.workspace = true -tokio = { workspace = true, features = ["sync"] } +tokio = { workspace = true, features = ["rt", "sync"] } [dev-dependencies] serde.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs index ad64b24d3c1..0b90e8151ea 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -14,6 +14,7 @@ use super::{ future::{ready_none, ready_value}, native::NativeResponseCache, request::{now, request, requests}, + semantic::{SemanticOperation, drive}, }; pub(super) enum CacheBinding { @@ -56,6 +57,11 @@ impl ResolvedCache { CacheBinding::Disabled => ready_none(py)?, CacheBinding::Native(service) => { let request = request(input)?; + if service.semantic_embedder().is_some() { + return Ok(ExecutionStep::Await( + drive(py, service.clone(), SemanticOperation::Lookup(request))?.unbind(), + )); + } let service = service.clone(); run_async( py, @@ -179,6 +185,13 @@ impl ResolvedCache { CacheBinding::Native(service) => { let request = self::request(request)?; let response: Value = from_py(response)?; + if service.semantic_embedder().is_some() { + return drive( + py, + service.clone(), + SemanticOperation::Store(request, response), + ); + } let service = service.clone(); run_async( py, @@ -240,7 +253,14 @@ impl ResolvedCache { "batch cache requests and responses must have equal lengths", )); } - let entries = requests.into_iter().zip(responses).collect(); + let entries = requests.into_iter().zip(responses).collect::>(); + if service.semantic_embedder().is_some() { + return drive( + py, + service.clone(), + SemanticOperation::StoreBatch(entries.into()), + ); + } let service = service.clone(); run_async( py, diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs index 63e078cd815..26edb26f428 100644 --- a/litellm-rust/crates/python-bridge/src/cache/embedder.rs +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -6,6 +6,17 @@ use litellm_host_python::to_py; use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict}; use serde_json::{Map, Value}; +tokio::task_local! { + static PREPARED_EMBEDDING: Result, Error>; +} + +pub(super) fn with_prepared_embedding( + vector: Result, Error>, + future: F, +) -> impl Future { + PREPARED_EMBEDDING.scope(vector, future) +} + pub(super) struct PythonEmbedder(Py); impl PythonEmbedder { @@ -34,7 +45,20 @@ impl PythonEmbedder { Ok(kwargs) } - fn extract(vector: Bound<'_, PyAny>) -> PyResult> { + pub(super) fn async_embedding_coroutine( + &self, + py: Python<'_>, + prompt: &str, + metadata: &Map, + ) -> PyResult> { + let kwargs = Self::metadata_kwargs(py, metadata)?; + self.0 + .bind(py) + .call_method("_get_async_embedding", (prompt,), Some(&kwargs)) + .map(Bound::unbind) + } + + pub(super) fn extract(vector: Bound<'_, PyAny>) -> PyResult> { Ok(vector .extract::>()? .into_iter() @@ -58,28 +82,36 @@ impl Embedder for PythonEmbedder { fn async_embed( &self, - prompt: &str, - metadata: &Map, + _prompt: &str, + _metadata: &Map, ) -> impl Future, Error>> + Send { - let coroutine = Python::attach(|py| { - let kwargs = Self::metadata_kwargs(py, metadata)?; - self.0 - .bind(py) - .call_method("_get_async_embedding", (prompt,), Some(&kwargs)) - .map(Bound::unbind) - }) - .map_err(|_| Error::Unavailable); - async move { - let coroutine = coroutine?; - let awaited = Python::attach(|py| { - pyo3_async_runtimes::tokio::into_future(coroutine.into_bound(py)) - }) - .map_err(|_| Error::Unavailable)? - .await - .map_err(|_| Error::Unavailable)?; - let vector = Python::attach(|py| awaited.extract::>(py)) - .map_err(|_| Error::Unavailable)?; - Ok(vector.into_iter().map(|value| value as f32).collect()) - } + let seeded = PREPARED_EMBEDDING + .try_with(Clone::clone) + .unwrap_or(Err(Error::Unavailable)); + std::future::ready(seeded) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn async_embed_returns_the_seeded_vector_or_unavailable() { + let embedder = Python::attach(|py| PythonEmbedder::new(py.None())); + let metadata = Map::new(); + let embedder_ref = &embedder; + let metadata_ref = &metadata; + assert_eq!( + with_prepared_embedding(Ok(vec![0.5f32, 0.25]), async move { + embedder_ref.async_embed("prompt", metadata_ref).await + }) + .await, + Ok(vec![0.5, 0.25]) + ); + assert_eq!( + embedder.async_embed("prompt", &metadata).await, + Err(Error::Unavailable) + ); } } diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 4cc87367d91..cd772d571cb 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -8,6 +8,7 @@ mod handle; mod native; mod request; mod resolver; +mod semantic; use litellm_cache::Error; use pyo3::{ diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 182010fab02..de9c4afa236 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -122,6 +122,13 @@ impl NativeResponseCache { } } + pub fn semantic_embedder(&self) -> Option<&PythonEmbedder> { + match self { + Self::RedisSemantic(cache) => Some(cache.backend().embedder()), + Self::Memory(_) | Self::Redis { .. } => None, + } + } + pub fn embedder_object(&self) -> Option<&Py> { match self { Self::RedisSemantic(cache) => Some(cache.backend().embedder().object()), diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 26e0fe4e62c..b06087bcc83 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -20,6 +20,7 @@ struct RequestInput { scope: Option, } +#[derive(Clone)] pub(super) struct CacheRequest { key: CacheKeyInput, controls: CacheControls, diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/semantic.rs new file mode 100644 index 00000000000..eb38b8b9c67 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/semantic.rs @@ -0,0 +1,165 @@ +use std::collections::VecDeque; + +use litellm_cache::Error; +use litellm_cache_redis_semantic::prompt_from_context; +use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; +use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use serde_json::Value; + +use super::{ + cache_error, + embedder::{PythonEmbedder, with_prepared_embedding}, + native::NativeResponseCache, + request::{CacheRequest, now}, +}; + +pub(super) enum SemanticOperation { + Lookup(CacheRequest), + Store(CacheRequest, Value), + StoreBatch(VecDeque<(CacheRequest, Value)>), +} + +enum Phase { + Start, + AwaitingEmbedding, + AwaitingBackend, +} + +pub(super) struct SemanticBody { + service: NativeResponseCache, + operation: SemanticOperation, + pending: Option<(CacheRequest, Option)>, + phase: Phase, +} + +impl SemanticBody { + pub(super) fn new(service: NativeResponseCache, operation: SemanticOperation) -> Self { + Self { + service, + operation, + pending: None, + phase: Phase::Start, + } + } + + fn backend_step( + &mut self, + py: Python<'_>, + seed: Result, Error>, + ) -> PyResult { + self.phase = Phase::AwaitingBackend; + let (request, response) = self.pending.take().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution resumed without a pending operation") + })?; + let service = self.service.clone(); + let future = async move { + match response { + None => service.async_lookup(&request, now()).await, + Some(response) => service + .async_store(&request, response, now()) + .await + .map(|_| None), + } + }; + let awaitable = run_async(py, with_prepared_embedding(seed, future), cache_error)?; + Ok(ExecutionStep::Await(awaitable.unbind())) + } +} + +impl ExecutionBody for SemanticBody { + fn resume(&mut self, mut result: Option>>) -> PyResult { + Python::attach(|py| { + loop { + match self.phase { + Phase::Start => { + if result.is_some() { + return Err(PyRuntimeError::new_err( + "semantic execution received a result before starting", + )); + } + if self.pending.is_none() { + match &mut self.operation { + SemanticOperation::Lookup(request) => { + self.pending = Some((request.clone(), None)); + } + SemanticOperation::Store(request, response) => { + let response = std::mem::replace(response, Value::Null); + self.pending = Some((request.clone(), Some(response))); + } + SemanticOperation::StoreBatch(queue) => { + let Some((request, response)) = queue.pop_front() else { + return Ok(ExecutionStep::Return(py.None())); + }; + self.pending = Some((request, Some(response))); + } + } + } + let (request, _) = self.pending.as_ref().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution has no pending operation") + })?; + let semantic = request.semantic(); + let Some(prompt) = prompt_from_context(&semantic.context) else { + return self.backend_step(py, Err(Error::Unavailable)); + }; + let embedder = self.service.semantic_embedder().ok_or_else(|| { + PyRuntimeError::new_err( + "semantic execution requires a redis-semantic backend", + ) + })?; + let coroutine = embedder.async_embedding_coroutine( + py, + &prompt, + &semantic.context.metadata, + )?; + self.phase = Phase::AwaitingEmbedding; + return Ok(ExecutionStep::Await(coroutine)); + } + Phase::AwaitingEmbedding => { + let result = result.take().ok_or_else(|| { + PyRuntimeError::new_err( + "semantic execution expected an embedding result", + ) + })?; + let seed = result + .and_then(|value| PythonEmbedder::extract(value.into_bound(py))) + .map_err(|_| Error::Unavailable); + return self.backend_step(py, seed); + } + Phase::AwaitingBackend => { + let result = result.take().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution expected a backend result") + })?; + let value = match result { + Ok(value) => value, + Err(error) => return Err(error), + }; + let more = matches!( + &self.operation, + SemanticOperation::StoreBatch(queue) if !queue.is_empty() + ); + if more { + self.phase = Phase::Start; + continue; + } + return Ok(ExecutionStep::Return(value)); + } + } + } + }) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.service.traverse(visit) + } +} + +pub(super) fn drive( + py: Python<'_>, + service: NativeResponseCache, + operation: SemanticOperation, +) -> PyResult> { + let execution = Py::new(py, Execution::new(SemanticBody::new(service, operation)))?; + py.import("litellm.rust_bridge.lifecycle")? + .getattr("drive")? + .call1((execution,)) +} diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index a60feb9973a..9312fdde075 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -479,6 +479,7 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n PARAPHRASE_MARKER: Final = " (paraphrase)" SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic" SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_" +SEMANTIC_CONTEXT: Final = contextvars.ContextVar("semantic_test_context", default="unset") def _normalized(vector: list[float]) -> list[float]: @@ -509,6 +510,7 @@ def _semantic_embedding(prompt: str) -> list[float]: class DeterministicEmbedding(litellm.CustomLLM): def __init__(self) -> None: self.calls: list[dict[str, object]] = [] + self.async_calls: list[dict[str, object]] = [] def _respond( self, @@ -553,6 +555,16 @@ class DeterministicEmbedding(litellm.CustomLLM): timeout: object = None, litellm_params: object = None, ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.async_calls.append( + { + "model": model, + "input": texts, + "task": asyncio.current_task(), + "context": SEMANTIC_CONTEXT.get(), + } + ) + SEMANTIC_CONTEXT.set("written-in-aembedding") return self._respond(model, input, model_response) @@ -755,6 +767,50 @@ async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( client.close() +async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + caller: Final = asyncio.current_task() + SEMANTIC_CONTEXT.set("caller-sentinel") + response: Final = {"choices": [{"text": "paris"}]} + + await binding.async_store( + semantic_request("inline", "what is the capital of france"), response + ) + assert ( + await binding.async_lookup( + semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}") + ) + == response + ) + assert await binding.async_lookup(semantic_request("inline", "python written prompt")) is None + assert SEMANTIC_CONTEXT.get() == "written-in-aembedding" + assert semantic_embedding.async_calls == [ + { + "model": "deterministic", + "input": ["what is the capital of france"], + "task": caller, + "context": "caller-sentinel", + }, + { + "model": "deterministic", + "input": [f"what is the capital of france{PARAPHRASE_MARKER}"], + "task": caller, + "context": "written-in-aembedding", + }, + { + "model": "deterministic", + "input": ["python written prompt"], + "task": caller, + "context": "written-in-aembedding", + }, + ], semantic_embedding.async_calls + + def test_redis_semantic_similarity_tag_and_threshold_boundaries( redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding ) -> None: From eeaf4f36e49d507cdfe0614c9bb374d5338b571d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:07:15 +0000 Subject: [PATCH 18/22] test(python-bridge): initialize the interpreter in the embedder seed test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/src/cache/embedder.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs index 26edb26f428..99c3de34e05 100644 --- a/litellm-rust/crates/python-bridge/src/cache/embedder.rs +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -98,6 +98,7 @@ mod tests { #[tokio::test] async fn async_embed_returns_the_seeded_vector_or_unavailable() { + Python::initialize(); let embedder = Python::attach(|py| PythonEmbedder::new(py.None())); let metadata = Map::new(); let embedder_ref = &embedder; From 2ab4b255883b660bbc88a94e682c5ecd9e4e9ccf Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:30:29 +0000 Subject: [PATCH 19/22] fix(python-bridge): update cache test handle stubs for merged backends Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/rust_bridge/_native.pyi | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 7eb266d5a09..baac21bb4bd 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -109,11 +109,14 @@ class _CacheTestHandle: *, ttl_seconds: float = 60.0, namespace: str | None = None, + startup_nodes: list[tuple[str, int]] | None = None, ) -> _CacheTestHandle: ... @staticmethod + def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ... + @staticmethod def redis_semantic(backend: object) -> _CacheTestHandle: ... @property - def backend(self) -> Literal["memory", "redis", "redis_semantic"]: ... + def backend(self) -> Literal["memory", "redis", "azure-blob", "redis_semantic"]: ... def _bind_facade(self, facade: object) -> None: ... @final From d38514dfc6be3ad591d3462ee57b1d7987a3c1fe Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 23:02:23 +0000 Subject: [PATCH 20/22] test(cache-redis-semantic): pin shared-index behavior across embedding dimensions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../cache-redis-semantic/tests/cache.rs | 148 ++++++++++++++++++ 1 file changed, 148 insertions(+) diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs index 85a35a033b2..fc37cf9f97f 100644 --- a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -817,6 +817,154 @@ async fn async_paths_embed_then_run_blocking_redis_work() { ); } +#[test] +fn shared_base_index_across_dimensions_replaces_the_isolated_index() { + // Pins parity with Python's `_isolated` + overwrite=True flow. + let prompt = "shared prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let value = entry(); + let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); + let store_hash = |index: &str, vector: &[f32]| { + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{index}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ) + }; + + let vector_a = vec![0.1f32; 8]; + let connection_a = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(INDEX, 8), Ok("OK")), + store_hash(INDEX, &vector_a), + ]) + .assert_all_commands_consumed(); + let (embedder_a, _) = FakeEmbedder::new(&[(prompt, &vector_a)]); + let worker_a = RedisSemanticCache::with_connection(connection_a, embedder_a, config()) + .with_clock(|| 1700000000.5); + worker_a.set_cache(tag, value.clone(), &context()).unwrap(); + + let vector_b = vec![0.2f32; 4]; + let connection_b = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(8))), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 4), Ok("OK")), + store_hash(&isolated, &vector_b), + MockCmd::new( + search_command(&isolated, tag, &vector_b), + Ok(search_result(hit_fields(tag, "0.0", encoded(&value)))), + ), + MockCmd::new( + search_command(&isolated, tag, &vector_b), + Err::(redis::RedisError::from(( + redis::ErrorKind::Extension, + "Vector dimension mismatch", + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder_b, _) = FakeEmbedder::new(&[(prompt, &vector_b)]); + let worker_b = RedisSemanticCache::with_connection(connection_b, embedder_b, config()) + .with_clock(|| 1700000000.5); + worker_b.set_cache(tag, value.clone(), &context()).unwrap(); + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap(), + Some(value.clone()) + ); + + let vector_c = vec![0.3f32; 16]; + let connection_c = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(8))), + MockCmd::new(redis::cmd("FT.INFO").arg(&isolated), Ok(compatible_info(4))), + MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")), + MockCmd::new(create_index_command(&isolated, 16), Ok("OK")), + store_hash(&isolated, &vector_c), + ]) + .assert_all_commands_consumed(); + let (embedder_c, _) = FakeEmbedder::new(&[(prompt, &vector_c)]); + let worker_c = RedisSemanticCache::with_connection(connection_c, embedder_c, config()) + .with_clock(|| 1700000000.5); + worker_c.set_cache(tag, value.clone(), &context()).unwrap(); + + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap_err(), + Error::Unavailable + ); +} + +#[test] +fn live_shared_index_is_replaced_across_dimensions() { + let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { + return; + }; + // Pins parity with Python's `_isolated` + overwrite=True flow. + let base = format!("rust_semantic_shared_{}", std::process::id()); + let isolated = format!("{base}_isolated"); + let prompt = "shared live prompt"; + let tag = "key1"; + let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); + let value = entry(); + let worker = |vector: Vec| { + let (embedder, _) = FakeEmbedder::new(&[(prompt, vector.as_slice())]); + RedisSemanticCache::new( + &url, + embedder, + RedisSemanticConfig { + index_name: base.clone(), + similarity_threshold: 0.9, + }, + ) + .unwrap() + }; + + let worker_a = worker(vec![0.1f32; 8]); + worker_a.set_cache(tag, value.clone(), &context()).unwrap(); + + let worker_b = worker(vec![0.2f32; 4]); + worker_b.set_cache(tag, value.clone(), &context()).unwrap(); + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap(), + Some(value.clone()) + ); + + let worker_c = worker(vec![0.3f32; 16]); + worker_c.set_cache(tag, value.clone(), &context()).unwrap(); + + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap_err(), + Error::Unavailable + ); + + let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); + for index in [&base, &isolated] { + let _: Result<(), _> = redis::cmd("FT.DROPINDEX") + .arg(index) + .arg("DD") + .query(&mut connection); + } +} + #[test] fn live_store_lookup_and_ttl_against_redis_stack() { let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { From 7282494c30e009ba655e6f1453e8c28f6e1e5f3c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 23:04:44 +0000 Subject: [PATCH 21/22] fix(python-bridge): propagate cancellation from semantic embedding awaits Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../python-bridge/src/cache/semantic.rs | 19 +++++++--- tests/test_litellm_rust/test_cache.py | 35 +++++++++++++++++++ 2 files changed, 50 insertions(+), 4 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/semantic.rs index eb38b8b9c67..f0f75de1edf 100644 --- a/litellm-rust/crates/python-bridge/src/cache/semantic.rs +++ b/litellm-rust/crates/python-bridge/src/cache/semantic.rs @@ -3,7 +3,11 @@ use std::collections::VecDeque; use litellm_cache::Error; use litellm_cache_redis_semantic::prompt_from_context; use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; -use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyException, PyRuntimeError}, + prelude::*, +}; use serde_json::Value; use super::{ @@ -120,9 +124,16 @@ impl ExecutionBody for SemanticBody { "semantic execution expected an embedding result", ) })?; - let seed = result - .and_then(|value| PythonEmbedder::extract(value.into_bound(py))) - .map_err(|_| Error::Unavailable); + let seed = match result { + Ok(value) => PythonEmbedder::extract(value.into_bound(py)) + .map_err(|_| Error::Unavailable), + Err(error) => { + if !error.is_instance_of::(py) { + return Err(error); + } + Err(Error::Unavailable) + } + }; return self.backend_step(py, seed); } Phase::AwaitingBackend => { diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 3f5449a1aa9..8a8c83cd82b 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -736,6 +736,8 @@ class DeterministicEmbedding(litellm.CustomLLM): def __init__(self) -> None: self.calls: list[dict[str, object]] = [] self.async_calls: list[dict[str, object]] = [] + self.entered = asyncio.Event() + self.gate: asyncio.Event | None = None def _respond( self, @@ -790,6 +792,9 @@ class DeterministicEmbedding(litellm.CustomLLM): } ) SEMANTIC_CONTEXT.set("written-in-aembedding") + self.entered.set() + if self.gate is not None: + await self.gate.wait() return self._respond(model, input, model_response) @@ -1036,6 +1041,36 @@ async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task( ], semantic_embedding.async_calls +async def test_native_semantic_cancellation_during_embedding_skips_the_backend( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + semantic_embedding.gate = asyncio.Event() + + async def lookup() -> object: + return await binding.async_lookup( + semantic_request("cancel", "cancelled prompt") + ) + + task: Final = asyncio.create_task(lookup()) + await semantic_embedding.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + semantic_embedding.gate.set() + + assert len(semantic_embedding.async_calls) == 1 + assert ( + await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "cancel", messages=semantic_messages("cancelled prompt") + ) + is None + ) + + def test_redis_semantic_similarity_tag_and_threshold_boundaries( redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding ) -> None: From 69b7224aef7e64e7a38eb624d1075c8e141335d7 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 23:56:40 +0000 Subject: [PATCH 22/22] test(python-bridge): initialize the interpreter in the embedder seed test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/src/cache/embedder.rs | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs index 10ccf396510..9398e5a862b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/embedder.rs +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -140,10 +140,8 @@ mod tests { #[tokio::test] async fn async_embed_returns_the_seeded_vector_or_unavailable() { - let object = Python::attach(|py| { - Python::initialize(); - py.None() - }); + Python::initialize(); + let object = Python::attach(|py| py.None()); let embedder = PythonEmbedder::new(object); let scoped_embedder = embedder.clone(); let scoped = with_prepared_embedding(Ok(vec![0.25]), async move {