diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f91deaaa401..b93a8827428 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3569,37 +3569,6 @@ dependencies = [ "uuid", ] -[[package]] -name = "litellm-cache-redis" -version = "0.1.0" -dependencies = [ - "litellm-cache", - "litellm-cache-testing", - "litellm-error", - "r2d2", - "redis", - "redis-test", - "rstest", - "serde_json", - "thiserror 2.0.19", - "tokio", -] - -[[package]] -name = "litellm-cache-redis-semantic" -version = "0.1.0" -dependencies = [ - "litellm-cache", - "litellm-cache-redis", - "litellm-cache-testing", - "redis", - "redis-test", - "rstest", - "serde_json", - "sha2 0.10.9", - "tokio", -] - [[package]] name = "litellm-cache-response" version = "0.1.0" @@ -3607,11 +3576,8 @@ dependencies = [ "litellm-cache", "litellm-cache-gcs", "litellm-cache-memory", - "litellm-cache-redis", "litellm-http", "py_literal", - "redis", - "redis-test", "rstest", "serde", "serde_json", @@ -3650,24 +3616,6 @@ dependencies = [ "litellm-cache", ] -[[package]] -name = "litellm-cache-valkey-semantic" -version = "0.1.0" -dependencies = [ - "litellm-cache", - "litellm-cache-redis", - "litellm-cache-testing", - "litellm-error", - "redis", - "redis-test", - "rstest", - "serde_json", - "sha2 0.10.9", - "thiserror 2.0.19", - "tokio", - "uuid", -] - [[package]] name = "litellm-callbacks-legacy-python" version = "0.1.0" @@ -4339,8 +4287,6 @@ dependencies = [ "litellm-auth", "litellm-auth-aws", "litellm-cache", - "litellm-cache-memory", - "litellm-cache-redis", "litellm-cache-response", "litellm-callbacks-legacy-python", "litellm-core-utils", @@ -4369,7 +4315,6 @@ dependencies = [ "prost", "pyo3", "pyo3-async-runtimes", - "redis", "reqwest", "rstest", "serde", @@ -5782,17 +5727,6 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" -[[package]] -name = "r2d2" -version = "0.8.10" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93" -dependencies = [ - "log", - "parking_lot", - "scheduled-thread-pool", -] - [[package]] name = "rand" version = "0.8.7" @@ -5966,18 +5900,6 @@ dependencies = [ "xxhash-rust", ] -[[package]] -name = "redis-test" -version = "1.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca" -dependencies = [ - "rand 0.9.5", - "redis", - "socket2 0.6.5", - "tempfile", -] - [[package]] name = "redox_syscall" version = "0.5.18" @@ -6454,15 +6376,6 @@ dependencies = [ "windows-sys 0.61.2", ] -[[package]] -name = "scheduled-thread-pool" -version = "0.2.7" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19" -dependencies = [ - "parking_lot", -] - [[package]] name = "schemars" version = "0.9.0" diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 75cac2a5d63..943e8576ae6 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -57,11 +57,9 @@ litellm-db-testing = { path = "crates/db-testing" } litellm-cache = { path = "crates/cache" } litellm-cache-azure-blob = { path = "crates/cache-azure-blob" } litellm-cache-memory = { path = "crates/cache-memory" } -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-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" } litellm-cache-testing = { path = "crates/cache-testing" } @@ -97,7 +95,6 @@ macro_rules_attribute = "0.2.3" schemars = "1" reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] } redis = "1.8.0" -redis-test = "1.0.4" testcontainers-modules = "0.15.0" qdrant-client = { version = "1.19.0", default-features = false } uuid = { version = "1", features = ["v4"] } diff --git a/litellm-rust/crates/cache-redis-semantic/Cargo.toml b/litellm-rust/crates/cache-redis-semantic/Cargo.toml deleted file mode 100644 index 6a5c794bde7..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/Cargo.toml +++ /dev/null @@ -1,19 +0,0 @@ -[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 -redis = { workspace = true, features = ["tls-rustls"] } -sha2.workspace = true - -[dev-dependencies] -litellm-cache-testing.workspace = true -redis-test.workspace = true -rstest.workspace = true -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 deleted file mode 100644 index 42fc6ae585f..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ /dev/null @@ -1,404 +0,0 @@ -use std::{ - sync::Arc, - time::{Duration, SystemTime, UNIX_EPOCH}, -}; - -use litellm_cache::{ - BaseCache, CacheCodec, Error, SemanticCacheContext, - semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context}, -}; -use litellm_cache_redis::{ - RedisTopology, - connection::{ConnectionRef, Connections}, -}; -use sha2::{Digest, Sha256}; - -use crate::{ - RedisSemanticConfig, - index::{CACHE_KEY_FIELD, Index, VECTOR_FIELD}, - reply::{bytes_field, first_document, number_field, string_field}, -}; - -struct Inner { - index: Index, - distance_threshold: f64, - clock: fn() -> f64, -} - -impl Inner { - fn new(config: RedisSemanticConfig, clock: fn() -> f64) -> Self { - Self { - index: Index::new(config.index_name), - distance_threshold: 1.0 - f64::from(config.similarity_threshold), - clock, - } - } - - fn store( - &self, - connection: &mut ConnectionRef<'_>, - tag: &str, - response: Vec, - prompt: &str, - vector: &[f32], - ttl: Option, - ) -> Result<(), Error> { - let index = self.index.ensure(connection, vector.len())?; - let entry_id = entry_id(prompt, tag); - let hash_key = format!("{index}:{entry_id}"); - 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.index.ensure(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(SemanticLookup::miss(Some(0.0))); - }; - if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) { - return Ok(SemanticLookup::miss(Some(0.0))); - } - // redisvl's range query only returns entries within the distance threshold, so a - // farther hit reads as no result. - let Some(distance) = number_field(fields, "vector_distance") - .filter(|distance| *distance <= self.distance_threshold) - else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let Some(response) = bytes_field(fields, "response") else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - Ok(SemanticLookup { - value: Some(response), - similarity: Some(1.0 - distance), - }) - } -} - -/// `RedisSemanticCache`: a redisvl-compatible semantic index on Redis Stack. Values go through -/// the injected codec, so the response layer decides what a cached entry is. -pub struct RedisSemanticCache { - connections: Arc>, - embedder: E, - codec: S, - inner: Arc, -} - -impl RedisSemanticCache { - pub fn new( - url: &str, - embedder: E, - codec: S, - config: RedisSemanticConfig, - ) -> Result { - Ok(Self { - connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), - embedder, - codec, - inner: Arc::new(Inner::new(config, timestamp)), - }) - } -} - -impl RedisSemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - pub fn with_connection( - connection: C, - embedder: E, - codec: S, - config: RedisSemanticConfig, - ) -> Self { - Self { - connections: Arc::new(Connections::fixed(connection)), - embedder, - codec, - inner: Arc::new(Inner::new(config, timestamp)), - } - } - - pub fn with_clock(self, clock: fn() -> f64) -> Self { - let config = RedisSemanticConfig { - index_name: self.index_name().to_owned(), - similarity_threshold: self.similarity_threshold(), - }; - Self { - inner: Arc::new(Inner::new(config, 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) - } - - fn decode(&self, lookup: SemanticLookup>) -> Result, Error> { - Ok(SemanticLookup { - value: lookup - .value - .map(|bytes| self.codec.decode(&bytes)) - .transpose()?, - similarity: lookup.similarity, - }) - } -} - -impl BaseCache for RedisSemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type Value = S::Value; - 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 response = self.codec.encode(&value)?; - let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; - let tag = Self::tag(key, context); - self.connections.execute(|connection| { - self.inner - .store(connection, tag, response, &prompt, &vector, context.ttl) - }) - } - - fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { - self.get_cache_with_similarity(key, context) - .map(|lookup| lookup.value) - } - - 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 response = self.codec.encode(&value)?; - let vector = self - .embedder - .async_embed(&prompt, context.metadata.as_ref()) - .await?; - let tag = Self::tag(key, &context).to_owned(); - let inner = Arc::clone(&self.inner); - Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - inner.store(connection, &tag, response, &prompt, &vector, context.ttl) - }) - .await - } - - async fn async_get_cache( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - self.async_get_cache_with_similarity(key, context) - .await - .map(|lookup| lookup.value) - } -} - -/// Python stamps a similarity of `0.0` when there is no prompt or no hit in the key's scope. -impl SemanticCache for RedisSemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn get_cache_with_similarity( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; - let tag = Self::tag(key, context); - let lookup = self - .connections - .execute(|connection| self.inner.lookup(connection, tag, &vector))?; - self.decode(lookup) - } - - async fn async_get_cache_with_similarity( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let vector = self - .embedder - .async_embed(&prompt, context.metadata.as_ref()) - .await?; - let tag = Self::tag(key, context).to_owned(); - let inner = Arc::clone(&self.inner); - let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - inner.lookup(connection, &tag, &vector) - }) - .await?; - self.decode(lookup) - } -} - -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 { - let mut escaped = String::with_capacity(value.len()); - for ch in value.chars() { - if matches!( - ch, - ',' | '.' - | '<' - | '>' - | '{' - | '}' - | '[' - | ']' - | '\\' - | '"' - | '\'' - | ':' - | ';' - | '!' - | '@' - | '#' - | '$' - | '%' - | '^' - | '&' - | '*' - | '(' - | ')' - | '-' - | '+' - | '=' - | '~' - | '|' - | '/' - | ' ' - | '?' - ) { - escaped.push('\\'); - } - escaped.push(ch); - } - escaped -} - -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/config.rs b/litellm-rust/crates/cache-redis-semantic/src/config.rs deleted file mode 100644 index 6b810628be7..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/src/config.rs +++ /dev/null @@ -1,8 +0,0 @@ -/// `RedisSemanticCache.DEFAULT_REDIS_INDEX_NAME`. -pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; - -#[derive(Clone, Debug)] -pub struct RedisSemanticConfig { - pub index_name: String, - pub similarity_threshold: f32, -} diff --git a/litellm-rust/crates/cache-redis-semantic/src/index.rs b/litellm-rust/crates/cache-redis-semantic/src/index.rs deleted file mode 100644 index 23401938a0e..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/src/index.rs +++ /dev/null @@ -1,209 +0,0 @@ -use std::sync::OnceLock; - -use litellm_cache::Error; -use litellm_cache_redis::connection::ConnectionRef; - -use crate::reply::{number_value, string_value}; - -pub(crate) const CACHE_KEY_FIELD: &str = "litellm_cache_key"; -pub(crate) const VECTOR_FIELD: &str = "prompt_vector"; - -/// The redisvl `SemanticCache` index, resolved once per cache: the configured name when its -/// schema fits, else `_isolated`, recreated when that one is stale too. -pub(crate) struct Index { - name: String, - resolved: OnceLock, -} - -impl Index { - pub(crate) fn new(name: String) -> Self { - Self { - name, - resolved: OnceLock::new(), - } - } - - pub(crate) fn name(&self) -> &str { - &self.name - } - - pub(crate) fn ensure( - &self, - connection: &mut ConnectionRef<'_>, - dims: usize, - ) -> Result { - if let Some(name) = self.resolved.get() { - return Ok(name.clone()); - } - let name = match index_compatible(connection, &self.name, dims)? { - Some(true) => self.name.clone(), - Some(false) => self.isolated(connection, dims)?, - None => match create_index(connection, &self.name, dims) { - Ok(()) => self.name.clone(), - Err(error) => match index_compatible(connection, &self.name, dims)? { - Some(true) => self.name.clone(), - Some(false) => self.isolated(connection, dims)?, - None => return Err(error), - }, - }, - }; - let _ = self.resolved.set(name.clone()); - Ok(name) - } - - fn isolated(&self, connection: &mut ConnectionRef<'_>, dims: usize) -> Result { - let name = format!("{}_isolated", self.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 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(error) => return Err(Error::unavailable(error)), - }; - Ok(Some(schema_compatible(&info, dims))) -} - -fn unknown_index(error: &redis::RedisError) -> bool { - error - .code() - .is_some_and(|code| code.eq_ignore_ascii_case("unknown")) - && error - .detail() - .is_some_and(|detail| detail.to_ascii_lowercase().starts_with("index")) -} - -struct Attribute { - name: Option, - field_type: Option, - dim: Option, - data_type: Option, - distance_metric: Option, -} - -fn attribute(value: &redis::Value) -> Option { - let redis::Value::Array(pairs) = value else { - return None; - }; - let mut attribute = Attribute { - name: None, - field_type: None, - dim: None, - data_type: None, - distance_metric: None, - }; - for pair in pairs.as_chunks::<2>().0 { - match string_value(&pair[0]).as_deref() { - Some("identifier") => attribute.name = string_value(&pair[1]), - Some("type") => attribute.field_type = string_value(&pair[1]), - Some("dim") => attribute.dim = number_value(&pair[1]), - Some("data_type") => attribute.data_type = string_value(&pair[1]), - Some("distance_metric") => attribute.distance_metric = string_value(&pair[1]), - _ => {} - } - } - Some(attribute) -} - -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().filter_map(attribute).collect::>(); - let has_field = |name: &str, field_type: &str| { - fields.iter().any(|field| { - field.name.as_deref() == Some(name) && field.field_type.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(|field| { - field.name.as_deref() == Some(VECTOR_FIELD) - && field.field_type.as_deref() == Some("VECTOR") - && field.dim == Some(dims as f64) - && field - .data_type - .as_deref() - .is_some_and(|data| data.eq_ignore_ascii_case("float32")) - && field - .distance_metric - .as_deref() - .is_some_and(|metric| metric.eq_ignore_ascii_case("cosine")) - }) -} diff --git a/litellm-rust/crates/cache-redis-semantic/src/lib.rs b/litellm-rust/crates/cache-redis-semantic/src/lib.rs deleted file mode 100644 index 251323df6cf..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/src/lib.rs +++ /dev/null @@ -1,7 +0,0 @@ -mod cache; -mod config; -mod index; -mod reply; - -pub use cache::RedisSemanticCache; -pub use config::{DEFAULT_INDEX_NAME, RedisSemanticConfig}; diff --git a/litellm-rust/crates/cache-redis-semantic/src/reply.rs b/litellm-rust/crates/cache-redis-semantic/src/reply.rs deleted file mode 100644 index 24cbd573cde..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/src/reply.rs +++ /dev/null @@ -1,57 +0,0 @@ -pub(crate) 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, - } -} - -pub(crate) 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()), - } -} - -pub(crate) 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]) -} - -pub(crate) fn string_field(fields: &[redis::Value], name: &str) -> Option { - field_value(fields, name).and_then(string_value) -} - -pub(crate) fn number_field(fields: &[redis::Value], name: &str) -> Option { - field_value(fields, name).and_then(number_value) -} - -pub(crate) 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/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs deleted file mode 100644 index cc42f4c3924..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs +++ /dev/null @@ -1,778 +0,0 @@ -mod support; - -use std::time::Duration; - -use litellm_cache::{ - BaseCache, Error, JsonCodec, SemanticCacheContext, - semantic::{SemanticCache, SemanticLookup}, -}; -use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig}; -use redis_test::{MockCmd, MockRedisConnection}; -use rstest::{fixture, rstest}; -use serde_json::{Value, json}; -use sha2::{Digest, Sha256}; -use support::FakeEmbedder; - -const INDEX: &str = DEFAULT_INDEX_NAME; -const PROMPT: &str = "hello prompt"; -const CLOCK: fn() -> f64 = || 1700000000.5; -const VECTOR: [f32; 3] = [0.1, 0.2, 0.3]; - -type MockCache = RedisSemanticCache, MockRedisConnection>; - -#[fixture] -fn config() -> RedisSemanticConfig { - RedisSemanticConfig { - index_name: INDEX.into(), - similarity_threshold: 0.9, - } -} - -#[fixture] -fn entry() -> Value { - json!({"timestamp": 1.0, "response": {"answer": "yes"}}) -} - -#[fixture] -fn context() -> SemanticCacheContext { - messages_context(vec![json!({"role": "user", "content": PROMPT})]) -} - -fn messages_context(messages: Vec) -> SemanticCacheContext { - SemanticCacheContext { - messages: Some(Value::Array(messages)), - ..Default::default() - } -} - -fn cache(commands: Vec, embedder: FakeEmbedder) -> MockCache { - RedisSemanticCache::with_connection( - MockRedisConnection::new(commands).assert_all_commands_consumed(), - embedder, - JsonCodec::new(), - config(), - ) - .with_clock(CLOCK) -} - -fn encoded(value: &Value) -> Vec { - serde_json::to_vec(value).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 { - support::server_error("Unknown index name") -} - -fn info_missing(index: &str) -> MockCmd { - MockCmd::new( - redis::cmd("FT.INFO").arg(index), - Err::(unknown_index_error()), - ) -} - -fn info(index: &str, value: redis::Value) -> MockCmd { - MockCmd::new(redis::cmd("FT.INFO").arg(index), Ok(value)) -} - -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 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_with(dims, "FLOAT32", "COSINE")) -} - -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_with(dims, "FLOAT32", "COSINE"), - ]) -} - -fn create_index(name: &str, dims: usize) -> MockCmd { - MockCmd::new(create_index_command(name, dims), Ok("OK")) -} - -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 hset(index: &str, prompt: &str, tag: &str, vector: &[f32], value: &Value) -> MockCmd { - 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), - ) -} - -fn search( - index: &str, - tag: &str, - vector: &[f32], - reply: redis::RedisResult, -) -> MockCmd { - 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)); - MockCmd::new(command, reply) -} - -fn hit(tag: &str, distance: &str, response: Vec) -> redis::Value { - redis::Value::Array(vec![ - redis::Value::Int(1), - s("litellm_semantic_cache_index:stored-id"), - redis::Value::Array(vec![ - s("entry_id"), - s("stored-id"), - s("prompt"), - s(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 empty_result() -> redis::Value { - redis::Value::Array(vec![redis::Value::Int(0)]) -} - -#[rstest] -#[case::creates_index_and_expires(false, Some(Duration::from_secs(5)), Some(5))] -#[case::existing_index_without_ttl(true, None, None)] -#[case::fractional_ttl_rounds_up(true, Some(Duration::from_millis(1500)), Some(2))] -fn store_writes_the_redisvl_hash( - #[case] index_exists: bool, - #[case] ttl: Option, - #[case] expire: Option, - entry: Value, - context: SemanticCacheContext, -) { - let hash_key = format!("{INDEX}:{}", entry_id(PROMPT, "key1")); - let mut commands = if index_exists { - vec![info(INDEX, compatible_info(3))] - } else { - vec![info_missing(INDEX), create_index(INDEX, 3)] - }; - commands.push(hset(INDEX, PROMPT, "key1", &VECTOR, &entry)); - commands.extend( - expire.map(|seconds| MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(seconds), Ok(1))), - ); - let cache = cache(commands, FakeEmbedder::new(&[])); - - cache - .set_cache("key1", entry, &SemanticCacheContext { ttl, ..context }) - .unwrap(); -} - -#[rstest] -#[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))] -#[case::above_distance_threshold("key1", "0.5", None, Ok(None))] -#[case::other_cache_key("other", "0.05", None, Ok(None))] -#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry(None)))] -fn lookup_applies_threshold_scope_and_codec( - #[case] stored_tag: &str, - #[case] distance: &str, - #[case] response: Option<&[u8]>, - #[case] expected: Result, Error>, - entry: Value, - context: SemanticCacheContext, -) { - let response = response.map_or_else(|| encoded(&entry), <[u8]>::to_vec); - let cache = cache( - vec![ - info(INDEX, compatible_info(3)), - search( - INDEX, - "key1", - &VECTOR, - Ok(hit(stored_tag, distance, response)), - ), - ], - FakeEmbedder::new(&[]), - ); - - let kind = |result: Result, Error>| { - result.map_err(|error| std::mem::discriminant(&error)) - }; - assert_eq!(kind(cache.get_cache("key1", &context)), kind(expected)); -} - -#[rstest] -#[case::hit(context(), Some(hit("key1", "0.05", encoded(&entry()))), Some(entry()), Some(1.0 - 0.05))] -#[case::beyond_distance_threshold(context(), Some(hit("key1", "0.5", encoded(&entry()))), None, Some(0.0))] -#[case::no_results(context(), Some(empty_result()), None, Some(0.0))] -#[case::other_cache_key(context(), Some(hit("other", "0.05", encoded(&entry()))), None, Some(0.0))] -#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))] -#[tokio::test] -async fn lookup_reports_python_semantic_similarity( - #[case] context: SemanticCacheContext, - #[case] reply: Option, - #[case] value: Option, - #[case] similarity: Option, - #[values(false, true)] use_async: bool, -) { - let commands = reply.map_or_else(Vec::new, |reply| { - vec![ - info(INDEX, compatible_info(3)), - search(INDEX, "key1", &VECTOR, Ok(reply)), - ] - }); - let cache = cache(commands, FakeEmbedder::new(&[])); - - let lookup = if use_async { - cache - .async_get_cache_with_similarity("key1", &context) - .await - } else { - cache.get_cache_with_similarity("key1", &context) - }; - - assert_eq!(lookup, Ok(SemanticLookup { value, similarity })); -} - -#[rstest] -#[tokio::test] -async fn missing_prompt_is_a_noop_that_never_embeds(entry: Value) { - let embedder = FakeEmbedder::new(&[]); - let calls = embedder.calls.clone(); - let cache = cache(Vec::new(), embedder); - let context = SemanticCacheContext::default(); - - cache.set_cache("key1", entry.clone(), &context).unwrap(); - assert_eq!(cache.get_cache("key1", &context).unwrap(), None); - cache - .async_set_cache("key1", entry, context.clone()) - .await - .unwrap(); - assert_eq!(cache.async_get_cache("key1", &context).await.unwrap(), None); - assert!(calls.lock().unwrap().is_empty()); -} - -#[rstest] -fn scope_overrides_key_as_filter_tag(entry: Value, context: SemanticCacheContext) { - let cache = cache( - vec![ - info(INDEX, compatible_info(3)), - hset(INDEX, PROMPT, "scope-a", &VECTOR, &entry), - search( - INDEX, - "scope\\-a", - &VECTOR, - Ok(hit("scope-a", "0.05", encoded(&entry))), - ), - ], - FakeEmbedder::new(&[]), - ); - let context = SemanticCacheContext { - scope: Some("scope-a".into()), - ..context - }; - - cache.set_cache("key1", entry.clone(), &context).unwrap(); - assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(entry)); -} - -#[rstest] -#[case::unscoped_schema(unscoped_info(3))] -#[case::wrong_distance_metric(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2")))] -#[case::wrong_data_type(info_with_vector(vector_attribute_with(3, "FLOAT64", "COSINE")))] -fn incompatible_schema_falls_back_to_isolated_index( - #[case] base_info: redis::Value, - entry: Value, - context: SemanticCacheContext, -) { - let isolated = format!("{INDEX}_isolated"); - let cache = cache( - vec![ - info(INDEX, base_info), - info_missing(&isolated), - create_index(&isolated, 3), - hset(&isolated, PROMPT, "key1", &VECTOR, &entry), - ], - FakeEmbedder::new(&[]), - ); - - cache.set_cache("key1", entry, &context).unwrap(); -} - -#[rstest] -fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticCacheContext) { - let cache = cache( - vec![ - info_missing(INDEX), - MockCmd::new( - create_index_command(INDEX, 3), - Err::<&str, _>(support::server_error("Index already exists")), - ), - info(INDEX, compatible_info(3)), - hset(INDEX, PROMPT, "key1", &VECTOR, &entry), - ], - FakeEmbedder::new(&[]), - ); - - cache.set_cache("key1", entry, &context).unwrap(); -} - -#[rstest] -#[case::punctuation_and_spaces("a:b, c|d", "a\\:b\\,\\ c\\|d")] -#[case::braces_and_dots("{x}.y", "\\{x\\}\\.y")] -#[case::plain("key1", "key1")] -fn tag_special_characters_are_escaped_in_search_filter( - #[case] tag: &str, - #[case] escaped: &str, - context: SemanticCacheContext, -) { - let cache = cache( - vec![ - info(INDEX, compatible_info(3)), - search(INDEX, escaped, &VECTOR, Ok(empty_result())), - ], - FakeEmbedder::new(&[]), - ); - - assert_eq!(cache.get_cache(tag, &context).unwrap(), None); -} - -#[rstest] -#[case::content_parts( - SemanticCacheContext { - messages: Some(json!([ - {"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]}, - {"role": "assistant", "content": "reply"}, - ])), - ..Default::default() - }, - "firstsecondreply" -)] -#[case::responses_string_input( - SemanticCacheContext { input: Some(json!(" plain input ")), ..Default::default() }, - "plain input" -)] -#[case::responses_nested_input( - SemanticCacheContext { - input: Some(json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"])), - ..Default::default() - }, - "nested\ntail" -)] -#[case::responses_output_text( - SemanticCacheContext { input: Some(json!({"output_text": " result text "})), ..Default::default() }, - "result text" -)] -#[case::search_results( - messages_context(vec![json!({ - "role": "user", - "content": "question", - "search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}], - })]), - "questionsrctfound{\"a\":1}" -)] -#[case::empty_messages_fall_back_to_input( - SemanticCacheContext { messages: Some(json!([])), input: Some(json!("fallback")), ..Default::default() }, - "fallback" -)] -fn prompt_extraction_matches_python_message_and_input_shapes( - #[case] context: SemanticCacheContext, - #[case] prompt: &str, -) { - let embedder = FakeEmbedder::new(&[]); - let calls = embedder.calls.clone(); - let cache = cache( - vec![ - info(INDEX, compatible_info(3)), - search(INDEX, "key1", &VECTOR, Ok(empty_result())), - ], - embedder, - ); - - cache.get_cache("key1", &context).unwrap(); - - assert_eq!(*calls.lock().unwrap(), vec![(prompt.to_owned(), None)]); -} - -#[rstest] -#[case(None)] -#[case(Some(Duration::from_secs(9)))] -fn ttl_passes_through_context_only(#[case] ttl: Option) { - let cache = cache(Vec::new(), FakeEmbedder::new(&[])); - - assert_eq!( - cache.get_ttl(&SemanticCacheContext { - ttl, - ..Default::default() - }), - ttl - ); -} - -#[rstest] -#[tokio::test] -async fn async_paths_embed_with_metadata_then_run_blocking_redis_work( - entry: Value, - context: SemanticCacheContext, -) { - let embedder = FakeEmbedder::new(&[]); - let calls = embedder.calls.clone(); - let cache = cache( - vec![ - info(INDEX, compatible_info(3)), - hset(INDEX, PROMPT, "key1", &VECTOR, &entry), - search( - INDEX, - "key1", - &VECTOR, - Ok(hit("key1", "0.05", encoded(&entry))), - ), - ], - embedder, - ); - let context = SemanticCacheContext { - metadata: Some(json!({"tenant": "team"})), - ..context - }; - - cache - .async_set_cache("key1", entry.clone(), context.clone()) - .await - .unwrap(); - assert_eq!( - cache.async_get_cache("key1", &context).await.unwrap(), - Some(entry) - ); - assert_eq!( - *calls.lock().unwrap(), - vec![(PROMPT.to_owned(), context.metadata.clone()); 2] - ); -} - -#[rstest] -fn accessors_report_the_config(config: RedisSemanticConfig) { - let cache = cache(Vec::new(), FakeEmbedder::new(&[])); - - assert_eq!(cache.index_name(), config.index_name); - assert!((cache.similarity_threshold() - config.similarity_threshold).abs() < 1e-6); -} - -#[rstest] -fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value) { - // Pins parity with Python's `_isolated` + overwrite=True flow. - let prompt = "shared prompt"; - let isolated = format!("{INDEX}_isolated"); - let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); - - let vector_a = vec![0.1f32; 8]; - let worker_a = cache( - vec![ - info_missing(INDEX), - create_index(INDEX, 8), - hset(INDEX, prompt, "key1", &vector_a, &entry), - ], - FakeEmbedder::new(&[(prompt, &vector_a)]), - ); - worker_a - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - - let vector_b = vec![0.2f32; 4]; - let worker_b = cache( - vec![ - info(INDEX, compatible_info(8)), - info_missing(&isolated), - create_index(&isolated, 4), - hset(&isolated, prompt, "key1", &vector_b, &entry), - search( - &isolated, - "key1", - &vector_b, - Ok(hit("key1", "0.0", encoded(&entry))), - ), - search( - &isolated, - "key1", - &vector_b, - Err(redis::RedisError::from(( - redis::ErrorKind::Extension, - "Vector dimension mismatch", - ))), - ), - ], - FakeEmbedder::new(&[(prompt, &vector_b)]), - ); - worker_b - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - assert_eq!( - worker_b.get_cache("key1", &context()).unwrap(), - Some(entry.clone()) - ); - - let vector_c = vec![0.3f32; 16]; - let worker_c = cache( - vec![ - info(INDEX, compatible_info(8)), - info(&isolated, compatible_info(4)), - MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")), - create_index(&isolated, 16), - hset(&isolated, prompt, "key1", &vector_c, &entry), - ], - FakeEmbedder::new(&[(prompt, &vector_c)]), - ); - worker_c - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - - assert!(matches!( - worker_b.get_cache("key1", &context()), - Err(Error::Unavailable(Some(_))) - )); -} - -#[fixture] -fn redis_stack_url() -> Option { - std::env::var("LITELLM_REDIS_STACK_URL").ok() -} - -fn live_cache( - url: &str, - index_name: &str, - prompt: &str, - vector: Vec, -) -> RedisSemanticCache> { - RedisSemanticCache::new( - url, - FakeEmbedder::new(&[(prompt, vector.as_slice())]), - JsonCodec::::new(), - RedisSemanticConfig { - index_name: index_name.to_owned(), - similarity_threshold: 0.9, - }, - ) - .unwrap() -} - -#[rstest] -fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option, entry: Value) { - let Some(url) = 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 context = || messages_context(vec![json!({"role": "user", "content": prompt})]); - - let worker_a = live_cache(&url, &base, prompt, vec![0.1f32; 8]); - worker_a - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - - let worker_b = live_cache(&url, &base, prompt, vec![0.2f32; 4]); - worker_b - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - assert_eq!( - worker_b.get_cache("key1", &context()).unwrap(), - Some(entry.clone()) - ); - - let worker_c = live_cache(&url, &base, prompt, vec![0.3f32; 16]); - worker_c - .set_cache("key1", entry.clone(), &context()) - .unwrap(); - - assert!(matches!( - worker_b.get_cache("key1", &context()), - Err(Error::Unavailable(Some(_))) - )); - - 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); - } -} - -#[rstest] -fn live_store_lookup_and_ttl_against_redis_stack(redis_stack_url: Option, entry: Value) { - let Some(url) = redis_stack_url else { - return; - }; - let prompt = "rust semantic cache live prompt"; - let index_name = format!("rust_semantic_test_{}", std::process::id()); - let cache = live_cache(&url, &index_name, prompt, vec![0.1, 0.2, 0.3, 0.4]); - let context = SemanticCacheContext { - ttl: Some(Duration::from_secs(120)), - ..messages_context(vec![json!({"role": "user", "content": prompt})]) - }; - - cache - .set_cache("live-key", entry.clone(), &context) - .unwrap(); - assert_eq!(cache.get_cache("live-key", &context).unwrap(), Some(entry)); - 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, "live-key")), - ) - .unwrap(); - assert!( - ttl > 0, - "expected stored hash to carry an expiry, got {ttl}" - ); -} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/contract.rs b/litellm-rust/crates/cache-redis-semantic/tests/contract.rs deleted file mode 100644 index fb3cda9616e..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/tests/contract.rs +++ /dev/null @@ -1,63 +0,0 @@ -mod support; - -use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding}; -use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig}; -use litellm_cache_testing as contract; -use rstest::{fixture, rstest}; -use serde_json::{Value, json}; -use support::FakeSearch; - -type Cache = RedisSemanticCache, FakeSearch>; - -const PREFIX: &str = "contract:"; - -#[fixture] -fn cache() -> Cache { - RedisSemanticCache::with_connection( - FakeSearch::default(), - PreparedEmbedding(vec![0.6, 0.8]), - JsonCodec::new(), - RedisSemanticConfig { - index_name: DEFAULT_INDEX_NAME.into(), - similarity_threshold: 0.9, - }, - ) -} - -#[fixture] -fn context() -> SemanticCacheContext { - SemanticCacheContext { - messages: Some(json!([{"role": "user", "content": "contract prompt"}])), - ..Default::default() - } -} - -#[rstest] -#[tokio::test] -async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) { - contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await; -} - -#[rstest] -#[tokio::test] -async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) { - contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await; -} - -#[rstest] -#[tokio::test] -async fn overwrite_replaces(cache: Cache, context: SemanticCacheContext) { - contract::overwrite_replaces(&cache, context, PREFIX, json!(1), json!({"b": 2})).await; -} - -#[rstest] -#[tokio::test] -async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) { - contract::pipeline_writes_every_entry( - &cache, - context, - PREFIX, - vec![json!("a"), json!(2), json!({"c": true})], - ) - .await; -} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs deleted file mode 100644 index f112dbe96a7..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs +++ /dev/null @@ -1,302 +0,0 @@ -#![allow(dead_code)] - -use std::{ - collections::{BTreeMap, HashMap}, - sync::{Arc, Mutex}, -}; - -use litellm_cache::{Error, semantic::Embedder}; -use serde_json::Value; - -pub type EmbedCalls = Arc)>>>; - -/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every -/// prompt with its metadata. -pub struct FakeEmbedder { - vectors: HashMap>, - pub calls: EmbedCalls, -} - -impl FakeEmbedder { - pub fn new(vectors: &[(&str, &[f32])]) -> Self { - Self { - vectors: vectors - .iter() - .map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec())) - .collect(), - calls: EmbedCalls::default(), - } - } -} - -impl Embedder for FakeEmbedder { - fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { - self.calls - .lock() - .unwrap() - .push((prompt.to_owned(), metadata.cloned())); - 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) - } -} - -struct FakeIndex { - prefix: Vec, - dims: usize, - vector_field: String, -} - -#[derive(Default)] -struct SearchState { - indexes: HashMap, - hashes: BTreeMap, BTreeMap>>, -} - -/// An in-memory Redis Stack speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache -/// sends, with exact cosine KNN over the hashes under an index prefix. -#[derive(Clone, Default)] -pub struct FakeSearch { - state: Arc>, -} - -impl FakeSearch { - fn run(&self, args: Vec>) -> redis::RedisResult { - let mut state = self.state.lock().unwrap(); - let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned(); - match text(0).to_uppercase().as_str() { - "FT.CREATE" => { - let name = text(1); - if state.indexes.contains_key(&name) { - return Err(server_error("Index already exists")); - } - let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes()); - let prefix = args[position("PREFIX").unwrap() + 2].clone(); - let dims = text(position("DIM").unwrap() + 1).parse().unwrap(); - let vector_field = text(position("VECTOR").unwrap() - 1); - state.indexes.insert( - name, - FakeIndex { - prefix, - dims, - vector_field, - }, - ); - Ok(redis::Value::Okay) - } - "FT.INFO" => { - let index = state - .indexes - .get(&text(1)) - .ok_or_else(|| server_error("Unknown index name"))?; - Ok(index_info(index)) - } - "FT.DROPINDEX" => { - state.indexes.remove(&text(1)); - Ok(redis::Value::Okay) - } - "HSET" => { - let hash = state.hashes.entry(args[1].clone()).or_default(); - for pair in args[2..].chunks(2) { - hash.insert( - String::from_utf8_lossy(&pair[0]).into_owned(), - pair[1].clone(), - ); - } - Ok(redis::Value::Int(((args.len() - 2) / 2) as i64)) - } - "EXPIRE" => Ok(redis::Value::Int(i64::from( - state.hashes.contains_key(&args[1]), - ))), - "FT.SEARCH" => { - let index = state - .indexes - .get(&text(1)) - .ok_or_else(|| server_error("no such index"))?; - let query = text(2); - let tag = query_tag(&query); - let params = args.iter().position(|arg| arg == b"PARAMS").unwrap(); - let vector = floats(&args[params + 3]); - let best = state - .hashes - .iter() - .filter(|(key, _)| key.starts_with(&index.prefix)) - .filter(|(_, fields)| { - fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes()) - }) - .filter_map(|(key, fields)| { - let stored = floats(fields.get(&index.vector_field)?); - (stored.len() == index.dims) - .then(|| (key, fields, 1.0 - cosine(&vector, &stored))) - }) - .min_by(|left, right| left.2.total_cmp(&right.2)); - let Some((key, fields, distance)) = best else { - return Ok(redis::Value::Array(vec![redis::Value::Int(0)])); - }; - let mut reply = fields - .iter() - .filter(|(name, _)| **name != index.vector_field) - .flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)]) - .collect::>(); - reply.extend([ - bulk(b"vector_distance"), - bulk(distance.to_string().as_bytes()), - ]); - Ok(redis::Value::Array(vec![ - redis::Value::Int(1), - bulk(key), - redis::Value::Array(reply), - ])) - } - "PING" => Ok(redis::Value::SimpleString("PONG".into())), - _ => Err(server_error("unsupported command")), - } - } -} - -impl redis::ConnectionLike for FakeSearch { - fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult { - let mut commands = parse_commands(command); - self.run(commands.remove(0)) - } - - fn req_packed_commands( - &mut self, - commands: &[u8], - offset: usize, - count: usize, - ) -> redis::RedisResult> { - let replies = parse_commands(commands) - .into_iter() - .map(|args| self.run(args)) - .collect::>>()?; - Ok(replies.into_iter().skip(offset).take(count).collect()) - } - - fn get_db(&self) -> i64 { - 0 - } - - fn check_connection(&mut self) -> bool { - true - } - - fn is_open(&self) -> bool { - true - } -} - -pub fn server_error(reply: &str) -> redis::RedisError { - redis::parse_redis_value(format!("-{reply}\r\n").as_bytes()) - .expect("error reply") - .extract_error() - .expect_err("server error") -} - -fn bulk(bytes: &[u8]) -> redis::Value { - redis::Value::BulkString(bytes.to_vec()) -} - -fn index_info(index: &FakeIndex) -> redis::Value { - let attribute = |name: &str, field_type: &str| { - redis::Value::Array(vec![ - bulk(b"identifier"), - bulk(name.as_bytes()), - bulk(b"type"), - bulk(field_type.as_bytes()), - ]) - }; - redis::Value::Array(vec![ - bulk(b"attributes"), - redis::Value::Array(vec![ - attribute("prompt", "TEXT"), - attribute("response", "TEXT"), - attribute("inserted_at", "NUMERIC"), - attribute("updated_at", "NUMERIC"), - attribute("litellm_cache_key", "TAG"), - redis::Value::Array(vec![ - bulk(b"identifier"), - bulk(index.vector_field.as_bytes()), - bulk(b"type"), - bulk(b"VECTOR"), - bulk(b"dim"), - redis::Value::Int(index.dims as i64), - bulk(b"data_type"), - bulk(b"FLOAT32"), - bulk(b"distance_metric"), - bulk(b"COSINE"), - ]), - ]), - ]) -} - -/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed. -fn query_tag(query: &str) -> String { - let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len(); - let mut tag = String::new(); - let mut characters = query[start..].chars(); - while let Some(character) = characters.next() { - match character { - '\\' => tag.extend(characters.next()), - '}' => break, - character => tag.push(character), - } - } - tag -} - -fn floats(bytes: &[u8]) -> Vec { - bytes - .as_chunks::<4>() - .0 - .iter() - .map(|chunk| f32::from_le_bytes(*chunk)) - .collect() -} - -fn cosine(left: &[f32], right: &[f32]) -> f64 { - let dot = left - .iter() - .zip(right) - .map(|(left, right)| f64::from(*left) * f64::from(*right)) - .sum::(); - let norm = |vector: &[f32]| { - vector - .iter() - .map(|value| f64::from(*value).powi(2)) - .sum::() - .sqrt() - }; - dot / (norm(left) * norm(right)) -} - -/// Splits a packed RESP request into each command's arguments. -fn parse_commands(mut bytes: &[u8]) -> Vec>> { - let line = |bytes: &mut &[u8]| { - let end = bytes - .windows(2) - .position(|window| window == b"\r\n") - .unwrap(); - let text = String::from_utf8(bytes[1..end].to_vec()).unwrap(); - *bytes = &bytes[end + 2..]; - text.parse::().unwrap() - }; - let mut commands = Vec::new(); - while !bytes.is_empty() { - let count = line(&mut bytes); - let mut args = Vec::with_capacity(count); - for _ in 0..count { - let length = line(&mut bytes); - args.push(bytes[..length].to_vec()); - bytes = &bytes[length + 2..]; - } - commands.push(args); - } - commands -} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml deleted file mode 100644 index f7621299506..00000000000 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ /dev/null @@ -1,20 +0,0 @@ -[package] -name = "litellm-cache-redis" -version = "0.1.0" -edition.workspace = true -license.workspace = true -repository.workspace = true - -[dependencies] -litellm-cache.workspace = true -litellm-error.workspace = true -redis = { workspace = true, features = ["cluster", "tls-rustls"] } -r2d2 = "0.8.10" -thiserror.workspace = true -tokio.workspace = true - -[dev-dependencies] -litellm-cache-testing.workspace = true -redis-test.workspace = true -rstest.workspace = true -serde_json.workspace = true diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs deleted file mode 100644 index b873a2c0282..00000000000 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ /dev/null @@ -1,192 +0,0 @@ -use std::{ - sync::{Arc, OnceLock}, - time::Duration, -}; - -use litellm_cache::{BatchEntry, CacheCodec, Error}; - -use crate::{ - connection::{ConnectionRef, Connections}, - topology::RedisTopology, -}; - -const DEFAULT_TTL: Duration = Duration::from_secs(600); - -pub struct RedisCache { - pub(crate) connections: Arc>, - pub(crate) default_ttl: Duration, - pub(crate) codec: S, - pub(crate) namespace: Option, - pub(crate) topology: RedisTopology, - /// The server's major version, read from `INFO` once, like Python's `redis_version`. - pub(crate) major_version: Arc>, -} - -impl RedisCache { - pub fn new(url: &str, default_ttl: Option, codec: S) -> Result { - Self::connect(url, &RedisTopology::Standalone, default_ttl, codec) - } - - pub fn connect( - url: &str, - topology: &RedisTopology, - default_ttl: Option, - codec: S, - ) -> Result { - let connections = Connections::open(url, topology)?; - Ok(Self { - connections: Arc::new(connections), - default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), - codec, - namespace: None, - topology: topology.clone(), - major_version: Arc::default(), - }) - } -} - -impl RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - pub fn with_connection(connection: C, default_ttl: Option, codec: S) -> Self { - Self { - connections: Arc::new(Connections::fixed(connection)), - default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), - codec, - namespace: None, - topology: RedisTopology::Standalone, - major_version: Arc::default(), - } - } - - pub fn with_namespace(self, namespace: Option) -> Self { - Self { - namespace: namespace.filter(|value| !value.is_empty()), - ..self - } - } - - pub fn namespace(&self) -> Option<&str> { - self.namespace.as_deref() - } - - pub fn topology(&self) -> &RedisTopology { - &self.topology - } - - pub(crate) fn namespaced_key(&self, key: &str) -> String { - namespaced_key(self.namespace.as_deref(), key) - } - - pub(crate) fn namespaced_keys(&self, keys: &[String]) -> Vec { - keys.iter().map(|key| self.namespaced_key(key)).collect() - } - - /// Whole seconds for `ttl`, falling back to the default TTL like Python's `get_ttl`. - pub(crate) fn ttl_or_default(&self, ttl: Option) -> u64 { - ttl_seconds(ttl.unwrap_or(self.default_ttl)) - } - - pub(crate) fn execute( - &self, - operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, - ) -> Result { - self.connections.execute(operation) - } - - /// `_parse_redis_major_version`: the major version from `INFO`, or - /// `DEFAULT_REDIS_MAJOR_VERSION` when `INFO` fails or its version does not parse. The first - /// answer is kept, as Python reads `redis_version` once at construction. - pub(crate) async fn major_version(&self) -> u32 { - if let Some(version) = self.major_version.get() { - return *version; - } - let info = self - .run(|connection| connection.node_text(&redis::cmd("INFO"))) - .await; - let version = info - .ok() - .and_then(|info| parse_major_version(&info)) - .unwrap_or_else(default_major_version); - *self.major_version.get_or_init(|| version) - } - - pub(crate) async fn run(&self, operation: F) -> Result - where - T: Send + 'static, - F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, - { - Connections::run_blocking(Arc::clone(&self.connections), operation).await - } - - pub(crate) fn namespaced_pattern(&self) -> Result { - let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?; - let escaped: String = namespace - .chars() - .flat_map(|ch| { - if matches!(ch, '*' | '?' | '[' | ']' | '\\') { - vec!['\\', ch] - } else { - vec![ch] - } - }) - .collect(); - Ok(format!("{escaped}:*")) - } - - pub(crate) fn decode_response(&self, value: redis::Value) -> Result, Error> { - match value { - redis::Value::Nil => Ok(None), - redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some), - redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some), - _ => Err(Error::InvalidEntry(None)), - } - } - - pub(crate) fn decode_batch_response( - &self, - value: redis::Value, - ) -> Result, Error> { - match self.decode_response(value) { - Ok(Some(value)) => Ok(BatchEntry::Hit(value)), - Ok(None) => Ok(BatchEntry::Miss), - Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid), - Err(error) => Err(error), - } - } -} - -pub(crate) fn namespaced_key(namespace: Option<&str>, key: &str) -> String { - match namespace { - Some(namespace) if !key.starts_with(&format!("{namespace}:")) => { - format!("{namespace}:{key}") - } - _ => key.into(), - } -} - -pub(crate) fn ttl_seconds(ttl: Duration) -> u64 { - ttl.as_secs() - .saturating_add(u64::from(ttl.subsec_nanos() > 0)) - .max(1) -} - -fn parse_major_version(info: &str) -> Option { - let version = info - .lines() - .find_map(|line| line.trim().strip_prefix("redis_version:"))? - .trim(); - match version.split_once('.') { - Some((major, _)) => major.parse().ok(), - None => version.parse::().ok().map(|major| major as u32), - } -} - -fn default_major_version() -> u32 { - std::env::var("DEFAULT_REDIS_MAJOR_VERSION") - .ok() - .and_then(|value| value.parse().ok()) - .unwrap_or(7) -} diff --git a/litellm-rust/crates/cache-redis/src/claim.rs b/litellm-rust/crates/cache-redis/src/claim.rs deleted file mode 100644 index c57a1b03395..00000000000 --- a/litellm-rust/crates/cache-redis/src/claim.rs +++ /dev/null @@ -1,109 +0,0 @@ -use litellm_cache::{CacheCodec, ClaimCache, Error, ExactCacheContext}; -use redis::Commands; - -use crate::{cache::RedisCache, connection::ConnectionRef}; - -const CLAIM_SCRIPT: &str = concat!( - "local current = redis.call('GET', KEYS[1]); ", - "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", - "elseif current ~= ARGV[1] then return 0; end; ", - "if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ", - "elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1" -); -const CLAIM_ATTEMPTS: usize = 8; - -fn stored_bytes(value: redis::Value) -> Result>, Error> { - match value { - redis::Value::Nil => Ok(None), - redis::Value::BulkString(bytes) => Ok(Some(bytes)), - redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())), - _ => Err(Error::InvalidEntry(None)), - } -} - -/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's -/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the -/// bytes that decision was made on, retried when another claimant wins the race. -fn claim( - connection: &mut ConnectionRef<'_>, - codec: &S, - key: &str, - candidate: S::Value, - eligible: &[S::Value], - ttl: u64, -) -> Result -where - S::Value: PartialEq, -{ - let payload = codec.encode(&candidate)?; - if payload.is_empty() { - return Err(Error::InvalidEntry(None)); - } - for _ in 0..CLAIM_ATTEMPTS { - let current = stored_bytes( - connection - .get::<_, redis::Value>(key) - .map_err(Error::unavailable)?, - )? - .filter(|bytes| !bytes.is_empty()); - let existing = current - .as_deref() - .and_then(|bytes| codec.decode(bytes).ok()) - .filter(|existing| eligible.is_empty() || eligible.contains(existing)); - let refresh = existing - .as_ref() - .is_some_and(|existing| !eligible.is_empty() || *existing == candidate); - let write: &[u8] = if existing.is_some() { b"" } else { &payload }; - let applied = redis::cmd("EVAL") - .arg(CLAIM_SCRIPT) - .arg(1) - .arg(key) - .arg(current.as_deref().unwrap_or_default()) - .arg(ttl) - .arg(write) - .arg(u8::from(refresh)) - .query::(connection) - .map_err(Error::unavailable)?; - if applied { - return Ok(existing.unwrap_or(candidate)); - } - } - Err(Error::unavailable(ClaimContended)) -} - -#[derive(Debug, thiserror::Error)] -#[error("claim kept losing the race after {CLAIM_ATTEMPTS} attempts")] -struct ClaimContended; - -impl ClaimCache for RedisCache -where - S: CacheCodec + Clone + 'static, - S::Value: PartialEq, - C: redis::ConnectionLike + Send + 'static, -{ - fn claim_cache( - &self, - key: &str, - candidate: S::Value, - eligible: &[S::Value], - context: ExactCacheContext, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(context.ttl); - self.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl)) - } - - async fn async_claim_cache( - &self, - key: &str, - candidate: S::Value, - eligible: Vec, - context: ExactCacheContext, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(context.ttl); - let codec = self.codec.clone(); - self.run(move |connection| claim(connection, &codec, &key, candidate, &eligible, ttl)) - .await - } -} diff --git a/litellm-rust/crates/cache-redis/src/connection.rs b/litellm-rust/crates/cache-redis/src/connection.rs deleted file mode 100644 index 99befa16ffa..00000000000 --- a/litellm-rust/crates/cache-redis/src/connection.rs +++ /dev/null @@ -1,513 +0,0 @@ -use std::{ - sync::{Arc, Mutex}, - time::Duration, -}; - -use litellm_cache::Error; -use redis::{ - ConnectionAddr, ConnectionInfo, IntoConnectionInfo, - cluster::{ - ClusterClient, ClusterClientBuilder, ClusterConnection, ClusterPipeline, NodeAddress, - }, - cluster_routing::{ - MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo, - }, -}; - -use crate::topology::{RedisNode, RedisTopology}; - -pub(crate) const REDIS_TIMEOUT: Duration = Duration::from_secs(5); -const REDIS_POOL_SIZE: u32 = 16; - -#[allow(private_interfaces)] -pub enum Connections { - Pool(r2d2::Pool), - Cluster(r2d2::Pool), - Fixed(Mutex), -} - -impl Connections -where - C: redis::ConnectionLike + Send + 'static, -{ - pub fn execute( - &self, - operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, - ) -> Result { - match self { - Self::Pool(pool) => { - let mut pooled = pool.get().map_err(Error::unavailable)?; - let result = operation(&mut ConnectionRef::Node(&mut pooled.connection)); - pooled.failed = matches!(result, Err(Error::Unavailable(_))); - result - } - Self::Cluster(pool) => { - let mut pooled = pool.get().map_err(Error::unavailable)?; - let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection)); - pooled.failed = matches!(result, Err(Error::Unavailable(_))); - result - } - Self::Fixed(connection) => { - let mut connection = connection.lock()?; - operation(&mut ConnectionRef::Node(&mut *connection)) - } - } - } - - pub async fn run_blocking(connections: Arc, operation: F) -> Result - where - T: Send + 'static, - F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, - { - tokio::task::spawn_blocking(move || connections.execute(operation)) - .await - .map_err(Error::unavailable)? - } - - pub fn fixed(connection: C) -> Self { - Self::Fixed(Mutex::new(connection)) - } - - pub fn open(url: &str, topology: &RedisTopology) -> Result { - match topology { - RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)), - RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool( - ClusterConnectionManager::open(url, startup_nodes)?, - )?)), - } - } - - /// Closes every idle pooled connection; the next operation opens a fresh one. Connections - /// checked out right now return to the pool, and a caller-owned connection stays open. - pub fn disconnect(&self) { - match self { - Self::Pool(pool) => close_idle(pool), - Self::Cluster(pool) => close_idle(pool), - Self::Fixed(_) => {} - } - } -} - -fn pool(manager: M) -> Result, Error> { - r2d2::Pool::builder() - .max_size(REDIS_POOL_SIZE) - .min_idle(Some(0)) - .connection_timeout(REDIS_TIMEOUT) - .test_on_check_out(false) - .build(manager) - .map_err(Error::unavailable) -} - -fn close_idle(pool: &r2d2::Pool) -where - M: r2d2::ManageConnection>, -{ - let mut idle = Vec::new(); - while let Some(mut connection) = pool.try_get() { - connection.failed = true; - idle.push(connection); - } -} - -pub(crate) struct PooledConnection { - connection: C, - failed: bool, -} - -/// 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(crate) struct ConnectionManager(redis::Client); - -impl ConnectionManager { - fn open(url: &str) -> Result { - redis::Client::open(url) - .map(Self) - .map_err(Error::unavailable) - } -} - -impl r2d2::ManageConnection for ConnectionManager { - type Connection = PooledConnection; - type Error = redis::RedisError; - - fn connect(&self) -> Result { - let connection = self.0.get_connection()?; - connection.set_read_timeout(Some(REDIS_TIMEOUT))?; - connection.set_write_timeout(Some(REDIS_TIMEOUT))?; - Ok(PooledConnection { - connection, - failed: false, - }) - } - - fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> { - redis::cmd("PING").query::(&mut connection.connection)?; - Ok(()) - } - - fn has_broken(&self, connection: &mut Self::Connection) -> bool { - connection.failed || !redis::ConnectionLike::is_open(&connection.connection) - } -} - -#[derive(Debug, thiserror::Error)] -enum TopologyFault { - #[error("Redis Cluster needs at least one startup node")] - NoStartupNodes, - #[error("Redis Cluster nodes need a TCP address")] - NonTcpAddress, - #[error("Redis Cluster answered a fan-out command without a per-node map")] - FanOutReply, - #[error("scan target does not match the connection topology")] - ScanTarget, -} - -pub(crate) struct ClusterConnectionManager(ClusterClient); - -impl ClusterConnectionManager { - fn open(url: &str, startup_nodes: &[RedisNode]) -> Result { - if startup_nodes.is_empty() { - return Err(Error::unavailable(TopologyFault::NoStartupNodes)); - } - let info = url.into_connection_info().map_err(Error::unavailable)?; - let nodes = startup_nodes - .iter() - .map(|node| node_info(&info, node)) - .collect::, _>>()?; - ClusterClientBuilder::new(nodes) - .connection_timeout(REDIS_TIMEOUT) - .response_timeout(REDIS_TIMEOUT) - .build() - .map(Self) - .map_err(Error::unavailable) - } -} - -fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result { - let addr = match info.addr() { - ConnectionAddr::Tcp(..) => ConnectionAddr::Tcp(node.host.clone(), node.port), - ConnectionAddr::TcpTls { - insecure, - tls_params, - .. - } => ConnectionAddr::TcpTls { - host: node.host.clone(), - port: node.port, - insecure: *insecure, - tls_params: tls_params.clone(), - }, - _ => return Err(Error::unavailable(TopologyFault::NonTcpAddress)), - }; - Ok(info.clone().set_addr(addr)) -} - -impl r2d2::ManageConnection for ClusterConnectionManager { - type Connection = PooledConnection; - type Error = redis::RedisError; - - fn connect(&self) -> Result { - let connection = self.0.get_connection()?; - connection.set_read_timeout(Some(REDIS_TIMEOUT))?; - connection.set_write_timeout(Some(REDIS_TIMEOUT))?; - Ok(PooledConnection { - connection, - failed: false, - }) - } - - fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> { - redis::cmd("PING").query::(&mut connection.connection)?; - Ok(()) - } - - fn has_broken(&self, connection: &mut Self::Connection) -> bool { - connection.failed || !redis::ConnectionLike::is_open(&connection.connection) - } -} - -pub enum ConnectionRef<'a> { - Node(&'a mut dyn redis::ConnectionLike), - Cluster(&'a mut ClusterConnection), -} - -impl redis::ConnectionLike for ConnectionRef<'_> { - fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { - match self { - Self::Node(connection) => connection.req_packed_command(cmd), - Self::Cluster(connection) => connection.req_packed_command(cmd), - } - } - - fn req_packed_commands( - &mut self, - cmd: &[u8], - offset: usize, - count: usize, - ) -> redis::RedisResult> { - match self { - Self::Node(connection) => connection.req_packed_commands(cmd, offset, count), - Self::Cluster(connection) => connection.req_packed_commands(cmd, offset, count), - } - } - - fn get_db(&self) -> i64 { - match self { - Self::Node(connection) => connection.get_db(), - Self::Cluster(connection) => redis::ConnectionLike::get_db(*connection), - } - } - - fn supports_pipelining(&self) -> bool { - match self { - Self::Node(connection) => connection.supports_pipelining(), - Self::Cluster(connection) => redis::ConnectionLike::supports_pipelining(*connection), - } - } - - fn check_connection(&mut self) -> bool { - match self { - Self::Node(connection) => connection.check_connection(), - Self::Cluster(connection) => connection.check_connection(), - } - } - - fn is_open(&self) -> bool { - match self { - Self::Node(connection) => connection.is_open(), - Self::Cluster(connection) => redis::ConnectionLike::is_open(*connection), - } - } -} - -impl ConnectionRef<'_> { - /// Runs `pipeline` and decodes its non-ignored replies as `T`. A cluster connection refuses - /// `Pipeline::query`, so there a transaction goes to its keys' slot as one MULTI/EXEC and - /// anything else is split per node by `ClusterPipeline`; either way the raw replies are - /// handed back to `pipeline` to decode. - pub(crate) fn query_pipeline( - &mut self, - pipeline: &redis::Pipeline, - ) -> Result { - match self { - Self::Node(connection) => pipeline.query(*connection), - Self::Cluster(connection) if pipeline.is_transaction() => { - redis::ConnectionLike::req_packed_commands( - *connection, - &pipeline.get_packed_pipeline(), - pipeline.len() + 1, - 1, - ) - .and_then(|replies| pipeline.query(&mut Replies(Some(replies)))) - } - Self::Cluster(connection) => { - let mut cluster = ClusterPipeline::with_capacity(pipeline.len()); - for command in pipeline.cmd_iter() { - cluster.add_command(command.clone()); - } - cluster - .query(connection) - .and_then(|replies| pipeline.query(&mut Replies(Some(replies)))) - } - } - .map_err(Error::unavailable) - } - - pub(crate) fn scan( - &mut self, - pattern: &str, - count: usize, - mut visit: impl FnMut(&mut Self, Vec) -> Result, - ) -> Result<(), Error> { - let pages = match self { - Self::Node(connection) => { - let page = scan_command(0, pattern, count) - .query::(*connection) - .map_err(Error::unavailable)?; - vec![(None, page)] - } - Self::Cluster(connection) => connection - .route_command( - &scan_command(0, pattern, count), - RoutingInfo::MultiNode(( - MultipleNodeRoutingInfo::AllMasters, - Some(ResponsePolicy::Special), - )), - ) - .map_err(Error::unavailable) - .and_then(primary_pages)? - .into_iter() - .map(|(node, page)| (Some(node), page)) - .collect(), - }; - for (node, (mut cursor, mut keys)) in pages { - loop { - if !visit(self, keys)? { - return Ok(()); - } - if cursor == 0 { - break; - } - (cursor, keys) = self.scan_page(node.as_ref(), cursor, pattern, count)?; - } - } - Ok(()) - } - - pub(crate) fn ping(&mut self) -> Result { - let command = redis::cmd("PING"); - match self { - Self::Node(connection) => command - .query::(*connection) - .map(|response| response == "PONG"), - Self::Cluster(connection) => connection - .route_command( - &command, - RoutingInfo::MultiNode(( - MultipleNodeRoutingInfo::AllNodes, - Some(ResponsePolicy::AllSucceeded), - )), - ) - .map(|_| true), - } - } - - pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result { - match self { - Self::Node(connection) => command.query(*connection).map_err(Error::unavailable), - Self::Cluster(connection) => { - let value = connection - .route_command( - command, - RoutingInfo::MultiNode(( - MultipleNodeRoutingInfo::AllNodes, - Some(ResponsePolicy::Special), - )), - ) - .map_err(Error::unavailable)?; - let redis::Value::Map(entries) = value else { - return Err(Error::unavailable(TopologyFault::FanOutReply)); - }; - let mut replies = entries - .into_iter() - .map(|(node, reply)| { - Ok(( - redis::from_redis_value::(node).map_err(Error::unavailable)?, - redis::from_redis_value::(reply).map_err(Error::unavailable)?, - )) - }) - .collect::, Error>>()?; - replies.sort(); - Ok(replies - .into_iter() - .map(|(_, reply)| reply) - .collect::>() - .join("\n")) - } - } - } - - pub(crate) fn flushall(&mut self) -> Result<(), Error> { - let command = redis::cmd("FLUSHALL"); - match self { - Self::Node(connection) => command.query(*connection).map_err(Error::unavailable), - Self::Cluster(connection) => connection - .route_command( - &command, - RoutingInfo::MultiNode(( - MultipleNodeRoutingInfo::AllMasters, - Some(ResponsePolicy::AllSucceeded), - )), - ) - .map(|_| ()) - .map_err(Error::unavailable), - } - } - - fn scan_page( - &mut self, - node: Option<&NodeAddress>, - cursor: u64, - pattern: &str, - count: usize, - ) -> Result { - let command = scan_command(cursor, pattern, count); - match (self, node) { - (Self::Node(connection), None) => { - command.query(*connection).map_err(Error::unavailable) - } - (Self::Cluster(connection), Some(node)) => connection - .route_command( - &command, - RoutingInfo::SingleNode(SingleNodeRoutingInfo::ByAddress { - host: node.host().to_string(), - port: node.port(), - }), - ) - .map_err(Error::unavailable) - .and_then(|value| redis::from_redis_value(value).map_err(Error::unavailable)), - _ => Err(Error::unavailable(TopologyFault::ScanTarget)), - } - } -} - -type ScanPage = (u64, Vec); - -fn primary_pages(value: redis::Value) -> Result, Error> { - let redis::Value::Map(entries) = value else { - return Err(Error::unavailable(TopologyFault::FanOutReply)); - }; - entries - .into_iter() - .map(|(node, page)| { - let node = redis::from_redis_value::(node).map_err(Error::unavailable)?; - let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?; - let page = redis::from_redis_value::(page).map_err(Error::unavailable)?; - Ok((node, page)) - }) - .collect() -} - -fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd { - let mut command = redis::cmd("SCAN"); - command - .cursor_arg(cursor) - .arg("MATCH") - .arg(pattern) - .arg("COUNT") - .arg(count); - command -} - -/// Hands already received pipeline replies to `Pipeline::query`, so it applies its own -/// ignore and error handling to replies a cluster pipeline gathered from several nodes. -struct Replies(Option>); - -impl redis::ConnectionLike for Replies { - fn req_packed_command(&mut self, _: &[u8]) -> redis::RedisResult { - Err((redis::ErrorKind::Client, "replies hold a pipeline only").into()) - } - - fn req_packed_commands( - &mut self, - _: &[u8], - _: usize, - _: usize, - ) -> redis::RedisResult> { - self.0 - .take() - .ok_or_else(|| (redis::ErrorKind::Client, "replies were already read").into()) - } - - fn get_db(&self) -> i64 { - 0 - } - - fn check_connection(&mut self) -> bool { - true - } - - fn is_open(&self) -> bool { - true - } -} diff --git a/litellm-rust/crates/cache-redis/src/counter.rs b/litellm-rust/crates/cache-redis/src/counter.rs deleted file mode 100644 index a4beb9f5712..00000000000 --- a/litellm-rust/crates/cache-redis/src/counter.rs +++ /dev/null @@ -1,205 +0,0 @@ -use std::time::Duration; - -use litellm_cache::{ - BoundedCounterCache, CacheCodec, CountReadCache, CounterCache, Error, ExactCacheContext, - IncrementOperation, -}; - -use crate::{ - cache::{RedisCache, ttl_seconds}, - connection::ConnectionRef, - store::mget, -}; - -const INCREMENT_SCRIPT: &str = concat!( - "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", - "if redis.call('TTL', KEYS[1]) == -1 then ", - "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" -); -const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!( - "local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ", - "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ", - "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", - "return count" -); -const SET_MAX_SCRIPT: &str = concat!( - "local current = redis.call('GET', KEYS[1]); ", - "if current == false or tonumber(current) < tonumber(ARGV[1]) then ", - "redis.call('SET', KEYS[1], ARGV[1]); ", - "if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", - "return ARGV[1]; end; return current" -); - -impl CounterCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn increment_cache( - &self, - key: &str, - amount: f64, - context: ExactCacheContext, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(context.ttl); - self.execute(|connection| increment(connection, key, amount, ttl, false)) - } - - /// Python `_incrbyfloat_with_ttl`: without `refresh_ttl` the TTL is set only on a key that - /// has none, in one atomic script; with it, every increment re-arms the TTL. - async fn async_increment( - &self, - key: &str, - amount: f64, - context: ExactCacheContext, - refresh_ttl: bool, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(context.ttl); - self.run(move |connection| increment(connection, key, amount, ttl, refresh_ttl)) - .await - } - - async fn async_increment_pipeline( - &self, - operations: Vec, - ) -> Result, Error> { - if operations.is_empty() { - return Ok(Vec::new()); - } - let mut pipeline = redis::pipe(); - for operation in operations { - let key = self.namespaced_key(&operation.key); - pipeline.cmd("INCRBYFLOAT").arg(&key).arg(operation.amount); - if let Some(ttl) = operation.ttl { - pipeline - .cmd("EXPIRE") - .arg(key) - .arg(ttl_seconds(ttl)) - .ignore(); - } - } - self.run(move |connection| connection.query_pipeline(&pipeline)) - .await - } -} - -fn increment( - connection: &mut ConnectionRef<'_>, - key: String, - amount: f64, - ttl: u64, - refresh_ttl: bool, -) -> Result { - if !refresh_ttl { - return redis::cmd("EVAL") - .arg(INCREMENT_SCRIPT) - .arg(1) - .arg(key) - .arg(amount) - .arg(ttl) - .query(connection) - .map_err(Error::unavailable); - } - connection - .query_pipeline( - redis::pipe() - .cmd("INCRBYFLOAT") - .arg(&key) - .arg(amount) - .cmd("EXPIRE") - .arg(&key) - .arg(ttl) - .ignore(), - ) - .map(|(value,)| value) -} - -impl CountReadCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn batch_get_counts(&self, keys: &[String]) -> Result>, Error> { - let keys = self.namespaced_keys(keys); - self.execute(|connection| mget(connection, keys))? - .into_iter() - .map(count) - .collect() - } - - async fn async_batch_get_counts(&self, keys: Vec) -> Result>, Error> { - let keys = self.namespaced_keys(&keys); - self.run(move |connection| mget(connection, keys)) - .await? - .into_iter() - .map(count) - .collect() - } -} - -fn count(value: redis::Value) -> Result, Error> { - redis::from_redis_value(value).map_err(Error::invalid_entry) -} - -impl BoundedCounterCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn increment_with_floor(&self, key: &str, amount: i64, ttl: Duration) -> Result { - let key = self.namespaced_key(key); - let ttl = ttl_seconds(ttl); - self.execute(|connection| increment_with_floor(connection, key, amount, ttl)) - } - - async fn async_increment_with_floor( - &self, - key: &str, - amount: i64, - ttl: Duration, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = ttl_seconds(ttl); - self.run(move |connection| increment_with_floor(connection, key, amount, ttl)) - .await - } - - async fn async_set_max( - &self, - key: &str, - value: f64, - ttl: Option, - ) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(ttl); - self.run(move |connection| { - redis::cmd("EVAL") - .arg(SET_MAX_SCRIPT) - .arg(1) - .arg(key) - .arg(value) - .arg(ttl) - .query(connection) - .map_err(Error::unavailable) - }) - .await - } -} - -fn increment_with_floor( - connection: &mut ConnectionRef<'_>, - key: String, - amount: i64, - ttl: u64, -) -> Result { - redis::cmd("EVAL") - .arg(INCREMENT_WITH_FLOOR_SCRIPT) - .arg(1) - .arg(key) - .arg(amount) - .arg(ttl) - .query(connection) - .map_err(Error::unavailable) -} diff --git a/litellm-rust/crates/cache-redis/src/keys.rs b/litellm-rust/crates/cache-redis/src/keys.rs deleted file mode 100644 index 2e8a9e60e65..00000000000 --- a/litellm-rust/crates/cache-redis/src/keys.rs +++ /dev/null @@ -1,63 +0,0 @@ -use std::time::Duration; - -use litellm_cache::{CacheCodec, Error, RefreshTtlCache, ScanCache, TtlCache}; - -use crate::cache::RedisCache; - -impl TtlCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn async_get_ttl(&self, key: &str) -> Result, Error> { - let key = self.namespaced_key(key); - let ttl = self - .run(move |connection| { - redis::cmd("TTL") - .arg(key) - .query::(connection) - .map_err(Error::unavailable) - }) - .await?; - Ok(u64::try_from(ttl).ok().map(Duration::from_secs)) - } -} - -impl RefreshTtlCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn async_refresh_ttl(&self, key: &str, ttl: Option) -> Result { - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(ttl); - self.run(move |connection| { - redis::cmd("EXPIRE") - .arg(key) - .arg(ttl) - .query(connection) - .map_err(Error::unavailable) - }) - .await - } -} - -impl ScanCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, Error> { - let pattern = format!("{}*", self.namespaced_key(pattern)); - self.run(move |connection| { - let mut matches = Vec::new(); - connection.scan(&pattern, count, |_, keys| { - matches.extend(keys); - Ok(matches.len() < count) - })?; - matches.truncate(count); - Ok(matches) - }) - .await - } -} diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs deleted file mode 100644 index 037e39e5d40..00000000000 --- a/litellm-rust/crates/cache-redis/src/lib.rs +++ /dev/null @@ -1,15 +0,0 @@ -mod cache; -mod claim; -pub mod connection; -mod counter; -mod keys; -mod lifecycle; -mod queue; -mod script; -mod store; -mod topology; - -pub use cache::RedisCache; -pub use queue::{RedisLpopOperation, RedisLpopResult, RedisRpushOperation}; -pub use script::{RedisArg, RedisScript}; -pub use topology::{RedisNode, RedisTopology}; diff --git a/litellm-rust/crates/cache-redis/src/lifecycle.rs b/litellm-rust/crates/cache-redis/src/lifecycle.rs deleted file mode 100644 index 54fa8ecc3bf..00000000000 --- a/litellm-rust/crates/cache-redis/src/lifecycle.rs +++ /dev/null @@ -1,86 +0,0 @@ -use litellm_cache::{ - CacheCodec, CacheConnectionResult, CacheConnectionStatus, ClientInfoCache, ConnectionCache, - DisconnectCache, Error, PingCache, -}; - -use crate::{cache::RedisCache, topology::RedisTopology}; - -impl PingCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn sync_ping(&self) -> Result { - self.execute(|connection| connection.ping().map_err(Error::unavailable)) - } - - async fn ping(&self) -> Result { - self.run(|connection| connection.ping().map_err(Error::unavailable)) - .await - } -} - -impl ConnectionCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - /// Python `RedisCache.test_connection`, or `RedisClusterCache.test_connection` for a - /// cluster topology, which differs only in its messages. - async fn test_connection(&self) -> Result { - let label = match self.topology { - RedisTopology::Standalone => "Redis", - RedisTopology::Cluster { .. } => "Redis Cluster", - }; - let ping = self - .run(|connection| Ok(connection.ping().map_err(Error::unavailable))) - .await - .and_then(|ping| ping) - .map_err(|error| litellm_error::Report::new(&error).to_string()); - Ok(match ping { - Ok(true) => CacheConnectionResult { - status: CacheConnectionStatus::Success, - message: format!("{label} connection test successful"), - error: None, - }, - Ok(false) => CacheConnectionResult { - status: CacheConnectionStatus::Failed, - message: format!("{label} ping returned False"), - error: None, - }, - Err(error) => CacheConnectionResult { - status: CacheConnectionStatus::Failed, - message: format!("{label} connection failed: {error}"), - error: Some(error), - }, - }) - } -} - -impl DisconnectCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn disconnect(&self) -> Result<(), Error> { - self.connections.disconnect(); - Ok(()) - } -} - -impl ClientInfoCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type ClientList = String; - type Info = String; - - fn client_list(&self) -> Result { - self.execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST"))) - } - - fn info(&self) -> Result { - self.execute(|connection| connection.node_text(&redis::cmd("INFO"))) - } -} diff --git a/litellm-rust/crates/cache-redis/src/queue.rs b/litellm-rust/crates/cache-redis/src/queue.rs deleted file mode 100644 index d55e4f49187..00000000000 --- a/litellm-rust/crates/cache-redis/src/queue.rs +++ /dev/null @@ -1,226 +0,0 @@ -use std::time::Duration; - -use litellm_cache::{CacheCodec, Error, PopOperation, PushOperation, QueueCache, SetCache}; - -use crate::{cache::RedisCache, script::RedisArg}; - -pub type RedisRpushOperation = PushOperation; -pub type RedisLpopOperation = PopOperation; - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum RedisLpopResult { - Missing, - Value(Vec), - Values(Vec>), -} - -impl SetCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type SetValue = RedisArg; - type SetResult = usize; - - async fn async_set_cache_sadd( - &self, - key: &str, - values: Vec, - ttl: Option, - ) -> Result { - if values.is_empty() { - return Err(Error::InvalidEntry(None)); - } - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(ttl); - let mut pipeline = redis::pipe(); - pipeline - .cmd("SADD") - .arg(&key) - .arg(values) - .cmd("EXPIRE") - .arg(&key) - .arg(ttl) - .ignore(); - self.run(move |connection| connection.query_pipeline(&pipeline)) - .await - .map(|(added,)| added) - } -} - -impl QueueCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type QueueValue = RedisArg; - type PopResult = RedisLpopResult; - - async fn async_rpush(&self, key: &str, values: Vec) -> Result { - if values.is_empty() { - return Err(Error::InvalidEntry(None)); - } - let key = self.namespaced_key(key); - self.run(move |connection| { - redis::cmd("RPUSH") - .arg(key) - .arg(values) - .query(connection) - .map_err(Error::unavailable) - }) - .await - } - - async fn async_rpush_and_trim( - &self, - key: &str, - values: Vec, - max_len: usize, - ) -> Result { - if values.is_empty() { - return Err(Error::InvalidEntry(None)); - } - let key = self.namespaced_key(key); - let start = i64::try_from(max_len).map_or(i64::MIN, |max_len| -max_len); - let mut pipeline = redis::pipe(); - pipeline - .atomic() - .cmd("RPUSH") - .arg(&key) - .arg(values) - .cmd("LTRIM") - .arg(&key) - .arg(start) - .arg(-1) - .ignore(); - self.run(move |connection| connection.query_pipeline(&pipeline)) - .await - .map(|(length,)| length) - } - - async fn async_rpush_pipeline( - &self, - operations: Vec, - ) -> Result, Error> { - if operations.is_empty() { - return Ok(Vec::new()); - } - let mut pipeline = redis::pipe(); - for operation in operations { - if operation.values.is_empty() { - return Err(Error::InvalidEntry(None)); - } - pipeline - .cmd("RPUSH") - .arg(self.namespaced_key(&operation.key)) - .arg(operation.values); - } - self.run(move |connection| connection.query_pipeline(&pipeline)) - .await - } - - async fn async_lpop(&self, key: &str, count: Option) -> Result { - if let Some(count) = count - && self.major_version().await < 7 - { - return self.lpop_one_at_a_time(key, count).await; - } - let command = lpop(self.namespaced_key(key), count); - let value = self - .run(move |connection| { - command - .query::(connection) - .map_err(Error::unavailable) - }) - .await?; - lpop_result(value, count.is_some()) - } - - async fn async_lpop_pipeline( - &self, - operations: Vec, - ) -> Result, Error> { - if operations.is_empty() { - return Ok(Vec::new()); - } - if operations.iter().any(|operation| operation.count.is_some()) - && self.major_version().await < 7 - { - let mut results = Vec::with_capacity(operations.len()); - for operation in &operations { - results.push(self.async_lpop(&operation.key, operation.count).await?); - } - return Ok(results); - } - let multiple = operations - .iter() - .map(|operation| operation.count.is_some()) - .collect::>(); - let mut pipeline = redis::pipe(); - for operation in operations { - pipeline.add_command(lpop(self.namespaced_key(&operation.key), operation.count)); - } - self.run(move |connection| connection.query_pipeline::>(&pipeline)) - .await? - .into_iter() - .zip(multiple) - .map(|(value, multiple)| lpop_result(value, multiple)) - .collect() - } -} - -fn lpop(key: String, count: Option) -> redis::Cmd { - let mut command = redis::cmd("LPOP"); - command.arg(key); - if let Some(count) = count { - command.arg(count); - } - command -} - -fn redis_bytes(value: redis::Value) -> Result, Error> { - match value { - redis::Value::BulkString(bytes) => Ok(bytes), - redis::Value::SimpleString(text) => Ok(text.into_bytes()), - _ => Err(Error::InvalidEntry(None)), - } -} - -impl RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - /// `handle_lpop_count_for_older_redis_versions`: `count` single-`LPOP` pipelines, keeping - /// only the values actually popped. - async fn lpop_one_at_a_time(&self, key: &str, count: usize) -> Result { - let key = self.namespaced_key(key); - let mut values = Vec::new(); - for _ in 0..count { - let mut pipeline = redis::pipe(); - pipeline.add_command(lpop(key.clone(), None)); - let replies = self - .run(move |connection| connection.query_pipeline::>(&pipeline)) - .await?; - for reply in replies { - if reply != redis::Value::Nil { - values.push(redis_bytes(reply)?); - } - } - } - Ok(RedisLpopResult::Values(values)) - } -} - -fn lpop_result(value: redis::Value, multiple: bool) -> Result { - match value { - redis::Value::Nil => Ok(RedisLpopResult::Missing), - redis::Value::Array(values) if multiple => values - .into_iter() - .map(redis_bytes) - .collect::, _>>() - .map(RedisLpopResult::Values), - value if !multiple => redis_bytes(value).map(RedisLpopResult::Value), - _ => Err(Error::InvalidEntry(None)), - } -} diff --git a/litellm-rust/crates/cache-redis/src/script.rs b/litellm-rust/crates/cache-redis/src/script.rs deleted file mode 100644 index 09987437ed4..00000000000 --- a/litellm-rust/crates/cache-redis/src/script.rs +++ /dev/null @@ -1,136 +0,0 @@ -use std::sync::Arc; - -use litellm_cache::{CacheCodec, CacheScript, Error, ScriptCache}; - -use crate::{ - cache::{RedisCache, namespaced_key}, - connection::Connections, -}; - -#[derive(Clone, Debug, PartialEq)] -pub enum RedisArg { - Bytes(Vec), - Integer(i64), - Float(f64), -} - -impl From<&str> for RedisArg { - fn from(value: &str) -> Self { - Self::Bytes(value.as_bytes().to_vec()) - } -} - -impl From for RedisArg { - fn from(value: String) -> Self { - Self::Bytes(value.into_bytes()) - } -} - -impl From> for RedisArg { - fn from(value: Vec) -> Self { - Self::Bytes(value) - } -} - -impl From for RedisArg { - fn from(value: i64) -> Self { - Self::Integer(value) - } -} - -impl From for RedisArg { - fn from(value: f64) -> Self { - Self::Float(value) - } -} - -impl redis::ToRedisArgs for RedisArg { - fn write_redis_args(&self, out: &mut W) - where - W: ?Sized + redis::RedisWrite, - { - match self { - Self::Bytes(value) => value.write_redis_args(out), - Self::Integer(value) => value.write_redis_args(out), - Self::Float(value) => value.write_redis_args(out), - } - } -} - -pub struct RedisScript { - connections: Arc>, - namespace: Option, - source: String, -} - -impl CacheScript for RedisScript -where - C: redis::ConnectionLike + Send + 'static, -{ - type Argument = RedisArg; - type Output = redis::Value; - - async fn invoke( - &self, - keys: Vec, - arguments: Vec, - ) -> Result { - let keys = keys - .into_iter() - .map(|key| namespaced_key(self.namespace.as_deref(), &key)) - .collect::>(); - let source = self.source.clone(); - Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - eval(connection, &source, keys, arguments) - }) - .await - } -} - -impl RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - pub async fn async_eval( - &self, - script: String, - keys: Vec, - arguments: Vec, - ) -> Result { - let keys = self.namespaced_keys(&keys); - self.run(move |connection| eval(connection, &script, keys, arguments)) - .await - } -} - -fn eval( - connection: &mut impl redis::ConnectionLike, - script: &str, - keys: Vec, - arguments: Vec, -) -> Result { - redis::cmd("EVAL") - .arg(script) - .arg(keys.len()) - .arg(keys) - .arg(arguments) - .query(connection) - .map_err(Error::unavailable) -} - -impl ScriptCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type Script = RedisScript; - - fn async_register_script(&self, source: String) -> Self::Script { - RedisScript { - connections: Arc::clone(&self.connections), - namespace: self.namespace.clone(), - source, - } - } -} diff --git a/litellm-rust/crates/cache-redis/src/store.rs b/litellm-rust/crates/cache-redis/src/store.rs deleted file mode 100644 index daadc940ca9..00000000000 --- a/litellm-rust/crates/cache-redis/src/store.rs +++ /dev/null @@ -1,232 +0,0 @@ -use std::time::Duration; - -use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheCodec, DeleteCache, Error, - ExactCacheContext, FlushAllCache, FlushCache, TtlPipelineCache, -}; -use redis::Commands; - -use crate::{cache::RedisCache, connection::ConnectionRef}; - -impl BaseCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type Value = S::Value; - type Context = ExactCacheContext; - - fn get_ttl(&self, context: &Self::Context) -> Option { - context.ttl.or(Some(self.default_ttl)) - } - - fn set_cache( - &self, - key: &str, - value: Self::Value, - context: &ExactCacheContext, - ) -> Result<(), Error> { - let payload = self.codec.encode(&value)?; - let ttl = self.ttl_or_default(context.ttl); - let key = self.namespaced_key(key); - self.execute(|connection| { - connection - .set_ex::<_, _, ()>(key, payload, ttl) - .map_err(Error::unavailable) - }) - } - - fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, Error> { - let key = self.namespaced_key(key); - let value = self.execute(|connection| { - connection - .get::<_, redis::Value>(key) - .map_err(Error::unavailable) - })?; - self.decode_response(value) - } - - async fn async_set_cache( - &self, - key: &str, - value: Self::Value, - context: ExactCacheContext, - ) -> Result<(), Error> { - let payload = self.codec.encode(&value)?; - let key = self.namespaced_key(key); - let ttl = self.ttl_or_default(context.ttl); - self.run(move |connection| { - connection - .set_ex::<_, _, ()>(key, payload, ttl) - .map_err(Error::unavailable) - }) - .await - } - - async fn async_get_cache( - &self, - key: &str, - _: &ExactCacheContext, - ) -> Result, Error> { - let key = self.namespaced_key(key); - let value = self - .run(move |connection| { - connection - .get::<_, redis::Value>(key) - .map_err(Error::unavailable) - }) - .await?; - self.decode_response(value) - } - - async fn async_set_cache_pipeline( - &self, - cache_list: Vec<(String, Self::Value)>, - context: ExactCacheContext, - ) -> Result<(), Error> { - self.async_set_cache_pipeline_with_ttls( - cache_list - .into_iter() - .map(|(key, value)| (key, value, context.ttl)) - .collect(), - ) - .await - } -} - -impl TtlPipelineCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn async_set_cache_pipeline_with_ttls( - &self, - entries: Vec<(String, Self::Value, Option)>, - ) -> Result<(), Error> { - if entries.is_empty() { - return Ok(()); - } - let mut pipeline = redis::pipe(); - for (key, value, ttl) in entries { - pipeline - .cmd("SETEX") - .arg(self.namespaced_key(&key)) - .arg(self.ttl_or_default(ttl)) - .arg(self.codec.encode(&value)?) - .ignore(); - } - self.run(move |connection| connection.query_pipeline(&pipeline)) - .await - } -} - -impl BatchCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn batch_get_cache( - &self, - keys: &[String], - _: &ExactCacheContext, - ) -> Result>, Error> { - let keys = self.namespaced_keys(keys); - self.execute(|connection| mget(connection, keys))? - .into_iter() - .map(|value| self.decode_batch_response(value)) - .collect() - } - - async fn async_batch_get_cache( - &self, - keys: Vec, - _: ExactCacheContext, - ) -> Result>, Error> { - let keys = self.namespaced_keys(&keys); - self.run(move |connection| mget(connection, keys)) - .await? - .into_iter() - .map(|value| self.decode_batch_response(value)) - .collect() - } -} - -pub(crate) fn mget( - connection: &mut ConnectionRef<'_>, - keys: Vec, -) -> Result, Error> { - redis::cmd("MGET") - .arg(keys) - .query(connection) - .map_err(Error::unavailable) -} - -impl DeleteCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn delete_cache(&self, key: &str) -> Result<(), Error> { - let key = self.namespaced_key(key); - self.execute(|connection| connection.del::<_, ()>(key).map_err(Error::unavailable)) - } - - async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { - let key = self.namespaced_key(key); - self.run(move |connection| connection.del::<_, ()>(key).map_err(Error::unavailable)) - .await - } -} - -impl BulkDeleteCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn delete_cache_keys(&self, keys: Vec) -> Result { - if keys.is_empty() { - return Ok(0); - } - let keys = self.namespaced_keys(&keys); - self.run(move |connection| connection.del(keys).map_err(Error::unavailable)) - .await - } -} - -impl FlushCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn flush_cache(&self) -> Result<(), Error> { - let pattern = self.namespaced_pattern()?; - self.execute(|connection| flush_matching(connection, &pattern)) - } - - async fn async_flush_cache(&self) -> Result<(), Error> { - let pattern = self.namespaced_pattern()?; - self.run(move |connection| flush_matching(connection, &pattern)) - .await - } -} - -impl FlushAllCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn flushall(&self) -> Result<(), Error> { - self.execute(|connection| connection.flushall()) - } -} - -fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> { - connection.scan(pattern, 1000, |connection, keys| { - if !keys.is_empty() { - connection - .del::<_, usize>(keys) - .map_err(Error::unavailable)?; - } - Ok(true) - }) -} diff --git a/litellm-rust/crates/cache-redis/src/topology.rs b/litellm-rust/crates/cache-redis/src/topology.rs deleted file mode 100644 index 7f4ee48b222..00000000000 --- a/litellm-rust/crates/cache-redis/src/topology.rs +++ /dev/null @@ -1,14 +0,0 @@ -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RedisNode { - pub host: String, - pub port: u16, -} - -#[derive(Clone, Debug, Default, PartialEq, Eq)] -pub enum RedisTopology { - #[default] - Standalone, - Cluster { - startup_nodes: Vec, - }, -} diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs deleted file mode 100644 index 9ec85251733..00000000000 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ /dev/null @@ -1,1164 +0,0 @@ -mod support; - -use std::time::Duration; - -use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, BoundedCounterCache, BulkDeleteCache, CacheCodec, - CacheConnectionStatus, CacheScript, ClaimCache, ClientInfoCache, ConnectionCache, - CountReadCache, CounterCache, DeleteCache, DisconnectCache, Error, ExactCacheContext, - FlushAllCache, FlushCache, IncrementOperation, JsonCodec, PingCache, QueueCache, - RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache, TtlPipelineCache, get_cache, - set_cache, -}; -use litellm_cache_redis::{ - RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, -}; -use redis_test::{MockCmd, MockRedisConnection}; -use rstest::{fixture, rstest}; -use serde_json::json; -use support::TaggedByteCodec; - -type Mocked = RedisCache; - -fn mock(commands: Vec) -> MockRedisConnection { - MockRedisConnection::new(commands).assert_all_commands_consumed() -} - -fn tagged(commands: Vec) -> Mocked { - RedisCache::with_connection(mock(commands), None, TaggedByteCodec(42)) -} - -fn json_cache(commands: Vec) -> Mocked> { - RedisCache::with_connection(mock(commands), None, JsonCodec::new()) -} - -fn team(commands: Vec) -> Mocked> { - json_cache(commands).with_namespace(Some("team".into())) -} - -#[fixture] -fn context() -> ExactCacheContext { - ExactCacheContext::default() -} - -fn scan(pattern: &str, cursor: u64, count: usize, reply: redis::Value) -> MockCmd { - MockCmd::new( - redis::cmd("SCAN") - .cursor_arg(cursor) - .arg("MATCH") - .arg(pattern) - .arg("COUNT") - .arg(count), - Ok(reply), - ) -} - -#[rstest] -fn constructor_rejects_invalid_urls() { - assert!(RedisCache::new("not a redis url", None, JsonCodec::::new()).is_err()); -} - -#[rstest] -#[case::zero_rounds_up_to_one(Some(Duration::ZERO), 1)] -#[case::fractions_round_up(Some(Duration::from_millis(1500)), 2)] -#[case::whole_seconds_are_kept(Some(Duration::from_secs(15)), 15)] -#[case::missing_ttl_uses_default(None, 600)] -fn writes_round_ttls_up_to_positive_seconds(#[case] ttl: Option, #[case] seconds: u64) { - let cache = tagged(vec![MockCmd::new( - redis::cmd("SETEX") - .arg("key") - .arg(seconds) - .arg([42u8, 7].as_slice()), - Ok("OK"), - )]); - cache - .set_cache("key", 7, &ExactCacheContext { ttl }) - .unwrap(); -} - -#[rstest] -fn generic_helpers_use_the_injected_codec_and_ttl() { - let cache = tagged(vec![ - MockCmd::new( - redis::cmd("SETEX") - .arg("counter") - .arg(2) - .arg([42u8, 7].as_slice()), - Ok("OK"), - ), - MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])), - ]); - let context = ExactCacheContext { - ttl: Some(Duration::from_millis(1500)), - }; - set_cache(&cache, "counter", 7, &context).unwrap(); - assert_eq!(get_cache(&cache, "counter", &context).unwrap(), Some(7)); -} - -#[rstest] -fn commands_round_trip_entries_and_delete_only_namespaced_keys(context: ExactCacheContext) { - let value = json!({"deployment": "model-a", "cooldown_seconds": 30}); - let payload = JsonCodec::::new() - .encode(&value) - .unwrap(); - let cache = json_cache(vec![ - MockCmd::new( - redis::cmd("SETEX") - .arg("litellm-cache:key") - .arg(600) - .arg(payload.clone()), - Ok("OK"), - ), - MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)), - MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)), - ]) - .with_namespace(Some("litellm-cache".into())); - - cache.set_cache("key", value.clone(), &context).unwrap(); - assert_eq!(cache.get_cache("key", &context).unwrap(), Some(value)); - cache.delete_cache("key").unwrap(); -} - -#[rstest] -#[tokio::test] -async fn async_operations_preserve_codec_ttl_and_missing_values(context: ExactCacheContext) { - let cache = RedisCache::with_connection( - mock(vec![ - MockCmd::new( - redis::cmd("SETEX") - .arg("counter") - .arg(9) - .arg([42u8, 7].as_slice()), - Ok("OK"), - ), - MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])), - MockCmd::new( - redis::cmd("SETEX") - .arg("batch") - .arg(2) - .arg([42u8, 8].as_slice()), - Ok("OK"), - ), - MockCmd::new(redis::cmd("DEL").arg("counter"), Ok(1u32)), - MockCmd::new(redis::cmd("GET").arg("counter"), Ok(redis::Value::Nil)), - ]), - Some(Duration::from_secs(9)), - TaggedByteCodec(42), - ); - cache - .batch_cache_write("counter", 7, context.clone()) - .await - .unwrap(); - assert_eq!( - cache.async_get_cache("counter", &context).await.unwrap(), - Some(7) - ); - cache - .async_set_cache_pipeline( - vec![("batch".into(), 8)], - ExactCacheContext { - ttl: Some(Duration::from_millis(1500)), - }, - ) - .await - .unwrap(); - cache.async_delete_cache("counter").await.unwrap(); - assert_eq!( - cache.async_get_cache("counter", &context).await.unwrap(), - None - ); -} - -#[rstest] -#[tokio::test] -async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCacheContext) { - let cache = tagged(vec![ - MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])), - MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])), - ]); - assert_eq!( - cache.set_cache("invalid", 255, &context), - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache.async_set_cache("invalid", 255, context.clone()).await, - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache - .async_set_cache_pipeline( - vec![("valid".into(), 7), ("invalid".into(), 255)], - context.clone(), - ) - .await, - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache - .async_set_cache_pipeline_with_ttls(vec![ - ("valid".into(), 7, None), - ("invalid".into(), 255, None), - ]) - .await, - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache.get_cache("invalid", &context), - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache.async_get_cache("invalid", &context).await, - Err(Error::InvalidEntry(None)) - ); -} - -#[rstest] -#[case::bare_key("key")] -#[case::already_prefixed("team:key")] -fn namespaces_are_added_once(#[case] key: &str, context: ExactCacheContext) { - let cache = team(vec![MockCmd::new( - redis::cmd("GET").arg("team:key"), - Ok(redis::Value::Nil), - )]); - assert_eq!(cache.get_cache(key, &context).unwrap(), None); -} - -#[rstest] -#[case::empty(Some(String::new()))] -#[case::missing(None)] -fn empty_namespaces_leave_keys_unprefixed( - #[case] namespace: Option, - context: ExactCacheContext, -) { - let cache = json_cache(vec![MockCmd::new( - redis::cmd("GET").arg("key"), - Ok(redis::Value::Nil), - )]) - .with_namespace(namespace); - assert_eq!(cache.namespace(), None); - assert_eq!(cache.get_cache("key", &context).unwrap(), None); -} - -#[rstest] -#[tokio::test] -async fn flush_requires_a_namespace() { - let unscoped = json_cache(Vec::new()); - assert_eq!(unscoped.flush_cache(), Err(Error::UnscopedFlush)); - assert_eq!( - unscoped.async_flush_cache().await, - Err(Error::UnscopedFlush) - ); -} - -#[rstest] -#[case::plain_namespace("litellm-cache", "litellm-cache:*", "litellm-cache:key")] -#[case::glob_metacharacters_are_escaped("team*", "team\\*:*", "team*:key")] -fn flush_scans_and_deletes_only_namespaced_keys( - #[case] namespace: &str, - #[case] pattern: &str, - #[case] key: &str, -) { - let cache = json_cache(vec![ - scan(pattern, 0, 1000, redis_test::redis_value!(["0", [key]])), - MockCmd::new(redis::cmd("DEL").arg(key), Ok(1u32)), - ]) - .with_namespace(Some(namespace.into())); - cache.flush_cache().unwrap(); -} - -#[rstest] -#[tokio::test] -async fn async_flush_deletes_each_scan_page_separately() { - let cache = team(vec![ - scan( - "team:*", - 0, - 1000, - redis_test::redis_value!(["7", ["team:a", "team:b"]]), - ), - MockCmd::new(redis::cmd("DEL").arg("team:a").arg("team:b"), Ok(2u32)), - scan( - "team:*", - 7, - 1000, - redis_test::redis_value!(["0", ["team:c"]]), - ), - MockCmd::new(redis::cmd("DEL").arg("team:c"), Ok(1u32)), - ]); - cache.async_flush_cache().await.unwrap(); -} - -#[rstest] -fn flushall_ignores_the_namespace() { - team(vec![MockCmd::new(redis::cmd("FLUSHALL"), Ok("OK"))]) - .flushall() - .unwrap(); -} - -#[rstest] -#[tokio::test] -async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries( - context: ExactCacheContext, -) { - let cache = tagged(vec![MockCmd::new( - redis::cmd("MGET").arg("hit").arg("miss").arg("invalid"), - Ok(vec![ - redis::Value::BulkString(vec![42, 7]), - redis::Value::Nil, - redis::Value::BulkString(vec![99, 7]), - ]), - )]); - - assert_eq!( - cache - .async_batch_get_cache(vec!["hit".into(), "miss".into(), "invalid".into()], context) - .await - .unwrap(), - vec![BatchEntry::Hit(7), BatchEntry::Miss, BatchEntry::Invalid] - ); -} - -#[rstest] -#[tokio::test] -async fn ttl_pipeline_keeps_each_entry_ttl_and_defaults_missing_ones() { - let mut pipeline = redis::pipe(); - pipeline - .cmd("SETEX") - .arg("ns:team_id:t1") - .arg(60u64) - .arg(r#"{"team_id":"t1"}"#) - .cmd("SETEX") - .arg("ns:u1") - .arg(7u64) - .arg(r#"{"user_id":"u1"}"#) - .cmd("SETEX") - .arg("ns:org_id:o1") - .arg(300u64) - .arg(r#"{"a":1}"#); - let cache = RedisCache::with_connection( - mock(vec![MockCmd::with_values( - pipeline, - Ok(vec!["OK", "OK", "OK"]), - )]), - Some(Duration::from_secs(300)), - JsonCodec::::new(), - ) - .with_namespace(Some("ns".into())); - - cache - .async_set_cache_pipeline_with_ttls(vec![ - ( - "team_id:t1".into(), - json!({"team_id": "t1"}), - Some(Duration::from_secs(60)), - ), - ( - "u1".into(), - json!({"user_id": "u1"}), - Some(Duration::from_secs(7)), - ), - ("org_id:o1".into(), json!({"a": 1}), None), - ]) - .await - .unwrap(); -} - -#[rstest] -#[tokio::test] -async fn empty_pipelines_skip_the_round_trip(context: ExactCacheContext) { - let cache = json_cache(Vec::new()); - cache - .async_set_cache_pipeline(Vec::new(), context) - .await - .unwrap(); - cache - .async_set_cache_pipeline_with_ttls(Vec::new()) - .await - .unwrap(); - assert_eq!(cache.delete_cache_keys(Vec::new()).await.unwrap(), 0); - assert_eq!( - cache.async_rpush_pipeline(Vec::new()).await.unwrap(), - Vec::::new() - ); - assert_eq!( - cache.async_lpop_pipeline(Vec::new()).await.unwrap(), - Vec::::new() - ); - assert_eq!( - cache.async_increment_pipeline(Vec::new()).await.unwrap(), - Vec::::new() - ); -} - -#[rstest] -#[tokio::test] -async fn count_reads_parse_integers_and_keep_missing_counters() { - let mget = || { - MockCmd::new( - redis::cmd("MGET").arg("team:count").arg("team:missing"), - Ok(redis_test::redis_value!(["7", nil])), - ) - }; - let cache = team(vec![mget(), mget()]); - let keys = vec!["count".to_string(), "missing".to_string()]; - - assert_eq!(cache.batch_get_counts(&keys).unwrap(), [Some(7), None]); - assert_eq!( - cache.async_batch_get_counts(keys).await.unwrap(), - [Some(7), None] - ); -} - -#[rstest] -#[tokio::test] -async fn pings_run_on_both_paths() { - let cache = team(vec![ - MockCmd::new(redis::cmd("PING"), Ok("PONG")), - MockCmd::new(redis::cmd("PING"), Ok("PONG")), - ]); - assert!(cache.sync_ping().unwrap()); - assert!(cache.ping().await.unwrap()); -} - -#[rstest] -#[case::remaining(12, Some(Duration::from_secs(12)))] -#[case::no_expiry(-1, None)] -#[case::missing(-2, None)] -#[tokio::test] -async fn ttl_reads_hide_negative_replies(#[case] reply: i64, #[case] ttl: Option) { - let cache = team(vec![MockCmd::new( - redis::cmd("TTL").arg("team:key"), - Ok(reply), - )]); - assert_eq!(cache.async_get_ttl("key").await.unwrap(), ttl); -} - -#[rstest] -#[case::explicit_ttl(Some(Duration::from_secs(30)), 30, 1, true)] -#[case::default_ttl(None, 600, 1, true)] -#[case::missing_key(Some(Duration::from_secs(30)), 30, 0, false)] -#[tokio::test] -async fn refresh_ttl_expires_existing_keys_only( - #[case] ttl: Option, - #[case] seconds: u64, - #[case] reply: i64, - #[case] refreshed: bool, -) { - let cache = team(vec![MockCmd::new( - redis::cmd("EXPIRE").arg("team:key").arg(seconds), - Ok(reply), - )]); - assert_eq!( - cache.async_refresh_ttl("key", ttl).await.unwrap(), - refreshed - ); -} - -#[rstest] -#[tokio::test] -async fn scan_stops_at_count_and_bulk_delete_reports_existing_keys() { - let cache = team(vec![ - scan( - "team:job-*", - 0, - 25, - redis_test::redis_value!(["4", ["team:job-a"]]), - ), - scan( - "team:job-*", - 4, - 25, - redis_test::redis_value!(["0", ["team:job-b"]]), - ), - MockCmd::new( - redis::cmd("DEL").arg("team:job-a").arg("team:job-b"), - Ok(2u32), - ), - ]); - assert_eq!( - cache.async_scan_iter("job-", 25).await.unwrap(), - ["team:job-a", "team:job-b"] - ); - assert_eq!( - cache - .delete_cache_keys(vec!["job-a".into(), "job-b".into()]) - .await - .unwrap(), - 2 - ); -} - -#[rstest] -#[tokio::test] -async fn sets_add_members_and_arm_the_default_ttl() { - let mut pipeline = redis::pipe(); - pipeline - .cmd("SADD") - .arg("team:members") - .arg("a") - .arg("b") - .cmd("EXPIRE") - .arg("team:members") - .arg(600u64) - .ignore(); - let cache = team(vec![MockCmd::with_values( - pipeline, - Ok(vec![redis::Value::Int(2), redis::Value::Int(1)]), - )]); - assert_eq!( - cache - .async_set_cache_sadd("members", vec!["a".into(), "b".into()], None) - .await - .unwrap(), - 2 - ); - assert_eq!( - cache - .async_set_cache_sadd("members", Vec::new(), None) - .await, - Err(Error::InvalidEntry(None)) - ); -} - -#[rstest] -#[tokio::test] -async fn queues_push_and_pop_namespaced_lists() { - let cache = team(vec![ - MockCmd::new( - redis::cmd("RPUSH").arg("team:queue").arg("a").arg("b"), - Ok(2u32), - ), - MockCmd::new( - redis::cmd("INFO"), - Ok("# Server\r\nredis_version:7.2.4\r\n"), - ), - MockCmd::new( - redis::cmd("LPOP").arg("team:queue").arg(2usize), - Ok(redis_test::redis_value!(["a", "b"])), - ), - MockCmd::new(redis::cmd("LPOP").arg("team:queue"), Ok("c")), - ]); - assert_eq!( - cache - .async_rpush("queue", vec!["a".into(), "b".into()]) - .await - .unwrap(), - 2 - ); - assert_eq!( - cache.async_rpush("queue", Vec::new()).await, - Err(Error::InvalidEntry(None)) - ); - assert_eq!( - cache.async_lpop("queue", Some(2)).await.unwrap(), - RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()]) - ); - assert_eq!( - cache.async_lpop("queue", None).await.unwrap(), - RedisLpopResult::Value(b"c".to_vec()) - ); -} - -/// Python `RedisCache.async_lpop` checks `redis_version` from `INFO` and, below major version 7, -/// pops a counted batch as `count` single-command `LPOP` pipelines, dropping `None` replies. -#[rstest] -#[tokio::test] -async fn counted_lpop_falls_back_to_single_pops_below_redis_7() { - let single_pop = || { - let mut pipeline = redis::pipe(); - pipeline.cmd("LPOP").arg("team:queue"); - pipeline - }; - let cache = team(vec![ - MockCmd::new( - redis::cmd("INFO"), - Ok("# Server\r\nredis_version:6.2.14\r\n"), - ), - MockCmd::with_values(single_pop(), Ok(vec![redis_test::redis_value!("a")])), - MockCmd::with_values(single_pop(), Ok(vec![redis_test::redis_value!("b")])), - MockCmd::with_values(single_pop(), Ok(vec![redis::Value::Nil])), - ]); - - assert_eq!( - cache.async_lpop("queue", Some(3)).await.unwrap(), - RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()]) - ); -} - -/// Python keeps `redis_version = "Unknown"` when `INFO` fails and then assumes -/// `DEFAULT_REDIS_MAJOR_VERSION` (7), so a counted pop is one `LPOP key count`. The version is read -/// once: the second pop sends no second `INFO`. -#[rstest] -#[tokio::test] -async fn counted_lpop_assumes_redis_7_when_info_fails() { - let cache = team(vec![ - MockCmd::new( - redis::cmd("INFO"), - Err::(redis::RedisError::from((redis::ErrorKind::Io, "down"))), - ), - MockCmd::new( - redis::cmd("LPOP").arg("team:queue").arg(2usize), - Ok(redis_test::redis_value!(["a", "b"])), - ), - MockCmd::new( - redis::cmd("LPOP").arg("team:queue").arg(1usize), - Ok(redis_test::redis_value!(["c"])), - ), - ]); - - assert_eq!( - cache.async_lpop("queue", Some(2)).await.unwrap(), - RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()]) - ); - assert_eq!( - cache.async_lpop("queue", Some(1)).await.unwrap(), - RedisLpopResult::Values(vec![b"c".to_vec()]) - ); -} - -fn push_and_trim(start: i64) -> redis::Pipeline { - let mut pipeline = redis::pipe(); - pipeline - .atomic() - .cmd("RPUSH") - .arg("ns:buf") - .arg("c") - .arg("d") - .cmd("LTRIM") - .arg("ns:buf") - .arg(start) - .arg(-1); - pipeline -} - -#[rstest] -#[case::keeps_newest_entries(3, -3)] -#[case::zero_keeps_everything(0, 0)] -#[tokio::test] -async fn rpush_and_trim_runs_push_and_trim_in_one_transaction( - #[case] max_len: usize, - #[case] start: i64, -) { - let cache = json_cache(vec![MockCmd::with_values( - push_and_trim(start), - Ok(vec![redis::Value::Array(vec![ - redis::Value::Int(4), - redis::Value::Okay, - ])]), - )]) - .with_namespace(Some("ns".into())); - - assert_eq!( - cache - .async_rpush_and_trim("buf", vec!["c".into(), "d".into()], max_len) - .await - .unwrap(), - 4 - ); -} - -#[rstest] -#[tokio::test] -async fn rpush_and_trim_raises_when_a_queued_command_fails() { - let wrong_type = redis::parse_redis_value( - b"-WRONGTYPE Operation against a key holding the wrong kind of value\r\n", - ) - .unwrap(); - let cache = json_cache(vec![MockCmd::with_values( - push_and_trim(-3), - Ok(vec![redis::Value::Array(vec![ - wrong_type, - redis::Value::Okay, - ])]), - )]) - .with_namespace(Some("ns".into())); - - assert!(matches!( - cache - .async_rpush_and_trim("buf", vec!["c".into(), "d".into()], 3) - .await, - Err(Error::Unavailable(Some(_))) - )); - assert_eq!( - cache.async_rpush_and_trim("buf", Vec::new(), 3).await, - Err(Error::InvalidEntry(None)) - ); -} - -#[rstest] -#[tokio::test] -async fn scripts_and_eval_namespace_their_keys() { - let eval = || { - MockCmd::new( - redis::cmd("EVAL") - .arg("return KEYS[1]") - .arg(1usize) - .arg("team:key"), - Ok("team:key"), - ) - }; - let cache = team(vec![eval(), eval()]); - assert_eq!( - cache - .async_eval("return KEYS[1]".into(), vec!["key".into()], Vec::new()) - .await - .unwrap(), - redis::Value::BulkString(b"team:key".to_vec()) - ); - assert_eq!( - cache - .async_register_script("return KEYS[1]".into()) - .invoke(vec!["key".into()], Vec::new()) - .await - .unwrap(), - redis::Value::BulkString(b"team:key".to_vec()) - ); -} - -#[rstest] -fn client_list_and_info_return_server_text() { - let cache = team(vec![ - MockCmd::new(redis::cmd("CLIENT").arg("LIST"), Ok("id=1")), - MockCmd::new(redis::cmd("INFO"), Ok("redis_version:7")), - ]); - assert_eq!(cache.client_list().unwrap(), "id=1"); - assert_eq!(cache.info().unwrap(), "redis_version:7"); -} - -#[rstest] -#[tokio::test] -async fn pipelines_preserve_operation_order() { - let mut rpush_pipeline = redis::pipe(); - rpush_pipeline - .cmd("RPUSH") - .arg("team:a") - .arg("one") - .cmd("RPUSH") - .arg("team:b") - .arg("two"); - let mut lpop_pipeline = redis::pipe(); - lpop_pipeline - .cmd("LPOP") - .arg("team:a") - .arg(2usize) - .cmd("LPOP") - .arg("team:b"); - let queue = team(vec![ - MockCmd::with_values( - rpush_pipeline, - Ok(vec![redis::Value::Int(1), redis::Value::Int(2)]), - ), - MockCmd::new( - redis::cmd("INFO"), - Ok("# Server\r\nredis_version:7.2.4\r\n"), - ), - MockCmd::with_values( - lpop_pipeline, - Ok(vec![redis_test::redis_value!(["one"]), redis::Value::Nil]), - ), - ]); - - assert_eq!( - queue - .async_rpush_pipeline(vec![ - RedisRpushOperation { - key: "a".into(), - values: vec![RedisArg::from("one")], - }, - RedisRpushOperation { - key: "b".into(), - values: vec![RedisArg::from("two")], - }, - ]) - .await - .unwrap(), - [1, 2] - ); - assert_eq!( - queue - .async_lpop_pipeline(vec![ - RedisLpopOperation { - key: "a".into(), - count: Some(2), - }, - RedisLpopOperation { - key: "b".into(), - count: None, - }, - ]) - .await - .unwrap(), - [ - RedisLpopResult::Values(vec![b"one".to_vec()]), - RedisLpopResult::Missing, - ] - ); -} - -/// Below major version 7 a counted `LPOP` is unsupported, so a pipeline that mixes counted and -/// plain pops runs each operation through `async_lpop`: `count` single-`LPOP` pipelines for the -/// counted ones, a bare `LPOP` for the rest. No `LPOP key count` reaches the connection. -#[rstest] -#[tokio::test] -async fn lpop_pipeline_pops_one_at_a_time_below_redis_7() { - let single_pop = |key: &str| { - let mut pipeline = redis::pipe(); - pipeline.cmd("LPOP").arg(key); - pipeline - }; - let cache = team(vec![ - MockCmd::new( - redis::cmd("INFO"), - Ok("# Server\r\nredis_version:6.2.14\r\n"), - ), - MockCmd::with_values( - single_pop("team:a"), - Ok(vec![redis_test::redis_value!("one")]), - ), - MockCmd::with_values( - single_pop("team:a"), - Ok(vec![redis_test::redis_value!("two")]), - ), - MockCmd::new(redis::cmd("LPOP").arg("team:b"), Ok(redis::Value::Nil)), - ]); - - assert_eq!( - cache - .async_lpop_pipeline(vec![ - RedisLpopOperation { - key: "a".into(), - count: Some(2), - }, - RedisLpopOperation { - key: "b".into(), - count: None, - }, - ]) - .await - .unwrap(), - [ - RedisLpopResult::Values(vec![b"one".to_vec(), b"two".to_vec()]), - RedisLpopResult::Missing, - ] - ); -} - -#[rstest] -#[tokio::test] -async fn increment_pipeline_expires_only_operations_with_a_ttl() { - let mut increment_pipeline = redis::pipe(); - increment_pipeline - .cmd("INCRBYFLOAT") - .arg("team:counter") - .arg(1.5f64) - .cmd("EXPIRE") - .arg("team:counter") - .arg(10u64) - .ignore() - .cmd("INCRBYFLOAT") - .arg("team:counter") - .arg(2.0f64); - let counters = team(vec![MockCmd::with_values( - increment_pipeline, - Ok(vec![ - redis::Value::BulkString(b"1.5".to_vec()), - redis::Value::Int(1), - redis::Value::BulkString(b"3.5".to_vec()), - ]), - )]); - assert_eq!( - counters - .async_increment_pipeline(vec![ - IncrementOperation { - key: "counter".into(), - amount: 1.5, - ttl: Some(Duration::from_secs(10)), - }, - IncrementOperation { - key: "counter".into(), - amount: 2.0, - ttl: None, - }, - ]) - .await - .unwrap(), - [1.5, 3.5] - ); -} - -const INCREMENT_SCRIPT: &str = concat!( - "local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ", - "if redis.call('TTL', KEYS[1]) == -1 then ", - "redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value" -); -const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!( - "local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ", - "if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ", - "if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", - "return count" -); -const SET_MAX_SCRIPT: &str = concat!( - "local current = redis.call('GET', KEYS[1]); ", - "if current == false or tonumber(current) < tonumber(ARGV[1]) then ", - "redis.call('SET', KEYS[1], ARGV[1]); ", - "if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ", - "return ARGV[1]; end; return current" -); - -fn increment_script(amount: f64) -> MockCmd { - MockCmd::new( - redis::cmd("EVAL") - .arg(INCREMENT_SCRIPT) - .arg(1) - .arg("counter") - .arg(amount) - .arg(600), - Ok("4.5"), - ) -} - -#[rstest] -#[tokio::test] -async fn increments_keep_an_existing_ttl_in_one_atomic_script(context: ExactCacheContext) { - let cache = RedisCache::with_connection( - mock(vec![increment_script(2.5), increment_script(2.5)]), - None, - JsonCodec::::new(), - ); - - assert_eq!( - cache - .increment_cache("counter", 2.5, context.clone()) - .unwrap(), - 4.5 - ); - assert_eq!( - cache - .async_increment("counter", 2.5, context, false) - .await - .unwrap(), - 4.5 - ); -} - -#[rstest] -#[case::explicit_ttl(Some(Duration::from_secs(60)), 60u64)] -#[case::default_ttl(None, 600u64)] -#[tokio::test] -async fn refresh_ttl_increments_rearm_the_ttl_in_the_same_round_trip( - #[case] ttl: Option, - #[case] seconds: u64, -) { - let mut pipeline = redis::pipe(); - pipeline - .cmd("INCRBYFLOAT") - .arg("ns:spend:key:k") - .arg(1.5f64) - .cmd("EXPIRE") - .arg("ns:spend:key:k") - .arg(seconds); - let cache = json_cache(vec![MockCmd::with_values( - pipeline, - Ok(vec![ - redis::Value::BulkString(b"1.5".to_vec()), - redis::Value::Int(1), - ]), - )]) - .with_namespace(Some("ns".into())); - - assert_eq!( - cache - .async_increment("spend:key:k", 1.5, ExactCacheContext { ttl }, true) - .await - .unwrap(), - 1.5 - ); -} - -#[rstest] -#[tokio::test] -async fn counter_repairs_are_atomic_and_use_default_ttl() { - let floor = || { - MockCmd::new( - redis::cmd("EVAL") - .arg(INCREMENT_WITH_FLOOR_SCRIPT) - .arg(1) - .arg("team:counter") - .arg(-2i64) - .arg(30u64), - Ok(0i64), - ) - }; - let cache = team(vec![ - floor(), - floor(), - MockCmd::new( - redis::cmd("EVAL") - .arg(SET_MAX_SCRIPT) - .arg(1) - .arg("team:counter") - .arg(4.5f64) - .arg(600u64), - Ok("4.5"), - ), - ]); - - assert_eq!( - cache - .increment_with_floor("counter", -2, Duration::from_secs(30)) - .unwrap(), - 0 - ); - assert_eq!( - cache - .async_increment_with_floor("counter", -2, Duration::from_secs(30)) - .await - .unwrap(), - 0 - ); - assert_eq!( - cache.async_set_max("counter", 4.5, None).await.unwrap(), - 4.5 - ); -} - -const CLAIM_SCRIPT: &str = concat!( - "local current = redis.call('GET', KEYS[1]); ", - "if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ", - "elseif current ~= ARGV[1] then return 0; end; ", - "if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ", - "elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1" -); - -fn claim_eval(expected: &str, write: &str, refresh: bool, applied: i64) -> MockCmd { - MockCmd::new( - redis::cmd("EVAL") - .arg(CLAIM_SCRIPT) - .arg(1) - .arg("pin") - .arg(expected) - .arg(600) - .arg(write) - .arg(u8::from(refresh)), - Ok(applied), - ) -} - -#[rstest] -#[tokio::test] -async fn claims_match_eligible_values_written_by_another_encoder(context: ExactCacheContext) { - let python_payload = r#"{"model_id": "a", "deployment": "east"}"#; - let stored = json!({"deployment": "east", "model_id": "a"}); - let cache = json_cache(vec![ - MockCmd::new(redis::cmd("GET").arg("pin"), Ok(python_payload)), - claim_eval(python_payload, "", true, 1), - ]); - - assert_eq!( - cache - .async_claim_cache( - "pin", - json!({"model_id": "b"}), - vec![stored.clone()], - context - ) - .await - .unwrap(), - stored - ); -} - -#[rstest] -fn claims_retry_when_the_key_changes_and_replace_ineligible_winners(context: ExactCacheContext) { - let candidate = json!({"model_id": "b"}); - let payload = r#"{"model_id":"b"}"#; - let cache = json_cache(vec![ - MockCmd::new(redis::cmd("GET").arg("pin"), Ok(redis::Value::Nil)), - claim_eval("", payload, false, 0), - MockCmd::new(redis::cmd("GET").arg("pin"), Ok(r#"{"model_id":"gone"}"#)), - claim_eval(r#"{"model_id":"gone"}"#, payload, false, 1), - ]); - - assert_eq!( - cache - .claim_cache( - "pin", - candidate.clone(), - &[json!({"model_id": "a"})], - context - ) - .unwrap(), - candidate - ); -} - -#[rstest] -fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl( - context: ExactCacheContext, -) { - let stored = r#"{"model_id": "a"}"#; - let cache = json_cache(vec![ - MockCmd::new(redis::cmd("GET").arg("pin"), Ok(stored)), - claim_eval(stored, "", false, 1), - ]); - - assert_eq!( - cache - .claim_cache("pin", json!({"model_id": "b"}), &[], context) - .unwrap(), - json!({"model_id": "a"}) - ); -} - -#[rstest] -#[tokio::test] -async fn test_connection_reports_success_with_the_python_message() { - let cache = team(vec![MockCmd::new(redis::cmd("PING"), Ok("PONG"))]); - - let result = cache.test_connection().await.unwrap(); - assert_eq!(result.status, CacheConnectionStatus::Success); - assert_eq!(result.message, "Redis connection test successful"); - assert_eq!(result.error, None); -} - -#[rstest] -#[case::unexpected_reply(Ok("NOPE"), "Redis ping returned False", None)] -#[case::connection_refused( - Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))), - "Redis connection failed:", - Some("connection refused") -)] -#[tokio::test] -async fn test_connection_failures_use_the_python_result_contract( - #[case] reply: redis::RedisResult<&'static str>, - #[case] message: &str, - #[case] cause: Option<&str>, -) { - let cache = json_cache(vec![MockCmd::new(redis::cmd("PING"), reply)]); - - let result = cache.test_connection().await.unwrap(); - assert_eq!(result.status, CacheConnectionStatus::Failed); - assert!(result.message.starts_with(message), "{}", result.message); - assert_eq!(result.error.is_some(), cause.is_some()); - assert!(cause.is_none_or(|cause| { - result - .error - .as_deref() - .is_some_and(|error| error.contains(cause)) - })); -} - -#[rstest] -#[tokio::test] -async fn disconnect_keeps_a_caller_owned_connection_usable() { - let cache = team(vec![MockCmd::new(redis::cmd("PING"), Ok("PONG"))]); - cache.disconnect().await.unwrap(); - assert!(cache.ping().await.unwrap()); -} - -#[rstest] -#[tokio::test] -async fn disconnect_drains_an_idle_pool_without_connecting() { - let cache = RedisCache::new("redis://127.0.0.1:1", None, JsonCodec::::new()).unwrap(); - cache.disconnect().await.unwrap(); -} diff --git a/litellm-rust/crates/cache-redis/tests/cluster.rs b/litellm-rust/crates/cache-redis/tests/cluster.rs deleted file mode 100644 index e1ba01bcb04..00000000000 --- a/litellm-rust/crates/cache-redis/tests/cluster.rs +++ /dev/null @@ -1,563 +0,0 @@ -//! Tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a comma -//! separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them. - -mod support; - -use std::{collections::HashSet, time::Duration}; - -use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, BoundedCounterCache, BulkDeleteCache, CacheConnectionStatus, - CacheScript, ClaimCache, ClientInfoCache, ConnectionCache, CounterCache, DeleteCache, - DisconnectCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec, - PingCache, QueueCache, RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache, - TtlPipelineCache, -}; -use litellm_cache_redis::{ - RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisNode, RedisRpushOperation, - RedisTopology, -}; -use redis::cluster_routing::Slot; -use rstest::{fixture, rstest}; -use serde_json::json; -use support::{JsonCache, cluster_cache, cluster_url}; - -type Counter = RedisCache>; - -#[fixture] -fn cache(#[default("cache")] label: &str) -> Option { - cluster_cache(label, Duration::from_secs(120), JsonCodec::new()) -} - -#[fixture] -fn counter(#[default("counter")] label: &str) -> Option { - cluster_cache(label, Duration::from_secs(60), JsonCodec::new()) -} - -#[fixture] -fn context() -> ExactCacheContext { - ExactCacheContext::default() -} - -fn multi_slot_keys(count: usize) -> Vec { - let keys: Vec = (0..count).map(|index| format!("key-{index}")).collect(); - let slots: HashSet = keys.iter().map(Slot::for_key).collect(); - assert!(slots.len() > 1, "keys must span multiple slots"); - keys -} - -fn seconds(seconds: u64) -> Option { - Some(Duration::from_secs(seconds)) -} - -#[rstest] -#[case::no_startup_nodes("redis://127.0.0.1:7000", Vec::new())] -#[case::unix_socket_url( - "redis+unix:///tmp/redis.sock", - vec![RedisNode { host: "127.0.0.1".into(), port: 7000 }] -)] -fn constructor_rejects_unusable_cluster_configs( - #[case] url: &str, - #[case] startup_nodes: Vec, -) { - let error = JsonCache::connect( - url, - &RedisTopology::Cluster { startup_nodes }, - None, - JsonCodec::new(), - ) - .err(); - assert!(matches!(error, Some(Error::Unavailable(_)))); -} - -#[rstest] -#[tokio::test] -async fn single_key_operations_round_trip_with_ttl_rounding( - #[with("single")] cache: Option, -) { - let Some(cache) = cache else { return }; - let context = ExactCacheContext { - ttl: Some(Duration::from_millis(1500)), - }; - let keys = multi_slot_keys(12); - for (index, key) in keys.iter().enumerate() { - cache - .set_cache(key, json!({ "index": index }), &context) - .unwrap(); - } - for (index, key) in keys.iter().enumerate() { - assert_eq!( - cache.get_cache(key, &context).unwrap(), - Some(json!({ "index": index })) - ); - } - assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(2)); - assert!( - cache - .async_refresh_ttl(&keys[0], seconds(40)) - .await - .unwrap() - ); - assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(40)); - cache.delete_cache(&keys[0]).unwrap(); - assert_eq!(cache.get_cache(&keys[0], &context).unwrap(), None); - assert!(!cache.async_refresh_ttl(&keys[0], None).await.unwrap()); - assert!(cache.sync_ping().unwrap()); - cache.async_flush_cache().await.unwrap(); -} - -#[rstest] -#[tokio::test] -async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries( - #[with("batch")] cache: Option, - context: ExactCacheContext, -) { - let Some(cache) = cache else { return }; - let keys = multi_slot_keys(40); - for (index, key) in keys.iter().enumerate() { - if index % 5 == 0 { - continue; - } - cache - .async_set_cache(key, json!(index), context.clone()) - .await - .unwrap(); - } - let mut raw = redis::cluster::ClusterClient::new(vec![cluster_url()]) - .unwrap() - .get_connection() - .unwrap(); - let malformed = format!("{}:{}", cache.namespace().unwrap(), keys[1]); - redis::cmd("SET") - .arg(&malformed) - .arg("not json") - .exec(&mut raw) - .unwrap(); - - let entries = cache - .async_batch_get_cache(keys.clone(), context.clone()) - .await - .unwrap(); - assert_eq!(entries.len(), keys.len()); - for (index, entry) in entries.iter().enumerate() { - let expected = if index == 1 { - BatchEntry::Invalid - } else if index % 5 == 0 { - BatchEntry::Miss - } else { - BatchEntry::Hit(json!(index)) - }; - assert_eq!(*entry, expected, "entry {index}"); - } - assert_eq!(cache.batch_get_cache(&keys, &context).unwrap(), entries); - - cache.delete_cache_keys(keys.clone()).await.unwrap(); - let entries = cache.async_batch_get_cache(keys, context).await.unwrap(); - assert!(entries.iter().all(|entry| *entry == BatchEntry::Miss)); -} - -#[rstest] -#[tokio::test] -async fn pipelines_group_by_slot_and_return_results_in_submission_order( - #[with("pipeline")] cache: Option, - counter: Option, - context: ExactCacheContext, -) { - let (Some(cache), Some(counter)) = (cache, counter) else { - return; - }; - let keys = multi_slot_keys(30); - cache - .async_set_cache_pipeline( - keys.iter() - .enumerate() - .map(|(index, key)| (key.clone(), json!(index))) - .collect(), - context.clone(), - ) - .await - .unwrap(); - let hits = cache - .async_batch_get_cache(keys.clone(), context.clone()) - .await - .unwrap(); - assert!( - hits.iter() - .enumerate() - .all(|(index, entry)| *entry == BatchEntry::Hit(json!(index))) - ); - - cache - .async_set_cache_pipeline_with_ttls( - keys.iter() - .enumerate() - .map(|(index, key)| (key.clone(), json!(index), seconds(index as u64 + 10))) - .collect(), - ) - .await - .unwrap(); - for (index, key) in keys.iter().enumerate() { - assert_eq!( - cache.async_get_ttl(key).await.unwrap(), - seconds(index as u64 + 10), - "{key}" - ); - } - - let queues: Vec = keys.iter().map(|key| format!("queue:{key}")).collect(); - let pushed = cache - .async_rpush_pipeline( - queues - .iter() - .enumerate() - .map(|(index, key)| RedisRpushOperation { - key: key.clone(), - values: (0..=index) - .map(|value| RedisArg::Integer(value as i64)) - .collect(), - }) - .collect(), - ) - .await - .unwrap(); - assert_eq!(pushed, (1..=keys.len()).collect::>()); - let popped = cache - .async_lpop_pipeline( - queues - .iter() - .enumerate() - .map(|(index, key)| RedisLpopOperation { - key: key.clone(), - count: (index % 2 == 0).then_some(2), - }) - .collect(), - ) - .await - .unwrap(); - for (index, result) in popped.into_iter().enumerate() { - match result { - RedisLpopResult::Value(value) => { - assert_eq!(index % 2, 1, "queue {index}"); - assert_eq!(value, b"0"); - } - RedisLpopResult::Values(values) => { - assert_eq!(index % 2, 0, "queue {index}"); - let expected: Vec> = (0..=index) - .take(2) - .map(|value| value.to_string().into_bytes()) - .collect(); - assert_eq!(values, expected); - } - other => panic!("queue {index}: {other:?}"), - } - } - - let counters: Vec = keys.iter().map(|key| format!("counter:{key}")).collect(); - let totals = counter - .async_increment_pipeline( - counters - .iter() - .enumerate() - .map(|(index, key)| IncrementOperation { - key: key.clone(), - amount: index as f64 + 0.5, - ttl: (index % 3 == 0).then_some(Duration::from_secs(30)), - }) - .collect(), - ) - .await - .unwrap(); - let expected: Vec = (0..keys.len()).map(|index| index as f64 + 0.5).collect(); - assert_eq!(totals, expected); - assert_eq!( - counter.async_get_ttl(&counters[0]).await.unwrap(), - seconds(30) - ); - assert_eq!(counter.async_get_ttl(&counters[1]).await.unwrap(), None); - counter.async_flush_cache().await.unwrap(); - cache.async_flush_cache().await.unwrap(); -} - -#[rstest] -#[tokio::test] -async fn rpush_and_trim_is_one_transaction_on_the_key_slot( - #[with("trim")] cache: Option, -) { - let Some(cache) = cache else { return }; - let values = |values: &[&str]| values.iter().map(|value| RedisArg::from(*value)).collect(); - assert_eq!( - cache - .async_rpush_and_trim("buf", values(&["a", "b"]), 3) - .await - .unwrap(), - 2 - ); - assert_eq!( - cache - .async_rpush_and_trim("buf", values(&["c", "d"]), 3) - .await - .unwrap(), - 4 - ); - assert_eq!( - cache.async_lpop("buf", Some(10)).await.unwrap(), - RedisLpopResult::Values(vec![b"b".to_vec(), b"c".to_vec(), b"d".to_vec()]) - ); - cache.async_flush_cache().await.unwrap(); -} - -#[rstest] -#[tokio::test] -async fn scan_and_scoped_flush_cover_every_primary( - #[with("flush")] cache: Option, - #[from(cache)] - #[with("other")] - other: Option, - context: ExactCacheContext, -) { - let (Some(cache), Some(other)) = (cache, other) else { - return; - }; - let keys = multi_slot_keys(60); - for key in &keys { - cache - .async_set_cache(key, json!(true), context.clone()) - .await - .unwrap(); - other - .async_set_cache(key, json!(true), context.clone()) - .await - .unwrap(); - } - let mut scanned = cache.async_scan_iter("key-", 1000).await.unwrap(); - scanned.sort(); - let mut expected: Vec = keys - .iter() - .map(|key| format!("{}:{key}", cache.namespace().unwrap())) - .collect(); - expected.sort(); - assert_eq!(scanned, expected); - assert_eq!(cache.async_scan_iter("key-", 7).await.unwrap().len(), 7); - - cache.flush_cache().unwrap(); - let flushed = cache - .async_batch_get_cache(keys.clone(), context.clone()) - .await - .unwrap(); - assert!(flushed.iter().all(|entry| *entry == BatchEntry::Miss)); - let kept = other.async_batch_get_cache(keys, context).await.unwrap(); - assert!( - kept.iter() - .all(|entry| *entry == BatchEntry::Hit(json!(true))) - ); - other.async_flush_cache().await.unwrap(); -} - -fn ping_calls_per_node(startup: &redis::Client) -> Vec<(String, u64)> { - let mut connection = startup.get_connection().unwrap(); - let nodes: String = redis::cmd("CLUSTER") - .arg("NODES") - .query(&mut connection) - .unwrap(); - let mut counts: Vec<(String, u64)> = nodes - .lines() - .map(|line| { - let address = line.split_whitespace().nth(1).unwrap(); - let address = address.split('@').next().unwrap(); - let mut node = redis::Client::open(format!("redis://{address}")) - .unwrap() - .get_connection() - .unwrap(); - let stats: String = redis::cmd("INFO") - .arg("commandstats") - .query(&mut node) - .unwrap(); - let calls = stats - .lines() - .find_map(|stat| stat.strip_prefix("cmdstat_ping:calls=")) - .and_then(|rest| rest.split(',').next()) - .map_or(0, |calls| calls.parse().unwrap()); - (address.to_string(), calls) - }) - .collect(); - counts.sort(); - counts -} - -#[rstest] -#[tokio::test] -async fn ping_reaches_every_node(#[with("ping")] cache: Option) { - let Some(cache) = cache else { return }; - let startup = redis::Client::open(cluster_url()).unwrap(); - let before = ping_calls_per_node(&startup); - assert!(before.len() >= 2, "{before:?}"); - assert!(cache.ping().await.unwrap()); - let after = ping_calls_per_node(&startup); - for ((node, calls_before), (_, calls_after)) in before.iter().zip(&after) { - assert!(calls_after > calls_before, "{node} was not pinged"); - } - assert!(cache.sync_ping().unwrap()); - let result = cache.test_connection().await.unwrap(); - assert_eq!(result.status, CacheConnectionStatus::Success); - assert_eq!(result.message, "Redis Cluster connection test successful"); -} - -#[rstest] -#[tokio::test] -async fn disconnect_closes_idle_connections_and_reconnects_on_demand( - #[with("disconnect")] cache: Option, -) { - let Some(cache) = cache else { return }; - assert!(cache.ping().await.unwrap()); - cache.disconnect().await.unwrap(); - assert!(cache.ping().await.unwrap()); -} - -#[rstest] -#[case::keep_existing_ttl(false)] -#[case::refresh_ttl(true)] -#[tokio::test] -async fn increments_refresh_the_ttl_only_when_asked( - counter: Option, - #[case] refresh_ttl: bool, -) { - let Some(counter) = counter else { return }; - let context = ExactCacheContext { ttl: seconds(60) }; - counter - .async_set_cache("spend", 0.0, ExactCacheContext { ttl: seconds(600) }) - .await - .unwrap(); - assert_eq!( - counter - .async_increment("spend", 1.5, context.clone(), refresh_ttl) - .await - .unwrap(), - 1.5 - ); - assert_eq!( - counter - .async_increment("spend", 2.0, context, refresh_ttl) - .await - .unwrap(), - 3.5 - ); - let ttl = counter.async_get_ttl("spend").await.unwrap().unwrap(); - assert_eq!(ttl <= Duration::from_secs(60), refresh_ttl, "{ttl:?}"); - counter.async_flush_cache().await.unwrap(); -} - -#[rstest] -#[tokio::test] -async fn counters_claims_scripts_and_sets_work_on_the_cluster( - counter: Option, - #[with("claim")] cache: Option, - context: ExactCacheContext, -) { - let (Some(counter), Some(cache)) = (counter, cache) else { - return; - }; - assert_eq!( - counter - .increment_cache("spend", 1.5, context.clone()) - .unwrap(), - 1.5 - ); - assert_eq!( - counter - .async_increment("spend", 2.0, context.clone(), false) - .await - .unwrap(), - 3.5 - ); - assert_eq!( - counter - .increment_with_floor("budget", -3, Duration::from_secs(30)) - .unwrap(), - 0 - ); - assert_eq!( - counter - .async_increment_with_floor("budget", 7, Duration::from_secs(30)) - .await - .unwrap(), - 7 - ); - assert_eq!(counter.async_set_max("peak", 4.0, None).await.unwrap(), 4.0); - assert_eq!(counter.async_set_max("peak", 2.0, None).await.unwrap(), 4.0); - counter.flush_cache().unwrap(); - - let owner = json!("owner-a"); - let rival = json!("owner-b"); - assert_eq!( - cache - .claim_cache("lock", owner.clone(), &[], context.clone()) - .unwrap(), - owner - ); - assert_eq!( - cache - .async_claim_cache("lock", rival.clone(), vec![owner.clone()], context.clone()) - .await - .unwrap(), - owner - ); - assert_eq!( - cache - .claim_cache("lock", rival.clone(), &[], context.clone()) - .unwrap(), - owner - ); - assert_eq!( - cache - .async_claim_cache("lock", rival.clone(), vec![rival.clone()], context.clone()) - .await - .unwrap(), - rival - ); - - let script = cache - .async_register_script("return redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])".into()); - let reply = script - .invoke( - vec!["scripted".into()], - vec![RedisArg::Bytes(b"payload".to_vec()), RedisArg::Integer(5)], - ) - .await - .unwrap(); - assert_eq!(reply, redis::Value::Okay); - assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), seconds(5)); - let evaluated = cache - .async_eval( - "return redis.call('GET', KEYS[1])".into(), - vec!["scripted".into()], - Vec::new(), - ) - .await - .unwrap(); - assert_eq!(evaluated, redis::Value::BulkString(b"payload".to_vec())); - - assert_eq!( - cache - .async_set_cache_sadd( - "members", - vec![ - RedisArg::Bytes(b"a".to_vec()), - RedisArg::Bytes(b"b".to_vec()) - ], - seconds(9), - ) - .await - .unwrap(), - 2 - ); - assert_eq!(cache.async_get_ttl("members").await.unwrap(), seconds(9)); - - let result = cache.test_connection().await.unwrap(); - assert_eq!(result.status, CacheConnectionStatus::Success); - assert!(cache.ping().await.unwrap()); - let info = cache.info().unwrap(); - assert!(info.matches("redis_version").count() > 1, "{info}"); - assert!(cache.client_list().unwrap().contains("id=")); - cache.async_flush_cache().await.unwrap(); - assert_eq!(cache.async_get_ttl("members").await.unwrap(), None); - assert_eq!(cache.get_cache("lock", &context).unwrap(), None); -} diff --git a/litellm-rust/crates/cache-redis/tests/contract.rs b/litellm-rust/crates/cache-redis/tests/contract.rs deleted file mode 100644 index 85d043b90b2..00000000000 --- a/litellm-rust/crates/cache-redis/tests/contract.rs +++ /dev/null @@ -1,95 +0,0 @@ -//! The shared cache contracts, run against the in-process fake connection and, when -//! `LITELLM_TEST_REDIS_CLUSTER_NODES` is set, against a live Redis Cluster. - -mod support; - -use std::time::Duration; - -use litellm_cache::{ExactCacheContext, JsonCodec}; -use litellm_cache_testing as contract; -use rstest::rstest; -use serde_json::json; -use support::{JsonCache, cluster_cache, fake_cache}; - -const PREFIX: &str = "contract:"; - -#[derive(Clone, Copy, Debug)] -enum Contract { - HitAndMiss, - SyncAsyncEquivalence, - OverwriteReplaces, - PipelineWritesEveryEntry, - BatchPreservesOrder, - DeleteRemovesKey, - FlushClears, - CounterAccumulates, -} - -#[derive(Clone, Copy, Debug)] -enum Server { - Fake, - Cluster, -} - -async fn check(contract: Contract, cache: &JsonCache) -where - C: redis::ConnectionLike + Send + 'static, -{ - let context = ExactCacheContext::default(); - match contract { - Contract::HitAndMiss => { - contract::hit_and_miss(cache, context, PREFIX, json!({"answer": 42})).await - } - Contract::SyncAsyncEquivalence => { - contract::sync_async_equivalence(cache, context, PREFIX, json!("first"), json!([2])) - .await - } - Contract::OverwriteReplaces => { - contract::overwrite_replaces(cache, context, PREFIX, json!(1), json!({"b": 2})).await - } - Contract::PipelineWritesEveryEntry => { - contract::pipeline_writes_every_entry( - cache, - context, - PREFIX, - vec![json!("a"), json!(2), json!({"c": true})], - ) - .await - } - Contract::BatchPreservesOrder => { - contract::batch_preserves_order(cache, context, PREFIX, json!("first"), json!(2)).await - } - Contract::DeleteRemovesKey => { - contract::delete_removes_key(cache, context, PREFIX, json!("value")).await - } - Contract::FlushClears => { - contract::flush_clears(cache, context, PREFIX, json!("value")).await - } - Contract::CounterAccumulates => contract::counter_accumulates(cache, context, PREFIX).await, - } -} - -#[rstest] -#[case::hit_and_miss(Contract::HitAndMiss)] -#[case::sync_async_equivalence(Contract::SyncAsyncEquivalence)] -#[case::overwrite_replaces(Contract::OverwriteReplaces)] -#[case::pipeline_writes_every_entry(Contract::PipelineWritesEveryEntry)] -#[case::batch_preserves_order(Contract::BatchPreservesOrder)] -#[case::delete_removes_key(Contract::DeleteRemovesKey)] -#[case::flush_clears(Contract::FlushClears)] -#[case::counter_accumulates(Contract::CounterAccumulates)] -#[tokio::test] -async fn redis_satisfies_the_cache_contract( - #[case] contract: Contract, - #[values(Server::Fake, Server::Cluster)] server: Server, -) { - match server { - Server::Fake => check(contract, &fake_cache("contract")).await, - Server::Cluster => { - let label = format!("{contract:?}"); - if let Some(cache) = cluster_cache(&label, Duration::from_secs(120), JsonCodec::new()) { - check(contract, &cache).await; - } - } - } -} diff --git a/litellm-rust/crates/cache-redis/tests/support/mod.rs b/litellm-rust/crates/cache-redis/tests/support/mod.rs deleted file mode 100644 index b04e1c29f3b..00000000000 --- a/litellm-rust/crates/cache-redis/tests/support/mod.rs +++ /dev/null @@ -1,231 +0,0 @@ -#![allow(dead_code)] - -use std::{ - collections::BTreeMap, - time::{Duration, SystemTime, UNIX_EPOCH}, -}; - -use litellm_cache::{CacheCodec, Error, JsonCodec}; -use litellm_cache_redis::{RedisCache, RedisNode, RedisTopology}; - -/// Encodes a byte behind a tag, so a value written with another tag decodes as invalid. -pub struct TaggedByteCodec(pub u8); - -impl CacheCodec for TaggedByteCodec { - type Value = u8; - - fn encode(&self, value: &u8) -> Result, Error> { - if *value > 127 { - return Err(Error::InvalidEntry(None)); - } - Ok(vec![self.0, *value]) - } - - fn decode(&self, bytes: &[u8]) -> Result { - match bytes { - [tag, value] if *tag == self.0 => Ok(*value), - _ => Err(Error::InvalidEntry(None)), - } - } -} - -/// A stateful in-process stand-in for a Redis server that understands the string commands the -/// shared contracts exercise, so they run without a live server. TTLs are accepted and ignored. -#[derive(Default)] -pub struct FakeRedis { - strings: BTreeMap, Vec>, -} - -impl FakeRedis { - fn run(&mut self, command: Vec>) -> redis::RedisResult { - let name = String::from_utf8_lossy(&command[0]).to_ascii_uppercase(); - let args = &command[1..]; - Ok(match name.as_str() { - "PING" => redis::Value::SimpleString("PONG".into()), - "SET" | "SETEX" => { - let value = if name == "SET" { &args[1] } else { &args[2] }; - self.strings.insert(args[0].clone(), value.clone()); - redis::Value::Okay - } - "GET" => self.get(&args[0]), - "MGET" => redis::Value::Array(args.iter().map(|key| self.get(key)).collect()), - "DEL" => { - let removed = args - .iter() - .filter(|key| self.strings.remove(*key).is_some()) - .count(); - redis::Value::Int(removed as i64) - } - "SCAN" => { - let pattern = &args[2]; - let keys = self - .strings - .keys() - .filter(|key| glob(pattern, key)) - .map(|key| redis::Value::BulkString(key.clone())) - .collect(); - redis::Value::Array(vec![ - redis::Value::BulkString(b"0".to_vec()), - redis::Value::Array(keys), - ]) - } - "EVAL" if args[0].windows(11).any(|window| window == b"INCRBYFLOAT") => { - self.increment_by_float(&args[2], &args[3]) - } - "INCRBYFLOAT" => self.increment_by_float(&args[0], &args[1]), - _ => { - return Err(redis::RedisError::from(( - redis::ErrorKind::Client, - "unsupported command", - name, - ))); - } - }) - } - - fn get(&self, key: &[u8]) -> redis::Value { - self.strings.get(key).map_or(redis::Value::Nil, |value| { - redis::Value::BulkString(value.clone()) - }) - } - - fn increment_by_float(&mut self, key: &[u8], amount: &[u8]) -> redis::Value { - let current = self - .strings - .get(key) - .map_or(0.0, |value| parse_float(value)); - let total = format!("{}", current + parse_float(amount)); - self.strings - .insert(key.to_vec(), total.clone().into_bytes()); - redis::Value::BulkString(total.into_bytes()) - } -} - -fn parse_float(bytes: &[u8]) -> f64 { - std::str::from_utf8(bytes).unwrap().parse().unwrap() -} - -/// Redis `MATCH` globbing for `*`, `?` and backslash escapes. -fn glob(pattern: &[u8], key: &[u8]) -> bool { - match pattern.split_first() { - None => key.is_empty(), - Some((b'*', rest)) => (0..=key.len()).any(|skip| glob(rest, &key[skip..])), - Some((b'?', rest)) => !key.is_empty() && glob(rest, &key[1..]), - Some((b'\\', [escaped, rest @ ..])) => { - key.first() == Some(escaped) && glob(rest, &key[1..]) - } - Some((literal, rest)) => key.first() == Some(literal) && glob(rest, &key[1..]), - } -} - -/// Splits RESP request bytes into the commands they carry. -fn commands(mut bytes: &[u8]) -> Vec>> { - fn line<'a>(bytes: &mut &'a [u8]) -> &'a [u8] { - let end = bytes - .windows(2) - .position(|window| window == b"\r\n") - .unwrap(); - let (line, rest) = bytes.split_at(end); - *bytes = &rest[2..]; - line - } - fn length(line: &[u8]) -> usize { - std::str::from_utf8(&line[1..]).unwrap().parse().unwrap() - } - let mut commands = Vec::new(); - while !bytes.is_empty() { - let count = length(line(&mut bytes)); - let command = (0..count) - .map(|_| { - let size = length(line(&mut bytes)); - let (argument, rest) = bytes.split_at(size); - bytes = &rest[2..]; - argument.to_vec() - }) - .collect(); - commands.push(command); - } - commands -} - -impl redis::ConnectionLike for FakeRedis { - fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult { - let command = commands(cmd).into_iter().next().unwrap(); - self.run(command) - } - - fn req_packed_commands( - &mut self, - cmd: &[u8], - offset: usize, - count: usize, - ) -> redis::RedisResult> { - let replies = commands(cmd) - .into_iter() - .map(|command| self.run(command)) - .collect::>>()?; - Ok(replies.into_iter().skip(offset).take(count).collect()) - } - - fn get_db(&self) -> i64 { - 0 - } - - fn check_connection(&mut self) -> bool { - true - } - - fn is_open(&self) -> bool { - true - } -} - -pub type JsonCache = RedisCache, C>; - -pub fn fake_cache(namespace: &str) -> JsonCache { - RedisCache::with_connection(FakeRedis::default(), None, JsonCodec::new()) - .with_namespace(Some(namespace.into())) -} - -/// Startup nodes from `LITELLM_TEST_REDIS_CLUSTER_NODES` (`host:port,host:port`); tests that -/// need a live cluster skip when it is unset. -pub fn cluster_topology() -> Option { - let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?; - let startup_nodes = nodes - .split(',') - .map(|node| { - let (host, port) = node.trim().rsplit_once(':').expect("host:port"); - RedisNode { - host: host.to_string(), - port: port.parse().expect("port"), - } - }) - .collect(); - Some(RedisTopology::Cluster { startup_nodes }) -} - -pub fn cluster_url() -> String { - std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL") - .unwrap_or_else(|_| "redis://127.0.0.1:7000".into()) -} - -pub fn unique_namespace(label: &str) -> String { - let nanos = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap() - .as_nanos(); - format!("cluster-test:{label}:{nanos}") -} - -pub fn cluster_cache( - label: &str, - default_ttl: Duration, - codec: S, -) -> Option> { - let topology = cluster_topology()?; - Some( - RedisCache::connect(&cluster_url(), &topology, Some(default_ttl), codec) - .expect("cluster connection") - .with_namespace(Some(unique_namespace(label))), - ) -} diff --git a/litellm-rust/crates/cache-response/Cargo.toml b/litellm-rust/crates/cache-response/Cargo.toml index 8f4570edc6f..46ef6ab92a8 100644 --- a/litellm-rust/crates/cache-response/Cargo.toml +++ b/litellm-rust/crates/cache-response/Cargo.toml @@ -16,9 +16,6 @@ sha2.workspace = true litellm-cache-gcs.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache-memory.workspace = true -litellm-cache-redis.workspace = true -redis.workspace = true -redis-test.workspace = true rstest.workspace = true tokio.workspace = true wiremock.workspace = true diff --git a/litellm-rust/crates/cache-response/tests/connection.rs b/litellm-rust/crates/cache-response/tests/connection.rs index bc24cf56846..f64057f6f1c 100644 --- a/litellm-rust/crates/cache-response/tests/connection.rs +++ b/litellm-rust/crates/cache-response/tests/connection.rs @@ -2,79 +2,11 @@ mod support; use std::{sync::Arc, time::Duration}; -use litellm_cache::CacheConnectionStatus; use litellm_cache_memory::InMemoryCache; -use litellm_cache_response::{ - CacheEntry, ConnectionProbe, ExactResponseCache, ResponseCache, ResponseCacheRequest, -}; -use redis_test::MockCmd; +use litellm_cache_response::{CacheEntry, ExactResponseCache, ResponseCache, ResponseCacheRequest}; use rstest::rstest; use serde_json::json; -use support::{keyed, memory, redis, request}; - -#[rstest] -#[case::reachable( - Ok("PONG"), - CacheConnectionStatus::Success, - "Redis connection test successful", - false -)] -#[case::unexpected_reply( - Ok("NOPE"), - CacheConnectionStatus::Failed, - "Redis ping returned False", - false -)] -#[case::connection_refused( - Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))), - CacheConnectionStatus::Failed, - "Redis connection failed:", - true -)] -#[tokio::test] -async fn connection_backends_are_reachable_as_a_probe( - #[case] reply: redis::RedisResult<&'static str>, - #[case] status: CacheConnectionStatus, - #[case] message: &str, - #[case] has_error: bool, -) { - let probe: Arc = - Arc::new(redis(vec![MockCmd::new(redis::cmd("PING"), reply)], None)); - - let result = probe.test_connection().await.unwrap(); - assert_eq!(result.status, status); - assert!(result.message.starts_with(message), "{}", result.message); - assert_eq!(result.error.is_some(), has_error); -} - -#[rstest] -#[tokio::test] -async fn one_service_serves_both_the_exact_cache_and_its_probe(request: ResponseCacheRequest) { - let service = Arc::new(redis( - vec![ - MockCmd::new(redis::cmd("PING"), Ok("PONG")), - MockCmd::new( - redis::cmd("GET").arg("tenant:key"), - Ok(br#"{"timestamp":100.0,"response":{"ok":true}}"#.to_vec()), - ), - ], - Some("tenant"), - )); - let probe: Arc = service.clone(); - let exact: Arc = service; - - assert_eq!( - probe.test_connection().await.unwrap().status, - CacheConnectionStatus::Success - ); - assert_eq!( - exact - .async_lookup(&request, Duration::from_secs(100)) - .await - .unwrap(), - Some(json!({"ok": true})) - ); -} +use support::{keyed, memory, request}; /// The in-memory backend has no `test_connection`, as in Python, and still serves every response /// operation. diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index 70ae259d33a..659ba065cf2 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -17,10 +17,9 @@ use litellm_cache_response::{ CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest, WriteBuffer, cache_key, }; -use redis_test::MockCmd; use rstest::{fixture, rstest}; use serde_json::{Value, json}; -use support::{keyed, memory, redis, request}; +use support::{keyed, memory, request}; type Memory = Arc>>; @@ -331,79 +330,6 @@ async fn directives_skip_io_and_keep_reads_and_writes_independent( ); } -#[rstest] -#[case::python_sync_literal( - br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.as_slice() -)] -#[case::python_async_json(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice())] -#[tokio::test] -async fn redis_consumer_reads_python_sync_and_async_envelopes( - request: ResponseCacheRequest, - #[case] stored: &[u8], - #[values(false, true)] asynchronous: bool, -) { - let cache = redis( - vec![MockCmd::new( - redis::cmd("GET").arg("tenant:key"), - Ok(stored.to_vec()), - )], - Some("tenant"), - ); - let now = Duration::from_secs(101); - let found = if asynchronous { - cache.async_lookup(&request, now).await.unwrap() - } else { - cache.lookup(&request, now).unwrap() - }; - assert_eq!(found, Some(json!({"ok": true, "text": "cached"}))); -} - -#[rstest] -#[case::object( - json!({"ok": true, "text": "cached"}), - br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice() -)] -#[case::array(json!([1, 2]), br#"{"timestamp":100.0,"response":"[1,2]"}"#.as_slice())] -#[tokio::test] -async fn redis_consumer_writes_python_compatible_json( - request: ResponseCacheRequest, - #[case] response: Value, - #[case] wire: &[u8], -) { - let cache = redis( - vec![MockCmd::new( - redis::cmd("SETEX").arg("tenant:key").arg(600).arg(wire), - Ok("OK"), - )], - Some("tenant"), - ); - cache - .async_store(&request, response, Duration::from_secs(100)) - .await - .unwrap(); -} - -#[rstest] -#[tokio::test] -async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis( - mut request: ResponseCacheRequest, -) { - let cache = redis( - vec![MockCmd::new( - redis::cmd("GET").arg("tenant:key"), - Ok(b"invalid".to_vec()), - )], - None, - ); - request.controls.no_cache = true; - assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None); - request.controls.no_cache = false; - assert_eq!( - cache.async_lookup(&request, Duration::ZERO).await.unwrap(), - None - ); -} - #[rstest] #[tokio::test] async fn captured_service_keeps_the_selected_backend_for_background_writes( @@ -541,27 +467,6 @@ async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses( assert_eq!(memory.lookup(&requests[2], now).unwrap(), None); } -#[rstest] -#[tokio::test] -async fn batch_lookup_with_no_readable_request_skips_the_backend( - #[values(false, true)] asynchronous: bool, -) { - let cache = redis(Vec::new(), None); - let mut request = keyed("key"); - request.controls.no_cache = true; - let requests = [request.clone(), request]; - let partial = if asynchronous { - cache - .async_lookup_batch(&requests, Duration::ZERO) - .await - .unwrap() - } else { - cache.lookup_batch(&requests, Duration::ZERO).unwrap() - }; - assert_eq!(partial.values, vec![None, None]); - assert_eq!(partial.missing_indices, vec![0, 1]); -} - #[rstest] #[tokio::test] async fn deferred_entries_keep_the_time_they_were_produced( diff --git a/litellm-rust/crates/cache-response/tests/support/mod.rs b/litellm-rust/crates/cache-response/tests/support/mod.rs index b992a937b70..f84da303659 100644 --- a/litellm-rust/crates/cache-response/tests/support/mod.rs +++ b/litellm-rust/crates/cache-response/tests/support/mod.rs @@ -1,15 +1,9 @@ use std::{sync::Arc, time::Duration}; use litellm_cache_memory::InMemoryCache; -use litellm_cache_redis::RedisCache; -use litellm_cache_response::{ - CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, -}; -use redis_test::{MockCmd, MockRedisConnection}; +use litellm_cache_response::{CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheRequest}; use rstest::fixture; -pub type MockedRedis = RedisCache; - #[fixture] pub fn memory() -> Arc>> { Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new( @@ -29,12 +23,3 @@ pub fn keyed(key: &str) -> ResponseCacheRequest { ..Default::default() }) } - -/// A Redis response cache that must receive exactly `commands`, in order. -pub fn redis(commands: Vec, namespace: Option<&str>) -> ResponseCache { - let connection = MockRedisConnection::new(commands).assert_all_commands_consumed(); - ResponseCache::new(Arc::new( - RedisCache::with_connection(connection, None, ResponseCacheCodec) - .with_namespace(namespace.map(str::to_owned)), - )) -} diff --git a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml deleted file mode 100644 index 5b4e99de382..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml +++ /dev/null @@ -1,22 +0,0 @@ -[package] -name = "litellm-cache-valkey-semantic" -version = "0.1.0" -edition.workspace = true -license.workspace = true -repository.workspace = true - -[dependencies] -litellm-cache.workspace = true -litellm-cache-redis.workspace = true -redis = { workspace = true, features = ["tls-rustls"] } -sha2.workspace = true -thiserror.workspace = true -uuid = { version = "1", features = ["v4"] } - -[dev-dependencies] -litellm-error.workspace = true -litellm-cache-testing.workspace = true -redis-test.workspace = true -rstest.workspace = true -serde_json.workspace = true -tokio.workspace = true diff --git a/litellm-rust/crates/cache-valkey-semantic/src/cache.rs b/litellm-rust/crates/cache-valkey-semantic/src/cache.rs deleted file mode 100644 index 947c0edc479..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/src/cache.rs +++ /dev/null @@ -1,245 +0,0 @@ -use std::{ - sync::{Arc, Mutex}, - time::Duration, -}; - -use litellm_cache::{ - BaseCache, CacheCodec, Error, SemanticCacheContext, - semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context}, -}; -use litellm_cache_redis::{RedisTopology, connection::Connections}; - -use crate::{ - ValkeySemanticConfig, - index::IndexState, - search::{embedding_bytes, scope_tag, search_document, write_document}, -}; - -/// `ValkeySemanticCache`: a semantic cache on valkey-search's TAG + VECTOR index. Values go -/// through the injected codec, so the response layer decides what a cached entry is. -pub struct ValkeySemanticCache { - connections: Arc>, - embedder: E, - codec: S, - config: ValkeySemanticConfig, - index_dimension: Arc>>, -} - -impl ValkeySemanticCache { - pub fn new( - url: &str, - embedder: E, - codec: S, - config: ValkeySemanticConfig, - ) -> Result { - Ok(Self { - connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), - embedder, - codec, - config, - index_dimension: Arc::new(Mutex::new(None)), - }) - } -} - -impl ValkeySemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - pub fn with_connection( - connection: C, - embedder: E, - codec: S, - config: ValkeySemanticConfig, - ) -> Self { - Self { - connections: Arc::new(Connections::fixed(connection)), - embedder, - codec, - config, - index_dimension: Arc::new(Mutex::new(None)), - } - } - - pub fn similarity_threshold(&self) -> f64 { - self.config.similarity_threshold - } - - pub fn index_name(&self) -> &str { - &self.config.index_name - } - - fn index_state(&self) -> IndexState { - IndexState { - name: self.config.index_name.clone(), - prefix: format!("{}:", self.config.index_name), - dimension: Arc::clone(&self.index_dimension), - similarity_threshold: self.config.similarity_threshold, - } - } - - fn decode(&self, lookup: SemanticLookup>) -> Result, Error> { - Ok(SemanticLookup { - value: lookup - .value - .map(|bytes| self.codec.decode(&bytes)) - .transpose()?, - similarity: lookup.similarity, - }) - } -} - -impl ValkeySemanticCache -where - E: Embedder, - S: CacheCodec + Clone, - C: redis::ConnectionLike + Send + 'static, -{ - /// The same index and connections behind a different embedder. - pub fn with_embedder(&self, embedder: E2) -> ValkeySemanticCache { - ValkeySemanticCache { - connections: Arc::clone(&self.connections), - embedder, - codec: self.codec.clone(), - config: self.config.clone(), - index_dimension: Arc::clone(&self.index_dimension), - } - } -} - -impl BaseCache for ValkeySemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type Value = S::Value; - type Context = SemanticCacheContext; - - fn get_ttl(&self, context: &Self::Context) -> Option { - context.ttl - } - - fn set_cache( - &self, - key: &str, - value: Self::Value, - context: &Self::Context, - ) -> Result<(), Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(()); - }; - let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; - let response = self.codec.encode(&value)?; - let index = self.index_state(); - self.connections.execute(|connection| { - write_document( - connection, - &index, - &scope_tag(key), - &prompt, - response, - embedding_bytes(&embedding), - self.get_ttl(context), - ) - }) - } - - fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { - self.get_cache_with_similarity(key, context) - .map(|lookup| lookup.value) - } - - 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 embedding = self - .embedder - .async_embed(&prompt, context.metadata.as_ref()) - .await?; - let response = self.codec.encode(&value)?; - let index = self.index_state(); - let scope = scope_tag(key); - Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - write_document( - connection, - &index, - &scope, - &prompt, - response, - embedding_bytes(&embedding), - context.ttl, - ) - }) - .await - } - - async fn async_get_cache( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - self.async_get_cache_with_similarity(key, context) - .await - .map(|lookup| lookup.value) - } -} - -/// Python stamps a similarity of `0.0` when there is no prompt or no document in the key's -/// scope, and the closest document's similarity even when it misses the threshold. -impl SemanticCache for ValkeySemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn get_cache_with_similarity( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; - let index = self.index_state(); - let lookup = self.connections.execute(|connection| { - search_document( - connection, - &index, - &scope_tag(key), - embedding_bytes(&embedding), - ) - })?; - self.decode(lookup) - } - - async fn async_get_cache_with_similarity( - &self, - key: &str, - context: &Self::Context, - ) -> Result, Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let embedding = self - .embedder - .async_embed(&prompt, context.metadata.as_ref()) - .await?; - let index = self.index_state(); - let scope = scope_tag(key); - let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - search_document(connection, &index, &scope, embedding_bytes(&embedding)) - }) - .await?; - self.decode(lookup) - } -} diff --git a/litellm-rust/crates/cache-valkey-semantic/src/config.rs b/litellm-rust/crates/cache-valkey-semantic/src/config.rs deleted file mode 100644 index c417652c27b..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/src/config.rs +++ /dev/null @@ -1,8 +0,0 @@ -/// `ValkeySemanticCache.DEFAULT_VALKEY_INDEX_NAME`. -pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; - -#[derive(Clone, Debug, PartialEq)] -pub struct ValkeySemanticConfig { - pub similarity_threshold: f64, - pub index_name: String, -} diff --git a/litellm-rust/crates/cache-valkey-semantic/src/index.rs b/litellm-rust/crates/cache-valkey-semantic/src/index.rs deleted file mode 100644 index 3dfff32ff7c..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/src/index.rs +++ /dev/null @@ -1,117 +0,0 @@ -use std::sync::{Arc, Mutex}; - -use litellm_cache::Error; -use litellm_cache_redis::connection::ConnectionRef; - -use crate::search::value_text; - -/// The valkey-search index one cache writes to, with the dimension it was last ensured for. -#[derive(Clone)] -pub(crate) struct IndexState { - pub(crate) name: String, - pub(crate) prefix: String, - pub(crate) dimension: Arc>>, - pub(crate) similarity_threshold: f64, -} - -/// `_ensure_index_sync` / `_ensure_index_async`: create the TAG + HNSW index once per dimension, -/// and accept an existing index unless it reports a different dimension. -pub(crate) fn ensure_index( - connection: &mut ConnectionRef<'_>, - index: &IndexState, - dimension: usize, -) -> Result<(), Error> { - if index - .dimension - .lock()? - .is_some_and(|existing| existing == dimension) - { - return Ok(()); - } - let create = redis::cmd("FT.CREATE") - .arg(&index.name) - .arg("ON") - .arg("HASH") - .arg("PREFIX") - .arg(1) - .arg(&index.prefix) - .arg("SCHEMA") - .arg("litellm_cache_key") - .arg("TAG") - .arg("embedding") - .arg("VECTOR") - .arg("HNSW") - .arg(6) - .arg("TYPE") - .arg("FLOAT32") - .arg("DIM") - .arg(dimension) - .arg("DISTANCE_METRIC") - .arg("COSINE") - .query::(connection) - .map(|_| ()); - if let Err(error) = create { - if !index_exists(&error) { - return Err(Error::unavailable(error)); - } - let info = redis::cmd("FT.INFO") - .arg(&index.name) - .query::(connection) - .map_err(Error::unavailable)?; - if let Some(existing) = - index_dimension_from_info(&info).filter(|existing| *existing != dimension) - { - return Err(Error::unavailable(DimensionMismatch { - existing, - dimension, - })); - } - } - *index.dimension.lock()? = Some(dimension); - Ok(()) -} - -fn index_exists(error: &redis::RedisError) -> bool { - error.code().is_some() - && error - .detail() - .is_some_and(|detail| detail.to_ascii_lowercase().contains("already exists")) -} - -#[derive(Debug, thiserror::Error)] -#[error("index exists with dimension {existing}, not {dimension}")] -struct DimensionMismatch { - existing: usize, - dimension: usize, -} - -/// `_extract_index_dim`: flatten each attribute one level and read the value after -/// `dimensions`. -fn index_dimension_from_info(value: &redis::Value) -> Option { - let redis::Value::Array(values) = value else { - return None; - }; - let attributes = values.windows(2).find_map(|pair| { - (value_text(&pair[0]).as_deref() == Some("attributes")).then_some(&pair[1]) - })?; - let redis::Value::Array(fields) = attributes else { - return None; - }; - fields.iter().find_map(|field| { - let redis::Value::Array(values) = field else { - return None; - }; - let values = values - .iter() - .flat_map(|value| match value { - redis::Value::Array(values) => values.as_slice(), - _ => std::slice::from_ref(value), - }) - .collect::>(); - values.windows(2).find_map(|pair| { - (value_text(pair[0]).as_deref() == Some("dimensions")) - .then(|| value_text(pair[1]).and_then(|value| value.parse().ok())) - .flatten() - }) - }) -} diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs deleted file mode 100644 index f2f62bc95bd..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs +++ /dev/null @@ -1,7 +0,0 @@ -mod cache; -mod config; -mod index; -mod search; - -pub use cache::ValkeySemanticCache; -pub use config::{DEFAULT_INDEX_NAME, ValkeySemanticConfig}; diff --git a/litellm-rust/crates/cache-valkey-semantic/src/search.rs b/litellm-rust/crates/cache-valkey-semantic/src/search.rs deleted file mode 100644 index c5e4e748b58..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/src/search.rs +++ /dev/null @@ -1,161 +0,0 @@ -use std::time::Duration; - -use litellm_cache::{Error, semantic::SemanticLookup}; -use litellm_cache_redis::connection::ConnectionRef; -use sha2::{Digest, Sha256}; -use uuid::Uuid; - -use crate::index::{IndexState, ensure_index}; - -/// `_scope_tag`: valkey-search TAG fields cannot match arbitrary keys verbatim, so scopes are -/// the key's lowercase SHA-256. -pub(crate) fn scope_tag(key: &str) -> String { - let digest = Sha256::digest(key.as_bytes()); - digest.iter().map(|byte| format!("{byte:02x}")).collect() -} - -pub(crate) fn embedding_bytes(embedding: &[f32]) -> Vec { - embedding - .iter() - .flat_map(|value| value.to_le_bytes()) - .collect() -} - -/// `HSET` a fresh `:` document, then `EXPIRE` it when a TTL is set. -pub(crate) fn write_document( - connection: &mut ConnectionRef<'_>, - index: &IndexState, - scope: &str, - prompt: &str, - response: Vec, - vector: Vec, - ttl: Option, -) -> Result<(), Error> { - ensure_index(connection, index, vector.len() / size_of::())?; - let document = format!("{}{scope}:{}", index.prefix, Uuid::new_v4()); - let mut pipeline = redis::pipe(); - pipeline - .cmd("HSET") - .arg(&document) - .arg("litellm_cache_key") - .arg(scope) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(response) - .arg("embedding") - .arg(vector) - .ignore(); - if let Some(ttl) = ttl { - pipeline - .cmd("EXPIRE") - .arg(&document) - .arg(ttl.as_secs()) - .ignore(); - } - pipeline.query::<()>(connection).map_err(Error::unavailable) -} - -/// The KNN-1 search within `scope`: the closest document's similarity, and its stored response -/// when that similarity reaches the threshold. No document reads as a similarity of `0.0`. -pub(crate) fn search_document( - connection: &mut ConnectionRef<'_>, - index: &IndexState, - scope: &str, - vector: Vec, -) -> Result>, Error> { - ensure_index(connection, index, vector.len() / size_of::())?; - let query = - format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]"); - let response = redis::cmd("FT.SEARCH") - .arg(&index.name) - .arg(query) - .arg("PARAMS") - .arg(2) - .arg("vec") - .arg(vector) - .arg("RETURN") - .arg(2) - .arg("response") - .arg("vector_distance") - .arg("DIALECT") - .arg(2) - .query::(connection) - .map_err(Error::unavailable)?; - let Some(fields) = search_fields(response)? else { - return Ok(SemanticLookup::miss(Some(0.0))); - }; - let field = |name: &str| { - fields - .iter() - .find_map(|(field, value)| (field == name).then(|| value.clone())) - .ok_or(Error::InvalidEntry(None)) - }; - let response = field("response")?; - let similarity = 1.0 - parse_f64(&field("vector_distance")?)?; - Ok(SemanticLookup { - value: (similarity >= index.similarity_threshold).then_some(response), - similarity: Some(similarity), - }) -} - -type SearchFields = Vec<(String, Vec)>; - -fn search_fields(value: redis::Value) -> Result, Error> { - let redis::Value::Array(values) = value else { - return Err(Error::InvalidEntry(None)); - }; - let total = parse_i64(values.first().ok_or(Error::InvalidEntry(None))?)?; - if total <= 0 || values.len() < 3 { - return Ok(None); - } - let redis::Value::Array(fields) = &values[2] else { - return Err(Error::InvalidEntry(None)); - }; - let (pairs, remainder) = fields.as_chunks::<2>(); - if !remainder.is_empty() { - return Err(Error::InvalidEntry(None)); - } - pairs - .iter() - .map(|pair| { - Ok(( - value_text(&pair[0]).ok_or(Error::InvalidEntry(None))?, - value_bytes(&pair[1])?, - )) - }) - .collect::, Error>>() - .map(Some) -} - -fn parse_i64(value: &redis::Value) -> Result { - value_text(value) - .ok_or(Error::InvalidEntry(None))? - .parse() - .map_err(Error::invalid_entry) -} - -fn parse_f64(value: &[u8]) -> Result { - std::str::from_utf8(value) - .map_err(Error::invalid_entry)? - .parse() - .map_err(Error::invalid_entry) -} - -pub(crate) fn value_text(value: &redis::Value) -> Option { - match value { - redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), - redis::Value::SimpleString(value) => Some(value.clone()), - redis::Value::Int(value) => Some(value.to_string()), - _ => None, - } -} - -fn value_bytes(value: &redis::Value) -> Result, Error> { - match value { - redis::Value::BulkString(bytes) => Ok(bytes.clone()), - redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()), - redis::Value::Int(value) => Ok(value.to_string().into_bytes()), - _ => Err(Error::InvalidEntry(None)), - } -} diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs b/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs deleted file mode 100644 index 10e32c0183b..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs +++ /dev/null @@ -1,430 +0,0 @@ -mod support; - -use std::{ - sync::{Arc, Mutex}, - time::Duration, -}; - -use litellm_cache::{ - BaseCache, Error, JsonCodec, SemanticCacheContext, - semantic::{PreparedEmbedding, SemanticCache, SemanticLookup}, -}; -use litellm_cache_valkey_semantic::{ - DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig, -}; -use redis_test::MockRedisConnection; -use rstest::{fixture, rstest}; -use serde_json::{Value, json}; -use support::{EmbedCalls, FakeEmbedder, RecordingConnection}; - -type Requests = Arc>>>; -type RecordingCache = ValkeySemanticCache, RecordingConnection>; - -const KEY_SCOPE: &str = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683"; - -struct Recording { - cache: RecordingCache, - requests: Requests, - calls: EmbedCalls, -} - -impl Recording { - fn text(&self) -> String { - self.requests - .lock() - .unwrap() - .iter() - .map(|request| String::from_utf8_lossy(request).into_owned()) - .collect::>() - .join("\n") - } -} - -fn config() -> ValkeySemanticConfig { - ValkeySemanticConfig { - similarity_threshold: 0.8, - index_name: "test".into(), - } -} - -fn recording(replies: impl IntoIterator>) -> Recording { - let connection = RecordingConnection::new(replies); - let requests = connection.requests(); - let embedder = FakeEmbedder::new(&[]); - let calls = Arc::clone(&embedder.calls); - Recording { - cache: ValkeySemanticCache::with_connection( - connection, - embedder, - JsonCodec::new(), - config(), - ), - requests, - calls, - } -} - -#[fixture] -fn entry() -> Value { - json!({"timestamp": 1.0, "response": {"answer": "ok"}}) -} - -#[fixture] -fn context() -> SemanticCacheContext { - SemanticCacheContext { - messages: Some(json!([{"role": "user", "content": "hello"}])), - metadata: Some(json!({"source": "test"})), - ..Default::default() - } -} - -fn ok() -> redis::RedisResult { - Ok(redis::Value::SimpleString("OK".into())) -} - -fn already_exists() -> redis::RedisResult { - Err(support::server_error("Index already exists")) -} - -fn bulk(value: &[u8]) -> redis::Value { - redis::Value::BulkString(value.to_vec()) -} - -fn search_hit(response: &[u8], distance: &str) -> redis::Value { - redis::Value::Array(vec![ - redis::Value::Int(1), - bulk(b"test:document"), - redis::Value::Array(vec![ - bulk(b"response"), - bulk(response), - bulk(b"vector_distance"), - bulk(distance.as_bytes()), - ]), - ]) -} - -fn encoded(value: &Value) -> Vec { - serde_json::to_vec(value).unwrap() -} - -/// `FT.INFO` with the vector field's dimension nested one level down, as valkey-search reports. -fn nested_dimension_info(dimension: i64) -> redis::Value { - redis::Value::Array(vec![ - redis::Value::SimpleString("attributes".into()), - redis::Value::Array(vec![redis::Value::Array(vec![ - redis::Value::SimpleString("embedding".into()), - redis::Value::Array(vec![ - redis::Value::SimpleString("dimensions".into()), - redis::Value::Int(dimension), - ]), - ])]), - ]) -} - -/// `FT.INFO` with `dimensions` as a sibling string of the identifier. -fn flat_dimension_info(dimension: &str) -> redis::Value { - redis::Value::Array(vec![ - redis::Value::SimpleString("attributes".into()), - redis::Value::Array(vec![redis::Value::Array(vec![ - redis::Value::SimpleString("identifier".into()), - redis::Value::SimpleString("embedding".into()), - redis::Value::Array(vec![ - redis::Value::SimpleString("dimensions".into()), - redis::Value::SimpleString(dimension.into()), - ]), - ])]), - ]) -} - -#[rstest] -#[case::string_content(json!([{"content": "hello"}]), None, Some("hello"))] -#[case::text_parts(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))] -#[case::non_object_parts_are_skipped(json!([{"content": ["raw", {"text": "hello"}]}]), None, Some("hello"))] -#[case::search_results(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))] -#[case::responses_string_input(json!([]), Some(json!(" hello ")), Some("hello"))] -#[case::responses_item_input(json!([]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))] -#[case::blank_input(json!([]), Some(json!(" ")), None)] -fn prompt_shapes_follow_redis_semantic_extraction( - #[case] messages: Value, - #[case] input: Option, - #[case] expected: Option<&str>, -) { - let recording = recording([ok(), Ok(redis::Value::Array(vec![redis::Value::Int(0)]))]); - let context = SemanticCacheContext { - messages: Some(messages), - input, - ..Default::default() - }; - - assert_eq!(recording.cache.get_cache("key", &context).unwrap(), None); - - let prompts = recording - .calls - .lock() - .unwrap() - .iter() - .map(|(prompt, _)| prompt.clone()) - .collect::>(); - assert_eq!( - prompts, - expected.into_iter().map(str::to_owned).collect::>() - ); - assert_eq!( - recording.requests.lock().unwrap().is_empty(), - expected.is_none() - ); -} - -#[rstest] -#[case::key("key", KEY_SCOPE)] -#[case::empty_key("", "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")] -fn documents_are_scoped_by_the_keys_sha256( - #[case] key: &str, - #[case] scope: &str, - entry: Value, - context: SemanticCacheContext, -) { - let recording = recording([ok()]); - - recording.cache.set_cache(key, entry, &context).unwrap(); - - let text = recording.text(); - assert!(text.contains(&format!("test:{scope}:"))); - assert!(text.contains(&format!("litellm_cache_key\r\n$64\r\n{scope}\r\n"))); -} - -#[rstest] -#[case::no_ttl(None, None)] -#[case::whole_seconds(Some(Duration::from_secs(5)), Some("5"))] -#[case::fractional_seconds_truncate(Some(Duration::from_millis(1900)), Some("1"))] -fn set_writes_hset_and_expires_only_with_a_ttl( - #[case] ttl: Option, - #[case] expire: Option<&str>, - entry: Value, - context: SemanticCacheContext, -) { - let recording = recording([ok()]); - - recording - .cache - .set_cache( - "key", - entry, - &SemanticCacheContext { - ttl, - ..context.clone() - }, - ) - .unwrap(); - - let text = recording.text(); - assert!(text.contains("FT.CREATE")); - assert!(text.contains("HSET")); - let expire_seconds = text - .split_once("EXPIRE\r\n") - .and_then(|(_, rest)| rest.split("\r\n").nth(3)); - assert_eq!(expire_seconds, expire); - assert_eq!( - *recording.calls.lock().unwrap(), - vec![("hello".to_owned(), context.metadata)] - ); -} - -#[rstest] -fn second_set_skips_create_after_dimension_is_cached(entry: Value, context: SemanticCacheContext) { - let recording = recording([ok()]); - - recording - .cache - .set_cache("key", entry.clone(), &context) - .unwrap(); - recording.cache.set_cache("key", entry, &context).unwrap(); - - let text = recording.text(); - assert_eq!(text.matches("FT.CREATE").count(), 1); - assert_eq!(text.matches("HSET").count(), 2); -} - -#[rstest] -#[case::nested_matching(nested_dimension_info(3), Ok(()))] -#[case::nested_mismatch( - nested_dimension_info(2), - Err("cache is unavailable: index exists with dimension 2, not 3") -)] -#[case::flat_matching(flat_dimension_info("3"), Ok(()))] -#[case::flat_mismatch( - flat_dimension_info("2"), - Err("cache is unavailable: index exists with dimension 2, not 3") -)] -#[case::unreported_dimension_is_accepted(redis::Value::Array(vec![]), Ok(()))] -fn existing_index_dimension_must_match_embedding( - #[case] info: redis::Value, - #[case] expected: Result<(), &str>, - entry: Value, - context: SemanticCacheContext, -) { - let recording = recording([already_exists(), Ok(info)]); - - let result = recording.cache.set_cache("key", entry, &context); - assert_eq!( - result.map_err(|error| litellm_error::Report::new(&error).to_string()), - expected.map_err(str::to_owned) - ); -} - -#[rstest] -#[case::create_failure(Err(redis::RedisError::from((redis::ErrorKind::Io, "boom"))))] -fn index_creation_failures_are_unavailable( - #[case] reply: redis::RedisResult, - entry: Value, - context: SemanticCacheContext, -) { - let recording = recording([reply]); - - assert!(matches!( - recording.cache.set_cache("key", entry, &context), - Err(Error::Unavailable(Some(_))) - )); -} - -#[rstest] -#[case::within_threshold(search_hit(&encoded(&entry()), "0.1"), Ok(Some(entry())))] -#[case::at_threshold(search_hit(&encoded(&entry()), "0.2"), Ok(Some(entry())))] -#[case::beyond_threshold(search_hit(&encoded(&entry()), "0.5"), Ok(None))] -#[case::zero_documents(redis::Value::Array(vec![redis::Value::Int(0)]), Ok(None))] -#[case::missing_response( - redis::Value::Array(vec![ - redis::Value::Int(1), - bulk(b"document"), - redis::Value::Array(vec![bulk(b"vector_distance"), bulk(b"0.1")]), - ]), - Err(Error::InvalidEntry(None)) -)] -#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry(None)))] -#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry(None)))] -fn get_applies_threshold_and_decodes_entry( - #[case] reply: redis::Value, - #[case] expected: Result, Error>, - context: SemanticCacheContext, -) { - let recording = recording([ok(), Ok(reply)]); - - let kind = |result: Result, Error>| { - result.map_err(|error| std::mem::discriminant(&error)) - }; - assert_eq!( - kind(recording.cache.get_cache("key", &context)), - kind(expected) - ); -} - -#[rstest] -#[case::hit(context(), Some(search_hit(&encoded(&entry()), "0.1")), Some(entry()), Some(1.0 - 0.1))] -#[case::below_threshold(context(), Some(search_hit(&encoded(&entry()), "0.5")), None, Some(1.0 - 0.5))] -#[case::no_results(context(), Some(redis::Value::Array(vec![redis::Value::Int(0)])), None, Some(0.0))] -#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))] -#[tokio::test] -async fn lookup_reports_python_semantic_similarity( - #[case] context: SemanticCacheContext, - #[case] reply: Option, - #[case] value: Option, - #[case] similarity: Option, - #[values(false, true)] use_async: bool, -) { - let searched = reply.is_some(); - let recording = recording(reply.map_or_else(Vec::new, |reply| vec![ok(), Ok(reply)])); - - let lookup = if use_async { - recording - .cache - .async_get_cache_with_similarity("key", &context) - .await - } else { - recording.cache.get_cache_with_similarity("key", &context) - }; - - assert_eq!(lookup, Ok(SemanticLookup { value, similarity })); - assert_eq!(recording.text().contains("FT.SEARCH"), searched); -} - -#[rstest] -#[tokio::test] -async fn missing_prompt_does_not_touch_valkey(entry: Value) { - let cache = ValkeySemanticCache::with_connection( - MockRedisConnection::new([]).assert_all_commands_consumed(), - FakeEmbedder::new(&[]), - JsonCodec::new(), - config(), - ); - let context = SemanticCacheContext::default(); - - cache.set_cache("key", entry.clone(), &context).unwrap(); - assert_eq!(cache.get_cache("key", &context).unwrap(), None); - cache - .async_set_cache("key", entry, context.clone()) - .await - .unwrap(); - assert_eq!(cache.async_get_cache("key", &context).await.unwrap(), None); - assert_eq!(cache.get_ttl(&context), None); -} - -#[rstest] -fn with_embedder_shares_index_state_and_connections(entry: Value, context: SemanticCacheContext) { - let recording = recording([ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]); - recording - .cache - .set_cache("key", entry.clone(), &context) - .unwrap(); - - let prepared = recording - .cache - .with_embedder(PreparedEmbedding(vec![0.1, 0.2, 0.3])); - - assert_eq!(prepared.get_cache("key", &context).unwrap(), Some(entry)); - assert_eq!(recording.text().matches("FT.CREATE").count(), 1); -} - -#[rstest] -fn accessors_report_the_config() { - let recording = recording([]); - - assert_eq!(recording.cache.index_name(), "test"); - assert_eq!(recording.cache.similarity_threshold(), 0.8); - assert_eq!(DEFAULT_INDEX_NAME, "litellm_semantic_cache_index"); -} - -#[rstest] -#[tokio::test] -async fn async_set_and_get_use_shared_document_helpers( - entry: Value, - context: SemanticCacheContext, -) { - let recording = recording([ok(), ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]); - let context = SemanticCacheContext { - ttl: Some(Duration::from_millis(1900)), - ..context - }; - - recording - .cache - .async_set_cache("key", entry.clone(), context.clone()) - .await - .unwrap(); - assert_eq!( - recording - .cache - .async_get_cache("key", &context) - .await - .unwrap(), - Some(entry) - ); - - let text = recording.text(); - assert!(text.contains("FT.CREATE")); - assert!(text.contains("HSET")); - assert!(text.contains("EXPIRE")); - assert_eq!( - *recording.calls.lock().unwrap(), - vec![("hello".to_owned(), context.metadata.clone()); 2] - ); -} diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/contract.rs b/litellm-rust/crates/cache-valkey-semantic/tests/contract.rs deleted file mode 100644 index 71081fca92b..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/tests/contract.rs +++ /dev/null @@ -1,62 +0,0 @@ -//! `overwrite_replaces` does not apply: like Python, every write is a new `:` -//! document, so a second write with the same prompt adds a tie instead of replacing the first. - -mod support; - -use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding}; -use litellm_cache_testing as contract; -use litellm_cache_valkey_semantic::{ - DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig, -}; -use rstest::{fixture, rstest}; -use serde_json::{Value, json}; -use support::FakeSearch; - -type Cache = ValkeySemanticCache, FakeSearch>; - -const PREFIX: &str = "contract:"; - -#[fixture] -fn cache() -> Cache { - ValkeySemanticCache::with_connection( - FakeSearch::default(), - PreparedEmbedding(vec![0.6, 0.8]), - JsonCodec::new(), - ValkeySemanticConfig { - similarity_threshold: 0.9, - index_name: DEFAULT_INDEX_NAME.into(), - }, - ) -} - -#[fixture] -fn context() -> SemanticCacheContext { - SemanticCacheContext { - messages: Some(json!([{"role": "user", "content": "contract prompt"}])), - ..Default::default() - } -} - -#[rstest] -#[tokio::test] -async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) { - contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await; -} - -#[rstest] -#[tokio::test] -async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) { - contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await; -} - -#[rstest] -#[tokio::test] -async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) { - contract::pipeline_writes_every_entry( - &cache, - context, - PREFIX, - vec![json!("a"), json!(2), json!({"c": true})], - ) - .await; -} diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs deleted file mode 100644 index 2af05943c98..00000000000 --- a/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs +++ /dev/null @@ -1,352 +0,0 @@ -#![allow(dead_code)] - -use std::{ - collections::{BTreeMap, HashMap, VecDeque}, - sync::{Arc, Mutex}, -}; - -use litellm_cache::{Error, semantic::Embedder}; -use serde_json::Value; - -pub type EmbedCalls = Arc)>>>; - -/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every -/// prompt with its metadata. -pub struct FakeEmbedder { - vectors: HashMap>, - pub calls: EmbedCalls, -} - -impl FakeEmbedder { - pub fn new(vectors: &[(&str, &[f32])]) -> Self { - Self { - vectors: vectors - .iter() - .map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec())) - .collect(), - calls: EmbedCalls::default(), - } - } -} - -impl Embedder for FakeEmbedder { - fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { - self.calls - .lock() - .unwrap() - .push((prompt.to_owned(), metadata.cloned())); - 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) - } -} - -struct FakeIndex { - prefix: Vec, - dims: usize, - vector_field: String, -} - -#[derive(Default)] -struct SearchState { - indexes: HashMap, - hashes: BTreeMap, BTreeMap>>, -} - -/// An in-memory valkey-search speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache -/// sends, with exact cosine KNN over the hashes under an index prefix. -#[derive(Clone, Default)] -pub struct FakeSearch { - state: Arc>, -} - -impl FakeSearch { - fn run(&self, args: Vec>) -> redis::RedisResult { - let mut state = self.state.lock().unwrap(); - let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned(); - match text(0).to_uppercase().as_str() { - "FT.CREATE" => { - let name = text(1); - if state.indexes.contains_key(&name) { - return Err(server_error("Index already exists")); - } - let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes()); - let prefix = args[position("PREFIX").unwrap() + 2].clone(); - let dims = text(position("DIM").unwrap() + 1).parse().unwrap(); - let vector_field = text(position("VECTOR").unwrap() - 1); - state.indexes.insert( - name, - FakeIndex { - prefix, - dims, - vector_field, - }, - ); - Ok(redis::Value::Okay) - } - "FT.INFO" => { - let index = state - .indexes - .get(&text(1)) - .ok_or_else(|| server_error("Unknown index name"))?; - Ok(index_info(index)) - } - "FT.DROPINDEX" => { - state.indexes.remove(&text(1)); - Ok(redis::Value::Okay) - } - "HSET" => { - let hash = state.hashes.entry(args[1].clone()).or_default(); - for pair in args[2..].chunks(2) { - hash.insert( - String::from_utf8_lossy(&pair[0]).into_owned(), - pair[1].clone(), - ); - } - Ok(redis::Value::Int(((args.len() - 2) / 2) as i64)) - } - "EXPIRE" => Ok(redis::Value::Int(i64::from( - state.hashes.contains_key(&args[1]), - ))), - "FT.SEARCH" => { - let index = state - .indexes - .get(&text(1)) - .ok_or_else(|| server_error("no such index"))?; - let query = text(2); - let tag = query_tag(&query); - let params = args.iter().position(|arg| arg == b"PARAMS").unwrap(); - let vector = floats(&args[params + 3]); - let best = state - .hashes - .iter() - .filter(|(key, _)| key.starts_with(&index.prefix)) - .filter(|(_, fields)| { - fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes()) - }) - .filter_map(|(key, fields)| { - let stored = floats(fields.get(&index.vector_field)?); - (stored.len() == index.dims) - .then(|| (key, fields, 1.0 - cosine(&vector, &stored))) - }) - .min_by(|left, right| left.2.total_cmp(&right.2)); - let Some((key, fields, distance)) = best else { - return Ok(redis::Value::Array(vec![redis::Value::Int(0)])); - }; - let mut reply = fields - .iter() - .filter(|(name, _)| **name != index.vector_field) - .flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)]) - .collect::>(); - reply.extend([ - bulk(b"vector_distance"), - bulk(distance.to_string().as_bytes()), - ]); - Ok(redis::Value::Array(vec![ - redis::Value::Int(1), - bulk(key), - redis::Value::Array(reply), - ])) - } - "PING" => Ok(redis::Value::SimpleString("PONG".into())), - _ => Err(server_error("unsupported command")), - } - } -} - -impl redis::ConnectionLike for FakeSearch { - fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult { - let mut commands = parse_commands(command); - self.run(commands.remove(0)) - } - - fn req_packed_commands( - &mut self, - commands: &[u8], - offset: usize, - count: usize, - ) -> redis::RedisResult> { - let replies = parse_commands(commands) - .into_iter() - .map(|args| self.run(args)) - .collect::>>()?; - Ok(replies.into_iter().skip(offset).take(count).collect()) - } - - fn get_db(&self) -> i64 { - 0 - } - - fn check_connection(&mut self) -> bool { - true - } - - fn is_open(&self) -> bool { - true - } -} - -pub fn server_error(reply: &str) -> redis::RedisError { - redis::parse_redis_value(format!("-{reply}\r\n").as_bytes()) - .expect("error reply") - .extract_error() - .expect_err("server error") -} - -fn bulk(bytes: &[u8]) -> redis::Value { - redis::Value::BulkString(bytes.to_vec()) -} - -fn index_info(index: &FakeIndex) -> redis::Value { - redis::Value::Array(vec![ - bulk(b"index_name"), - bulk(b"fake"), - bulk(b"attributes"), - redis::Value::Array(vec![ - redis::Value::Array(vec![ - bulk(b"identifier"), - bulk(b"litellm_cache_key"), - bulk(b"type"), - bulk(b"TAG"), - ]), - redis::Value::Array(vec![ - bulk(b"identifier"), - bulk(index.vector_field.as_bytes()), - bulk(b"type"), - bulk(b"VECTOR"), - bulk(b"index"), - redis::Value::Array(vec![ - bulk(b"dimensions"), - redis::Value::Int(index.dims as i64), - ]), - ]), - ]), - ]) -} - -/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed. -fn query_tag(query: &str) -> String { - let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len(); - let mut tag = String::new(); - let mut characters = query[start..].chars(); - while let Some(character) = characters.next() { - match character { - '\\' => tag.extend(characters.next()), - '}' => break, - character => tag.push(character), - } - } - tag -} - -fn floats(bytes: &[u8]) -> Vec { - bytes - .as_chunks::<4>() - .0 - .iter() - .map(|chunk| f32::from_le_bytes(*chunk)) - .collect() -} - -fn cosine(left: &[f32], right: &[f32]) -> f64 { - let dot = left - .iter() - .zip(right) - .map(|(left, right)| f64::from(*left) * f64::from(*right)) - .sum::(); - let norm = |vector: &[f32]| { - vector - .iter() - .map(|value| f64::from(*value).powi(2)) - .sum::() - .sqrt() - }; - dot / (norm(left) * norm(right)) -} - -/// Splits a packed RESP request into each command's arguments. -fn parse_commands(mut bytes: &[u8]) -> Vec>> { - let line = |bytes: &mut &[u8]| { - let end = bytes - .windows(2) - .position(|window| window == b"\r\n") - .unwrap(); - let text = String::from_utf8(bytes[1..end].to_vec()).unwrap(); - *bytes = &bytes[end + 2..]; - text.parse::().unwrap() - }; - let mut commands = Vec::new(); - while !bytes.is_empty() { - let count = line(&mut bytes); - let mut args = Vec::with_capacity(count); - for _ in 0..count { - let length = line(&mut bytes); - args.push(bytes[..length].to_vec()); - bytes = &bytes[length + 2..]; - } - commands.push(args); - } - commands -} - -/// Records every packed request and answers from a script, `OK` once the script runs out. -pub struct RecordingConnection { - requests: Arc>>>, - replies: Mutex>>, -} - -impl RecordingConnection { - pub fn new(replies: impl IntoIterator>) -> Self { - Self { - requests: Arc::default(), - replies: Mutex::new(replies.into_iter().collect()), - } - } - - pub fn requests(&self) -> Arc>>> { - Arc::clone(&self.requests) - } - - fn reply(&self) -> redis::RedisResult { - self.replies - .lock() - .unwrap() - .pop_front() - .unwrap_or_else(|| Ok(redis::Value::SimpleString("OK".into()))) - } -} - -impl redis::ConnectionLike for RecordingConnection { - fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult { - self.requests.lock().unwrap().push(command.to_vec()); - self.reply() - } - - fn req_packed_commands( - &mut self, - command: &[u8], - _offset: usize, - count: usize, - ) -> redis::RedisResult> { - self.requests.lock().unwrap().push(command.to_vec()); - (0..count).map(|_| self.reply()).collect() - } - - fn get_db(&self) -> i64 { - 0 - } - - fn check_connection(&mut self) -> bool { - true - } - - fn is_open(&self) -> bool { - true - } -} diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index e3ef45f7876..fd826f549e5 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -29,8 +29,6 @@ bytes.workspace = true futures-util.workspace = true litellm-cache.workspace = true litellm-error.workspace = true -litellm-cache-memory.workspace = true -litellm-cache-redis.workspace = true litellm-cache-response.workspace = true serde.workspace = true litellm-auth.workspace = true @@ -56,7 +54,6 @@ pyo3.workspace = true prost.workspace = true pyo3-async-runtimes.workspace = true reqwest.workspace = true -redis = { workspace = true, features = ["tls-rustls"] } serde_json.workspace = true strum.workspace = true veil.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md index 41c82187373..e79562b8762 100644 --- a/litellm-rust/crates/python-bridge/src/cache/AGENTS.md +++ b/litellm-rust/crates/python-bridge/src/cache/AGENTS.md @@ -1,7 +1,7 @@ # Cache boundary -This folder owns how Rust inference reaches the selected cache: global cache selection, route admission, delegation to a Python cache and the experimental V2 native handles. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates +This folder owns how Rust inference reaches the selected Python cache. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates -`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache, and owns route admission and inference protocol composition for both adapters +`mod.rs` exposes the cache boundary to routes. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache and owns route admission and inference protocol composition -`python/` delegates operations to the selected Python cache without discovering configuration. `native/` owns the experimental V2 native cache handles. Neither adapter depends on shared selection or the other adapter. Shared composition depends on the adapters, and routes use only the parent module's exports +`python/` delegates operations to the selected Python cache without discovering configuration. Routes use only the parent module's exports diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 87dba3efcb1..4e3847951d8 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,32 +1,15 @@ -mod native; mod python; mod selection; -pub(crate) use native::NativeCacheHandle; pub(crate) use python::{CacheCall, PythonCache}; -pub(crate) use selection::{Cached, Selection, configure, configured_native}; +pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled}; -use litellm_cache::Error; -use pyo3::{ - exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, - prelude::*, -}; +use pyo3::{exceptions::PyRuntimeError, prelude::*}; -/// A cache call the bridge rejects itself, before or after the cache ran. #[derive(Debug, thiserror::Error)] pub(crate) enum HandleError { - #[error("recreate the v2 cache after fork")] - Forked, - #[error("cache limits must be positive")] - Limits, - #[error("namespace and a positive cache limit are required")] - RedisLimits, - #[error("cache durations must be finite and positive")] - Duration(#[source] Option), - #[error("unsupported v2 cache control: {0}")] - Control(String), - #[error("native inference needs a v2 cache, the catalog keeps Python caches on Python")] - NotNative, + #[error("this inference route does not support configured caches")] + Unsupported, #[error("cache operation without configured cache")] Unconfigured, #[error("cache arguments unavailable")] @@ -39,66 +22,10 @@ impl From for PyErr { fn from(error: HandleError) -> Self { let message = litellm_error::Report::new(&error).to_string(); match error { - HandleError::Limits - | HandleError::RedisLimits - | HandleError::Duration(_) - | HandleError::Control(_) => PyValueError::new_err(message), - HandleError::Forked - | HandleError::NotNative + HandleError::Unsupported | HandleError::Unconfigured | HandleError::MissingArguments | HandleError::NothingPending => PyRuntimeError::new_err(message), } } } - -fn cache_error(error: Error) -> PyErr { - let message = litellm_error::Report::new(&error).to_string(); - match error { - Error::InvalidEntry(_) => PyValueError::new_err(message), - Error::UnsupportedOperation => PyNotImplementedError::new_err(message), - Error::Unavailable(_) | Error::UnscopedFlush | Error::MissingPrompt => { - PyRuntimeError::new_err(message) - } - } -} - -#[cfg(test)] -mod tests { - use pyo3::{ - exceptions::{PyRuntimeError, PyValueError}, - prelude::*, - }; - use rstest::rstest; - - use super::{Error, cache_error}; - - fn refused() -> std::io::Error { - std::io::Error::other("connection refused") - } - - #[rstest] - #[case::unavailable( - Error::unavailable(refused()), - false, - "cache is unavailable: connection refused" - )] - #[case::invalid_entry( - Error::invalid_entry(refused()), - true, - "invalid cache entry: connection refused" - )] - fn python_sees_the_class_and_the_whole_chain( - #[case] error: Error, - #[case] caller_fault: bool, - #[case] message: &str, - ) { - Python::initialize(); - Python::attach(|py| { - let failure = cache_error(error); - assert_eq!(failure.is_instance_of::(py), caller_fault); - assert_eq!(failure.is_instance_of::(py), !caller_fault); - assert_eq!(failure.value(py).to_string(), message); - }); - } -} diff --git a/litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md b/litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md deleted file mode 100644 index 09a34d188cd..00000000000 --- a/litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md +++ /dev/null @@ -1,9 +0,0 @@ -# Native cache bindings - -This directory exposes the experimental V2 native cache handles (`NativeCacheHandle`) to Python. Cache algorithms and storage protocols remain in their cache crates - -Accept the cache object selected by the parent module. Do not read global `litellm.cache`, decide route admission, or select the Python cache adapter here - -Keep Python-facing cache classes and method signatures stable when reorganizing modules. Native internals stay private to this directory unless the shared cache boundary or Python module registration needs them - -Verify changes with the V2 cache tests using a freshly built extension. Test cache behavior, not module paths or file structure diff --git a/litellm-rust/crates/python-bridge/src/cache/native/mod.rs b/litellm-rust/crates/python-bridge/src/cache/native/mod.rs deleted file mode 100644 index da1460f90ed..00000000000 --- a/litellm-rust/crates/python-bridge/src/cache/native/mod.rs +++ /dev/null @@ -1,3 +0,0 @@ -pub(super) mod v2; - -pub(crate) use v2::NativeCacheHandle; diff --git a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs b/litellm-rust/crates/python-bridge/src/cache/native/v2.rs deleted file mode 100644 index c54cb641843..00000000000 --- a/litellm-rust/crates/python-bridge/src/cache/native/v2.rs +++ /dev/null @@ -1,336 +0,0 @@ -use crate::cache::{HandleError, cache_error}; -use std::{ - sync::Arc, - time::{Duration, SystemTime, UNIX_EPOCH}, -}; - -use litellm_cache::{DeleteCache, DisconnectCache, PingCache}; -use litellm_host_python::{from_py, release_gil, to_py}; -use serde_json::Value; - -use litellm_cache_memory::InMemoryCache; -use litellm_cache_redis::{RedisCache, RedisTopology}; -use litellm_cache_response::{ - CacheEntry, CacheKeyInput, ExactResponseCache, ResponseCache, ResponseCacheCodec, - ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService, -}; -use pyo3::{prelude::*, types::PyDict}; - -#[pyclass( - frozen, - name = "NativeCacheHandle", - module = "litellm.rust_bridge._native" -)] -pub(crate) struct NativeCacheHandle { - service: Arc, - backend: Arc, - storage: Storage, - pid: u32, -} - -#[derive(Clone)] -enum Storage { - Memory(Arc>), - Redis(Arc>), -} - -impl NativeCacheHandle { - fn check_process(&self) -> PyResult<()> { - if self.pid != std::process::id() { - return Err(HandleError::Forked.into()); - } - Ok(()) - } -} - -fn request(key: String, ttl: Option) -> PyResult { - let mut request: ResponseCacheRequest = ResponseCacheRequest::new(CacheKeyInput { - preset: Some(key), - ..Default::default() - }); - request.context.ttl = ttl.map(duration).transpose()?; - Ok(request) -} - -#[pymethods] -impl NativeCacheHandle { - #[staticmethod] - #[pyo3(signature = (*, ttl=600.0, capacity=200, max_entry_bytes=4194304))] - fn memory(ttl: f64, capacity: usize, max_entry_bytes: usize) -> PyResult { - let ttl = duration(ttl)?; - if capacity == 0 || max_entry_bytes == 0 { - return Err(HandleError::Limits.into()); - } - let storage = Arc::new(InMemoryCache::new(Some(capacity), Some(ttl))); - let backend = Arc::new(ResponseCache::new(storage.clone()).with_config( - ResponseCacheConfig { - namespace: "sdk".into(), - max_entry_bytes, - }, - )); - Ok(Self { - service: backend.clone(), - backend, - storage: Storage::Memory(storage), - pid: std::process::id(), - }) - } - - #[staticmethod] - #[pyo3(signature = (url, *, namespace, ttl=600.0, max_entry_bytes=4194304))] - fn redis( - py: Python<'_>, - url: &str, - namespace: String, - ttl: f64, - max_entry_bytes: usize, - ) -> PyResult { - let ttl = duration(ttl)?; - if namespace.is_empty() || max_entry_bytes == 0 { - return Err(HandleError::RedisLimits.into()); - } - let storage = Arc::new( - release_gil(py, || { - RedisCache::connect( - url, - &RedisTopology::Standalone, - Some(ttl), - ResponseCacheCodec, - ) - }) - .map_err(cache_error)? - .with_namespace(Some(namespace.clone())), - ); - let backend = Arc::new(ResponseCache::new(storage.clone()).with_config( - ResponseCacheConfig { - namespace, - max_entry_bytes, - }, - )); - Ok(Self { - service: backend.clone(), - backend, - storage: Storage::Redis(storage), - pid: std::process::id(), - }) - } - fn get(&self, py: Python<'_>, key: String) -> PyResult> { - self.check_process()?; - let request = request(key, None)?; - let value = - release_gil(py, || self.backend.lookup(&request, now())).map_err(cache_error)?; - to_py(py, &value) - } - - #[pyo3(signature = (key, value, *, ttl=None))] - fn set( - &self, - py: Python<'_>, - key: String, - value: &Bound<'_, PyAny>, - ttl: Option, - ) -> PyResult<()> { - self.check_process()?; - let request = request(key, ttl)?; - let value: Value = from_py(value)?; - release_gil(py, || self.backend.store(&request, value, now())).map_err(cache_error) - } - - fn async_get<'py>(&self, py: Python<'py>, key: String) -> PyResult> { - self.check_process()?; - let request = request(key, None)?; - let backend = self.backend.clone(); - crate::execution::run_async( - py, - async move { backend.async_lookup(&request, now()).await }, - cache_error, - ) - } - - #[pyo3(signature = (key, value, *, ttl=None))] - fn async_set<'py>( - &self, - py: Python<'py>, - key: String, - value: &Bound<'_, PyAny>, - ttl: Option, - ) -> PyResult> { - self.check_process()?; - let request = request(key, ttl)?; - let value: Value = from_py(value)?; - let backend = self.backend.clone(); - crate::execution::run_async( - py, - async move { backend.async_store(&request, value, now()).await }, - cache_error, - ) - } - - #[pyo3(signature = (entries, *, ttl=None))] - fn async_set_many<'py>( - &self, - py: Python<'py>, - entries: &Bound<'_, PyAny>, - ttl: Option, - ) -> PyResult> { - self.check_process()?; - let entries: Vec<(String, Value)> = from_py(entries)?; - let entries = entries - .into_iter() - .map(|(key, value)| Ok((request(key, ttl)?, value))) - .collect::>>()?; - let backend = self.backend.clone(); - crate::execution::run_async( - py, - async move { backend.async_store_batch(entries, now()).await }, - cache_error, - ) - } - - fn flush(&self, py: Python<'_>) -> PyResult> { - self.check_process()?; - let backend = self.backend.clone(); - crate::execution::run_sync(py, async move { backend.async_flush().await }, cache_error) - } - - fn async_flush<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - let backend = self.backend.clone(); - crate::execution::run_async(py, async move { backend.async_flush().await }, cache_error) - } - - fn ping<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - let storage = self.storage.clone(); - crate::execution::run_async( - py, - async move { - match storage { - Storage::Memory(_) => Ok(true), - Storage::Redis(cache) => cache.ping().await, - } - }, - cache_error, - ) - } - - fn disconnect<'py>(&self, py: Python<'py>) -> PyResult> { - self.check_process()?; - let storage = self.storage.clone(); - crate::execution::run_async( - py, - async move { - match storage { - Storage::Memory(cache) => cache.disconnect().await, - Storage::Redis(cache) => cache.disconnect().await, - } - }, - cache_error, - ) - } - - fn delete<'py>(&self, py: Python<'py>, keys: Vec) -> PyResult> { - self.check_process()?; - let storage = self.storage.clone(); - crate::execution::run_async( - py, - async move { - for key in keys { - match &storage { - Storage::Memory(cache) => cache.async_delete_cache(&key).await?, - Storage::Redis(cache) => cache.async_delete_cache(&key).await?, - } - } - Ok(()) - }, - cache_error, - ) - } -} - -fn duration(seconds: f64) -> Result { - match Duration::try_from_secs_f64(seconds) { - Ok(duration) if !duration.is_zero() => Ok(duration), - Ok(_) => Err(HandleError::Duration(None)), - Err(error) => Err(HandleError::Duration(Some(error))), - } -} - -pub(in crate::cache) fn native_handle<'py>( - configured: &Bound<'py, PyAny>, -) -> PyResult>> { - Ok(configured - .getattr_opt("cache")? - .map(|backend| backend.getattr_opt("native_handle")) - .transpose()? - .flatten() - .filter(|handle| handle.is_instance_of::())) -} - -pub(in crate::cache) fn configured( - configured: &Bound<'_, PyAny>, - kwargs: &Bound<'_, PyDict>, -) -> PyResult<( - Option>, - litellm_cache_response::CacheOptions, -)> { - let handle = native_handle(configured)?.ok_or(HandleError::NotNative)?; - let cache = handle.extract::>()?; - cache.check_process()?; - let controls = kwargs.get_item("cache")?.filter(|value| !value.is_none()); - let controls = controls - .as_ref() - .map(|value| value.cast::()) - .transpose()?; - if let Some(controls) = controls { - for name in controls.keys() { - let name = name.extract::()?; - if !matches!( - name.as_str(), - "no-cache" | "no-store" | "ttl" | "s-maxage" | "s-max-age" | "use-cache" - ) { - return Err(HandleError::Control(name).into()); - } - } - } - let boolean = |name: &str| -> PyResult { - controls - .map(|values| values.get_item(name)) - .transpose()? - .flatten() - .map(|value| value.extract()) - .transpose() - .map(|value| value.unwrap_or(false)) - }; - let seconds = |name: &str| -> PyResult> { - controls - .map(|values| values.get_item(name)) - .transpose()? - .flatten() - .map(|value| Ok(duration(value.extract()?)?)) - .transpose() - }; - Ok(( - Some(cache.service.clone()), - litellm_cache_response::CacheOptions { - policy: litellm_cache_response::CachePolicy { - caching: kwargs - .get_item("caching")? - .filter(|value| !value.is_none()) - .map(|value| value.extract()) - .transpose()?, - no_cache: boolean("no-cache")?, - no_store: boolean("no-store")?, - ttl: seconds("ttl")?, - max_age: seconds("s-max-age")?.or(seconds("s-maxage")?), - }, - scope: litellm_cache_response::CacheScope::Shared, - }, - )) -} - -fn now() -> Duration { - SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() -} diff --git a/litellm-rust/crates/python-bridge/src/cache/selection.rs b/litellm-rust/crates/python-bridge/src/cache/selection.rs index 00fc7cd822f..e96f079e5da 100644 --- a/litellm-rust/crates/python-bridge/src/cache/selection.rs +++ b/litellm-rust/crates/python-bridge/src/cache/selection.rs @@ -1,14 +1,10 @@ -use super::{native, python}; -use litellm_cache_response::{ - CacheOptions, CachePolicy, CacheScope, ResponseCacheService, ScopedCache, -}; +use super::{HandleError, python}; +use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope, ScopedCache}; use litellm_host::{ machine::{HostServices, MachineFault}, protocol::Protocol, }; use pyo3::{prelude::*, types::PyDict}; -use std::sync::Arc; - const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache"; pub(crate) struct Cached

(std::marker::PhantomData

); @@ -24,7 +20,6 @@ impl Protocol for Cached

{ enum Backend { Disabled, - Native(Arc), Python { namespace: String }, } @@ -43,7 +38,6 @@ impl Selection { { let service = match self.backend { Backend::Disabled => None, - Backend::Native(service) => Some(service), Backend::Python { namespace } => Some(python::service(services, namespace)), }; ( @@ -76,37 +70,21 @@ fn selected_cache<'py>( Ok(Some(configured)) } -pub(crate) fn configured_native( +pub(crate) fn ensure_cache_disabled( py: Python<'_>, kwargs: &Bound<'_, PyDict>, call_type: &str, -) -> PyResult<( - Option>, - litellm_cache_response::CacheOptions, -)> { +) -> PyResult<()> { let Some(configured) = selected_cache(py, kwargs, call_type)? else { - return Ok(( - None, - litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared), - )); + return Ok(()); }; - native_configuration(&configured, kwargs) -} - -fn native_configuration( - configured: &Bound<'_, PyAny>, - kwargs: &Bound<'_, PyDict>, -) -> PyResult<(Option>, CacheOptions)> { - if !configured + if configured .call_method("should_use_cache", (), Some(kwargs))? .extract::()? { - return Ok(( - None, - litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared), - )); + return Err(HandleError::Unsupported.into()); } - native::v2::configured(configured, kwargs) + Ok(()) } pub(crate) fn configure( @@ -122,13 +100,6 @@ pub(crate) fn configure( options: CacheOptions::new(CacheScope::Shared), }); }; - if native::v2::native_handle(&cache)?.is_some() { - let (native, options) = native_configuration(&cache, arguments)?; - return Ok(Selection { - backend: native.map_or(Backend::Disabled, Backend::Native), - options, - }); - } let enabled = cache .call_method("should_use_cache", (), Some(arguments))? .extract::()?; diff --git a/litellm-rust/crates/python-bridge/src/lib.rs b/litellm-rust/crates/python-bridge/src/lib.rs index fdcb969a423..de35419308a 100644 --- a/litellm-rust/crates/python-bridge/src/lib.rs +++ b/litellm-rust/crates/python-bridge/src/lib.rs @@ -65,10 +65,6 @@ mod _native { fn init(module: &Bound<'_, PyModule>) -> PyResult<()> { let py = module.py(); let dict = module.dict(); - dict.set_item( - "NativeCacheHandle", - py.get_type::(), - )?; dict.set_item( "_SecretManagerRuntime", py.get_type::(), @@ -92,7 +88,6 @@ mod tests { Python::initialize(); Python::attach(|py| { let mut expected = vec![ - "NativeCacheHandle", "RustBridgeDeclined", "RustUpstreamError", "ForkedAfterNativeRuntimeStarted", diff --git a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs index d78d141f0f2..1bcdc18a765 100644 --- a/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use host::ChatCompletionsPythonHost; use litellm_auth::AuthServices; -use litellm_cache_response::{CachePolicy, ScopedCache}; +use litellm_cache_response::CachePolicy; use litellm_callbacks_legacy_python::LoggingOperation; use litellm_host::{call::HostedMachine, protocol::Protocol}; use litellm_inference_chat::{ChatCompletionsRoute, route::ChatCompletions}; @@ -57,10 +57,6 @@ impl InferenceRoute for ChatCompletionsRoute { Self::new(http, auth, secrets) } - fn with_cache(self, cache: ScopedCache) -> Self { - self.with_cache(cache) - } - fn machine( self, call: ::Request, diff --git a/litellm-rust/crates/python-bridge/src/routes/inference.rs b/litellm-rust/crates/python-bridge/src/routes/inference.rs index ca30c94b02f..1245c126c4e 100644 --- a/litellm-rust/crates/python-bridge/src/routes/inference.rs +++ b/litellm-rust/crates/python-bridge/src/routes/inference.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use litellm_auth::AuthServices; -use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache}; +use litellm_cache_response::CachePolicy; use litellm_callbacks_legacy_python::LoggingOperation; use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider; use litellm_host::{call::HostedMachine, protocol::Protocol}; @@ -142,7 +142,6 @@ pub(super) trait InferenceRoute: Sized + 'static { auth: Arc, secrets: Arc, ) -> Self; - fn with_cache(self, cache: ScopedCache) -> Self; fn machine( self, call: ::Request, @@ -176,12 +175,8 @@ where crate::http::resources().auth.clone(), crate::secrets::source(py)?, ); - let (cache, cache_options) = crate::cache::configured_native(py, arguments, call_type)?; - let route = match cache { - Some(cache) => route.with_cache(ScopedCache::new(cache, CacheScope::Shared)), - None => route, - }; - Ok(route.machine(request, cache_options.policy)) + crate::cache::ensure_cache_disabled(py, arguments, call_type)?; + Ok(route.machine(request, CachePolicy::default())) }, host, hooks, diff --git a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs index 9c045d65c13..cc4d85b4d26 100644 --- a/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs +++ b/litellm-rust/crates/python-bridge/src/routes/responses/mod.rs @@ -5,7 +5,7 @@ use std::sync::Arc; use host::ResponsesPythonHost; use litellm_auth::AuthServices; -use litellm_cache_response::{CachePolicy, ScopedCache}; +use litellm_cache_response::CachePolicy; use litellm_callbacks_legacy_python::LoggingOperation; use litellm_host::{call::HostedMachine, protocol::Protocol}; use litellm_inference_responses::{ResponsesRoute, route::Responses}; @@ -49,10 +49,6 @@ impl InferenceRoute for ResponsesRoute { Self::new(http, auth, secrets) } - fn with_cache(self, cache: ScopedCache) -> Self { - self.with_cache(cache) - } - fn machine( self, call: ::Request, diff --git a/litellm/_v2/AGENTS.md b/litellm/_v2/AGENTS.md deleted file mode 100644 index 55998814a7f..00000000000 --- a/litellm/_v2/AGENTS.md +++ /dev/null @@ -1,3 +0,0 @@ -Everything here is experimental and should not be documented - -Use this directory to explore alternative APIs where the Rust migration makes backward compatibility difficult. The gateway can use them for performance, but keep them behind the v2 flag to avoid breaking SDK users diff --git a/litellm/_v2/__init__.py b/litellm/_v2/__init__.py deleted file mode 100644 index 79156cf0f80..00000000000 --- a/litellm/_v2/__init__.py +++ /dev/null @@ -1,3 +0,0 @@ -from litellm._v2.cache import Cache - -__all__ = ("Cache",) diff --git a/litellm/_v2/cache/AGENTS.md b/litellm/_v2/cache/AGENTS.md deleted file mode 100644 index 308419b7253..00000000000 --- a/litellm/_v2/cache/AGENTS.md +++ /dev/null @@ -1,11 +0,0 @@ -# Python v2 cache - -Keep `litellm._v2.cache.Cache` import-compatible when reorganizing this package. This package owns Python cache factories and the adapter between the existing `BaseCache` interface and `NativeCacheHandle` - -Construct native handles at this boundary and inject the adapter through the existing cache facade's `_backend` parameter. Keep the facade's `type`, namespace, and TTL consistent with the configured native backend - -Keep storage implementation in the Rust storage crates and response-cache policy in `litellm-cache-response` and core. Do not duplicate cache-key generation, freshness rules, response encoding, or inference orchestration here - -Preserve synchronous and asynchronous cache operations, including TTL forwarding and lifecycle methods. Validate Python values before passing them to typed native interfaces. Keep native extension imports lazy so importing the package does not require loading the extension - -Extend the existing v2 cache tests in `tests/test_litellm_rust/test_v2.py` for behavioral changes, following that directory's `AGENTS.md`. Test observable cache behavior rather than package layout or implementation structure diff --git a/litellm/_v2/cache/__init__.py b/litellm/_v2/cache/__init__.py deleted file mode 100644 index 5a3860b7a6e..00000000000 --- a/litellm/_v2/cache/__init__.py +++ /dev/null @@ -1,72 +0,0 @@ -from __future__ import annotations - -from collections.abc import Sequence -from typing import TYPE_CHECKING, Final - -from pydantic import TypeAdapter - -from litellm.caching.base_cache import BaseCache -from litellm.caching.caching import Cache as CacheFacade -from litellm.types.caching import LiteLLMCacheType - -if TYPE_CHECKING: - from litellm.rust_bridge._native import NativeCacheHandle - -_DURATION: Final[TypeAdapter[float | None]] = TypeAdapter(float | None) - - -class NativeBackend(BaseCache): - def __init__(self, handle: NativeCacheHandle) -> None: - self.native_handle = handle - - def get_cache(self, key: str, **kwargs: object) -> object: - return self.native_handle.get(key) - - async def async_get_cache(self, key: str, **kwargs: object) -> object: - return await self.native_handle.async_get(key) - - def set_cache(self, key: str, value: object, **kwargs: object) -> None: - self.native_handle.set(key, value, ttl=_DURATION.validate_python(kwargs.get("ttl"))) - - async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None: - await self.native_handle.async_set(key, value, ttl=_DURATION.validate_python(kwargs.get("ttl"))) - - async def async_set_cache_pipeline(self, cache_list: Sequence[tuple[str, object]], **kwargs: object) -> None: - await self.native_handle.async_set_many(cache_list, ttl=_DURATION.validate_python(kwargs.get("ttl"))) - - async def batch_cache_write(self, key: str, value: object, **kwargs: object) -> None: - await self.async_set_cache(key, value, **kwargs) - - def flush_cache(self) -> None: - self.native_handle.flush() - - async def async_flush_cache(self) -> None: - await self.native_handle.async_flush() - - async def ping(self) -> bool: - return await self.native_handle.ping() - - async def disconnect(self) -> None: - await self.native_handle.disconnect() - - async def delete_cache_keys(self, keys: Sequence[str]) -> None: - await self.native_handle.delete(keys) - - async def test_connection(self) -> dict[str, str]: - return {"status": "success" if await self.ping() else "failed"} - - -class Cache: - @staticmethod - def memory(*, ttl: float = 600, capacity: int = 200, max_entry_bytes: int = 4194304) -> CacheFacade: - from litellm.rust_bridge._native import NativeCacheHandle - - handle: Final = NativeCacheHandle.memory(ttl=ttl, capacity=capacity, max_entry_bytes=max_entry_bytes) - return CacheFacade(type=LiteLLMCacheType.LOCAL, ttl=ttl, _backend=NativeBackend(handle)) - - @staticmethod - def redis(url: str, *, namespace: str, ttl: float = 600, max_entry_bytes: int = 4194304) -> CacheFacade: - from litellm.rust_bridge._native import NativeCacheHandle - - handle: Final = NativeCacheHandle.redis(url, namespace=namespace, ttl=ttl, max_entry_bytes=max_entry_bytes) - return CacheFacade(type=LiteLLMCacheType.REDIS, namespace=namespace, ttl=ttl, _backend=NativeBackend(handle)) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 312bdbc55ab..7e990e3c71b 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -254,7 +254,9 @@ async def anthropic_response( model_id=model_id, version=version, response_cost=( - failure_row_response_cost(base_llm_response_processor.data) if is_agentic_followup_failure(e) else 0 + failure_row_response_cost(base_llm_response_processor.data) + if is_agentic_followup_failure(answered) + else 0 ), model_region=getattr(user_api_key_dict, "allowed_model_region", ""), request_data=data, diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 8177e4a5cc9..954815f6a4f 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -296,32 +296,3 @@ class _SecretManagerRuntime: timeout: float | httpx.Timeout | None = None, primary_secret_name: str | None = None, ) -> Future[JsonValue]: ... - -@final -class NativeCacheHandle: - def __new__(cls, _uninstantiable: Never, /) -> Never: ... - @staticmethod - def memory( - *, - ttl: float = 600.0, - capacity: int = 200, - max_entry_bytes: int = 4194304, - ) -> NativeCacheHandle: ... - @staticmethod - def redis( - url: str, - *, - namespace: str, - ttl: float = 600.0, - max_entry_bytes: int = 4194304, - ) -> NativeCacheHandle: ... - def get(self, key: str) -> object: ... - def set(self, key: str, value: object, *, ttl: float | None = None) -> None: ... - def async_get(self, key: str) -> Future[object]: ... - def async_set(self, key: str, value: object, *, ttl: float | None = None) -> Future[None]: ... - def async_set_many(self, entries: Sequence[tuple[str, object]], *, ttl: float | None = None) -> Future[None]: ... - def flush(self) -> None: ... - def async_flush(self) -> Future[None]: ... - def ping(self) -> Future[bool]: ... - def disconnect(self) -> Future[None]: ... - def delete(self, keys: Sequence[str]) -> Future[None]: ... diff --git a/tests/test_litellm_rust/cache/test_python_cache.py b/tests/test_litellm_rust/cache/test_python_cache.py index beb818c7162..c9ac7af1a78 100644 --- a/tests/test_litellm_rust/cache/test_python_cache.py +++ b/tests/test_litellm_rust/cache/test_python_cache.py @@ -5,13 +5,12 @@ import pytest from pydantic import TypeAdapter import litellm -from litellm import _v2 from litellm.caching.caching import Cache, CacheMode from litellm.caching.caching_handler import ( _PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request ) from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict -from litellm.proxy._types import Litellm_EntityType, UserAPIKeyAuth +from litellm.proxy._types import Litellm_EntityType from litellm.proxy.hooks.model_max_budget_limiter import ( PROXY_VirtualKeyModelMaxBudgetLimiter, model_budget_spend_cache_key, @@ -19,10 +18,8 @@ from litellm.proxy.hooks.model_max_budget_limiter import ( from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.utils import InternalUsageCache from litellm.rust_bridge.catalog import Route -from litellm.types.caching import CachingSupportedCallTypes from tests.test_litellm_rust.cache.support import cache_key, collect, invoke, payload from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_EVENTS, MESSAGES_MODEL, MESSAGES_RESPONSE -from tests.test_litellm_rust.responses.data import RESPONSES_RESPONSE from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging, request_body from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.routing import native_route @@ -32,27 +29,18 @@ pytestmark = pytest.mark.requires_rust_extension @pytest.mark.asyncio @pytest.mark.parametrize( - ("route", "stream", "legacy"), - ( - ("chat", False, False), - ("messages", False, False), - ("responses", False, False), - ("messages", True, False), - ("messages", False, True), - ("messages", True, True), - ), + "stream", + (False, True), ) @pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) -async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage( +async def test_legacy_cache_hit_keeps_model_budget_spend_but_accounts_for_usage( recording_server: RecordingServer, - route: Literal["chat", "messages", "responses"], stream: bool, native: bool, monkeypatch: pytest.MonkeyPatch, - legacy: bool, ) -> None: monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") - litellm.cache = Cache() if legacy else _v2.Cache.memory() + litellm.cache = Cache() counters: Final = litellm.DualCache() budget: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(counters) limiter: Final = PROXY_MaxParallelRequestsHandler_v3( @@ -72,7 +60,7 @@ async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage( } spend_key: Final = model_budget_spend_cache_key(Litellm_EntityType.KEY, key_hash, "cached-model", "1h") token_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "tokens") - first: Final = await invoke(route, recording_server, options, native=native) + first: Final = await invoke("messages", recording_server, options, native=native) if stream: await collect(first) await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) @@ -81,13 +69,13 @@ async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage( first_log: Final = TypeAdapter(dict[str, object]).validate_python(first_events[0].kwargs) first_payload: Final = TypeAdapter(dict[str, object]).validate_python(first_log["standard_logging_object"]) expected_cost: Final = TypeAdapter(float).validate_python(first_log["response_cost"]) - usage: Final = RESPONSES_RESPONSE["usage"] if route == "responses" else MESSAGES_RESPONSE["usage"] + usage: Final = MESSAGES_RESPONSE["usage"] expected_tokens: Final = usage["input_tokens"] + usage["output_tokens"] assert expected_cost > 0 assert counters.get_cache(spend_key) == pytest.approx(expected_cost) assert counters.get_cache(token_key) == first_payload["total_tokens"] == expected_tokens - second: Final = await invoke(route, recording_server, options, native=native) + second: Final = await invoke("messages", recording_server, options, native=native) if stream: await collect(second) await drain_logging() @@ -101,7 +89,7 @@ async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage( assert cached_payload["cache_hit"] is True assert cached_payload["id"] != first_payload["id"] assert cached_payload["custom_llm_provider"] == first_payload["custom_llm_provider"] - assert cached_payload["custom_llm_provider"] == ("openai" if route == "responses" else "anthropic"), { + assert cached_payload["custom_llm_provider"] == "anthropic", { "miss_provider": first_log.get("custom_llm_provider"), "hit_provider": cached_log.get("custom_llm_provider"), } @@ -111,110 +99,22 @@ async def test_cache_hit_keeps_model_budget_spend_but_accounts_for_usage( @pytest.mark.asyncio -@pytest.mark.parametrize("native", (False, True), ids=("python", "rust")) -@pytest.mark.parametrize("backend", ("disabled", "memory", "redis")) -async def test_response_cache_backend_does_not_control_coordination( - recording_server: RecordingServer, - native: bool, - monkeypatch: pytest.MonkeyPatch, - backend: Literal["disabled", "memory", "redis"], - redis_url: str, -) -> None: - monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") - litellm.cache = ( - None - if backend == "disabled" - else _v2.Cache.memory() - if backend == "memory" - else _v2.Cache.redis(redis_url, namespace="independent-coordination") - ) - recording_server.expected_requests = 2 if backend == "disabled" else 1 - counters: Final = litellm.DualCache() - budget: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(counters) - limiter: Final = PROXY_MaxParallelRequestsHandler_v3( - InternalUsageCache(counters), model_group_resolver=lambda model: model - ) - key_hash: Final = "b" * 64 - identity: Final = UserAPIKeyAuth(api_key=key_hash, rpm_limit=2, tpm_limit=1000, max_parallel_requests=1) - spend_key: Final = model_budget_spend_cache_key(Litellm_EntityType.KEY, key_hash, "cached-model", "1h") - request_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "requests") - token_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "tokens") - parallel_key: Final = limiter.create_rate_limit_keys("api_key", key_hash, "max_parallel_requests") - expected_tokens: Final = MESSAGES_RESPONSE["usage"]["input_tokens"] + MESSAGES_RESPONSE["usage"]["output_tokens"] - recorder: Final = RecordingLogger() - - async def request(call_id: str, successes: int) -> object: - data: Final = { - "model": MESSAGES_MODEL, - "messages": list(MESSAGES), - "litellm_call_id": call_id, - "max_tokens": 32, - "metadata": { - "user_api_key": key_hash, - "model_group": "cached-model", - "user_api_key_model_max_budget": {"cached-model": {"max_budget": 1, "budget_duration": "1h"}}, - }, - } - await limiter.async_pre_call_hook(identity, counters, data, "acompletion") - assert len(TypeAdapter(dict[str, float]).validate_python(counters.get_cache(parallel_key))) == 1 - response: Final = await invoke( - "chat", recording_server, {**data, "callbacks": [budget, limiter, recorder]}, native=native - ) - await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) - await recorder.wait_for_async("async_log_success_event", count=successes) - return response - - await asyncio.create_task(request("cache-miss", 1)) - assert counters.get_cache(parallel_key) == {} - assert counters.get_cache(request_key) == 1 - assert counters.get_cache(token_key) == expected_tokens - first_events: Final = await recorder.wait_for_async("async_log_success_event") - first_cost: Final = TypeAdapter(float).validate_python(first_events[0].kwargs["response_cost"]) - assert first_cost > 0 - assert counters.get_cache(spend_key) == pytest.approx(first_cost) - await asyncio.create_task(request("cache-hit", 2)) - expected_spend: Final = first_cost * recording_server.expected_requests - assert counters.get_cache(spend_key) == pytest.approx(expected_spend) - assert len(recording_server.requests) == recording_server.expected_requests - assert counters.get_cache(parallel_key) == {} - assert counters.get_cache(request_key) == 2 - assert counters.get_cache(token_key) == 2 * expected_tokens - with pytest.raises(litellm.RateLimitError): - await asyncio.create_task(request("over-rpm-limit", 3)) - assert len(recording_server.requests) == recording_server.expected_requests - assert counters.get_cache(parallel_key) == {} - assert counters.get_cache(token_key) == 2 * expected_tokens - - assert counters.get_cache(spend_key) == pytest.approx(expected_spend) - if litellm.cache is not None: - await litellm.cache.disconnect() - - -@pytest.mark.asyncio -@pytest.mark.parametrize( - ("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True)) -) -async def test_cache_controls_and_backend_credential_key_semantics( - recording_server: RecordingServer, - route: Literal["chat", "messages", "responses"], - legacy: bool, -) -> None: - recording_server.expected_requests = 3 if legacy else 4 - litellm.cache = Cache() if legacy else _v2.Cache.memory() - await invoke(route, recording_server, {"cache": {"no-store": True}}) - await invoke(route, recording_server, {}) - await invoke(route, recording_server, {}) +async def test_legacy_cache_controls_and_backend_credential_key_semantics(recording_server: RecordingServer) -> None: + recording_server.expected_requests = 3 + litellm.cache = Cache() + await invoke("messages", recording_server, {"cache": {"no-store": True}}) + await invoke("messages", recording_server, {}) + await invoke("messages", recording_server, {}) assert len(recording_server.requests) == 2 - await invoke(route, recording_server, {"cache": {"no-cache": True}}) - await invoke(route, recording_server, {"api_key": "another-key"}) + await invoke("messages", recording_server, {"cache": {"no-cache": True}}) + await invoke("messages", recording_server, {"api_key": "another-key"}) assert len(recording_server.requests) == recording_server.expected_requests @pytest.mark.asyncio -@pytest.mark.parametrize("legacy", (False, True)) -async def test_v2_messages_replays_a_completed_stream(recording_server: RecordingServer, legacy: bool) -> None: +async def test_legacy_messages_replays_a_completed_stream(recording_server: RecordingServer) -> None: recording_server.default_response = ResponseSpec(body=None, events=MESSAGES_EVENTS) - litellm.cache = Cache() if legacy else _v2.Cache.memory() + litellm.cache = Cache() recorder: Final = RecordingLogger() parameters: Final = { "model": MESSAGES_MODEL, @@ -280,74 +180,30 @@ async def test_rust_messages_uses_a_legacy_cache_without_python_inference( @pytest.mark.asyncio -@pytest.mark.parametrize("route", ("chat", "messages", "responses")) @pytest.mark.parametrize("native", (False, True)) -@pytest.mark.parametrize("excluded", (None, [], ["embedding"])) -async def test_v2_cache_honors_supported_call_types_for_reads_and_writes( +async def test_legacy_cache_default_off_requires_opt_in_even_for_existing_entries( recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch, - route: Literal["chat", "messages", "responses"], native: bool, - excluded: list[CachingSupportedCallTypes] | None, ) -> None: monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") - litellm.cache = _v2.Cache.memory() - call_type: Final[CachingSupportedCallTypes] = ( - "acompletion" if route == "chat" else "anthropic_messages" if route == "messages" else "aresponses" - ) - recording_server.expected_requests = 4 - litellm.cache.supported_call_types = excluded - await invoke(route, recording_server, {}, native=native) - await invoke(route, recording_server, {}, native=native) - assert len(recording_server.requests) == 2 - litellm.cache.supported_call_types = [call_type] - await invoke(route, recording_server, {}, native=native) - assert len(recording_server.requests) == 3 - await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) - await invoke(route, recording_server, {}, native=native) - assert len(recording_server.requests) == 3 - litellm.cache.supported_call_types = excluded - await invoke(route, recording_server, {}, native=native) - assert len(recording_server.requests) == 4 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("native", (False, True)) -@pytest.mark.parametrize( - ("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True)) -) -async def test_v2_default_off_requires_opt_in_even_for_existing_entries( - recording_server: RecordingServer, - monkeypatch: pytest.MonkeyPatch, - route: Literal["chat", "messages", "responses"], - native: bool, - legacy: bool, -) -> None: - monkeypatch.setenv("LITELLM_RUST", "1" if native else "0") - litellm.cache = Cache() if legacy else _v2.Cache.memory() + litellm.cache = Cache() litellm.cache.mode = CacheMode.default_off recording_server.expected_requests = 4 - await invoke(route, recording_server, {}, native=native) - await invoke(route, recording_server, {}, native=native) + await invoke("messages", recording_server, {}, native=native) + await invoke("messages", recording_server, {}, native=native) assert len(recording_server.requests) == 2 - await invoke(route, recording_server, {"cache": {"use-cache": True}}, native=native) + await invoke("messages", recording_server, {"cache": {"use-cache": True}}, native=native) assert len(recording_server.requests) == 3 await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) - await invoke(route, recording_server, {"cache": {"use-cache": True}}, native=native) + await invoke("messages", recording_server, {"cache": {"use-cache": True}}, native=native) assert len(recording_server.requests) == 3 - await invoke(route, recording_server, {}, native=native) + await invoke("messages", recording_server, {}, native=native) assert len(recording_server.requests) == 4 @pytest.mark.asyncio -@pytest.mark.parametrize( - ("route", "legacy"), (("chat", False), ("messages", False), ("responses", False), ("messages", True)) -) -async def test_cache_lookup_uses_backend_request_callback_semantics( - recording_server: RecordingServer, - route: Literal["chat", "messages", "responses"], - legacy: bool, -) -> None: +async def test_legacy_cache_lookup_uses_backend_request_callback_semantics(recording_server: RecordingServer) -> None: class Rewrite(RecordingLogger): temperature = 0.1 @@ -356,22 +212,40 @@ async def test_cache_lookup_uses_backend_request_callback_semantics( super().log_pre_api_call(model, messages, kwargs) logger: Final = Rewrite() - litellm.cache = Cache() if legacy else _v2.Cache.memory() - recording_server.expected_requests = 1 if legacy else 2 - await invoke(route, recording_server, {"callbacks": [logger]}) - first_hit: Final = await invoke(route, recording_server, {"callbacks": [logger]}) + litellm.cache = Cache() + recording_server.expected_requests = 1 + await invoke("messages", recording_server, {"callbacks": [logger]}) + first_hit: Final = await invoke("messages", recording_server, {"callbacks": [logger]}) logger.temperature = 0.8 - await invoke(route, recording_server, {"callbacks": [logger]}) - second_hit: Final = await invoke(route, recording_server, {"callbacks": [logger]}) + await invoke("messages", recording_server, {"callbacks": [logger]}) + second_hit: Final = await invoke("messages", recording_server, {"callbacks": [logger]}) assert logger.names.count("log_pre_api_call") == 4 assert len(recording_server.requests) == recording_server.expected_requests assert recording_server.requests[0].body["temperature"] == 0.1 - if not legacy: - assert recording_server.requests[1].body["temperature"] == 0.8 assert isinstance(cache_key(first_hit), str) assert isinstance(cache_key(second_hit), str) - if not legacy: - assert cache_key(first_hit) != cache_key(second_hit) + + +@pytest.mark.asyncio +async def test_rust_chat_completions_rejects_a_configured_python_cache( + recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + recording_server.expected_requests = 0 + monkeypatch.setattr(litellm, "cache", Cache()) + with pytest.raises(RuntimeError, match="this inference route does not support configured caches"): + await invoke("chat", recording_server, {}) + assert recording_server.requests == [] + + +@pytest.mark.asyncio +async def test_rust_responses_rejects_a_configured_python_cache( + recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch +) -> None: + recording_server.expected_requests = 0 + monkeypatch.setattr(litellm, "cache", Cache()) + with pytest.raises(RuntimeError, match="this inference route does not support configured caches"): + await invoke("responses", recording_server, {}) + assert recording_server.requests == [] @pytest.mark.asyncio diff --git a/tests/test_litellm_rust/cache/v2/__init__.py b/tests/test_litellm_rust/cache/v2/__init__.py deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/tests/test_litellm_rust/cache/v2/test_cache.py b/tests/test_litellm_rust/cache/v2/test_cache.py deleted file mode 100644 index 1461e6eac61..00000000000 --- a/tests/test_litellm_rust/cache/v2/test_cache.py +++ /dev/null @@ -1,207 +0,0 @@ -import asyncio -from collections.abc import AsyncIterator -from typing import Final, Literal - -import pytest -from pydantic import TypeAdapter - -import litellm -from litellm import _v2 -from litellm._v2.cache import NativeBackend -from litellm.caching.caching_handler import ( - _PENDING_CACHE_WRITES, # pyright: ignore[reportPrivateUsage] # await the existing background cache writer before the next request -) -from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict -from litellm.types.utils import ModelResponse -from tests.test_litellm_rust.messages.data import MESSAGES, MESSAGES_MODEL, MESSAGES_RESPONSE -from tests.test_litellm_rust.responses.data import RESPONSES_MODEL, RESPONSES_RESPONSE, RESPONSES_TEXT -from tests.test_litellm_rust.cache.support import cache_key, invoke, payload -from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging -from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec - -pytestmark = pytest.mark.requires_rust_extension - - -@pytest.mark.asyncio -@pytest.mark.parametrize("route", ("chat", "messages", "responses")) -@pytest.mark.parametrize("backend", ("memory", "redis")) -async def test_v2_cache_skips_provider_and_reports_one_success_per_call( - recording_server: RecordingServer, - route: Literal["chat", "messages", "responses"], - backend: Literal["memory", "redis"], - redis_url: str, -) -> None: - recording_server.expected_requests = 2 - litellm.cache = _v2.Cache.memory() if backend == "memory" else _v2.Cache.redis(redis_url, namespace="headers") - recorder: Final = RecordingLogger() - first: Final = await invoke(route, recording_server, {"callbacks": [recorder]}) - await recorder.wait_for_async("async_log_success_event") - second: Final = await invoke(route, recording_server, {"callbacks": [recorder]}) - assert payload(first) == payload(second) - assert cache_key(first) is None - key: Final = cache_key(second) - assert isinstance(key, str) - assert key == get_hidden_params_dict(second)["cache_key"] - assert len(recording_server.requests) == 1 - await drain_logging() - successes: Final = await recorder.wait_for_async("async_log_success_event", count=2) - assert len(successes) == 2 - cached_log: Final = TypeAdapter(dict[str, object]).validate_python(successes[-1].kwargs) - assert cached_log["cache_hit"] is True - assert cached_log["response_cost"] == 0 - await litellm.cache.delete_cache_keys([key]) - refreshed: Final = await invoke(route, recording_server, {"callbacks": [recorder]}) - assert cache_key(refreshed) is None - assert len(recording_server.requests) == 2 - assert len(await recorder.wait_for_async("async_log_success_event", count=3)) == 3 - await litellm.cache.disconnect() - - -@pytest.mark.asyncio -async def test_v2_global_cache_leaves_legacy_only_calls_usable() -> None: - litellm.cache = _v2.Cache.memory() - response: Final = await litellm.aembedding( - model="openai/cache-test-embedding", - input=["hello"], - api_key="test-key", - mock_response=[0.25, 0.75], - ) - assert response.model_dump(include={"data"}) == { - "data": [{"embedding": [0.25, 0.75], "index": 0, "object": "embedding"}] - } - - -@pytest.mark.parametrize("route", ("chat", "responses")) -def test_v2_cache_works_through_python_inference( - recording_server: RecordingServer, route: Literal["chat", "messages", "responses"], monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setenv("LITELLM_RUST", "0") - litellm.cache = _v2.Cache.memory() - common: Final = {"api_key": "test-key", "api_base": recording_server.base_url} - if route == "responses": - recording_server.default_response = ResponseSpec(body=RESPONSES_RESPONSE) - parameters: Final = {"model": RESPONSES_MODEL, "input": "hello", **common} - first: Final = litellm.responses(**parameters) - second: Final = litellm.responses(**parameters) - assert payload(first) == payload(second) - else: - recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) - arguments: Final = {"model": MESSAGES_MODEL, "messages": list(MESSAGES), "max_tokens": 32, **common} - initial: Final = litellm.completion(**arguments) - cached: Final = litellm.completion(**arguments) - assert isinstance(initial, ModelResponse) and isinstance(cached, ModelResponse) - assert ( - initial.choices[0].message.content - == cached.choices[0].message.content - == MESSAGES_RESPONSE["content"][0]["text"] - ) - assert len(recording_server.requests) == 1 - - -@pytest.mark.asyncio -async def test_v2_facade_and_backend_share_storage_and_management() -> None: - cache: Final = _v2.Cache.memory() - await cache.async_add_cache({"answer": 7}, cache_key="shared") - assert cache.get_cache(cache_key="shared") == {"answer": 7} - assert await cache.ping() is True - await cache.delete_cache_keys(["shared"]) - assert await cache.async_get_cache(cache_key="shared") is None - cache.add_cache({"answer": 8}, cache_key="flush") - backend: Final = cache.cache - assert isinstance(backend, NativeBackend) - backend.flush_cache() - assert cache.get_cache(cache_key="flush") is None - await cache.disconnect() - - -@pytest.mark.asyncio -@pytest.mark.parametrize("control", ("s-maxage", "s-max-age")) -async def test_v2_native_cache_accepts_existing_freshness_aliases( - recording_server: RecordingServer, control: str -) -> None: - litellm.cache = _v2.Cache.memory() - first: Final = await invoke("responses", recording_server, {}) - second: Final = await invoke("responses", recording_server, {"cache": {control: 600}}) - assert payload(first) == payload(second) - assert len(recording_server.requests) == 1 - - -@pytest.mark.asyncio -async def test_v2_cache_does_not_force_native_responses_streaming( - recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setenv("LITELLM_RUST", "0") - litellm.cache = _v2.Cache.memory() - recording_server.default_response = ResponseSpec( - body=None, - events=( - ("response.created", {"type": "response.created", "sequence_number": 0, "response": RESPONSES_RESPONSE}), - ( - "response.completed", - {"type": "response.completed", "sequence_number": 1, "response": RESPONSES_RESPONSE}, - ), - ), - ) - response: Final = await litellm.aresponses( - model=RESPONSES_MODEL, - input="hello", - stream=True, - caching=False, - api_key="test-key", - api_base=recording_server.base_url, - ) - assert isinstance(response, AsyncIterator) - chunks: Final = [chunk async for chunk in response] - assert chunks[-1].type == "response.completed" - assert chunks[-1].response.output[0].content[0].text == RESPONSES_TEXT - - -@pytest.mark.asyncio -async def test_v2_cache_works_through_python_messages( - recording_server: RecordingServer, monkeypatch: pytest.MonkeyPatch -) -> None: - monkeypatch.setenv("LITELLM_RUST", "0") - litellm.cache = _v2.Cache.memory() - recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) - parameters: Final = { - "model": MESSAGES_MODEL, - "messages": list(MESSAGES), - "max_tokens": 32, - "api_key": "test-key", - "api_base": recording_server.base_url, - } - first: Final = await litellm.anthropic_messages(**parameters) - await asyncio.gather(*tuple(_PENDING_CACHE_WRITES)) - second: Final = await litellm.anthropic_messages(**parameters) - assert payload(first) == payload(second) - assert len(recording_server.requests) == 1 - - -@pytest.mark.asyncio -@pytest.mark.parametrize("asynchronous", (False, True)) -async def test_v2_redis_flush_only_removes_its_namespace( - redis_url: str, recording_server: RecordingServer, asynchronous: bool -) -> None: - own: Final = _v2.Cache.redis(redis_url, namespace="flush-own") - other: Final = _v2.Cache.redis(redis_url, namespace="flush-other") - litellm.cache = own - recording_server.expected_requests = 2 - await own.async_add_cache({"answer": "own"}, cache_key="shared") - await other.async_add_cache({"answer": "other"}, cache_key="shared") - await invoke("responses", recording_server, {}) - hit: Final = await invoke("responses", recording_server, {}) - assert isinstance(cache_key(hit), str) - assert await own.async_get_cache(cache_key="shared") == {"answer": "own"} - backend: Final = own.cache - assert isinstance(backend, NativeBackend) - if asynchronous: - await backend.async_flush_cache() - else: - backend.flush_cache() - assert await own.async_get_cache(cache_key="shared") is None - assert await other.async_get_cache(cache_key="shared") == {"answer": "other"} - refreshed: Final = await invoke("responses", recording_server, {}) - assert cache_key(refreshed) is None - assert len(recording_server.requests) == 2 - await own.disconnect() - await other.disconnect() diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 3336cce9aaf..ceb9ef38887 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -25,7 +25,7 @@ _REPO: Final = Path(__file__).resolve().parents[2] _REDIS_PRODUCER_ROOTS: Final = ("litellm", "enterprise") # The cache implementations and facades: they emit the service events, their callers declare the family. -_CACHE_LAYER_DIRS: Final = ("litellm/caching", "litellm/_v2/cache") +_CACHE_LAYER_DIRS: Final = ("litellm/caching",) # Helpers that act on a cache handed in by the declaring caller, or forward to the response-cache facade. _CACHE_PARAMETER_HELPERS: Final = frozenset( {