diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 29e661dea9c..66270501312 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -40,6 +40,12 @@ dependencies = [ "cc", ] +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + [[package]] name = "android_system_properties" version = "0.1.6" @@ -230,6 +236,7 @@ dependencies = [ "aws-credential-types", "aws-sigv4", "aws-smithy-async", + "aws-smithy-eventstream", "aws-smithy-http", "aws-smithy-runtime", "aws-smithy-runtime-api", @@ -238,7 +245,9 @@ dependencies = [ "bytes", "bytes-utils", "fastrand", + "http 0.2.12", "http 1.4.2", + "http-body 0.4.6", "http-body 1.1.0", "percent-encoding", "pin-project-lite", @@ -272,6 +281,43 @@ dependencies = [ "tracing", ] +[[package]] +name = "aws-sdk-s3" +version = "1.146.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2cd651b4400d4011b8927b83a9552bf90ff11e6e5da0b9f0a7583247aceec971" +dependencies = [ + "arc-swap", + "aws-credential-types", + "aws-runtime", + "aws-sigv4", + "aws-smithy-async", + "aws-smithy-checksums", + "aws-smithy-eventstream", + "aws-smithy-http", + "aws-smithy-json", + "aws-smithy-observability", + "aws-smithy-runtime", + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "aws-smithy-xml 0.62.1", + "aws-types", + "bytes", + "fastrand", + "hex", + "hmac", + "http 0.2.12", + "http 1.4.2", + "http-body 1.1.0", + "lru", + "percent-encoding", + "regex-lite", + "sha2 0.11.0", + "tracing", + "url", +] + [[package]] name = "aws-sdk-secretsmanager" version = "1.117.0" @@ -316,7 +362,7 @@ dependencies = [ "aws-smithy-runtime-api", "aws-smithy-schema", "aws-smithy-types", - "aws-smithy-xml", + "aws-smithy-xml 0.61.1", "aws-types", "fastrand", "http 0.2.12", @@ -332,6 +378,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "31d955e76ff96acd555bf06fa0fa6d5bf9335fa84ae7c64481b20ae61d231f70" dependencies = [ "aws-credential-types", + "aws-smithy-eventstream", "aws-smithy-http", "aws-smithy-runtime-api", "aws-smithy-types", @@ -359,10 +406,31 @@ dependencies = [ ] [[package]] -name = "aws-smithy-eventstream" -version = "0.61.1" +name = "aws-smithy-checksums" +version = "0.65.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5a9381123ab62d20c13082b151f30f962a3b112b727345394536dfa39a482944" +checksum = "b67ecd999972b58e67cab052f5129906c08c25883bd0788ceefc55ef97d61307" +dependencies = [ + "aws-smithy-http", + "aws-smithy-types", + "bytes", + "crc-fast", + "hex", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "md-5", + "pin-project-lite", + "sha1 0.11.0", + "sha2 0.11.0", + "tracing", +] + +[[package]] +name = "aws-smithy-eventstream" +version = "0.61.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80c2051c2f1016fb8e6548dd07b8bc2ac9c3fe583721444b92f515e856d31609" dependencies = [ "aws-smithy-types", "bytes", @@ -375,6 +443,7 @@ version = "0.64.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37843d9add67c3aff5856f409c6dc315d3cdff60f9c0cb5b670dab1e9920306d" dependencies = [ + "aws-smithy-eventstream", "aws-smithy-runtime-api", "aws-smithy-types", "bytes", @@ -554,6 +623,18 @@ dependencies = [ "xmlparser", ] +[[package]] +name = "aws-smithy-xml" +version = "0.62.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b932c8d6dc127fc980eecd78f8694ae9b9551b69a93a7def2a199c1c0033daf" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", + "xmlparser", +] + [[package]] name = "aws-types" version = "1.6.0" @@ -980,6 +1061,16 @@ dependencies = [ "libc", ] +[[package]] +name = "crc-fast" +version = "1.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e75b2483e97a5a7da73ac68a05b629f9c53cff58d8ed1c77866079e18b00dba5" +dependencies = [ + "digest 0.10.7", + "spin", +] + [[package]] name = "crc16" version = "0.4.0" @@ -1348,7 +1439,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -1925,6 +2016,8 @@ version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" dependencies = [ + "allocator-api2", + "equivalent", "foldhash", ] @@ -2148,7 +2241,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.6.5", + "socket2 0.5.10", "tokio", "tower-service", "tracing", @@ -2679,6 +2772,37 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-s3" +version = "0.1.0" +dependencies = [ + "aws-credential-types", + "aws-sdk-s3", + "aws-smithy-types", + "aws-types", + "litellm-auth-aws", + "litellm-cache", + "serde_json", + "tokio", + "wiremock", +] + +[[package]] +name = "litellm-cache-valkey-semantic" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "litellm-cache-redis", + "litellm-cache-response", + "redis", + "redis-test", + "rstest", + "serde_json", + "sha2 0.10.9", + "tokio", + "uuid", +] + [[package]] name = "litellm-callbacks-legacy-python" version = "0.1.0" @@ -2841,6 +2965,7 @@ dependencies = [ "criterion", "futures-util", "litellm-auth", + "litellm-auth-aws", "litellm-auth-gcp", "litellm-cache", "litellm-cache-azure-blob", @@ -2850,6 +2975,8 @@ dependencies = [ "litellm-cache-redis", "litellm-cache-redis-semantic", "litellm-cache-response", + "litellm-cache-s3", + "litellm-cache-valkey-semantic", "litellm-callbacks-legacy-python", "litellm-core", "litellm-core-utils", @@ -2860,10 +2987,12 @@ dependencies = [ "litellm-types", "pyo3", "pyo3-async-runtimes", + "redis", "rstest", "serde", "serde_json", "serde_with", + "sha2 0.10.9", "tokio", "tokio-tungstenite", ] @@ -3091,6 +3220,15 @@ version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" +[[package]] +name = "lru" +version = "0.18.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff9840bcc50b71349309900da0ce7279aa336ae71d73250b07998932c7d97c25" +dependencies = [ + "hashbrown 0.17.1", +] + [[package]] name = "lru-slab" version = "0.1.2" @@ -3113,6 +3251,16 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c" +[[package]] +name = "md-5" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" +dependencies = [ + "cfg-if", + "digest 0.11.3", +] + [[package]] name = "memchr" version = "2.8.3" @@ -3755,7 +3903,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls 0.23.42", - "socket2 0.6.5", + "socket2 0.5.10", "thiserror 2.0.19", "tokio", "tracing", @@ -3794,9 +3942,9 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.6.5", + "socket2 0.5.10", "tracing", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4282,7 +4430,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4353,7 +4501,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -4626,6 +4774,17 @@ dependencies = [ "digest 0.10.7", ] +[[package]] +name = "sha1" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "digest 0.11.3", +] + [[package]] name = "sha1_smol" version = "1.0.1" @@ -4742,6 +4901,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "spin" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "023a211cb3138dbc438680b32560ad89f699977624c9f8dbb95a47d5b4c07dd3" + [[package]] name = "spm_precompiled" version = "0.1.4" @@ -4908,10 +5073,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.4.3", + "getrandom 0.3.4", "once_cell", "rustix", - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] @@ -5375,7 +5540,7 @@ dependencies = [ "rand 0.8.7", "rustls 0.23.42", "rustls-pki-types", - "sha1", + "sha1 0.10.7", "thiserror 1.0.69", "utf-8", ] @@ -5784,7 +5949,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.61.2", + "windows-sys 0.52.0", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 3e28c11347b..4e370c13766 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -33,6 +33,7 @@ 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" } diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs index 4181ce719e3..e0ac31f3630 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -8,29 +8,26 @@ use litellm_cache::{ BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error, SemanticCacheContext, }; -use litellm_cache_redis::connection::{ConnectionRef, Connections, ttl_seconds}; +use litellm_cache_redis::{ + RedisTopology, + connection::{ConnectionRef, Connections}, +}; use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; use serde_json::Value; use sha2::{Digest, Sha256}; use crate::prompt::prompt_from_context; -const REDIS_TIMEOUT: Duration = Duration::from_secs(5); -const REDIS_POOL_SIZE: u32 = 16; const CACHE_KEY_FIELD: &str = "litellm_cache_key"; const VECTOR_FIELD: &str = "prompt_vector"; pub trait Embedder: Send + Sync + 'static { - fn embed( - &self, - prompt: &str, - metadata: &serde_json::Map, - ) -> Result, Error>; + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error>; fn async_embed( &self, prompt: &str, - metadata: &serde_json::Map, + metadata: Option<&Value>, ) -> impl Future, Error>> + Send; } @@ -212,7 +209,7 @@ pub struct RedisSemanticCache { impl RedisSemanticCache { pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result { Ok(Self { - connections: Arc::new(Connections::pooled(url, REDIS_TIMEOUT, REDIS_POOL_SIZE)?), + connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), embedder, inner: Arc::new(Inner::new(config)), }) @@ -277,7 +274,7 @@ impl BaseCache let Some(prompt) = prompt_from_context(context) else { return Ok(()); }; - let vector = self.embedder.embed(&prompt, &context.metadata)?; + let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; let tag = Self::tag(key, context).to_string(); self.connections.execute(|connection| { self.inner @@ -289,7 +286,7 @@ impl BaseCache let Some(prompt) = prompt_from_context(context) else { return Ok(None); }; - let vector = self.embedder.embed(&prompt, &context.metadata)?; + let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; let tag = Self::tag(key, context).to_string(); self.connections .execute(|connection| self.inner.lookup(connection, &tag, &vector)) @@ -306,7 +303,7 @@ impl BaseCache }; let vector = self .embedder - .async_embed(&prompt, &context.metadata) + .async_embed(&prompt, context.metadata.as_ref()) .await?; let tag = Self::tag(key, &context).to_string(); let inner = Arc::clone(&self.inner); @@ -326,7 +323,7 @@ impl BaseCache }; let vector = self .embedder - .async_embed(&prompt, &context.metadata) + .async_embed(&prompt, context.metadata.as_ref()) .await?; let tag = Self::tag(key, context).to_string(); let inner = Arc::clone(&self.inner); @@ -613,3 +610,9 @@ fn bytes_field(fields: &[redis::Value], name: &str) -> Option> { _ => None, } } + +fn ttl_seconds(ttl: Duration) -> u64 { + ttl.as_secs() + .saturating_add(u64::from(ttl.subsec_nanos() > 0)) + .max(1) +} diff --git a/litellm-rust/crates/cache-redis-semantic/src/prompt.rs b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs index fc99898d847..b9c38e98d77 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/prompt.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs @@ -2,8 +2,10 @@ use litellm_cache::SemanticCacheContext; use serde_json::Value; pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { - if !context.messages.is_empty() { - return Some(messages_text(&context.messages)); + if let Some(messages) = context.messages.as_ref().and_then(Value::as_array) + && !messages.is_empty() + { + return Some(messages_text(messages)); } let input = context.input.as_ref()?; let mut parts = Vec::new(); diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs index fc37cf9f97f..233b87ec52f 100644 --- a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -35,7 +35,7 @@ impl FakeEmbedder { } impl Embedder for FakeEmbedder { - fn embed(&self, prompt: &str, _: &serde_json::Map) -> Result, Error> { + fn embed(&self, prompt: &str, _: Option<&Value>) -> Result, Error> { self.calls.lock().unwrap().push(prompt.to_string()); Ok(self @@ -45,11 +45,7 @@ impl Embedder for FakeEmbedder { .unwrap_or_else(|| vec![0.1, 0.2, 0.3])) } - async fn async_embed( - &self, - prompt: &str, - metadata: &serde_json::Map, - ) -> Result, Error> { + async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { self.embed(prompt, metadata) } } @@ -63,7 +59,7 @@ fn config() -> RedisSemanticConfig { fn messages_context(messages: Vec) -> SemanticCacheContext { SemanticCacheContext { - messages, + messages: Some(Value::Array(messages)), ..Default::default() } } diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 35ad659dc63..24399c9b2f9 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -40,6 +40,7 @@ const CLAIM_SCRIPT: &str = concat!( ); const CLAIM_ATTEMPTS: usize = 8; +#[allow(private_interfaces)] pub enum Connections { Pool(r2d2::Pool), Cluster(r2d2::Pool), @@ -50,21 +51,6 @@ impl Connections where C: redis::ConnectionLike + Send + 'static, { - pub fn pooled(url: &str, timeout: Duration, pool_size: u32) -> Result { - let pool = r2d2::Pool::builder() - .max_size(pool_size) - .min_idle(Some(0)) - .connection_timeout(timeout) - .test_on_check_out(false) - .build(ConnectionManager::open(url)?) - .map_err(|_| Error::Unavailable)?; - Ok(Self::Pool(pool)) - } - - pub fn fixed(connection: C) -> Self { - Self::Fixed(Mutex::new(connection)) - } - pub fn execute( &self, operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, @@ -98,6 +84,19 @@ where .await .map_err(|_| Error::Unavailable)? } + + pub fn fixed(connection: C) -> Self { + Self::Fixed(Mutex::new(connection)) + } + + pub fn open(url: &str, topology: &RedisTopology) -> Result { + match topology { + RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)), + RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool( + ClusterConnectionManager::open(url, startup_nodes)?, + )?)), + } + } } pub struct RedisCache { @@ -119,12 +118,7 @@ impl RedisCache { default_ttl: Option, codec: S, ) -> Result { - let connections = match topology { - RedisTopology::Standalone => Connections::Pool(pool(ConnectionManager::open(url)?)?), - RedisTopology::Cluster { startup_nodes } => { - Connections::Cluster(pool(ClusterConnectionManager::open(url, startup_nodes)?)?) - } - }; + let connections = Connections::open(url, topology)?; Ok(Self { connections: Arc::new(connections), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), @@ -224,22 +218,10 @@ where } fn ttl_seconds(ttl: Duration) -> u64 { - ttl_seconds(ttl) + ttl.as_secs() + .saturating_add(u64::from(ttl.subsec_nanos() > 0)) + .max(1) } - - async fn run_blocking(connections: Arc>, operation: F) -> Result - where - T: Send + 'static, - F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, - { - Connections::run_blocking(connections, operation).await - } -} - -pub fn ttl_seconds(ttl: Duration) -> u64 { - ttl.as_secs() - .saturating_add(u64::from(ttl.subsec_nanos() > 0)) - .max(1) } fn namespaced_key(namespace: Option<&str>, key: &str) -> String { @@ -338,7 +320,7 @@ where if entries.is_empty() { return Ok(()); } - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let commands = entries .into_iter() .map(|(key, payload)| { @@ -357,7 +339,7 @@ where } async fn test_connection(&self) -> Result { - match Self::run_blocking(Arc::clone(&self.connections), |connection| { + match Connections::run_blocking(Arc::clone(&self.connections), |connection| { Ok(match connection.ping() { Ok(_) => CacheConnectionResult { status: CacheConnectionStatus::Success, diff --git a/litellm-rust/crates/cache-redis/src/cache/operations.rs b/litellm-rust/crates/cache-redis/src/cache/operations.rs index 4da190efda8..4345ee879b3 100644 --- a/litellm-rust/crates/cache-redis/src/cache/operations.rs +++ b/litellm-rust/crates/cache-redis/src/cache/operations.rs @@ -188,7 +188,7 @@ where } pub async fn ping(&self) -> Result { - Self::run_blocking(Arc::clone(&self.connections), |connection| { + Connections::run_blocking(Arc::clone(&self.connections), |connection| { connection.ping().map_err(|_| Error::Unavailable) }) .await @@ -208,7 +208,7 @@ where pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, Error> { let pattern = format!("{}*", self.namespaced_key(pattern)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut matches = Vec::new(); connection.scan(&pattern, count, |_, keys| { matches.extend(keys); @@ -231,7 +231,7 @@ where } let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut sadd = redis::cmd("SADD"); sadd.arg(&key).arg(values); let mut expire = redis::cmd("EXPIRE"); @@ -279,7 +279,7 @@ where if operations.is_empty() { return Ok(Vec::new()); } - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let commands = operations .into_iter() .map(|(key, values)| { @@ -333,7 +333,7 @@ where .iter() .map(|(_, count)| count.is_some()) .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let commands = operations .into_iter() .map(|(key, count)| { @@ -426,7 +426,7 @@ where if operations.is_empty() { return Ok(Vec::new()); } - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut commands = Vec::with_capacity(operations.len() * 2); let mut increments = Vec::with_capacity(operations.len()); for (key, amount, ttl) in operations { diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs index ea75906e9c9..efb0db931ac 100644 --- a/litellm-rust/crates/cache-redis/src/lib.rs +++ b/litellm-rust/crates/cache-redis/src/lib.rs @@ -2,7 +2,7 @@ mod cache; mod topology; pub mod connection { - pub use crate::cache::{ConnectionRef, Connections, ttl_seconds}; + pub use crate::cache::{ConnectionRef, Connections}; } pub use cache::{ diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 56c1646d343..46e561ddad1 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -58,4 +58,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths -Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees +Redis cluster, disk, and cloud stores remain follow-ups. Semantic backends plug in through `SemanticCacheContext`, which carries the prompt inputs and metadata alongside the cache TTL. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index a27cc1967d5..2f949d511de 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,22 +1,21 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, - ExactCacheContext, FlushCache, + BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, FlushCache, }; use serde_json::Value; use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; #[derive(Clone)] -pub struct ResponseCacheRequest { +pub struct ResponseCacheRequest { pub key: CacheKeyInput, pub controls: CacheControls, pub context: C, pub max_age: Option, } -impl ResponseCacheRequest { +impl ResponseCacheRequest { pub fn new(key: CacheKeyInput) -> Self { Self { key, @@ -27,17 +26,24 @@ impl ResponseCacheRequest { default_on: true, ..Default::default() }, - context: ExactCacheContext::default(), + context: C::default(), max_age: None, } } } -pub struct ResponseCache> { +pub struct ResponseCache> +where + B::Context: Default + PartialEq, +{ backend: Arc, } -impl> ResponseCache { +impl ResponseCache +where + B: BaseCache, + B::Context: Default + PartialEq, +{ pub fn new(backend: Arc) -> Self { Self { backend } } @@ -46,10 +52,11 @@ impl> ResponseCache { &self.backend } - pub fn default_ttl(&self) -> Option - where - B::Context: Default, - { + pub fn backend_arc(&self) -> &Arc { + &self.backend + } + + pub fn default_ttl(&self) -> Option { self.backend.get_ttl(&B::Context::default()) } @@ -199,10 +206,7 @@ impl> ResponseCache { &self, entries: Vec<(ResponseCacheRequest, Value)>, now: Duration, - ) -> Result<(), Error> - where - B::Context: PartialEq, - { + ) -> Result<(), Error> { self.async_store_entries( entries .into_iter() @@ -217,10 +221,7 @@ impl> ResponseCache { pub async fn async_store_entries( &self, entries: Vec<(ResponseCacheRequest, Value, Duration)>, - ) -> Result<(), Error> - where - B::Context: PartialEq, - { + ) -> Result<(), Error> { let writable = entries .into_iter() .filter(|(request, _, _)| request.controls.writes()) diff --git a/litellm-rust/crates/cache-s3/Cargo.toml b/litellm-rust/crates/cache-s3/Cargo.toml new file mode 100644 index 00000000000..cdc17e732cb --- /dev/null +++ b/litellm-rust/crates/cache-s3/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "litellm-cache-s3" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-auth-aws.workspace = true +aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] } +aws-credential-types = "1.3.0" +aws-smithy-types = "1.6.0" +aws-types = "1.6.0" +tokio.workspace = true + +[dev-dependencies] +wiremock = "0.6.5" +serde_json.workspace = true +tokio = { workspace = true, features = ["macros", "rt-multi-thread"] } diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs new file mode 100644 index 00000000000..b7ca722cea3 --- /dev/null +++ b/litellm-rust/crates/cache-s3/src/auth.rs @@ -0,0 +1,101 @@ +use aws_credential_types::{ + Credentials as AwsCredentials, + provider::{ProvideCredentials, error::CredentialsError, future}, +}; +use litellm_auth_aws::{AwsAuthConfig, resolve_credentials}; + +#[derive(Clone)] +pub(crate) struct Credentials { + config: AwsAuthConfig, + env: fn(&str) -> Option, +} + +impl Credentials { + pub(crate) fn new(config: AwsAuthConfig) -> Self { + Self::with_env(config, |name| std::env::var(name).ok()) + } + + pub(crate) fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option) -> Self { + Self { config, env } + } +} + +impl ProvideCredentials for Credentials { + fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a> + where + Self: 'a, + { + future::ProvideCredentials::new(async { + if let (Some(access_key_id), Some(secret_access_key)) = ( + self.config.access_key_id.clone(), + self.config.secret_access_key.clone(), + ) { + return Ok(AwsCredentials::new( + access_key_id, + secret_access_key, + self.config.session_token.clone(), + None, + "litellm-s3-cache", + )); + } + resolve_credentials(self.config.clone(), &self.env) + .await + .map_err(|_| CredentialsError::provider_error("S3 cache authentication failed")) + }) + } +} + +impl std::fmt::Debug for Credentials { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Credentials").finish_non_exhaustive() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn explicit_keys_ignore_an_ambient_session_token() { + let provider = Credentials::with_env( + AwsAuthConfig { + access_key_id: Some("key".to_string()), + secret_access_key: Some("secret".to_string()), + region_name: Some("us-east-1".to_string()), + ..Default::default() + }, + |name| (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()), + ); + let credentials = provider.provide_credentials().await.unwrap(); + assert_eq!(credentials.access_key_id(), "key"); + assert_eq!(credentials.secret_access_key(), "secret"); + assert_eq!(credentials.session_token(), None); + } + + #[tokio::test] + async fn explicit_keys_keep_their_session_token() { + let provider = Credentials::new(AwsAuthConfig { + access_key_id: Some("key".to_string()), + secret_access_key: Some("secret".to_string()), + session_token: Some("t".to_string()), + region_name: Some("us-east-1".to_string()), + ..Default::default() + }); + let credentials = provider.provide_credentials().await.unwrap(); + assert_eq!(credentials.session_token(), Some("t")); + } + + #[tokio::test] + async fn environment_keys_resolve_with_their_session_token() { + let provider = Credentials::with_env(AwsAuthConfig::default(), |name| match name { + "AWS_ACCESS_KEY_ID" => Some("env-key".to_string()), + "AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()), + "AWS_SESSION_TOKEN" => Some("env-token".to_string()), + _ => None, + }); + let credentials = provider.provide_credentials().await.unwrap(); + assert_eq!(credentials.access_key_id(), "env-key"); + assert_eq!(credentials.secret_access_key(), "env-secret"); + assert_eq!(credentials.session_token(), Some("env-token")); + } +} diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs new file mode 100644 index 00000000000..9c791f42c3b --- /dev/null +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -0,0 +1,220 @@ +use std::{ + future::Future, + sync::Arc, + time::{Duration, SystemTime}, +}; + +use aws_sdk_s3::{ + config::{BehaviorVersion, Region, RequestChecksumCalculation, ResponseChecksumValidation}, + error::SdkError, + primitives::ByteStream, +}; +use aws_smithy_types::{DateTime, date_time::Format}; +use litellm_auth_aws::AwsAuthConfig; +use litellm_cache::{ + BaseCache, BatchCache, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache, +}; +use tokio::runtime::Handle; + +use crate::auth::Credentials; + +pub struct S3Endpoint { + pub url: String, +} + +pub struct S3CacheConfig { + pub bucket: String, + pub key_prefix: String, + pub region: String, + pub endpoint: Option, + pub auth: AwsAuthConfig, +} + +pub struct S3Cache { + client: aws_sdk_s3::Client, + codec: C, + runtime: Handle, + bucket: Arc, + key_prefix: Arc, + region: Arc, + endpoint: Option>, +} + +impl S3Cache { + pub fn new(config: S3CacheConfig, codec: C, runtime: Handle) -> Self { + let endpoint_url: Option = config.endpoint.map(|endpoint| endpoint.url); + let base = aws_sdk_s3::Config::builder() + .behavior_version(BehaviorVersion::latest()) + .region(Region::new(config.region.clone())) + .credentials_provider(Credentials::new(config.auth)) + .request_checksum_calculation(RequestChecksumCalculation::WhenRequired) + .response_checksum_validation(ResponseChecksumValidation::WhenRequired); + let builder = match &endpoint_url { + Some(url) => base.endpoint_url(url).force_path_style(true), + None => base, + }; + Self { + client: aws_sdk_s3::Client::from_conf(builder.build()), + codec, + runtime, + bucket: config.bucket.into(), + key_prefix: config.key_prefix.into(), + region: config.region.into(), + endpoint: endpoint_url.map(Into::into), + } + } + + pub fn bucket(&self) -> &str { + &self.bucket + } + + pub fn key_prefix(&self) -> &str { + &self.key_prefix + } + + pub fn region(&self) -> &str { + &self.region + } + + pub fn endpoint(&self) -> Option<&str> { + self.endpoint.as_deref() + } + + pub fn to_s3_key(&self, key: &str) -> String { + format!("{}{}", self.key_prefix, key.replace(':', "/")) + } + + fn block_on(&self, future: F) -> F::Output { + if Handle::try_current().is_ok() { + tokio::task::block_in_place(|| self.runtime.block_on(future)) + } else { + self.runtime.block_on(future) + } + } + + async fn put( + &self, + key: &str, + value: C::Value, + context: &ExactCacheContext, + ) -> Result<(), Error> { + let s3_key = self.to_s3_key(key); + let body = self.codec.encode(&value)?; + let request = self + .client + .put_object() + .bucket(self.bucket.as_ref()) + .key(&s3_key) + .body(ByteStream::from(body)) + .content_type("application/json") + .content_language("en") + .content_disposition(format!("inline; filename=\"{s3_key}.json\"")); + let request = match context.ttl { + Some(ttl) => { + let seconds = ttl.as_secs_f64(); + request + .cache_control(format!("immutable, max-age={seconds}, s-maxage={seconds}")) + .expires(DateTime::from(SystemTime::now() + ttl)) + } + None => request.cache_control("immutable, max-age=31536000, s-maxage=31536000"), + }; + request.send().await.map_err(|_| Error::Unavailable)?; + Ok(()) + } + + async fn get(&self, key: &str) -> Result, Error> { + let output = match self + .client + .get_object() + .bucket(self.bucket.as_ref()) + .key(self.to_s3_key(key)) + .send() + .await + { + Ok(output) => output, + Err(error) => { + if let SdkError::ServiceError(service) = &error { + let status = error + .raw_response() + .map(|response| response.status().as_u16()); + let not_found = service.err().is_no_such_key() + || service.err().meta().code() == Some("AccessDenied") + || status == Some(404) + || status == Some(403); + if not_found { + return Ok(None); + } + } + return Err(Error::Unavailable); + } + }; + if let Some(expires) = output.expires_string() + && let Ok(expires) = DateTime::from_str(expires, Format::HttpDate) + && expires < DateTime::from(SystemTime::now()) + { + return Ok(None); + } + let bytes = output + .body + .collect() + .await + .map_err(|_| Error::Unavailable)? + .into_bytes(); + self.codec.decode(&bytes).map(Some) + } +} + +impl BaseCache for S3Cache { + type Value = C::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + self.block_on(self.put(key, value, context)) + } + + fn get_cache(&self, key: &str, _context: &Self::Context) -> Result, Error> { + self.block_on(self.get(key)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> Result<(), Error> { + self.put(key, value, &context).await + } + + async fn async_get_cache( + &self, + key: &str, + _context: &Self::Context, + ) -> Result, Error> { + self.get(key).await + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Err(Error::UnsupportedOperation) + } +} + +impl BatchCache for S3Cache {} + +impl FlushCache for S3Cache { + fn flush_cache(&self) -> Result<(), Error> { + Ok(()) + } +} diff --git a/litellm-rust/crates/cache-s3/src/lib.rs b/litellm-rust/crates/cache-s3/src/lib.rs new file mode 100644 index 00000000000..f6126dfa908 --- /dev/null +++ b/litellm-rust/crates/cache-s3/src/lib.rs @@ -0,0 +1,4 @@ +mod auth; +mod cache; + +pub use cache::{S3Cache, S3CacheConfig, S3Endpoint}; diff --git a/litellm-rust/crates/cache-s3/tests/cache.rs b/litellm-rust/crates/cache-s3/tests/cache.rs new file mode 100644 index 00000000000..9a71656286b --- /dev/null +++ b/litellm-rust/crates/cache-s3/tests/cache.rs @@ -0,0 +1,278 @@ +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +use litellm_auth_aws::AwsAuthConfig; +use litellm_cache::{ + BaseCache, BatchCache, BatchEntry, Error, ExactCacheContext, FlushCache, JsonCodec, +}; +use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint}; +use serde_json::{Value, json}; +use tokio::runtime::Handle; +use wiremock::{ + Mock, MockServer, ResponseTemplate, + matchers::{method, path}, +}; + +fn config(endpoint: String) -> S3CacheConfig { + S3CacheConfig { + bucket: "cache-bucket".to_string(), + key_prefix: "team/".to_string(), + region: "us-east-1".to_string(), + endpoint: Some(S3Endpoint { url: endpoint }), + auth: AwsAuthConfig { + access_key_id: Some("key".to_string()), + secret_access_key: Some("secret".to_string()), + region_name: Some("us-east-1".to_string()), + ..Default::default() + }, + } +} + +fn cache(endpoint: &str) -> S3Cache> { + S3Cache::new( + config(endpoint.to_string()), + JsonCodec::::new(), + Handle::current(), + ) +} + +async fn mock_server() -> MockServer { + let server = MockServer::start().await; + Mock::given(method("PUT")) + .respond_with(ResponseTemplate::new(200).insert_header("etag", "\"etag\"")) + .mount(&server) + .await; + server +} + +fn http_date_from(headers: &wiremock::http::HeaderMap, name: &str) -> Option { + use aws_smithy_types::{DateTime, date_time::Format}; + headers + .get(name) + .and_then(|value| DateTime::from_str(value.to_str().ok()?, Format::HttpDate).ok()) + .map(|date| UNIX_EPOCH + Duration::new(date.secs() as u64, date.subsec_nanos())) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn set_writes_python_metadata_with_and_without_ttl() { + let server = mock_server().await; + let cache = cache(&server.uri()); + let context = ExactCacheContext { + ttl: Some(Duration::from_secs(90)), + }; + cache + .set_cache("alpha:beta", json!({"answer": 1}), &context) + .unwrap(); + cache + .set_cache("plain", json!({"answer": 2}), &ExactCacheContext::default()) + .unwrap(); + + let requests = server.received_requests().await.unwrap(); + let ttl_request = requests + .iter() + .find(|request| request.url.path() == "/cache-bucket/team/alpha/beta") + .expect("ttl write should hit the converted S3 key"); + assert_eq!( + ttl_request.headers["cache-control"].to_str().unwrap(), + "immutable, max-age=90, s-maxage=90" + ); + assert_eq!( + ttl_request.headers["content-type"].to_str().unwrap(), + "application/json" + ); + assert_eq!( + ttl_request.headers["content-language"].to_str().unwrap(), + "en" + ); + assert_eq!( + ttl_request.headers["content-disposition"].to_str().unwrap(), + "inline; filename=\"team/alpha/beta.json\"" + ); + let expires = http_date_from(&ttl_request.headers, "expires").expect("ttl write sets Expires"); + let remaining = expires.duration_since(SystemTime::now()).unwrap(); + assert!(remaining > Duration::from_secs(60) && remaining <= Duration::from_secs(91)); + assert_eq!( + serde_json::from_slice::(&ttl_request.body).unwrap(), + json!({"answer": 1}) + ); + + let plain = requests + .iter() + .find(|request| request.url.path() == "/cache-bucket/team/plain") + .expect("no-ttl write should hit the converted S3 key"); + assert_eq!( + plain.headers["cache-control"].to_str().unwrap(), + "immutable, max-age=31536000, s-maxage=31536000" + ); + assert!(plain.headers.get("expires").is_none()); + assert_eq!( + plain.headers["content-disposition"].to_str().unwrap(), + "inline; filename=\"team/plain.json\"" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn get_hit_miss_expired_and_invalid_entries() { + let server = mock_server().await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/hit")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 3}))) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/missing")) + .respond_with( + ResponseTemplate::new(404).set_body_string("NoSuchKey"), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/denied")) + .respond_with( + ResponseTemplate::new(403).set_body_string("AccessDenied"), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/expired")) + .respond_with( + ResponseTemplate::new(200) + .insert_header("expires", "Thu, 01 Jan 1970 00:00:00 GMT") + .set_body_json(json!({"answer": 4})), + ) + .mount(&server) + .await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/malformed")) + .respond_with(ResponseTemplate::new(200).set_body_string("not a cache entry")) + .mount(&server) + .await; + let cache = cache(&server.uri()); + let context = ExactCacheContext::default(); + + assert_eq!( + cache.get_cache("hit", &context).unwrap(), + Some(json!({"answer": 3})) + ); + assert_eq!(cache.get_cache("missing", &context).unwrap(), None); + assert_eq!(cache.get_cache("denied", &context).unwrap(), None); + assert_eq!(cache.get_cache("expired", &context).unwrap(), None); + assert_eq!( + cache.get_cache("malformed", &context), + Err(Error::InvalidEntry) + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn batch_get_preserves_order_with_hits_misses_and_invalid() { + let server = mock_server().await; + for (key, status, body) in [ + ("first", 200, "{\"answer\": 1}"), + ("invalid", 200, "garbage"), + ] { + Mock::given(method("GET")) + .and(path(format!("/cache-bucket/team/{key}"))) + .respond_with(ResponseTemplate::new(status).set_body_string(body)) + .mount(&server) + .await; + } + Mock::given(method("GET")) + .and(path("/cache-bucket/team/miss")) + .respond_with(ResponseTemplate::new(404)) + .mount(&server) + .await; + let cache = cache(&server.uri()); + let context = ExactCacheContext::default(); + let keys = vec![ + "first".to_string(), + "miss".to_string(), + "invalid".to_string(), + ]; + + let entries = cache.batch_get_cache(&keys, &context).unwrap(); + + assert_eq!( + entries, + vec![ + BatchEntry::Hit(json!({"answer": 1})), + BatchEntry::Miss, + BatchEntry::Invalid, + ] + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unsupported_and_noop_capabilities_match_python() { + let server = mock_server().await; + let cache = cache(&server.uri()); + + assert_eq!( + cache.test_connection().await, + Err(Error::UnsupportedOperation) + ); + cache.flush_cache().unwrap(); + cache.disconnect().await.unwrap(); + assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None); + assert_eq!( + cache.get_ttl(&ExactCacheContext { + ttl: Some(Duration::from_secs(45)), + }), + Some(Duration::from_secs(45)) + ); + assert!(server.received_requests().await.unwrap().is_empty()); +} + +#[test] +fn key_conversion_prefixes_and_splits_colons() { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(1) + .enable_all() + .build() + .unwrap(); + let _guard = runtime.enter(); + let cache = S3Cache::new( + S3CacheConfig { + key_prefix: "team/".to_string(), + ..config("http://localhost".to_string()) + }, + JsonCodec::::new(), + runtime.handle().clone(), + ); + + assert_eq!(cache.bucket(), "cache-bucket"); + assert_eq!(cache.key_prefix(), "team/"); + assert_eq!(cache.to_s3_key("a:b:c"), "team/a/b/c"); + assert_eq!(cache.to_s3_key("plain"), "team/plain"); + + let unprefixed = S3Cache::new( + S3CacheConfig { + key_prefix: String::new(), + ..config("http://localhost".to_string()) + }, + JsonCodec::::new(), + runtime.handle().clone(), + ); + assert_eq!(unprefixed.to_s3_key("a:b"), "a/b"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn sync_methods_block_inside_and_outside_the_runtime() { + let server = mock_server().await; + Mock::given(method("GET")) + .and(path("/cache-bucket/team/key")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 9}))) + .mount(&server) + .await; + let uri = server.uri(); + let cache = tokio::task::spawn_blocking(move || { + let cache = cache(&uri); + let context = ExactCacheContext::default(); + cache + .set_cache("key", json!({"answer": 9}), &context) + .unwrap(); + cache.get_cache("key", &context).unwrap() + }) + .await + .unwrap(); + + assert_eq!(cache, Some(json!({"answer": 9}))); +} diff --git a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml new file mode 100644 index 00000000000..f98bb5a5fa8 --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml @@ -0,0 +1,20 @@ +[package] +name = "litellm-cache-valkey-semantic" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-cache-redis.workspace = true +litellm-cache-response.workspace = true +redis = { version = "1.7.0", features = ["tls-rustls"] } +serde_json.workspace = true +sha2.workspace = true +tokio.workspace = true +uuid = { version = "1", features = ["v4"] } + +[dev-dependencies] +redis-test = "1.0.4" +rstest.workspace = true diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs new file mode 100644 index 00000000000..6062ccc842c --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs @@ -0,0 +1,1153 @@ +use std::{ + future::Future, + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext}; +use litellm_cache_redis::{ + RedisTopology, + connection::{ConnectionRef, Connections}, +}; +use litellm_cache_response::CacheEntry; +use serde_json::Value; +use sha2::{Digest, Sha256}; +use uuid::Uuid; + +pub trait Embedder: Send + Sync + 'static { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error>; + + fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> impl Future, Error>> + Send; +} + +pub struct PreparedEmbedding(pub Vec); + +impl Embedder for PreparedEmbedding { + fn embed(&self, _prompt: &str, _metadata: Option<&Value>) -> Result, Error> { + Ok(self.0.clone()) + } + + async fn async_embed( + &self, + _prompt: &str, + _metadata: Option<&Value>, + ) -> Result, Error> { + Ok(self.0.clone()) + } +} + +#[derive(Clone, Debug, PartialEq)] +pub struct ValkeySemanticConfig { + pub similarity_threshold: f64, + pub index_name: String, +} + +pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; + +#[derive(Clone)] +struct IndexState { + name: String, + prefix: String, + dimension: Arc>>, + similarity_threshold: f64, +} + +pub struct ValkeySemanticCache< + E: Embedder, + S: CacheCodec, + C = redis::Connection, +> { + connections: Arc>, + embedder: E, + codec: S, + config: ValkeySemanticConfig, + index_dimension: Arc>>, +} + +impl ValkeySemanticCache +where + E: Embedder, + S: CacheCodec, +{ + pub fn new( + url: &str, + embedder: E, + codec: S, + config: ValkeySemanticConfig, + ) -> Result { + 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, + } + } +} + +impl ValkeySemanticCache +where + E: Embedder, + S: CacheCodec + Clone, + C: redis::ConnectionLike + Send + 'static, +{ + pub fn with_embedder(&self, embedder: E2) -> ValkeySemanticCache { + ValkeySemanticCache { + connections: Arc::clone(&self.connections), + embedder, + codec: self.codec.clone(), + config: self.config.clone(), + index_dimension: Arc::clone(&self.index_dimension), + } + } +} + +impl BaseCache for ValkeySemanticCache +where + E: Embedder, + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(()); + }; + let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; + let scope = scope_tag(key); + let response = self.codec.encode(&value)?; + let vector = embedding_bytes(&embedding); + let index = self.index_state(); + self.connections.execute(|connection| { + write_document( + connection, + &index, + &scope, + &prompt, + response, + vector, + self.get_ttl(context), + ) + }) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; + let scope = scope_tag(key); + let vector = embedding_bytes(&embedding); + let index = self.index_state(); + let response = self.connections.execute(|connection| { + search_document(connection, &index, &scope, vector, embedding.len()) + })?; + let Some(response) = response else { + return Ok(None); + }; + self.codec.decode(&response).map(Some) + } + + fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> impl Future> + Send { + let key = key.to_owned(); + let prompt = prompt_from_context(&context); + let metadata = context.metadata.clone(); + async move { + let Some(prompt) = prompt else { + return Ok(()); + }; + let embedding = self + .embedder + .async_embed(&prompt, metadata.as_ref()) + .await?; + let connections = Arc::clone(&self.connections); + let index = self.index_state(); + let response = self.codec.encode(&value)?; + let vector = embedding_bytes(&embedding); + let scope = scope_tag(&key); + let ttl = context.ttl; + Connections::run_blocking(connections, move |connection| { + write_document(connection, &index, &scope, &prompt, response, vector, ttl) + }) + .await + } + } + + fn async_get_cache( + &self, + key: &str, + context: &Self::Context, + ) -> impl Future, Error>> + Send { + let key = key.to_owned(); + let prompt = prompt_from_context(context); + let metadata = context.metadata.clone(); + async move { + let Some(prompt) = prompt else { + return Ok(None); + }; + let embedding = self + .embedder + .async_embed(&prompt, metadata.as_ref()) + .await?; + let connections = Arc::clone(&self.connections); + let index = self.index_state(); + Connections::run_blocking(connections, move |connection| { + let scope = scope_tag(&key); + let vector = embedding_bytes(&embedding); + search_document(connection, &index, &scope, vector, embedding.len()) + }) + .await + .and_then(|response| response.map(|bytes| self.codec.decode(&bytes)).transpose()) + } + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Err(Error::UnsupportedOperation) + } +} + +pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { + if let Some(Value::Array(messages)) = context.messages.as_ref() + && !messages.is_empty() + { + return messages + .iter() + .filter_map(Value::as_object) + .map(message_text) + .collect(); + } + let input = context.input.as_ref()?; + let mut parts = Vec::new(); + collect_input_text(input, &mut parts); + let prompt = parts.join("\n").trim().to_owned(); + (!prompt.is_empty()).then_some(prompt) +} + +fn message_text(message: &serde_json::Map) -> Option { + let content = match message.get("content") { + Some(Value::String(value)) => value.clone(), + Some(Value::Array(parts)) => { + let mut content = String::new(); + for part in parts { + let part = part.as_object()?; + if let Some(text) = part.get("text").and_then(Value::as_str) { + content.push_str(text); + } + } + content + } + _ => String::new(), + }; + Some(format!( + "{content}{}", + search_results_text(message.get("search_results")) + )) +} + +fn search_results_text(value: Option<&Value>) -> String { + let Some(Value::Array(results)) = value else { + return String::new(); + }; + results + .iter() + .filter_map(Value::as_object) + .map(|result| { + let source = result.get("source").and_then(Value::as_str).unwrap_or(""); + let title = result.get("title").and_then(Value::as_str).unwrap_or(""); + let content = result + .get("content") + .and_then(Value::as_array) + .map(|blocks| { + blocks + .iter() + .filter_map(Value::as_object) + .filter_map(|block| block.get("text").and_then(Value::as_str)) + .collect::() + }) + .unwrap_or_default(); + let citations = result + .get("citations") + .filter(|value| !value.is_null()) + .and_then(|value| serde_json::to_string(value).ok()) + .unwrap_or_default(); + format!("{source}{title}{content}{citations}") + }) + .collect() +} + +fn collect_input_text(value: &Value, parts: &mut Vec) { + match value { + Value::String(value) => { + let value = value.trim(); + if !value.is_empty() { + parts.push(value.to_owned()); + } + } + Value::Array(values) => values + .iter() + .for_each(|value| collect_input_text(value, parts)), + Value::Object(object) => { + if let Some(content) = object.get("content").filter(|value| !value.is_null()) { + collect_input_text(content, parts); + return; + } + for key in ["text", "output", "input_text", "output_text"] { + if let Some(Value::String(value)) = object.get(key) { + let value = value.trim(); + if !value.is_empty() { + parts.push(value.to_owned()); + return; + } + } + } + } + _ => {} + } +} + +fn scope_tag(key: &str) -> String { + let digest = Sha256::digest(key.as_bytes()); + digest.iter().map(|byte| format!("{byte:02x}")).collect() +} + +fn embedding_bytes(embedding: &[f32]) -> Vec { + embedding + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() +} + +fn write_document( + connection: &mut ConnectionRef<'_>, + index: &IndexState, + scope: &str, + prompt: &str, + response: Vec, + vector: Vec, + ttl: Option, +) -> Result<(), Error> { + let dimension = vector.len() / std::mem::size_of::(); + ensure_index( + connection, + &index.name, + &index.prefix, + &index.dimension, + dimension, + )?; + 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) +} + +fn search_document( + connection: &mut ConnectionRef<'_>, + index: &IndexState, + scope: &str, + vector: Vec, + dimension: usize, +) -> Result>, Error> { + ensure_index( + connection, + &index.name, + &index.prefix, + &index.dimension, + dimension, + )?; + 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(None); + }; + let response = fields + .iter() + .find_map(|(name, value)| (name == "response").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = fields + .iter() + .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) + .ok_or(Error::InvalidEntry)?; + let distance = parse_f64(&distance)?; + if 1.0 - distance < index.similarity_threshold { + return Ok(None); + } + Ok(Some(response)) +} + +fn ensure_index( + connection: &mut ConnectionRef<'_>, + index_name: &str, + prefix: &str, + index_dimension: &Mutex>, + dimension: usize, +) -> Result<(), Error> { + if index_dimension + .lock() + .map_err(|_| Error::Unavailable)? + .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(prefix) + .arg("SCHEMA") + .arg("litellm_cache_key") + .arg("TAG") + .arg("embedding") + .arg("VECTOR") + .arg("HNSW") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dimension) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .query::(connection) + .map(|_| ()) + .map_err(|error| error.to_string()); + if let Err(message) = create { + if !message.to_ascii_lowercase().contains("already exists") { + return Err(Error::Unavailable); + } + let info = redis::cmd("FT.INFO") + .arg(index_name) + .query::(connection) + .map_err(|_| Error::Unavailable)?; + let existing = index_dimension_from_info(&info).ok_or(Error::Unavailable)?; + if existing != dimension { + return Err(Error::Unavailable); + } + } + *index_dimension.lock().map_err(|_| Error::Unavailable)? = Some(dimension); + Ok(()) +} + +fn index_dimension_from_info(value: &redis::Value) -> Option { + let redis::Value::Array(values) = value else { + return None; + }; + let attributes = values.windows(2).find_map(|pair| { + (value_text(&pair[0]).as_deref() == Some("attributes")).then_some(&pair[1]) + })?; + let redis::Value::Array(fields) = attributes else { + return None; + }; + fields.iter().find_map(|field| { + let redis::Value::Array(values) = field else { + return None; + }; + let flattened = values.iter().flat_map(|value| match value { + redis::Value::Array(values) => values.as_slice(), + _ => std::slice::from_ref(value), + }); + let values = flattened.collect::>(); + values.windows(2).find_map(|pair| { + if value_text(pair[0]).as_deref() == Some("dimensions") { + return value_text(pair[1]).and_then(|value| value.parse().ok()); + } + None + }) + }) +} + +type SearchFields = Vec<(String, Vec)>; + +fn search_fields(value: redis::Value) -> Result, Error> { + let redis::Value::Array(values) = value else { + return Err(Error::InvalidEntry); + }; + let total = parse_i64(values.first().ok_or(Error::InvalidEntry)?)?; + if total <= 0 || values.len() < 3 { + return Ok(None); + } + let redis::Value::Array(fields) = &values[2] else { + return Err(Error::InvalidEntry); + }; + let (pairs, remainder) = fields.as_chunks::<2>(); + if !remainder.is_empty() { + return Err(Error::InvalidEntry); + } + let pairs = pairs + .iter() + .map(|pair| { + Ok(( + value_text(&pair[0]).ok_or(Error::InvalidEntry)?, + value_bytes(&pair[1])?, + )) + }) + .collect::, Error>>()?; + Ok(Some(pairs)) +} + +fn parse_i64(value: &redis::Value) -> Result { + value_text(value) + .ok_or(Error::InvalidEntry)? + .parse() + .map_err(|_| Error::InvalidEntry) +} + +fn parse_f64(value: &[u8]) -> Result { + std::str::from_utf8(value) + .map_err(|_| Error::InvalidEntry)? + .parse() + .map_err(|_| Error::InvalidEntry) +} + +fn value_text(value: &redis::Value) -> Option { + match value { + redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), + redis::Value::SimpleString(value) => Some(value.clone()), + redis::Value::Int(value) => Some(value.to_string()), + _ => None, + } +} + +fn value_bytes(value: &redis::Value) -> Result, Error> { + match value { + redis::Value::BulkString(bytes) => Ok(bytes.clone()), + redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()), + redis::Value::Int(value) => Ok(value.to_string().into_bytes()), + _ => Err(Error::InvalidEntry), + } +} + +#[cfg(test)] +mod tests { + use std::{ + collections::VecDeque, + sync::{Arc, Mutex}, + time::Duration, + }; + + use litellm_cache::{BaseCache, CacheCodec}; + use litellm_cache_response::{ + CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, + }; + use redis_test::MockRedisConnection; + use rstest::rstest; + use serde_json::{Value, json}; + + use super::{ + Embedder, PreparedEmbedding, ValkeySemanticCache, ValkeySemanticConfig, + index_dimension_from_info, prompt_from_context, scope_tag, + }; + + #[derive(Clone)] + struct FixedEmbedder { + vector: Vec, + calls: EmbedderCalls, + } + + type EmbedderCalls = Arc)>>>; + type RecordingCache = + ValkeySemanticCache; + type RecordingSetup = (RecordingCache, Arc>>>, EmbedderCalls); + + impl Embedder for FixedEmbedder { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, super::Error> { + self.calls + .lock() + .unwrap() + .push((prompt.into(), metadata.cloned())); + Ok(self.vector.clone()) + } + + async fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> Result, super::Error> { + self.embed(prompt, metadata) + } + } + + struct RecordingConnection { + requests: Arc>>>, + replies: Mutex>>, + } + + impl RecordingConnection { + fn new(replies: impl IntoIterator>) -> Self { + Self { + requests: Arc::default(), + replies: Mutex::new(replies.into_iter().collect()), + } + } + + fn requests(&self) -> Arc>>> { + Arc::clone(&self.requests) + } + + fn reply(&self) -> redis::RedisResult { + self.replies + .lock() + .unwrap() + .pop_front() + .unwrap_or_else(|| Ok(redis::Value::SimpleString("OK".into()))) + } + } + + impl redis::ConnectionLike for RecordingConnection { + fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult { + self.requests.lock().unwrap().push(command.to_vec()); + self.reply() + } + + fn req_packed_commands( + &mut self, + command: &[u8], + _offset: usize, + count: usize, + ) -> redis::RedisResult> { + self.requests.lock().unwrap().push(command.to_vec()); + (0..count).map(|_| self.reply()).collect() + } + + fn get_db(&self) -> i64 { + 0 + } + + fn check_connection(&mut self) -> bool { + true + } + + fn is_open(&self) -> bool { + true + } + } + + fn context( + messages: Option, + input: Option, + ) -> litellm_cache::SemanticCacheContext { + litellm_cache::SemanticCacheContext { + messages, + input, + ..Default::default() + } + } + + #[rstest] + #[case(json!([{"content": "hello"}]), None, Some("hello"))] + #[case(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))] + #[case(json!([{"content": ["raw", {"text": "hello"}]}]), None, None)] + #[case(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))] + #[case(Value::Array(vec![]), Some(json!(" hello ")), Some("hello"))] + #[case(Value::Array(vec![]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))] + #[case(Value::Array(vec![]), Some(json!(" ")), None)] + fn prompt_shapes( + #[case] messages: Value, + #[case] input: Option, + #[case] expected: Option<&str>, + ) { + assert_eq!( + prompt_from_context(&context(Some(messages), input)), + expected.map(str::to_owned) + ); + } + + #[test] + fn scope_tags_are_lowercase_sha256() { + assert_eq!( + scope_tag("key"), + "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683" + ); + } + + #[test] + fn existing_index_dimension_is_read_from_attributes() { + let info = redis::Value::Array(vec![ + redis::Value::SimpleString("attributes".into()), + redis::Value::Array(vec![redis::Value::Array(vec![ + redis::Value::SimpleString("identifier".into()), + redis::Value::SimpleString("embedding".into()), + redis::Value::Array(vec![ + redis::Value::SimpleString("dimensions".into()), + redis::Value::SimpleString("2".into()), + ]), + ])]), + ]); + assert_eq!(index_dimension_from_info(&info), Some(2)); + } + + #[tokio::test] + async fn unsupported_connection_test_is_reported() { + let cache = ValkeySemanticCache::with_connection( + MockRedisConnection::new([]).assert_all_commands_consumed(), + FixedEmbedder { + vector: vec![1.0, 0.0], + calls: Arc::default(), + }, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold: 0.8, + index_name: "test".into(), + }, + ); + assert_eq!( + cache.test_connection().await, + Err(super::Error::UnsupportedOperation) + ); + } + + #[tokio::test] + async fn prepared_embedding_returns_its_vector_for_any_prompt() { + let embedding = PreparedEmbedding(vec![1.0, 2.0]); + assert_eq!( + embedding + .async_embed("different prompt", None) + .await + .unwrap(), + vec![1.0, 2.0] + ); + } + + #[test] + fn with_embedder_shares_index_state_and_connections() { + let entry = CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "ok"}), + }; + let encoded = ResponseCacheCodec.encode(&entry).unwrap(); + let cache = ValkeySemanticCache::with_connection( + RecordingConnection::new([ok(), ok(), Ok(search_hit(encoded, "0.1"))]), + FixedEmbedder { + vector: vec![1.0, 0.0], + calls: Arc::default(), + }, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold: 0.8, + index_name: "test".into(), + }, + ); + cache + .set_cache("key", entry.clone(), &semantic_context(None)) + .unwrap(); + let prepared = cache.with_embedder(PreparedEmbedding(vec![1.0, 0.0])); + assert_eq!( + prepared.get_cache("key", &semantic_context(None)).unwrap(), + Some(entry) + ); + } + + #[test] + fn missing_prompt_does_not_touch_redis() { + let cache = ValkeySemanticCache::with_connection( + MockRedisConnection::new([]).assert_all_commands_consumed(), + FixedEmbedder { + vector: vec![1.0, 0.0], + calls: Arc::default(), + }, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold: 0.8, + index_name: "test".into(), + }, + ); + assert_eq!(cache.get_cache("key", &context(None, None)).unwrap(), None); + assert_eq!(cache.get_ttl(&context(None, None)), None); + } + + fn semantic_context(ttl: Option) -> litellm_cache::SemanticCacheContext { + litellm_cache::SemanticCacheContext { + messages: Some(json!([{"role": "user", "content": "hello"}])), + metadata: Some(json!({"source": "test"})), + ttl, + ..Default::default() + } + } + + fn cache_with_recording( + replies: impl IntoIterator>, + vector: Vec, + threshold: f64, + ) -> RecordingSetup { + let connection = RecordingConnection::new(replies); + let requests = connection.requests(); + let calls: EmbedderCalls = Arc::default(); + let cache = ValkeySemanticCache::with_connection( + connection, + FixedEmbedder { + vector, + calls: Arc::clone(&calls), + }, + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold: threshold, + index_name: "test".into(), + }, + ); + (cache, requests, calls) + } + + fn ok() -> redis::RedisResult { + Ok(redis::Value::SimpleString("OK".into())) + } + + fn already_exists() -> redis::RedisResult { + Err(redis::RedisError::from(( + redis::ErrorKind::Io, + "already exists", + ))) + } + + fn info_dimension(dimension: usize) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::SimpleString("attributes".into()), + redis::Value::Array(vec![redis::Value::Array(vec![ + redis::Value::SimpleString("embedding".into()), + redis::Value::Array(vec![ + redis::Value::SimpleString("dimensions".into()), + redis::Value::Int(dimension as i64), + ]), + ])]), + ]) + } + + fn search_hit(response: Vec, distance: &str) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"test:document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"response".to_vec()), + redis::Value::BulkString(response), + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(distance.as_bytes().to_vec()), + ]), + ]) + } + + fn requests_text(requests: &Arc>>>) -> String { + requests + .lock() + .unwrap() + .iter() + .map(|request| String::from_utf8_lossy(request)) + .collect::>() + .join("\n") + } + + #[test] + fn set_without_ttl_writes_hset_without_expire() { + let (cache, requests, calls) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); + cache + .set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(None), + ) + .unwrap(); + let text = requests_text(&requests); + assert!(text.contains("FT.CREATE")); + assert!(text.contains("HSET")); + assert!( + text.contains("test:2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683:") + ); + assert!(!text.contains("EXPIRE")); + assert_eq!( + *calls.lock().unwrap(), + vec![("hello".into(), Some(json!({"source": "test"})))] + ); + } + + #[test] + fn set_with_ttl_truncates_expire_seconds() { + let (cache, requests, _) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); + cache + .set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(Some(Duration::from_millis(1900))), + ) + .unwrap(); + let text = requests_text(&requests); + assert!(text.contains("EXPIRE")); + assert!(text.contains("\r\n$1\r\n1\r\n")); + } + + #[test] + fn second_set_skips_create_after_dimension_is_cached() { + let (cache, requests, _) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); + let context = semantic_context(None); + let entry = CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }; + cache.set_cache("key", entry.clone(), &context).unwrap(); + cache.set_cache("key", entry, &context).unwrap(); + let text = requests_text(&requests); + assert_eq!(text.matches("FT.CREATE").count(), 1); + assert_eq!(text.matches("HSET").count(), 2); + } + + #[test] + fn existing_index_dimension_must_match_embedding() { + let (cache, _, _) = cache_with_recording( + [already_exists(), Ok(info_dimension(2))], + vec![1.0, 0.0], + 0.8, + ); + cache + .set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(None), + ) + .unwrap(); + + let (cache, _, _) = cache_with_recording( + [already_exists(), Ok(info_dimension(3))], + vec![1.0, 0.0], + 0.8, + ); + assert_eq!( + cache.set_cache( + "key", + CacheEntry { + timestamp: None, + response: json!({"answer": "ok"}), + }, + &semantic_context(None), + ), + Err(super::Error::Unavailable) + ); + } + + #[test] + fn get_applies_threshold_and_decodes_entry() { + let entry = CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "ok"}), + }; + let encoded = ResponseCacheCodec.encode(&entry).unwrap(); + let (cache, _, _) = cache_with_recording( + [ok(), Ok(search_hit(encoded.clone(), "0.1"))], + vec![1.0, 0.0], + 0.8, + ); + assert_eq!( + cache.get_cache("key", &semantic_context(None)).unwrap(), + Some(entry) + ); + + let (cache, _, _) = + cache_with_recording([ok(), Ok(search_hit(encoded, "0.5"))], vec![1.0, 0.0], 0.8); + assert_eq!( + cache.get_cache("key", &semantic_context(None)).unwrap(), + None + ); + } + + #[test] + fn get_zero_docs_is_a_miss() { + let (cache, _, _) = cache_with_recording( + [ok(), Ok(redis::Value::Array(vec![redis::Value::Int(0)]))], + vec![1.0, 0.0], + 0.8, + ); + assert_eq!( + cache.get_cache("key", &semantic_context(None)).unwrap(), + None + ); + } + + #[rstest] + #[case(redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(b"0.1".to_vec()), + ]), + ]))] + #[case(redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"response".to_vec()), + redis::Value::BulkString(b"not-json".to_vec()), + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(b"abc".to_vec()), + ]), + ]))] + fn malformed_entries_are_invalid(#[case] search: redis::Value) { + let (cache, _, _) = cache_with_recording([ok(), Ok(search)], vec![1.0, 0.0], 0.8); + assert_eq!( + cache.get_cache("key", &semantic_context(None)), + Err(super::Error::InvalidEntry) + ); + } + + #[test] + fn response_cache_turns_invalid_entries_into_misses() { + let (cache, _, _) = cache_with_recording( + [ + ok(), + Ok(redis::Value::Array(vec![ + redis::Value::Int(1), + redis::Value::BulkString(b"document".to_vec()), + redis::Value::Array(vec![ + redis::Value::BulkString(b"response".to_vec()), + redis::Value::BulkString(b"not-json".to_vec()), + redis::Value::BulkString(b"vector_distance".to_vec()), + redis::Value::BulkString(b"0.1".to_vec()), + ]), + ])), + ], + vec![1.0, 0.0], + 0.8, + ); + let service = ResponseCache::new(Arc::new(cache)); + let request = ResponseCacheRequest { + key: CacheKeyInput { + preset: Some("key".into()), + ..Default::default() + }, + context: semantic_context(None), + ..ResponseCacheRequest::new(CacheKeyInput::default()) + }; + assert_eq!(service.lookup(&request, Duration::ZERO).unwrap(), None); + } + + #[tokio::test] + async fn async_set_and_get_use_shared_document_helpers() { + let entry = CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "ok"}), + }; + let encoded = ResponseCacheCodec.encode(&entry).unwrap(); + let (cache, requests, calls) = cache_with_recording( + [ok(), ok(), ok(), Ok(search_hit(encoded, "0.1"))], + vec![1.0, 0.0], + 0.8, + ); + let context = semantic_context(Some(Duration::from_millis(1900))); + cache + .async_set_cache("key", entry.clone(), context.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache("key", &context).await.unwrap(), + Some(entry) + ); + let text = requests_text(&requests); + assert!(text.contains("FT.CREATE")); + assert!(text.contains("HSET")); + assert!(text.contains("EXPIRE")); + assert_eq!(calls.lock().unwrap().len(), 2); + } +} diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 6f961798795..5c10e7fd5c3 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -35,8 +35,8 @@ impl CacheContext for ExactCacheContext { #[derive(Clone, Debug, Default, PartialEq)] pub struct SemanticCacheContext { pub input: Option, - pub messages: Vec, - pub metadata: serde_json::Map, + pub messages: Option, + pub metadata: Option, pub scope: Option, pub ttl: Option, } @@ -48,11 +48,8 @@ impl CacheContext for SemanticCacheContext { fn with_ttl(&self, ttl: Option) -> Self { Self { - input: self.input.clone(), - messages: self.messages.clone(), - metadata: self.metadata.clone(), - scope: self.scope.clone(), ttl, + ..self.clone() } } } @@ -130,3 +127,31 @@ pub trait BaseCache: Send + Sync { fn test_connection(&self) -> impl Future> + Send; } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use serde_json::json; + + use super::{CacheContext, SemanticCacheContext}; + + #[test] + fn semantic_context_with_ttl_only_replaces_ttl() { + let context = SemanticCacheContext { + input: Some(json!({"input": "hello"})), + messages: Some(json!([{"role": "user", "content": "hello"}])), + metadata: Some(json!({"tenant": "team"})), + scope: Some("scope".into()), + ttl: Some(Duration::from_secs(10)), + }; + + let updated = context.with_ttl(Some(Duration::from_secs(20))); + + assert_eq!(updated.ttl, Some(Duration::from_secs(20))); + assert_eq!(updated.input, context.input); + assert_eq!(updated.messages, context.messages); + assert_eq!(updated.metadata, context.metadata); + assert_eq!(updated.scope, context.scope); + } +} diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 33171f36f46..36307ac9b33 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -131,8 +131,8 @@ fn associated_context_preserves_backend_specific_lookup_inputs() { fn semantic_context_with_ttl_preserves_lookup_inputs() { let context = SemanticCacheContext { input: Some(serde_json::json!("text")), - messages: vec![serde_json::json!({"role": "user", "content": "hi"})], - metadata: serde_json::Map::from_iter([("key".into(), serde_json::json!("value"))]), + messages: Some(serde_json::json!([{"role": "user", "content": "hi"}])), + metadata: Some(serde_json::json!({"key": "value"})), scope: Some("scope".into()), ttl: None, }; diff --git a/litellm-rust/crates/framer/Cargo.toml b/litellm-rust/crates/framer/Cargo.toml index d22502f871a..62bfcc7da3d 100644 --- a/litellm-rust/crates/framer/Cargo.toml +++ b/litellm-rust/crates/framer/Cargo.toml @@ -11,7 +11,7 @@ aws = ["dep:aws-smithy-eventstream", "dep:aws-smithy-types"] sse = ["dep:sse-stream"] [dependencies] -aws-smithy-eventstream = { version = "=0.61.1", optional = true } +aws-smithy-eventstream = { version = "=0.61.4", optional = true } aws-smithy-types = { version = "1.6.1", optional = true } bytes = "1" futures-util.workspace = true diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 0cc7af1836f..f04b78feee1 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -34,7 +34,7 @@ tokio = { workspace = true, features = ["sync"] } url.workspace = true [dev-dependencies] -aws-smithy-eventstream = "=0.61.1" +aws-smithy-eventstream = "=0.61.4" aws-smithy-types = "1.6.1" rstest.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 0080a972481..bbe727f9345 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -24,12 +24,15 @@ litellm-cache.workspace = true litellm-cache-azure-blob.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true +litellm-cache-s3.workspace = true litellm-cache-gcs.workspace = true litellm-cache-disk.workspace = true litellm-cache-redis-semantic.workspace = true litellm-cache-response.workspace = true +litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" } serde.workspace = true litellm-auth.workspace = true +litellm-auth-aws.workspace = true litellm-callbacks-legacy-python.workspace = true litellm-core.workspace = true litellm-core-utils.workspace = true @@ -41,6 +44,7 @@ litellm-host-python.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true +redis = { version = "1.7.0", features = ["tls-rustls"] } serde_json.workspace = true tokio = { workspace = true, features = ["rt", "sync"] } @@ -50,6 +54,7 @@ serde_with.workspace = true criterion.workspace = true futures-util.workspace = true rstest.workspace = true +sha2.workspace = true tokio-tungstenite.workspace = true [[bench]] diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs index 0b90e8151ea..2ff73238202 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -14,7 +14,6 @@ use super::{ future::{ready_none, ready_value}, native::NativeResponseCache, request::{now, request, requests}, - semantic::{SemanticOperation, drive}, }; pub(super) enum CacheBinding { @@ -57,17 +56,7 @@ impl ResolvedCache { CacheBinding::Disabled => ready_none(py)?, CacheBinding::Native(service) => { let request = request(input)?; - if service.semantic_embedder().is_some() { - return Ok(ExecutionStep::Await( - drive(py, service.clone(), SemanticOperation::Lookup(request))?.unbind(), - )); - } - let service = service.clone(); - run_async( - py, - async move { service.async_lookup(&request, now()).await }, - cache_error, - )? + service.async_lookup_py(py, request)? } CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?, }; @@ -185,19 +174,7 @@ impl ResolvedCache { CacheBinding::Native(service) => { let request = self::request(request)?; let response: Value = from_py(response)?; - if service.semantic_embedder().is_some() { - return drive( - py, - service.clone(), - SemanticOperation::Store(request, response), - ); - } - let service = service.clone(); - run_async( - py, - async move { service.async_store(&request, response, now()).await }, - cache_error, - ) + service.async_store_py(py, request, response) } CacheBinding::PythonCallback(callback) => { callback.async_store(py, response, callback_kwargs) @@ -253,20 +230,8 @@ impl ResolvedCache { "batch cache requests and responses must have equal lengths", )); } - let entries = requests.into_iter().zip(responses).collect::>(); - if service.semantic_embedder().is_some() { - return drive( - py, - service.clone(), - SemanticOperation::StoreBatch(entries.into()), - ); - } - let service = service.clone(); - run_async( - py, - async move { service.async_store_batch(entries, now()).await }, - cache_error, - ) + let entries = requests.into_iter().zip(responses).collect(); + service.async_store_batch_py(py, entries) } CacheBinding::PythonCallback(callback) => { callback.async_store_batch(py, callback_result, callback_kwargs) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 3a7ae5d618f..ef2a03ac9ba 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -1,11 +1,13 @@ use std::{path::PathBuf, time::Duration}; +use litellm_auth_aws::AwsAuthConfig; use litellm_cache::CacheType; use litellm_cache_redis::{RedisNode, RedisTopology}; +use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; use pyo3::{ - exceptions::{PyTypeError, PyValueError}, + exceptions::{PyAttributeError, PyTypeError, PyValueError}, prelude::*, - types::{PyAny, PyDict, PyList, PyString}, + types::{PyAny, PyBool, PyDict, PyList, PyString}, }; use super::{native::NativeResponseCache, request::duration}; @@ -115,10 +117,20 @@ struct RedisClientProjection<'py> { const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31; +#[allow(dead_code, reason = "consumed by the cache activation follow-up")] +pub(super) struct ValkeySemanticCacheConfig { + pub(super) similarity_threshold: f64, + pub(super) index_name: String, + pub(super) embedding_model: String, + pub(super) connection: RedisConnectionConfig, +} + pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + S3(Box), Gcs(GcsCacheConfig), + ValkeySemantic(Box), Disk(DiskCacheConfig), AzureBlob(AzureBlobCacheConfig), RedisSemantic(Box), @@ -136,6 +148,9 @@ pub(super) enum UnsupportedCacheConfig { RedisCredentials, RedisConnection, RedisOption, + S3Client, + S3Credentials, + S3Option, GcsBucket, DiskStore, } @@ -148,6 +163,9 @@ impl UnsupportedCacheConfig { Self::RedisCredentials => "native Redis credentials require Python", Self::RedisConnection => "native Redis connection type is not implemented", Self::RedisOption => "native Redis configuration requires Python", + Self::S3Client => "native S3 client type is not implemented", + Self::S3Credentials => "native S3 credentials require Python", + Self::S3Option => "native S3 configuration requires Python", Self::GcsBucket => "native GCS cache requires a configured bucket name", Self::DiskStore => "native disk cache requires the built-in diskcache store", } @@ -192,6 +210,13 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::S3) => match project_s3(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::S3(Box::new(backend)), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, Some(CacheType::Gcs) => match project_gcs(&backend)? { Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { policy, @@ -199,6 +224,13 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::ValkeySemantic) => match project_valkey_semantic(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::ValkeySemantic(Box::new(backend)), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, Some(CacheType::Disk) => match project_disk(&backend)? { Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { policy, @@ -218,11 +250,9 @@ impl NativeCacheConfig { backend: CacheBackendConfig::RedisSemantic(Box::new(backend)), })) }), - Some(CacheType::ValkeySemantic | CacheType::S3 | CacheType::QdrantSemantic) | None => { - Ok(CacheConfigProjection::Unsupported( - UnsupportedCacheConfig::Backend, - )) - } + Some(CacheType::QdrantSemantic) | None => Ok(CacheConfigProjection::Unsupported( + UnsupportedCacheConfig::Backend, + )), } } @@ -230,12 +260,16 @@ impl NativeCacheConfig { let default_ttl = match &self.backend { CacheBackendConfig::Memory(config) => Some(config.default_ttl), CacheBackendConfig::Redis(config) => Some(config.default_ttl), + CacheBackendConfig::S3(_) => None, + CacheBackendConfig::ValkeySemantic(_) => Some(Duration::ZERO), CacheBackendConfig::Disk(_) | CacheBackendConfig::AzureBlob(_) | CacheBackendConfig::Gcs(_) | CacheBackendConfig::RedisSemantic(_) => None, }; - if service.default_ttl() != default_ttl { + if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_)) + && service.default_ttl() != default_ttl + { return Some("facade and native backend default TTLs must match"); } match &self.backend { @@ -260,6 +294,30 @@ impl NativeCacheConfig { CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) .then_some("facade and native backend namespaces must match"), + CacheBackendConfig::S3(_) if service.kind() != "s3" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::S3(config) if service.bucket() != Some(config.bucket.as_str()) => { + Some("facade and native backend buckets must match") + } + CacheBackendConfig::S3(config) + if service.key_prefix() != Some(config.key_prefix.as_str()) => + { + Some("facade and native backend key prefixes must match") + } + CacheBackendConfig::S3(config) if service.region() != Some(config.region.as_str()) => { + Some("facade and native backend regions must match") + } + CacheBackendConfig::S3(config) + if service.endpoint() + != config + .endpoint + .as_ref() + .map(|endpoint| endpoint.url.as_str()) => + { + Some("facade and native backend endpoints must match") + } + CacheBackendConfig::S3(_) => None, CacheBackendConfig::Gcs(_) if service.kind() != "gcs" => { Some("facade and native backend types must match") } @@ -285,6 +343,16 @@ impl NativeCacheConfig { Some("facade and native backend credentials must match") } CacheBackendConfig::Gcs(_) => None, + CacheBackendConfig::ValkeySemantic(config) => { + if service.kind() != "valkey-semantic" { + return Some("facade and native backend types must match"); + } + let Some((threshold, index_name)) = service.semantic_config() else { + return Some("facade and native backend types must match"); + }; + (threshold != config.similarity_threshold || index_name != config.index_name) + .then_some("facade and native semantic settings must match") + } CacheBackendConfig::Disk(_) if service.kind() != "disk" => { Some("facade and native backend types must match") } @@ -497,6 +565,77 @@ fn project_redis( })) } +#[inline(never)] +fn project_s3( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let client = backend.getattr("s3_client")?; + if !instance_class_is(&client, "botocore.client", "S3")? { + return Ok(Err(UnsupportedCacheConfig::S3Client)); + } + let meta = client.getattr("meta")?; + let Some(region) = optional_string(meta.getattr("region_name")?)? else { + return Ok(Err(UnsupportedCacheConfig::S3Option)); + }; + let Some(endpoint_url) = optional_string(meta.getattr("endpoint_url")?)? else { + return Ok(Err(UnsupportedCacheConfig::S3Option)); + }; + let client_config = meta.getattr("config")?; + for name in ["s3", "proxies", "client_cert"] { + if optional_attribute(&client_config, name)?.is_some_and(|value| !value.is_none()) { + return Ok(Err(UnsupportedCacheConfig::S3Option)); + } + } + let signature = match optional_attribute(&client_config, "signature_version")? { + Some(value) => value.extract::>()?, + None => None, + }; + if signature.as_deref() != Some("s3v4") { + return Ok(Err(UnsupportedCacheConfig::S3Option)); + } + let insecure = endpoint_url.starts_with("http://"); + let verify = optional_attribute_chain(&client, &["_endpoint", "http_session", "_verify"])?; + let verified = verify + .and_then(|value| value.cast::().ok().map(|value| value.is_true())) + .unwrap_or(false); + if !verified && !insecure { + return Ok(Err(UnsupportedCacheConfig::S3Option)); + } + let credentials = optional_attribute_chain(&client, &["_request_signer", "_credentials"])? + .ok_or(UnsupportedCacheConfig::S3Credentials); + let credentials = match credentials { + Ok(credentials) if !credentials.is_none() => credentials, + _ => return Ok(Err(UnsupportedCacheConfig::S3Credentials)), + }; + let auth = if credentials.getattr("method")?.extract::()?.as_str() == "explicit" { + AwsAuthConfig { + access_key_id: credentials + .getattr("access_key")? + .extract::>()?, + secret_access_key: credentials + .getattr("secret_key")? + .extract::>()?, + session_token: credentials.getattr("token")?.extract::>()?, + region_name: Some(region.clone()), + ..Default::default() + } + } else { + AwsAuthConfig { + region_name: Some(region.clone()), + ..Default::default() + } + }; + let default_endpoint = endpoint_url == format!("https://s3.{region}.amazonaws.com") + || (region == "us-east-1" && endpoint_url == "https://s3.amazonaws.com"); + Ok(Ok(S3CacheConfig { + bucket: backend.getattr("bucket_name")?.extract::()?, + key_prefix: backend.getattr("key_prefix")?.extract::()?, + region, + endpoint: (!default_endpoint).then_some(S3Endpoint { url: endpoint_url }), + auth, + })) +} + #[inline(never)] fn project_standalone_client<'py>( client: &Bound<'py, PyAny>, @@ -531,7 +670,7 @@ fn project_standalone_client<'py>( #[inline(never)] fn project_cluster_client<'py>( - source: &Bound<'py, PyDict>, + source: &Bound<'_, PyDict>, client: &Bound<'py, PyAny>, ) -> PyResult, UnsupportedCacheConfig>> { let Some(startup_nodes) = startup_nodes(source)? else { @@ -619,6 +758,71 @@ fn port(value: i64) -> PyResult { u16::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis port")) } +#[inline(never)] +fn project_valkey_semantic( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let client = backend.getattr("sync_client")?; + let pool = client.getattr("connection_pool")?; + let Ok((resolved, is_tls)) = project_connection_pool(&pool)? else { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + }; + for key in ["credential_provider", "redis_connect_func"] { + if has_value(&resolved, key)? { + return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); + } + } + if is_tls { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + let connection = RedisConnectionConfig { + host: required_string(&resolved, "host")?, + port: u16::try_from(required_i64(&resolved, "port")?) + .map_err(|_| PyValueError::new_err("invalid Redis port"))?, + database: optional_i64(&resolved, "db")?.unwrap_or(0), + username: optional_dict_string(&resolved, "username")?, + password: optional_dict_string(&resolved, "password")?, + protocol: RedisProtocol::Resp2, + pool_size: pool.getattr("max_connections")?.extract::()?, + read_timeout: None, + connect_timeout: None, + socket_keepalive: None, + health_check_interval: Duration::ZERO, + client_name: None, + tls: None, + }; + if connection.host.is_empty() { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + Ok(Ok(ValkeySemanticCacheConfig { + similarity_threshold: backend.getattr("similarity_threshold")?.extract()?, + index_name: backend.getattr("index_name")?.extract()?, + embedding_model: backend.getattr("embedding_model")?.extract()?, + connection, + })) +} + +#[inline(never)] +fn project_connection_pool<'py>( + pool: &Bound<'py, PyAny>, +) -> PyResult, bool), UnsupportedCacheConfig>> { + if !instance_class_is(pool, "redis.connection", "ConnectionPool")? { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + } + let resolved = pool.getattr("connection_kwargs")?.cast_into::()?; + let connection_class = resolved + .get_item("connection_class")? + .unwrap_or(pool.getattr("connection_class")?); + let is_tls = if class_is(&connection_class, "redis.connection", "Connection")? { + false + } else if class_is(&connection_class, "redis.connection", "SSLConnection")? { + true + } else { + return Ok(Err(UnsupportedCacheConfig::RedisConnection)); + }; + Ok(Ok((resolved, is_tls))) +} + #[inline(never)] fn project_tls(values: &Bound<'_, PyDict>) -> PyResult { Ok(RedisTlsConfig { @@ -700,6 +904,31 @@ fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult( + value: &Bound<'py, PyAny>, + name: &str, +) -> PyResult>> { + match value.getattr(name) { + Ok(value) => Ok(Some(value)), + Err(error) if error.is_instance_of::(value.py()) => Ok(None), + Err(error) => Err(error), + } +} + +#[inline(never)] +fn optional_attribute_chain<'py>( + value: &Bound<'py, PyAny>, + names: &[&str], +) -> PyResult>> { + names + .iter() + .try_fold(Some(value.clone()), |current, name| match current { + Some(current) => optional_attribute(¤t, name), + None => Ok(None), + }) +} + #[inline(never)] fn optional_string(value: Bound<'_, PyAny>) -> PyResult> { Ok(value @@ -790,8 +1019,8 @@ mod tests { use litellm_cache_redis::{RedisNode, RedisTopology}; use super::{ - CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement, - DiskCacheConfig, GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig, + CacheBackendConfig, CacheConfigProjection, CertificateRequirement, GcsCacheConfig, + NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig, }; use crate::cache::native::NativeResponseCache; @@ -866,6 +1095,194 @@ mod tests { }); } + #[test] + fn projects_resolved_redis_tls_configuration() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "pool = ConnectionPool()\n\ + pool.connection_class = SSLConnection\n\ + pool.max_connections = 29\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6380, 'db': 4, 'username': 'user', 'password': 'secret', 'protocol': 3, 'socket_timeout': 7.5, 'socket_connect_timeout': 2, 'socket_keepalive': True, 'health_check_interval': 15, 'client_name': 'litellm', 'ssl_cert_reqs': 'optional', 'ssl_check_hostname': True, 'ssl_ca_certs': '/ca.pem', 'ssl_ca_data': 'CA DATA', 'ssl_certfile': '/client.pem', 'ssl_keyfile': '/client.key'}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\ + facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("Redis cache should be supported"); + }; + let CacheBackendConfig::Redis(redis) = config.backend else { + panic!("expected Redis configuration"); + }; + assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777)); + assert_eq!(redis.namespace.as_deref(), Some("team")); + assert_eq!(redis.flush_size, 31); + assert_eq!(redis.connection.host, "cache.internal"); + assert_eq!(redis.connection.port, 6380); + assert_eq!(redis.connection.database, 4); + assert_eq!(redis.connection.protocol, RedisProtocol::Resp3); + assert_eq!(redis.connection.pool_size, 29); + let tls = redis.connection.tls.unwrap(); + assert_eq!( + tls.certificate_requirement, + CertificateRequirement::Optional + ); + assert!(tls.check_hostname); + assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem")); + assert_eq!(tls.ca_data.as_deref(), Some("CA DATA")); + assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem")); + assert_eq!(tls.client_key.as_deref(), Some("/client.key")); + }); + } + + #[test] + fn projects_valkey_semantic_configuration() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "pool = ConnectionPool()\n\ + pool.connection_class = Connection\n\ + pool.max_connections = 12\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'db': 2}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ + facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("Valkey semantic cache should be supported"); + }; + let CacheBackendConfig::ValkeySemantic(valkey) = config.backend else { + panic!("expected Valkey semantic configuration"); + }; + assert_eq!(valkey.similarity_threshold, 0.85); + assert_eq!(valkey.index_name, "semantic_idx"); + assert_eq!(valkey.embedding_model, "text-embedding-3-small"); + assert_eq!(valkey.connection.host, "cache.internal"); + assert_eq!(valkey.connection.port, 6390); + assert_eq!(valkey.connection.database, 2); + assert_eq!(valkey.connection.pool_size, 12); + assert_eq!(valkey.connection.protocol, RedisProtocol::Resp2); + assert!(valkey.connection.tls.is_none()); + }); + } + + #[test] + fn valkey_semantic_tls_stays_on_python() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "pool = ConnectionPool()\n\ + pool.connection_class = SSLConnection\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ + facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("TLS Valkey semantic cache should stay on Python"); + }; + assert_eq!( + reason.message(), + "native Redis connection type is not implemented" + ); + }); + } + + #[test] + fn valkey_semantic_dynamic_auth_stays_on_python() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "pool = ConnectionPool()\n\ + pool.connection_class = Connection\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'credential_provider': object()}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ + facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("dynamic Valkey authentication must stay on Python"); + }; + assert_eq!(reason.message(), "native Redis credentials require Python"); + }); + } + + #[test] + fn dynamic_redis_auth_stays_on_python() { + Python::initialize(); + Python::attach(|py| { + let facade = facade( + py, + "backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\ + facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("dynamic authentication must stay on Python"); + }; + assert_eq!(reason.message(), "native Redis credentials require Python"); + }); + } + + #[test] + fn projects_cluster_startup_nodes_as_redis_topology() { + Python::initialize(); + Python::attach(|py| { + let facade = cluster_facade( + py, + "[{'host': 'node-a', 'port': 7000}, {'host': 'node-b', 'port': 7001}]", + "client.on_connect", + ); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("cluster startup nodes should project natively"); + }; + let CacheBackendConfig::Redis(redis) = &config.backend else { + panic!("expected Redis configuration"); + }; + let expected = RedisTopology::Cluster { + startup_nodes: vec![ + RedisNode { + host: "node-a".into(), + port: 7000, + }, + RedisNode { + host: "node-b".into(), + port: 7001, + }, + ], + }; + assert_eq!(redis.topology, expected); + assert_eq!(redis.connection.host, "node-a"); + assert_eq!(redis.connection.port, 7000); + assert_eq!(redis.connection.password.as_deref(), Some("secret")); + assert_eq!(redis.connection.protocol, RedisProtocol::Resp3); + assert_eq!( + redis + .connection + .tls + .as_ref() + .unwrap() + .certificate_requirement, + CertificateRequirement::None + ); + }); + } + #[test] fn projects_gcs_configuration() { Python::initialize(); @@ -931,192 +1348,6 @@ mod tests { }); } - #[test] - fn projects_resolved_redis_tls_configuration() { - Python::initialize(); - Python::attach(|py| { - let facade = facade( - py, - "pool = ConnectionPool()\n\ - pool.connection_class = SSLConnection\n\ - pool.max_connections = 29\n\ - pool.connection_kwargs = {'host': 'cache.internal', 'port': 6380, 'db': 4, 'username': 'user', 'password': 'secret', 'protocol': 3, 'socket_timeout': 7.5, 'socket_connect_timeout': 2, 'socket_keepalive': True, 'health_check_interval': 15, 'client_name': 'litellm', 'ssl_cert_reqs': 'optional', 'ssl_check_hostname': True, 'ssl_ca_certs': '/ca.pem', 'ssl_ca_data': 'CA DATA', 'ssl_certfile': '/client.pem', 'ssl_keyfile': '/client.key'}\n\ - client = SimpleNamespace(connection_pool=pool)\n\ - backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\ - facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)", - ); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("Redis cache should be supported"); - }; - let CacheBackendConfig::Redis(redis) = config.backend else { - panic!("expected Redis configuration"); - }; - assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777)); - assert_eq!(redis.namespace.as_deref(), Some("team")); - assert_eq!(redis.flush_size, 31); - assert_eq!(redis.connection.host, "cache.internal"); - assert_eq!(redis.connection.port, 6380); - assert_eq!(redis.connection.database, 4); - assert_eq!(redis.connection.protocol, RedisProtocol::Resp3); - assert_eq!(redis.connection.pool_size, 29); - let tls = redis.connection.tls.unwrap(); - assert_eq!( - tls.certificate_requirement, - CertificateRequirement::Optional - ); - assert!(tls.check_hostname); - assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem")); - assert_eq!(tls.ca_data.as_deref(), Some("CA DATA")); - assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem")); - assert_eq!(tls.client_key.as_deref(), Some("/client.key")); - }); - } - - #[test] - fn dynamic_redis_auth_stays_on_python() { - Python::initialize(); - Python::attach(|py| { - let facade = facade( - py, - "backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\ - facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)", - ); - let CacheConfigProjection::Unsupported(reason) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("dynamic authentication must stay on Python"); - }; - assert_eq!(reason.message(), "native Redis credentials require Python"); - }); - } - #[test] - fn projects_builtin_disk_configuration_and_rejects_custom_stores() { - Python::initialize(); - Python::attach(|py| { - let root = - std::env::temp_dir().join(format!("litellm-disk-config-{}", std::process::id())); - let directory = root.to_string_lossy(); - let disk_facade = facade( - py, - &format!( - "Cache = type('Cache', (), {{'__module__': 'diskcache.core'}})\n\ - Disk = type('Disk', (), {{'__module__': 'diskcache.core'}})\n\ - store = Cache()\n\ - store._disk = Disk()\n\ - store.directory = {directory:?}\n\ - backend = SimpleNamespace(disk_cache=store)\n\ - facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)" - ), - ); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&disk_facade).unwrap() - else { - panic!("disk cache should be supported"); - }; - let CacheBackendConfig::Disk(disk) = config.backend else { - panic!("expected disk configuration"); - }; - assert_eq!(disk.directory, root); - let matching = NativeResponseCache::disk(&directory).unwrap(); - assert_eq!( - (NativeCacheConfig { - policy: config.policy, - backend: CacheBackendConfig::Disk(disk), - }) - .service_mismatch(&matching), - None - ); - let other = NativeResponseCache::disk(&root.join("other").to_string_lossy()).unwrap(); - let mismatch = NativeCacheConfig { - policy: CachePolicy { - mode: "default-on".into(), - ttl: None, - namespace: None, - supported_call_types: None, - redis_flush_size: None, - semantic_cache_scope: "key".into(), - }, - backend: CacheBackendConfig::Disk(DiskCacheConfig { - directory: root.clone(), - }), - }; - assert_eq!( - mismatch.service_mismatch(&other), - Some("facade and native backend directories must match") - ); - - let custom = facade( - py, - &format!( - "CustomCache = type('CustomCache', (), {{'__module__': 'mypkg'}})\n\ - CustomDisk = type('CustomDisk', (), {{'__module__': 'mypkg'}})\n\ - store = CustomCache()\n\ - store._disk = CustomDisk()\n\ - store.directory = {directory:?}\n\ - backend = SimpleNamespace(disk_cache=store)\n\ - facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)" - ), - ); - let CacheConfigProjection::Unsupported(reason) = - NativeCacheConfig::project(&custom).unwrap() - else { - panic!("custom disk store must stay on Python"); - }; - assert_eq!( - reason.message(), - "native disk cache requires the built-in diskcache store" - ); - }); - } - - #[test] - fn projects_cluster_startup_nodes_as_redis_topology() { - Python::initialize(); - Python::attach(|py| { - let facade = cluster_facade( - py, - "[{'host': 'node-a', 'port': 7000}, {'host': 'node-b', 'port': 7001}]", - "client.on_connect", - ); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("cluster startup nodes should project natively"); - }; - let CacheBackendConfig::Redis(redis) = &config.backend else { - panic!("expected Redis configuration"); - }; - let expected = RedisTopology::Cluster { - startup_nodes: vec![ - RedisNode { - host: "node-a".into(), - port: 7000, - }, - RedisNode { - host: "node-b".into(), - port: 7001, - }, - ], - }; - assert_eq!(redis.topology, expected); - assert_eq!(redis.connection.host, "node-a"); - assert_eq!(redis.connection.port, 7000); - assert_eq!(redis.connection.password.as_deref(), Some("secret")); - assert_eq!(redis.connection.protocol, RedisProtocol::Resp3); - assert_eq!( - redis - .connection - .tls - .as_ref() - .unwrap() - .certificate_requirement, - CertificateRequirement::None - ); - }); - } - #[test] fn malformed_startup_nodes_and_foreign_connect_hooks_stay_on_python() { Python::initialize(); diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs index 99c3de34e05..10ccf396510 100644 --- a/litellm-rust/crates/python-bridge/src/cache/embedder.rs +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -1,10 +1,9 @@ use std::future::Future; use litellm_cache::Error; -use litellm_cache_redis_semantic::Embedder; use litellm_host_python::to_py; use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict}; -use serde_json::{Map, Value}; +use serde_json::Value; tokio::task_local! { static PREPARED_EMBEDDING: Result, Error>; @@ -19,11 +18,21 @@ pub(super) fn with_prepared_embedding( pub(super) struct PythonEmbedder(Py); +impl Clone for PythonEmbedder { + fn clone(&self) -> Self { + Python::attach(|py| Self(self.0.clone_ref(py))) + } +} + impl PythonEmbedder { pub(super) fn new(object: Py) -> Self { Self(object) } + pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult { + Ok(Self(backend.clone().unbind())) + } + pub(super) fn object(&self) -> &Py { &self.0 } @@ -32,16 +41,24 @@ impl PythonEmbedder { visit.call(&self.0) } + pub(super) fn async_embed_awaitable<'py>( + &self, + py: Python<'py>, + prompt: &str, + metadata: &Option, + ) -> PyResult> { + let metadata = to_py(py, metadata)?; + self.0 + .bind(py) + .call_method1("_get_async_embedding", (prompt, metadata)) + } + fn metadata_kwargs<'py>( py: Python<'py>, - metadata: &Map, + metadata: Option<&Value>, ) -> PyResult> { let kwargs = PyDict::new(py); - if metadata.is_empty() { - kwargs.set_item("metadata", py.None())?; - } else { - kwargs.set_item("metadata", to_py(py, metadata)?)?; - } + kwargs.set_item("metadata", to_py(py, &metadata)?)?; Ok(kwargs) } @@ -49,7 +66,7 @@ impl PythonEmbedder { &self, py: Python<'_>, prompt: &str, - metadata: &Map, + metadata: Option<&Value>, ) -> PyResult> { let kwargs = Self::metadata_kwargs(py, metadata)?; self.0 @@ -67,8 +84,33 @@ impl PythonEmbedder { } } -impl Embedder for PythonEmbedder { - fn embed(&self, prompt: &str, metadata: &Map) -> Result, Error> { +impl litellm_cache_valkey_semantic::Embedder for PythonEmbedder { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + let result = Python::attach(|py| -> PyResult> { + let metadata = to_py(py, &metadata)?; + self.0 + .bind(py) + .call_method1("_get_embedding", (prompt, metadata))? + .extract() + }) + .map_err(|_| Error::Unavailable)?; + Ok(result.into_iter().map(|value| value as f32).collect()) + } + + fn async_embed( + &self, + _prompt: &str, + _metadata: Option<&Value>, + ) -> impl Future, Error>> + Send { + let seeded = PREPARED_EMBEDDING + .try_with(Clone::clone) + .unwrap_or(Err(Error::Unavailable)); + std::future::ready(seeded) + } +} + +impl litellm_cache_redis_semantic::Embedder for PythonEmbedder { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { Python::attach(|py| { let kwargs = Self::metadata_kwargs(py, metadata)?; Self::extract(self.0.bind(py).call_method( @@ -83,7 +125,7 @@ impl Embedder for PythonEmbedder { fn async_embed( &self, _prompt: &str, - _metadata: &Map, + _metadata: Option<&Value>, ) -> impl Future, Error>> + Send { let seeded = PREPARED_EMBEDDING .try_with(Clone::clone) @@ -98,21 +140,19 @@ mod tests { #[tokio::test] async fn async_embed_returns_the_seeded_vector_or_unavailable() { - Python::initialize(); - let embedder = Python::attach(|py| PythonEmbedder::new(py.None())); - let metadata = Map::new(); - let embedder_ref = &embedder; - let metadata_ref = &metadata; - assert_eq!( - with_prepared_embedding(Ok(vec![0.5f32, 0.25]), async move { - embedder_ref.async_embed("prompt", metadata_ref).await - }) - .await, - Ok(vec![0.5, 0.25]) - ); - assert_eq!( - embedder.async_embed("prompt", &metadata).await, - Err(Error::Unavailable) - ); + let object = Python::attach(|py| { + Python::initialize(); + py.None() + }); + let embedder = PythonEmbedder::new(object); + let scoped_embedder = embedder.clone(); + let scoped = with_prepared_embedding(Ok(vec![0.25]), async move { + litellm_cache_redis_semantic::Embedder::async_embed(&scoped_embedder, "prompt", None) + .await + }); + assert_eq!(scoped.await, Ok(vec![0.25])); + let unscoped = + litellm_cache_redis_semantic::Embedder::async_embed(&embedder, "prompt", None).await; + assert_eq!(unscoped, Err(Error::Unavailable)); } } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index c4be7a9254a..f1389b745e9 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -31,6 +31,7 @@ struct RedisPoolGuard { connection_class: Py, connection_kwargs: Py, max_connections: Option, + client_name: &'static str, attributes: RedisPoolAttributes, } @@ -46,11 +47,18 @@ struct AzureBlobClientGuard { container_name: String, } +struct S3ClientGuard { + reference: Py, +} + enum ConnectionGuard { None, RedisPool(RedisPoolGuard), AzureBlob(AzureBlobClientGuard), + S3(S3ClientGuard), } + +#[derive(Clone, Copy)] struct RedisPoolAttributes { pool: &'static str, connection_class: &'static str, @@ -69,6 +77,8 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes { max_connections: None, }; +const VALKEY_POOL: RedisPoolAttributes = STANDALONE_POOL; + pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, @@ -155,10 +165,9 @@ impl ObjectGuard { return Ok(false); } for (name, value) in &expected.attributes { - if !attributes.get_item(name)?.is(value.bind(py)) { - return Ok(false); - } - if instance.contains(name)? && !self.config_names.contains(&name.as_str()) { + if (instance.contains(name)? && !self.config_names.contains(&name.as_str())) + || !attributes.get_item(name)?.is(value.bind(py)) + { return Ok(false); } } @@ -179,8 +188,12 @@ impl ObjectGuard { } impl RedisPoolGuard { - fn capture(backend: &Bound<'_, PyAny>, attributes: RedisPoolAttributes) -> PyResult { - let pool = backend.getattr("redis_client")?.getattr(attributes.pool)?; + fn capture( + backend: &Bound<'_, PyAny>, + client_name: &'static str, + attributes: RedisPoolAttributes, + ) -> PyResult { + let pool = backend.getattr(client_name)?.getattr(attributes.pool)?; Ok(Self { reference: pool.clone().unbind(), connection_class: pool.getattr(attributes.connection_class)?.unbind(), @@ -188,31 +201,30 @@ impl RedisPoolGuard { .getattr("connection_kwargs")? .call_method0("copy")? .unbind(), - max_connections: Self::max_connections(&pool, &attributes)?, + max_connections: attributes + .max_connections + .map(|name| pool.getattr(name)?.extract::()) + .transpose()?, + client_name, attributes, }) } - fn max_connections( - pool: &Bound<'_, PyAny>, - attributes: &RedisPoolAttributes, - ) -> PyResult> { - attributes - .max_connections - .map(|name| pool.getattr(name)?.extract::()) - .transpose() - } - fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { let pool = backend - .getattr("redis_client")? + .getattr(self.client_name)? .getattr(self.attributes.pool)?; Ok(self.reference.bind(py).is(&pool) && self .connection_class .bind(py) .is(&pool.getattr(self.attributes.connection_class)?) - && self.max_connections == Self::max_connections(&pool, &self.attributes)? + && self.max_connections + == self + .attributes + .max_connections + .map(|name| pool.getattr(name)?.extract::()) + .transpose()? && self .connection_kwargs .bind(py) @@ -274,12 +286,43 @@ impl AzureBlobClientGuard { } } +impl S3ClientGuard { + fn capture(backend: &Bound<'_, PyAny>) -> PyResult { + Ok(Self { + reference: backend.getattr("s3_client")?.unbind(), + }) + } + + fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult { + Ok(self.reference.bind(py).is(&backend.getattr("s3_client")?)) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.reference) + } +} + impl ConnectionGuard { fn capture(kind: &str, cluster: bool, backend: &Bound<'_, PyAny>) -> PyResult { Ok(match (kind, cluster) { - ("redis", false) => Self::RedisPool(RedisPoolGuard::capture(backend, STANDALONE_POOL)?), - ("redis", true) => Self::RedisPool(RedisPoolGuard::capture(backend, CLUSTER_POOL)?), + ("redis", false) => Self::RedisPool(RedisPoolGuard::capture( + backend, + "redis_client", + STANDALONE_POOL, + )?), + ("redis", true) => Self::RedisPool(RedisPoolGuard::capture( + backend, + "redis_client", + CLUSTER_POOL, + )?), + ("valkey-semantic", _) => Self::RedisPool(RedisPoolGuard::capture( + backend, + "sync_client", + VALKEY_POOL, + )?), + ("disk", _) => Self::None, ("azure-blob", _) => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?), + ("s3", _) => Self::S3(S3ClientGuard::capture(backend)?), _ => Self::None, }) } @@ -289,6 +332,7 @@ impl ConnectionGuard { Self::None => Ok(true), Self::RedisPool(guard) => guard.matches(py, backend), Self::AzureBlob(guard) => guard.matches(py, backend), + Self::S3(guard) => guard.matches(py, backend), } } @@ -297,6 +341,7 @@ impl ConnectionGuard { Self::None => Ok(()), Self::RedisPool(guard) => guard.traverse(visit), Self::AzureBlob(guard) => guard.traverse(visit), + Self::S3(guard) => guard.traverse(visit), } } } @@ -329,12 +374,18 @@ impl FacadeGuard { "redis", ), ("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"), + ("valkey-semantic", false) => ( + "litellm.caching.valkey_semantic_cache", + "ValkeySemanticCache", + "valkey-semantic", + ), ("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"), ("azure-blob", _) => ( "litellm.caching.azure_blob_cache", "AzureBlobCache", "azure-blob", ), + ("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -394,6 +445,11 @@ impl FacadeGuard { "embedding_timeout", "_index_name", "_redis_url", + "similarity_threshold", + "embedding_model", + "index_name", + "embedding_max_input_tokens", + "embedding_timeout", "bucket_name", "key_prefix", "path_service_account", diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 8c64e5bbb19..e6edf48d72f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,5 +1,8 @@ +use litellm_auth_aws::AwsAuthConfig; +use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; use litellm_cache_redis::{RedisNode, RedisTopology}; use litellm_cache_redis_semantic::RedisSemanticConfig; +use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; use litellm_host_python::{release_gil, run_sync_value}; use pyo3::{ PyTraverseError, PyVisit, @@ -7,8 +10,6 @@ use pyo3::{ prelude::*, }; -use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; - use super::{ cache_error, config::project_redis_semantic, embedder::PythonEmbedder, facade::FacadeGuard, native::NativeResponseCache, request::duration, @@ -74,6 +75,40 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (bucket, *, region, endpoint_url=None, key_prefix="", access_key_id=None, secret_access_key=None, session_token=None))] + fn s3( + py: Python<'_>, + bucket: String, + region: String, + endpoint_url: Option, + key_prefix: &str, + access_key_id: Option, + secret_access_key: Option, + session_token: Option, + ) -> PyResult { + let config = S3CacheConfig { + bucket, + key_prefix: key_prefix.to_string(), + region: region.clone(), + endpoint: endpoint_url.map(|url| S3Endpoint { url }), + auth: AwsAuthConfig { + access_key_id, + secret_access_key, + session_token, + region_name: Some(region), + ..Default::default() + }, + }; + let service = run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) })?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[staticmethod] #[pyo3(signature = (bucket_name, *, gcs_path=None, path_service_account=None, endpoint=None, token=None))] fn gcs( @@ -111,6 +146,29 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[pyo3(signature = (url, similarity_threshold, index_name, embedder))] + fn valkey_semantic( + url: String, + similarity_threshold: f64, + index_name: String, + embedder: &Bound<'_, PyAny>, + ) -> PyResult { + let python_embedder = PythonEmbedder::from_backend(embedder)?; + let service = NativeResponseCache::valkey_semantic( + &url, + similarity_threshold, + index_name, + python_embedder, + ) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[staticmethod] #[pyo3(signature = (account_url, container))] fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult { @@ -164,11 +222,17 @@ impl CacheTestHandle { fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> { let service = self.service()?; let guard = FacadeGuard::capture(py, facade, &service)?; - let service = service.with_redis_flush_size( - facade - .getattr("redis_flush_size")? - .extract::>()?, - ); + let service = service + .with_scope( + facade + .getattr("semantic_cache_scope")? + .extract::()?, + ) + .with_redis_flush_size( + facade + .getattr("redis_flush_size")? + .extract::>()?, + ); let handle = Py::new( py, Self { diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index cd772d571cb..fa028518559 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -9,6 +9,7 @@ mod native; mod request; mod resolver; mod semantic; +mod semantic_step; use litellm_cache::Error; use pyo3::{ diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index d2bced66f2c..acc457d8c9c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,6 +1,8 @@ use std::{path::Path, sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, CacheConnectionResult, Error, ExactCacheContext}; +use litellm_cache::{ + CacheCodec, CacheConnectionResult, Error, ExactCacheContext, SemanticCacheContext, +}; use litellm_cache_azure_blob::AzureBlobCache; use litellm_cache_disk::DiskCache; use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource}; @@ -8,12 +10,69 @@ use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::{RedisCache, RedisTopology}; use litellm_cache_redis_semantic::{RedisSemanticCache, RedisSemanticConfig}; use litellm_cache_response::{ - CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, + CacheEntry, CacheKeyField, PartialHits, ResponseCache, ResponseCacheCodec, + ResponseCacheRequest, WriteBuffer, }; -use pyo3::{Py, PyAny, PyTraverseError, PyVisit}; +use litellm_cache_s3::{S3Cache, S3CacheConfig}; +use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig}; +use pyo3::{PyTraverseError, PyVisit, prelude::*}; use serde_json::Value; -use super::{embedder::PythonEmbedder, request::CacheRequest}; +use super::{ + embedder::PythonEmbedder, + request::NativeRequest, + semantic::{SemanticBody, SemanticOperation, drive}, + semantic_step::{SemanticEmbedExecution, drive_semantic}, +}; + +fn semantic_key(request: &NativeRequest, scope: &str) -> litellm_cache_response::CacheKeyInput { + let mut key = request.key.clone(); + if key.preset.is_some() { + return key; + } + key.fields + .retain(|field| !matches!(field.name.as_str(), "messages" | "prompt" | "input")); + const TENANT: [&str; 3] = [ + "user_api_key", + "user_api_key_team_id", + "user_api_key_org_id", + ]; + let end_user = (scope == "end_user").then_some("user_api_key_end_user_id"); + for name in TENANT.into_iter().chain(end_user) { + let sources = [ + request.metadata.as_ref(), + request.litellm_metadata.as_ref(), + request + .litellm_params + .as_ref() + .and_then(|params| params.get("metadata")), + request + .litellm_params + .as_ref() + .and_then(|params| params.get("litellm_metadata")), + ]; + let Some(value) = sources.into_iter().flatten().find_map(|source| { + source + .as_object() + .and_then(|values| values.get(name)) + .filter(|value| !value.is_null()) + }) else { + continue; + }; + let value = match value { + Value::Null => continue, + Value::String(text) => text.clone(), + other => other.to_string(), + }; + key.fields.push(CacheKeyField { + name: name.to_owned(), + value: Some(value), + api_parameter: true, + internal_parameter: false, + }); + } + key +} #[derive(Clone)] pub(super) enum NativeResponseCache { @@ -22,10 +81,19 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + S3(Arc>>), Gcs(Arc>>), + ValkeySemantic { + cache: Arc>>, + embedder: PythonEmbedder, + scope: String, + }, + RedisSemantic { + cache: Arc>>, + embedder: PythonEmbedder, + }, Disk(Arc>>), AzureBlob(Arc>>), - RedisSemantic(Arc>>), } impl NativeResponseCache { @@ -56,9 +124,36 @@ impl NativeResponseCache { buffer: None, }) } - pub fn disk(directory: &str) -> Result { - let cache = DiskCache::open(directory, ResponseCacheCodec)?; - Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache))))) + + pub async fn s3(config: S3CacheConfig) -> Self { + let runtime = tokio::runtime::Handle::current(); + Self::S3(Arc::new(ResponseCache::new(Arc::new(S3Cache::new( + config, + ResponseCacheCodec, + runtime, + ))))) + } + + pub fn valkey_semantic( + url: &str, + similarity_threshold: f64, + index_name: String, + embedder: PythonEmbedder, + ) -> Result { + let backend = ValkeySemanticCache::new( + url, + embedder.clone(), + ResponseCacheCodec, + ValkeySemanticConfig { + similarity_threshold, + index_name, + }, + )?; + Ok(Self::ValkeySemantic { + cache: Arc::new(ResponseCache::new(Arc::new(backend))), + embedder, + scope: String::from("key"), + }) } pub fn redis_semantic( @@ -66,10 +161,16 @@ impl NativeResponseCache { embedder: PythonEmbedder, config: RedisSemanticConfig, ) -> Result { - let backend = RedisSemanticCache::new(url, embedder, config)?; - Ok(Self::RedisSemantic(Arc::new(ResponseCache::new(Arc::new( - backend, - ))))) + let backend = RedisSemanticCache::new(url, embedder.clone(), config)?; + Ok(Self::RedisSemantic { + cache: Arc::new(ResponseCache::new(Arc::new(backend))), + embedder, + }) + } + + pub fn disk(directory: &str) -> Result { + let cache = DiskCache::open(directory, ResponseCacheCodec)?; + Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache))))) } pub fn gcs(config: GcsConfig, token: Option) -> Result { @@ -105,21 +206,90 @@ impl NativeResponseCache { )), Self::Memory(_) | Self::Redis { .. } + | Self::S3(_) + | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::Disk(_) - | Self::Gcs(_) - | Self::RedisSemantic(_) => None, + | Self::Gcs(_) => None, + } + } + + fn exact(request: &NativeRequest) -> ResponseCacheRequest { + ResponseCacheRequest { + key: request.key.clone(), + controls: request.controls, + context: ExactCacheContext { ttl: request.ttl }, + max_age: request.max_age, + } + } + + pub(super) fn redis_semantic_request( + request: &NativeRequest, + ) -> ResponseCacheRequest { + ResponseCacheRequest { + key: request.key.clone(), + controls: request.controls, + context: SemanticCacheContext { + input: request.input.clone(), + messages: request.messages.clone(), + metadata: request.metadata.clone(), + scope: request.scope.clone(), + ttl: request.ttl, + }, + max_age: request.max_age, + } + } + + fn semantic( + request: &NativeRequest, + scope: &str, + ) -> ResponseCacheRequest { + ResponseCacheRequest { + key: semantic_key(request, scope), + controls: request.controls, + context: SemanticCacheContext { + input: request.input.clone(), + messages: request.messages.clone(), + metadata: request.metadata.clone(), + scope: Some(scope.to_owned()), + ttl: request.ttl, + }, + max_age: request.max_age, + } + } + + pub fn with_redis_flush_size(self, flush_size: Option) -> Self { + match self { + Self::Redis { cache, .. } => Self::Redis { + cache, + buffer: flush_size.map(|size| Arc::new(WriteBuffer::new(size))), + }, + value => value, + } + } + + pub fn with_scope(self, scope: String) -> Self { + match self { + Self::ValkeySemantic { + cache, embedder, .. + } => Self::ValkeySemantic { + cache, + embedder, + scope, + }, + value => value, } } -} -impl NativeResponseCache { pub fn kind(&self) -> &'static str { match self { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", + Self::S3(_) => "s3", Self::Gcs(_) => "gcs", + Self::ValkeySemantic { .. } => "valkey-semantic", + Self::RedisSemantic { .. } => "redis_semantic", Self::Disk(_) => "disk", - Self::RedisSemantic(_) => "redis_semantic", Self::AzureBlob(_) => "azure-blob", } } @@ -128,28 +298,65 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), + Self::S3(cache) => cache.default_ttl(), Self::Gcs(cache) => cache.default_ttl(), + Self::ValkeySemantic { cache, .. } => cache.default_ttl(), + Self::RedisSemantic { cache, .. } => cache.default_ttl(), Self::Disk(cache) => cache.default_ttl(), - Self::RedisSemantic(cache) => cache.default_ttl(), Self::AzureBlob(cache) => cache.default_ttl(), } } + pub fn bucket(&self) -> Option<&str> { + match self { + Self::S3(cache) => Some(cache.backend().bucket()), + _ => None, + } + } + + pub fn key_prefix(&self) -> Option<&str> { + match self { + Self::S3(cache) => Some(cache.backend().key_prefix()), + _ => None, + } + } + + pub fn region(&self) -> Option<&str> { + match self { + Self::S3(cache) => Some(cache.backend().region()), + _ => None, + } + } + + pub fn endpoint(&self) -> Option<&str> { + match self { + Self::S3(cache) => cache.backend().endpoint(), + _ => None, + } + } + pub fn namespace(&self) -> Option<&str> { match self { - Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::RedisSemantic(_) => None, + Self::Memory(_) + | Self::S3(_) + | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } + | Self::Disk(_) + | Self::AzureBlob(_) + | Self::Gcs(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), - Self::Gcs(_) => None, } } pub fn topology(&self) -> Option<&RedisTopology> { match self { Self::Memory(_) + | Self::S3(_) + | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::Disk(_) | Self::AzureBlob(_) - | Self::Gcs(_) - | Self::RedisSemantic(_) => None, + | Self::Gcs(_) => None, Self::Redis { cache, .. } => Some(cache.backend().topology()), } } @@ -158,10 +365,12 @@ impl NativeResponseCache { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), Self::Redis { .. } + | Self::S3(_) + | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::Disk(_) | Self::AzureBlob(_) - | Self::Gcs(_) - | Self::RedisSemantic(_) => None, + | Self::Gcs(_) => None, } } @@ -169,282 +378,454 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), Self::Redis { .. } + | Self::S3(_) + | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::Disk(_) | Self::AzureBlob(_) - | Self::Gcs(_) - | Self::RedisSemantic(_) => None, - } - } - - pub fn index_name(&self) -> Option<&str> { - match self { - Self::RedisSemantic(cache) => Some(cache.backend().index_name()), - Self::Memory(_) - | Self::Redis { .. } - | Self::AzureBlob(_) - | Self::Disk(_) | Self::Gcs(_) => None, } } - pub fn similarity_threshold(&self) -> Option { - match self { - Self::RedisSemantic(cache) => Some(cache.backend().similarity_threshold()), - Self::Memory(_) - | Self::Redis { .. } - | Self::AzureBlob(_) - | Self::Disk(_) - | Self::Gcs(_) => None, - } - } - - pub fn semantic_embedder(&self) -> Option<&PythonEmbedder> { - match self { - Self::RedisSemantic(cache) => Some(cache.backend().embedder()), - Self::Memory(_) - | Self::Redis { .. } - | Self::AzureBlob(_) - | Self::Disk(_) - | Self::Gcs(_) => None, - } - } - - pub fn embedder_object(&self) -> Option<&Py> { - match self { - Self::RedisSemantic(cache) => Some(cache.backend().embedder().object()), - Self::Memory(_) - | Self::Redis { .. } - | Self::AzureBlob(_) - | Self::Disk(_) - | Self::Gcs(_) => None, - } - } - - pub fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - if let Self::RedisSemantic(cache) = self { - cache.backend().embedder().traverse(visit)?; - } - Ok(()) - } - - pub fn with_redis_flush_size(self, flush_size: Option) -> Self { - match self { - Self::Redis { cache, .. } => Self::Redis { - cache, - buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))), - }, - other => other, - } - } - pub fn directory(&self) -> Option<&Path> { match self { Self::Disk(cache) => Some(cache.backend().directory()), Self::Memory(_) | Self::Redis { .. } + | Self::S3(_) + | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::AzureBlob(_) - | Self::Gcs(_) - | Self::RedisSemantic(_) => None, + | Self::Gcs(_) => None, } } - fn exact_requests(requests: &[CacheRequest]) -> Vec> { - requests.iter().map(CacheRequest::exact).collect() + pub fn semantic_config(&self) -> Option<(f64, &str)> { + match self { + Self::ValkeySemantic { cache, .. } => Some(( + cache.backend().similarity_threshold(), + cache.backend().index_name(), + )), + Self::RedisSemantic { cache, .. } => Some(( + f64::from(cache.backend().similarity_threshold()), + cache.backend().index_name(), + )), + _ => None, + } } - pub fn lookup(&self, request: &CacheRequest, now: Duration) -> Result, Error> { + pub fn index_name(&self) -> Option<&str> { match self { - Self::Memory(cache) => cache.lookup(&request.exact(), now), - Self::Redis { cache, .. } => cache.lookup(&request.exact(), now), - Self::Gcs(cache) => cache.lookup(&request.exact(), now), - Self::Disk(cache) => cache.lookup(&request.exact(), now), - Self::AzureBlob(cache) => cache.lookup(&request.exact(), now), - Self::RedisSemantic(cache) => cache.lookup(&request.semantic(), now), + Self::RedisSemantic { cache, .. } => Some(cache.backend().index_name()), + _ => None, + } + } + + pub fn similarity_threshold(&self) -> Option { + match self { + Self::RedisSemantic { cache, .. } => Some(cache.backend().similarity_threshold()), + _ => None, + } + } + + pub fn semantic_embedder(&self) -> Option<&PythonEmbedder> { + match self { + Self::RedisSemantic { embedder, .. } => Some(embedder), + _ => None, + } + } + + pub fn embedder_object(&self) -> Option<&Py> { + match self { + Self::RedisSemantic { embedder, .. } => Some(embedder.object()), + _ => None, + } + } + + pub fn lookup(&self, request: &NativeRequest, now: Duration) -> Result, Error> { + match self { + Self::Memory(cache) => cache.lookup(&Self::exact(request), now), + Self::Redis { cache, .. } => cache.lookup(&Self::exact(request), now), + Self::S3(cache) => cache.lookup(&Self::exact(request), now), + Self::ValkeySemantic { cache, scope, .. } => { + cache.lookup(&Self::semantic(request, scope), now) + } + Self::RedisSemantic { cache, .. } => { + cache.lookup(&Self::redis_semantic_request(request), now) + } + Self::Gcs(cache) => cache.lookup(&Self::exact(request), now), + Self::Disk(cache) => cache.lookup(&Self::exact(request), now), + Self::AzureBlob(cache) => cache.lookup(&Self::exact(request), now), } } pub fn store( &self, - request: &CacheRequest, + request: &NativeRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.store(&request.exact(), response, now), - Self::Redis { cache, .. } => cache.store(&request.exact(), response, now), - Self::Gcs(cache) => cache.store(&request.exact(), response, now), - Self::Disk(cache) => cache.store(&request.exact(), response, now), - Self::AzureBlob(cache) => cache.store(&request.exact(), response, now), - Self::RedisSemantic(cache) => cache.store(&request.semantic(), response, now), + Self::Memory(cache) => cache.store(&Self::exact(request), response, now), + Self::Redis { cache, .. } => cache.store(&Self::exact(request), response, now), + Self::S3(cache) => cache.store(&Self::exact(request), response, now), + Self::ValkeySemantic { cache, scope, .. } => { + cache.store(&Self::semantic(request, scope), response, now) + } + Self::RedisSemantic { cache, .. } => { + cache.store(&Self::redis_semantic_request(request), response, now) + } + Self::Gcs(cache) => cache.store(&Self::exact(request), response, now), + Self::Disk(cache) => cache.store(&Self::exact(request), response, now), + Self::AzureBlob(cache) => cache.store(&Self::exact(request), response, now), } } pub fn lookup_batch( &self, - requests: &[CacheRequest], + requests: &[NativeRequest], now: Duration, ) -> Result { match self { - Self::Memory(cache) => cache.lookup_batch(&Self::exact_requests(requests), now), - Self::Redis { cache, .. } => cache.lookup_batch(&Self::exact_requests(requests), now), - Self::Gcs(cache) => cache.lookup_batch(&Self::exact_requests(requests), now), - Self::Disk(cache) => cache.lookup_batch(&Self::exact_requests(requests), now), - Self::AzureBlob(cache) => cache.lookup_batch(&Self::exact_requests(requests), now), - Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), + Self::Memory(cache) => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.lookup_batch(&requests, now) + } + Self::Redis { cache, .. } => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.lookup_batch(&requests, now) + } + Self::S3(cache) => { + cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) + } + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + Err(Error::UnsupportedOperation) + } + Self::Gcs(cache) => { + cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) + } + Self::Disk(cache) => { + cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) + } + Self::AzureBlob(cache) => { + cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) + } } } pub async fn async_lookup( &self, - request: &CacheRequest, + request: &NativeRequest, now: Duration, ) -> Result, Error> { match self { - Self::Memory(cache) => cache.async_lookup(&request.exact(), now).await, - Self::Redis { cache, .. } => cache.async_lookup(&request.exact(), now).await, - Self::Gcs(cache) => cache.async_lookup(&request.exact(), now).await, - Self::Disk(cache) => cache.async_lookup(&request.exact(), now).await, - Self::AzureBlob(cache) => cache.async_lookup(&request.exact(), now).await, - Self::RedisSemantic(cache) => cache.async_lookup(&request.semantic(), now).await, + Self::Memory(cache) => cache.async_lookup(&Self::exact(request), now).await, + Self::Redis { cache, .. } => cache.async_lookup(&Self::exact(request), now).await, + Self::S3(cache) => cache.async_lookup(&Self::exact(request), now).await, + Self::ValkeySemantic { cache, scope, .. } => { + cache + .async_lookup(&Self::semantic(request, scope), now) + .await + } + Self::RedisSemantic { cache, .. } => { + cache + .async_lookup(&Self::redis_semantic_request(request), now) + .await + } + Self::Gcs(cache) => cache.async_lookup(&Self::exact(request), now).await, + Self::Disk(cache) => cache.async_lookup(&Self::exact(request), now).await, + Self::AzureBlob(cache) => cache.async_lookup(&Self::exact(request), now).await, + } + } + + pub(super) fn async_lookup_py<'py>( + &self, + py: Python<'py>, + request: NativeRequest, + ) -> PyResult> { + match self { + Self::Memory(_) + | Self::Redis { .. } + | Self::S3(_) + | Self::Disk(_) + | Self::AzureBlob(_) + | Self::Gcs(_) => { + let service = self.clone(); + litellm_host_python::run_async( + py, + async move { service.async_lookup(&request, super::request::now()).await }, + super::cache_error, + ) + } + Self::ValkeySemantic { + cache, + embedder, + scope, + } => drive_semantic( + py, + SemanticEmbedExecution::lookup( + Arc::clone(cache.backend_arc()), + embedder.clone(), + Self::semantic(&request, scope), + ), + ), + Self::RedisSemantic { .. } => drive( + py, + SemanticBody::new(self.clone(), SemanticOperation::Lookup(request)), + ), } } pub async fn async_store( &self, - request: &CacheRequest, + request: &NativeRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_store(&request.exact(), response, now).await, + Self::Memory(cache) => { + cache + .async_store(&Self::exact(request), response, now) + .await + } Self::Redis { cache, buffer: None, - } => cache.async_store(&request.exact(), response, now).await, + } => { + cache + .async_store(&Self::exact(request), response, now) + .await + } Self::Redis { cache, buffer: Some(buffer), } => { buffer - .async_store(cache, &request.exact(), response, now) + .async_store(cache, &Self::exact(request), response, now) .await } - Self::Gcs(cache) => cache.async_store(&request.exact(), response, now).await, - Self::Disk(cache) => cache.async_store(&request.exact(), response, now).await, - Self::AzureBlob(cache) => cache.async_store(&request.exact(), response, now).await, - Self::RedisSemantic(cache) => { - cache.async_store(&request.semantic(), response, now).await + Self::S3(cache) => { + cache + .async_store(&Self::exact(request), response, now) + .await } + Self::ValkeySemantic { cache, scope, .. } => { + cache + .async_store(&Self::semantic(request, scope), response, now) + .await + } + Self::RedisSemantic { cache, .. } => { + cache + .async_store(&Self::redis_semantic_request(request), response, now) + .await + } + Self::Gcs(cache) => { + cache + .async_store(&Self::exact(request), response, now) + .await + } + Self::Disk(cache) => { + cache + .async_store(&Self::exact(request), response, now) + .await + } + Self::AzureBlob(cache) => { + cache + .async_store(&Self::exact(request), response, now) + .await + } + } + } + + pub(super) fn async_store_py<'py>( + &self, + py: Python<'py>, + request: NativeRequest, + response: Value, + ) -> PyResult> { + match self { + Self::Memory(_) + | Self::Redis { .. } + | Self::S3(_) + | Self::Disk(_) + | Self::AzureBlob(_) + | Self::Gcs(_) => { + let service = self.clone(); + litellm_host_python::run_async( + py, + async move { + service + .async_store(&request, response, super::request::now()) + .await + }, + super::cache_error, + ) + } + Self::ValkeySemantic { + cache, + embedder, + scope, + } => drive_semantic( + py, + SemanticEmbedExecution::store( + Arc::clone(cache.backend_arc()), + embedder.clone(), + Self::semantic(&request, scope), + response, + ), + ), + Self::RedisSemantic { .. } => drive( + py, + SemanticBody::new(self.clone(), SemanticOperation::Store(request, response)), + ), } } pub async fn async_lookup_batch( &self, - requests: &[CacheRequest], + requests: &[NativeRequest], now: Duration, ) -> Result { match self { Self::Memory(cache) => { - cache - .async_lookup_batch(&Self::exact_requests(requests), now) - .await + let requests = requests.iter().map(Self::exact).collect::>(); + cache.async_lookup_batch(&requests, now).await } Self::Redis { cache, .. } => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.async_lookup_batch(&requests, now).await + } + Self::S3(cache) => { cache - .async_lookup_batch(&Self::exact_requests(requests), now) + .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) .await } + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + Err(Error::UnsupportedOperation) + } Self::Gcs(cache) => { cache - .async_lookup_batch(&Self::exact_requests(requests), now) + .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) .await } Self::Disk(cache) => { cache - .async_lookup_batch(&Self::exact_requests(requests), now) + .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) .await } Self::AzureBlob(cache) => { cache - .async_lookup_batch(&Self::exact_requests(requests), now) + .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) .await } - Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), } } pub async fn async_store_batch( &self, - entries: Vec<(CacheRequest, Value)>, + entries: Vec<(NativeRequest, Value)>, now: Duration, ) -> Result<(), Error> { match self { Self::Memory(cache) => { - cache - .async_store_batch( - entries - .into_iter() - .map(|(request, value)| (request.exact(), value)) - .collect(), - now, - ) - .await + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await } Self::Redis { cache, .. } => { - cache - .async_store_batch( - entries - .into_iter() - .map(|(request, value)| (request.exact(), value)) - .collect(), - now, - ) - .await + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await } + Self::S3(cache) => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + Self::ValkeySemantic { cache, scope, .. } => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::semantic(&request, scope), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + Self::RedisSemantic { .. } => Err(Error::UnsupportedOperation), Self::Gcs(cache) => { - cache - .async_store_batch( - entries - .into_iter() - .map(|(request, value)| (request.exact(), value)) - .collect(), - now, - ) - .await + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await } Self::Disk(cache) => { - cache - .async_store_batch( - entries - .into_iter() - .map(|(request, value)| (request.exact(), value)) - .collect(), - now, - ) - .await + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await } Self::AzureBlob(cache) => { - cache - .async_store_batch( - entries - .into_iter() - .map(|(request, value)| (request.exact(), value)) - .collect(), - now, - ) - .await + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await } - Self::RedisSemantic(cache) => { - cache - .async_store_batch( - entries - .into_iter() - .map(|(request, value)| (request.semantic(), value)) - .collect(), - now, - ) - .await + } + } + + pub(super) fn async_store_batch_py<'py>( + &self, + py: Python<'py>, + entries: Vec<(NativeRequest, Value)>, + ) -> PyResult> { + match self { + Self::Memory(_) + | Self::Redis { .. } + | Self::S3(_) + | Self::Disk(_) + | Self::AzureBlob(_) + | Self::Gcs(_) => { + let service = self.clone(); + litellm_host_python::run_async( + py, + async move { + service + .async_store_batch(entries, super::request::now()) + .await + }, + super::cache_error, + ) } + Self::ValkeySemantic { + cache, + embedder, + scope, + } => { + let (requests, responses): (Vec<_>, Vec<_>) = entries + .into_iter() + .map(|(request, response)| (Self::semantic(&request, scope), response)) + .unzip(); + drive_semantic( + py, + SemanticEmbedExecution::store_batch( + Arc::clone(cache.backend_arc()), + embedder.clone(), + requests, + responses, + ), + ) + } + Self::RedisSemantic { .. } => drive( + py, + SemanticBody::new(self.clone(), SemanticOperation::StoreBatch(entries.into())), + ), } } @@ -457,9 +838,12 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::S3(cache) => cache.async_flush().await, + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + Err(Error::UnsupportedOperation) + } Self::Gcs(cache) => cache.async_flush().await, Self::Disk(cache) => cache.async_flush().await, - Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), Self::AzureBlob(cache) => cache.async_flush().await, } } @@ -468,13 +852,24 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::S3(cache) => cache.test_connection().await, + Self::ValkeySemantic { cache, .. } => cache.test_connection().await, + Self::RedisSemantic { .. } => Err(Error::UnsupportedOperation), Self::Gcs(cache) => cache.test_connection().await, Self::Disk(cache) => cache.test_connection().await, - Self::RedisSemantic(_) => Err(Error::UnsupportedOperation), Self::AzureBlob(cache) => cache.test_connection().await, } } + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + match self { + Self::ValkeySemantic { embedder, .. } => embedder.traverse(visit)?, + Self::RedisSemantic { embedder, .. } => embedder.traverse(visit)?, + _ => {} + } + Ok(()) + } + pub fn gcs_backend(&self) -> Option<&GcsCache> { match self { Self::Gcs(cache) => Some(cache.backend()), @@ -482,3 +877,80 @@ impl NativeResponseCache { } } } + +#[cfg(test)] +mod tests { + use litellm_cache_response::{CacheControls, CacheKeyInput, cache_key}; + use serde_json::json; + use sha2::{Digest, Sha256}; + + use super::*; + + fn native_request(key: CacheKeyInput, metadata: Value) -> NativeRequest { + NativeRequest { + key, + controls: CacheControls::default(), + ttl: None, + max_age: None, + messages: Some(json!([{"role": "user", "content": "prompt"}])), + input: None, + metadata: Some(metadata), + litellm_metadata: None, + litellm_params: None, + scope: None, + } + } + + #[test] + fn semantic_key_matches_python_scope_material() { + let key = CacheKeyInput { + fields: vec![ + CacheKeyField { + name: "model".to_owned(), + value: Some("gpt-4.1".to_owned()), + api_parameter: true, + internal_parameter: false, + }, + CacheKeyField { + name: "messages".to_owned(), + value: Some("prompt".to_owned()), + api_parameter: true, + internal_parameter: false, + }, + ], + ..Default::default() + }; + let request = native_request( + key, + json!({"user_api_key": "k1", "user_api_key_team_id": null}), + ); + let expected = format!("{:x}", Sha256::digest(b"model: gpt-4.1user_api_key: k1")); + assert_eq!(cache_key(&semantic_key(&request, "key")), expected); + + let end_user_request = native_request( + request.key.clone(), + json!({"user_api_key": "k1", "user_api_key_end_user_id": "u1"}), + ); + let expected = format!( + "{:x}", + Sha256::digest(b"model: gpt-4.1user_api_key: k1user_api_key_end_user_id: u1") + ); + assert_eq!( + cache_key(&semantic_key(&end_user_request, "end_user")), + expected + ); + + let preset_request = native_request( + CacheKeyInput { + preset: Some("preset-key".to_owned()), + ..Default::default() + }, + json!({"user_api_key": "k1"}), + ); + assert_eq!( + semantic_key(&preset_request, "end_user").preset.as_deref(), + Some("preset-key") + ); + assert!(semantic_key(&preset_request, "end_user").fields.is_empty()); + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index b06087bcc83..3b4b910c1f0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -1,11 +1,11 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::{ExactCacheContext, SemanticCacheContext}; +use litellm_cache::ExactCacheContext; use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; use serde::Deserialize; -use serde_json::{Map, Value}; +use serde_json::Value; #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -14,69 +14,52 @@ struct RequestInput { controls: Option, ttl_seconds: Option, max_age_seconds: Option, + messages: Option, input: Option, - messages: Option>, - metadata: Option>, + metadata: Option, + litellm_metadata: Option, + litellm_params: Option, scope: Option, } #[derive(Clone)] -pub(super) struct CacheRequest { - key: CacheKeyInput, - controls: CacheControls, - ttl: Option, - max_age: Option, - input: Option, - messages: Vec, - metadata: Map, - scope: Option, +pub(super) struct NativeRequest { + pub(super) key: CacheKeyInput, + pub(super) controls: CacheControls, + pub(super) ttl: Option, + pub(super) max_age: Option, + pub(super) messages: Option, + pub(super) input: Option, + pub(super) metadata: Option, + pub(super) litellm_metadata: Option, + pub(super) litellm_params: Option, + pub(super) scope: Option, } -impl CacheRequest { - pub(super) fn exact(&self) -> ResponseCacheRequest { - let mut request = ResponseCacheRequest::new(self.key.clone()); - request.controls = self.controls; - request.context.ttl = self.ttl; - request.max_age = self.max_age; - request - } - - pub(super) fn semantic(&self) -> ResponseCacheRequest { - ResponseCacheRequest { - key: self.key.clone(), - controls: self.controls, - context: SemanticCacheContext { - input: self.input.clone(), - messages: self.messages.clone(), - metadata: self.metadata.clone(), - scope: self.scope.clone(), - ttl: self.ttl, - }, - max_age: self.max_age, - } - } -} - -pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { +pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { let input: RequestInput = from_py(value)?; request_input(input) } -fn request_input(input: RequestInput) -> PyResult { - let defaults = ResponseCacheRequest::::new(input.key.clone()); - Ok(CacheRequest { +fn request_input(input: RequestInput) -> PyResult { + let controls = input.controls.unwrap_or_else(|| { + ResponseCacheRequest::::new(input.key.clone()).controls + }); + Ok(NativeRequest { key: input.key, - controls: input.controls.unwrap_or(defaults.controls), + controls, ttl: input.ttl_seconds.map(duration).transpose()?, max_age: input.max_age_seconds.map(duration).transpose()?, + messages: input.messages, input: input.input, - messages: input.messages.unwrap_or_default(), - metadata: input.metadata.unwrap_or_default(), + metadata: input.metadata, + litellm_metadata: input.litellm_metadata, + litellm_params: input.litellm_params, scope: input.scope, }) } -pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { +pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { from_py::>(value)? .into_iter() .map(request_input) diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/semantic.rs index f0f75de1edf..7661598c5f9 100644 --- a/litellm-rust/crates/python-bridge/src/cache/semantic.rs +++ b/litellm-rust/crates/python-bridge/src/cache/semantic.rs @@ -14,13 +14,13 @@ use super::{ cache_error, embedder::{PythonEmbedder, with_prepared_embedding}, native::NativeResponseCache, - request::{CacheRequest, now}, + request::{NativeRequest, now}, }; pub(super) enum SemanticOperation { - Lookup(CacheRequest), - Store(CacheRequest, Value), - StoreBatch(VecDeque<(CacheRequest, Value)>), + Lookup(NativeRequest), + Store(NativeRequest, Value), + StoreBatch(VecDeque<(NativeRequest, Value)>), } enum Phase { @@ -32,7 +32,7 @@ enum Phase { pub(super) struct SemanticBody { service: NativeResponseCache, operation: SemanticOperation, - pending: Option<(CacheRequest, Option)>, + pending: Option<(NativeRequest, Option)>, phase: Phase, } @@ -101,7 +101,7 @@ impl ExecutionBody for SemanticBody { let (request, _) = self.pending.as_ref().ok_or_else(|| { PyRuntimeError::new_err("semantic execution has no pending operation") })?; - let semantic = request.semantic(); + let semantic = NativeResponseCache::redis_semantic_request(request); let Some(prompt) = prompt_from_context(&semantic.context) else { return self.backend_step(py, Err(Error::Unavailable)); }; @@ -113,7 +113,7 @@ impl ExecutionBody for SemanticBody { let coroutine = embedder.async_embedding_coroutine( py, &prompt, - &semantic.context.metadata, + semantic.context.metadata.as_ref(), )?; self.phase = Phase::AwaitingEmbedding; return Ok(ExecutionStep::Await(coroutine)); @@ -160,16 +160,15 @@ impl ExecutionBody for SemanticBody { } fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { - self.service.traverse(visit) + if let Some(embedder) = self.service.semantic_embedder() { + embedder.traverse(visit)?; + } + Ok(()) } } -pub(super) fn drive( - py: Python<'_>, - service: NativeResponseCache, - operation: SemanticOperation, -) -> PyResult> { - let execution = Py::new(py, Execution::new(SemanticBody::new(service, operation)))?; +pub(super) fn drive(py: Python<'_>, body: SemanticBody) -> PyResult> { + let execution = Py::new(py, Execution::new(body))?; py.import("litellm.rust_bridge.lifecycle")? .getattr("drive")? .call1((execution,)) diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs new file mode 100644 index 00000000000..24caf3374d6 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/semantic_step.rs @@ -0,0 +1,249 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_cache::SemanticCacheContext; +use litellm_cache_response::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest}; +use litellm_cache_valkey_semantic::{PreparedEmbedding, ValkeySemanticCache, prompt_from_context}; +use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; +use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use serde_json::Value; + +use super::{cache_error, embedder::PythonEmbedder}; + +pub(super) enum Op { + Lookup, + Store(Value), + StoreBatch(Vec), +} + +#[derive(Clone, Copy)] +enum State { + Start, + AwaitingEmbedding, + AwaitingStorage, + Done, +} + +pub(super) struct SemanticEmbedExecution { + backend: Arc>, + embedder: PythonEmbedder, + requests: Vec>, + op: Op, + now: Option, + prepared: Vec>>, + index: usize, + state: State, +} + +impl SemanticEmbedExecution { + pub(super) fn lookup( + backend: Arc>, + embedder: PythonEmbedder, + request: ResponseCacheRequest, + ) -> Self { + Self { + backend, + embedder, + requests: vec![request], + op: Op::Lookup, + now: None, + prepared: vec![None], + index: 0, + state: State::Start, + } + } + + pub(super) fn store( + backend: Arc>, + embedder: PythonEmbedder, + request: ResponseCacheRequest, + response: Value, + ) -> Self { + Self { + backend, + embedder, + requests: vec![request], + op: Op::Store(response), + now: None, + prepared: vec![None], + index: 0, + state: State::Start, + } + } + + pub(super) fn store_batch( + backend: Arc>, + embedder: PythonEmbedder, + requests: Vec>, + responses: Vec, + ) -> Self { + Self { + backend, + embedder, + prepared: vec![None; requests.len()], + requests, + op: Op::StoreBatch(responses), + now: None, + index: 0, + state: State::Start, + } + } + + fn start(&mut self, py: Python<'_>) -> PyResult { + if self.now.is_none() { + self.now = Some(super::request::now()); + } + while self.index < self.requests.len() { + let request = &self.requests[self.index]; + let enabled = match &self.op { + Op::Lookup => request.controls.reads(), + Op::Store(_) | Op::StoreBatch(_) => request.controls.writes(), + }; + if !enabled { + self.index += 1; + continue; + } + let Some(prompt) = prompt_from_context(&request.context) else { + self.index += 1; + continue; + }; + let metadata = request.context.metadata.clone(); + let awaitable = self + .embedder + .async_embed_awaitable(py, &prompt, &metadata)?; + self.state = State::AwaitingEmbedding; + return Ok(ExecutionStep::Await(awaitable.unbind())); + } + self.state = State::AwaitingStorage; + self.storage_step(py) + } + + fn storage_step(&self, py: Python<'_>) -> PyResult { + let requests = self.requests.clone(); + let prepared = self.prepared.clone(); + let backend = Arc::clone(&self.backend); + let now = self + .now + .ok_or_else(|| PyRuntimeError::new_err("semantic cache timestamp is unavailable"))?; + let awaitable = match &self.op { + Op::Lookup => { + let Some(request) = requests.into_iter().next() else { + return Err(PyRuntimeError::new_err( + "semantic lookup requires one request", + )); + }; + match prepared.into_iter().next().flatten() { + Some(values) => { + let backend = backend.with_embedder(PreparedEmbedding(values)); + let cache = Arc::new(ResponseCache::new(Arc::new(backend))); + run_async( + py, + async move { cache.async_lookup(&request, now).await }, + cache_error, + )? + } + None => { + let cache = Arc::new(ResponseCache::new(backend)); + run_async( + py, + async move { cache.async_lookup(&request, now).await }, + cache_error, + )? + } + } + } + Op::Store(response) => { + let Some(request) = requests.into_iter().next() else { + return Err(PyRuntimeError::new_err( + "semantic store requires one request", + )); + }; + let response = response.clone(); + match prepared.into_iter().next().flatten() { + Some(values) => { + let backend = backend.with_embedder(PreparedEmbedding(values)); + let cache = Arc::new(ResponseCache::new(Arc::new(backend))); + run_async( + py, + async move { cache.async_store(&request, response, now).await }, + cache_error, + )? + } + None => { + let cache = Arc::new(ResponseCache::new(backend)); + run_async( + py, + async move { cache.async_store(&request, response, now).await }, + cache_error, + )? + } + } + } + Op::StoreBatch(responses) => { + let responses = responses.clone(); + run_async( + py, + async move { + for ((request, response), prepared) in + requests.into_iter().zip(responses).zip(prepared) + { + let Some(values) = prepared else { + continue; + }; + let backend = backend.with_embedder(PreparedEmbedding(values)); + let cache = ResponseCache::new(Arc::new(backend)); + cache.async_store(&request, response, now).await?; + } + Ok(()) + }, + cache_error, + )? + } + }; + Ok(ExecutionStep::Await(awaitable.unbind())) + } + + fn resume_py( + &mut self, + py: Python<'_>, + result: Option>>, + ) -> PyResult { + match (self.state, result) { + (State::Start, None) => self.start(py), + (State::AwaitingEmbedding, Some(Ok(value))) => { + let values = value.bind(py).extract::>()?; + self.prepared[self.index] = + Some(values.into_iter().map(|value| value as f32).collect()); + self.index += 1; + self.start(py) + } + (State::AwaitingStorage, Some(Ok(value))) => { + self.state = State::Done; + Ok(ExecutionStep::Return(value)) + } + (_, Some(Err(error))) => Err(error), + _ => Err(PyRuntimeError::new_err( + "invalid semantic cache execution state", + )), + } + } +} + +impl ExecutionBody for SemanticEmbedExecution { + fn resume(&mut self, result: Option>>) -> PyResult { + Python::attach(|py| self.resume_py(py, result)) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + self.embedder.traverse(visit) + } +} + +pub(super) fn drive_semantic<'py>( + py: Python<'py>, + body: SemanticEmbedExecution, +) -> PyResult> { + let execution = Py::new(py, Execution::new(body))?; + py.import("litellm.rust_bridge.lifecycle")? + .getattr("drive")? + .call1((execution,)) +} diff --git a/litellm/litellm_core_utils/logging_worker.py b/litellm/litellm_core_utils/logging_worker.py index 5ccc5632646..2f8e7bdccea 100644 --- a/litellm/litellm_core_utils/logging_worker.py +++ b/litellm/litellm_core_utils/logging_worker.py @@ -484,9 +484,13 @@ class LoggingWorker: so it correctly handles items that have been dequeued but whose callback hasn't finished yet — ``queue.empty()`` would return True in that window and cause us to skip the wait. + + ``start()`` runs first so a queue left behind by a previous event loop + is carried onto this one and drained here instead of joined forever. """ if self._queue is None: return + self.start() await self._queue.join() async def clear_queue(self): diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index beb90c2fb8c..5cf5ffc102e 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -43037,21 +43037,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.92272e-07, + "input_cost_per_token": 8.87226e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.784544e-06, + "output_cost_per_token": 1.774452e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.4356e-08, + "cache_read_input_token_cost": 7.39355e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -73794,6 +73794,7 @@ "supports_web_search": false }, "openrouter/bytedance-seed/seed-1.6": { + "deprecation_date": "2026-11-11", "input_cost_per_token": 2.5e-07, "input_cost_per_token_above_128k_tokens": 5e-07, "litellm_provider": "openrouter", @@ -73815,6 +73816,7 @@ "supports_web_search": false }, "openrouter/bytedance-seed/seed-1.6-flash": { + "deprecation_date": "2026-11-11", "input_cost_per_token": 7.5e-08, "input_cost_per_token_above_128k_tokens": 1e-07, "litellm_provider": "openrouter", @@ -73855,6 +73857,7 @@ "supports_web_search": false }, "openrouter/bytedance-seed/seed-2.0-code": { + "deprecation_date": "2026-11-11", "input_cost_per_token": 5e-07, "input_cost_per_token_above_128k_tokens": 1e-06, "litellm_provider": "openrouter", @@ -76876,6 +76879,26 @@ "supports_vision": true, "supports_web_search": true }, + "moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "global.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, @@ -76975,5 +76998,51 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false + }, + "xiaomi_mimo/mimo-v2.6-pro": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 4.35e-07, + "litellm_provider": "xiaomi_mimo", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8.7e-07, + "source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "xiaomi_mimo/mimo-v2.6-flash": { + "cache_read_input_token_cost": 2.8e-09, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "xiaomi_mimo", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true } } diff --git a/litellm/proxy/auth/handle_jwt.py b/litellm/proxy/auth/handle_jwt.py index 996911cdfaa..010a1b4536e 100644 --- a/litellm/proxy/auth/handle_jwt.py +++ b/litellm/proxy/auth/handle_jwt.py @@ -14,7 +14,7 @@ import hashlib import os import re import time -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Collection, Mapping, Sequence from dataclasses import dataclass from typing import Any, Final, Literal, NoReturn, Protocol, TypeVar, cast @@ -1602,6 +1602,7 @@ class JWTAuthManager: team_object: LiteLLM_TeamTable | None, route: str, request_method: str | None = None, + team_allowed_routes: Collection[str] = (), ) -> bool: normalized_request_method: Final = request_method.upper() if isinstance(request_method, str) else None if not RouteChecks.is_auth_enforced_pass_through_route( @@ -1610,8 +1611,11 @@ class JWTAuthManager: ): return True + if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=team_allowed_routes): + return True + # JWT team selection is team-scoped; key metadata is not available here, - # so passthrough access is granted only by the selected team's metadata. + # so beyond the JWT config grant above, only the selected team's metadata grants access. return RouteChecks.check_passthrough_route_access( route=route, user_api_key_dict=UserAPIKeyAuth(team_metadata=(team_object.metadata or {}) if team_object else {}), @@ -1689,6 +1693,7 @@ class JWTAuthManager: team_object=team_object, route=route, request_method=request_method, + team_allowed_routes=jwt_handler.litellm_jwtauth.team_allowed_routes, ): is_allowed = False denied_auth_enforced_pass_through_route = True @@ -2584,6 +2589,7 @@ class JWTAuthManager: team_object=team_object, route=route, request_method=request_method, + team_allowed_routes=handler.litellm_jwtauth.team_allowed_routes, ): JWTAuthManager._raise_team_passthrough_route_denial(route=route) @@ -2653,6 +2659,7 @@ class JWTAuthManager: team_object=team_object, route=route, request_method=request_method, + team_allowed_routes=handler.litellm_jwtauth.team_allowed_routes, ): JWTAuthManager._raise_team_passthrough_route_denial(route=route) elif team_id is None: diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index afddb4866a9..38189a2d07b 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -278,7 +278,11 @@ class RouteChecks: route=route, method=RouteChecks._get_request_method(request=request), ): - RouteChecks._require_auth_pass_through_access(route=route, valid_token=valid_token) + RouteChecks._require_auth_pass_through_access( + route=route, + valid_token=valid_token, + jwt_team_allowed_routes=RouteChecks._jwt_team_allowed_routes(valid_token=valid_token), + ) elif RouteChecks.is_llm_api_route(route=route): pass elif RouteChecks.is_info_route(route=route): @@ -689,16 +693,43 @@ class RouteChecks: ), ) + @staticmethod + def jwt_team_routes_grant_pass_through(route: str, team_allowed_routes: Collection[str]) -> bool: + """ + Explicit paths and trailing-wildcard prefixes grant auth=true pass-through. Blanket grants never do: + a named route group like ``openai_routes`` is only ever compared as a path, and an entry that names + no path segment (``*``, ``/*``) is skipped. + """ + return any( + RouteChecks.route_matches_wildcard_pattern(route=route, pattern=allowed_route) + for allowed_route in team_allowed_routes + if allowed_route.rstrip("*").strip("/") + ) + + @staticmethod + def _jwt_team_allowed_routes(valid_token: UserAPIKeyAuth) -> Collection[str]: + """``team_allowed_routes`` for team tokens built by JWT auth; JWT-mapped virtual keys stay key-scoped.""" + if valid_token.jwt_claims is None or valid_token.token is not None or valid_token.team_id is None: + return () + + from litellm.proxy.proxy_server import jwt_handler + + return jwt_handler.litellm_jwtauth.team_allowed_routes + @staticmethod def _require_auth_pass_through_access( route: str, valid_token: UserAPIKeyAuth, + jwt_team_allowed_routes: Collection[str] = (), ) -> None: """ - Require an explicit ``allowed_passthrough_routes`` match for auth=true pass-through. + Require an explicit grant for auth=true pass-through: ``allowed_passthrough_routes`` on the + key or team, or an explicit JWT ``team_allowed_routes`` entry. """ if RouteChecks.check_passthrough_route_access(route=route, user_api_key_dict=valid_token): return + if RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=jwt_team_allowed_routes): + return raise RouteChecks._auth_pass_through_denied_exception(route=route) @staticmethod diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index c1388e8bb81..ee0e4db291b 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -1078,6 +1078,9 @@ if MCP_AVAILABLE: return {"servers": registry_servers} ## FastAPI Routes + def _mcp_server_display_order(server: LiteLLM_MCPServerTable) -> tuple[str, str]: + return ((server.server_name or server.alias or server.server_id).lower(), server.server_id) + def _get_user_mcp_management_mode() -> UserMCPManagementMode: from litellm.proxy.proxy_server import ( general_settings as proxy_general_settings, @@ -1228,10 +1231,12 @@ if MCP_AVAILABLE: detail="You do not have permission to view MCP servers for this team.", ) - redacted_mcp_servers = await _get_team_scoped_mcp_server_list(sanitized_team_id) + redacted_mcp_servers = sorted( + await _get_team_scoped_mcp_server_list(sanitized_team_id), key=_mcp_server_display_order + ) else: servers: Final = await _resolve_accessible_mcp_servers(user_api_key_dict) - redacted_mcp_servers = _redact_mcp_credentials_list(servers) + redacted_mcp_servers = sorted(_redact_mcp_credentials_list(servers), key=_mcp_server_display_order) if connected_app_view is True and is_ui_session_credential(user_api_key_dict): reachable_ids: Final = await _connected_app_reachable_server_ids(user_api_key_dict) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index f823dcfb135..c0bb93daadc 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -1,6 +1,6 @@ from asyncio import Future from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence -from typing import Literal, Never, final +from typing import Never, final from litellm.llms.base_llm.ocr.transformation import OCRResponse from litellm.rust_bridge.messages.entrypoints import LiteLLMMessagesRequest @@ -93,6 +93,59 @@ class ResponsesWebSocketConnection: def recv_text(self) -> Future[str | None]: ... def close(self) -> Future[None]: ... +@final +class _CacheTestBinding: + @property + def kind(self) -> str: ... + def lookup( + self, + request: object, + *, + callback_kwargs: Mapping[str, object] | Sequence[object] | None = None, + ) -> object: ... + def store( + self, + request: object, + response: object, + *, + callback_kwargs: Mapping[str, object] | None = None, + ) -> None: ... + def lookup_batch( + self, + requests: Sequence[object], + *, + callback_kwargs: Sequence[object] | None = None, + ) -> object: ... + def async_lookup( + self, + request: object, + *, + callback_kwargs: Mapping[str, object] | None = None, + ) -> Future[object]: ... + def async_store( + self, + request: object, + response: object, + *, + callback_kwargs: Mapping[str, object] | None = None, + ) -> Future[None]: ... + def async_lookup_batch( + self, + requests: Sequence[object], + *, + callback_kwargs: Sequence[object] | None = None, + ) -> Future[object]: ... + def async_store_batch( + self, + requests: Sequence[object], + responses: Sequence[object], + *, + callback_result: object = None, + callback_kwargs: Mapping[str, object] | None = None, + ) -> Future[object]: ... + def async_flush(self) -> Future[None]: ... + def ping(self) -> Future[object]: ... + @final class _CacheTestHandle: def __new__(cls, _uninstantiable: Never, /) -> Never: ... @@ -109,11 +162,22 @@ class _CacheTestHandle: *, ttl_seconds: float = 60.0, namespace: str | None = None, - startup_nodes: list[tuple[str, int]] | None = None, + startup_nodes: Sequence[tuple[str, int]] | None = None, ) -> _CacheTestHandle: ... @staticmethod def disk(directory: str) -> _CacheTestHandle: ... @staticmethod + def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ... + @staticmethod + def redis_semantic(backend: object) -> _CacheTestHandle: ... + @staticmethod + def valkey_semantic( + url: str, + similarity_threshold: float, + index_name: str, + embedder: object, + ) -> _CacheTestHandle: ... + @staticmethod def gcs( bucket_name: str, *, @@ -123,48 +187,19 @@ class _CacheTestHandle: token: str | None = None, ) -> _CacheTestHandle: ... @staticmethod - def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ... - @staticmethod - def redis_semantic(backend: object) -> _CacheTestHandle: ... - @property - def backend( - self, - ) -> Literal["memory", "redis", "gcs", "disk", "azure-blob", "redis_semantic"]: ... - def _bind_facade(self, facade: object) -> None: ... - -@final -class _CacheTestBinding: - def __new__(cls, _uninstantiable: Never, /) -> Never: ... - @property - def kind(self) -> Literal["disabled", "native", "python_callback"]: ... - def lookup( - self, request: object, *, callback_kwargs: object = None - ) -> object: ... - def store( - self, request: object, response: object, *, callback_kwargs: object = None - ) -> None: ... - def lookup_batch( - self, requests: object, *, callback_kwargs: object = None - ) -> object: ... - def async_lookup( - self, request: object, *, callback_kwargs: object = None - ) -> Future[object]: ... - def async_store( - self, request: object, response: object, *, callback_kwargs: object = None - ) -> Future[object]: ... - def async_lookup_batch( - self, requests: object, *, callback_kwargs: object = None - ) -> Future[object]: ... - def async_store_batch( - self, - requests: object, - responses: object, + def s3( + bucket: str, *, - callback_result: object = None, - callback_kwargs: object = None, - ) -> Future[object]: ... - def async_flush(self) -> Future[object]: ... - def ping(self) -> Future[object]: ... + region: str, + endpoint_url: str | None = None, + key_prefix: str = "", + access_key_id: str | None = None, + secret_access_key: str | None = None, + session_token: str | None = None, + ) -> _CacheTestHandle: ... + @property + def backend(self) -> str: ... + def _bind_facade(self, facade: object) -> None: ... @final class _CacheTestResolver: diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index beb90c2fb8c..5cf5ffc102e 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -43037,21 +43037,21 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro": { - "input_cost_per_token": 8.92272e-07, + "input_cost_per_token": 8.87226e-07, "input_cost_per_token_cache_hit": 4.4e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 1.784544e-06, + "output_cost_per_token": 1.774452e-06, "source": "https://openrouter.ai/api/v1/models", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_tool_choice": true, - "cache_read_input_token_cost": 7.4356e-08, + "cache_read_input_token_cost": 7.39355e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -73794,6 +73794,7 @@ "supports_web_search": false }, "openrouter/bytedance-seed/seed-1.6": { + "deprecation_date": "2026-11-11", "input_cost_per_token": 2.5e-07, "input_cost_per_token_above_128k_tokens": 5e-07, "litellm_provider": "openrouter", @@ -73815,6 +73816,7 @@ "supports_web_search": false }, "openrouter/bytedance-seed/seed-1.6-flash": { + "deprecation_date": "2026-11-11", "input_cost_per_token": 7.5e-08, "input_cost_per_token_above_128k_tokens": 1e-07, "litellm_provider": "openrouter", @@ -73855,6 +73857,7 @@ "supports_web_search": false }, "openrouter/bytedance-seed/seed-2.0-code": { + "deprecation_date": "2026-11-11", "input_cost_per_token": 5e-07, "input_cost_per_token_above_128k_tokens": 1e-06, "litellm_provider": "openrouter", @@ -76876,6 +76879,26 @@ "supports_vision": true, "supports_web_search": true }, + "moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "global.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, @@ -76975,5 +76998,51 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": false + }, + "xiaomi_mimo/mimo-v2.6-pro": { + "cache_read_input_token_cost": 3.6e-09, + "input_cost_per_token": 4.35e-07, + "litellm_provider": "xiaomi_mimo", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 8.7e-07, + "source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true + }, + "xiaomi_mimo/mimo-v2.6-flash": { + "cache_read_input_token_cost": 2.8e-09, + "input_cost_per_token": 1.4e-07, + "litellm_provider": "xiaomi_mimo", + "max_input_tokens": 1048576, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 2.8e-07, + "source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md", + "supported_endpoints": [ + "/v1/chat/completions" + ], + "supports_audio_input": true, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_video_input": true, + "supports_vision": true } } diff --git a/tests/e2e/coverage_registry/llm_conversational.yaml b/tests/e2e/coverage_registry/llm_conversational.yaml index 49d4d92ff0b..64335aa560c 100644 --- a/tests/e2e/coverage_registry/llm_conversational.yaml +++ b/tests/e2e/coverage_registry/llm_conversational.yaml @@ -89,6 +89,9 @@ - {id: llm.chat_completions.together_ai.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool result round trip"} - {id: llm.chat_completions.together_ai.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together cost header and spend row match the registry price"} - {id: llm.chat_completions.together_ai.thinking.nonstream.effort_none_disables, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [effort_none_disables], source: "llm_translation/test_together_ai_e2e.py", rationale: "reasoning_effort=none maps to Together's reasoning disable toggle on hybrid models"} +- {id: llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "Native MiMo v2.6 rows price the cost header and spend row from the cost map"} +- {id: llm.chat_completions.xiaomi_mimo.thinking.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: thinking, streaming: stream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo reasoning deltas stream as reasoning_content"} +- {id: llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo tool calls are not dropped"} - {id: llm.chat_completions.together_ai.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: structured_output, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "response_format json_schema reaches Together and constrains the reply"} - {id: llm.chat_completions.together_ai.prompt_cache_5m.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: prompt_cache_5m, streaming: nonstream, assertions: [cache_hit, cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together prefix-cache reads bill at cache_read_input_token_cost, not full input price"} - {id: llm.messages.together_ai.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together over /v1/messages streaming"} diff --git a/tests/e2e/coverage_registry/schema.py b/tests/e2e/coverage_registry/schema.py index 8b0d38a083c..d9c20d5c588 100644 --- a/tests/e2e/coverage_registry/schema.py +++ b/tests/e2e/coverage_registry/schema.py @@ -58,6 +58,7 @@ LlmRoute = Literal[ "openai", "together_ai", "vertex", + "xiaomi_mimo", ] LlmCapability = Literal[ diff --git a/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py b/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py new file mode 100644 index 00000000000..efca216634b --- /dev/null +++ b/tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py @@ -0,0 +1,214 @@ +"""Live e2e: Xiaomi MiMo v2.6 through the gateway on /chat/completions. + +Both native ``xiaomi_mimo/`` v2.6 rows (pro and flash) are registered via +``/model/new`` and driven against Xiaomi's own endpoint. What the gateway owes +us is that the reasoning chain surfaces as ``reasoning_content``, tool calls +survive translation, and the cost header plus spend row follow the proxy's own +cost-map price for the row (read back from ``/model/info``, never pinned here). +Requires XIAOMI_MIMO_API_KEY on the proxy; no skip gate. +""" + +from __future__ import annotations + +from typing import Final + +import pytest +from e2e_config import unique_marker +from e2e_http import StreamingResponse, require_successful_call, unwrap +from lifecycle import ResourceManager +from models import ( + ChatBody, + ChatMessage, + ChatResponse, + ChatTool, + ChatToolFunction, + CostMapEntry, + LiteLLMParamsBody, + OutMessage, + SpendLogRow, +) +from passthrough_client import PassthroughClient +from pydantic import BaseModel + +pytestmark = pytest.mark.e2e + +BACKENDS: Final = ("xiaomi_mimo/mimo-v2.6-pro", "xiaomi_mimo/mimo-v2.6-flash") +ARITHMETIC_PROMPT = "What is 17 + 26? Answer with just the number." +WEATHER_PROMPT = "What is the weather in Paris? Use the tool." +COUNTING_PROMPT = "Count from 1 to 50, one number per line." + +WEATHER_TOOL = ChatTool( + function=ChatToolFunction( + name="get_weather", + description="Get the current weather for a location.", + parameters={ + "type": "object", + "properties": {"location": {"type": "string"}}, + "required": ["location"], + }, + ) +) + + +class _WeatherArgs(BaseModel): + location: str + + +class _StreamDelta(BaseModel): + content: str | None = None + reasoning_content: str | None = None + + +class _StreamChoice(BaseModel): + delta: _StreamDelta | None = None + + +class _StreamChunk(BaseModel): + choices: list[_StreamChoice] = [] + + +def _approx_equal(actual: float, expected: float) -> bool: + return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2) + + +@pytest.fixture(scope="module") +def registry(client: PassthroughClient) -> dict[str, CostMapEntry]: + return client.proxy.model_cost_map() + + +def _register(client: PassthroughClient, resources: ResourceManager, backend: str) -> tuple[str, str]: + model = f"e2e-xiaomi-{unique_marker()}" + model_id = client.proxy.create_model( + model, LiteLLMParamsBody(model=backend, api_key="os.environ/XIAOMI_MIMO_API_KEY") + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model, resources.key() + + +def _message(response: ChatResponse) -> OutMessage: + assert response.choices, f"Xiaomi returned no choices: {response}" + message = response.choices[0].message + assert message is not None, f"Xiaomi choice has no message: {response}" + return message + + +def _deltas(result: StreamingResponse) -> list[_StreamDelta]: + require_successful_call(result) + assert result.is_streaming, f"response was not streamed: {result.headers}" + assert not result.stream_error, f"stream errored: {result.stream_error}" + assert result.stream_done, f"stream never reached [DONE]: {result.stream_events[-3:]}" + return [ + choice.delta + for event in result.stream_events + for choice in _StreamChunk.model_validate_json(event).choices + if choice.delta is not None + ] + + +@pytest.mark.parametrize("backend", BACKENDS) +class TestXiaomiMimoChatCompletions: + @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged") + def test_cost_header_and_spend_row_match_the_registry_price( + self, + client: PassthroughClient, + resources: ResourceManager, + registry: dict[str, CostMapEntry], + backend: str, + ) -> None: + price = registry.get(backend) + assert price is not None, f"{backend} has no row in the proxy's cost map, so native calls would bill $0" + assert price.litellm_provider == "xiaomi_mimo", f"{backend} is filed under the wrong provider: {price}" + assert price.input_cost_per_token and price.output_cost_per_token, f"{backend} carries no price: {price}" + model, key = _register(client, resources, backend) + + result = client.proxy.transport.send( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=f"{ARITHMETIC_PROMPT} {unique_marker()}")], + max_tokens=1024, + ), + ) + require_successful_call(result) + response = ChatResponse.model_validate_json(result.body) + message = _message(response) + assert message.content and "43" in message.content, f"answer lost: {message}" + assert message.reasoning_content, f"{backend} reasons, but no reasoning_content came back: {message}" + + usage = response.usage + assert usage is not None and usage.prompt_tokens and usage.completion_tokens, ( + f"response carries no usage, so the cost cannot be real: {result.body[:300]}" + ) + header_cost = result.response_cost + assert header_cost is not None and header_cost > 0, ( + f"x-litellm-response-cost header missing or non-positive: {result.headers}" + ) + cached = (usage.prompt_tokens_details.cached_tokens or 0) if usage.prompt_tokens_details else 0 + expected = ( + (usage.prompt_tokens - cached) * price.input_cost_per_token + + cached * (price.cache_read_input_token_cost or 0.0) + + usage.completion_tokens * price.output_cost_per_token + ) + assert _approx_equal(header_cost, expected), ( + f"header cost {header_cost} disagrees with the registry price for {backend} at {usage}: expected {expected}" + ) + + def _priced(rows: list[SpendLogRow]) -> bool: + return any(row.spend is not None and row.spend > 0 for row in rows) + + rows = client.proxy.poll_logs_for_key(key, predicate=_priced) + priced = [row for row in rows if row.spend is not None and row.spend > 0] + assert priced, f"no priced spend row landed for key {key}; got {rows}" + row = priced[0] + assert row.custom_llm_provider == "xiaomi_mimo", f"spend row misattributed: {row}" + assert row.spend is not None and _approx_equal(row.spend, header_cost), ( + f"logged spend {row.spend} disagrees with the x-litellm-response-cost header {header_cost}" + ) + + @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.thinking.stream.works") + def test_reasoning_and_answer_stream_as_deltas( + self, client: PassthroughClient, resources: ResourceManager, backend: str + ) -> None: + model, key = _register(client, resources, backend) + + deltas = _deltas( + client.proxy.chat_stream( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=COUNTING_PROMPT)], + max_tokens=2048, + stream=True, + ), + ) + ) + reasoning = "".join(delta.reasoning_content or "" for delta in deltas) + content = "".join(delta.content or "" for delta in deltas) + assert reasoning, f"stream carried no reasoning_content deltas: {deltas[:5]}" + assert "50" in content, f"streamed answer lost: {content[:300]!r}" + + @pytest.mark.covers("llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works") + def test_tool_call_is_returned(self, client: PassthroughClient, resources: ResourceManager, backend: str) -> None: + model, key = _register(client, resources, backend) + + message = _message( + unwrap( + client.proxy.chat( + key, + ChatBody( + model=model, + messages=[ChatMessage(role="user", content=WEATHER_PROMPT)], + tools=[WEATHER_TOOL], + max_tokens=1024, + ), + ) + ) + ) + assert message.tool_calls, f"{backend} dropped the tool call: {message}" + call = message.tool_calls[0] + assert call.id, f"tool call carries no id, so a tool result cannot answer it: {call}" + assert call.function.name == "get_weather", f"wrong tool called: {call}" + assert call.function.arguments, f"tool call carries no arguments: {call}" + args = _WeatherArgs.model_validate_json(call.function.arguments) + assert "paris" in args.location.lower(), f"tool arguments lost the location: {args}" diff --git a/tests/test_litellm/litellm_core_utils/test_logging_worker.py b/tests/test_litellm/litellm_core_utils/test_logging_worker.py index 1553e788472..2bb93a58531 100644 --- a/tests/test_litellm/litellm_core_utils/test_logging_worker.py +++ b/tests/test_litellm/litellm_core_utils/test_logging_worker.py @@ -180,6 +180,53 @@ class TestLoggingWorker: assert sorted(fired) == ["first", "second"] + @pytest.mark.parametrize("stranded", ["still_queued", "dequeued_never_started"]) + def test_flush_on_new_loop_drains_tasks_stranded_on_previous_loop(self, stranded): + """ + Regression: ``flush()`` from a new event loop used to ``join()`` the queue bound to the + previous loop, whose unfinished counter nothing on the new loop ever decrements. The first + such flush hung until pytest-timeout killed it and every later one raised + ``RuntimeError: ... is bound to a different event loop`` from the queue's Event. + """ + worker = LoggingWorker(timeout=1.0, max_queue_size=10) + callback = AsyncMock() + + async def enqueue_on_first_loop(): + if stranded == "still_queued": + worker._ensure_queue() + worker.enqueue(callback()) + return + worker.ensure_initialized_and_enqueue(callback()) + + asyncio.run(enqueue_on_first_loop()) + assert worker._queue is not None + expected_shape = (1, 0) if stranded == "still_queued" else (0, 1) + assert (worker._queue.qsize(), len(worker._unstarted_dequeued_tasks())) == expected_shape + assert callback.await_count == 0, "precondition: the callback never ran before the first loop closed" + + async def flush_twice_on_second_loop(): + await asyncio.wait_for(worker.flush(), timeout=5) + await asyncio.wait_for(worker.flush(), timeout=5) + + asyncio.run(flush_twice_on_second_loop()) + + assert callback.await_count == 1 + + def test_flush_starts_a_worker_when_the_queue_has_none(self): + """``flush()`` must drain a queue that exists on the current loop without a running worker.""" + worker = LoggingWorker(timeout=1.0, max_queue_size=10) + callback = AsyncMock() + + async def enqueue_then_flush(): + worker._ensure_queue() + worker.enqueue(callback()) + assert worker._worker_task is None, "precondition: nothing is draining the queue yet" + await asyncio.wait_for(worker.flush(), timeout=3) + + asyncio.run(enqueue_then_flush()) + + assert callback.await_count == 1 + def test_flush_on_exit_swallows_cancellation_and_drains_remaining(self): """A callback raising CancelledError must not abort the atexit flush of later events.""" worker = LoggingWorker(timeout=1.0, max_queue_size=10) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index e768139f04a..0969a913605 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -291,6 +291,106 @@ async def test_find_team_with_model_access_uses_request_method_for_passthrough_a assert "allowed_passthrough_routes" in exc_info.value.detail +_AUTH_ENFORCED_MODEL_HOST_ROUTES: Final = { + "test-uuid-1:subpath:/model-host/v1/extractor:GET,POST": { + "endpoint_id": "test-uuid-1", + "path": "/model-host/v1/extractor", + "type": "subpath", + "auth": True, + }, +} + + +@pytest.mark.asyncio +async def test_find_team_with_model_access_team_allowed_routes_wildcard_grants_auth_passthrough(): + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_allowed_routes=["openai_routes", "/model-host/*"]) + team_without_passthrough_allowlist = LiteLLM_TeamTable(team_id="team-a", models=["all-proxy-models"], metadata={}) + + with ( + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + return_value=team_without_passthrough_allowlist, + ), + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + _AUTH_ENFORCED_MODEL_HOST_ROUTES, + ), + patch("litellm.proxy.utils.get_server_root_path", return_value="/"), + ): + team_id, team_obj = await JWTAuthManager.find_team_with_model_access( + team_ids={"team-a"}, + requested_model=None, + route="/model-host/v1/extractor/predict", + jwt_handler=jwt_handler, + prisma_client=None, + user_api_key_cache=MagicMock(), + parent_otel_span=None, + proxy_logging_obj=MagicMock(), + request_method="POST", + ) + + assert team_id == "team-a" + assert team_obj == team_without_passthrough_allowlist + + +@pytest.mark.asyncio +async def test_auth_builder_header_team_allows_auth_passthrough_for_team_allowed_routes_wildcard(): + from litellm.proxy.utils import ProxyLogging + + jwt_handler = JWTHandler() + user_api_key_cache = DualCache() + jwt_handler.update_environment( + prisma_client=None, + user_api_key_cache=user_api_key_cache, + litellm_jwtauth=LiteLLM_JWTAuth( + team_ids_jwt_field="groups", + user_id_jwt_field="sub", + team_allowed_routes=["openai_routes", "/model-host/*"], + ), + ) + + with ( + patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock) as mock_auth_jwt, + patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), + patch.object(JWTAuthManager, "check_admin_access", new_callable=AsyncMock, return_value=None), + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + return_value=LiteLLM_TeamTable(team_id="team-2", metadata={}), + ), + patch.object( + JWTAuthManager, + "get_objects", + new_callable=AsyncMock, + return_value=(None, None, None, None, "user-1"), + ), + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + _AUTH_ENFORCED_MODEL_HOST_ROUTES, + ), + patch("litellm.proxy.utils.get_server_root_path", return_value="/"), + ): + mock_auth_jwt.return_value = {"sub": "user-1", "scope": "", "groups": ["team-1", "team-2"]} + + result = await JWTAuthManager.auth_builder( + api_key="jwt-token", + jwt_handler=jwt_handler, + request_data={}, + general_settings={}, + route="/model-host/v1/extractor/predict", + prisma_client=None, + user_api_key_cache=user_api_key_cache, + parent_otel_span=None, + proxy_logging_obj=ProxyLogging(user_api_key_cache=user_api_key_cache), + request_headers={"x-litellm-team-id": "team-2"}, + request_method="POST", + ) + + assert result["team_id"] == "team-2" + + @pytest.mark.asyncio async def test_auth_builder_proxy_admin_user_role(): """Test that is_proxy_admin is True when user_object.user_role is PROXY_ADMIN""" @@ -6463,6 +6563,90 @@ async def test_auth_builder_db_fallback_enforces_passthrough_route_access(): assert "passthrough route" in exc_info.value.detail +async def _auth_builder_via_db_team_fallback(team_allowed_routes: list[str]): + user_id = "u_passthrough" + user_object = LiteLLM_UserTable( + user_id=user_id, + user_role=LitellmUserRoles.INTERNAL_USER, + teams=["team_no_passthrough"], + ) + jwt_handler = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(fallback_to_db_teams=True, team_allowed_routes=team_allowed_routes) + + async def fake_get_team(team_id, **kwargs): + return LiteLLM_TeamTable(team_id=team_id, metadata={}) + + with ( + patch.object(jwt_handler, "auth_jwt", new_callable=AsyncMock, return_value={"sub": user_id, "scope": ""}), + patch.object(JWTAuthManager, "check_rbac_role", new_callable=AsyncMock), + patch.object(jwt_handler, "get_rbac_role", return_value=None), + patch.object(jwt_handler, "get_scopes", return_value=[]), + patch.object(jwt_handler, "get_object_id", return_value=None), + patch.object( + JWTAuthManager, + "get_user_info", + new_callable=AsyncMock, + return_value=(user_id, "u@example.com", True), + ), + patch.object(jwt_handler, "get_org_id", return_value=None), + patch.object(jwt_handler, "get_end_user_id", return_value=None), + patch.object(JWTAuthManager, "check_admin_access", new_callable=AsyncMock, return_value=None), + patch.object( + JWTAuthManager, + "get_objects", + new_callable=AsyncMock, + return_value=(user_object, None, None, None, user_id), + ), + patch.object(JWTAuthManager, "map_user_to_teams", new_callable=AsyncMock), + patch.object(JWTAuthManager, "validate_object_id", return_value=True), + patch.object(JWTAuthManager, "sync_user_role_and_teams", new_callable=AsyncMock), + patch( + "litellm.proxy.auth.handle_jwt.get_team_object", + new_callable=AsyncMock, + side_effect=fake_get_team, + ), + patch( + "litellm.proxy.auth.handle_jwt.get_team_membership", + new_callable=AsyncMock, + return_value=None, + ), + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + _AUTH_ENFORCED_MODEL_HOST_ROUTES, + ), + patch("litellm.proxy.utils.get_server_root_path", return_value="/"), + ): + return await JWTAuthManager.auth_builder( + api_key="test_jwt_token", + jwt_handler=jwt_handler, + request_data={}, + general_settings={"enforce_rbac": False}, + route="/model-host/v1/extractor/predict", + prisma_client=None, + user_api_key_cache=None, + parent_otel_span=None, + proxy_logging_obj=None, + request_headers=None, + request_method="POST", + ) + + +@pytest.mark.asyncio +async def test_auth_builder_db_fallback_team_allowed_routes_wildcard_grants_auth_passthrough(): + result = await _auth_builder_via_db_team_fallback(team_allowed_routes=["openai_routes", "/model-host/*"]) + + assert result["team_id"] == "team_no_passthrough" + + +@pytest.mark.asyncio +async def test_auth_builder_db_fallback_route_groups_alone_do_not_grant_auth_passthrough(): + with pytest.raises(HTTPException) as exc_info: + await _auth_builder_via_db_team_fallback(team_allowed_routes=["openai_routes", "mapped_pass_through_routes"]) + + assert exc_info.value.status_code == 403, exc_info.value.detail + assert "allowed_passthrough_routes" in exc_info.value.detail + + @pytest.mark.asyncio async def test_sync_user_role_and_teams_singular_claim_reconciles_memberships(): """When fallback_to_db_teams is on but the JWT carries a singular team claim diff --git a/tests/test_litellm/proxy/auth/test_route_checks.py b/tests/test_litellm/proxy/auth/test_route_checks.py index c9b8d6764c9..8da93ba341b 100644 --- a/tests/test_litellm/proxy/auth/test_route_checks.py +++ b/tests/test_litellm/proxy/auth/test_route_checks.py @@ -1249,6 +1249,122 @@ def test_non_proxy_admin_allows_auth_pass_through_with_team_allowlist(): ) +@pytest.mark.parametrize( + "route, team_allowed_routes, expected", + [ + ("/model-host/v1/extractor/predict", ["/model-host/*"], True), + ("/model-host", ["/model-host/*"], False), + ("/model-host/v1/extractor", ["/model-host/v1/extractor"], True), + ("/model-host/v1/extractor/predict", ["/model-host/v1/extractor"], False), + ("/other/v1/extractor", ["/model-host/*"], False), + ("/model-host/v1/extractor", ["openai_routes", "llm_api_routes", "mapped_pass_through_routes"], False), + ("/model-host/v1/extractor", ["*"], False), + ("/model-host/v1/extractor", ["/*"], False), + ("/model-host/v1/extractor", [], False), + ], +) +def test_jwt_team_routes_grant_pass_through_only_for_explicit_paths(route, team_allowed_routes, expected): + assert ( + RouteChecks.jwt_team_routes_grant_pass_through(route=route, team_allowed_routes=team_allowed_routes) + is expected + ) + + +_AUTH_ENFORCED_MODEL_HOST_ROUTES: Final = { + "test-uuid-1:subpath:/model-host/v1/extractor:GET,POST": { + "endpoint_id": "test-uuid-1", + "path": "/model-host/v1/extractor", + "type": "subpath", + "auth": True, + }, +} + + +def _jwt_handler_with_team_allowed_routes(team_allowed_routes: list[str]): + from litellm.proxy._types import LiteLLM_JWTAuth + from litellm.proxy.auth.handle_jwt import JWTHandler + + jwt_handler: Final = JWTHandler() + jwt_handler.litellm_jwtauth = LiteLLM_JWTAuth(team_allowed_routes=team_allowed_routes) + return jwt_handler + + +def _check_model_host_route_as(valid_token: UserAPIKeyAuth, team_allowed_routes: list[str]) -> None: + with ( + patch( + "litellm.proxy.pass_through_endpoints.pass_through_endpoints._registered_pass_through_routes", + _AUTH_ENFORCED_MODEL_HOST_ROUTES, + ), + patch("litellm.proxy.utils.get_server_root_path", return_value="/"), + patch( + "litellm.proxy.proxy_server.jwt_handler", + _jwt_handler_with_team_allowed_routes(team_allowed_routes), + ), + ): + RouteChecks.non_proxy_admin_allowed_routes_check( + user_obj=None, + _user_role=LitellmUserRoles.INTERNAL_USER.value, + route="/model-host/v1/extractor/predict", + request=MagicMock(spec=Request), + valid_token=valid_token, + request_data={}, + ) + + +def test_non_proxy_admin_allows_auth_pass_through_for_jwt_team_allowed_routes_wildcard(): + jwt_token: Final = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + team_id="team-a", + jwt_claims={"sub": "test_user"}, + ) + + _check_model_host_route_as(jwt_token, team_allowed_routes=["openai_routes", "/model-host/*"]) + + +def test_non_proxy_admin_denies_auth_pass_through_for_jwt_when_only_route_groups_configured(): + jwt_token: Final = UserAPIKeyAuth( + user_id="test_user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + team_id="team-a", + jwt_claims={"sub": "test_user"}, + ) + + with pytest.raises(HTTPException) as exc_info: + _check_model_host_route_as(jwt_token, team_allowed_routes=["openai_routes", "mapped_pass_through_routes"]) + + assert exc_info.value.status_code == 403, exc_info.value.detail + assert "allowed_passthrough_routes" in exc_info.value.detail + + +@pytest.mark.parametrize( + "api_key, team_id, jwt_claims", + [ + ("sk-test-key", "team-a", None), + ("sk-test-key", "team-a", {"sub": "test_user"}), + (None, "team-a", None), + (None, None, {"sub": "test_user"}), + ], + ids=["plain_virtual_key", "jwt_mapped_virtual_key", "keyless_non_jwt_caller", "jwt_without_team"], +) +def test_non_proxy_admin_jwt_team_allowed_routes_grant_pass_through_only_to_jwt_team_callers( + api_key, team_id, jwt_claims +): + caller: Final = UserAPIKeyAuth( + api_key=api_key, + user_id="test_user", + user_role=LitellmUserRoles.INTERNAL_USER.value, + team_id=team_id, + jwt_claims=jwt_claims, + ) + + with pytest.raises(HTTPException) as exc_info: + _check_model_host_route_as(caller, team_allowed_routes=["openai_routes", "/model-host/*"]) + + assert exc_info.value.status_code == 403, exc_info.value.detail + assert "allowed_passthrough_routes" in exc_info.value.detail + + def test_virtual_key_without_llm_api_routes_cannot_access_pass_through(): """ Test that virtual keys without llm_api_routes permission cannot access registered pass-through endpoints. diff --git a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py index 80773f314d8..c66d490f122 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_mcp_management_endpoints.py @@ -6,7 +6,7 @@ import logging from contextlib import ExitStack from datetime import datetime, timedelta from types import SimpleNamespace -from typing import List, Optional, cast +from typing import Final, List, Optional, cast from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1614,6 +1614,60 @@ class TestTeamScopedMCPServerAccess: result = await fetch_all_mcp_servers(user_api_key_dict=mock_user_auth, team_id="any-team-id") assert len(result) == 1 + +class TestFetchAllMCPServersOrdering: + def test_display_order_is_case_insensitive_name_then_id(self) -> None: + servers: Final = ( + LiteLLM_MCPServerTable(server_id="s-2", server_name="GitHub", alias="aaa", transport=MCPTransport.http), + LiteLLM_MCPServerTable(server_id="s-1", alias="github", transport=MCPTransport.http), + LiteLLM_MCPServerTable(server_id="s-0", server_name="Slack", alias="zzz", transport=MCPTransport.http), + LiteLLM_MCPServerTable(server_id="confluence", server_name="", alias="", transport=MCPTransport.http), + ) + + ordered: Final = sorted(servers, key=mgmt_endpoints._mcp_server_display_order) + assert [s.server_id for s in ordered] == ["confluence", "s-1", "s-2", "s-0"] + + @pytest.mark.parametrize("team_id", [None, "team-1"]) + @pytest.mark.parametrize("reverse", [False, True]) + @pytest.mark.asyncio + async def test_list_is_sorted_by_display_name_regardless_of_resolution_order( + self, team_id: str | None, reverse: bool + ) -> None: + mock_user_auth: Final = generate_mock_user_api_key_auth( + user_role=LitellmUserRoles.PROXY_ADMIN, + user_id="admin_user", + ) + servers: Final = ( + generate_mock_mcp_server_db_record(server_id="s-zeta", alias="zeta"), + generate_mock_mcp_server_db_record(server_id="s-alpha", alias="Alpha"), + generate_mock_mcp_server_db_record(server_id="s-mid", alias="mid"), + ) + resolved: Final = list(reversed(servers) if reverse else servers) + mock_manager: Final = MagicMock() + mock_manager.get_all_allowed_mcp_servers = AsyncMock(return_value=resolved) + with ( + patch( # test-quality-ok: the route reads a module-global manager with no injection seam + "litellm.proxy.management_endpoints.mcp_management_endpoints.global_mcp_server_manager", + mock_manager, + ), + patch( # test-quality-ok: admin view is derived from module-global proxy settings + "litellm.proxy.management_endpoints.mcp_management_endpoints._user_has_admin_view", + return_value=True, + ), + patch( # test-quality-ok: auth contexts need a live prisma client + "litellm.proxy.management_endpoints.mcp_management_endpoints.build_effective_auth_contexts", + AsyncMock(return_value=[mock_user_auth]), + ), + patch( # test-quality-ok: isolate the route's ordering from team database resolution + "litellm.proxy.management_endpoints.mcp_management_endpoints._get_team_scoped_mcp_server_list", + AsyncMock(return_value=resolved), + ), + ): + result: Final = await mgmt_endpoints.fetch_all_mcp_servers( + user_api_key_dict=mock_user_auth, team_id=team_id + ) + assert [s.server_id for s in result] == ["s-alpha", "s-mid", "s-zeta"] + @pytest.mark.asyncio async def test_restricted_virtual_key_cannot_use_team_id_filter(self): """Restricted virtual keys must not bypass access limits via team_id.""" diff --git a/tests/test_litellm_rust/support/s3_stub.py b/tests/test_litellm_rust/support/s3_stub.py new file mode 100644 index 00000000000..5a683fb78f3 --- /dev/null +++ b/tests/test_litellm_rust/support/s3_stub.py @@ -0,0 +1,112 @@ +"""In-process path-style S3 stub for native cache parity tests.""" + +import threading +from dataclasses import dataclass, field +from email.utils import parsedate_to_datetime +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Final +from urllib.parse import unquote, urlsplit + +_STORED_HEADERS: Final = ( + "cache-control", + "content-type", + "content-language", + "content-disposition", + "expires", +) + + +@dataclass +class S3Object: + body: bytes + headers: dict[str, str] = field(default_factory=dict) + + +class S3Stub: + """Minimal path-style S3 endpoint serving PUT and GET object operations.""" + + def __init__(self) -> None: + self._objects: dict[str, S3Object] = {} + stub: Final = self + + class Handler(BaseHTTPRequestHandler): + def _key(self) -> str: + parts: Final = urlsplit(self.path).path.lstrip("/").split("/", 1) + return unquote(parts[1]) if len(parts) == 2 else "" + + def _read_body(self) -> bytes: + transfer: Final = self.headers.get("transfer-encoding", "") + if "chunked" not in transfer: + return self.rfile.read(int(self.headers.get("content-length", 0))) + chunks: Final = bytearray() + while True: + size = int(self.rfile.readline().split(b";")[0].strip(), 16) + if size == 0: + while self.rfile.readline().strip(): + pass + return bytes(chunks) + chunks.extend(self.rfile.read(size)) + self.rfile.readline() + + def do_PUT(self) -> None: + body: Final = self._read_body() + headers: Final = {name: self.headers[name] for name in _STORED_HEADERS if name in self.headers} + stub._objects = {**stub._objects, self._key(): S3Object(body=body, headers=headers)} + self.send_response(200) + self.send_header("ETag", '"stub"') + self.send_header("Content-Length", "0") + self.end_headers() + + def do_HEAD(self) -> None: + self._object(send_body=False) + + def do_GET(self) -> None: + self._object(send_body=True) + + def _object(self, send_body: bool) -> None: + entry: Final = stub._objects.get(self._key()) + if entry is None: + self.send_response(404) + self.send_header("Content-Type", "application/xml") + body: Final = b'NoSuchKey' + self.send_header("Content-Length", str(len(body))) + self.end_headers() + if send_body: + self.wfile.write(body) + return + self.send_response(200) + for name, value in entry.headers.items(): + self.send_header(name, value) + self.send_header("ETag", '"stub"') + self.send_header("Content-Length", str(len(entry.body))) + self.end_headers() + if send_body: + self.wfile.write(entry.body) + + def log_message(self, format: str, *args: object) -> None: + pass + + self._server: Final = ThreadingHTTPServer(("127.0.0.1", 0), Handler) + self._worker: Final = threading.Thread(target=self._server.serve_forever, daemon=True) + self._worker.start() + + @property + def url(self) -> str: + host, port = self._server.server_address[:2] + return f"http://{host}:{port}" + + @property + def objects(self) -> dict[str, S3Object]: + return self._objects + + def put_object(self, key: str, body: bytes, headers: dict[str, str] | None = None) -> None: + self._objects = {**self._objects, key: S3Object(body=body, headers=headers or {})} + + def expires(self, key: str) -> object: + header: Final = self._objects[key].headers.get("expires") + return parsedate_to_datetime(header) if header else None + + def close(self) -> None: + self._server.shutdown() + self._server.server_close() + self._worker.join(timeout=5) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index c11a43abc80..45f7214ecbf 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -11,12 +11,16 @@ import uuid import weakref from collections.abc import Callable, Generator from contextlib import ExitStack +from datetime import datetime from pathlib import Path from types import SimpleNamespace from typing import Final, Protocol, cast +from unittest.mock import Mock from urllib.parse import urlparse from uuid import uuid4 +import boto3 +import botocore.config import diskcache import fakeredis import pytest @@ -26,17 +30,19 @@ from azure.storage.blob import ContainerClient import litellm from litellm.caching.azure_blob_cache import AzureBlobCache from litellm.caching.caching import Cache, disable_cache, enable_cache, update_cache -from litellm.caching.gcs_cache import GCSCache from litellm.caching.disk_cache import DiskCache +from litellm.caching.gcs_cache import GCSCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.caching.redis_semantic_cache import RedisSemanticCache +from litellm.caching.s3_cache import S3Cache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType from litellm.types.llms.custom_llm import CustomLLMItem from litellm.types.utils import EmbeddingResponse from tests.test_litellm_rust.support.fake_gcs import FakeGcs from tests.test_litellm_rust.support.isolation import rebound +from tests.test_litellm_rust.support.s3_stub import S3Stub _CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name _CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name @@ -436,15 +442,20 @@ def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure assert handle.backend == "azure-blob" account_url: Final = backend.container_client.url.removesuffix(f"/{backend.container_client.container_name}") with pytest.raises(TypeError, match="containers must match"): - _native._CacheTestHandle.azure_blob(account_url, f"{backend.container_client.container_name}-other")._bind_facade( - azure_blob_facade - ) + _native._CacheTestHandle.azure_blob( + account_url, f"{backend.container_client.container_name}-other" + )._bind_facade(azure_blob_facade) handle._bind_facade(azure_blob_facade) resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=azure_blob_facade)) native: Final = resolver.resolve() assert native.kind == "native" - response: Final = {"choices": [{"text": "caf\u00e9 \u2603"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} + response: Final = { + "choices": [{"text": "caf\u00e9 \u2603"}], + "usage": {"total_tokens": 3}, + "flag": True, + "empty": None, + } native.store({**request("sync"), "ttl_seconds": 0.001}, response) native.store(request("sync"), {"choices": [{"text": "second"}]}) time.sleep(0.01) @@ -499,7 +510,9 @@ async def test_azure_blob_native_async_writes_overwrite_batch_and_flush_like_pyt await binding.async_store({**request("async"), "ttl_seconds": 0.001}, {"value": 2}) time.sleep(0.01) assert await binding.async_lookup(request("async")) == {"value": 2} - assert await backend.async_get_cache("async") == json.loads(backend.container_client.download_blob("async").readall()) + assert await backend.async_get_cache("async") == json.loads( + backend.container_client.download_blob("async").readall() + ) assert cast(CacheLookup, azure_blob_facade).get_cache(cache_key="async") == {"value": 2} await binding.async_store_batch([request("first"), request("second")], [{"value": 3}, {"value": 4}]) @@ -544,6 +557,8 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: assert client.get("second") is not None await facade.cache.disconnect() client.close() + + async def test_disk_reads_python_entries_and_python_reads_native_entries(tmp_path: Path) -> None: disk_cache: Final = DiskCache(disk_cache_dir=str(tmp_path)) response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} @@ -652,6 +667,183 @@ async def test_disk_native_batch_lookup_and_store_report_partial_hits(tmp_path: } +@pytest.fixture +def s3_stub() -> Generator[S3Stub]: + stub: Final = S3Stub() + try: + yield stub + finally: + stub.close() + + +def python_s3(url: str) -> S3Cache: + return S3Cache( + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + + +async def test_s3_reads_python_entries_and_writes_with_python_metadata(s3_stub: S3Stub) -> None: + python_cache: Final = python_s3(s3_stub.url) + response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}} + python_cache.set_cache("sync:key", {"timestamp": time.time(), "response": response}, ttl=90) + python_cache.set_cache("plain", {"timestamp": time.time(), "response": response}) + s3_stub.put_object("team/malformed", b"not a cache entry") + s3_stub.put_object( + "team/expired", + json.dumps({"timestamp": time.time(), "response": response}).encode(), + {"expires": "Thu, 01 Jan 1970 00:00:00 GMT"}, + ) + binding: Final = _native._CacheTestResolver( + SimpleNamespace( + cache=_native._CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + ) + ).resolve() + + assert binding.lookup(request("sync:key")) == response + assert await binding.async_lookup(request("plain")) == response + assert binding.lookup(request("malformed")) is None + assert binding.lookup(request("expired")) is None + assert binding.lookup(request("absent")) is None + + binding.store({**request("native:key"), "ttl_seconds": 90.0}, response) + await binding.async_store(request("no_ttl"), response) + stored: Final = s3_stub.objects["team/native/key"] + assert stored.headers["content-type"] == "application/json" + assert stored.headers["content-language"] == "en" + assert stored.headers["content-disposition"] == 'inline; filename="team/native/key.json"' + assert stored.headers["cache-control"] == "immutable, max-age=90, s-maxage=90" + expires: Final = cast(datetime, s3_stub.expires("team/native/key")) + remaining: Final = (expires - datetime.now(expires.tzinfo)).total_seconds() + assert 60 < remaining <= 91 + no_ttl: Final = s3_stub.objects["team/no_ttl"] + assert no_ttl.headers["cache-control"] == "immutable, max-age=31536000, s-maxage=31536000" + assert "expires" not in no_ttl.headers + assert python_cache.get_cache("native:key")["response"] == response + + partial: Final = await binding.async_lookup_batch([request("native:key"), request("absent"), request("malformed")]) + assert partial == {"values": [response, None, None], "missing_indices": [1, 2]} + + +def test_s3_facade_binds_only_exact_configuration_and_falls_back_on_mutation(s3_stub: S3Stub) -> None: + facade: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + handle: Final = _native._CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + with pytest.raises(TypeError, match="buckets must match"): + _native._CacheTestHandle.s3("other", region="us-east-1", endpoint_url=s3_stub.url)._bind_facade(facade) + with pytest.raises(TypeError, match="key prefixes must match"): + _native._CacheTestHandle.s3( + "cache-bucket", region="us-east-1", endpoint_url=s3_stub.url, key_prefix="other/" + )._bind_facade(facade) + handle._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + binding: Final = resolver.resolve() + assert binding.kind == "native" + + handler: Final = Mock() + facade.cache.s3_client.meta.events.register("before-call.s3.*", handler) + binding.store(request("native"), {"answer": 1}) + assert binding.lookup(request("native")) == {"answer": 1} + assert handler.call_count == 0 + assert "team/native" in s3_stub.objects + + with rebound(facade.cache, "bucket_name", "other"): + assert resolver.resolve().kind == "python_callback" + other_client: Final = boto3.client( + "s3", + region_name="us-east-1", + endpoint_url=s3_stub.url, + aws_access_key_id="key", + aws_secret_access_key="secret", + ) + with rebound(facade.cache, "s3_client", other_client): + assert resolver.resolve().kind == "python_callback" + + class CustomS3Cache(S3Cache): + pass + + subclassed: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + subclassed.cache = CustomS3Cache( + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + with pytest.raises(TypeError): + handle._bind_facade(subclassed) + assert _native._CacheTestResolver(SimpleNamespace(cache=subclassed)).resolve().kind == "python_callback" + + +def test_s3_facade_rejects_configurations_that_require_python(s3_stub: S3Stub) -> None: + handle: Final = _native._CacheTestHandle.s3( + "cache-bucket", + region="us-east-1", + endpoint_url=s3_stub.url, + key_prefix="team/", + access_key_id="key", + secret_access_key="secret", + ) + unverified: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url="https://s3.example.test", + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + s3_verify=False, + ) + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(unverified) + proxied: Final = Cache( + type=LiteLLMCacheType.S3, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=s3_stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + s3_config=botocore.config.Config(proxies={"https": "http://proxy.test"}), + ) + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(proxied) + + async def test_gcs_reads_python_entries_and_writes_python_compatible_objects( fake_gcs: FakeGcs, monkeypatch: pytest.MonkeyPatch ) -> None: @@ -787,6 +979,7 @@ async def test_gcs_facade_binds_only_exact_matching_configuration( assert resolver.resolve().kind == "python_callback" with rebound(facade.cache, "path_service_account", "sa.json"): assert resolver.resolve().kind == "python_callback" + def no_get_cache(*args: object, **kwargs: object) -> None: return None @@ -921,7 +1114,9 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n await binding.async_flush() - remaining: Final = tuple(sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node))) + remaining: Final = tuple( + sorted(key for node in client.get_primaries() for key in client.keys("parity:*", target_nodes=node)) + ) assert remaining == (), remaining assert client.get("unscoped") == b"stays" client.delete("unscoped") diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py new file mode 100644 index 00000000000..c87a9f86a80 --- /dev/null +++ b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py @@ -0,0 +1,599 @@ +import asyncio +import contextvars +import hashlib +import os +import struct +import threading +import time +from collections.abc import Generator, Mapping +from types import SimpleNamespace +from typing import Final, cast +from uuid import uuid4 + +import pytest +import redis + +from litellm.caching.caching import Cache +from litellm.caching.valkey_semantic_cache import ValkeySemanticCache +from litellm.rust_bridge import _native +from litellm.types.caching import LiteLLMCacheType + +pytestmark: Final = pytest.mark.requires_rust_extension +embedding_context: Final = contextvars.ContextVar("embedding_context") + + +@pytest.fixture +def valkey_url() -> str: + url: Final = os.environ.get("LITELLM_TEST_VALKEY_URL") + if url is None: + pytest.skip("LITELLM_TEST_VALKEY_URL is not set") + return url + + +@pytest.fixture +def index_name(valkey_url: str) -> Generator[str]: + index: Final = f"litellm_test_{uuid4().hex}" + yield index + client: Final = redis.Redis.from_url(valkey_url) + try: + client.ft(index).dropindex(delete_documents=True) + except redis.ResponseError: + pass + finally: + client.close() + + +def _request(prompt: str = "semantic cache prompt") -> dict[str, object]: + return { + "key": {"preset": "key"}, + "messages": [{"role": "user", "content": prompt}], + } + + +def _field_request( + prompt: str, + metadata: Mapping[str, object], + *, + namespace: str | None = None, + litellm_metadata: Mapping[str, object] | None = None, + litellm_params: Mapping[str, object] | None = None, +) -> dict[str, object]: + request: Final = { + "key": { + "fields": [ + { + "name": "model", + "value": "gpt-4.1", + "api_parameter": True, + "internal_parameter": False, + }, + { + "name": "messages", + "value": prompt, + "api_parameter": True, + "internal_parameter": False, + }, + ], + "namespace": namespace, + }, + "messages": [{"role": "user", "content": prompt}], + "metadata": dict(metadata), + } + if litellm_metadata is not None: + request["litellm_metadata"] = dict(litellm_metadata) + if litellm_params is not None: + request["litellm_params"] = dict(litellm_params) + return request + + +def _facade( + url: str, + index_name: str, + embeddings: Mapping[str, list[float]], + *, + namespace: str | None = None, +) -> Cache: + facade: Final = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + valkey_semantic_cache_index_name=index_name, + namespace=namespace, + ) + vectors: Final = embeddings + + def embed(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]: + return vectors[prompt] + + async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]: + return vectors[prompt] + + facade.cache._get_embedding = embed + facade.cache._get_async_embedding = async_embedding + return facade + + +def _backend( + url: str, + index_name: str, + embeddings: Mapping[str, list[float]] | None = None, +) -> ValkeySemanticCache: + vectors: Final = embeddings or {"semantic cache prompt": [1.0, 0.0]} + backend: Final = ValkeySemanticCache( + redis_url=url, + similarity_threshold=0.8, + index_name=index_name, + ) + + def embed(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]: + return vectors[prompt] + + async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]: + return vectors[prompt] + + backend._get_embedding = embed + backend._get_async_embedding = async_embedding + return backend + + +def test_python_write_native_read( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + response: Final = {"answer": "python"} + backend.set_cache("key", response, messages=_request()["messages"]) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + assert binding.lookup(_request()) == response + + +def test_native_write_python_read( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + response: Final = {"answer": "native"} + binding.store({**_request(), "ttl_seconds": 2.0}, response) + cached: Final = cast(Mapping[str, object], backend.get_cache("key", messages=_request()["messages"])) + assert cached["response"] == response + + +async def test_async_lookup_and_store( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + request: Final = {**_request(), "ttl_seconds": 2.0} + await binding.async_store(request, {"answer": "async"}) + assert await binding.async_lookup(request) == {"answer": "async"} + + +async def test_disabled_cache_controls_skip_async_embedding( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + calls: Final = [] + + async def fail_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]: + calls.append(prompt) + raise AssertionError("embedding must not run") + + backend._get_async_embedding = fail_embedding + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + controls: Final = { + "supported_call_type": True, + "configured": True, + "native_backend": True, + "default_on": True, + "caching": True, + "no_cache": False, + "no_store": False, + "use_cache": True, + } + no_read_request: Final = {**_request(), "controls": {**controls, "no_cache": True}} + assert await binding.async_lookup(no_read_request) is None + no_write_request: Final = {**_request(), "controls": {**controls, "no_store": True}} + await binding.async_store(no_write_request, {"answer": "blocked"}) + assert calls == [] + client: Final = redis.Redis.from_url(valkey_url) + assert list(client.scan_iter(f"{index_name}:*")) == [] + client.close() + + +async def test_async_embedding_runs_inline_in_caller_task( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + observed: dict[str, object] = {} + + async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]: + observed["context"] = embedding_context.get("missing") + observed["task"] = asyncio.current_task() + observed["thread"] = threading.get_ident() + embedding_context.set("embedder") + return [1.0, 0.0] + + backend._get_async_embedding = async_embedding + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + request: Final = {**_request(), "ttl_seconds": 2.0} + caller_task: Final = asyncio.current_task() + caller_thread: Final = threading.get_ident() + token: Final = embedding_context.set("caller") + try: + await binding.async_store(request, {"answer": "inline"}) + assert observed["context"] == "caller" + assert observed["task"] is caller_task + assert observed["thread"] == caller_thread + assert embedding_context.get() == "embedder" + assert await binding.async_lookup(request) == {"answer": "inline"} + finally: + embedding_context.reset(token) + + +def test_facade_activation_and_mutation_fallback( + valkey_url: str, + index_name: str, +) -> None: + facade: Final = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + redis_url=valkey_url, + similarity_threshold=0.8, + valkey_semantic_cache_index_name=index_name, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + facade.cache, + ) + handle._bind_facade(facade) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + facade.cache.similarity_threshold = 0.7 + assert resolver.resolve().kind == "python_callback" + + +def test_batch_lookup_is_unsupported( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + backend, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + with pytest.raises(NotImplementedError): + binding.lookup_batch([_request()]) + + +def test_ttl_expiry( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store({**_request(), "ttl_seconds": 1.0}, {"answer": "expires"}) + client: Final = redis.Redis.from_url(valkey_url) + documents: Final = list(client.scan_iter(f"{index_name}:*")) + assert len(documents) == 1 + assert client.ttl(documents[0]) > 0 + time.sleep(1.5) + assert binding.lookup(_request()) is None + + +def test_no_ttl_is_persistent_and_python_reads_native_value( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + response: Final = {"answer": "persistent"} + binding.store(_request(), response) + client: Final = redis.Redis.from_url(valkey_url) + documents: Final = list(client.scan_iter(f"{index_name}:*")) + assert len(documents) == 1 + assert client.ttl(documents[0]) == -1 + cached: Final = cast(Mapping[str, object], backend.get_cache("key", messages=_request()["messages"])) + assert cached["response"] == response + + +def test_below_threshold_misses_on_native_and_python( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend( + valkey_url, + index_name, + {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store(_request("prompt A"), {"answer": "A"}) + assert binding.lookup(_request("prompt B")) is None + assert backend.get_cache("key", messages=_request("prompt B")["messages"]) is None + + +def test_malformed_entry_is_a_miss_on_native_and_python( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + client: Final = redis.Redis.from_url(valkey_url) + scope: Final = hashlib.sha256(b"key").hexdigest() + document: Final = f"{index_name}:{scope}:{uuid4().hex}" + client.hset( + document, + mapping={ + "litellm_cache_key": scope, + "prompt": "semantic cache prompt", + "response": "not json", + "embedding": struct.pack("<2f", 1.0, 0.0), + }, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + assert binding.lookup(_request()) is None + assert backend.get_cache("key", messages=_request()["messages"]) is None + + +def test_mixed_content_parts_match_python_semantic_behavior( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + messages: Final = [{"role": "user", "content": ["raw", {"text": "hello"}]}] + backend.set_cache("key", {"answer": "mixed"}, messages=messages) + assert backend.get_cache("key", messages=messages) is None + + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + request: Final = {**_request(), "messages": messages} + binding.store(request, {"answer": "mixed"}) + assert binding.lookup(request) is None + client: Final = redis.Redis.from_url(valkey_url) + assert list(client.scan_iter(f"{index_name}:*")) == [] + client.close() + + +async def test_async_store_batch_and_lookup( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend( + valkey_url, + index_name, + {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}, + ) + sync_calls: Final = [] + async_tasks: Final = [] + + def sync_embedding(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]: + sync_calls.append(prompt) + return {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}[prompt] + + async def async_embedding( + prompt: str, + metadata: dict[str, object] | None = None, + ) -> list[float]: + async_tasks.append(asyncio.current_task()) + return {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}[prompt] + + backend._get_embedding = sync_embedding + backend._get_async_embedding = async_embedding + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + requests: Final = [_request("prompt A"), _request("prompt B")] + responses: Final = [{"answer": "A"}, {"answer": "B"}] + caller_task: Final = asyncio.current_task() + await binding.async_store_batch(requests, responses) + assert sync_calls == [] + assert async_tasks + assert all(task is caller_task for task in async_tasks) + assert await binding.async_lookup(requests[0]) == responses[0] + assert await binding.async_lookup(requests[1]) == responses[1] + + +def test_subclass_backend_falls_back_to_python( + valkey_url: str, + index_name: str, +) -> None: + class Custom(ValkeySemanticCache): + pass + + facade: Final = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + redis_url=valkey_url, + similarity_threshold=0.8, + valkey_semantic_cache_index_name=index_name, + ) + facade.cache = Custom(redis_url=valkey_url, similarity_threshold=0.8, index_name=index_name) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "python_callback" + + +def test_field_key_matches_python_semantic_scope( + valkey_url: str, + index_name: str, +) -> None: + facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]}) + metadata: Final = {"user_api_key": "k1"} + expected: Final = facade.get_cache_key( + model="gpt-4.1", + messages=[{"role": "user", "content": "semantic cache prompt"}], + metadata=metadata, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + facade.cache, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store(_field_request("semantic cache prompt", metadata), {"answer": "scoped"}) + client: Final = redis.Redis.from_url(valkey_url) + documents: Final = list(client.scan_iter(f"{index_name}:*")) + assert len(documents) == 1 + document_parts: Final = documents[0].decode().split(":") + assert document_parts[1] == hashlib.sha256(expected.encode()).hexdigest() + client.close() + + +def test_field_key_reads_all_python_tenant_metadata_sources( + valkey_url: str, + index_name: str, +) -> None: + facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]}) + params_metadata: Final = {"user_api_key_team_id": "team-from-params"} + expected: Final = facade.get_cache_key( + model="gpt-4.1", + messages=[{"role": "user", "content": "semantic cache prompt"}], + metadata={}, + litellm_params={"metadata": params_metadata}, + ) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + facade.cache, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store( + _field_request( + "semantic cache prompt", + {}, + litellm_params={"metadata": params_metadata}, + ), + {"answer": "params"}, + ) + client: Final = redis.Redis.from_url(valkey_url) + documents: Final = list(client.scan_iter(f"{index_name}:*")) + assert len(documents) == 1 + document_parts: Final = documents[0].decode().split(":") + assert document_parts[1] == hashlib.sha256(expected.encode()).hexdigest() + client.close() + + assert ( + binding.lookup( + _field_request( + "semantic cache prompt", + {}, + litellm_metadata={"user_api_key_team_id": "team-from-litellm"}, + ) + ) + is None + ) + + +def test_namespace_isolates_semantic_entries( + valkey_url: str, + index_name: str, +) -> None: + facade: Final = _facade( + valkey_url, + index_name, + {"semantic cache prompt": [1.0, 0.0]}, + namespace="team-a", + ) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + facade.cache, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + team_a: Final = _field_request("semantic cache prompt", {}, namespace="team-a") + team_b: Final = _field_request("semantic cache prompt", {}, namespace="team-b") + binding.store(team_a, {"answer": "team-a"}) + assert binding.lookup(team_b) is None + assert binding.lookup(team_a) == {"answer": "team-a"} + cached: Final = cast( + Mapping[str, object], + facade.get_cache( + model="gpt-4.1", + messages=[{"role": "user", "content": "semantic cache prompt"}], + ), + ) + assert cached == {"answer": "team-a"} + + +def test_field_key_isolates_tenant_scope( + valkey_url: str, + index_name: str, +) -> None: + facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]}) + handle: Final = _native._CacheTestHandle.valkey_semantic( + valkey_url, + 0.8, + index_name, + facade.cache, + ) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + binding.store( + _field_request("semantic cache prompt", {"user_api_key": "k1"}), + {"answer": "tenant one"}, + ) + assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k2"})) is None + assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k1"})) == {"answer": "tenant one"} + + +def test_tls_valkey_facade_falls_back_to_python( + index_name: str, +) -> None: + facade: Final = Cache( + type=LiteLLMCacheType.VALKEY_SEMANTIC, + redis_url="rediss://127.0.0.1:6390/0", + similarity_threshold=0.8, + valkey_semantic_cache_index_name=index_name, + ) + resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "python_callback" + + +async def test_ping_maps_unsupported_native_operation_to_not_implemented( + valkey_url: str, + index_name: str, +) -> None: + backend: Final = _backend(valkey_url, index_name) + handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + with pytest.raises(NotImplementedError): + await binding.ping() diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx index 091b2f1403f..1217d878489 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.test.tsx @@ -3,7 +3,8 @@ import { render, waitFor, screen, act, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, it, expect, vi, beforeEach } from "vitest"; import { QueryClient, QueryClientProvider } from "@tanstack/react-query"; -import MCPServers from "./mcp_servers"; +import MCPServers, { compareServers, type SortKey } from "./mcp_servers"; +import type { MCPServer } from "@/components/mcp_tools/types"; import * as networking from "@/components/networking"; // Mock the networking module @@ -31,6 +32,98 @@ const createQueryClient = () => }, }); +describe("compareServers", () => { + const server = (server_id: string, name: string, created_at = ""): MCPServer => ({ + server_id, + server_name: name, + created_at, + updated_at: created_at, + created_by: "user", + updated_by: "user", + }); + + const shuffled = [server("c", "github"), server("a", "slack"), server("b", "Jira")]; + + it("orders servers without timestamps by name so config.yaml servers render in a stable order", () => { + const byCreated = [...shuffled].sort((a, b) => compareServers(a, b, "created_desc")).map((s) => s.server_id); + const byUpdated = [...shuffled].sort((a, b) => compareServers(a, b, "updated_desc")).map((s) => s.server_id); + const byHealth = [...shuffled].sort((a, b) => compareServers(a, b, "health")).map((s) => s.server_id); + + expect(byCreated).toEqual(["c", "b", "a"]); + expect(byUpdated).toEqual(["c", "b", "a"]); + expect(byHealth).toEqual(["c", "b", "a"]); + }); + + it("keeps newest-first when timestamps differ", () => { + const newest = server("new", "zzz", "2026-02-01T00:00:00Z"); + const oldest = server("old", "aaa", "2026-01-01T00:00:00Z"); + expect([oldest, newest].sort((a, b) => compareServers(a, b, "created_desc")).map((s) => s.server_id)).toEqual([ + "new", + "old", + ]); + }); + + it.each(["created_desc", "updated_desc", "name_asc", "health"])( + "breaks equal timestamps and names by ID for %s regardless of input order", + (sort) => { + const servers = [ + server("b", "GitHub", "2026-01-01T00:00:00Z"), + server("c", "Slack", "2026-01-01T00:00:00Z"), + server("a", "github", "2026-01-01T00:00:00Z"), + ]; + for (const input of [servers, [...servers].reverse()]) { + expect([...input].sort((a, b) => compareServers(a, b, sort)).map((s) => s.server_id)).toEqual(["a", "b", "c"]); + } + }, + ); + + it("uses the display name before alias, then falls back to alias and ID", () => { + const servers: MCPServer[] = [ + { ...server("s-slack", "Slack"), alias: "aaa" }, + { ...server("s-github", ""), server_name: null, alias: "GitHub" }, + { ...server("confluence", ""), alias: "" }, + ]; + for (const input of [servers, [...servers].reverse()]) { + expect([...input].sort((a, b) => compareServers(a, b, "name_asc")).map((s) => s.server_id)).toEqual([ + "confluence", + "s-github", + "s-slack", + ]); + } + }); + + it.each(["created_desc", "updated_desc", "health"])( + "keeps timestamped servers before missing timestamps for %s", + (sort) => { + const servers = [ + server("config", "aaa"), + server("older", "bbb", "2026-01-01T00:00:00Z"), + server("newer", "zzz", "2026-02-01T00:00:00Z"), + ]; + for (const input of [servers, [...servers].reverse()]) { + expect([...input].sort((a, b) => compareServers(a, b, sort)).map((s) => s.server_id)).toEqual([ + "newer", + "older", + "config", + ]); + } + }, + ); + + it("sorts health before recency and display name", () => { + const servers: MCPServer[] = [ + { ...server("healthy", "aaa", "2026-03-01T00:00:00Z"), status: "healthy" }, + { ...server("unknown", "bbb", "2026-02-01T00:00:00Z"), status: "unknown" }, + { ...server("unhealthy", "zzz", "2026-01-01T00:00:00Z"), status: "unhealthy" }, + ]; + expect(servers.sort((a, b) => compareServers(a, b, "health")).map((s) => s.server_id)).toEqual([ + "unhealthy", + "unknown", + "healthy", + ]); + }); +}); + describe("MCPServers", () => { const defaultProps = { accessToken: "123", @@ -74,6 +167,134 @@ describe("MCPServers", () => { const myConnections = await screen.findByRole("link", { name: "My Connections" }); expect(myConnections).toBeVisible(); expect(myConnections).toHaveAttribute("href", "/ui/connect"); + for (const name of ["Semantic Filter", "Tool Search", "Network Settings", "Submitted MCPs"]) { + const tab = screen.queryByRole("tab", { name }); + if (userRole === "Admin") { + expect(tab).toBeVisible(); + } else { + expect(tab).not.toBeInTheDocument(); + } + } + expect( + screen.getByRole("button", { + name: userRole === "Admin" ? "+ Add New MCP Server" : "+ Submit MCP Server", + }), + ).toBeVisible(); + }); + + it.each(["cancel", "success", "failure", "unnamed"])("preserves delete confirmation on %s", async (outcome) => { + const server: MCPServer = { + created_at: "", + updated_at: "", + server_id: "delete-server", + server_name: outcome === "unnamed" ? null : "Delete fixture", + alias: "delete-alias", + url: outcome === "unnamed" ? null : "https://example.com/mcp", + created_by: "user", + updated_by: "user", + }; + vi.mocked(networking.fetchMCPServers).mockResolvedValue([server]); + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]); + let finishDelete: () => void = () => {}; + vi.mocked(networking.deleteMCPServer).mockImplementation( + () => + new Promise((resolve, reject) => { + finishDelete = () => (outcome === "failure" ? reject(new Error("Delete failed")) : resolve(undefined)); + }), + ); + render( + + + , + ); + await userEvent.click(await screen.findByRole("button", { name: "Server actions" })); + await userEvent.click(await screen.findByRole("menuitem", { name: "Delete" })); + const dialog = await screen.findByRole("alertdialog", { name: "Delete MCP Server?" }); + expect(within(dialog).getByText("delete-server")).toBeVisible(); + if (outcome === "unnamed") { + expect(within(dialog).queryByText("Name")).not.toBeInTheDocument(); + expect(within(dialog).queryByText("URL")).not.toBeInTheDocument(); + } else { + expect(within(dialog).getByText("Delete fixture")).toBeVisible(); + expect(within(dialog).getByText("https://example.com/mcp")).toBeVisible(); + } + if (outcome === "cancel") { + await userEvent.click(within(dialog).getByRole("button", { name: "Cancel" })); + expect(networking.deleteMCPServer).not.toHaveBeenCalled(); + } else { + await userEvent.click(within(dialog).getByRole("button", { name: "Delete" })); + expect(within(dialog).getByRole("button", { name: "Deleting..." })).toBeDisabled(); + expect(within(dialog).getByRole("button", { name: "Cancel" })).toBeDisabled(); + expect(networking.deleteMCPServer).toHaveBeenCalledWith("123", "delete-server"); + await act(async () => finishDelete()); + } + await waitFor(() => expect(screen.queryByRole("alertdialog")).not.toBeInTheDocument()); + }); + + it("filters servers by access group", async () => { + const server = { created_by: "user", updated_by: "user" }; + vi.mocked(networking.fetchMCPServers).mockResolvedValue([ + { + ...server, + server_id: "string-group", + server_name: "String group", + alias: "string-alias", + mcp_access_groups: ["shared"], + }, + { + ...server, + server_id: "legacy-group", + server_name: "Legacy group", + alias: "legacy-alias", + mcp_access_groups: ["shared"], + }, + { + ...server, + server_id: "other-group", + server_name: "Other group", + alias: "other-alias", + mcp_access_groups: ["different"], + }, + ]); + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]); + render( + + + , + ); + await screen.findByText("String group"); + await userEvent.click(screen.getByRole("combobox", { name: "Access Group" })); + await userEvent.click(await screen.findByRole("option", { name: "shared" })); + expect(screen.getByText("String group")).toBeVisible(); + expect(screen.getByText("Legacy group")).toBeVisible(); + expect(screen.queryByText("Other group")).not.toBeInTheDocument(); + }); + + it.each(["server_name", "alias", "url", "server_id"] as const)("searches by %s case-insensitively", async (field) => { + const server: MCPServer = { + created_at: "", + updated_at: "", + server_id: "search-server", + server_name: "Search fixture", + created_by: "user", + updated_by: "user", + [field]: "Needle", + }; + vi.mocked(networking.fetchMCPServers).mockResolvedValue([server]); + vi.mocked(networking.fetchMCPServerHealth).mockResolvedValue([]); + render( + + + , + ); + await screen.findByTestId("mcp-servers-grid"); + const search = screen.getByPlaceholderText("Search by name, alias, URL, or ID"); + await userEvent.type(search, " NEEDLE "); + expect(screen.getByTestId("mcp-servers-grid")).toBeVisible(); + await userEvent.clear(search); + await userEvent.type(search, "no-match"); + expect(screen.queryByTestId("mcp-servers-grid")).not.toBeInTheDocument(); + expect(screen.getByText("No servers match the current filters or search.")).toBeVisible(); }); it("should render mocked MCP servers data in the table", async () => { @@ -316,9 +537,7 @@ describe("MCPServers", () => { expect(screen.getByText("Team B Server")).toBeInTheDocument(); expect(screen.getByText("Team A Server 2")).toBeInTheDocument(); - // Find the team select by its "Team" label, then the combobox it labels - const teamLabel = screen.getByText("Team"); - const teamSelect = within(teamLabel.parentElement!).getByRole("combobox"); + const teamSelect = screen.getByRole("combobox", { name: "Team" }); await userEvent.click(teamSelect); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx index 2df304eff96..818b5150650 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/mcp-servers/_components/mcp_servers.tsx @@ -49,7 +49,7 @@ import { cn } from "@/lib/cva.config"; import UserEnvVarsModal from "./UserEnvVarsModal"; import { listMCPUserEnvVarStatus } from "@/components/networking"; -type SortKey = "created_desc" | "updated_desc" | "name_asc" | "health"; +export type SortKey = "created_desc" | "updated_desc" | "name_asc" | "health"; const SORT_OPTIONS: { value: SortKey; label: string }[] = [ { value: "created_desc", label: "Recently created" }, @@ -64,32 +64,33 @@ const HEALTH_RANK: Record = { healthy: 2, }; -const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => { +const compareByName = (a: MCPServer, b: MCPServer): number => { + const nameA = (a.server_name || a.alias || a.server_id).toLowerCase(); + const nameB = (b.server_name || b.alias || b.server_id).toLowerCase(); + return nameA.localeCompare(nameB) || a.server_id.localeCompare(b.server_id); +}; + +const compareByTimestampDesc = (a: string | null | undefined, b: string | null | undefined): number => { + const ta = a ? new Date(a).getTime() : 0; + const tb = b ? new Date(b).getTime() : 0; + return tb - ta; +}; + +export const compareServers = (a: MCPServer, b: MCPServer, sort: SortKey): number => { switch (sort) { - case "name_asc": { - const nameA = (a.server_name || a.alias || a.server_id).toLowerCase(); - const nameB = (b.server_name || b.alias || b.server_id).toLowerCase(); - return nameA.localeCompare(nameB); - } - case "updated_desc": { - const ta = a.updated_at ? new Date(a.updated_at).getTime() : 0; - const tb = b.updated_at ? new Date(b.updated_at).getTime() : 0; - return tb - ta; - } + case "name_asc": + return compareByName(a, b); + case "updated_desc": + return compareByTimestampDesc(a.updated_at, b.updated_at) || compareByName(a, b); case "health": { const ra = HEALTH_RANK[a.status ?? "unknown"] ?? 1; const rb = HEALTH_RANK[b.status ?? "unknown"] ?? 1; if (ra !== rb) return ra - rb; - const ta = a.created_at ? new Date(a.created_at).getTime() : 0; - const tb = b.created_at ? new Date(b.created_at).getTime() : 0; - return tb - ta; + return compareByTimestampDesc(a.created_at, b.created_at) || compareByName(a, b); } case "created_desc": - default: { - const ta = a.created_at ? new Date(a.created_at).getTime() : 0; - const tb = b.created_at ? new Date(b.created_at).getTime() : 0; - return tb - ta; - } + default: + return compareByTimestampDesc(a.created_at, b.created_at) || compareByName(a, b); } }; @@ -112,6 +113,62 @@ const readToolsOAuthServerId = (): string | null => { } }; +function DeleteServerDialog({ + open, + onOpenChange, + server, + isDeleting, + onConfirm, +}: { + open: boolean; + onOpenChange: (open: boolean) => void; + server: MCPServer | undefined; + isDeleting: boolean; + onConfirm: () => Promise; +}) { + return ( + + + + Delete MCP Server? + +
+

