diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index c02f460f22e..66270501312 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2741,6 +2741,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "litellm-cache-redis", + "litellm-cache-response", + "r2d2", + "redis", + "redis-test", + "serde_json", + "sha2 0.10.9", + "tokio", +] + [[package]] name = "litellm-cache-response" version = "0.1.0" @@ -2958,6 +2973,7 @@ dependencies = [ "litellm-cache-gcs", "litellm-cache-memory", "litellm-cache-redis", + "litellm-cache-redis-semantic", "litellm-cache-response", "litellm-cache-s3", "litellm-cache-valkey-semantic", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index be86240ac46..4e370c13766 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -36,6 +36,7 @@ litellm-cache-redis = { path = "crates/cache-redis" } litellm-cache-s3 = { path = "crates/cache-s3" } litellm-cache-gcs = { path = "crates/cache-gcs" } litellm-cache-disk = { path = "crates/cache-disk" } +litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" } litellm-cache-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } diff --git a/litellm-rust/crates/cache-redis-semantic/Cargo.toml b/litellm-rust/crates/cache-redis-semantic/Cargo.toml new file mode 100644 index 00000000000..9a8755a189e --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-cache-redis.workspace = true +litellm-cache-response.workspace = true +redis = { version = "1.7.0", features = ["tls-rustls"] } +r2d2 = "0.8.10" +serde_json.workspace = true +sha2.workspace = true +tokio.workspace = true + +[dev-dependencies] +redis-test = "1.0.4" +serde_json.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs new file mode 100644 index 00000000000..e0ac31f3630 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -0,0 +1,618 @@ +use std::{ + future::Future, + sync::{Arc, OnceLock}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_cache::{ + BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error, + SemanticCacheContext, +}; +use litellm_cache_redis::{ + RedisTopology, + connection::{ConnectionRef, Connections}, +}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +use crate::prompt::prompt_from_context; + +const CACHE_KEY_FIELD: &str = "litellm_cache_key"; +const VECTOR_FIELD: &str = "prompt_vector"; + +pub trait Embedder: Send + Sync + 'static { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error>; + + fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> impl Future, Error>> + Send; +} + +#[derive(Clone, Debug)] +pub struct RedisSemanticConfig { + pub index_name: String, + pub similarity_threshold: f32, +} + +struct Inner { + index_name: String, + distance_threshold: f64, + resolved_index: OnceLock, + codec: ResponseCacheCodec, + clock: fn() -> f64, +} + +impl Inner { + fn new(config: RedisSemanticConfig) -> Self { + Self { + index_name: config.index_name, + distance_threshold: 1.0 - f64::from(config.similarity_threshold), + resolved_index: OnceLock::new(), + codec: ResponseCacheCodec, + clock: timestamp, + } + } + + fn ensure_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + if let Some(name) = self.resolved_index.get() { + return Ok(name.clone()); + } + let name = match index_compatible(connection, &self.index_name, dims)? { + Some(true) => self.index_name.clone(), + Some(false) => self.isolated_index(connection, dims)?, + None => match create_index(connection, &self.index_name, dims) { + Ok(()) => self.index_name.clone(), + Err(_) => match index_compatible(connection, &self.index_name, dims)? { + Some(true) => self.index_name.clone(), + Some(false) => self.isolated_index(connection, dims)?, + None => return Err(Error::Unavailable), + }, + }, + }; + let _ = self.resolved_index.set(name.clone()); + Ok(name) + } + + fn isolated_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + let name = format!("{}_isolated", self.index_name); + match index_compatible(connection, &name, dims)? { + Some(true) => Ok(name), + Some(false) => { + redis::cmd("FT.DROPINDEX") + .arg(&name) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + create_index(connection, &name, dims)?; + Ok(name) + } + None => { + create_index(connection, &name, dims)?; + Ok(name) + } + } + } + + fn store( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + value: &CacheEntry, + prompt: &str, + vector: &[f32], + ttl: Option, + ) -> Result<(), Error> { + let index = self.ensure_index(connection, vector.len())?; + let entry_id = entry_id(prompt, tag); + let hash_key = format!("{index}:{entry_id}"); + let response = self.codec.encode(value)?; + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(&entry_id) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg(VECTOR_FIELD) + .arg(vector_buffer(vector)) + .arg("inserted_at") + .arg(format!("{}", (self.clock)())) + .arg("updated_at") + .arg(format!("{}", (self.clock)())) + .arg(CACHE_KEY_FIELD) + .arg(tag) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + if let Some(ttl) = ttl { + redis::cmd("EXPIRE") + .arg(&hash_key) + .arg(ttl_seconds(ttl)) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + } + Ok(()) + } + + fn lookup( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + vector: &[f32], + ) -> Result, Error> { + let index = self.ensure_index(connection, vector.len())?; + let query = format!( + "(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]", + escape_tag(tag) + ); + let result = redis::cmd("FT.SEARCH") + .arg(&index) + .arg(query) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg(CACHE_KEY_FIELD) + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_buffer(vector)) + .query::(connection) + .map_err(|_| Error::Unavailable)?; + let Some(fields) = first_document(&result) else { + return Ok(None); + }; + if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) { + return Ok(None); + } + if number_field(fields, "vector_distance") + .is_none_or(|distance| distance > self.distance_threshold) + { + return Ok(None); + } + let Some(response) = bytes_field(fields, "response") else { + return Ok(None); + }; + self.codec.decode(&response).map(Some) + } +} + +pub struct RedisSemanticCache { + connections: Arc>, + embedder: E, + inner: Arc, +} + +impl RedisSemanticCache { + pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result { + Ok(Self { + connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), + embedder, + inner: Arc::new(Inner::new(config)), + }) + } +} + +impl RedisSemanticCache { + pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self { + Self { + connections: Arc::new(Connections::fixed(connection)), + embedder, + inner: Arc::new(Inner::new(config)), + } + } + + pub fn with_clock(self, clock: fn() -> f64) -> Self { + Self { + inner: Arc::new(Inner { + index_name: self.inner.index_name.clone(), + distance_threshold: self.inner.distance_threshold, + resolved_index: OnceLock::new(), + codec: self.inner.codec, + clock, + }), + ..self + } + } + + pub fn embedder(&self) -> &E { + &self.embedder + } + + pub fn index_name(&self) -> &str { + &self.inner.index_name + } + + pub fn similarity_threshold(&self) -> f32 { + (1.0 - self.inner.distance_threshold) as f32 + } + + fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str { + context.scope.as_deref().unwrap_or(key) + } +} + +impl BaseCache + for RedisSemanticCache +{ + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(()); + }; + let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; + let tag = Self::tag(key, context).to_string(); + self.connections.execute(|connection| { + self.inner + .store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; + let tag = Self::tag(key, context).to_string(); + self.connections + .execute(|connection| self.inner.lookup(connection, &tag, &vector)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(&context) else { + return Ok(()); + }; + let vector = self + .embedder + .async_embed(&prompt, context.metadata.as_ref()) + .await?; + let tag = Self::tag(key, &context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + .await + } + + async fn async_get_cache( + &self, + key: &str, + context: &Self::Context, + ) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self + .embedder + .async_embed(&prompt, context.metadata.as_ref()) + .await?; + let tag = Self::tag(key, context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.lookup(connection, &tag, &vector) + }) + .await + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + match Connections::run_blocking(Arc::clone(&self.connections), |connection| { + Ok(match redis::cmd("PING").query::(connection) { + Ok(_) => CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, + }, + Err(error) => CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }, + }) + }) + .await + { + Ok(result) => Ok(result), + Err(error) => Ok(CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }), + } + } +} + +fn timestamp() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or_default() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(CACHE_KEY_FIELD.as_bytes()); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn vector_buffer(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn escape_tag(value: &str) -> String { + value + .chars() + .flat_map(|ch| { + if matches!( + ch, + ',' | '.' + | '<' + | '>' + | '{' + | '}' + | '[' + | ']' + | '\\' + | '"' + | '\'' + | ':' + | ';' + | '!' + | '@' + | '#' + | '$' + | '%' + | '^' + | '&' + | '*' + | '(' + | ')' + | '-' + | '+' + | '=' + | '~' + | '|' + | '/' + | ' ' + | '?' + ) { + vec!['\\', ch] + } else { + vec![ch] + } + }) + .collect() +} + +fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> { + redis::cmd("FT.CREATE") + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg(VECTOR_FIELD) + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg(CACHE_KEY_FIELD) + .arg("TAG") + .arg("SEPARATOR") + .arg(",") + .query::<()>(connection) + .map_err(|_| Error::Unavailable) +} + +fn index_compatible( + connection: &mut ConnectionRef<'_>, + name: &str, + dims: usize, +) -> Result, Error> { + let info = match redis::cmd("FT.INFO") + .arg(name) + .query::(connection) + { + Ok(info) => info, + Err(error) if unknown_index(&error) => return Ok(None), + Err(_) => return Err(Error::Unavailable), + }; + Ok(Some(schema_compatible(&info, dims))) +} + +fn unknown_index(error: &redis::RedisError) -> bool { + let message = error.to_string().to_lowercase(); + message.contains("unknown") && message.contains("index") +} + +fn schema_compatible(info: &redis::Value, dims: usize) -> bool { + let redis::Value::Array(entries) = info else { + return false; + }; + let attributes = entries + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some("attributes")) + .map(|pair| &pair[1]); + let Some(redis::Value::Array(attributes)) = attributes else { + return false; + }; + let fields = attributes + .iter() + .map(|attribute| { + let redis::Value::Array(attribute) = attribute else { + return (None, None, None, None, None); + }; + let mut name = None; + let mut field_type = None; + let mut dim = None; + let mut data_type = None; + let mut distance_metric = None; + for pair in attribute.as_chunks::<2>().0 { + match string_value(&pair[0]).as_deref() { + Some("identifier") => name = string_value(&pair[1]), + Some("type") => field_type = string_value(&pair[1]), + Some("dim") => dim = number_value(&pair[1]), + Some("data_type") => data_type = string_value(&pair[1]), + Some("distance_metric") => distance_metric = string_value(&pair[1]), + _ => {} + } + } + (name, field_type, dim, data_type, distance_metric) + }) + .collect::>(); + let has_field = |name: &str, field_type: &str| { + fields + .iter() + .any(|(n, t, ..)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type)) + }; + has_field("prompt", "TEXT") + && has_field("response", "TEXT") + && has_field("inserted_at", "NUMERIC") + && has_field("updated_at", "NUMERIC") + && has_field(CACHE_KEY_FIELD, "TAG") + && fields.iter().any(|(n, t, d, data, metric)| { + n.as_deref() == Some(VECTOR_FIELD) + && t.as_deref() == Some("VECTOR") + && *d == Some(dims as f64) + && data + .as_deref() + .is_some_and(|data| data.eq_ignore_ascii_case("float32")) + && metric + .as_deref() + .is_some_and(|metric| metric.eq_ignore_ascii_case("cosine")) + }) +} + +fn string_value(value: &redis::Value) -> Option { + match value { + redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), + redis::Value::SimpleString(text) => Some(text.clone()), + redis::Value::VerbatimString { text, .. } => Some(text.clone()), + _ => None, + } +} + +fn number_value(value: &redis::Value) -> Option { + match value { + redis::Value::Int(number) => Some(*number as f64), + redis::Value::Double(number) => Some(*number), + _ => string_value(value).and_then(|text| text.parse().ok()), + } +} + +fn first_document(result: &redis::Value) -> Option<&[redis::Value]> { + let redis::Value::Array(items) = result else { + return None; + }; + let [count, _document_id, fields, ..] = items.as_slice() else { + return None; + }; + if !matches!(count, redis::Value::Int(count) if *count > 0) { + return None; + } + match fields { + redis::Value::Array(fields) => Some(fields.as_slice()), + _ => None, + } +} + +fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> { + fields + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some(name)) + .map(|pair| &pair[1]) +} + +fn string_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(string_value) +} + +fn number_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(number_value) +} + +fn bytes_field(fields: &[redis::Value], name: &str) -> Option> { + match field_value(fields, name)? { + redis::Value::BulkString(bytes) => Some(bytes.clone()), + redis::Value::SimpleString(text) => Some(text.clone().into_bytes()), + _ => None, + } +} + +fn ttl_seconds(ttl: Duration) -> u64 { + ttl.as_secs() + .saturating_add(u64::from(ttl.subsec_nanos() > 0)) + .max(1) +} diff --git a/litellm-rust/crates/cache-redis-semantic/src/lib.rs b/litellm-rust/crates/cache-redis-semantic/src/lib.rs new file mode 100644 index 00000000000..51d0b4ba5f3 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/lib.rs @@ -0,0 +1,5 @@ +mod cache; +mod prompt; + +pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig}; +pub use prompt::prompt_from_context; diff --git a/litellm-rust/crates/cache-redis-semantic/src/prompt.rs b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs new file mode 100644 index 00000000000..b9c38e98d77 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs @@ -0,0 +1,97 @@ +use litellm_cache::SemanticCacheContext; +use serde_json::Value; + +pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { + if let Some(messages) = context.messages.as_ref().and_then(Value::as_array) + && !messages.is_empty() + { + return Some(messages_text(messages)); + } + let input = context.input.as_ref()?; + let mut parts = Vec::new(); + collect_input_text(input, &mut parts); + let prompt = parts.join("\n").trim().to_string(); + (!prompt.is_empty()).then_some(prompt) +} + +fn messages_text(messages: &[Value]) -> String { + let mut text = String::new(); + for message in messages { + let Some(message) = message.as_object() else { + continue; + }; + match message.get("content") { + Some(Value::String(content)) => text.push_str(content), + Some(Value::Array(parts)) => { + for part in parts { + if let Some(text_content) = part.get("text").and_then(Value::as_str) { + text.push_str(text_content); + } + } + } + _ => {} + } + text.push_str(&search_results_text(message.get("search_results"))); + } + text +} + +fn search_results_text(search_results: Option<&Value>) -> String { + let Some(Value::Array(results)) = search_results else { + return String::new(); + }; + let mut text = String::new(); + for result in results { + let Some(result) = result.as_object() else { + continue; + }; + for key in ["source", "title"] { + if let Some(value) = result.get(key).and_then(Value::as_str) { + text.push_str(value); + } + } + if let Some(Value::Array(content)) = result.get("content") { + for block in content { + if let Some(value) = block.get("text").and_then(Value::as_str) { + text.push_str(value); + } + } + } + if let Some(citations) = result.get("citations") { + text.push_str(&citations.to_string()); + } + } + text +} + +fn collect_input_text(value: &Value, parts: &mut Vec) { + match value { + Value::String(text) => { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + } + } + Value::Array(items) => { + for item in items { + collect_input_text(item, parts); + } + } + Value::Object(map) => { + if let Some(content) = map.get("content").filter(|content| !content.is_null()) { + collect_input_text(content, parts); + return; + } + for key in ["text", "output", "input_text", "output_text"] { + if let Some(Value::String(text)) = map.get(key) { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + return; + } + } + } + } + _ => {} + } +} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs new file mode 100644 index 00000000000..233b87ec52f --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -0,0 +1,1003 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{BaseCache, CacheCodec, Error, SemanticCacheContext}; +use litellm_cache_redis_semantic::{Embedder, RedisSemanticCache, RedisSemanticConfig}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use redis_test::{MockCmd, MockRedisConnection}; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; + +const INDEX: &str = "litellm_semantic_cache_index"; + +struct FakeEmbedder { + vectors: HashMap>, + calls: Arc>>, +} + +impl FakeEmbedder { + fn new(vectors: &[(&str, &[f32])]) -> (Self, Arc>>) { + let calls = Arc::new(Mutex::new(Vec::new())); + ( + Self { + vectors: vectors + .iter() + .map(|(prompt, vector)| (prompt.to_string(), vector.to_vec())) + .collect(), + calls: Arc::clone(&calls), + }, + calls, + ) + } +} + +impl Embedder for FakeEmbedder { + fn embed(&self, prompt: &str, _: Option<&Value>) -> Result, Error> { + self.calls.lock().unwrap().push(prompt.to_string()); + + Ok(self + .vectors + .get(prompt) + .cloned() + .unwrap_or_else(|| vec![0.1, 0.2, 0.3])) + } + + async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + self.embed(prompt, metadata) + } +} + +fn config() -> RedisSemanticConfig { + RedisSemanticConfig { + index_name: INDEX.into(), + similarity_threshold: 0.9, + } +} + +fn messages_context(messages: Vec) -> SemanticCacheContext { + SemanticCacheContext { + messages: Some(Value::Array(messages)), + ..Default::default() + } +} + +fn entry() -> CacheEntry { + CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "yes"}), + } +} + +fn encoded(entry: &CacheEntry) -> Vec { + ResponseCacheCodec.encode(entry).unwrap() +} + +fn vector_bytes(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(b"litellm_cache_key"); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn s(value: &str) -> redis::Value { + redis::Value::BulkString(value.as_bytes().to_vec()) +} + +fn unknown_index_error() -> redis::RedisError { + redis::RedisError::from((redis::ErrorKind::Extension, "Unknown index name")) +} + +fn attribute(name: &str, field_type: &str, extra: Vec) -> redis::Value { + let mut parts = vec![ + s("identifier"), + s(name), + s("attribute"), + s(name), + s("type"), + s(field_type), + ]; + parts.extend(extra); + redis::Value::Array(parts) +} + +fn index_info(attributes: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("index_name"), + s(INDEX), + s("attributes"), + redis::Value::Array(attributes), + ]) +} + +fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> redis::Value { + attribute( + "prompt_vector", + "VECTOR", + vec![ + s("algorithm"), + s("FLAT"), + s("data_type"), + s(data_type), + s("dim"), + redis::Value::Int(dims), + s("distance_metric"), + s(distance_metric), + ], + ) +} + +fn vector_attribute(dims: i64) -> redis::Value { + vector_attribute_with(dims, "FLOAT32", "COSINE") +} + +fn info_with_vector(vector: redis::Value) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector, + attribute("litellm_cache_key", "TAG", vec![]), + ]) +} + +fn compatible_info(dims: i64) -> redis::Value { + info_with_vector(vector_attribute(dims)) +} + +fn unscoped_info(dims: i64) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector_attribute(dims), + ]) +} + +fn create_index_command(name: &str, dims: usize) -> redis::Cmd { + let mut command = redis::cmd("FT.CREATE"); + command + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg("prompt_vector") + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg("litellm_cache_key") + .arg("TAG") + .arg("SEPARATOR") + .arg(","); + command +} + +fn search_command(index: &str, tag: &str, vector: &[f32]) -> redis::Cmd { + let mut command = redis::cmd("FT.SEARCH"); + command + .arg(index) + .arg(format!( + "(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]" + )) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg("litellm_cache_key") + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_bytes(vector)); + command +} + +fn hit_fields(tag: &str, distance: &str, response: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("entry_id"), + s("stored-id"), + s("prompt"), + s("hello prompt"), + s("response"), + redis::Value::BulkString(response), + s("inserted_at"), + s("1700000000.5"), + s("updated_at"), + s("1700000000.5"), + s("litellm_cache_key"), + s(tag), + s("vector_distance"), + s(distance), + ]) +} + +fn search_result(fields: redis::Value) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::Int(1), + s("litellm_semantic_cache_index:stored-id"), + fields, + ]) +} + +fn empty_result() -> redis::Value { + redis::Value::Array(vec![redis::Value::Int(0)]) +} + +#[test] +fn store_creates_index_and_writes_hash_with_expire() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(INDEX, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(5), Ok(1)), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(5)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + cache.set_cache(tag, value, &context).unwrap(); +} + +#[test] +fn store_without_ttl_skips_expire() { + let prompt = "hello prompt"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "key1"))) + .arg("entry_id") + .arg(entry_id(prompt, "key1")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("key1"), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + "key1", + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn lookup_returns_hit_below_distance_threshold() { + let vector = vec![0.1f32, 0.2, 0.3]; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let hit = cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]), + ) + .unwrap(); + assert_eq!(hit, Some(value)); +} + +#[test] +fn lookup_misses_above_distance_threshold_and_on_tag_mismatch() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.5", encoded(&entry())))), + ), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "other", + "0.05", + encoded(&entry()), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + let context = messages_context(vec![json!({"role": "user", "content": "hello prompt"})]); + + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); +} + +#[test] +fn lookup_returns_invalid_entry_on_malformed_response() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "key1", + "0.05", + b"not json!".to_vec(), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[test] +fn missing_prompt_is_noop_and_never_embeds() { + let connection = MockRedisConnection::new(Vec::::new()).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let context = SemanticCacheContext::default(); + cache.set_cache("key1", entry(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert!(calls.lock().unwrap().is_empty()); +} + +#[test] +fn scope_overrides_key_as_filter_tag() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "scope-a"))) + .arg("entry_id") + .arg(entry_id(prompt, "scope-a")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("scope-a"), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, "scope\\-a", &vector), + Ok(search_result(hit_fields( + "scope-a", + "0.05", + encoded(&value), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = SemanticCacheContext { + scope: Some("scope-a".into()), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + + cache.set_cache("key1", value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(value)); +} + +#[test] +fn incompatible_schema_falls_back_to_isolated_index() { + let prompt = "hello prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(unscoped_info(3))), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{isolated}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn create_index_race_rechecks_schema_and_stores() { + let prompt = "hello prompt"; + let tag = "key1"; + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new( + create_index_command(INDEX, 3), + Err::<&str, _>(redis::RedisError::from(( + redis::ErrorKind::Extension, + "Index already exists", + ))), + ), + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn wrong_distance_metric_falls_back_to_isolated_index() { + let prompt = "hello prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Ok(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2"))), + ), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{isolated}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn tag_special_characters_are_escaped_in_search_filter() { + let vector = vec![0.1f32, 0.2, 0.3]; + let tag = "a:b, c|d"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "a\\:b\\,\\ c\\|d", &vector), + Ok(empty_result()), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + tag, + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap(), + None + ); +} + +#[test] +fn prompt_extraction_matches_python_message_and_input_shapes() { + let vector = vec![0.1f32, 0.2, 0.3]; + let lookups = 5; + let mut commands = vec![MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Ok(compatible_info(3)), + )]; + for _ in 0..lookups { + commands.push(MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(empty_result()), + )); + } + let connection = MockRedisConnection::new(commands).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + cache + .get_cache( + "key1", + &messages_context(vec![ + json!({"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]}), + json!({"role": "assistant", "content": "reply"}), + ]), + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!(" plain input ")), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some( + json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"]), + ), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!({"output_text": " result text "})), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &messages_context(vec![json!({ + "role": "user", + "content": "question", + "search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}], + })]), + ) + .unwrap(); + + assert_eq!( + *calls.lock().unwrap(), + vec![ + "firstsecondreply", + "plain input", + "nested\ntail", + "result text", + "questionsrctfound{\"a\":1}", + ] + ); +} + +#[test] +fn ttl_passes_through_context_only() { + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection( + MockRedisConnection::new(Vec::::new()), + embedder, + config(), + ); + assert_eq!(cache.get_ttl(&SemanticCacheContext::default()), None); + assert_eq!( + cache.get_ttl(&SemanticCacheContext { + ttl: Some(Duration::from_secs(9)), + ..Default::default() + }), + Some(Duration::from_secs(9)) + ); +} + +#[tokio::test] +async fn async_paths_embed_then_run_blocking_redis_work() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, tag, &vector), + Ok(search_result(hit_fields(tag, "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = messages_context(vec![json!({"role": "user", "content": prompt})]); + + cache + .async_set_cache(tag, value.clone(), context.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache(tag, &context).await.unwrap(), + Some(value) + ); +} + +#[test] +fn shared_base_index_across_dimensions_replaces_the_isolated_index() { + // Pins parity with Python's `_isolated` + overwrite=True flow. + let prompt = "shared prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let value = entry(); + let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); + let store_hash = |index: &str, vector: &[f32]| { + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{index}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ) + }; + + let vector_a = vec![0.1f32; 8]; + let connection_a = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(INDEX, 8), Ok("OK")), + store_hash(INDEX, &vector_a), + ]) + .assert_all_commands_consumed(); + let (embedder_a, _) = FakeEmbedder::new(&[(prompt, &vector_a)]); + let worker_a = RedisSemanticCache::with_connection(connection_a, embedder_a, config()) + .with_clock(|| 1700000000.5); + worker_a.set_cache(tag, value.clone(), &context()).unwrap(); + + let vector_b = vec![0.2f32; 4]; + let connection_b = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(8))), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 4), Ok("OK")), + store_hash(&isolated, &vector_b), + MockCmd::new( + search_command(&isolated, tag, &vector_b), + Ok(search_result(hit_fields(tag, "0.0", encoded(&value)))), + ), + MockCmd::new( + search_command(&isolated, tag, &vector_b), + Err::(redis::RedisError::from(( + redis::ErrorKind::Extension, + "Vector dimension mismatch", + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder_b, _) = FakeEmbedder::new(&[(prompt, &vector_b)]); + let worker_b = RedisSemanticCache::with_connection(connection_b, embedder_b, config()) + .with_clock(|| 1700000000.5); + worker_b.set_cache(tag, value.clone(), &context()).unwrap(); + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap(), + Some(value.clone()) + ); + + let vector_c = vec![0.3f32; 16]; + let connection_c = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(8))), + MockCmd::new(redis::cmd("FT.INFO").arg(&isolated), Ok(compatible_info(4))), + MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")), + MockCmd::new(create_index_command(&isolated, 16), Ok("OK")), + store_hash(&isolated, &vector_c), + ]) + .assert_all_commands_consumed(); + let (embedder_c, _) = FakeEmbedder::new(&[(prompt, &vector_c)]); + let worker_c = RedisSemanticCache::with_connection(connection_c, embedder_c, config()) + .with_clock(|| 1700000000.5); + worker_c.set_cache(tag, value.clone(), &context()).unwrap(); + + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap_err(), + Error::Unavailable + ); +} + +#[test] +fn live_shared_index_is_replaced_across_dimensions() { + let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { + return; + }; + // Pins parity with Python's `_isolated` + overwrite=True flow. + let base = format!("rust_semantic_shared_{}", std::process::id()); + let isolated = format!("{base}_isolated"); + let prompt = "shared live prompt"; + let tag = "key1"; + let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); + let value = entry(); + let worker = |vector: Vec| { + let (embedder, _) = FakeEmbedder::new(&[(prompt, vector.as_slice())]); + RedisSemanticCache::new( + &url, + embedder, + RedisSemanticConfig { + index_name: base.clone(), + similarity_threshold: 0.9, + }, + ) + .unwrap() + }; + + let worker_a = worker(vec![0.1f32; 8]); + worker_a.set_cache(tag, value.clone(), &context()).unwrap(); + + let worker_b = worker(vec![0.2f32; 4]); + worker_b.set_cache(tag, value.clone(), &context()).unwrap(); + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap(), + Some(value.clone()) + ); + + let worker_c = worker(vec![0.3f32; 16]); + worker_c.set_cache(tag, value.clone(), &context()).unwrap(); + + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap_err(), + Error::Unavailable + ); + + let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); + for index in [&base, &isolated] { + let _: Result<(), _> = redis::cmd("FT.DROPINDEX") + .arg(index) + .arg("DD") + .query(&mut connection); + } +} + +#[test] +fn live_store_lookup_and_ttl_against_redis_stack() { + let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { + return; + }; + let vector = vec![0.1f32, 0.2, 0.3, 0.4]; + let prompt = "rust semantic cache live prompt"; + let tag = "live-key"; + let index_name = format!("rust_semantic_test_{}", std::process::id()); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::new( + &url, + embedder, + RedisSemanticConfig { + index_name: index_name.clone(), + similarity_threshold: 0.9, + }, + ) + .unwrap(); + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(120)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + let value = entry(); + + cache.set_cache(tag, value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache(tag, &context).unwrap(), Some(value)); + assert_eq!(cache.get_cache("other-key", &context).unwrap(), None); + + let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); + let ttl: i64 = redis::Commands::ttl( + &mut connection, + format!("{index_name}:{}", entry_id(prompt, tag)), + ) + .unwrap(); + assert!( + ttl > 0, + "expected stored hash to carry an expiry, got {ttl}" + ); +} diff --git a/litellm-rust/crates/cache-redis/src/cache/connection.rs b/litellm-rust/crates/cache-redis/src/cache/connection.rs index 06364296992..013bf055f89 100644 --- a/litellm-rust/crates/cache-redis/src/cache/connection.rs +++ b/litellm-rust/crates/cache-redis/src/cache/connection.rs @@ -12,7 +12,7 @@ use redis::{ use super::REDIS_TIMEOUT; use crate::topology::RedisNode; -pub(super) struct PooledConnection { +pub struct PooledConnection { pub(super) connection: C, pub(super) failed: bool, } @@ -20,7 +20,7 @@ pub(super) 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. -pub(super) struct ConnectionManager(redis::Client); +pub struct ConnectionManager(redis::Client); impl ConnectionManager { pub(super) fn open(url: &str) -> Result { @@ -54,7 +54,7 @@ impl r2d2::ManageConnection for ConnectionManager { } } -pub(super) struct ClusterConnectionManager(ClusterClient); +pub struct ClusterConnectionManager(ClusterClient); impl ClusterConnectionManager { pub(super) fn open(url: &str, startup_nodes: &[RedisNode]) -> Result { diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 9180ee9d0dc..36307ac9b33 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,7 +1,8 @@ use std::{sync::Mutex, time::Duration}; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache, + BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, SemanticCacheContext, + get_cache, }; struct TestCache { @@ -126,6 +127,24 @@ fn associated_context_preserves_backend_specific_lookup_inputs() { ); } +#[test] +fn semantic_context_with_ttl_preserves_lookup_inputs() { + let context = SemanticCacheContext { + input: Some(serde_json::json!("text")), + messages: Some(serde_json::json!([{"role": "user", "content": "hi"}])), + metadata: Some(serde_json::json!({"key": "value"})), + scope: Some("scope".into()), + ttl: None, + }; + let updated = context.with_ttl(Some(Duration::from_secs(30))); + assert_eq!(updated.ttl(), Some(Duration::from_secs(30))); + assert_eq!(updated.input, context.input); + assert_eq!(updated.messages, context.messages); + assert_eq!(updated.metadata, context.metadata); + assert_eq!(updated.scope, context.scope); + assert_eq!(context.with_ttl(None).ttl(), None); +} + #[tokio::test] async fn default_batch_operations_use_async_writes_and_stop_on_failure() { let cache = TestCache { diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index d07a9839ebf..bbe727f9345 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -27,6 +27,7 @@ litellm-cache-redis.workspace = true litellm-cache-s3.workspace = true litellm-cache-gcs.workspace = true litellm-cache-disk.workspace = true +litellm-cache-redis-semantic.workspace = true litellm-cache-response.workspace = true litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" } serde.workspace = true @@ -45,7 +46,7 @@ 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"] } +tokio = { workspace = true, features = ["rt", "sync"] } [dev-dependencies] serde.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 9bc666c4f2d..ef2a03ac9ba 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -88,6 +88,24 @@ pub(super) struct GcsCacheConfig { pub(super) path_service_account: Option, } +pub(super) struct AzureBlobCacheConfig { + pub(super) account_url: String, + pub(super) container: String, +} + +#[allow( + dead_code, + reason = "embedding settings are projected so drift falls back to Python" +)] +pub(super) struct RedisSemanticCacheConfig { + pub(super) redis_url: String, + pub(super) index_name: String, + pub(super) similarity_threshold: f64, + pub(super) embedding_model: String, + pub(super) embedding_max_input_tokens: Option, + pub(super) embedding_timeout: Option, +} + struct RedisClientProjection<'py> { topology: RedisTopology, host: String, @@ -107,11 +125,6 @@ pub(super) struct ValkeySemanticCacheConfig { pub(super) connection: RedisConnectionConfig, } -pub(super) struct AzureBlobCacheConfig { - pub(super) account_url: String, - pub(super) container: String, -} - pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), @@ -120,6 +133,7 @@ pub(super) enum CacheBackendConfig { ValkeySemantic(Box), Disk(DiskCacheConfig), AzureBlob(AzureBlobCacheConfig), + RedisSemantic(Box), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -230,9 +244,15 @@ impl NativeCacheConfig { backend: CacheBackendConfig::AzureBlob(backend), })) }), - Some(CacheType::RedisSemantic | CacheType::QdrantSemantic) | None => Ok( - CacheConfigProjection::Unsupported(UnsupportedCacheConfig::Backend), - ), + Some(CacheType::RedisSemantic) => project_redis_semantic(&backend).map(|backend| { + CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::RedisSemantic(Box::new(backend)), + })) + }), + Some(CacheType::QdrantSemantic) | None => Ok(CacheConfigProjection::Unsupported( + UnsupportedCacheConfig::Backend, + )), } } @@ -244,7 +264,8 @@ impl NativeCacheConfig { CacheBackendConfig::ValkeySemantic(_) => Some(Duration::ZERO), CacheBackendConfig::Disk(_) | CacheBackendConfig::AzureBlob(_) - | CacheBackendConfig::Gcs(_) => None, + | CacheBackendConfig::Gcs(_) + | CacheBackendConfig::RedisSemantic(_) => None, }; if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_)) && service.default_ttl() != default_ttl @@ -343,6 +364,20 @@ impl NativeCacheConfig { let facade = std::fs::canonicalize(&config.directory).ok(); (native != facade).then_some("facade and native backend directories must match") } + CacheBackendConfig::RedisSemantic(_) if service.kind() != "redis_semantic" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::RedisSemantic(config) + if service.index_name() != Some(config.index_name.as_str()) => + { + Some("facade and native backend index names must match") + } + CacheBackendConfig::RedisSemantic(config) + if service.similarity_threshold() != Some(config.similarity_threshold as f32) => + { + Some("facade and native backend similarity thresholds must match") + } + CacheBackendConfig::RedisSemantic(_) => None, CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() { None => Some("facade and native backend types must match"), Some((account_url, container)) @@ -371,6 +406,27 @@ fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult, +) -> PyResult { + Ok(RedisSemanticCacheConfig { + redis_url: backend.getattr("_redis_url")?.extract::()?, + index_name: backend + .getattr("_index_name")? + .extract::>()? + .unwrap_or_else(|| "litellm_semantic_cache_index".into()), + similarity_threshold: backend.getattr("similarity_threshold")?.extract::()?, + embedding_model: backend.getattr("embedding_model")?.extract::()?, + embedding_max_input_tokens: backend + .getattr("embedding_max_input_tokens")? + .extract::>()?, + embedding_timeout: backend + .getattr("embedding_timeout")? + .extract::>()?, + }) +} + #[inline(never)] fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { let max_size_kib = backend.getattr("max_size_per_item")?.extract::()?; diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs index 3de0ceb3b67..9398e5a862b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/embedder.rs +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -1,23 +1,44 @@ -use std::{future::Future, sync::Arc}; +use std::future::Future; use litellm_cache::Error; -use litellm_cache_valkey_semantic::Embedder; use litellm_host_python::to_py; -use pyo3::{PyTraverseError, PyVisit, prelude::*}; +use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict}; use serde_json::Value; -#[derive(Clone)] -pub(super) struct PythonEmbedder { - sync_embed: Arc>, - async_embed_callable: Arc>, +tokio::task_local! { + static PREPARED_EMBEDDING: Result, Error>; +} + +pub(super) fn with_prepared_embedding( + vector: Result, Error>, + future: F, +) -> impl Future { + PREPARED_EMBEDDING.scope(vector, future) +} + +pub(super) struct PythonEmbedder(Py); + +impl Clone for PythonEmbedder { + fn clone(&self) -> Self { + Python::attach(|py| Self(self.0.clone_ref(py))) + } } impl PythonEmbedder { + pub(super) fn new(object: Py) -> Self { + Self(object) + } + pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult { - Ok(Self { - sync_embed: Arc::new(backend.getattr("_get_embedding")?.unbind()), - async_embed_callable: Arc::new(backend.getattr("_get_async_embedding")?.unbind()), - }) + Ok(Self(backend.clone().unbind())) + } + + pub(super) fn object(&self) -> &Py { + &self.0 + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) } pub(super) fn async_embed_awaitable<'py>( @@ -27,37 +48,109 @@ impl PythonEmbedder { metadata: &Option, ) -> PyResult> { let metadata = to_py(py, metadata)?; - self.async_embed_callable.bind(py).call1((prompt, metadata)) + self.0 + .bind(py) + .call_method1("_get_async_embedding", (prompt, metadata)) } - pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - visit.call(&*self.sync_embed)?; - visit.call(&*self.async_embed_callable) + fn metadata_kwargs<'py>( + py: Python<'py>, + metadata: Option<&Value>, + ) -> PyResult> { + let kwargs = PyDict::new(py); + kwargs.set_item("metadata", to_py(py, &metadata)?)?; + Ok(kwargs) + } + + pub(super) fn async_embedding_coroutine( + &self, + py: Python<'_>, + prompt: &str, + metadata: Option<&Value>, + ) -> PyResult> { + let kwargs = Self::metadata_kwargs(py, metadata)?; + self.0 + .bind(py) + .call_method("_get_async_embedding", (prompt,), Some(&kwargs)) + .map(Bound::unbind) + } + + pub(super) fn extract(vector: Bound<'_, PyAny>) -> PyResult> { + Ok(vector + .extract::>()? + .into_iter() + .map(|value| value as f32) + .collect()) } } -impl Embedder for PythonEmbedder { +impl litellm_cache_valkey_semantic::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 + self.0 .bind(py) - .call1((prompt, metadata))? + .call_method1("_get_embedding", (prompt, metadata))? .extract() }) .map_err(|_| Error::Unavailable)?; 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>, ) -> impl Future, Error>> + Send { - async { Err(Error::Unavailable) } + let seeded = PREPARED_EMBEDDING + .try_with(Clone::clone) + .unwrap_or(Err(Error::Unavailable)); + std::future::ready(seeded) + } +} + +impl litellm_cache_redis_semantic::Embedder for PythonEmbedder { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + Python::attach(|py| { + let kwargs = Self::metadata_kwargs(py, metadata)?; + Self::extract(self.0.bind(py).call_method( + "_get_embedding", + (prompt,), + Some(&kwargs), + )?) + }) + .map_err(|_| Error::Unavailable) + } + + fn async_embed( + &self, + _prompt: &str, + _metadata: Option<&Value>, + ) -> impl Future, Error>> + Send { + let seeded = PREPARED_EMBEDDING + .try_with(Clone::clone) + .unwrap_or(Err(Error::Unavailable)); + std::future::ready(seeded) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn async_embed_returns_the_seeded_vector_or_unavailable() { + Python::initialize(); + let object = Python::attach(|py| py.None()); + let embedder = PythonEmbedder::new(object); + let scoped_embedder = embedder.clone(); + let scoped = with_prepared_embedding(Ok(vec![0.25]), async move { + litellm_cache_redis_semantic::Embedder::async_embed(&scoped_embedder, "prompt", None) + .await + }); + assert_eq!(scoped.await, Ok(vec![0.25])); + let unscoped = + litellm_cache_redis_semantic::Embedder::async_embed(&embedder, "prompt", None).await; + assert_eq!(unscoped, Err(Error::Unavailable)); } } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 4b4e3255cb4..f1389b745e9 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -363,6 +363,11 @@ impl FacadeGuard { let (module, name, cache_kind) = match (kind, cluster) { ("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), ("redis", false) => ("litellm.caching.redis_cache", "RedisCache", "redis"), + ("redis_semantic", _) => ( + "litellm.caching.redis_semantic_cache", + "RedisSemanticCache", + "redis-semantic", + ), ("redis", true) => ( "litellm.caching.redis_cluster_cache", "RedisClusterCache", @@ -400,6 +405,15 @@ impl FacadeGuard { if let Some(message) = config.service_mismatch(service) { return Err(PyTypeError::new_err(message)); } + if kind == "redis_semantic" + && service + .embedder_object() + .is_none_or(|embedder| !backend.is(embedder.bind(py))) + { + return Err(PyTypeError::new_err( + "facade backend must be the native embedder", + )); + } Ok(Self { outer: ObjectGuard::capture( py, @@ -425,6 +439,13 @@ impl FacadeGuard { "redis_kwargs", "redis_flush_size", "similarity_threshold", + "distance_threshold", + "embedding_model", + "embedding_max_input_tokens", + "embedding_timeout", + "_index_name", + "_redis_url", + "similarity_threshold", "embedding_model", "index_name", "embedding_max_input_tokens", diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 769ad3548be..e6edf48d72f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,13 +1,18 @@ use litellm_auth_aws::AwsAuthConfig; use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; use litellm_cache_redis::{RedisNode, RedisTopology}; +use litellm_cache_redis_semantic::RedisSemanticConfig; use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; use litellm_host_python::{release_gil, run_sync_value}; -use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyRuntimeError, PyTypeError}, + prelude::*, +}; use super::{ - cache_error, embedder::PythonEmbedder, facade::FacadeGuard, native::NativeResponseCache, - request::duration, + cache_error, config::project_redis_semantic, embedder::PythonEmbedder, facade::FacadeGuard, + native::NativeResponseCache, request::duration, }; #[pyclass(frozen, name = "_CacheTestHandle")] @@ -179,6 +184,36 @@ impl CacheTestHandle { }) } + #[staticmethod] + fn redis_semantic(py: Python<'_>, backend: Bound<'_, PyAny>) -> PyResult { + let class = py + .import("litellm.caching.redis_semantic_cache")? + .getattr("RedisSemanticCache")?; + if !backend.get_type().is(&class) { + return Err(PyTypeError::new_err( + "native redis-semantic handles require the built-in RedisSemanticCache", + )); + } + let config = project_redis_semantic(&backend)?; + let embedder = PythonEmbedder::new(backend.unbind()); + let service = release_gil(py, move || { + NativeResponseCache::redis_semantic( + &config.redis_url, + embedder, + RedisSemanticConfig { + index_name: config.index_name, + similarity_threshold: config.similarity_threshold as f32, + }, + ) + }) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() @@ -210,6 +245,7 @@ impl CacheTestHandle { } fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + self.service.traverse(&visit)?; if let Some(guard) = &self.guard { guard.traverse(visit)?; } diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 278d3da1ff9..fa028518559 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; mod semantic_step; use litellm_cache::Error; diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 575dae45833..acc457d8c9c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -8,18 +8,20 @@ use litellm_cache_disk::DiskCache; use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::{RedisCache, RedisTopology}; +use litellm_cache_redis_semantic::{RedisSemanticCache, RedisSemanticConfig}; use litellm_cache_response::{ CacheEntry, CacheKeyField, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, }; use litellm_cache_s3::{S3Cache, S3CacheConfig}; use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig}; -use pyo3::prelude::*; +use pyo3::{PyTraverseError, PyVisit, prelude::*}; use serde_json::Value; use super::{ embedder::PythonEmbedder, request::NativeRequest, + semantic::{SemanticBody, SemanticOperation, drive}, semantic_step::{SemanticEmbedExecution, drive_semantic}, }; @@ -86,6 +88,10 @@ pub(super) enum NativeResponseCache { embedder: PythonEmbedder, scope: String, }, + RedisSemantic { + cache: Arc>>, + embedder: PythonEmbedder, + }, Disk(Arc>>), AzureBlob(Arc>>), } @@ -150,6 +156,18 @@ impl NativeResponseCache { }) } + pub fn redis_semantic( + url: &str, + embedder: PythonEmbedder, + config: RedisSemanticConfig, + ) -> Result { + let backend = RedisSemanticCache::new(url, embedder.clone(), config)?; + Ok(Self::RedisSemantic { + cache: Arc::new(ResponseCache::new(Arc::new(backend))), + embedder, + }) + } + pub fn disk(directory: &str) -> Result { let cache = DiskCache::open(directory, ResponseCacheCodec)?; Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache))))) @@ -190,6 +208,7 @@ impl NativeResponseCache { | Self::Redis { .. } | Self::S3(_) | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::Disk(_) | Self::Gcs(_) => None, } @@ -204,6 +223,23 @@ impl NativeResponseCache { } } + pub(super) fn redis_semantic_request( + request: &NativeRequest, + ) -> ResponseCacheRequest { + ResponseCacheRequest { + key: request.key.clone(), + controls: request.controls, + context: SemanticCacheContext { + input: request.input.clone(), + messages: request.messages.clone(), + metadata: request.metadata.clone(), + scope: request.scope.clone(), + ttl: request.ttl, + }, + max_age: request.max_age, + } + } + fn semantic( request: &NativeRequest, scope: &str, @@ -252,6 +288,7 @@ impl NativeResponseCache { Self::S3(_) => "s3", Self::Gcs(_) => "gcs", Self::ValkeySemantic { .. } => "valkey-semantic", + Self::RedisSemantic { .. } => "redis_semantic", Self::Disk(_) => "disk", Self::AzureBlob(_) => "azure-blob", } @@ -264,6 +301,7 @@ impl NativeResponseCache { Self::S3(cache) => cache.default_ttl(), Self::Gcs(cache) => cache.default_ttl(), Self::ValkeySemantic { cache, .. } => cache.default_ttl(), + Self::RedisSemantic { cache, .. } => cache.default_ttl(), Self::Disk(cache) => cache.default_ttl(), Self::AzureBlob(cache) => cache.default_ttl(), } @@ -302,6 +340,7 @@ impl NativeResponseCache { Self::Memory(_) | Self::S3(_) | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, @@ -314,6 +353,7 @@ impl NativeResponseCache { Self::Memory(_) | Self::S3(_) | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, @@ -327,6 +367,7 @@ impl NativeResponseCache { Self::Redis { .. } | Self::S3(_) | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, @@ -339,6 +380,7 @@ impl NativeResponseCache { Self::Redis { .. } | Self::S3(_) | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, @@ -352,6 +394,7 @@ impl NativeResponseCache { | Self::Redis { .. } | Self::S3(_) | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::AzureBlob(_) | Self::Gcs(_) => None, } @@ -363,6 +406,38 @@ impl NativeResponseCache { cache.backend().similarity_threshold(), cache.backend().index_name(), )), + Self::RedisSemantic { cache, .. } => Some(( + f64::from(cache.backend().similarity_threshold()), + cache.backend().index_name(), + )), + _ => None, + } + } + + pub fn index_name(&self) -> Option<&str> { + match self { + Self::RedisSemantic { cache, .. } => Some(cache.backend().index_name()), + _ => None, + } + } + + pub fn similarity_threshold(&self) -> Option { + match self { + Self::RedisSemantic { cache, .. } => Some(cache.backend().similarity_threshold()), + _ => None, + } + } + + pub fn semantic_embedder(&self) -> Option<&PythonEmbedder> { + match self { + Self::RedisSemantic { embedder, .. } => Some(embedder), + _ => None, + } + } + + pub fn embedder_object(&self) -> Option<&Py> { + match self { + Self::RedisSemantic { embedder, .. } => Some(embedder.object()), _ => None, } } @@ -375,6 +450,9 @@ impl NativeResponseCache { Self::ValkeySemantic { cache, scope, .. } => { cache.lookup(&Self::semantic(request, scope), now) } + Self::RedisSemantic { cache, .. } => { + cache.lookup(&Self::redis_semantic_request(request), now) + } Self::Gcs(cache) => cache.lookup(&Self::exact(request), now), Self::Disk(cache) => cache.lookup(&Self::exact(request), now), Self::AzureBlob(cache) => cache.lookup(&Self::exact(request), now), @@ -394,6 +472,9 @@ impl NativeResponseCache { Self::ValkeySemantic { cache, scope, .. } => { cache.store(&Self::semantic(request, scope), response, now) } + Self::RedisSemantic { cache, .. } => { + cache.store(&Self::redis_semantic_request(request), response, now) + } Self::Gcs(cache) => cache.store(&Self::exact(request), response, now), Self::Disk(cache) => cache.store(&Self::exact(request), response, now), Self::AzureBlob(cache) => cache.store(&Self::exact(request), response, now), @@ -417,7 +498,9 @@ impl NativeResponseCache { Self::S3(cache) => { cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) } - Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + Err(Error::UnsupportedOperation) + } Self::Gcs(cache) => { cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) } @@ -444,6 +527,11 @@ impl NativeResponseCache { .async_lookup(&Self::semantic(request, scope), now) .await } + Self::RedisSemantic { cache, .. } => { + cache + .async_lookup(&Self::redis_semantic_request(request), now) + .await + } Self::Gcs(cache) => cache.async_lookup(&Self::exact(request), now).await, Self::Disk(cache) => cache.async_lookup(&Self::exact(request), now).await, Self::AzureBlob(cache) => cache.async_lookup(&Self::exact(request), now).await, @@ -481,6 +569,10 @@ impl NativeResponseCache { Self::semantic(&request, scope), ), ), + Self::RedisSemantic { .. } => drive( + py, + SemanticBody::new(self.clone(), SemanticOperation::Lookup(request)), + ), } } @@ -522,6 +614,11 @@ impl NativeResponseCache { .async_store(&Self::semantic(request, scope), response, now) .await } + Self::RedisSemantic { cache, .. } => { + cache + .async_store(&Self::redis_semantic_request(request), response, now) + .await + } Self::Gcs(cache) => { cache .async_store(&Self::exact(request), response, now) @@ -577,6 +674,10 @@ impl NativeResponseCache { response, ), ), + Self::RedisSemantic { .. } => drive( + py, + SemanticBody::new(self.clone(), SemanticOperation::Store(request, response)), + ), } } @@ -599,7 +700,9 @@ impl NativeResponseCache { .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) .await } - Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + Err(Error::UnsupportedOperation) + } Self::Gcs(cache) => { cache .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) @@ -652,6 +755,7 @@ impl NativeResponseCache { .collect(); cache.async_store_batch(entries, now).await } + Self::RedisSemantic { .. } => Err(Error::UnsupportedOperation), Self::Gcs(cache) => { let entries = entries .into_iter() @@ -718,6 +822,10 @@ impl NativeResponseCache { ), ) } + Self::RedisSemantic { .. } => drive( + py, + SemanticBody::new(self.clone(), SemanticOperation::StoreBatch(entries.into())), + ), } } @@ -731,7 +839,9 @@ impl NativeResponseCache { cache.async_flush().await } Self::S3(cache) => cache.async_flush().await, - Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation), + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + Err(Error::UnsupportedOperation) + } Self::Gcs(cache) => cache.async_flush().await, Self::Disk(cache) => cache.async_flush().await, Self::AzureBlob(cache) => cache.async_flush().await, @@ -744,12 +854,22 @@ impl NativeResponseCache { Self::Redis { cache, .. } => cache.test_connection().await, Self::S3(cache) => cache.test_connection().await, Self::ValkeySemantic { cache, .. } => cache.test_connection().await, + Self::RedisSemantic { .. } => Err(Error::UnsupportedOperation), Self::Gcs(cache) => cache.test_connection().await, Self::Disk(cache) => cache.test_connection().await, Self::AzureBlob(cache) => cache.test_connection().await, } } + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + match self { + Self::ValkeySemantic { embedder, .. } => embedder.traverse(visit)?, + Self::RedisSemantic { embedder, .. } => embedder.traverse(visit)?, + _ => {} + } + Ok(()) + } + pub fn gcs_backend(&self) -> Option<&GcsCache> { match self { Self::Gcs(cache) => Some(cache.backend()), @@ -777,6 +897,7 @@ mod tests { metadata: Some(metadata), litellm_metadata: None, litellm_params: None, + scope: None, } } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 036951891a1..3b4b910c1f0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -19,8 +19,10 @@ struct RequestInput { metadata: Option, litellm_metadata: Option, litellm_params: Option, + scope: Option, } +#[derive(Clone)] pub(super) struct NativeRequest { pub(super) key: CacheKeyInput, pub(super) controls: CacheControls, @@ -31,6 +33,7 @@ pub(super) struct NativeRequest { pub(super) metadata: Option, pub(super) litellm_metadata: Option, pub(super) litellm_params: Option, + pub(super) scope: Option, } pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { @@ -52,6 +55,7 @@ fn request_input(input: RequestInput) -> PyResult { metadata: input.metadata, litellm_metadata: input.litellm_metadata, litellm_params: input.litellm_params, + scope: input.scope, }) } diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/semantic.rs new file mode 100644 index 00000000000..7661598c5f9 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/semantic.rs @@ -0,0 +1,175 @@ +use std::collections::VecDeque; + +use litellm_cache::Error; +use litellm_cache_redis_semantic::prompt_from_context; +use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyException, PyRuntimeError}, + prelude::*, +}; +use serde_json::Value; + +use super::{ + cache_error, + embedder::{PythonEmbedder, with_prepared_embedding}, + native::NativeResponseCache, + request::{NativeRequest, now}, +}; + +pub(super) enum SemanticOperation { + Lookup(NativeRequest), + Store(NativeRequest, Value), + StoreBatch(VecDeque<(NativeRequest, Value)>), +} + +enum Phase { + Start, + AwaitingEmbedding, + AwaitingBackend, +} + +pub(super) struct SemanticBody { + service: NativeResponseCache, + operation: SemanticOperation, + pending: Option<(NativeRequest, Option)>, + phase: Phase, +} + +impl SemanticBody { + pub(super) fn new(service: NativeResponseCache, operation: SemanticOperation) -> Self { + Self { + service, + operation, + pending: None, + phase: Phase::Start, + } + } + + fn backend_step( + &mut self, + py: Python<'_>, + seed: Result, Error>, + ) -> PyResult { + self.phase = Phase::AwaitingBackend; + let (request, response) = self.pending.take().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution resumed without a pending operation") + })?; + let service = self.service.clone(); + let future = async move { + match response { + None => service.async_lookup(&request, now()).await, + Some(response) => service + .async_store(&request, response, now()) + .await + .map(|_| None), + } + }; + let awaitable = run_async(py, with_prepared_embedding(seed, future), cache_error)?; + Ok(ExecutionStep::Await(awaitable.unbind())) + } +} + +impl ExecutionBody for SemanticBody { + fn resume(&mut self, mut result: Option>>) -> PyResult { + Python::attach(|py| { + loop { + match self.phase { + Phase::Start => { + if result.is_some() { + return Err(PyRuntimeError::new_err( + "semantic execution received a result before starting", + )); + } + if self.pending.is_none() { + match &mut self.operation { + SemanticOperation::Lookup(request) => { + self.pending = Some((request.clone(), None)); + } + SemanticOperation::Store(request, response) => { + let response = std::mem::replace(response, Value::Null); + self.pending = Some((request.clone(), Some(response))); + } + SemanticOperation::StoreBatch(queue) => { + let Some((request, response)) = queue.pop_front() else { + return Ok(ExecutionStep::Return(py.None())); + }; + self.pending = Some((request, Some(response))); + } + } + } + let (request, _) = self.pending.as_ref().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution has no pending operation") + })?; + let semantic = NativeResponseCache::redis_semantic_request(request); + let Some(prompt) = prompt_from_context(&semantic.context) else { + return self.backend_step(py, Err(Error::Unavailable)); + }; + let embedder = self.service.semantic_embedder().ok_or_else(|| { + PyRuntimeError::new_err( + "semantic execution requires a redis-semantic backend", + ) + })?; + let coroutine = embedder.async_embedding_coroutine( + py, + &prompt, + semantic.context.metadata.as_ref(), + )?; + self.phase = Phase::AwaitingEmbedding; + return Ok(ExecutionStep::Await(coroutine)); + } + Phase::AwaitingEmbedding => { + let result = result.take().ok_or_else(|| { + PyRuntimeError::new_err( + "semantic execution expected an embedding result", + ) + })?; + let seed = match result { + Ok(value) => PythonEmbedder::extract(value.into_bound(py)) + .map_err(|_| Error::Unavailable), + Err(error) => { + if !error.is_instance_of::(py) { + return Err(error); + } + Err(Error::Unavailable) + } + }; + return self.backend_step(py, seed); + } + Phase::AwaitingBackend => { + let result = result.take().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution expected a backend result") + })?; + let value = match result { + Ok(value) => value, + Err(error) => return Err(error), + }; + let more = matches!( + &self.operation, + SemanticOperation::StoreBatch(queue) if !queue.is_empty() + ); + if more { + self.phase = Phase::Start; + continue; + } + return Ok(ExecutionStep::Return(value)); + } + } + } + }) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Some(embedder) = self.service.semantic_embedder() { + embedder.traverse(visit)?; + } + Ok(()) + } +} + +pub(super) fn drive(py: Python<'_>, body: SemanticBody) -> PyResult> { + let execution = Py::new(py, Execution::new(body))?; + py.import("litellm.rust_bridge.lifecycle")? + .getattr("drive")? + .call1((execution,)) +} diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 978e770b4be..c0bb93daadc 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -169,6 +169,8 @@ class _CacheTestHandle: @staticmethod def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ... @staticmethod + def redis_semantic(backend: object) -> _CacheTestHandle: ... + @staticmethod def valkey_semantic( url: str, similarity_threshold: float, diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index e0d92a2e957..45f7214ecbf 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -1,19 +1,23 @@ import asyncio import contextvars import gc +import hashlib import json +import math import os import threading import time import uuid import weakref -from collections.abc import Generator +from collections.abc import Callable, Generator +from contextlib import ExitStack from datetime import datetime from pathlib import Path from types import SimpleNamespace from typing import Final, Protocol, cast from unittest.mock import Mock from urllib.parse import urlparse +from uuid import uuid4 import boto3 import botocore.config @@ -30,13 +34,19 @@ from litellm.caching.disk_cache import DiskCache from litellm.caching.gcs_cache import GCSCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.caching.redis_semantic_cache import RedisSemanticCache from litellm.caching.s3_cache import S3Cache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType +from litellm.types.llms.custom_llm import CustomLLMItem +from litellm.types.utils import EmbeddingResponse from tests.test_litellm_rust.support.fake_gcs import FakeGcs from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.s3_stub import S3Stub +_CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name +_CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name + pytestmark: Final = pytest.mark.requires_rust_extension @@ -114,14 +124,14 @@ def test_existing_constructor_and_global_are_unchanged() -> None: assert type(facade.cache) is InMemoryCache assert "_native_cache_handle" not in vars(facade) with rebound(litellm, "cache", facade): - resolver: Final = _native._CacheTestResolver(litellm) + resolver: Final = _CacheTestResolver(litellm) assert resolver.resolve().kind == "python_callback" resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None: - resolver: Final = _native._CacheTestResolver(litellm) + resolver: Final = _CacheTestResolver(litellm) enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30) enabled: Final = litellm.cache @@ -144,13 +154,13 @@ def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> Non async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: - namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.memory()) - resolver: Final = _native._CacheTestResolver(namespace) + namespace: Final = SimpleNamespace(cache=_CacheTestHandle.memory()) + resolver: Final = _CacheTestResolver(namespace) selected: Final = resolver.resolve() assert selected.kind == "native" selected.store(request(), {"answer": 1}) assert await selected.async_lookup(request()) == {"answer": 1} - with rebound(namespace, "cache", _native._CacheTestHandle.memory()): + with rebound(namespace, "cache", _CacheTestHandle.memory()): replacement: Final = resolver.resolve() await selected.async_store(request(), {"answer": 2}) assert replacement.lookup(request()) is None @@ -183,7 +193,7 @@ async def test_python_callback_preserves_identity_caller_task_context_and_errors raise failure namespace: Final = SimpleNamespace(cache=CustomCache()) - binding: Final = _native._CacheTestResolver(namespace).resolve() + binding: Final = _CacheTestResolver(namespace).resolve() assert binding.kind == "python_callback" assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel assert context.get() == "callback" @@ -204,7 +214,7 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: finally: finished.set() - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() async def lookup() -> object: return await binding.async_lookup(None, callback_kwargs={}) @@ -219,9 +229,9 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle: Final = _native._CacheTestHandle.memory() + handle: Final = _CacheTestHandle.memory() handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) native: Final = resolver.resolve() assert native.kind == "native" native.store(request(), {"source": "native"}) @@ -252,12 +262,12 @@ def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not class CustomCache(Cache): pass - handle: Final = _native._CacheTestHandle.memory() + handle: Final = _CacheTestHandle.memory() with pytest.raises(TypeError): handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) facade: Final = Cache(type=LiteLLMCacheType.LOCAL) handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) with rebound(facade, "cache", InMemoryCache()): assert resolver.resolve().kind == "python_callback" with rebound(facade, "ttl", 12): @@ -282,7 +292,7 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: def cyclic_reference() -> weakref.ReferenceType[CustomCache]: callback: Final = CustomCache() namespace: Final = SimpleNamespace(cache=callback) - binding: Final = _native._CacheTestResolver(namespace).resolve() + binding: Final = _CacheTestResolver(namespace).resolve() setattr(callback, "binding", binding) return weakref.ref(callback) @@ -293,8 +303,8 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: client: Final = redis.Redis.from_url(redis_url) - namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url, namespace="team")) - binding: Final = _native._CacheTestResolver(namespace).resolve() + namespace: Final = SimpleNamespace(cache=_CacheTestHandle.redis(redis_url, namespace="team")) + binding: Final = _CacheTestResolver(namespace).resolve() response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} client.set("team:sync", str(envelope)) @@ -316,33 +326,33 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd def test_invalid_duration_and_request_shape_fail_before_storage() -> None: - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() for seconds in (-1.0, float("nan"), float("inf")): with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) assert binding.lookup(request()) is None with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): - _native._CacheTestHandle.memory(ttl_seconds=-1) + _CacheTestHandle.memory(ttl_seconds=-1) async def test_memory_size_policy_is_applied_by_the_native_host() -> None: - handle: Final = _native._CacheTestHandle.memory(capacity=2, max_entry_bytes=128) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + handle: Final = _CacheTestHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=handle)).resolve() small: Final = {"answer": "ok"} binding.store(request("small"), small) assert await binding.async_lookup(request("small")) == small await binding.async_store(request("large"), {"answer": "x" * 256}) assert binding.lookup(request("large")) is None assert binding.lookup(request("small")) == small - disabled: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.memory(capacity=0)) + disabled: Final = _CacheTestResolver( + SimpleNamespace(cache=_CacheTestHandle.memory(capacity=0)) ).resolve() await disabled.async_store(request(), small) assert await disabled.async_lookup(request()) is None async def test_native_batch_lookup_and_store_report_partial_hits() -> None: - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() requests: Final = [request("hit"), request("miss"), request("disabled")] requests[2]["controls"] = { "supported_call_type": True, @@ -380,7 +390,7 @@ async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: ) -> object: return result, kwargs - binding: Final = _native._CacheTestResolver( + binding: Final = _CacheTestResolver( SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL)) ).resolve() assert binding.kind == "python_callback" @@ -410,7 +420,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) cache.cache.set_cache("key", "value") - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=cache)).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=cache)).resolve() assert binding.kind == "python_callback" setattr(cache.cache, "ping", ping) @@ -422,7 +432,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: def test_facade_registration_rejects_mismatched_capacity() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) with pytest.raises(TypeError, match="capacities must match"): - _native._CacheTestHandle.memory(capacity=7)._bind_facade(facade) + _CacheTestHandle.memory(capacity=7)._bind_facade(facade) def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure_blob_facade: Cache) -> None: @@ -525,19 +535,19 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: redis_flush_size=2, ) with pytest.raises(TypeError, match="default TTLs must match"): - _native._CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) + _CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) with pytest.raises(TypeError, match="namespaces must match"): - _native._CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) - _native._CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + _CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) + _CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() client: Final = redis.Redis.from_url(redis_url) with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" pool: Final = facade.cache.redis_client.connection_pool with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}): - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" await binding.async_store(request("first"), {"value": 1}) assert client.get("first") is None @@ -1112,3 +1122,572 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n client.delete("unscoped") client.close() facade.cache.redis_client.close() + + +PARAPHRASE_MARKER: Final = " (paraphrase)" +SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic" +SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_" +SEMANTIC_CONTEXT: Final = contextvars.ContextVar("semantic_test_context", default="unset") + + +def _normalized(vector: list[float]) -> list[float]: + norm: Final = math.sqrt(sum(component * component for component in vector)) + return [component / norm for component in vector] + + +def _base_embedding(prompt: str) -> list[float]: + digest: Final = hashlib.sha256(prompt.encode("utf-8")).digest() + return _normalized([float(digest[index] + 1) for index in range(8)]) + + +def _semantic_embedding(prompt: str) -> list[float]: + if PARAPHRASE_MARKER not in prompt: + return _base_embedding(prompt) + base: Final = _base_embedding(prompt.replace(PARAPHRASE_MARKER, "").strip()) + pivot: Final = min(range(8), key=lambda index: abs(base[index])) + direction: Final = _normalized( + [ + (1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot] + for index in range(8) + ] + ) + # Rotating an orthogonal unit direction by 0.329 produces ~0.05 cosine distance + return _normalized([base[index] + 0.329 * direction[index] for index in range(8)]) + + +class DeterministicEmbedding(litellm.CustomLLM): + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + self.async_calls: list[dict[str, object]] = [] + self.entered = asyncio.Event() + self.gate: asyncio.Event | None = None + + def _respond( + self, + model: str, + input: object, + model_response: EmbeddingResponse, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.calls.append({"model": model, "input": texts}) + model_response.model = model + model_response.data = [ + {"object": "embedding", "index": index, "embedding": _semantic_embedding(str(text))} + for index, text in enumerate(texts) + ] + return model_response + + def embedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + return self._respond(model, input, model_response) + + async def aembedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.async_calls.append( + { + "model": model, + "input": texts, + "task": asyncio.current_task(), + "context": SEMANTIC_CONTEXT.get(), + } + ) + SEMANTIC_CONTEXT.set("written-in-aembedding") + self.entered.set() + if self.gate is not None: + await self.gate.wait() + return self._respond(model, input, model_response) + + +@pytest.fixture +def semantic_embedding() -> Generator[DeterministicEmbedding]: + handler: Final = DeterministicEmbedding() + with ExitStack() as stack: + stack.enter_context( + rebound( + litellm, + "custom_provider_map", + [ + *litellm.custom_provider_map, + cast( + CustomLLMItem, + {"provider": "semantic-test", "custom_handler": handler}, + ), + ], + ) + ) + stack.enter_context( + rebound( + litellm, + "_custom_providers", # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + [*litellm._custom_providers, "semantic-test"], # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + ) + ) + stack.enter_context( + rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"]) + ) + yield handler + + +@pytest.fixture +def redis_stack() -> Generator[tuple[str, str]]: + url: Final = os.environ.get("LITELLM_REDIS_STACK_URL") + if url is None: + pytest.skip("LITELLM_REDIS_STACK_URL is not set") + index: Final = f"{SEMANTIC_INDEX_PREFIX}{uuid4().hex}" + yield url, index + client: Final = redis.Redis.from_url(url) + try: + client.execute_command("FT.DROPINDEX", index, "DD") # pyright: ignore[reportUnknownMemberType] # redis-py leaves execute_command partially unknown + except redis.RedisError: + pass + client.close() + + +def semantic_request(key: str, prompt: str, **extra: object) -> dict[str, object]: + return { + "key": {"preset": key}, + "messages": [{"role": "user", "content": prompt}], + **extra, + } + + +def semantic_messages(prompt: str) -> list[dict[str, object]]: + return [{"role": "user", "content": prompt}] + + +def semantic_entry_id(prompt: str, tag: str) -> str: + return hashlib.sha256(f"{prompt}litellm_cache_key{tag}".encode()).hexdigest() + + +def semantic_facade(url: str, index: str, *, similarity_threshold: float = 0.8) -> Cache: + facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=similarity_threshold, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(facade) + return facade + + +def test_redis_semantic_constructor_identity_and_provenance( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + backend: Final = cast(RedisSemanticCache, facade.cache) + assert backend.__class__.__module__ == "litellm.caching.redis_semantic_cache" + assert type(backend) is RedisSemanticCache + assert backend._redis_url == url # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend._index_name == index # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend.similarity_threshold == 0.8 + assert backend.embedding_model == SEMANTIC_EMBEDDING_MODEL + handle: Final = cast(object, getattr(facade, "_native_cache_handle")) + assert isinstance(handle, _CacheTestHandle) + assert handle.backend == "redis_semantic" + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + + +def test_redis_semantic_native_and_python_sync_entries_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + response: Final = {"choices": [{"text": "paris"}], "usage": {"total_tokens": 2}} + + binding.store(semantic_request("geo", "what is the capital of france"), response) + + native_hash_key: Final = f"{index}:{semantic_entry_id('what is the capital of france', 'geo')}" + stored: Final = client.hgetall(native_hash_key) + assert set(stored) == { + b"entry_id", + b"prompt", + b"response", + b"prompt_vector", + b"inserted_at", + b"updated_at", + b"litellm_cache_key", + }, stored + assert stored[b"entry_id"].decode() == native_hash_key.split(":", 1)[1] + assert stored[b"prompt"] == b"what is the capital of france" + assert stored[b"litellm_cache_key"] == b"geo" + assert len(stored[b"prompt_vector"]) == 32 + decoded: Final = cast(dict[str, object], json.loads(stored[b"response"])) + assert decoded["response"] == response + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "geo", messages=semantic_messages("what is the capital of france") + ) + == decoded + ) + assert semantic_embedding.calls == [ + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["dimension test"]}, + ] + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "math", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": 42}}), + messages=semantic_messages("what is 6 times 7"), + ) + python_hash_key: Final = f"{index}:{semantic_entry_id('what is 6 times 7', 'math')}" + assert json.loads(cast(bytes, client.hget(python_hash_key, "response"))) == { + "timestamp": 1700000000.0, + "response": {"answer": 42}, + } + assert binding.lookup(semantic_request("math", "what is 6 times 7")) == {"answer": 42} + client.close() + + +async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + await binding.async_store( + semantic_request("async", "name a primary color"), {"answer": "blue"} + ) + hash_key: Final = f"{index}:{semantic_entry_id('name a primary color', 'async')}" + decoded: Final = cast(dict[str, object], json.loads(cast(bytes, client.hget(hash_key, "response")))) + python_read: Final = await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async", messages=semantic_messages("name a primary color") + ) + assert python_read == decoded + + await binding.async_store_batch( + [ + semantic_request("batch-one", "first batch prompt"), + semantic_request("batch-two", "second batch prompt"), + ], + [{"answer": 1}, {"answer": 2}], + ) + expected: Final = { + key: json.loads( + cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response")) + ) + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ) + } + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ): + assert cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + key, messages=semantic_messages(prompt) + ) == expected[key], key + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async-python", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": "python"}}), + messages=semantic_messages("python written prompt"), + ) + assert await binding.async_lookup( + semantic_request("async-python", "python written prompt") + ) == {"answer": "python"} + client.close() + + +async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + caller: Final = asyncio.current_task() + SEMANTIC_CONTEXT.set("caller-sentinel") + response: Final = {"choices": [{"text": "paris"}]} + + await binding.async_store( + semantic_request("inline", "what is the capital of france"), response + ) + assert ( + await binding.async_lookup( + semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}") + ) + == response + ) + assert await binding.async_lookup(semantic_request("inline", "python written prompt")) is None + assert SEMANTIC_CONTEXT.get() == "written-in-aembedding" + assert semantic_embedding.async_calls == [ + { + "model": "deterministic", + "input": ["what is the capital of france"], + "task": caller, + "context": "caller-sentinel", + }, + { + "model": "deterministic", + "input": [f"what is the capital of france{PARAPHRASE_MARKER}"], + "task": caller, + "context": "written-in-aembedding", + }, + { + "model": "deterministic", + "input": ["python written prompt"], + "task": caller, + "context": "written-in-aembedding", + }, + ], semantic_embedding.async_calls + + +async def test_native_semantic_cancellation_during_embedding_skips_the_backend( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + semantic_embedding.gate = asyncio.Event() + + async def lookup() -> object: + return await binding.async_lookup( + semantic_request("cancel", "cancelled prompt") + ) + + task: Final = asyncio.create_task(lookup()) + await semantic_embedding.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + semantic_embedding.gate.set() + + assert len(semantic_embedding.async_calls) == 1 + assert ( + await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "cancel", messages=semantic_messages("cancelled prompt") + ) + is None + ) + + +def test_redis_semantic_similarity_tag_and_threshold_boundaries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + binding.store(semantic_request("sim", "tell me a joke"), {"answer": "haha"}) + paraphrase: Final = f"tell me a joke{PARAPHRASE_MARKER}" + assert binding.lookup(semantic_request("sim", paraphrase)) == {"answer": "haha"} + assert binding.lookup(semantic_request("sim", "an unrelated question about spreadsheets")) is None + assert binding.lookup(semantic_request("other-key", "tell me a joke")) is None + + strict: Final = semantic_facade(url, index, similarity_threshold=0.99) + strict_binding: Final = _CacheTestResolver(SimpleNamespace(cache=strict)).resolve() + assert strict_binding.lookup(semantic_request("sim", paraphrase)) is None + assert strict_binding.lookup(semantic_request("sim", "tell me a joke")) == {"answer": "haha"} + + +def test_redis_semantic_ttl_is_written_only_when_requested( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store( + {**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1} + ) + expiring: Final = f"{index}:{semantic_entry_id('ttl prompt', 'ttl')}" + assert 0 < client.ttl(expiring) <= 12 + + binding.store(semantic_request("ttl-none", "untimed prompt"), {"answer": 2}) + persistent: Final = f"{index}:{semantic_entry_id('untimed prompt', 'ttl-none')}" + assert client.ttl(persistent) == -1 + + binding.store( + {**semantic_request("ttl-fraction", "fractional prompt"), "ttl_seconds": 1.5}, + {"answer": 3}, + ) + fractional: Final = f"{index}:{semantic_entry_id('fractional prompt', 'ttl-fraction')}" + assert client.ttl(fractional) == 2 + client.close() + + +def test_redis_semantic_malformed_response_is_a_miss_for_both_readers( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(semantic_request("bad", "corrupt me"), {"answer": 1}) + hash_key: Final = f"{index}:{semantic_entry_id('corrupt me', 'bad')}" + client.hset(hash_key, "response", b"{not json") + assert binding.lookup(semantic_request("bad", "corrupt me")) is None + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "bad", messages=semantic_messages("corrupt me") + ) + is None + ) + client.close() + + +async def test_redis_semantic_unsupported_operations_raise_not_implemented( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + with pytest.raises(NotImplementedError): + binding.lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_flush() + with pytest.raises(NotImplementedError): + await binding.ping() + + +def test_redis_semantic_requests_without_prompt_are_noops( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(request("plain"), {"answer": 1}) + assert binding.lookup(request("plain")) is None + assert semantic_embedding.calls == [] + assert client.keys(f"{index}:*") == [] + client.close() + + +def test_redis_semantic_scope_overrides_the_tag_and_isolates_entries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + scoped: Final = {**semantic_request("scoped", "scoped prompt"), "scope": "team-a"} + binding.store(scoped, {"answer": "kept"}) + hash_key: Final = f"{index}:{semantic_entry_id('scoped prompt', 'team-a')}" + assert client.hget(hash_key, "litellm_cache_key") == b"team-a" + assert binding.lookup(scoped) == {"answer": "kept"} + assert binding.lookup(semantic_request("scoped", "scoped prompt")) is None + assert binding.lookup({**scoped, "scope": "team-b"}) is None + client.close() + + +def test_redis_semantic_configuration_drift_falls_back_to_python( + redis_stack: tuple[str, str], + semantic_embedding: DeterministicEmbedding, + monkeypatch: pytest.MonkeyPatch, +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + + with rebound(facade.cache, "similarity_threshold", 0.5): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "semantic_cache_scope", "end_user"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "embedding_model", "other-model"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "_index_name", "other-index"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "CACHE_KEY_FIELD_NAME", "other-field"): + assert resolver.resolve().kind == "python_callback" + + def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]: + return _semantic_embedding(prompt) + + monkeypatch.setattr(RedisSemanticCache, "_get_embedding", patched_embedding) + assert resolver.resolve().kind == "python_callback" + + +def test_redis_semantic_handle_rejects_wrong_backends( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + + class CustomSemanticCache(RedisSemanticCache): + pass + + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + _CacheTestHandle.redis_semantic(object()) + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + _CacheTestHandle.redis_semantic( + CustomSemanticCache( + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=f"{index}_subclass", + ) + ) + + facade: Final = semantic_facade(url, index) + with pytest.raises(TypeError, match="backend types must match"): + _CacheTestHandle.redis(url)._bind_facade(facade) + + subclassed_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + subclassed_facade.cache = CustomSemanticCache( # pyright: ignore[reportAttributeAccessIssue] # facade backend slot is not declared + + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=index, + ) + with pytest.raises(TypeError): + _CacheTestHandle.redis_semantic( + subclassed_facade.cache + )._bind_facade(subclassed_facade) + + replacement_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + with pytest.raises(TypeError, match="must be the native embedder"): + _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade)