From d2f457f144a430ad2848b33839d298d9312c8d4e Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:20:31 +0000 Subject: [PATCH 01/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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 212ab630b81957610b9708af7fb85b827538979e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:44:26 -0700 Subject: [PATCH 18/24] fix(logging_worker): make flush() survive an event loop change flush() awaited join() on whatever queue the worker held, even one bound to an event loop that has since closed. Its unfinished counter is never decremented on the new loop, so the first flush() after a loop change hung until pytest-timeout killed it and every later one raised "is bound to a different event loop" from the queue's Event. The CircleCI unit job has been red on every branch since the first tests that flush without enqueueing landed, and an SDK script that flushes from a second asyncio.run() hangs the same way. flush() now goes through start() first, which carries the tasks stranded on the previous loop onto the current one and guarantees a worker there to drain them, the same loop-change handling every other entry point already had. --- litellm/litellm_core_utils/logging_worker.py | 6 ++++ .../litellm_core_utils/test_logging_worker.py | 34 +++++++++++++++++++ 2 files changed, 40 insertions(+) diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 5ccc5632646..cb3d8bf4fe5 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -484,9 +484,15 @@ class LoggingWorker: so it correctly handles items that have been dequeued but whose callback hasn't finished yet — ``queue.empty()`` would return True in that window and cause us to skip the wait. + + ``start()`` runs first so that, after an event loop change, the tasks + still on the previous loop's queue move onto this loop and a worker + here drains them; joining the old queue directly would wait on a + counter nothing on this loop ever decrements. """ if self._queue is None: return + self.start() await self._queue.join() async def clear_queue(self): diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/test_litellm/litellm_core_utils/test_logging_worker.py index 1553e788472..891bd0686b2 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -180,6 +180,40 @@ class TestLoggingWorker: assert sorted(fired) == ["first", "second"] + @pytest.mark.parametrize("stranded", ["still_queued", "dequeued_never_started"]) + def test_flush_on_new_loop_drains_tasks_stranded_on_previous_loop(self, stranded): + """ + Regression: ``flush()`` from a new event loop used to ``join()`` the queue bound to the + previous loop, whose unfinished counter nothing on the new loop ever decrements. The first + such flush hung until pytest-timeout killed it and every later one raised + ``RuntimeError: ... is bound to a different event loop`` from the queue's Event. + """ + worker = LoggingWorker(timeout=1.0, max_queue_size=10) + fired = [] + + async def marker(): + fired.append(True) + + async def enqueue_on_first_loop(): + if stranded == "still_queued": + worker._ensure_queue() + worker.enqueue(marker()) + return + worker.ensure_initialized_and_enqueue(marker()) + + asyncio.run(enqueue_on_first_loop()) + assert worker._queue is not None + expected_shape = (1, 0) if stranded == "still_queued" else (0, 1) + assert (worker._queue.qsize(), len(worker._unstarted_dequeued_tasks())) == expected_shape + assert fired == [], "precondition: the callback never ran before the first loop closed" + + async def flush_on_second_loop(): + await asyncio.wait_for(worker.flush(), timeout=5) + + asyncio.run(flush_on_second_loop()) + + assert fired == [True] + def test_flush_on_exit_swallows_cancellation_and_drains_remaining(self): """A callback raising CancelledError must not abort the atexit flush of later events.""" worker = LoggingWorker(timeout=1.0, max_queue_size=10) From e86ba8bbebfad37f558b91bb67e5fe55a71493f2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 15:56:09 -0700 Subject: [PATCH 19/24] test(logging_worker): cover a same-loop flush and a repeated flush after a loop change --- litellm/litellm_core_utils/logging_worker.py | 6 ++--- .../litellm_core_utils/test_logging_worker.py | 23 +++++++++++++++++-- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index cb3d8bf4fe5..2f8e7bdccea 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -485,10 +485,8 @@ class LoggingWorker: callback hasn't finished yet — ``queue.empty()`` would return True in that window and cause us to skip the wait. - ``start()`` runs first so that, after an event loop change, the tasks - still on the previous loop's queue move onto this loop and a worker - here drains them; joining the old queue directly would wait on a - counter nothing on this loop ever decrements. + ``start()`` runs first so a queue left behind by a previous event loop + is carried onto this one and drained here instead of joined forever. """ if self._queue is None: return diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/test_litellm/litellm_core_utils/test_logging_worker.py index 891bd0686b2..336067976fa 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -207,10 +207,29 @@ class TestLoggingWorker: assert (worker._queue.qsize(), len(worker._unstarted_dequeued_tasks())) == expected_shape assert fired == [], "precondition: the callback never ran before the first loop closed" - async def flush_on_second_loop(): + async def flush_twice_on_second_loop(): + await asyncio.wait_for(worker.flush(), timeout=5) await asyncio.wait_for(worker.flush(), timeout=5) - asyncio.run(flush_on_second_loop()) + asyncio.run(flush_twice_on_second_loop()) + + assert fired == [True] + + def test_flush_starts_a_worker_when_the_queue_has_none(self): + """``flush()`` must drain a queue that exists on the current loop without a running worker.""" + worker = LoggingWorker(timeout=1.0, max_queue_size=10) + fired = [] + + async def marker(): + fired.append(True) + + async def enqueue_then_flush(): + worker._ensure_queue() + worker.enqueue(marker()) + assert worker._worker_task is None, "precondition: nothing is draining the queue yet" + await asyncio.wait_for(worker.flush(), timeout=3) + + asyncio.run(enqueue_then_flush()) assert fired == [True] From 9fa1683ed97f29f4593b28259c0250a9f771f2c8 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 23:09:03 +0000 Subject: [PATCH 20/24] 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, } From 47a1053065116e54f78f287ba50c0d7756a195bc Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 23:12:57 +0000 Subject: [PATCH 21/24] feat(xiaomi_mimo): add mimo-v2.6-pro and mimo-v2.6-flash cost map rows with live e2e coverage Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 46 ++++ model_prices_and_context_window.json | 46 ++++ .../coverage_registry/llm_conversational.yaml | 3 + tests/e2e/coverage_registry/schema.py | 1 + .../llm_translation/test_xiaomi_mimo_e2e.py | 214 ++++++++++++++++++ 5 files changed, 310 insertions(+) create mode 100644 tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index beb90c2fb8c..2d8be7f1e87 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -76975,5 +76975,51 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false + }, + "xiaomi_mimo/mimo-v2.6-pro": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 4.35e-07, + "litellm_provider": "xiaomi_mimo", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8.7e-07, + "source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "xiaomi_mimo/mimo-v2.6-flash": { + "cache_read_input_token_cost": 2.8e-09, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "xiaomi_mimo", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index beb90c2fb8c..2d8be7f1e87 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -76975,5 +76975,51 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false + }, + "xiaomi_mimo/mimo-v2.6-pro": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 4.35e-07, + "litellm_provider": "xiaomi_mimo", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8.7e-07, + "source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "xiaomi_mimo/mimo-v2.6-flash": { + "cache_read_input_token_cost": 2.8e-09, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "xiaomi_mimo", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true } } diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 49d4d92ff0b..64335aa560c 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -89,6 +89,9 @@ - {id: llm.chat_completions.together_ai.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool result round trip"} - {id: llm.chat_completions.together_ai.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together cost header and spend row match the registry price"} - {id: llm.chat_completions.together_ai.thinking.nonstream.effort_none_disables, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [effort_none_disables], source: "llm_translation/test_together_ai_e2e.py", rationale: "reasoning_effort=none maps to Together's reasoning disable toggle on hybrid models"} +- {id: llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "Native MiMo v2.6 rows price the cost header and spend row from the cost map"} +- {id: llm.chat_completions.xiaomi_mimo.thinking.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: thinking, streaming: stream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo reasoning deltas stream as reasoning_content"} +- {id: llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo tool calls are not dropped"} - {id: llm.chat_completions.together_ai.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: structured_output, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "response_format json_schema reaches Together and constrains the reply"} - {id: llm.chat_completions.together_ai.prompt_cache_5m.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: prompt_cache_5m, streaming: nonstream, assertions: [cache_hit, cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together prefix-cache reads bill at cache_read_input_token_cost, not full input price"} - {id: llm.messages.together_ai.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together over /v1/messages streaming"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 8b0d38a083c..d9c20d5c588 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -58,6 +58,7 @@ LlmRoute = Literal[ "openai", "together_ai", "vertex", + "xiaomi_mimo", ] LlmCapability = Literal[ diff --git a/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py b/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py new file mode 100644 index 00000000000..efca216634b --- /dev/null +++ b/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py @@ -0,0 +1,214 @@ +"""Live e2e: Xiaomi MiMo v2.6 through the gateway on /chat/completions. + +Both native ``xiaomi_mimo/`` v2.6 rows (pro and flash) are registered via +``/model/new`` and driven against Xiaomi's own endpoint. What the gateway owes +us is that the reasoning chain surfaces as ``reasoning_content``, tool calls +survive translation, and the cost header plus spend row follow the proxy's own +cost-map price for the row (read back from ``/model/info``, never pinned here). +Requires XIAOMI_MIMO_API_KEY on the proxy; no skip gate. +""" + +from __future__ import annotations + +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call, unwrap +from lifecycle import ResourceManager +from models import ( + ChatBody, + ChatMessage, + ChatResponse, + ChatTool, + ChatToolFunction, + CostMapEntry, + LiteLLMParamsBody, + OutMessage, + SpendLogRow, +) +from passthrough_client import PassthroughClient +from pydantic import BaseModel + +pytestmark = pytest.mark.e2e + +BACKENDS: Final = ("xiaomi_mimo/mimo-v2.6-pro", "xiaomi_mimo/mimo-v2.6-flash") +ARITHMETIC_PROMPT = "What is 17 + 26? Answer with just the number." +WEATHER_PROMPT = "What is the weather in Paris? Use the tool." +COUNTING_PROMPT = "Count from 1 to 50, one number per line." + +WEATHER_TOOL = ChatTool( + function=ChatToolFunction( + name="get_weather", + description="Get the current weather for a location.", + parameters={ + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + ) +) + + +class _WeatherArgs(BaseModel): + location: str + + +class _StreamDelta(BaseModel): + content: str | None = None + reasoning_content: str | None = None + + +class _StreamChoice(BaseModel): + delta: _StreamDelta | None = None + + +class _StreamChunk(BaseModel): + choices: list[_StreamChoice] = [] + + +def _approx_equal(actual: float, expected: float) -> bool: + return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) + + +@pytest.fixture(scope="module") +def registry(client: PassthroughClient) -> dict[str, CostMapEntry]: + return client.proxy.model_cost_map() + + +def _register(client: PassthroughClient, resources: ResourceManager, backend: str) -> tuple[str, str]: + model = f"e2e-xiaomi-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=backend, api_key="os.environ/XIAOMI_MIMO_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model, resources.key() + + +def _message(response: ChatResponse) -> OutMessage: + assert response.choices, f"Xiaomi returned no choices: {response}" + message = response.choices[0].message + assert message is not None, f"Xiaomi choice has no message: {response}" + return message + + +def _deltas(result: StreamingResponse) -> list[_StreamDelta]: + require_successful_call(result) + assert result.is_streaming, f"response was not streamed: {result.headers}" + assert not result.stream_error, f"stream errored: {result.stream_error}" + assert result.stream_done, f"stream never reached [DONE]: {result.stream_events[-3:]}" + return [ + choice.delta + for event in result.stream_events + for choice in _StreamChunk.model_validate_json(event).choices + if choice.delta is not None + ] + + +@pytest.mark.parametrize("backend", BACKENDS) +class TestXiaomiMimoChatCompletions: + @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged") + def test_cost_header_and_spend_row_match_the_registry_price( + self, + client: PassthroughClient, + resources: ResourceManager, + registry: dict[str, CostMapEntry], + backend: str, + ) -> None: + price = registry.get(backend) + assert price is not None, f"{backend} has no row in the proxy's cost map, so native calls would bill $0" + assert price.litellm_provider == "xiaomi_mimo", f"{backend} is filed under the wrong provider: {price}" + assert price.input_cost_per_token and price.output_cost_per_token, f"{backend} carries no price: {price}" + model, key = _register(client, resources, backend) + + result = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"{ARITHMETIC_PROMPT} {unique_marker()}")], + max_tokens=1024, + ), + ) + require_successful_call(result) + response = ChatResponse.model_validate_json(result.body) + message = _message(response) + assert message.content and "43" in message.content, f"answer lost: {message}" + assert message.reasoning_content, f"{backend} reasons, but no reasoning_content came back: {message}" + + usage = response.usage + assert usage is not None and usage.prompt_tokens and usage.completion_tokens, ( + f"response carries no usage, so the cost cannot be real: {result.body[:300]}" + ) + header_cost = result.response_cost + assert header_cost is not None and header_cost > 0, ( + f"x-litellm-response-cost header missing or non-positive: {result.headers}" + ) + cached = (usage.prompt_tokens_details.cached_tokens or 0) if usage.prompt_tokens_details else 0 + expected = ( + (usage.prompt_tokens - cached) * price.input_cost_per_token + + cached * (price.cache_read_input_token_cost or 0.0) + + usage.completion_tokens * price.output_cost_per_token + ) + assert _approx_equal(header_cost, expected), ( + f"header cost {header_cost} disagrees with the registry price for {backend} at {usage}: expected {expected}" + ) + + def _priced(rows: list[SpendLogRow]) -> bool: + return any(row.spend is not None and row.spend > 0 for row in rows) + + rows = client.proxy.poll_logs_for_key(key, predicate=_priced) + priced = [row for row in rows if row.spend is not None and row.spend > 0] + assert priced, f"no priced spend row landed for key {key}; got {rows}" + row = priced[0] + assert row.custom_llm_provider == "xiaomi_mimo", f"spend row misattributed: {row}" + assert row.spend is not None and _approx_equal(row.spend, header_cost), ( + f"logged spend {row.spend} disagrees with the x-litellm-response-cost header {header_cost}" + ) + + @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.thinking.stream.works") + def test_reasoning_and_answer_stream_as_deltas( + self, client: PassthroughClient, resources: ResourceManager, backend: str + ) -> None: + model, key = _register(client, resources, backend) + + deltas = _deltas( + client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=COUNTING_PROMPT)], + max_tokens=2048, + stream=True, + ), + ) + ) + reasoning = "".join(delta.reasoning_content or "" for delta in deltas) + content = "".join(delta.content or "" for delta in deltas) + assert reasoning, f"stream carried no reasoning_content deltas: {deltas[:5]}" + assert "50" in content, f"streamed answer lost: {content[:300]!r}" + + @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works") + def test_tool_call_is_returned(self, client: PassthroughClient, resources: ResourceManager, backend: str) -> None: + model, key = _register(client, resources, backend) + + message = _message( + unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=WEATHER_PROMPT)], + tools=[WEATHER_TOOL], + max_tokens=1024, + ), + ) + ) + ) + assert message.tool_calls, f"{backend} dropped the tool call: {message}" + call = message.tool_calls[0] + assert call.id, f"tool call carries no id, so a tool result cannot answer it: {call}" + assert call.function.name == "get_weather", f"wrong tool called: {call}" + assert call.function.arguments, f"tool call carries no arguments: {call}" + args = _WeatherArgs.model_validate_json(call.function.arguments) + assert "paris" in args.location.lower(), f"tool arguments lost the location: {args}" From 9388602f467b4c1d3ce2f137191c4a8638a8a5d1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 16:13:09 -0700 Subject: [PATCH 22/24] test(logging_worker): track callback runs with AsyncMock instead of a mutated list --- .../litellm_core_utils/test_logging_worker.py | 22 +++++++------------ 1 file changed, 8 insertions(+), 14 deletions(-) diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/test_litellm/litellm_core_utils/test_logging_worker.py index 336067976fa..2bb93a58531 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -189,23 +189,20 @@ class TestLoggingWorker: ``RuntimeError: ... is bound to a different event loop`` from the queue's Event. """ worker = LoggingWorker(timeout=1.0, max_queue_size=10) - fired = [] - - async def marker(): - fired.append(True) + callback = AsyncMock() async def enqueue_on_first_loop(): if stranded == "still_queued": worker._ensure_queue() - worker.enqueue(marker()) + worker.enqueue(callback()) return - worker.ensure_initialized_and_enqueue(marker()) + worker.ensure_initialized_and_enqueue(callback()) asyncio.run(enqueue_on_first_loop()) assert worker._queue is not None expected_shape = (1, 0) if stranded == "still_queued" else (0, 1) assert (worker._queue.qsize(), len(worker._unstarted_dequeued_tasks())) == expected_shape - assert fired == [], "precondition: the callback never ran before the first loop closed" + assert callback.await_count == 0, "precondition: the callback never ran before the first loop closed" async def flush_twice_on_second_loop(): await asyncio.wait_for(worker.flush(), timeout=5) @@ -213,25 +210,22 @@ class TestLoggingWorker: asyncio.run(flush_twice_on_second_loop()) - assert fired == [True] + assert callback.await_count == 1 def test_flush_starts_a_worker_when_the_queue_has_none(self): """``flush()`` must drain a queue that exists on the current loop without a running worker.""" worker = LoggingWorker(timeout=1.0, max_queue_size=10) - fired = [] - - async def marker(): - fired.append(True) + callback = AsyncMock() async def enqueue_then_flush(): worker._ensure_queue() - worker.enqueue(marker()) + worker.enqueue(callback()) assert worker._worker_task is None, "precondition: nothing is draining the queue yet" await asyncio.wait_for(worker.flush(), timeout=3) asyncio.run(enqueue_then_flush()) - assert fired == [True] + assert callback.await_count == 1 def test_flush_on_exit_swallows_cancellation_and_drains_remaining(self): """A callback raising CancelledError must not abort the atexit flush of later events.""" From 9af3363c2fd646bc5fc8945c22258c9e8fb368dd Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 23:18:18 +0000 Subject: [PATCH 23/24] fix(bedrock): add bare moonshotai.kimi-k3 cost map entry mirroring the global inference profile Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- model_prices_and_context_window.json | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 6a414a0908e..b9b47fd45b5 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -76879,6 +76879,26 @@ "supports_vision": true, "supports_web_search": true }, + "moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "global.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, From cbed8cd0d81c351cb368ab90cb72d432a046322d Mon Sep 17 00:00:00 2001 From: kerry Date: Mon, 21 Sep 2026 23:20:26 +0000 Subject: [PATCH 24/24] fix(bedrock): sync moonshotai.kimi-k3 entry into cost map backup file Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 6a414a0908e..b9b47fd45b5 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -76879,6 +76879,26 @@ "supports_vision": true, "supports_web_search": true }, + "moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "global.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07,