From 6ab121a3e7364178544bcbee529f6c57b3115a0a Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:29:49 +0000 Subject: [PATCH] feat(cache-redis-semantic): add native Redis Semantic cache backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 15 + litellm-rust/Cargo.toml | 1 + .../crates/cache-redis-semantic/Cargo.toml | 21 + .../crates/cache-redis-semantic/src/cache.rs | 599 ++++++++++++++ .../crates/cache-redis-semantic/src/lib.rs | 4 + .../crates/cache-redis-semantic/src/prompt.rs | 95 +++ .../cache-redis-semantic/tests/cache.rs | 751 ++++++++++++++++++ litellm-rust/crates/cache/tests/caching.rs | 9 +- 8 files changed, 1489 insertions(+), 6 deletions(-) create mode 100644 litellm-rust/crates/cache-redis-semantic/Cargo.toml create mode 100644 litellm-rust/crates/cache-redis-semantic/src/cache.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/src/lib.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/src/prompt.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/tests/cache.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ed4ae4e3353..e911d0d9c45 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2486,6 +2486,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "litellm-cache-redis", + "litellm-cache-response", + "r2d2", + "redis", + "redis-test", + "serde_json", + "sha2 0.10.9", + "tokio", +] + [[package]] name = "litellm-cache-response" version = "0.1.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 570d0dd3568..05eea6bc299 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -29,6 +29,7 @@ litellm-core-utils = { path = "crates/core-utils" } litellm-cache = { path = "crates/cache" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-cache-redis = { path = "crates/cache-redis" } +litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" } litellm-cache-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } diff --git a/litellm-rust/crates/cache-redis-semantic/Cargo.toml b/litellm-rust/crates/cache-redis-semantic/Cargo.toml new file mode 100644 index 00000000000..9a8755a189e --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-cache-redis.workspace = true +litellm-cache-response.workspace = true +redis = { version = "1.7.0", features = ["tls-rustls"] } +r2d2 = "0.8.10" +serde_json.workspace = true +sha2.workspace = true +tokio.workspace = true + +[dev-dependencies] +redis-test = "1.0.4" +serde_json.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs new file mode 100644 index 00000000000..26d19d34670 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -0,0 +1,599 @@ +use std::{ + future::Future, + sync::{Arc, OnceLock}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_cache::{ + BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error, + SemanticCacheContext, +}; +use litellm_cache_redis::connection::{ConnectionRef, Connections, ttl_seconds}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +use crate::prompt::prompt_from_context; + +const REDIS_TIMEOUT: Duration = Duration::from_secs(5); +const REDIS_POOL_SIZE: u32 = 16; +const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; +const CACHE_KEY_FIELD: &str = "litellm_cache_key"; +const VECTOR_FIELD: &str = "prompt_vector"; + +pub trait Embedder: Send + Sync + 'static { + fn embed( + &self, + prompt: &str, + metadata: &serde_json::Map, + ) -> Result, Error>; + + fn async_embed( + &self, + prompt: &str, + metadata: &serde_json::Map, + ) -> impl Future, Error>> + Send; +} + +#[derive(Clone, Debug)] +pub struct RedisSemanticConfig { + pub index_name: String, + pub similarity_threshold: f32, +} + +impl Default for RedisSemanticConfig { + fn default() -> Self { + Self { + index_name: DEFAULT_INDEX_NAME.into(), + similarity_threshold: 0.9, + } + } +} + +struct Inner { + index_name: String, + distance_threshold: f64, + resolved_index: OnceLock, + codec: ResponseCacheCodec, + clock: fn() -> f64, +} + +impl Inner { + fn new(config: RedisSemanticConfig) -> Self { + Self { + index_name: config.index_name, + distance_threshold: 1.0 - f64::from(config.similarity_threshold), + resolved_index: OnceLock::new(), + codec: ResponseCacheCodec, + clock: timestamp, + } + } + + fn ensure_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + if let Some(name) = self.resolved_index.get() { + return Ok(name.clone()); + } + let name = match index_compatible(connection, &self.index_name, dims)? { + Some(true) => self.index_name.clone(), + Some(false) => self.isolated_index(connection, dims)?, + None => { + create_index(connection, &self.index_name, dims)?; + self.index_name.clone() + } + }; + let _ = self.resolved_index.set(name.clone()); + Ok(name) + } + + fn isolated_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + let name = format!("{}_isolated", self.index_name); + match index_compatible(connection, &name, dims)? { + Some(true) => Ok(name), + Some(false) => { + redis::cmd("FT.DROPINDEX") + .arg(&name) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + create_index(connection, &name, dims)?; + Ok(name) + } + None => { + create_index(connection, &name, dims)?; + Ok(name) + } + } + } + + fn store( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + value: &CacheEntry, + prompt: &str, + vector: &[f32], + ttl: Option, + ) -> Result<(), Error> { + let index = self.ensure_index(connection, vector.len())?; + let entry_id = entry_id(prompt, tag); + let hash_key = format!("{index}:{entry_id}"); + let response = self.codec.encode(value)?; + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(&entry_id) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg(VECTOR_FIELD) + .arg(vector_buffer(vector)) + .arg("inserted_at") + .arg(format!("{}", (self.clock)())) + .arg("updated_at") + .arg(format!("{}", (self.clock)())) + .arg(CACHE_KEY_FIELD) + .arg(tag) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + if let Some(ttl) = ttl { + redis::cmd("EXPIRE") + .arg(&hash_key) + .arg(ttl_seconds(ttl)) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + } + Ok(()) + } + + fn lookup( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + vector: &[f32], + ) -> Result, Error> { + let index = self.ensure_index(connection, vector.len())?; + let query = format!( + "(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]", + escape_tag(tag) + ); + let result = redis::cmd("FT.SEARCH") + .arg(&index) + .arg(query) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg(CACHE_KEY_FIELD) + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_buffer(vector)) + .query::(connection) + .map_err(|_| Error::Unavailable)?; + let Some(fields) = first_document(&result) else { + return Ok(None); + }; + if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) { + return Ok(None); + } + if number_field(fields, "vector_distance") + .is_none_or(|distance| distance > self.distance_threshold) + { + return Ok(None); + } + let Some(response) = bytes_field(fields, "response") else { + return Ok(None); + }; + self.codec.decode(&response).map(Some) + } +} + +pub struct RedisSemanticCache { + connections: Arc>, + embedder: E, + inner: Arc, +} + +impl RedisSemanticCache { + pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result { + Ok(Self { + connections: Arc::new(Connections::pooled(url, REDIS_TIMEOUT, REDIS_POOL_SIZE)?), + embedder, + inner: Arc::new(Inner::new(config)), + }) + } +} + +impl RedisSemanticCache { + pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self { + Self { + connections: Arc::new(Connections::fixed(connection)), + embedder, + inner: Arc::new(Inner::new(config)), + } + } + + pub fn with_clock(self, clock: fn() -> f64) -> Self { + Self { + inner: Arc::new(Inner { + index_name: self.inner.index_name.clone(), + distance_threshold: self.inner.distance_threshold, + resolved_index: OnceLock::new(), + codec: self.inner.codec, + clock, + }), + ..self + } + } + + fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str { + context.scope.as_deref().unwrap_or(key) + } +} + +impl BaseCache + for RedisSemanticCache +{ + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(()); + }; + let vector = self.embedder.embed(&prompt, &context.metadata)?; + let tag = Self::tag(key, context).to_string(); + self.connections.execute(|connection| { + self.inner + .store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self.embedder.embed(&prompt, &context.metadata)?; + let tag = Self::tag(key, context).to_string(); + self.connections + .execute(|connection| self.inner.lookup(connection, &tag, &vector)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(&context) else { + return Ok(()); + }; + let vector = self + .embedder + .async_embed(&prompt, &context.metadata) + .await?; + let tag = Self::tag(key, &context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + .await + } + + async fn async_get_cache( + &self, + key: &str, + context: &Self::Context, + ) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self + .embedder + .async_embed(&prompt, &context.metadata) + .await?; + let tag = Self::tag(key, context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.lookup(connection, &tag, &vector) + }) + .await + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + match Connections::run_blocking(Arc::clone(&self.connections), |connection| { + Ok(match redis::cmd("PING").query::(connection) { + Ok(_) => CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, + }, + Err(error) => CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }, + }) + }) + .await + { + Ok(result) => Ok(result), + Err(error) => Ok(CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }), + } + } +} + +fn timestamp() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or_default() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(CACHE_KEY_FIELD.as_bytes()); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn vector_buffer(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn escape_tag(value: &str) -> String { + value + .chars() + .flat_map(|ch| { + if matches!( + ch, + ',' | '.' + | '<' + | '>' + | '{' + | '}' + | '[' + | ']' + | '\\' + | '"' + | '\'' + | ':' + | ';' + | '!' + | '@' + | '#' + | '$' + | '%' + | '^' + | '&' + | '*' + | '(' + | ')' + | '-' + | '+' + | '=' + | '~' + | '|' + | '/' + | ' ' + | '?' + ) { + vec!['\\', ch] + } else { + vec![ch] + } + }) + .collect() +} + +fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> { + redis::cmd("FT.CREATE") + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg(VECTOR_FIELD) + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg(CACHE_KEY_FIELD) + .arg("TAG") + .arg("SEPARATOR") + .arg(",") + .query::<()>(connection) + .map_err(|_| Error::Unavailable) +} + +fn index_compatible( + connection: &mut ConnectionRef<'_>, + name: &str, + dims: usize, +) -> Result, Error> { + let info = match redis::cmd("FT.INFO") + .arg(name) + .query::(connection) + { + Ok(info) => info, + Err(error) if unknown_index(&error) => return Ok(None), + Err(_) => return Err(Error::Unavailable), + }; + Ok(Some(schema_compatible(&info, dims))) +} + +fn unknown_index(error: &redis::RedisError) -> bool { + let message = error.to_string().to_lowercase(); + message.contains("unknown") && message.contains("index") +} + +fn schema_compatible(info: &redis::Value, dims: usize) -> bool { + let redis::Value::Array(entries) = info else { + return false; + }; + let attributes = entries + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some("attributes")) + .map(|pair| &pair[1]); + let Some(redis::Value::Array(attributes)) = attributes else { + return false; + }; + let fields = attributes + .iter() + .map(|attribute| { + let redis::Value::Array(attribute) = attribute else { + return (None, None, None); + }; + let mut name = None; + let mut field_type = None; + let mut dim = None; + for pair in attribute.as_chunks::<2>().0 { + match string_value(&pair[0]).as_deref() { + Some("identifier") => name = string_value(&pair[1]), + Some("type") => field_type = string_value(&pair[1]), + Some("dim") => dim = number_value(&pair[1]), + _ => {} + } + } + (name, field_type, dim) + }) + .collect::>(); + let has_field = |name: &str, field_type: &str| { + fields + .iter() + .any(|(n, t, _)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type)) + }; + has_field("prompt", "TEXT") + && has_field("response", "TEXT") + && has_field("inserted_at", "NUMERIC") + && has_field("updated_at", "NUMERIC") + && has_field(CACHE_KEY_FIELD, "TAG") + && fields.iter().any(|(n, t, d)| { + n.as_deref() == Some(VECTOR_FIELD) + && t.as_deref() == Some("VECTOR") + && *d == Some(dims as f64) + }) +} + +fn string_value(value: &redis::Value) -> Option { + match value { + redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), + redis::Value::SimpleString(text) => Some(text.clone()), + redis::Value::VerbatimString { text, .. } => Some(text.clone()), + _ => None, + } +} + +fn number_value(value: &redis::Value) -> Option { + match value { + redis::Value::Int(number) => Some(*number as f64), + redis::Value::Double(number) => Some(*number), + _ => string_value(value).and_then(|text| text.parse().ok()), + } +} + +fn first_document(result: &redis::Value) -> Option<&[redis::Value]> { + let redis::Value::Array(items) = result else { + return None; + }; + let [count, _document_id, fields, ..] = items.as_slice() else { + return None; + }; + if !matches!(count, redis::Value::Int(count) if *count > 0) { + return None; + } + match fields { + redis::Value::Array(fields) => Some(fields.as_slice()), + _ => None, + } +} + +fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> { + fields + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some(name)) + .map(|pair| &pair[1]) +} + +fn string_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(string_value) +} + +fn number_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(number_value) +} + +fn bytes_field(fields: &[redis::Value], name: &str) -> Option> { + match field_value(fields, name)? { + redis::Value::BulkString(bytes) => Some(bytes.clone()), + redis::Value::SimpleString(text) => Some(text.clone().into_bytes()), + _ => None, + } +} diff --git a/litellm-rust/crates/cache-redis-semantic/src/lib.rs b/litellm-rust/crates/cache-redis-semantic/src/lib.rs new file mode 100644 index 00000000000..a34603cd18f --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/lib.rs @@ -0,0 +1,4 @@ +mod cache; +mod prompt; + +pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig}; diff --git a/litellm-rust/crates/cache-redis-semantic/src/prompt.rs b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs new file mode 100644 index 00000000000..fc99898d847 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs @@ -0,0 +1,95 @@ +use litellm_cache::SemanticCacheContext; +use serde_json::Value; + +pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { + if !context.messages.is_empty() { + return Some(messages_text(&context.messages)); + } + let input = context.input.as_ref()?; + let mut parts = Vec::new(); + collect_input_text(input, &mut parts); + let prompt = parts.join("\n").trim().to_string(); + (!prompt.is_empty()).then_some(prompt) +} + +fn messages_text(messages: &[Value]) -> String { + let mut text = String::new(); + for message in messages { + let Some(message) = message.as_object() else { + continue; + }; + match message.get("content") { + Some(Value::String(content)) => text.push_str(content), + Some(Value::Array(parts)) => { + for part in parts { + if let Some(text_content) = part.get("text").and_then(Value::as_str) { + text.push_str(text_content); + } + } + } + _ => {} + } + text.push_str(&search_results_text(message.get("search_results"))); + } + text +} + +fn search_results_text(search_results: Option<&Value>) -> String { + let Some(Value::Array(results)) = search_results else { + return String::new(); + }; + let mut text = String::new(); + for result in results { + let Some(result) = result.as_object() else { + continue; + }; + for key in ["source", "title"] { + if let Some(value) = result.get(key).and_then(Value::as_str) { + text.push_str(value); + } + } + if let Some(Value::Array(content)) = result.get("content") { + for block in content { + if let Some(value) = block.get("text").and_then(Value::as_str) { + text.push_str(value); + } + } + } + if let Some(citations) = result.get("citations") { + text.push_str(&citations.to_string()); + } + } + text +} + +fn collect_input_text(value: &Value, parts: &mut Vec) { + match value { + Value::String(text) => { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + } + } + Value::Array(items) => { + for item in items { + collect_input_text(item, parts); + } + } + Value::Object(map) => { + if let Some(content) = map.get("content").filter(|content| !content.is_null()) { + collect_input_text(content, parts); + return; + } + for key in ["text", "output", "input_text", "output_text"] { + if let Some(Value::String(text)) = map.get(key) { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + return; + } + } + } + } + _ => {} + } +} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs new file mode 100644 index 00000000000..77b057ae3b9 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -0,0 +1,751 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{BaseCache, CacheCodec, Error, SemanticCacheContext}; +use litellm_cache_redis_semantic::{Embedder, RedisSemanticCache, RedisSemanticConfig}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use redis_test::{MockCmd, MockRedisConnection}; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; + +const INDEX: &str = "litellm_semantic_cache_index"; + +struct FakeEmbedder { + vectors: HashMap>, + calls: Arc>>, +} + +impl FakeEmbedder { + fn new(vectors: &[(&str, &[f32])]) -> (Self, Arc>>) { + let calls = Arc::new(Mutex::new(Vec::new())); + ( + Self { + vectors: vectors + .iter() + .map(|(prompt, vector)| (prompt.to_string(), vector.to_vec())) + .collect(), + calls: Arc::clone(&calls), + }, + calls, + ) + } +} + +impl Embedder for FakeEmbedder { + fn embed(&self, prompt: &str, _: &serde_json::Map) -> Result, Error> { + self.calls.lock().unwrap().push(prompt.to_string()); + + Ok(self + .vectors + .get(prompt) + .cloned() + .unwrap_or_else(|| vec![0.1, 0.2, 0.3])) + } + + async fn async_embed( + &self, + prompt: &str, + metadata: &serde_json::Map, + ) -> Result, Error> { + self.embed(prompt, metadata) + } +} + +fn config() -> RedisSemanticConfig { + RedisSemanticConfig { + index_name: INDEX.into(), + similarity_threshold: 0.9, + } +} + +fn messages_context(messages: Vec) -> SemanticCacheContext { + SemanticCacheContext { + messages, + ..Default::default() + } +} + +fn entry() -> CacheEntry { + CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "yes"}), + } +} + +fn encoded(entry: &CacheEntry) -> Vec { + ResponseCacheCodec.encode(entry).unwrap() +} + +fn vector_bytes(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(b"litellm_cache_key"); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn s(value: &str) -> redis::Value { + redis::Value::BulkString(value.as_bytes().to_vec()) +} + +fn unknown_index_error() -> redis::RedisError { + redis::RedisError::from((redis::ErrorKind::Extension, "Unknown index name")) +} + +fn attribute(name: &str, field_type: &str, extra: Vec) -> redis::Value { + let mut parts = vec![ + s("identifier"), + s(name), + s("attribute"), + s(name), + s("type"), + s(field_type), + ]; + parts.extend(extra); + redis::Value::Array(parts) +} + +fn index_info(attributes: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("index_name"), + s(INDEX), + s("attributes"), + redis::Value::Array(attributes), + ]) +} + +fn vector_attribute(dims: i64) -> redis::Value { + attribute( + "prompt_vector", + "VECTOR", + vec![ + s("algorithm"), + s("FLAT"), + s("data_type"), + s("FLOAT32"), + s("dim"), + redis::Value::Int(dims), + s("distance_metric"), + s("COSINE"), + ], + ) +} + +fn compatible_info(dims: i64) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector_attribute(dims), + attribute("litellm_cache_key", "TAG", vec![]), + ]) +} + +fn unscoped_info(dims: i64) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector_attribute(dims), + ]) +} + +fn create_index_command(name: &str, dims: usize) -> redis::Cmd { + let mut command = redis::cmd("FT.CREATE"); + command + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg("prompt_vector") + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg("litellm_cache_key") + .arg("TAG") + .arg("SEPARATOR") + .arg(","); + command +} + +fn search_command(index: &str, tag: &str, vector: &[f32]) -> redis::Cmd { + let mut command = redis::cmd("FT.SEARCH"); + command + .arg(index) + .arg(format!( + "(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]" + )) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg("litellm_cache_key") + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_bytes(vector)); + command +} + +fn hit_fields(tag: &str, distance: &str, response: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("entry_id"), + s("stored-id"), + s("prompt"), + s("hello prompt"), + s("response"), + redis::Value::BulkString(response), + s("inserted_at"), + s("1700000000.5"), + s("updated_at"), + s("1700000000.5"), + s("litellm_cache_key"), + s(tag), + s("vector_distance"), + s(distance), + ]) +} + +fn search_result(fields: redis::Value) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::Int(1), + s("litellm_semantic_cache_index:stored-id"), + fields, + ]) +} + +fn empty_result() -> redis::Value { + redis::Value::Array(vec![redis::Value::Int(0)]) +} + +#[test] +fn store_creates_index_and_writes_hash_with_expire() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(INDEX, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(5), Ok(1)), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(5)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + cache.set_cache(tag, value, &context).unwrap(); +} + +#[test] +fn store_without_ttl_skips_expire() { + let prompt = "hello prompt"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "key1"))) + .arg("entry_id") + .arg(entry_id(prompt, "key1")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("key1"), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + "key1", + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn lookup_returns_hit_below_distance_threshold() { + let vector = vec![0.1f32, 0.2, 0.3]; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let hit = cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]), + ) + .unwrap(); + assert_eq!(hit, Some(value)); +} + +#[test] +fn lookup_misses_above_distance_threshold_and_on_tag_mismatch() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.5", encoded(&entry())))), + ), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "other", + "0.05", + encoded(&entry()), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + let context = messages_context(vec![json!({"role": "user", "content": "hello prompt"})]); + + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); +} + +#[test] +fn lookup_returns_invalid_entry_on_malformed_response() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "key1", + "0.05", + b"not json!".to_vec(), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[test] +fn missing_prompt_is_noop_and_never_embeds() { + let connection = MockRedisConnection::new(Vec::::new()).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let context = SemanticCacheContext::default(); + cache.set_cache("key1", entry(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert!(calls.lock().unwrap().is_empty()); +} + +#[test] +fn scope_overrides_key_as_filter_tag() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "scope-a"))) + .arg("entry_id") + .arg(entry_id(prompt, "scope-a")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("scope-a"), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, "scope\\-a", &vector), + Ok(search_result(hit_fields( + "scope-a", + "0.05", + encoded(&value), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = SemanticCacheContext { + scope: Some("scope-a".into()), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + + cache.set_cache("key1", value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(value)); +} + +#[test] +fn incompatible_schema_falls_back_to_isolated_index() { + let prompt = "hello prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(unscoped_info(3))), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{isolated}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn tag_special_characters_are_escaped_in_search_filter() { + let vector = vec![0.1f32, 0.2, 0.3]; + let tag = "a:b, c|d"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "a\\:b\\,\\ c\\|d", &vector), + Ok(empty_result()), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + tag, + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap(), + None + ); +} + +#[test] +fn prompt_extraction_matches_python_message_and_input_shapes() { + let vector = vec![0.1f32, 0.2, 0.3]; + let lookups = 5; + let mut commands = vec![MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Ok(compatible_info(3)), + )]; + for _ in 0..lookups { + commands.push(MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(empty_result()), + )); + } + let connection = MockRedisConnection::new(commands).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + cache + .get_cache( + "key1", + &messages_context(vec![ + json!({"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]}), + json!({"role": "assistant", "content": "reply"}), + ]), + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!(" plain input ")), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some( + json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"]), + ), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!({"output_text": " result text "})), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &messages_context(vec![json!({ + "role": "user", + "content": "question", + "search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}], + })]), + ) + .unwrap(); + + assert_eq!( + *calls.lock().unwrap(), + vec![ + "firstsecondreply", + "plain input", + "nested\ntail", + "result text", + "questionsrctfound{\"a\":1}", + ] + ); +} + +#[test] +fn ttl_passes_through_context_only() { + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection( + MockRedisConnection::new(Vec::::new()), + embedder, + config(), + ); + assert_eq!(cache.get_ttl(&SemanticCacheContext::default()), None); + assert_eq!( + cache.get_ttl(&SemanticCacheContext { + ttl: Some(Duration::from_secs(9)), + ..Default::default() + }), + Some(Duration::from_secs(9)) + ); +} + +#[tokio::test] +async fn async_paths_embed_then_run_blocking_redis_work() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, tag, &vector), + Ok(search_result(hit_fields(tag, "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = messages_context(vec![json!({"role": "user", "content": prompt})]); + + cache + .async_set_cache(tag, value.clone(), context.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache(tag, &context).await.unwrap(), + Some(value) + ); +} + +#[test] +fn live_store_lookup_and_ttl_against_redis_stack() { + let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { + return; + }; + let vector = vec![0.1f32, 0.2, 0.3, 0.4]; + let prompt = "rust semantic cache live prompt"; + let tag = "live-key"; + let index_name = format!("rust_semantic_test_{}", std::process::id()); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::new( + &url, + embedder, + RedisSemanticConfig { + index_name: index_name.clone(), + similarity_threshold: 0.9, + }, + ) + .unwrap(); + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(120)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + let value = entry(); + + cache.set_cache(tag, value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache(tag, &context).unwrap(), Some(value)); + assert_eq!(cache.get_cache("other-key", &context).unwrap(), None); + + let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); + let ttl: i64 = redis::Commands::ttl( + &mut connection, + format!("{index_name}:{}", entry_id(prompt, tag)), + ) + .unwrap(); + assert!( + ttl > 0, + "expected stored hash to carry an expiry, got {ttl}" + ); +} diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 2e65b4eeae5..33171f36f46 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,8 +1,8 @@ use std::{sync::Mutex, time::Duration}; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, - SemanticCacheContext, get_cache, + BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, SemanticCacheContext, + get_cache, }; struct TestCache { @@ -132,10 +132,7 @@ fn semantic_context_with_ttl_preserves_lookup_inputs() { let context = SemanticCacheContext { input: Some(serde_json::json!("text")), messages: vec![serde_json::json!({"role": "user", "content": "hi"})], - metadata: serde_json::Map::from_iter([( - "key".into(), - serde_json::json!("value"), - )]), + metadata: serde_json::Map::from_iter([("key".into(), serde_json::json!("value"))]), scope: Some("scope".into()), ttl: None, };