diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql new file mode 100644 index 00000000000..d594b0056df --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260919000000_add_autorouter_user_session_rollup/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260919000000_add_autorouter_user_session_rollup/migration.sql new file mode 100644 index 00000000000..2b864131ab2 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260919000000_add_autorouter_user_session_rollup/migration.sql @@ -0,0 +1,42 @@ +CREATE TABLE IF NOT EXISTS "LiteLLM_AutoRouterUserSession" ( + "user_id" TEXT NOT NULL, + "api_key" TEXT NOT NULL, + "session_id" TEXT NOT NULL, + "router_name" TEXT NOT NULL, + "router_type" TEXT NOT NULL, + "first_turn_at" TIMESTAMP(3) NOT NULL, + "last_turn_at" TIMESTAMP(3) NOT NULL, + "last_model" TEXT NOT NULL, + "models" JSONB NOT NULL DEFAULT '{}', + "turns" INTEGER NOT NULL DEFAULT 0, + "unordered_turns" INTEGER NOT NULL DEFAULT 0, + "covered_turns" INTEGER NOT NULL DEFAULT 0, + "cache_hits" INTEGER NOT NULL DEFAULT 0, + "same_model_turns" INTEGER NOT NULL DEFAULT 0, + "same_model_hits" INTEGER NOT NULL DEFAULT 0, + "first_visit_turns" INTEGER NOT NULL DEFAULT 0, + "first_visit_hits" INTEGER NOT NULL DEFAULT 0, + "return_turns" INTEGER NOT NULL DEFAULT 0, + "return_hits" INTEGER NOT NULL DEFAULT 0, + "return_expired_misses" INTEGER NOT NULL DEFAULT 0, + "return_within_ttl_misses" INTEGER NOT NULL DEFAULT 0, + "ttl_5m_turns" INTEGER NOT NULL DEFAULT 0, + "ttl_1h_turns" INTEGER NOT NULL DEFAULT 0, + "total_tokens" BIGINT NOT NULL DEFAULT 0, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0, + "saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0, + "savings_estimated_turns" INTEGER NOT NULL DEFAULT 0, + "savings_estimated_actual_spend" DOUBLE PRECISION NOT NULL DEFAULT 0, + "savings_estimated_saved_spend" DOUBLE PRECISION NOT NULL DEFAULT 0, + "savings_estimated_baseline_models" JSONB NOT NULL DEFAULT '{}', + "classifier_cost" DOUBLE PRECISION NOT NULL DEFAULT 0, + "classifier_cost_recorded_turns" INTEGER NOT NULL DEFAULT 0, + "tier_turns" JSONB NOT NULL DEFAULT '{}', + "baseline_models" JSONB NOT NULL DEFAULT '{}', + + CONSTRAINT "LiteLLM_AutoRouterUserSession_pkey" PRIMARY KEY ("user_id", "api_key", "session_id", "router_name") +); + +CREATE INDEX IF NOT EXISTS "idx_autorouter_user_session_last_turn" ON "LiteLLM_AutoRouterUserSession"("last_turn_at"); + +CREATE INDEX IF NOT EXISTS "idx_autorouter_user_session_user_last_turn" ON "LiteLLM_AutoRouterUserSession"("user_id", "last_turn_at"); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 368864f9cd1..85996430bc5 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -73,6 +73,7 @@ model LiteLLM_AgentsTable { static_headers Json? @default("{}") extra_headers String[] @default([]) agent_access_groups String[] @default([]) + access_group_ids String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) @@ -1623,6 +1624,47 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } +model LiteLLM_AutoRouterUserSession { + user_id String + api_key String + session_id String + router_name String + router_type String + first_turn_at DateTime + last_turn_at DateTime + last_model String + models Json @default("{}") + turns Int @default(0) + unordered_turns Int @default(0) + covered_turns Int @default(0) + cache_hits Int @default(0) + same_model_turns Int @default(0) + same_model_hits Int @default(0) + first_visit_turns Int @default(0) + first_visit_hits Int @default(0) + return_turns Int @default(0) + return_hits Int @default(0) + return_expired_misses Int @default(0) + return_within_ttl_misses Int @default(0) + ttl_5m_turns Int @default(0) + ttl_1h_turns Int @default(0) + total_tokens BigInt @default(0) + spend Float @default(0) + saved_spend Float @default(0) + savings_estimated_turns Int @default(0) + savings_estimated_actual_spend Float @default(0) + savings_estimated_saved_spend Float @default(0) + savings_estimated_baseline_models Json @default("{}") + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) + tier_turns Json @default("{}") + baseline_models Json @default("{}") + + @@id([user_id, api_key, session_id, router_name]) + @@index([last_turn_at], map: "idx_autorouter_user_session_last_turn") + @@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn") +} + // Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in // either direction. forward duplicates the requests the keys did not route through the // router through it, answering whether they should adopt it; reverse duplicates the diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index de234fa3afa..66270501312 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2741,6 +2741,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +dependencies = [ + "litellm-cache", + "litellm-cache-redis", + "litellm-cache-response", + "r2d2", + "redis", + "redis-test", + "serde_json", + "sha2 0.10.9", + "tokio", +] + [[package]] name = "litellm-cache-response" version = "0.1.0" @@ -2772,6 +2787,22 @@ dependencies = [ "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" @@ -2942,8 +2973,10 @@ dependencies = [ "litellm-cache-gcs", "litellm-cache-memory", "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", @@ -2954,10 +2987,12 @@ dependencies = [ "litellm-types", "pyo3", "pyo3-async-runtimes", + "redis", "rstest", "serde", "serde_json", "serde_with", + "sha2 0.10.9", "tokio", "tokio-tungstenite", ] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index be86240ac46..4e370c13766 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -36,6 +36,7 @@ litellm-cache-redis = { path = "crates/cache-redis" } litellm-cache-s3 = { path = "crates/cache-s3" } litellm-cache-gcs = { path = "crates/cache-gcs" } litellm-cache-disk = { path = "crates/cache-disk" } +litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" } litellm-cache-response = { path = "crates/cache-response" } litellm-token-counter = { path = "crates/token-counter" } litellm-token-counter-fast = { path = "crates/token-counter-fast" } diff --git a/litellm-rust/crates/cache-redis-semantic/Cargo.toml b/litellm-rust/crates/cache-redis-semantic/Cargo.toml new file mode 100644 index 00000000000..9a8755a189e --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "litellm-cache-redis-semantic" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +litellm-cache.workspace = true +litellm-cache-redis.workspace = true +litellm-cache-response.workspace = true +redis = { version = "1.7.0", features = ["tls-rustls"] } +r2d2 = "0.8.10" +serde_json.workspace = true +sha2.workspace = true +tokio.workspace = true + +[dev-dependencies] +redis-test = "1.0.4" +serde_json.workspace = true +tokio.workspace = true diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs new file mode 100644 index 00000000000..e0ac31f3630 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -0,0 +1,618 @@ +use std::{ + future::Future, + sync::{Arc, OnceLock}, + time::{Duration, SystemTime, UNIX_EPOCH}, +}; + +use litellm_cache::{ + BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error, + SemanticCacheContext, +}; +use litellm_cache_redis::{ + RedisTopology, + connection::{ConnectionRef, Connections}, +}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use serde_json::Value; +use sha2::{Digest, Sha256}; + +use crate::prompt::prompt_from_context; + +const CACHE_KEY_FIELD: &str = "litellm_cache_key"; +const VECTOR_FIELD: &str = "prompt_vector"; + +pub trait Embedder: Send + Sync + 'static { + fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error>; + + fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> impl Future, Error>> + Send; +} + +#[derive(Clone, Debug)] +pub struct RedisSemanticConfig { + pub index_name: String, + pub similarity_threshold: f32, +} + +struct Inner { + index_name: String, + distance_threshold: f64, + resolved_index: OnceLock, + codec: ResponseCacheCodec, + clock: fn() -> f64, +} + +impl Inner { + fn new(config: RedisSemanticConfig) -> Self { + Self { + index_name: config.index_name, + distance_threshold: 1.0 - f64::from(config.similarity_threshold), + resolved_index: OnceLock::new(), + codec: ResponseCacheCodec, + clock: timestamp, + } + } + + fn ensure_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + if let Some(name) = self.resolved_index.get() { + return Ok(name.clone()); + } + let name = match index_compatible(connection, &self.index_name, dims)? { + Some(true) => self.index_name.clone(), + Some(false) => self.isolated_index(connection, dims)?, + None => match create_index(connection, &self.index_name, dims) { + Ok(()) => self.index_name.clone(), + Err(_) => match index_compatible(connection, &self.index_name, dims)? { + Some(true) => self.index_name.clone(), + Some(false) => self.isolated_index(connection, dims)?, + None => return Err(Error::Unavailable), + }, + }, + }; + let _ = self.resolved_index.set(name.clone()); + Ok(name) + } + + fn isolated_index( + &self, + connection: &mut ConnectionRef<'_>, + dims: usize, + ) -> Result { + let name = format!("{}_isolated", self.index_name); + match index_compatible(connection, &name, dims)? { + Some(true) => Ok(name), + Some(false) => { + redis::cmd("FT.DROPINDEX") + .arg(&name) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + create_index(connection, &name, dims)?; + Ok(name) + } + None => { + create_index(connection, &name, dims)?; + Ok(name) + } + } + } + + fn store( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + value: &CacheEntry, + prompt: &str, + vector: &[f32], + ttl: Option, + ) -> Result<(), Error> { + let index = self.ensure_index(connection, vector.len())?; + let entry_id = entry_id(prompt, tag); + let hash_key = format!("{index}:{entry_id}"); + let response = self.codec.encode(value)?; + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(&entry_id) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(response) + .arg(VECTOR_FIELD) + .arg(vector_buffer(vector)) + .arg("inserted_at") + .arg(format!("{}", (self.clock)())) + .arg("updated_at") + .arg(format!("{}", (self.clock)())) + .arg(CACHE_KEY_FIELD) + .arg(tag) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + if let Some(ttl) = ttl { + redis::cmd("EXPIRE") + .arg(&hash_key) + .arg(ttl_seconds(ttl)) + .query::<()>(connection) + .map_err(|_| Error::Unavailable)?; + } + Ok(()) + } + + fn lookup( + &self, + connection: &mut ConnectionRef<'_>, + tag: &str, + vector: &[f32], + ) -> Result, Error> { + let index = self.ensure_index(connection, vector.len())?; + let query = format!( + "(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]", + escape_tag(tag) + ); + let result = redis::cmd("FT.SEARCH") + .arg(&index) + .arg(query) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg(CACHE_KEY_FIELD) + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_buffer(vector)) + .query::(connection) + .map_err(|_| Error::Unavailable)?; + let Some(fields) = first_document(&result) else { + return Ok(None); + }; + if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) { + return Ok(None); + } + if number_field(fields, "vector_distance") + .is_none_or(|distance| distance > self.distance_threshold) + { + return Ok(None); + } + let Some(response) = bytes_field(fields, "response") else { + return Ok(None); + }; + self.codec.decode(&response).map(Some) + } +} + +pub struct RedisSemanticCache { + connections: Arc>, + embedder: E, + inner: Arc, +} + +impl RedisSemanticCache { + pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result { + Ok(Self { + connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), + embedder, + inner: Arc::new(Inner::new(config)), + }) + } +} + +impl RedisSemanticCache { + pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self { + Self { + connections: Arc::new(Connections::fixed(connection)), + embedder, + inner: Arc::new(Inner::new(config)), + } + } + + pub fn with_clock(self, clock: fn() -> f64) -> Self { + Self { + inner: Arc::new(Inner { + index_name: self.inner.index_name.clone(), + distance_threshold: self.inner.distance_threshold, + resolved_index: OnceLock::new(), + codec: self.inner.codec, + clock, + }), + ..self + } + } + + pub fn embedder(&self) -> &E { + &self.embedder + } + + pub fn index_name(&self) -> &str { + &self.inner.index_name + } + + pub fn similarity_threshold(&self) -> f32 { + (1.0 - self.inner.distance_threshold) as f32 + } + + fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str { + context.scope.as_deref().unwrap_or(key) + } +} + +impl BaseCache + for RedisSemanticCache +{ + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(()); + }; + let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; + let tag = Self::tag(key, context).to_string(); + self.connections.execute(|connection| { + self.inner + .store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?; + let tag = Self::tag(key, context).to_string(); + self.connections + .execute(|connection| self.inner.lookup(connection, &tag, &vector)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> Result<(), Error> { + let Some(prompt) = prompt_from_context(&context) else { + return Ok(()); + }; + let vector = self + .embedder + .async_embed(&prompt, context.metadata.as_ref()) + .await?; + let tag = Self::tag(key, &context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.store(connection, &tag, &value, &prompt, &vector, context.ttl) + }) + .await + } + + async fn async_get_cache( + &self, + key: &str, + context: &Self::Context, + ) -> Result, Error> { + let Some(prompt) = prompt_from_context(context) else { + return Ok(None); + }; + let vector = self + .embedder + .async_embed(&prompt, context.metadata.as_ref()) + .await?; + let tag = Self::tag(key, context).to_string(); + let inner = Arc::clone(&self.inner); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + inner.lookup(connection, &tag, &vector) + }) + .await + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + match Connections::run_blocking(Arc::clone(&self.connections), |connection| { + Ok(match redis::cmd("PING").query::(connection) { + Ok(_) => CacheConnectionResult { + status: CacheConnectionStatus::Success, + message: "Redis cache connection test successful".into(), + error: None, + }, + Err(error) => CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }, + }) + }) + .await + { + Ok(result) => Ok(result), + Err(error) => Ok(CacheConnectionResult { + status: CacheConnectionStatus::Failed, + message: format!("Redis connection failed: {error}"), + error: Some(error.to_string()), + }), + } + } +} + +fn timestamp() -> f64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|duration| duration.as_secs_f64()) + .unwrap_or_default() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(CACHE_KEY_FIELD.as_bytes()); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn vector_buffer(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn escape_tag(value: &str) -> String { + value + .chars() + .flat_map(|ch| { + if matches!( + ch, + ',' | '.' + | '<' + | '>' + | '{' + | '}' + | '[' + | ']' + | '\\' + | '"' + | '\'' + | ':' + | ';' + | '!' + | '@' + | '#' + | '$' + | '%' + | '^' + | '&' + | '*' + | '(' + | ')' + | '-' + | '+' + | '=' + | '~' + | '|' + | '/' + | ' ' + | '?' + ) { + vec!['\\', ch] + } else { + vec![ch] + } + }) + .collect() +} + +fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> { + redis::cmd("FT.CREATE") + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg(VECTOR_FIELD) + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg(CACHE_KEY_FIELD) + .arg("TAG") + .arg("SEPARATOR") + .arg(",") + .query::<()>(connection) + .map_err(|_| Error::Unavailable) +} + +fn index_compatible( + connection: &mut ConnectionRef<'_>, + name: &str, + dims: usize, +) -> Result, Error> { + let info = match redis::cmd("FT.INFO") + .arg(name) + .query::(connection) + { + Ok(info) => info, + Err(error) if unknown_index(&error) => return Ok(None), + Err(_) => return Err(Error::Unavailable), + }; + Ok(Some(schema_compatible(&info, dims))) +} + +fn unknown_index(error: &redis::RedisError) -> bool { + let message = error.to_string().to_lowercase(); + message.contains("unknown") && message.contains("index") +} + +fn schema_compatible(info: &redis::Value, dims: usize) -> bool { + let redis::Value::Array(entries) = info else { + return false; + }; + let attributes = entries + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some("attributes")) + .map(|pair| &pair[1]); + let Some(redis::Value::Array(attributes)) = attributes else { + return false; + }; + let fields = attributes + .iter() + .map(|attribute| { + let redis::Value::Array(attribute) = attribute else { + return (None, None, None, None, None); + }; + let mut name = None; + let mut field_type = None; + let mut dim = None; + let mut data_type = None; + let mut distance_metric = None; + for pair in attribute.as_chunks::<2>().0 { + match string_value(&pair[0]).as_deref() { + Some("identifier") => name = string_value(&pair[1]), + Some("type") => field_type = string_value(&pair[1]), + Some("dim") => dim = number_value(&pair[1]), + Some("data_type") => data_type = string_value(&pair[1]), + Some("distance_metric") => distance_metric = string_value(&pair[1]), + _ => {} + } + } + (name, field_type, dim, data_type, distance_metric) + }) + .collect::>(); + let has_field = |name: &str, field_type: &str| { + fields + .iter() + .any(|(n, t, ..)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type)) + }; + has_field("prompt", "TEXT") + && has_field("response", "TEXT") + && has_field("inserted_at", "NUMERIC") + && has_field("updated_at", "NUMERIC") + && has_field(CACHE_KEY_FIELD, "TAG") + && fields.iter().any(|(n, t, d, data, metric)| { + n.as_deref() == Some(VECTOR_FIELD) + && t.as_deref() == Some("VECTOR") + && *d == Some(dims as f64) + && data + .as_deref() + .is_some_and(|data| data.eq_ignore_ascii_case("float32")) + && metric + .as_deref() + .is_some_and(|metric| metric.eq_ignore_ascii_case("cosine")) + }) +} + +fn string_value(value: &redis::Value) -> Option { + match value { + redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(), + redis::Value::SimpleString(text) => Some(text.clone()), + redis::Value::VerbatimString { text, .. } => Some(text.clone()), + _ => None, + } +} + +fn number_value(value: &redis::Value) -> Option { + match value { + redis::Value::Int(number) => Some(*number as f64), + redis::Value::Double(number) => Some(*number), + _ => string_value(value).and_then(|text| text.parse().ok()), + } +} + +fn first_document(result: &redis::Value) -> Option<&[redis::Value]> { + let redis::Value::Array(items) = result else { + return None; + }; + let [count, _document_id, fields, ..] = items.as_slice() else { + return None; + }; + if !matches!(count, redis::Value::Int(count) if *count > 0) { + return None; + } + match fields { + redis::Value::Array(fields) => Some(fields.as_slice()), + _ => None, + } +} + +fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> { + fields + .as_chunks::<2>() + .0 + .iter() + .find(|pair| string_value(&pair[0]).as_deref() == Some(name)) + .map(|pair| &pair[1]) +} + +fn string_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(string_value) +} + +fn number_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(number_value) +} + +fn bytes_field(fields: &[redis::Value], name: &str) -> Option> { + match field_value(fields, name)? { + redis::Value::BulkString(bytes) => Some(bytes.clone()), + redis::Value::SimpleString(text) => Some(text.clone().into_bytes()), + _ => None, + } +} + +fn ttl_seconds(ttl: Duration) -> u64 { + ttl.as_secs() + .saturating_add(u64::from(ttl.subsec_nanos() > 0)) + .max(1) +} diff --git a/litellm-rust/crates/cache-redis-semantic/src/lib.rs b/litellm-rust/crates/cache-redis-semantic/src/lib.rs new file mode 100644 index 00000000000..51d0b4ba5f3 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/lib.rs @@ -0,0 +1,5 @@ +mod cache; +mod prompt; + +pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig}; +pub use prompt::prompt_from_context; diff --git a/litellm-rust/crates/cache-redis-semantic/src/prompt.rs b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs new file mode 100644 index 00000000000..b9c38e98d77 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs @@ -0,0 +1,97 @@ +use litellm_cache::SemanticCacheContext; +use serde_json::Value; + +pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { + if let Some(messages) = context.messages.as_ref().and_then(Value::as_array) + && !messages.is_empty() + { + return Some(messages_text(messages)); + } + let input = context.input.as_ref()?; + let mut parts = Vec::new(); + collect_input_text(input, &mut parts); + let prompt = parts.join("\n").trim().to_string(); + (!prompt.is_empty()).then_some(prompt) +} + +fn messages_text(messages: &[Value]) -> String { + let mut text = String::new(); + for message in messages { + let Some(message) = message.as_object() else { + continue; + }; + match message.get("content") { + Some(Value::String(content)) => text.push_str(content), + Some(Value::Array(parts)) => { + for part in parts { + if let Some(text_content) = part.get("text").and_then(Value::as_str) { + text.push_str(text_content); + } + } + } + _ => {} + } + text.push_str(&search_results_text(message.get("search_results"))); + } + text +} + +fn search_results_text(search_results: Option<&Value>) -> String { + let Some(Value::Array(results)) = search_results else { + return String::new(); + }; + let mut text = String::new(); + for result in results { + let Some(result) = result.as_object() else { + continue; + }; + for key in ["source", "title"] { + if let Some(value) = result.get(key).and_then(Value::as_str) { + text.push_str(value); + } + } + if let Some(Value::Array(content)) = result.get("content") { + for block in content { + if let Some(value) = block.get("text").and_then(Value::as_str) { + text.push_str(value); + } + } + } + if let Some(citations) = result.get("citations") { + text.push_str(&citations.to_string()); + } + } + text +} + +fn collect_input_text(value: &Value, parts: &mut Vec) { + match value { + Value::String(text) => { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + } + } + Value::Array(items) => { + for item in items { + collect_input_text(item, parts); + } + } + Value::Object(map) => { + if let Some(content) = map.get("content").filter(|content| !content.is_null()) { + collect_input_text(content, parts); + return; + } + for key in ["text", "output", "input_text", "output_text"] { + if let Some(Value::String(text)) = map.get(key) { + let trimmed = text.trim(); + if !trimmed.is_empty() { + parts.push(trimmed.to_string()); + return; + } + } + } + } + _ => {} + } +} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs new file mode 100644 index 00000000000..233b87ec52f --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -0,0 +1,1003 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::{BaseCache, CacheCodec, Error, SemanticCacheContext}; +use litellm_cache_redis_semantic::{Embedder, RedisSemanticCache, RedisSemanticConfig}; +use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use redis_test::{MockCmd, MockRedisConnection}; +use serde_json::{Value, json}; +use sha2::{Digest, Sha256}; + +const INDEX: &str = "litellm_semantic_cache_index"; + +struct FakeEmbedder { + vectors: HashMap>, + calls: Arc>>, +} + +impl FakeEmbedder { + fn new(vectors: &[(&str, &[f32])]) -> (Self, Arc>>) { + let calls = Arc::new(Mutex::new(Vec::new())); + ( + Self { + vectors: vectors + .iter() + .map(|(prompt, vector)| (prompt.to_string(), vector.to_vec())) + .collect(), + calls: Arc::clone(&calls), + }, + calls, + ) + } +} + +impl Embedder for FakeEmbedder { + fn embed(&self, prompt: &str, _: Option<&Value>) -> Result, Error> { + self.calls.lock().unwrap().push(prompt.to_string()); + + Ok(self + .vectors + .get(prompt) + .cloned() + .unwrap_or_else(|| vec![0.1, 0.2, 0.3])) + } + + async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { + self.embed(prompt, metadata) + } +} + +fn config() -> RedisSemanticConfig { + RedisSemanticConfig { + index_name: INDEX.into(), + similarity_threshold: 0.9, + } +} + +fn messages_context(messages: Vec) -> SemanticCacheContext { + SemanticCacheContext { + messages: Some(Value::Array(messages)), + ..Default::default() + } +} + +fn entry() -> CacheEntry { + CacheEntry { + timestamp: Some(1.0), + response: json!({"answer": "yes"}), + } +} + +fn encoded(entry: &CacheEntry) -> Vec { + ResponseCacheCodec.encode(entry).unwrap() +} + +fn vector_bytes(vector: &[f32]) -> Vec { + vector + .iter() + .flat_map(|component| component.to_le_bytes()) + .collect() +} + +fn entry_id(prompt: &str, tag: &str) -> String { + let mut digest = Sha256::new(); + digest.update(prompt.as_bytes()); + digest.update(b"litellm_cache_key"); + digest.update(tag.as_bytes()); + format!("{:x}", digest.finalize()) +} + +fn s(value: &str) -> redis::Value { + redis::Value::BulkString(value.as_bytes().to_vec()) +} + +fn unknown_index_error() -> redis::RedisError { + redis::RedisError::from((redis::ErrorKind::Extension, "Unknown index name")) +} + +fn attribute(name: &str, field_type: &str, extra: Vec) -> redis::Value { + let mut parts = vec![ + s("identifier"), + s(name), + s("attribute"), + s(name), + s("type"), + s(field_type), + ]; + parts.extend(extra); + redis::Value::Array(parts) +} + +fn index_info(attributes: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("index_name"), + s(INDEX), + s("attributes"), + redis::Value::Array(attributes), + ]) +} + +fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> redis::Value { + attribute( + "prompt_vector", + "VECTOR", + vec![ + s("algorithm"), + s("FLAT"), + s("data_type"), + s(data_type), + s("dim"), + redis::Value::Int(dims), + s("distance_metric"), + s(distance_metric), + ], + ) +} + +fn vector_attribute(dims: i64) -> redis::Value { + vector_attribute_with(dims, "FLOAT32", "COSINE") +} + +fn info_with_vector(vector: redis::Value) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector, + attribute("litellm_cache_key", "TAG", vec![]), + ]) +} + +fn compatible_info(dims: i64) -> redis::Value { + info_with_vector(vector_attribute(dims)) +} + +fn unscoped_info(dims: i64) -> redis::Value { + index_info(vec![ + attribute("prompt", "TEXT", vec![]), + attribute("response", "TEXT", vec![]), + attribute("inserted_at", "NUMERIC", vec![]), + attribute("updated_at", "NUMERIC", vec![]), + vector_attribute(dims), + ]) +} + +fn create_index_command(name: &str, dims: usize) -> redis::Cmd { + let mut command = redis::cmd("FT.CREATE"); + command + .arg(name) + .arg("ON") + .arg("HASH") + .arg("PREFIX") + .arg(1) + .arg(name) + .arg("SCORE") + .arg(1.0) + .arg("SCHEMA") + .arg("prompt") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("response") + .arg("TEXT") + .arg("WEIGHT") + .arg(1) + .arg("inserted_at") + .arg("NUMERIC") + .arg("updated_at") + .arg("NUMERIC") + .arg("prompt_vector") + .arg("VECTOR") + .arg("FLAT") + .arg(6) + .arg("TYPE") + .arg("FLOAT32") + .arg("DIM") + .arg(dims) + .arg("DISTANCE_METRIC") + .arg("COSINE") + .arg("litellm_cache_key") + .arg("TAG") + .arg("SEPARATOR") + .arg(","); + command +} + +fn search_command(index: &str, tag: &str, vector: &[f32]) -> redis::Cmd { + let mut command = redis::cmd("FT.SEARCH"); + command + .arg(index) + .arg(format!( + "(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]" + )) + .arg("RETURN") + .arg(8) + .arg("entry_id") + .arg("prompt") + .arg("response") + .arg("inserted_at") + .arg("updated_at") + .arg("metadata") + .arg("litellm_cache_key") + .arg("vector_distance") + .arg("SORTBY") + .arg("vector_distance") + .arg("ASC") + .arg("DIALECT") + .arg(2) + .arg("LIMIT") + .arg(0) + .arg(1) + .arg("PARAMS") + .arg(2) + .arg("vector") + .arg(vector_bytes(vector)); + command +} + +fn hit_fields(tag: &str, distance: &str, response: Vec) -> redis::Value { + redis::Value::Array(vec![ + s("entry_id"), + s("stored-id"), + s("prompt"), + s("hello prompt"), + s("response"), + redis::Value::BulkString(response), + s("inserted_at"), + s("1700000000.5"), + s("updated_at"), + s("1700000000.5"), + s("litellm_cache_key"), + s(tag), + s("vector_distance"), + s(distance), + ]) +} + +fn search_result(fields: redis::Value) -> redis::Value { + redis::Value::Array(vec![ + redis::Value::Int(1), + s("litellm_semantic_cache_index:stored-id"), + fields, + ]) +} + +fn empty_result() -> redis::Value { + redis::Value::Array(vec![redis::Value::Int(0)]) +} + +#[test] +fn store_creates_index_and_writes_hash_with_expire() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(INDEX, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(5), Ok(1)), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(5)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + cache.set_cache(tag, value, &context).unwrap(); +} + +#[test] +fn store_without_ttl_skips_expire() { + let prompt = "hello prompt"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "key1"))) + .arg("entry_id") + .arg(entry_id(prompt, "key1")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("key1"), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + "key1", + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn lookup_returns_hit_below_distance_threshold() { + let vector = vec![0.1f32, 0.2, 0.3]; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let hit = cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]), + ) + .unwrap(); + assert_eq!(hit, Some(value)); +} + +#[test] +fn lookup_misses_above_distance_threshold_and_on_tag_mismatch() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields("key1", "0.5", encoded(&entry())))), + ), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "other", + "0.05", + encoded(&entry()), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + let context = messages_context(vec![json!({"role": "user", "content": "hello prompt"})]); + + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); +} + +#[test] +fn lookup_returns_invalid_entry_on_malformed_response() { + let vector = vec![0.1f32, 0.2, 0.3]; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(search_result(hit_fields( + "key1", + "0.05", + b"not json!".to_vec(), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + "key1", + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap_err(), + Error::InvalidEntry + ); +} + +#[test] +fn missing_prompt_is_noop_and_never_embeds() { + let connection = MockRedisConnection::new(Vec::::new()).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + let context = SemanticCacheContext::default(); + cache.set_cache("key1", entry(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + assert!(calls.lock().unwrap().is_empty()); +} + +#[test] +fn scope_overrides_key_as_filter_tag() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, "scope-a"))) + .arg("entry_id") + .arg(entry_id(prompt, "scope-a")) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg("scope-a"), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, "scope\\-a", &vector), + Ok(search_result(hit_fields( + "scope-a", + "0.05", + encoded(&value), + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = SemanticCacheContext { + scope: Some("scope-a".into()), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + + cache.set_cache("key1", value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(value)); +} + +#[test] +fn incompatible_schema_falls_back_to_isolated_index() { + let prompt = "hello prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(unscoped_info(3))), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{isolated}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn create_index_race_rechecks_schema_and_stores() { + let prompt = "hello prompt"; + let tag = "key1"; + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new( + create_index_command(INDEX, 3), + Err::<&str, _>(redis::RedisError::from(( + redis::ErrorKind::Extension, + "Index already exists", + ))), + ), + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{INDEX}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn wrong_distance_metric_falls_back_to_isolated_index() { + let prompt = "hello prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let connection = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Ok(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2"))), + ), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{isolated}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&entry())) + .arg("prompt_vector") + .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + + cache + .set_cache( + tag, + entry(), + &messages_context(vec![json!({"role": "user", "content": prompt})]), + ) + .unwrap(); +} + +#[test] +fn tag_special_characters_are_escaped_in_search_filter() { + let vector = vec![0.1f32, 0.2, 0.3]; + let tag = "a:b, c|d"; + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + search_command(INDEX, "a\\:b\\,\\ c\\|d", &vector), + Ok(empty_result()), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + assert_eq!( + cache + .get_cache( + tag, + &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) + ) + .unwrap(), + None + ); +} + +#[test] +fn prompt_extraction_matches_python_message_and_input_shapes() { + let vector = vec![0.1f32, 0.2, 0.3]; + let lookups = 5; + let mut commands = vec![MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Ok(compatible_info(3)), + )]; + for _ in 0..lookups { + commands.push(MockCmd::new( + search_command(INDEX, "key1", &vector), + Ok(empty_result()), + )); + } + let connection = MockRedisConnection::new(commands).assert_all_commands_consumed(); + let (embedder, calls) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()); + + cache + .get_cache( + "key1", + &messages_context(vec![ + json!({"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]}), + json!({"role": "assistant", "content": "reply"}), + ]), + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!(" plain input ")), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some( + json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"]), + ), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &SemanticCacheContext { + input: Some(json!({"output_text": " result text "})), + ..Default::default() + }, + ) + .unwrap(); + cache + .get_cache( + "key1", + &messages_context(vec![json!({ + "role": "user", + "content": "question", + "search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}], + })]), + ) + .unwrap(); + + assert_eq!( + *calls.lock().unwrap(), + vec![ + "firstsecondreply", + "plain input", + "nested\ntail", + "result text", + "questionsrctfound{\"a\":1}", + ] + ); +} + +#[test] +fn ttl_passes_through_context_only() { + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection( + MockRedisConnection::new(Vec::::new()), + embedder, + config(), + ); + assert_eq!(cache.get_ttl(&SemanticCacheContext::default()), None); + assert_eq!( + cache.get_ttl(&SemanticCacheContext { + ttl: Some(Duration::from_secs(9)), + ..Default::default() + }), + Some(Duration::from_secs(9)) + ); +} + +#[tokio::test] +async fn async_paths_embed_then_run_blocking_redis_work() { + let vector = vec![0.1f32, 0.2, 0.3]; + let prompt = "hello prompt"; + let tag = "key1"; + let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); + let value = entry(); + let connection = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), + MockCmd::new( + redis::cmd("HSET") + .arg(&hash_key) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(&vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ), + MockCmd::new( + search_command(INDEX, tag, &vector), + Ok(search_result(hit_fields(tag, "0.05", encoded(&value)))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder, _) = FakeEmbedder::new(&[]); + let cache = RedisSemanticCache::with_connection(connection, embedder, config()) + .with_clock(|| 1700000000.5); + let context = messages_context(vec![json!({"role": "user", "content": prompt})]); + + cache + .async_set_cache(tag, value.clone(), context.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache(tag, &context).await.unwrap(), + Some(value) + ); +} + +#[test] +fn shared_base_index_across_dimensions_replaces_the_isolated_index() { + // Pins parity with Python's `_isolated` + overwrite=True flow. + let prompt = "shared prompt"; + let tag = "key1"; + let isolated = format!("{INDEX}_isolated"); + let value = entry(); + let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); + let store_hash = |index: &str, vector: &[f32]| { + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{index}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(&value)) + .arg("prompt_vector") + .arg(vector_bytes(vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ) + }; + + let vector_a = vec![0.1f32; 8]; + let connection_a = MockRedisConnection::new([ + MockCmd::new( + redis::cmd("FT.INFO").arg(INDEX), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(INDEX, 8), Ok("OK")), + store_hash(INDEX, &vector_a), + ]) + .assert_all_commands_consumed(); + let (embedder_a, _) = FakeEmbedder::new(&[(prompt, &vector_a)]); + let worker_a = RedisSemanticCache::with_connection(connection_a, embedder_a, config()) + .with_clock(|| 1700000000.5); + worker_a.set_cache(tag, value.clone(), &context()).unwrap(); + + let vector_b = vec![0.2f32; 4]; + let connection_b = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(8))), + MockCmd::new( + redis::cmd("FT.INFO").arg(&isolated), + Err::(unknown_index_error()), + ), + MockCmd::new(create_index_command(&isolated, 4), Ok("OK")), + store_hash(&isolated, &vector_b), + MockCmd::new( + search_command(&isolated, tag, &vector_b), + Ok(search_result(hit_fields(tag, "0.0", encoded(&value)))), + ), + MockCmd::new( + search_command(&isolated, tag, &vector_b), + Err::(redis::RedisError::from(( + redis::ErrorKind::Extension, + "Vector dimension mismatch", + ))), + ), + ]) + .assert_all_commands_consumed(); + let (embedder_b, _) = FakeEmbedder::new(&[(prompt, &vector_b)]); + let worker_b = RedisSemanticCache::with_connection(connection_b, embedder_b, config()) + .with_clock(|| 1700000000.5); + worker_b.set_cache(tag, value.clone(), &context()).unwrap(); + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap(), + Some(value.clone()) + ); + + let vector_c = vec![0.3f32; 16]; + let connection_c = MockRedisConnection::new([ + MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(8))), + MockCmd::new(redis::cmd("FT.INFO").arg(&isolated), Ok(compatible_info(4))), + MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")), + MockCmd::new(create_index_command(&isolated, 16), Ok("OK")), + store_hash(&isolated, &vector_c), + ]) + .assert_all_commands_consumed(); + let (embedder_c, _) = FakeEmbedder::new(&[(prompt, &vector_c)]); + let worker_c = RedisSemanticCache::with_connection(connection_c, embedder_c, config()) + .with_clock(|| 1700000000.5); + worker_c.set_cache(tag, value.clone(), &context()).unwrap(); + + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap_err(), + Error::Unavailable + ); +} + +#[test] +fn live_shared_index_is_replaced_across_dimensions() { + let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { + return; + }; + // Pins parity with Python's `_isolated` + overwrite=True flow. + let base = format!("rust_semantic_shared_{}", std::process::id()); + let isolated = format!("{base}_isolated"); + let prompt = "shared live prompt"; + let tag = "key1"; + let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); + let value = entry(); + let worker = |vector: Vec| { + let (embedder, _) = FakeEmbedder::new(&[(prompt, vector.as_slice())]); + RedisSemanticCache::new( + &url, + embedder, + RedisSemanticConfig { + index_name: base.clone(), + similarity_threshold: 0.9, + }, + ) + .unwrap() + }; + + let worker_a = worker(vec![0.1f32; 8]); + worker_a.set_cache(tag, value.clone(), &context()).unwrap(); + + let worker_b = worker(vec![0.2f32; 4]); + worker_b.set_cache(tag, value.clone(), &context()).unwrap(); + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap(), + Some(value.clone()) + ); + + let worker_c = worker(vec![0.3f32; 16]); + worker_c.set_cache(tag, value.clone(), &context()).unwrap(); + + assert_eq!( + worker_b.get_cache(tag, &context()).unwrap_err(), + Error::Unavailable + ); + + let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); + for index in [&base, &isolated] { + let _: Result<(), _> = redis::cmd("FT.DROPINDEX") + .arg(index) + .arg("DD") + .query(&mut connection); + } +} + +#[test] +fn live_store_lookup_and_ttl_against_redis_stack() { + let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { + return; + }; + let vector = vec![0.1f32, 0.2, 0.3, 0.4]; + let prompt = "rust semantic cache live prompt"; + let tag = "live-key"; + let index_name = format!("rust_semantic_test_{}", std::process::id()); + let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); + let cache = RedisSemanticCache::new( + &url, + embedder, + RedisSemanticConfig { + index_name: index_name.clone(), + similarity_threshold: 0.9, + }, + ) + .unwrap(); + let context = SemanticCacheContext { + ttl: Some(Duration::from_secs(120)), + ..messages_context(vec![json!({"role": "user", "content": prompt})]) + }; + let value = entry(); + + cache.set_cache(tag, value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache(tag, &context).unwrap(), Some(value)); + assert_eq!(cache.get_cache("other-key", &context).unwrap(), None); + + let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); + let ttl: i64 = redis::Commands::ttl( + &mut connection, + format!("{index_name}:{}", entry_id(prompt, tag)), + ) + .unwrap(); + assert!( + ttl > 0, + "expected stored hash to carry an expiry, got {ttl}" + ); +} diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index e2e2656fcbb..24399c9b2f9 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -14,7 +14,7 @@ use crate::topology::RedisTopology; mod connection; mod operations; -pub(crate) use connection::ConnectionRef; +pub use connection::ConnectionRef; use connection::{ClusterConnectionManager, ConnectionManager}; pub use operations::{ @@ -40,7 +40,8 @@ const CLAIM_SCRIPT: &str = concat!( ); const CLAIM_ATTEMPTS: usize = 8; -enum Connections { +#[allow(private_interfaces)] +pub enum Connections { Pool(r2d2::Pool), Cluster(r2d2::Pool), Fixed(Mutex), @@ -50,7 +51,7 @@ impl Connections where C: redis::ConnectionLike + Send + 'static, { - fn execute( + pub fn execute( &self, operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, ) -> Result { @@ -73,6 +74,29 @@ where } } } + + pub async fn run_blocking(connections: Arc, operation: F) -> Result + where + T: Send + 'static, + F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, + { + tokio::task::spawn_blocking(move || connections.execute(operation)) + .await + .map_err(|_| Error::Unavailable)? + } + + pub fn fixed(connection: C) -> Self { + Self::Fixed(Mutex::new(connection)) + } + + pub fn open(url: &str, topology: &RedisTopology) -> Result { + match topology { + RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)), + RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool( + ClusterConnectionManager::open(url, startup_nodes)?, + )?)), + } + } } pub struct RedisCache { @@ -94,12 +118,7 @@ impl RedisCache { default_ttl: Option, codec: S, ) -> Result { - let connections = match topology { - RedisTopology::Standalone => Connections::Pool(pool(ConnectionManager::open(url)?)?), - RedisTopology::Cluster { startup_nodes } => { - Connections::Cluster(pool(ClusterConnectionManager::open(url, startup_nodes)?)?) - } - }; + let connections = Connections::open(url, topology)?; Ok(Self { connections: Arc::new(connections), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), @@ -127,7 +146,7 @@ where { pub fn with_connection(connection: C, default_ttl: Option, codec: S) -> Self { Self { - connections: Arc::new(Connections::Fixed(Mutex::new(connection))), + connections: Arc::new(Connections::fixed(connection)), default_ttl: default_ttl.unwrap_or(DEFAULT_TTL), codec, namespace: None, @@ -203,16 +222,6 @@ where .saturating_add(u64::from(ttl.subsec_nanos() > 0)) .max(1) } - - async fn run_blocking(connections: Arc>, operation: F) -> Result - where - T: Send + 'static, - F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, - { - tokio::task::spawn_blocking(move || connections.execute(operation)) - .await - .map_err(|_| Error::Unavailable)? - } } fn namespaced_key(namespace: Option<&str>, key: &str) -> String { @@ -271,7 +280,7 @@ where let payload = self.codec.encode(&value)?; let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection .set_ex::<_, _, ()>(key, payload, ttl) .map_err(|_| Error::Unavailable) @@ -285,7 +294,7 @@ where _: &ExactCacheContext, ) -> Result, Error> { let key = self.namespaced_key(key); - let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection .get::<_, redis::Value>(key) .map_err(|_| Error::Unavailable) @@ -311,7 +320,7 @@ where if entries.is_empty() { return Ok(()); } - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let commands = entries .into_iter() .map(|(key, payload)| { @@ -330,7 +339,7 @@ where } async fn test_connection(&self) -> Result { - match Self::run_blocking(Arc::clone(&self.connections), |connection| { + match Connections::run_blocking(Arc::clone(&self.connections), |connection| { Ok(match connection.ping() { Ok(_) => CacheConnectionResult { status: CacheConnectionStatus::Success, @@ -391,7 +400,7 @@ where .iter() .map(|key| self.namespaced_key(key)) .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("MGET") .arg(keys) .query::>(connection) @@ -418,7 +427,7 @@ where async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { let key = self.namespaced_key(key); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection.del::<_, ()>(key).map_err(|_| Error::Unavailable) }) .await @@ -438,7 +447,7 @@ where async fn async_flush_cache(&self) -> Result<(), Error> { let pattern = self.namespaced_pattern()?; - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { Self::flush_matching(connection, &pattern) }) .await @@ -470,7 +479,7 @@ where ) -> Result { let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { increment(connection, key, amount, ttl) }) .await @@ -581,7 +590,7 @@ where let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl)); let codec = self.codec.clone(); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { claim(connection, &codec, &key, candidate, &eligible, ttl) }) .await diff --git a/litellm-rust/crates/cache-redis/src/cache/connection.rs b/litellm-rust/crates/cache-redis/src/cache/connection.rs index 1834f1d94e5..013bf055f89 100644 --- a/litellm-rust/crates/cache-redis/src/cache/connection.rs +++ b/litellm-rust/crates/cache-redis/src/cache/connection.rs @@ -12,7 +12,7 @@ use redis::{ use super::REDIS_TIMEOUT; use crate::topology::RedisNode; -pub(super) struct PooledConnection { +pub struct PooledConnection { pub(super) connection: C, pub(super) failed: bool, } @@ -20,7 +20,7 @@ pub(super) struct PooledConnection { /// Pools connections without a checkout PING, which would double every operation's round trips. /// A timed-out command leaves its reply on the socket while redis still reports the connection /// open, so any connection whose operation failed is discarded instead of being reused. -pub(super) struct ConnectionManager(redis::Client); +pub struct ConnectionManager(redis::Client); impl ConnectionManager { pub(super) fn open(url: &str) -> Result { @@ -54,7 +54,7 @@ impl r2d2::ManageConnection for ConnectionManager { } } -pub(super) struct ClusterConnectionManager(ClusterClient); +pub struct ClusterConnectionManager(ClusterClient); impl ClusterConnectionManager { pub(super) fn open(url: &str, startup_nodes: &[RedisNode]) -> Result { @@ -117,7 +117,7 @@ impl r2d2::ManageConnection for ClusterConnectionManager { } } -pub(crate) enum ConnectionRef<'a> { +pub enum ConnectionRef<'a> { Node(&'a mut dyn redis::ConnectionLike), Cluster(&'a mut ClusterConnection), } diff --git a/litellm-rust/crates/cache-redis/src/cache/operations.rs b/litellm-rust/crates/cache-redis/src/cache/operations.rs index 9a7023338bf..4345ee879b3 100644 --- a/litellm-rust/crates/cache-redis/src/cache/operations.rs +++ b/litellm-rust/crates/cache-redis/src/cache/operations.rs @@ -144,7 +144,7 @@ where .into_iter() .map(|key| self.namespaced_key(&key)) .collect::>(); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { connection.del(keys).map_err(|_| Error::Unavailable) }) .await @@ -172,7 +172,7 @@ where .iter() .map(|key| self.namespaced_key(key)) .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("MGET") .arg(keys) .query::>(connection) @@ -188,7 +188,7 @@ where } pub async fn ping(&self) -> Result { - Self::run_blocking(Arc::clone(&self.connections), |connection| { + Connections::run_blocking(Arc::clone(&self.connections), |connection| { connection.ping().map_err(|_| Error::Unavailable) }) .await @@ -196,7 +196,7 @@ where pub async fn async_get_ttl(&self, key: &str) -> Result, Error> { let key = self.namespaced_key(key); - let ttl = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let ttl = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("TTL") .arg(key) .query::(connection) @@ -208,7 +208,7 @@ where pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, Error> { let pattern = format!("{}*", self.namespaced_key(pattern)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut matches = Vec::new(); connection.scan(&pattern, count, |_, keys| { matches.extend(keys); @@ -231,7 +231,7 @@ where } let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut sadd = redis::cmd("SADD"); sadd.arg(&key).arg(values); let mut expire = redis::cmd("EXPIRE"); @@ -253,7 +253,7 @@ where return Err(Error::InvalidEntry); } let key = self.namespaced_key(key); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("RPUSH") .arg(key) .arg(values) @@ -279,7 +279,7 @@ where if operations.is_empty() { return Ok(Vec::new()); } - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let commands = operations .into_iter() .map(|(key, values)| { @@ -304,7 +304,7 @@ where ) -> Result { let key = self.namespaced_key(key); let multiple = count.is_some(); - let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut command = redis::cmd("LPOP"); command.arg(key); if let Some(count) = count { @@ -333,7 +333,7 @@ where .iter() .map(|(_, count)| count.is_some()) .collect::>(); - let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| { + let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let commands = operations .into_iter() .map(|(key, count)| { @@ -365,7 +365,7 @@ where .into_iter() .map(|key| self.namespaced_key(&key)) .collect::>(); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("EVAL") .arg(script) .arg(keys.len()) @@ -426,7 +426,7 @@ where if operations.is_empty() { return Ok(Vec::new()); } - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { let mut commands = Vec::with_capacity(operations.len() * 2); let mut increments = Vec::with_capacity(operations.len()); for (key, amount, ttl) in operations { @@ -460,7 +460,7 @@ where ) -> Result { let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(ttl); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { increment_with_floor(connection, key, amount, ttl) }) .await @@ -474,7 +474,7 @@ where ) -> Result { let key = self.namespaced_key(key); let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl)); - Self::run_blocking(Arc::clone(&self.connections), move |connection| { + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { redis::cmd("EVAL") .arg(SET_MAX_SCRIPT) .arg(1) diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs index 98f6bfd8ce5..efb0db931ac 100644 --- a/litellm-rust/crates/cache-redis/src/lib.rs +++ b/litellm-rust/crates/cache-redis/src/lib.rs @@ -1,6 +1,10 @@ mod cache; mod topology; +pub mod connection { + pub use crate::cache::{ConnectionRef, Connections}; +} + pub use cache::{ RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript, }; diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index 56c1646d343..46e561ddad1 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -58,4 +58,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths -Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees +Redis cluster, disk, and cloud stores remain follow-ups. Semantic backends plug in through `SemanticCacheContext`, which carries the prompt inputs and metadata alongside the cache TTL. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index e50e68cdabb..2f949d511de 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,21 +1,21 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache, + BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, FlushCache, }; use serde_json::Value; use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; #[derive(Clone)] -pub struct ResponseCacheRequest { +pub struct ResponseCacheRequest { pub key: CacheKeyInput, pub controls: CacheControls, - pub context: ExactCacheContext, + pub context: C, pub max_age: Option, } -impl ResponseCacheRequest { +impl ResponseCacheRequest { pub fn new(key: CacheKeyInput) -> Self { Self { key, @@ -26,17 +26,24 @@ impl ResponseCacheRequest { default_on: true, ..Default::default() }, - context: ExactCacheContext::default(), + context: C::default(), max_age: None, } } } -pub struct ResponseCache> { +pub struct ResponseCache> +where + B::Context: Default + PartialEq, +{ backend: Arc, } -impl> ResponseCache { +impl ResponseCache +where + B: BaseCache, + B::Context: Default + PartialEq, +{ pub fn new(backend: Arc) -> Self { Self { backend } } @@ -45,8 +52,12 @@ impl> ResponseCach &self.backend } + pub fn backend_arc(&self) -> &Arc { + &self.backend + } + pub fn default_ttl(&self) -> Option { - self.backend.get_ttl(&ExactCacheContext::default()) + self.backend.get_ttl(&B::Context::default()) } pub async fn async_flush(&self) -> Result<(), Error> @@ -62,7 +73,7 @@ impl> ResponseCach pub fn lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -81,7 +92,7 @@ impl> ResponseCach pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -101,7 +112,7 @@ impl> ResponseCach pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -126,7 +137,7 @@ impl> ResponseCach pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -153,7 +164,7 @@ impl> ResponseCach pub fn store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -172,7 +183,7 @@ impl> ResponseCach pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -193,7 +204,7 @@ impl> ResponseCach pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(ResponseCacheRequest, Value)>, now: Duration, ) -> Result<(), Error> { self.async_store_entries( @@ -209,7 +220,7 @@ impl> ResponseCach /// the freshness of its original response. pub async fn async_store_entries( &self, - entries: Vec<(ResponseCacheRequest, Value, Duration)>, + entries: Vec<(ResponseCacheRequest, Value, Duration)>, ) -> Result<(), Error> { let writable = entries .into_iter() @@ -249,8 +260,8 @@ impl> ResponseCach } fn partial_hits( - requests: &[ResponseCacheRequest], - readable: Vec<(usize, &ResponseCacheRequest)>, + requests: &[ResponseCacheRequest], + readable: Vec<(usize, &ResponseCacheRequest)>, entries: Vec>, now: Duration, ) -> Result { diff --git a/litellm-rust/crates/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 8bd69ba5ad6..5c10e7fd5c3 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -32,6 +32,28 @@ impl CacheContext for ExactCacheContext { } } +#[derive(Clone, Debug, Default, PartialEq)] +pub struct SemanticCacheContext { + pub input: Option, + pub messages: Option, + pub metadata: Option, + pub scope: Option, + pub ttl: Option, +} + +impl CacheContext for SemanticCacheContext { + fn ttl(&self) -> Option { + self.ttl + } + + fn with_ttl(&self, ttl: Option) -> Self { + Self { + ttl, + ..self.clone() + } + } +} + #[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)] #[serde(rename_all = "lowercase")] pub enum CacheConnectionStatus { @@ -105,3 +127,31 @@ pub trait BaseCache: Send + Sync { fn test_connection(&self) -> impl Future> + Send; } + +#[cfg(test)] +mod tests { + use std::time::Duration; + + use serde_json::json; + + use super::{CacheContext, SemanticCacheContext}; + + #[test] + fn semantic_context_with_ttl_only_replaces_ttl() { + let context = SemanticCacheContext { + input: Some(json!({"input": "hello"})), + messages: Some(json!([{"role": "user", "content": "hello"}])), + metadata: Some(json!({"tenant": "team"})), + scope: Some("scope".into()), + ttl: Some(Duration::from_secs(10)), + }; + + let updated = context.with_ttl(Some(Duration::from_secs(20))); + + assert_eq!(updated.ttl, Some(Duration::from_secs(20))); + assert_eq!(updated.input, context.input); + assert_eq!(updated.messages, context.messages); + assert_eq!(updated.metadata, context.metadata); + assert_eq!(updated.scope, context.scope); + } +} diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index ce9f93b6dc4..8364c635e3a 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -8,7 +8,7 @@ mod error; pub use base_cache::{ BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext, - ExactCacheContext, + ExactCacheContext, SemanticCacheContext, }; pub use cache_type::CacheType; pub use caching::{Cache, CacheBackend, get_cache, set_cache}; diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 9180ee9d0dc..36307ac9b33 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,7 +1,8 @@ use std::{sync::Mutex, time::Duration}; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache, + BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, SemanticCacheContext, + get_cache, }; struct TestCache { @@ -126,6 +127,24 @@ fn associated_context_preserves_backend_specific_lookup_inputs() { ); } +#[test] +fn semantic_context_with_ttl_preserves_lookup_inputs() { + let context = SemanticCacheContext { + input: Some(serde_json::json!("text")), + messages: Some(serde_json::json!([{"role": "user", "content": "hi"}])), + metadata: Some(serde_json::json!({"key": "value"})), + scope: Some("scope".into()), + ttl: None, + }; + let updated = context.with_ttl(Some(Duration::from_secs(30))); + assert_eq!(updated.ttl(), Some(Duration::from_secs(30))); + assert_eq!(updated.input, context.input); + assert_eq!(updated.messages, context.messages); + assert_eq!(updated.metadata, context.metadata); + assert_eq!(updated.scope, context.scope); + assert_eq!(context.with_ttl(None).ttl(), None); +} + #[tokio::test] async fn default_batch_operations_use_async_writes_and_stop_on_failure() { let cache = TestCache { diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 9191e61a25c..bbe727f9345 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -27,7 +27,9 @@ 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 @@ -42,8 +44,9 @@ litellm-host-python.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true +redis = { version = "1.7.0", features = ["tls-rustls"] } serde_json.workspace = true -tokio = { workspace = true, features = ["sync"] } +tokio = { workspace = true, features = ["rt", "sync"] } [dev-dependencies] serde.workspace = true @@ -51,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 ad64b24d3c1..2ff73238202 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -56,12 +56,7 @@ impl ResolvedCache { CacheBinding::Disabled => ready_none(py)?, CacheBinding::Native(service) => { let request = request(input)?; - let service = service.clone(); - run_async( - py, - async move { service.async_lookup(&request, now()).await }, - cache_error, - )? + service.async_lookup_py(py, request)? } CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?, }; @@ -179,12 +174,7 @@ impl ResolvedCache { CacheBinding::Native(service) => { let request = self::request(request)?; let response: Value = from_py(response)?; - let service = service.clone(); - run_async( - py, - async move { service.async_store(&request, response, now()).await }, - cache_error, - ) + service.async_store_py(py, request, response) } CacheBinding::PythonCallback(callback) => { callback.async_store(py, response, callback_kwargs) @@ -241,12 +231,7 @@ impl ResolvedCache { )); } let entries = requests.into_iter().zip(responses).collect(); - let service = service.clone(); - run_async( - py, - async move { service.async_store_batch(entries, now()).await }, - cache_error, - ) + service.async_store_batch_py(py, entries) } CacheBinding::PythonCallback(callback) => { callback.async_store_batch(py, callback_result, callback_kwargs) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 7316b3e259c..ef2a03ac9ba 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -93,6 +93,19 @@ pub(super) struct AzureBlobCacheConfig { pub(super) container: String, } +#[allow( + dead_code, + reason = "embedding settings are projected so drift falls back to Python" +)] +pub(super) struct RedisSemanticCacheConfig { + pub(super) redis_url: String, + pub(super) index_name: String, + pub(super) similarity_threshold: f64, + pub(super) embedding_model: String, + pub(super) embedding_max_input_tokens: Option, + pub(super) embedding_timeout: Option, +} + struct RedisClientProjection<'py> { topology: RedisTopology, host: String, @@ -104,13 +117,23 @@ 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), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -201,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, @@ -214,10 +244,13 @@ impl NativeCacheConfig { backend: CacheBackendConfig::AzureBlob(backend), })) }), - Some( - CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::QdrantSemantic, - ) - | None => Ok(CacheConfigProjection::Unsupported( + Some(CacheType::RedisSemantic) => project_redis_semantic(&backend).map(|backend| { + CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::RedisSemantic(Box::new(backend)), + })) + }), + Some(CacheType::QdrantSemantic) | None => Ok(CacheConfigProjection::Unsupported( UnsupportedCacheConfig::Backend, )), } @@ -228,11 +261,15 @@ impl NativeCacheConfig { 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(_) => None, + | 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 { @@ -306,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") } @@ -317,6 +364,20 @@ impl NativeCacheConfig { let facade = std::fs::canonicalize(&config.directory).ok(); (native != facade).then_some("facade and native backend directories must match") } + CacheBackendConfig::RedisSemantic(_) if service.kind() != "redis_semantic" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::RedisSemantic(config) + if service.index_name() != Some(config.index_name.as_str()) => + { + Some("facade and native backend index names must match") + } + CacheBackendConfig::RedisSemantic(config) + if service.similarity_threshold() != Some(config.similarity_threshold as f32) => + { + Some("facade and native backend similarity thresholds must match") + } + CacheBackendConfig::RedisSemantic(_) => None, CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() { None => Some("facade and native backend types must match"), Some((account_url, container)) @@ -345,6 +406,27 @@ fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult, +) -> PyResult { + Ok(RedisSemanticCacheConfig { + redis_url: backend.getattr("_redis_url")?.extract::()?, + index_name: backend + .getattr("_index_name")? + .extract::>()? + .unwrap_or_else(|| "litellm_semantic_cache_index".into()), + similarity_threshold: backend.getattr("similarity_threshold")?.extract::()?, + embedding_model: backend.getattr("embedding_model")?.extract::()?, + embedding_max_input_tokens: backend + .getattr("embedding_max_input_tokens")? + .extract::>()?, + embedding_timeout: backend + .getattr("embedding_timeout")? + .extract::>()?, + }) +} + #[inline(never)] fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { let max_size_kib = backend.getattr("max_size_per_item")?.extract::()?; @@ -588,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 { @@ -676,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 { @@ -869,16 +1016,13 @@ mod tests { use pyo3::{prelude::*, types::PyDict}; - use litellm_auth_aws::AwsAuthConfig; - use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; - use litellm_host_python::run_sync_value; + 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; - use litellm_cache_redis::{RedisNode, RedisTopology}; fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> { facade( @@ -951,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(); @@ -1016,375 +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" - ); - }); - } - - fn s3_facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> { - let locals = PyDict::new(py); - py.run( - &CString::new(format!( - "from types import SimpleNamespace\n\ - S3Client = type('S3', (), {{'__module__': 'botocore.client'}})\n\ - client = S3Client()\n\ - client.meta = SimpleNamespace(region_name='us-east-1', endpoint_url='https://example.test', config=SimpleNamespace(s3=None, proxies=None, client_cert=None, signature_version='s3v4'))\n\ - client._endpoint = SimpleNamespace(http_session=SimpleNamespace(_verify=True))\n\ - client._request_signer = SimpleNamespace(_credentials=SimpleNamespace(method='explicit', access_key='key', secret_key='secret', token='token'))\n\ - backend = SimpleNamespace(bucket_name='bucket', key_prefix='team/', s3_client=client)\n\ - facade = SimpleNamespace(type='s3', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)\n\ - {body}" - )) - .unwrap(), - None, - Some(&locals), - ) - .unwrap(); - locals.get_item("facade").unwrap().unwrap() - } - - #[test] - fn projects_s3_configuration_with_explicit_credentials_and_custom_endpoint() { - Python::initialize(); - Python::attach(|py| { - let facade = s3_facade(py, ""); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("S3 cache should be supported"); - }; - let CacheBackendConfig::S3(s3) = config.backend else { - panic!("expected S3 configuration"); - }; - assert_eq!(s3.bucket, "bucket"); - assert_eq!(s3.key_prefix, "team/"); - assert_eq!(s3.region, "us-east-1"); - assert_eq!( - s3.endpoint.map(|endpoint| endpoint.url).as_deref(), - Some("https://example.test") - ); - assert_eq!(s3.auth.access_key_id.as_deref(), Some("key")); - assert_eq!(s3.auth.secret_access_key.as_deref(), Some("secret")); - assert_eq!(s3.auth.session_token.as_deref(), Some("token")); - assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1")); - }); - } - - #[test] - fn default_s3_endpoint_projects_no_custom_endpoint() { - Python::initialize(); - Python::attach(|py| { - let facade = s3_facade( - py, - "facade.cache.s3_client.meta.endpoint_url = 'https://s3.us-east-1.amazonaws.com'", - ); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("S3 cache should be supported"); - }; - let CacheBackendConfig::S3(s3) = config.backend else { - panic!("expected S3 configuration"); - }; - assert!(s3.endpoint.is_none()); - }); - } - - #[test] - fn non_sigv4_proxies_and_disabled_verification_stay_on_python() { - Python::initialize(); - Python::attach(|py| { - for (body, message) in [ - ( - "facade.cache.s3_client.meta.config.signature_version = 's3'", - "native S3 configuration requires Python", - ), - ( - "facade.cache.s3_client.meta.config.proxies = {'https': 'proxy'}", - "native S3 configuration requires Python", - ), - ( - "facade.cache.s3_client._endpoint.http_session._verify = False", - "native S3 configuration requires Python", - ), - ( - "del facade.cache.s3_client._endpoint.http_session._verify", - "native S3 configuration requires Python", - ), - ] { - let facade = s3_facade(py, body); - let CacheConfigProjection::Unsupported(reason) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("{body} must stay on Python"); - }; - assert_eq!(reason.message(), message); - } - let facade = s3_facade(py, "facade.cache.s3_client = SimpleNamespace()"); - let CacheConfigProjection::Unsupported(reason) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("non-botocore client must stay on Python"); - }; - assert_eq!(reason.message(), "native S3 client type is not implemented"); - let facade = s3_facade( - py, - "facade.cache.s3_client._request_signer._credentials = None", - ); - let CacheConfigProjection::Unsupported(reason) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("missing credentials must stay on Python"); - }; - assert_eq!(reason.message(), "native S3 credentials require Python"); - }); - } - - #[test] - fn non_explicit_s3_credentials_use_the_default_chain() { - Python::initialize(); - Python::attach(|py| { - let facade = s3_facade( - py, - "facade.cache.s3_client._request_signer._credentials = SimpleNamespace(method='sso', access_key=None, secret_key=None, token=None)", - ); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("default-chain credentials should be supported"); - }; - let CacheBackendConfig::S3(s3) = config.backend else { - panic!("expected S3 configuration"); - }; - assert_eq!(s3.auth.access_key_id, None); - assert_eq!(s3.auth.secret_access_key, None); - assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1")); - }); - } - - fn s3_service(py: Python<'_>, region: &str, endpoint: Option<&str>) -> NativeResponseCache { - let config = S3CacheConfig { - bucket: "bucket".to_string(), - key_prefix: "team/".to_string(), - region: region.to_string(), - endpoint: endpoint.map(|url| S3Endpoint { - url: url.to_string(), - }), - auth: AwsAuthConfig::default(), - }; - run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) }).unwrap() - } - - #[test] - fn s3_binding_rejects_region_and_endpoint_mismatches() { - Python::initialize(); - Python::attach(|py| { - let facade = s3_facade(py, ""); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("S3 cache should be supported"); - }; - assert_eq!( - config.service_mismatch(&s3_service(py, "us-east-1", Some("https://example.test"))), - None - ); - assert_eq!( - config.service_mismatch(&s3_service(py, "us-west-2", Some("https://example.test"))), - Some("facade and native backend regions must match") - ); - assert_eq!( - config.service_mismatch(&s3_service(py, "us-east-1", Some("https://other.test"))), - Some("facade and native backend endpoints must match") - ); - assert_eq!( - config.service_mismatch(&s3_service(py, "us-east-1", None)), - Some("facade and native backend endpoints must match") - ); - }); - } - - #[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 new file mode 100644 index 00000000000..9398e5a862b --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -0,0 +1,156 @@ +use std::future::Future; + +use litellm_cache::Error; +use litellm_host_python::to_py; +use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict}; +use serde_json::Value; + +tokio::task_local! { + static PREPARED_EMBEDDING: Result, Error>; +} + +pub(super) fn with_prepared_embedding( + vector: Result, Error>, + future: F, +) -> impl Future { + PREPARED_EMBEDDING.scope(vector, future) +} + +pub(super) struct PythonEmbedder(Py); + +impl Clone for PythonEmbedder { + fn clone(&self) -> Self { + Python::attach(|py| Self(self.0.clone_ref(py))) + } +} + +impl PythonEmbedder { + pub(super) fn new(object: Py) -> Self { + Self(object) + } + + pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult { + Ok(Self(backend.clone().unbind())) + } + + pub(super) fn object(&self) -> &Py { + &self.0 + } + + pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + visit.call(&self.0) + } + + pub(super) fn async_embed_awaitable<'py>( + &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: Option<&Value>, + ) -> PyResult> { + let kwargs = PyDict::new(py); + kwargs.set_item("metadata", to_py(py, &metadata)?)?; + Ok(kwargs) + } + + pub(super) fn async_embedding_coroutine( + &self, + py: Python<'_>, + prompt: &str, + metadata: Option<&Value>, + ) -> PyResult> { + let kwargs = Self::metadata_kwargs(py, metadata)?; + self.0 + .bind(py) + .call_method("_get_async_embedding", (prompt,), Some(&kwargs)) + .map(Bound::unbind) + } + + pub(super) fn extract(vector: Bound<'_, PyAny>) -> PyResult> { + Ok(vector + .extract::>()? + .into_iter() + .map(|value| value as f32) + .collect()) + } +} + +impl 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( + "_get_embedding", + (prompt,), + Some(&kwargs), + )?) + }) + .map_err(|_| Error::Unavailable) + } + + fn async_embed( + &self, + _prompt: &str, + _metadata: Option<&Value>, + ) -> impl Future, Error>> + Send { + let seeded = PREPARED_EMBEDDING + .try_with(Clone::clone) + .unwrap_or(Err(Error::Unavailable)); + std::future::ready(seeded) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn async_embed_returns_the_seeded_vector_or_unavailable() { + Python::initialize(); + let object = Python::attach(|py| py.None()); + let embedder = PythonEmbedder::new(object); + let scoped_embedder = embedder.clone(); + let scoped = with_prepared_embedding(Ok(vec![0.25]), async move { + litellm_cache_redis_semantic::Embedder::async_embed(&scoped_embedder, "prompt", None) + .await + }); + assert_eq!(scoped.await, Ok(vec![0.25])); + let unscoped = + litellm_cache_redis_semantic::Embedder::async_embed(&embedder, "prompt", None).await; + assert_eq!(unscoped, Err(Error::Unavailable)); + } +} diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index e89d0fc9715..f1389b745e9 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -31,17 +31,15 @@ struct RedisPoolGuard { connection_class: Py, connection_kwargs: Py, max_connections: Option, + client_name: &'static str, attributes: RedisPoolAttributes, } -struct S3ClientGuard { - reference: Py, -} - struct DiskStoreGuard { reference: Py, directory: String, } + struct AzureBlobClientGuard { sync_client: Py, async_client: Py, @@ -49,12 +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, @@ -72,6 +76,9 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes { connection_class: "connection_pool_class", max_connections: None, }; + +const VALKEY_POOL: RedisPoolAttributes = STANDALONE_POOL; + pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, @@ -158,7 +165,9 @@ impl ObjectGuard { return Ok(false); } for (name, value) in &expected.attributes { - if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) { + if (instance.contains(name)? && !self.config_names.contains(&name.as_str())) + || !attributes.get_item(name)?.is(value.bind(py)) + { return Ok(false); } } @@ -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) @@ -226,22 +238,6 @@ impl RedisPoolGuard { } } -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 DiskStoreGuard { fn capture(backend: &Bound<'_, PyAny>) -> PyResult { let store = backend.getattr("disk_cache")?; @@ -261,6 +257,7 @@ impl DiskStoreGuard { visit.call(&self.reference) } } + impl AzureBlobClientGuard { fn capture(backend: &Bound<'_, PyAny>) -> PyResult { let sync_client = backend.getattr("container_client")?; @@ -289,11 +286,41 @@ 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, @@ -318,6 +345,7 @@ impl ConnectionGuard { } } } + impl FacadeGuard { pub(super) fn capture( py: Python<'_>, @@ -335,19 +363,29 @@ impl FacadeGuard { let (module, name, cache_kind) = match (kind, cluster) { ("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"), ("redis", false) => ("litellm.caching.redis_cache", "RedisCache", "redis"), + ("redis_semantic", _) => ( + "litellm.caching.redis_semantic_cache", + "RedisSemanticCache", + "redis-semantic", + ), ("redis", true) => ( "litellm.caching.redis_cluster_cache", "RedisClusterCache", "redis", ), - ("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"), ("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")?; @@ -367,6 +405,15 @@ impl FacadeGuard { if let Some(message) = config.service_mismatch(service) { return Err(PyTypeError::new_err(message)); } + if kind == "redis_semantic" + && service + .embedder_object() + .is_none_or(|embedder| !backend.is(embedder.bind(py))) + { + return Err(PyTypeError::new_err( + "facade backend must be the native embedder", + )); + } Ok(Self { outer: ObjectGuard::capture( py, @@ -391,6 +438,18 @@ impl FacadeGuard { "max_size_per_item", "redis_kwargs", "redis_flush_size", + "similarity_threshold", + "distance_threshold", + "embedding_model", + "embedding_max_input_tokens", + "embedding_timeout", + "_index_name", + "_redis_url", + "similarity_threshold", + "embedding_model", + "index_name", + "embedding_max_input_tokens", + "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 d6e622a5c49..e6edf48d72f 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,12 +1,19 @@ use litellm_auth_aws::AwsAuthConfig; +use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; use litellm_cache_redis::{RedisNode, RedisTopology}; +use litellm_cache_redis_semantic::RedisSemanticConfig; use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; use litellm_host_python::{release_gil, run_sync_value}; -use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyRuntimeError, PyTypeError}, + prelude::*, +}; -use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; - -use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; +use super::{ + cache_error, config::project_redis_semantic, embedder::PythonEmbedder, facade::FacadeGuard, + native::NativeResponseCache, request::duration, +}; #[pyclass(frozen, name = "_CacheTestHandle")] pub(crate) struct CacheTestHandle { @@ -139,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 { @@ -153,6 +183,37 @@ impl CacheTestHandle { pid: std::process::id(), }) } + + #[staticmethod] + fn redis_semantic(py: Python<'_>, backend: Bound<'_, PyAny>) -> PyResult { + let class = py + .import("litellm.caching.redis_semantic_cache")? + .getattr("RedisSemanticCache")?; + if !backend.get_type().is(&class) { + return Err(PyTypeError::new_err( + "native redis-semantic handles require the built-in RedisSemanticCache", + )); + } + let config = project_redis_semantic(&backend)?; + let embedder = PythonEmbedder::new(backend.unbind()); + let service = release_gil(py, move || { + NativeResponseCache::redis_semantic( + &config.redis_url, + embedder, + RedisSemanticConfig { + index_name: config.index_name, + similarity_threshold: config.similarity_threshold as f32, + }, + ) + }) + .map_err(cache_error)?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() @@ -161,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 { @@ -178,6 +245,7 @@ impl CacheTestHandle { } fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> { + self.service.traverse(&visit)?; if let Some(guard) = &self.guard { guard.traverse(visit)?; } diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index ac4494150d9..fa028518559 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,12 +1,15 @@ mod binding; mod callback; mod config; +mod embedder; mod facade; mod future; mod handle; 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 39496890009..acc457d8c9c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,17 +1,79 @@ use std::{path::Path, sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; +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}; 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 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::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 { Memory(Arc>>), @@ -21,6 +83,15 @@ pub(super) enum NativeResponseCache { }, S3(Arc>>), Gcs(Arc>>), + ValkeySemantic { + cache: Arc>>, + embedder: PythonEmbedder, + scope: String, + }, + RedisSemantic { + cache: Arc>>, + embedder: PythonEmbedder, + }, Disk(Arc>>), AzureBlob(Arc>>), } @@ -53,6 +124,7 @@ impl NativeResponseCache { buffer: None, }) } + pub async fn s3(config: S3CacheConfig) -> Self { let runtime = tokio::runtime::Handle::current(); Self::S3(Arc::new(ResponseCache::new(Arc::new(S3Cache::new( @@ -61,6 +133,41 @@ impl NativeResponseCache { 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( + url: &str, + embedder: PythonEmbedder, + config: RedisSemanticConfig, + ) -> Result { + let backend = RedisSemanticCache::new(url, embedder.clone(), config)?; + Ok(Self::RedisSemantic { + cache: Arc::new(ResponseCache::new(Arc::new(backend))), + embedder, + }) + } + pub fn disk(directory: &str) -> Result { let cache = DiskCache::open(directory, ResponseCacheCodec)?; Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache))))) @@ -97,20 +204,91 @@ impl NativeResponseCache { cache.backend().account_url(), cache.backend().container_name(), )), - Self::Memory(_) | Self::Redis { .. } | Self::S3(_) | Self::Disk(_) | Self::Gcs(_) => { - None - } + Self::Memory(_) + | Self::Redis { .. } + | Self::S3(_) + | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } + | Self::Disk(_) + | 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::AzureBlob(_) => "azure-blob", } @@ -122,6 +300,8 @@ impl NativeResponseCache { 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::AzureBlob(cache) => cache.default_ttl(), } @@ -157,17 +337,27 @@ impl NativeResponseCache { pub fn namespace(&self) -> Option<&str> { match self { - Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) => None, + Self::Memory(_) + | Self::S3(_) + | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } + | Self::Disk(_) + | Self::AzureBlob(_) + | Self::Gcs(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), - Self::S3(_) | Self::Gcs(_) => None, } } pub fn topology(&self) -> Option<&RedisTopology> { match self { - Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, + Self::Memory(_) + | Self::S3(_) + | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } + | Self::Disk(_) + | Self::AzureBlob(_) + | Self::Gcs(_) => None, Self::Redis { cache, .. } => Some(cache.backend().topology()), - Self::S3(_) => None, } } @@ -176,6 +366,8 @@ impl NativeResponseCache { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), Self::Redis { .. } | Self::S3(_) + | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, @@ -187,144 +379,453 @@ impl NativeResponseCache { Self::Memory(cache) => cache.backend().max_entry_bytes(), Self::Redis { .. } | Self::S3(_) + | Self::ValkeySemantic { .. } + | Self::RedisSemantic { .. } | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, } } - 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(_) => None, } } - pub fn lookup( - &self, - request: &ResponseCacheRequest, - now: Duration, - ) -> Result, Error> { + pub fn semantic_config(&self) -> Option<(f64, &str)> { match self { - Self::Memory(cache) => cache.lookup(request, now), - Self::Redis { cache, .. } => cache.lookup(request, now), - Self::S3(cache) => cache.lookup(request, now), - Self::Gcs(cache) => cache.lookup(request, now), - Self::Disk(cache) => cache.lookup(request, now), - Self::AzureBlob(cache) => cache.lookup(request, now), + 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 index_name(&self) -> Option<&str> { + match self { + Self::RedisSemantic { cache, .. } => Some(cache.backend().index_name()), + _ => None, + } + } + + pub fn similarity_threshold(&self) -> Option { + match self { + Self::RedisSemantic { cache, .. } => Some(cache.backend().similarity_threshold()), + _ => None, + } + } + + pub fn semantic_embedder(&self) -> Option<&PythonEmbedder> { + match self { + Self::RedisSemantic { embedder, .. } => Some(embedder), + _ => None, + } + } + + pub fn embedder_object(&self) -> Option<&Py> { + match self { + Self::RedisSemantic { embedder, .. } => Some(embedder.object()), + _ => None, + } + } + + 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: &ResponseCacheRequest, + request: &NativeRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.store(request, response, now), - Self::Redis { cache, .. } => cache.store(request, response, now), - Self::S3(cache) => cache.store(request, response, now), - Self::Gcs(cache) => cache.store(request, response, now), - Self::Disk(cache) => cache.store(request, response, now), - Self::AzureBlob(cache) => cache.store(request, response, now), + Self::Memory(cache) => cache.store(&Self::exact(request), response, now), + Self::Redis { cache, .. } => cache.store(&Self::exact(request), response, now), + Self::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: &[ResponseCacheRequest], + requests: &[NativeRequest], now: Duration, ) -> Result { match self { - Self::Memory(cache) => cache.lookup_batch(requests, now), - Self::Redis { cache, .. } => cache.lookup_batch(requests, now), - Self::S3(cache) => cache.lookup_batch(requests, now), - Self::Gcs(cache) => cache.lookup_batch(requests, now), - Self::Disk(cache) => cache.lookup_batch(requests, now), - Self::AzureBlob(cache) => cache.lookup_batch(requests, now), + Self::Memory(cache) => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.lookup_batch(&requests, now) + } + Self::Redis { cache, .. } => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.lookup_batch(&requests, now) + } + Self::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: &ResponseCacheRequest, + request: &NativeRequest, now: Duration, ) -> Result, Error> { match self { - Self::Memory(cache) => cache.async_lookup(request, now).await, - Self::Redis { cache, .. } => cache.async_lookup(request, now).await, - Self::S3(cache) => cache.async_lookup(request, now).await, - Self::Gcs(cache) => cache.async_lookup(request, now).await, - Self::Disk(cache) => cache.async_lookup(request, now).await, - Self::AzureBlob(cache) => cache.async_lookup(request, now).await, + Self::Memory(cache) => cache.async_lookup(&Self::exact(request), now).await, + Self::Redis { cache, .. } => cache.async_lookup(&Self::exact(request), now).await, + Self::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: &ResponseCacheRequest, + request: &NativeRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_store(request, response, now).await, + Self::Memory(cache) => { + cache + .async_store(&Self::exact(request), response, now) + .await + } Self::Redis { cache, buffer: None, - } => cache.async_store(request, response, now).await, + } => { + cache + .async_store(&Self::exact(request), response, now) + .await + } Self::Redis { cache, buffer: Some(buffer), - } => buffer.async_store(cache, request, response, now).await, - Self::S3(cache) => cache.async_store(request, response, now).await, - Self::Gcs(cache) => cache.async_store(request, response, now).await, - Self::Disk(cache) => cache.async_store(request, response, now).await, - Self::AzureBlob(cache) => cache.async_store(request, response, now).await, + } => { + buffer + .async_store(cache, &Self::exact(request), 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: &[ResponseCacheRequest], + requests: &[NativeRequest], now: Duration, ) -> Result { match self { - Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, - Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, - Self::S3(cache) => cache.async_lookup_batch(requests, now).await, - Self::Gcs(cache) => cache.async_lookup_batch(requests, now).await, - Self::Disk(cache) => cache.async_lookup_batch(requests, now).await, - Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await, + Self::Memory(cache) => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.async_lookup_batch(&requests, now).await + } + Self::Redis { cache, .. } => { + let requests = requests.iter().map(Self::exact).collect::>(); + cache.async_lookup_batch(&requests, now).await + } + Self::S3(cache) => { + cache + .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(&requests.iter().map(Self::exact).collect::>(), now) + .await + } + Self::Disk(cache) => { + cache + .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) + .await + } + Self::AzureBlob(cache) => { + cache + .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) + .await + } } } pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(NativeRequest, Value)>, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_store_batch(entries, now).await, - Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, - Self::S3(cache) => cache.async_store_batch(entries, now).await, - Self::Gcs(cache) => cache.async_store_batch(entries, now).await, - Self::Disk(cache) => cache.async_store_batch(entries, now).await, - Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await, + Self::Memory(cache) => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + Self::Redis { cache, .. } => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + Self::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) => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + Self::Disk(cache) => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + Self::AzureBlob(cache) => { + let entries = entries + .into_iter() + .map(|(request, value)| (Self::exact(&request), value)) + .collect(); + cache.async_store_batch(entries, now).await + } + } + } + + pub(super) fn async_store_batch_py<'py>( + &self, + py: Python<'py>, + entries: Vec<(NativeRequest, Value)>, + ) -> PyResult> { + match self { + Self::Memory(_) + | Self::Redis { .. } + | 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())), + ), } } @@ -338,6 +839,9 @@ 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::AzureBlob(cache) => cache.async_flush().await, @@ -349,12 +853,23 @@ impl NativeResponseCache { 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::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()), @@ -362,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 0c5343a63d0..3b4b910c1f0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -1,9 +1,11 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use litellm_cache::ExactCacheContext; use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; use serde::Deserialize; +use serde_json::Value; #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -12,24 +14,52 @@ struct RequestInput { controls: Option, ttl_seconds: Option, max_age_seconds: Option, + messages: Option, + input: Option, + metadata: Option, + litellm_metadata: Option, + litellm_params: Option, + scope: Option, } -pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { +#[derive(Clone)] +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, +} + +pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { let input: RequestInput = from_py(value)?; request_input(input) } -fn request_input(input: RequestInput) -> PyResult { - let mut request = ResponseCacheRequest::new(input.key); - if let Some(controls) = input.controls { - request.controls = controls; - } - request.context.ttl = input.ttl_seconds.map(duration).transpose()?; - request.max_age = input.max_age_seconds.map(duration).transpose()?; - Ok(request) +fn request_input(input: RequestInput) -> PyResult { + let controls = input.controls.unwrap_or_else(|| { + ResponseCacheRequest::::new(input.key.clone()).controls + }); + Ok(NativeRequest { + key: input.key, + controls, + ttl: input.ttl_seconds.map(duration).transpose()?, + max_age: input.max_age_seconds.map(duration).transpose()?, + messages: input.messages, + input: input.input, + metadata: input.metadata, + 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 new file mode 100644 index 00000000000..7661598c5f9 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/semantic.rs @@ -0,0 +1,175 @@ +use std::collections::VecDeque; + +use litellm_cache::Error; +use litellm_cache_redis_semantic::prompt_from_context; +use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async}; +use pyo3::{ + PyTraverseError, PyVisit, + exceptions::{PyException, PyRuntimeError}, + prelude::*, +}; +use serde_json::Value; + +use super::{ + cache_error, + embedder::{PythonEmbedder, with_prepared_embedding}, + native::NativeResponseCache, + request::{NativeRequest, now}, +}; + +pub(super) enum SemanticOperation { + Lookup(NativeRequest), + Store(NativeRequest, Value), + StoreBatch(VecDeque<(NativeRequest, Value)>), +} + +enum Phase { + Start, + AwaitingEmbedding, + AwaitingBackend, +} + +pub(super) struct SemanticBody { + service: NativeResponseCache, + operation: SemanticOperation, + pending: Option<(NativeRequest, Option)>, + phase: Phase, +} + +impl SemanticBody { + pub(super) fn new(service: NativeResponseCache, operation: SemanticOperation) -> Self { + Self { + service, + operation, + pending: None, + phase: Phase::Start, + } + } + + fn backend_step( + &mut self, + py: Python<'_>, + seed: Result, Error>, + ) -> PyResult { + self.phase = Phase::AwaitingBackend; + let (request, response) = self.pending.take().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution resumed without a pending operation") + })?; + let service = self.service.clone(); + let future = async move { + match response { + None => service.async_lookup(&request, now()).await, + Some(response) => service + .async_store(&request, response, now()) + .await + .map(|_| None), + } + }; + let awaitable = run_async(py, with_prepared_embedding(seed, future), cache_error)?; + Ok(ExecutionStep::Await(awaitable.unbind())) + } +} + +impl ExecutionBody for SemanticBody { + fn resume(&mut self, mut result: Option>>) -> PyResult { + Python::attach(|py| { + loop { + match self.phase { + Phase::Start => { + if result.is_some() { + return Err(PyRuntimeError::new_err( + "semantic execution received a result before starting", + )); + } + if self.pending.is_none() { + match &mut self.operation { + SemanticOperation::Lookup(request) => { + self.pending = Some((request.clone(), None)); + } + SemanticOperation::Store(request, response) => { + let response = std::mem::replace(response, Value::Null); + self.pending = Some((request.clone(), Some(response))); + } + SemanticOperation::StoreBatch(queue) => { + let Some((request, response)) = queue.pop_front() else { + return Ok(ExecutionStep::Return(py.None())); + }; + self.pending = Some((request, Some(response))); + } + } + } + let (request, _) = self.pending.as_ref().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution has no pending operation") + })?; + let semantic = NativeResponseCache::redis_semantic_request(request); + let Some(prompt) = prompt_from_context(&semantic.context) else { + return self.backend_step(py, Err(Error::Unavailable)); + }; + let embedder = self.service.semantic_embedder().ok_or_else(|| { + PyRuntimeError::new_err( + "semantic execution requires a redis-semantic backend", + ) + })?; + let coroutine = embedder.async_embedding_coroutine( + py, + &prompt, + semantic.context.metadata.as_ref(), + )?; + self.phase = Phase::AwaitingEmbedding; + return Ok(ExecutionStep::Await(coroutine)); + } + Phase::AwaitingEmbedding => { + let result = result.take().ok_or_else(|| { + PyRuntimeError::new_err( + "semantic execution expected an embedding result", + ) + })?; + let seed = match result { + Ok(value) => PythonEmbedder::extract(value.into_bound(py)) + .map_err(|_| Error::Unavailable), + Err(error) => { + if !error.is_instance_of::(py) { + return Err(error); + } + Err(Error::Unavailable) + } + }; + return self.backend_step(py, seed); + } + Phase::AwaitingBackend => { + let result = result.take().ok_or_else(|| { + PyRuntimeError::new_err("semantic execution expected a backend result") + })?; + let value = match result { + Ok(value) => value, + Err(error) => return Err(error), + }; + let more = matches!( + &self.operation, + SemanticOperation::StoreBatch(queue) if !queue.is_empty() + ); + if more { + self.phase = Phase::Start; + continue; + } + return Ok(ExecutionStep::Return(value)); + } + } + } + }) + } + + fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> { + if let Some(embedder) = self.service.semantic_embedder() { + embedder.traverse(visit)?; + } + Ok(()) + } +} + +pub(super) fn drive(py: Python<'_>, body: SemanticBody) -> PyResult> { + let execution = Py::new(py, Execution::new(body))?; + py.import("litellm.rust_bridge.lifecycle")? + .getattr("drive")? + .call1((execution,)) +} diff --git a/litellm-rust/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/constants.py b/litellm/constants.py index 5e86927ef56..1012ca831f2 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1749,6 +1749,12 @@ SPEND_LOG_CLEANUP_BATCH_FAILURE_BACKOFF_SECONDS: Final = float( SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_RUN_BUDGET_SECONDS", "300")) SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS: Final = float(os.getenv("SPEND_LOG_CLEANUP_BATCH_TIMEOUT_SECONDS", "30")) SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP: Final = int(os.getenv("SPEND_LOG_CLEANUP_REMAINING_COUNT_CAP", "100000")) +SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS: Final = float( + os.getenv("SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS", "5") +) +SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS: Final = float( + os.getenv("SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS", "5") +) TOOL_SPEND_TOP_TOOLS: Final = 100 SPEND_LOG_PARTITION_INTERVAL: Final = os.getenv("SPEND_LOG_PARTITION_INTERVAL", "day") SPEND_LOG_PARTITION_PRECREATE_AHEAD: Final = int(os.getenv("SPEND_LOG_PARTITION_PRECREATE_AHEAD", 7)) 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/llms/azure/azure.py b/litellm/llms/azure/azure.py index 449319c1b95..2e7e7bb0c9d 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -481,7 +481,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): additional_args={"complete_input_dict": data}, original_response=str(e), ) - raise AzureOpenAIError(status_code=500, message=str(e)) + raise except Exception as e: message: Final = getattr(e, "message", str(e)) body: Final = getattr(e, "body", None) diff --git a/litellm/llms/openai/videos/guardrail_translation/__init__.py b/litellm/llms/openai/videos/guardrail_translation/__init__.py new file mode 100644 index 00000000000..7bd869612d6 --- /dev/null +++ b/litellm/llms/openai/videos/guardrail_translation/__init__.py @@ -0,0 +1,23 @@ +"""OpenAI Video Generation handler for Unified Guardrails.""" + +from typing import Final + +from litellm.llms.openai.videos.guardrail_translation.handler import ( + OpenAIVideoGenerationHandler, +) +from litellm.types.utils import CallTypes + +guardrail_translation_mappings: Final = { # mutable-ok: discover_guardrail_translation_mappings only accepts isinstance(mappings, dict) + CallTypes.video_generation: OpenAIVideoGenerationHandler, + CallTypes.avideo_generation: OpenAIVideoGenerationHandler, + CallTypes.create_video: OpenAIVideoGenerationHandler, + CallTypes.acreate_video: OpenAIVideoGenerationHandler, + CallTypes.video_remix: OpenAIVideoGenerationHandler, + CallTypes.avideo_remix: OpenAIVideoGenerationHandler, + CallTypes.video_edit: OpenAIVideoGenerationHandler, + CallTypes.avideo_edit: OpenAIVideoGenerationHandler, + CallTypes.video_extension: OpenAIVideoGenerationHandler, + CallTypes.avideo_extension: OpenAIVideoGenerationHandler, +} + +__all__ = ("OpenAIVideoGenerationHandler", "guardrail_translation_mappings") diff --git a/litellm/llms/openai/videos/guardrail_translation/handler.py b/litellm/llms/openai/videos/guardrail_translation/handler.py new file mode 100644 index 00000000000..49a8d05100c --- /dev/null +++ b/litellm/llms/openai/videos/guardrail_translation/handler.py @@ -0,0 +1,48 @@ +from typing import TYPE_CHECKING, Final + +from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth + + +class OpenAIVideoGenerationHandler(BaseTranslation): + async def process_input_messages( + self, + data: dict[str, object], # mutable-ok: BaseTranslation contract passes the proxy's request dict through + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> dict[str, object]: # mutable-ok: BaseTranslation contract returns the proxy's request dict + prompt: Final = data.get("prompt") + if not isinstance(prompt, str): + return data + + model: Final = data.get("model") + texts: Final = [prompt] # mutable-ok: GenericGuardrailAPIInputs.texts is declared list[str] + inputs: Final = ( + GenericGuardrailAPIInputs(texts=texts, model=model) + if isinstance(model, str) + else GenericGuardrailAPIInputs(texts=texts) + ) + guardrailed_inputs: Final = await guardrail_to_apply.apply_guardrail( # pyright: ignore[reportUnknownMemberType] # request_data is a bare dict + inputs=inputs, + request_data=data, + input_type="request", + logging_obj=litellm_logging_obj, + ) + guardrailed_texts: Final = guardrailed_inputs.get("texts") + guardrailed_prompt: Final = guardrailed_texts[0] if guardrailed_texts else prompt + return {**data, "prompt": guardrailed_prompt} # mutable-ok: BaseTranslation contract returns a dict + + async def process_output_response( + self, + response: object, + guardrail_to_apply: "CustomGuardrail", + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, + request_data: dict[str, object] | None = None, # mutable-ok: BaseTranslation contract + ) -> object: + return response diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 35c61871e09..80bb2bf9bd2 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.87226e-07, + "input_cost_per_token": 9.5526e-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.774452e-06, + "output_cost_per_token": 1.91052e-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.39355e-08, + "cache_read_input_token_cost": 7.9605e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -43079,22 +43079,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 5.58624e-07, "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.675872e-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": 4.4e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "cache_read_input_token_cost": 1.86208e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.58624e-7,"output_cost_per_token":0.000001675872,"cache_read_input_token_cost":1.86208e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -68941,9 +68941,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 5.544e-08, - "output_cost_per_token": 1.1088e-07, - "cache_read_input_token_cost": 1.1088e-08, + "input_cost_per_token": 8.8606e-08, + "output_cost_per_token": 1.77212e-07, + "cache_read_input_token_cost": 1.77212e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -70299,8 +70299,8 @@ "supports_web_search": true }, "openrouter/meta-llama/llama-4-maverick": { - "input_cost_per_token": 2e-07, - "output_cost_per_token": 8e-07, + "input_cost_per_token": 1.875e-07, + "output_cost_per_token": 6.525e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 16384, @@ -72999,15 +72999,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 4.4e-08, - "input_cost_per_token": 1.32e-06, + "cache_read_input_token_cost": 1.86208e-08, + "input_cost_per_token": 5.58624e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, - "output_cost_per_token": 3.96e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.58624e-7,"output_cost_per_token":0.000001675872,"cache_read_input_token_cost":1.86208e-8}, + "output_cost_per_token": 1.675872e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -76879,6 +76879,26 @@ "supports_vision": true, "supports_web_search": true }, + "moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "global.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index b0640e4f0dd..2e5cad29b0a 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -44,6 +44,11 @@ from litellm.proxy._types import ( UserAPIKeyAuth, user_api_key_has_admin_view, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + CeilingResolver, + resolve_agent_access_group_ceiling, +) +from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.auth.user_api_key_auth import ( _get_bearer_token_or_received_api_key, # pyright: ignore[reportPrivateUsage] # shared x-litellm-api-key parser lives with user_api_key_auth @@ -184,6 +189,21 @@ def _has_client_supplied_mcp_auth( return bool(mcp_auth_header) or bool(mcp_server_auth_headers) +def _agent_capped_servers( + allowed_mcp_servers: Sequence[str], + agent_servers: Sequence[str], + agent_access_group_servers: frozenset[str] | None, +) -> tuple[str, ...] | None: + if not agent_servers and agent_access_group_servers is None: + return None + return tuple( + s + for s in allowed_mcp_servers + if (not agent_servers or s in agent_servers) + and (agent_access_group_servers is None or s in agent_access_group_servers) + ) + + def _is_mcp_admitted_user_subject(user_api_key_auth: UserAPIKeyAuth | None) -> bool: """True when this auth is a keyless subject admitted by the gateway session / bridge user path, as opposed to a JWT or other keyless auth that merely lacks a ``team_id``. @@ -1546,25 +1566,33 @@ class MCPRequestHandler: # Check agent permissions if agent_id is set on the key ######################################################### if user_api_key_auth and user_api_key_auth.agent_id: - allowed_mcp_servers_for_agent: Final = await MCPRequestHandler._get_allowed_mcp_servers_for_agent( - user_api_key_auth + agent_capped: Final = _agent_capped_servers( + allowed_mcp_servers, + await MCPRequestHandler._get_allowed_mcp_servers_for_agent(user_api_key_auth), + await MCPRequestHandler._get_agent_access_group_server_ceiling(user_api_key_auth), ) - if len(allowed_mcp_servers_for_agent) > 0: + if agent_capped is not None: has_lower_level_mcp_restrictions = True - # Intersect: agent can only use servers allowed by BOTH key/team AND agent config - allowed_mcp_servers = [s for s in allowed_mcp_servers if s in allowed_mcp_servers_for_agent] + allowed_mcp_servers = list(agent_capped) verbose_logger.debug( "Applied agent intersection filter. Final allowed servers: %s", allowed_mcp_servers ) + ######################################################### + # Cap an agent key at what the user and team that invoked the agent may reach + ######################################################### + caller_capped, caller_restricts = await MCPRequestHandler._apply_agent_caller_ceiling( + allowed_mcp_servers, user_api_key_auth + ) + ######################################################### # Apply the internal user's own ceiling (the entitlement attached to the human) ######################################################### capped, user_restricts = await MCPRequestHandler._apply_user_server_ceiling( - allowed_mcp_servers, user_api_key_auth, keyless_source=keyless_source + caller_capped, user_api_key_auth, keyless_source=keyless_source ) allowed_mcp_servers = list(capped) - has_lower_level_mcp_restrictions = has_lower_level_mcp_restrictions or user_restricts + has_lower_level_mcp_restrictions = has_lower_level_mcp_restrictions or caller_restricts or user_restricts ######################################################### # Apply org-level ceiling if org_id is set @@ -2907,6 +2935,28 @@ class MCPRequestHandler: verbose_logger.debug("Applied user ceiling filter. Final allowed servers: %s", capped) return capped, True + @staticmethod + async def _apply_agent_caller_ceiling( + allowed_mcp_servers: Sequence[str], + user_api_key_auth: UserAPIKeyAuth | None = None, + ) -> tuple[tuple[str, ...], bool]: + """Narrow an agent key's servers to those the invoking user and team (echoed back by the agent + as ``x-litellm-user-id`` / ``x-litellm-team-id``) may reach: the echoed team's grants when it + names any, then the echoed user's own entitlement. Raises like the user ceiling when that + entitlement is known but unreadable, so the resolver denies rather than widens.""" + caller_auth: Final = agent_caller_auth(user_api_key_auth) if user_api_key_auth else None + if caller_auth is None: + return tuple(allowed_mcp_servers), False + team_servers: Final = frozenset(await MCPRequestHandler._get_allowed_mcp_servers_for_team(caller_auth)) + team_capped: Final = ( + tuple(server for server in allowed_mcp_servers if server in team_servers) + if team_servers + else tuple(allowed_mcp_servers) + ) + user_capped, user_restricts = await MCPRequestHandler._apply_user_server_ceiling(team_capped, caller_auth) + verbose_logger.debug("Applied agent caller ceiling. Final allowed servers: %s", user_capped) + return user_capped, bool(team_servers) or user_restricts + @staticmethod async def _user_places_mcp_ceiling(user_api_key_auth: UserAPIKeyAuth | None = None) -> bool: """Whether this human's own entitlement bounds their MCP access at all. @@ -3137,6 +3187,27 @@ class MCPRequestHandler: verbose_logger.warning("Failed to get allowed MCP servers for agent: %s", e) return [] + @staticmethod + async def _get_agent_access_group_server_ceiling( + user_api_key_auth: UserAPIKeyAuth, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, + ) -> frozenset[str] | None: + """ + Server IDs the agent's attached unified access groups (``LiteLLM_AgentsTable.access_group_ids``) + allow, or None when the agent has none attached. Unlike the object_permission path above, an + attached group set that names no servers is an empty ceiling and denies every server. + """ + from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( + global_mcp_server_manager, + ) + + if not user_api_key_auth.agent_id: + return None + ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id) + if ceiling is None: + return None + return frozenset(global_mcp_server_manager.expand_permission_list(sorted(ceiling.mcp_server_ids))) + @staticmethod async def _get_agent_tool_permissions_for_server( server_id: str, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 6f9a2d8c96d..1293a05e4a6 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -2357,6 +2357,20 @@ }, "AgentConfig": { "properties": { + "access_group_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Access Group Ids" + }, "agent_card_params": { "$ref": "#/components/schemas/AgentCard" }, @@ -2683,6 +2697,20 @@ }, "AgentResponse": { "properties": { + "access_group_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Access Group Ids" + }, "agent_card_params": { "additionalProperties": true, "title": "Agent Card Params", @@ -3506,6 +3534,20 @@ }, "PatchAgentRequest": { "properties": { + "access_group_ids": { + "anyOf": [ + { + "items": { + "type": "string" + }, + "type": "array" + }, + { + "type": "null" + } + ], + "title": "Access Group Ids" + }, "agent_card_params": { "$ref": "#/components/schemas/AgentCard" }, diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 209d5267dec..c8e8314e966 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -28,6 +28,7 @@ from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( validate_langfuse_span_scope_value, validate_no_callback_env_reference, ) +from litellm.types.agents import AgentCaller from litellm.types.integrations.compression_interception import ( CompressionSavingsMetadata, ) @@ -3341,6 +3342,15 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob "user id." ), ) + agent_caller: AgentCaller | None = Field( + default=None, + exclude=True, + description=( + "Set per request from the x-litellm-user-id / x-litellm-team-id headers an agent echoes back on " + "calls made with its own key. Every check treats it as a ceiling, so a forged value can only " + "narrow the agent's access." + ), + ) budget_reservation: dict[str, Any] | None = Field(default=None, exclude=True) team_budget_snapshot: TeamBudgetSnapshot | None = Field(default=None, exclude=True) user_budget_snapshot: UserBudgetSnapshot | None = Field(default=None, exclude=True) @@ -3373,6 +3383,7 @@ class UserAPIKeyAuth(LiteLLM_VerificationTokenView): # the expected response ob values.pop("mcp_session_resource_server_id", None) values.pop("mcp_toolset_id", None) values.pop("via_virtual_key", None) + values.pop("agent_caller", None) if values.get("api_key") is not None: values.update({"token": cls._safe_hash_litellm_api_key(values.get("api_key"))}) if isinstance(values.get("api_key"), str): @@ -4328,6 +4339,11 @@ class ProxyErrorTypes(str, enum.Enum): Project does not have access to the model """ + agent_model_access_denied = "agent_model_access_denied" + """ + The agent behind the key does not have access to the model + """ + model_cost_map_missing = "model_cost_map_missing" expired_key = "expired_key" @@ -4402,7 +4418,7 @@ class ProxyErrorTypes(str, enum.Enum): @classmethod def get_model_access_error_type_for_object( - cls, object_type: Literal["key", "user", "team", "org", "project"] + cls, object_type: Literal["key", "user", "team", "org", "project", "agent"] ) -> "ProxyErrorTypes": """ Get the model access error type for object_type @@ -4417,6 +4433,8 @@ class ProxyErrorTypes(str, enum.Enum): return cls.org_model_access_denied elif object_type == "project": return cls.project_model_access_denied + elif object_type == "agent": + return cls.agent_model_access_denied @classmethod def get_vector_store_access_error_type_for_object( diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 834c16ba6dc..2a189a76545 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -146,12 +146,17 @@ def _validate_push_notification_url(url: str) -> None: def _caller_identity_headers(user_api_key_dict: UserAPIKeyAuth) -> Mapping[str, str]: + """The human behind this call. An agent key acting for an invoking user forwards that user, not + itself, so a chain of agents stays capped at what the original caller may reach.""" + caller: Final = user_api_key_dict.agent_caller + user_id: Final = caller.user_id if caller is not None else user_api_key_dict.user_id + team_id: Final = caller.team_id if caller is not None else user_api_key_dict.team_id return MappingProxyType( { name: value for name, value in ( - ("X-LiteLLM-User-Id", user_api_key_dict.user_id), - ("X-LiteLLM-Team-Id", user_api_key_dict.team_id), + ("X-LiteLLM-User-Id", user_id), + ("X-LiteLLM-Team-Id", team_id), ) if value } diff --git a/litellm/proxy/agent_endpoints/agent_registry.py b/litellm/proxy/agent_endpoints/agent_registry.py index c7b6bca72cf..d6b12e830e1 100644 --- a/litellm/proxy/agent_endpoints/agent_registry.py +++ b/litellm/proxy/agent_endpoints/agent_registry.py @@ -7,6 +7,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, NamedTuple, Protocol, TypedDict from pydantic import TypeAdapter, ValidationError +from typing_extensions import ReadOnly import litellm from litellm.constants import REDACTED_BY_LITELM_STRING @@ -37,6 +38,7 @@ class AgentRecordDump(TypedDict): agent_card_params: dict[str, object] static_headers: dict[str, str] | None extra_headers: list[str] | None + access_group_ids: ReadOnly[Sequence[str] | None] object_permission: dict[str, object] | None spend: float tpm_limit: int | None @@ -65,6 +67,9 @@ class AgentRecord(Protocol): @property def object_permission(self) -> AgentObjectPermissionRecord | None: ... + @property + def access_group_ids(self) -> Sequence[str] | None: ... + @property def spend(self) -> float: ... @@ -284,6 +289,12 @@ def _resolved_agent_param_value( return _MISSING_AGENT_PARAM +def _patched_access_group_ids(agent: PatchAgentRequest) -> Mapping[str, object]: + if "access_group_ids" not in agent: + return MappingProxyType({}) + return MappingProxyType({"access_group_ids": tuple(dict.fromkeys(agent.get("access_group_ids") or ()))}) + + def _restore_redacted_litellm_params( incoming: Mapping[str, object], existing: Mapping[str, object], @@ -516,6 +527,7 @@ class AgentRegistry: static_headers_val: Final[str | None] = safe_dumps(dict(static_headers_obj)) if static_headers_obj else None extra_headers_val: Final = agent.get("extra_headers") + access_group_ids_val: Final = agent.get("access_group_ids") create_data: Final[dict[str, object]] = { "agent_name": agent_name, @@ -532,6 +544,8 @@ class AgentRegistry: create_data["static_headers"] = static_headers_val if extra_headers_val is not None: create_data["extra_headers"] = extra_headers_val + if access_group_ids_val is not None: + create_data["access_group_ids"] = tuple(dict.fromkeys(access_group_ids_val)) if object_permission_id is not None: create_data["object_permission_id"] = object_permission_id @@ -601,7 +615,7 @@ class AgentRegistry: existing_agent: Final[Mapping[str, object]] = dict(existing_record) augment_agent: Final = {**existing_agent, **agent} - update_data: Final[dict[str, object]] = {} + update_data: Final[dict[str, object]] = {**_patched_access_group_ids(agent)} if augment_agent.get("agent_name"): update_data["agent_name"] = augment_agent.get("agent_name") if "litellm_params" in agent: @@ -703,6 +717,7 @@ class AgentRegistry: safe_dumps(dict(static_headers_obj_u)) if static_headers_obj_u is not None else safe_dumps({}) ) extra_headers_val_u: Final = agent.get("extra_headers") or [] + access_group_ids_val_u: Final = tuple(dict.fromkeys(agent.get("access_group_ids") or ())) update_data: Final[dict[str, object]] = { "agent_name": agent_name, @@ -710,6 +725,7 @@ class AgentRegistry: "agent_card_params": agent_card_params, "static_headers": static_headers_val_u, "extra_headers": extra_headers_val_u, + "access_group_ids": access_group_ids_val_u, "updated_by": updated_by, "updated_at": datetime.now(timezone.utc), } diff --git a/litellm/proxy/agent_endpoints/auth/agent_access_groups.py b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py new file mode 100644 index 00000000000..49e5407ff88 --- /dev/null +++ b/litellm/proxy/agent_endpoints/auth/agent_access_groups.py @@ -0,0 +1,75 @@ +import asyncio +from collections.abc import Awaitable, Callable +from dataclasses import dataclass +from typing import Final, TypeAlias + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_AccessGroupTable + +AccessGroupIds: TypeAlias = tuple[str, ...] +AccessGroupIdsLoader: TypeAlias = Callable[[str], Awaitable[AccessGroupIds]] # mutable-ok: Callable params +LoadedAccessGroup: TypeAlias = LiteLLM_AccessGroupTable | None +AccessGroupLoader: TypeAlias = Callable[[str], Awaitable[LoadedAccessGroup]] # mutable-ok: Callable parameter syntax + + +@dataclass(frozen=True, slots=True) +class AgentAccessGroupCeiling: + """Everything the agent's attached access groups allow. An empty set denies that resource kind.""" + + access_group_ids: AccessGroupIds + models: frozenset[str] + mcp_server_ids: frozenset[str] + agent_ids: frozenset[str] + + +CeilingResolver: TypeAlias = Callable[[str], Awaitable[AgentAccessGroupCeiling | None]] # mutable-ok: Callable params + + +async def _registry_access_group_ids(agent_id: str) -> AccessGroupIds: + from litellm.proxy.common_utils.registry_read_through import get_agent_with_read_through + + agent: Final = await get_agent_with_read_through(agent_id) + return tuple(agent.access_group_ids or ()) if agent is not None else () + + +async def _load_access_group(access_group_id: str) -> LoadedAccessGroup: + from litellm.proxy.auth.auth_checks import get_access_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + if prisma_client is None: + verbose_proxy_logger.warning("Agent access group %s cannot be loaded without a DB", access_group_id) + return None + try: + return await get_access_object( + access_group_id=access_group_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException as e: + verbose_proxy_logger.warning( + "Agent access group %s could not be loaded, treating it as empty: %s", access_group_id, e.detail + ) + return None + + +async def resolve_agent_access_group_ceiling( + agent_id: str, + load_access_group_ids: AccessGroupIdsLoader = _registry_access_group_ids, + load_access_group: AccessGroupLoader = _load_access_group, +) -> AgentAccessGroupCeiling | None: + """``None`` when the agent has no access groups attached, so nothing is capped.""" + access_group_ids: Final = await load_access_group_ids(agent_id) + if not access_group_ids: + return None + + loaded: Final = await asyncio.gather(*(load_access_group(group_id) for group_id in access_group_ids)) + groups: Final = tuple(group for group in loaded if group is not None) + return AgentAccessGroupCeiling( + access_group_ids=access_group_ids, + models=frozenset(model for group in groups for model in group.access_model_names), + mcp_server_ids=frozenset(server_id for group in groups for server_id in group.access_mcp_server_ids), + agent_ids=frozenset(target_id for group in groups for target_id in group.access_agent_ids), + ) diff --git a/litellm/proxy/agent_endpoints/auth/agent_caller.py b/litellm/proxy/agent_endpoints/auth/agent_caller.py new file mode 100644 index 00000000000..47d43e8f71b --- /dev/null +++ b/litellm/proxy/agent_endpoints/auth/agent_caller.py @@ -0,0 +1,87 @@ +"""The human behind an agent's own proxy calls. + +``/a2a/{agent}`` forwards the invoking key's ``X-LiteLLM-User-Id`` / ``X-LiteLLM-Team-Id`` to the +agent backend. When the agent echoes them back on requests made with its own key, the proxy caps +that key at what the invoking user and team may reach. The cap is intersected with, never +substituted for, the agent key's own grants and the agent's access group ceiling, so the headers +can only narrow access and need no trust. +""" + +from collections.abc import Mapping +from typing import Final + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import LiteLLM_TeamTable, LiteLLM_UserTable, UserAPIKeyAuth +from litellm.types.agents import ( + AGENT_CALLER_TEAM_ID_HEADER, + AGENT_CALLER_USER_ID_HEADER, + AgentCaller, +) + + +def _header(headers: Mapping[str, str], name: str) -> str | None: + value: Final = next((raw for key, raw in headers.items() if key.lower() == name), None) + return value.strip() or None if value is not None else None + + +def agent_caller_from_headers(headers: Mapping[str, str], user_api_key_auth: UserAPIKeyAuth) -> AgentCaller | None: + """The caller an agent key is acting for, or ``None`` when the key is not an agent's or no id was echoed.""" + if not user_api_key_auth.agent_id: + return None + user_id: Final = _header(headers, AGENT_CALLER_USER_ID_HEADER) + team_id: Final = _header(headers, AGENT_CALLER_TEAM_ID_HEADER) + if user_id is None and team_id is None: + return None + return AgentCaller(user_id=user_id, team_id=team_id) + + +def agent_caller_auth(user_api_key_auth: UserAPIKeyAuth) -> UserAPIKeyAuth | None: + """A minimal auth context standing for the invoking user and team, so the shared key/team/user + resolvers can be reused unchanged to compute what the caller may reach.""" + caller: Final = user_api_key_auth.agent_caller + if caller is None: + return None + return UserAPIKeyAuth( + user_id=caller.user_id, + team_id=caller.team_id, + parent_otel_span=user_api_key_auth.parent_otel_span, + ) + + +async def load_agent_caller_team(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: + """The invoking team's row, or ``None`` when no team id was echoed. Raises when the id names a team + that cannot be loaded, since a caller we cannot resolve must not be treated as unrestricted.""" + from litellm.proxy.auth.auth_checks import get_team_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + caller: Final = user_api_key_auth.agent_caller + if caller is None or caller.team_id is None: + return None + return await get_team_object( + team_id=caller.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + + +async def load_agent_caller_user(user_api_key_auth: UserAPIKeyAuth) -> LiteLLM_UserTable | None: + """The invoking user's row, or ``None`` when no user id was echoed or the row does not exist.""" + from litellm.proxy.auth.auth_checks import get_user_object + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + caller: Final = user_api_key_auth.agent_caller + if caller is None or caller.user_id is None: + return None + user_object: Final = await get_user_object( + user_id=caller.user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + user_id_upsert=False, + parent_otel_span=user_api_key_auth.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + if user_object is None: + verbose_proxy_logger.debug("agent caller user %r not found; no user ceiling applied", caller.user_id) + return user_object diff --git a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py index e4dd77e2f82..9fe74bfee3f 100644 --- a/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py +++ b/litellm/proxy/agent_endpoints/auth/agent_permission_handler.py @@ -19,6 +19,11 @@ from litellm.proxy._types import ( LitellmUserRoles, UserAPIKeyAuth, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + CeilingResolver, + resolve_agent_access_group_ceiling, +) +from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth from litellm.repositories.table_repositories import AgentsRepository from litellm.types.agents import AgentResponse @@ -44,6 +49,22 @@ def _to_stable_ids(agent_ids: frozenset[str]) -> frozenset[str]: return frozenset(global_agent_registry.stable_agent_id(agent_id) for agent_id in agent_ids) +def _restricted_ids(access: AgentAccess) -> frozenset[str] | None: + if isinstance(access, UnrestrictedAgentAccess): + return None + return _to_stable_ids(access.agent_ids) + + +def _intersect_agent_access(key_access: AgentAccess, team_access: AgentAccess) -> AgentAccess: + key_ids: Final = _restricted_ids(key_access) + team_ids: Final = _restricted_ids(team_access) + if key_ids is None: + return UnrestrictedAgentAccess() if team_ids is None else RestrictedAgentAccess(team_ids) + if team_ids is None: + return RestrictedAgentAccess(key_ids) + return RestrictedAgentAccess(key_ids & team_ids) + + class AgentRequestHandler: """ Class to handle agent permission checking, including: @@ -61,35 +82,56 @@ class AgentRequestHandler: @staticmethod async def resolve_agent_access( user_api_key_auth: UserAPIKeyAuth | None = None, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> AgentAccess: - """ - Resolve the agents the given user/key may reach. + """Agents the key may reach: key and team grants, intersected with the agent's access group ceiling + and, for an agent key acting on behalf of an invoking user, with that user's team grants.""" + key_team_access: Final = await AgentRequestHandler._resolve_key_team_agent_access(user_api_key_auth) + caller_access: Final = await AgentRequestHandler._agent_caller_access(user_api_key_auth) + own_access: Final = _intersect_agent_access(key_team_access, caller_access) + agent_ceiling: Final = await AgentRequestHandler._agent_access_group_ceiling(user_api_key_auth, resolve_ceiling) + if agent_ceiling is None: + return own_access + if isinstance(own_access, UnrestrictedAgentAccess): + return RestrictedAgentAccess(agent_ceiling) + return RestrictedAgentAccess(own_access.agent_ids & agent_ceiling) - ``UnrestrictedAgentAccess`` is only returned when neither the key nor its team - carries any grant. Grants that intersect to nothing stay restricted, so - narrowing a caller can never widen what it reaches. - """ + @staticmethod + async def _agent_caller_access(user_api_key_auth: UserAPIKeyAuth | None) -> AgentAccess: + caller_auth: Final = agent_caller_auth(user_api_key_auth) if user_api_key_auth else None + if caller_auth is None: + return UnrestrictedAgentAccess() + return await AgentRequestHandler._get_allowed_agents_for_team(caller_auth) + + @staticmethod + async def _resolve_key_team_agent_access( + user_api_key_auth: UserAPIKeyAuth | None, + ) -> AgentAccess: try: key_access: Final = await AgentRequestHandler._get_allowed_agents_for_key(user_api_key_auth) team_access: Final = await AgentRequestHandler._get_allowed_agents_for_team(user_api_key_auth) - - match (key_access, team_access): - case (UnrestrictedAgentAccess(), UnrestrictedAgentAccess()): - return UnrestrictedAgentAccess() - case (UnrestrictedAgentAccess(), RestrictedAgentAccess(team_ids)): - return RestrictedAgentAccess(_to_stable_ids(team_ids)) - case (RestrictedAgentAccess(key_ids), UnrestrictedAgentAccess()): - return RestrictedAgentAccess(_to_stable_ids(key_ids)) - case (RestrictedAgentAccess(key_ids), RestrictedAgentAccess(team_ids)): - return RestrictedAgentAccess(_to_stable_ids(key_ids) & _to_stable_ids(team_ids)) except Exception as e: verbose_logger.warning("Failed to get allowed agents: %s", e) return UnrestrictedAgentAccess() + return _intersect_agent_access(key_access, team_access) + + @staticmethod + async def _agent_access_group_ceiling( + user_api_key_auth: UserAPIKeyAuth | None, + resolve_ceiling: CeilingResolver, + ) -> frozenset[str] | None: + if user_api_key_auth is None or not user_api_key_auth.agent_id: + return None + ceiling: Final = await resolve_ceiling(user_api_key_auth.agent_id) + if ceiling is None: + return None + return _to_stable_ids(ceiling.agent_ids) @staticmethod async def is_agent_allowed( agent_id: str, user_api_key_auth: UserAPIKeyAuth | None = None, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> bool: """ Check if a specific agent is allowed for the given user/key. @@ -103,7 +145,7 @@ class AgentRequestHandler: """ from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry - match await AgentRequestHandler.resolve_agent_access(user_api_key_auth): + match await AgentRequestHandler.resolve_agent_access(user_api_key_auth, resolve_ceiling): case UnrestrictedAgentAccess(): return True case RestrictedAgentAccess(allowed_agent_ids): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index c2279fb2fe1..b3efde7f05f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -15,7 +15,7 @@ import re import time from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeAlias from fastapi import HTTPException, Request, status from pydantic import BaseModel, TypeAdapter @@ -68,6 +68,15 @@ from litellm.proxy._types import ( SpecialModelNames, UserAPIKeyAuth, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + CeilingResolver, + resolve_agent_access_group_ceiling, +) +from litellm.proxy.agent_endpoints.auth.agent_caller import ( + agent_caller_auth, + load_agent_caller_team, + load_agent_caller_user, +) from litellm.proxy.auth.budget_throttle import ( budget_throttle_percentage, should_throttle_budget_exceeded, @@ -1006,6 +1015,16 @@ async def common_checks( code=status.HTTP_400_BAD_REQUEST, ) + await _check_agent_access_group_model_access(model=_model, valid_token=valid_token, llm_router=llm_router) + await _check_agent_caller_model_access( + model=_model, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + ## 2.1 If user can call model (if personal key) if _model and team_object is None and user_object is not None: with tracer.trace("litellm.proxy.auth.common_checks.can_user_call_model"): @@ -4251,7 +4270,7 @@ def _can_object_call_model( models: list[str], team_model_aliases: dict[str, str] | None = None, team_id: str | None = None, - object_type: Literal["user", "team", "key", "org", "project"] = "user", + object_type: Literal["user", "team", "key", "org", "project", "agent"] = "user", fallback_depth: int = 0, ) -> Literal[True]: """ @@ -4317,6 +4336,82 @@ def _can_object_call_model( ) +async def _check_agent_access_group_model_access( + model: str | list[str] | None, # mutable-ok: _can_object_call_model and the client message helper take list[str] + valid_token: UserAPIKeyAuth | None, + llm_router: Router | None, + resolve_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, +) -> Literal[True]: + """Attached groups naming no model deny every model; the empty allowlist in ``_can_object_call_model`` allows.""" + if not model or valid_token is None or not valid_token.agent_id: + return True + ceiling: Final = await resolve_ceiling(valid_token.agent_id) + if ceiling is None: + return True + if not ceiling.models: + raise ModelAccessDeniedProxyException( + message=model_access_denied_client_message(model=model), + internal_message=f"agent {valid_token.agent_id} access groups {ceiling.access_group_ids} grant no models", + type=ProxyErrorTypes.agent_model_access_denied, + param="model", + code=status.HTTP_403_FORBIDDEN, + ) + return _can_object_call_model( + model=model, + llm_router=llm_router, + models=sorted(ceiling.models), + team_id=valid_token.team_id, + object_type="agent", + ) + + +LoadedCallerTeam: TypeAlias = LiteLLM_TeamTable | None +LoadedCallerUser: TypeAlias = LiteLLM_UserTable | None +CallerTeamLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerTeam]] # mutable-ok: Callable params +CallerUserLoader: TypeAlias = Callable[[UserAPIKeyAuth], Awaitable[LoadedCallerUser]] # mutable-ok: Callable params + + +async def _check_agent_caller_model_access( + model: str | list[str] | None, # mutable-ok: the model checks it delegates to take list[str] + valid_token: UserAPIKeyAuth | None, + llm_router: Router | None, + prisma_client: Optional["PrismaClient"], + user_api_key_cache: UserApiKeyCache, + proxy_logging_obj: ProxyLogging, + load_team: CallerTeamLoader = load_agent_caller_team, + load_user: CallerUserLoader = load_agent_caller_user, +) -> None: + """An agent key acting for an invoking user may call only what that user's own key could: the + invoking team's models (and per-member scope) when a team was echoed, else the user's models.""" + if not model or valid_token is None: + return + caller_auth: Final = agent_caller_auth(valid_token) + if caller_auth is None: + return + caller_team: Final = await load_team(valid_token) + if caller_team is not None: + await can_team_access_model( + model=model, + team_object=caller_team, + llm_router=llm_router, + prisma_client=prisma_client, + ) + await _check_team_member_model_access( + model=model, + team_object=caller_team, + valid_token=caller_auth, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + return + caller_user: Final = await load_user(valid_token) + if caller_user is None: + return + await can_user_call_model(model=model, llm_router=llm_router, user_object=caller_user) + + def _model_in_team_aliases(model: str, team_model_aliases: dict[str, str] | None = None) -> bool: """ Returns True if `model` being accessed is an alias of a team model 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/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 4371ce4fda8..08f64e610bd 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -39,6 +39,7 @@ from litellm.integrations.otel.runtime import phase_span, seed_request_identity from litellm.litellm_core_utils.dd_tracing import tracer from litellm.litellm_core_utils.dot_notation_indexing import get_nested_value from litellm.proxy._types import * +from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_from_headers from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, TeamNotFoundError, @@ -3330,6 +3331,9 @@ async def user_api_key_auth( raise body_parse_exception raise user_api_key_auth_obj.budget_reservation = None + user_api_key_auth_obj.agent_caller = agent_caller_from_headers( + _safe_get_request_headers(request), user_api_key_auth_obj + ) _seed_request_destinations(user_api_key_auth_obj, request) # A body that never parsed is authenticated (so the trace carries identity diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index f4bebc4a4cb..b3fdc4695cb 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -22,8 +22,12 @@ from pathlib import Path from types import MappingProxyType from typing import Final, TypeAlias +import click +from filelock import FileLock +from packaging.version import InvalidVersion, Version from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError +from litellm._version import version as litellm_version from litellm.litellm_core_utils.private_json import ( commit_staged_json, discard_staged_json, @@ -75,6 +79,7 @@ BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" AUTOROUTE_BACKUP_PATH: Final = Path.home() / ".litellm" / "autorouter" / "claude_settings_backup.json" CONFIGURE_STATE_PATH: Final = Path.home() / ".litellm" / "claude_configure_state.json" STATUSLINE_SCRIPT_PATH: Final = Path.home() / ".litellm" / "statusline.py" +STATUSLINE_VERSION_PREFIX: Final = b"# litellm-statusline-version: " @dataclass(frozen=True, slots=True) @@ -305,11 +310,54 @@ def statusline_command(script_path: Path, platform: str = sys.platform) -> str: return " ".join(quote(token) for token in (sys.executable, str(script_path))) -def install_statusline_script(script_path: Path | None = None) -> str: +def _statusline_version(value: str) -> Version | None: + try: + return Version(value) + except InvalidVersion: + return None + + +def _installed_statusline_version(target: Path) -> Version | None: + try: + with target.open("rb") as script: + header: Final = script.readline(256) + except FileNotFoundError: + return None + if not header.startswith(STATUSLINE_VERSION_PREFIX): + return None + try: + return _statusline_version(header.removeprefix(STATUSLINE_VERSION_PREFIX).decode("ascii").strip()) + except UnicodeDecodeError: + return None + + +def install_statusline_script( + script_path: Path | None = None, + *, + package_version: str = litellm_version, + write: Callable[[str, bytes], None] = write_private_bytes, +) -> str: target: Final = script_path or STATUSLINE_SCRIPT_PATH try: ensure_private_dir(target.parent) - write_private_bytes(str(target), Path(statusline_script.__file__).read_bytes()) + bundled_version: Final = _statusline_version(package_version) + with FileLock(str(target) + ".lock", timeout=10, mode=0o600): + installed_version: Final = _installed_statusline_version(target) + if installed_version is not None and (bundled_version is None or installed_version > bundled_version): + cli_version: Final = str(bundled_version) if bundled_version is not None else "unknown" + click.echo( + f"Keeping the status line from LiteLLM {installed_version}; this CLI is {cli_version}. " + "Upgrade the CLI to refresh it.", + err=True, + ) + return statusline_command(target) + source: Final = Path(statusline_script.__file__).read_bytes() + header: Final = ( + STATUSLINE_VERSION_PREFIX + str(bundled_version).encode("ascii") + b"\n" + if bundled_version is not None + else b"" + ) + write(str(target), header + source) except OSError as e: raise ClaudeSettingsError(f"Could not install the status line script at {target}: {e}") from e return statusline_command(target) diff --git a/litellm/proxy/client/cli/commands/statusline_script.py b/litellm/proxy/client/cli/commands/statusline_script.py index 09dd062c888..d16160b1ab8 100644 --- a/litellm/proxy/client/cli/commands/statusline_script.py +++ b/litellm/proxy/client/cli/commands/statusline_script.py @@ -1,7 +1,7 @@ """Claude Code status line and Codex Stop hook for auto-routed sessions. -`lite` copies this file verbatim to ~/.litellm/statusline.py and registers it as Claude -Code's `statusLine` command and as Codex's `[[hooks.Stop]]` command, so it must stay +`lite` copies this file to ~/.litellm/statusline.py with a CLI version header when known and registers +it as Claude Code's `statusLine` command and as Codex's `[[hooks.Stop]]` command, so it must stay standard-library only and must never import litellm. Claude Code re-runs it on every status refresh (about every 300ms while typing), so the proxy is asked at most once per TTL per session and every other refresh is served from a small on-disk cache that holds diff --git a/litellm/proxy/config_resolvers/__init__.py b/litellm/proxy/config_resolvers/__init__.py index eee760df458..f031a8d74b9 100644 --- a/litellm/proxy/config_resolvers/__init__.py +++ b/litellm/proxy/config_resolvers/__init__.py @@ -5,6 +5,17 @@ from litellm.proxy.config_resolvers._descriptors import ( FieldSource, resolve_fields, ) -from litellm.proxy.config_resolvers.settings_store import SettingsStore, config_ownership_message +from litellm.proxy.config_resolvers.settings_store import ( + SettingsStore, + config_ownership_message, + source_for, +) -__all__ = ("FieldDescriptor", "FieldSource", "SettingsStore", "config_ownership_message", "resolve_fields") +__all__ = ( + "FieldDescriptor", + "FieldSource", + "SettingsStore", + "config_ownership_message", + "resolve_fields", + "source_for", +) diff --git a/litellm/proxy/config_resolvers/settings_store.py b/litellm/proxy/config_resolvers/settings_store.py index 291000b3b6a..f05af3de03a 100644 --- a/litellm/proxy/config_resolvers/settings_store.py +++ b/litellm/proxy/config_resolvers/settings_store.py @@ -176,3 +176,10 @@ class SettingsStore(MutableMapping[str, JsonValue]): def _resolution_for(self, key: str) -> Resolved: yaml_value: Final[SettingValue] = self._yaml_values.get(key, ABSENT) return resolve(yaml_value, self._db_value(key)) + + +def source_for(settings: SettingsStore, key: str, default: object = None) -> FieldSource: + source: Final = settings.source(key) + if source == "unset": + return "default" if default is not None else "unset" + return source diff --git a/litellm/proxy/db/autorouter_session_rollup.py b/litellm/proxy/db/autorouter_session_rollup.py index 0d812ee812a..dd08cfd1bef 100644 --- a/litellm/proxy/db/autorouter_session_rollup.py +++ b/litellm/proxy/db/autorouter_session_rollup.py @@ -4,7 +4,7 @@ Per-session auto-router benchmarks rollup. At request time the spend writer builds one AutoRouterTurnTransaction per successful auto-routed request (a request whose metadata carries a routing_decision) and queues it on the prisma client. The spend-log flush job drains the queue into -LiteLLM_AutoRouterSession with one conditional upsert per turn: the statement classifies +key and user session rollups with one atomic statement per turn: each upsert classifies the turn (same model, first visit, return to a model the session already used, out of order) against the row's own columns, so nothing is read before the write and concurrent pods compose. The benchmarks endpoint aggregates these rows and never touches @@ -35,10 +35,27 @@ if TYPE_CHECKING: CACHE_TTL_5M_SECONDS: Final = 300 CACHE_TTL_1H_SECONDS: Final = 3600 -AUTOROUTER_BENCHMARKS_SQL: Final = """ +_SESSION_COLUMNS: Final = """ + api_key, session_id, router_name, router_type, first_turn_at, last_turn_at, + last_model, models, turns, unordered_turns, covered_turns, cache_hits, + same_model_turns, same_model_hits, first_visit_turns, first_visit_hits, + return_turns, return_hits, return_expired_misses, return_within_ttl_misses, + ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns, + baseline_models, savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend, + savings_estimated_baseline_models +""" + +AUTOROUTER_BENCHMARKS_SQL: Final = f""" WITH windowed AS ( - SELECT * FROM "LiteLLM_AutoRouterSession" - WHERE last_turn_at >= $1::timestamp + SELECT {_SESSION_COLUMNS} FROM "LiteLLM_AutoRouterSession" + WHERE $4::text IS NULL + AND last_turn_at >= $1::timestamp + AND first_turn_at < $2::timestamp + AND ($3::text IS NULL OR api_key = $3::text) + UNION ALL + SELECT {_SESSION_COLUMNS} FROM "LiteLLM_AutoRouterUserSession" + WHERE (($4::text IS NOT NULL AND user_id = $4::text) OR ($4::text IS NULL AND api_key = '')) + AND last_turn_at >= $1::timestamp AND first_turn_at < $2::timestamp AND ($3::text IS NULL OR api_key = $3::text) ), @@ -53,7 +70,7 @@ tier_maps AS ( ) SELECT agg.*, - COALESCE(tier_maps.tier_turns, '{}'::jsonb) AS tier_turns + COALESCE(tier_maps.tier_turns, '{{}}'::jsonb) AS tier_turns FROM ( SELECT router_name, @@ -111,6 +128,7 @@ class AutoRouterTurnTransaction: savings_estimated_turns: int = 0 savings_estimated_actual_spend: float = 0.0 savings_estimated_saved_spend: float = 0.0 + user_id: str = "" class TurnCacheFacts(NamedTuple): @@ -214,10 +232,11 @@ def build_autorouter_turn_transaction( if not isinstance(routing_decision, Mapping) or not routing_decision: return None router_name: Final = routing_decision.get("router_model_name") or payload.get("model_group") - api_key: Final = payload.get("api_key") + api_key: Final = payload.get("api_key") or "" + user_id: Final = payload.get("user") or "" session_id: Final = payload.get("session_id") model: Final = payload.get("model") - if not (isinstance(router_name, str) and router_name and api_key and session_id and model): + if not (isinstance(router_name, str) and router_name and (api_key or user_id) and session_id and model): return None turn_at: Final = _turn_time_utc(str(payload.get("startTime") or "")) if turn_at is None: @@ -236,6 +255,7 @@ def build_autorouter_turn_transaction( estimated_savings: Final = recorded_estimated_autorouter_savings(metadata) return AutoRouterTurnTransaction( api_key=api_key, + user_id=user_id, session_id=bounded_session_id(session_id), router_name=router_name, router_type=str(routing_decision.get("router_type") or "unknown"), @@ -293,18 +313,18 @@ _RETURN_MISS: Final = ( _IDLE_SECONDS: Final = f"EXTRACT(EPOCH FROM {_TURN_AT}::timestamp) - (t.models -> {_MODEL} ->> 'at')::float8" _CACHE_TOUCHED: Final = f"{_TOUCHED}::int = 1" -UPSERT_AUTOROUTER_SESSION_SQL: Final = f""" -INSERT INTO "LiteLLM_AutoRouterSession" AS t ( - api_key, session_id, router_name, router_type, first_turn_at, last_turn_at, - last_model, models, turns, unordered_turns, covered_turns, cache_hits, - same_model_turns, same_model_hits, first_visit_turns, first_visit_hits, - return_turns, return_hits, return_expired_misses, return_within_ttl_misses, - ttl_5m_turns, ttl_1h_turns, total_tokens, spend, saved_spend, classifier_cost, classifier_cost_recorded_turns, tier_turns, - baseline_models, savings_estimated_turns, savings_estimated_actual_spend, savings_estimated_saved_spend, - savings_estimated_baseline_models + +def _session_upsert_sql(*, user_scoped: bool) -> str: + table_name: Final = "LiteLLM_AutoRouterUserSession" if user_scoped else "LiteLLM_AutoRouterSession" + user_column: Final = "user_id, " if user_scoped else "" + user_value: Final = f"{_p('user_id')}::text, " if user_scoped else "" + required_identity: Final = _p("user_id" if user_scoped else "api_key") + return f""" +INSERT INTO "{table_name}" AS t ( + {user_column}{_SESSION_COLUMNS} ) -VALUES ( - {_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp, +SELECT + {user_value}{_p("api_key")}, {_p("session_id")}, {_p("router_name")}, {_p("router_type")}, {_TURN_AT}::timestamp, {_TURN_AT}::timestamp, {_MODEL}, jsonb_build_object({_MODEL}, jsonb_build_object('at', EXTRACT(EPOCH FROM {_TURN_AT}::timestamp), 'ttl', {_CACHE_TTL}::int)), 1, 0, {_COVERED}::int, {_CACHE_HIT}::int, 0, 0, 1, {_CACHE_HIT}::int, @@ -315,8 +335,8 @@ VALUES ( {_p("classifier_cost")}::float8, 1, {_TIER_DELTA}, {_BASELINE_DELTA}, {_p("savings_estimated_turns")}::int, {_p("savings_estimated_actual_spend")}::float8, {_p("savings_estimated_saved_spend")}::float8, {_ESTIMATED_BASELINE_DELTA} -) -ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET +WHERE {required_identity}::text <> '' +ON CONFLICT ({user_column}api_key, session_id, router_name) DO UPDATE SET turns = t.turns + 1, total_tokens = t.total_tokens + EXCLUDED.total_tokens, spend = t.spend + EXCLUDED.spend, @@ -365,6 +385,17 @@ ON CONFLICT (api_key, session_id, router_name) DO UPDATE SET """ +UPSERT_AUTOROUTER_SESSION_SQL: Final = f""" +WITH key_rollup AS ( + {_session_upsert_sql(user_scoped=False)} + RETURNING 1 +) +{_session_upsert_sql(user_scoped=True)} +""" + +UPSERT_AUTOROUTER_USER_SESSION_SQL: Final = _session_upsert_sql(user_scoped=True) + + def _as_sql_param(value: str | float | bool | datetime | None) -> str | float | None: if isinstance(value, bool): return int(value) @@ -377,18 +408,23 @@ def _upsert_params(transaction: AutoRouterTurnTransaction) -> tuple[str | float return tuple(_as_sql_param(getattr(transaction, name)) for name in _UPSERT_PARAM_FIELDS) -async def write_autorouter_turn(db: SupportsExecuteRaw, transaction: AutoRouterTurnTransaction) -> None: - await db.execute_raw(UPSERT_AUTOROUTER_SESSION_SQL, *_upsert_params(transaction)) +async def write_autorouter_turn( + db: SupportsExecuteRaw, + transaction: AutoRouterTurnTransaction, + statement: str = UPSERT_AUTOROUTER_SESSION_SQL, +) -> None: + await db.execute_raw(statement, *_upsert_params(transaction)) async def _upsert_turn_with_retry( prisma_client: PrismaClient, transaction: AutoRouterTurnTransaction, n_retry_times: int, + statement: str, ) -> None: for attempt in range(n_retry_times + 1): try: - await write_autorouter_turn(prisma_client.db, transaction) + await write_autorouter_turn(prisma_client.db, transaction, statement) except DB_RETRY_SAFE_ERROR_TYPES: if attempt >= n_retry_times: raise @@ -397,6 +433,58 @@ async def _upsert_turn_with_retry( return +def _session_partition(transaction: AutoRouterTurnTransaction) -> tuple[str, str, str, str]: + identity: Final = ("key", transaction.api_key) if transaction.api_key else ("user", transaction.user_id) + return (*identity, transaction.session_id, transaction.router_name) + + +async def _drain_session_partition( + prisma_client: PrismaClient, + transactions: tuple[AutoRouterTurnTransaction, ...], + n_retry_times: int, + statement: str, +) -> tuple[AutoRouterTurnTransaction, ...]: + for position, transaction in enumerate(transactions): + try: + await _upsert_turn_with_retry(prisma_client, transaction, n_retry_times, statement) + except Exception as flush_err: # noqa: BLE001 # stop dependent turns without retrying an ambiguous write + verbose_proxy_logger.error( + "Spend tracking - auto-router session rollup flush failed for router %s; " + "%s of %s turn writes stopped in this partition: %s", + transaction.router_name, + len(transactions) - position, + len(transactions), + flush_err, + ) + return transactions[position:] + return () + + +async def _flush_session_partition( + prisma_client: PrismaClient, + transactions: tuple[AutoRouterTurnTransaction, ...], + n_retry_times: int, +) -> None: + failed_suffix: Final = await _drain_session_partition( + prisma_client, transactions, n_retry_times, UPSERT_AUTOROUTER_SESSION_SQL + ) + if not failed_suffix or not failed_suffix[0].api_key: + return + failed_user: Final = failed_suffix[0].user_id + other_users: Final = sorted( + ( + transaction + for transaction in failed_suffix[1:] + if transaction.user_id and transaction.user_id != failed_user + ), + key=lambda transaction: transaction.user_id, + ) + for _, user_turns in groupby(other_users, key=lambda transaction: transaction.user_id): + await _drain_session_partition( + prisma_client, tuple(user_turns), n_retry_times, UPSERT_AUTOROUTER_USER_SESSION_SQL + ) + + async def flush_autorouter_turn_transactions( prisma_client: PrismaClient, transactions: Sequence[AutoRouterTurnTransaction], @@ -407,38 +495,20 @@ async def flush_autorouter_turn_transactions( Statements run sequentially in per-session event order: a turn's classification depends on the turns before it, and Postgres rejects one multi-row INSERT touching the same key twice. Only ConnectError is retried, per statement, because it proves - that statement never reached the database. Any other failure drops the remaining - turns of THAT session only, with an error log, and the flush continues with the - next session: sessions are independent state machines, so one poisoned statement - must not discard unrelated sessions, and a repeated increment is worse than an - undercount. Callers must not add their own retry around this function. + that statement never reached the database. A failed write stops its key and user + histories for this batch. Other users sharing that key can still advance their + independent user histories, with the key projection disabled and the real key + identity preserved. The failed turn is never replayed. Callers must not add their + own retry around this function. """ if not transactions: return ordered: Final = sorted( transactions, - key=lambda transaction: ( - transaction.api_key, - transaction.session_id, - transaction.router_name, - transaction.turn_at, - ), + key=lambda transaction: (*_session_partition(transaction), transaction.turn_at), ) - for session_key, session_group in groupby( + for _, session_group in groupby( ordered, - key=lambda transaction: (transaction.api_key, transaction.session_id, transaction.router_name), + key=_session_partition, ): - session_turns = tuple(session_group) - for position, transaction in enumerate(session_turns): - try: - await _upsert_turn_with_retry(prisma_client, transaction, n_retry_times) - except Exception as flush_err: # noqa: BLE001 # a statement failure drops only its session's remainder by design - verbose_proxy_logger.error( - "Spend tracking - auto-router session rollup flush failed for router %s; " - "%s of %s turn transactions dropped for one session: %s", - session_key[2], - len(session_turns) - position, - len(session_turns), - flush_err, - ) - break + await _flush_session_partition(prisma_client, tuple(session_group), n_retry_times) diff --git a/litellm/proxy/db/baseline_accounting.py b/litellm/proxy/db/baseline_accounting.py index 8622cb9e481..4219102d9aa 100644 --- a/litellm/proxy/db/baseline_accounting.py +++ b/litellm/proxy/db/baseline_accounting.py @@ -171,6 +171,7 @@ class _Change(BaseModel): request_id: str publication: BaselinePublication api_key: str + user_id: str = "" session_id: str router_name: str baseline_model: str @@ -256,42 +257,54 @@ SET publication = x.publication::text FROM jsonb_to_recordset($1::jsonb) AS x(request_id text, publication jsonb) WHERE observations.request_id = x.request_id """ -_UPDATE_SESSIONS: Final = """ + + +def _session_correction_sql(*, user_scoped: bool) -> str: + table_name: Final = "LiteLLM_AutoRouterUserSession" if user_scoped else "LiteLLM_AutoRouterSession" + identity_columns: Final = ("user_id, " if user_scoped else "") + "api_key, session_id, router_name" + user_filter: Final = "WHERE user_id <> ''" if user_scoped else "" + user_match: Final = "session.user_id = totals.user_id AND " if user_scoped else "" + return f""" WITH changes AS ( SELECT * FROM jsonb_to_recordset($1::jsonb) AS x( - api_key text, session_id text, router_name text, baseline_model text, + user_id text, api_key text, session_id text, router_name text, baseline_model text, covered_delta int, actual_delta float8, savings_delta float8 ) + {user_filter} ), totals AS ( - SELECT api_key, session_id, router_name, SUM(covered_delta)::int AS covered_delta, + SELECT {identity_columns}, SUM(covered_delta)::int AS covered_delta, SUM(actual_delta) AS actual_delta, SUM(savings_delta) AS savings_delta - FROM changes GROUP BY api_key, session_id, router_name + FROM changes GROUP BY {identity_columns} ), models AS ( - SELECT api_key, session_id, router_name, jsonb_object_agg(baseline_model, delta) AS deltas + SELECT {identity_columns}, jsonb_object_agg(baseline_model, delta) AS deltas FROM ( - SELECT api_key, session_id, router_name, baseline_model, SUM(covered_delta)::int AS delta - FROM changes GROUP BY api_key, session_id, router_name, baseline_model - ) grouped GROUP BY api_key, session_id, router_name + SELECT {identity_columns}, baseline_model, SUM(covered_delta)::int AS delta + FROM changes GROUP BY {identity_columns}, baseline_model + ) grouped GROUP BY {identity_columns} ) -UPDATE "LiteLLM_AutoRouterSession" AS session +UPDATE "{table_name}" AS session SET saved_spend = session.saved_spend + totals.savings_delta, savings_estimated_turns = session.savings_estimated_turns + totals.covered_delta, savings_estimated_actual_spend = session.savings_estimated_actual_spend + totals.actual_delta, savings_estimated_saved_spend = session.savings_estimated_saved_spend + totals.savings_delta, savings_estimated_baseline_models = ( - SELECT COALESCE(jsonb_object_agg(key, value), '{}'::jsonb) FROM ( + SELECT COALESCE(jsonb_object_agg(key, value), '{{}}'::jsonb) FROM ( SELECT key, SUM(value::int)::int AS value FROM ( SELECT * FROM jsonb_each_text(session.savings_estimated_baseline_models) UNION ALL SELECT * FROM jsonb_each_text(models.deltas) ) combined GROUP BY key HAVING SUM(value::int) > 0 ) counts ) -FROM totals JOIN models USING (api_key, session_id, router_name) -WHERE session.api_key = totals.api_key AND session.session_id = totals.session_id +FROM totals JOIN models USING ({identity_columns}) +WHERE {user_match}session.api_key = totals.api_key AND session.session_id = totals.session_id AND session.router_name = totals.router_name """ +_UPDATE_SESSIONS: Final = _session_correction_sql(user_scoped=False) +_UPDATE_USER_SESSIONS: Final = _session_correction_sql(user_scoped=True) + + def _primary_transaction(client: PrismaClient) -> _TransactionManager: primary: Final = cast(_TransactionalDatabase, writer_wrapper(client.db)) return primary.tx(timeout=_TRANSACTION_TIMEOUT) @@ -308,6 +321,7 @@ def _change(record: BaselineAccountingRecord, old: BaselinePublication | None, n request_id=record.observation.request_id, publication=new, api_key=record.api_key, + user_id=record.turn.user_id if record.turn is not None else "", session_id=record.session_id, router_name=record.router_name, baseline_model=record.baseline_model, @@ -357,6 +371,8 @@ async def _publish(db: SupportsRawQueries, changes: Sequence[_Change]) -> None: serialized: Final = json.dumps(tuple(change.model_dump(mode="json") for change in changes), separators=(",", ":")) await db.execute_raw(_UPDATE_LOGS, serialized) await db.execute_raw(_UPDATE_SESSIONS, serialized) + if any(change.user_id for change in changes): + await db.execute_raw(_UPDATE_USER_SESSIONS, serialized) for entity, table in DAILY_SPEND_TABLES.items(): if adjustments := tuple( change.daily.adjustment(target, change.savings_delta, change.request_id) diff --git a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py index b28a653c9aa..85e19fa8a32 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py +++ b/litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py @@ -1,5 +1,6 @@ import asyncio import time +from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Final, Literal, TypeAlias @@ -40,6 +41,28 @@ class TableCleanupResult: stop_reason: StopReason +class _RunProgress: + """How far one cleanup run has got, reported if that run is cancelled""" + + def __init__(self) -> None: + self.rows_deleted: int = 0 + self.batches: int = 0 + + def record_batch(self, rows_deleted: int) -> None: + self.rows_deleted += rows_deleted + self.batches += 1 + + +_run_progress: ContextVar[_RunProgress] = ContextVar("spend_log_cleanup_run_progress") + + +def _record_run_batch(rows_deleted: int) -> None: + """Count a batch towards the run in progress, if a run is what issued it""" + progress: Final = _run_progress.get(None) + if progress is not None: + progress.record_batch(rows_deleted) + + class _RemainingRow(BaseModel): """One row of the capped outstanding-rows probe, validated out of prisma's untyped result.""" @@ -422,6 +445,7 @@ class SpendLogCleanup: total_deleted += deleted_count run_count += 1 + _record_run_batch(deleted_count) # Add a small sleep to prevent overwhelming the database await asyncio.sleep(0.1) @@ -492,6 +516,18 @@ class SpendLogCleanup: deadline=deadline, ) + async def _delete_old_autorouter_user_session_rows( + self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float + ) -> TableCleanupResult: + return await self._delete_old_rows_batched( + prisma_client, + cutoff_date, + table_name="LiteLLM_AutoRouterUserSession", + key_columns=("user_id", "api_key", "session_id", "router_name"), + time_column="last_turn_at", + deadline=deadline, + ) + async def _delete_old_health_check_rows( self, prisma_client: PrismaClient, cutoff_date: datetime, deadline: float ) -> TableCleanupResult: @@ -560,9 +596,17 @@ class SpendLogCleanup: ) except Exception: # noqa: BLE001 # retained observations are retried by the next cleanup job verbose_proxy_logger.warning("Auto-router baseline retention remains pending") - sessions_result: Final = await self._delete_old_autorouter_session_rows(prisma_client, session_cutoff, deadline) + sessions_result: Final = await self._delete_old_autorouter_session_rows( + prisma_client, session_cutoff, self._group_deadline(deadline, 2) + ) verbose_proxy_logger.info("Deleted %s expired auto-router session rollup rows", sessions_result.rows_deleted) - return (sessions_result,) + user_sessions_result: Final = await self._delete_old_autorouter_user_session_rows( + prisma_client, session_cutoff, deadline + ) + verbose_proxy_logger.info( + "Deleted %s expired auto-router user session rollup rows", user_sessions_result.rows_deleted + ) + return (sessions_result, user_sessions_result) async def _clean_health_checks( self, prisma_client: PrismaClient, retention_seconds: int, deadline: float @@ -601,6 +645,9 @@ class SpendLogCleanup: If no pod_lock_manager, runs cleanup without distributed locking. """ lock_acquired = False + run_started_at: Final = time.monotonic() + progress: Final = _RunProgress() + progress_token: Final = _run_progress.set(progress) try: verbose_proxy_logger.info("Cleanup job triggered at %s", datetime.now()) self._refresh_bounds() @@ -681,6 +728,15 @@ class SpendLogCleanup: self._run_outcome(spend_log_results + session_results + health_check_results) ) + except asyncio.CancelledError: + verbose_proxy_logger.error( + "Spend log cleanup cancelled after %.2fs (rows_deleted=%d, batches=%d); the next run resumes from here", + time.monotonic() - run_started_at, + progress.rows_deleted, + progress.batches, + ) + SpendLogCleanupMetrics.record_run("aborted") + raise except Exception as e: # .exception() captures the traceback; str(e) alone on a Prisma/DB # timeout is often empty and gives operators no signal to diagnose. @@ -692,6 +748,7 @@ class SpendLogCleanup: SpendLogCleanupMetrics.record_run("aborted") return # Return after error handling finally: + _run_progress.reset(progress_token) # Only release the lock if it was actually acquired if lock_acquired and self.pod_lock_manager and self.pod_lock_manager.redis_cache: await self.pod_lock_manager.release_lock(cronjob_id=SPEND_LOG_CLEANUP_JOB_NAME) diff --git a/litellm/proxy/management_endpoints/access_group_endpoints.py b/litellm/proxy/management_endpoints/access_group_endpoints.py index a6cc5140b15..b4923b0a2dc 100644 --- a/litellm/proxy/management_endpoints/access_group_endpoints.py +++ b/litellm/proxy/management_endpoints/access_group_endpoints.py @@ -5,6 +5,7 @@ from types import MappingProxyType from typing import Final, Protocol from fastapi import APIRouter, Depends, HTTPException, status +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager @@ -109,6 +110,36 @@ class _KeyTable(Protocol): async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... +class _AgentRecord(Protocol): + @property + def agent_id(self) -> str: ... + + @property + def access_group_ids(self) -> Sequence[str] | None: ... + + +class _AgentTable(Protocol): + async def find_many(self, where: Mapping[str, object]) -> Sequence[_AgentRecord]: ... + + async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + + +class _HasSomeFilter(TypedDict): + hasSome: ReadOnly[Sequence[str]] + + +class _AgentAccessGroupsWhere(TypedDict): + access_group_ids: ReadOnly[_HasSomeFilter] + + +class _AgentIdWhere(TypedDict): + agent_id: ReadOnly[str] + + +class _AgentAccessGroupsData(TypedDict): + access_group_ids: ReadOnly[Sequence[str]] + + class _AccessGroupTx(Protocol): @property def litellm_accessgrouptable(self) -> _AccessGroupTable: ... @@ -119,6 +150,9 @@ class _AccessGroupTx(Protocol): @property def litellm_verificationtoken(self) -> _KeyTable: ... + @property + def litellm_agentstable(self) -> _AgentTable: ... + def _require_proxy_admin(user_api_key_dict: UserAPIKeyAuth) -> None: if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN: @@ -324,6 +358,41 @@ async def _sync_remove_access_group_from_keys(tx: _AccessGroupTx, key_tokens: li ) +def _without_access_group(access_group_ids: Sequence[str] | None, access_group_id: str) -> tuple[str, ...]: + return tuple(ag for ag in (access_group_ids or ()) if ag != access_group_id) + + +async def _detach_access_group_from_agents(tx: _AccessGroupTx, access_group_id: str) -> tuple[str, ...]: + agents_with_group: Final = await tx.litellm_agentstable.find_many( + where=_AgentAccessGroupsWhere(access_group_ids=_HasSomeFilter(hasSome=(access_group_id,))) + ) + for agent in agents_with_group: + await tx.litellm_agentstable.update( + where=_AgentIdWhere(agent_id=agent.agent_id), + data=_AgentAccessGroupsData( + access_group_ids=_without_access_group(agent.access_group_ids, access_group_id) + ), + ) + return tuple(agent.agent_id for agent in agents_with_group) + + +def _detach_access_group_from_agent_registry(agent_ids: Sequence[str], access_group_id: str) -> None: + registered: Final = tuple( + agent + for agent in (global_agent_registry.get_agent_by_id(agent_id) for agent_id in agent_ids) + if agent is not None + ) + for agent in registered: + global_agent_registry.deregister_agent(agent_name=agent.agent_name) + global_agent_registry.register_agent( + agent_config=agent.model_copy( + update=_AgentAccessGroupsData( + access_group_ids=_without_access_group(agent.access_group_ids, access_group_id) + ) + ) + ) + + # --------------------------------------------------------------------------- # Cache patch helpers # --------------------------------------------------------------------------- @@ -705,11 +774,14 @@ async def delete_access_group( out_of_sync_key_tokens: Final = set(existing.assigned_key_ids or []) - {k.token for k in keys_with_group} await _sync_remove_access_group_from_keys(tx, list(out_of_sync_key_tokens), access_group_id) + detached_agent_ids: Final = await _detach_access_group_from_agents(tx, access_group_id) + await tx.litellm_accessgrouptable.delete(where={"access_group_id": access_group_id}) from litellm.proxy.proxy_server import proxy_logging_obj, user_api_key_cache await invalidate_access_group_cache(access_group_id) + _detach_access_group_from_agent_registry(detached_agent_ids, access_group_id) await _patch_team_caches_remove_access_group( affected_team_ids, access_group_id, user_api_key_cache, proxy_logging_obj ) diff --git a/litellm/proxy/management_endpoints/auto_router_endpoints.py b/litellm/proxy/management_endpoints/auto_router_endpoints.py index d8054a8dc4a..03fc58622ce 100644 --- a/litellm/proxy/management_endpoints/auto_router_endpoints.py +++ b/litellm/proxy/management_endpoints/auto_router_endpoints.py @@ -789,14 +789,18 @@ async def get_auto_router_benchmarks( ] = None, end_date: Annotated[str | None, Query(description="YYYY-MM-DD UTC, inclusive (defaults to today)")] = None, api_key: Annotated[str | None, Query(description="Filter to one virtual key token hash")] = None, + user_id: Annotated[ + str | None, Query(min_length=1, description="Filter to one canonical internal user recorded on each turn") + ] = None, ) -> AutoRouterBenchmarksResponse: """ Benchmarks for the auto-router dashboard: session shape, savings against the configured baseline, and prompt-caching behaviour bucketed by what the router did. - Reads the LiteLLM_AutoRouterSession rollup, folded once per request at spend-write time, - so this endpoint never scans LiteLLM_SpendLogs. A session is in the window when it - overlaps it: its last turn is on or after start_date and its first turn is on or before + Reads session rollups folded once per request at spend-write time, so this endpoint + never scans LiteLLM_SpendLogs. A user filter selects only turns attributed to that + internal user when written; older key-only history remains outside user views. A session + is in the window when it overlaps it: its last turn is on or after start_date and its first turn is on or before end_date. Overall hit rate is over telemetry-bearing turns; each bucket's hit rate is over that bucket's turns. @@ -826,6 +830,7 @@ async def get_auto_router_benchmarks( start_day.isoformat(), (end_day + timedelta(days=1)).isoformat(), api_key, + user_id, ) rows: Final = _SESSION_AGG_ROWS.validate_python(raw_rows or ()) groups: Final = ( diff --git a/litellm/proxy/management_endpoints/router_settings_endpoints.py b/litellm/proxy/management_endpoints/router_settings_endpoints.py index fa5eca7ef32..7a3b091551b 100644 --- a/litellm/proxy/management_endpoints/router_settings_endpoints.py +++ b/litellm/proxy/management_endpoints/router_settings_endpoints.py @@ -8,6 +8,8 @@ GET /router/fields - Get router settings field definitions without values (for U """ import inspect +from collections.abc import Mapping +from types import MappingProxyType from typing import Any, Final, get_args from fastapi import APIRouter, Depends @@ -16,6 +18,7 @@ from pydantic import BaseModel, Field from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.config_resolvers import FieldSource, SettingsStore, source_for from litellm.router import Router from litellm.types.management_endpoints import ( ROUTER_SETTINGS_FIELDS, @@ -30,6 +33,7 @@ class RouterSettingsResponse(BaseModel): fields: list[RouterSettingsField] = Field(description="List of all configurable router settings with metadata") current_values: dict[str, Any] = Field(description="Current values of router settings") routing_strategy_descriptions: dict[str, str] = Field(description="Descriptions for each routing strategy option") + source: dict[str, FieldSource] = Field(description="Source of each current router setting") class RouterFieldsResponse(BaseModel): @@ -39,6 +43,18 @@ class RouterFieldsResponse(BaseModel): routing_strategy_descriptions: dict[str, str] = Field(description="Descriptions for each routing strategy option") +def _router_setting_source( + settings: SettingsStore, + key: str, + current_value: object, + field_default: object, +) -> FieldSource: + source: Final = source_for(settings, key, field_default) + if source != "unset": + return source + return "default" if current_value is not None else "unset" + + def _get_routing_strategies_from_router_class() -> list[str]: """ Dynamically extract routing strategies from the Router class __init__ method. @@ -109,15 +125,29 @@ async def get_router_settings( # Merge with config values (config takes precedence) current_values.update(router_settings_from_config) - # Update field values with current values for field in router_fields: if field.field_name in current_values: field.field_value = current_values[field.field_name] + field_defaults: Final[Mapping[str, object]] = MappingProxyType( + {field.field_name: field.field_default for field in router_fields} + ) + source: Final[Mapping[str, FieldSource]] = MappingProxyType( + { + key: _router_setting_source( + proxy_config.router_settings, + key, + current_values[key], + field_defaults.get(key), + ) + for key in current_values + } + ) return RouterSettingsResponse( fields=router_fields, current_values=current_values, routing_strategy_descriptions=ROUTING_STRATEGY_DESCRIPTIONS, + source=source, ) except Exception as e: verbose_proxy_logger.error("Error fetching router settings: %s", e) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 011d20f212b..3ca15f3f17a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -51,6 +51,7 @@ import anyio import websockets import websockets.exceptions from pydantic import BaseModel, Json, JsonValue, TypeAdapter, ValidationError +from pydantic.fields import FieldInfo, PydanticUndefined from typing_extensions import NotRequired, ReadOnly, assert_never from litellm._uuid import uuid @@ -460,7 +461,13 @@ from litellm.proxy.common_utils.user_api_key_cache import ( project_spend_counter_key, tag_cache_key, ) -from litellm.proxy.config_resolvers import SettingsStore, config_ownership_message, resolve_fields +from litellm.proxy.config_resolvers import ( + FieldSource, + SettingsStore, + config_ownership_message, + resolve_fields, + source_for, +) from litellm.proxy.config_resolvers.alerting import ( EMAIL_DESCRIPTORS, MS_TEAMS_DESCRIPTORS, @@ -719,6 +726,11 @@ from litellm.proxy.route_llm_request import route_request from litellm.proxy.route_priority import hot_routes_first from litellm.proxy.search_endpoints.endpoints import router as search_router from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager +from litellm.proxy.shutdown.scheduled_jobs import ( + AwaitableAsyncIOExecutor, + pause_scheduled_jobs, + stop_in_flight_scheduler_jobs, +) from litellm.proxy.spend_tracking.budget_reservation import get_budget_window_start from litellm.proxy.spend_tracking.daily_global_spend_rollup import ( run_scheduled_daily_global_spend_reconcile, @@ -1492,6 +1504,10 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: if model_info_scheduler is not scheduler: model_info_scheduler.shutdown(wait=False) + # Shutdown event - stop starting scheduled jobs; the ones already running keep the drain window + if scheduler is not None: + pause_scheduled_jobs(scheduler) + # Shutdown event - drain in-flight requests before tearing down dependencies # so SIGTERM (rolling update, scale-down, liveness kill) doesn't drop them. GracefulShutdownManager.start_shutdown() @@ -1531,6 +1547,13 @@ async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[None, None]: await _drain_spend_event_producer_on_shutdown() + # Shutdown event - finish or cancel in-flight scheduled jobs before the shutdown flushes and the DB disconnect + if scheduler is not None and scheduler_executor is not None: + try: + await stop_in_flight_scheduler_jobs(scheduler, scheduler_executor) + except Exception as e: + verbose_proxy_logger.error("Error stopping in-flight scheduled jobs: %s", e) + await flush_spend_counters_on_shutdown() await _flush_spend_logs_queue_on_shutdown() @@ -2533,6 +2556,7 @@ celery_app_conn: Final = None celery_fn: Final = None # Redis Queue for handling requests scheduler = None +scheduler_executor: AwaitableAsyncIOExecutor | None = None # rebind-ok: bound once the scheduler is built at startup # Global variable for anthropic beta headers reload scheduling last_anthropic_beta_headers_reload = None @@ -4958,6 +4982,12 @@ def _resolve_env_params_from_os(params: Mapping[str, object]) -> dict[str, objec } +def _get_field_default(field_info: FieldInfo) -> JsonValue: + if field_info.default is PydanticUndefined: + return None + return cast(JsonValue, field_info.default) # cast-ok: Pydantic field defaults are JSON values at runtime + + def _bind_general_settings_store(settings: SettingsStore) -> None: global general_settings general_settings = settings # pyright: ignore[reportAssignmentType] # legacy global accepts mappings @@ -10121,7 +10151,7 @@ class ProxyStartupEvent: proxy_logging_obj: ProxyLogging, ) -> ProxyWorkerHeartbeat: """Initializes scheduled background jobs""" - global heuristic_v1_tuning_baselines, store_model_in_db, scheduler # rebind-ok: startup publishes the one read-only baseline snapshot + global heuristic_v1_tuning_baselines, store_model_in_db, scheduler, scheduler_executor # rebind-ok: startup publishes the one read-only baseline snapshot # MEMORY LEAK FIX: Configure scheduler with optimized settings # Memray analysis showed APScheduler's normalize() and _apply_jitter() causing @@ -10130,9 +10160,9 @@ class ProxyStartupEvent: # 1. Remove/minimize jitter to avoid normalize() memory explosion # 2. Use larger misfire_grace_time to prevent backlog calculations # 3. Set replace_existing=True to avoid duplicate jobs - from apscheduler.executors.asyncio import AsyncIOExecutor from apscheduler.jobstores.memory import MemoryJobStore + scheduler_executor = AwaitableAsyncIOExecutor() # rebind-ok: shutdown awaits the jobs this executor runs scheduler = AsyncIOScheduler( job_defaults={ "coalesce": APSCHEDULER_COALESCE, @@ -10145,7 +10175,7 @@ class ProxyStartupEvent: jobstores={"default": MemoryJobStore()}, # explicitly use memory job store # Use simple executor to minimize overhead executors={ - "default": AsyncIOExecutor(), + "default": scheduler_executor, }, # Disable timezone awareness to reduce computation timezone=None, @@ -16046,6 +16076,22 @@ async def model_settings(): #### ALERTING MANAGEMENT ENDPOINTS #### +def _nested_setting_source( + settings: SettingsStore, + db_values: Mapping[str, JsonValue], + parent_key: str, + field_name: str, + field_default: JsonValue, +) -> FieldSource: + unset_source: Final[FieldSource] = "default" if field_default is not None else "unset" + parent_value: Final = settings.config_value(parent_key) + if isinstance(parent_value, Mapping) and field_name in parent_value: + return "config" + if settings.owned_by_config(parent_key): + return unset_source + return "db" if field_name in db_values else unset_source + + @router.get( "/alerting/settings", description="Return the configurable alerting param, description, and current value", @@ -16083,17 +16129,20 @@ async def alerting_settings( where={"param_name": "general_settings"} ) - if db_general_settings is not None and db_general_settings.param_value is not None: - db_general_settings_dict: Final = dict(db_general_settings.param_value) - alerting_args_dict: dict = cast( # cast-ok: ConfigGeneralSettings validates alerting_args as a dict on write - dict[str, JsonValue], db_general_settings_dict.get("alerting_args", {}) - ) - alerting_values: list | None = cast( # cast-ok: ConfigGeneralSettings validates alerting as a list on write - list[JsonValue] | None, db_general_settings_dict.get("alerting") - ) - else: - alerting_args_dict = {} - alerting_values = None + db_general_settings_dict: Final[Mapping[str, JsonValue]] = MappingProxyType( + dict(db_general_settings.param_value) # mutable-ok: Prisma returns the JSON column as a plain dict + if db_general_settings is not None and db_general_settings.param_value is not None + else {} + ) + alerting_args_value: Final = db_general_settings_dict.get("alerting_args") + alerting_args_dict: Final[Mapping[str, JsonValue]] = MappingProxyType( + alerting_args_value if isinstance(alerting_args_value, dict) else {} + ) + alerting_values: Final = cast( # cast-ok: alerting is stored as a JSON list when present + list[JsonValue] | None, db_general_settings_dict.get("alerting") + ) + + settings: Final = proxy_config.settings allowed_args: Final = MappingProxyType( { @@ -16122,9 +16171,9 @@ async def alerting_settings( is_slack_enabled = False - if general_settings.get("alerting") and isinstance(general_settings["alerting"], list): - if "slack" in general_settings["alerting"]: - is_slack_enabled = True + alerting: Final = settings.get("alerting") + if isinstance(alerting, list) and "slack" in alerting: + is_slack_enabled = True _response_obj = ConfigList( field_name="slack_alerting", @@ -16132,6 +16181,7 @@ async def alerting_settings( field_description="Enable slack alerting for monitoring proxy in production: llm outages, budgets, spend tracking failures.", field_value=is_slack_enabled, stored_in_db=True if alerting_values is not None else False, + source=source_for(settings, "alerting"), field_default_value=None, premium_field=False, ) @@ -16139,6 +16189,7 @@ async def alerting_settings( for field_name, field_info in SlackAlertingArgs.model_fields.items(): if field_name in allowed_args: + field_default: JsonValue = _get_field_default(field_info) _stored_in_db: bool | None = None if field_name in alerting_args_dict: _stored_in_db = True @@ -16149,9 +16200,16 @@ async def alerting_settings( field_name=field_name, field_type=allowed_args[field_name], field_description=field_info.description or "", - field_value=_slack_alerting_args_dict.get(field_name, None), + field_value=_slack_alerting_args_dict.get(field_name, field_default), stored_in_db=_stored_in_db, - field_default_value=field_info.default, + source=_nested_setting_source( + settings, + alerting_args_dict, + "alerting_args", + field_name, + field_default, + ), + field_default_value=field_default, premium_field=(True if field_name == "region_outage_alert_ttl" else False), ) return_val.append(_response_obj) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 368864f9cd1..85996430bc5 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -73,6 +73,7 @@ model LiteLLM_AgentsTable { static_headers Json? @default("{}") extra_headers String[] @default([]) agent_access_groups String[] @default([]) + access_group_ids String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) @@ -1623,6 +1624,47 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } +model LiteLLM_AutoRouterUserSession { + user_id String + api_key String + session_id String + router_name String + router_type String + first_turn_at DateTime + last_turn_at DateTime + last_model String + models Json @default("{}") + turns Int @default(0) + unordered_turns Int @default(0) + covered_turns Int @default(0) + cache_hits Int @default(0) + same_model_turns Int @default(0) + same_model_hits Int @default(0) + first_visit_turns Int @default(0) + first_visit_hits Int @default(0) + return_turns Int @default(0) + return_hits Int @default(0) + return_expired_misses Int @default(0) + return_within_ttl_misses Int @default(0) + ttl_5m_turns Int @default(0) + ttl_1h_turns Int @default(0) + total_tokens BigInt @default(0) + spend Float @default(0) + saved_spend Float @default(0) + savings_estimated_turns Int @default(0) + savings_estimated_actual_spend Float @default(0) + savings_estimated_saved_spend Float @default(0) + savings_estimated_baseline_models Json @default("{}") + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) + tier_turns Json @default("{}") + baseline_models Json @default("{}") + + @@id([user_id, api_key, session_id, router_name]) + @@index([last_turn_at], map: "idx_autorouter_user_session_last_turn") + @@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn") +} + // Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in // either direction. forward duplicates the requests the keys did not route through the // router through it, answering whether they should adopt it; reverse duplicates the diff --git a/litellm/proxy/shutdown/scheduled_jobs.py b/litellm/proxy/shutdown/scheduled_jobs.py new file mode 100644 index 00000000000..7889c35cf4e --- /dev/null +++ b/litellm/proxy/shutdown/scheduled_jobs.py @@ -0,0 +1,79 @@ +# pyright: reportMissingTypeStubs=false # apscheduler ships no type information + +import asyncio +from collections.abc import Collection +from typing import Final, Protocol + +from apscheduler.executors.asyncio import AsyncIOExecutor + +from litellm._logging import verbose_proxy_logger +from litellm.constants import ( + SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, + SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, +) + + +class StoppableScheduler(Protocol): + """The slice of ``AsyncIOScheduler`` shutdown uses, which ships no type information""" + + @property + def running(self) -> bool: ... + + def pause(self) -> None: ... + + def shutdown(self, wait: bool = ...) -> None: ... + + +class AwaitableAsyncIOExecutor(AsyncIOExecutor): # pyright: ignore[reportUntypedBaseClass] # apscheduler ships no type information and is absent from the type-check env + """``AsyncIOExecutor`` whose in-flight job tasks can be awaited after ``shutdown`` cancels them""" + + _pending_futures: Collection["asyncio.Future[object]"] + + def in_flight_jobs(self) -> tuple["asyncio.Future[object]", ...]: + """The job tasks that are running right now, as a snapshot""" + return tuple(future for future in self._pending_futures if not future.done()) + + +def pause_scheduled_jobs(scheduler: StoppableScheduler) -> None: + """Stop the scheduler from starting jobs that shutdown would only cancel; running jobs continue""" + if scheduler.running: + scheduler.pause() + + +async def stop_in_flight_scheduler_jobs( + scheduler: StoppableScheduler, + executor: AwaitableAsyncIOExecutor, + *, + finish_timeout_seconds: float = SCHEDULED_JOB_SHUTDOWN_FINISH_TIMEOUT_SECONDS, + cancel_timeout_seconds: float = SCHEDULED_JOB_SHUTDOWN_CANCEL_TIMEOUT_SECONDS, +) -> None: + """ + Let in-flight jobs finish for up to finish_timeout_seconds, then stop the scheduler and wait, bounded by + cancel_timeout_seconds, for the jobs it cancels. + + Must run before the database is disconnected: a write job that finishes needs its connection, + and a job's cancellation handler is what records the run's outcome. + """ + if not scheduler.running: + return + in_flight: Final = executor.in_flight_jobs() + if in_flight: + verbose_proxy_logger.info( + "Waiting up to %ss for %d in-flight scheduled job(s) to finish", + finish_timeout_seconds, + len(in_flight), + ) + still_running: Final = ( + (await asyncio.wait(in_flight, timeout=finish_timeout_seconds))[1] if in_flight else frozenset() + ) + scheduler.shutdown(wait=False) + if not still_running: + return + verbose_proxy_logger.info("Cancelling %d in-flight scheduled job(s) for shutdown", len(still_running)) + _done, pending = await asyncio.wait(still_running, timeout=cancel_timeout_seconds) + if pending: + verbose_proxy_logger.warning( + "%d scheduled job(s) did not finish within %ss of cancellation; giving up on them", + len(pending), + cancel_timeout_seconds, + ) diff --git a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py index 5e24738d4ed..98718580e7f 100644 --- a/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py +++ b/litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py @@ -15,8 +15,8 @@ from typing import ( from urllib.parse import urlparse from fastapi import APIRouter, Body, Depends, File, HTTPException, UploadFile -from pydantic import ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model -from pydantic.fields import FieldInfo +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, create_model +from pydantic.fields import FieldInfo, PydanticUndefined from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm @@ -25,6 +25,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import mask_sensitive_keys from litellm.proxy._experimental.mcp_server.tool_search import MCP_TOOL_SEARCH_SETTINGS_KEY from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.config_resolvers import FieldSource, SettingsStore, source_for from litellm.proxy.config_resolvers.settings_store import ConfigOwnedKeyError from litellm.proxy.config_resolvers.sso import ( SSO_FIELD_ENV_VARS, @@ -35,7 +36,10 @@ from litellm.proxy.management_endpoints.team_admin_field_permissions import ( SUPPORTED_TEAM_ADMIN_PERMISSIONS, TEAM_ADMIN_EDITABLE_TEAM_FIELDS_SETTING, ) -from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled +from litellm.proxy.spend_tracking.ptu_feature_flag import ( + PTU_COST_ATTRIBUTION_ENV_VAR, + is_ptu_cost_attribution_enabled, +) from litellm.proxy.utils import invalidate_config_param from litellm.repositories.config_repository import ConfigRepository from litellm.repositories.organization_repository import OrganizationRepository @@ -45,6 +49,7 @@ from litellm.repositories.table_repositories import ( UISettingsRepository, ) from litellm.repositories.team_repository import TeamRepository +from litellm.secret_managers.main import get_secret from litellm.types.mcp import MCPToolSearchSettings from litellm.types.proxy.management_endpoints.ui_sso import ( DefaultTeamSSOParams, @@ -199,6 +204,11 @@ class SettingsResponse(BaseModel): """Schema information including descriptions and property types for UI display""" +class _SettingsWithSchema(BaseModel): + values: dict[str, object] + field_schema: dict[str, object] + + class SSOSettingsResponse(SettingsResponse): """Response model for SSO settings""" @@ -330,6 +340,8 @@ class UISettings(BaseModel): class UISettingsResponse(SettingsResponse): """Response model for UI settings""" + source: dict[str, FieldSource] + # Allowlist of UI settings that can be stored ALLOWED_UI_SETTINGS_FIELDS: Final = { @@ -755,6 +767,25 @@ def _root_schema(settings_class: type[BaseModel]) -> _RootSchema: ) +def _model_field_default(settings_class: type[BaseModel], field_name: str) -> object: + field_info: Final = settings_class.model_fields.get(field_name) + if field_info is None or field_info.default is PydanticUndefined: + return None + return cast(object, field_info.default) # cast-ok: Pydantic field defaults are untyped + + +def _ui_setting_source( + key: str, + value: object, + settings: SettingsStore, + settings_class: type[BaseModel], +) -> FieldSource: + if key == ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: + configured_value: Final = get_secret(PTU_COST_ATTRIBUTION_ENV_VAR, None) + return "config" if configured_value is not None or value is True else "default" + return source_for(settings, key, _model_field_default(settings_class, key)) + + async def _get_settings_with_schema( settings_key: str, settings_class: type[BaseModel], @@ -1706,7 +1737,7 @@ async def get_ui_settings(): Get UI-specific configuration flags. All authenticated users can fetch these settings for client-side behavior. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import prisma_client, proxy_config if prisma_client is None: raise HTTPException( @@ -1731,22 +1762,43 @@ async def get_ui_settings(): await user_api_key_cache.async_set_cache(key=UI_SETTINGS_CACHE_KEY, value=ui_settings, ttl=UI_SETTINGS_CACHE_TTL) - # Build config-like object for schema helper - config: Final[Mapping[str, object]] = { - "litellm_settings": {"ui_settings": ui_settings} - } # mutable-ok: schema helper only reads it - - settings: Final = await _get_settings_with_schema( - settings_key="ui_settings", - settings_class=_get_effective_ui_settings_class(), - config=config, + effective_ui_settings: Final[Mapping[str, object]] = MappingProxyType( + { + **ui_settings, + **{key: proxy_config.settings[key] for key in ALLOWED_UI_SETTINGS_FIELDS if key in proxy_config.settings}, + } + ) + config: Final[Mapping[str, object]] = MappingProxyType( + {"litellm_settings": MappingProxyType({"ui_settings": effective_ui_settings})} + ) + settings_class: Final = _get_effective_ui_settings_class() + resolved_settings: Final = _SettingsWithSchema.model_validate( + await _get_settings_with_schema( + settings_key="ui_settings", + settings_class=settings_class, + config=config, + ) + ) + values: Final[Mapping[str, object]] = MappingProxyType( + { + **resolved_settings.values, + ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(), + } + ) + source: Final[Mapping[str, FieldSource]] = MappingProxyType( + { + key: ( + _ui_setting_source(key, values[key], proxy_config.settings, settings_class) + if key in proxy_config.settings or key not in ui_settings + else "db" + ) + for key in values + } ) return UISettingsResponse( - values={ - **settings["values"], - ENABLE_PTU_COST_ATTRIBUTION_UI_SETTING: is_ptu_cost_attribution_enabled(), - }, - field_schema=settings["field_schema"], + values=values, + field_schema=resolved_settings.field_schema, + source=source, ) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index de5f545f109..bfd412db1a7 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -148,6 +148,7 @@ from litellm.proxy._types import ( Member, UserAPIKeyAuth, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change @@ -8261,6 +8262,51 @@ async def _get_access_group_models( return tuple(dict.fromkeys((*team_group_models, *key_group_models))) +async def _agent_access_group_visible_models( + user_api_key_dict: "UserAPIKeyAuth", + llm_router: "Router | None", + include_model_access_groups: bool, + return_wildcard_routes: bool, + team_id: str | None, + resolve_agent_ceiling: CeilingResolver, +) -> frozenset[str] | None: + """Models an agent key may still list once its attached access groups cap it, ``None`` when + nothing caps it, so ``/v1/models`` never advertises a model the same key would be denied on.""" + from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models + + if not user_api_key_dict.agent_id: + return None + ceiling: Final = await resolve_agent_ceiling(user_api_key_dict.agent_id) + if ceiling is None: + return None + if llm_router is None: + return ceiling.models + proxy_model_list: Final = llm_router.get_model_names() + model_access_groups: Final = llm_router.get_model_access_groups() + granted: Final = get_team_models( + team_models=sorted(ceiling.models), + proxy_model_list=proxy_model_list, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + ) + if not granted: + return frozenset() + return frozenset( + get_complete_model_list( + key_models=granted, + team_models=(), + proxy_model_list=proxy_model_list, + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=return_wildcard_routes, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + team_id=team_id, + ) + ) + + async def get_available_models_for_user( user_api_key_dict: "UserAPIKeyAuth", llm_router: Optional["Router"], @@ -8273,6 +8319,7 @@ async def get_available_models_for_user( only_model_access_groups: bool = False, return_wildcard_routes: bool = False, user_api_key_cache: Optional["UserApiKeyCache"] = None, + resolve_agent_ceiling: CeilingResolver = resolve_agent_access_group_ceiling, ) -> list[str]: """ Get the list of models available to a user based on their API key and team permissions. @@ -8376,7 +8423,18 @@ async def get_available_models_for_user( team_id=effective_team_id, ) - return all_models + agent_visible: Final = await _agent_access_group_visible_models( + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + include_model_access_groups=include_model_access_groups, + return_wildcard_routes=return_wildcard_routes, + team_id=effective_team_id, + resolve_agent_ceiling=resolve_agent_ceiling, + ) + if agent_visible is None: + return all_models + capped: Final = [m for m in all_models if m in agent_visible] # mutable-ok: callers expect the list all_models is + return capped def _safe_get_model_info(model: str, get_model_info: Callable[[str], ModelInfo]) -> ModelInfo | None: diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 1537e3a540c..c4dbe68ebe2 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -1321,7 +1321,7 @@ class ComplexityRouterConfig(BaseModel): ) enable_context_window_escalation: bool = Field( - default=True, + default=False, description=( "Escalate a request off a tier whose models provably cannot hold its prompt, before " "dispatch. The classifier scores complexity and never prompt size, so a long agentic " @@ -1331,7 +1331,8 @@ class ComplexityRouterConfig(BaseModel): "moves to the lowest configured tier with a model whose declared window fits; when " "only some of the tier's models fit, the pick is restricted to those and the tier " "keeps the request. Models with no resolvable window are never escalated away from " - "and never escalated onto. Set false to dispatch on complexity alone, as before." + "and never escalated onto. Disabled by default: omit or set false to dispatch on " + "complexity alone; set true to enable context-window escalation." ), ) context_window_escalation_buffer: float = Field( diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 90d876b71e6..c0bb93daadc 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -169,6 +169,15 @@ class _CacheTestHandle: @staticmethod def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ... @staticmethod + def redis_semantic(backend: object) -> _CacheTestHandle: ... + @staticmethod + def valkey_semantic( + url: str, + similarity_threshold: float, + index_name: str, + embedder: object, + ) -> _CacheTestHandle: ... + @staticmethod def gcs( bucket_name: str, *, diff --git a/litellm/types/agents.py b/litellm/types/agents.py index dbaaab62d86..7f8d8c6af66 100644 --- a/litellm/types/agents.py +++ b/litellm/types/agents.py @@ -2,7 +2,7 @@ from collections.abc import Mapping, Sequence from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal -from pydantic import BaseModel, PrivateAttr, StrictInt +from pydantic import BaseModel, ConfigDict, PrivateAttr, StrictInt from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMPydanticObjectBase @@ -189,6 +189,7 @@ class AgentConfig(TypedDict, total=False): session_rpm_limit: int | None static_headers: dict[str, str] | None extra_headers: list[str] | None + access_group_ids: ReadOnly[Sequence[str] | None] class PatchAgentRequest(TypedDict, total=False): @@ -202,6 +203,21 @@ class PatchAgentRequest(TypedDict, total=False): session_rpm_limit: int | None static_headers: dict[str, str] | None extra_headers: list[str] | None + access_group_ids: ReadOnly[Sequence[str] | None] + + +AGENT_CALLER_USER_ID_HEADER: Final = "x-litellm-user-id" +AGENT_CALLER_TEAM_ID_HEADER: Final = "x-litellm-team-id" + + +class AgentCaller(BaseModel): + """The user and team that invoked an agent, echoed back by the agent on its own proxy calls. + Only ever narrows what the agent's key may do.""" + + model_config = ConfigDict(frozen=True) + + user_id: str | None = None + team_id: str | None = None # Request/Response models for CRUD endpoints @@ -226,6 +242,7 @@ class AgentResponse(BaseModel): session_rpm_limit: int | None = None static_headers: dict[str, str] | None = None extra_headers: list[str] | None = None + access_group_ids: Sequence[str] | None = None keys: list[AgentKeySummary] | None = None search_score: float | None = None created_at: datetime | None = None diff --git a/litellm/types/utils.py b/litellm/types/utils.py index e1d43b7fccb..e23f329ee83 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -459,6 +459,8 @@ class CallTypes(str, Enum): ######################################################### create_video = "create_video" acreate_video = "acreate_video" + video_generation = "video_generation" + avideo_generation = "avideo_generation" avideo_retrieve = "avideo_retrieve" video_retrieve = "video_retrieve" avideo_content = "avideo_content" diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 35c61871e09..80bb2bf9bd2 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.87226e-07, + "input_cost_per_token": 9.5526e-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.774452e-06, + "output_cost_per_token": 1.91052e-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.39355e-08, + "cache_read_input_token_cost": 7.9605e-08, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -43079,22 +43079,22 @@ "supports_web_search": false }, "openrouter/deepseek/deepseek-v4-pro-0813": { - "input_cost_per_token": 1.32e-06, + "input_cost_per_token": 5.58624e-07, "input_cost_per_token_cache_hit": 1.9272e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "output_cost_per_token": 3.96e-06, + "output_cost_per_token": 1.675872e-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": 4.4e-08, - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, + "cache_read_input_token_cost": 1.86208e-08, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.58624e-7,"output_cost_per_token":0.000001675872,"cache_read_input_token_cost":1.86208e-8}, "supports_audio_input": false, "supports_pdf_input": false, "supports_vision": false, @@ -68941,9 +68941,9 @@ "supports_web_search": true }, "openrouter/deepseek/deepseek-v4-flash": { - "input_cost_per_token": 5.544e-08, - "output_cost_per_token": 1.1088e-07, - "cache_read_input_token_cost": 1.1088e-08, + "input_cost_per_token": 8.8606e-08, + "output_cost_per_token": 1.77212e-07, + "cache_read_input_token_cost": 1.77212e-08, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, @@ -70299,8 +70299,8 @@ "supports_web_search": true }, "openrouter/meta-llama/llama-4-maverick": { - "input_cost_per_token": 2e-07, - "output_cost_per_token": 8e-07, + "input_cost_per_token": 1.875e-07, + "output_cost_per_token": 6.525e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 16384, @@ -72999,15 +72999,15 @@ "supports_web_search": false }, "openrouter/~deepseek/deepseek-pro-latest": { - "cache_read_input_token_cost": 4.4e-08, - "input_cost_per_token": 1.32e-06, + "cache_read_input_token_cost": 1.86208e-08, + "input_cost_per_token": 5.58624e-07, "litellm_provider": "openrouter", "max_input_tokens": 1048576, "max_output_tokens": 384000, "max_tokens": 384000, "mode": "chat", - "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":6.6e-7,"output_cost_per_token":0.00000198,"cache_read_input_token_cost":2.2e-8}, - "output_cost_per_token": 3.96e-06, + "off_peak_pricing": {"windows":[{"weekdays":["saturday","sunday"],"hours_utc":"00:00-00:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"00:00-01:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"04:00-06:00"},{"weekdays":["monday","tuesday","wednesday","thursday","friday"],"hours_utc":"10:00-00:00"}],"input_cost_per_token":5.58624e-7,"output_cost_per_token":0.000001675872,"cache_read_input_token_cost":1.86208e-8}, + "output_cost_per_token": 1.675872e-06, "source": "https://openrouter.ai/api/v1/models", "supports_audio_input": false, "supports_function_calling": true, @@ -76879,6 +76879,26 @@ "supports_vision": true, "supports_web_search": true }, + "moonshotai.kimi-k3": { + "cache_creation_input_token_cost": 3.75e-06, + "cache_read_input_token_cost": 3e-07, + "input_cost_per_token": 3e-06, + "litellm_provider": "bedrock_converse", + "max_input_tokens": 1000000, + "max_output_tokens": 131072, + "max_tokens": 131072, + "mode": "chat", + "output_cost_per_token": 1.5e-05, + "source": "https://aws.amazon.com/bedrock/pricing/", + "supports_audio_input": false, + "supports_function_calling": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_vision": true + }, "global.moonshotai.kimi-k3": { "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, diff --git a/pyproject.toml b/pyproject.toml index 95da93df41e..a1b276e4e8c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -15,12 +15,14 @@ dependencies = [ # When changing a floor, verify it installs + imports on every supported # Python with: `uv pip install --resolution=lowest-direct .` "fastuuid>=0.14.0,<1.0", + "filelock>=3.16.1,<4.0", "httpx[http2]>=0.28.0,<1.0", "openai>=2.20.0,<3.0.0", "python-dotenv>=1.0.0,<2.0", "tiktoken>=0.8.0,<1.0; python_version < '3.14'", "tiktoken>=0.12.0,<1.0; python_version >= '3.14'", "importlib-metadata>=8.0.0,<9.0", + "packaging>=24.0", "tokenizers>=0.21.0,<1.0", "click>=8.0.0,<9.0", "jinja2>=3.1.6,<4.0", diff --git a/schema.prisma b/schema.prisma index 368864f9cd1..85996430bc5 100644 --- a/schema.prisma +++ b/schema.prisma @@ -73,6 +73,7 @@ model LiteLLM_AgentsTable { static_headers Json? @default("{}") extra_headers String[] @default([]) agent_access_groups String[] @default([]) + access_group_ids String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) spend Float @default(0.0) @@ -1623,6 +1624,47 @@ model LiteLLM_AutoRouterSession { @@index([last_turn_at], map: "idx_autorouter_session_last_turn") } +model LiteLLM_AutoRouterUserSession { + user_id String + api_key String + session_id String + router_name String + router_type String + first_turn_at DateTime + last_turn_at DateTime + last_model String + models Json @default("{}") + turns Int @default(0) + unordered_turns Int @default(0) + covered_turns Int @default(0) + cache_hits Int @default(0) + same_model_turns Int @default(0) + same_model_hits Int @default(0) + first_visit_turns Int @default(0) + first_visit_hits Int @default(0) + return_turns Int @default(0) + return_hits Int @default(0) + return_expired_misses Int @default(0) + return_within_ttl_misses Int @default(0) + ttl_5m_turns Int @default(0) + ttl_1h_turns Int @default(0) + total_tokens BigInt @default(0) + spend Float @default(0) + saved_spend Float @default(0) + savings_estimated_turns Int @default(0) + savings_estimated_actual_spend Float @default(0) + savings_estimated_saved_spend Float @default(0) + savings_estimated_baseline_models Json @default("{}") + classifier_cost Float @default(0) + classifier_cost_recorded_turns Int @default(0) + tier_turns Json @default("{}") + baseline_models Json @default("{}") + + @@id([user_id, api_key, session_id, router_name]) + @@index([last_turn_at], map: "idx_autorouter_user_session_last_turn") + @@index([user_id, last_turn_at], map: "idx_autorouter_user_session_user_last_turn") +} + // Shadow eval: evaluation of an auto-router against one or more keys' live traffic, in // either direction. forward duplicates the requests the keys did not route through the // router through it, answering whether they should adopt it; reverse duplicates the diff --git a/tests/e2e/coverage_registry/guardrail.yaml b/tests/e2e/coverage_registry/guardrail.yaml index 81832bebf49..86eb44f6cb1 100644 --- a/tests/e2e/coverage_registry/guardrail.yaml +++ b/tests/e2e/coverage_registry/guardrail.yaml @@ -6,6 +6,7 @@ - {id: guardrail.presidio.pre_call.logs_masked_entities, module: guardrail, tier: P0, hook_point: pre_call, assertions: [logs_masked_entities], exercised_on: [chat_completions], source: "guardrail_hooks/presidio.py", rationale: "A masking run must record itself on the spend log: the dashboard's guardrail panel renders the masked-entity counts and per-entity scores straight off metadata.guardrail_information, so a run that masks but records nothing leaves an operator unable to audit it"} - {id: guardrail.bedrock.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_hooks/bedrock_guardrails.py", rationale: "AWS content guardrail blocks harmful input"} - {id: guardrail.litellm_content_filter.pre_call.blocks, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Local content-filter default-on blocks banned keyword pre-call"} +- {id: guardrail.litellm_content_filter.pre_call.blocks_video, module: guardrail, tier: P0, hook_point: pre_call, assertions: [blocks], exercised_on: [videos], source: "test_key_guardrail_video_e2e.py", fail_before_fix: proven, rationale: "A content-filter guardrail attached to a key (metadata.guardrails) blocks a banned prompt on POST /v1/videos before the provider is called; before the fix the route's call type was unknown to the unified guardrail hook and the prompt went to the provider unscanned (LIT-6685)"} - {id: guardrail.litellm_content_filter.pre_call.allows, module: guardrail, tier: P0, hook_point: pre_call, assertions: [allows], exercised_on: [chat_completions], source: "test_team_disable_global_guardrail_e2e.py", rationale: "Team disable_global_guardrails bypasses default-on content filter"} - {id: guardrail.litellm_content_filter.apply_endpoint.blocks, module: guardrail, tier: P0, hook_point: apply_endpoint, assertions: [blocks], exercised_on: [chat_completions], source: "guardrail_endpoints.py:apply_guardrail", rationale: "POST /guardrails/apply_guardrail blocks banned content for customers that call the apply surface directly"} - {id: guardrail.litellm_content_filter.apply_endpoint.allows, module: guardrail, tier: P0, hook_point: apply_endpoint, assertions: [allows], exercised_on: [chat_completions], source: "guardrail_endpoints.py:apply_guardrail", rationale: "POST /guardrails/apply_guardrail returns clean text for allowed input"} diff --git a/tests/e2e/coverage_registry/reliability.yaml b/tests/e2e/coverage_registry/reliability.yaml index 65354100f58..334780eda53 100644 --- a/tests/e2e/coverage_registry/reliability.yaml +++ b/tests/e2e/coverage_registry/reliability.yaml @@ -12,6 +12,7 @@ - {id: reliability.cooldown.429.trips_then_recovers, module: reliability, tier: P0, behavior: cooldown, variant: "429", assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:69", rationale: "Cools on 429, avoids hammering exhausted provider"} - {id: reliability.cooldown.auth.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: auth, assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:74", rationale: "Cools on 401 auth error"} - {id: reliability.cooldown.timeout.trips_then_recovers, module: reliability, tier: P1, behavior: cooldown, variant: timeout, assertions: [trips_then_recovers], exercised_on: [chat_completions], source: "cooldown_handlers.py:77", rationale: "Cools on 408 timeout"} +- {id: reliability.cooldown.client_disconnect.stays_healthy, module: reliability, tier: P0, behavior: cooldown, variant: client_disconnect, assertions: [stays_healthy], exercised_on: [chat_completions], source: "llms/azure/azure.py:484", fail_before_fix: proven, rationale: "A client hanging up mid-request under cancel_on_disconnect never benches the Azure deployment it was talking to: the cancellation used to surface as a fake 500 that tripped the cooldown and sent every caller behind it to billed fallbacks (GitHub issues #35329 and #42222)"} - {id: reliability.routing.simple_shuffle.picks_healthy_deployment, module: reliability, tier: P1, behavior: routing, variant: simple_shuffle, assertions: [picks_healthy_deployment], exercised_on: [chat_completions], source: "router_strategy/simple_shuffle.py", rationale: "Baseline weighted/uniform pick"} - {id: reliability.routing.latency_based.picks_lowest_latency, module: reliability, tier: P1, behavior: routing, variant: latency_based, assertions: [picks_lowest_latency], exercised_on: [chat_completions], source: "router_strategy/lowest_latency.py", rationale: "Routes to lowest-latency deployment"} - {id: reliability.routing.cost_based.picks_lowest_cost, module: reliability, tier: P1, behavior: routing, variant: cost_based, assertions: [picks_lowest_cost], exercised_on: [chat_completions], source: "router_strategy/lowest_cost.py", rationale: "Spend-aware routing"} diff --git a/tests/e2e/e2e_http.py b/tests/e2e/e2e_http.py index d4978601b20..97f1e1671f8 100644 --- a/tests/e2e/e2e_http.py +++ b/tests/e2e/e2e_http.py @@ -676,6 +676,35 @@ def send( return streaming_outcome(resp, stream, sent_at=sent_at) +class AbandonedRequest(BaseModel): + """A non-streaming request whose socket the client closed ``after`` seconds in, + before the proxy had answered.""" + + kind: Literal["abandoned"] = "abandoned" + after: float + + +def abandon( + url: URL, *, headers: BaseModel, json: BaseModel, after: float, connect_timeout: float = 10.0 +) -> AbandonedRequest | StreamingResponse: + """POST and close the connection ``after`` seconds if no response head has arrived + by then; returns the response instead when the proxy answered first.""" + sent_at: Final = time.monotonic() + session: Final = requests.Session() + try: + resp = session.post( + str(url), + headers=_headers(headers), + json=wire_body(json), + timeout=(connect_timeout, after), + ) + except requests.exceptions.ReadTimeout: + return AbandonedRequest(after=after) + finally: + session.close() + return streaming_outcome(resp, False, sent_at=sent_at) + + def stream(url: URL, *, headers: BaseModel, json: BaseModel, timeout: float = 60.0) -> StreamingResponse: """Streaming (SSE) call: consumes the stream counting events, and captures the x-litellm-call-id + content-type headers. Body is elided.""" diff --git a/tests/e2e/gateway/stage_mirror_ci_config.yml b/tests/e2e/gateway/stage_mirror_ci_config.yml index 352caddf588..2d02fedddae 100644 --- a/tests/e2e/gateway/stage_mirror_ci_config.yml +++ b/tests/e2e/gateway/stage_mirror_ci_config.yml @@ -10,6 +10,7 @@ general_settings: store_prompts_in_spend_logs: true database_connection_pool_limit: 10 forward_client_headers_to_llm_api: false + cancel_on_disconnect: true maximum_spend_logs_retention_period: "60d" maximum_spend_logs_cleanup_cron: "0 1 * * *" proxy_budget_rescheduler_min_time: 15 diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index ed112a79b9b..3ceac737399 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -20,6 +20,7 @@ from models import ( ChatResponse, ChatTool, KeyGenerateBody, + KeyMetadata, LiteLLMParamsBody, TeamDeleteBody, TeamInfoParams, @@ -27,6 +28,8 @@ from models import ( TeamMetadata, TeamNewBody, TeamNewResponse, + VideoCreateBody, + VideoCreateResponse, ) from proxy_client import ProxyClient from pydantic import BaseModel @@ -151,12 +154,12 @@ class _ResponsesGuardrailBody(BaseModel): class GuardrailsClient: proxy: ProxyClient - def create_content_filter_guardrail(self, name: str, blocked_keyword: str) -> str: + def create_content_filter_guardrail(self, name: str, blocked_keyword: str, *, default_on: bool = True) -> str: return self.register( name, ContentFilterParamsBody( mode="pre_call", - default_on=True, + default_on=default_on, blocked_words=[BlockedWordBody(keyword=blocked_keyword, action="BLOCK")], ), ) @@ -266,6 +269,21 @@ class GuardrailsClient: def create_key_in_team(self, team_id: str) -> str: return self.proxy.generate_key(KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user")) + def create_key_with_guardrails(self, resources: ResourceManager, guardrails: list[str]) -> str: + key = self.proxy.generate_key( + KeyGenerateBody(user_id="e2e-guardrails-user", metadata=KeyMetadata(guardrails=guardrails)) + ) + resources.defer(lambda: self.proxy.delete_key(key)) + return key + + def create_video(self, key: str, model: str, prompt: str) -> Result[VideoCreateResponse]: + return self.proxy.transport.post( + "/v1/videos", + headers=self.proxy.transport.bearer(key), + json=VideoCreateBody(model=model, prompt=prompt, seconds="4"), + response_type=VideoCreateResponse, + ) + def chat( self, key: str, diff --git a/tests/e2e/guardrails/test_key_guardrail_video_e2e.py b/tests/e2e/guardrails/test_key_guardrail_video_e2e.py new file mode 100644 index 00000000000..5f318e141a5 --- /dev/null +++ b/tests/e2e/guardrails/test_key_guardrail_video_e2e.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +import pytest +from e2e_config import unique_marker +from e2e_http import Success, UnknownApiError +from guardrails_client import GuardrailsClient, poll_until_blocked +from lifecycle import ResourceManager +from models import LiteLLMParamsBody + +pytestmark = pytest.mark.e2e + +CHAT_MODEL = "gemini-2.5-flash" +VIDEO_BACKEND = "vertex_ai/veo-3.1-fast-generate-001" + + +def _video_prompt_with(banned_keyword: str) -> str: + return f"A short clip of a paper boat floating down a stream. {banned_keyword}" + + +def _create_video_model(client: GuardrailsClient, resources: ResourceManager) -> str: + model_name = f"e2e-guard-video-{unique_marker()}" + model_id = client.proxy.create_model( + model_name, + LiteLLMParamsBody( + model=VIDEO_BACKEND, + vertex_project="os.environ/VERTEXAI_PROJECT", + vertex_location="os.environ/VERTEXAI_LOCATION", + vertex_credentials="os.environ/VERTEXAI_CREDENTIALS", + ), + provider_live=True, + ) + resources.defer(lambda: client.proxy.delete_model(model_id)) + return model_name + + +class TestKeyAttachedGuardrailOnVideos: + @pytest.mark.covers( + "guardrail.litellm_content_filter.pre_call.blocks_video", + exercised_on=["videos"], + ) + def test_key_attached_content_filter_blocks_banned_video_prompt( + self, client: GuardrailsClient, resources: ResourceManager + ) -> None: + banned = unique_marker() + guardrail_name = f"e2e-video-filter-{banned}" + guardrail_id = client.create_content_filter_guardrail(guardrail_name, banned, default_on=False) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + key = client.create_key_with_guardrails(resources, [guardrail_name]) + model = _create_video_model(client, resources) + + synced = poll_until_blocked(lambda: client.chat(key, CHAT_MODEL, _video_prompt_with(banned))) + assert isinstance(synced, UnknownApiError) and synced.status_code == 400, ( + f"key guardrail {guardrail_name!r} never synced to the proxy on /chat/completions: {synced}" + ) + + result = client.create_video(key, model, _video_prompt_with(banned)) + match result: + case UnknownApiError(status_code=status, body=body): + assert status == 400, f"expected a 400 guardrail block, got {status}: {body[:300]}" + assert "content blocked" in body.lower() or banned in body, ( + f"block response missing content-filter reason: {body[:300]}" + ) + case Success(data=video): + pytest.fail( + f"key-attached guardrail {guardrail_name!r} was skipped on /v1/videos: " + f"the banned prompt reached the provider and started video job {video.id}" + ) + case _: + pytest.fail(f"unexpected /v1/videos outcome for a banned prompt: {result}") diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 355329585fb..2608da2e4fc 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -60,6 +60,7 @@ class KeyMetadata(BaseModel): priority: str | None = None batch_enqueued_token_limit: int | None = None tag: str | None = None + guardrails: list[str] | None = None class ObjectPermission(BaseModel): @@ -697,6 +698,20 @@ class EmbedResponse(BaseModel): model: str | None = None +# ---------- videos ---------- + + +class VideoCreateBody(BaseModel): + model: str + prompt: str + seconds: str | None = None + + +class VideoCreateResponse(BaseModel): + id: str + status: str | None = None + + # ---------- rerank ---------- @@ -936,6 +951,23 @@ class RouterSettingsResponse(BaseModel): current_values: RouterCurrentValues +class ConfigListParams(BaseModel): + config_type: Literal["general_settings"] + + +class ConfigField(BaseModel): + """One row of GET /config/list: a general_settings field and the value the + proxy is running with, the two fields a test preconditions on.""" + + model_config = ConfigDict(extra="ignore") + field_name: str + field_value: JsonValue = None + + +class ConfigFieldList(RootModel[tuple[ConfigField, ...]]): + """GET /config/list answers with a bare array of general_settings fields.""" + + class CostMapEntry(BaseModel): model_config = ConfigDict(extra="ignore") litellm_provider: str | None = None @@ -1051,6 +1083,7 @@ class ModelInfoBody(BaseModel): mode: ModelMode | None = None access_groups: list[str] | None = None team_id: str | None = None + allowed_fails: int | None = None allowed_fails_policy: dict[str, int] | None = None diff --git a/tests/e2e/proxy_client.py b/tests/e2e/proxy_client.py index 2f32361e083..76c57452c69 100644 --- a/tests/e2e/proxy_client.py +++ b/tests/e2e/proxy_client.py @@ -47,6 +47,8 @@ from models import ( AnthropicMessagesResponse, ChatBody, ChatResponse, + ConfigFieldList, + ConfigListParams, CostMap, CostMapEntry, CountTokensBody, @@ -630,6 +632,19 @@ class ProxyClient: provider_live=provider_live, ) + def general_setting_enabled(self, field_name: str) -> bool: + """Whether the proxy is running with the named general_settings flag on, for + a test whose behavior only exists under a config flag the stack has to carry.""" + fields = unwrap( + self.transport.get( + "/config/list", + headers=self.transport.master, + params=ConfigListParams(config_type="general_settings"), + response_type=ConfigFieldList, + ) + ).root + return any(entry.field_name == field_name and entry.field_value is True for entry in fields) + def register_model( self, body: ModelNewBody, listed_for: str | None = None, *, provider_live: bool = False ) -> str: diff --git a/tests/e2e/router/reliability_support.py b/tests/e2e/router/reliability_support.py index 976c05ffceb..3d5b76f6408 100644 --- a/tests/e2e/router/reliability_support.py +++ b/tests/e2e/router/reliability_support.py @@ -42,7 +42,7 @@ REAL_KEY = "os.environ/OPENAI_API_KEY" CACHING_MODEL = "anthropic/claude-haiku-4-5" CACHING_KEY = "os.environ/ANTHROPIC_API_KEY" -CONTENT_FILTERED_MODEL = "azure/gpt-5.4-nano" +AZURE_MODEL = "azure/gpt-5.4-nano" AZURE_KEY = "os.environ/AZURE_API_KEY" AZURE_BASE = "os.environ/AZURE_API_BASE" AZURE_API_VERSION = "2024-10-21" @@ -53,6 +53,7 @@ CONTENT_POLICY_PROMPT = ( ) COOLDOWN_SECONDS = 30.0 +REPLICA_PROPAGATION_SECONDS = 15.0 # The smallest-context chat model OpenAI still serves (16385 tokens). A prompt # past that limit comes back as a real `context_length_exceeded` 400, which is @@ -111,7 +112,7 @@ def create_content_filtered_deployment(proxy: ProxyClient, name: str) -> str: return proxy.create_model( name, LiteLLMParamsBody( - model=CONTENT_FILTERED_MODEL, + model=AZURE_MODEL, api_key=AZURE_KEY, api_base=AZURE_BASE, api_version=AZURE_API_VERSION, @@ -120,6 +121,26 @@ def create_content_filtered_deployment(proxy: ProxyClient, name: str) -> str: ) +def create_azure_benched_on_first_failure_deployment(proxy: ProxyClient, name: str, cooldown_time: float) -> str: + """The live Azure OpenAI deployment holding all of the group's shuffle weight, + benched on its first failure of any class, with the client's own retries off.""" + return proxy.register_model( + ModelNewBody( + model_name=name, + litellm_params=LiteLLMParamsBody( + model=AZURE_MODEL, + api_key=AZURE_KEY, + api_base=AZURE_BASE, + api_version=AZURE_API_VERSION, + max_retries=0, + weight=1, + cooldown_time=cooldown_time, + ), + model_info=ModelInfoBody(allowed_fails=0), + ) + ) + + def create_caching_deployment(proxy: ProxyClient, name: str) -> str: """Register the Anthropic deployment whose prompt cache the affinity check pins to.""" return proxy.create_model(name, LiteLLMParamsBody(model=CACHING_MODEL, api_key=CACHING_KEY, weight=1)) diff --git a/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py b/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py new file mode 100644 index 00000000000..06174e97d20 --- /dev/null +++ b/tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py @@ -0,0 +1,139 @@ +"""Live e2e: a client hanging up mid-request under cancel_on_disconnect never +benches the deployment it was talking to. + +The group is the cooldown suite's pair: the live Azure deployment holding all of +the shuffle weight, benched on its first failure of any class with a cooldown that +outlasts the test, plus a healthy backup at weight 0 the shuffle only reaches once +the Azure deployment is benched. A cheap call first proves the Azure deployment +answers the key and warms its auth path. The test then asks for an answer far +longer than CLIENT_HANGS_UP_AFTER_SECONDS of generation, retries off, and hangs up +that many seconds in: late enough that the proxy has handed the call to Azure (a +hang-up before the provider call is in flight cancels nothing the router could +bench, so the cell would pass vacuously). An answer that comes back inside the +window proves nothing and benches nothing either, since a success never counts +against the deployment, so the cell asks again up to HANG_UP_ATTEMPTS times and +fails out loud naming the window only when every ask came back early. After the +cooldown suite's replica propagation window, every one of the next calls has to +come back 200 from the Azure deployment itself, named in x-litellm-model-id; a +single answer from the backup means the hang-up was booked as a failure. + +The test reads `cancel_on_disconnect` back from the proxy first: without the flag +the hang-up cancels nothing and the cell would pass vacuously. +""" + +from __future__ import annotations + +import time + +import pytest +from complexity_router_client import ComplexityRouterClient +from e2e_config import unique_marker +from e2e_http import AbandonedRequest, StreamingResponse +from lifecycle import ResourceManager +from models import ChatMessage, ReliabilityChatBody, RouterSettingsOverride +from reliability_support import ( + REPLICA_PROPAGATION_SECONDS, + chat_override, + create_azure_benched_on_first_failure_deployment, + create_zero_weight_backup_deployment, + model_id_of, +) + +pytestmark = pytest.mark.e2e + +CLIENT_HANGS_UP_AFTER_SECONDS = 5.0 +HANG_UP_ATTEMPTS = 3 +LONG_ANSWER_MAX_TOKENS = 16384 +BENCH_OUTLASTS_TEST_SECONDS = 300.0 +CALLS_AFTER_HANGUP = 6 + + +def _say_hi(client: ComplexityRouterClient, key: str, group: str) -> StreamingResponse: + return chat_override( + client.proxy, + key, + group, + f"say hi {unique_marker()}", + override=RouterSettingsOverride(num_retries=0), + ) + + +def _ask_for_a_long_answer_then_hang_up( + client: ComplexityRouterClient, key: str, group: str +) -> AbandonedRequest | StreamingResponse: + return client.proxy.transport.abandon( + "/chat/completions", + headers=client.proxy.transport.bearer(key), + json=ReliabilityChatBody( + model=group, + messages=[ + ChatMessage( + role="user", + content=( + "Write an essay on the history of the telegraph with one section per decade from the 1830s " + f"to the 2020s, each section at least 300 words. {unique_marker()}" + ), + ) + ], + max_tokens=LONG_ANSWER_MAX_TOKENS, + router_settings_override=RouterSettingsOverride(num_retries=0), + ), + after=CLIENT_HANGS_UP_AFTER_SECONDS, + ) + + +def _hang_up_mid_answer(client: ComplexityRouterClient, key: str, group: str) -> None: + for attempt in range(1, HANG_UP_ATTEMPTS + 1): + match _ask_for_a_long_answer_then_hang_up(client, key, group): + case AbandonedRequest(): + return + case StreamingResponse(status_code=200): + continue + case StreamingResponse(status_code=status_code, body=body): + pytest.fail( + f"hang-up attempt {attempt} should have found the long answer still in flight after " + f"{CLIENT_HANGS_UP_AFTER_SECONDS:.0f}s, but the proxy answered {status_code}: {body[:300]}" + ) + pytest.fail( + f"the proxy answered all {HANG_UP_ATTEMPTS} long asks within {CLIENT_HANGS_UP_AFTER_SECONDS:.0f}s, so the " + "client never hung up with a call still in flight and the bench this cell guards against could not happen" + ) + + +class TestReliabilityCancelOnDisconnect: + @pytest.mark.covers("reliability.cooldown.client_disconnect.stays_healthy") + def test_client_hanging_up_never_benches_the_deployment( + self, client: ComplexityRouterClient, resources: ResourceManager, scoped_key: str + ) -> None: + assert client.proxy.general_setting_enabled("cancel_on_disconnect"), ( + "this cell needs general_settings.cancel_on_disconnect: true in the proxy config; without it the " + "hang-up cancels nothing and the bench it guards against can never happen" + ) + + group = f"reliability-cooldown-disconnect-{unique_marker()}" + azure = create_azure_benched_on_first_failure_deployment( + client.proxy, group, cooldown_time=BENCH_OUTLASTS_TEST_SECONDS + ) + resources.defer(lambda: client.proxy.delete_model(azure)) + backup = create_zero_weight_backup_deployment(client.proxy, group) + resources.defer(lambda: client.proxy.delete_model(backup)) + + warm_up = _say_hi(client, scoped_key, group) + assert warm_up.status_code == 200 and model_id_of(warm_up) == azure, ( + f"before any hang-up the Azure deployment {azure} should answer the group, got {warm_up.status_code} " + f"from {model_id_of(warm_up)!r}: {warm_up.body[:300]}" + ) + + _hang_up_mid_answer(client, scoped_key, group) + time.sleep(REPLICA_PROPAGATION_SECONDS) + + for call in range(1, CALLS_AFTER_HANGUP + 1): + resp = _say_hi(client, scoped_key, group) + assert resp.status_code == 200, ( + f"call {call} after the hang-up should have been a plain 200 from the group, got " + f"{resp.status_code}: {resp.body[:300]}" + ) + assert model_id_of(resp) == azure, ( + f"call {call} after the hang-up should have been served by the Azure deployment {azure}, the proxy " + f"named {model_id_of(resp)!r}: the cancelled call was booked as a failure and benched it" + ) diff --git a/tests/e2e/router/test_reliability_cooldowns_e2e.py b/tests/e2e/router/test_reliability_cooldowns_e2e.py index 5b5cec09f06..769971e1533 100644 --- a/tests/e2e/router/test_reliability_cooldowns_e2e.py +++ b/tests/e2e/router/test_reliability_cooldowns_e2e.py @@ -43,6 +43,7 @@ from lifecycle import ResourceManager from models import KeyGenerateBody, RouterSettingsOverride from reliability_support import ( COOLDOWN_SECONDS, + REPLICA_PROPAGATION_SECONDS, chat_override, create_always_5xx_deployment, create_always_rate_limited_deployment, @@ -57,7 +58,6 @@ from reliability_support import ( pytestmark = pytest.mark.e2e RECOVERY_GRACE_SECONDS = 10 -REPLICA_PROPAGATION_SECONDS = 15.0 PROPAGATION_POLL_SECONDS = 0.25 BENCH_MARGIN_SECONDS = 4.0 diff --git a/tests/e2e/transport.py b/tests/e2e/transport.py index 0022c0c4355..a3eec815441 100644 --- a/tests/e2e/transport.py +++ b/tests/e2e/transport.py @@ -13,6 +13,7 @@ from typing import Protocol import e2e_http from e2e_http import ( URL, + AbandonedRequest, AuthHeaders, BinaryStream, NetworkError, @@ -58,6 +59,10 @@ class Transport(Protocol): stream: bool = False, ) -> StreamingResponse: ... + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: ... + def get[R: BaseModel]( self, path: str, @@ -243,6 +248,11 @@ class HttpTransport: timeout=self.request_timeout, ) + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: + return e2e_http.abandon(self._url(path), headers=headers, json=json, after=after) + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: return e2e_http.probe( self._url(path), @@ -420,6 +430,11 @@ class SplitTransport: ) -> StreamingResponse: return self._route(path).send(path, headers=headers, json=json, params=params, stream=stream) + def abandon( + self, path: str, *, headers: BaseModel, json: BaseModel, after: float + ) -> AbandonedRequest | StreamingResponse: + return self._route(path).abandon(path, headers=headers, json=json, after=after) + def probe(self, path: str, *, params: BaseModel, headers: BaseModel | None = None) -> ProbeResult: return self._route(path).probe(path, params=params, headers=headers) diff --git a/tests/e2e/ui/helpers/roundTrip.ts b/tests/e2e/ui/helpers/roundTrip.ts index 8d6e264e622..91e48b7087c 100644 --- a/tests/e2e/ui/helpers/roundTrip.ts +++ b/tests/e2e/ui/helpers/roundTrip.ts @@ -12,17 +12,111 @@ export async function captureRequestBody( match: { method: string; urlIncludes: string }, action: () => Promise, ): Promise> { - const pending = page.waitForRequest((req) => req.method() === match.method && req.url().includes(match.urlIncludes)); + const pending = page.waitForRequest( + (req) => + req.method() === match.method && req.url().includes(match.urlIncludes), + ); await action(); const request = await pending; return JSON.parse(request.postData() ?? "{}") as Record; } /** Reads an endpoint as the master key, so a failure is bad data and not an expired UI token. */ -export async function readBack(page: Page, endpoint: string): Promise { +export async function readBack( + page: Page, + endpoint: string, +): Promise { const res = await page.request.get(endpoint, { headers: { Authorization: `Bearer ${masterKey()}` }, }); expect(res.ok(), `GET ${endpoint}`).toBe(true); return (await res.json()) as T; } + +type OperationOutcome = + | { readonly status: "success" } + | { readonly status: "failure"; readonly error: unknown }; + +type RunFailure = + | { readonly status: "action_failure"; readonly error: unknown } + | { readonly status: "cleanup_failure"; readonly error: unknown } + | { + readonly status: "action_and_cleanup_failure"; + readonly actionError: unknown; + readonly cleanupError: unknown; + }; + +function toRunFailure( + actionOutcome: OperationOutcome, + cleanupOutcome: OperationOutcome, +): RunFailure | null { + if ( + actionOutcome.status === "failure" && + cleanupOutcome.status === "failure" + ) { + return { + status: "action_and_cleanup_failure", + actionError: actionOutcome.error, + cleanupError: cleanupOutcome.error, + }; + } + if (actionOutcome.status === "failure") { + return { status: "action_failure", error: actionOutcome.error }; + } + if (cleanupOutcome.status === "failure") { + return { status: "cleanup_failure", error: cleanupOutcome.error }; + } + return null; +} + +function raiseRunFailure(failure: RunFailure): never { + switch (failure.status) { + case "action_failure": + throw failure.error; + case "cleanup_failure": + throw failure.error; + case "action_and_cleanup_failure": + throw new AggregateError( + [failure.actionError, failure.cleanupError], + "Action and cleanup failed", + ); + } +} + +async function runAction( + action: () => void | Promise, +): Promise { + return Promise.resolve() + .then(action) + .then( + () => ({ status: "success" as const }), + (error: unknown) => ({ status: "failure" as const, error }), + ); +} + +async function runCleanup( + cleanup: () => boolean | Promise, +): Promise { + return Promise.resolve() + .then(cleanup) + .then( + (succeeded) => + succeeded + ? { status: "success" as const } + : { + status: "failure" as const, + error: new Error("Failed to clean up UI E2E resource"), + }, + (error: unknown) => ({ status: "failure" as const, error }), + ); +} + +export async function runWithCleanup( + action: () => void | Promise, + cleanup: () => boolean | Promise, +): Promise { + const actionOutcome = await runAction(action); + const cleanupOutcome = await runCleanup(cleanup); + const failure = toRunFailure(actionOutcome, cleanupOutcome); + if (failure !== null) raiseRunFailure(failure); +} diff --git a/tests/e2e/ui/tests/prompts/addPrompt.spec.ts b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts new file mode 100644 index 00000000000..9d85236c4a6 --- /dev/null +++ b/tests/e2e/ui/tests/prompts/addPrompt.spec.ts @@ -0,0 +1,75 @@ +import { test, expect } from "@playwright/test"; + +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page as DashboardPage } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { runWithCleanup } from "../../helpers/roundTrip"; +import { masterKey, uniqueSuffix } from "../../helpers/traffic"; + +test.use({ storageState: ADMIN_STORAGE_PATH }); + +test.describe("Prompt upload form", () => { + test("uploads a prompt file and reads the created prompt back", async ({ + page, + }) => { + const promptId = `e2e-prompt-${uniqueSuffix()}`; + const promptContent = "Hello {{name}}"; + + await runWithCleanup( + async () => { + await navigateToPage(page, DashboardPage.Prompts); + await page.getByRole("button", { name: "Upload .prompt File" }).click(); + await expect( + page.getByRole("dialog", { name: "Add New Prompt" }), + ).toBeVisible(); + await page.getByLabel("Prompt ID").fill(promptId); + await page.locator('input[type="file"]').setInputFiles({ + name: "e2e.prompt", + mimeType: "text/plain", + buffer: Buffer.from( + `---\nmodel: fake-openai-gpt-4\n---\n${promptContent}\n`, + ), + }); + await expect(page.getByText("Selected: e2e.prompt")).toBeVisible(); + await page.getByRole("button", { name: "Create Prompt" }).click(); + + await expect + .poll(async () => { + const response = await page.request.get( + `/prompts/${encodeURIComponent(promptId)}/info?environment=development`, + { + headers: { Authorization: `Bearer ${masterKey()}` }, + }, + ); + return response.ok(); + }) + .toBe(true); + await expect + .poll(async () => { + const response = await page.request.get( + `/prompts/${encodeURIComponent(promptId)}/info?environment=development`, + { + headers: { Authorization: `Bearer ${masterKey()}` }, + }, + ); + if (!response.ok()) return undefined; + const promptInfo = (await response.json()) as { + raw_prompt_template?: { content?: string }; + }; + return promptInfo.raw_prompt_template?.content; + }) + .toBe(promptContent); + await expect(page.getByText(promptId, { exact: true })).toBeVisible(); + }, + async () => { + const response = await page.request.delete( + `/prompts/${encodeURIComponent(promptId)}?environment=development`, + { + headers: { Authorization: `Bearer ${masterKey()}` }, + }, + ); + return response.ok(); + }, + ); + }); +}); diff --git a/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts new file mode 100644 index 00000000000..fe659080eab --- /dev/null +++ b/tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts @@ -0,0 +1,84 @@ +import { test, expect } from "@playwright/test"; + +import { ADMIN_STORAGE_PATH } from "../../constants"; +import { Page as DashboardPage } from "../../fixtures/pages"; +import { navigateToPage } from "../../helpers/navigation"; +import { + captureRequestBody, + readBack, + runWithCleanup, +} from "../../helpers/roundTrip"; +import { masterKey, uniqueSuffix } from "../../helpers/traffic"; + +test.use({ storageState: ADMIN_STORAGE_PATH }); + +test.describe("Tag management", () => { + test("creates, edits, reopens, and reads back a tag", async ({ page }) => { + const tagName = `e2e-tag-${uniqueSuffix()}`; + const description = "synthetic tag description"; + const updatedDescription = `${description} updated`; + + await runWithCleanup( + async () => { + await navigateToPage(page, DashboardPage.TagManagement); + await page.getByRole("button", { name: "+ Create New Tag" }).click(); + await expect( + page.getByRole("dialog", { name: "Create New Tag" }), + ).toBeVisible(); + await page.getByLabel("Tag Name").fill(tagName); + await page.getByLabel("Description").fill(description); + await page.getByRole("button", { name: "Create Tag" }).click(); + + await expect + .poll(async () => { + const response = await readBack>( + page, + "/tag/list", + ); + return response.some((tag) => tag.name === tagName); + }) + .toBe(true); + await expect(page.getByText(tagName, { exact: true })).toBeVisible(); + + await page.getByText(tagName, { exact: true }).click(); + await expect(page.getByText("Tag Name:")).toBeVisible(); + await page.getByRole("button", { name: "Edit Tag" }).click(); + await page.getByLabel("Description").fill(updatedDescription); + const updateBody = await captureRequestBody( + page, + { method: "POST", urlIncludes: "/tag/update" }, + () => page.getByRole("button", { name: "Save Changes" }).click(), + ); + expect(updateBody).toMatchObject({ + name: tagName, + description: updatedDescription, + }); + + await expect + .poll(async () => { + const infoResponse = await page.request.post("/tag/info", { + headers: { Authorization: `Bearer ${masterKey()}` }, + data: { names: [tagName] }, + }); + expect(infoResponse.ok()).toBe(true); + const info = (await infoResponse.json()) as Record< + string, + { description?: string } + >; + return info[tagName]?.description; + }) + .toBe(updatedDescription); + }, + async () => { + const response = await page.request.post("/tag/delete", { + headers: { + Authorization: `Bearer ${masterKey()}`, + "Content-Type": "application/json", + }, + data: { name: tagName }, + }); + return response.ok(); + }, + ); + }); +}); diff --git a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py index f3c68b489a5..77549b527d8 100644 --- a/tests/proxy_behavior/spend/test_autorouter_session_rollup.py +++ b/tests/proxy_behavior/spend/test_autorouter_session_rollup.py @@ -6,17 +6,24 @@ tests/test_litellm/proxy/db/test_autorouter_session_rollup.py. """ import asyncio +import time import uuid from datetime import datetime, timedelta, timezone -from typing import Final +from types import SimpleNamespace +from typing import Final, TypedDict, cast import pytest from prisma import Prisma +from prisma.errors import RawQueryError +from typing_extensions import ReadOnly from litellm.proxy.db.autorouter_session_rollup import ( AUTOROUTER_BENCHMARKS_SQL, UPSERT_AUTOROUTER_SESSION_SQL, + AutoRouterTurnTransaction, + flush_autorouter_turn_transactions, ) +from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup pytestmark = pytest.mark.asyncio(loop_scope="session") @@ -45,6 +52,7 @@ async def _turn( tier: "str | None" = None, baseline: "str | None" = None, estimated: bool = True, + user_id: str = "", ) -> None: touched: Final = 1 if (hit or ttl is not None or not covered) else 0 await db.execute_raw( @@ -68,6 +76,7 @@ async def _turn( int(estimated), spend if estimated else 0.0, saved if estimated else 0.0, + user_id, ) @@ -217,7 +226,7 @@ async def test_subtotal_coverage_survives_legacy_and_rolling_writers(db, writers assert row["savings_estimated_actual_spend"] == pytest.approx(0.01 * sum(writers)) assert row["savings_estimated_saved_spend"] == pytest.approx(0.02 * sum(writers)) groups: Final = await db.query_raw( - AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key + AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key, None ) assert len(groups) == 1 assert groups[0]["classifier_cost"] == row["classifier_cost"] @@ -242,7 +251,7 @@ async def test_unknown_and_legacy_turns_preserve_actual_spend_without_entering_t assert row["saved_spend"] == pytest.approx(-0.03) assert row["savings_estimated_baseline_models"] == {"opus": 1} groups: Final = await db.query_raw( - AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key + AUTOROUTER_BENCHMARKS_SQL, T0.isoformat(), (T0 + timedelta(days=1)).isoformat(), key, None ) assert len(groups) == 1 for actual in (row, groups[0]): @@ -277,6 +286,7 @@ async def test_the_benchmarks_aggregate_reads_only_overlapping_sessions(db): (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), None, + None, ) matching = [row for row in rows if row["router_name"] == router] assert len(matching) == 1 @@ -304,6 +314,7 @@ async def test_the_benchmarks_aggregate_can_filter_to_one_key(db): (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), first_key, + None, ) matching = [row for row in rows if row["router_name"] == router] assert len(matching) == 1 @@ -317,10 +328,160 @@ async def test_the_benchmarks_aggregate_can_filter_to_one_key(db): (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), f"k-{uuid.uuid4()}", + None, ) assert [row for row in unknown_key_rows if row["router_name"] == router] == [] +class _BenchmarkRow(TypedDict): + sessions: ReadOnly[int] + turns: ReadOnly[int] + same_model_turns: ReadOnly[int] + first_visit_turns: ReadOnly[int] + spend: ReadOnly[float] + saved_spend: ReadOnly[float] + tier_turns: ReadOnly[dict[str, int]] + cache_hits: ReadOnly[int] + savings_estimated_turns: ReadOnly[int] + savings_estimated_actual_spend: ReadOnly[float] + savings_estimated_saved_spend: ReadOnly[float] + + +async def _scoped_benchmarks( + db: Prisma, router: str, user_id: str | None = None, key: str | None = None +) -> tuple[_BenchmarkRow, ...]: + rows: Final = await db.query_raw( + AUTOROUTER_BENCHMARKS_SQL, + (T0 - timedelta(days=1)).isoformat(), + (T0 + timedelta(days=1)).isoformat(), + key, + user_id, + ) + return tuple(cast(_BenchmarkRow, row) for row in rows if row["router_name"] == router) + + +async def test_users_keep_written_identity_across_shared_keys_and_keyless_sessions(db: Prisma) -> None: + router: Final = f"r-{uuid.uuid4()}" + alice: Final = f"u-{uuid.uuid4()}" + bob: Final = f"u-{uuid.uuid4()}" + first_key: Final = f"k-{uuid.uuid4()}" + second_key: Final = f"k-{uuid.uuid4()}" + await _legacy_turn(db, first_key, T0, router=router) + await _turn(db, first_key, "A", T0 + timedelta(seconds=10), router=router, user_id=alice, tier="simple") + await _turn( + db, first_key, "B", T0 + timedelta(seconds=20), router=router, user_id=bob, spend=0.03, saved=0.06, tier="complex" + ) + await _turn(db, second_key, "C", T0, router=router, user_id=alice, spend=0.02, saved=0.04) + await _turn(db, "", "A", T0, router=router, user_id=alice, ttl=300) + await _turn(db, "", "A", T0 + timedelta(seconds=1), router=router, user_id=alice, hit=1) + await _turn(db, "", "B", T0, router=router, user_id=bob, spend=0.04, saved=0.08) + await _turn(db, second_key, "C", T0 - timedelta(days=40), router=router, user_id=alice, session_id="expired") + + alice_rows: Final = await _scoped_benchmarks(db, router, user_id=alice) + bob_rows: Final = await _scoped_benchmarks(db, router, user_id=bob) + global_rows: Final = await _scoped_benchmarks(db, router) + key_rows: Final = await _scoped_benchmarks(db, router, key=first_key) + intersection: Final = await _scoped_benchmarks(db, router, user_id=alice, key=first_key) + assert len(alice_rows) == len(bob_rows) == len(global_rows) == len(key_rows) == len(intersection) == 1 + assert (alice_rows[0]["sessions"], alice_rows[0]["turns"], alice_rows[0]["same_model_turns"]) == (3, 4, 1) + assert (bob_rows[0]["sessions"], bob_rows[0]["turns"], bob_rows[0]["first_visit_turns"]) == (2, 2, 2) + assert alice_rows[0]["spend"] == pytest.approx(0.05) + assert bob_rows[0]["spend"] == pytest.approx(0.07) + assert alice_rows[0]["tier_turns"] == {"simple": 1} + assert bob_rows[0]["tier_turns"] == {"complex": 1} + assert (alice_rows[0]["cache_hits"], bob_rows[0]["cache_hits"]) == (1, 0) + assert (global_rows[0]["sessions"], global_rows[0]["turns"]) == (4, 7) + assert (alice_rows[0]["savings_estimated_turns"], bob_rows[0]["savings_estimated_turns"]) == (4, 2) + assert global_rows[0]["savings_estimated_turns"] == 6 + for scoped in (alice_rows[0], bob_rows[0]): + assert scoped["savings_estimated_actual_spend"] == pytest.approx(scoped["spend"]) + assert scoped["savings_estimated_saved_spend"] == pytest.approx(scoped["saved_spend"]) + assert global_rows[0]["spend"] == pytest.approx(alice_rows[0]["spend"] + bob_rows[0]["spend"] + 0.01) + assert global_rows[0]["saved_spend"] == pytest.approx(alice_rows[0]["saved_spend"] + bob_rows[0]["saved_spend"] + 0.02) + assert global_rows[0]["tier_turns"] == {"simple": 1, "complex": 1} + assert (key_rows[0]["sessions"], key_rows[0]["turns"]) == (1, 3) + assert key_rows[0]["spend"] == pytest.approx(0.05) + assert (intersection[0]["sessions"], intersection[0]["turns"]) == (1, 1) + assert intersection[0]["spend"] == pytest.approx(0.01) + assert await _scoped_benchmarks(db, router, user_id=bob, key=second_key) == () + assert await _scoped_benchmarks(db, router, user_id=f"u-{uuid.uuid4()}") == () + assert await _scoped_benchmarks(db, router, user_id="") == () + + +async def test_a_failed_user_projection_rolls_back_the_keys_increment(db: Prisma) -> None: + key: Final = f"k-{uuid.uuid4()}" + user_id: Final = "".join(str(uuid.uuid4()) for _ in range(200)) + await _turn(db, key, "A", T0) + before: Final = await _row(db, key) + + with pytest.raises(RawQueryError, match=r"index row (requires|size)"): + await _turn(db, key, "B", T0 + timedelta(seconds=1), user_id=user_id) + + assert await _row(db, key) == before + assert await db.query_raw('SELECT user_id FROM "LiteLLM_AutoRouterUserSession" WHERE user_id = $1', user_id) == [] + + first_user: Final = f"u-{uuid.uuid4()}" + second_user: Final = f"u-{uuid.uuid4()}" + turns: Final = tuple( + AutoRouterTurnTransaction( + api_key=key, + user_id=user, + session_id="s1", + router_name="auto-1", + router_type="complexity", + model=model, + turn_at=T0 + timedelta(seconds=second), + total_tokens=100, + spend=0.01, + saved_spend=0.02, + classifier_cost=0.0, + covered=True, + cache_hit=False, + cache_ttl_seconds=None, + cache_touched=False, + ) + for user, model, second in ( + (first_user, "A", 1), + (user_id, "B", 2), + (first_user, "B", 3), + (second_user, "C", 4), + (first_user, "B", 5), + (second_user, "C", 6), + (user_id, "A", 7), + ) + ) + await flush_autorouter_turn_transactions(SimpleNamespace(db=db), tuple(reversed(turns)), n_retry_times=0) + + key_row: Final = await _row(db, key) + assert (key_row["turns"], key_row["last_model"], key_row["unordered_turns"]) == (2, "A", 0) + assert key_row["spend"] == pytest.approx(0.02) + user_rows: Final = await db.query_raw('SELECT * FROM "LiteLLM_AutoRouterUserSession" WHERE api_key = $1', key) + by_user: Final = {row["user_id"]: row for row in user_rows} + assert set(by_user) == {first_user, second_user} + for user, count, model in ((first_user, 3, "B"), (second_user, 2, "C")): + row: Final = by_user[user] + assert (row["turns"], row["same_model_turns"], row["unordered_turns"], row["last_model"]) == (count, 1, 0, model) + assert row["spend"] == pytest.approx(count * 0.01) + assert row["saved_spend"] == pytest.approx(count * 0.02) + + +async def test_user_session_cleanup_keeps_another_users_recent_keyless_session(db: Prisma) -> None: + router: Final = f"r-{uuid.uuid4()}" + expired_user: Final = f"u-{uuid.uuid4()}" + recent_user: Final = f"u-{uuid.uuid4()}" + await _turn(db, "", "A", T0 - timedelta(days=1), router=router, user_id=expired_user) + await _turn(db, "", "A", T0 + timedelta(days=1), router=router, user_id=recent_user) + cleaner: Final = SpendLogCleanup(general_settings={}) + + await cleaner._delete_old_autorouter_user_session_rows( + SimpleNamespace(db=db), T0.replace(tzinfo=timezone.utc), time.monotonic() + 60 + ) + + assert await db.query_raw( + 'SELECT user_id, turns FROM "LiteLLM_AutoRouterUserSession" WHERE router_name = $1', router + ) == [{"user_id": recent_user, "turns": 1}] + + async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db): key = f"k-{uuid.uuid4()}" router = f"r-{uuid.uuid4()}" @@ -334,6 +495,7 @@ async def test_a_reconfigured_alias_reports_each_router_type_as_its_own_group(db (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), None, + None, ) matching = sorted( (row for row in rows if row["router_name"] == router), @@ -418,6 +580,7 @@ async def test_the_benchmarks_aggregate_sums_tier_turns_across_sessions(db): (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), None, + None, ) grouped = next(row for row in rows if row["router_name"] == router) assert grouped["tier_turns"] == {"simple": 2, "complex": 1} @@ -446,6 +609,7 @@ async def test_tier_maps_stay_separate_per_router_type_on_a_reconfigured_alias(d (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), None, + None, ) by_type = {row["router_type"]: row["tier_turns"] for row in rows if row["router_name"] == router} assert by_type == {"complexity": {"medium": 1}, "quality": {"2": 1}} @@ -461,6 +625,7 @@ async def test_a_window_with_no_tiered_turns_aggregates_to_an_empty_map(db): (T0 - timedelta(days=1)).isoformat(), (T0 + timedelta(days=1)).isoformat(), None, + None, ) grouped = next(row for row in rows if row["router_name"] == router) assert grouped["tier_turns"] == {} diff --git a/tests/proxy_behavior/spend/test_baseline_accounting.py b/tests/proxy_behavior/spend/test_baseline_accounting.py index e187a44c29d..3504751d132 100644 --- a/tests/proxy_behavior/spend/test_baseline_accounting.py +++ b/tests/proxy_behavior/spend/test_baseline_accounting.py @@ -56,7 +56,9 @@ def record() -> Callable[..., BaselineAccountingRecord]: }, ) - def create(label: str = "first", started: float = 10000.0, identical: bool = True) -> BaselineAccountingRecord: + def create( + label: str = "first", started: float = 10000.0, identical: bool = True, user_id: str = "" + ) -> BaselineAccountingRecord: return BaselineAccountingRecord( scope="autorouter-baseline:v3:" + run * 2, api_key=run, session_id=run, router_name="test-router", baseline_model="anthropic/claude-opus-5", @@ -76,6 +78,7 @@ def record() -> Callable[..., BaselineAccountingRecord]: total_tokens=6230, spend=0.17, saved_spend=0.0, classifier_cost=0.0, covered=True, cache_hit=False, cache_ttl_seconds=3600, cache_touched=True, baseline_model="anthropic/claude-opus-5", + user_id=user_id, ), daily=DailyBaselineAttribution( date="2026-09-15", api_key=run, model="claude-opus-5", custom_llm_provider="anthropic", @@ -99,21 +102,36 @@ async def _session(db: Prisma, record: BaselineAccountingRecord): return rows[0] +async def _user_sessions(db: Prisma, record: BaselineAccountingRecord) -> dict[str, dict[str, object]]: + rows: Final = await db.query_raw('SELECT * FROM "LiteLLM_AutoRouterUserSession" WHERE api_key=$1', record.api_key) + return {str(row["user_id"]): row for row in rows} + + async def test_late_replay_updates_all_projections_without_rebilling(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: store: Final = _store(db) - late: Final = record("late", 10001.0) - early: Final = record("early", identical=False) + late: Final = record("late", 10001.0, user_id="late-user") + early: Final = record("early", identical=False, user_id="early-user") await _log(db, late) assert await store.append(late) == "recorded" assert await store.project(late.scope) == "published" before: Final = await _session(db, late) assert before["savings_estimated_actual_spend"] == before["spend"] == 0.17 assert before["saved_spend"] == 0.0 + before_users: Final = await _user_sessions(db, late) + assert set(before_users) == {"late-user"} + assert before_users["late-user"]["savings_estimated_turns"] == 1 + assert before_users["late-user"]["savings_estimated_baseline_models"] == {late.baseline_model: 1} await _log(db, early) assert await store.append(early) == "recorded" pending: Final = await _session(db, late) assert pending["spend"] == 0.34 and pending["savings_estimated_turns"] == 0 assert pending["saved_spend"] == pending["savings_estimated_actual_spend"] == 0.0 + pending_users: Final = await _user_sessions(db, late) + assert set(pending_users) == {"late-user", "early-user"} + for user in pending_users.values(): + assert user["turns"] == 1 and user["spend"] == 0.17 + assert user["savings_estimated_turns"] == user["savings_estimated_actual_spend"] == user["saved_spend"] == 0 + assert user["savings_estimated_baseline_models"] == {} waiting: Final = await db.query_raw('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=$1', late.observation.request_id) assert waiting[0]["metadata"]["autorouter_savings"] is None assert waiting[0]["metadata"]["autorouter_savings_estimate"]["reason"] == "pending_projection" @@ -125,37 +143,69 @@ async def test_late_replay_updates_all_projections_without_rebilling(db: Prisma, assert logs[0]["spend"] == 0.17 assert logs[0]["metadata"]["autorouter_savings_estimate"]["provenance"] == "modeled" assert after["saved_spend"] == pytest.approx(logs[0]["metadata"]["autorouter_savings"]) + after_users: Final = await _user_sessions(db, late) + assert after_users["early-user"] == pending_users["early-user"] + for field in ( + "saved_spend", "savings_estimated_turns", "savings_estimated_actual_spend", + "savings_estimated_saved_spend", "savings_estimated_baseline_models", + ): + assert after_users["late-user"][field] == after[field] + assert after_users["late-user"]["turns"] == 1 and after_users["late-user"]["spend"] == 0.17 for table in ("DailyUserSpend", "DailyTeamSpend", "DailyOrganizationSpend", "DailyEndUserSpend", "DailyAgentSpend", "DailyTagSpend"): rows: Final = await db.query_raw(f'SELECT spend,api_requests,autorouter_savings_spend FROM "LiteLLM_{table}" WHERE api_key=$1', late.api_key) assert rows[0]["spend"] == rows[0]["api_requests"] == 0 assert rows[0]["autorouter_savings_spend"] == pytest.approx(after["saved_spend"]) -async def test_commit_ack_loss_and_concurrent_duplicate_delivery_are_idempotent(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: - event: Final = record() +@pytest.mark.parametrize("attributed", [True, False]) +async def test_commit_ack_loss_and_concurrent_duplicate_delivery_are_idempotent( + db: Prisma, record: Callable[..., BaselineAccountingRecord], attributed: bool +) -> None: + event: Final = record(user_id="first-user" if attributed else "") + other: Final = record("other", 10001.0, user_id="second-user" if attributed else "") await _log(db, event) assert await _store(db, after_commit=True).append(event) == "unavailable" store: Final = _store(db) assert set(await asyncio.gather(*(store.append(event) for _ in range(4)))) == {"recorded"} + await _log(db, other) + assert await store.append(other) == "recorded" + if not attributed: + await db.execute_raw( + 'UPDATE "LiteLLM_AutoRouterBaselineObservation" SET data=(data::jsonb #- \'{turn,user_id}\')::text WHERE scope=$1', + event.scope, + ) assert await store.project(event.scope) == "published" assert await store.project(event.scope) == "unchanged" session: Final = await _session(db, event) - assert session["turns"] == session["savings_estimated_turns"] == 1 - assert session["spend"] == session["savings_estimated_actual_spend"] == 0.17 + assert session["turns"] == session["savings_estimated_turns"] == 2 + assert session["spend"] == session["savings_estimated_actual_spend"] == 0.34 + users: Final = await _user_sessions(db, event) + assert set(users) == ({"first-user", "second-user"} if attributed else set()) + for user in users.values(): + assert user["turns"] == user["savings_estimated_turns"] == 1 + assert user["spend"] == user["savings_estimated_actual_spend"] == 0.17 + assert user["savings_estimated_baseline_models"] == {event.baseline_model: 1} async def test_publication_rollback_keeps_dirty_revision_for_retry(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: - event: Final = record() + event: Final = record(user_id="rollback-user") await _log(db, event) store: Final = _store(db) assert await store.append(event) == "recorded" assert await _store(db, before_commit=True).project(event.scope) == "unavailable" session: Final = await _session(db, event) assert session["spend"] == 0.17 and session["savings_estimated_turns"] == 0 + before_users: Final = await _user_sessions(db, event) + assert before_users["rollback-user"]["spend"] == 0.17 + assert before_users["rollback-user"]["savings_estimated_turns"] == 0 + assert before_users["rollback-user"]["savings_estimated_baseline_models"] == {} revisions: Final = await db.query_raw('SELECT revision,published_revision FROM "LiteLLM_AutoRouterBaselineComparison" WHERE scope=$1', event.scope) assert revisions[0]["revision"] > revisions[0]["published_revision"] assert await store.project(event.scope) == "published" assert (await _session(db, event))["savings_estimated_turns"] == 1 + after_users: Final = await _user_sessions(db, event) + assert after_users["rollback-user"]["turns"] == after_users["rollback-user"]["savings_estimated_turns"] == 1 + assert after_users["rollback-user"]["spend"] == after_users["rollback-user"]["savings_estimated_actual_spend"] == 0.17 async def test_conflicting_duplicate_cannot_restore_an_observed_estimate(db: Prisma, record: Callable[..., BaselineAccountingRecord]) -> None: @@ -196,6 +246,7 @@ async def test_native_observation_enters_spend_pipeline_once_with_shared_daily_a db: Prisma, record: Callable[..., BaselineAccountingRecord], monkeypatch: pytest.MonkeyPatch, ) -> None: import os + from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.hooks.autorouter_baseline_cache import CapturedBaselineObservation 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/llms/azure/test_azure.py b/tests/test_litellm/llms/azure/test_azure.py index 6b6832f623c..86065c7adc6 100644 --- a/tests/test_litellm/llms/azure/test_azure.py +++ b/tests/test_litellm/llms/azure/test_azure.py @@ -1,10 +1,13 @@ """Tests for litellm/llms/azure/azure.py AzureChatCompletion handler behaviour.""" +import asyncio import time from typing import Final -from openai import AzureOpenAI +import pytest +from openai import AsyncAzureOpenAI, AzureOpenAI +import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.azure import AzureChatCompletion @@ -52,3 +55,25 @@ def test_sync_streaming_stamps_response_headers_on_the_logging_obj() -> None: ) assert logging_obj.model_call_details["response_headers"] == {"x-ms-is-spilled-over": "true"} + + +class _CancelledRawCompletions: + async def create(self, **kwargs): + raise asyncio.CancelledError() + + +@pytest.mark.asyncio +async def test_acompletion_propagates_cancelled_error() -> None: + client = AsyncAzureOpenAI( + api_key="fake-key", + api_version="2024-02-01", + azure_endpoint="https://fake-resource.openai.azure.com", + ) + client.chat.completions.with_raw_response = _CancelledRawCompletions() + + with pytest.raises(asyncio.CancelledError): + await litellm.acompletion( + model="azure/fake-deployment", + messages=[{"role": "user", "content": "hi"}], + client=client, + ) diff --git a/tests/test_litellm/llms/azure/test_azure_common_utils.py b/tests/test_litellm/llms/azure/test_azure_common_utils.py index 83ec85f1176..caf941ebd19 100644 --- a/tests/test_litellm/llms/azure/test_azure_common_utils.py +++ b/tests/test_litellm/llms/azure/test_azure_common_utils.py @@ -597,7 +597,8 @@ async def test_ensure_initialize_azure_sdk_client_always_used(call_type): "litellm.files.main.azure_files_instance.initialize_azure_sdk_client" ) elif ( - call_type == CallTypes.avideo_content + call_type == CallTypes.avideo_generation + or call_type == CallTypes.avideo_content or call_type == CallTypes.avideo_list or call_type == CallTypes.avideo_remix or call_type == CallTypes.avideo_create_character diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 087c5a03498..5c885a3168d 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -12,6 +12,7 @@ from starlette.datastructures import Headers from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( MCPRequestHandler, UnloadableEntitlementError, + _agent_capped_servers, _is_mcp_admitted_user_subject, ) from litellm.proxy._types import ( @@ -21,6 +22,7 @@ from litellm.proxy._types import ( SpecialMCPServerNames, UserAPIKeyAuth, ) +from litellm.types.agents import AgentCaller @pytest.mark.asyncio @@ -4169,10 +4171,114 @@ async def test_get_allowed_mcp_servers_for_key_prefers_in_memory_permission(): global_mcp_server_manager.registry.pop("direct-server", None) +@pytest.mark.parametrize( + ("agent_servers", "group_ceiling", "expected"), + [ + ([], frozenset({"server_1"}), ("server_1",)), + ([], frozenset({"server_1", "server_2", "server_3"}), ("server_1", "server_2")), + ([], frozenset(), ()), + (["server_2"], frozenset({"server_1", "server_2"}), ("server_2",)), + (["server_1"], frozenset({"server_2"}), ()), + (["server_1"], None, ("server_1",)), + ], +) +def test_agent_capped_servers_intersects_agent_config_and_access_groups(agent_servers, group_ceiling, expected): + """The agent's attached access groups cap the key/team servers alongside its own + object_permission; groups naming no server deny all.""" + assert _agent_capped_servers(["server_1", "server_2"], agent_servers, group_ceiling) == expected + + +def test_agent_capped_servers_without_agent_restrictions_is_uncapped(): + assert _agent_capped_servers(["server_1", "server_2"], [], None) is None + + @pytest.mark.asyncio class TestAgentMCPPermissions: """Test agent-level MCP server and tool permission intersection.""" + @staticmethod + def _agent_key_acting_for(user_id: str, team_id: str | None) -> UserAPIKeyAuth: + agent_key = UserAPIKeyAuth(api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1") + agent_key.agent_caller = AgentCaller(user_id=user_id, team_id=team_id) + return agent_key + + @staticmethod + def _team_servers(grants: dict[str, list[str]]) -> AsyncMock: + async def by_team(user_api_key_auth: UserAPIKeyAuth | None = None) -> list[str]: + assert user_api_key_auth is not None + return grants.get(user_api_key_auth.team_id or "", []) + + return AsyncMock(side_effect=by_team) + + @staticmethod + def _user_servers(grants: dict[str, list[str] | None]) -> AsyncMock: + async def by_user(user_api_key_auth: UserAPIKeyAuth | None = None) -> list[str] | None: + assert user_api_key_auth is not None + return grants.get(user_api_key_auth.user_id or "", []) + + return AsyncMock(side_effect=by_user) + + async def test_agent_key_acting_for_a_user_is_capped_at_the_invoking_teams_servers(self): + """LIT-8014: the agent's own key reaches server_1 and server_2, but the human who invoked it + belongs to a team granted only server_2, so on their behalf the agent reaches only server_2.""" + agent_key = self._agent_key_acting_for(user_id="alice", team_id="callers") + + with ( + patch.object( # test-quality-ok: the level resolvers read proxy_server globals with no injection seam + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server_1", "server_2"]) + ), + patch.object( # test-quality-ok: same seam, keyed by which team is being asked about + MCPRequestHandler, + "_get_allowed_mcp_servers_for_team", + self._team_servers({"callers": ["server_2", "server_3"]}), + ), + patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here + MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) + ), + patch.object( # test-quality-ok: neither the agent's owner nor the caller has a personal grant + MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({}) + ), + ): + assert await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=agent_key) == ["server_2"] + + async def test_agent_key_acting_for_a_teamless_user_is_capped_at_that_users_servers(self): + agent_key = self._agent_key_acting_for(user_id="alice", team_id=None) + + with ( + patch.object( # test-quality-ok: the level resolvers read proxy_server globals with no injection seam + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server_1", "server_2"]) + ), + patch.object( # test-quality-ok: same seam + MCPRequestHandler, "_get_allowed_mcp_servers_for_team", self._team_servers({}) + ), + patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here + MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) + ), + patch.object( # test-quality-ok: same seam, keyed by which user is being asked about + MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({"alice": ["server_1"]}) + ), + ): + assert await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=agent_key) == ["server_1"] + + async def test_agent_key_acting_for_a_caller_whose_entitlement_is_unreadable_reaches_nothing(self): + agent_key = self._agent_key_acting_for(user_id="alice", team_id=None) + + with ( + patch.object( # test-quality-ok: the level resolvers read proxy_server globals with no injection seam + MCPRequestHandler, "_get_allowed_mcp_servers_for_key", AsyncMock(return_value=["server_1"]) + ), + patch.object( # test-quality-ok: same seam + MCPRequestHandler, "_get_allowed_mcp_servers_for_team", self._team_servers({}) + ), + patch.object( # test-quality-ok: agent object_permission lookup hits the DB, not under test here + MCPRequestHandler, "_get_allowed_mcp_servers_for_agent", AsyncMock(return_value=[]) + ), + patch.object( # test-quality-ok: None is the resolver's own "entitlement unresolvable" signal + MCPRequestHandler, "_get_allowed_mcp_servers_for_user", self._user_servers({"alice": None}) + ), + ): + assert await MCPRequestHandler.get_allowed_mcp_servers(user_api_key_auth=agent_key) == [] + async def test_get_allowed_mcp_servers_agent_intersection(self): """Key/team allow [server_1, server_2]; agent allows [server_1]. Result = [server_1].""" user_api_key_auth = UserAPIKeyAuth( @@ -4208,6 +4314,46 @@ class TestAgentMCPPermissions: assert sorted(result) == ["server_1", "server_2"] mock_agent.assert_called_once_with(user_api_key_auth) + async def test_agent_access_group_server_ceiling_expands_group_servers(self): + from litellm.proxy._experimental.mcp_server.mcp_server_manager import global_mcp_server_manager + from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling + from litellm.types.mcp import MCPTransport + from litellm.types.mcp_server.mcp_server_manager import MCPServer + + asked: list[str] = [] + + async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None: + asked.append(agent_id) + return AgentAccessGroupCeiling( + access_group_ids=("ag-1",), + models=frozenset(), + mcp_server_ids=frozenset({"aliased-server"}), + agent_ids=frozenset(), + ) + + global_mcp_server_manager.registry["ag-server-id"] = MCPServer( + server_id="ag-server-id", + name="ag-server", + server_name="ag-server", + alias="aliased-server", + url="https://ag-server.example.com", + transport=MCPTransport.http, + ) + try: + result = await MCPRequestHandler._get_agent_access_group_server_ceiling( + UserAPIKeyAuth(api_key="test-key", agent_id="agent-ag"), resolve + ) + finally: + global_mcp_server_manager.registry.pop("ag-server-id", None) + + assert result == frozenset({"ag-server-id"}) + assert asked == ["agent-ag"] + assert ( + await MCPRequestHandler._get_agent_access_group_server_ceiling(UserAPIKeyAuth(api_key="k"), resolve) + is None + ) + assert asked == ["agent-ag"] + async def test_get_allowed_mcp_servers_key_team_agent_intersection(self): """Key allows [1, 2], agent allows [2, 3]. Result = [2].""" user_api_key_auth = UserAPIKeyAuth( diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py new file mode 100644 index 00000000000..e744e84d671 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py @@ -0,0 +1,146 @@ +from typing import Final + +import pytest +from fastapi import HTTPException + +from litellm.models.access_group import LiteLLM_AccessGroupTable +from litellm.proxy.agent_endpoints.auth.agent_access_groups import ( + AgentAccessGroupCeiling, + resolve_agent_access_group_ceiling, +) +from litellm.types.agents import AgentResponse + +_CARD: Final = {"name": "agent", "url": "http://localhost:9999", "version": "1.0.0"} + + +def _agent(access_group_ids: list[str] | None) -> AgentResponse: + return AgentResponse( + agent_id="agent-1", agent_name="agent", agent_card_params=_CARD, access_group_ids=access_group_ids + ) + + +def _group( + group_id: str, + models: tuple[str, ...] = (), + mcp_servers: tuple[str, ...] = (), + agents: tuple[str, ...] = (), +) -> LiteLLM_AccessGroupTable: + return LiteLLM_AccessGroupTable( + access_group_id=group_id, + access_group_name=group_id, + access_model_names=list(models), + access_mcp_server_ids=list(mcp_servers), + access_agent_ids=list(agents), + ) + + +def _loaders(agent: AgentResponse | None, groups: dict[str, LiteLLM_AccessGroupTable]): + async def load_agent(agent_id: str) -> tuple[str, ...]: + return tuple(agent.access_group_ids or ()) if agent is not None else () + + async def load_group(group_id: str) -> LiteLLM_AccessGroupTable | None: + return groups.get(group_id) + + return load_agent, load_group + + +@pytest.mark.asyncio +@pytest.mark.parametrize("access_group_ids", [None, []]) +async def test_agent_without_access_groups_has_no_ceiling(access_group_ids: list[str] | None): + load_agent, load_group = _loaders(_agent(access_group_ids), {"g1": _group("g1", models=("gpt-5",))}) + + assert await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) is None + + +@pytest.mark.asyncio +async def test_unknown_agent_has_no_ceiling(): + load_agent, load_group = _loaders(None, {}) + + assert await resolve_agent_access_group_ceiling("missing", load_agent, load_group) is None + + +@pytest.mark.asyncio +async def test_ceiling_is_the_union_of_every_attached_group(): + load_agent, load_group = _loaders( + _agent(["g1", "g2"]), + { + "g1": _group("g1", models=("gpt-5",), mcp_servers=("mcp-a",), agents=("agent-b",)), + "g2": _group("g2", models=("claude-sonnet",), mcp_servers=("mcp-b",), agents=("agent-c",)), + }, + ) + + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) + + assert ceiling == AgentAccessGroupCeiling( + access_group_ids=("g1", "g2"), + models=frozenset({"gpt-5", "claude-sonnet"}), + mcp_server_ids=frozenset({"mcp-a", "mcp-b"}), + agent_ids=frozenset({"agent-b", "agent-c"}), + ) + + +@pytest.mark.asyncio +async def test_unloadable_group_contributes_nothing_but_the_ceiling_still_applies(): + load_agent, load_group = _loaders(_agent(["g1", "gone"]), {"g1": _group("g1", models=("gpt-5",))}) + + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) + + assert ceiling == AgentAccessGroupCeiling( + access_group_ids=("g1", "gone"), + models=frozenset({"gpt-5"}), + mcp_server_ids=frozenset(), + agent_ids=frozenset(), + ) + + +@pytest.mark.asyncio +async def test_only_unloadable_groups_is_an_empty_ceiling_not_unrestricted(): + load_agent, load_group = _loaders(_agent(["gone"]), {}) + + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_agent, load_group) + + assert ceiling is not None + assert ceiling.models == frozenset() + assert ceiling.mcp_server_ids == frozenset() + assert ceiling.agent_ids == frozenset() + + +@pytest.mark.asyncio +async def test_default_agent_loader_reads_the_attached_groups_from_the_registry(): + from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry + + _, load_group = _loaders(None, {"g1": _group("g1", models=("gpt-5",))}) + global_agent_registry.register_agent(_agent(["g1"])) + try: + ceiling: Final = await resolve_agent_access_group_ceiling("agent-1", load_access_group=load_group) + finally: + global_agent_registry.deregister_agent("agent") + + assert ceiling == AgentAccessGroupCeiling( + access_group_ids=("g1",), models=frozenset({"gpt-5"}), mcp_server_ids=frozenset(), agent_ids=frozenset() + ) + + +@pytest.mark.asyncio +async def test_default_loader_treats_a_missing_group_as_unreadable(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints.auth.agent_access_groups import _load_access_group + from litellm.proxy.auth import auth_checks + + async def missing_group(**_: object) -> LiteLLM_AccessGroupTable: + raise HTTPException(status_code=404, detail={"error": "Access group doesn't exist in db."}) + + monkeypatch.setattr(proxy_server, "prisma_client", object()) + monkeypatch.setattr(auth_checks, "get_access_object", missing_group) + + assert await _load_access_group("gone") is None + + +@pytest.mark.asyncio +async def test_default_loader_returns_nothing_without_a_db(monkeypatch: pytest.MonkeyPatch): + from litellm.proxy import proxy_server + from litellm.proxy.agent_endpoints.auth.agent_access_groups import _load_access_group + + monkeypatch.setattr(proxy_server, "prisma_client", None) + + assert await _load_access_group("ag-1") is None diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py new file mode 100644 index 00000000000..b08964503c8 --- /dev/null +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py @@ -0,0 +1,57 @@ +from typing import Final + +import pytest + +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.agent_endpoints.auth.agent_caller import agent_caller_auth, agent_caller_from_headers +from litellm.types.agents import AgentCaller + +_AGENT_KEY: Final = UserAPIKeyAuth(api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1") + + +def test_agent_key_echoing_both_ids_acts_for_that_user_and_team() -> None: + headers: Final = {"X-LiteLLM-User-Id": " alice ", "x-litellm-team-id": "callers"} + + assert agent_caller_from_headers(headers, _AGENT_KEY) == AgentCaller(user_id="alice", team_id="callers") + + +def test_agent_key_echoing_only_a_user_id_acts_for_a_teamless_user() -> None: + assert agent_caller_from_headers({"x-litellm-user-id": "alice"}, _AGENT_KEY) == AgentCaller(user_id="alice") + + +@pytest.mark.parametrize("headers", [{}, {"x-litellm-user-id": " ", "x-litellm-team-id": ""}]) +def test_agent_key_echoing_no_caller_acts_for_itself(headers: dict[str, str]) -> None: + assert agent_caller_from_headers(headers, _AGENT_KEY) is None + + +def test_caller_headers_on_a_key_without_an_agent_are_ignored() -> None: + plain_key: Final = UserAPIKeyAuth(api_key="plain-key", user_id="bob") + + assert agent_caller_from_headers({"x-litellm-user-id": "alice", "x-litellm-team-id": "callers"}, plain_key) is None + + +def test_caller_auth_stands_for_the_invoking_user_not_the_agent() -> None: + agent_key: Final = UserAPIKeyAuth( + api_key="agent-key", user_id="agent-owner", team_id="agent-team", agent_id="agent-1" + ) + agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers") + + caller_auth: Final = agent_caller_auth(agent_key) + + assert caller_auth is not None + assert (caller_auth.user_id, caller_auth.team_id, caller_auth.agent_id, caller_auth.api_key) == ( + "alice", + "callers", + None, + None, + ) + assert agent_caller_auth(_AGENT_KEY) is None + + +def test_agent_caller_cannot_be_set_from_a_request_payload() -> None: + forged: Final = UserAPIKeyAuth.model_validate( + {"api_key": "agent-key", "agent_id": "agent-1", "agent_caller": {"user_id": "alice", "team_id": "callers"}} + ) + + assert forged.agent_caller is None + assert "agent_caller" not in forged.model_dump() diff --git a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py index 383b72e5c58..a87716375e8 100644 --- a/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py +++ b/tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py @@ -9,10 +9,10 @@ from unittest.mock import AsyncMock, patch import pytest - from litellm.constants import UI_SESSION_TOKEN_TEAM_ID -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy._types import LiteLLM_ObjectPermissionTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry +from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( AgentAccess, AgentRequestHandler, @@ -20,6 +20,7 @@ from litellm.proxy.agent_endpoints.auth.agent_permission_handler import ( UnrestrictedAgentAccess, accessible_agents, ) +from litellm.types.agents import AgentCaller def _registry_with(*agent_names: str) -> AgentRegistry: @@ -157,6 +158,130 @@ class TestAgentRequestHandler: is False ), agent_id + @staticmethod + def _ceiling_resolver(agent_ids: frozenset[str] | None) -> tuple[CeilingResolver, list[str]]: + """A resolver that records the agent ids it was asked about and answers with a fixed + ceiling, or None when the agent has no access groups attached.""" + asked: Final[list[str]] = [] + + async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None: + asked.append(agent_id) + if agent_ids is None: + return None + return AgentAccessGroupCeiling( + access_group_ids=("ag-1",), models=frozenset(), mcp_server_ids=frozenset(), agent_ids=agent_ids + ) + + return resolve, asked + + @staticmethod + def _key_granting(agent_ids: list[str], agent_id: str | None) -> UserAPIKeyAuth: + return UserAPIKeyAuth( + api_key="test-key", + user_id="test-user", + agent_id=agent_id, + object_permission=LiteLLM_ObjectPermissionTable(object_permission_id="obj-1", agents=agent_ids), + ) + + async def test_agent_access_groups_cap_an_otherwise_unrestricted_key(self): + """A key with no agent grant of its own may still only reach the agents its + agent's attached access groups name.""" + agent_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="caller-agent") + resolve, asked = self._ceiling_resolver(frozenset({"agent-beta"})) + + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset({"agent-beta"}) + ) + assert await AgentRequestHandler.is_agent_allowed("agent-beta", agent_key, resolve) is True + assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key, resolve) is False + assert asked == ["caller-agent"] * 3 + + @staticmethod + def _team_grants(grants: dict[str, AgentAccess]) -> AsyncMock: + async def by_team(user_api_key_auth: UserAPIKeyAuth | None = None) -> AgentAccess: + assert user_api_key_auth is not None + return grants.get(user_api_key_auth.team_id or "", UnrestrictedAgentAccess()) + + return AsyncMock(side_effect=by_team) + + async def test_agent_key_acting_for_a_user_is_capped_at_the_invoking_teams_agents(self): + """LIT-8014: the agent's key and access groups reach alpha and beta, but the human who + invoked it belongs to a team granted only beta, so on their behalf the agent reaches only beta.""" + agent_key: Final = self._key_granting(["agent-alpha", "agent-beta"], agent_id="caller-agent") + agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers") + resolve, _ = self._ceiling_resolver(frozenset({"agent-alpha", "agent-beta", "agent-gamma"})) + + with patch.object( # test-quality-ok: the team resolver reads proxy_server globals with no injection seam + AgentRequestHandler, + "_get_allowed_agents_for_team", + self._team_grants({"callers": RestrictedAgentAccess(frozenset({"agent-beta", "agent-gamma"}))}), + ) as mock_team: + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset({"agent-beta"}) + ) + assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key, resolve) is False + + assert {call.args[0].team_id for call in mock_team.call_args_list} == {None, "callers"} + + async def test_agent_key_acting_for_a_user_whose_team_grants_no_agent_reaches_none(self): + agent_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="caller-agent") + agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers") + resolve, _ = self._ceiling_resolver(None) + + with patch.object( # test-quality-ok: the team resolver reads proxy_server globals with no injection seam + AgentRequestHandler, + "_get_allowed_agents_for_team", + self._team_grants({"callers": RestrictedAgentAccess(frozenset())}), + ): + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset() + ) + + async def test_agent_key_acting_for_an_ungranted_caller_keeps_its_own_agents(self): + agent_key: Final = self._key_granting(["agent-alpha"], agent_id="caller-agent") + agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers") + resolve, _ = self._ceiling_resolver(None) + + with patch.object( # test-quality-ok: the team resolver reads proxy_server globals with no injection seam + AgentRequestHandler, "_get_allowed_agents_for_team", self._team_grants({}) + ): + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset({"agent-alpha"}) + ) + + + async def test_agent_access_groups_intersect_with_key_grants(self): + agent_key: Final = self._key_granting(["agent-alpha", "agent-beta"], agent_id="caller-agent") + resolve, _ = self._ceiling_resolver(frozenset({"agent-beta", "agent-gamma"})) + + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset({"agent-beta"}) + ) + assert await AgentRequestHandler.is_agent_allowed("agent-gamma", agent_key, resolve) is False + + async def test_agent_access_groups_naming_no_agent_deny_every_agent(self): + agent_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user", agent_id="caller-agent") + resolve, _ = self._ceiling_resolver(frozenset()) + + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess(frozenset()) + assert await AgentRequestHandler.is_agent_allowed("agent-alpha", agent_key, resolve) is False + + async def test_agent_without_access_groups_keeps_key_grants(self): + agent_key: Final = self._key_granting(["agent-alpha"], agent_id="caller-agent") + resolve, asked = self._ceiling_resolver(None) + + assert await AgentRequestHandler.resolve_agent_access(agent_key, resolve) == RestrictedAgentAccess( + frozenset({"agent-alpha"}) + ) + assert asked == ["caller-agent"] + + async def test_key_without_agent_never_consults_agent_access_groups(self): + plain_key: Final = UserAPIKeyAuth(api_key="test-key", user_id="test-user") + resolve, asked = self._ceiling_resolver(frozenset()) + + assert await AgentRequestHandler.resolve_agent_access(plain_key, resolve) == UnrestrictedAgentAccess() + assert asked == [] + async def test_empty_access_group_denies_every_agent(self): """LIT-5143: a key restricted to an access group that resolves to no agents is restricted to nothing, not unrestricted. A failed group lookup still fails open.""" diff --git a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py index 441e9640ef9..b9a260f5b14 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py @@ -16,6 +16,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.agents import AgentCaller AddLiteLLMData = Callable[..., Awaitable[dict[str, object]]] @@ -511,6 +512,24 @@ async def test_message_methods_forward_caller_identity_headers(method: str): assert forwarded_headers.get("X-LiteLLM-Team-Id") == "team-xyz" +@pytest.mark.asyncio +@pytest.mark.parametrize("method", ["message/send", "message/stream"]) +async def test_agent_calling_another_agent_forwards_the_human_who_invoked_it(method: str): + """LIT-8014: an agent acting for alice calls a second agent through the proxy. That hop must + carry alice, not the first agent's owner, so the chain stays capped at what alice may reach.""" + mock_request = _make_request_mock(method, _HELLO_MESSAGE_PARAMS) + agent_key = UserAPIKeyAuth(api_key="sk-agent", user_id="agent-owner", team_id="agent-team", agent_id="agent-1") + agent_key.agent_caller = AgentCaller(user_id="alice", team_id="callers") + + captured = await _invoke_message_method(method, mock_request, agent_key) + + forwarded_headers = captured.agent_extra_headers or {} + assert (forwarded_headers.get("X-LiteLLM-User-Id"), forwarded_headers.get("X-LiteLLM-Team-Id")) == ( + "alice", + "callers", + ) + + @pytest.mark.asyncio @pytest.mark.parametrize("method", ["message/send", "message/stream"]) async def test_message_methods_send_the_entra_bearer_for_azure_agents(method: str): diff --git a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py index 231626c7eb5..b036e0dac4d 100644 --- a/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py +++ b/tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py @@ -15,6 +15,7 @@ from litellm.proxy.agent_endpoints.agent_registry import ( _restore_redacted_litellm_params, redact_sensitive_agent_litellm_params, ) +from litellm.types.agents import PatchAgentRequest # Obviously-fake stand-ins for a real AWS credential pair (LIT-6736 regression # fixtures) -- never a real key shape, and must never appear in any response. @@ -990,3 +991,138 @@ async def test_patch_agent_in_db_preserves_secret_when_echoed_back_redacted(): stored_params: Final = json.loads(mock_update.call_args.kwargs["data"]["litellm_params"]) assert stored_params["aws_secret_access_key"] == SENTINEL_AWS_SECRET_ACCESS_KEY assert stored_params["is_public"] is True + + +def _agent_row_mock(access_group_ids: list[str]) -> MagicMock: + row: Final = MagicMock() + row.model_dump.return_value = { + "agent_id": "agent-123", + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {}, + "object_permission": None, + "access_group_ids": access_group_ids, + } + row.object_permission = None + return row + + +@pytest.mark.asyncio +async def test_add_agent_to_db_persists_deduplicated_access_group_ids(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_create = AsyncMock(return_value=_agent_row_mock(["ag-1", "ag-2"])) + mock_prisma.db.litellm_agentstable.create = mock_create + + result: Final = await registry.add_agent_to_db( + agent={ + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "access_group_ids": ["ag-1", "ag-2", "ag-1"], + }, + prisma_client=mock_prisma, + created_by="test-user", + ) + + assert tuple(mock_create.call_args.kwargs["data"]["access_group_ids"]) == ("ag-1", "ag-2") + assert result.access_group_ids == ["ag-1", "ag-2"] + + +@pytest.mark.asyncio +async def test_add_agent_to_db_without_access_group_ids_leaves_column_to_its_default(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_create = AsyncMock(return_value=_agent_row_mock([])) + mock_prisma.db.litellm_agentstable.create = mock_create + + await registry.add_agent_to_db( + agent={"agent_name": "Test Agent", "agent_card_params": _sample_agent_card_params()}, + prisma_client=mock_prisma, + created_by="test-user", + ) + + assert "access_group_ids" not in mock_create.call_args.kwargs["data"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("patch_body", "expected"), + [ + ({"access_group_ids": ["ag-2", "ag-3"]}, ["ag-2", "ag-3"]), + ({"access_group_ids": []}, []), + ({"access_group_ids": None}, []), + ], +) +async def test_patch_agent_in_db_replaces_access_group_ids_when_provided( + patch_body: PatchAgentRequest, expected: list[str] +): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Test Agent", + "litellm_params": {}, + "object_permission_id": None, + "access_group_ids": ["ag-1"], + } + ) + mock_update = AsyncMock(return_value=_agent_row_mock(expected)) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", agent=patch_body, prisma_client=mock_prisma, updated_by="test-user" + ) + + assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected) + + +@pytest.mark.asyncio +async def test_patch_agent_in_db_keeps_access_group_ids_when_omitted(): + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value={ + "agent_id": "agent-123", + "agent_name": "Old Name", + "litellm_params": {}, + "object_permission_id": None, + "access_group_ids": ["ag-1"], + } + ) + mock_update = AsyncMock(return_value=_agent_row_mock(["ag-1"])) + mock_prisma.db.litellm_agentstable.update = mock_update + + await registry.patch_agent_in_db( + agent_id="agent-123", agent={"agent_name": "New Name"}, prisma_client=mock_prisma, updated_by="test-user" + ) + + assert "access_group_ids" not in mock_update.call_args.kwargs["data"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("body_access_group_ids", "expected"), + [(["ag-9", "ag-9"], ["ag-9"]), (None, []), ("omitted", [])], +) +async def test_update_agent_in_db_always_writes_access_group_ids(body_access_group_ids, expected: list[str]): + """PUT is a full replacement: omitting the field clears any previously attached groups.""" + registry: Final = AgentRegistry() + mock_prisma: Final = MagicMock() + mock_prisma.db.litellm_agentstable.find_unique = AsyncMock( + return_value=SimpleNamespace(litellm_params={}, object_permission_id=None, access_group_ids=["ag-1"]) + ) + mock_update = AsyncMock(return_value=_agent_row_mock(expected)) + mock_prisma.db.litellm_agentstable.update = mock_update + body: Final = { + "agent_name": "Test Agent", + "agent_card_params": _sample_agent_card_params(), + "litellm_params": {"model": "bedrock/agentcore/my-agent"}, + **({} if body_access_group_ids == "omitted" else {"access_group_ids": body_access_group_ids}), + } + + await registry.update_agent_in_db( + agent_id="agent-123", agent=body, prisma_client=mock_prisma, updated_by="test-user" + ) + + assert tuple(mock_update.call_args.kwargs["data"]["access_group_ids"]) == tuple(expected) diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 44c1d6a3c6b..0cab6535f3d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -34,20 +34,26 @@ from litellm.proxy._types import ( UserAPIKeyAuth, WebhookEvent, ) +from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling, CeilingResolver +from litellm.types.agents import AgentCaller from litellm.proxy.auth.auth_checks import ( ExperimentalUIJWTToken, _cache_management_object, _can_object_call_model, _can_object_call_vector_stores, + _check_agent_access_group_model_access, _check_end_user_budget, _check_team_member_budget, _fetch_key_object_from_db_with_reconnect, _get_fuzzy_user_object, + CallerTeamLoader, + CallerUserLoader, _get_team_db_check, _log_budget_lookup_failure, _tag_max_budget_check, _team_max_budget_check, _virtual_key_max_budget_alert_check, + _check_agent_caller_model_access, _virtual_key_max_budget_check, _virtual_key_soft_budget_check, get_key_object, @@ -8930,6 +8936,69 @@ def test_request_skips_budget_checks_extends_route_rule_with_zero_cost_models() assert request_skips_budget_checks(route="/v1/chat/completions", model=None, llm_router=None) is False +def _agent_model_ceiling_resolver( + models: frozenset[str] | None, +) -> tuple[CeilingResolver, list[str]]: + """Resolver that records the agent ids it was asked about and answers with a fixed model + ceiling, or None when the agent has no access groups attached.""" + asked: Final[list[str]] = [] + + async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None: + asked.append(agent_id) + if models is None: + return None + return AgentAccessGroupCeiling( + access_group_ids=("ag-1",), models=models, mcp_server_ids=frozenset(), agent_ids=frozenset() + ) + + return resolve, asked + + +@pytest.mark.asyncio +async def test_agent_access_groups_cap_models_even_when_key_allows_them(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) + resolve, asked = _agent_model_ceiling_resolver(frozenset({"gpt-5"})) + + assert await _check_agent_access_group_model_access("gpt-5", agent_key, None, resolve) is True + + with pytest.raises(ProxyException) as exc_info: + await _check_agent_access_group_model_access("claude-sonnet", agent_key, None, resolve) + + assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied + assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN) + assert asked == ["agent-1", "agent-1"] + + +@pytest.mark.asyncio +async def test_agent_access_groups_naming_no_model_deny_every_model(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=[]) + resolve, _ = _agent_model_ceiling_resolver(frozenset()) + + with pytest.raises(ProxyException) as exc_info: + await _check_agent_access_group_model_access("gpt-5", agent_key, None, resolve) + + assert exc_info.value.type == ProxyErrorTypes.agent_model_access_denied + + +@pytest.mark.asyncio +async def test_agent_without_access_groups_adds_no_model_ceiling(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) + resolve, asked = _agent_model_ceiling_resolver(None) + + assert await _check_agent_access_group_model_access("gpt-5", agent_key, None, resolve) is True + assert await _check_agent_access_group_model_access("claude-sonnet", agent_key, None, resolve) is True + assert asked == ["agent-1", "agent-1"] + + +@pytest.mark.asyncio +async def test_key_without_agent_never_consults_agent_access_groups(): + plain_key: Final = UserAPIKeyAuth(token="plain-token", models=["gpt-5"]) + resolve, asked = _agent_model_ceiling_resolver(frozenset()) + + assert await _check_agent_access_group_model_access("gpt-5", plain_key, None, resolve) is True + assert asked == [] + + @pytest.mark.asyncio async def test_team_member_budget_check_temp_budget_increase_extends_cap(): """Spend above max_budget but below max_budget + active temp increase @@ -9086,3 +9155,117 @@ async def test_team_member_budget_check_adds_temp_increase_to_live_team_default( proxy_logging_obj=ProxyLogging(user_api_key_cache=None), ) assert exc_info.value.max_budget == expected_cap + + +def _agent_key_acting_for(user_id: str | None, team_id: str | None) -> UserAPIKeyAuth: + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) + agent_key.agent_caller = AgentCaller(user_id=user_id, team_id=team_id) + return agent_key + + +def _caller_loaders( + team: LiteLLM_TeamTable | None, + user: LiteLLM_UserTable | None, +) -> tuple[CallerTeamLoader, CallerUserLoader, list[str]]: + """Loaders that hand back fixed caller rows and record the agent_caller they were asked about.""" + asked: Final[list[str]] = [] + + async def load_team(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: + asked.append(f"team:{valid_token.agent_caller.team_id if valid_token.agent_caller else None}") + return team + + async def load_user(valid_token: UserAPIKeyAuth) -> LiteLLM_UserTable | None: + asked.append(f"user:{valid_token.agent_caller.user_id if valid_token.agent_caller else None}") + return user + + return load_team, load_user, asked + + +async def _cache_with_membership(user_id: str, team_id: str, allowed_models: list[str] | None) -> UserApiKeyCache: + from litellm.proxy._types import LiteLLM_TeamMembership + from litellm.proxy.common_utils.user_api_key_cache import team_membership_reservation_cache_key + + cache: Final = UserApiKeyCache() + await cache.async_set_cache( + key=team_membership_reservation_cache_key(user_id=user_id, team_id=team_id), + value=LiteLLM_TeamMembership( + user_id=user_id, + team_id=team_id, + litellm_budget_table=LiteLLM_BudgetTable(allowed_models=allowed_models) if allowed_models else None, + ), + model_type=LiteLLM_TeamMembership, + ) + return cache + + +async def _check_caller_models( + agent_key: UserAPIKeyAuth, + model: str, + load_team: CallerTeamLoader, + load_user: CallerUserLoader, + cache: UserApiKeyCache | None = None, +) -> None: + await _check_agent_caller_model_access( + model=model, + valid_token=agent_key, + llm_router=None, + prisma_client=None, + user_api_key_cache=cache or UserApiKeyCache(), + proxy_logging_obj=MagicMock(), + load_team=load_team, + load_user=load_user, + ) + + +@pytest.mark.asyncio +async def test_agent_key_acting_for_a_team_is_capped_at_that_teams_models(): + """LIT-8014: the invoking team may only call gpt-5, so the agent's own claude grant does not help.""" + agent_key: Final = _agent_key_acting_for(user_id="alice", team_id="team-a") + load_team, load_user, asked = _caller_loaders(LiteLLM_TeamTable(team_id="team-a", models=["gpt-5"]), None) + cache: Final = await _cache_with_membership("alice", "team-a", allowed_models=None) + + await _check_caller_models(agent_key, "gpt-5", load_team, load_user, cache) + with pytest.raises(ProxyException) as exc_info: + await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user, cache) + + assert exc_info.value.type == ProxyErrorTypes.team_model_access_denied + assert exc_info.value.code == str(status.HTTP_403_FORBIDDEN) + assert asked == ["team:team-a", "team:team-a"] + + +@pytest.mark.asyncio +async def test_agent_key_acting_for_a_team_member_is_capped_at_the_members_scope(): + agent_key: Final = _agent_key_acting_for(user_id="alice", team_id="team-a") + load_team, load_user, _ = _caller_loaders( + LiteLLM_TeamTable(team_id="team-a", models=["gpt-5", "claude-sonnet"]), None + ) + cache: Final = await _cache_with_membership("alice", "team-a", allowed_models=["gpt-5"]) + + await _check_caller_models(agent_key, "gpt-5", load_team, load_user, cache) + with pytest.raises(ProxyException) as exc_info: + await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user, cache) + + assert "User=alice, Team=team-a" in exc_info.value.internal_message + + +@pytest.mark.asyncio +async def test_agent_key_acting_for_a_teamless_user_is_capped_at_that_users_models(): + agent_key: Final = _agent_key_acting_for(user_id="alice", team_id=None) + load_team, load_user, asked = _caller_loaders(None, LiteLLM_UserTable(user_id="alice", models=["gpt-5"])) + + await _check_caller_models(agent_key, "gpt-5", load_team, load_user) + with pytest.raises(ProxyException) as exc_info: + await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user) + + assert exc_info.value.type == ProxyErrorTypes.user_model_access_denied + assert asked == ["team:None", "user:alice", "team:None", "user:alice"] + + +@pytest.mark.asyncio +async def test_agent_key_without_an_echoed_caller_keeps_its_own_models(): + agent_key: Final = UserAPIKeyAuth(token="agent-token", agent_id="agent-1", models=["gpt-5", "claude-sonnet"]) + load_team, load_user, asked = _caller_loaders(LiteLLM_TeamTable(team_id="team-a", models=[]), None) + + await _check_caller_models(agent_key, "claude-sonnet", load_team, load_user) + + assert asked == [] 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/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index a48c64eb4a0..9353c149d15 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -5,13 +5,16 @@ import shlex import stat import sys import time +from concurrent.futures import ThreadPoolExecutor, TimeoutError as FutureTimeoutError from pathlib import Path +from threading import Event +from typing import Final from unittest.mock import patch import pytest from click.testing import CliRunner -from litellm.litellm_core_utils.private_json import commit_staged_json +from litellm.litellm_core_utils.private_json import commit_staged_json, write_private_bytes from litellm.proxy.client.cli.commands.claude_settings import ( ANTHROPIC_DEFAULT_MODEL_ENV_KEYS, AUTOROUTE_BACKUP_PATH, @@ -773,7 +776,7 @@ class TestStatusLine: script = tmp_path / "lite" / "statusline.py" command = install_statusline_script(script) - assert script.read_bytes() == pathlib.Path(statusline_script.__file__).read_bytes() + assert script.read_bytes().split(b"\n", 1)[1] == pathlib.Path(statusline_script.__file__).read_bytes() assert shlex.split(command) == [sys.executable, str(script)] assert command == statusline_command(script) assert stat.S_IMODE(script.stat().st_mode) == 0o600 @@ -783,15 +786,13 @@ class TestStatusLine: def test_a_reinstall_replaces_the_script_in_one_step_and_a_refused_one_leaves_the_old_script_whole(self, tmp_path): # Claude Code may be running the script at the moment `lite` reinstalls it; the file it has open # must stay complete, and a reinstall that cannot land must not leave a truncated script behind. - from litellm.proxy.client.cli.commands import statusline_script - script = tmp_path / "lite" / "statusline.py" install_statusline_script(script) - bundled = pathlib.Path(statusline_script.__file__).read_bytes() + bundled = script.read_bytes() with script.open("rb") as running: install_statusline_script(script) assert running.read() == bundled - assert [child.name for child in script.parent.iterdir()] == ["statusline.py"] + assert {child.name for child in script.parent.iterdir()} <= {"statusline.py", "statusline.py.lock"} if os.geteuid() != 0: script.parent.chmod(0o500) @@ -802,6 +803,141 @@ class TestStatusLine: script.parent.chmod(0o700) assert script.read_bytes() == bundled + @pytest.mark.parametrize( + ("installed_version", "older_version"), + (("2.10.0", "2.9.0"), ("2.1.0", "2.1.0rc1"), ("2.1.0rc1", "2.1.0.dev2"), ("2.1.0.post1", "2.1.0")), + ) + def test_an_older_cli_preserves_the_newer_footer( + self, tmp_path: Path, capsys: pytest.CaptureFixture[str], installed_version: str, older_version: str + ) -> None: + script: Final = tmp_path / "statusline.py" + command: Final = install_statusline_script(script, package_version=installed_version) + installed: Final = script.read_bytes() + modified: Final = script.stat().st_mtime_ns + + assert install_statusline_script(script, package_version=older_version) == command + + assert script.read_bytes() == installed + assert script.stat().st_mtime_ns == modified + assert f"Keeping the status line from LiteLLM {installed_version}" in capsys.readouterr().err + + @pytest.mark.parametrize( + "old_header", (b"", b"# litellm-statusline-version: invalid\n", b"# litellm-statusline-version: \xff\n") + ) + def test_a_legacy_or_damaged_version_marker_is_repaired(self, tmp_path: Path, old_header: bytes) -> None: + from litellm.proxy.client.cli.commands import statusline_script + + script: Final = tmp_path / "statusline.py" + script.write_bytes(old_header + b"print('old footer')\n") + + install_statusline_script(script, package_version="2.1.0") + + assert script.read_bytes() == ( + b"# litellm-statusline-version: 2.1.0\n" + Path(statusline_script.__file__).read_bytes() + ) + + @pytest.mark.parametrize("next_version", ("2.1.0", "2.2.0")) + def test_an_equal_or_newer_cli_refreshes_the_footer(self, tmp_path: Path, next_version: str) -> None: + from litellm.proxy.client.cli.commands import statusline_script + + script: Final = tmp_path / "statusline.py" + script.write_bytes(b"# litellm-statusline-version: 2.1.0\nprint('old footer')\n") + + install_statusline_script(script, package_version=next_version) + + assert script.read_bytes() == ( + f"# litellm-statusline-version: {next_version}\n".encode() + Path(statusline_script.__file__).read_bytes() + ) + + def test_configure_keeps_a_newer_footer_while_updating_settings( + self, tmp_path: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + script: Final = tmp_path / "statusline.py" + script.write_bytes(b"# litellm-statusline-version: 999999.0.0\nprint('newer footer')\n") + installed: Final = script.read_bytes() + rig: Final = _Rig(tmp_path, {"theme": "dark"}) + + rig.configure(script_path=script) + + assert script.read_bytes() == installed + assert rig.read()["statusLine"]["command"] == statusline_command(script) + assert rig.read()["env"]["ANTHROPIC_BASE_URL"] == PROXY + assert "Keeping the status line" in capsys.readouterr().err + + @pytest.mark.parametrize("package_version", ("unknown", "", "invalid-version")) + @pytest.mark.parametrize("existing", (None, b"print('legacy footer')\n", b"# litellm-statusline-version: invalid\n")) + def test_an_unknown_cli_version_can_install_and_refresh_an_unversioned_footer( + self, tmp_path: Path, package_version: str, existing: bytes | None + ) -> None: + from litellm.proxy.client.cli.commands import statusline_script + + script: Final = tmp_path / "statusline.py" + if existing is not None: + script.write_bytes(existing) + + assert install_statusline_script(script, package_version=package_version) == statusline_command(script) + assert script.read_bytes() == Path(statusline_script.__file__).read_bytes() + + def test_an_unknown_cli_version_preserves_a_versioned_footer( + self, tmp_path: Path, capsys: pytest.CaptureFixture[str] + ) -> None: + script: Final = tmp_path / "statusline.py" + command: Final = install_statusline_script(script, package_version="2.1.0") + installed: Final = script.read_bytes() + + assert install_statusline_script(script, package_version="unknown") == command + assert script.read_bytes() == installed + assert "Keeping the status line from LiteLLM 2.1.0" in capsys.readouterr().err + + @pytest.mark.parametrize(("first_version", "second_version"), (("2.0", "3.0"), ("3.0", "2.0"))) + def test_overlapping_installs_keep_the_newest_footer( + self, tmp_path: Path, first_version: str, second_version: str + ) -> None: + from litellm.proxy.client.cli.commands import statusline_script + + script: Final = tmp_path / "statusline.py" + first_writing: Final = Event() + release_first: Final = Event() + second_started: Final = Event() + + def paused_write(path: str, data: bytes) -> None: + first_writing.set() + assert release_first.wait(5), "First installer was never released" + write_private_bytes(path, data) + + def second_install() -> str: + second_started.set() + return install_statusline_script(script, package_version=second_version) + + with ThreadPoolExecutor(max_workers=2) as pool: + first: Final = pool.submit(install_statusline_script, script, package_version=first_version, write=paused_write) + try: + assert first_writing.wait(5), "First installer did not reach the write" + second: Final = pool.submit(second_install) + assert second_started.wait(5), "Second installer did not start" + with pytest.raises(FutureTimeoutError): + second.result(timeout=0.5) + finally: + release_first.set() + assert first.result(timeout=5) == statusline_command(script) + assert second.result(timeout=5) == statusline_command(script) + + assert script.read_bytes() == b"# litellm-statusline-version: 3.0\n" + Path(statusline_script.__file__).read_bytes() + + def test_a_failed_install_keeps_the_footer_and_releases_the_lock(self, tmp_path: Path) -> None: + script: Final = tmp_path / "statusline.py" + install_statusline_script(script, package_version="2.0") + installed: Final = script.read_bytes() + + def failed_write(path: str, data: bytes) -> None: + raise OSError("disk full") + + with pytest.raises(ClaudeSettingsError, match="disk full"): + install_statusline_script(script, package_version="3.0", write=failed_write) + assert script.read_bytes() == installed + assert install_statusline_script(script, package_version="3.0") == statusline_command(script) + assert script.read_bytes().startswith(b"# litellm-statusline-version: 3.0\n") + def test_configure_installs_it_and_unconfigure_removes_only_ours(self, tmp_path): rig = _Rig(tmp_path, {"theme": "dark"}) script = tmp_path / "statusline.py" diff --git a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py index acd3dc18b54..c61a489f894 100644 --- a/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py +++ b/tests/test_litellm/proxy/db/test_autorouter_session_rollup.py @@ -56,6 +56,31 @@ def _build(payload: dict | None = None, metadata: dict | None = None): class TestBuildTransaction: + @pytest.mark.parametrize( + "api_key, user_id, included", + [ + ("hashed-key", "canonical-user", True), + ("hashed-key", None, True), + ("hashed-key", "", True), + ("", "canonical-user", True), + ("", None, False), + ("", "", False), + ], + ) + def test_attribution_uses_the_canonical_user_even_without_a_key( + self, api_key: str, user_id: str | None, included: bool + ) -> None: + transaction: Final = _build( + payload=_payload(api_key=api_key, user=user_id), + metadata=_metadata(user="client-user", user_api_key_user_id="metadata-user"), + ) + if not included: + assert transaction is None + return + assert transaction is not None + assert transaction.api_key == api_key + assert transaction.user_id == (user_id or "") + def test_successful_auto_routed_turn_builds_every_field(self): transaction = _build( metadata=_metadata( @@ -205,23 +230,43 @@ class TestBuildTransaction: class _FakeDB: - def __init__(self, failures: "list[Exception] | None" = None, poison_session: str | None = None): + def __init__( + self, + failures: "list[Exception] | None" = None, + poison_session: str | None = None, + poison_user: str | None = None, + commit_then_error_users: frozenset[str] = frozenset(), + ): self.calls: list[tuple] = [] + self.attempts: list[tuple[str, tuple[object, ...]]] = [] self._failures = list(failures or []) self._poison_session = poison_session + self._poison_user = poison_user + self._commit_then_error_users = commit_then_error_users async def execute_raw(self, sql: str, *params: object) -> int: + self.attempts.append((sql, params)) if self._poison_session is not None and params[1] == self._poison_session: raise RuntimeError("index row size exceeds btree maximum") + if self._poison_user is not None and params[19] == self._poison_user: + raise RuntimeError("index row size exceeds btree maximum") if self._failures: raise self._failures.pop(0) self.calls.append((sql, params)) + if params[19] in self._commit_then_error_users: + raise RuntimeError("commit succeeded but acknowledgement was lost") return 1 class _FakeClient: - def __init__(self, failures: "list[Exception] | None" = None, poison_session: str | None = None): - self.db = _FakeDB(failures, poison_session) + def __init__( + self, + failures: "list[Exception] | None" = None, + poison_session: str | None = None, + poison_user: str | None = None, + commit_then_error_users: frozenset[str] = frozenset(), + ): + self.db = _FakeDB(failures, poison_session, poison_user, commit_then_error_users) def _transaction( @@ -229,9 +274,11 @@ def _transaction( at: datetime = datetime(2026, 8, 1, 12, 0, 0), tier: str | None = "medium", baseline_model: str | None = "anthropic/claude-opus-5", + api_key: str = "k1", + user_id: str = "", ) -> AutoRouterTurnTransaction: return AutoRouterTurnTransaction( - api_key="k1", + api_key=api_key, session_id=session_id, router_name="live-auto", router_type="complexity", @@ -247,6 +294,7 @@ def _transaction( cache_touched=False, tier=tier, baseline_model=baseline_model, + user_id=user_id, ) @@ -261,7 +309,7 @@ class TestFlush: def test_params_marshal_in_statement_order(self): client = _FakeClient() - asyncio.run(flush_autorouter_turn_transactions(client, [_transaction()])) + asyncio.run(flush_autorouter_turn_transactions(client, [_transaction(user_id="canonical-user")])) sql, params = client.db.calls[0] assert sql == UPSERT_AUTOROUTER_SESSION_SQL assert params == ( @@ -284,8 +332,65 @@ class TestFlush: 0, 0.0, 0.0, + "canonical-user", ) + def test_a_keys_turns_stay_chronological_when_its_canonical_user_changes(self) -> None: + client: Final = _FakeClient() + earlier: Final = _transaction(user_id="z-user", at=datetime(2026, 8, 1, 12, 0, 0)) + later: Final = _transaction(user_id="a-user", at=datetime(2026, 8, 1, 12, 0, 10)) + asyncio.run(flush_autorouter_turn_transactions(client, [later, earlier])) + assert [(params[5], params[19]) for _, params in client.db.calls] == [ + ("2026-08-01T12:00:00", "z-user"), + ("2026-08-01T12:00:10", "a-user"), + ] + + def test_one_keyless_users_failed_session_does_not_drop_another_users_turn(self) -> None: + client: Final = _FakeClient(poison_user="a-user") + failed: Final = _transaction(api_key="", user_id="a-user") + other: Final = _transaction(api_key="", user_id="b-user", at=datetime(2026, 8, 1, 12, 0, 10)) + asyncio.run(flush_autorouter_turn_transactions(client, [other, failed])) + assert [(params[0], params[1], params[19]) for _, params in client.db.calls] == [("", "s1", "b-user")] + + def test_uncertain_commits_quarantine_only_the_key_and_each_failed_user(self) -> None: + client: Final = _FakeClient(commit_then_error_users=frozenset({"a-failed", "c-failed"})) + turns: Final = tuple( + _transaction(user_id=user, at=datetime(2026, 8, 1, 12, 0, second), api_key=key) + for user, second, key in ( + ("b-healthy", 0, "k1"), + ("a-failed", 1, "k1"), + ("b-healthy", 2, "k1"), + ("c-failed", 3, "k1"), + ("b-healthy", 4, "k1"), + ("d-healthy", 5, "k1"), + ("c-failed", 6, "k1"), + ("d-healthy", 7, "k1"), + ("a-failed", 8, "k1"), + ("", 9, "k1"), + ("z-other", 10, "k2"), + ) + ) + asyncio.run(flush_autorouter_turn_transactions(client, tuple(reversed(turns)))) + + assert client.db.attempts == client.db.calls + assert [ + (params[0], params[19], params[5]) + for sql, params in client.db.calls + if sql == UPSERT_AUTOROUTER_SESSION_SQL + ] == [ + ("k1", "b-healthy", "2026-08-01T12:00:00"), + ("k1", "a-failed", "2026-08-01T12:00:01"), + ("k2", "z-other", "2026-08-01T12:00:10"), + ] + assert [params[19] for _, params in client.db.attempts].count("a-failed") == 1 + assert [params[19] for _, params in client.db.attempts].count("c-failed") == 1 + for user, seconds in (("b-healthy", (2, 4)), ("c-failed", (3,)), ("d-healthy", (5, 7))): + assert [ + (params[0], params[5]) + for sql, params in client.db.calls + if sql != UPSERT_AUTOROUTER_SESSION_SQL and params[19] == user + ] == [("k1", f"2026-08-01T12:00:{second:02d}") for second in seconds] + def test_a_connect_error_retries_the_same_statement(self): client = _FakeClient(failures=[httpx.ConnectError("boom")]) asyncio.run(flush_autorouter_turn_transactions(client, [_transaction()])) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index d1d22d0d7c2..c90f88ec110 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -2,7 +2,7 @@ import logging from types import SimpleNamespace -from typing import Final +from typing import TYPE_CHECKING, Final, Literal import pytest @@ -13,7 +13,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route -from litellm.llms import load_guardrail_translation_mappings +from litellm.llms import discover_guardrail_translation_mappings, load_guardrail_translation_mappings from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.llms.base_llm.guardrail_translation.utils import ( effective_skip_system_message_for_guardrail, @@ -41,7 +41,10 @@ from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrai ) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponsesAPIResponse -from litellm.types.utils import CallTypes, Delta, ModelResponseStream, StreamingChoices +from litellm.types.utils import CallTypes, Delta, GenericGuardrailAPIInputs, ModelResponseStream, StreamingChoices + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj class RecordingGuardrail(CustomGuardrail): @@ -61,6 +64,18 @@ class RecordingGuardrail(CustomGuardrail): return {"texts": inputs.get("texts", [])} +class RewritingGuardrail(RecordingGuardrail): + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, # mutable-ok: CustomGuardrail.apply_guardrail contract + input_type: Literal["request", "response"], + logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> GenericGuardrailAPIInputs: + recorded: Final = await super().apply_guardrail(inputs, request_data, input_type, logging_obj=logging_obj) + return GenericGuardrailAPIInputs(texts=[f"{text} [GUARDRAILED]" for text in recorded["texts"]]) + + class _NoopTranslation(BaseTranslation): """Test translation handler that simply echoes input/output.""" @@ -115,9 +130,7 @@ class TestUnifiedLLMGuardrails: assert msgs[0]["content"] == "sys" def test_effective_skip_respects_per_guardrail_over_global(self, monkeypatch): - monkeypatch.setattr( - litellm, "skip_system_message_in_guardrail", True, raising=False - ) + monkeypatch.setattr(litellm, "skip_system_message_in_guardrail", True, raising=False) class G: skip_system_message_in_guardrail = False @@ -130,21 +143,15 @@ class TestUnifiedLLMGuardrails: assert effective_skip_system_message_for_guardrail(G2()) is True @pytest.mark.asyncio - async def test_openai_handler_skips_system_in_guardrail_inputs( - self, monkeypatch - ): - monkeypatch.setattr( - litellm, "skip_system_message_in_guardrail", True, raising=False - ) + async def test_openai_handler_skips_system_in_guardrail_inputs(self, monkeypatch): + monkeypatch.setattr(litellm, "skip_system_message_in_guardrail", True, raising=False) captured = {} class MockGuardrail: skip_system_message_in_guardrail = None - async def apply_guardrail( - self, inputs, request_data, input_type, logging_obj=None - ): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): captured["inputs"] = inputs return inputs @@ -169,21 +176,15 @@ class TestUnifiedLLMGuardrails: assert data["messages"][0]["content"] == "secret system" @pytest.mark.asyncio - async def test_openai_handler_per_guardrail_skip_false_overrides_global( - self, monkeypatch - ): - monkeypatch.setattr( - litellm, "skip_system_message_in_guardrail", True, raising=False - ) + async def test_openai_handler_per_guardrail_skip_false_overrides_global(self, monkeypatch): + monkeypatch.setattr(litellm, "skip_system_message_in_guardrail", True, raising=False) captured = {} class MockGuardrail: skip_system_message_in_guardrail = False - async def apply_guardrail( - self, inputs, request_data, input_type, logging_obj=None - ): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): captured["inputs"] = inputs return inputs @@ -201,10 +202,7 @@ class TestUnifiedLLMGuardrails: ) assert "sys" in captured["inputs"]["texts"] - roles = { - m.get("role") - for m in (captured["inputs"].get("structured_messages") or []) - } + roles = {m.get("role") for m in (captured["inputs"].get("structured_messages") or [])} assert "system" in roles class TestSkipToolMessageForChatCompletions: @@ -229,12 +227,8 @@ class TestUnifiedLLMGuardrails: assert all(m["role"] != "tool" for m in out) assert msgs[2]["content"] == "tool result" - def test_effective_skip_tool_respects_per_guardrail_over_global( - self, monkeypatch - ): - monkeypatch.setattr( - litellm, "skip_tool_message_in_guardrail", True, raising=False - ) + def test_effective_skip_tool_respects_per_guardrail_over_global(self, monkeypatch): + monkeypatch.setattr(litellm, "skip_tool_message_in_guardrail", True, raising=False) class G: skip_tool_message_in_guardrail = False @@ -248,18 +242,14 @@ class TestUnifiedLLMGuardrails: @pytest.mark.asyncio async def test_openai_handler_skips_tool_in_guardrail_inputs(self, monkeypatch): - monkeypatch.setattr( - litellm, "skip_tool_message_in_guardrail", True, raising=False - ) + monkeypatch.setattr(litellm, "skip_tool_message_in_guardrail", True, raising=False) captured = {} class MockGuardrail: skip_tool_message_in_guardrail = None - async def apply_guardrail( - self, inputs, request_data, input_type, logging_obj=None - ): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): captured["inputs"] = inputs return inputs @@ -299,21 +289,15 @@ class TestUnifiedLLMGuardrails: assert data["messages"][2]["content"] == "secret tool result" @pytest.mark.asyncio - async def test_openai_handler_per_guardrail_skip_tool_false_overrides_global( - self, monkeypatch - ): - monkeypatch.setattr( - litellm, "skip_tool_message_in_guardrail", True, raising=False - ) + async def test_openai_handler_per_guardrail_skip_tool_false_overrides_global(self, monkeypatch): + monkeypatch.setattr(litellm, "skip_tool_message_in_guardrail", True, raising=False) captured = {} class MockGuardrail: skip_tool_message_in_guardrail = False - async def apply_guardrail( - self, inputs, request_data, input_type, logging_obj=None - ): + async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None): captured["inputs"] = inputs return inputs @@ -331,10 +315,7 @@ class TestUnifiedLLMGuardrails: ) assert "tr" in captured["inputs"]["texts"] - roles = { - m.get("role") - for m in (captured["inputs"].get("structured_messages") or []) - } + roles = {m.get("role") for m in (captured["inputs"].get("structured_messages") or [])} assert "tool" in roles class TestAsyncPreCallHook: @@ -360,6 +341,38 @@ class TestUnifiedLLMGuardrails: assert guardrail.event_history == [GuardrailEventHooks.pre_mcp_call] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "call_type", + ["avideo_generation", "acreate_video", "avideo_remix", "avideo_edit", "avideo_extension"], + ) + async def test_video_routes_scan_prompt_and_keep_rewrite(self, monkeypatch, call_type: str) -> None: + """LIT-6685: /v1/videos dispatches call_type="avideo_generation", which the + hook once swallowed as an unknown CallTypes value and returned unscanned. + Runs against the discovered handler map so the video package must really exist.""" + _patch_translation_mappings(monkeypatch, discover_guardrail_translation_mappings()) + handler = UnifiedLLMGuardrails() + guardrail = RewritingGuardrail() + data = { + "guardrail_to_apply": guardrail, + "model": "veo-3.1-fast", + "prompt": "a paper boat on a stream", + "seconds": "4", + } + + result = await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + cache=DualCache(), + data=data, + call_type=call_type, + ) + + assert guardrail.event_history == [GuardrailEventHooks.pre_call] + assert [call["inputs"]["texts"] for call in guardrail.apply_calls] == [["a paper boat on a stream"]] + assert guardrail.apply_calls[0]["inputs"]["model"] == "veo-3.1-fast" + assert result["prompt"] == "a paper boat on a stream [GUARDRAILED]" + assert result["seconds"] == "4" + class TestAsyncModerationHook: @pytest.mark.asyncio async def test_uses_mcp_event_type(self): @@ -424,7 +437,9 @@ class TestUnifiedLLMGuardrails: async def process_input_messages(self, data, guardrail_to_apply, litellm_logging_obj=None): # type: ignore[override] return data - async def process_output_response(self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None): # type: ignore[override] + async def process_output_response( + self, response, guardrail_to_apply, litellm_logging_obj=None, user_api_key_dict=None + ): # type: ignore[override] return response async def process_output_streaming_response( @@ -493,9 +508,7 @@ class TestUnifiedLLMGuardrails: response=mock_stream(), request_data=request_data, ): - content = ( - item.choices[0].delta.content if item.choices[0].delta else None - ) + content = item.choices[0].delta.content if item.choices[0].delta else None yielded_contents.append(content) # Every chunk should have non-empty content @@ -546,23 +559,18 @@ class TestUnifiedLLMGuardrails: ], ) @pytest.mark.asyncio - async def test_post_call_scans_output_on_every_registered_alias( - self, request_route: str - ) -> None: + async def test_post_call_scans_output_on_every_registered_alias(self, request_route: str) -> None: handler = UnifiedLLMGuardrails() guardrail = RecordingGuardrail() await handler.async_post_call_success_hook( data={"guardrail_to_apply": guardrail, "model": "gpt-4o"}, - user_api_key_dict=UserAPIKeyAuth( - api_key="test-key", request_route=request_route - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route=request_route), response=self._responses_api_response(), ) assert guardrail.apply_calls, ( - f"guardrail never ran for request_route={request_route!r}; model " - f"output reached the client unscanned" + f"guardrail never ran for request_route={request_route!r}; model output reached the client unscanned" ) assert guardrail.apply_calls[0]["input_type"] == "response" assert guardrail.apply_calls[0]["inputs"]["texts"] == ["Paris"] @@ -592,18 +600,14 @@ class TestUnifiedLLMGuardrails: assert CallTypes.responses in mappings @pytest.mark.asyncio - async def test_unresolvable_route_skips_scanning_and_says_so( - self, caplog: pytest.LogCaptureFixture - ) -> None: + async def test_unresolvable_route_skips_scanning_and_says_so(self, caplog: pytest.LogCaptureFixture) -> None: handler = UnifiedLLMGuardrails() guardrail = RecordingGuardrail() with caplog.at_level(logging.WARNING): result = await handler.async_post_call_success_hook( data={"guardrail_to_apply": guardrail, "model": "gpt-4o"}, - user_api_key_dict=UserAPIKeyAuth( - api_key="test-key", request_route="/cursor/chat/completions" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/cursor/chat/completions"), response=self._responses_api_response(), ) @@ -622,9 +626,7 @@ class TestUnifiedLLMGuardrails: with caplog.at_level(logging.WARNING): await handler.async_post_call_success_hook( data={"guardrail_to_apply": guardrail, "model": "gpt-4o"}, - user_api_key_dict=UserAPIKeyAuth( - api_key="test-key", request_route="/v1/chat/completions" - ), + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", request_route="/v1/chat/completions"), response=self._responses_api_response(), ) @@ -734,15 +736,10 @@ class TestUnifiedLLMGuardrails: assert guardrail.event_history == [GuardrailEventHooks.pre_call] assert len(guardrail.apply_calls) == 1 assert guardrail.apply_calls[0]["input_type"] == "request" - assert ( - "https://arxiv.org/pdf/2201.04234" - in guardrail.apply_calls[0]["inputs"]["texts"] - ) + assert "https://arxiv.org/pdf/2201.04234" in guardrail.apply_calls[0]["inputs"]["texts"] # Data should be returned with document intact - assert ( - result["document"]["document_url"] == "https://arxiv.org/pdf/2201.04234" - ) + assert result["document"]["document_url"] == "https://arxiv.org/pdf/2201.04234" @pytest.mark.asyncio async def test_moderation_hook_invokes_ocr_handler(self): @@ -770,10 +767,7 @@ class TestUnifiedLLMGuardrails: assert guardrail.event_history == [GuardrailEventHooks.during_call] assert len(guardrail.apply_calls) == 1 - assert ( - "https://example.com/scan.png" - in guardrail.apply_calls[0]["inputs"]["texts"] - ) + assert "https://example.com/scan.png" in guardrail.apply_calls[0]["inputs"]["texts"] @pytest.mark.asyncio async def test_post_call_success_hook_guardrails_ocr_output(self): @@ -789,9 +783,7 @@ class TestUnifiedLLMGuardrails: def should_run_guardrail(self, data, event_type): # type: ignore[override] return True - async def apply_guardrail( - self, inputs, request_data, input_type, **kwargs - ): + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): texts = inputs.get("texts", []) return {"texts": [t.replace("SECRET", "[REDACTED]") for t in texts]} @@ -1538,9 +1530,7 @@ class TestStreamingTransform: # And the redacted text ("SECRET") reached the wire on some non-tool # chunk (i.e. the text terminator). transformed = "".join( - item.choices[0].delta.content or "" - for item in out - if item.choices and not item.choices[0].delta.tool_calls + item.choices[0].delta.content or "" for item in out if item.choices and not item.choices[0].delta.tool_calls ) assert "SECRET" in transformed assert "secret" not in transformed @@ -1685,7 +1675,9 @@ class TestStreamingTransform: _stream_chunk("went home."), ModelResponseStream( choices=[ - StreamingChoices(index=0, delta=Delta(content=None, role="assistant", tool_calls=None), finish_reason=None), + StreamingChoices( + index=0, delta=Delta(content=None, role="assistant", tool_calls=None), finish_reason=None + ), StreamingChoices( index=1, delta=Delta( @@ -1985,9 +1977,7 @@ class TestStreamingHttpErrorFrames: guardrail = _EosHttpBlockingGuardrail() chunks = _anthropic_message_chunks(["hello ", "world"]) - out = await _drive_stream( - UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages" - ) + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages") raw = b"".join(c for c in out if isinstance(c, bytes)).decode() assert "hello " in raw @@ -2011,9 +2001,7 @@ class TestStreamingHttpErrorFrames: }, ] - out = await _drive_stream( - UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses" - ) + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses") assert chunks[0] in out and chunks[1] in out assert chunks[2] not in out @@ -2076,9 +2064,7 @@ class TestStreamingGuardrailInformationBucket: for chunk in chunks: yield chunk - user_api_key_dict = UserAPIKeyAuth( - api_key="test-key", user_id="user-1", request_route="/v1/chat/completions" - ) + user_api_key_dict = UserAPIKeyAuth(api_key="test-key", user_id="user-1", request_route="/v1/chat/completions") request_data = {"guardrail_to_apply": guardrail, "model": "gpt-4", "metadata": {}} out = [] @@ -2407,7 +2393,5 @@ class TestTranslationMappingsAreReadLive: assert len(guardrail.apply_calls) == 1 assert not [ - name - for name, value in vars(unified_module).items() - if isinstance(value, dict) and CallTypes.aocr in value + name for name, value in vars(unified_module).items() if isinstance(value, dict) and CallTypes.aocr in value ] diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py index d687f8d1c8d..13e57408bd0 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_endpoints.py @@ -17,7 +17,9 @@ from litellm.proxy._types import ( LitellmUserRoles, UserAPIKeyAuth, ) +from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry from litellm.proxy.proxy_server import app +from litellm.types.agents import AgentResponse def _make_access_group_record( @@ -126,6 +128,7 @@ def client_and_mocks(monkeypatch): mock_agents_table = MagicMock() mock_agents_table.find_many = AsyncMock(return_value=[]) + mock_agents_table.update = AsyncMock(return_value=None) @asynccontextmanager async def mock_tx(): @@ -133,6 +136,7 @@ def client_and_mocks(monkeypatch): litellm_accessgrouptable=mock_access_group_table, litellm_teamtable=mock_team_table, litellm_verificationtoken=mock_key_table, + litellm_agentstable=mock_agents_table, ) yield tx @@ -158,15 +162,9 @@ def client_and_mocks(monkeypatch): mock_proxy_logging = MagicMock() mock_proxy_logging.internal_usage_cache = MagicMock() mock_proxy_logging.internal_usage_cache.dual_cache = MagicMock() - mock_proxy_logging.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock( - return_value=None - ) - mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock( - return_value=None - ) - mock_proxy_logging.internal_usage_cache.dual_cache.async_set_cache = AsyncMock( - return_value=None - ) + mock_proxy_logging.internal_usage_cache.dual_cache.async_delete_cache = AsyncMock(return_value=None) + mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=None) + mock_proxy_logging.internal_usage_cache.dual_cache.async_set_cache = AsyncMock(return_value=None) monkeypatch.setattr(ps, "proxy_logging_obj", mock_proxy_logging) admin_user = UserAPIKeyAuth( @@ -239,9 +237,7 @@ def test_create_access_group_duplicate_name_conflict(client_and_mocks): "unique constraint violation", ], ) -def test_create_access_group_race_condition_returns_409( - client_and_mocks, error_message -): +def test_create_access_group_race_condition_returns_409(client_and_mocks, error_message): """Create race condition: Prisma unique constraint surfaces as 409, not 500.""" client, _, mock_table, *_ = client_and_mocks @@ -288,9 +284,7 @@ def test_create_access_group_500_on_non_constraint_prisma_error(client_and_mocks # Use raise_server_exceptions=False so unhandled exceptions become 500 responses test_client = TestClient(app, raise_server_exceptions=False) - resp = test_client.post( - "/v1/access_group", json={"access_group_name": "test-group"} - ) + resp = test_client.post("/v1/access_group", json={"access_group_name": "test-group"}) assert resp.status_code == 500 @@ -558,9 +552,7 @@ def test_update_access_group_empty_body(client_and_mocks): """Update with empty body succeeds; only updated_by is set.""" client, _, mock_table, *_ = client_and_mocks - existing = _make_access_group_record( - access_group_id="ag-update", access_group_name="unchanged" - ) + existing = _make_access_group_record(access_group_id="ag-update", access_group_name="unchanged") mock_table.find_unique = AsyncMock(return_value=existing) resp = client.put("/v1/access_group/ag-update", json={}) @@ -576,14 +568,10 @@ def test_update_access_group_name_success(client_and_mocks): """Update access_group_name succeeds when new name is unique.""" client, _, mock_table, *_ = client_and_mocks - existing = _make_access_group_record( - access_group_id="ag-update", access_group_name="old-name" - ) + existing = _make_access_group_record(access_group_id="ag-update", access_group_name="old-name") mock_table.find_unique = AsyncMock(return_value=existing) - resp = client.put( - "/v1/access_group/ag-update", json={"access_group_name": "new-name"} - ) + resp = client.put("/v1/access_group/ag-update", json={"access_group_name": "new-name"}) assert resp.status_code == 200 mock_table.update.assert_awaited_once() call_kwargs = mock_table.update.call_args.kwargs @@ -594,19 +582,13 @@ def test_update_access_group_name_duplicate_conflict(client_and_mocks): """Update access_group_name to existing name returns 409 (unique constraint).""" client, _, mock_table, *_ = client_and_mocks - existing = _make_access_group_record( - access_group_id="ag-update", access_group_name="old-name" - ) + existing = _make_access_group_record(access_group_id="ag-update", access_group_name="old-name") mock_table.find_unique = AsyncMock(return_value=existing) mock_table.update = AsyncMock( - side_effect=Exception( - "Unique constraint failed on the fields: (`access_group_name`)" - ) + side_effect=Exception("Unique constraint failed on the fields: (`access_group_name`)") ) - resp = client.put( - "/v1/access_group/ag-update", json={"access_group_name": "taken-name"} - ) + resp = client.put("/v1/access_group/ag-update", json={"access_group_name": "taken-name"}) assert resp.status_code == 409 assert "already exists" in resp.json()["detail"] mock_table.update.assert_awaited_once() @@ -620,21 +602,15 @@ def test_update_access_group_name_duplicate_conflict(client_and_mocks): "unique constraint violation", ], ) -def test_update_access_group_name_unique_constraint_returns_409( - client_and_mocks, error_message -): +def test_update_access_group_name_unique_constraint_returns_409(client_and_mocks, error_message): """Update access_group_name: Prisma unique constraint surfaces as 409.""" client, _, mock_table, *_ = client_and_mocks - existing = _make_access_group_record( - access_group_id="ag-update", access_group_name="old-name" - ) + existing = _make_access_group_record(access_group_id="ag-update", access_group_name="old-name") mock_table.find_unique = AsyncMock(return_value=existing) mock_table.update = AsyncMock(side_effect=Exception(error_message)) - resp = client.put( - "/v1/access_group/ag-update", json={"access_group_name": "race-name"} - ) + resp = client.put("/v1/access_group/ag-update", json={"access_group_name": "race-name"}) assert resp.status_code == 409 assert "already exists" in resp.json()["detail"] @@ -690,9 +666,7 @@ def test_delete_access_group_forbidden_non_admin(client_and_mocks, user_role): def test_delete_access_group_cleans_up_teams_and_keys(client_and_mocks): """Delete removes access_group_id from teams and keys before deleting the group.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable mock_key_table = mock_prisma.db.litellm_verificationtoken @@ -722,10 +696,61 @@ def test_delete_access_group_cleans_up_teams_and_keys(client_and_mocks): where={"token": "key-token-1"}, data={"access_group_ids": []}, ) - mock_access_group_table.delete.assert_awaited_once_with( - where={"access_group_id": "ag-to-delete"} + mock_access_group_table.delete.assert_awaited_once_with(where={"access_group_id": "ag-to-delete"}) + + +def test_delete_access_group_detaches_group_from_agents(client_and_mocks): + """Delete strips the group from every agent that had it attached, so agents are not left + pointing at a group that no longer exists (which would deny them every model, server and agent).""" + client, mock_prisma, mock_access_group_table, _mock_cache, _mock_proxy_logging = client_and_mocks + mock_agents_table = mock_prisma.db.litellm_agentstable + + existing = _make_access_group_record(access_group_id="ag-to-delete") + mock_access_group_table.find_unique = AsyncMock(return_value=existing) + + agent_with_group = MagicMock() + agent_with_group.agent_id = "agent-1" + agent_with_group.access_group_ids = ["ag-keep", "ag-to-delete"] + mock_agents_table.find_many = AsyncMock(return_value=[agent_with_group]) + global_agent_registry.register_agent( + AgentResponse( + agent_id="agent-1", + agent_name="detach-test-agent", + agent_card_params={"name": "detach-test-agent", "url": "http://localhost:9", "version": "1"}, + access_group_ids=["ag-keep", "ag-to-delete"], + ) ) + try: + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 204 + + mock_agents_table.update.assert_awaited_once_with( + where={"agent_id": "agent-1"}, + data={"access_group_ids": ("ag-keep",)}, + ) + mock_access_group_table.delete.assert_awaited_once_with(where={"access_group_id": "ag-to-delete"}) + registered = global_agent_registry.get_agent_by_id("agent-1") + assert registered is not None + assert tuple(registered.access_group_ids or ()) == ("ag-keep",) + finally: + global_agent_registry.deregister_agent("detach-test-agent") + + +def test_delete_access_group_without_attached_agents_leaves_agents_untouched(client_and_mocks): + client, mock_prisma, mock_access_group_table, _mock_cache, _mock_proxy_logging = client_and_mocks + mock_agents_table = mock_prisma.db.litellm_agentstable + + mock_access_group_table.find_unique = AsyncMock( + return_value=_make_access_group_record(access_group_id="ag-to-delete") + ) + + resp = client.delete("/v1/access_group/ag-to-delete") + assert resp.status_code == 204 + + mock_agents_table.find_many.assert_awaited_once_with(where={"access_group_ids": {"hasSome": ("ag-to-delete",)}}) + mock_agents_table.update.assert_not_awaited() + @pytest.mark.parametrize( "team_cache_group_ids,key_cache_group_ids,expected_team_ids_after,expected_key_ids_after", @@ -792,9 +817,7 @@ def test_delete_access_group_patches_cached_team_and_key( """Delete patches cached team/key objects to remove the deleted access_group_id.""" from litellm.proxy._types import LiteLLM_TeamTableCachedObj - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable mock_key_table = mock_prisma.db.litellm_verificationtoken @@ -820,13 +843,9 @@ def test_delete_access_group_patches_cached_team_and_key( team_id="team-1", access_group_ids=list(team_cache_group_ids), ) - mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock( - return_value=cached_team - ) + mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=cached_team) else: - mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock( - return_value=None - ) + mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=None) # user_api_key_cache is queried both for teams (fallback after dual_cache) and # hashed keys — return the right stub per ``key``. A single AsyncMock(return_value=key) @@ -834,9 +853,7 @@ def test_delete_access_group_patches_cached_team_and_key( # Use a synchronous side_effect (not async def): AsyncMock awaits coroutine side_effects # inconsistently across Python/unittest versions; sync returns are awaited as immediate results. def user_cache_get_side_effect(*args, **kwargs): - cache_key = ( - kwargs.get("key") if "key" in kwargs else (args[0] if args else None) - ) + cache_key = kwargs.get("key") if "key" in kwargs else (args[0] if args else None) if cache_key == "team_id:team-1": if team_cache_group_ids is None: return None @@ -868,14 +885,11 @@ def test_delete_access_group_patches_cached_team_and_key( team_set_calls = [ c for c in mock_cache.async_set_cache.call_args_list - if c.kwargs.get("key", "") == "team_id:team-1" - or (len(c.args) >= 1 and c.args[0] == "team_id:team-1") + if c.kwargs.get("key", "") == "team_id:team-1" or (len(c.args) >= 1 and c.args[0] == "team_id:team-1") ] assert len(team_set_calls) >= 1, "Expected team cache to be patched" # The cached team object should have the updated access_group_ids - written_team = ( - team_set_calls[0].kwargs.get("value") or team_set_calls[0].args[1] - ) + written_team = team_set_calls[0].kwargs.get("value") or team_set_calls[0].args[1] if isinstance(written_team, LiteLLM_TeamTableCachedObj): assert written_team.access_group_ids == expected_team_ids_after else: @@ -883,8 +897,7 @@ def test_delete_access_group_patches_cached_team_and_key( team_set_calls = [ c for c in mock_cache.async_set_cache.call_args_list - if c.kwargs.get("key", "") == "team_id:team-1" - or (len(c.args) >= 1 and c.args[0] == "team_id:team-1") + if c.kwargs.get("key", "") == "team_id:team-1" or (len(c.args) >= 1 and c.args[0] == "team_id:team-1") ] assert len(team_set_calls) == 0, "Should not patch team cache when not cached" @@ -892,8 +905,7 @@ def test_delete_access_group_patches_cached_team_and_key( key_set_calls = [ c for c in mock_cache.async_set_cache.call_args_list - if c.kwargs.get("key", "") == "hashed-key-1" - or (len(c.args) >= 1 and c.args[0] == "hashed-key-1") + if c.kwargs.get("key", "") == "hashed-key-1" or (len(c.args) >= 1 and c.args[0] == "hashed-key-1") ] assert len(key_set_calls) >= 1, "Expected key cache to be patched" written_key = key_set_calls[0].kwargs.get("value") or key_set_calls[0].args[1] @@ -903,17 +915,14 @@ def test_delete_access_group_patches_cached_team_and_key( key_set_calls = [ c for c in mock_cache.async_set_cache.call_args_list - if c.kwargs.get("key", "") == "hashed-key-1" - or (len(c.args) >= 1 and c.args[0] == "hashed-key-1") + if c.kwargs.get("key", "") == "hashed-key-1" or (len(c.args) >= 1 and c.args[0] == "hashed-key-1") ] assert len(key_set_calls) == 0, "Should not patch key cache when not cached" def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): """Delete patches key cache — mock returns UserAPIKeyAuth (what UserApiKeyCache emits after deserialize).""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable mock_key_table = mock_prisma.db.litellm_verificationtoken @@ -929,9 +938,7 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): mock_key_table.find_unique = AsyncMock(return_value=key_with_group) # No team in cache - mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock( - return_value=None - ) + mock_proxy_logging.internal_usage_cache.dual_cache.async_get_cache = AsyncMock(return_value=None) # Serialized shape from Redis dict; UserApiKeyCache.async_get_cache(model_type=...) yields a model — simulate that. cached_key_payload = { @@ -940,18 +947,14 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): } def user_cache_get_dict_when_key_matches(*args, **kwargs): - cache_key = ( - kwargs.get("key") if "key" in kwargs else (args[0] if args else None) - ) + cache_key = kwargs.get("key") if "key" in kwargs else (args[0] if args else None) if cache_key == "team_id:team-1": return None if cache_key == "hashed-key-dict": return UserAPIKeyAuth.model_validate(cached_key_payload) return None - mock_cache.async_get_cache = AsyncMock( - side_effect=user_cache_get_dict_when_key_matches - ) + mock_cache.async_get_cache = AsyncMock(side_effect=user_cache_get_dict_when_key_matches) resp = client.delete("/v1/access_group/ag-to-delete") assert resp.status_code == 204 @@ -960,8 +963,7 @@ def test_delete_access_group_patches_key_cached_as_dict(client_and_mocks): key_set_calls = [ c for c in mock_cache.async_set_cache.call_args_list - if c.kwargs.get("key", "") == "hashed-key-dict" - or (len(c.args) >= 1 and c.args[0] == "hashed-key-dict") + if c.kwargs.get("key", "") == "hashed-key-dict" or (len(c.args) >= 1 and c.args[0] == "hashed-key-dict") ] assert len(key_set_calls) >= 1, "Expected key cache to be patched" written_key = key_set_calls[0].kwargs.get("value") or key_set_calls[0].args[1] @@ -988,9 +990,7 @@ def test_delete_access_group_404_on_p2025_or_record_not_found(client_and_mocks): existing = _make_access_group_record(access_group_id="ag-to-delete") mock_table.find_unique = AsyncMock(return_value=existing) - mock_table.delete = AsyncMock( - side_effect=Exception("P2025: Record to delete does not exist") - ) + mock_table.delete = AsyncMock(side_effect=Exception("P2025: Record to delete does not exist")) resp = client.delete("/v1/access_group/ag-to-delete") assert resp.status_code == 404 @@ -1039,9 +1039,7 @@ def test_delete_access_group_500_on_generic_exception(client_and_mocks): ("delete", "/v1/unified_access_group/ag-123", lambda: {}), ], ) -def test_access_group_endpoints_db_not_connected( - client_and_mocks, monkeypatch, method, url, factory -): +def test_access_group_endpoints_db_not_connected(client_and_mocks, monkeypatch, method, url, factory): """All endpoints return 500 when DB is not connected.""" client, *_ = client_and_mocks @@ -1049,9 +1047,7 @@ def test_access_group_endpoints_db_not_connected( resp = getattr(client, method)(url, **factory()) assert resp.status_code == 500 - assert ( - resp.json()["detail"]["error"] == CommonProxyErrors.db_not_connected_error.value - ) + assert resp.json()["detail"]["error"] == CommonProxyErrors.db_not_connected_error.value # --------------------------------------------------------------------------- @@ -1107,9 +1103,7 @@ def test_attached_team_ids_by_group_keeps_column_order_then_appends_unmirrored_t def test_create_access_group_syncs_assigned_teams(client_and_mocks): """Create adds access_group_id to each assigned team's access_group_ids in DB.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable team_record = _make_team_record("team-1") @@ -1132,9 +1126,7 @@ def test_create_access_group_syncs_assigned_teams(client_and_mocks): def test_create_access_group_syncs_assigned_keys(client_and_mocks): """Create adds access_group_id to each assigned key's access_group_ids in DB.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_key_table = mock_prisma.db.litellm_verificationtoken key_record = MagicMock() @@ -1148,9 +1140,7 @@ def test_create_access_group_syncs_assigned_keys(client_and_mocks): ) assert resp.status_code == 201 - mock_key_table.find_unique.assert_awaited_once_with( - where={"token": "hashed-token-1"} - ) + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "hashed-token-1"}) mock_key_table.update.assert_awaited_once() call_kwargs = mock_key_table.update.call_args.kwargs assert call_kwargs["where"] == {"token": "hashed-token-1"} @@ -1200,14 +1190,10 @@ def test_create_access_group_idempotent_team_sync(client_and_mocks): def test_update_access_group_syncs_added_teams(client_and_mocks): """Update adds access_group_id to newly assigned teams.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable - existing = _make_access_group_record( - access_group_id="ag-update", assigned_team_ids=["team-existing"] - ) + existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-existing"]) mock_access_group_table.find_unique = AsyncMock(return_value=existing) team_record = _make_team_record("team-new") @@ -1248,14 +1234,10 @@ def test_update_access_group_rejects_nonexistent_team(client_and_mocks): def test_update_access_group_syncs_removed_teams(client_and_mocks): """Update removes access_group_id from de-assigned teams.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable - existing = _make_access_group_record( - access_group_id="ag-update", assigned_team_ids=["team-keep", "team-remove"] - ) + existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-keep", "team-remove"]) mock_access_group_table.find_unique = AsyncMock(return_value=existing) team_to_remove = _make_team_record("team-remove", ["ag-update"]) @@ -1268,9 +1250,7 @@ def test_update_access_group_syncs_removed_teams(client_and_mocks): ) assert resp.status_code == 200 - mock_team_table.find_unique.assert_awaited_once_with( - where={"team_id": "team-remove"} - ) + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-remove"}) mock_team_table.update.assert_awaited_once() call_kwargs = mock_team_table.update.call_args.kwargs assert call_kwargs["where"] == {"team_id": "team-remove"} @@ -1296,19 +1276,15 @@ def test_update_access_group_detaches_team_the_mirror_missed(client_and_mocks): mock_team_table.update.assert_awaited_once() call_kwargs = mock_team_table.update.call_args.kwargs assert call_kwargs["where"] == {"team_id": "team-unmirrored"} - assert call_kwargs["data"]["access_group_ids"] == ["ag-other"] + assert tuple(call_kwargs["data"]["access_group_ids"]) == ("ag-other",) def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_mocks): """Update does not sync teams when assigned_team_ids is absent from the payload.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable - existing = _make_access_group_record( - access_group_id="ag-update", assigned_team_ids=["team-1"] - ) + existing = _make_access_group_record(access_group_id="ag-update", assigned_team_ids=["team-1"]) mock_access_group_table.find_unique = AsyncMock(return_value=existing) resp = client.put("/v1/access_group/ag-update", json={"description": "new desc"}) @@ -1320,14 +1296,10 @@ def test_update_access_group_no_team_sync_when_ids_not_in_payload(client_and_moc def test_update_access_group_syncs_added_keys(client_and_mocks): """Update adds access_group_id to newly assigned keys.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_key_table = mock_prisma.db.litellm_verificationtoken - existing = _make_access_group_record( - access_group_id="ag-update", assigned_key_ids=["old-token"] - ) + existing = _make_access_group_record(access_group_id="ag-update", assigned_key_ids=["old-token"]) mock_access_group_table.find_unique = AsyncMock(return_value=existing) key_record = MagicMock() @@ -1350,14 +1322,10 @@ def test_update_access_group_syncs_added_keys(client_and_mocks): def test_update_access_group_syncs_removed_keys(client_and_mocks): """Update removes access_group_id from de-assigned keys.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_key_table = mock_prisma.db.litellm_verificationtoken - existing = _make_access_group_record( - access_group_id="ag-update", assigned_key_ids=["keep-token", "remove-token"] - ) + existing = _make_access_group_record(access_group_id="ag-update", assigned_key_ids=["keep-token", "remove-token"]) mock_access_group_table.find_unique = AsyncMock(return_value=existing) key_to_remove = MagicMock() @@ -1385,9 +1353,7 @@ def test_update_access_group_syncs_removed_keys(client_and_mocks): def test_delete_access_group_handles_out_of_sync_assigned_teams(client_and_mocks): """Delete includes teams from assigned_team_ids even when not found by hasSome query.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_team_table = mock_prisma.db.litellm_teamtable # Access group has assigned_team_ids but the team's access_group_ids is not synced @@ -1409,18 +1375,14 @@ def test_delete_access_group_handles_out_of_sync_assigned_teams(client_and_mocks assert resp.status_code == 204 # find_unique is called for the out-of-sync team (included via union with assigned_team_ids) - mock_team_table.find_unique.assert_awaited_once_with( - where={"team_id": "team-out-of-sync"} - ) + mock_team_table.find_unique.assert_awaited_once_with(where={"team_id": "team-out-of-sync"}) # No update needed since team's access_group_ids doesn't contain "ag-to-delete" mock_team_table.update.assert_not_awaited() def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks): """Delete includes keys from assigned_key_ids even when not found by hasSome query.""" - client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = ( - client_and_mocks - ) + client, mock_prisma, mock_access_group_table, mock_cache, mock_proxy_logging = client_and_mocks mock_key_table = mock_prisma.db.litellm_verificationtoken existing = _make_access_group_record( @@ -1439,9 +1401,7 @@ def test_delete_access_group_handles_out_of_sync_assigned_keys(client_and_mocks) resp = client.delete("/v1/access_group/ag-to-delete") assert resp.status_code == 204 - mock_key_table.find_unique.assert_awaited_once_with( - where={"token": "token-out-of-sync"} - ) + mock_key_table.find_unique.assert_awaited_once_with(where={"token": "token-out-of-sync"}) mock_key_table.update.assert_not_awaited() @@ -1536,10 +1496,16 @@ def test_list_access_groups_resolves_names_with_one_query_per_table(client_and_m mock_table.find_many = AsyncMock( return_value=[ _make_access_group_record( - access_group_id="ag-1", access_mcp_server_ids=["mcp-a"], access_agent_ids=["agent-a"], assigned_key_ids=["key-a"] + access_group_id="ag-1", + access_mcp_server_ids=["mcp-a"], + access_agent_ids=["agent-a"], + assigned_key_ids=["key-a"], ), _make_access_group_record( - access_group_id="ag-2", access_mcp_server_ids=["mcp-b"], access_agent_ids=["agent-b"], assigned_key_ids=["key-b"] + access_group_id="ag-2", + access_mcp_server_ids=["mcp-b"], + access_agent_ids=["agent-b"], + assigned_key_ids=["key-b"], ), ] ) @@ -1573,7 +1539,10 @@ def test_list_access_groups_skips_lookups_when_nothing_to_resolve(client_and_moc """Groups with no MCP servers, agents or keys must not trigger an empty IN () query per table.""" client, mock_prisma, mock_table, *_ = client_and_mocks mock_table.find_many = AsyncMock( - return_value=[_make_access_group_record(access_group_id="ag-1"), _make_access_group_record(access_group_id="ag-2")] + return_value=[ + _make_access_group_record(access_group_id="ag-1"), + _make_access_group_record(access_group_id="ag-2"), + ] ) resp = client.get("/v1/access_group") diff --git a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py index 931531441d3..a282ae731dd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_auto_router_endpoints.py @@ -5,6 +5,7 @@ Unit tests for auto router management endpoints from collections.abc import Mapping, Sequence from functools import partial from pathlib import Path +from types import SimpleNamespace from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -781,17 +782,43 @@ class TestAutoRouterBenchmarks: assert _summed_agg_row([complexity, quality]).tier_turns == {} @pytest.mark.asyncio - async def test_non_admin_roles_cannot_read_benchmarks(self): + @pytest.mark.parametrize("user_id", [None, "own-user", "other-user"]) + async def test_non_admin_roles_cannot_read_benchmarks(self, user_id: str | None): from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks with pytest.raises(HTTPException) as err: await get_auto_router_benchmarks( - user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-x"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, api_key="sk-x", user_id="own-user" + ), start_date="2026-08-01", end_date="2026-08-02", + user_id=user_id, ) assert err.value.status_code == 403 + @pytest.mark.asyncio + async def test_an_empty_user_filter_is_rejected_before_querying_deployment_data( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + import httpx + from fastapi import FastAPI + + from litellm.proxy import proxy_server + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks + + query: Final = AsyncMock(return_value=[]) + monkeypatch.setattr(proxy_server, "prisma_client", SimpleNamespace(db=SimpleNamespace(query_raw=query))) + app: Final = FastAPI() + app.get("/auto_router/benchmarks")(get_auto_router_benchmarks) + app.dependency_overrides[user_api_key_auth] = lambda: ADMIN + async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: + response: Final = await client.get("/auto_router/benchmarks", params={"user_id": ""}) + + assert response.status_code == 422 + query.assert_not_awaited() + @pytest.mark.asyncio async def test_a_reversed_window_is_rejected(self, monkeypatch: pytest.MonkeyPatch): from litellm.proxy import proxy_server @@ -807,7 +834,11 @@ class TestAutoRouterBenchmarks: assert err.value.status_code == 400 @pytest.mark.asyncio - async def test_endpoint_returns_groups_and_totals_from_the_rollup(self, monkeypatch: pytest.MonkeyPatch): + @pytest.mark.parametrize("role", [LitellmUserRoles.PROXY_ADMIN, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY]) + @pytest.mark.parametrize("user_id", [None, "selected-user"]) + async def test_endpoint_returns_groups_and_totals_from_the_rollup( + self, monkeypatch: pytest.MonkeyPatch, role: LitellmUserRoles, user_id: str | None + ): from litellm.proxy import proxy_server from litellm.proxy.management_endpoints.auto_router_endpoints import get_auto_router_benchmarks @@ -822,12 +853,13 @@ class TestAutoRouterBenchmarks: monkeypatch.setattr(proxy_server, "prisma_client", type("P", (), {"db": _DB()})()) response = await get_auto_router_benchmarks( - user_api_key_dict=ADMIN, + user_api_key_dict=UserAPIKeyAuth(user_role=role, api_key="sk-admin", user_id="viewer"), start_date="2026-07-01", end_date="2026-08-01", api_key="key-hash", + user_id=user_id, ) - assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00", "key-hash") + assert captured["params"] == ("2026-07-01T00:00:00", "2026-08-02T00:00:00", "key-hash", user_id) assert response.routers_in_scope == 1 assert response.groups[0].router_name == "live-auto" assert response.groups[0].saved_pct == response.totals.saved_pct == 75.0 diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index 1d7f44b2099..95bf404d651 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -1285,11 +1285,11 @@ class TestUpdateModel: "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), ), - patch( + patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", side_effect=lambda value: value, ), - patch( + patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", new=AsyncMock( return_value=ReconcileOutcome(still_desired=None, live_after=None) @@ -4616,7 +4616,7 @@ class TestPatchModelBlockedAuthGate: "litellm.proxy.management_endpoints.model_management_endpoints.ModelManagementAuthChecks.can_user_make_model_call", new=AsyncMock(return_value=None), ), - patch( + patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", new=AsyncMock( return_value=ReconcileOutcome(still_desired=None, live_after=None) @@ -7285,3 +7285,337 @@ class TestTeamMemberAutoRouterWrites: assert json.loads(written["model_info"])["member_auto_router"] is True assert appended.await_args.kwargs["data"].models == ["new-personal-router"] assert appended.await_args.kwargs["data"].team_id == "member-team" + + +class TestModelManagementActorEdges: + @pytest.mark.asyncio + async def test_add_model_rejects_non_team_internal_user(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + actor: Final = UserAPIKeyAuth(user_id="internal-user", user_role=LitellmUserRoles.INTERNAL_USER) + prisma: Final = MagicMock() + deployment: Final = Deployment( + model_name="internal-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id="internal-model-id"), + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model(model_params=deployment, user_api_key_dict=actor) + + assert str(exc_info.value.code) == "403" + assert "permission" in str(exc_info.value).lower() + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + @pytest.mark.asyncio + async def test_add_model_rejects_proxy_admin_viewer(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + actor: Final = UserAPIKeyAuth( + user_id="view-only-user", user_role=LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY + ) + prisma: Final = MagicMock() + deployment: Final = Deployment( + model_name="view-only-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id="view-only-model-id"), + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model(model_params=deployment, user_api_key_dict=actor) + + assert str(exc_info.value.code) == "403" + assert "view-only" in str(exc_info.value).lower() + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + @pytest.mark.asyncio + async def test_add_model_requires_database_storage(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import add_new_model + + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + prisma: Final = MagicMock() + deployment: Final = Deployment( + model_name="database-disabled-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id="database-disabled-model-id"), + ) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", False), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] endpoint reads proxy-server state through its only test seam + ): + with pytest.raises(ProxyException) as exc_info: + await add_new_model(model_params=deployment, user_api_key_dict=actor) + + assert str(exc_info.value.code) == "500" + assert "STORE_MODEL_IN_DB" in str(exc_info.value) + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + @pytest.mark.asyncio + async def test_legacy_model_update_persists_changed_field(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_model + + model_id: Final = "legacy-update-model-id" + existing_row: Final = MagicMock() + existing_row.litellm_params = {"model": "openai/test-model", "timeout": 30} + existing_row.model_dump.return_value = { + "model_name": "legacy-update-model", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": model_id}, + } + existing_row.model_dump_json.return_value = "{}" + updated_row: Final = MagicMock() + updated_row.model_dump_json.return_value = "{}" + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) + router: Final = MagicMock() + router.get_model_ids.return_value = [model_id] + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value: value, + ), + patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams(timeout=42), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=actor, + ) + + written: Final = json.loads( + prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"] + ) + assert written["timeout"] == 42 + assert written["model"] == "openai/test-model" + + @pytest.mark.asyncio + async def test_legacy_model_update_explicit_null_preserves_existing_field(self): + from litellm.proxy.management_endpoints.model_management_endpoints import update_model + + model_id: Final = "legacy-null-model-id" + existing_row: Final = MagicMock() + existing_row.litellm_params = {"model": "openai/test-model", "timeout": 30} + existing_row.model_dump.return_value = { + "model_name": "legacy-null-model", + "litellm_params": existing_row.litellm_params, + "model_info": {"id": model_id}, + } + existing_row.model_dump_json.return_value = "{}" + updated_row: Final = MagicMock() + updated_row.model_dump_json.return_value = "{}" + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) + router: Final = MagicMock() + router.get_model_ids.return_value = [model_id] + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] update endpoint reads proxy-server state through its only test seam + patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value: value, + ), + patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + await update_model( + model_params=updateDeployment( + litellm_params=updateLiteLLMParams(timeout=None), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=actor, + ) + + written: Final = json.loads( + prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"] + ) + assert written["timeout"] == 30 + + @pytest.mark.asyncio + async def test_patch_model_rejects_config_file_model(self): + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.model_management_endpoints import patch_model + + model_id: Final = "config-model-id" + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_proxymodeltable.update = AsyncMock() + router: Final = MagicMock() + router.get_deployment.return_value = Deployment( + model_name="config-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id=model_id), + ) + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] patch endpoint reads proxy-server state through its only test seam + ): + with pytest.raises(ProxyException) as exc_info: + await patch_model( + model_id=model_id, + patch_data=updateDeployment( + litellm_params=updateLiteLLMParams(timeout=42), + model_info=ModelInfo(id=model_id), + ), + user_api_key_dict=actor, + ) + + assert str(exc_info.value.code) == "400" + assert "Cannot edit config-based model" in str(exc_info.value) + prisma.db.litellm_proxymodeltable.update.assert_not_awaited() + + @contextlib.contextmanager + def _client_for(self, actor: UserAPIKeyAuth) -> Iterator[TestClient]: + import litellm.proxy.proxy_server as proxy_server + from litellm.proxy.proxy_server import app + + app.dependency_overrides[proxy_server.user_api_key_auth] = lambda: actor + try: + yield TestClient(app) + finally: + app.dependency_overrides.pop(proxy_server.user_api_key_auth, None) + + def test_post_model_new_binds_to_actor_guard(self): + actor: Final = UserAPIKeyAuth(user_id="internal-user", user_role=LitellmUserRoles.INTERNAL_USER) + prisma: Final = MagicMock() + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + self._client_for(actor) as client, + ): + response: Final = client.post( + "/model/new", + json={ + "model_name": "internal-model", + "litellm_params": {"model": "openai/test-model"}, + "model_info": {"id": "internal-model-id"}, + }, + ) + + assert response.status_code == 403 + assert "permission" in response.text.lower() + prisma.db.litellm_proxymodeltable.create.assert_not_called() + + def test_post_legacy_model_update_binds_to_persistence(self): + model_id: Final = "legacy-route-model-id" + existing_row: Final = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name="legacy-route-model", + litellm_params={"model": "openai/test-model", "timeout": 30}, + model_info={"id": model_id}, + created_by="admin", + updated_by="admin", + ) + updated_row: Final = LiteLLM_ProxyModelTable( + model_id=model_id, + model_name="legacy-route-model", + litellm_params={"model": "openai/test-model", "timeout": 42}, + model_info={"id": model_id}, + created_by="admin", + updated_by="admin", + ) + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=existing_row) + prisma.db.litellm_proxymodeltable.update = AsyncMock(return_value=updated_row) + router: Final = MagicMock() + router.get_model_ids.return_value = [model_id] + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch( # test-quality-ok: [TQ008] isolate persistence from encryption implementation + "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper", + side_effect=lambda value: value, + ), + patch( # test-quality-ok: [TQ008] isolate persistence from router reload implementation + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + patch( # test-quality-ok: [TQ008] audit logging is outside the persistence contract + "litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log", + new=AsyncMock(return_value=None), + ), + self._client_for(actor) as client, + ): + response: Final = client.post( + "/model/update", + json={ + "litellm_params": {"timeout": 42}, + "model_info": {"id": model_id}, + }, + ) + + assert response.status_code == 200, response.text + written: Final = json.loads( + prisma.db.litellm_proxymodeltable.update.await_args.kwargs["data"]["litellm_params"] + ) + assert written["timeout"] == 42 + + def test_patch_config_model_binds_to_patch_route(self): + model_id: Final = "config-route-model-id" + prisma: Final = MagicMock() + prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=None) + prisma.db.litellm_proxymodeltable.update = AsyncMock() + router: Final = MagicMock() + router.get_deployment.return_value = Deployment( + model_name="config-route-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id=model_id), + ) + actor: Final = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN) + with ( + patch("litellm.proxy.proxy_server.prisma_client", prisma), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.llm_router", router), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.store_model_in_db", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + patch("litellm.proxy.proxy_server.premium_user", True), # test-quality-ok: [TQ008] route reads proxy-server state through its only test seam + self._client_for(actor) as client, + ): + response: Final = client.patch( + f"/model/{model_id}/update", + json={ + "litellm_params": {"timeout": 42}, + "model_info": {"id": model_id}, + }, + ) + + assert response.status_code == 400 + assert "Cannot edit config-based model" in response.text + prisma.db.litellm_proxymodeltable.update.assert_not_awaited() diff --git a/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py index bdfaa00b5e5..74ed43041a3 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_router_settings_endpoints.py @@ -4,14 +4,15 @@ Tests for router settings management endpoints. Tests the GET endpoints for router settings and router fields. """ -from unittest.mock import AsyncMock, MagicMock, patch +from collections.abc import Mapping +from typing import Any, Final import pytest from fastapi.testclient import TestClient - from litellm.proxy import proxy_server -from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth, ProxyRuntimeConfig +from litellm.proxy._types import LitellmUserRoles, ProxyRuntimeConfig, UserAPIKeyAuth +from litellm.proxy.config_resolvers import SettingsStore from litellm.proxy.management_endpoints.router_settings_endpoints import ( get_router_settings, ) @@ -21,6 +22,16 @@ from litellm.router import Router client = TestClient(app) +class _StubProxyConfig: + def __init__(self, router_settings: SettingsStore, config_router_settings: Mapping[str, Any]) -> None: + self.router_settings: Final = router_settings + self._config_router_settings: Final = dict(config_router_settings) + + async def get_config(self, config_file_path: str | None = None) -> ProxyRuntimeConfig: + del config_file_path + return ProxyRuntimeConfig.from_resolved({"router_settings": dict(self._config_router_settings)}) + + class TestRouterSettingsEndpoints: """Test suite for router settings endpoints""" @@ -75,6 +86,31 @@ class TestRouterSettingsEndpoints: assert isinstance(routing_strategy_field["options"], list) assert len(routing_strategy_field["options"]) > 0 + @pytest.mark.asyncio + async def test_get_router_settings_reports_sources( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + store = SettingsStore("router_settings") + store.load_yaml({"routing_strategy": "simple-shuffle"}) + store.apply_db_row("router_settings", {"num_retries": 3}) + monkeypatch.setattr( + proxy_server, + "proxy_config", + _StubProxyConfig( + store, + {"routing_strategy": "simple-shuffle", "num_retries": 3}, + ), + ) + monkeypatch.setattr(proxy_server, "llm_router", None) + + admin_user = UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-x" + ) + response = await get_router_settings(user_api_key_dict=admin_user) + + assert response.source["routing_strategy"] == "config" + assert response.source["num_retries"] == "db" + @pytest.mark.asyncio async def test_get_router_settings_includes_routing_groups_from_live_router( self, monkeypatch @@ -102,12 +138,10 @@ class TestRouterSettingsEndpoints: ) monkeypatch.setattr(proxy_server, "llm_router", llm_router) - - async def fake_get_config(self, config_file_path=None): - return ProxyRuntimeConfig.from_resolved({}) - monkeypatch.setattr( - proxy_server.ProxyConfig, "get_config", fake_get_config, raising=True + proxy_server, + "proxy_config", + _StubProxyConfig(SettingsStore("router_settings"), {}), ) admin_user = UserAPIKeyAuth( @@ -116,6 +150,8 @@ class TestRouterSettingsEndpoints: response = await get_router_settings(user_api_key_dict=admin_user) assert response.current_values.get("routing_groups") == groups + assert response.current_values["timeout"] is not None + assert response.source["timeout"] == "default" rg_field = next(f for f in response.fields if f.field_name == "routing_groups") assert rg_field.field_value == groups diff --git a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py index 246e2cbba54..b5536b7618c 100644 --- a/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py +++ b/tests/test_litellm/proxy/proxy_server/test_routes_model_metrics.py @@ -11,13 +11,18 @@ Pins (PR2): from __future__ import annotations +from collections.abc import Callable, Mapping +from contextlib import AbstractContextManager from unittest.mock import AsyncMock, MagicMock import pytest +from fastapi.testclient import TestClient import litellm from litellm.proxy import proxy_server from litellm.proxy._types import LitellmUserRoles +from litellm.proxy.config_resolvers.settings_rules import JsonValue +from litellm.proxy.config_resolvers.settings_store import SettingsStore from .conftest import normalize # type: ignore[import-not-found] @@ -179,6 +184,177 @@ def test_model_settings_method_not_allowed(client, auth_as): # --------------------------------------------------------------------------- +def _alerting_client( + monkeypatch: pytest.MonkeyPatch, + *, + yaml_values: Mapping[str, JsonValue], + db_row: Mapping[str, JsonValue], + live_args: Mapping[str, JsonValue], +) -> "SettingsStore": + pc = MagicMock() + row = MagicMock() + row.param_value = db_row + pc.db.litellm_config.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(proxy_server, "prisma_client", pc) + + logging_obj = MagicMock() + args_model = MagicMock() + args_model.model_dump = MagicMock(return_value=live_args) + logging_obj.slack_alerting_instance.alerting_args = args_model + monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj) + + store = SettingsStore("general_settings") + store.load_yaml(yaml_values) + store.apply_db_row("general_settings", db_row) + monkeypatch.setattr(proxy_server.proxy_config, "settings", store) + monkeypatch.setattr(proxy_server, "general_settings", store) + return store + + +def test_alerting_settings_reports_sources( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + _alerting_client( + monkeypatch, + yaml_values={ + "alerting": ["slack"], + "alerting_args": {"daily_report_frequency": 3, "report_check_interval": 300}, + }, + db_row={"alerting_args": {"daily_report_frequency": 7, "outage_alert_ttl": 4242}}, + live_args={"daily_report_frequency": 3}, + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/alerting/settings") + + assert response.status_code == 200 + by_name = {entry["field_name"]: entry for entry in response.json()} + + assert by_name["slack_alerting"]["source"] == "config" + assert by_name["daily_report_frequency"]["source"] == "config" + assert by_name["report_check_interval"]["source"] == "config" + assert by_name["outage_alert_ttl"]["source"] == "default" + assert by_name["budget_alert_ttl"]["source"] == "default" + + +def test_alerting_settings_reports_db_source_when_the_file_omits_alerting_args( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + store = _alerting_client( + monkeypatch, + yaml_values={"alerting": ["slack"]}, + db_row={ + "alerting_args": { + "outage_alert_ttl": 4242, + "region_outage_alert_ttl": [], + "report_check_interval": None, + } + }, + live_args={"outage_alert_ttl": 4242}, + ) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/alerting/settings") + + assert response.status_code == 200 + by_name = {entry["field_name"]: entry for entry in response.json()} + + assert store.source("alerting_args") == "db" + assert by_name["outage_alert_ttl"]["source"] == "db" + assert by_name["region_outage_alert_ttl"]["source"] == "db" + assert by_name["report_check_interval"]["source"] == "db" + assert by_name["budget_alert_ttl"]["source"] == "default" + + +def test_alerting_settings_reports_config_source_when_db_disagrees( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.config_resolvers import SettingsStore + + db_alerting_args = {"daily_report_frequency": 7} + + pc = MagicMock() + row = MagicMock() + row.param_value = {"alerting_args": db_alerting_args} + pc.db.litellm_config.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(proxy_server, "prisma_client", pc) + + logging_obj = MagicMock() + args_model = MagicMock() + args_model.model_dump = MagicMock(return_value={"daily_report_frequency": 3}) + logging_obj.slack_alerting_instance.alerting_args = args_model + monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj) + + store = SettingsStore("general_settings") + store.load_yaml({"alerting_args": {"daily_report_frequency": 3}}) + store.apply_db_row("general_settings", {"alerting_args": db_alerting_args}) + monkeypatch.setattr(proxy_server.proxy_config, "settings", store) + monkeypatch.setattr(proxy_server, "general_settings", store) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/alerting/settings") + + assert response.status_code == 200 + by_name = {entry["field_name"]: entry for entry in response.json()} + assert store.source("alerting_args") == "config" + assert by_name["daily_report_frequency"]["field_value"] == 3 + assert by_name["daily_report_frequency"]["source"] == "config" + + +@pytest.mark.parametrize("db_alerting_args", [None, []]) +def test_alerting_settings_handles_empty_db_args( + client: TestClient, + auth_as: Callable[..., AbstractContextManager[None]], + monkeypatch: pytest.MonkeyPatch, + db_alerting_args: JsonValue, +) -> None: + from litellm.proxy.config_resolvers import SettingsStore + + pc = MagicMock() + row = MagicMock() + row.param_value = {"alerting_args": db_alerting_args} + pc.db.litellm_config.find_first = AsyncMock(return_value=row) + monkeypatch.setattr(proxy_server, "prisma_client", pc) + + logging_obj = MagicMock() + args_model = MagicMock() + args_model.model_dump = MagicMock(return_value={}) + logging_obj.slack_alerting_instance.alerting_args = args_model + monkeypatch.setattr(proxy_server, "proxy_logging_obj", logging_obj) + + store = SettingsStore("general_settings") + store.load_yaml({"alerting_args": {"report_check_interval": 300}}) + monkeypatch.setattr(proxy_server.proxy_config, "settings", store) + monkeypatch.setattr(proxy_server, "general_settings", store) + + with auth_as(LitellmUserRoles.PROXY_ADMIN): + response = client.get("/alerting/settings") + + assert response.status_code == 200 + by_name = {entry["field_name"]: entry for entry in response.json()} + assert by_name["report_check_interval"]["source"] == "config" + assert by_name["budget_alert_ttl"]["source"] == "default" + + +@pytest.mark.parametrize( + ("field_default", "expected"), + [(43200, "default"), (None, "unset")], +) +def test_nested_setting_source_without_a_config_or_db_value(field_default: JsonValue, expected: str) -> None: + store = SettingsStore("general_settings") + store.load_yaml({}) + + assert ( + proxy_server._nested_setting_source(store, {}, "alerting_args", "budget_alert_ttl", field_default) == expected + ) + + def test_alerting_settings_no_db_error(client, auth_as, no_prisma): """Pins ``GET /alerting/settings`` (error: db not connected).""" with auth_as(LitellmUserRoles.PROXY_ADMIN): diff --git a/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py new file mode 100644 index 00000000000..fbce38db39f --- /dev/null +++ b/tests/test_litellm/proxy/shutdown/test_scheduled_jobs.py @@ -0,0 +1,158 @@ +import asyncio +import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager +from datetime import datetime, timedelta + +import pytest +from apscheduler.schedulers.asyncio import AsyncIOScheduler + +from litellm.proxy.shutdown.scheduled_jobs import ( + AwaitableAsyncIOExecutor, + pause_scheduled_jobs, + stop_in_flight_scheduler_jobs, +) + + +class _Job: + """A scheduled job that blocks until cancelled, or for ``work_seconds``, and records what it observed""" + + def __init__(self, swallow_cancellation: bool = False, work_seconds: float | None = None) -> None: + self.started = asyncio.Event() + self.events: list[str] = [] + self.swallow_cancellation = swallow_cancellation + self.work_seconds = work_seconds + + async def run(self) -> None: + self.started.set() + try: + if self.work_seconds is None: + await asyncio.Event().wait() + else: + await asyncio.sleep(self.work_seconds) + self.events.append("committed") + except asyncio.CancelledError: + self.events.append("cancelled") + if self.swallow_cancellation: + await asyncio.Event().wait() + raise + finally: + self.events.append("finished") + + +@asynccontextmanager +async def _running_scheduler(*jobs: _Job) -> AsyncIterator[tuple[AsyncIOScheduler, AwaitableAsyncIOExecutor]]: + """A started scheduler with every job in flight, stopped on the way out whatever the test did""" + executor = AwaitableAsyncIOExecutor() + scheduler = AsyncIOScheduler(executors={"default": executor}) + for index, job in enumerate(jobs): + scheduler.add_job(job.run, id=f"job-{index}", next_run_time=datetime.now()) + scheduler.start() + try: + for job in jobs: + await asyncio.wait_for(job.started.wait(), timeout=5) + yield scheduler, executor + finally: + if scheduler.running: + scheduler.shutdown(wait=False) + stragglers = executor.in_flight_jobs() + for straggler in stragglers: + straggler.cancel() + await asyncio.gather(*stragglers, return_exceptions=True) + + +@pytest.mark.asyncio +async def test_in_flight_jobs_observe_cancellation_before_shutdown_returns(): + """The job's own CancelledError handler records how a run ended, so shutdown must wait for it""" + job = _Job() + async with _running_scheduler(job) as (scheduler, executor): + await stop_in_flight_scheduler_jobs(scheduler, executor) + + assert job.events == ["cancelled", "finished"] + assert scheduler.running is False + assert executor.in_flight_jobs() == () + + +@pytest.mark.asyncio +async def test_a_job_that_is_finishing_is_allowed_to_finish_rather_than_cancelled(): + """A spend write cancelled mid-commit drops the rows it popped, so short jobs get to finish first""" + write = _Job(work_seconds=0.2) + stuck = _Job() + async with _running_scheduler(write, stuck) as (scheduler, executor): + await stop_in_flight_scheduler_jobs(scheduler, executor, finish_timeout_seconds=2.0) + + assert write.events == ["committed", "finished"] + assert stuck.events == ["cancelled", "finished"] + assert scheduler.running is False + + +@pytest.mark.asyncio +async def test_every_in_flight_job_is_cancelled_not_only_the_first(): + first, second = _Job(), _Job() + async with _running_scheduler(first, second) as (scheduler, executor): + await stop_in_flight_scheduler_jobs(scheduler, executor) + + assert first.events == ["cancelled", "finished"] + assert second.events == ["cancelled", "finished"] + + +@pytest.mark.asyncio +async def test_a_job_that_ignores_cancellation_is_abandoned_after_the_timeout(caplog): + """A job that swallows CancelledError must not hold the pod past its termination grace period""" + job = _Job(swallow_cancellation=True) + async with _running_scheduler(job) as (scheduler, executor): + with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): + await stop_in_flight_scheduler_jobs(scheduler, executor, cancel_timeout_seconds=0.05) + + assert job.events == ["cancelled"] + assert "1 scheduled job(s) did not finish within 0.05s of cancellation" in caplog.text + + +@pytest.mark.asyncio +async def test_shutdown_with_nothing_in_flight_still_stops_the_scheduler(): + async with _running_scheduler() as (scheduler, executor): + await stop_in_flight_scheduler_jobs(scheduler, executor) + await asyncio.sleep(0) + + assert scheduler.running is False + + +@pytest.mark.asyncio +async def test_a_scheduler_that_never_started_is_left_alone(): + """The proxy runs without a scheduler when it has no database""" + executor = AwaitableAsyncIOExecutor() + scheduler = AsyncIOScheduler(executors={"default": executor}) + + await stop_in_flight_scheduler_jobs(scheduler, executor) + + assert scheduler.running is False + + +@pytest.mark.asyncio +async def test_pausing_stops_new_jobs_from_starting_but_leaves_running_ones_alone(): + """A job due during the shutdown drain would only be cancelled, so it must not start at all""" + running = _Job() + async with _running_scheduler(running) as (scheduler, executor): + late = _Job() + scheduler.add_job(late.run, id="late", next_run_time=datetime.now() + timedelta(seconds=0.1)) + + pause_scheduled_jobs(scheduler) + await asyncio.sleep(0.3) + + assert late.started.is_set() is False + assert running.events == [] + assert scheduler.running is True + + await stop_in_flight_scheduler_jobs(scheduler, executor) + + assert running.events == ["cancelled", "finished"] + assert late.started.is_set() is False + + +@pytest.mark.asyncio +async def test_pausing_a_scheduler_that_never_started_is_a_no_op(): + scheduler = AsyncIOScheduler(executors={"default": AwaitableAsyncIOExecutor()}) + + pause_scheduled_jobs(scheduler) + + assert scheduler.running is False diff --git a/tests/test_litellm/proxy/test_spend_log_cleanup.py b/tests/test_litellm/proxy/test_spend_log_cleanup.py index bf1538183ab..e333da03950 100644 --- a/tests/test_litellm/proxy/test_spend_log_cleanup.py +++ b/tests/test_litellm/proxy/test_spend_log_cleanup.py @@ -7,6 +7,7 @@ import math import time from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest @@ -793,18 +794,20 @@ async def test_spend_logs_retention_alone_does_not_touch_the_session_rollup(): tables = [call[0][0] for call in client.db.execute_raw.call_args_list] assert any('"LiteLLM_SpendLogs"' in sql for sql in tables) assert not any('"LiteLLM_AutoRouterSession"' in sql for sql in tables) + assert not any('"LiteLLM_AutoRouterUserSession"' in sql for sql in tables) assert not any('"LiteLLM_HealthCheckTable"' in sql for sql in tables) @pytest.mark.asyncio -async def test_session_retention_alone_cleans_only_the_session_rollup(): - client = _mock_prisma_for_retention([0]) +async def test_session_retention_alone_cleans_both_session_rollups(): + client = _mock_prisma_for_retention([0, 0]) cleaner = SpendLogCleanup(general_settings={"maximum_autorouter_session_retention_period": "365d"}) cleaner.pod_lock_manager = None await cleaner.cleanup_old_spend_logs(client) tables = [call[0][0] for call in client.db.execute_raw.call_args_list] - assert len(tables) == 1 + assert len(tables) == 2 assert '"LiteLLM_AutoRouterSession"' in tables[0] + assert '"LiteLLM_AutoRouterUserSession"' in tables[1] @pytest.mark.asyncio @@ -825,7 +828,7 @@ async def test_health_check_retention_alone_cleans_only_the_health_check_table() @pytest.mark.asyncio async def test_each_retention_key_cuts_off_at_its_own_horizon(): - client = _mock_prisma_for_retention([0, 0, 0, 0]) + client = _mock_prisma_for_retention([0, 0, 0, 0, 0]) cleaner = SpendLogCleanup( general_settings={ "maximum_spend_logs_retention_period": "7d", @@ -839,6 +842,8 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon(): ( "LiteLLM_AutoRouterSession" if '"LiteLLM_AutoRouterSession"' in call[0][0] + else "LiteLLM_AutoRouterUserSession" + if '"LiteLLM_AutoRouterUserSession"' in call[0][0] else "LiteLLM_HealthCheckTable" if '"LiteLLM_HealthCheckTable"' in call[0][0] else "logs" @@ -848,6 +853,7 @@ async def test_each_retention_key_cuts_off_at_its_own_horizon(): now = datetime.now(timezone.utc) assert (now - cutoffs["logs"]).days == 7 assert (now - cutoffs["LiteLLM_AutoRouterSession"]).days == 365 + assert cutoffs["LiteLLM_AutoRouterUserSession"] == cutoffs["LiteLLM_AutoRouterSession"] assert (now - cutoffs["LiteLLM_HealthCheckTable"]).days == 30 @@ -1417,3 +1423,125 @@ def test_the_reported_run_outcome_is_the_most_significant_reason_in_any_order(st """ results = tuple(TableCleanupResult(rows_deleted=0, stop_reason=reason) for reason in stop_reasons) assert SpendLogCleanup._run_outcome(results) == expected + + +_OTHER_OUTCOMES: Final = ("completed", "budget_exhausted", "batch_cap_reached", "skipped_locked", "skipped_disabled") + + +def _runs_recorded(outcome: str) -> float: + """The real ``litellm_spend_log_cleanup_runs_total`` sample for one outcome, 0 when unset""" + from prometheus_client import REGISTRY + + return REGISTRY.get_sample_value("litellm_spend_log_cleanup_runs_total", {"outcome": outcome}) or 0.0 + + +@pytest.mark.asyncio +async def test_a_cancelled_run_records_aborted_and_logs_its_progress_before_re_raising(monkeypatch): + """A run cut short by shutdown must leave its outcome and how far it got behind""" + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + mock_logger = MagicMock() + monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger) + aborted_runs_before = _runs_recorded("aborted") + other_runs_before = {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES} + + third_batch_reached = asyncio.Event() + + async def _execute_raw(sql, *args): + if third_batch_reached.is_set(): + raise AssertionError("no batch may be issued after the cancelled one") + if _execute_raw.calls < 2: + _execute_raw.calls += 1 + return 150 + third_batch_reached.set() + await asyncio.Event().wait() + + _execute_raw.calls = 0 + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + mock_prisma_client.db.execute_raw = _execute_raw + + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = MagicMock() + cleaner.pod_lock_manager.redis_cache = MagicMock() + cleaner.pod_lock_manager.acquire_lock = AsyncMock(return_value=True) + cleaner.pod_lock_manager.release_lock = AsyncMock() + + run = asyncio.ensure_future(cleaner.cleanup_old_spend_logs(mock_prisma_client)) + await asyncio.wait_for(third_batch_reached.wait(), timeout=5) + run.cancel() + with pytest.raises(asyncio.CancelledError): + await run + + assert _runs_recorded("aborted") == aborted_runs_before + 1 + assert {outcome: _runs_recorded(outcome) for outcome in _OTHER_OUTCOMES} == other_runs_before + cleaner.pod_lock_manager.release_lock.assert_awaited_once() + mock_logger.exception.assert_not_called() + (error_call,) = mock_logger.error.call_args_list + rendered = error_call[0][0] % error_call[0][1:] + assert rendered.startswith("Spend log cleanup cancelled after ") + assert "s (rows_deleted=300, batches=2)" in rendered + + +@pytest.mark.asyncio +async def test_progress_reported_for_a_cancelled_run_is_that_run_only(monkeypatch): + """The scheduler holds one cleaner for the life of the process, so progress must not carry over""" + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + mock_logger = MagicMock() + monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger) + + mock_prisma_client = MagicMock() + _wire_tx(mock_prisma_client.db) + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, 0, 0]) + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = None + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + mock_prisma_client.db.execute_raw = AsyncMock(side_effect=[150, asyncio.CancelledError()]) + with pytest.raises(asyncio.CancelledError): + await cleaner.cleanup_old_spend_logs(mock_prisma_client) + + (error_call,) = mock_logger.error.call_args_list + rendered = error_call[0][0] % error_call[0][1:] + assert "(rows_deleted=150, batches=1)" in rendered + + +@pytest.mark.asyncio +async def test_progress_reported_by_an_overlapping_run_is_its_own(monkeypatch): + """With APSCHEDULER_MAX_INSTANCES above one, two runs share the cleaner but not their progress""" + import litellm.proxy.db.db_transaction_queue.spend_log_cleanup as cleanup_module + + mock_logger = MagicMock() + monkeypatch.setattr(cleanup_module, "verbose_proxy_logger", mock_logger) + + first_batch_done = asyncio.Event() + second_run_done = asyncio.Event() + + async def _slow_execute_raw(sql, *args): + first_batch_done.set() + await second_run_done.wait() + return 100 + + slow_client = MagicMock() + _wire_tx(slow_client.db) + slow_client.db.execute_raw = _slow_execute_raw + fast_client = MagicMock() + _wire_tx(fast_client.db) + fast_client.db.execute_raw = AsyncMock(side_effect=[150, 150, 0, 0]) + + cleaner = SpendLogCleanup(general_settings={"maximum_spend_logs_retention_period": "7d"}) + cleaner.pod_lock_manager = None + + slow_run = asyncio.ensure_future(cleaner.cleanup_old_spend_logs(slow_client)) + await asyncio.wait_for(first_batch_done.wait(), timeout=5) + await cleaner.cleanup_old_spend_logs(fast_client) + second_run_done.set() + await asyncio.sleep(0) + slow_run.cancel() + with pytest.raises(asyncio.CancelledError): + await slow_run + + (error_call,) = mock_logger.error.call_args_list + rendered = error_call[0][0] % error_call[0][1:] + assert "(rows_deleted=100, batches=1)" in rendered diff --git a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py index c562646a881..ff271baee9b 100644 --- a/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py +++ b/tests/test_litellm/proxy/ui_crud_endpoints/test_proxy_setting_endpoints.py @@ -1342,6 +1342,48 @@ class TestProxySettingEndpoints: where={"id": "ui_settings"} ) + def test_get_ui_settings_reports_sources(self, monkeypatch: pytest.MonkeyPatch) -> None: + from unittest.mock import AsyncMock, MagicMock + + from litellm.proxy import proxy_server + from litellm.proxy.config_resolvers import SettingsStore + + mock_prisma = MagicMock() + mock_db_record = MagicMock() + mock_db_record.ui_settings = { + "disable_model_add_for_internal_users": True, + "require_auth_for_public_ai_hub": True, + } + mock_prisma.db.litellm_uisettings.find_unique = AsyncMock( + return_value=mock_db_record + ) + monkeypatch.setattr(proxy_server, "prisma_client", mock_prisma) + + store = SettingsStore("general_settings") + store.load_yaml( + { + "disable_model_add_for_internal_users": False, + "forward_client_headers_to_llm_api": True, + } + ) + store.apply_db_row( + "ui_settings", + {"disable_model_add_for_internal_users": True}, + ) + monkeypatch.setattr(proxy_server.proxy_config, "settings", store) + monkeypatch.setattr(proxy_server, "general_settings", store) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + data = response.json() + assert data["values"]["disable_model_add_for_internal_users"] is False + assert data["values"]["forward_client_headers_to_llm_api"] is True + assert data["values"]["require_auth_for_public_ai_hub"] is True + assert data["source"]["disable_model_add_for_internal_users"] == "config" + assert data["source"]["forward_client_headers_to_llm_api"] == "config" + assert data["source"]["require_auth_for_public_ai_hub"] == "db" + def test_get_ui_settings_schema_description_preserved_with_extensions( self, mock_auth, monkeypatch ): @@ -3482,6 +3524,7 @@ class TestPtuCostAttributionUISetting: assert response.status_code == 200 assert response.json()["values"]["enable_ptu_cost_attribution"] is False + assert response.json()["source"]["enable_ptu_cost_attribution"] == "default" def test_reported_true_once_the_env_var_is_set(self, mock_auth, monkeypatch): from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR @@ -3493,6 +3536,47 @@ class TestPtuCostAttributionUISetting: assert response.status_code == 200 assert response.json()["values"]["enable_ptu_cost_attribution"] is True + assert response.json()["source"]["enable_ptu_cost_attribution"] == "config" + + def test_reported_config_when_secret_manager_enables_the_flag( + self, mock_auth: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + monkeypatch.setattr( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.is_ptu_cost_attribution_enabled", + lambda: True, + ) + self._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_ptu_cost_attribution"] is True + assert response.json()["source"]["enable_ptu_cost_attribution"] == "config" + + def test_reported_config_when_secret_manager_disables_the_flag( + self, mock_auth: None, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.proxy.spend_tracking.ptu_feature_flag import PTU_COST_ATTRIBUTION_ENV_VAR + + monkeypatch.delenv(PTU_COST_ATTRIBUTION_ENV_VAR, raising=False) + monkeypatch.setattr( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.is_ptu_cost_attribution_enabled", + lambda: False, + ) + monkeypatch.setattr( + "litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints.get_secret", + lambda *_args: False, + ) + self._mock_prisma(monkeypatch) + + response = client.get("/get/ui_settings") + + assert response.status_code == 200 + assert response.json()["values"]["enable_ptu_cost_attribution"] is False + assert response.json()["source"]["enable_ptu_cost_attribution"] == "config" def test_a_persisted_true_cannot_forge_the_derived_value(self, mock_auth, monkeypatch): """A row written before the allowlist existed must not be able to turn the feature on.""" diff --git a/tests/test_litellm/proxy/utils/helpers/test_model_access.py b/tests/test_litellm/proxy/utils/helpers/test_model_access.py index 5fb4392eec6..7f77f938323 100644 --- a/tests/test_litellm/proxy/utils/helpers/test_model_access.py +++ b/tests/test_litellm/proxy/utils/helpers/test_model_access.py @@ -110,8 +110,7 @@ def test_create_model_info_response_happy_path_no_metadata(): "owned_by": result["owned_by"], "created_is_int": isinstance(result["created"], int), "metadata_absent": "metadata" not in result, - "max_input_tokens_positive_int": isinstance(result["max_input_tokens"], int) - and result["max_input_tokens"] > 0, + "max_input_tokens_positive_int": isinstance(result["max_input_tokens"], int) and result["max_input_tokens"] > 0, "max_output_tokens_positive_int": isinstance(result["max_output_tokens"], int) and result["max_output_tokens"] > 0, } @@ -205,9 +204,7 @@ def test_validate_model_access_happy_path_single_model_in_list(): def test_validate_model_access_happy_path_batch_all_accessible(): summary = { - "result": validate_model_access( - "gpt-4o,claude-haiku", ["gpt-4o", "claude-haiku", "gemini"] - ), + "result": validate_model_access("gpt-4o,claude-haiku", ["gpt-4o", "claude-haiku", "gemini"]), "input": "gpt-4o,claude-haiku", "available": ["gpt-4o", "claude-haiku", "gemini"], } @@ -389,9 +386,7 @@ async def test_get_available_models_for_user_error_path_complete_list_raises( def _boom(**_kwargs): raise RuntimeError("downstream failure") - monkeypatch.setattr( - "litellm.proxy.auth.model_checks.get_complete_model_list", _boom - ) + monkeypatch.setattr("litellm.proxy.auth.model_checks.get_complete_model_list", _boom) user_api_key_dict = UserAPIKeyAuth( api_key="sk-test-key", user_id="user-1", @@ -481,6 +476,7 @@ async def test_get_available_models_for_user_without_access_groups_grants_nothin ) assert result == [] + @pytest.mark.asyncio async def test_get_available_models_for_user_resolves_key_access_group_models( monkeypatch, @@ -521,3 +517,78 @@ async def test_get_available_models_for_user_resolves_key_access_group_models( user_api_key_cache=MagicMock(), ) assert result == ["model-b"] + + +def _agent_ceiling(models: frozenset[str] | None): + from litellm.proxy.agent_endpoints.auth.agent_access_groups import AgentAccessGroupCeiling + + async def resolve(agent_id: str) -> AgentAccessGroupCeiling | None: + if models is None: + return None + return AgentAccessGroupCeiling( + access_group_ids=("ag-agent",), models=models, mcp_server_ids=frozenset(), agent_ids=frozenset() + ) + + return resolve + + +def _agent_key(models: list[str]) -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-agent-key", user_id="user-1", agent_id="agent-1", models=models) + + +@pytest.mark.asyncio +async def test_agent_key_listing_is_capped_to_its_access_groups(): + result = await get_available_models_for_user( + user_api_key_dict=_agent_key(["model-a", "model-b", "model-c"]), + llm_router=_router_with_models(["model-a", "model-b", "model-c"]), + general_settings={}, + user_model=None, + resolve_agent_ceiling=_agent_ceiling(frozenset({"model-b", "model-d"})), + ) + assert result == ["model-b"] + + +@pytest.mark.asyncio +async def test_agent_key_listing_is_empty_when_its_groups_grant_no_model(): + result = await get_available_models_for_user( + user_api_key_dict=_agent_key(["model-a"]), + llm_router=_router_with_models(["model-a"]), + general_settings={}, + user_model=None, + resolve_agent_ceiling=_agent_ceiling(frozenset()), + ) + assert result == [] + + +@pytest.mark.asyncio +async def test_agent_ceiling_expands_a_model_access_group_name_for_listing(): + router = _router_with_models(["model-a", "model-b"]) + router.get_model_access_groups.return_value = {"fast-models": ["model-b"]} + result = await get_available_models_for_user( + user_api_key_dict=_agent_key(["model-a", "model-b"]), + llm_router=router, + general_settings={}, + user_model=None, + resolve_agent_ceiling=_agent_ceiling(frozenset({"fast-models"})), + ) + assert result == ["model-b"] + + +@pytest.mark.asyncio +async def test_listing_is_unchanged_without_an_agent_or_without_attached_groups(): + router = _router_with_models(["model-a", "model-b"]) + plain_key = await get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-plain", user_id="user-1", models=["model-a", "model-b"]), + llm_router=router, + general_settings={}, + user_model=None, + resolve_agent_ceiling=_agent_ceiling(frozenset({"model-a"})), + ) + agent_without_groups = await get_available_models_for_user( + user_api_key_dict=_agent_key(["model-a", "model-b"]), + llm_router=router, + general_settings={}, + user_model=None, + resolve_agent_ceiling=_agent_ceiling(None), + ) + assert (plain_key, agent_without_groups) == (["model-a", "model-b"], ["model-a", "model-b"]) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index 83f30dc52a4..97e7a799053 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -13,6 +13,7 @@ import time from collections.abc import AsyncIterator, Mapping, Sequence from copy import deepcopy from functools import partial +from types import MappingProxyType from typing import Dict, Final, List, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -90,6 +91,7 @@ from litellm.types.router import ( TaggedPreRoutingStrategy, ) from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.management_endpoints.auto_router_endpoints import RequestComplexityRouterConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -6681,6 +6683,7 @@ class TestTierModelAffinity: litellm_router_instance=_windowed_router(_SMALL, _BIG), complexity_router_config={ "tiers": {"SIMPLE": ["small-model", "big-model"]}, + "enable_context_window_escalation": True, "adaptive": adaptive, "deployment_affinity": True, "session_affinity": False, @@ -13878,8 +13881,12 @@ _CJK_TURNS = [ ] -def _tier_config(**overrides) -> Dict: - return {"tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}, **overrides} +def _tier_config(**overrides: object) -> dict[str, object]: + return { + "tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}, + "enable_context_window_escalation": True, + **overrides, + } class TestContextWindowEscalation: @@ -13938,7 +13945,7 @@ class TestContextWindowEscalation: router = ComplexityRouter( model_name="test-router", litellm_router_instance=_windowed_router(_SMALL, ("mid-model", "openai/gpt-4o-mini", 200000), _BIG), - complexity_router_config={"tiers": {"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}}, + complexity_router_config=_tier_config(tiers={"SIMPLE": ["small-model", "mid-model"], "COMPLEX": "big-model"}), ) result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) @@ -13975,7 +13982,7 @@ class TestContextWindowEscalation: }, ] ), - complexity_router_config={"tiers": {"SIMPLE": "mixed-pool", "COMPLEX": "big-model"}}, + complexity_router_config=_tier_config(tiers={"SIMPLE": "mixed-pool", "COMPLEX": "big-model"}), ) result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) @@ -14026,7 +14033,7 @@ class TestContextWindowEscalation: router = ComplexityRouter( model_name="test-router", litellm_router_instance=_windowed_router(*deployments), - complexity_router_config={"tiers": tiers}, + complexity_router_config=_tier_config(tiers=tiers), ) result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) @@ -14035,19 +14042,37 @@ class TestContextWindowEscalation: assert result.model == expected_model @pytest.mark.asyncio - async def test_the_disabled_gate_dispatches_on_complexity_alone(self): - """The escape hatch: enable_context_window_escalation false restores today's behavior.""" - router = ComplexityRouter( + @pytest.mark.parametrize("enabled", (None, False, True), ids=("omitted", "disabled", "enabled")) + @pytest.mark.parametrize("serialized", (False, True), ids=("config", "http-json")) + async def test_context_window_escalation_requires_opt_in(self, enabled: bool | None, serialized: bool) -> None: + setting: Final = ( + MappingProxyType({"enable_context_window_escalation": enabled}) + if enabled is not None + else MappingProxyType({}) + ) + raw_config: Final = RequestComplexityRouterConfig.model_validate( + MappingProxyType( + {"tiers": MappingProxyType({"SIMPLE": "small-model", "COMPLEX": "big-model"}), **setting} + ) + ) + config: Final = ( + RequestComplexityRouterConfig.model_validate_json(raw_config.model_dump_json()) + if serialized + else raw_config + ) + router: Final = ComplexityRouter( model_name="test-router", litellm_router_instance=_windowed_router(_SMALL, _BIG), - complexity_router_config=_tier_config(enable_context_window_escalation=False), + complexity_router_config=config.model_dump(exclude_unset=not serialized, exclude_none=True), ) - result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) + result: Final = await router.async_pre_routing_hook( + model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS + ) assert result is not None - assert result.model == "small-model" - assert "context_escalated" not in result.routing_decision + assert result.model == ("big-model" if enabled else "small-model") + assert result.routing_decision.get("context_escalated", False) is (enabled is True) @pytest.mark.asyncio async def test_out_of_band_system_and_tools_count_against_the_window(self): @@ -14156,7 +14181,7 @@ class TestContextWindowEscalation: }, ] ), - complexity_router_config={"adaptive": True, "tiers": {"SIMPLE": ["small-model", "mid-model"]}}, + complexity_router_config=_tier_config(adaptive=True, tiers={"SIMPLE": ["small-model", "mid-model"]}), ) result = await router.async_pre_routing_hook(model="test-router", request_kwargs={}, messages=_OVERSIZED_TURNS) @@ -14186,7 +14211,7 @@ class TestContextWindowEscalation: }, ] ), - complexity_router_config={"tiers": {"SIMPLE": "cop-pool", "COMPLEX": "big-model"}}, + complexity_router_config=_tier_config(tiers={"SIMPLE": "cop-pool", "COMPLEX": "big-model"}), ) real_get_llm_provider = litellm.get_llm_provider copilot_resolutions: List = [] @@ -14219,7 +14244,7 @@ class TestContextWindowEscalation: "model_name": "smart-router", "litellm_params": { "model": "auto_router/complexity_router", - "complexity_router_config": {"tiers": {"SIMPLE": "small-model", "COMPLEX": "big-model"}}, + "complexity_router_config": _tier_config(), }, }, { @@ -15139,7 +15164,12 @@ class TestHealthFallbackDispatch: ) -> None: from litellm.types.router import RouterRateLimitError - router: Final = self._router(config={"tiers": {"SIMPLE": "primary", "MEDIUM": "peer", "COMPLEX": "large"}}) + router: Final = self._router( + config={ + "tiers": {"SIMPLE": "primary", "MEDIUM": "peer", "COMPLEX": "large"}, + "enable_context_window_escalation": True, + } + ) router.add_deployment( Deployment( model_name="large", @@ -15216,7 +15246,13 @@ class TestHealthFallbackDispatch: @pytest.mark.asyncio @pytest.mark.parametrize("default_fits", [True, False]) async def test_modality_default_must_also_fit_context(self, default_fits: bool) -> None: - router: Final = self._router(config={"modality_routing": True, "tiers": {"SIMPLE": "primary"}}) + router: Final = self._router( + config={ + "modality_routing": True, + "tiers": {"SIMPLE": "primary"}, + "enable_context_window_escalation": True, + } + ) for deployment in router.model_list: deployment["model_info"]["supports_vision"] = deployment["model_name"] == "fallback" deployment["model_info"]["max_input_tokens"] = 10000 if default_fits else 10 diff --git a/tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py b/tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py new file mode 100644 index 00000000000..7b13b196d5b --- /dev/null +++ b/tests/test_litellm/router_strategy/test_lowest_tpm_rpm.py @@ -0,0 +1,54 @@ +from datetime import datetime, timedelta +from typing import Final + +from litellm import Router +from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict + +MODEL_GROUP: Final = "lowest-tpm-router" +HIGH_USAGE_DEPLOYMENT_ID: Final = "highest-usage" +LOW_USAGE_DEPLOYMENT_ID: Final = "lowest-usage" + + +def _deployment(deployment_id: str) -> DeploymentTypedDict: + params: LiteLLMParamsTypedDict = { + "model": "gpt-4o", + "api_key": "key", + "mock_response": f"from {deployment_id}", + } + return { + "model_name": MODEL_GROUP, + "litellm_params": params, + "model_info": {"id": deployment_id}, + } + + +def test_usage_based_routing_v1_selects_the_lowest_recorded_tpm() -> None: + router: Final = Router( + model_list=[ + _deployment(HIGH_USAGE_DEPLOYMENT_ID), + _deployment(LOW_USAGE_DEPLOYMENT_ID), + ], + routing_strategy="usage-based-routing", + num_retries=0, + ) + usage_by_deployment: Final = { + HIGH_USAGE_DEPLOYMENT_ID: 100, + LOW_USAGE_DEPLOYMENT_ID: 1, + } + now: Final = datetime.now() + cache_keys: Final = tuple( + f"{MODEL_GROUP}:tpm:{(now + timedelta(minutes=offset)).strftime('%H-%M')}" + for offset in range(60) + ) + + for cache_key in cache_keys: + router.cache.set_cache( + key=cache_key, value=usage_by_deployment, ttl=float("inf") + ) + + deployment: Final = router.get_available_deployment( + model=MODEL_GROUP, + messages=[{"role": "user", "content": "test"}], + ) + + assert deployment["model_info"]["id"] == LOW_USAGE_DEPLOYMENT_ID diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index e0d92a2e957..45f7214ecbf 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -1,19 +1,23 @@ import asyncio import contextvars import gc +import hashlib import json +import math import os import threading import time import uuid import weakref -from collections.abc import Generator +from collections.abc import Callable, Generator +from contextlib import ExitStack from datetime import datetime from pathlib import Path from types import SimpleNamespace from typing import Final, Protocol, cast from unittest.mock import Mock from urllib.parse import urlparse +from uuid import uuid4 import boto3 import botocore.config @@ -30,13 +34,19 @@ from litellm.caching.disk_cache import DiskCache from litellm.caching.gcs_cache import GCSCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cluster_cache import RedisClusterCache +from litellm.caching.redis_semantic_cache import RedisSemanticCache from litellm.caching.s3_cache import S3Cache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType +from litellm.types.llms.custom_llm import CustomLLMItem +from litellm.types.utils import EmbeddingResponse from tests.test_litellm_rust.support.fake_gcs import FakeGcs from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.s3_stub import S3Stub +_CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name +_CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name + pytestmark: Final = pytest.mark.requires_rust_extension @@ -114,14 +124,14 @@ def test_existing_constructor_and_global_are_unchanged() -> None: assert type(facade.cache) is InMemoryCache assert "_native_cache_handle" not in vars(facade) with rebound(litellm, "cache", facade): - resolver: Final = _native._CacheTestResolver(litellm) + resolver: Final = _CacheTestResolver(litellm) assert resolver.resolve().kind == "python_callback" resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"}) assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7} def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None: - resolver: Final = _native._CacheTestResolver(litellm) + resolver: Final = _CacheTestResolver(litellm) enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30) enabled: Final = litellm.cache @@ -144,13 +154,13 @@ def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> Non async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None: - namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.memory()) - resolver: Final = _native._CacheTestResolver(namespace) + namespace: Final = SimpleNamespace(cache=_CacheTestHandle.memory()) + resolver: Final = _CacheTestResolver(namespace) selected: Final = resolver.resolve() assert selected.kind == "native" selected.store(request(), {"answer": 1}) assert await selected.async_lookup(request()) == {"answer": 1} - with rebound(namespace, "cache", _native._CacheTestHandle.memory()): + with rebound(namespace, "cache", _CacheTestHandle.memory()): replacement: Final = resolver.resolve() await selected.async_store(request(), {"answer": 2}) assert replacement.lookup(request()) is None @@ -183,7 +193,7 @@ async def test_python_callback_preserves_identity_caller_task_context_and_errors raise failure namespace: Final = SimpleNamespace(cache=CustomCache()) - binding: Final = _native._CacheTestResolver(namespace).resolve() + binding: Final = _CacheTestResolver(namespace).resolve() assert binding.kind == "python_callback" assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel assert context.get() == "callback" @@ -204,7 +214,7 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: finally: finished.set() - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve() async def lookup() -> object: return await binding.async_lookup(None, callback_kwargs={}) @@ -219,9 +229,9 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None: def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) - handle: Final = _native._CacheTestHandle.memory() + handle: Final = _CacheTestHandle.memory() handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) native: Final = resolver.resolve() assert native.kind == "native" native.store(request(), {"source": "native"}) @@ -252,12 +262,12 @@ def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not class CustomCache(Cache): pass - handle: Final = _native._CacheTestHandle.memory() + handle: Final = _CacheTestHandle.memory() with pytest.raises(TypeError): handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL)) facade: Final = Cache(type=LiteLLMCacheType.LOCAL) handle._bind_facade(facade) - resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) with rebound(facade, "cache", InMemoryCache()): assert resolver.resolve().kind == "python_callback" with rebound(facade, "ttl", 12): @@ -282,7 +292,7 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: def cyclic_reference() -> weakref.ReferenceType[CustomCache]: callback: Final = CustomCache() namespace: Final = SimpleNamespace(cache=callback) - binding: Final = _native._CacheTestResolver(namespace).resolve() + binding: Final = _CacheTestResolver(namespace).resolve() setattr(callback, "binding", binding) return weakref.ref(callback) @@ -293,8 +303,8 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None: async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None: client: Final = redis.Redis.from_url(redis_url) - namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url, namespace="team")) - binding: Final = _native._CacheTestResolver(namespace).resolve() + namespace: Final = SimpleNamespace(cache=_CacheTestHandle.redis(redis_url, namespace="team")) + binding: Final = _CacheTestResolver(namespace).resolve() response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None} envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)} client.set("team:sync", str(envelope)) @@ -316,33 +326,33 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd def test_invalid_duration_and_request_shape_fail_before_storage() -> None: - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() for seconds in (-1.0, float("nan"), float("inf")): with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1}) assert binding.lookup(request()) is None with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"): - _native._CacheTestHandle.memory(ttl_seconds=-1) + _CacheTestHandle.memory(ttl_seconds=-1) async def test_memory_size_policy_is_applied_by_the_native_host() -> None: - handle: Final = _native._CacheTestHandle.memory(capacity=2, max_entry_bytes=128) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() + handle: Final = _CacheTestHandle.memory(capacity=2, max_entry_bytes=128) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=handle)).resolve() small: Final = {"answer": "ok"} binding.store(request("small"), small) assert await binding.async_lookup(request("small")) == small await binding.async_store(request("large"), {"answer": "x" * 256}) assert binding.lookup(request("large")) is None assert binding.lookup(request("small")) == small - disabled: Final = _native._CacheTestResolver( - SimpleNamespace(cache=_native._CacheTestHandle.memory(capacity=0)) + disabled: Final = _CacheTestResolver( + SimpleNamespace(cache=_CacheTestHandle.memory(capacity=0)) ).resolve() await disabled.async_store(request(), small) assert await disabled.async_lookup(request()) is None async def test_native_batch_lookup_and_store_report_partial_hits() -> None: - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve() requests: Final = [request("hit"), request("miss"), request("disabled")] requests[2]["controls"] = { "supported_call_type": True, @@ -380,7 +390,7 @@ async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None: ) -> object: return result, kwargs - binding: Final = _native._CacheTestResolver( + binding: Final = _CacheTestResolver( SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL)) ).resolve() assert binding.kind == "python_callback" @@ -410,7 +420,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: cache: Final = Cache(type=LiteLLMCacheType.LOCAL) cache.cache.set_cache("key", "value") - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=cache)).resolve() + binding: Final = _CacheTestResolver(SimpleNamespace(cache=cache)).resolve() assert binding.kind == "python_callback" setattr(cache.cache, "ping", ping) @@ -422,7 +432,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None: def test_facade_registration_rejects_mismatched_capacity() -> None: facade: Final = Cache(type=LiteLLMCacheType.LOCAL) with pytest.raises(TypeError, match="capacities must match"): - _native._CacheTestHandle.memory(capacity=7)._bind_facade(facade) + _CacheTestHandle.memory(capacity=7)._bind_facade(facade) def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure_blob_facade: Cache) -> None: @@ -525,19 +535,19 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None: redis_flush_size=2, ) with pytest.raises(TypeError, match="default TTLs must match"): - _native._CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) + _CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade) with pytest.raises(TypeError, match="namespaces must match"): - _native._CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) - _native._CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) - binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + _CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade) + _CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() client: Final = redis.Redis.from_url(redis_url) with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}): - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" pool: Final = facade.cache.redis_client.connection_pool with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}): - assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" await binding.async_store(request("first"), {"value": 1}) assert client.get("first") is None @@ -1112,3 +1122,572 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n client.delete("unscoped") client.close() facade.cache.redis_client.close() + + +PARAPHRASE_MARKER: Final = " (paraphrase)" +SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic" +SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_" +SEMANTIC_CONTEXT: Final = contextvars.ContextVar("semantic_test_context", default="unset") + + +def _normalized(vector: list[float]) -> list[float]: + norm: Final = math.sqrt(sum(component * component for component in vector)) + return [component / norm for component in vector] + + +def _base_embedding(prompt: str) -> list[float]: + digest: Final = hashlib.sha256(prompt.encode("utf-8")).digest() + return _normalized([float(digest[index] + 1) for index in range(8)]) + + +def _semantic_embedding(prompt: str) -> list[float]: + if PARAPHRASE_MARKER not in prompt: + return _base_embedding(prompt) + base: Final = _base_embedding(prompt.replace(PARAPHRASE_MARKER, "").strip()) + pivot: Final = min(range(8), key=lambda index: abs(base[index])) + direction: Final = _normalized( + [ + (1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot] + for index in range(8) + ] + ) + # Rotating an orthogonal unit direction by 0.329 produces ~0.05 cosine distance + return _normalized([base[index] + 0.329 * direction[index] for index in range(8)]) + + +class DeterministicEmbedding(litellm.CustomLLM): + def __init__(self) -> None: + self.calls: list[dict[str, object]] = [] + self.async_calls: list[dict[str, object]] = [] + self.entered = asyncio.Event() + self.gate: asyncio.Event | None = None + + def _respond( + self, + model: str, + input: object, + model_response: EmbeddingResponse, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.calls.append({"model": model, "input": texts}) + model_response.model = model + model_response.data = [ + {"object": "embedding", "index": index, "embedding": _semantic_embedding(str(text))} + for index, text in enumerate(texts) + ] + return model_response + + def embedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + return self._respond(model, input, model_response) + + async def aembedding( + self, + model: str, + input: list[object], + model_response: EmbeddingResponse, + print_verbose: Callable[..., object], + logging_obj: object, + optional_params: dict[str, object], + api_key: object = None, + api_base: object = None, + timeout: object = None, + litellm_params: object = None, + ) -> EmbeddingResponse: + texts: Final = cast(list[object], input if isinstance(input, list) else [input]) + self.async_calls.append( + { + "model": model, + "input": texts, + "task": asyncio.current_task(), + "context": SEMANTIC_CONTEXT.get(), + } + ) + SEMANTIC_CONTEXT.set("written-in-aembedding") + self.entered.set() + if self.gate is not None: + await self.gate.wait() + return self._respond(model, input, model_response) + + +@pytest.fixture +def semantic_embedding() -> Generator[DeterministicEmbedding]: + handler: Final = DeterministicEmbedding() + with ExitStack() as stack: + stack.enter_context( + rebound( + litellm, + "custom_provider_map", + [ + *litellm.custom_provider_map, + cast( + CustomLLMItem, + {"provider": "semantic-test", "custom_handler": handler}, + ), + ], + ) + ) + stack.enter_context( + rebound( + litellm, + "_custom_providers", # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + [*litellm._custom_providers, "semantic-test"], # pyright: ignore[reportPrivateUsage] # no public provider-registration hook + ) + ) + stack.enter_context( + rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"]) + ) + yield handler + + +@pytest.fixture +def redis_stack() -> Generator[tuple[str, str]]: + url: Final = os.environ.get("LITELLM_REDIS_STACK_URL") + if url is None: + pytest.skip("LITELLM_REDIS_STACK_URL is not set") + index: Final = f"{SEMANTIC_INDEX_PREFIX}{uuid4().hex}" + yield url, index + client: Final = redis.Redis.from_url(url) + try: + client.execute_command("FT.DROPINDEX", index, "DD") # pyright: ignore[reportUnknownMemberType] # redis-py leaves execute_command partially unknown + except redis.RedisError: + pass + client.close() + + +def semantic_request(key: str, prompt: str, **extra: object) -> dict[str, object]: + return { + "key": {"preset": key}, + "messages": [{"role": "user", "content": prompt}], + **extra, + } + + +def semantic_messages(prompt: str) -> list[dict[str, object]]: + return [{"role": "user", "content": prompt}] + + +def semantic_entry_id(prompt: str, tag: str) -> str: + return hashlib.sha256(f"{prompt}litellm_cache_key{tag}".encode()).hexdigest() + + +def semantic_facade(url: str, index: str, *, similarity_threshold: float = 0.8) -> Cache: + facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=similarity_threshold, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(facade) + return facade + + +def test_redis_semantic_constructor_identity_and_provenance( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + backend: Final = cast(RedisSemanticCache, facade.cache) + assert backend.__class__.__module__ == "litellm.caching.redis_semantic_cache" + assert type(backend) is RedisSemanticCache + assert backend._redis_url == url # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend._index_name == index # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config + assert backend.similarity_threshold == 0.8 + assert backend.embedding_model == SEMANTIC_EMBEDDING_MODEL + handle: Final = cast(object, getattr(facade, "_native_cache_handle")) + assert isinstance(handle, _CacheTestHandle) + assert handle.backend == "redis_semantic" + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + + +def test_redis_semantic_native_and_python_sync_entries_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + response: Final = {"choices": [{"text": "paris"}], "usage": {"total_tokens": 2}} + + binding.store(semantic_request("geo", "what is the capital of france"), response) + + native_hash_key: Final = f"{index}:{semantic_entry_id('what is the capital of france', 'geo')}" + stored: Final = client.hgetall(native_hash_key) + assert set(stored) == { + b"entry_id", + b"prompt", + b"response", + b"prompt_vector", + b"inserted_at", + b"updated_at", + b"litellm_cache_key", + }, stored + assert stored[b"entry_id"].decode() == native_hash_key.split(":", 1)[1] + assert stored[b"prompt"] == b"what is the capital of france" + assert stored[b"litellm_cache_key"] == b"geo" + assert len(stored[b"prompt_vector"]) == 32 + decoded: Final = cast(dict[str, object], json.loads(stored[b"response"])) + assert decoded["response"] == response + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "geo", messages=semantic_messages("what is the capital of france") + ) + == decoded + ) + assert semantic_embedding.calls == [ + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["what is the capital of france"]}, + {"model": "deterministic", "input": ["dimension test"]}, + ] + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "math", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": 42}}), + messages=semantic_messages("what is 6 times 7"), + ) + python_hash_key: Final = f"{index}:{semantic_entry_id('what is 6 times 7', 'math')}" + assert json.loads(cast(bytes, client.hget(python_hash_key, "response"))) == { + "timestamp": 1700000000.0, + "response": {"answer": 42}, + } + assert binding.lookup(semantic_request("math", "what is 6 times 7")) == {"answer": 42} + client.close() + + +async def test_redis_semantic_async_paths_and_store_batch_share_one_layout( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + await binding.async_store( + semantic_request("async", "name a primary color"), {"answer": "blue"} + ) + hash_key: Final = f"{index}:{semantic_entry_id('name a primary color', 'async')}" + decoded: Final = cast(dict[str, object], json.loads(cast(bytes, client.hget(hash_key, "response")))) + python_read: Final = await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async", messages=semantic_messages("name a primary color") + ) + assert python_read == decoded + + await binding.async_store_batch( + [ + semantic_request("batch-one", "first batch prompt"), + semantic_request("batch-two", "second batch prompt"), + ], + [{"answer": 1}, {"answer": 2}], + ) + expected: Final = { + key: json.loads( + cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response")) + ) + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ) + } + for key, prompt in ( + ("batch-one", "first batch prompt"), + ("batch-two", "second batch prompt"), + ): + assert cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + key, messages=semantic_messages(prompt) + ) == expected[key], key + + cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "async-python", + json.dumps({"timestamp": 1700000000.0, "response": {"answer": "python"}}), + messages=semantic_messages("python written prompt"), + ) + assert await binding.async_lookup( + semantic_request("async-python", "python written prompt") + ) == {"answer": "python"} + client.close() + + +async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + caller: Final = asyncio.current_task() + SEMANTIC_CONTEXT.set("caller-sentinel") + response: Final = {"choices": [{"text": "paris"}]} + + await binding.async_store( + semantic_request("inline", "what is the capital of france"), response + ) + assert ( + await binding.async_lookup( + semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}") + ) + == response + ) + assert await binding.async_lookup(semantic_request("inline", "python written prompt")) is None + assert SEMANTIC_CONTEXT.get() == "written-in-aembedding" + assert semantic_embedding.async_calls == [ + { + "model": "deterministic", + "input": ["what is the capital of france"], + "task": caller, + "context": "caller-sentinel", + }, + { + "model": "deterministic", + "input": [f"what is the capital of france{PARAPHRASE_MARKER}"], + "task": caller, + "context": "written-in-aembedding", + }, + { + "model": "deterministic", + "input": ["python written prompt"], + "task": caller, + "context": "written-in-aembedding", + }, + ], semantic_embedding.async_calls + + +async def test_native_semantic_cancellation_during_embedding_skips_the_backend( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + semantic_embedding.gate = asyncio.Event() + + async def lookup() -> object: + return await binding.async_lookup( + semantic_request("cancel", "cancelled prompt") + ) + + task: Final = asyncio.create_task(lookup()) + await semantic_embedding.entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + semantic_embedding.gate.set() + + assert len(semantic_embedding.async_calls) == 1 + assert ( + await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "cancel", messages=semantic_messages("cancelled prompt") + ) + is None + ) + + +def test_redis_semantic_similarity_tag_and_threshold_boundaries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + binding.store(semantic_request("sim", "tell me a joke"), {"answer": "haha"}) + paraphrase: Final = f"tell me a joke{PARAPHRASE_MARKER}" + assert binding.lookup(semantic_request("sim", paraphrase)) == {"answer": "haha"} + assert binding.lookup(semantic_request("sim", "an unrelated question about spreadsheets")) is None + assert binding.lookup(semantic_request("other-key", "tell me a joke")) is None + + strict: Final = semantic_facade(url, index, similarity_threshold=0.99) + strict_binding: Final = _CacheTestResolver(SimpleNamespace(cache=strict)).resolve() + assert strict_binding.lookup(semantic_request("sim", paraphrase)) is None + assert strict_binding.lookup(semantic_request("sim", "tell me a joke")) == {"answer": "haha"} + + +def test_redis_semantic_ttl_is_written_only_when_requested( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store( + {**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1} + ) + expiring: Final = f"{index}:{semantic_entry_id('ttl prompt', 'ttl')}" + assert 0 < client.ttl(expiring) <= 12 + + binding.store(semantic_request("ttl-none", "untimed prompt"), {"answer": 2}) + persistent: Final = f"{index}:{semantic_entry_id('untimed prompt', 'ttl-none')}" + assert client.ttl(persistent) == -1 + + binding.store( + {**semantic_request("ttl-fraction", "fractional prompt"), "ttl_seconds": 1.5}, + {"answer": 3}, + ) + fractional: Final = f"{index}:{semantic_entry_id('fractional prompt', 'ttl-fraction')}" + assert client.ttl(fractional) == 2 + client.close() + + +def test_redis_semantic_malformed_response_is_a_miss_for_both_readers( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(semantic_request("bad", "corrupt me"), {"answer": 1}) + hash_key: Final = f"{index}:{semantic_entry_id('corrupt me', 'bad')}" + client.hset(hash_key, "response", b"{not json") + assert binding.lookup(semantic_request("bad", "corrupt me")) is None + assert ( + cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class + "bad", messages=semantic_messages("corrupt me") + ) + is None + ) + client.close() + + +async def test_redis_semantic_unsupported_operations_raise_not_implemented( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + + with pytest.raises(NotImplementedError): + binding.lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_lookup_batch([semantic_request("batch", "prompt one")]) + with pytest.raises(NotImplementedError): + await binding.async_flush() + with pytest.raises(NotImplementedError): + await binding.ping() + + +def test_redis_semantic_requests_without_prompt_are_noops( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + binding.store(request("plain"), {"answer": 1}) + assert binding.lookup(request("plain")) is None + assert semantic_embedding.calls == [] + assert client.keys(f"{index}:*") == [] + client.close() + + +def test_redis_semantic_scope_overrides_the_tag_and_isolates_entries( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + client: Final = redis.Redis.from_url(url) + + scoped: Final = {**semantic_request("scoped", "scoped prompt"), "scope": "team-a"} + binding.store(scoped, {"answer": "kept"}) + hash_key: Final = f"{index}:{semantic_entry_id('scoped prompt', 'team-a')}" + assert client.hget(hash_key, "litellm_cache_key") == b"team-a" + assert binding.lookup(scoped) == {"answer": "kept"} + assert binding.lookup(semantic_request("scoped", "scoped prompt")) is None + assert binding.lookup({**scoped, "scope": "team-b"}) is None + client.close() + + +def test_redis_semantic_configuration_drift_falls_back_to_python( + redis_stack: tuple[str, str], + semantic_embedding: DeterministicEmbedding, + monkeypatch: pytest.MonkeyPatch, +) -> None: + url, index = redis_stack + facade: Final = semantic_facade(url, index) + resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade)) + assert resolver.resolve().kind == "native" + + with rebound(facade.cache, "similarity_threshold", 0.5): + assert resolver.resolve().kind == "python_callback" + with rebound(facade, "semantic_cache_scope", "end_user"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "embedding_model", "other-model"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "_index_name", "other-index"): + assert resolver.resolve().kind == "python_callback" + with rebound(facade.cache, "CACHE_KEY_FIELD_NAME", "other-field"): + assert resolver.resolve().kind == "python_callback" + + def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]: + return _semantic_embedding(prompt) + + monkeypatch.setattr(RedisSemanticCache, "_get_embedding", patched_embedding) + assert resolver.resolve().kind == "python_callback" + + +def test_redis_semantic_handle_rejects_wrong_backends( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding +) -> None: + url, index = redis_stack + + class CustomSemanticCache(RedisSemanticCache): + pass + + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + _CacheTestHandle.redis_semantic(object()) + with pytest.raises(TypeError, match="built-in RedisSemanticCache"): + _CacheTestHandle.redis_semantic( + CustomSemanticCache( + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=f"{index}_subclass", + ) + ) + + facade: Final = semantic_facade(url, index) + with pytest.raises(TypeError, match="backend types must match"): + _CacheTestHandle.redis(url)._bind_facade(facade) + + subclassed_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + subclassed_facade.cache = CustomSemanticCache( # pyright: ignore[reportAttributeAccessIssue] # facade backend slot is not declared + + redis_url=url, + similarity_threshold=0.8, + embedding_model=SEMANTIC_EMBEDDING_MODEL, + index_name=index, + ) + with pytest.raises(TypeError): + _CacheTestHandle.redis_semantic( + subclassed_facade.cache + )._bind_facade(subclassed_facade) + + replacement_facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + with pytest.raises(TypeError, match="must be the native embedder"): + _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade) 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)/agents/_components/AgentFormKit.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx index 15b85001cdc..8e100d0c3ed 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/AgentFormKit.tsx @@ -90,6 +90,7 @@ export interface AgentFormValues { guardrails?: string[]; entitlement_models?: string[]; entitlement_agents?: string[]; + access_group_ids?: string[]; allowed_mcp_servers_and_groups?: McpServerSelection; mcp_tool_permissions?: Record; defaultInputModes?: string[]; @@ -121,6 +122,7 @@ export interface AgentRequestPayload { agent_card_params?: Record; litellm_params?: Record; object_permission?: Record; + access_group_ids?: string[]; } interface AgentFormFieldProps { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx index ebc97891744..457ee656415 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.integration.test.tsx @@ -21,6 +21,9 @@ vi.mock("./agent_card_discovery", () => ({ default: () =>
({ default: () =>
})); vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({ default: () =>
})); vi.mock("@/components/guardrails/GuardrailSelector", () => ({ default: () =>
})); +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: () => ({ data: [], isLoading: false, isError: false }), +})); vi.mock("@/components/common_components/team_dropdown", () => ({ default: () =>
})); const a2aInfo: AgentCreateInfo = { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx index ccac244f019..fbf5cf8c1fb 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.test.tsx @@ -44,6 +44,14 @@ vi.mock("@/components/mcp_server_management/MCPToolPermissions", () => ({ default: () => null, })); +vi.mock("@/components/common_components/AccessGroupSelector", () => ({ + default: ({ onChange }: { onChange: (value: string[]) => void }) => ( + + ), +})); + vi.mock("@/components/common_components/team_dropdown", () => ({ default: () => null, })); @@ -141,5 +149,29 @@ describe("AddAgentForm logos", () => { await vi.waitFor(() => expect(networking.createAgentCall).toHaveBeenCalled()); const [, payload] = vi.mocked(networking.createAgentCall).mock.calls[0]; expect(payload.object_permission).toEqual({ mcp_toolsets: ["ts-1"] }); + expect(payload).not.toHaveProperty("access_group_ids"); + }); + + it("includes selected access groups in the create payload", async () => { + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + vi.mocked(networking.createAgentCall) + .mockReset() + .mockResolvedValue({ + agent_id: "agent-1", + agent_name: "Test Agent", + } as never); + vi.mocked(networking.keyListCall).mockResolvedValue({ keys: [] }); + + renderForm(); + await user.click(screen.getByRole("button", { name: "Next →" })); + await user.click(screen.getByTestId("select-access-group")); + await user.click(screen.getByRole("button", { name: "Next →" })); + await user.click(screen.getByRole("button", { name: "Next →" })); + await user.click(screen.getByText(/Skip for now/)); + await user.click(screen.getByRole("button", { name: "Create Agent →" })); + + await vi.waitFor(() => expect(networking.createAgentCall).toHaveBeenCalled()); + const [, payload] = vi.mocked(networking.createAgentCall).mock.calls[0]; + expect(payload.access_group_ids).toEqual(["ag-1", "ag-2"]); }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx index e71fed40209..5bd6ea9b83a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/add_agent_form.tsx @@ -50,6 +50,7 @@ import { } from "./AgentFormKit"; import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; +import AccessGroupSelector from "@/components/common_components/AccessGroupSelector"; import GuardrailSelector from "@/components/guardrails/GuardrailSelector"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; @@ -113,6 +114,7 @@ const SHARED_INITIAL_VALUES: AgentFormValues = { mcp_tool_permissions: {}, entitlement_models: [], entitlement_agents: [], + access_group_ids: [], guardrails: [], }; @@ -374,6 +376,9 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok if (Object.keys(objectPermission).length > 0) { agentData.object_permission = objectPermission; } + if (values.access_group_ids?.length) { + agentData.access_group_ids = values.access_group_ids; + } // Wire trace-id flags and budget controls into agent litellm_params (before create call) if (requireTraceIdInbound || requireTraceIdOutbound) { @@ -494,6 +499,22 @@ const AddAgentForm: React.FC = ({ visible, onClose, accessTok )} + + {({ value, onChange }) => ( + + )} + + { return agentData; }; +export const parseAccessGroupIdsForForm = (agent: { access_group_ids?: string[] | null }) => ({ + access_group_ids: agent.access_group_ids ?? [], +}); + export const parseMcpPermissionsForForm = (agent: any) => ({ allowed_mcp_servers_and_groups: { servers: agent.object_permission?.mcp_servers ?? [], @@ -377,5 +381,6 @@ export const parseAgentForForm = (agent: any) => { // extra_headers: already an array of strings extra_headers: agent.extra_headers ?? [], ...parseMcpPermissionsForForm(agent), + ...parseAccessGroupIdsForForm(agent), }; }; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx index 357f924cbe7..37e00766a75 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.integration.test.tsx @@ -25,6 +25,10 @@ vi.mock("@/app/(dashboard)/hooks/keys/useKeys", () => ({ vi.mock("./agent_card_discovery", () => ({ default: () =>
})); +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: () => ({ data: [], isLoading: false, isError: false }), +})); + const A2A_AGENT = { agent_id: "agent-1", agent_name: "my-agent", @@ -176,6 +180,7 @@ describe("AgentInfoView update payload", () => { session_tpm_limit: 333, session_rpm_limit: 444, object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, + access_group_ids: [], }); }); @@ -217,6 +222,7 @@ describe("AgentInfoView update payload", () => { session_tpm_limit: 333, session_rpm_limit: 444, object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, + access_group_ids: [], }); }); @@ -295,6 +301,7 @@ describe("AgentInfoView update payload", () => { model: "langgraph/asst_1", }, object_permission: { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }, + access_group_ids: [], }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx index 29d97a20afe..7e6c7c0e05c 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.test.tsx @@ -28,6 +28,28 @@ vi.mock("@/app/(dashboard)/hooks/mcpServers/useMCPServers", () => ({ useMCPServers: () => ({ data: [{ server_id: "srv-1", server_name: "github" }] }), })); +vi.mock("@/app/(dashboard)/hooks/accessGroups/useAccessGroups", () => ({ + useAccessGroups: () => ({ + data: [{ access_group_id: "ag-1", access_group_name: "support-tools" }], + isLoading: false, + isError: false, + }), +})); + +vi.mock("@/components/common_components/AccessGroupSelector", () => ({ + default: ({ value, onChange }: { value?: string[]; onChange: (value: string[]) => void }) => ( +
+ {(value ?? []).join(",")} + + +
+ ), +})); + vi.mock("@/components/mcp_server_management/MCPServerSelector", () => ({ default: () =>
, })); @@ -76,6 +98,38 @@ describe("AgentInfoView settings", () => { expect(payload.tpm_limit).toBe(42); const clearedMcpGrants = { mcp_servers: [], mcp_access_groups: [], mcp_toolsets: [], mcp_tool_permissions: {} }; expect(payload.object_permission).toEqual(clearedMcpGrants); + expect(payload.access_group_ids).toEqual([]); + }); + + it("sends the newly attached access group in the update payload", async () => { + render(); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(screen.getByRole("button", { name: "Edit Settings" })); + fireEvent.click(await screen.findByRole("button", { name: "Attach ag-1" })); + expect(screen.getByTestId("selected-access-groups")).toHaveTextContent("ag-1"); + + fireEvent.click(screen.getByRole("button", { name: /Save Changes/ })); + + await waitFor(() => expect(networking.patchAgentCall).toHaveBeenCalledTimes(1)); + const [, , payload] = vi.mocked(networking.patchAgentCall).mock.calls[0]; + expect(payload.access_group_ids).toEqual(["ag-1"]); + }); + + it("loads the attached access groups into the editor and sends an empty list once detached", async () => { + vi.mocked(networking.getAgentInfo).mockResolvedValue({ ...agent, access_group_ids: ["ag-1"] }); + render(); + + fireEvent.click(await screen.findByRole("tab", { name: "Settings" })); + fireEvent.click(screen.getByRole("button", { name: "Edit Settings" })); + expect(await screen.findByTestId("selected-access-groups")).toHaveTextContent("ag-1"); + + fireEvent.click(screen.getByRole("button", { name: "Detach all access groups" })); + fireEvent.click(screen.getByRole("button", { name: /Save Changes/ })); + + await waitFor(() => expect(networking.patchAgentCall).toHaveBeenCalledTimes(1)); + const [, , payload] = vi.mocked(networking.patchAgentCall).mock.calls[0]; + expect(payload.access_group_ids).toEqual([]); }); it("shows MCP grants with server names on the overview tab", async () => { @@ -88,4 +142,20 @@ describe("AgentInfoView settings", () => { expect(await screen.findByText("github (srv-1)")).toBeInTheDocument(); }); + + it("shows attached access groups with their names on the overview tab", async () => { + vi.mocked(networking.getAgentInfo).mockResolvedValue({ ...agent, access_group_ids: ["ag-1", "ag-unknown"] }); + + render(); + + expect(await screen.findByText("support-tools (ag-1)")).toBeInTheDocument(); + expect(screen.getByText("ag-unknown")).toBeInTheDocument(); + }); + + it("shows None when the agent has no access groups attached", async () => { + render(); + + expect(await screen.findByText("Access Groups")).toBeInTheDocument(); + expect(screen.getByText("None")).toBeInTheDocument(); + }); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx index 6cb99e9692f..adac456f232 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/agents/_components/agent_info.tsx @@ -16,6 +16,8 @@ import { Agent } from "@/components/agents/types"; import { KeyResponse } from "@/components/key_team_helpers/key_list"; import { useKeys } from "@/app/(dashboard)/hooks/keys/useKeys"; import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers"; +import { useAccessGroups } from "@/app/(dashboard)/hooks/accessGroups/useAccessGroups"; +import AccessGroupSelector from "@/components/common_components/AccessGroupSelector"; import KeyInfoView from "@/components/templates/key_info_view"; import MCPServerSelector from "@/components/mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "@/components/mcp_server_management/MCPToolPermissions"; @@ -26,6 +28,7 @@ import { AGENT_FORM_CONFIG, buildAgentDataFromForm, buildMcpObjectPermission, + parseAccessGroupIdsForForm, parseAgentForForm, parseMcpPermissionsForForm, } from "./agent_config"; @@ -122,7 +125,11 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT } else { const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType); if (typeInfo) { - form.reset({ ...parseDynamicAgentForForm(data, typeInfo), ...parseMcpPermissionsForForm(data) }); + form.reset({ + ...parseDynamicAgentForForm(data, typeInfo), + ...parseMcpPermissionsForForm(data), + ...parseAccessGroupIdsForForm(data), + }); } else { form.reset(parseAgentForForm(data)); } @@ -142,7 +149,11 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT if (agentType !== "a2a") { const typeInfo = agentTypeMetadata.find((t) => t.agent_type === agentType); if (typeInfo) { - form.reset({ ...parseDynamicAgentForForm(agent, typeInfo), ...parseMcpPermissionsForForm(agent) }); + form.reset({ + ...parseDynamicAgentForForm(agent, typeInfo), + ...parseMcpPermissionsForForm(agent), + ...parseAccessGroupIdsForForm(agent), + }); } } } @@ -153,12 +164,18 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT const mcpSelection = useWatch({ control: form.control, name: "allowed_mcp_servers_and_groups" }); const mcpToolPermissions = useWatch({ control: form.control, name: "mcp_tool_permissions" }); const { data: mcpServers = [] } = useMCPServers(); + const { data: accessGroups = [] } = useAccessGroups(); const mcpServerLabel = (serverId: string) => { const server = mcpServers.find((s) => s.server_id === serverId); return server?.server_name ? `${server.server_name} (${serverId})` : serverId; }; + const accessGroupLabel = (accessGroupId: string) => { + const group = accessGroups.find((g) => g.access_group_id === accessGroupId); + return group ? `${group.access_group_name} (${accessGroupId})` : accessGroupId; + }; + const discoveryRequest = useMemo( () => buildDiscoveryRequest(detectedAgentType, watchedFormValues || {}, selectedAgentTypeInfo), [watchedFormValues, selectedAgentTypeInfo, detectedAgentType], @@ -221,6 +238,7 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT await patchAgentCall(accessToken, agentId, { ...updateData, object_permission: buildMcpObjectPermission(values), + access_group_ids: values.access_group_ids ?? [], }); toast.success("Agent updated successfully"); setIsEditing(false); @@ -350,6 +368,17 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {agent.rpm_limit ?? "Unlimited"} {agent.session_tpm_limit ?? "Unlimited"} {agent.session_rpm_limit ?? "Unlimited"} + + {agent.access_group_ids?.length ? ( +
+ {agent.access_group_ids.map((accessGroupId) => ( +
{accessGroupLabel(accessGroupId)}
+ ))} +
+ ) : ( + "None" + )} + {formatDate(agent.created_at)} {formatDate(agent.updated_at)} @@ -489,6 +518,26 @@ const AgentInfoView: React.FC = ({ agentId, onClose, accessT {rateLimitField("session_rpm_limit", "Session RPM Limit")}
+ +

Access Groups

+ + + {({ value, onChange }) => ( + + )} + + +

MCP Servers

diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx index 5c7453c1394..a144630cdd0 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.test.tsx @@ -434,7 +434,7 @@ describe("AutoRouterBenchmarksTab", () => { mockHook({ data: response([group()]) }); const { dateValue, onDateChange } = renderTab(); - expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, undefined); + expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, undefined, undefined); expect(screen.getByText("Jul 6 – Aug 5 (UTC)")).toBeInTheDocument(); fireEvent.click(screen.getByTestId("date-picker")); @@ -460,7 +460,7 @@ describe("AutoRouterBenchmarksTab", () => { , ); - expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, "key-hash-1"); + expect(vi.mocked(useAutoRouterBenchmarks)).toHaveBeenCalledWith("sk-test", dateValue, "key-hash-1", undefined); expect(screen.getByText("Total estimated savings")).toBeInTheDocument(); expect(screen.queryByRole("tab", { name: "Shadow Evals" })).not.toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx index ce55b633b60..063598bd46e 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab.tsx @@ -312,8 +312,7 @@ const BenchmarksBody: React.FC = ({ isPending, error, data, length. Total actual spend includes every turn; savings and baseline spend include only turns with a current estimate, including turns with zero savings. Savings are net of recorded LLM classification cost. Classification cost per 1K turns is averaged over all auto-router turns, including those that skip classification. The range - counts whole sessions that overlap it, so totals can differ slightly from the Overall tab, which buckets savings - by UTC day. + counts whole sessions that overlap it, so totals can differ from savings views that group usage by UTC day.

@@ -333,11 +332,17 @@ interface AutoRouterBenchmarksTabProps { accessToken: string | null; activity: Pick; apiKey?: string; + userId?: string; } -export const AutoRouterUsageView: React.FC = ({ accessToken, activity, apiKey }) => { +export const AutoRouterUsageView: React.FC = ({ + accessToken, + activity, + apiKey, + userId, +}) => { const { dateValue, onDateChange } = activity; - const { data, isPending, error } = useAutoRouterBenchmarks(accessToken, dateValue, apiKey); + const { data, isPending, error } = useAutoRouterBenchmarks(accessToken, dateValue, apiKey, userId); const [selectedKey, setSelectedKey] = useState(ALL_ROUTERS); const { data: autoRouters } = useAutoRouters(); @@ -372,6 +377,12 @@ export const AutoRouterUsageView: React.FC = ({ ac
+ {userId && ( +

+ Usage for this user across API keys and JWT-authenticated requests. Older sessions recorded without a user ID + are not included. +

+ )} +export const useAutoRouterBenchmarks = ( + accessToken: string | null, + range: DateRange, + apiKey?: string, + userId?: string, +) => $api.useQuery( "get", "/auto_router/benchmarks", - { params: { query: { ...benchmarksWindow(range, new Date()), api_key: apiKey } } }, + { params: { query: { ...benchmarksWindow(range, new Date()), api_key: apiKey, user_id: userId } } }, { enabled: Boolean(accessToken && range.from && range.to), retry: false }, ); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx index 4059303d5a5..e501cf00b90 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.test.tsx @@ -15,6 +15,8 @@ vi.mock("@/app/(dashboard)/usage/_components/hooks/usePaginatedDailyActivity", ( isFetchingMore: false, progress: { currentPage: 4, totalPages: 9 }, cancelled: false, + failed: false, + coversRange: true, cancel: mockCancel, }; }, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts index 92dd24b8d6d..4eb9f257d30 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/useDailyActivityRange.ts @@ -67,14 +67,16 @@ export const useScopedDailyActivityRange = ( args: [accessToken, startTime, endTime, userId, true, apiKey], enabled: !!accessToken && !!startTime && !!endTime, }; - const { data, loading, isFetchingMore, progress, cancelled, failed, cancel } = + const { data, loading, isFetchingMore, progress, cancelled, failed, coversRange, cancel } = usePaginatedDailyActivity(activityQueryOptions); + const readUnavailable = failed || cancelled; + const waitingForRange = activityQueryOptions.enabled && !coversRange && !readUnavailable; return { dateValue, onDateChange, results: data.results as DailyData[], - loading, + loading: loading || waitingForRange, isFetchingMore, progress, cancelled, diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx index 0f1a44851c7..6a0e55a6dda 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.integration.test.tsx @@ -1,7 +1,17 @@ -import { fireEvent, renderWithProviders as render, screen, waitFor } from "../../../../../../tests/test-utils"; +import { + act, + fireEvent, + renderWithProviders as render, + screen, + testQueryClient, + waitFor, +} from "../../../../../../tests/test-utils"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; -import { describe, expect, it, vi, beforeEach } from "vitest"; +import { describe, expect, it, vi, beforeEach, afterEach } from "vitest"; +import { Profiler } from "react"; import UserInfoView from "./user_info_view"; +import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; +import type { AutoRouterBenchmarksResponse } from "@/app/(dashboard)/cost-optimization/_components/autoRouterBenchmarks"; const mockTeamMemberAddCall = vi.fn(); const mockTeamMemberDeleteCall = vi.fn(); @@ -11,6 +21,8 @@ const mockTeamInfoCall = vi.fn(); const mockUserUpdateUserCall = vi.fn(); const mockFetchMCPServers = vi.fn(); const mockListMCPTools = vi.fn(); +const mockUserDailyActivityCall = vi.fn(); +const mockUserDailyActivityAggregatedCall = vi.fn(); const MCP_SERVER = { server_id: "srv-1", server_name: "GitHub MCP", alias: "GitHub MCP" }; @@ -47,13 +59,18 @@ vi.mock("next/navigation", () => ({ useSearchParams: () => new URLSearchParams(window.location.search), })); -vi.mock("@/components/networking", () => { +vi.mock("@/components/networking", async (importOriginal) => { + const original = await importOriginal(); return { + formatDate: original.formatDate, serverRootPath: "/", userGetInfoV2: (...args: unknown[]) => mockUserGetInfoV2(...args), + userDailyActivityCall: (...args: unknown[]) => mockUserDailyActivityCall(...args), + userDailyActivityAggregatedCall: (...args: unknown[]) => mockUserDailyActivityAggregatedCall(...args), userDeleteCall: vi.fn(), userUpdateUserCall: (...args: unknown[]) => mockUserUpdateUserCall(...args), modelAvailableCall: vi.fn().mockResolvedValue({ data: [] }), + modelInfoCall: vi.fn().mockResolvedValue({ data: [], total_pages: 1 }), invitationCreateCall: vi.fn(), teamInfoCall: (...args: unknown[]) => mockTeamInfoCall(...args), teamListCall: (...args: unknown[]) => mockTeamListCall(...args), @@ -291,3 +308,337 @@ describe("UserInfoView add-to-team form", () => { expect(screen.getByText("Add User to Team")).toBeInTheDocument(); }); }); + +const savingsDay = (date: string, metrics: Partial): DailyData => ({ + date, + metrics: { + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 1, + successful_requests: 1, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + ...metrics, + }, + breakdown: { models: {}, model_groups: {}, mcp_servers: {}, providers: {}, api_keys: {}, entities: {} }, +}); + +const savingsResponse = (results: DailyData[]) => ({ + results, + metadata: { total_pages: 1, has_more: false, page: 1 }, +}); + +const routerUsageResponse = (saved: number): AutoRouterBenchmarksResponse => ({ + start_date: "2026-09-01", + end_date: "2026-09-19", + routers_in_scope: 0, + groups: [], + totals: { + sessions: 2, + turns: 2, + avg_turns_per_session: 1, + avg_session_seconds: 0, + avg_tokens_per_session: 100, + spend: 10, + savings_estimated_turns: 2, + savings_estimated_actual_spend: 10, + classifier_cost: 0, + saved_spend: saved, + baseline_spend: 10 + saved, + saved_pct: (100 * saved) / (10 + saved), + saved_per_session: saved / 2, + cache: { + coverage_pct: 100, + hit_rate_pct: 0, + same_model: { turns: 0, hits: 0, hit_rate_pct: 0 }, + first_visit: { turns: 2, hits: 0, hit_rate_pct: 0 }, + return_to_tier: { turns: 0, hits: 0, hit_rate_pct: 0 }, + unordered_turns: 0, + return_misses_expired: 0, + return_misses_within_ttl: 0, + return_misses_unknown: 0, + ttl_5m_turns: 0, + ttl_1h_turns: 0, + }, + }, +}); + +describe("UserInfoView auto-router usage", () => { + const props = { + userId: "user-123", + onClose: vi.fn(), + accessToken: "admin-token", + userRole: "proxy_admin", + possibleUIRoles: null, + }; + const mockFetch = vi.fn(); + + beforeEach(() => { + testQueryClient.clear(); + vi.clearAllMocks(); + mockUserGetInfoV2.mockImplementation((_token: string, userId: string) => + Promise.resolve({ ...MOCK_USER_DATA_NO_TEAMS, user_id: userId }), + ); + mockFetch.mockReset().mockResolvedValue(Response.json(routerUsageResponse(42))); + vi.stubGlobal("fetch", mockFetch); + }); + + afterEach(() => { + testQueryClient.clear(); + vi.unstubAllGlobals(); + }); + + it.each(["proxy_admin", "proxy_admin_viewer"])( + "loads selected-user usage lazily for %s without a key filter", + async (userRole) => { + const user = userEvent.setup(); + render(); + const tab = await screen.findByRole("tab", { name: "Auto-router usage" }); + expect(mockFetch).not.toHaveBeenCalled(); + await user.click(tab); + + expect(await screen.findByText("$42.00")).toBeInTheDocument(); + const request = mockFetch.mock.calls[0][0] as Request; + const params = new URL(request.url).searchParams; + expect(params.get("user_id")).toBe("user-123"); + expect(params.has("api_key")).toBe(false); + expect(screen.getByText(/Older sessions recorded without a user ID are not included/)).toBeInTheDocument(); + }, + ); + + it("switches query scope without displaying the previous user's usage", async () => { + const nextUser = Promise.withResolvers(); + mockFetch.mockResolvedValueOnce(Response.json(routerUsageResponse(42))).mockReturnValue(nextUser.promise); + const user = userEvent.setup(); + const { rerender } = render(); + await user.click(await screen.findByRole("tab", { name: "Auto-router usage" })); + expect(await screen.findByText("$42.00")).toBeInTheDocument(); + + rerender(); + expect(screen.getByText("Loading auto-router usage...")).toBeInTheDocument(); + expect(screen.queryByText("$42.00")).not.toBeInTheDocument(); + await act(async () => nextUser.resolve(Response.json(routerUsageResponse(-7)))); + expect(await screen.findByText("-$7.00")).toBeInTheDocument(); + expect( + mockFetch.mock.calls.map(([request]) => new URL((request as Request).url).searchParams.get("user_id")), + ).toEqual(["user-123", "user-456"]); + }); + + it.each(["internal_user", "org_admin", null])("keeps the admin-only tab unavailable to %s", async (userRole) => { + render(); + await screen.findByRole("tab", { name: "Overview" }); + expect(screen.queryByRole("tab", { name: "Auto-router usage" })).not.toBeInTheDocument(); + expect(mockFetch).not.toHaveBeenCalled(); + }); + + it("never turns an absent user ID into a deployment-wide request", async () => { + const user = userEvent.setup(); + render(); + await user.click(await screen.findByRole("tab", { name: "Auto-router usage" })); + expect(screen.getByRole("alert")).toHaveTextContent("this user has no ID"); + expect(mockFetch).not.toHaveBeenCalled(); + }); +}); + +describe("UserInfoView savings", () => { + const props = { + userId: "user-123", + onClose: vi.fn(), + accessToken: "admin-token", + userRole: "proxy_admin", + possibleUIRoles: null, + }; + + beforeEach(() => { + vi.clearAllMocks(); + mockUserGetInfoV2.mockImplementation((_token: string, userId: string) => + Promise.resolve({ ...MOCK_USER_DATA_NO_TEAMS, user_id: userId }), + ); + mockUserDailyActivityAggregatedCall.mockReset().mockResolvedValue(savingsResponse([])); + mockUserDailyActivityCall.mockReset().mockResolvedValue(savingsResponse([])); + }); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it.each(["internal_user", "org_admin", "team_admin"])( + "only offers self savings to %s and stops querying after switching to another user", + async (userRole) => { + const user = userEvent.setup(); + const { rerender } = render(); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + expect(await screen.findByText("No usage recorded for this user in this range.")).toBeInTheDocument(); + expect(mockUserDailyActivityAggregatedCall.mock.calls[0][3]).toBe("user-1"); + + mockUserDailyActivityAggregatedCall.mockClear(); + mockUserDailyActivityCall.mockClear(); + rerender(); + await screen.findAllByText("another-user"); + expect(screen.getByRole("tab", { name: "Overview" })).toHaveAttribute("aria-selected", "true"); + expect(screen.queryByRole("tab", { name: "Savings" })).not.toBeInTheDocument(); + expect(screen.queryByText("No usage recorded for this user in this range.")).not.toBeInTheDocument(); + expect(mockUserDailyActivityAggregatedCall).not.toHaveBeenCalled(); + expect(mockUserDailyActivityCall).not.toHaveBeenCalled(); + }, + ); + + it("loads selected user savings without a key filter, including losses", async () => { + const firstDay: Partial = { + compression_savings_spend: 1.5, + gateway_injected_caching_savings_spend: 0.1, + prompt_caching_savings_spend: 0.25, + autorouter_savings_spend: -1, + }; + const secondDay: Partial = { + compression_savings_spend: 0.5, + gateway_injected_caching_savings_spend: 0.3, + prompt_caching_savings_spend: 0.75, + autorouter_savings_spend: -2, + }; + mockUserDailyActivityAggregatedCall.mockResolvedValue( + savingsResponse([savingsDay("2026-09-18", firstDay), savingsDay("2026-09-19", secondDay)]), + ); + const user = userEvent.setup(); + render(); + const savingsTab = await screen.findByRole("tab", { name: "Savings" }); + expect(mockUserDailyActivityAggregatedCall).not.toHaveBeenCalled(); + expect(mockUserDailyActivityCall).not.toHaveBeenCalled(); + + await user.click(savingsTab); + + expect(await screen.findByTestId("summary-card-total-recorded-savings")).toHaveTextContent("-$0.6000"); + expect(screen.getByTestId("summary-card-compression-savings")).toHaveTextContent("$2.00"); + expect(screen.getByTestId("summary-card-prompt-caching-savings")).toHaveTextContent("$0.4000"); + expect(screen.getByTestId("summary-card-prompt-caching-savings")).toHaveTextContent("$1.00Total"); + expect(screen.getByTestId("summary-card-auto-router-savings")).toHaveTextContent("-$3.00"); + expect(mockUserDailyActivityAggregatedCall).toHaveBeenCalledExactlyOnceWith( + "admin-token", + expect.any(Date), + expect.any(Date), + "user-123", + true, + null, + ); + expect(screen.getByTestId("user-savings-scope-note")).toHaveTextContent("JWT-authenticated requests"); + await user.click(screen.getByRole("tab", { name: "Per day" })); + expect(screen.getByRole("tab", { name: "Per day" })).toHaveAttribute("aria-selected", "true"); + expect(screen.getByTestId("summary-card-total-recorded-savings")).toHaveTextContent("-$0.6000"); + }); + + it("removes the prior user's savings while the newly selected user's results are loading", async () => { + const nextUser = Promise.withResolvers>(); + mockUserDailyActivityAggregatedCall + .mockResolvedValueOnce(savingsResponse([savingsDay("2026-09-19", { compression_savings_spend: 42 })])) + .mockReturnValueOnce(nextUser.promise); + const user = userEvent.setup(); + const { rerender } = render(); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + expect(await screen.findByTestId("summary-card-total-recorded-savings")).toHaveTextContent("$42.00"); + + rerender(); + + expect(await screen.findByTestId("user-savings-empty")).toHaveTextContent("Loading savings"); + expect(screen.queryByTestId("summary-card-total-recorded-savings")).not.toBeInTheDocument(); + expect(mockUserDailyActivityAggregatedCall).toHaveBeenLastCalledWith( + "admin-token", + expect.any(Date), + expect.any(Date), + "user-456", + true, + null, + ); + await act(async () => { + nextUser.resolve(savingsResponse([savingsDay("2026-09-19", { autorouter_savings_spend: -7 })])); + }); + expect(await screen.findByTestId("summary-card-total-recorded-savings")).toHaveTextContent("-$7.00"); + expect(screen.queryByText("$42.00")).not.toBeInTheDocument(); + }); + + it("never commits the previous range's savings under the newly selected dates", async () => { + vi.stubGlobal("requestIdleCallback", (callback: IdleRequestCallback) => + window.setTimeout(() => callback({ didTimeout: false, timeRemaining: () => 0 }), 0), + ); + const nextRange = Promise.withResolvers>(); + mockUserDailyActivityAggregatedCall + .mockResolvedValueOnce(savingsResponse([savingsDay("2026-09-19", { compression_savings_spend: 42 })])) + .mockReturnValue(nextRange.promise); + const committedTotals: Array = []; + const captureNewRange = () => { + if (screen.queryByText("Running total saved · Sep 1 – Sep 2 (UTC)")) { + committedTotals.push(screen.queryByTestId("summary-card-total-recorded-savings")?.textContent ?? null); + } + }; + const user = userEvent.setup(); + render( + + + , + ); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + expect(await screen.findByTestId("summary-card-total-recorded-savings")).toHaveTextContent("$42.00"); + + await user.click(screen.getByRole("button", { name: / - / })); + const [startDateInput, endDateInput] = screen.getAllByDisplayValue(/^\d{4}-\d{2}-\d{2}$/); + fireEvent.change(startDateInput, { target: { value: "2026-09-01" } }); + fireEvent.change(endDateInput, { target: { value: "2026-09-02" } }); + await user.click(screen.getByRole("button", { name: "Apply" })); + + expect(committedTotals.length).toBeGreaterThan(0); + expect(committedTotals.every((total) => total === null)).toBe(true); + expect(screen.getByTestId("user-savings-empty")).toHaveTextContent("Loading savings"); + await act(async () => { + nextRange.resolve(savingsResponse([savingsDay("2026-09-02", { autorouter_savings_spend: -7 })])); + }); + expect(await screen.findByTestId("summary-card-total-recorded-savings")).toHaveTextContent("-$7.00"); + }); + + it("reports an incomplete paginated read as unavailable instead of displaying a partial savings total", async () => { + mockUserDailyActivityAggregatedCall.mockRejectedValue(new Error("aggregated unavailable")); + mockUserDailyActivityCall + .mockResolvedValueOnce({ + results: [savingsDay("2026-09-19", { compression_savings_spend: 42 })], + metadata: { total_pages: 2, has_more: true, page: 1 }, + }) + .mockRejectedValueOnce(new Error("next page unavailable")); + const user = userEvent.setup(); + render(); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + + expect(await screen.findByRole("alert")).toHaveTextContent("Savings are unavailable for this range"); + expect(mockUserDailyActivityCall).toHaveBeenLastCalledWith( + "admin-token", + expect.any(Date), + expect.any(Date), + 2, + "user-123", + true, + null, + ); + expect(screen.queryByTestId("summary-card-total-recorded-savings")).not.toBeInTheDocument(); + expect(screen.queryByText(/No usage recorded/)).not.toBeInTheDocument(); + }); + + it("distinguishes a user with no usage from an unavailable read", async () => { + const user = userEvent.setup(); + render(); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + + expect(await screen.findByTestId("user-savings-empty")).toHaveTextContent("No usage recorded for this user"); + expect(screen.getByTestId("summary-card-total-recorded-savings")).toHaveTextContent("$0.00"); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); + + it.each(["", " "])("never queries an absent selected user ID (%j)", async (userId) => { + const user = userEvent.setup(); + render(); + await user.click(await screen.findByRole("tab", { name: "Savings" })); + + expect(screen.getByRole("alert")).toHaveTextContent("this user has no ID"); + expect(mockUserDailyActivityAggregatedCall).not.toHaveBeenCalled(); + expect(mockUserDailyActivityCall).not.toHaveBeenCalled(); + }); +}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx index e083e549552..c95badc587a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/users/_components/view_users/user_info_view.tsx @@ -28,7 +28,7 @@ import { ComboboxList, } from "@/components/ui/combobox"; import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; -import { rolesWithWriteAccess } from "@/utils/roles"; +import { hasProxyWideSpendView, rolesWithWriteAccess } from "@/utils/roles"; import { teamDetailHref } from "@/utils/entityLinks"; import { BadgeLink } from "@/components/shared/BadgeLink"; import { UserEditView } from "../user_edit_view"; @@ -44,6 +44,9 @@ import { useMCPServers } from "@/app/(dashboard)/hooks/mcpServers/useMCPServers" import { useMCPToolsets } from "@/app/(dashboard)/hooks/mcpServers/useMCPToolsets"; import { extractMcpEntitlement } from "@/components/mcp_server_management/mcpEntitlement"; import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/ui/dialog"; +import ScopedSavingsTab from "@/components/shared/ScopedSavingsTab"; +import { AutoRouterUsageView } from "@/app/(dashboard)/cost-optimization/_components/AutoRouterBenchmarksTab"; +import { useActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; interface UserInfoViewProps { userId: string; @@ -85,7 +88,10 @@ export default function UserInfoView({ initialTab = 0, startInEditMode = false, }: UserInfoViewProps) { - const { premiumUser } = useAuthorized(); + const { premiumUser, userId: signedInUserId } = useAuthorized(); + const canViewAutoRouterUsage = hasProxyWideSpendView(userRole); + const canViewSavings = canViewAutoRouterUsage || (Boolean(userId.trim()) && userId === signedInUserId); + const activityDateRange = useActivityDateRange(); const [userData, setUserData] = useState(null); const [teamDetails, setTeamDetails] = useState([]); const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false); @@ -97,6 +103,8 @@ export default function UserInfoView({ const [invitationLinkData, setInvitationLinkData] = useState(null); const [baseUrl, setBaseUrl] = useState(null); const [activeTab, setActiveTab] = useState(initialTab === 1 ? "details" : "overview"); + const hiddenSavingsTab = activeTab === "savings" && !canViewSavings; + const hiddenRouterTab = activeTab === "auto-router-usage" && !canViewAutoRouterUsage; const [copiedStates, setCopiedStates] = useState>({}); const [isTeamsExpanded, setIsTeamsExpanded] = useState(false); const [isAddTeamModalOpen, setIsAddTeamModalOpen] = useState(false); @@ -467,7 +475,11 @@ export default function UserInfoView({ confirmLoading={isDeletingUser} /> - setActiveTab(String(v))} className="gap-0"> + setActiveTab(String(v))} + className="gap-0" + > Overview @@ -475,6 +487,16 @@ export default function UserInfoView({ Details + {canViewSavings && ( + + Savings + + )} + {canViewAutoRouterUsage && ( + + Auto-router usage + + )} {/* Overview Panel */} @@ -685,6 +707,38 @@ export default function UserInfoView({ )} + {canViewSavings && ( + + {activeTab === "savings" && + (userId.trim() ? ( + + ) : ( +

Savings are unavailable because this user has no ID.

+ ))} + + )} + {canViewAutoRouterUsage && ( + + {activeTab === "auto-router-usage" && + (userId.trim() ? ( + + ) : ( +

Auto-router usage is unavailable because this user has no ID.

+ ))} + + )} void; -}> = ({ value, classifierType, onTypeChange }) => { - const scorerLocked = Boolean(value.custom_tier_set); - const scorerLockedReason = restrictedBy(value, "heuristicClassifier")?.reason; - return ( - onTypeChange(classifierType as ClassifierType)} - className="w-full" - > -
- - - - - - - - - - - - - - -
- - ); -}; - const ClassificationMethodConfig: React.FC = ({ value, onChange, @@ -467,6 +391,10 @@ const ClassificationMethodConfig: React.FC = ({ <> + {classifierType === "custom" && ( + + )} + {classifierType === "heuristic_v2" && (
diff --git a/ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx b/ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx index c0a65076d20..ad09efd8059 100644 --- a/ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ContextWindowEscalationConfig.tsx @@ -7,7 +7,7 @@ const ContextWindowEscalationConfig: React.FC<{ value: ComplexityRouterConfigValue; onChange: (value: ComplexityRouterConfigValue) => void; }> = ({ value, onChange }) => { - const enabled = value.enable_context_window_escalation ?? true; + const enabled = value.enable_context_window_escalation ?? false; // A number input renders Number("0.") as "0", so a decimal cannot be typed without a local draft. const [bufferDraft, setBufferDraft] = React.useState(null); const commitBuffer = (raw: string) => { @@ -32,7 +32,8 @@ const ContextWindowEscalationConfig: React.FC<{
When a prompt provably cannot fit the decided tier's context windows, route it to the lowest tier whose - window holds it instead of letting the provider reject it. Off means requests dispatch on complexity alone. + window holds it instead of letting the provider reject it. Disabled by default. Off means requests dispatch on + complexity alone. {enabled && (
diff --git a/ui/litellm-dashboard/src/components/add_model/HeuristicKeywordOverrides.tsx b/ui/litellm-dashboard/src/components/add_model/HeuristicKeywordOverrides.tsx new file mode 100644 index 00000000000..185f187bbfc --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/HeuristicKeywordOverrides.tsx @@ -0,0 +1,42 @@ +import React from "react"; +import { MultiSelect } from "@/components/shared/MultiSelect"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; + +const fields = [ + ["code_keywords", "Code keywords"], + ["reasoning_keywords", "Reasoning keywords"], + ["technical_keywords", "Technical keywords"], + ["simple_keywords", "Simple keywords"], +] as const; + +const HeuristicKeywordOverrides: React.FC<{ + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; +}> = ({ value, onChange }) => ( +
+

+ Each list replaces the built-in keyword list of the same name for the heuristic scorer. Leave a list empty to keep + the built-in one. To add technical terms without replacing the list, use custom technical keywords under + Classification Method. +

+ {fields.map(([key, label]) => { + const keywords = value[key] ?? []; + return ( +
+ {label} + ({ label: keyword, value: keyword }))} + value={keywords} + onValueChange={(next) => onChange({ ...value, [key]: next.length > 0 ? next : undefined })} + placeholder={`Add ${label.toLowerCase()}`} + emptyText="Type to add a keyword" + allowCustomValues + className="w-full" + /> +
+ ); + })} +
+); + +export default HeuristicKeywordOverrides; diff --git a/ui/litellm-dashboard/src/components/add_model/HousekeepingRoutingControls.tsx b/ui/litellm-dashboard/src/components/add_model/HousekeepingRoutingControls.tsx new file mode 100644 index 00000000000..2b6caba6291 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/HousekeepingRoutingControls.tsx @@ -0,0 +1,44 @@ +import React from "react"; +import { MultiSelect } from "@/components/shared/MultiSelect"; +import { Switch } from "@/components/ui/switch"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; + +const HousekeepingRoutingControls: React.FC<{ + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; +}> = ({ value, onChange }) => { + const enabled = value.route_housekeeping_to_cheapest_tier ?? true; + const patterns = value.housekeeping_patterns ?? []; + return ( + <> +
+ onChange({ ...value, route_housekeeping_to_cheapest_tier: next })} + aria-label="Route housekeeping calls to the cheapest tier" + /> + Route housekeeping calls to the cheapest tier +
+ + Conversation-title style calls skip the classifier and go to the cheapest tier. + + Additional housekeeping sentinels + ({ label: pattern, value: pattern }))} + value={patterns} + onValueChange={(next) => onChange({ ...value, housekeeping_patterns: next.length > 0 ? next : undefined })} + placeholder="e.g., conversation title" + emptyText="Type to add a sentinel" + allowCustomValues + disabled={!enabled} + className="w-full" + /> + + Case-sensitive literal strings added to the built-in conversation-title sentinels. + {!enabled && " Turn housekeeping routing on for these to take effect."} + + + ); +}; + +export default HousekeepingRoutingControls; diff --git a/ui/litellm-dashboard/src/components/add_model/PlanModeOverrideControls.tsx b/ui/litellm-dashboard/src/components/add_model/PlanModeOverrideControls.tsx index 83fd0c42d16..e195fb62222 100644 --- a/ui/litellm-dashboard/src/components/add_model/PlanModeOverrideControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/PlanModeOverrideControls.tsx @@ -1,5 +1,6 @@ import React from "react"; import { Switch } from "@/components/ui/switch"; +import { MultiSelect } from "@/components/shared/MultiSelect"; import TierRowSelect from "./TierRowSelect"; import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; @@ -38,6 +39,23 @@ const PlanModeOverrideControls: React.FC<{ />
)} +
+ Additional plan-mode sentinels + ({ label: pattern, value: pattern }))} + value={value.plan_mode_patterns ?? []} + onValueChange={(patterns) => + onChange({ ...value, plan_mode_patterns: patterns.length > 0 ? patterns : undefined }) + } + placeholder="e.g., enter plan mode" + emptyText="Type to add a sentinel" + allowCustomValues + className="w-full" + /> + + Case-sensitive literal strings added to the built-in Claude Code and Copilot plan-mode markers. + +
); diff --git a/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx b/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx new file mode 100644 index 00000000000..c7f9ee48e0b --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/ReminderMarkers.tsx @@ -0,0 +1,82 @@ +import React from "react"; +import { Plus, Trash2 } from "lucide-react"; +import { Button } from "@/components/ui/button"; +import { Input } from "@/components/ui/input"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import { getReminderMarkersError, type ReminderMarkerPair } from "./build_complexity_router_config"; + +const ReminderMarkers: React.FC<{ + value: ComplexityRouterConfigValue; + onChange: (value: ComplexityRouterConfigValue) => void; + showValidationErrors?: boolean; +}> = ({ value, onChange, showValidationErrors = false }) => { + const markers = value.reminder_markers ?? []; + const update = (index: number, patch: Partial) => + onChange({ + ...value, + reminder_markers: markers.map((marker, markerIndex) => + markerIndex === index ? { ...marker, ...patch } : marker, + ), + }); + const remove = (index: number) => { + const next = markers.filter((_, markerIndex) => markerIndex !== index); + onChange({ ...value, reminder_markers: next.length > 0 ? next : undefined }); + }; + const error = getReminderMarkersError(value.reminder_markers); + return ( +
+

+ Delimiter pairs that wrap harness-injected reminder blocks, which are stripped before classification. Setting + any pair replaces the built-in pairs, so list every pair your harness emits. Matching is case-insensitive and + values are saved lowercased. +

+
+ {markers.map((marker, index) => ( +
+
+ + update(index, { open: event.target.value })} + /> +
+
+ + update(index, { close: event.target.value })} + /> +
+ +
+ ))} +
+ + {showValidationErrors && error &&

{error}

} +
+ ); +}; + +export default ReminderMarkers; diff --git a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx index 68dd880a684..6e7b1489a4c 100644 --- a/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/ResponseFormatControls.tsx @@ -18,6 +18,18 @@ const ResponseFormatControls: React.FC<{ Return the resolved underlying model name in responses instead of the autorouter alias. +
+ onChange({ ...value, max_tokens_from_tier_model: enabled })} + aria-label="Cap max_tokens at the tier model's output ceiling" + /> + Cap max_tokens at the tier model's output ceiling +
+ + Replace the caller's max_tokens with the routed tier model's output ceiling so one client value fits + every tier. Off forwards the caller's value unchanged. + ); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx index fdff7ad001c..242893e2ae8 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.test.tsx @@ -744,7 +744,7 @@ describe("AddAutoRouterTab", () => { ); }); - it("carries a context-window escalation opt-out through to the create payload", async () => { + it("starts context-window escalation disabled and carries an explicit opt-in to the create payload", async () => { const user = userEvent.setup(); vi.mocked(getMissingTiersError).mockReturnValue(null); @@ -754,14 +754,15 @@ describe("AddAutoRouterTab", () => { expandDetailedConfiguration(); await user.click(screen.getByText("Advanced: Context Window Escalation")); const toggle = await screen.findByRole("switch", { name: "Escalate oversized prompts to a tier that fits" }); - expect(toggle).toBeChecked(); + expect(toggle).not.toBeChecked(); + expect(screen.queryByLabelText("Window fit buffer")).not.toBeInTheDocument(); await user.click(toggle); await user.click(screen.getByRole("button", { name: /add auto router/i })); await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); expect(vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config).toMatchObject({ - enable_context_window_escalation: false, + enable_context_window_escalation: true, }); }); @@ -774,6 +775,7 @@ describe("AddAutoRouterTab", () => { await user.type(screen.getByPlaceholderText(/smart_router/i), "ctx-buffer-router"); expandDetailedConfiguration(); await user.click(screen.getByText("Advanced: Context Window Escalation")); + await user.click(screen.getByRole("switch", { name: "Escalate oversized prompts to a tier that fits" })); const buffer = await screen.findByLabelText("Window fit buffer"); fireEvent.change(buffer, { target: { value: "1.5" } }); fireEvent.blur(buffer, { target: { value: "1.5" } }); @@ -783,7 +785,7 @@ describe("AddAutoRouterTab", () => { await waitFor(() => expect(handleAddAutoRouterSubmit).toHaveBeenCalled()); const config = vi.mocked(handleAddAutoRouterSubmit).mock.calls.at(-1)?.[0].complexity_router_config; expect(config).toMatchObject({ context_window_escalation_buffer: 1 }); - expect(config).not.toHaveProperty("enable_context_window_escalation"); + expect(config).toHaveProperty("enable_context_window_escalation", true); }); it("clearing the buffer removes it from the payload so the router tracks the backend default", async () => { @@ -795,6 +797,7 @@ describe("AddAutoRouterTab", () => { await user.type(screen.getByPlaceholderText(/smart_router/i), "ctx-clear-router"); expandDetailedConfiguration(); await user.click(screen.getByText("Advanced: Context Window Escalation")); + await user.click(screen.getByRole("switch", { name: "Escalate oversized prompts to a tier that fits" })); const buffer = await screen.findByLabelText("Window fit buffer"); fireEvent.change(buffer, { target: { value: "0.8" } }); fireEvent.blur(buffer, { target: { value: "0.8" } }); diff --git a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx index 84e44fee9c3..9be49edf08d 100644 --- a/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx +++ b/ui/litellm-dashboard/src/components/add_model/add_auto_router_tab.tsx @@ -28,10 +28,6 @@ import ComplexityRouterConfig, { effectiveClassifierType, usesLlmClassifier, heuristicScoringRole, - DEFAULT_ADAPTIVE_WEIGHTS, - DEFAULT_SESSION_AFFINITY, - DEFAULT_DEPLOYMENT_AFFINITY, - DEFAULT_TIER_DISTANCE_PENALTY, } from "./ComplexityRouterConfig"; import { KeywordTierRule } from "./KeywordTierRules"; import { customDimensionsError } from "./custom_dimensions"; @@ -48,6 +44,8 @@ import { getKeywordTierRulesError, getClassifierModelError, getHeuristicV2SuccessThresholdError, + getReminderMarkersError, + getClassifierPluginTimeoutError, getClassifierReasoningEffortError, getMissingTiersError, getPlanModeTierError, @@ -55,6 +53,7 @@ import { getTierLabelsError, dryRunRejection, } from "./build_complexity_router_config"; +import { builderParamsFromValue } from "./complexity_router_builder_params"; import { activeTierName, activeTierRows, getCustomTierRowsError, resolveComplexityDefaultModel } from "./tier_rows"; import { tierRowLabel } from "./complexity_router_tiers"; import { buildAutoRouterTestTargets, AutoRouterTestTarget } from "./build_auto_router_test_targets"; @@ -152,6 +151,8 @@ export const getSubmitBlockedReason = ( getKeywordTierRulesError(keywordTierRules, activeTierRows(config)) ?? getClassifierModelError(config) ?? getHeuristicV2SuccessThresholdError(config.heuristic_v2_success_threshold) ?? + getReminderMarkersError(config.reminder_markers) ?? + getClassifierPluginTimeoutError(config.classifier_type, config.classifier_plugin_timeout_ms) ?? (heuristicScoringRole(config) === "decides" ? customDimensionsError(config.custom_dimensions) : null) ?? getClassifierReasoningEffortError(config, modelInfo) ?? getReferencedModelsError(referencedModelsParams, availability) @@ -399,56 +400,25 @@ const AddAutoRouterTab: React.FC = ({ ); const complexityRouterConfigParams: BuildComplexityRouterConfigParams = { - tiers: complexityRouterConfig.tiers, - enableNonReasoningTier: complexityRouterConfig.enable_non_reasoning_tier, - customTierSet: complexityRouterConfig.custom_tier_set, - defaultModel: complexityRouterConfig.default_model, - planModeMinTier: complexityRouterConfig.plan_mode_min_tier, - classificationPrompt: complexityRouterConfig.classification_prompt, - classificationExamples: complexityRouterConfig.classification_examples, - heuristicFirstMaxTier: complexityRouterConfig.heuristic_first_max_tier, - hybridBoundaryMargin: complexityRouterConfig.hybrid_boundary_margin, - classificationMode: complexityRouterConfig.classification_mode, - tierLabels: complexityRouterConfig.tier_labels, - classifierType: complexityRouterConfig.classifier_type, - jevClassifierConfig: complexityRouterConfig.jev_classifier_config, - heuristicV2SuccessThreshold: complexityRouterConfig.heuristic_v2_success_threshold, - capabilityClassifierConfig: complexityRouterConfig.capability_classifier_config, - llmV2Config: complexityRouterConfig.llm_v2_config, - classifierLlmConfig: complexityRouterConfig.classifier_llm_config, - classifierContextWindowSize: complexityRouterConfig.classifier_context_window_size, - classifierContextBudgetChars: complexityRouterConfig.classifier_context_budget_chars, - classifierContextPerTurnChars: complexityRouterConfig.classifier_context_per_turn_chars, - classifierContextIncludeAssistantTurns: complexityRouterConfig.classifier_context_include_assistant_turns, - classifierFallback: complexityRouterConfig.classifier_fallback, - sessionAffinity: complexityRouterConfig.session_affinity ?? DEFAULT_SESSION_AFFINITY, - modalityRouting: complexityRouterConfig.modality_routing ?? false, - modalityPinOverride: complexityRouterConfig.modality_pin_override ?? false, - deploymentAffinity: complexityRouterConfig.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, + ...builderParamsFromValue(complexityRouterConfig), customTechnicalKeywords, keywordTierRules, semanticMatchingEnabled, embeddingModel, matchThreshold, escalationKeywords, - stallEscalationEnabled: complexityRouterConfig.stall_escalation_enabled, - stallEscalationWindow: complexityRouterConfig.stall_escalation_window, - stallEscalationRepeatThreshold: complexityRouterConfig.stall_escalation_repeat_threshold, - adaptive: complexityRouterConfig.adaptive ?? false, - adaptiveWeights: complexityRouterConfig.adaptive_weights ?? DEFAULT_ADAPTIVE_WEIGHTS, - tierDistancePenalty: complexityRouterConfig.tier_distance_penalty ?? DEFAULT_TIER_DISTANCE_PENALTY, - adaptiveEligible: complexityRouterConfig.adaptive_eligible ?? "all", - returnRawModelName: complexityRouterConfig.return_raw_model_name ?? false, - tierModelParams: complexityRouterConfig.tier_model_params, - tierBoundaries: complexityRouterConfig.tier_boundaries, - tokenThresholds: complexityRouterConfig.token_thresholds, - dimensionWeights: complexityRouterConfig.dimension_weights, - customDimensions: complexityRouterConfig.custom_dimensions, - reasoningOverrideMinScore: complexityRouterConfig.reasoning_override_min_score, - enableContextWindowEscalation: complexityRouterConfig.enable_context_window_escalation, - contextWindowEscalationBuffer: complexityRouterConfig.context_window_escalation_buffer, - sessionAffinityTtlSeconds: complexityRouterConfig.session_affinity_ttl_seconds, }; + const jevRequestParams = + effectiveClassifierType(complexityRouterConfig) === "jev" + ? { + prompt: JEV_CONNECTION_TEST_PROMPT, + config: buildComplexityRouterConfig(complexityRouterConfigParams), + defaultModel: resolveComplexityDefaultModel(complexityRouterConfig, complexityRouterConfig.default_model), + routerName: watchedName, + teamId: requiresTeamScope ? watchedTeamId ?? undefined : undefined, + } + : undefined; + const jevRequest = jevRequestParams ? buildAutoRouterRoutingTestRequest(jevRequestParams) : undefined; const submitRecommendedRouter = async (name: string) => { // The one answer the submit button reads, so a disabled button and a refused submit cannot @@ -857,20 +827,7 @@ const AddAutoRouterTab: React.FC = ({ testId={connectionTestId} accessToken={accessToken} targets={testTargets} - jevRequest={ - effectiveClassifierType(complexityRouterConfig) === "jev" - ? buildAutoRouterRoutingTestRequest({ - prompt: JEV_CONNECTION_TEST_PROMPT, - config: buildComplexityRouterConfig(complexityRouterConfigParams), - defaultModel: resolveComplexityDefaultModel( - complexityRouterConfig, - complexityRouterConfig.default_model, - ), - routerName: watchedName, - teamId: requiresTeamScope ? watchedTeamId ?? undefined : undefined, - }) - : undefined - } + jevRequest={jevRequest} onTestComplete={() => setIsTestingConnection(false)} /> diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts index 054aba7a6aa..6d458d61797 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.test.ts @@ -6,6 +6,8 @@ import { getKeywordTierRulesError, getClassifierModelError, getHeuristicV2SuccessThresholdError, + getReminderMarkersError, + getClassifierPluginTimeoutError, getClassifierReasoningEffortError, getMissingTiersError, hydrateCustomTierSet, @@ -166,7 +168,7 @@ describe("buildComplexityRouterConfig", () => { }); it.each(["capability", "llm_v2", "heuristic"] as const)( - "disables the removed overrides only for forecast creates: %s", + "preserves explicit context-window opt-in beside forecast restrictions: %s", (classifierType) => { const forecast = classifierType !== "heuristic"; const params = { @@ -178,14 +180,10 @@ describe("buildComplexityRouterConfig", () => { }; const config = buildComplexityRouterConfig(params); expect(config.adaptive).toBe(!forecast); - expect(config.enable_context_window_escalation).toBe(!forecast); + expect(config.enable_context_window_escalation).toBe(true); + expect(config.context_window_escalation_buffer).toBe(0.9); expect(config.escalation_keywords).toEqual(forecast ? [] : ["LITELLM ESCALATE"]); - for (const key of [ - "adaptive_weights", - "adaptive_eligible", - "tier_distance_penalty", - "context_window_escalation_buffer", - ]) { + for (const key of ["adaptive_weights", "adaptive_eligible", "tier_distance_penalty"]) { expect(Object.hasOwn(config, key)).toBe(!forecast); } if (forecast) { @@ -224,13 +222,14 @@ describe("buildComplexityRouterConfig", () => { expect(config).toEqual(expected); }); - it("carries an explicit context-window escalation opt-out and buffer, false included", () => { + it.each([undefined, false, true])("preserves the context-window escalation setting: %s", (enabled) => { const config = buildComplexityRouterConfig({ ...baseParams, - enableContextWindowEscalation: false, + enableContextWindowEscalation: enabled, contextWindowEscalationBuffer: 0.9, }); - expect(config.enable_context_window_escalation).toBe(false); + expect(config.enable_context_window_escalation).toBe(enabled); + expect(Object.hasOwn(config, "enable_context_window_escalation")).toBe(enabled !== undefined); expect(config.context_window_escalation_buffer).toBe(0.9); }); @@ -1482,3 +1481,75 @@ describe("classifier vision wire payload", () => { expect(payload.classifier_llm_config).not.toHaveProperty("vision"); }); }); + +describe("advanced complexity router fields", () => { + it("normalizes lists, reminder markers, and explicit false values", () => { + const payload = buildComplexityRouterConfig({ + ...baseParams, + codeKeywords: [" async ", " "], + reasoningKeywords: ["prove"], + technicalKeywords: ["api"], + simpleKeywords: ["hello"], + planModePatterns: [" plan "], + routeHousekeepingToCheapestTier: false, + housekeepingPatterns: [" title "], + reminderMarkers: [{ open: " ", close: " " }], + maxTokensFromTierModel: false, + classifierType: "custom", + classifierPluginTimeoutMs: 3000, + }); + expect(payload).toMatchObject({ + code_keywords: ["async"], + reasoning_keywords: ["prove"], + technical_keywords: ["api"], + simple_keywords: ["hello"], + plan_mode_patterns: ["plan"], + route_housekeeping_to_cheapest_tier: false, + housekeeping_patterns: ["title"], + reminder_markers: [{ open: "", close: "" }], + max_tokens_from_tier_model: false, + classifier_plugin_timeout_ms: 3000, + }); + }); + + it("omits defaults, empty lists, and timeout values for non-custom classifiers", () => { + const payload = buildComplexityRouterConfig({ + ...baseParams, + codeKeywords: [" ", ""], + reminderMarkers: [], + routeHousekeepingToCheapestTier: true, + maxTokensFromTierModel: true, + classifierPluginTimeoutMs: 3000, + }); + expect(payload).not.toHaveProperty("code_keywords"); + expect(payload).not.toHaveProperty("reminder_markers"); + expect(payload).not.toHaveProperty("route_housekeeping_to_cheapest_tier"); + expect(payload).not.toHaveProperty("max_tokens_from_tier_model"); + expect(payload).not.toHaveProperty("classifier_plugin_timeout_ms"); + }); + + it.each([ + "code_keywords", + "reasoning_keywords", + "technical_keywords", + "simple_keywords", + "plan_mode_patterns", + "route_housekeeping_to_cheapest_tier", + "housekeeping_patterns", + "reminder_markers", + "max_tokens_from_tier_model", + "classifier_plugin_timeout_ms", + ])("omits unset advanced field %s", (key) => { + const payload = buildComplexityRouterConfig(baseParams); + expect(payload).not.toHaveProperty(key); + }); + + it("validates marker pairs and custom classifier timeout", () => { + expect(getReminderMarkersError([{ open: " ", close: " " }])).toContain("different"); + expect(getReminderMarkersError([{ open: "", close: "" }])).toContain("needs both"); + expect(getReminderMarkersError([{ open: "", close: "" }])).toBeNull(); + expect(getClassifierPluginTimeoutError("custom", 0)).toContain("whole number"); + expect(getClassifierPluginTimeoutError("custom", 3000)).toBeNull(); + expect(getClassifierPluginTimeoutError("heuristic", 0)).toBeNull(); + }); +}); diff --git a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts index 15ae2b4c0b0..9b4f1980f62 100644 --- a/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts +++ b/ui/litellm-dashboard/src/components/add_model/build_complexity_router_config.ts @@ -54,6 +54,10 @@ import { export type ClassifierVisionConfig = { enabled?: boolean; max_images?: number }; export type ClassifierLLMConfigWire = ClassifierLLMConfig & { vision?: ClassifierVisionConfig }; +export interface ReminderMarkerPair { + open: string; + close: string; +} /** * Drop an empty system_prompt so the payload carries an override only when there is one. The @@ -181,6 +185,16 @@ export interface StoredComplexityRouterConfig { stall_escalation_enabled?: unknown; stall_escalation_window?: unknown; stall_escalation_repeat_threshold?: unknown; + code_keywords?: unknown; + reasoning_keywords?: unknown; + technical_keywords?: unknown; + simple_keywords?: unknown; + plan_mode_patterns?: unknown; + route_housekeeping_to_cheapest_tier?: unknown; + housekeeping_patterns?: unknown; + reminder_markers?: unknown; + max_tokens_from_tier_model?: unknown; + classifier_plugin_timeout_ms?: unknown; } export interface BuildComplexityRouterConfigParams { @@ -233,6 +247,16 @@ export interface BuildComplexityRouterConfigParams { enableContextWindowEscalation?: boolean; contextWindowEscalationBuffer?: number; sessionAffinityTtlSeconds?: number; + codeKeywords?: string[]; + reasoningKeywords?: string[]; + technicalKeywords?: string[]; + simpleKeywords?: string[]; + planModePatterns?: string[]; + routeHousekeepingToCheapestTier?: boolean; + housekeepingPatterns?: string[]; + reminderMarkers?: ReminderMarkerPair[]; + maxTokensFromTierModel?: boolean; + classifierPluginTimeoutMs?: number; } /** @@ -302,6 +326,16 @@ export interface ComplexityRouterConfigPayload { enable_context_window_escalation?: boolean; context_window_escalation_buffer?: number; tier_model_configs?: Record; + code_keywords?: string[]; + reasoning_keywords?: string[]; + technical_keywords?: string[]; + simple_keywords?: string[]; + plan_mode_patterns?: string[]; + route_housekeeping_to_cheapest_tier?: boolean; + housekeeping_patterns?: string[]; + reminder_markers?: ReminderMarkerPair[]; + max_tokens_from_tier_model?: boolean; + classifier_plugin_timeout_ms?: number; } export const serializeTierLabels = (tierLabels: ComplexityTierLabels | undefined): ComplexityTierLabels | undefined => { @@ -376,6 +410,26 @@ export const getHeuristicV2SuccessThresholdError = (threshold: number | undefine return validProbability ? null : "Success threshold must be a number between 0 and 1"; }; +export const getReminderMarkersError = (pairs: ReminderMarkerPair[] | undefined): string | null => { + for (const [index, pair] of (pairs ?? []).entries()) { + const open = pair.open.trim().toLowerCase(); + const close = pair.close.trim().toLowerCase(); + if (!open || !close) return `Reminder marker pair ${index + 1} needs both an opening and a closing delimiter`; + if (open === close) return `Reminder marker pair ${index + 1} must use different opening and closing delimiters`; + } + return null; +}; + +export const getClassifierPluginTimeoutError = ( + classifierType: ClassifierType, + timeoutMs: number | undefined, +): string | null => { + if (classifierType !== "custom" || timeoutMs === undefined) return null; + return Number.isInteger(timeoutMs) && timeoutMs > 0 + ? null + : "Classifier plugin timeout must be a whole number of milliseconds greater than 0"; +}; + export const getClassifierModelError = ( config: Pick< ComplexityRouterConfigValue, @@ -640,6 +694,16 @@ export const buildComplexityRouterConfig = ({ enableContextWindowEscalation, contextWindowEscalationBuffer, sessionAffinityTtlSeconds, + codeKeywords, + reasoningKeywords, + technicalKeywords, + simpleKeywords, + planModePatterns, + routeHousekeepingToCheapestTier, + housekeepingPatterns, + reminderMarkers, + maxTokensFromTierModel, + classifierPluginTimeoutMs, }: BuildComplexityRouterConfigParams): ComplexityRouterConfigPayload => { const serializedTierModelConfigs = customTierSet ? serializeTierModelConfigs( @@ -672,6 +736,29 @@ export const buildComplexityRouterConfig = ({ }; const effectiveType = effectiveClassifierType({ custom_tier_set: customTierSet, classifier_type: classifierType }); const forecast = isForecastClassifier(effectiveType); + const preserveContextWindowBuffer = !forecast || enableContextWindowEscalation === true; + const cleanList = (items: string[] | undefined): string[] | undefined => { + const cleaned = (items ?? []).map((item) => item.trim()).filter(Boolean); + return cleaned.length > 0 ? cleaned : undefined; + }; + const cleanedReminderMarkers = reminderMarkers?.map(({ open, close }) => ({ + open: open.trim().toLowerCase(), + close: close.trim().toLowerCase(), + })); + const cleanedListValues = { + code_keywords: cleanList(codeKeywords), + reasoning_keywords: cleanList(reasoningKeywords), + technical_keywords: cleanList(technicalKeywords), + simple_keywords: cleanList(simpleKeywords), + plan_mode_patterns: cleanList(planModePatterns), + housekeeping_patterns: cleanList(housekeepingPatterns), + }; + const cleanedLists = Object.fromEntries(Object.entries(cleanedListValues).filter(([, list]) => list !== undefined)); + const hasValidCustomClassifierTimeout = + classifierType === "custom" && + classifierPluginTimeoutMs !== undefined && + Number.isInteger(classifierPluginTimeoutMs) && + classifierPluginTimeoutMs > 0; const supportsOpeningPrompt = !customTierSet && !forecast && usesLlmClassifier(effectiveType); const payload: ComplexityRouterConfigPayload = { @@ -729,17 +816,21 @@ export const buildComplexityRouterConfig = ({ adaptive_eligible: adaptiveEligible, }), ...(returnRawModelName && { return_raw_model_name: true }), - // Omission enables the backend default, so hidden forecast controls need an explicit opt-out. ...((forecast || enableContextWindowEscalation !== undefined) && { - enable_context_window_escalation: forecast ? false : enableContextWindowEscalation, + enable_context_window_escalation: enableContextWindowEscalation ?? false, }), - ...(!forecast && + ...(preserveContextWindowBuffer && contextWindowEscalationBuffer !== undefined && { context_window_escalation_buffer: contextWindowEscalationBuffer, }), ...(sessionAffinityTtlSeconds !== undefined && { session_affinity_ttl_seconds: sessionAffinityTtlSeconds, }), + ...cleanedLists, + ...(routeHousekeepingToCheapestTier === false && { route_housekeeping_to_cheapest_tier: false }), + ...(cleanedReminderMarkers && cleanedReminderMarkers.length > 0 && { reminder_markers: cleanedReminderMarkers }), + ...(maxTokensFromTierModel === false && { max_tokens_from_tier_model: false }), + ...(hasValidCustomClassifierTimeout && { classifier_plugin_timeout_ms: classifierPluginTimeoutMs }), ...scorerKnobs, }; if (!customTierSet) return payload; diff --git a/ui/litellm-dashboard/src/components/add_model/classifier_types.ts b/ui/litellm-dashboard/src/components/add_model/classifier_types.ts index ec88166ed2e..aa9d5619052 100644 --- a/ui/litellm-dashboard/src/components/add_model/classifier_types.ts +++ b/ui/litellm-dashboard/src/components/add_model/classifier_types.ts @@ -6,7 +6,8 @@ export type ClassifierType = | "heuristic_first" | "hybrid" | "capability" - | "llm_v2"; + | "llm_v2" + | "custom"; export const usesLlmClassifier = (classifierType: ClassifierType): boolean => (["llm", "heuristic_first", "hybrid", "capability", "llm_v2"] as const).some((type) => type === classifierType); diff --git a/ui/litellm-dashboard/src/components/add_model/complexity_router_builder_params.ts b/ui/litellm-dashboard/src/components/add_model/complexity_router_builder_params.ts new file mode 100644 index 00000000000..124a85ce9a3 --- /dev/null +++ b/ui/litellm-dashboard/src/components/add_model/complexity_router_builder_params.ts @@ -0,0 +1,74 @@ +import type { BuildComplexityRouterConfigParams } from "./build_complexity_router_config"; +import type { ComplexityRouterConfigValue } from "./ComplexityRouterConfig"; +import { + DEFAULT_ADAPTIVE_WEIGHTS, + DEFAULT_DEPLOYMENT_AFFINITY, + DEFAULT_SESSION_AFFINITY, + DEFAULT_TIER_DISTANCE_PENALTY, +} from "./ComplexityRouterConfig"; + +export const builderParamsFromValue = ( + value: ComplexityRouterConfigValue, +): Omit< + BuildComplexityRouterConfigParams, + | "customTechnicalKeywords" + | "keywordTierRules" + | "semanticMatchingEnabled" + | "embeddingModel" + | "matchThreshold" + | "escalationKeywords" +> => ({ + tiers: value.tiers, + enableNonReasoningTier: value.enable_non_reasoning_tier, + customTierSet: value.custom_tier_set, + defaultModel: value.default_model, + planModeMinTier: value.plan_mode_min_tier, + classificationPrompt: value.classification_prompt, + classificationExamples: value.classification_examples, + heuristicFirstMaxTier: value.heuristic_first_max_tier, + hybridBoundaryMargin: value.hybrid_boundary_margin, + classificationMode: value.classification_mode, + tierLabels: value.tier_labels, + classifierType: value.classifier_type, + jevClassifierConfig: value.jev_classifier_config, + heuristicV2SuccessThreshold: value.heuristic_v2_success_threshold, + capabilityClassifierConfig: value.capability_classifier_config, + llmV2Config: value.llm_v2_config, + classifierLlmConfig: value.classifier_llm_config, + classifierContextWindowSize: value.classifier_context_window_size, + classifierContextBudgetChars: value.classifier_context_budget_chars, + classifierContextPerTurnChars: value.classifier_context_per_turn_chars, + classifierContextIncludeAssistantTurns: value.classifier_context_include_assistant_turns, + classifierFallback: value.classifier_fallback, + sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY, + sessionAffinityTtlSeconds: value.session_affinity_ttl_seconds, + modalityRouting: value.modality_routing ?? false, + modalityPinOverride: value.modality_pin_override ?? false, + deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, + adaptive: value.adaptive ?? false, + adaptiveWeights: value.adaptive_weights ?? DEFAULT_ADAPTIVE_WEIGHTS, + tierDistancePenalty: value.tier_distance_penalty ?? DEFAULT_TIER_DISTANCE_PENALTY, + adaptiveEligible: value.adaptive_eligible ?? "all", + returnRawModelName: value.return_raw_model_name ?? false, + tierBoundaries: value.tier_boundaries, + tokenThresholds: value.token_thresholds, + dimensionWeights: value.dimension_weights, + customDimensions: value.custom_dimensions, + reasoningOverrideMinScore: value.reasoning_override_min_score, + tierModelParams: value.tier_model_params, + enableContextWindowEscalation: value.enable_context_window_escalation, + contextWindowEscalationBuffer: value.context_window_escalation_buffer, + stallEscalationEnabled: value.stall_escalation_enabled, + stallEscalationWindow: value.stall_escalation_window, + stallEscalationRepeatThreshold: value.stall_escalation_repeat_threshold, + codeKeywords: value.code_keywords, + reasoningKeywords: value.reasoning_keywords, + technicalKeywords: value.technical_keywords, + simpleKeywords: value.simple_keywords, + planModePatterns: value.plan_mode_patterns, + routeHousekeepingToCheapestTier: value.route_housekeeping_to_cheapest_tier, + housekeepingPatterns: value.housekeeping_patterns, + reminderMarkers: value.reminder_markers, + maxTokensFromTierModel: value.max_tokens_from_tier_model, + classifierPluginTimeoutMs: value.classifier_plugin_timeout_ms, +}); diff --git a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx index 9d792480c9f..923cf2aa0e3 100644 --- a/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/handle_add_model_submit.test.tsx @@ -101,4 +101,23 @@ describe("prepareModelAddRequest", () => { expect(deployment.litellmParamsObj.litellm_credential_name).toBe("from-json"); expect(deployment.litellmParamsObj.timeout).toBe(5); }); + + it.each([ + ["OpenAI", "openai/*"], + ["Azure_AI_Studio", "azure_ai/*"], + ["Petals", "petals/*"], + ])("composes wildcard names for the all-model selection", async (custom_llm_provider, wildcardModel) => { + const formValues = { + model_mappings: [], + model: "all-wildcard", + custom_llm_provider, + }; + + const deployments = await prepareModelAddRequest({ ...formValues }, "token", null); + + expect(deployments).toHaveLength(1); + const [deployment] = deployments!; + expect(deployment.modelName).toBe(wildcardModel); + expect(deployment.litellmParamsObj.model).toBe(wildcardModel); + }); }); diff --git a/ui/litellm-dashboard/src/components/agents/types.ts b/ui/litellm-dashboard/src/components/agents/types.ts index 24ff0c0e12c..6adb3fa9dda 100644 --- a/ui/litellm-dashboard/src/components/agents/types.ts +++ b/ui/litellm-dashboard/src/components/agents/types.ts @@ -21,6 +21,7 @@ export interface Agent { [key: string]: any; }; object_permission?: AgentObjectPermission; + access_group_ids?: string[] | null; keys?: AgentAttachedKey[] | null; spend?: number; tpm_limit?: number | null; diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts index 2450f7bce27..162414b7814 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts +++ b/ui/litellm-dashboard/src/components/edit_auto_router/build_updated_complexity_router_config.test.ts @@ -189,14 +189,10 @@ describe("buildUpdatedComplexityRouterConfig keyword matching", () => { const saved = buildUpdatedComplexityRouterConfig(stored, value, undefined, keywordState); const forecast = classifier_type !== "heuristic"; expect(saved.adaptive).toBe(!forecast); - expect(saved.enable_context_window_escalation).toBe(!forecast); + expect(saved.enable_context_window_escalation).toBe(true); + expect(saved.context_window_escalation_buffer).toBe(0.9); expect(saved.escalation_keywords).toEqual(forecast ? [] : stored.escalation_keywords); - for (const key of [ - "adaptive_weights", - "adaptive_eligible", - "tier_distance_penalty", - "context_window_escalation_buffer", - ]) { + for (const key of ["adaptive_weights", "adaptive_eligible", "tier_distance_penalty"]) { expect(Object.hasOwn(saved, key)).toBe(!forecast); } expect(saved.keyword_tier_rules).toEqual(STORED.keyword_tier_rules); @@ -208,6 +204,19 @@ describe("buildUpdatedComplexityRouterConfig keyword matching", () => { }, ); + it.each([undefined, false, true])("preserves stored context-window escalation on save: %s", (enabled) => { + const stored = { + ...STORED, + ...(enabled !== undefined && { enable_context_window_escalation: enabled }), + }; + const value = hydrateComplexityRouterConfig(stored, undefined); + const saved = buildUpdatedComplexityRouterConfig(stored, value, undefined, hydratedState); + const serialized: typeof saved = JSON.parse(JSON.stringify(saved)); + expect(value.enable_context_window_escalation).toBe(enabled); + expect(serialized.enable_context_window_escalation).toBe(enabled); + expect(Object.hasOwn(serialized, "enable_context_window_escalation")).toBe(enabled !== undefined); + }); + it("round-trips an untouched edit without changing any keyword-matching value", () => { // Opening the modal hydrates state from STORED; saving with nothing changed must be a // no-op. These keys are now MANAGED, so a hydration bug silently wipes them. @@ -854,6 +863,16 @@ describe("managed keys survive an untouched open-and-save", () => { reasoning_override_min_score: 0.3, enable_context_window_escalation: false, context_window_escalation_buffer: 0.9, + code_keywords: ["async", "await"], + reasoning_keywords: ["prove"], + technical_keywords: ["api"], + simple_keywords: ["hello"], + plan_mode_patterns: ["plan now"], + route_housekeeping_to_cheapest_tier: false, + housekeeping_patterns: ["conversation title"], + reminder_markers: [{ open: "", close: "" }], + max_tokens_from_tier_model: false, + classifier_plugin_timeout_ms: 3000, }; // tier_definitions and fallback_tier cannot sit beside heuristic_first, which this fixture uses, @@ -864,6 +883,7 @@ describe("managed keys survive an untouched open-and-save", () => { "fallback_tier", "hybrid_boundary_margin", "jev_classifier_config", + "classifier_plugin_timeout_ms", ]); // The stall keys are rejected beside the session pinning and user-turn classification this @@ -894,6 +914,12 @@ describe("managed keys survive an untouched open-and-save", () => { expect(dropped).toEqual([]); }); + it("keeps the custom classifier plugin timeout through an untouched save", () => { + const stored = { ...STORED_ALL_MANAGED, classifier_type: "custom", classifier_plugin_timeout_ms: 3000 }; + const hydrated = hydrateComplexityRouterConfig(stored, undefined); + expect(buildUpdatedComplexityRouterConfig(stored, hydrated).classifier_plugin_timeout_ms).toBe(3000); + }); + it("carries an enabled non-reasoning tier and its models through their own round trip", () => { // `tiers` is rewritten wholesale on save, so this is the regression that matters: opening an // enabled router and saving an unrelated edit must not delete the tier or its pool. diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx index 34db61483cf..aa6cf92ceb9 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.integration.test.tsx @@ -347,6 +347,77 @@ describe("EditAutoRouterModal keyword matching", () => { }); }); +describe("EditAutoRouterModal advanced field round trips", () => { + const storedAdvancedConfig = { + ...STORED_CONFIG, + route_housekeeping_to_cheapest_tier: false, + housekeeping_patterns: ["conversation title"], + reminder_markers: [{ open: "", close: "" }], + max_tokens_from_tier_model: false, + }; + + const renderAdvancedModal = (props: Partial> = {}) => + renderModal({ + modelData: { + ...MODEL_DATA, + litellm_params: { ...MODEL_DATA.litellm_params, complexity_router_config: storedAdvancedConfig }, + }, + ...props, + }); + + beforeEach(() => { + modelPatchUpdateCall.mockClear(); + }); + + it("hydrates housekeeping and reminder fields, then omits the default max-token value after editing", async () => { + const user = userEvent.setup(); + renderAdvancedModal(); + + await user.click(await screen.findByText("Advanced: Housekeeping Routing")); + expect(screen.getByRole("switch", { name: "Route housekeeping calls to the cheapest tier" })).not.toBeChecked(); + expect(screen.getByRole("combobox", { name: "e.g., conversation title" })).toHaveValue(""); + + await user.click(screen.getByText("Advanced: Reminder Markers")); + expect(screen.getByLabelText("Opening delimiter")).toHaveValue(""); + expect(screen.getByLabelText("Closing delimiter")).toHaveValue(""); + + await user.click(screen.getByText("Advanced: Response Format")); + const maxTokensSwitch = screen.getByRole("switch", { name: "Cap max_tokens at the tier model's output ceiling" }); + await user.click(maxTokensSwitch); + await user.click(screen.getByRole("button", { name: /save changes/i })); + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalledOnce()); + + expect(savedConfig()).not.toHaveProperty("max_tokens_from_tier_model"); + expect(savedConfig()).toMatchObject({ + route_housekeeping_to_cheapest_tier: false, + housekeeping_patterns: ["conversation title"], + reminder_markers: [{ open: "", close: "" }], + }); + }); + + it("does not PATCH when the edit is cancelled", async () => { + const user = userEvent.setup(); + const onCancel = vi.fn(); + renderAdvancedModal({ onCancel }); + await user.click(screen.getByRole("button", { name: /cancel/i })); + expect(onCancel).toHaveBeenCalledOnce(); + expect(modelPatchUpdateCall).not.toHaveBeenCalled(); + }); + + it("preserves all stored advanced fields through an untouched save", async () => { + const user = userEvent.setup(); + renderAdvancedModal(); + await user.click(screen.getByRole("button", { name: /save changes/i })); + await waitFor(() => expect(modelPatchUpdateCall).toHaveBeenCalledOnce()); + expect(savedConfig()).toMatchObject({ + route_housekeeping_to_cheapest_tier: false, + housekeeping_patterns: ["conversation title"], + reminder_markers: [{ open: "", close: "" }], + max_tokens_from_tier_model: false, + }); + }); +}); + describe("EditAutoRouterModal classifier context window", () => { beforeEach(() => { modelPatchUpdateCall.mockClear(); diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx index c88bbb101f7..0dafa9b330a 100644 --- a/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx +++ b/ui/litellm-dashboard/src/components/edit_auto_router/edit_auto_router_modal.tsx @@ -1,14 +1,7 @@ import AutoRouterClassifierTabs from "../add_model/AutoRouterClassifierTabs"; import { usesClassifierContext } from "../add_model/classifier_types"; -import { defaultJevClassifierConfig, jevClassifierConfigSchema } from "../add_model/jev_classifier_config"; -import type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config"; export type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config"; -import { - getForecastConfigError, - isForecastClassifier, - capabilitySettingsSchema, - fuseSettingsSchema, -} from "../add_model/forecast_classifier_config"; +import { getForecastConfigError, isForecastClassifier } from "../add_model/forecast_classifier_config"; import React, { useEffect, useMemo, useState } from "react"; import { complexityRouterSchema, @@ -30,13 +23,10 @@ import ModelChoiceCombobox, { type ModelChoice } from "../add_model/ModelChoiceC import { modelAvailableCall, modelPatchUpdateCall, validateAutoRouterConfig } from "../networking"; import { fetchAutoRouterModels, fetchAvailableModels, ModelGroup } from "@/components/llm_calls/fetch_models"; import RouterConfigBuilder, { type RouterConfig, serializeRouterConfig } from "../add_model/RouterConfigBuilder"; -import { hydrateTierModelParams } from "../add_model/complexity_router_tiers"; import { - type ActiveTierSet, CUSTOM_TIER_OMITTED_KEYS, activeTierRows, getCustomTierRowsError, - tierParamsByRowId, resolveComplexityDefaultModel, } from "../add_model/tier_rows"; import { isComplexityRouter } from "../add_model/auto_router_strategies"; @@ -45,16 +35,14 @@ import { buildComplexityRouterConfig, getClassifierModelError, getHeuristicV2SuccessThresholdError, + getReminderMarkersError, + getClassifierPluginTimeoutError, getClassifierReasoningEffortError, getKeywordTierRulesError, getMissingTiersError, getSemanticConfigError, getPlanModeTierError, getTierLabelsError, - hydrateBuiltInTiers, - hydrateCustomTierSet, - hydratePlanModeMinTier, - hydrateTierLabels, dryRunRejection, } from "../add_model/build_complexity_router_config"; import { KeywordTierRule } from "../add_model/KeywordTierRules"; @@ -66,22 +54,14 @@ import { hydrateAutoRouterCompression, } from "../add_model/buildAutoRouterCompression"; import { hydrateKeywordTierRules } from "../add_model/complexity_router_keywords"; -import { customDimensionsError, hydrateCustomDimensions } from "../add_model/custom_dimensions"; -import { - hydrateDimensionWeights, - hydrateReasoningOverrideMinScore, - hydrateTierBoundaries, - hydrateTokenThresholds, -} from "../add_model/heuristic_scoring_knobs"; +import { customDimensionsError } from "../add_model/custom_dimensions"; import ComplexityRouterConfig, { ComplexityRouterConfigValue, effectiveClassifierType, heuristicScoringRole, - DEFAULT_ADAPTIVE_WEIGHTS, - DEFAULT_SESSION_AFFINITY, - DEFAULT_DEPLOYMENT_AFFINITY, - DEFAULT_TIER_DISTANCE_PENALTY, } from "../add_model/ComplexityRouterConfig"; +import { builderParamsFromValue } from "../add_model/complexity_router_builder_params"; +import { hydrateComplexityRouterConfig, hydratePinnedDefaultModel } from "./hydrate_complexity_router_config"; import { Dialog, DialogContent, @@ -104,124 +84,7 @@ interface EditAutoRouterModalProps { // Keys this modal rewrites from its own form state on save. Anything absent from this set is // carried through untouched from the stored config, so a key only belongs here once the modal // actually renders a control that can set it. - -/** - * The stored complexity_router_config as form state. Every key in MANAGED_COMPLEXITY_ROUTER_KEYS is - * rewritten from this state on save, so a key missing here is silently dropped from the saved config. - */ -export const hydrateComplexityRouterConfig = ( - parsedConfig: StoredComplexityRouterConfig, - complexityRouterDefaultModel: string | null | undefined, -): ComplexityRouterConfigValue => { - const builtIn = hydrateBuiltInTiers(parsedConfig.tiers, parsedConfig.enable_non_reasoning_tier); - const { tiers: hydratedTiers, enable_non_reasoning_tier } = builtIn; - const custom_tier_set = hydrateCustomTierSet(parsedConfig); - const activeTiers = { ...builtIn, custom_tier_set }; - - return { - tiers: hydratedTiers, - enable_non_reasoning_tier, - custom_tier_set, - tier_model_params: tierParamsByRowId( - hydrateTierModelParams(parsedConfig.tiers, parsedConfig.tier_model_configs), - activeTierRows(activeTiers), - ), - default_model: hydratePinnedDefaultModel(parsedConfig.default_model, complexityRouterDefaultModel, activeTiers), - plan_mode_min_tier: hydratePlanModeMinTier(parsedConfig.plan_mode_min_tier, custom_tier_set), - tier_labels: hydrateTierLabels(parsedConfig.tier_labels), - classifier_type: parsedConfig.classifier_type || "heuristic", - heuristic_v2_success_threshold: - typeof parsedConfig.heuristic_v2_success_threshold === "number" - ? parsedConfig.heuristic_v2_success_threshold - : undefined, - capability_classifier_config: capabilitySettingsSchema.safeParse(parsedConfig.capability_classifier_config).data, - llm_v2_config: fuseSettingsSchema.safeParse(parsedConfig.llm_v2_config).data, - classifier_llm_config: parsedConfig.classifier_type === "jev" ? undefined : parsedConfig.classifier_llm_config, - jev_classifier_config: - parsedConfig.classifier_type === "jev" - ? jevClassifierConfigSchema.safeParse(parsedConfig.jev_classifier_config ?? {}).data ?? - defaultJevClassifierConfig() - : undefined, - classifier_context_window_size: - typeof parsedConfig.classifier_context_window_size === "number" - ? parsedConfig.classifier_context_window_size - : undefined, - classifier_context_budget_chars: - typeof parsedConfig.classifier_context_budget_chars === "number" - ? parsedConfig.classifier_context_budget_chars - : undefined, - classifier_context_per_turn_chars: - typeof parsedConfig.classifier_context_per_turn_chars === "number" - ? parsedConfig.classifier_context_per_turn_chars - : undefined, - classifier_context_include_assistant_turns: - typeof parsedConfig.classifier_context_include_assistant_turns === "boolean" - ? parsedConfig.classifier_context_include_assistant_turns - : undefined, - classifier_fallback: - parsedConfig.classifier_fallback === "default_model" || parsedConfig.classifier_fallback === "heuristic" - ? parsedConfig.classifier_fallback - : undefined, - classification_prompt: - typeof parsedConfig.classification_prompt === "string" && parsedConfig.classification_prompt.trim() !== "" - ? parsedConfig.classification_prompt - : undefined, - classification_examples: - typeof parsedConfig.classification_examples === "string" && parsedConfig.classification_examples.trim() !== "" - ? parsedConfig.classification_examples - : undefined, - heuristic_first_max_tier: - typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== "" - ? parsedConfig.heuristic_first_max_tier - : undefined, - hybrid_boundary_margin: - typeof parsedConfig.hybrid_boundary_margin === "number" ? parsedConfig.hybrid_boundary_margin : undefined, - classification_mode: - parsedConfig.classification_mode === "user_turn" || parsedConfig.classification_mode === "every_request" - ? parsedConfig.classification_mode - : undefined, - tier_boundaries: hydrateTierBoundaries(parsedConfig.tier_boundaries), - token_thresholds: hydrateTokenThresholds(parsedConfig.token_thresholds), - dimension_weights: hydrateDimensionWeights(parsedConfig.dimension_weights), - custom_dimensions: hydrateCustomDimensions(parsedConfig.custom_dimensions), - reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score), - session_affinity: - typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY, - session_affinity_ttl_seconds: - typeof parsedConfig.session_affinity_ttl_seconds === "number" && - Number.isFinite(parsedConfig.session_affinity_ttl_seconds) - ? parsedConfig.session_affinity_ttl_seconds - : undefined, - modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false, - modality_pin_override: - typeof parsedConfig.modality_pin_override === "boolean" ? parsedConfig.modality_pin_override : false, - deployment_affinity: - typeof parsedConfig.deployment_affinity === "boolean" - ? parsedConfig.deployment_affinity - : DEFAULT_DEPLOYMENT_AFFINITY, - adaptive: parsedConfig.adaptive || false, - adaptive_weights: parsedConfig.adaptive_weights, - tier_distance_penalty: parsedConfig.tier_distance_penalty, - adaptive_eligible: parsedConfig.adaptive_eligible || "all", - return_raw_model_name: parsedConfig.return_raw_model_name || false, - enable_context_window_escalation: - typeof parsedConfig.enable_context_window_escalation === "boolean" - ? parsedConfig.enable_context_window_escalation - : undefined, - context_window_escalation_buffer: - typeof parsedConfig.context_window_escalation_buffer === "number" - ? parsedConfig.context_window_escalation_buffer - : undefined, - stall_escalation_enabled: parsedConfig.stall_escalation_enabled === true || undefined, - stall_escalation_window: - typeof parsedConfig.stall_escalation_window === "number" ? parsedConfig.stall_escalation_window : undefined, - stall_escalation_repeat_threshold: - typeof parsedConfig.stall_escalation_repeat_threshold === "number" - ? parsedConfig.stall_escalation_repeat_threshold - : undefined, - }; -}; - +export { hydrateComplexityRouterConfig, hydratePinnedDefaultModel }; export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "tiers", "enable_non_reasoning_tier", @@ -266,6 +129,16 @@ export const MANAGED_COMPLEXITY_ROUTER_KEYS = new Set([ "stall_escalation_enabled", "stall_escalation_window", "stall_escalation_repeat_threshold", + "code_keywords", + "reasoning_keywords", + "technical_keywords", + "simple_keywords", + "plan_mode_patterns", + "route_housekeeping_to_cheapest_tier", + "housekeeping_patterns", + "reminder_markers", + "max_tokens_from_tier_model", + "classifier_plugin_timeout_ms", ]); // Managed only when the caller passes the corresponding state. A caller that does not render @@ -285,24 +158,6 @@ const toRecord = (value: unknown): Record => { : {}; }; -// A pin lives in two places: complexity_router_config.default_model (this UI's own marker, added -// by PR #36615) and litellm_params.complexity_router_default_model (what the backend reads). Only -// the marker proves an operator picked it, because before #36615 every save wrote a tier-derived -// value into litellm_params. So with no marker, a litellm_params value counts as a pin only when -// it diverges from what the tiers alone derive; a match stays unpinned and keeps tracking tiers. -export const hydratePinnedDefaultModel = ( - storedConfigDefaultModel: unknown, - litellmParamsDefaultModel: string | null | undefined, - activeTiers: ActiveTierSet, -): string | undefined => { - if (typeof storedConfigDefaultModel === "string" && storedConfigDefaultModel.trim()) { - return storedConfigDefaultModel; - } - const tierDerived = resolveComplexityDefaultModel(activeTiers); - const externalOverride = litellmParamsDefaultModel?.trim(); - return externalOverride && externalOverride !== tierDerived ? externalOverride : undefined; -}; - export interface KeywordMatchingState { keywordTierRules: KeywordTierRule[]; escalationKeywords: string[]; @@ -338,55 +193,13 @@ export const buildUpdatedComplexityRouterConfig = ( ); const builderParams: BuildComplexityRouterConfigParams = { - tiers: value.tiers, - enableNonReasoningTier: value.enable_non_reasoning_tier, - customTierSet: value.custom_tier_set, - defaultModel: value.default_model, - planModeMinTier: value.plan_mode_min_tier, - classificationPrompt: value.classification_prompt, - classificationExamples: value.classification_examples, - heuristicFirstMaxTier: value.heuristic_first_max_tier, - hybridBoundaryMargin: value.hybrid_boundary_margin, - classificationMode: value.classification_mode, - tierLabels: value.tier_labels, - classifierType: value.classifier_type, - jevClassifierConfig: value.jev_classifier_config, - heuristicV2SuccessThreshold: value.heuristic_v2_success_threshold, - capabilityClassifierConfig: value.capability_classifier_config, - llmV2Config: value.llm_v2_config, - classifierLlmConfig: value.classifier_llm_config, - classifierContextWindowSize: value.classifier_context_window_size, - classifierContextBudgetChars: value.classifier_context_budget_chars, - classifierContextPerTurnChars: value.classifier_context_per_turn_chars, - classifierContextIncludeAssistantTurns: value.classifier_context_include_assistant_turns, - classifierFallback: value.classifier_fallback, - sessionAffinity: value.session_affinity ?? DEFAULT_SESSION_AFFINITY, - sessionAffinityTtlSeconds: value.session_affinity_ttl_seconds, - modalityRouting: value.modality_routing ?? false, - modalityPinOverride: value.modality_pin_override ?? false, - deploymentAffinity: value.deployment_affinity ?? DEFAULT_DEPLOYMENT_AFFINITY, + ...builderParamsFromValue(value), customTechnicalKeywords: customTechnicalKeywords ?? [], keywordTierRules: keywordMatching?.keywordTierRules ?? [], semanticMatchingEnabled: keywordMatching?.semanticMatchingEnabled ?? false, embeddingModel: keywordMatching?.embeddingModel, matchThreshold: keywordMatching?.matchThreshold ?? DEFAULT_MATCH_THRESHOLD, escalationKeywords: keywordMatching?.escalationKeywords ?? [], - adaptive: value.adaptive ?? false, - adaptiveWeights: value.adaptive_weights ?? DEFAULT_ADAPTIVE_WEIGHTS, - tierDistancePenalty: value.tier_distance_penalty ?? DEFAULT_TIER_DISTANCE_PENALTY, - adaptiveEligible: value.adaptive_eligible ?? "all", - returnRawModelName: value.return_raw_model_name ?? false, - tierBoundaries: value.tier_boundaries, - tokenThresholds: value.token_thresholds, - dimensionWeights: value.dimension_weights, - customDimensions: value.custom_dimensions, - reasoningOverrideMinScore: value.reasoning_override_min_score, - tierModelParams: value.tier_model_params, - enableContextWindowEscalation: value.enable_context_window_escalation, - contextWindowEscalationBuffer: value.context_window_escalation_buffer, - stallEscalationEnabled: value.stall_escalation_enabled, - stallEscalationWindow: value.stall_escalation_window, - stallEscalationRepeatThreshold: value.stall_escalation_repeat_threshold, }; const built = buildComplexityRouterConfig(builderParams); @@ -557,12 +370,13 @@ const EditAutoRouterModal: React.FC = ({ setRouterConfig(parsedConfig); // Set form values - form.reset({ + const routerFormValues = { auto_router_name: modelData.model_name, auto_router_default_model: modelData.litellm_params?.auto_router_default_model || null, auto_router_embedding_model: modelData.litellm_params?.auto_router_embedding_model || null, model_access_group: modelData.model_info?.access_groups || [], - }); + }; + form.reset(routerFormValues); } catch (error) { console.error("Error parsing auto router config:", error); toast.fromError("Error loading auto router configuration"); @@ -585,6 +399,11 @@ const EditAutoRouterModal: React.FC = ({ const classifierError = getClassifierModelError(complexityRouterConfig) ?? getHeuristicV2SuccessThresholdError(complexityRouterConfig.heuristic_v2_success_threshold) ?? + getReminderMarkersError(complexityRouterConfig.reminder_markers) ?? + getClassifierPluginTimeoutError( + complexityRouterConfig.classifier_type, + complexityRouterConfig.classifier_plugin_timeout_ms, + ) ?? getForecastConfigError(complexityRouterConfig) ?? (heuristicScoringRole(complexityRouterConfig) === "decides" ? customDimensionsError(complexityRouterConfig.custom_dimensions) @@ -635,11 +454,18 @@ const EditAutoRouterModal: React.FC = ({ // Dual write: complexity_router_config.default_model (the pin marker hydratePinnedDefaultModel // reads back) and complexity_router_default_model (what the backend routes on) must always be // written together from the same value. Same pairing in add_auto_router_tab.tsx. + const keywordMatching = { + keywordTierRules, + escalationKeywords, + semanticMatchingEnabled, + embeddingModel, + matchThreshold, + }; const updatedConfig = buildUpdatedComplexityRouterConfig( modelData.litellm_params?.complexity_router_config, complexityRouterConfig, customTechnicalKeywords, - { keywordTierRules, escalationKeywords, semanticMatchingEnabled, embeddingModel, matchThreshold }, + keywordMatching, ); const serverVerdict = await validateAutoRouterConfig(accessToken, updatedConfig, modelData?.model_info?.team_id); const dryRunError = dryRunRejection(serverVerdict); @@ -676,12 +502,13 @@ const EditAutoRouterModal: React.FC = ({ ); toast.success("Auto router configuration updated successfully"); - onSuccess({ + const updatedModelData = { ...modelData, model_name: values.auto_router_name, litellm_params: updatedLitellmParams, model_info: updatedModelInfo, - }); + }; + onSuccess(updatedModelData); onCancel(); return; } diff --git a/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts new file mode 100644 index 00000000000..6dbd2b19b52 --- /dev/null +++ b/ui/litellm-dashboard/src/components/edit_auto_router/hydrate_complexity_router_config.ts @@ -0,0 +1,186 @@ +import { defaultJevClassifierConfig, jevClassifierConfigSchema } from "../add_model/jev_classifier_config"; +import { capabilitySettingsSchema, fuseSettingsSchema } from "../add_model/forecast_classifier_config"; +import type { StoredComplexityRouterConfig } from "../add_model/build_complexity_router_config"; +import { + hydrateBuiltInTiers, + hydrateCustomTierSet, + hydratePlanModeMinTier, + hydrateTierLabels, +} from "../add_model/build_complexity_router_config"; +import { hydrateTierModelParams } from "../add_model/complexity_router_tiers"; +import { hydrateCustomDimensions } from "../add_model/custom_dimensions"; +import { + hydrateDimensionWeights, + hydrateReasoningOverrideMinScore, + hydrateTierBoundaries, + hydrateTokenThresholds, +} from "../add_model/heuristic_scoring_knobs"; +import type { ComplexityRouterConfigValue } from "../add_model/ComplexityRouterConfig"; +import { DEFAULT_DEPLOYMENT_AFFINITY, DEFAULT_SESSION_AFFINITY } from "../add_model/ComplexityRouterConfig"; +import { + type ActiveTierSet, + activeTierRows, + tierParamsByRowId, + resolveComplexityDefaultModel, +} from "../add_model/tier_rows"; + +const isReminderMarkerPair = (input: unknown): input is { open: string; close: string } => { + if (typeof input !== "object" || input === null) { + return false; + } + if (!("open" in input) || !("close" in input)) { + return false; + } + return typeof input.open === "string" && typeof input.close === "string"; +}; + +const stringList = (input: unknown): string[] | undefined => + Array.isArray(input) ? input.filter((item): item is string => typeof item === "string") : undefined; + +export const hydratePinnedDefaultModel = ( + storedConfigDefaultModel: unknown, + litellmParamsDefaultModel: string | null | undefined, + activeTiers: ActiveTierSet, +): string | undefined => { + if (typeof storedConfigDefaultModel === "string" && storedConfigDefaultModel.trim()) { + return storedConfigDefaultModel; + } + const tierDerived = resolveComplexityDefaultModel(activeTiers); + const externalOverride = litellmParamsDefaultModel?.trim(); + return externalOverride && externalOverride !== tierDerived ? externalOverride : undefined; +}; + +export const hydrateComplexityRouterConfig = ( + parsedConfig: StoredComplexityRouterConfig, + complexityRouterDefaultModel: string | null | undefined, +): ComplexityRouterConfigValue => { + const builtIn = hydrateBuiltInTiers(parsedConfig.tiers, parsedConfig.enable_non_reasoning_tier); + const { tiers: hydratedTiers, enable_non_reasoning_tier } = builtIn; + const custom_tier_set = hydrateCustomTierSet(parsedConfig); + const activeTiers = { ...builtIn, custom_tier_set }; + + return { + tiers: hydratedTiers, + enable_non_reasoning_tier, + custom_tier_set, + tier_model_params: tierParamsByRowId( + hydrateTierModelParams(parsedConfig.tiers, parsedConfig.tier_model_configs), + activeTierRows(activeTiers), + ), + default_model: hydratePinnedDefaultModel(parsedConfig.default_model, complexityRouterDefaultModel, activeTiers), + plan_mode_min_tier: hydratePlanModeMinTier(parsedConfig.plan_mode_min_tier, custom_tier_set), + tier_labels: hydrateTierLabels(parsedConfig.tier_labels), + classifier_type: parsedConfig.classifier_type || "heuristic", + heuristic_v2_success_threshold: + typeof parsedConfig.heuristic_v2_success_threshold === "number" + ? parsedConfig.heuristic_v2_success_threshold + : undefined, + capability_classifier_config: capabilitySettingsSchema.safeParse(parsedConfig.capability_classifier_config).data, + llm_v2_config: fuseSettingsSchema.safeParse(parsedConfig.llm_v2_config).data, + classifier_llm_config: parsedConfig.classifier_type === "jev" ? undefined : parsedConfig.classifier_llm_config, + jev_classifier_config: + parsedConfig.classifier_type === "jev" + ? jevClassifierConfigSchema.safeParse(parsedConfig.jev_classifier_config ?? {}).data ?? + defaultJevClassifierConfig() + : undefined, + classifier_context_window_size: + typeof parsedConfig.classifier_context_window_size === "number" + ? parsedConfig.classifier_context_window_size + : undefined, + classifier_context_budget_chars: + typeof parsedConfig.classifier_context_budget_chars === "number" + ? parsedConfig.classifier_context_budget_chars + : undefined, + classifier_context_per_turn_chars: + typeof parsedConfig.classifier_context_per_turn_chars === "number" + ? parsedConfig.classifier_context_per_turn_chars + : undefined, + classifier_context_include_assistant_turns: + typeof parsedConfig.classifier_context_include_assistant_turns === "boolean" + ? parsedConfig.classifier_context_include_assistant_turns + : undefined, + classifier_fallback: + parsedConfig.classifier_fallback === "default_model" || parsedConfig.classifier_fallback === "heuristic" + ? parsedConfig.classifier_fallback + : undefined, + classification_prompt: + typeof parsedConfig.classification_prompt === "string" && parsedConfig.classification_prompt.trim() !== "" + ? parsedConfig.classification_prompt + : undefined, + classification_examples: + typeof parsedConfig.classification_examples === "string" && parsedConfig.classification_examples.trim() !== "" + ? parsedConfig.classification_examples + : undefined, + heuristic_first_max_tier: + typeof parsedConfig.heuristic_first_max_tier === "string" && parsedConfig.heuristic_first_max_tier.trim() !== "" + ? parsedConfig.heuristic_first_max_tier + : undefined, + hybrid_boundary_margin: + typeof parsedConfig.hybrid_boundary_margin === "number" ? parsedConfig.hybrid_boundary_margin : undefined, + classification_mode: + parsedConfig.classification_mode === "user_turn" || parsedConfig.classification_mode === "every_request" + ? parsedConfig.classification_mode + : undefined, + tier_boundaries: hydrateTierBoundaries(parsedConfig.tier_boundaries), + token_thresholds: hydrateTokenThresholds(parsedConfig.token_thresholds), + dimension_weights: hydrateDimensionWeights(parsedConfig.dimension_weights), + custom_dimensions: hydrateCustomDimensions(parsedConfig.custom_dimensions), + reasoning_override_min_score: hydrateReasoningOverrideMinScore(parsedConfig.reasoning_override_min_score), + session_affinity: + typeof parsedConfig.session_affinity === "boolean" ? parsedConfig.session_affinity : DEFAULT_SESSION_AFFINITY, + session_affinity_ttl_seconds: + typeof parsedConfig.session_affinity_ttl_seconds === "number" && + Number.isFinite(parsedConfig.session_affinity_ttl_seconds) + ? parsedConfig.session_affinity_ttl_seconds + : undefined, + modality_routing: typeof parsedConfig.modality_routing === "boolean" ? parsedConfig.modality_routing : false, + modality_pin_override: + typeof parsedConfig.modality_pin_override === "boolean" ? parsedConfig.modality_pin_override : false, + deployment_affinity: + typeof parsedConfig.deployment_affinity === "boolean" + ? parsedConfig.deployment_affinity + : DEFAULT_DEPLOYMENT_AFFINITY, + adaptive: parsedConfig.adaptive || false, + adaptive_weights: parsedConfig.adaptive_weights, + tier_distance_penalty: parsedConfig.tier_distance_penalty, + adaptive_eligible: parsedConfig.adaptive_eligible || "all", + return_raw_model_name: parsedConfig.return_raw_model_name || false, + enable_context_window_escalation: + typeof parsedConfig.enable_context_window_escalation === "boolean" + ? parsedConfig.enable_context_window_escalation + : undefined, + context_window_escalation_buffer: + typeof parsedConfig.context_window_escalation_buffer === "number" + ? parsedConfig.context_window_escalation_buffer + : undefined, + stall_escalation_enabled: parsedConfig.stall_escalation_enabled === true || undefined, + stall_escalation_window: + typeof parsedConfig.stall_escalation_window === "number" ? parsedConfig.stall_escalation_window : undefined, + stall_escalation_repeat_threshold: + typeof parsedConfig.stall_escalation_repeat_threshold === "number" + ? parsedConfig.stall_escalation_repeat_threshold + : undefined, + code_keywords: stringList(parsedConfig.code_keywords), + reasoning_keywords: stringList(parsedConfig.reasoning_keywords), + technical_keywords: stringList(parsedConfig.technical_keywords), + simple_keywords: stringList(parsedConfig.simple_keywords), + plan_mode_patterns: stringList(parsedConfig.plan_mode_patterns), + route_housekeeping_to_cheapest_tier: + typeof parsedConfig.route_housekeeping_to_cheapest_tier === "boolean" + ? parsedConfig.route_housekeeping_to_cheapest_tier + : undefined, + housekeeping_patterns: stringList(parsedConfig.housekeeping_patterns), + reminder_markers: Array.isArray(parsedConfig.reminder_markers) + ? parsedConfig.reminder_markers.filter(isReminderMarkerPair) + : undefined, + max_tokens_from_tier_model: + typeof parsedConfig.max_tokens_from_tier_model === "boolean" + ? parsedConfig.max_tokens_from_tier_model + : undefined, + classifier_plugin_timeout_ms: + typeof parsedConfig.classifier_plugin_timeout_ms === "number" && + Number.isFinite(parsedConfig.classifier_plugin_timeout_ms) + ? parsedConfig.classifier_plugin_timeout_ms + : undefined, + }; +}; diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index f4011a653ab..82399792674 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -6211,6 +6211,7 @@ export const patchAgentCall = async ( rpm_limit?: number | null; session_tpm_limit?: number | null; session_rpm_limit?: number | null; + access_group_ids?: string[]; }, ) => { try { diff --git a/ui/litellm-dashboard/src/components/shared/ScopedSavingsTab.tsx b/ui/litellm-dashboard/src/components/shared/ScopedSavingsTab.tsx new file mode 100644 index 00000000000..7a0455f9d7e --- /dev/null +++ b/ui/litellm-dashboard/src/components/shared/ScopedSavingsTab.tsx @@ -0,0 +1,133 @@ +"use client"; + +import { useMemo, useState } from "react"; + +import { AreaChart, BarChart, CustomLegend } from "@/components/shared/charts"; +import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; +import SavingsTiles from "@/components/shared/SavingsTiles"; +import { Card, CardAction, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; +import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { + formatRangeLabel, + localIsoDay, + MAX_POINTS_WITH_DOTS, + SAVINGS_COLORS, + SAVINGS_SERIES, + SavingsAccumulation, + SavingsPoint, + savingsSeriesOf, + shortDate, + toCumulative, + usd, + withStartAnchor, +} from "@/app/(dashboard)/cost-optimization/_components/costOptimizationUtils"; +import { + useScopedDailyActivityRange, + type ActivityDateRange, + type DailyActivityScope, +} from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; + +interface ScopedSavingsTabProps { + accessToken: string | null; + scope: DailyActivityScope; + activity: ActivityDateRange; + entityType: "key" | "user"; + scopeNote?: string; +} + +const ScopedSavingsTab = ({ accessToken, scope, activity, entityType, scopeNote }: ScopedSavingsTabProps) => { + const { dateValue, onDateChange, results, loading, isFetchingMore, failed, cancelled } = useScopedDailyActivityRange( + accessToken, + scope, + activity, + ); + const startTime = dateValue.from; + const endTime = dateValue.to; + + const [accumulation, setAccumulation] = useState("cumulative"); + + const perInterval = useMemo(() => savingsSeriesOf(results), [results]); + + const overTime = useMemo(() => { + if (accumulation !== "cumulative") return perInterval; + const startLabel = startTime ? shortDate(localIsoDay(startTime)) : ""; + return withStartAnchor(toCumulative(perInterval), startLabel); + }, [accumulation, perInterval, startTime]); + + const intervalLabel = "Per day"; + const rangeLabel = formatRangeLabel(startTime, endTime); + const savingsSubtitle = [ + accumulation === "cumulative" ? "Running total saved" : `Saved ${intervalLabel.toLowerCase()}`, + rangeLabel && `${rangeLabel} (UTC)`, + ] + .filter(Boolean) + .join(" · "); + + const isLoading = loading || isFetchingMore; + const unavailable = failed || cancelled; + const showResults = !isLoading && !unavailable; + const hasRows = results.length > 0; + const showEmpty = !unavailable && (isLoading || !hasRows); + const showChart = showResults && hasRows; + const chartProps = { + data: overTime, + index: "date", + categories: SAVINGS_SERIES, + colors: SAVINGS_COLORS, + valueFormatter: usd, + showLegend: false, + }; + + return ( +
+
+ Spend is bucketed by UTC day + +
+ + {scopeNote && ( +

+ {scopeNote} +

+ )} + + {unavailable && ( +

+ Savings are unavailable for this range. Try another date range or reopen this tab. +

+ )} + {showResults && } + + + + Savings + {savingsSubtitle} + + + setAccumulation(value as SavingsAccumulation)}> + + Cumulative + {intervalLabel} + + + + + + {showEmpty && ( +

+ {isLoading ? "Loading savings..." : `No usage recorded for this ${entityType} in this range.`} +

+ )} + {showChart && + (accumulation === "cumulative" ? ( + + ) : ( + + ))} + + +
+ ); +}; + +export default ScopedSavingsTab; diff --git a/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx index c981010dff9..9f320049b09 100644 --- a/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/shared/SearchSelect.integration.test.tsx @@ -111,4 +111,36 @@ describe("SearchSelect", () => { expect(screen.queryByText("Growth")).not.toBeInTheDocument(); expect(onValueChange).not.toHaveBeenCalled(); }); + + it("supports keyboard select, clear, and reselect", async () => { + const onValueChange = vi.fn(); + const user = userEvent.setup(); + function Controlled() { + const [value, setValue] = useState(null); + return ( + { + setValue(next); + onValueChange(next); + }} + /> + ); + } + + render(); + const input = screen.getByRole("combobox"); + await user.tab(); + await user.keyboard("{Enter}"); + await user.keyboard("{ArrowDown}{Enter}"); + expect(onValueChange).toHaveBeenLastCalledWith("team-1"); + const clear = screen.getByRole("button", { name: "Clear" }); + clear.focus(); + await user.keyboard("{Enter}"); + expect(onValueChange).toHaveBeenLastCalledWith(null); + input.focus(); + await user.keyboard("{Enter}{ArrowDown}{Enter}"); + expect(onValueChange).toHaveBeenLastCalledWith("team-1"); + }); }); diff --git a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx index c33529eb042..d1395e0be51 100644 --- a/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeySavingsTab.tsx @@ -1,132 +1,29 @@ "use client"; -import React, { useMemo, useState } from "react"; - -import { AreaChart, BarChart, CustomLegend } from "@/components/shared/charts"; -import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; -import SavingsTiles from "@/components/shared/SavingsTiles"; -import { Card, CardAction, CardContent, CardDescription, CardHeader, CardTitle } from "@/components/ui/card"; -import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import ScopedSavingsTab from "@/components/shared/ScopedSavingsTab"; import { hasProxyWideSpendView, spendScopeUserId } from "@/utils/roles"; -import { - formatRangeLabel, - localIsoDay, - MAX_POINTS_WITH_DOTS, - SAVINGS_COLORS, - SAVINGS_SERIES, - SavingsAccumulation, - SavingsPoint, - savingsSeriesOf, - shortDate, - toCumulative, - usd, - withStartAnchor, -} from "@/app/(dashboard)/cost-optimization/_components/costOptimizationUtils"; -import { - useScopedDailyActivityRange, - type ActivityDateRange, -} from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; +import type { ActivityDateRange } from "@/app/(dashboard)/cost-optimization/_components/useDailyActivityRange"; interface KeySavingsTabProps { accessToken: string | null; - /** The key's token hash — what spend rows are keyed by, not the one-time plaintext secret. */ keyToken: string; userId: string | null; userRole: string; activity: ActivityDateRange; } -const KeySavingsTab: React.FC = ({ accessToken, keyToken, userId, userRole, activity }) => { - // Proxy admins read the whole key. For anyone else the endpoint applies the caller's own user_id - // alongside the key filter, so the figures cover only that viewer's requests on this key -- said - // plainly in the scope note below rather than left to be misread as the key's total. - const readsWholeKey = hasProxyWideSpendView(userRole); - const { dateValue, onDateChange, results, loading, isFetchingMore } = useScopedDailyActivityRange( - accessToken, - { userId: spendScopeUserId(userRole, userId), apiKey: keyToken }, - activity, - ); - const startTime = dateValue.from ?? null; - const endTime = dateValue.to ?? null; - - const [accumulation, setAccumulation] = useState("cumulative"); - - const perInterval = useMemo(() => savingsSeriesOf(results), [results]); - - const overTime = useMemo(() => { - if (accumulation !== "cumulative") return perInterval; - const startLabel = startTime ? shortDate(localIsoDay(startTime)) : ""; - return withStartAnchor(toCumulative(perInterval), startLabel); - }, [accumulation, perInterval, startTime]); - - const intervalLabel = "Per day"; - const rangeLabel = formatRangeLabel(startTime ?? undefined, endTime ?? undefined); - const savingsSubtitle = [ - accumulation === "cumulative" ? "Running total saved" : `Saved ${intervalLabel.toLowerCase()}`, - rangeLabel && `${rangeLabel} (UTC)`, - ] - .filter(Boolean) - .join(" · "); - - const isLoading = loading || isFetchingMore; - const hasRows = results.length > 0; - const chartProps = { - data: overTime, - index: "date", - categories: SAVINGS_SERIES, - colors: SAVINGS_COLORS, - valueFormatter: usd, - showLegend: false, - }; - - return ( -
-
- Spend is bucketed by UTC day - -
- - {!readsWholeKey && ( -

- Showing your own requests on this key. A key shared across a team will have spend from other members that is - not counted here. -

- )} - - - - - - Savings - {savingsSubtitle} - - - setAccumulation(value as SavingsAccumulation)}> - - Cumulative - {intervalLabel} - - - - - - {/* Distinguishes "still fetching" from "this key genuinely had no traffic": an empty - chart alone reads as a broken panel, and a $0.00 tile reads as a real zero. */} - {!hasRows && ( -

- {isLoading ? "Loading savings..." : "No usage recorded for this key in this range."} -

- )} - {hasRows && accumulation === "cumulative" && ( - - )} - {/* Not stacked: auto-router can go negative on a cold-cache write, and stacking would - draw that below the axis while the rest of the bar still read as the total. */} - {hasRows && accumulation !== "cumulative" && } - - -
- ); -}; +const KeySavingsTab = ({ accessToken, keyToken, userId, userRole, activity }: KeySavingsTabProps) => ( + +); export default KeySavingsTab; diff --git a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx index 5a35c7ae16b..8d1847e0121 100644 --- a/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/RequestLogsFilters.test.tsx @@ -344,4 +344,15 @@ describe("RequestLogsFilters", () => { expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.CACHE_STATUS, undefined); }); + + it("clears the raw Error Code combobox through the undefined filter contract", async () => { + const user = userEvent.setup(); + const { set } = renderFilters({ [LOG_FILTER_IDS.ERROR_CODE]: "429" }); + const input = await screen.findByPlaceholderText("Select or type an error code"); + + await user.click(input); + await user.click(screen.getByRole("button", { name: "Clear", hidden: true })); + + expect(set).toHaveBeenCalledWith(LOG_FILTER_IDS.ERROR_CODE, undefined); + }); }); diff --git a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts index e1d9e6ab5ac..87efb2b1d5e 100644 --- a/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts +++ b/ui/litellm-dashboard/src/lib/autorouter_presets.test.ts @@ -746,7 +746,7 @@ describe("autorouter_presets", () => { expect(prefill.escalationKeywords).toEqual([]); }); - it("carries a preset's context-window escalation opt-out and buffer through the prefill", () => { + it.each([undefined, false, true])("preserves a preset's context-window escalation setting: %s", (enabled) => { const prefill = buildPresetPrefill( { tiers: { SIMPLE: ["gpt-5-nano"], MEDIUM: [], COMPLEX: [], REASONING: [] }, @@ -754,12 +754,12 @@ describe("autorouter_presets", () => { classification_mode: "every_request", session_affinity: false, deployment_affinity: true, - enable_context_window_escalation: false, + enable_context_window_escalation: enabled, context_window_escalation_buffer: 0.9, }, groupsOnly(["gpt-5-nano"]), ); - expect(prefill.complexityRouterConfig.enable_context_window_escalation).toBe(false); + expect(prefill.complexityRouterConfig.enable_context_window_escalation).toBe(enabled); expect(prefill.complexityRouterConfig.context_window_escalation_buffer).toBe(0.9); }); diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 84e3270c48b..a0f6d2c8d8b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -1246,9 +1246,10 @@ export interface paths { * @description Benchmarks for the auto-router dashboard: session shape, savings against the configured * baseline, and prompt-caching behaviour bucketed by what the router did. * - * Reads the LiteLLM_AutoRouterSession rollup, folded once per request at spend-write time, - * so this endpoint never scans LiteLLM_SpendLogs. A session is in the window when it - * overlaps it: its last turn is on or after start_date and its first turn is on or before + * Reads session rollups folded once per request at spend-write time, so this endpoint + * never scans LiteLLM_SpendLogs. A user filter selects only turns attributed to that + * internal user when written; older key-only history remains outside user views. A session + * is in the window when it overlaps it: its last turn is on or after start_date and its first turn is on or before * end_date. Overall hit rate is over telemetry-bearing turns; each bucket's hit rate is * over that bucket's turns. * @@ -23639,6 +23640,8 @@ export interface components { }; /** AgentConfig */ AgentConfig: { + /** Access Group Ids */ + access_group_ids?: string[] | null; agent_card_params: components["schemas"]["AgentCard"]; /** Agent Name */ agent_name: string; @@ -23786,6 +23789,8 @@ export interface components { }; /** AgentResponse */ AgentResponse: { + /** Access Group Ids */ + access_group_ids?: string[] | null; /** Agent Card Params */ agent_card_params: { [key: string]: unknown; @@ -25830,7 +25835,7 @@ export interface components { * CallTypes * @enum {string} */ - CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; + CallTypes: "embedding" | "aembedding" | "completion" | "acompletion" | "atext_completion" | "text_completion" | "image_generation" | "aimage_generation" | "image_edit" | "aimage_edit" | "moderation" | "amoderation" | "atranscription" | "transcription" | "aspeech" | "speech" | "rerank" | "arerank" | "search" | "asearch" | "_arealtime" | "_aresponses_websocket" | "create_batch" | "acreate_batch" | "aretrieve_batch" | "retrieve_batch" | "acancel_batch" | "cancel_batch" | "pass_through_endpoint" | "anthropic_messages" | "aanthropic_messages" | "get_assistants" | "aget_assistants" | "create_assistants" | "acreate_assistants" | "delete_assistant" | "adelete_assistant" | "acreate_thread" | "create_thread" | "aget_thread" | "get_thread" | "a_add_message" | "add_message" | "aget_messages" | "get_messages" | "arun_thread" | "run_thread" | "arun_thread_stream" | "run_thread_stream" | "afile_retrieve" | "file_retrieve" | "afile_delete" | "file_delete" | "afile_list" | "file_list" | "acreate_file" | "create_file" | "afile_content" | "file_content" | "create_fine_tuning_job" | "acreate_fine_tuning_job" | "create_video" | "acreate_video" | "video_generation" | "avideo_generation" | "avideo_retrieve" | "video_retrieve" | "avideo_content" | "video_content" | "video_remix" | "avideo_remix" | "video_list" | "avideo_list" | "video_retrieve_job" | "avideo_retrieve_job" | "video_delete" | "avideo_delete" | "video_create_character" | "avideo_create_character" | "video_get_character" | "avideo_get_character" | "video_edit" | "avideo_edit" | "video_extension" | "avideo_extension" | "vector_store_file_create" | "avector_store_file_create" | "vector_store_file_list" | "avector_store_file_list" | "vector_store_file_retrieve" | "avector_store_file_retrieve" | "vector_store_file_content" | "avector_store_file_content" | "vector_store_file_update" | "avector_store_file_update" | "vector_store_file_delete" | "avector_store_file_delete" | "vector_store_create" | "avector_store_create" | "vector_store_search" | "avector_store_search" | "ingest" | "aingest" | "query" | "aquery" | "create_interaction" | "acreate_interaction" | "create_container" | "acreate_container" | "list_containers" | "alist_containers" | "retrieve_container" | "aretrieve_container" | "delete_container" | "adelete_container" | "list_container_files" | "alist_container_files" | "upload_container_file" | "aupload_container_file" | "create_sandbox" | "acreate_sandbox" | "delete_sandbox" | "adelete_sandbox" | "run_code" | "arun_code" | "code_interpreter_tool" | "acode_interpreter_tool" | "acancel_fine_tuning_job" | "cancel_fine_tuning_job" | "alist_fine_tuning_jobs" | "list_fine_tuning_jobs" | "aretrieve_fine_tuning_job" | "retrieve_fine_tuning_job" | "responses" | "aresponses" | "alist_input_items" | "llm_passthrough_route" | "allm_passthrough_route" | "generate_content" | "agenerate_content" | "generate_content_stream" | "agenerate_content_stream" | "ocr" | "aocr" | "call_mcp_tool" | "list_mcp_tools" | "asend_message" | "send_message" | "acreate_skill"; /** CallbackDelete */ CallbackDelete: { /** Callback Name */ @@ -34982,6 +34987,8 @@ export interface components { }; /** PatchAgentRequest */ PatchAgentRequest: { + /** Access Group Ids */ + access_group_ids?: string[] | null; agent_card_params?: components["schemas"]["AgentCard"]; /** Agent Name */ agent_name?: string; @@ -36884,8 +36891,8 @@ export interface components { embedding_model?: string | null; /** * Enable Context Window Escalation - * @description Escalate a request off a tier whose models provably cannot hold its prompt, before dispatch. The classifier scores complexity and never prompt size, so a long agentic session whose newest ask is trivial lands on a small-window tier and the provider rejects it with a context-window 400 that nothing retries. When every model of the decided tier has a declared window smaller than the estimated prompt, the request moves to the lowest configured tier with a model whose declared window fits; when only some of the tier's models fit, the pick is restricted to those and the tier keeps the request. Models with no resolvable window are never escalated away from and never escalated onto. Set false to dispatch on complexity alone, as before. - * @default true + * @description Escalate a request off a tier whose models provably cannot hold its prompt, before dispatch. The classifier scores complexity and never prompt size, so a long agentic session whose newest ask is trivial lands on a small-window tier and the provider rejects it with a context-window 400 that nothing retries. When every model of the decided tier has a declared window smaller than the estimated prompt, the request moves to the lowest configured tier with a model whose declared window fits; when only some of the tier's models fit, the pick is restricted to those and the tier keeps the request. Models with no resolvable window are never escalated away from and never escalated onto. Disabled by default: omit or set false to dispatch on complexity alone; set true to enable context-window escalation. + * @default false */ enable_context_window_escalation: boolean; /** @@ -37244,6 +37251,13 @@ export interface components { routing_strategy_descriptions: { [key: string]: string; }; + /** + * Source + * @description Source of each current router setting + */ + source: { + [key: string]: "config" | "db" | "env" | "default" | "unset"; + }; }; /** * RoutingGroup @@ -39707,6 +39721,10 @@ export interface components { field_schema: { [key: string]: unknown; }; + /** Source */ + source: { + [key: string]: "config" | "db" | "env" | "default" | "unset"; + }; /** Values */ values: { [key: string]: unknown; @@ -43881,6 +43899,8 @@ export interface operations { end_date?: string | null; /** @description Filter to one virtual key token hash */ api_key?: string | null; + /** @description Filter to one canonical internal user recorded on each turn */ + user_id?: string | null; }; header?: never; path?: never; diff --git a/uv.lock b/uv.lock index 543581cbc23..18d57aa991b 100644 --- a/uv.lock +++ b/uv.lock @@ -4516,11 +4516,13 @@ dependencies = [ { name = "boto3" }, { name = "click" }, { name = "fastuuid" }, + { name = "filelock" }, { name = "httpx", extra = ["http2"] }, { name = "importlib-metadata" }, { name = "jinja2" }, { name = "jsonschema" }, { name = "openai" }, + { name = "packaging" }, { name = "pydantic" }, { name = "pydantic-settings" }, { name = "python-dotenv" }, @@ -4769,6 +4771,7 @@ requires-dist = [ { name = "fastapi", marker = "extra == 'proxy'", specifier = ">=0.136.3,<1.0" }, { name = "fastapi-sso", marker = "extra == 'proxy'", specifier = ">=0.19.0,<1.0" }, { name = "fastuuid", specifier = ">=0.14.0,<1.0" }, + { name = "filelock", specifier = ">=3.16.1,<4.0" }, { name = "google-cloud-aiplatform", marker = "extra == 'google'", specifier = ">=1.133.0,<2.0" }, { name = "google-cloud-aiplatform", marker = "extra == 'proxy-runtime'", specifier = ">=1.133.0,<2.0" }, { name = "google-cloud-iam", marker = "extra == 'extra-proxy'", specifier = ">=2.19.1,<3.0" }, @@ -4802,6 +4805,7 @@ requires-dist = [ { name = "numpydoc", marker = "extra == 'utils'", specifier = ">=1.8.0,<2.0" }, { name = "nvidia-riva-client", marker = "extra == 'stt-nvidia-riva'", specifier = ">=2.15.0" }, { name = "openai", specifier = ">=2.20.0,<3.0.0" }, + { name = "packaging", specifier = ">=24.0" }, { name = "opentelemetry-api", marker = "extra == 'proxy-runtime'", specifier = "==1.28.0" }, { name = "opentelemetry-exporter-otlp", marker = "extra == 'proxy-runtime'", specifier = "==1.28.0" }, { name = "opentelemetry-instrumentation-fastapi", marker = "extra == 'proxy-runtime'", specifier = "==0.49b0" },