From d2f457f144a430ad2848b33839d298d9312c8d4e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:20:31 +0000 Subject: [PATCH 01/18] feat(cache): add semantic cache context and unsupported operation error Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache/src/base_cache.rs | 50 +++++++++++++++++++++ litellm-rust/crates/cache/src/error.rs | 2 + litellm-rust/crates/cache/src/lib.rs | 2 +- 3 files changed, 53 insertions(+), 1 deletion(-) diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 8bd69ba5ad6..5c10e7fd5c3 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -32,6 +32,28 @@ impl CacheContext for ExactCacheContext { } } +#[derive(Clone, Debug, Default, PartialEq)] +pub struct SemanticCacheContext { + pub input: Option, + pub messages: Option, + pub metadata: Option, + pub scope: Option, + pub ttl: Option, +} + +impl CacheContext for SemanticCacheContext { + fn ttl(&self) -> Option { + self.ttl + } + + fn with_ttl(&self, ttl: Option) -> Self { + Self { + ttl, + ..self.clone() + } + } +} + #[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum CacheConnectionStatus { @@ -105,3 +127,31 @@ pub trait BaseCache: Send + Sync { fn test_connection(&self) -> impl Future> + Send; } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use serde_json::json; + + use super::{CacheContext, SemanticCacheContext}; + + #[test] + fn semantic_context_with_ttl_only_replaces_ttl() { + let context = SemanticCacheContext { + input: Some(json!({"input": "hello"})), + messages: Some(json!([{"role": "user", "content": "hello"}])), + metadata: Some(json!({"tenant": "team"})), + scope: Some("scope".into()), + ttl: Some(Duration::from_secs(10)), + }; + + let updated = context.with_ttl(Some(Duration::from_secs(20))); + + assert_eq!(updated.ttl, Some(Duration::from_secs(20))); + assert_eq!(updated.input, context.input); + assert_eq!(updated.messages, context.messages); + assert_eq!(updated.metadata, context.metadata); + assert_eq!(updated.scope, context.scope); + } +} diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index ff3ff6572d4..51e4fe2d66a 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 operation is not supported by this backend")] + UnsupportedOperation, } 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}; From df97b274fc78ab261064f0a691016488c6c709dc Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:21:19 +0000 Subject: [PATCH 02/18] refactor(cache-response): generalize ResponseCache over the backend context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache-response/README.md | 2 +- .../crates/cache-response/src/response.rs | 43 +++++++++++-------- .../crates/python-bridge/src/cache/request.rs | 3 +- 3 files changed, 28 insertions(+), 20 deletions(-) diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 56c1646d343..46e561ddad1 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -58,4 +58,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths -Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees +Redis cluster, disk, and cloud stores remain follow-ups. Semantic backends plug in through `SemanticCacheContext`, which carries the prompt inputs and metadata alongside the cache TTL. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index e50e68cdabb..e70e07a5d26 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,21 +1,21 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache, + BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, 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, @@ -26,17 +26,24 @@ impl ResponseCacheRequest { default_on: true, ..Default::default() }, - context: ExactCacheContext::default(), + context: C::default(), max_age: None, } } } -pub struct ResponseCache> { +pub struct ResponseCache> +where + B::Context: Default + PartialEq, +{ backend: Arc, } -impl> ResponseCache { +impl ResponseCache +where + B: BaseCache, + B::Context: Default + PartialEq, +{ pub fn new(backend: Arc) -> Self { Self { backend } } @@ -46,7 +53,7 @@ impl> ResponseCach } pub fn default_ttl(&self) -> Option { - self.backend.get_ttl(&ExactCacheContext::default()) + self.backend.get_ttl(&B::Context::default()) } pub async fn async_flush(&self) -> Result<(), Error> @@ -62,7 +69,7 @@ impl> ResponseCach pub fn lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -81,7 +88,7 @@ impl> ResponseCach pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -101,7 +108,7 @@ impl> ResponseCach pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -126,7 +133,7 @@ impl> ResponseCach pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -153,7 +160,7 @@ impl> ResponseCach pub fn store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -172,7 +179,7 @@ impl> ResponseCach pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -193,7 +200,7 @@ impl> ResponseCach pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(ResponseCacheRequest, Value)>, now: Duration, ) -> Result<(), Error> { self.async_store_entries( @@ -209,7 +216,7 @@ impl> ResponseCach /// the freshness of its original response. pub async fn async_store_entries( &self, - entries: Vec<(ResponseCacheRequest, Value, Duration)>, + entries: Vec<(ResponseCacheRequest, Value, Duration)>, ) -> Result<(), Error> { let writable = entries .into_iter() @@ -249,8 +256,8 @@ impl> ResponseCach } fn partial_hits( - requests: &[ResponseCacheRequest], - readable: Vec<(usize, &ResponseCacheRequest)>, + requests: &[ResponseCacheRequest], + readable: Vec<(usize, &ResponseCacheRequest)>, entries: Vec>, now: Duration, ) -> Result { diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 0c5343a63d0..52a5f7d9055 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -1,5 +1,6 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use litellm_cache::ExactCacheContext; use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; @@ -20,7 +21,7 @@ pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult PyResult { - let mut request = ResponseCacheRequest::new(input.key); + let mut request = ResponseCacheRequest::::new(input.key); if let Some(controls) = input.controls { request.controls = controls; } From 1f86bb8e4640fd7e758e10f66106fd7c1bda01de Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:24:34 +0000 Subject: [PATCH 03/18] feat(cache-valkey-semantic): add native Valkey semantic cache backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 16 + .../crates/cache-valkey-semantic/Cargo.toml | 20 + .../crates/cache-valkey-semantic/src/lib.rs | 844 ++++++++++++++++++ 3 files changed, 880 insertions(+) create mode 100644 litellm-rust/crates/cache-valkey-semantic/Cargo.toml create mode 100644 litellm-rust/crates/cache-valkey-semantic/src/lib.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..5daf691d4c3 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2502,6 +2502,22 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-valkey-semantic" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "litellm-cache-response", + "r2d2", + "redis", + "redis-test", + "rstest", + "serde_json", + "sha2 0.10.9", + "tokio", + "uuid", +] + [[package]] name = "litellm-callbacks-legacy-python" version = "0.1.0" diff --git a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml new file mode 100644 index 00000000000..9a0a566ca3b --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "litellm-cache-valkey-semantic" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-cache-response.workspace = true +r2d2 = "0.8.10" +redis = { version = "1.7.0", features = ["tls-rustls"] } +serde_json.workspace = true +sha2.workspace = true +tokio.workspace = true +uuid = { version = "1", features = ["v4"] } + +[dev-dependencies] +redis-test = "1.0.4" +rstest.workspace = true diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs new file mode 100644 index 00000000000..85c4c9af15c --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs @@ -0,0 +1,844 @@ +use std::{ + future::Future, + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext}; +use litellm_cache_response::CacheEntry; +use serde_json::Value; +use sha2::{Digest, Sha256}; +use uuid::Uuid; + +pub trait Embedder: Send + Sync + 'static { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error>; + + fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> impl Future, Error>> + Send; +} + +#[derive(Clone, Debug, PartialEq)] +pub struct ValkeySemanticConfig { + pub similarity_threshold: f64, + pub index_name: String, +} + +pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; + +struct PooledConnection { + connection: redis::Connection, + failed: bool, +} + +struct ConnectionManager(redis::Client); + +impl r2d2::ManageConnection for ConnectionManager { + type Connection = PooledConnection; + type Error = redis::RedisError; + + fn connect(&self) -> Result { + let connection = self.0.get_connection()?; + Ok(PooledConnection { + connection, + failed: false, + }) + } + + fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), Self::Error> { + redis::cmd("PING").query::(&mut connection.connection)?; + Ok(()) + } + + fn has_broken(&self, connection: &mut Self::Connection) -> bool { + connection.failed || !redis::ConnectionLike::is_open(&connection.connection) + } +} + +enum Connections { + Pool(r2d2::Pool), + Fixed(Mutex), +} + +struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); + +impl redis::ConnectionLike for ConnectionRef<'_> { + fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { + self.0.req_packed_command(cmd) + } + + fn req_packed_commands( + &mut self, + cmd: &[u8], + offset: usize, + count: usize, + ) -> redis::RedisResult> { + self.0.req_packed_commands(cmd, offset, count) + } + + fn get_db(&self) -> i64 { + self.0.get_db() + } + + fn supports_pipelining(&self) -> bool { + self.0.supports_pipelining() + } + + fn check_connection(&mut self) -> bool { + self.0.check_connection() + } + + fn is_open(&self) -> bool { + self.0.is_open() + } +} + +impl Connections +where + C: redis::ConnectionLike + Send + 'static, +{ + fn execute( + &self, + operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, + ) -> Result { + match self { + Self::Pool(pool) => { + let mut pooled = pool.get().map_err(|_| Error::Unavailable)?; + let result = operation(&mut ConnectionRef(&mut pooled.connection)); + pooled.failed = matches!(result, Err(Error::Unavailable)); + result + } + Self::Fixed(connection) => { + let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; + operation(&mut ConnectionRef(&mut *connection)) + } + } + } +} + +pub struct ValkeySemanticCache< + E: Embedder, + S: CacheCodec, + C = redis::Connection, +> { + connections: Arc>, + embedder: E, + codec: S, + config: ValkeySemanticConfig, + index_dimension: Arc>>, +} + +impl ValkeySemanticCache +where + E: Embedder, + S: CacheCodec, +{ + pub fn new( + url: &str, + embedder: E, + codec: S, + config: ValkeySemanticConfig, + ) -> Result { + let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; + let pool = r2d2::Pool::builder() + .max_size(16) + .min_idle(Some(0)) + .test_on_check_out(false) + .build(ConnectionManager(client)) + .map_err(|_| Error::Unavailable)?; + Ok(Self { + connections: Arc::new(Connections::Pool(pool)), + embedder, + codec, + config, + index_dimension: Arc::new(Mutex::new(None)), + }) + } +} + +impl ValkeySemanticCache +where + E: Embedder, + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + pub fn with_connection( + connection: C, + embedder: E, + codec: S, + config: ValkeySemanticConfig, + ) -> Self { + Self { + connections: Arc::new(Connections::Fixed(Mutex::new(connection))), + embedder, + codec, + config, + index_dimension: Arc::new(Mutex::new(None)), + } + } + + pub fn similarity_threshold(&self) -> f64 { + self.config.similarity_threshold + } + + pub fn index_name(&self) -> &str { + &self.config.index_name + } + + fn key_prefix(&self) -> String { + format!("{}:", self.config.index_name) + } + + fn ensure_index(&self, dimension: usize) -> Result<(), Error> { + ensure_index( + &self.connections, + &self.config.index_name, + &self.key_prefix(), + &self.index_dimension, + dimension, + ) + } +} + +impl BaseCache for ValkeySemanticCache +where + E: Embedder, + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + 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 embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; + self.ensure_index(embedding.len())?; + let scope = scope_tag(key); + let document = format!("{}{}:{}", self.key_prefix(), scope, Uuid::new_v4()); + let response = self.codec.encode(&value)?; + let vector = embedding_bytes(&embedding); + let ttl = self.get_ttl(context); + self.connections.execute(|connection| { + let mut pipeline = redis::pipe(); + pipeline + .cmd("HSET") + .arg(&document) + .arg("litellm_cache_key") + .arg(&scope) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg("embedding") + .arg(vector) + .ignore(); + if let Some(ttl) = ttl { + pipeline + .cmd("EXPIRE") + .arg(&document) + .arg(ttl.as_secs()) + .ignore(); + } + pipeline + .query::<()>(connection) + .map_err(|_| Error::Unavailable) + }) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; + self.ensure_index(embedding.len())?; + let scope = scope_tag(key); + let query = + format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]"); + let vector = embedding_bytes(&embedding); + let response = self.connections.execute(|connection| { + redis::cmd("FT.SEARCH") + .arg(&self.config.index_name) + .arg(query) + .arg("PARAMS") + .arg(2) + .arg("vec") + .arg(vector) + .arg("RETURN") + .arg(2) + .arg("response") + .arg("vector_distance") + .arg("DIALECT") + .arg(2) + .query::(connection) + .map_err(|_| Error::Unavailable) + })?; + let Some(fields) = search_fields(response)? else { + return Ok(None); + }; + let response = fields + .iter() + .find_map(|(name, value)| (name == "response").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = fields + .iter() + .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = parse_f64(&distance)?; + if 1.0 - distance < self.config.similarity_threshold { + return Ok(None); + } + self.codec.decode(&response).map(Some) + } + + fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> impl Future> + Send { + let key = key.to_owned(); + let prompt = prompt_from_context(&context); + let metadata = context.metadata.clone(); + async move { + let Some(prompt) = prompt else { + return Ok(()); + }; + let embedding = self + .embedder + .async_embed(&prompt, metadata.as_ref()) + .await?; + let connections = Arc::clone(&self.connections); + let config = self.config.clone(); + let index_dimension = Arc::clone(&self.index_dimension); + let response = self.codec.encode(&value)?; + let vector = embedding_bytes(&embedding); + let prefix = format!("{}:", config.index_name); + let scope = scope_tag(&key); + let document = format!("{prefix}{scope}:{}", Uuid::new_v4()); + let ttl = context.ttl; + tokio::task::spawn_blocking(move || { + ensure_index( + &connections, + &config.index_name, + &prefix, + &index_dimension, + embedding.len(), + )?; + connections.execute(|connection| { + let mut pipeline = redis::pipe(); + pipeline + .cmd("HSET") + .arg(&document) + .arg("litellm_cache_key") + .arg(&scope) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg("embedding") + .arg(vector) + .ignore(); + if let Some(ttl) = ttl { + pipeline + .cmd("EXPIRE") + .arg(&document) + .arg(ttl.as_secs()) + .ignore(); + } + pipeline + .query::<()>(connection) + .map_err(|_| Error::Unavailable) + }) + }) + .await + .map_err(|_| Error::Unavailable)? + } + } + + fn async_get_cache( + &self, + key: &str, + context: &Self::Context, + ) -> impl Future, Error>> + Send { + let key = key.to_owned(); + let prompt = prompt_from_context(context); + let metadata = context.metadata.clone(); + async move { + let Some(prompt) = prompt else { + return Ok(None); + }; + let embedding = self + .embedder + .async_embed(&prompt, metadata.as_ref()) + .await?; + let connections = Arc::clone(&self.connections); + let config = self.config.clone(); + let index_dimension = Arc::clone(&self.index_dimension); + let threshold = config.similarity_threshold; + tokio::task::spawn_blocking(move || { + let prefix = format!("{}:", config.index_name); + ensure_index( + &connections, + &config.index_name, + &prefix, + &index_dimension, + embedding.len(), + )?; + let scope = scope_tag(&key); + let query = format!( + "(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]" + ); + let vector = embedding_bytes(&embedding); + let response = connections.execute(|connection| { + redis::cmd("FT.SEARCH") + .arg(&config.index_name) + .arg(query) + .arg("PARAMS") + .arg(2) + .arg("vec") + .arg(vector) + .arg("RETURN") + .arg(2) + .arg("response") + .arg("vector_distance") + .arg("DIALECT") + .arg(2) + .query::(connection) + .map_err(|_| Error::Unavailable) + })?; + let Some(fields) = search_fields(response)? else { + return Ok(None); + }; + let response = fields + .iter() + .find_map(|(name, value)| (name == "response").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = fields + .iter() + .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = parse_f64(&distance)?; + if 1.0 - distance < threshold { + return Ok(None); + } + Ok(Some(response)) + }) + .await + .map_err(|_| Error::Unavailable)? + .and_then(|response| response.map(|bytes| self.codec.decode(&bytes)).transpose()) + } + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Err(Error::UnsupportedOperation) + } +} + +pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { + if let Some(Value::Array(messages)) = context.messages.as_ref() + && !messages.is_empty() + { + return Some( + messages + .iter() + .filter_map(Value::as_object) + .map(message_text) + .collect(), + ); + } + let input = context.input.as_ref()?; + let mut parts = Vec::new(); + collect_input_text(input, &mut parts); + let prompt = parts.join("\n").trim().to_owned(); + (!prompt.is_empty()).then_some(prompt) +} + +fn message_text(message: &serde_json::Map) -> String { + let content = match message.get("content") { + Some(Value::String(value)) => value.clone(), + Some(Value::Array(parts)) => parts + .iter() + .filter_map(Value::as_object) + .filter_map(|part| part.get("text").and_then(Value::as_str)) + .filter(|text| !text.is_empty()) + .collect(), + _ => String::new(), + }; + format!( + "{content}{}", + search_results_text(message.get("search_results")) + ) +} + +fn search_results_text(value: Option<&Value>) -> String { + let Some(Value::Array(results)) = value else { + return String::new(); + }; + results + .iter() + .filter_map(Value::as_object) + .map(|result| { + let source = result.get("source").and_then(Value::as_str).unwrap_or(""); + let title = result.get("title").and_then(Value::as_str).unwrap_or(""); + let content = result + .get("content") + .and_then(Value::as_array) + .map(|blocks| { + blocks + .iter() + .filter_map(Value::as_object) + .filter_map(|block| block.get("text").and_then(Value::as_str)) + .collect::() + }) + .unwrap_or_default(); + let citations = result + .get("citations") + .filter(|value| !value.is_null()) + .and_then(|value| serde_json::to_string(value).ok()) + .unwrap_or_default(); + format!("{source}{title}{content}{citations}") + }) + .collect() +} + +fn collect_input_text(value: &Value, parts: &mut Vec) { + match value { + Value::String(value) => { + let value = value.trim(); + if !value.is_empty() { + parts.push(value.to_owned()); + } + } + Value::Array(values) => values + .iter() + .for_each(|value| collect_input_text(value, parts)), + Value::Object(object) => { + if let Some(content) = object.get("content").filter(|value| !value.is_null()) { + collect_input_text(content, parts); + return; + } + for key in ["text", "output", "input_text", "output_text"] { + if let Some(Value::String(value)) = object.get(key) { + let value = value.trim(); + if !value.is_empty() { + parts.push(value.to_owned()); + return; + } + } + } + } + _ => {} + } +} + +fn scope_tag(key: &str) -> String { + let digest = Sha256::digest(key.as_bytes()); + digest.iter().map(|byte| format!("{byte:02x}")).collect() +} + +fn embedding_bytes(embedding: &[f32]) -> Vec { + embedding + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() +} + +fn ensure_index( + connections: &Connections, + index_name: &str, + prefix: &str, + index_dimension: &Mutex>, + dimension: usize, +) -> Result<(), Error> +where + C: redis::ConnectionLike + Send + 'static, +{ + if index_dimension + .lock() + .map_err(|_| Error::Unavailable)? + .is_some_and(|existing| existing == dimension) + { + return Ok(()); + } + let create = connections.execute(|connection| { + Ok(redis::cmd("FT.CREATE") + .arg(index_name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(prefix) + .arg("SCHEMA") + .arg("litellm_cache_key") + .arg("TAG") + .arg("embedding") + .arg("VECTOR") + .arg("HNSW") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dimension) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .query::(connection) + .map(|_| ()) + .map_err(|error| error.to_string())) + })?; + if let Err(message) = create { + if !message.to_ascii_lowercase().contains("already exists") { + return Err(Error::Unavailable); + } + let info = connections.execute(|connection| { + redis::cmd("FT.INFO") + .arg(index_name) + .query::(connection) + .map_err(|_| Error::Unavailable) + })?; + let existing = index_dimension_from_info(&info).ok_or(Error::Unavailable)?; + if existing != dimension { + return Err(Error::Unavailable); + } + } + *index_dimension.lock().map_err(|_| Error::Unavailable)? = Some(dimension); + Ok(()) +} + +fn index_dimension_from_info(value: &redis::Value) -> Option { + let redis::Value::Array(values) = value else { + return None; + }; + let attributes = values.windows(2).find_map(|pair| { + (value_text(&pair[0]).as_deref() == Some("attributes")).then_some(&pair[1]) + })?; + let redis::Value::Array(fields) = attributes else { + return None; + }; + fields.iter().find_map(|field| { + let redis::Value::Array(values) = field else { + return None; + }; + let flattened = values.iter().flat_map(|value| match value { + redis::Value::Array(values) => values.as_slice(), + _ => std::slice::from_ref(value), + }); + let values = flattened.collect::>(); + values.windows(2).find_map(|pair| { + if value_text(pair[0]).as_deref() == Some("dimensions") { + return value_text(pair[1]).and_then(|value| value.parse().ok()); + } + None + }) + }) +} + +type SearchFields = Vec<(String, Vec)>; + +fn search_fields(value: redis::Value) -> Result, Error> { + let redis::Value::Array(values) = value else { + return Err(Error::InvalidEntry); + }; + let total = parse_i64(values.first().ok_or(Error::InvalidEntry)?)?; + if total <= 0 || values.len() < 3 { + return Ok(None); + } + let redis::Value::Array(fields) = &values[2] else { + return Err(Error::InvalidEntry); + }; + let (pairs, remainder) = fields.as_chunks::<2>(); + if !remainder.is_empty() { + return Err(Error::InvalidEntry); + } + let pairs = pairs + .iter() + .map(|pair| { + Ok(( + value_text(&pair[0]).ok_or(Error::InvalidEntry)?, + value_bytes(&pair[1])?, + )) + }) + .collect::, Error>>()?; + Ok(Some(pairs)) +} + +fn parse_i64(value: &redis::Value) -> Result { + value_text(value) + .ok_or(Error::InvalidEntry)? + .parse() + .map_err(|_| Error::InvalidEntry) +} + +fn parse_f64(value: &[u8]) -> Result { + std::str::from_utf8(value) + .map_err(|_| Error::InvalidEntry)? + .parse() + .map_err(|_| Error::InvalidEntry) +} + +fn value_text(value: &redis::Value) -> Option { + match value { + redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), + redis::Value::SimpleString(value) => Some(value.clone()), + redis::Value::Int(value) => Some(value.to_string()), + _ => None, + } +} + +fn value_bytes(value: &redis::Value) -> Result, Error> { + match value { + redis::Value::BulkString(bytes) => Ok(bytes.clone()), + redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()), + redis::Value::Int(value) => Ok(value.to_string().into_bytes()), + _ => Err(Error::InvalidEntry), + } +} + +#[cfg(test)] +mod tests { + use std::sync::{Arc, Mutex}; + + use litellm_cache::BaseCache; + use litellm_cache_response::ResponseCacheCodec; + use redis_test::MockRedisConnection; + use rstest::rstest; + use serde_json::{Value, json}; + + use super::{ + Embedder, ValkeySemanticCache, ValkeySemanticConfig, index_dimension_from_info, + prompt_from_context, scope_tag, + }; + + #[derive(Clone)] + struct FixedEmbedder { + vector: Vec, + calls: EmbedderCalls, + } + + type EmbedderCalls = Arc)>>>; + + impl Embedder for FixedEmbedder { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, super::Error> { + self.calls + .lock() + .unwrap() + .push((prompt.into(), metadata.cloned())); + Ok(self.vector.clone()) + } + + async fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> Result, super::Error> { + self.embed(prompt, metadata) + } + } + + fn context( + messages: Option, + input: Option, + ) -> litellm_cache::SemanticCacheContext { + litellm_cache::SemanticCacheContext { + messages, + input, + ..Default::default() + } + } + + #[rstest] + #[case(json!([{"content": "hello"}]), None, Some("hello"))] + #[case(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))] + #[case(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))] + #[case(Value::Array(vec![]), Some(json!(" hello ")), Some("hello"))] + #[case(Value::Array(vec![]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))] + #[case(Value::Array(vec![]), Some(json!(" ")), None)] + fn prompt_shapes( + #[case] messages: Value, + #[case] input: Option, + #[case] expected: Option<&str>, + ) { + assert_eq!( + prompt_from_context(&context(Some(messages), input)), + expected.map(str::to_owned) + ); + } + + #[test] + fn scope_tags_are_lowercase_sha256() { + assert_eq!( + scope_tag("key"), + "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683" + ); + } + + #[test] + fn existing_index_dimension_is_read_from_attributes() { + let info = redis::Value::Array(vec![ + redis::Value::SimpleString("attributes".into()), + redis::Value::Array(vec![redis::Value::Array(vec![ + redis::Value::SimpleString("identifier".into()), + redis::Value::SimpleString("embedding".into()), + redis::Value::Array(vec![ + redis::Value::SimpleString("dimensions".into()), + redis::Value::SimpleString("2".into()), + ]), + ])]), + ]); + assert_eq!(index_dimension_from_info(&info), Some(2)); + } + + #[tokio::test] + async fn unsupported_connection_test_is_reported() { + let cache = ValkeySemanticCache::with_connection( + MockRedisConnection::new([]).assert_all_commands_consumed(), + FixedEmbedder { + vector: vec![1.0, 0.0], + calls: Arc::default(), + }, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold: 0.8, + index_name: "test".into(), + }, + ); + assert_eq!( + cache.test_connection().await, + Err(super::Error::UnsupportedOperation) + ); + } + + #[test] + fn missing_prompt_does_not_touch_redis() { + let cache = ValkeySemanticCache::with_connection( + MockRedisConnection::new([]).assert_all_commands_consumed(), + FixedEmbedder { + vector: vec![1.0, 0.0], + calls: Arc::default(), + }, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold: 0.8, + index_name: "test".into(), + }, + ); + assert_eq!(cache.get_cache("key", &context(None, None)).unwrap(), None); + assert_eq!(cache.get_ttl(&context(None, None)), None); + } +} From 4db34812449038fed2c724a3ef099fefb3198ac2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:33:32 +0000 Subject: [PATCH 04/18] feat(python-bridge): serve ValkeySemanticCache natively Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 2 + litellm-rust/crates/python-bridge/Cargo.toml | 2 + .../crates/python-bridge/src/cache/config.rs | 84 +++++- .../python-bridge/src/cache/embedder.rs | 58 ++++ .../crates/python-bridge/src/cache/facade.rs | 58 +++- .../crates/python-bridge/src/cache/handle.rs | 44 ++- .../crates/python-bridge/src/cache/mod.rs | 4 +- .../crates/python-bridge/src/cache/native.rs | 253 +++++++++++++----- .../crates/python-bridge/src/cache/request.rs | 39 ++- .../test_valkey_semantic_cache.py | 149 +++++++++++ 10 files changed, 598 insertions(+), 95 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/cache/embedder.rs create mode 100644 tests/test_litellm_rust/test_valkey_semantic_cache.py diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 5daf691d4c3..e4c9c385f1f 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2685,6 +2685,7 @@ dependencies = [ "litellm-cache-memory", "litellm-cache-redis", "litellm-cache-response", + "litellm-cache-valkey-semantic", "litellm-callbacks-legacy-python", "litellm-core", "litellm-core-utils", @@ -2695,6 +2696,7 @@ dependencies = [ "litellm-types", "pyo3", "pyo3-async-runtimes", + "redis", "rstest", "serde", "serde_json", diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1eb2ec28036..ce2405f33c4 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -24,6 +24,7 @@ litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true litellm-cache-response.workspace = true +litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" } serde.workspace = true litellm-auth.workspace = true litellm-callbacks-legacy-python.workspace = true @@ -37,6 +38,7 @@ litellm-host-python.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true +redis = { version = "1.7.0", features = ["tls-rustls"] } serde_json.workspace = true tokio = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 0e7d6aee11d..7218805b8df 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -73,9 +73,18 @@ pub(super) struct RedisCacheConfig { pub(super) connection: RedisConnectionConfig, } +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] +pub(super) struct ValkeySemanticCacheConfig { + pub(super) similarity_threshold: f64, + pub(super) index_name: String, + pub(super) embedding_model: String, + pub(super) connection: RedisConnectionConfig, +} + pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + ValkeySemantic(Box), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -142,9 +151,15 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::ValkeySemantic) => match project_valkey_semantic(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::ValkeySemantic(Box::new(backend)), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, Some( CacheType::RedisSemantic - | CacheType::ValkeySemantic | CacheType::S3 | CacheType::Disk | CacheType::QdrantSemantic @@ -158,11 +173,13 @@ 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, - }) + if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_)) + && service.default_ttl() + != Some(match &self.backend { + CacheBackendConfig::Memory(config) => config.default_ttl, + CacheBackendConfig::Redis(config) => config.default_ttl, + CacheBackendConfig::ValkeySemantic(_) => Duration::ZERO, + }) { return Some("facade and native backend default TTLs must match"); } @@ -185,6 +202,16 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), + CacheBackendConfig::ValkeySemantic(config) => { + if service.kind() != "valkey-semantic" { + return Some("facade and native backend types must match"); + } + let Some((threshold, index_name)) = service.semantic_config() else { + return Some("facade and native backend types must match"); + }; + (threshold != config.similarity_threshold || index_name != config.index_name) + .then_some("facade and native semantic settings must match") + } } } } @@ -299,6 +326,51 @@ fn project_redis( })) } +#[inline(never)] +fn project_valkey_semantic( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let client = backend.getattr("sync_client")?; + let pool = client.getattr("connection_pool")?; + if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; + let connection_class = resolved + .get_item("connection_class")? + .unwrap_or(pool.getattr("connection_class")?); + if !class_is(&connection_class, "redis.connection", "Connection")? + && !class_is(&connection_class, "redis.connection", "SSLConnection")? + { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + let connection = RedisConnectionConfig { + host: required_string(&resolved, "host")?, + port: u16::try_from(required_i64(&resolved, "port")?) + .map_err(|_| PyValueError::new_err("invalid Redis port"))?, + database: optional_i64(&resolved, "db")?.unwrap_or(0), + username: optional_dict_string(&resolved, "username")?, + password: optional_dict_string(&resolved, "password")?, + protocol: RedisProtocol::Resp2, + pool_size: pool.getattr("max_connections")?.extract::()?, + read_timeout: None, + connect_timeout: None, + socket_keepalive: None, + health_check_interval: Duration::ZERO, + client_name: None, + tls: None, + }; + if connection.host.is_empty() { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + Ok(Ok(ValkeySemanticCacheConfig { + similarity_threshold: backend.getattr("similarity_threshold")?.extract()?, + index_name: backend.getattr("index_name")?.extract()?, + embedding_model: backend.getattr("embedding_model")?.extract()?, + connection, + })) +} + #[inline(never)] fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { Ok(RedisTlsConfig { 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..d240d9d019e --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -0,0 +1,58 @@ +use std::{future::Future, sync::Arc}; + +use litellm_cache::Error; +use litellm_cache_valkey_semantic::Embedder; +use litellm_host_python::to_py; +use pyo3::prelude::*; +use serde_json::Value; + +#[derive(Clone)] +pub(super) struct PythonEmbedder { + sync_embed: Arc>, + async_embed: Arc>, +} + +impl PythonEmbedder { + pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult { + Ok(Self { + sync_embed: Arc::new(backend.getattr("_get_embedding")?.unbind()), + async_embed: Arc::new(backend.getattr("_get_async_embedding")?.unbind()), + }) + } +} + +impl Embedder for PythonEmbedder { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + let result = Python::attach(|py| -> PyResult> { + let metadata = to_py(py, &metadata)?; + self.sync_embed + .bind(py) + .call1((prompt, metadata))? + .extract() + }) + .map_err(|_| Error::Unavailable)?; + Ok(result.into_iter().map(|value| value as f32).collect()) + } + + fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> impl Future, Error>> + Send { + let callable = Arc::clone(&self.async_embed); + let prompt = prompt.to_owned(); + let metadata = metadata.cloned(); + async move { + let future = Python::attach(|py| -> PyResult<_> { + let metadata = to_py(py, &metadata)?; + let awaitable = callable.bind(py).call1((prompt, metadata))?; + pyo3_async_runtimes::tokio::into_future(awaitable) + }) + .map_err(|_| Error::Unavailable)?; + let result = future.await.map_err(|_| Error::Unavailable)?; + let result = Python::attach(|py| result.bind(py).extract::>()) + .map_err(|_| Error::Unavailable)?; + Ok(result.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..eb76d22097d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -36,6 +36,7 @@ pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, redis_pool: Option, + redis_client_name: Option<&'static str>, } impl ObjectGuard { @@ -117,7 +118,9 @@ impl ObjectGuard { return Ok(false); } for (name, value) in &expected.attributes { - if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) { + if (instance.contains(name)? && !self.config_names.contains(&name.as_str())) + || !attributes.get_item(name)?.is(value.bind(py)) + { return Ok(false); } } @@ -138,10 +141,8 @@ impl ObjectGuard { } impl RedisPoolGuard { - fn capture(backend: &Bound<'_, PyAny>) -> PyResult { - let pool = backend - .getattr("redis_client")? - .getattr("connection_pool")?; + fn capture(backend: &Bound<'_, PyAny>, client_name: &str) -> PyResult { + let pool = backend.getattr(client_name)?.getattr("connection_pool")?; Ok(Self { reference: pool.clone().unbind(), connection_class: pool.getattr("connection_class")?.unbind(), @@ -153,10 +154,13 @@ impl RedisPoolGuard { }) } - fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { - let pool = backend - .getattr("redis_client")? - .getattr("connection_pool")?; + fn matches( + &self, + py: Python<'_>, + backend: &Bound<'_, PyAny>, + client_name: &str, + ) -> PyResult { + let pool = backend.getattr(client_name)?.getattr("connection_pool")?; Ok(self.reference.bind(py).is(&pool) && self .connection_class @@ -192,6 +196,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"), + "valkey-semantic" => ( + "litellm.caching.valkey_semantic_cache", + "ValkeySemanticCache", + "valkey-semantic", + ), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -235,11 +244,32 @@ impl FacadeGuard { "max_size_per_item", "redis_kwargs", "redis_flush_size", + "similarity_threshold", + "embedding_model", + "index_name", + "embedding_max_input_tokens", + "embedding_timeout", ], )?, - redis_pool: (kind == "redis") - .then(|| RedisPoolGuard::capture(&backend)) + redis_pool: (kind == "redis" || kind == "valkey-semantic") + .then(|| { + RedisPoolGuard::capture( + &backend, + if kind == "redis" { + "redis_client" + } else { + "sync_client" + }, + ) + }) .transpose()?, + redis_client_name: (kind == "redis" || kind == "valkey-semantic").then_some( + if kind == "redis" { + "redis_client" + } else { + "sync_client" + }, + ), }) } @@ -252,7 +282,11 @@ impl FacadeGuard { return Ok(false); } match &self.redis_pool { - Some(guard) => guard.matches(py, &backend), + Some(guard) => guard.matches( + py, + &backend, + self.redis_client_name.unwrap_or("redis_client"), + ), None => Ok(true), } } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 8251b3df06c..119bd35cd25 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,7 +1,10 @@ use litellm_host_python::release_gil; use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; -use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; +use super::{ + cache_error, embedder::PythonEmbedder, facade::FacadeGuard, native::NativeResponseCache, + request::duration, +}; #[pyclass(frozen, name = "_CacheTestHandle")] pub(crate) struct CacheTestHandle { @@ -51,6 +54,29 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[pyo3(signature = (url, similarity_threshold, index_name, embedder))] + fn valkey_semantic( + url: String, + similarity_threshold: f64, + index_name: String, + embedder: &Bound<'_, PyAny>, + ) -> PyResult { + let python_embedder = PythonEmbedder::from_backend(embedder)?; + let service = NativeResponseCache::valkey_semantic( + &url, + similarity_threshold, + index_name, + python_embedder, + ) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() @@ -59,11 +85,17 @@ impl CacheTestHandle { fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { let service = self.service()?; let guard = FacadeGuard::capture(py, facade, &service)?; - let service = service.with_redis_flush_size( - facade - .getattr("redis_flush_size")? - .extract::>()?, - ); + let service = service + .with_scope( + facade + .getattr("semantic_cache_scope")? + .extract::()?, + ) + .with_redis_flush_size( + facade + .getattr("redis_flush_size")? + .extract::>()?, + ); let handle = Py::new( py, Self { 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..d314cd41ac5 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,13 +1,18 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; +use litellm_cache::{ + CacheCodec, CacheConnectionResult, Error, ExactCacheContext, SemanticCacheContext, +}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, }; +use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig}; use serde_json::Value; +use super::{embedder::PythonEmbedder, request::NativeRequest}; + #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), @@ -15,6 +20,10 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + ValkeySemantic { + cache: Arc>>, + scope: String, + }, } impl NativeResponseCache { @@ -43,41 +52,52 @@ impl NativeResponseCache { buffer: None, }) } -} -impl NativeResponseCache { - pub fn kind(&self) -> &'static str { - match self { - Self::Memory(_) => "memory", - Self::Redis { .. } => "redis", + pub fn valkey_semantic( + url: &str, + similarity_threshold: f64, + index_name: String, + embedder: PythonEmbedder, + ) -> Result { + let backend = ValkeySemanticCache::new( + url, + embedder, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold, + index_name, + }, + )?; + Ok(Self::ValkeySemantic { + cache: Arc::new(ResponseCache::new(Arc::new(backend))), + scope: String::from("key"), + }) + } + + fn exact(request: &NativeRequest) -> ResponseCacheRequest { + ResponseCacheRequest { + key: request.key.clone(), + controls: request.controls, + context: ExactCacheContext { ttl: request.ttl }, + max_age: request.max_age, } } - pub fn default_ttl(&self) -> Option { - match self { - Self::Memory(cache) => cache.default_ttl(), - Self::Redis { cache, .. } => cache.default_ttl(), - } - } - - pub fn namespace(&self) -> Option<&str> { - match self { - Self::Memory(_) => None, - Self::Redis { cache, .. } => cache.backend().namespace(), - } - } - - pub fn capacity(&self) -> Option { - match self { - Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), - Self::Redis { .. } => None, - } - } - - pub fn max_entry_bytes(&self) -> Option { - match self { - Self::Memory(cache) => cache.backend().max_entry_bytes(), - Self::Redis { .. } => None, + fn semantic( + request: &NativeRequest, + scope: &str, + ) -> ResponseCacheRequest { + ResponseCacheRequest { + key: request.key.clone(), + controls: request.controls, + context: SemanticCacheContext { + input: request.input.clone(), + messages: request.messages.clone(), + metadata: request.metadata.clone(), + scope: Some(scope.to_owned()), + ttl: request.ttl, + }, + max_age: request.max_age, } } @@ -85,95 +105,206 @@ impl NativeResponseCache { match self { Self::Redis { cache, .. } => Self::Redis { cache, - buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))), + buffer: flush_size.map(|size| Arc::new(WriteBuffer::new(size))), }, - memory => memory, + value => value, } } - pub fn lookup( - &self, - request: &ResponseCacheRequest, - now: Duration, - ) -> Result, Error> { + pub fn with_scope(self, scope: String) -> Self { match self { - Self::Memory(cache) => cache.lookup(request, now), - Self::Redis { cache, .. } => cache.lookup(request, now), + Self::ValkeySemantic { cache, .. } => Self::ValkeySemantic { cache, scope }, + value => value, + } + } + + pub fn kind(&self) -> &'static str { + match self { + Self::Memory(_) => "memory", + Self::Redis { .. } => "redis", + Self::ValkeySemantic { .. } => "valkey-semantic", + } + } + + pub fn default_ttl(&self) -> Option { + match self { + Self::Memory(cache) => cache.default_ttl(), + Self::Redis { cache, .. } => cache.default_ttl(), + Self::ValkeySemantic { cache, .. } => cache.default_ttl(), + } + } + + pub fn namespace(&self) -> Option<&str> { + match self { + Self::Memory(_) | Self::ValkeySemantic { .. } => None, + Self::Redis { cache, .. } => cache.backend().namespace(), + } + } + + pub fn capacity(&self) -> Option { + match self { + Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), + Self::Redis { .. } | Self::ValkeySemantic { .. } => None, + } + } + + pub fn max_entry_bytes(&self) -> Option { + match self { + Self::Memory(cache) => cache.backend().max_entry_bytes(), + Self::Redis { .. } | Self::ValkeySemantic { .. } => None, + } + } + + pub fn semantic_config(&self) -> Option<(f64, &str)> { + match self { + Self::ValkeySemantic { cache, .. } => Some(( + cache.backend().similarity_threshold(), + cache.backend().index_name(), + )), + _ => None, + } + } + + pub fn lookup(&self, request: &NativeRequest, now: Duration) -> Result, Error> { + match self { + Self::Memory(cache) => cache.lookup(&Self::exact(request), now), + Self::Redis { cache, .. } => cache.lookup(&Self::exact(request), now), + Self::ValkeySemantic { cache, scope } => { + cache.lookup(&Self::semantic(request, scope), now) + } } } pub fn store( &self, - request: &ResponseCacheRequest, + request: &NativeRequest, 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(&Self::exact(request), response, now), + Self::Redis { cache, .. } => cache.store(&Self::exact(request), response, now), + Self::ValkeySemantic { cache, scope } => { + cache.store(&Self::semantic(request, scope), response, now) + } } } pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[NativeRequest], now: Duration, ) -> Result { match self { - Self::Memory(cache) => cache.lookup_batch(requests, now), - Self::Redis { cache, .. } => cache.lookup_batch(requests, now), + Self::Memory(cache) => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.lookup_batch(&requests, now) + } + Self::Redis { cache, .. } => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.lookup_batch(&requests, now) + } + Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), } } pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &NativeRequest, 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(&Self::exact(request), now).await, + Self::Redis { cache, .. } => cache.async_lookup(&Self::exact(request), now).await, + Self::ValkeySemantic { cache, scope } => { + cache + .async_lookup(&Self::semantic(request, scope), now) + .await + } } } pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &NativeRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_store(request, response, now).await, + Self::Memory(cache) => { + cache + .async_store(&Self::exact(request), response, now) + .await + } Self::Redis { cache, buffer: None, - } => cache.async_store(request, response, now).await, + } => { + cache + .async_store(&Self::exact(request), response, now) + .await + } Self::Redis { cache, buffer: Some(buffer), - } => buffer.async_store(cache, request, response, now).await, + } => { + buffer + .async_store(cache, &Self::exact(request), response, now) + .await + } + Self::ValkeySemantic { cache, scope } => { + cache + .async_store(&Self::semantic(request, scope), response, now) + .await + } } } pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[NativeRequest], 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) => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.async_lookup_batch(&requests, now).await + } + Self::Redis { cache, .. } => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.async_lookup_batch(&requests, now).await + } + Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), } } pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(NativeRequest, 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) => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + Self::Redis { cache, .. } => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + Self::ValkeySemantic { cache, scope } => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::semantic(&request, scope), value)) + .collect(); + cache.async_store_batch(entries, now).await + } } } @@ -186,6 +317,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), } } @@ -193,6 +325,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::ValkeySemantic { cache, .. } => cache.test_connection().await, } } } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 52a5f7d9055..3e19e7fdc22 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -5,6 +5,7 @@ use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest} use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; use serde::Deserialize; +use serde_json::Value; #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -13,24 +14,42 @@ struct RequestInput { controls: Option, ttl_seconds: Option, max_age_seconds: Option, + messages: Option, + input: Option, + metadata: Option, } -pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { +pub(super) struct NativeRequest { + pub(super) key: CacheKeyInput, + pub(super) controls: CacheControls, + pub(super) ttl: Option, + pub(super) max_age: Option, + pub(super) messages: Option, + pub(super) input: Option, + pub(super) metadata: Option, +} + +pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { let input: RequestInput = from_py(value)?; request_input(input) } -fn request_input(input: RequestInput) -> PyResult { - let mut request = ResponseCacheRequest::::new(input.key); - if let Some(controls) = input.controls { - request.controls = controls; - } - request.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 controls = input.controls.unwrap_or_else(|| { + ResponseCacheRequest::::new(input.key.clone()).controls + }); + Ok(NativeRequest { + key: input.key, + controls, + ttl: input.ttl_seconds.map(duration).transpose()?, + max_age: input.max_age_seconds.map(duration).transpose()?, + messages: input.messages, + input: input.input, + metadata: input.metadata, + }) } -pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { +pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { from_py::>(value)? .into_iter() .map(request_input) diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache.py b/tests/test_litellm_rust/test_valkey_semantic_cache.py new file mode 100644 index 00000000000..00037e6e29f --- /dev/null +++ b/tests/test_litellm_rust/test_valkey_semantic_cache.py @@ -0,0 +1,149 @@ +import os +from collections.abc import Generator, Mapping +from types import SimpleNamespace +from typing import Final, cast +from uuid import uuid4 + +import pytest +import redis + +from litellm.caching.caching import Cache +from litellm.caching.valkey_semantic_cache import ValkeySemanticCache +from litellm.rust_bridge import _native +from litellm.types.caching import LiteLLMCacheType + +pytestmark: Final = pytest.mark.requires_rust_extension + + +@pytest.fixture +def valkey_url() -> str: + url: Final = os.environ.get("LITELLM_TEST_VALKEY_URL") + if url is None: + pytest.skip("LITELLM_TEST_VALKEY_URL is not set") + return url + + +@pytest.fixture +def index_name(valkey_url: str) -> Generator[str]: + index: Final = f"litellm_test_{uuid4().hex}" + yield index + client: Final = redis.Redis.from_url(valkey_url) + try: + client.ft(index).dropindex(delete_documents=True) + except redis.ResponseError: + pass + finally: + client.close() + + +def _request() -> dict[str, object]: + return { + "key": {"preset": "key"}, + "messages": [{"role": "user", "content": "semantic cache prompt"}], + } + + +def _backend(url: str, index_name: str) -> ValkeySemanticCache: + backend: Final = ValkeySemanticCache( + redis_url=url, + similarity_threshold=0.8, + index_name=index_name, + ) + backend._get_embedding = lambda prompt, metadata=None: [1.0, 0.0] + + async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]: + return [1.0, 0.0] + + backend._get_async_embedding = async_embedding + return backend + + +def test_python_write_native_read( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + response: Final = {"answer": "python"} + backend.set_cache("key", response, messages=_request()["messages"]) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + assert binding.lookup(_request()) == response + + +def test_native_write_python_read( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + response: Final = {"answer": "native"} + binding.store({**_request(), "ttl_seconds": 2.0}, response) + cached: Final = cast(Mapping[str, object], backend.get_cache("key", messages=_request()["messages"])) + assert cached["response"] == response + + +async def test_async_lookup_and_store( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + request: Final = {**_request(), "ttl_seconds": 2.0} + await binding.async_store(request, {"answer": "async"}) + assert await binding.async_lookup(request) == {"answer": "async"} + + +def test_facade_activation_and_mutation_fallback( + valkey_url: str, + index_name: str, +) -> None: + facade: Final = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + redis_url=valkey_url, + similarity_threshold=0.8, + valkey_semantic_cache_index_name=index_name, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + facade.cache, + ) + handle._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + facade.cache.similarity_threshold = 0.7 + assert resolver.resolve().kind == "python_callback" + + +def test_batch_lookup_is_unsupported( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + with pytest.raises(NotImplementedError): + binding.lookup_batch([_request()]) From f6db876a3d72e143fd6638f226ac4c01a34ca088 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:34:21 +0000 Subject: [PATCH 05/18] test(cache): disambiguate Valkey semantic test module Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...key_semantic_cache.py => test_valkey_semantic_cache_native.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/test_litellm_rust/{test_valkey_semantic_cache.py => test_valkey_semantic_cache_native.py} (100%) diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache.py b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py similarity index 100% rename from tests/test_litellm_rust/test_valkey_semantic_cache.py rename to tests/test_litellm_rust/test_valkey_semantic_cache_native.py From 56237af7a9dbfcdb20b41b3959620128c8fbccd2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:42:41 +0000 Subject: [PATCH 06/18] test(cache): add Valkey semantic contract coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-valkey-semantic/src/lib.rs | 659 +++++++++++++----- .../crates/python-bridge/src/cache/config.rs | 84 ++- .../test_valkey_semantic_cache_native.py | 140 +++- 3 files changed, 695 insertions(+), 188 deletions(-) diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs index 85c4c9af15c..ef028f0a7b2 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs @@ -62,6 +62,14 @@ enum Connections { Fixed(Mutex), } +#[derive(Clone)] +struct IndexState { + name: String, + prefix: String, + dimension: Arc>>, + similarity_threshold: f64, +} + struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); impl redis::ConnectionLike for ConnectionRef<'_> { @@ -187,18 +195,13 @@ where &self.config.index_name } - fn key_prefix(&self) -> String { - format!("{}:", self.config.index_name) - } - - fn ensure_index(&self, dimension: usize) -> Result<(), Error> { - ensure_index( - &self.connections, - &self.config.index_name, - &self.key_prefix(), - &self.index_dimension, - dimension, - ) + fn index_state(&self) -> IndexState { + IndexState { + name: self.config.index_name.clone(), + prefix: format!("{}:", self.config.index_name), + dimension: Arc::clone(&self.index_dimension), + similarity_threshold: self.config.similarity_threshold, + } } } @@ -225,37 +228,19 @@ where return Ok(()); }; let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; - self.ensure_index(embedding.len())?; let scope = scope_tag(key); - let document = format!("{}{}:{}", self.key_prefix(), scope, Uuid::new_v4()); let response = self.codec.encode(&value)?; let vector = embedding_bytes(&embedding); - let ttl = self.get_ttl(context); - self.connections.execute(|connection| { - let mut pipeline = redis::pipe(); - pipeline - .cmd("HSET") - .arg(&document) - .arg("litellm_cache_key") - .arg(&scope) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(response) - .arg("embedding") - .arg(vector) - .ignore(); - if let Some(ttl) = ttl { - pipeline - .cmd("EXPIRE") - .arg(&document) - .arg(ttl.as_secs()) - .ignore(); - } - pipeline - .query::<()>(connection) - .map_err(|_| Error::Unavailable) - }) + let index = self.index_state(); + write_document( + &self.connections, + &index, + &scope, + &prompt, + response, + vector, + self.get_ttl(context), + ) } fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { @@ -263,43 +248,14 @@ where return Ok(None); }; let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; - self.ensure_index(embedding.len())?; let scope = scope_tag(key); - let query = - format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]"); let vector = embedding_bytes(&embedding); - let response = self.connections.execute(|connection| { - redis::cmd("FT.SEARCH") - .arg(&self.config.index_name) - .arg(query) - .arg("PARAMS") - .arg(2) - .arg("vec") - .arg(vector) - .arg("RETURN") - .arg(2) - .arg("response") - .arg("vector_distance") - .arg("DIALECT") - .arg(2) - .query::(connection) - .map_err(|_| Error::Unavailable) - })?; - let Some(fields) = search_fields(response)? else { + let index = self.index_state(); + let Some(response) = + search_document(&self.connections, &index, &scope, vector, embedding.len())? + else { return Ok(None); }; - let response = fields - .iter() - .find_map(|(name, value)| (name == "response").then(|| value.clone())) - .ok_or(Error::InvalidEntry)?; - let distance = fields - .iter() - .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) - .ok_or(Error::InvalidEntry)?; - let distance = parse_f64(&distance)?; - if 1.0 - distance < self.config.similarity_threshold { - return Ok(None); - } self.codec.decode(&response).map(Some) } @@ -321,47 +277,13 @@ where .async_embed(&prompt, metadata.as_ref()) .await?; let connections = Arc::clone(&self.connections); - let config = self.config.clone(); - let index_dimension = Arc::clone(&self.index_dimension); + let index = self.index_state(); let response = self.codec.encode(&value)?; let vector = embedding_bytes(&embedding); - let prefix = format!("{}:", config.index_name); let scope = scope_tag(&key); - let document = format!("{prefix}{scope}:{}", Uuid::new_v4()); let ttl = context.ttl; tokio::task::spawn_blocking(move || { - ensure_index( - &connections, - &config.index_name, - &prefix, - &index_dimension, - embedding.len(), - )?; - connections.execute(|connection| { - let mut pipeline = redis::pipe(); - pipeline - .cmd("HSET") - .arg(&document) - .arg("litellm_cache_key") - .arg(&scope) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(response) - .arg("embedding") - .arg(vector) - .ignore(); - if let Some(ttl) = ttl { - pipeline - .cmd("EXPIRE") - .arg(&document) - .arg(ttl.as_secs()) - .ignore(); - } - pipeline - .query::<()>(connection) - .map_err(|_| Error::Unavailable) - }) + write_document(&connections, &index, &scope, &prompt, response, vector, ttl) }) .await .map_err(|_| Error::Unavailable)? @@ -385,56 +307,11 @@ where .async_embed(&prompt, metadata.as_ref()) .await?; let connections = Arc::clone(&self.connections); - let config = self.config.clone(); - let index_dimension = Arc::clone(&self.index_dimension); - let threshold = config.similarity_threshold; + let index = self.index_state(); tokio::task::spawn_blocking(move || { - let prefix = format!("{}:", config.index_name); - ensure_index( - &connections, - &config.index_name, - &prefix, - &index_dimension, - embedding.len(), - )?; let scope = scope_tag(&key); - let query = format!( - "(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]" - ); let vector = embedding_bytes(&embedding); - let response = connections.execute(|connection| { - redis::cmd("FT.SEARCH") - .arg(&config.index_name) - .arg(query) - .arg("PARAMS") - .arg(2) - .arg("vec") - .arg(vector) - .arg("RETURN") - .arg(2) - .arg("response") - .arg("vector_distance") - .arg("DIALECT") - .arg(2) - .query::(connection) - .map_err(|_| Error::Unavailable) - })?; - let Some(fields) = search_fields(response)? else { - return Ok(None); - }; - let response = fields - .iter() - .find_map(|(name, value)| (name == "response").then(|| value.clone())) - .ok_or(Error::InvalidEntry)?; - let distance = fields - .iter() - .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) - .ok_or(Error::InvalidEntry)?; - let distance = parse_f64(&distance)?; - if 1.0 - distance < threshold { - return Ok(None); - } - Ok(Some(response)) + search_document(&connections, &index, &scope, vector, embedding.len()) }) .await .map_err(|_| Error::Unavailable)? @@ -560,6 +437,108 @@ fn embedding_bytes(embedding: &[f32]) -> Vec { .collect() } +fn write_document( + connections: &Connections, + index: &IndexState, + scope: &str, + prompt: &str, + response: Vec, + vector: Vec, + ttl: Option, +) -> Result<(), Error> +where + C: redis::ConnectionLike + Send + 'static, +{ + let dimension = vector.len() / std::mem::size_of::(); + ensure_index( + connections, + &index.name, + &index.prefix, + &index.dimension, + dimension, + )?; + let document = format!("{}{scope}:{}", index.prefix, Uuid::new_v4()); + connections.execute(|connection| { + let mut pipeline = redis::pipe(); + pipeline + .cmd("HSET") + .arg(&document) + .arg("litellm_cache_key") + .arg(scope) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg("embedding") + .arg(vector) + .ignore(); + if let Some(ttl) = ttl { + pipeline + .cmd("EXPIRE") + .arg(&document) + .arg(ttl.as_secs()) + .ignore(); + } + pipeline + .query::<()>(connection) + .map_err(|_| Error::Unavailable) + }) +} + +fn search_document( + connections: &Connections, + index: &IndexState, + scope: &str, + vector: Vec, + dimension: usize, +) -> Result>, Error> +where + C: redis::ConnectionLike + Send + 'static, +{ + ensure_index( + connections, + &index.name, + &index.prefix, + &index.dimension, + dimension, + )?; + let query = + format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]"); + let response = connections.execute(|connection| { + redis::cmd("FT.SEARCH") + .arg(&index.name) + .arg(query) + .arg("PARAMS") + .arg(2) + .arg("vec") + .arg(vector) + .arg("RETURN") + .arg(2) + .arg("response") + .arg("vector_distance") + .arg("DIALECT") + .arg(2) + .query::(connection) + .map_err(|_| Error::Unavailable) + })?; + let Some(fields) = search_fields(response)? else { + return Ok(None); + }; + let response = fields + .iter() + .find_map(|(name, value)| (name == "response").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = fields + .iter() + .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = parse_f64(&distance)?; + if 1.0 - distance < index.similarity_threshold { + return Ok(None); + } + Ok(Some(response)) +} + fn ensure_index( connections: &Connections, index_name: &str, @@ -712,10 +691,16 @@ fn value_bytes(value: &redis::Value) -> Result, Error> { #[cfg(test)] mod tests { - use std::sync::{Arc, Mutex}; + use std::{ + collections::VecDeque, + sync::{Arc, Mutex}, + time::Duration, + }; - use litellm_cache::BaseCache; - use litellm_cache_response::ResponseCacheCodec; + use litellm_cache::{BaseCache, CacheCodec}; + use litellm_cache_response::{ + CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, + }; use redis_test::MockRedisConnection; use rstest::rstest; use serde_json::{Value, json}; @@ -732,6 +717,9 @@ mod tests { } type EmbedderCalls = Arc)>>>; + type RecordingCache = + ValkeySemanticCache; + type RecordingSetup = (RecordingCache, Arc>>>, EmbedderCalls); impl Embedder for FixedEmbedder { fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, super::Error> { @@ -751,6 +739,61 @@ mod tests { } } + struct RecordingConnection { + requests: Arc>>>, + replies: Mutex>>, + } + + impl RecordingConnection { + fn new(replies: impl IntoIterator>) -> Self { + Self { + requests: Arc::default(), + replies: Mutex::new(replies.into_iter().collect()), + } + } + + fn requests(&self) -> Arc>>> { + Arc::clone(&self.requests) + } + + fn reply(&self) -> redis::RedisResult { + self.replies + .lock() + .unwrap() + .pop_front() + .unwrap_or_else(|| Ok(redis::Value::SimpleString("OK".into()))) + } + } + + impl redis::ConnectionLike for RecordingConnection { + fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult { + self.requests.lock().unwrap().push(command.to_vec()); + self.reply() + } + + fn req_packed_commands( + &mut self, + command: &[u8], + _offset: usize, + count: usize, + ) -> redis::RedisResult> { + self.requests.lock().unwrap().push(command.to_vec()); + (0..count).map(|_| self.reply()).collect() + } + + fn get_db(&self) -> i64 { + 0 + } + + fn check_connection(&mut self) -> bool { + true + } + + fn is_open(&self) -> bool { + true + } + } + fn context( messages: Option, input: Option, @@ -841,4 +884,302 @@ mod tests { assert_eq!(cache.get_cache("key", &context(None, None)).unwrap(), None); assert_eq!(cache.get_ttl(&context(None, None)), None); } + + fn semantic_context(ttl: Option) -> litellm_cache::SemanticCacheContext { + litellm_cache::SemanticCacheContext { + messages: Some(json!([{"role": "user", "content": "hello"}])), + metadata: Some(json!({"source": "test"})), + ttl, + ..Default::default() + } + } + + fn cache_with_recording( + replies: impl IntoIterator>, + vector: Vec, + threshold: f64, + ) -> RecordingSetup { + let connection = RecordingConnection::new(replies); + let requests = connection.requests(); + let calls: EmbedderCalls = Arc::default(); + let cache = ValkeySemanticCache::with_connection( + connection, + FixedEmbedder { + vector, + calls: Arc::clone(&calls), + }, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold: threshold, + index_name: "test".into(), + }, + ); + (cache, requests, calls) + } + + fn ok() -> redis::RedisResult { + Ok(redis::Value::SimpleString("OK".into())) + } + + fn already_exists() -> redis::RedisResult { + Err(redis::RedisError::from(( + redis::ErrorKind::Io, + "already exists", + ))) + } + + fn info_dimension(dimension: usize) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::SimpleString("attributes".into()), + redis::Value::Array(vec![redis::Value::Array(vec![ + redis::Value::SimpleString("embedding".into()), + redis::Value::Array(vec![ + redis::Value::SimpleString("dimensions".into()), + redis::Value::Int(dimension as i64), + ]), + ])]), + ]) + } + + fn search_hit(response: Vec, distance: &str) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"test:document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"response".to_vec()), + redis::Value::BulkString(response), + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(distance.as_bytes().to_vec()), + ]), + ]) + } + + fn requests_text(requests: &Arc>>>) -> String { + requests + .lock() + .unwrap() + .iter() + .map(|request| String::from_utf8_lossy(request)) + .collect::>() + .join("\n") + } + + #[test] + fn set_without_ttl_writes_hset_without_expire() { + let (cache, requests, calls) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); + cache + .set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(None), + ) + .unwrap(); + let text = requests_text(&requests); + assert!(text.contains("FT.CREATE")); + assert!(text.contains("HSET")); + assert!( + text.contains("test:2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683:") + ); + assert!(!text.contains("EXPIRE")); + assert_eq!( + *calls.lock().unwrap(), + vec![("hello".into(), Some(json!({"source": "test"})))] + ); + } + + #[test] + fn set_with_ttl_truncates_expire_seconds() { + let (cache, requests, _) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); + cache + .set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(Some(Duration::from_millis(1900))), + ) + .unwrap(); + let text = requests_text(&requests); + assert!(text.contains("EXPIRE")); + assert!(text.contains("\r\n$1\r\n1\r\n")); + } + + #[test] + fn second_set_skips_create_after_dimension_is_cached() { + let (cache, requests, _) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); + let context = semantic_context(None); + let entry = CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }; + cache.set_cache("key", entry.clone(), &context).unwrap(); + cache.set_cache("key", entry, &context).unwrap(); + let text = requests_text(&requests); + assert_eq!(text.matches("FT.CREATE").count(), 1); + assert_eq!(text.matches("HSET").count(), 2); + } + + #[test] + fn existing_index_dimension_must_match_embedding() { + let (cache, _, _) = cache_with_recording( + [already_exists(), Ok(info_dimension(2))], + vec![1.0, 0.0], + 0.8, + ); + cache + .set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(None), + ) + .unwrap(); + + let (cache, _, _) = cache_with_recording( + [already_exists(), Ok(info_dimension(3))], + vec![1.0, 0.0], + 0.8, + ); + assert_eq!( + cache.set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(None), + ), + Err(super::Error::Unavailable) + ); + } + + #[test] + fn get_applies_threshold_and_decodes_entry() { + let entry = CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "ok"}), + }; + let encoded = ResponseCacheCodec.encode(&entry).unwrap(); + let (cache, _, _) = cache_with_recording( + [ok(), Ok(search_hit(encoded.clone(), "0.1"))], + vec![1.0, 0.0], + 0.8, + ); + assert_eq!( + cache.get_cache("key", &semantic_context(None)).unwrap(), + Some(entry) + ); + + let (cache, _, _) = + cache_with_recording([ok(), Ok(search_hit(encoded, "0.5"))], vec![1.0, 0.0], 0.8); + assert_eq!( + cache.get_cache("key", &semantic_context(None)).unwrap(), + None + ); + } + + #[test] + fn get_zero_docs_is_a_miss() { + let (cache, _, _) = cache_with_recording( + [ok(), Ok(redis::Value::Array(vec![redis::Value::Int(0)]))], + vec![1.0, 0.0], + 0.8, + ); + assert_eq!( + cache.get_cache("key", &semantic_context(None)).unwrap(), + None + ); + } + + #[rstest] + #[case(redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(b"0.1".to_vec()), + ]), + ]))] + #[case(redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"response".to_vec()), + redis::Value::BulkString(b"not-json".to_vec()), + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(b"abc".to_vec()), + ]), + ]))] + fn malformed_entries_are_invalid(#[case] search: redis::Value) { + let (cache, _, _) = cache_with_recording([ok(), Ok(search)], vec![1.0, 0.0], 0.8); + assert_eq!( + cache.get_cache("key", &semantic_context(None)), + Err(super::Error::InvalidEntry) + ); + } + + #[test] + fn response_cache_turns_invalid_entries_into_misses() { + let (cache, _, _) = cache_with_recording( + [ + ok(), + Ok(redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"response".to_vec()), + redis::Value::BulkString(b"not-json".to_vec()), + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(b"0.1".to_vec()), + ]), + ])), + ], + vec![1.0, 0.0], + 0.8, + ); + let service = ResponseCache::new(Arc::new(cache)); + let request = ResponseCacheRequest { + key: CacheKeyInput { + preset: Some("key".into()), + ..Default::default() + }, + context: semantic_context(None), + ..ResponseCacheRequest::new(CacheKeyInput::default()) + }; + assert_eq!(service.lookup(&request, Duration::ZERO).unwrap(), None); + } + + #[tokio::test] + async fn async_set_and_get_use_shared_document_helpers() { + let entry = CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "ok"}), + }; + let encoded = ResponseCacheCodec.encode(&entry).unwrap(); + let (cache, requests, calls) = cache_with_recording( + [ok(), ok(), ok(), Ok(search_hit(encoded, "0.1"))], + vec![1.0, 0.0], + 0.8, + ); + let context = semantic_context(Some(Duration::from_millis(1900))); + cache + .async_set_cache("key", entry.clone(), context.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache("key", &context).await.unwrap(), + Some(entry) + ); + let text = requests_text(&requests); + assert!(text.contains("FT.CREATE")); + assert!(text.contains("HSET")); + assert!(text.contains("EXPIRE")); + assert_eq!(calls.lock().unwrap().len(), 2); + } } diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 7218805b8df..e074c5e2f5d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -276,25 +276,15 @@ fn project_redis( let client = backend.getattr("redis_client")?; let pool = client.getattr("connection_pool")?; - if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? { + let Ok((resolved, is_tls)) = project_connection_pool(&pool)? else { return Ok(Err(UnsupportedCacheConfig::RedisConnection)); - } - let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; + }; for key in ["credential_provider", "redis_connect_func"] { if has_value(&resolved, key)? { return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); } } - let connection_class = resolved - .get_item("connection_class")? - .unwrap_or(pool.getattr("connection_class")?); - let tls = if class_is(&connection_class, "redis.connection", "Connection")? { - None - } else if class_is(&connection_class, "redis.connection", "SSLConnection")? { - Some(project_tls(&resolved)?) - } else { - return Ok(Err(UnsupportedCacheConfig::RedisConnection)); - }; + let tls = is_tls.then(|| project_tls(&resolved)).transpose()?; let protocol = match optional_i64(&resolved, "protocol")?.unwrap_or(2) { 2 => RedisProtocol::Resp2, @@ -332,18 +322,9 @@ fn project_valkey_semantic( ) -> PyResult> { let client = backend.getattr("sync_client")?; let pool = client.getattr("connection_pool")?; - if !instance_class_is(&pool, "redis.connection", "ConnectionPool")? { + let Ok((resolved, _is_tls)) = project_connection_pool(&pool)? else { return Ok(Err(UnsupportedCacheConfig::RedisConnection)); - } - let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; - let connection_class = resolved - .get_item("connection_class")? - .unwrap_or(pool.getattr("connection_class")?); - if !class_is(&connection_class, "redis.connection", "Connection")? - && !class_is(&connection_class, "redis.connection", "SSLConnection")? - { - return Ok(Err(UnsupportedCacheConfig::RedisConnection)); - } + }; let connection = RedisConnectionConfig { host: required_string(&resolved, "host")?, port: u16::try_from(required_i64(&resolved, "port")?) @@ -371,6 +352,27 @@ fn project_valkey_semantic( })) } +#[inline(never)] +fn project_connection_pool<'py>( + pool: &Bound<'py, PyAny>, +) -> PyResult, bool), UnsupportedCacheConfig>> { + if !instance_class_is(pool, "redis.connection", "ConnectionPool")? { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; + let connection_class = resolved + .get_item("connection_class")? + .unwrap_or(pool.getattr("connection_class")?); + let is_tls = if class_is(&connection_class, "redis.connection", "Connection")? { + false + } else if class_is(&connection_class, "redis.connection", "SSLConnection")? { + true + } else { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + }; + Ok(Ok((resolved, is_tls))) +} + #[inline(never)] fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { Ok(RedisTlsConfig { @@ -646,6 +648,40 @@ mod tests { }); } + #[test] + fn projects_valkey_semantic_configuration() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "pool = ConnectionPool()\n\ + pool.connection_class = Connection\n\ + pool.max_connections = 12\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'db': 2}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ + facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("Valkey semantic cache should be supported"); + }; + let CacheBackendConfig::ValkeySemantic(valkey) = config.backend else { + panic!("expected Valkey semantic configuration"); + }; + assert_eq!(valkey.similarity_threshold, 0.85); + assert_eq!(valkey.index_name, "semantic_idx"); + assert_eq!(valkey.embedding_model, "text-embedding-3-small"); + assert_eq!(valkey.connection.host, "cache.internal"); + assert_eq!(valkey.connection.port, 6390); + assert_eq!(valkey.connection.database, 2); + assert_eq!(valkey.connection.pool_size, 12); + assert_eq!(valkey.connection.protocol, RedisProtocol::Resp2); + assert!(valkey.connection.tls.is_none()); + }); + } + #[test] fn dynamic_redis_auth_stays_on_python() { Python::initialize(); diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py index 00037e6e29f..dc4ede8ea30 100644 --- a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py +++ b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py @@ -1,4 +1,7 @@ +import hashlib import os +import struct +import time from collections.abc import Generator, Mapping from types import SimpleNamespace from typing import Final, cast @@ -36,24 +39,32 @@ def index_name(valkey_url: str) -> Generator[str]: client.close() -def _request() -> dict[str, object]: +def _request(prompt: str = "semantic cache prompt") -> dict[str, object]: return { "key": {"preset": "key"}, - "messages": [{"role": "user", "content": "semantic cache prompt"}], + "messages": [{"role": "user", "content": prompt}], } -def _backend(url: str, index_name: str) -> ValkeySemanticCache: +def _backend( + url: str, + index_name: str, + embeddings: Mapping[str, list[float]] | None = None, +) -> ValkeySemanticCache: + vectors: Final = embeddings or {"semantic cache prompt": [1.0, 0.0]} backend: Final = ValkeySemanticCache( redis_url=url, similarity_threshold=0.8, index_name=index_name, ) - backend._get_embedding = lambda prompt, metadata=None: [1.0, 0.0] + + def embed(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]: + return vectors[prompt] async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]: - return [1.0, 0.0] + return vectors[prompt] + backend._get_embedding = embed backend._get_async_embedding = async_embedding return backend @@ -147,3 +158,122 @@ def test_batch_lookup_is_unsupported( binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() with pytest.raises(NotImplementedError): binding.lookup_batch([_request()]) + + +def test_ttl_expiry( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store({**_request(), "ttl_seconds": 1.0}, {"answer": "expires"}) + client: Final = redis.Redis.from_url(valkey_url) + documents: Final = list(client.scan_iter(f"{index_name}:*")) + assert len(documents) == 1 + assert client.ttl(documents[0]) > 0 + time.sleep(1.5) + assert binding.lookup(_request()) is None + + +def test_no_ttl_is_persistent_and_python_reads_native_value( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + response: Final = {"answer": "persistent"} + binding.store(_request(), response) + client: Final = redis.Redis.from_url(valkey_url) + documents: Final = list(client.scan_iter(f"{index_name}:*")) + assert len(documents) == 1 + assert client.ttl(documents[0]) == -1 + cached: Final = cast(Mapping[str, object], backend.get_cache("key", messages=_request()["messages"])) + assert cached["response"] == response + + +def test_below_threshold_misses_on_native_and_python( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend( + valkey_url, + index_name, + {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store(_request("prompt A"), {"answer": "A"}) + assert binding.lookup(_request("prompt B")) is None + assert backend.get_cache("key", messages=_request("prompt B")["messages"]) is None + + +def test_malformed_entry_is_a_miss_on_native_and_python( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + client: Final = redis.Redis.from_url(valkey_url) + scope: Final = hashlib.sha256(b"key").hexdigest() + document: Final = f"{index_name}:{scope}:{uuid4().hex}" + client.hset( + document, + mapping={ + "litellm_cache_key": scope, + "prompt": "semantic cache prompt", + "response": "not json", + "embedding": struct.pack("<2f", 1.0, 0.0), + }, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + assert binding.lookup(_request()) is None + assert backend.get_cache("key", messages=_request()["messages"]) is None + + +async def test_async_store_batch_and_lookup( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend( + valkey_url, + index_name, + {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + requests: Final = [_request("prompt A"), _request("prompt B")] + responses: Final = [{"answer": "A"}, {"answer": "B"}] + await binding.async_store_batch(requests, responses) + assert await binding.async_lookup(requests[0]) == responses[0] + assert await binding.async_lookup(requests[1]) == responses[1] + + +def test_subclass_backend_falls_back_to_python( + valkey_url: str, + index_name: str, +) -> None: + class Custom(ValkeySemanticCache): + pass + + facade: Final = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + redis_url=valkey_url, + similarity_threshold=0.8, + valkey_semantic_cache_index_name=index_name, + ) + facade.cache = Custom(redis_url=valkey_url, similarity_threshold=0.8, index_name=index_name) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "python_callback" + + +async def test_ping_maps_unsupported_native_operation_to_not_implemented( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + with pytest.raises(NotImplementedError): + await binding.ping() From 0d09e9d8929f0bac68932ec59dfbe465a805f2c1 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:04:57 +0000 Subject: [PATCH 07/18] fix(rust-wheel): reduce native extension size 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 570d0dd3568..8725b25cfc7 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -73,7 +73,7 @@ fancy-regex = "0.19.2" veil = "0.3.0" [profile.release] -opt-level = 3 +opt-level = 2 lto = "thin" codegen-units = 1 panic = "unwind" From 4509eb991453a89ba45c13b500bc5533d4801c9a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:22:01 +0000 Subject: [PATCH 08/18] fix(cache): align native semantic cache scope keys 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 | 57 ++++++++- .../crates/python-bridge/src/cache/native.rs | 115 +++++++++++++++++- .../test_valkey_semantic_cache_native.py | 115 ++++++++++++++++++ 5 files changed, 286 insertions(+), 3 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e4c9c385f1f..6ca4ff69648 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2700,6 +2700,7 @@ dependencies = [ "rstest", "serde", "serde_json", + "sha2 0.10.9", "tokio", "tokio-tungstenite", ] diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index ce2405f33c4..afc4f8833b5 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -46,6 +46,7 @@ tokio = { workspace = true, features = ["sync"] } criterion.workspace = true futures-util.workspace = true rstest.workspace = true +sha2.workspace = true tokio-tungstenite.workspace = true [[bench]] diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index e074c5e2f5d..169bb5accda 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -322,9 +322,17 @@ fn project_valkey_semantic( ) -> PyResult> { let client = backend.getattr("sync_client")?; let pool = client.getattr("connection_pool")?; - let Ok((resolved, _is_tls)) = project_connection_pool(&pool)? else { + let Ok((resolved, is_tls)) = project_connection_pool(&pool)? else { return Ok(Err(UnsupportedCacheConfig::RedisConnection)); }; + for key in ["credential_provider", "redis_connect_func"] { + if has_value(&resolved, key)? { + return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); + } + } + if is_tls { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } let connection = RedisConnectionConfig { host: required_string(&resolved, "host")?, port: u16::try_from(required_i64(&resolved, "port")?) @@ -682,6 +690,53 @@ mod tests { }); } + #[test] + fn valkey_semantic_tls_stays_on_python() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "pool = ConnectionPool()\n\ + pool.connection_class = SSLConnection\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ + facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("TLS Valkey semantic cache should stay on Python"); + }; + assert_eq!( + reason.message(), + "native Redis connection type is not implemented" + ); + }); + } + + #[test] + fn valkey_semantic_dynamic_auth_stays_on_python() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "pool = ConnectionPool()\n\ + pool.connection_class = Connection\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'credential_provider': object()}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ + facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("dynamic Valkey authentication must stay on Python"); + }; + assert_eq!(reason.message(), "native Redis credentials require Python"); + }); + } + #[test] fn dynamic_redis_auth_stays_on_python() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index d314cd41ac5..e7c875579fd 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -6,13 +6,50 @@ use litellm_cache::{ use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ - CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, + CacheEntry, CacheKeyField, PartialHits, ResponseCache, ResponseCacheCodec, + ResponseCacheRequest, WriteBuffer, }; use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig}; use serde_json::Value; use super::{embedder::PythonEmbedder, request::NativeRequest}; +fn semantic_key(request: &NativeRequest, scope: &str) -> litellm_cache_response::CacheKeyInput { + let mut key = request.key.clone(); + if key.preset.is_some() { + return key; + } + key.fields + .retain(|field| !matches!(field.name.as_str(), "messages" | "prompt" | "input")); + const TENANT: [&str; 3] = [ + "user_api_key", + "user_api_key_team_id", + "user_api_key_org_id", + ]; + let end_user = (scope == "end_user").then_some("user_api_key_end_user_id"); + for name in TENANT.into_iter().chain(end_user) { + let Some(value) = request + .metadata + .as_ref() + .and_then(|metadata| metadata.get(name)) + else { + continue; + }; + let value = match value { + Value::Null => continue, + Value::String(text) => text.clone(), + other => other.to_string(), + }; + key.fields.push(CacheKeyField { + name: name.to_owned(), + value: Some(value), + api_parameter: true, + internal_parameter: false, + }); + } + key +} + #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), @@ -88,7 +125,7 @@ impl NativeResponseCache { scope: &str, ) -> ResponseCacheRequest { ResponseCacheRequest { - key: request.key.clone(), + key: semantic_key(request, scope), controls: request.controls, context: SemanticCacheContext { input: request.input.clone(), @@ -329,3 +366,77 @@ impl NativeResponseCache { } } } + +#[cfg(test)] +mod tests { + use litellm_cache_response::{CacheControls, CacheKeyInput, cache_key}; + use serde_json::json; + use sha2::{Digest, Sha256}; + + use super::*; + + fn native_request(key: CacheKeyInput, metadata: Value) -> NativeRequest { + NativeRequest { + key, + controls: CacheControls::default(), + ttl: None, + max_age: None, + messages: Some(json!([{"role": "user", "content": "prompt"}])), + input: None, + metadata: Some(metadata), + } + } + + #[test] + fn semantic_key_matches_python_scope_material() { + let key = CacheKeyInput { + fields: vec![ + CacheKeyField { + name: "model".to_owned(), + value: Some("gpt-4.1".to_owned()), + api_parameter: true, + internal_parameter: false, + }, + CacheKeyField { + name: "messages".to_owned(), + value: Some("prompt".to_owned()), + api_parameter: true, + internal_parameter: false, + }, + ], + ..Default::default() + }; + let request = native_request( + key, + json!({"user_api_key": "k1", "user_api_key_team_id": null}), + ); + let expected = format!("{:x}", Sha256::digest(b"model: gpt-4.1user_api_key: k1")); + assert_eq!(cache_key(&semantic_key(&request, "key")), expected); + + let end_user_request = native_request( + request.key.clone(), + json!({"user_api_key": "k1", "user_api_key_end_user_id": "u1"}), + ); + let expected = format!( + "{:x}", + Sha256::digest(b"model: gpt-4.1user_api_key: k1user_api_key_end_user_id: u1") + ); + assert_eq!( + cache_key(&semantic_key(&end_user_request, "end_user")), + expected + ); + + let preset_request = native_request( + CacheKeyInput { + preset: Some("preset-key".to_owned()), + ..Default::default() + }, + json!({"user_api_key": "k1"}), + ); + assert_eq!( + semantic_key(&preset_request, "end_user").preset.as_deref(), + Some("preset-key") + ); + assert!(semantic_key(&preset_request, "end_user").fields.is_empty()); + } +} diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py index dc4ede8ea30..81fcf00ffc0 100644 --- a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py +++ b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py @@ -46,6 +46,56 @@ def _request(prompt: str = "semantic cache prompt") -> dict[str, object]: } +def _field_request( + prompt: str, + metadata: Mapping[str, object], +) -> dict[str, object]: + return { + "key": { + "fields": [ + { + "name": "model", + "value": "gpt-4.1", + "api_parameter": True, + "internal_parameter": False, + }, + { + "name": "messages", + "value": prompt, + "api_parameter": True, + "internal_parameter": False, + }, + ] + }, + "messages": [{"role": "user", "content": prompt}], + "metadata": dict(metadata), + } + + +def _facade( + url: str, + index_name: str, + embeddings: Mapping[str, list[float]], +) -> Cache: + facade: Final = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + valkey_semantic_cache_index_name=index_name, + ) + vectors: Final = embeddings + + def embed(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]: + return vectors[prompt] + + async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]: + return vectors[prompt] + + facade.cache._get_embedding = embed + facade.cache._get_async_embedding = async_embedding + return facade + + def _backend( url: str, index_name: str, @@ -268,6 +318,71 @@ def test_subclass_backend_falls_back_to_python( assert resolver.resolve().kind == "python_callback" +def test_field_key_matches_python_semantic_scope( + valkey_url: str, + index_name: str, +) -> None: + facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]}) + metadata: Final = {"user_api_key": "k1"} + expected: Final = facade.get_cache_key( + model="gpt-4.1", + messages=[{"role": "user", "content": "semantic cache prompt"}], + metadata=metadata, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + facade.cache, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store(_field_request("semantic cache prompt", metadata), {"answer": "scoped"}) + client: Final = redis.Redis.from_url(valkey_url) + documents: Final = list(client.scan_iter(f"{index_name}:*")) + assert len(documents) == 1 + document_parts: Final = documents[0].decode().split(":") + assert document_parts[1] == hashlib.sha256(expected.encode()).hexdigest() + client.close() + + +def test_field_key_isolates_tenant_scope( + valkey_url: str, + index_name: str, +) -> None: + facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]}) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + facade.cache, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store( + _field_request("semantic cache prompt", {"user_api_key": "k1"}), + {"answer": "tenant one"}, + ) + assert ( + binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k2"})) + is None + ) + assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k1"})) == { + "answer": "tenant one" + } + + +def test_tls_valkey_facade_falls_back_to_python( + index_name: str, +) -> None: + facade: Final = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + redis_url="rediss://127.0.0.1:6390/0", + similarity_threshold=0.8, + valkey_semantic_cache_index_name=index_name, + ) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "python_callback" + + async def test_ping_maps_unsupported_native_operation_to_not_implemented( valkey_url: str, index_name: str, From 769917a7e8d9ffccca42f084dddfab94f299e246 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:22:15 +0000 Subject: [PATCH 09/18] revert(rust-wheel): restore release optimization 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 8725b25cfc7..570d0dd3568 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -73,7 +73,7 @@ fancy-regex = "0.19.2" veil = "0.3.0" [profile.release] -opt-level = 2 +opt-level = 3 lto = "thin" codegen-units = 1 panic = "unwind" From 2b257bd9a3f09faedecb700cffa00a175df46d49 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:24:23 +0000 Subject: [PATCH 10/18] 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 a5571333fb95d4989d81f7e0ab55b69d0260db22 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:25:22 +0000 Subject: [PATCH 11/18] refactor(cache-valkey-semantic): reuse cache-redis connection layer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 2 +- .../crates/cache-valkey-semantic/Cargo.toml | 2 +- .../crates/cache-valkey-semantic/src/lib.rs | 303 ++++++------------ 3 files changed, 99 insertions(+), 208 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 6ca4ff69648..425509c6c1a 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2507,8 +2507,8 @@ name = "litellm-cache-valkey-semantic" version = "0.1.0" dependencies = [ "litellm-cache", + "litellm-cache-redis", "litellm-cache-response", - "r2d2", "redis", "redis-test", "rstest", diff --git a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml index 9a0a566ca3b..f98bb5a5fa8 100644 --- a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml @@ -7,8 +7,8 @@ repository.workspace = true [dependencies] litellm-cache.workspace = true +litellm-cache-redis.workspace = true litellm-cache-response.workspace = true -r2d2 = "0.8.10" redis = { version = "1.7.0", features = ["tls-rustls"] } serde_json.workspace = true sha2.workspace = true diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs index ef028f0a7b2..dfc6c82596c 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs @@ -5,6 +5,7 @@ use std::{ }; use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext}; +use litellm_cache_redis::connection::{ConnectionRef, Connections}; use litellm_cache_response::CacheEntry; use serde_json::Value; use sha2::{Digest, Sha256}; @@ -28,40 +29,6 @@ pub struct ValkeySemanticConfig { pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; -struct PooledConnection { - connection: redis::Connection, - failed: bool, -} - -struct ConnectionManager(redis::Client); - -impl r2d2::ManageConnection for ConnectionManager { - type Connection = PooledConnection; - type Error = redis::RedisError; - - fn connect(&self) -> Result { - let connection = self.0.get_connection()?; - Ok(PooledConnection { - connection, - failed: false, - }) - } - - fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), Self::Error> { - redis::cmd("PING").query::(&mut connection.connection)?; - Ok(()) - } - - fn has_broken(&self, connection: &mut Self::Connection) -> bool { - connection.failed || !redis::ConnectionLike::is_open(&connection.connection) - } -} - -enum Connections { - Pool(r2d2::Pool), - Fixed(Mutex), -} - #[derive(Clone)] struct IndexState { name: String, @@ -70,61 +37,8 @@ struct IndexState { similarity_threshold: f64, } -struct ConnectionRef<'a>(&'a mut dyn redis::ConnectionLike); - -impl redis::ConnectionLike for ConnectionRef<'_> { - fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { - self.0.req_packed_command(cmd) - } - - fn req_packed_commands( - &mut self, - cmd: &[u8], - offset: usize, - count: usize, - ) -> redis::RedisResult> { - self.0.req_packed_commands(cmd, offset, count) - } - - fn get_db(&self) -> i64 { - self.0.get_db() - } - - fn supports_pipelining(&self) -> bool { - self.0.supports_pipelining() - } - - fn check_connection(&mut self) -> bool { - self.0.check_connection() - } - - fn is_open(&self) -> bool { - self.0.is_open() - } -} - -impl Connections -where - C: redis::ConnectionLike + Send + 'static, -{ - fn execute( - &self, - operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, - ) -> Result { - match self { - Self::Pool(pool) => { - let mut pooled = pool.get().map_err(|_| Error::Unavailable)?; - let result = operation(&mut ConnectionRef(&mut pooled.connection)); - pooled.failed = matches!(result, Err(Error::Unavailable)); - result - } - Self::Fixed(connection) => { - let mut connection = connection.lock().map_err(|_| Error::Unavailable)?; - operation(&mut ConnectionRef(&mut *connection)) - } - } - } -} +const REDIS_TIMEOUT: Duration = Duration::from_secs(5); +const REDIS_POOL_SIZE: u32 = 16; pub struct ValkeySemanticCache< E: Embedder, @@ -149,15 +63,8 @@ where codec: S, config: ValkeySemanticConfig, ) -> Result { - let client = redis::Client::open(url).map_err(|_| Error::Unavailable)?; - let pool = r2d2::Pool::builder() - .max_size(16) - .min_idle(Some(0)) - .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)?), embedder, codec, config, @@ -179,7 +86,7 @@ where config: ValkeySemanticConfig, ) -> Self { Self { - connections: Arc::new(Connections::Fixed(Mutex::new(connection))), + connections: Arc::new(Connections::fixed(connection)), embedder, codec, config, @@ -232,15 +139,17 @@ where let response = self.codec.encode(&value)?; let vector = embedding_bytes(&embedding); let index = self.index_state(); - write_document( - &self.connections, - &index, - &scope, - &prompt, - response, - vector, - self.get_ttl(context), - ) + self.connections.execute(|connection| { + write_document( + connection, + &index, + &scope, + &prompt, + response, + vector, + self.get_ttl(context), + ) + }) } fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { @@ -251,9 +160,10 @@ where let scope = scope_tag(key); let vector = embedding_bytes(&embedding); let index = self.index_state(); - let Some(response) = - search_document(&self.connections, &index, &scope, vector, embedding.len())? - else { + let response = self.connections.execute(|connection| { + search_document(connection, &index, &scope, vector, embedding.len()) + })?; + let Some(response) = response else { return Ok(None); }; self.codec.decode(&response).map(Some) @@ -282,11 +192,10 @@ where let vector = embedding_bytes(&embedding); let scope = scope_tag(&key); let ttl = context.ttl; - tokio::task::spawn_blocking(move || { - write_document(&connections, &index, &scope, &prompt, response, vector, ttl) + Connections::run_blocking(connections, move |connection| { + write_document(connection, &index, &scope, &prompt, response, vector, ttl) }) .await - .map_err(|_| Error::Unavailable)? } } @@ -308,13 +217,12 @@ where .await?; let connections = Arc::clone(&self.connections); let index = self.index_state(); - tokio::task::spawn_blocking(move || { + Connections::run_blocking(connections, move |connection| { let scope = scope_tag(&key); let vector = embedding_bytes(&embedding); - search_document(&connections, &index, &scope, vector, embedding.len()) + search_document(connection, &index, &scope, vector, embedding.len()) }) .await - .map_err(|_| Error::Unavailable)? .and_then(|response| response.map(|bytes| self.codec.decode(&bytes)).transpose()) } } @@ -437,66 +345,58 @@ fn embedding_bytes(embedding: &[f32]) -> Vec { .collect() } -fn write_document( - connections: &Connections, +fn write_document( + connection: &mut ConnectionRef<'_>, index: &IndexState, scope: &str, prompt: &str, response: Vec, vector: Vec, ttl: Option, -) -> Result<(), Error> -where - C: redis::ConnectionLike + Send + 'static, -{ +) -> Result<(), Error> { let dimension = vector.len() / std::mem::size_of::(); ensure_index( - connections, + connection, &index.name, &index.prefix, &index.dimension, dimension, )?; let document = format!("{}{scope}:{}", index.prefix, Uuid::new_v4()); - connections.execute(|connection| { - let mut pipeline = redis::pipe(); + let mut pipeline = redis::pipe(); + pipeline + .cmd("HSET") + .arg(&document) + .arg("litellm_cache_key") + .arg(scope) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg("embedding") + .arg(vector) + .ignore(); + if let Some(ttl) = ttl { pipeline - .cmd("HSET") + .cmd("EXPIRE") .arg(&document) - .arg("litellm_cache_key") - .arg(scope) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(response) - .arg("embedding") - .arg(vector) + .arg(ttl.as_secs()) .ignore(); - if let Some(ttl) = ttl { - pipeline - .cmd("EXPIRE") - .arg(&document) - .arg(ttl.as_secs()) - .ignore(); - } - pipeline - .query::<()>(connection) - .map_err(|_| Error::Unavailable) - }) + } + pipeline + .query::<()>(connection) + .map_err(|_| Error::Unavailable) } -fn search_document( - connections: &Connections, +fn search_document( + connection: &mut ConnectionRef<'_>, index: &IndexState, scope: &str, vector: Vec, dimension: usize, -) -> Result>, Error> -where - C: redis::ConnectionLike + Send + 'static, -{ +) -> Result>, Error> { ensure_index( - connections, + connection, &index.name, &index.prefix, &index.dimension, @@ -504,23 +404,21 @@ where )?; let query = format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]"); - let response = connections.execute(|connection| { - redis::cmd("FT.SEARCH") - .arg(&index.name) - .arg(query) - .arg("PARAMS") - .arg(2) - .arg("vec") - .arg(vector) - .arg("RETURN") - .arg(2) - .arg("response") - .arg("vector_distance") - .arg("DIALECT") - .arg(2) - .query::(connection) - .map_err(|_| Error::Unavailable) - })?; + let response = redis::cmd("FT.SEARCH") + .arg(&index.name) + .arg(query) + .arg("PARAMS") + .arg(2) + .arg("vec") + .arg(vector) + .arg("RETURN") + .arg(2) + .arg("response") + .arg("vector_distance") + .arg("DIALECT") + .arg(2) + .query::(connection) + .map_err(|_| Error::Unavailable)?; let Some(fields) = search_fields(response)? else { return Ok(None); }; @@ -539,16 +437,13 @@ where Ok(Some(response)) } -fn ensure_index( - connections: &Connections, +fn ensure_index( + connection: &mut ConnectionRef<'_>, index_name: &str, prefix: &str, index_dimension: &Mutex>, dimension: usize, -) -> Result<(), Error> -where - C: redis::ConnectionLike + Send + 'static, -{ +) -> Result<(), Error> { if index_dimension .lock() .map_err(|_| Error::Unavailable)? @@ -556,41 +451,37 @@ where { return Ok(()); } - let create = connections.execute(|connection| { - Ok(redis::cmd("FT.CREATE") - .arg(index_name) - .arg("ON") - .arg("HASH") - .arg("PREFIX") - .arg(1) - .arg(prefix) - .arg("SCHEMA") - .arg("litellm_cache_key") - .arg("TAG") - .arg("embedding") - .arg("VECTOR") - .arg("HNSW") - .arg(6) - .arg("TYPE") - .arg("FLOAT32") - .arg("DIM") - .arg(dimension) - .arg("DISTANCE_METRIC") - .arg("COSINE") - .query::(connection) - .map(|_| ()) - .map_err(|error| error.to_string())) - })?; + let create = redis::cmd("FT.CREATE") + .arg(index_name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(prefix) + .arg("SCHEMA") + .arg("litellm_cache_key") + .arg("TAG") + .arg("embedding") + .arg("VECTOR") + .arg("HNSW") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dimension) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .query::(connection) + .map(|_| ()) + .map_err(|error| error.to_string()); if let Err(message) = create { if !message.to_ascii_lowercase().contains("already exists") { return Err(Error::Unavailable); } - let info = connections.execute(|connection| { - redis::cmd("FT.INFO") - .arg(index_name) - .query::(connection) - .map_err(|_| Error::Unavailable) - })?; + let info = redis::cmd("FT.INFO") + .arg(index_name) + .query::(connection) + .map_err(|_| Error::Unavailable)?; let existing = index_dimension_from_info(&info).ok_or(Error::Unavailable)?; if existing != dimension { return Err(Error::Unavailable); From 2a5112d146bf83f8fe70473d75bc11b259e139f7 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:32:40 +0000 Subject: [PATCH 12/18] 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 | 67 +++++++++++-------- .../cache-redis/src/cache/connection.rs | 2 +- .../cache-redis/src/cache/operations.rs | 28 ++++---- litellm-rust/crates/cache-redis/src/lib.rs | 4 ++ 4 files changed, 57 insertions(+), 44 deletions(-) diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index e2e2656fcbb..24399c9b2f9 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -14,7 +14,7 @@ use crate::topology::RedisTopology; mod connection; mod operations; -pub(crate) use connection::ConnectionRef; +pub use connection::ConnectionRef; use connection::{ClusterConnectionManager, ConnectionManager}; pub use operations::{ @@ -40,7 +40,8 @@ const CLAIM_SCRIPT: &str = concat!( ); const CLAIM_ATTEMPTS: usize = 8; -enum Connections { +#[allow(private_interfaces)] +pub enum Connections { Pool(r2d2::Pool), Cluster(r2d2::Pool), Fixed(Mutex), @@ -50,7 +51,7 @@ impl Connections where C: redis::ConnectionLike + Send + 'static, { - fn execute( + pub fn execute( &self, operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, ) -> Result { @@ -73,6 +74,29 @@ 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 fn fixed(connection: C) -> Self { + Self::Fixed(Mutex::new(connection)) + } + + pub fn open(url: &str, topology: &RedisTopology) -> Result { + match topology { + RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)), + RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool( + ClusterConnectionManager::open(url, startup_nodes)?, + )?)), + } + } } pub struct RedisCache { @@ -94,12 +118,7 @@ impl RedisCache { default_ttl: Option, codec: S, ) -> Result { - let connections = match topology { - RedisTopology::Standalone => Connections::Pool(pool(ConnectionManager::open(url)?)?), - RedisTopology::Cluster { startup_nodes } => { - Connections::Cluster(pool(ClusterConnectionManager::open(url, startup_nodes)?)?) - } - }; + let connections = Connections::open(url, topology)?; Ok(Self { connections: Arc::new(connections), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), @@ -127,7 +146,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, @@ -203,16 +222,6 @@ where .saturating_add(u64::from(ttl.subsec_nanos() > 0)) .max(1) } - - 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)? - } } fn namespaced_key(namespace: Option<&str>, key: &str) -> String { @@ -271,7 +280,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) @@ -285,7 +294,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) @@ -311,7 +320,7 @@ where if entries.is_empty() { return Ok(()); } - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let commands = entries .into_iter() .map(|(key, payload)| { @@ -330,7 +339,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 connection.ping() { Ok(_) => CacheConnectionResult { status: CacheConnectionStatus::Success, @@ -391,7 +400,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) @@ -418,7 +427,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 @@ -438,7 +447,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 @@ -470,7 +479,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 @@ -581,7 +590,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/connection.rs b/litellm-rust/crates/cache-redis/src/cache/connection.rs index 1834f1d94e5..06364296992 100644 --- a/litellm-rust/crates/cache-redis/src/cache/connection.rs +++ b/litellm-rust/crates/cache-redis/src/cache/connection.rs @@ -117,7 +117,7 @@ impl r2d2::ManageConnection for ClusterConnectionManager { } } -pub(crate) enum ConnectionRef<'a> { +pub enum ConnectionRef<'a> { Node(&'a mut dyn redis::ConnectionLike), Cluster(&'a mut ClusterConnection), } diff --git a/litellm-rust/crates/cache-redis/src/cache/operations.rs b/litellm-rust/crates/cache-redis/src/cache/operations.rs index 9a7023338bf..4345ee879b3 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) @@ -188,7 +188,7 @@ where } pub async fn ping(&self) -> Result { - Self::run_blocking(Arc::clone(&self.connections), |connection| { + Connections::run_blocking(Arc::clone(&self.connections), |connection| { connection.ping().map_err(|_| Error::Unavailable) }) .await @@ -196,7 +196,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) @@ -208,7 +208,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 matches = Vec::new(); connection.scan(&pattern, count, |_, keys| { matches.extend(keys); @@ -231,7 +231,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 sadd = redis::cmd("SADD"); sadd.arg(&key).arg(values); let mut expire = redis::cmd("EXPIRE"); @@ -253,7 +253,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) @@ -279,7 +279,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 commands = operations .into_iter() .map(|(key, values)| { @@ -304,7 +304,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 { @@ -333,7 +333,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 commands = operations .into_iter() .map(|(key, count)| { @@ -365,7 +365,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()) @@ -426,7 +426,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 commands = Vec::with_capacity(operations.len() * 2); let mut increments = Vec::with_capacity(operations.len()); for (key, amount, ttl) in operations { @@ -460,7 +460,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 @@ -474,7 +474,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..efb0db931ac 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}; +} + pub use cache::{ RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript, }; From b6b0e58ba399d5e8a01c4fc13de4f4c2c6af7fd6 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:32:42 +0000 Subject: [PATCH 13/18] refactor(cache-valkey-semantic): reuse cache-redis connection layer Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache-valkey-semantic/src/lib.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs index dfc6c82596c..2bd4bd71bce 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs @@ -5,7 +5,10 @@ use std::{ }; use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext}; -use litellm_cache_redis::connection::{ConnectionRef, Connections}; +use litellm_cache_redis::{ + RedisTopology, + connection::{ConnectionRef, Connections}, +}; use litellm_cache_response::CacheEntry; use serde_json::Value; use sha2::{Digest, Sha256}; @@ -37,9 +40,6 @@ struct IndexState { similarity_threshold: f64, } -const REDIS_TIMEOUT: Duration = Duration::from_secs(5); -const REDIS_POOL_SIZE: u32 = 16; - pub struct ValkeySemanticCache< E: Embedder, S: CacheCodec, @@ -64,7 +64,7 @@ where config: ValkeySemanticConfig, ) -> Result { Ok(Self { - connections: Arc::new(Connections::pooled(url, REDIS_TIMEOUT, REDIS_POOL_SIZE)?), + connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), embedder, codec, config, From cf7234c6f43ea220be426503bc1a1e69becbb62b Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:48:40 +0000 Subject: [PATCH 14/18] chore(rust): refresh workspace lockfile Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 7bd166f48b6..097884c8ecc 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2707,6 +2707,7 @@ dependencies = [ "serde", "serde_json", "serde_with", + "sha2 0.10.9", "tokio", "tokio-tungstenite", ] From 847f732f5d6e4c55d1b925bd50f480af4f858440 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:01:13 +0000 Subject: [PATCH 15/18] fix(cache): await valkey semantic embeddings inline on the caller loop Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-response/src/response.rs | 4 + .../crates/cache-valkey-semantic/src/lib.rs | 78 ++++++++- .../crates/python-bridge/src/cache/binding.rs | 14 +- .../python-bridge/src/cache/embedder.rs | 45 ++--- .../crates/python-bridge/src/cache/mod.rs | 1 + .../crates/python-bridge/src/cache/native.rs | 103 ++++++++++-- .../python-bridge/src/cache/semantic_step.rs | 154 ++++++++++++++++++ .../test_valkey_semantic_cache_native.py | 41 +++++ 8 files changed, 393 insertions(+), 47 deletions(-) create mode 100644 litellm-rust/crates/python-bridge/src/cache/semantic_step.rs diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index e70e07a5d26..2f949d511de 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -52,6 +52,10 @@ where &self.backend } + pub fn backend_arc(&self) -> &Arc { + &self.backend + } + pub fn default_ttl(&self) -> Option { self.backend.get_ttl(&B::Context::default()) } diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs index 2bd4bd71bce..e3a3e6094c4 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs @@ -24,6 +24,22 @@ pub trait Embedder: Send + Sync + 'static { ) -> impl Future, Error>> + Send; } +pub struct PreparedEmbedding(pub Vec); + +impl Embedder for PreparedEmbedding { + fn embed(&self, _prompt: &str, _metadata: Option<&Value>) -> Result, Error> { + Ok(self.0.clone()) + } + + async fn async_embed( + &self, + _prompt: &str, + _metadata: Option<&Value>, + ) -> Result, Error> { + Ok(self.0.clone()) + } +} + #[derive(Clone, Debug, PartialEq)] pub struct ValkeySemanticConfig { pub similarity_threshold: f64, @@ -112,6 +128,23 @@ where } } +impl ValkeySemanticCache +where + E: Embedder, + S: CacheCodec + Clone, + C: redis::ConnectionLike + Send + 'static, +{ + pub fn with_embedder(&self, embedder: E2) -> ValkeySemanticCache { + ValkeySemanticCache { + connections: Arc::clone(&self.connections), + embedder, + codec: self.codec.clone(), + config: self.config.clone(), + index_dimension: Arc::clone(&self.index_dimension), + } + } +} + impl BaseCache for ValkeySemanticCache where E: Embedder, @@ -597,8 +630,8 @@ mod tests { use serde_json::{Value, json}; use super::{ - Embedder, ValkeySemanticCache, ValkeySemanticConfig, index_dimension_from_info, - prompt_from_context, scope_tag, + Embedder, PreparedEmbedding, ValkeySemanticCache, ValkeySemanticConfig, + index_dimension_from_info, prompt_from_context, scope_tag, }; #[derive(Clone)] @@ -758,6 +791,47 @@ mod tests { ); } + #[tokio::test] + async fn prepared_embedding_returns_its_vector_for_any_prompt() { + let embedding = PreparedEmbedding(vec![1.0, 2.0]); + assert_eq!( + embedding + .async_embed("different prompt", None) + .await + .unwrap(), + vec![1.0, 2.0] + ); + } + + #[test] + fn with_embedder_shares_index_state_and_connections() { + let entry = CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "ok"}), + }; + let encoded = ResponseCacheCodec.encode(&entry).unwrap(); + let cache = ValkeySemanticCache::with_connection( + RecordingConnection::new([ok(), ok(), Ok(search_hit(encoded, "0.1"))]), + FixedEmbedder { + vector: vec![1.0, 0.0], + calls: Arc::default(), + }, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold: 0.8, + index_name: "test".into(), + }, + ); + cache + .set_cache("key", entry.clone(), &semantic_context(None)) + .unwrap(); + let prepared = cache.with_embedder(PreparedEmbedding(vec![1.0, 0.0])); + assert_eq!( + prepared.get_cache("key", &semantic_context(None)).unwrap(), + Some(entry) + ); + } + #[test] fn missing_prompt_does_not_touch_redis() { let cache = ValkeySemanticCache::with_connection( diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs index ad64b24d3c1..a8881f19e45 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -56,12 +56,7 @@ impl ResolvedCache { CacheBinding::Disabled => ready_none(py)?, CacheBinding::Native(service) => { let request = request(input)?; - let service = service.clone(); - run_async( - py, - async move { service.async_lookup(&request, now()).await }, - cache_error, - )? + service.async_lookup_py(py, request)? } CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?, }; @@ -179,12 +174,7 @@ impl ResolvedCache { CacheBinding::Native(service) => { let request = self::request(request)?; let response: Value = from_py(response)?; - let service = service.clone(); - run_async( - py, - async move { service.async_store(&request, response, now()).await }, - cache_error, - ) + service.async_store_py(py, request, response) } CacheBinding::PythonCallback(callback) => { callback.async_store(py, response, callback_kwargs) diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs index d240d9d019e..3de0ceb3b67 100644 --- a/litellm-rust/crates/python-bridge/src/cache/embedder.rs +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -3,22 +3,37 @@ use std::{future::Future, sync::Arc}; use litellm_cache::Error; use litellm_cache_valkey_semantic::Embedder; use litellm_host_python::to_py; -use pyo3::prelude::*; +use pyo3::{PyTraverseError, PyVisit, prelude::*}; use serde_json::Value; #[derive(Clone)] pub(super) struct PythonEmbedder { sync_embed: Arc>, - async_embed: Arc>, + async_embed_callable: Arc>, } impl PythonEmbedder { pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult { Ok(Self { sync_embed: Arc::new(backend.getattr("_get_embedding")?.unbind()), - async_embed: Arc::new(backend.getattr("_get_async_embedding")?.unbind()), + async_embed_callable: Arc::new(backend.getattr("_get_async_embedding")?.unbind()), }) } + + pub(super) fn async_embed_awaitable<'py>( + &self, + py: Python<'py>, + prompt: &str, + metadata: &Option, + ) -> PyResult> { + let metadata = to_py(py, metadata)?; + self.async_embed_callable.bind(py).call1((prompt, metadata)) + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&*self.sync_embed)?; + visit.call(&*self.async_embed_callable) + } } impl Embedder for PythonEmbedder { @@ -34,25 +49,15 @@ impl Embedder for PythonEmbedder { Ok(result.into_iter().map(|value| value as f32).collect()) } + #[expect( + clippy::manual_async_fn, + reason = "the shared Embedder trait uses an impl Future return" + )] fn async_embed( &self, - prompt: &str, - metadata: Option<&Value>, + _prompt: &str, + _metadata: Option<&Value>, ) -> impl Future, Error>> + Send { - let callable = Arc::clone(&self.async_embed); - let prompt = prompt.to_owned(); - let metadata = metadata.cloned(); - async move { - let future = Python::attach(|py| -> PyResult<_> { - let metadata = to_py(py, &metadata)?; - let awaitable = callable.bind(py).call1((prompt, metadata))?; - pyo3_async_runtimes::tokio::into_future(awaitable) - }) - .map_err(|_| Error::Unavailable)?; - let result = future.await.map_err(|_| Error::Unavailable)?; - let result = Python::attach(|py| result.bind(py).extract::>()) - .map_err(|_| Error::Unavailable)?; - Ok(result.into_iter().map(|value| value as f32).collect()) - } + async { 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..278d3da1ff9 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_step; 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 3380a914c41..7340ef6cce0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -10,9 +10,14 @@ use litellm_cache_response::{ ResponseCacheRequest, WriteBuffer, }; use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig}; +use pyo3::prelude::*; use serde_json::Value; -use super::{embedder::PythonEmbedder, request::NativeRequest}; +use super::{ + embedder::PythonEmbedder, + request::NativeRequest, + semantic_step::{SemanticEmbedExecution, drive_semantic}, +}; fn semantic_key(request: &NativeRequest, scope: &str) -> litellm_cache_response::CacheKeyInput { let mut key = request.key.clone(); @@ -59,6 +64,7 @@ pub(super) enum NativeResponseCache { }, ValkeySemantic { cache: Arc>>, + embedder: PythonEmbedder, scope: String, }, } @@ -100,7 +106,7 @@ impl NativeResponseCache { ) -> Result { let backend = ValkeySemanticCache::new( url, - embedder, + embedder.clone(), ResponseCacheCodec, ValkeySemanticConfig { similarity_threshold, @@ -109,6 +115,7 @@ impl NativeResponseCache { )?; Ok(Self::ValkeySemantic { cache: Arc::new(ResponseCache::new(Arc::new(backend))), + embedder, scope: String::from("key"), }) } @@ -152,7 +159,13 @@ impl NativeResponseCache { pub fn with_scope(self, scope: String) -> Self { match self { - Self::ValkeySemantic { cache, .. } => Self::ValkeySemantic { cache, scope }, + Self::ValkeySemantic { + cache, embedder, .. + } => Self::ValkeySemantic { + cache, + embedder, + scope, + }, value => value, } } @@ -215,7 +228,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.lookup(&Self::exact(request), now), Self::Redis { cache, .. } => cache.lookup(&Self::exact(request), now), - Self::ValkeySemantic { cache, scope } => { + Self::ValkeySemantic { cache, scope, .. } => { cache.lookup(&Self::semantic(request, scope), now) } } @@ -230,7 +243,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.store(&Self::exact(request), response, now), Self::Redis { cache, .. } => cache.store(&Self::exact(request), response, now), - Self::ValkeySemantic { cache, scope } => { + Self::ValkeySemantic { cache, scope, .. } => { cache.store(&Self::semantic(request, scope), response, now) } } @@ -262,7 +275,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.async_lookup(&Self::exact(request), now).await, Self::Redis { cache, .. } => cache.async_lookup(&Self::exact(request), now).await, - Self::ValkeySemantic { cache, scope } => { + Self::ValkeySemantic { cache, scope, .. } => { cache .async_lookup(&Self::semantic(request, scope), now) .await @@ -270,6 +283,36 @@ impl NativeResponseCache { } } + pub(super) fn async_lookup_py<'py>( + &self, + py: Python<'py>, + request: NativeRequest, + ) -> PyResult> { + match self { + Self::Memory(_) | Self::Redis { .. } => { + let service = self.clone(); + litellm_host_python::run_async( + py, + async move { service.async_lookup(&request, super::request::now()).await }, + super::cache_error, + ) + } + Self::ValkeySemantic { + cache, + embedder, + scope, + } => drive_semantic( + py, + SemanticEmbedExecution::lookup( + Arc::clone(cache.backend_arc()), + embedder.clone(), + Self::semantic(&request, scope), + super::request::now(), + ), + ), + } + } + pub async fn async_store( &self, request: &NativeRequest, @@ -298,7 +341,7 @@ impl NativeResponseCache { .async_store(cache, &Self::exact(request), response, now) .await } - Self::ValkeySemantic { cache, scope } => { + Self::ValkeySemantic { cache, scope, .. } => { cache .async_store(&Self::semantic(request, scope), response, now) .await @@ -306,6 +349,42 @@ impl NativeResponseCache { } } + pub(super) fn async_store_py<'py>( + &self, + py: Python<'py>, + request: NativeRequest, + response: Value, + ) -> PyResult> { + match self { + Self::Memory(_) | Self::Redis { .. } => { + let service = self.clone(); + litellm_host_python::run_async( + py, + async move { + service + .async_store(&request, response, super::request::now()) + .await + }, + super::cache_error, + ) + } + Self::ValkeySemantic { + cache, + embedder, + scope, + } => drive_semantic( + py, + SemanticEmbedExecution::store( + Arc::clone(cache.backend_arc()), + embedder.clone(), + Self::semantic(&request, scope), + response, + super::request::now(), + ), + ), + } + } + pub async fn async_lookup_batch( &self, requests: &[NativeRequest], @@ -344,12 +423,10 @@ impl NativeResponseCache { .collect(); cache.async_store_batch(entries, now).await } - Self::ValkeySemantic { cache, scope } => { - let entries = entries - .into_iter() - .map(|(request, value)| (Self::semantic(&request, scope), value)) - .collect(); - cache.async_store_batch(entries, now).await + Self::ValkeySemantic { cache, scope, .. } => { + entries.into_iter().try_for_each(|(request, value)| { + cache.store(&Self::semantic(&request, scope), value, now) + }) } } } diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs new file mode 100644 index 00000000000..c62cdb1d9a6 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs @@ -0,0 +1,154 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::SemanticCacheContext; +use litellm_cache_response::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; +use litellm_cache_valkey_semantic::{ + Embedder, PreparedEmbedding, ValkeySemanticCache, 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}; + +pub(super) enum Op { + Lookup, + Store(Value), +} + +#[derive(Clone, Copy)] +enum State { + Start, + AwaitingEmbedding, + AwaitingStorage, + Done, +} + +pub(super) struct SemanticEmbedExecution { + backend: Arc>, + embedder: PythonEmbedder, + request: ResponseCacheRequest, + op: Op, + now: Duration, + state: State, +} + +impl SemanticEmbedExecution { + pub(super) fn lookup( + backend: Arc>, + embedder: PythonEmbedder, + request: ResponseCacheRequest, + now: Duration, + ) -> Self { + Self { + backend, + embedder, + request, + op: Op::Lookup, + now, + state: State::Start, + } + } + + pub(super) fn store( + backend: Arc>, + embedder: PythonEmbedder, + request: ResponseCacheRequest, + response: Value, + now: Duration, + ) -> Self { + Self { + backend, + embedder, + request, + op: Op::Store(response), + now, + state: State::Start, + } + } + + fn start(&mut self, py: Python<'_>) -> PyResult { + let Some(prompt) = prompt_from_context(&self.request.context) else { + let cache = Arc::new(ResponseCache::new(Arc::clone(&self.backend))); + self.state = State::AwaitingStorage; + return storage_step(py, cache, self.request.clone(), &self.op, self.now); + }; + let awaitable = + self.embedder + .async_embed_awaitable(py, &prompt, &self.request.context.metadata)?; + self.state = State::AwaitingEmbedding; + Ok(ExecutionStep::Await(awaitable.unbind())) + } + + fn resume_py( + &mut self, + py: Python<'_>, + result: Option>>, + ) -> PyResult { + match (self.state, result) { + (State::Start, None) => self.start(py), + (State::AwaitingEmbedding, Some(Ok(value))) => { + let values = value.bind(py).extract::>()?; + let backend = self.backend.with_embedder(PreparedEmbedding( + values.into_iter().map(|value| value as f32).collect(), + )); + let cache = Arc::new(ResponseCache::new(Arc::new(backend))); + self.state = State::AwaitingStorage; + storage_step(py, cache, self.request.clone(), &self.op, self.now) + } + (State::AwaitingStorage, Some(Ok(value))) => { + self.state = State::Done; + Ok(ExecutionStep::Return(value)) + } + (_, Some(Err(error))) => Err(error), + _ => Err(PyRuntimeError::new_err( + "invalid semantic cache execution state", + )), + } + } +} + +impl ExecutionBody for SemanticEmbedExecution { + fn resume(&mut self, result: Option>>) -> PyResult { + Python::attach(|py| self.resume_py(py, result)) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.embedder.traverse(visit) + } +} + +fn storage_step( + py: Python<'_>, + cache: Arc>>, + request: ResponseCacheRequest, + op: &Op, + now: Duration, +) -> PyResult { + let awaitable = match op { + Op::Lookup => run_async( + py, + async move { cache.async_lookup(&request, now).await }, + cache_error, + )?, + Op::Store(response) => { + let response = response.clone(); + run_async( + py, + async move { cache.async_store(&request, response, now).await }, + cache_error, + )? + } + }; + Ok(ExecutionStep::Await(awaitable.unbind())) +} + +pub(super) fn drive_semantic<'py>( + py: Python<'py>, + body: SemanticEmbedExecution, +) -> PyResult> { + let execution = Py::new(py, Execution::new(body))?; + py.import("litellm.rust_bridge.lifecycle")? + .getattr("drive")? + .call1((execution,)) +} diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py index 81fcf00ffc0..e2bd3dcb10a 100644 --- a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py +++ b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py @@ -1,6 +1,9 @@ +import asyncio +import contextvars import hashlib import os import struct +import threading import time from collections.abc import Generator, Mapping from types import SimpleNamespace @@ -16,6 +19,7 @@ from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType pytestmark: Final = pytest.mark.requires_rust_extension +embedding_context: Final = contextvars.ContextVar("embedding_context") @pytest.fixture @@ -171,6 +175,43 @@ async def test_async_lookup_and_store( assert await binding.async_lookup(request) == {"answer": "async"} +async def test_async_embedding_runs_inline_in_caller_task( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + observed: dict[str, object] = {} + + async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]: + observed["context"] = embedding_context.get("missing") + observed["task"] = asyncio.current_task() + observed["thread"] = threading.get_ident() + embedding_context.set("embedder") + return [1.0, 0.0] + + backend._get_async_embedding = async_embedding + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + request: Final = {**_request(), "ttl_seconds": 2.0} + caller_task: Final = asyncio.current_task() + caller_thread: Final = threading.get_ident() + token: Final = embedding_context.set("caller") + try: + await binding.async_store(request, {"answer": "inline"}) + assert observed["context"] == "caller" + assert observed["task"] is caller_task + assert observed["thread"] == caller_thread + assert embedding_context.get() == "embedder" + assert await binding.async_lookup(request) == {"answer": "inline"} + finally: + embedding_context.reset(token) + + def test_facade_activation_and_mutation_fallback( valkey_url: str, index_name: str, From c233377d9c073fd9e2262c58ea14b2ca75aadf0a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:04:46 +0000 Subject: [PATCH 16/18] fix(cache): await valkey semantic batch embeddings inline Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/python-bridge/src/cache/binding.rs | 7 +- .../crates/python-bridge/src/cache/native.rs | 49 ++++- .../python-bridge/src/cache/semantic_step.rs | 175 +++++++++++++----- .../test_valkey_semantic_cache_native.py | 20 ++ 4 files changed, 197 insertions(+), 54 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs index a8881f19e45..2ff73238202 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -231,12 +231,7 @@ impl ResolvedCache { )); } let entries = requests.into_iter().zip(responses).collect(); - let service = service.clone(); - run_async( - py, - async move { service.async_store_batch(entries, now()).await }, - cache_error, - ) + service.async_store_batch_py(py, entries) } CacheBinding::PythonCallback(callback) => { callback.async_store_batch(py, callback_result, callback_kwargs) diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 7340ef6cce0..260dc8f349e 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -424,9 +424,52 @@ impl NativeResponseCache { cache.async_store_batch(entries, now).await } Self::ValkeySemantic { cache, scope, .. } => { - entries.into_iter().try_for_each(|(request, value)| { - cache.store(&Self::semantic(&request, scope), value, now) - }) + let entries = entries + .into_iter() + .map(|(request, value)| (Self::semantic(&request, scope), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + } + } + + pub(super) fn async_store_batch_py<'py>( + &self, + py: Python<'py>, + entries: Vec<(NativeRequest, Value)>, + ) -> PyResult> { + match self { + Self::Memory(_) | Self::Redis { .. } => { + let service = self.clone(); + litellm_host_python::run_async( + py, + async move { + service + .async_store_batch(entries, super::request::now()) + .await + }, + super::cache_error, + ) + } + Self::ValkeySemantic { + cache, + embedder, + scope, + } => { + let (requests, responses): (Vec<_>, Vec<_>) = entries + .into_iter() + .map(|(request, response)| (Self::semantic(&request, scope), response)) + .unzip(); + drive_semantic( + py, + SemanticEmbedExecution::store_batch( + Arc::clone(cache.backend_arc()), + embedder.clone(), + requests, + responses, + super::request::now(), + ), + ) } } } diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs index c62cdb1d9a6..d8f3ab297f8 100644 --- a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs +++ b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs @@ -2,9 +2,7 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::SemanticCacheContext; use litellm_cache_response::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; -use litellm_cache_valkey_semantic::{ - Embedder, PreparedEmbedding, ValkeySemanticCache, prompt_from_context, -}; +use litellm_cache_valkey_semantic::{PreparedEmbedding, ValkeySemanticCache, prompt_from_context}; use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; use serde_json::Value; @@ -14,6 +12,7 @@ use super::{cache_error, embedder::PythonEmbedder}; pub(super) enum Op { Lookup, Store(Value), + StoreBatch(Vec), } #[derive(Clone, Copy)] @@ -27,9 +26,11 @@ enum State { pub(super) struct SemanticEmbedExecution { backend: Arc>, embedder: PythonEmbedder, - request: ResponseCacheRequest, + requests: Vec>, op: Op, now: Duration, + prepared: Vec>>, + index: usize, state: State, } @@ -43,9 +44,11 @@ impl SemanticEmbedExecution { Self { backend, embedder, - request, + requests: vec![request], op: Op::Lookup, now, + prepared: vec![None], + index: 0, state: State::Start, } } @@ -60,23 +63,132 @@ impl SemanticEmbedExecution { Self { backend, embedder, - request, + requests: vec![request], op: Op::Store(response), now, + prepared: vec![None], + index: 0, + state: State::Start, + } + } + + pub(super) fn store_batch( + backend: Arc>, + embedder: PythonEmbedder, + requests: Vec>, + responses: Vec, + now: Duration, + ) -> Self { + Self { + backend, + embedder, + prepared: vec![None; requests.len()], + requests, + op: Op::StoreBatch(responses), + now, + index: 0, state: State::Start, } } fn start(&mut self, py: Python<'_>) -> PyResult { - let Some(prompt) = prompt_from_context(&self.request.context) else { - let cache = Arc::new(ResponseCache::new(Arc::clone(&self.backend))); - self.state = State::AwaitingStorage; - return storage_step(py, cache, self.request.clone(), &self.op, self.now); + while self.index < self.requests.len() { + let request = &self.requests[self.index]; + let Some(prompt) = prompt_from_context(&request.context) else { + self.index += 1; + continue; + }; + let metadata = request.context.metadata.clone(); + let awaitable = self + .embedder + .async_embed_awaitable(py, &prompt, &metadata)?; + self.state = State::AwaitingEmbedding; + return Ok(ExecutionStep::Await(awaitable.unbind())); + } + self.state = State::AwaitingStorage; + self.storage_step(py) + } + + fn storage_step(&self, py: Python<'_>) -> PyResult { + let requests = self.requests.clone(); + let prepared = self.prepared.clone(); + let backend = Arc::clone(&self.backend); + let now = self.now; + let awaitable = match &self.op { + Op::Lookup => { + let Some(request) = requests.into_iter().next() else { + return Err(PyRuntimeError::new_err( + "semantic lookup requires one request", + )); + }; + match prepared.into_iter().next().flatten() { + Some(values) => { + let backend = backend.with_embedder(PreparedEmbedding(values)); + let cache = Arc::new(ResponseCache::new(Arc::new(backend))); + run_async( + py, + async move { cache.async_lookup(&request, now).await }, + cache_error, + )? + } + None => { + let cache = Arc::new(ResponseCache::new(backend)); + run_async( + py, + async move { cache.async_lookup(&request, now).await }, + cache_error, + )? + } + } + } + Op::Store(response) => { + let Some(request) = requests.into_iter().next() else { + return Err(PyRuntimeError::new_err( + "semantic store requires one request", + )); + }; + let response = response.clone(); + match prepared.into_iter().next().flatten() { + Some(values) => { + let backend = backend.with_embedder(PreparedEmbedding(values)); + let cache = Arc::new(ResponseCache::new(Arc::new(backend))); + run_async( + py, + async move { cache.async_store(&request, response, now).await }, + cache_error, + )? + } + None => { + let cache = Arc::new(ResponseCache::new(backend)); + run_async( + py, + async move { cache.async_store(&request, response, now).await }, + cache_error, + )? + } + } + } + Op::StoreBatch(responses) => { + let responses = responses.clone(); + run_async( + py, + async move { + for ((request, response), prepared) in + requests.into_iter().zip(responses).zip(prepared) + { + let Some(values) = prepared else { + continue; + }; + let backend = backend.with_embedder(PreparedEmbedding(values)); + let cache = ResponseCache::new(Arc::new(backend)); + cache.async_store(&request, response, now).await?; + } + Ok(()) + }, + cache_error, + )? + } }; - let awaitable = - self.embedder - .async_embed_awaitable(py, &prompt, &self.request.context.metadata)?; - self.state = State::AwaitingEmbedding; Ok(ExecutionStep::Await(awaitable.unbind())) } @@ -89,12 +201,10 @@ impl SemanticEmbedExecution { (State::Start, None) => self.start(py), (State::AwaitingEmbedding, Some(Ok(value))) => { let values = value.bind(py).extract::>()?; - let backend = self.backend.with_embedder(PreparedEmbedding( - values.into_iter().map(|value| value as f32).collect(), - )); - let cache = Arc::new(ResponseCache::new(Arc::new(backend))); - self.state = State::AwaitingStorage; - storage_step(py, cache, self.request.clone(), &self.op, self.now) + self.prepared[self.index] = + Some(values.into_iter().map(|value| value as f32).collect()); + self.index += 1; + self.start(py) } (State::AwaitingStorage, Some(Ok(value))) => { self.state = State::Done; @@ -118,31 +228,6 @@ impl ExecutionBody for SemanticEmbedExecution { } } -fn storage_step( - py: Python<'_>, - cache: Arc>>, - request: ResponseCacheRequest, - op: &Op, - now: Duration, -) -> PyResult { - let awaitable = match op { - Op::Lookup => run_async( - py, - async move { cache.async_lookup(&request, now).await }, - cache_error, - )?, - Op::Store(response) => { - let response = response.clone(); - run_async( - py, - async move { cache.async_store(&request, response, now).await }, - cache_error, - )? - } - }; - Ok(ExecutionStep::Await(awaitable.unbind())) -} - pub(super) fn drive_semantic<'py>( py: Python<'py>, body: SemanticEmbedExecution, diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py index e2bd3dcb10a..fb5c5965333 100644 --- a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py +++ b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py @@ -332,11 +332,31 @@ async def test_async_store_batch_and_lookup( index_name, {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}, ) + sync_calls: Final = [] + async_tasks: Final = [] + + def sync_embedding(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]: + sync_calls.append(prompt) + return {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}[prompt] + + async def async_embedding( + prompt: str, + metadata: dict[str, object] | None = None, + ) -> list[float]: + async_tasks.append(asyncio.current_task()) + return {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}[prompt] + + backend._get_embedding = sync_embedding + backend._get_async_embedding = async_embedding handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() requests: Final = [_request("prompt A"), _request("prompt B")] responses: Final = [{"answer": "A"}, {"answer": "B"}] + caller_task: Final = asyncio.current_task() await binding.async_store_batch(requests, responses) + assert sync_calls == [] + assert async_tasks + assert all(task is caller_task for task in async_tasks) assert await binding.async_lookup(requests[0]) == responses[0] assert await binding.async_lookup(requests[1]) == responses[1] From 21d1604e64eb77a430bb95c4f80d3e44f94b1e1f Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:33:11 +0000 Subject: [PATCH 17/18] fix(cache): honor controls and tenant metadata in valkey semantic bridge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-valkey-semantic/src/lib.rs | 35 ++-- .../crates/python-bridge/src/cache/native.rs | 28 ++- .../crates/python-bridge/src/cache/request.rs | 6 + .../python-bridge/src/cache/semantic_step.rs | 26 ++- .../test_valkey_semantic_cache_native.py | 162 +++++++++++++++++- 5 files changed, 216 insertions(+), 41 deletions(-) diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs index e3a3e6094c4..6062ccc842c 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs @@ -273,13 +273,11 @@ pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { if let Some(Value::Array(messages)) = context.messages.as_ref() && !messages.is_empty() { - return Some( - messages - .iter() - .filter_map(Value::as_object) - .map(message_text) - .collect(), - ); + return messages + .iter() + .filter_map(Value::as_object) + .map(message_text) + .collect(); } let input = context.input.as_ref()?; let mut parts = Vec::new(); @@ -288,21 +286,25 @@ pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { (!prompt.is_empty()).then_some(prompt) } -fn message_text(message: &serde_json::Map) -> String { +fn message_text(message: &serde_json::Map) -> Option { let content = match message.get("content") { Some(Value::String(value)) => value.clone(), - Some(Value::Array(parts)) => parts - .iter() - .filter_map(Value::as_object) - .filter_map(|part| part.get("text").and_then(Value::as_str)) - .filter(|text| !text.is_empty()) - .collect(), + Some(Value::Array(parts)) => { + let mut content = String::new(); + for part in parts { + let part = part.as_object()?; + if let Some(text) = part.get("text").and_then(Value::as_str) { + content.push_str(text); + } + } + content + } _ => String::new(), }; - format!( + Some(format!( "{content}{}", search_results_text(message.get("search_results")) - ) + )) } fn search_results_text(value: Option<&Value>) -> String { @@ -732,6 +734,7 @@ mod tests { #[rstest] #[case(json!([{"content": "hello"}]), None, Some("hello"))] #[case(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))] + #[case(json!([{"content": ["raw", {"text": "hello"}]}]), None, None)] #[case(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))] #[case(Value::Array(vec![]), Some(json!(" hello ")), Some("hello"))] #[case(Value::Array(vec![]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))] diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index fefc09c1321..d8f7a693108 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -34,11 +34,24 @@ fn semantic_key(request: &NativeRequest, scope: &str) -> litellm_cache_response: ]; let end_user = (scope == "end_user").then_some("user_api_key_end_user_id"); for name in TENANT.into_iter().chain(end_user) { - let Some(value) = request - .metadata - .as_ref() - .and_then(|metadata| metadata.get(name)) - else { + let sources = [ + request.metadata.as_ref(), + request.litellm_metadata.as_ref(), + request + .litellm_params + .as_ref() + .and_then(|params| params.get("metadata")), + request + .litellm_params + .as_ref() + .and_then(|params| params.get("litellm_metadata")), + ]; + let Some(value) = sources.into_iter().flatten().find_map(|source| { + source + .as_object() + .and_then(|values| values.get(name)) + .filter(|value| !value.is_null()) + }) else { continue; }; let value = match value { @@ -340,7 +353,6 @@ impl NativeResponseCache { Arc::clone(cache.backend_arc()), embedder.clone(), Self::semantic(&request, scope), - super::request::now(), ), ), } @@ -417,7 +429,6 @@ impl NativeResponseCache { embedder.clone(), Self::semantic(&request, scope), response, - super::request::now(), ), ), } @@ -517,7 +528,6 @@ impl NativeResponseCache { embedder.clone(), requests, responses, - super::request::now(), ), ) } @@ -565,6 +575,8 @@ mod tests { messages: Some(json!([{"role": "user", "content": "prompt"}])), input: None, metadata: Some(metadata), + litellm_metadata: None, + litellm_params: None, } } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 3e19e7fdc22..036951891a1 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -17,6 +17,8 @@ struct RequestInput { messages: Option, input: Option, metadata: Option, + litellm_metadata: Option, + litellm_params: Option, } pub(super) struct NativeRequest { @@ -27,6 +29,8 @@ pub(super) struct NativeRequest { pub(super) messages: Option, pub(super) input: Option, pub(super) metadata: Option, + pub(super) litellm_metadata: Option, + pub(super) litellm_params: Option, } pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { @@ -46,6 +50,8 @@ fn request_input(input: RequestInput) -> PyResult { messages: input.messages, input: input.input, metadata: input.metadata, + litellm_metadata: input.litellm_metadata, + litellm_params: input.litellm_params, }) } diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs index d8f3ab297f8..24caf3374d6 100644 --- a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs +++ b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs @@ -28,7 +28,7 @@ pub(super) struct SemanticEmbedExecution { embedder: PythonEmbedder, requests: Vec>, op: Op, - now: Duration, + now: Option, prepared: Vec>>, index: usize, state: State, @@ -39,14 +39,13 @@ impl SemanticEmbedExecution { backend: Arc>, embedder: PythonEmbedder, request: ResponseCacheRequest, - now: Duration, ) -> Self { Self { backend, embedder, requests: vec![request], op: Op::Lookup, - now, + now: None, prepared: vec![None], index: 0, state: State::Start, @@ -58,14 +57,13 @@ impl SemanticEmbedExecution { embedder: PythonEmbedder, request: ResponseCacheRequest, response: Value, - now: Duration, ) -> Self { Self { backend, embedder, requests: vec![request], op: Op::Store(response), - now, + now: None, prepared: vec![None], index: 0, state: State::Start, @@ -77,7 +75,6 @@ impl SemanticEmbedExecution { embedder: PythonEmbedder, requests: Vec>, responses: Vec, - now: Duration, ) -> Self { Self { backend, @@ -85,15 +82,26 @@ impl SemanticEmbedExecution { prepared: vec![None; requests.len()], requests, op: Op::StoreBatch(responses), - now, + now: None, index: 0, state: State::Start, } } fn start(&mut self, py: Python<'_>) -> PyResult { + if self.now.is_none() { + self.now = Some(super::request::now()); + } while self.index < self.requests.len() { let request = &self.requests[self.index]; + let enabled = match &self.op { + Op::Lookup => request.controls.reads(), + Op::Store(_) | Op::StoreBatch(_) => request.controls.writes(), + }; + if !enabled { + self.index += 1; + continue; + } let Some(prompt) = prompt_from_context(&request.context) else { self.index += 1; continue; @@ -113,7 +121,9 @@ impl SemanticEmbedExecution { let requests = self.requests.clone(); let prepared = self.prepared.clone(); let backend = Arc::clone(&self.backend); - let now = self.now; + let now = self + .now + .ok_or_else(|| PyRuntimeError::new_err("semantic cache timestamp is unavailable"))?; let awaitable = match &self.op { Op::Lookup => { let Some(request) = requests.into_iter().next() else { diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py index fb5c5965333..c87a9f86a80 100644 --- a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py +++ b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py @@ -53,8 +53,12 @@ def _request(prompt: str = "semantic cache prompt") -> dict[str, object]: def _field_request( prompt: str, metadata: Mapping[str, object], + *, + namespace: str | None = None, + litellm_metadata: Mapping[str, object] | None = None, + litellm_params: Mapping[str, object] | None = None, ) -> dict[str, object]: - return { + request: Final = { "key": { "fields": [ { @@ -69,23 +73,32 @@ def _field_request( "api_parameter": True, "internal_parameter": False, }, - ] + ], + "namespace": namespace, }, "messages": [{"role": "user", "content": prompt}], "metadata": dict(metadata), } + if litellm_metadata is not None: + request["litellm_metadata"] = dict(litellm_metadata) + if litellm_params is not None: + request["litellm_params"] = dict(litellm_params) + return request def _facade( url: str, index_name: str, embeddings: Mapping[str, list[float]], + *, + namespace: str | None = None, ) -> Cache: facade: Final = Cache( type=LiteLLMCacheType.VALKEY_SEMANTIC, redis_url=url, similarity_threshold=0.8, valkey_semantic_cache_index_name=index_name, + namespace=namespace, ) vectors: Final = embeddings @@ -175,6 +188,45 @@ async def test_async_lookup_and_store( assert await binding.async_lookup(request) == {"answer": "async"} +async def test_disabled_cache_controls_skip_async_embedding( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + calls: Final = [] + + async def fail_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]: + calls.append(prompt) + raise AssertionError("embedding must not run") + + backend._get_async_embedding = fail_embedding + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + controls: Final = { + "supported_call_type": True, + "configured": True, + "native_backend": True, + "default_on": True, + "caching": True, + "no_cache": False, + "no_store": False, + "use_cache": True, + } + no_read_request: Final = {**_request(), "controls": {**controls, "no_cache": True}} + assert await binding.async_lookup(no_read_request) is None + no_write_request: Final = {**_request(), "controls": {**controls, "no_store": True}} + await binding.async_store(no_write_request, {"answer": "blocked"}) + assert calls == [] + client: Final = redis.Redis.from_url(valkey_url) + assert list(client.scan_iter(f"{index_name}:*")) == [] + client.close() + + async def test_async_embedding_runs_inline_in_caller_task( valkey_url: str, index_name: str, @@ -323,6 +375,25 @@ def test_malformed_entry_is_a_miss_on_native_and_python( assert backend.get_cache("key", messages=_request()["messages"]) is None +def test_mixed_content_parts_match_python_semantic_behavior( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + messages: Final = [{"role": "user", "content": ["raw", {"text": "hello"}]}] + backend.set_cache("key", {"answer": "mixed"}, messages=messages) + assert backend.get_cache("key", messages=messages) is None + + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + request: Final = {**_request(), "messages": messages} + binding.store(request, {"answer": "mixed"}) + assert binding.lookup(request) is None + client: Final = redis.Redis.from_url(valkey_url) + assert list(client.scan_iter(f"{index_name}:*")) == [] + client.close() + + async def test_async_store_batch_and_lookup( valkey_url: str, index_name: str, @@ -406,6 +477,84 @@ def test_field_key_matches_python_semantic_scope( client.close() +def test_field_key_reads_all_python_tenant_metadata_sources( + valkey_url: str, + index_name: str, +) -> None: + facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]}) + params_metadata: Final = {"user_api_key_team_id": "team-from-params"} + expected: Final = facade.get_cache_key( + model="gpt-4.1", + messages=[{"role": "user", "content": "semantic cache prompt"}], + metadata={}, + litellm_params={"metadata": params_metadata}, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + facade.cache, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store( + _field_request( + "semantic cache prompt", + {}, + litellm_params={"metadata": params_metadata}, + ), + {"answer": "params"}, + ) + client: Final = redis.Redis.from_url(valkey_url) + documents: Final = list(client.scan_iter(f"{index_name}:*")) + assert len(documents) == 1 + document_parts: Final = documents[0].decode().split(":") + assert document_parts[1] == hashlib.sha256(expected.encode()).hexdigest() + client.close() + + assert ( + binding.lookup( + _field_request( + "semantic cache prompt", + {}, + litellm_metadata={"user_api_key_team_id": "team-from-litellm"}, + ) + ) + is None + ) + + +def test_namespace_isolates_semantic_entries( + valkey_url: str, + index_name: str, +) -> None: + facade: Final = _facade( + valkey_url, + index_name, + {"semantic cache prompt": [1.0, 0.0]}, + namespace="team-a", + ) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + facade.cache, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + team_a: Final = _field_request("semantic cache prompt", {}, namespace="team-a") + team_b: Final = _field_request("semantic cache prompt", {}, namespace="team-b") + binding.store(team_a, {"answer": "team-a"}) + assert binding.lookup(team_b) is None + assert binding.lookup(team_a) == {"answer": "team-a"} + cached: Final = cast( + Mapping[str, object], + facade.get_cache( + model="gpt-4.1", + messages=[{"role": "user", "content": "semantic cache prompt"}], + ), + ) + assert cached == {"answer": "team-a"} + + def test_field_key_isolates_tenant_scope( valkey_url: str, index_name: str, @@ -422,13 +571,8 @@ def test_field_key_isolates_tenant_scope( _field_request("semantic cache prompt", {"user_api_key": "k1"}), {"answer": "tenant one"}, ) - assert ( - binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k2"})) - is None - ) - assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k1"})) == { - "answer": "tenant one" - } + assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k2"})) is None + assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k1"})) == {"answer": "tenant one"} def test_tls_valkey_facade_falls_back_to_python( From 9fa1683ed97f29f4593b28259c0250a9f771f2c8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 23:09:03 +0000 Subject: [PATCH 18/18] chore(cache): keep main's UnsupportedOperation message 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, 1 insertion(+), 1 deletion(-) diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index 51e4fe2d66a..1a381d0afd8 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -6,6 +6,6 @@ pub enum Error { InvalidEntry, #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, - #[error("cache operation is not supported by this backend")] + #[error("operation is not supported by this cache")] UnsupportedOperation, }