scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 98

merge: resolve conflicts with main

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
mateo 2026-09-22 00:29:42 +00:00
commit 539e027b71
164 changed files with 13498 additions and 1826 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
litellm-proxy-extras/litellm_proxy_extras/migrations/20260917000000_add_agent_access_group_ids/migration.sql Normal file
View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_AgentsTable" ADD COLUMN IF NOT EXISTS "access_group_ids" TEXT[] DEFAULT ARRAY[]::TEXT[];

42
litellm-proxy-extras/litellm_proxy_extras/migrations/20260919000000_add_autorouter_user_session_rollup/migration.sql Normal file
View file

@ -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");

42
litellm-proxy-extras/litellm_proxy_extras/schema.prisma
View file

@ -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

35
litellm-rust/Cargo.lock generated
View file

@ -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",
]

1
litellm-rust/Cargo.toml
View file

@ -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" }

21
litellm-rust/crates/cache-redis-semantic/Cargo.toml Normal file
View file

@ -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

618
litellm-rust/crates/cache-redis-semantic/src/cache.rs Normal file
View file

@ -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<Vec<f32>, Error>;
fn async_embed(
&self,
prompt: &str,
metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, 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<String>,
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<String, Error> {
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<String, Error> {
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<Duration>,
) -> 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<Option<CacheEntry>, 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::<redis::Value>(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<E: Embedder, C = redis::Connection> {
connections: Arc<Connections<C>>,
embedder: E,
inner: Arc<Inner>,
}
impl<E: Embedder> RedisSemanticCache<E> {
pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
inner: Arc::new(Inner::new(config)),
})
}
}
impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<E, C> {
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<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
for RedisSemanticCache<E, C>
{
type Value = CacheEntry;
type Context = SemanticCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
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<Option<Self::Value>, 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<Option<Self::Value>, 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<CacheConnectionResult, Error> {
match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
Ok(match redis::cmd("PING").query::<String>(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<u8> {
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<Option<bool>, Error> {
let info = match redis::cmd("FT.INFO")
.arg(name)
.query::<redis::Value>(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::<Vec<_>>();
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<String> {
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<f64> {
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<String> {
field_value(fields, name).and_then(string_value)
}
fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
field_value(fields, name).and_then(number_value)
}
fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
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)
}

5
litellm-rust/crates/cache-redis-semantic/src/lib.rs Normal file
View file

@ -0,0 +1,5 @@
mod cache;
mod prompt;
pub use cache::{Embedder, RedisSemanticCache, RedisSemanticConfig};
pub use prompt::prompt_from_context;

97
litellm-rust/crates/cache-redis-semantic/src/prompt.rs Normal file
View file

@ -0,0 +1,97 @@
use litellm_cache::SemanticCacheContext;
use serde_json::Value;
pub fn prompt_from_context(context: &SemanticCacheContext) -> Option<String> {
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<String>) {
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;
}
}
}
}
_ => {}
}
}

1003
litellm-rust/crates/cache-redis-semantic/tests/cache.rs Normal file
View file

File diff suppressed because it is too large Load diff

67
litellm-rust/crates/cache-redis/src/cache.rs
View file

@ -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<C> {
#[allow(private_interfaces)]
pub enum Connections<C> {
Pool(r2d2::Pool<ConnectionManager>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
@ -50,7 +51,7 @@ impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
fn execute<T>(
pub fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
@ -73,6 +74,29 @@ where
}
}
}
pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + 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<Self, Error> {
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<S, C = redis::Connection> {
@ -94,12 +118,7 @@ impl<S: CacheCodec> RedisCache<S> {
default_ttl: Option<Duration>,
codec: S,
) -> Result<Self, Error> {
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<Duration>, 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<T, F>(connections: Arc<Connections<C>>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + 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<Option<Self::Value>, 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<CacheConnectionResult, Error> {
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::<Vec<_>>();
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::<Vec<redis::Value>>(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<f64, Error> {
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

8
litellm-rust/crates/cache-redis/src/cache/connection.rs vendored
View file

@ -12,7 +12,7 @@ use redis::{
use super::REDIS_TIMEOUT;
use crate::topology::RedisNode;
pub(super) struct PooledConnection<C> {
pub struct PooledConnection<C> {
pub(super) connection: C,
pub(super) failed: bool,
}
@ -20,7 +20,7 @@ pub(super) struct PooledConnection<C> {
/// 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<Self, Error> {
@ -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<Self, Error> {
@ -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),
}

28
litellm-rust/crates/cache-redis/src/cache/operations.rs vendored
View file

@ -144,7 +144,7 @@ where
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
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::<Vec<_>>();
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::<Vec<redis::Value>>(connection)
@ -188,7 +188,7 @@ where
}
pub async fn ping(&self) -> Result<bool, Error> {
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<Option<i64>, 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::<i64>(connection)
@ -208,7 +208,7 @@ where
pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, 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<RedisLpopResult, Error> {
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::<Vec<_>>();
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::<Vec<_>>();
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<i64, Error> {
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<f64, Error> {
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)

4
litellm-rust/crates/cache-redis/src/lib.rs
View file

@ -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,
};

2
litellm-rust/crates/cache-response/README.md
View file

@ -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

47
litellm-rust/crates/cache-response/src/response.rs
View file

@ -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<C: CacheContext = litellm_cache::ExactCacheContext> {
pub key: CacheKeyInput,
pub controls: CacheControls,
pub context: ExactCacheContext,
pub context: C,
pub max_age: Option<Duration>,
}
impl ResponseCacheRequest {
impl<C: CacheContext + Default> ResponseCacheRequest<C> {
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<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>>
where
B::Context: Default + PartialEq,
{
backend: Arc<B>,
}
impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCache<B> {
impl<B> ResponseCache<B>
where
B: BaseCache<Value = CacheEntry>,
B::Context: Default + PartialEq,
{
pub fn new(backend: Arc<B>) -> Self {
Self { backend }
}
@ -45,8 +52,12 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
&self.backend
}
pub fn backend_arc(&self) -> &Arc<B> {
&self.backend
}
pub fn default_ttl(&self) -> Option<Duration> {
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<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn lookup(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
now: Duration,
) -> Result<Option<Value>, Error> {
if !request.controls.reads() {
@ -81,7 +92,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_lookup(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
now: Duration,
) -> Result<Option<Value>, Error> {
if !request.controls.reads() {
@ -101,7 +112,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[ResponseCacheRequest<B::Context>],
now: Duration,
) -> Result<PartialHits, Error>
where
@ -126,7 +137,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[ResponseCacheRequest<B::Context>],
now: Duration,
) -> Result<PartialHits, Error>
where
@ -153,7 +164,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn store(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
response: Value,
now: Duration,
) -> Result<(), Error> {
@ -172,7 +183,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_store(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
response: Value,
now: Duration,
) -> Result<(), Error> {
@ -193,7 +204,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_store_batch(
&self,
entries: Vec<(ResponseCacheRequest, Value)>,
entries: Vec<(ResponseCacheRequest<B::Context>, Value)>,
now: Duration,
) -> Result<(), Error> {
self.async_store_entries(
@ -209,7 +220,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
/// the freshness of its original response.
pub async fn async_store_entries(
&self,
entries: Vec<(ResponseCacheRequest, Value, Duration)>,
entries: Vec<(ResponseCacheRequest<B::Context>, Value, Duration)>,
) -> Result<(), Error> {
let writable = entries
.into_iter()
@ -249,8 +260,8 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
}
fn partial_hits(
requests: &[ResponseCacheRequest],
readable: Vec<(usize, &ResponseCacheRequest)>,
requests: &[ResponseCacheRequest<B::Context>],
readable: Vec<(usize, &ResponseCacheRequest<B::Context>)>,
entries: Vec<BatchEntry<CacheEntry>>,
now: Duration,
) -> Result<PartialHits, Error> {

20
litellm-rust/crates/cache-valkey-semantic/Cargo.toml Normal file
View file

@ -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

1153
litellm-rust/crates/cache-valkey-semantic/src/lib.rs Normal file
View file

File diff suppressed because it is too large Load diff

50
litellm-rust/crates/cache/src/base_cache.rs vendored
View file

@ -32,6 +32,28 @@ impl CacheContext for ExactCacheContext {
}
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct SemanticCacheContext {
pub input: Option<serde_json::Value>,
pub messages: Option<serde_json::Value>,
pub metadata: Option<serde_json::Value>,
pub scope: Option<String>,
pub ttl: Option<Duration>,
}
impl CacheContext for SemanticCacheContext {
fn ttl(&self) -> Option<Duration> {
self.ttl
}
fn with_ttl(&self, ttl: Option<Duration>) -> 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<Output = Result<CacheConnectionResult, Error>> + 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);
}
}

2
litellm-rust/crates/cache/src/lib.rs vendored
View file

@ -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};

21
litellm-rust/crates/cache/tests/caching.rs vendored
View file

@ -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 {

6
litellm-rust/crates/python-bridge/Cargo.toml
View file

@ -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]]

21
litellm-rust/crates/python-bridge/src/cache/binding.rs vendored
View file

@ -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)

727
litellm-rust/crates/python-bridge/src/cache/config.rs vendored
View file

@ -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<u64>,
pub(super) embedding_timeout: Option<f64>,
}
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<RedisCacheConfig>),
S3(Box<S3CacheConfig>),
Gcs(GcsCacheConfig),
ValkeySemantic(Box<ValkeySemanticCacheConfig>),
Disk(DiskCacheConfig),
AzureBlob(AzureBlobCacheConfig),
RedisSemantic(Box<RedisSemanticCacheConfig>),
}
#[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<AzureBlobCacheConf
})
}
#[inline(never)]
pub(super) fn project_redis_semantic(
backend: &Bound<'_, PyAny>,
) -> PyResult<RedisSemanticCacheConfig> {
Ok(RedisSemanticCacheConfig {
redis_url: backend.getattr("_redis_url")?.extract::<String>()?,
index_name: backend
.getattr("_index_name")?
.extract::<Option<String>>()?
.unwrap_or_else(|| "litellm_semantic_cache_index".into()),
similarity_threshold: backend.getattr("similarity_threshold")?.extract::<f64>()?,
embedding_model: backend.getattr("embedding_model")?.extract::<String>()?,
embedding_max_input_tokens: backend
.getattr("embedding_max_input_tokens")?
.extract::<Option<u64>>()?,
embedding_timeout: backend
.getattr("embedding_timeout")?
.extract::<Option<f64>>()?,
})
}
#[inline(never)]
fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
let max_size_kib = backend.getattr("max_size_per_item")?.extract::<usize>()?;
@ -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<Result<RedisClientProjection<'py>, UnsupportedCacheConfig>> {
let Some(startup_nodes) = startup_nodes(source)? else {
@ -676,6 +758,71 @@ fn port(value: i64) -> PyResult<u16> {
u16::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis port"))
}
#[inline(never)]
fn project_valkey_semantic(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<ValkeySemanticCacheConfig, UnsupportedCacheConfig>> {
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::<usize>()?,
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<Result<(Bound<'py, PyDict>, bool), UnsupportedCacheConfig>> {
if !instance_class_is(pool, "redis.connection", "ConnectionPool")? {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let resolved = pool.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
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<RedisTlsConfig> {
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();

156
litellm-rust/crates/python-bridge/src/cache/embedder.rs vendored Normal file
View file

@ -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<Vec<f32>, Error>;
}
pub(super) fn with_prepared_embedding<F: Future>(
vector: Result<Vec<f32>, Error>,
future: F,
) -> impl Future<Output = F::Output> {
PREPARED_EMBEDDING.scope(vector, future)
}
pub(super) struct PythonEmbedder(Py<PyAny>);
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<PyAny>) -> Self {
Self(object)
}
pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self(backend.clone().unbind()))
}
pub(super) fn object(&self) -> &Py<PyAny> {
&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<Value>,
) -> PyResult<Bound<'py, PyAny>> {
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<Bound<'py, PyDict>> {
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<Py<PyAny>> {
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<Vec<f32>> {
Ok(vector
.extract::<Vec<f64>>()?
.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<Vec<f32>, Error> {
let result = Python::attach(|py| -> PyResult<Vec<f64>> {
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<Output = Result<Vec<f32>, 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<Vec<f32>, 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<Output = Result<Vec<f32>, 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));
}
}

137
litellm-rust/crates/python-bridge/src/cache/facade.rs vendored
View file

@ -31,17 +31,15 @@ struct RedisPoolGuard {
connection_class: Py<PyAny>,
connection_kwargs: Py<PyAny>,
max_connections: Option<usize>,
client_name: &'static str,
attributes: RedisPoolAttributes,
}
struct S3ClientGuard {
reference: Py<PyAny>,
}
struct DiskStoreGuard {
reference: Py<PyAny>,
directory: String,
}
struct AzureBlobClientGuard {
sync_client: Py<PyAny>,
async_client: Py<PyAny>,
@ -49,12 +47,18 @@ struct AzureBlobClientGuard {
container_name: String,
}
struct S3ClientGuard {
reference: Py<PyAny>,
}
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<Self> {
let pool = backend.getattr("redis_client")?.getattr(attributes.pool)?;
fn capture(
backend: &Bound<'_, PyAny>,
client_name: &'static str,
attributes: RedisPoolAttributes,
) -> PyResult<Self> {
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::<usize>())
.transpose()?,
client_name,
attributes,
})
}
fn max_connections(
pool: &Bound<'_, PyAny>,
attributes: &RedisPoolAttributes,
) -> PyResult<Option<usize>> {
attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
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::<usize>())
.transpose()?
&& self
.connection_kwargs
.bind(py)
@ -226,22 +238,6 @@ impl RedisPoolGuard {
}
}
impl S3ClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
reference: backend.getattr("s3_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
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<Self> {
let store = backend.getattr("disk_cache")?;
@ -261,6 +257,7 @@ impl DiskStoreGuard {
visit.call(&self.reference)
}
}
impl AzureBlobClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let sync_client = backend.getattr("container_client")?;
@ -289,11 +286,41 @@ impl AzureBlobClientGuard {
}
}
impl S3ClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
reference: backend.getattr("s3_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
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<Self> {
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",

86
litellm-rust/crates/python-bridge/src/cache/handle.rs vendored
View file

@ -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<Self> {
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<Self> {
@ -153,6 +183,37 @@ impl CacheTestHandle {
pid: std::process::id(),
})
}
#[staticmethod]
fn redis_semantic(py: Python<'_>, backend: Bound<'_, PyAny>) -> PyResult<Self> {
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::<Option<usize>>()?,
);
let service = service
.with_scope(
facade
.getattr("semantic_cache_scope")?
.extract::<String>()?,
)
.with_redis_flush_size(
facade
.getattr("redis_flush_size")?
.extract::<Option<usize>>()?,
);
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)?;
}

3
litellm-rust/crates/python-bridge/src/cache/mod.rs vendored
View file

@ -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::{

742
litellm-rust/crates/python-bridge/src/cache/native.rs vendored
View file

@ -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<ResponseCache<InMemoryCache<CacheEntry>>>),
@ -21,6 +83,15 @@ pub(super) enum NativeResponseCache {
},
S3(Arc<ResponseCache<S3Cache<ResponseCacheCodec>>>),
Gcs(Arc<ResponseCache<GcsCache<ResponseCacheCodec>>>),
ValkeySemantic {
cache: Arc<ResponseCache<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>>,
embedder: PythonEmbedder,
scope: String,
},
RedisSemantic {
cache: Arc<ResponseCache<RedisSemanticCache<PythonEmbedder>>>,
embedder: PythonEmbedder,
},
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
}
@ -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<Self, Error> {
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<Self, Error> {
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<Self, Error> {
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<ExactCacheContext> {
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<SemanticCacheContext> {
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<SemanticCacheContext> {
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<usize>) -> 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<usize>) -> 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<Option<Value>, 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<f32> {
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<PyAny>> {
match self {
Self::RedisSemantic { embedder, .. } => Some(embedder.object()),
_ => None,
}
}
pub fn lookup(&self, request: &NativeRequest, now: Duration) -> Result<Option<Value>, 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<PartialHits, Error> {
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::<Vec<_>>();
cache.lookup_batch(&requests, now)
}
Self::Redis { cache, .. } => {
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
cache.lookup_batch(&requests, now)
}
Self::S3(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => {
Err(Error::UnsupportedOperation)
}
Self::Gcs(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
Self::Disk(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
Self::AzureBlob(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
}
}
pub async fn async_lookup(
&self,
request: &ResponseCacheRequest,
request: &NativeRequest,
now: Duration,
) -> Result<Option<Value>, 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<Bound<'py, PyAny>> {
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<Bound<'py, PyAny>> {
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<PartialHits, Error> {
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::<Vec<_>>();
cache.async_lookup_batch(&requests, now).await
}
Self::Redis { cache, .. } => {
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
cache.async_lookup_batch(&requests, now).await
}
Self::S3(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => {
Err(Error::UnsupportedOperation)
}
Self::Gcs(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
Self::Disk(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
Self::AzureBlob(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), 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<Bound<'py, PyAny>> {
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<ResponseCacheCodec>> {
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());
}
}

50
litellm-rust/crates/python-bridge/src/cache/request.rs vendored
View file

@ -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<CacheControls>,
ttl_seconds: Option<f64>,
max_age_seconds: Option<f64>,
messages: Option<Value>,
input: Option<Value>,
metadata: Option<Value>,
litellm_metadata: Option<Value>,
litellm_params: Option<Value>,
scope: Option<String>,
}
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<ResponseCacheRequest> {
#[derive(Clone)]
pub(super) struct NativeRequest {
pub(super) key: CacheKeyInput,
pub(super) controls: CacheControls,
pub(super) ttl: Option<Duration>,
pub(super) max_age: Option<Duration>,
pub(super) messages: Option<Value>,
pub(super) input: Option<Value>,
pub(super) metadata: Option<Value>,
pub(super) litellm_metadata: Option<Value>,
pub(super) litellm_params: Option<Value>,
pub(super) scope: Option<String>,
}
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<NativeRequest> {
let input: RequestInput = from_py(value)?;
request_input(input)
}
fn request_input(input: RequestInput) -> PyResult<ResponseCacheRequest> {
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<NativeRequest> {
let controls = input.controls.unwrap_or_else(|| {
ResponseCacheRequest::<ExactCacheContext>::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<Vec<ResponseCacheRequest>> {
pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult<Vec<NativeRequest>> {
from_py::<Vec<RequestInput>>(value)?
.into_iter()
.map(request_input)

175
litellm-rust/crates/python-bridge/src/cache/semantic.rs vendored Normal file
View file

@ -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<Value>)>,
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<Vec<f32>, Error>,
) -> PyResult<ExecutionStep> {
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<Py<PyAny>>>) -> PyResult<ExecutionStep> {
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::<PyException>(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<Bound<'_, PyAny>> {
let execution = Py::new(py, Execution::new(body))?;
py.import("litellm.rust_bridge.lifecycle")?
.getattr("drive")?
.call1((execution,))
}

249
litellm-rust/crates/python-bridge/src/cache/semantic_step.rs vendored Normal file
View file

@ -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<Value>),
}
#[derive(Clone, Copy)]
enum State {
Start,
AwaitingEmbedding,
AwaitingStorage,
Done,
}
pub(super) struct SemanticEmbedExecution {
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
op: Op,
now: Option<Duration>,
prepared: Vec<Option<Vec<f32>>>,
index: usize,
state: State,
}
impl SemanticEmbedExecution {
pub(super) fn lookup(
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
request: ResponseCacheRequest<SemanticCacheContext>,
) -> 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<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
request: ResponseCacheRequest<SemanticCacheContext>,
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<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
responses: Vec<Value>,
) -> 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<ExecutionStep> {
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<ExecutionStep> {
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<Py<PyAny>>>,
) -> PyResult<ExecutionStep> {
match (self.state, result) {
(State::Start, None) => self.start(py),
(State::AwaitingEmbedding, Some(Ok(value))) => {
let values = value.bind(py).extract::<Vec<f64>>()?;
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<Py<PyAny>>>) -> PyResult<ExecutionStep> {
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<Bound<'py, PyAny>> {
let execution = Py::new(py, Execution::new(body))?;
py.import("litellm.rust_bridge.lifecycle")?
.getattr("drive")?
.call1((execution,))
}

6
litellm/constants.py
View file

@ -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))

4
litellm/litellm_core_utils/logging_worker.py
View file

@ -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):

2
litellm/llms/azure/azure.py
View file

@ -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)

23
litellm/llms/openai/videos/guardrail_translation/__init__.py Normal file
View file

@ -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")

48
litellm/llms/openai/videos/guardrail_translation/handler.py Normal file
View file

@ -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

52
litellm/model_prices_and_context_window_backup.json
View file

@ -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,

85
litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py
View file

@ -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,

42
litellm/proxy/_lazy_openapi_snapshot.json
View file

@ -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"
},

20
litellm/proxy/_types.py
View file

@ -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(

9
litellm/proxy/agent_endpoints/a2a_endpoints.py
View file

@ -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
}

18
litellm/proxy/agent_endpoints/agent_registry.py
View file

@ -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),
}

75
litellm/proxy/agent_endpoints/auth/agent_access_groups.py Normal file
View file

@ -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),
)

87
litellm/proxy/agent_endpoints/auth/agent_caller.py Normal file
View file

@ -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

76
litellm/proxy/agent_endpoints/auth/agent_permission_handler.py
View file

@ -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):

99
litellm/proxy/auth/auth_checks.py
View file

@ -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

11
litellm/proxy/auth/handle_jwt.py
View file

@ -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:

35
litellm/proxy/auth/route_checks.py
View file

@ -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

4
litellm/proxy/auth/user_api_key_auth.py
View file

@ -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

52
litellm/proxy/client/cli/commands/claude_settings.py
View file

@ -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)

4
litellm/proxy/client/cli/commands/statusline_script.py
View file

@ -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

15
litellm/proxy/config_resolvers/__init__.py
View file

@ -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",
)

7
litellm/proxy/config_resolvers/settings_store.py
View file

@ -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

170
litellm/proxy/db/autorouter_session_rollup.py
View file

@ -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)

40
litellm/proxy/db/baseline_accounting.py
View file

@ -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)

61
litellm/proxy/db/db_transaction_queue/spend_log_cleanup.py
View file

@ -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)

72
litellm/proxy/management_endpoints/access_group_endpoints.py
View file

@ -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
)

11
litellm/proxy/management_endpoints/auto_router_endpoints.py
View file

@ -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 = (

32
litellm/proxy/management_endpoints/router_settings_endpoints.py
View file

@ -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)

98
litellm/proxy/proxy_server.py
View file

@ -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)

42
litellm/proxy/schema.prisma
View file

@ -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

79
litellm/proxy/shutdown/scheduled_jobs.py Normal file
View file

@ -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,
)

88
litellm/proxy/ui_crud_endpoints/proxy_setting_endpoints.py
View file

@ -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,
)

60
litellm/proxy/utils.py
View file

@ -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:

5
litellm/router_strategy/complexity_router/config.py
View file

@ -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(

9
litellm/rust_bridge/_native.pyi
View file

@ -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,
*,

19
litellm/types/agents.py
View file

@ -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

2
litellm/types/utils.py
View file

@ -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"

52
model_prices_and_context_window.json
View file

@ -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,

2
pyproject.toml
View file

@ -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",

42
schema.prisma
View file

@ -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

1
tests/e2e/coverage_registry/guardrail.yaml
View file

@ -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"}

1
tests/e2e/coverage_registry/reliability.yaml
View file

@ -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"}

29
tests/e2e/e2e_http.py
View file

@ -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."""

1
tests/e2e/gateway/stage_mirror_ci_config.yml
View file

@ -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

22
tests/e2e/guardrails/guardrails_client.py
View file

@ -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,

69
tests/e2e/guardrails/test_key_guardrail_video_e2e.py Normal file
View file

@ -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}")

33
tests/e2e/models.py
View file

@ -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

15
tests/e2e/proxy_client.py
View file

@ -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:

25
tests/e2e/router/reliability_support.py
View file

@ -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))

139
tests/e2e/router/test_reliability_cancel_on_disconnect_e2e.py Normal file
View file

@ -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"
)

2
tests/e2e/router/test_reliability_cooldowns_e2e.py
View file

@ -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

15
tests/e2e/transport.py
View file

@ -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)

98
tests/e2e/ui/helpers/roundTrip.ts
View file

@ -12,17 +12,111 @@ export async function captureRequestBody(
match: { method: string; urlIncludes: string },
action: () => Promise<void>,
): Promise<Record<string, any>> {
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<string, any>;
}
/** Reads an endpoint as the master key, so a failure is bad data and not an expired UI token. */
export async function readBack<T = any>(page: Page, endpoint: string): Promise<T> {
export async function readBack<T = any>(
page: Page,
endpoint: string,
): Promise<T> {
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<void>,
): Promise<OperationOutcome> {
return Promise.resolve()
.then(action)
.then(
() => ({ status: "success" as const }),
(error: unknown) => ({ status: "failure" as const, error }),
);
}
async function runCleanup(
cleanup: () => boolean | Promise<boolean>,
): Promise<OperationOutcome> {
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<void>,
cleanup: () => boolean | Promise<boolean>,
): Promise<void> {
const actionOutcome = await runAction(action);
const cleanupOutcome = await runCleanup(cleanup);
const failure = toRunFailure(actionOutcome, cleanupOutcome);
if (failure !== null) raiseRunFailure(failure);
}

75
tests/e2e/ui/tests/prompts/addPrompt.spec.ts Normal file
View file

@ -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();
},
);
});
});

84
tests/e2e/ui/tests/tagManagement/tagManagement.spec.ts Normal file
View file

@ -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<Array<{ name: string }>>(
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();
},
);
});
});

171
tests/proxy_behavior/spend/test_autorouter_session_rollup.py
View file

@ -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"] == {}

67
tests/proxy_behavior/spend/test_baseline_accounting.py
View file

@ -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

47
tests/test_litellm/litellm_core_utils/test_logging_worker.py
View file

@ -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)

27
tests/test_litellm/llms/azure/test_azure.py
View file

@ -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,
)

3
tests/test_litellm/llms/azure/test_azure_common_utils.py
View file

@ -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

146
tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py
View file

@ -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(

146
tests/test_litellm/proxy/agent_endpoints/auth/test_agent_access_groups.py Normal file
View file

@ -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

57
tests/test_litellm/proxy/agent_endpoints/auth/test_agent_caller.py Normal file
View file

@ -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()

129
tests/test_litellm/proxy/agent_endpoints/auth/test_agent_permission_handler.py
View file

@ -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."""

19
tests/test_litellm/proxy/agent_endpoints/test_a2a_endpoints.py
View file

@ -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):

136
tests/test_litellm/proxy/agent_endpoints/test_agent_registry.py
View file

@ -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)

183
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -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 == []

184
tests/test_litellm/proxy/auth/test_handle_jwt.py
View file

@ -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

Some files were not shown because too many files have changed in this diff Show more