+ This action is permanent and cannot be undone. All associated configurations will be removed. +

+ + {server && ( +
+ {server.server_name && ( +
+
Name
+
{server.server_name}
+
+ )} +
+
ID
+
{server.server_id}
+
+ {server.url && ( +
+
URL
+
{server.url}
+
+ )} +
+ )} +
+ + Cancel + + + + + ); +} + const MCPServers: React.FC = ({ accessToken, userRole, userID, isViewOnly = false }) => { const { data: mcpServers, isLoading: isLoadingServers, refetch } = useMCPServers(); @@ -298,16 +355,12 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID, i } if (group !== "all") { filtered = filtered.filter((server) => - server.mcp_access_groups?.some((g: any) => (typeof g === "string" ? g === group : g && g.name === group)), + server.mcp_access_groups?.some((g: string | { name?: string } | null) => + typeof g === "string" ? g === group : g?.name === group, + ), ); } - const sorted = [...filtered].sort((a, b) => { - if (!a.created_at && !b.created_at) return 0; - if (!a.created_at) return 1; - if (!b.created_at) return -1; - return new Date(b.created_at).getTime() - new Date(a.created_at).getTime(); - }); - setFilteredServers(sorted); + setFilteredServers(filtered); }, [serversWithHealth], ); @@ -338,7 +391,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID, i const alias = (s.alias || "").toLowerCase(); const url = (s.url || "").toLowerCase(); const id = s.server_id.toLowerCase(); - return name.includes(q) || alias.includes(q) || url.includes(q) || id.includes(q); + return [name, alias, url, id].some((value) => value.includes(q)); }) : filteredServers; return [...matches].sort((a, b) => compareServers(a, b, sortKey)); @@ -381,9 +434,7 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID, i }; // Find the server to delete from the servers list - const serverToDelete = serverIdToDelete - ? (mcpServers || []).find((server) => server.server_id === serverIdToDelete) - : null; + const serverToDelete = mcpServers?.find((server) => server.server_id === serverIdToDelete); const handleCreateSuccess = (newMcpServer: MCPServer) => { setFilteredServers((prev) => [...prev, newMcpServer]); @@ -425,45 +476,13 @@ const MCPServers: React.FC = ({ accessToken, userRole, userID, i return (
- !open && cancelDelete()}> - - - Delete MCP Server? - -
-

- This action is permanent and cannot be undone. All associated configurations will be removed. -

- - {serverToDelete && ( -
- {serverToDelete.server_name && ( -
-
Name
-
{serverToDelete.server_name}
-
- )} -
-
ID
-
{serverToDelete.server_id}
-
- {serverToDelete.url && ( -
-
URL
-
{serverToDelete.url}
-
- )} -
- )} -
- - Cancel - - - - + !open && cancelDelete()} + server={serverToDelete} + isDeleting={isDeletingServer} + onConfirm={confirmDelete} + /> = ({ accessToken, userRole, userID, i My Connections - {isAdminRole(userRole) && ( + {isAdminRole(userRole) ? ( <> - )} - {!isAdminRole(userRole) && ( + ) : (