scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Merge remote-tracking branch 'upstream/main' into litellm_tts_true_streaming

Browse source
This commit is contained in:
Gabriel Creti 2026-11-04 12:09:27 +01:00
commit 4163aae38e
No known key found for this signature in database
GPG key ID: E3015A42D3CA9D22
468 changed files with 27265 additions and 15701 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/osv-scan.yml vendored
View file

@ -5,6 +5,7 @@ on:
branches:
- main
- "litellm_**"
- "rc/**"
schedule:
- cron: "23 6 * * *"
workflow_dispatch:

9
litellm-proxy-extras/litellm_proxy_extras/migrations/20261104000000_skill_zip_upload/migration.sql Normal file
View file

@ -0,0 +1,9 @@
SET LOCAL lock_timeout = '10s';
CREATE TABLE IF NOT EXISTS "LiteLLM_ClaudeCodePluginUpload" (
"plugin_name" TEXT NOT NULL,
"content" BYTEA NOT NULL,
CONSTRAINT "LiteLLM_ClaudeCodePluginUpload_pkey" PRIMARY KEY ("plugin_name"),
CONSTRAINT "LiteLLM_ClaudeCodePluginUpload_plugin_name_fkey" FOREIGN KEY ("plugin_name")
REFERENCES "LiteLLM_ClaudeCodePluginTable"("name") ON DELETE CASCADE ON UPDATE CASCADE
);

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

@ -1615,9 +1615,18 @@ model LiteLLM_ClaudeCodePluginTable {
updated_at DateTime? @default(now()) @updatedAt
created_by String?
uploaded_archive LiteLLM_ClaudeCodePluginUpload?
@@map("LiteLLM_ClaudeCodePluginTable")
}
model LiteLLM_ClaudeCodePluginUpload {
plugin_name String @id
content Bytes
plugin LiteLLM_ClaudeCodePluginTable @relation(fields: [plugin_name], references: [name], onDelete: Cascade, onUpdate: Cascade)
}
// User/team-scoped memory store with a GLOBAL unique key.
// `value` is a string (typically markdown/text meant for LLM context);
// `metadata` is an optional JSON envelope for structured tags without schema changes.

87
litellm-rust/Cargo.lock generated
View file

@ -3569,37 +3569,6 @@ dependencies = [
"uuid",
]
[[package]]
name = "litellm-cache-redis"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-testing",
"litellm-error",
"r2d2",
"redis",
"redis-test",
"rstest",
"serde_json",
"thiserror 2.0.19",
"tokio",
]
[[package]]
name = "litellm-cache-redis-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-testing",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"tokio",
]
[[package]]
name = "litellm-cache-response"
version = "0.1.0"
@ -3607,11 +3576,8 @@ dependencies = [
"litellm-cache",
"litellm-cache-gcs",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-http",
"py_literal",
"redis",
"redis-test",
"rstest",
"serde",
"serde_json",
@ -3650,24 +3616,6 @@ dependencies = [
"litellm-cache",
]
[[package]]
name = "litellm-cache-valkey-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-testing",
"litellm-error",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"thiserror 2.0.19",
"tokio",
"uuid",
]
[[package]]
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
@ -4339,8 +4287,6 @@ dependencies = [
"litellm-auth",
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-cache-response",
"litellm-callbacks-legacy-python",
"litellm-core-utils",
@ -4369,7 +4315,6 @@ dependencies = [
"prost",
"pyo3",
"pyo3-async-runtimes",
"redis",
"reqwest",
"rstest",
"serde",
@ -5782,17 +5727,6 @@ version = "6.0.0"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf"
[[package]]
name = "r2d2"
version = "0.8.10"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "51de85fb3fb6524929c8a2eb85e6b6d363de4e8c48f9e2c2eac4944abc181c93"
dependencies = [
"log",
"parking_lot",
"scheduled-thread-pool",
]
[[package]]
name = "rand"
version = "0.8.7"
@ -5966,18 +5900,6 @@ dependencies = [
"xxhash-rust",
]
[[package]]
name = "redis-test"
version = "1.0.4"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "804d36862e4323b69f96440cbb13c9894fc90176abdeaf91264e21d5d77f6aca"
dependencies = [
"rand 0.9.5",
"redis",
"socket2 0.6.5",
"tempfile",
]
[[package]]
name = "redox_syscall"
version = "0.5.18"
@ -6454,15 +6376,6 @@ dependencies = [
"windows-sys 0.61.2",
]
[[package]]
name = "scheduled-thread-pool"
version = "0.2.7"
source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "3cbc66816425a074528352f5789333ecff06ca41b36b0b0efdfbb29edc391a19"
dependencies = [
"parking_lot",
]
[[package]]
name = "schemars"
version = "0.9.0"

3
litellm-rust/Cargo.toml
View file

@ -57,11 +57,9 @@ litellm-db-testing = { path = "crates/db-testing" }
litellm-cache = { path = "crates/cache" }
litellm-cache-azure-blob = { path = "crates/cache-azure-blob" }
litellm-cache-memory = { path = "crates/cache-memory" }
litellm-cache-redis = { path = "crates/cache-redis" }
litellm-cache-s3 = { path = "crates/cache-s3" }
litellm-cache-gcs = { path = "crates/cache-gcs" }
litellm-cache-disk = { path = "crates/cache-disk" }
litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" }
litellm-cache-response = { path = "crates/cache-response" }
litellm-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" }
litellm-cache-testing = { path = "crates/cache-testing" }
@ -97,7 +95,6 @@ macro_rules_attribute = "0.2.3"
schemars = "1"
reqwest = { version = "0.13", default-features = false, features = ["json", "multipart", "rustls", "http2", "stream", "query"] }
redis = "1.8.0"
redis-test = "1.0.4"
testcontainers-modules = "0.15.0"
qdrant-client = { version = "1.19.0", default-features = false }
uuid = { version = "1", features = ["v4"] }

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

@ -1,19 +0,0 @@
[package]
name = "litellm-cache-redis-semantic"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-cache-redis.workspace = true
redis = { workspace = true, features = ["tls-rustls"] }
sha2.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
redis-test.workspace = true
rstest.workspace = true
serde_json.workspace = true
tokio.workspace = true

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

@ -1,404 +0,0 @@
use std::{
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, CacheCodec, Error, SemanticCacheContext,
semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
};
use litellm_cache_redis::{
RedisTopology,
connection::{ConnectionRef, Connections},
};
use sha2::{Digest, Sha256};
use crate::{
RedisSemanticConfig,
index::{CACHE_KEY_FIELD, Index, VECTOR_FIELD},
reply::{bytes_field, first_document, number_field, string_field},
};
struct Inner {
index: Index,
distance_threshold: f64,
clock: fn() -> f64,
}
impl Inner {
fn new(config: RedisSemanticConfig, clock: fn() -> f64) -> Self {
Self {
index: Index::new(config.index_name),
distance_threshold: 1.0 - f64::from(config.similarity_threshold),
clock,
}
}
fn store(
&self,
connection: &mut ConnectionRef<'_>,
tag: &str,
response: Vec<u8>,
prompt: &str,
vector: &[f32],
ttl: Option<Duration>,
) -> Result<(), Error> {
let index = self.index.ensure(connection, vector.len())?;
let entry_id = entry_id(prompt, tag);
let hash_key = format!("{index}:{entry_id}");
redis::cmd("HSET")
.arg(&hash_key)
.arg("entry_id")
.arg(&entry_id)
.arg("prompt")
.arg(prompt)
.arg("response")
.arg(response)
.arg(VECTOR_FIELD)
.arg(vector_buffer(vector))
.arg("inserted_at")
.arg(format!("{}", (self.clock)()))
.arg("updated_at")
.arg(format!("{}", (self.clock)()))
.arg(CACHE_KEY_FIELD)
.arg(tag)
.query::<()>(connection)
.map_err(Error::unavailable)?;
if let Some(ttl) = ttl {
redis::cmd("EXPIRE")
.arg(&hash_key)
.arg(ttl_seconds(ttl))
.query::<()>(connection)
.map_err(Error::unavailable)?;
}
Ok(())
}
fn lookup(
&self,
connection: &mut ConnectionRef<'_>,
tag: &str,
vector: &[f32],
) -> Result<SemanticLookup<Vec<u8>>, Error> {
let index = self.index.ensure(connection, vector.len())?;
let query = format!(
"(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]",
escape_tag(tag)
);
let result = redis::cmd("FT.SEARCH")
.arg(&index)
.arg(query)
.arg("RETURN")
.arg(8)
.arg("entry_id")
.arg("prompt")
.arg("response")
.arg("inserted_at")
.arg("updated_at")
.arg("metadata")
.arg(CACHE_KEY_FIELD)
.arg("vector_distance")
.arg("SORTBY")
.arg("vector_distance")
.arg("ASC")
.arg("DIALECT")
.arg(2)
.arg("LIMIT")
.arg(0)
.arg(1)
.arg("PARAMS")
.arg(2)
.arg("vector")
.arg(vector_buffer(vector))
.query::<redis::Value>(connection)
.map_err(Error::unavailable)?;
let Some(fields) = first_document(&result) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) {
return Ok(SemanticLookup::miss(Some(0.0)));
}
// redisvl's range query only returns entries within the distance threshold, so a
// farther hit reads as no result.
let Some(distance) = number_field(fields, "vector_distance")
.filter(|distance| *distance <= self.distance_threshold)
else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let Some(response) = bytes_field(fields, "response") else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
Ok(SemanticLookup {
value: Some(response),
similarity: Some(1.0 - distance),
})
}
}
/// `RedisSemanticCache`: a redisvl-compatible semantic index on Redis Stack. Values go through
/// the injected codec, so the response layer decides what a cached entry is.
pub struct RedisSemanticCache<E, S, C = redis::Connection> {
connections: Arc<Connections<C>>,
embedder: E,
codec: S,
inner: Arc<Inner>,
}
impl<E: Embedder, S: CacheCodec> RedisSemanticCache<E, S> {
pub fn new(
url: &str,
embedder: E,
codec: S,
config: RedisSemanticConfig,
) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
codec,
inner: Arc::new(Inner::new(config, timestamp)),
})
}
}
impl<E, S, C> RedisSemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub fn with_connection(
connection: C,
embedder: E,
codec: S,
config: RedisSemanticConfig,
) -> Self {
Self {
connections: Arc::new(Connections::fixed(connection)),
embedder,
codec,
inner: Arc::new(Inner::new(config, timestamp)),
}
}
pub fn with_clock(self, clock: fn() -> f64) -> Self {
let config = RedisSemanticConfig {
index_name: self.index_name().to_owned(),
similarity_threshold: self.similarity_threshold(),
};
Self {
inner: Arc::new(Inner::new(config, clock)),
..self
}
}
pub fn embedder(&self) -> &E {
&self.embedder
}
pub fn index_name(&self) -> &str {
self.inner.index.name()
}
pub fn similarity_threshold(&self) -> f32 {
(1.0 - self.inner.distance_threshold) as f32
}
fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str {
context.scope.as_deref().unwrap_or(key)
}
fn decode(&self, lookup: SemanticLookup<Vec<u8>>) -> Result<SemanticLookup<S::Value>, Error> {
Ok(SemanticLookup {
value: lookup
.value
.map(|bytes| self.codec.decode(&bytes))
.transpose()?,
similarity: lookup.similarity,
})
}
}
impl<E, S, C> BaseCache for RedisSemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = S::Value;
type Context = SemanticCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<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 response = self.codec.encode(&value)?;
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context);
self.connections.execute(|connection| {
self.inner
.store(connection, tag, response, &prompt, &vector, context.ttl)
})
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.get_cache_with_similarity(key, context)
.map(|lookup| lookup.value)
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> Result<(), Error> {
let Some(prompt) = prompt_from_context(&context) else {
return Ok(());
};
let response = self.codec.encode(&value)?;
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, &context).to_owned();
let inner = Arc::clone(&self.inner);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
inner.store(connection, &tag, response, &prompt, &vector, context.ttl)
})
.await
}
async fn async_get_cache(
&self,
key: &str,
context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.async_get_cache_with_similarity(key, context)
.await
.map(|lookup| lookup.value)
}
}
/// Python stamps a similarity of `0.0` when there is no prompt or no hit in the key's scope.
impl<E, S, C> SemanticCache for RedisSemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context);
let lookup = self
.connections
.execute(|connection| self.inner.lookup(connection, tag, &vector))?;
self.decode(lookup)
}
async fn async_get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, context).to_owned();
let inner = Arc::clone(&self.inner);
let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
inner.lookup(connection, &tag, &vector)
})
.await?;
self.decode(lookup)
}
}
fn timestamp() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs_f64())
.unwrap_or_default()
}
fn entry_id(prompt: &str, tag: &str) -> String {
let mut digest = Sha256::new();
digest.update(prompt.as_bytes());
digest.update(CACHE_KEY_FIELD.as_bytes());
digest.update(tag.as_bytes());
format!("{:x}", digest.finalize())
}
fn vector_buffer(vector: &[f32]) -> Vec<u8> {
vector
.iter()
.flat_map(|component| component.to_le_bytes())
.collect()
}
fn escape_tag(value: &str) -> String {
let mut escaped = String::with_capacity(value.len());
for ch in value.chars() {
if matches!(
ch,
',' | '.'
| '<'
| '>'
| '{'
| '}'
| '['
| ']'
| '\\'
| '"'
| '\''
| ':'
| ';'
| '!'
| '@'
| '#'
| '$'
| '%'
| '^'
| '&'
| '*'
| '('
| ')'
| '-'
| '+'
| '='
| '~'
| '|'
| '/'
| ' '
| '?'
) {
escaped.push('\\');
}
escaped.push(ch);
}
escaped
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}

8
litellm-rust/crates/cache-redis-semantic/src/config.rs
View file

@ -1,8 +0,0 @@
/// `RedisSemanticCache.DEFAULT_REDIS_INDEX_NAME`.
pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index";
#[derive(Clone, Debug)]
pub struct RedisSemanticConfig {
pub index_name: String,
pub similarity_threshold: f32,
}

209
litellm-rust/crates/cache-redis-semantic/src/index.rs
View file

@ -1,209 +0,0 @@
use std::sync::OnceLock;
use litellm_cache::Error;
use litellm_cache_redis::connection::ConnectionRef;
use crate::reply::{number_value, string_value};
pub(crate) const CACHE_KEY_FIELD: &str = "litellm_cache_key";
pub(crate) const VECTOR_FIELD: &str = "prompt_vector";
/// The redisvl `SemanticCache` index, resolved once per cache: the configured name when its
/// schema fits, else `<name>_isolated`, recreated when that one is stale too.
pub(crate) struct Index {
name: String,
resolved: OnceLock<String>,
}
impl Index {
pub(crate) fn new(name: String) -> Self {
Self {
name,
resolved: OnceLock::new(),
}
}
pub(crate) fn name(&self) -> &str {
&self.name
}
pub(crate) fn ensure(
&self,
connection: &mut ConnectionRef<'_>,
dims: usize,
) -> Result<String, Error> {
if let Some(name) = self.resolved.get() {
return Ok(name.clone());
}
let name = match index_compatible(connection, &self.name, dims)? {
Some(true) => self.name.clone(),
Some(false) => self.isolated(connection, dims)?,
None => match create_index(connection, &self.name, dims) {
Ok(()) => self.name.clone(),
Err(error) => match index_compatible(connection, &self.name, dims)? {
Some(true) => self.name.clone(),
Some(false) => self.isolated(connection, dims)?,
None => return Err(error),
},
},
};
let _ = self.resolved.set(name.clone());
Ok(name)
}
fn isolated(&self, connection: &mut ConnectionRef<'_>, dims: usize) -> Result<String, Error> {
let name = format!("{}_isolated", self.name);
match index_compatible(connection, &name, dims)? {
Some(true) => Ok(name),
Some(false) => {
redis::cmd("FT.DROPINDEX")
.arg(&name)
.query::<()>(connection)
.map_err(Error::unavailable)?;
create_index(connection, &name, dims)?;
Ok(name)
}
None => {
create_index(connection, &name, dims)?;
Ok(name)
}
}
}
}
fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> {
redis::cmd("FT.CREATE")
.arg(name)
.arg("ON")
.arg("HASH")
.arg("PREFIX")
.arg(1)
.arg(name)
.arg("SCORE")
.arg(1.0)
.arg("SCHEMA")
.arg("prompt")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("response")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("inserted_at")
.arg("NUMERIC")
.arg("updated_at")
.arg("NUMERIC")
.arg(VECTOR_FIELD)
.arg("VECTOR")
.arg("FLAT")
.arg(6)
.arg("TYPE")
.arg("FLOAT32")
.arg("DIM")
.arg(dims)
.arg("DISTANCE_METRIC")
.arg("COSINE")
.arg(CACHE_KEY_FIELD)
.arg("TAG")
.arg("SEPARATOR")
.arg(",")
.query::<()>(connection)
.map_err(Error::unavailable)
}
fn index_compatible(
connection: &mut ConnectionRef<'_>,
name: &str,
dims: usize,
) -> Result<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(error) => return Err(Error::unavailable(error)),
};
Ok(Some(schema_compatible(&info, dims)))
}
fn unknown_index(error: &redis::RedisError) -> bool {
error
.code()
.is_some_and(|code| code.eq_ignore_ascii_case("unknown"))
&& error
.detail()
.is_some_and(|detail| detail.to_ascii_lowercase().starts_with("index"))
}
struct Attribute {
name: Option<String>,
field_type: Option<String>,
dim: Option<f64>,
data_type: Option<String>,
distance_metric: Option<String>,
}
fn attribute(value: &redis::Value) -> Option<Attribute> {
let redis::Value::Array(pairs) = value else {
return None;
};
let mut attribute = Attribute {
name: None,
field_type: None,
dim: None,
data_type: None,
distance_metric: None,
};
for pair in pairs.as_chunks::<2>().0 {
match string_value(&pair[0]).as_deref() {
Some("identifier") => attribute.name = string_value(&pair[1]),
Some("type") => attribute.field_type = string_value(&pair[1]),
Some("dim") => attribute.dim = number_value(&pair[1]),
Some("data_type") => attribute.data_type = string_value(&pair[1]),
Some("distance_metric") => attribute.distance_metric = string_value(&pair[1]),
_ => {}
}
}
Some(attribute)
}
fn schema_compatible(info: &redis::Value, dims: usize) -> bool {
let redis::Value::Array(entries) = info else {
return false;
};
let attributes = entries
.as_chunks::<2>()
.0
.iter()
.find(|pair| string_value(&pair[0]).as_deref() == Some("attributes"))
.map(|pair| &pair[1]);
let Some(redis::Value::Array(attributes)) = attributes else {
return false;
};
let fields = attributes.iter().filter_map(attribute).collect::<Vec<_>>();
let has_field = |name: &str, field_type: &str| {
fields.iter().any(|field| {
field.name.as_deref() == Some(name) && field.field_type.as_deref() == Some(field_type)
})
};
has_field("prompt", "TEXT")
&& has_field("response", "TEXT")
&& has_field("inserted_at", "NUMERIC")
&& has_field("updated_at", "NUMERIC")
&& has_field(CACHE_KEY_FIELD, "TAG")
&& fields.iter().any(|field| {
field.name.as_deref() == Some(VECTOR_FIELD)
&& field.field_type.as_deref() == Some("VECTOR")
&& field.dim == Some(dims as f64)
&& field
.data_type
.as_deref()
.is_some_and(|data| data.eq_ignore_ascii_case("float32"))
&& field
.distance_metric
.as_deref()
.is_some_and(|metric| metric.eq_ignore_ascii_case("cosine"))
})
}

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

@ -1,7 +0,0 @@
mod cache;
mod config;
mod index;
mod reply;
pub use cache::RedisSemanticCache;
pub use config::{DEFAULT_INDEX_NAME, RedisSemanticConfig};

57
litellm-rust/crates/cache-redis-semantic/src/reply.rs
View file

@ -1,57 +0,0 @@
pub(crate) 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,
}
}
pub(crate) 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()),
}
}
pub(crate) fn first_document(result: &redis::Value) -> Option<&[redis::Value]> {
let redis::Value::Array(items) = result else {
return None;
};
let [count, _document_id, fields, ..] = items.as_slice() else {
return None;
};
if !matches!(count, redis::Value::Int(count) if *count > 0) {
return None;
}
match fields {
redis::Value::Array(fields) => Some(fields.as_slice()),
_ => None,
}
}
fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> {
fields
.as_chunks::<2>()
.0
.iter()
.find(|pair| string_value(&pair[0]).as_deref() == Some(name))
.map(|pair| &pair[1])
}
pub(crate) fn string_field(fields: &[redis::Value], name: &str) -> Option<String> {
field_value(fields, name).and_then(string_value)
}
pub(crate) fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
field_value(fields, name).and_then(number_value)
}
pub(crate) 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,
}
}

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

@ -1,778 +0,0 @@
mod support;
use std::time::Duration;
use litellm_cache::{
BaseCache, Error, JsonCodec, SemanticCacheContext,
semantic::{SemanticCache, SemanticLookup},
};
use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig};
use redis_test::{MockCmd, MockRedisConnection};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use sha2::{Digest, Sha256};
use support::FakeEmbedder;
const INDEX: &str = DEFAULT_INDEX_NAME;
const PROMPT: &str = "hello prompt";
const CLOCK: fn() -> f64 = || 1700000000.5;
const VECTOR: [f32; 3] = [0.1, 0.2, 0.3];
type MockCache = RedisSemanticCache<FakeEmbedder, JsonCodec<Value>, MockRedisConnection>;
#[fixture]
fn config() -> RedisSemanticConfig {
RedisSemanticConfig {
index_name: INDEX.into(),
similarity_threshold: 0.9,
}
}
#[fixture]
fn entry() -> Value {
json!({"timestamp": 1.0, "response": {"answer": "yes"}})
}
#[fixture]
fn context() -> SemanticCacheContext {
messages_context(vec![json!({"role": "user", "content": PROMPT})])
}
fn messages_context(messages: Vec<Value>) -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(Value::Array(messages)),
..Default::default()
}
}
fn cache(commands: Vec<MockCmd>, embedder: FakeEmbedder) -> MockCache {
RedisSemanticCache::with_connection(
MockRedisConnection::new(commands).assert_all_commands_consumed(),
embedder,
JsonCodec::new(),
config(),
)
.with_clock(CLOCK)
}
fn encoded(value: &Value) -> Vec<u8> {
serde_json::to_vec(value).unwrap()
}
fn vector_bytes(vector: &[f32]) -> Vec<u8> {
vector
.iter()
.flat_map(|component| component.to_le_bytes())
.collect()
}
fn entry_id(prompt: &str, tag: &str) -> String {
let mut digest = Sha256::new();
digest.update(prompt.as_bytes());
digest.update(b"litellm_cache_key");
digest.update(tag.as_bytes());
format!("{:x}", digest.finalize())
}
fn s(value: &str) -> redis::Value {
redis::Value::BulkString(value.as_bytes().to_vec())
}
fn unknown_index_error() -> redis::RedisError {
support::server_error("Unknown index name")
}
fn info_missing(index: &str) -> MockCmd {
MockCmd::new(
redis::cmd("FT.INFO").arg(index),
Err::<redis::Value, _>(unknown_index_error()),
)
}
fn info(index: &str, value: redis::Value) -> MockCmd {
MockCmd::new(redis::cmd("FT.INFO").arg(index), Ok(value))
}
fn attribute(name: &str, field_type: &str, extra: Vec<redis::Value>) -> redis::Value {
let mut parts = vec![
s("identifier"),
s(name),
s("attribute"),
s(name),
s("type"),
s(field_type),
];
parts.extend(extra);
redis::Value::Array(parts)
}
fn index_info(attributes: Vec<redis::Value>) -> redis::Value {
redis::Value::Array(vec![
s("index_name"),
s(INDEX),
s("attributes"),
redis::Value::Array(attributes),
])
}
fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> redis::Value {
attribute(
"prompt_vector",
"VECTOR",
vec![
s("algorithm"),
s("FLAT"),
s("data_type"),
s(data_type),
s("dim"),
redis::Value::Int(dims),
s("distance_metric"),
s(distance_metric),
],
)
}
fn info_with_vector(vector: redis::Value) -> redis::Value {
index_info(vec![
attribute("prompt", "TEXT", vec![]),
attribute("response", "TEXT", vec![]),
attribute("inserted_at", "NUMERIC", vec![]),
attribute("updated_at", "NUMERIC", vec![]),
vector,
attribute("litellm_cache_key", "TAG", vec![]),
])
}
fn compatible_info(dims: i64) -> redis::Value {
info_with_vector(vector_attribute_with(dims, "FLOAT32", "COSINE"))
}
fn unscoped_info(dims: i64) -> redis::Value {
index_info(vec![
attribute("prompt", "TEXT", vec![]),
attribute("response", "TEXT", vec![]),
attribute("inserted_at", "NUMERIC", vec![]),
attribute("updated_at", "NUMERIC", vec![]),
vector_attribute_with(dims, "FLOAT32", "COSINE"),
])
}
fn create_index(name: &str, dims: usize) -> MockCmd {
MockCmd::new(create_index_command(name, dims), Ok("OK"))
}
fn create_index_command(name: &str, dims: usize) -> redis::Cmd {
let mut command = redis::cmd("FT.CREATE");
command
.arg(name)
.arg("ON")
.arg("HASH")
.arg("PREFIX")
.arg(1)
.arg(name)
.arg("SCORE")
.arg(1.0)
.arg("SCHEMA")
.arg("prompt")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("response")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("inserted_at")
.arg("NUMERIC")
.arg("updated_at")
.arg("NUMERIC")
.arg("prompt_vector")
.arg("VECTOR")
.arg("FLAT")
.arg(6)
.arg("TYPE")
.arg("FLOAT32")
.arg("DIM")
.arg(dims)
.arg("DISTANCE_METRIC")
.arg("COSINE")
.arg("litellm_cache_key")
.arg("TAG")
.arg("SEPARATOR")
.arg(",");
command
}
fn hset(index: &str, prompt: &str, tag: &str, vector: &[f32], value: &Value) -> MockCmd {
MockCmd::new(
redis::cmd("HSET")
.arg(format!("{index}:{}", entry_id(prompt, tag)))
.arg("entry_id")
.arg(entry_id(prompt, tag))
.arg("prompt")
.arg(prompt)
.arg("response")
.arg(encoded(value))
.arg("prompt_vector")
.arg(vector_bytes(vector))
.arg("inserted_at")
.arg("1700000000.5")
.arg("updated_at")
.arg("1700000000.5")
.arg("litellm_cache_key")
.arg(tag),
Ok(7),
)
}
fn search(
index: &str,
tag: &str,
vector: &[f32],
reply: redis::RedisResult<redis::Value>,
) -> MockCmd {
let mut command = redis::cmd("FT.SEARCH");
command
.arg(index)
.arg(format!(
"(@litellm_cache_key:{{{tag}}})=>[KNN 1 @prompt_vector $vector AS vector_distance]"
))
.arg("RETURN")
.arg(8)
.arg("entry_id")
.arg("prompt")
.arg("response")
.arg("inserted_at")
.arg("updated_at")
.arg("metadata")
.arg("litellm_cache_key")
.arg("vector_distance")
.arg("SORTBY")
.arg("vector_distance")
.arg("ASC")
.arg("DIALECT")
.arg(2)
.arg("LIMIT")
.arg(0)
.arg(1)
.arg("PARAMS")
.arg(2)
.arg("vector")
.arg(vector_bytes(vector));
MockCmd::new(command, reply)
}
fn hit(tag: &str, distance: &str, response: Vec<u8>) -> redis::Value {
redis::Value::Array(vec![
redis::Value::Int(1),
s("litellm_semantic_cache_index:stored-id"),
redis::Value::Array(vec![
s("entry_id"),
s("stored-id"),
s("prompt"),
s(PROMPT),
s("response"),
redis::Value::BulkString(response),
s("inserted_at"),
s("1700000000.5"),
s("updated_at"),
s("1700000000.5"),
s("litellm_cache_key"),
s(tag),
s("vector_distance"),
s(distance),
]),
])
}
fn empty_result() -> redis::Value {
redis::Value::Array(vec![redis::Value::Int(0)])
}
#[rstest]
#[case::creates_index_and_expires(false, Some(Duration::from_secs(5)), Some(5))]
#[case::existing_index_without_ttl(true, None, None)]
#[case::fractional_ttl_rounds_up(true, Some(Duration::from_millis(1500)), Some(2))]
fn store_writes_the_redisvl_hash(
#[case] index_exists: bool,
#[case] ttl: Option<Duration>,
#[case] expire: Option<u64>,
entry: Value,
context: SemanticCacheContext,
) {
let hash_key = format!("{INDEX}:{}", entry_id(PROMPT, "key1"));
let mut commands = if index_exists {
vec![info(INDEX, compatible_info(3))]
} else {
vec![info_missing(INDEX), create_index(INDEX, 3)]
};
commands.push(hset(INDEX, PROMPT, "key1", &VECTOR, &entry));
commands.extend(
expire.map(|seconds| MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(seconds), Ok(1))),
);
let cache = cache(commands, FakeEmbedder::new(&[]));
cache
.set_cache("key1", entry, &SemanticCacheContext { ttl, ..context })
.unwrap();
}
#[rstest]
#[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))]
#[case::above_distance_threshold("key1", "0.5", None, Ok(None))]
#[case::other_cache_key("other", "0.05", None, Ok(None))]
#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry(None)))]
fn lookup_applies_threshold_scope_and_codec(
#[case] stored_tag: &str,
#[case] distance: &str,
#[case] response: Option<&[u8]>,
#[case] expected: Result<Option<Value>, Error>,
entry: Value,
context: SemanticCacheContext,
) {
let response = response.map_or_else(|| encoded(&entry), <[u8]>::to_vec);
let cache = cache(
vec![
info(INDEX, compatible_info(3)),
search(
INDEX,
"key1",
&VECTOR,
Ok(hit(stored_tag, distance, response)),
),
],
FakeEmbedder::new(&[]),
);
let kind = |result: Result<Option<Value>, Error>| {
result.map_err(|error| std::mem::discriminant(&error))
};
assert_eq!(kind(cache.get_cache("key1", &context)), kind(expected));
}
#[rstest]
#[case::hit(context(), Some(hit("key1", "0.05", encoded(&entry()))), Some(entry()), Some(1.0 - 0.05))]
#[case::beyond_distance_threshold(context(), Some(hit("key1", "0.5", encoded(&entry()))), None, Some(0.0))]
#[case::no_results(context(), Some(empty_result()), None, Some(0.0))]
#[case::other_cache_key(context(), Some(hit("other", "0.05", encoded(&entry()))), None, Some(0.0))]
#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))]
#[tokio::test]
async fn lookup_reports_python_semantic_similarity(
#[case] context: SemanticCacheContext,
#[case] reply: Option<redis::Value>,
#[case] value: Option<Value>,
#[case] similarity: Option<f64>,
#[values(false, true)] use_async: bool,
) {
let commands = reply.map_or_else(Vec::new, |reply| {
vec![
info(INDEX, compatible_info(3)),
search(INDEX, "key1", &VECTOR, Ok(reply)),
]
});
let cache = cache(commands, FakeEmbedder::new(&[]));
let lookup = if use_async {
cache
.async_get_cache_with_similarity("key1", &context)
.await
} else {
cache.get_cache_with_similarity("key1", &context)
};
assert_eq!(lookup, Ok(SemanticLookup { value, similarity }));
}
#[rstest]
#[tokio::test]
async fn missing_prompt_is_a_noop_that_never_embeds(entry: Value) {
let embedder = FakeEmbedder::new(&[]);
let calls = embedder.calls.clone();
let cache = cache(Vec::new(), embedder);
let context = SemanticCacheContext::default();
cache.set_cache("key1", entry.clone(), &context).unwrap();
assert_eq!(cache.get_cache("key1", &context).unwrap(), None);
cache
.async_set_cache("key1", entry, context.clone())
.await
.unwrap();
assert_eq!(cache.async_get_cache("key1", &context).await.unwrap(), None);
assert!(calls.lock().unwrap().is_empty());
}
#[rstest]
fn scope_overrides_key_as_filter_tag(entry: Value, context: SemanticCacheContext) {
let cache = cache(
vec![
info(INDEX, compatible_info(3)),
hset(INDEX, PROMPT, "scope-a", &VECTOR, &entry),
search(
INDEX,
"scope\\-a",
&VECTOR,
Ok(hit("scope-a", "0.05", encoded(&entry))),
),
],
FakeEmbedder::new(&[]),
);
let context = SemanticCacheContext {
scope: Some("scope-a".into()),
..context
};
cache.set_cache("key1", entry.clone(), &context).unwrap();
assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(entry));
}
#[rstest]
#[case::unscoped_schema(unscoped_info(3))]
#[case::wrong_distance_metric(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2")))]
#[case::wrong_data_type(info_with_vector(vector_attribute_with(3, "FLOAT64", "COSINE")))]
fn incompatible_schema_falls_back_to_isolated_index(
#[case] base_info: redis::Value,
entry: Value,
context: SemanticCacheContext,
) {
let isolated = format!("{INDEX}_isolated");
let cache = cache(
vec![
info(INDEX, base_info),
info_missing(&isolated),
create_index(&isolated, 3),
hset(&isolated, PROMPT, "key1", &VECTOR, &entry),
],
FakeEmbedder::new(&[]),
);
cache.set_cache("key1", entry, &context).unwrap();
}
#[rstest]
fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticCacheContext) {
let cache = cache(
vec![
info_missing(INDEX),
MockCmd::new(
create_index_command(INDEX, 3),
Err::<&str, _>(support::server_error("Index already exists")),
),
info(INDEX, compatible_info(3)),
hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
],
FakeEmbedder::new(&[]),
);
cache.set_cache("key1", entry, &context).unwrap();
}
#[rstest]
#[case::punctuation_and_spaces("a:b, c|d", "a\\:b\\,\\ c\\|d")]
#[case::braces_and_dots("{x}.y", "\\{x\\}\\.y")]
#[case::plain("key1", "key1")]
fn tag_special_characters_are_escaped_in_search_filter(
#[case] tag: &str,
#[case] escaped: &str,
context: SemanticCacheContext,
) {
let cache = cache(
vec![
info(INDEX, compatible_info(3)),
search(INDEX, escaped, &VECTOR, Ok(empty_result())),
],
FakeEmbedder::new(&[]),
);
assert_eq!(cache.get_cache(tag, &context).unwrap(), None);
}
#[rstest]
#[case::content_parts(
SemanticCacheContext {
messages: Some(json!([
{"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]},
{"role": "assistant", "content": "reply"},
])),
..Default::default()
},
"firstsecondreply"
)]
#[case::responses_string_input(
SemanticCacheContext { input: Some(json!(" plain input ")), ..Default::default() },
"plain input"
)]
#[case::responses_nested_input(
SemanticCacheContext {
input: Some(json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"])),
..Default::default()
},
"nested\ntail"
)]
#[case::responses_output_text(
SemanticCacheContext { input: Some(json!({"output_text": " result text "})), ..Default::default() },
"result text"
)]
#[case::search_results(
messages_context(vec![json!({
"role": "user",
"content": "question",
"search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}],
})]),
"questionsrctfound{\"a\":1}"
)]
#[case::empty_messages_fall_back_to_input(
SemanticCacheContext { messages: Some(json!([])), input: Some(json!("fallback")), ..Default::default() },
"fallback"
)]
fn prompt_extraction_matches_python_message_and_input_shapes(
#[case] context: SemanticCacheContext,
#[case] prompt: &str,
) {
let embedder = FakeEmbedder::new(&[]);
let calls = embedder.calls.clone();
let cache = cache(
vec![
info(INDEX, compatible_info(3)),
search(INDEX, "key1", &VECTOR, Ok(empty_result())),
],
embedder,
);
cache.get_cache("key1", &context).unwrap();
assert_eq!(*calls.lock().unwrap(), vec![(prompt.to_owned(), None)]);
}
#[rstest]
#[case(None)]
#[case(Some(Duration::from_secs(9)))]
fn ttl_passes_through_context_only(#[case] ttl: Option<Duration>) {
let cache = cache(Vec::new(), FakeEmbedder::new(&[]));
assert_eq!(
cache.get_ttl(&SemanticCacheContext {
ttl,
..Default::default()
}),
ttl
);
}
#[rstest]
#[tokio::test]
async fn async_paths_embed_with_metadata_then_run_blocking_redis_work(
entry: Value,
context: SemanticCacheContext,
) {
let embedder = FakeEmbedder::new(&[]);
let calls = embedder.calls.clone();
let cache = cache(
vec![
info(INDEX, compatible_info(3)),
hset(INDEX, PROMPT, "key1", &VECTOR, &entry),
search(
INDEX,
"key1",
&VECTOR,
Ok(hit("key1", "0.05", encoded(&entry))),
),
],
embedder,
);
let context = SemanticCacheContext {
metadata: Some(json!({"tenant": "team"})),
..context
};
cache
.async_set_cache("key1", entry.clone(), context.clone())
.await
.unwrap();
assert_eq!(
cache.async_get_cache("key1", &context).await.unwrap(),
Some(entry)
);
assert_eq!(
*calls.lock().unwrap(),
vec![(PROMPT.to_owned(), context.metadata.clone()); 2]
);
}
#[rstest]
fn accessors_report_the_config(config: RedisSemanticConfig) {
let cache = cache(Vec::new(), FakeEmbedder::new(&[]));
assert_eq!(cache.index_name(), config.index_name);
assert!((cache.similarity_threshold() - config.similarity_threshold).abs() < 1e-6);
}
#[rstest]
fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value) {
// Pins parity with Python's `_isolated` + overwrite=True flow.
let prompt = "shared prompt";
let isolated = format!("{INDEX}_isolated");
let context = || messages_context(vec![json!({"role": "user", "content": prompt})]);
let vector_a = vec![0.1f32; 8];
let worker_a = cache(
vec![
info_missing(INDEX),
create_index(INDEX, 8),
hset(INDEX, prompt, "key1", &vector_a, &entry),
],
FakeEmbedder::new(&[(prompt, &vector_a)]),
);
worker_a
.set_cache("key1", entry.clone(), &context())
.unwrap();
let vector_b = vec![0.2f32; 4];
let worker_b = cache(
vec![
info(INDEX, compatible_info(8)),
info_missing(&isolated),
create_index(&isolated, 4),
hset(&isolated, prompt, "key1", &vector_b, &entry),
search(
&isolated,
"key1",
&vector_b,
Ok(hit("key1", "0.0", encoded(&entry))),
),
search(
&isolated,
"key1",
&vector_b,
Err(redis::RedisError::from((
redis::ErrorKind::Extension,
"Vector dimension mismatch",
))),
),
],
FakeEmbedder::new(&[(prompt, &vector_b)]),
);
worker_b
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert_eq!(
worker_b.get_cache("key1", &context()).unwrap(),
Some(entry.clone())
);
let vector_c = vec![0.3f32; 16];
let worker_c = cache(
vec![
info(INDEX, compatible_info(8)),
info(&isolated, compatible_info(4)),
MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")),
create_index(&isolated, 16),
hset(&isolated, prompt, "key1", &vector_c, &entry),
],
FakeEmbedder::new(&[(prompt, &vector_c)]),
);
worker_c
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert!(matches!(
worker_b.get_cache("key1", &context()),
Err(Error::Unavailable(Some(_)))
));
}
#[fixture]
fn redis_stack_url() -> Option<String> {
std::env::var("LITELLM_REDIS_STACK_URL").ok()
}
fn live_cache(
url: &str,
index_name: &str,
prompt: &str,
vector: Vec<f32>,
) -> RedisSemanticCache<FakeEmbedder, JsonCodec<Value>> {
RedisSemanticCache::new(
url,
FakeEmbedder::new(&[(prompt, vector.as_slice())]),
JsonCodec::<Value>::new(),
RedisSemanticConfig {
index_name: index_name.to_owned(),
similarity_threshold: 0.9,
},
)
.unwrap()
}
#[rstest]
fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option<String>, entry: Value) {
let Some(url) = redis_stack_url else {
return;
};
// Pins parity with Python's `_isolated` + overwrite=True flow.
let base = format!("rust_semantic_shared_{}", std::process::id());
let isolated = format!("{base}_isolated");
let prompt = "shared live prompt";
let context = || messages_context(vec![json!({"role": "user", "content": prompt})]);
let worker_a = live_cache(&url, &base, prompt, vec![0.1f32; 8]);
worker_a
.set_cache("key1", entry.clone(), &context())
.unwrap();
let worker_b = live_cache(&url, &base, prompt, vec![0.2f32; 4]);
worker_b
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert_eq!(
worker_b.get_cache("key1", &context()).unwrap(),
Some(entry.clone())
);
let worker_c = live_cache(&url, &base, prompt, vec![0.3f32; 16]);
worker_c
.set_cache("key1", entry.clone(), &context())
.unwrap();
assert!(matches!(
worker_b.get_cache("key1", &context()),
Err(Error::Unavailable(Some(_)))
));
let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
for index in [&base, &isolated] {
let _: Result<(), _> = redis::cmd("FT.DROPINDEX")
.arg(index)
.arg("DD")
.query(&mut connection);
}
}
#[rstest]
fn live_store_lookup_and_ttl_against_redis_stack(redis_stack_url: Option<String>, entry: Value) {
let Some(url) = redis_stack_url else {
return;
};
let prompt = "rust semantic cache live prompt";
let index_name = format!("rust_semantic_test_{}", std::process::id());
let cache = live_cache(&url, &index_name, prompt, vec![0.1, 0.2, 0.3, 0.4]);
let context = SemanticCacheContext {
ttl: Some(Duration::from_secs(120)),
..messages_context(vec![json!({"role": "user", "content": prompt})])
};
cache
.set_cache("live-key", entry.clone(), &context)
.unwrap();
assert_eq!(cache.get_cache("live-key", &context).unwrap(), Some(entry));
assert_eq!(cache.get_cache("other-key", &context).unwrap(), None);
let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap();
let ttl: i64 = redis::Commands::ttl(
&mut connection,
format!("{index_name}:{}", entry_id(prompt, "live-key")),
)
.unwrap();
assert!(
ttl > 0,
"expected stored hash to carry an expiry, got {ttl}"
);
}

63
litellm-rust/crates/cache-redis-semantic/tests/contract.rs
View file

@ -1,63 +0,0 @@
mod support;
use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig};
use litellm_cache_testing as contract;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::FakeSearch;
type Cache = RedisSemanticCache<PreparedEmbedding, JsonCodec<Value>, FakeSearch>;
const PREFIX: &str = "contract:";
#[fixture]
fn cache() -> Cache {
RedisSemanticCache::with_connection(
FakeSearch::default(),
PreparedEmbedding(vec![0.6, 0.8]),
JsonCodec::new(),
RedisSemanticConfig {
index_name: DEFAULT_INDEX_NAME.into(),
similarity_threshold: 0.9,
},
)
}
#[fixture]
fn context() -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(json!([{"role": "user", "content": "contract prompt"}])),
..Default::default()
}
}
#[rstest]
#[tokio::test]
async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) {
contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await;
}
#[rstest]
#[tokio::test]
async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) {
contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await;
}
#[rstest]
#[tokio::test]
async fn overwrite_replaces(cache: Cache, context: SemanticCacheContext) {
contract::overwrite_replaces(&cache, context, PREFIX, json!(1), json!({"b": 2})).await;
}
#[rstest]
#[tokio::test]
async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) {
contract::pipeline_writes_every_entry(
&cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
}

302
litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs
View file

@ -1,302 +0,0 @@
#![allow(dead_code)]
use std::{
collections::{BTreeMap, HashMap},
sync::{Arc, Mutex},
};
use litellm_cache::{Error, semantic::Embedder};
use serde_json::Value;
pub type EmbedCalls = Arc<Mutex<Vec<(String, Option<Value>)>>>;
/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every
/// prompt with its metadata.
pub struct FakeEmbedder {
vectors: HashMap<String, Vec<f32>>,
pub calls: EmbedCalls,
}
impl FakeEmbedder {
pub fn new(vectors: &[(&str, &[f32])]) -> Self {
Self {
vectors: vectors
.iter()
.map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec()))
.collect(),
calls: EmbedCalls::default(),
}
}
}
impl Embedder for FakeEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
self.calls
.lock()
.unwrap()
.push((prompt.to_owned(), metadata.cloned()));
Ok(self
.vectors
.get(prompt)
.cloned()
.unwrap_or_else(|| vec![0.1, 0.2, 0.3]))
}
async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
self.embed(prompt, metadata)
}
}
struct FakeIndex {
prefix: Vec<u8>,
dims: usize,
vector_field: String,
}
#[derive(Default)]
struct SearchState {
indexes: HashMap<String, FakeIndex>,
hashes: BTreeMap<Vec<u8>, BTreeMap<String, Vec<u8>>>,
}
/// An in-memory Redis Stack speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache
/// sends, with exact cosine KNN over the hashes under an index prefix.
#[derive(Clone, Default)]
pub struct FakeSearch {
state: Arc<Mutex<SearchState>>,
}
impl FakeSearch {
fn run(&self, args: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
let mut state = self.state.lock().unwrap();
let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned();
match text(0).to_uppercase().as_str() {
"FT.CREATE" => {
let name = text(1);
if state.indexes.contains_key(&name) {
return Err(server_error("Index already exists"));
}
let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes());
let prefix = args[position("PREFIX").unwrap() + 2].clone();
let dims = text(position("DIM").unwrap() + 1).parse().unwrap();
let vector_field = text(position("VECTOR").unwrap() - 1);
state.indexes.insert(
name,
FakeIndex {
prefix,
dims,
vector_field,
},
);
Ok(redis::Value::Okay)
}
"FT.INFO" => {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| server_error("Unknown index name"))?;
Ok(index_info(index))
}
"FT.DROPINDEX" => {
state.indexes.remove(&text(1));
Ok(redis::Value::Okay)
}
"HSET" => {
let hash = state.hashes.entry(args[1].clone()).or_default();
for pair in args[2..].chunks(2) {
hash.insert(
String::from_utf8_lossy(&pair[0]).into_owned(),
pair[1].clone(),
);
}
Ok(redis::Value::Int(((args.len() - 2) / 2) as i64))
}
"EXPIRE" => Ok(redis::Value::Int(i64::from(
state.hashes.contains_key(&args[1]),
))),
"FT.SEARCH" => {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| server_error("no such index"))?;
let query = text(2);
let tag = query_tag(&query);
let params = args.iter().position(|arg| arg == b"PARAMS").unwrap();
let vector = floats(&args[params + 3]);
let best = state
.hashes
.iter()
.filter(|(key, _)| key.starts_with(&index.prefix))
.filter(|(_, fields)| {
fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes())
})
.filter_map(|(key, fields)| {
let stored = floats(fields.get(&index.vector_field)?);
(stored.len() == index.dims)
.then(|| (key, fields, 1.0 - cosine(&vector, &stored)))
})
.min_by(|left, right| left.2.total_cmp(&right.2));
let Some((key, fields, distance)) = best else {
return Ok(redis::Value::Array(vec![redis::Value::Int(0)]));
};
let mut reply = fields
.iter()
.filter(|(name, _)| **name != index.vector_field)
.flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)])
.collect::<Vec<_>>();
reply.extend([
bulk(b"vector_distance"),
bulk(distance.to_string().as_bytes()),
]);
Ok(redis::Value::Array(vec![
redis::Value::Int(1),
bulk(key),
redis::Value::Array(reply),
]))
}
"PING" => Ok(redis::Value::SimpleString("PONG".into())),
_ => Err(server_error("unsupported command")),
}
}
}
impl redis::ConnectionLike for FakeSearch {
fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
let mut commands = parse_commands(command);
self.run(commands.remove(0))
}
fn req_packed_commands(
&mut self,
commands: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
let replies = parse_commands(commands)
.into_iter()
.map(|args| self.run(args))
.collect::<redis::RedisResult<Vec<_>>>()?;
Ok(replies.into_iter().skip(offset).take(count).collect())
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}
pub fn server_error(reply: &str) -> redis::RedisError {
redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
.expect("error reply")
.extract_error()
.expect_err("server error")
}
fn bulk(bytes: &[u8]) -> redis::Value {
redis::Value::BulkString(bytes.to_vec())
}
fn index_info(index: &FakeIndex) -> redis::Value {
let attribute = |name: &str, field_type: &str| {
redis::Value::Array(vec![
bulk(b"identifier"),
bulk(name.as_bytes()),
bulk(b"type"),
bulk(field_type.as_bytes()),
])
};
redis::Value::Array(vec![
bulk(b"attributes"),
redis::Value::Array(vec![
attribute("prompt", "TEXT"),
attribute("response", "TEXT"),
attribute("inserted_at", "NUMERIC"),
attribute("updated_at", "NUMERIC"),
attribute("litellm_cache_key", "TAG"),
redis::Value::Array(vec![
bulk(b"identifier"),
bulk(index.vector_field.as_bytes()),
bulk(b"type"),
bulk(b"VECTOR"),
bulk(b"dim"),
redis::Value::Int(index.dims as i64),
bulk(b"data_type"),
bulk(b"FLOAT32"),
bulk(b"distance_metric"),
bulk(b"COSINE"),
]),
]),
])
}
/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed.
fn query_tag(query: &str) -> String {
let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len();
let mut tag = String::new();
let mut characters = query[start..].chars();
while let Some(character) = characters.next() {
match character {
'\\' => tag.extend(characters.next()),
'}' => break,
character => tag.push(character),
}
}
tag
}
fn floats(bytes: &[u8]) -> Vec<f32> {
bytes
.as_chunks::<4>()
.0
.iter()
.map(|chunk| f32::from_le_bytes(*chunk))
.collect()
}
fn cosine(left: &[f32], right: &[f32]) -> f64 {
let dot = left
.iter()
.zip(right)
.map(|(left, right)| f64::from(*left) * f64::from(*right))
.sum::<f64>();
let norm = |vector: &[f32]| {
vector
.iter()
.map(|value| f64::from(*value).powi(2))
.sum::<f64>()
.sqrt()
};
dot / (norm(left) * norm(right))
}
/// Splits a packed RESP request into each command's arguments.
fn parse_commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
let line = |bytes: &mut &[u8]| {
let end = bytes
.windows(2)
.position(|window| window == b"\r\n")
.unwrap();
let text = String::from_utf8(bytes[1..end].to_vec()).unwrap();
*bytes = &bytes[end + 2..];
text.parse::<usize>().unwrap()
};
let mut commands = Vec::new();
while !bytes.is_empty() {
let count = line(&mut bytes);
let mut args = Vec::with_capacity(count);
for _ in 0..count {
let length = line(&mut bytes);
args.push(bytes[..length].to_vec());
bytes = &bytes[length + 2..];
}
commands.push(args);
}
commands
}

20
litellm-rust/crates/cache-redis/Cargo.toml
View file

@ -1,20 +0,0 @@
[package]
name = "litellm-cache-redis"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-error.workspace = true
redis = { workspace = true, features = ["cluster", "tls-rustls"] }
r2d2 = "0.8.10"
thiserror.workspace = true
tokio.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
redis-test.workspace = true
rstest.workspace = true
serde_json.workspace = true

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

@ -1,192 +0,0 @@
use std::{
sync::{Arc, OnceLock},
time::Duration,
};
use litellm_cache::{BatchEntry, CacheCodec, Error};
use crate::{
connection::{ConnectionRef, Connections},
topology::RedisTopology,
};
const DEFAULT_TTL: Duration = Duration::from_secs(600);
pub struct RedisCache<S, C = redis::Connection> {
pub(crate) connections: Arc<Connections<C>>,
pub(crate) default_ttl: Duration,
pub(crate) codec: S,
pub(crate) namespace: Option<String>,
pub(crate) topology: RedisTopology,
/// The server's major version, read from `INFO` once, like Python's `redis_version`.
pub(crate) major_version: Arc<OnceLock<u32>>,
}
impl<S: CacheCodec> RedisCache<S> {
pub fn new(url: &str, default_ttl: Option<Duration>, codec: S) -> Result<Self, Error> {
Self::connect(url, &RedisTopology::Standalone, default_ttl, codec)
}
pub fn connect(
url: &str,
topology: &RedisTopology,
default_ttl: Option<Duration>,
codec: S,
) -> Result<Self, Error> {
let connections = Connections::open(url, topology)?;
Ok(Self {
connections: Arc::new(connections),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
topology: topology.clone(),
major_version: Arc::default(),
})
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub fn with_connection(connection: C, default_ttl: Option<Duration>, codec: S) -> Self {
Self {
connections: Arc::new(Connections::fixed(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
topology: RedisTopology::Standalone,
major_version: Arc::default(),
}
}
pub fn with_namespace(self, namespace: Option<String>) -> Self {
Self {
namespace: namespace.filter(|value| !value.is_empty()),
..self
}
}
pub fn namespace(&self) -> Option<&str> {
self.namespace.as_deref()
}
pub fn topology(&self) -> &RedisTopology {
&self.topology
}
pub(crate) fn namespaced_key(&self, key: &str) -> String {
namespaced_key(self.namespace.as_deref(), key)
}
pub(crate) fn namespaced_keys(&self, keys: &[String]) -> Vec<String> {
keys.iter().map(|key| self.namespaced_key(key)).collect()
}
/// Whole seconds for `ttl`, falling back to the default TTL like Python's `get_ttl`.
pub(crate) fn ttl_or_default(&self, ttl: Option<Duration>) -> u64 {
ttl_seconds(ttl.unwrap_or(self.default_ttl))
}
pub(crate) fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
self.connections.execute(operation)
}
/// `_parse_redis_major_version`: the major version from `INFO`, or
/// `DEFAULT_REDIS_MAJOR_VERSION` when `INFO` fails or its version does not parse. The first
/// answer is kept, as Python reads `redis_version` once at construction.
pub(crate) async fn major_version(&self) -> u32 {
if let Some(version) = self.major_version.get() {
return *version;
}
let info = self
.run(|connection| connection.node_text(&redis::cmd("INFO")))
.await;
let version = info
.ok()
.and_then(|info| parse_major_version(&info))
.unwrap_or_else(default_major_version);
*self.major_version.get_or_init(|| version)
}
pub(crate) async fn run<T, F>(&self, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
Connections::run_blocking(Arc::clone(&self.connections), operation).await
}
pub(crate) fn namespaced_pattern(&self) -> Result<String, Error> {
let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?;
let escaped: String = namespace
.chars()
.flat_map(|ch| {
if matches!(ch, '*' | '?' | '[' | ']' | '\\') {
vec!['\\', ch]
} else {
vec![ch]
}
})
.collect();
Ok(format!("{escaped}:*"))
}
pub(crate) fn decode_response(&self, value: redis::Value) -> Result<Option<S::Value>, Error> {
match value {
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
redis::Value::SimpleString(text) => self.codec.decode(text.as_bytes()).map(Some),
_ => Err(Error::InvalidEntry(None)),
}
}
pub(crate) fn decode_batch_response(
&self,
value: redis::Value,
) -> Result<BatchEntry<S::Value>, Error> {
match self.decode_response(value) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
Err(Error::InvalidEntry(_)) => Ok(BatchEntry::Invalid),
Err(error) => Err(error),
}
}
}
pub(crate) fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
match namespace {
Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
format!("{namespace}:{key}")
}
_ => key.into(),
}
}
pub(crate) fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
fn parse_major_version(info: &str) -> Option<u32> {
let version = info
.lines()
.find_map(|line| line.trim().strip_prefix("redis_version:"))?
.trim();
match version.split_once('.') {
Some((major, _)) => major.parse().ok(),
None => version.parse::<f64>().ok().map(|major| major as u32),
}
}
fn default_major_version() -> u32 {
std::env::var("DEFAULT_REDIS_MAJOR_VERSION")
.ok()
.and_then(|value| value.parse().ok())
.unwrap_or(7)
}

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

@ -1,109 +0,0 @@
use litellm_cache::{CacheCodec, ClaimCache, Error, ExactCacheContext};
use redis::Commands;
use crate::{cache::RedisCache, connection::ConnectionRef};
const CLAIM_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
"elseif current ~= ARGV[1] then return 0; end; ",
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
);
const CLAIM_ATTEMPTS: usize = 8;
fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
match value {
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => Ok(Some(bytes)),
redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
_ => Err(Error::InvalidEntry(None)),
}
}
/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's
/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the
/// bytes that decision was made on, retried when another claimant wins the race.
fn claim<S: CacheCodec>(
connection: &mut ConnectionRef<'_>,
codec: &S,
key: &str,
candidate: S::Value,
eligible: &[S::Value],
ttl: u64,
) -> Result<S::Value, Error>
where
S::Value: PartialEq,
{
let payload = codec.encode(&candidate)?;
if payload.is_empty() {
return Err(Error::InvalidEntry(None));
}
for _ in 0..CLAIM_ATTEMPTS {
let current = stored_bytes(
connection
.get::<_, redis::Value>(key)
.map_err(Error::unavailable)?,
)?
.filter(|bytes| !bytes.is_empty());
let existing = current
.as_deref()
.and_then(|bytes| codec.decode(bytes).ok())
.filter(|existing| eligible.is_empty() || eligible.contains(existing));
let refresh = existing
.as_ref()
.is_some_and(|existing| !eligible.is_empty() || *existing == candidate);
let write: &[u8] = if existing.is_some() { b"" } else { &payload };
let applied = redis::cmd("EVAL")
.arg(CLAIM_SCRIPT)
.arg(1)
.arg(key)
.arg(current.as_deref().unwrap_or_default())
.arg(ttl)
.arg(write)
.arg(u8::from(refresh))
.query::<bool>(connection)
.map_err(Error::unavailable)?;
if applied {
return Ok(existing.unwrap_or(candidate));
}
}
Err(Error::unavailable(ClaimContended))
}
#[derive(Debug, thiserror::Error)]
#[error("claim kept losing the race after {CLAIM_ATTEMPTS} attempts")]
struct ClaimContended;
impl<S, C> ClaimCache for RedisCache<S, C>
where
S: CacheCodec + Clone + 'static,
S::Value: PartialEq,
C: redis::ConnectionLike + Send + 'static,
{
fn claim_cache(
&self,
key: &str,
candidate: S::Value,
eligible: &[S::Value],
context: ExactCacheContext,
) -> Result<S::Value, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(context.ttl);
self.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl))
}
async fn async_claim_cache(
&self,
key: &str,
candidate: S::Value,
eligible: Vec<S::Value>,
context: ExactCacheContext,
) -> Result<S::Value, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(context.ttl);
let codec = self.codec.clone();
self.run(move |connection| claim(connection, &codec, &key, candidate, &eligible, ttl))
.await
}
}

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

@ -1,513 +0,0 @@
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::Error;
use redis::{
ConnectionAddr, ConnectionInfo, IntoConnectionInfo,
cluster::{
ClusterClient, ClusterClientBuilder, ClusterConnection, ClusterPipeline, NodeAddress,
},
cluster_routing::{
MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo,
},
};
use crate::topology::{RedisNode, RedisTopology};
pub(crate) const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
const REDIS_POOL_SIZE: u32 = 16;
#[allow(private_interfaces)]
pub enum Connections<C> {
Pool(r2d2::Pool<ConnectionManager>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
}
impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
pub fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
match self {
Self::Pool(pool) => {
let mut pooled = pool.get().map_err(Error::unavailable)?;
let result = operation(&mut ConnectionRef::Node(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
result
}
Self::Cluster(pool) => {
let mut pooled = pool.get().map_err(Error::unavailable)?;
let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable(_)));
result
}
Self::Fixed(connection) => {
let mut connection = connection.lock()?;
operation(&mut ConnectionRef::Node(&mut *connection))
}
}
}
pub async fn run_blocking<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)?,
)?)),
}
}
/// Closes every idle pooled connection; the next operation opens a fresh one. Connections
/// checked out right now return to the pool, and a caller-owned connection stays open.
pub fn disconnect(&self) {
match self {
Self::Pool(pool) => close_idle(pool),
Self::Cluster(pool) => close_idle(pool),
Self::Fixed(_) => {}
}
}
}
fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
r2d2::Pool::builder()
.max_size(REDIS_POOL_SIZE)
.min_idle(Some(0))
.connection_timeout(REDIS_TIMEOUT)
.test_on_check_out(false)
.build(manager)
.map_err(Error::unavailable)
}
fn close_idle<M, T>(pool: &r2d2::Pool<M>)
where
M: r2d2::ManageConnection<Connection = PooledConnection<T>>,
{
let mut idle = Vec::new();
while let Some(mut connection) = pool.try_get() {
connection.failed = true;
idle.push(connection);
}
}
pub(crate) struct PooledConnection<C> {
connection: C,
failed: bool,
}
/// Pools connections without a checkout PING, which would double every operation's round trips.
/// A timed-out command leaves its reply on the socket while redis still reports the connection
/// open, so any connection whose operation failed is discarded instead of being reused.
pub(crate) struct ConnectionManager(redis::Client);
impl ConnectionManager {
fn open(url: &str) -> Result<Self, Error> {
redis::Client::open(url)
.map(Self)
.map_err(Error::unavailable)
}
}
impl r2d2::ManageConnection for ConnectionManager {
type Connection = PooledConnection<redis::Connection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
let connection = self.0.get_connection()?;
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
Ok(PooledConnection {
connection,
failed: false,
})
}
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
redis::cmd("PING").query::<String>(&mut connection.connection)?;
Ok(())
}
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
}
}
#[derive(Debug, thiserror::Error)]
enum TopologyFault {
#[error("Redis Cluster needs at least one startup node")]
NoStartupNodes,
#[error("Redis Cluster nodes need a TCP address")]
NonTcpAddress,
#[error("Redis Cluster answered a fan-out command without a per-node map")]
FanOutReply,
#[error("scan target does not match the connection topology")]
ScanTarget,
}
pub(crate) struct ClusterConnectionManager(ClusterClient);
impl ClusterConnectionManager {
fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
if startup_nodes.is_empty() {
return Err(Error::unavailable(TopologyFault::NoStartupNodes));
}
let info = url.into_connection_info().map_err(Error::unavailable)?;
let nodes = startup_nodes
.iter()
.map(|node| node_info(&info, node))
.collect::<Result<Vec<_>, _>>()?;
ClusterClientBuilder::new(nodes)
.connection_timeout(REDIS_TIMEOUT)
.response_timeout(REDIS_TIMEOUT)
.build()
.map(Self)
.map_err(Error::unavailable)
}
}
fn node_info(info: &ConnectionInfo, node: &RedisNode) -> Result<ConnectionInfo, Error> {
let addr = match info.addr() {
ConnectionAddr::Tcp(..) => ConnectionAddr::Tcp(node.host.clone(), node.port),
ConnectionAddr::TcpTls {
insecure,
tls_params,
..
} => ConnectionAddr::TcpTls {
host: node.host.clone(),
port: node.port,
insecure: *insecure,
tls_params: tls_params.clone(),
},
_ => return Err(Error::unavailable(TopologyFault::NonTcpAddress)),
};
Ok(info.clone().set_addr(addr))
}
impl r2d2::ManageConnection for ClusterConnectionManager {
type Connection = PooledConnection<ClusterConnection>;
type Error = redis::RedisError;
fn connect(&self) -> Result<Self::Connection, redis::RedisError> {
let connection = self.0.get_connection()?;
connection.set_read_timeout(Some(REDIS_TIMEOUT))?;
connection.set_write_timeout(Some(REDIS_TIMEOUT))?;
Ok(PooledConnection {
connection,
failed: false,
})
}
fn is_valid(&self, connection: &mut Self::Connection) -> Result<(), redis::RedisError> {
redis::cmd("PING").query::<String>(&mut connection.connection)?;
Ok(())
}
fn has_broken(&self, connection: &mut Self::Connection) -> bool {
connection.failed || !redis::ConnectionLike::is_open(&connection.connection)
}
}
pub enum ConnectionRef<'a> {
Node(&'a mut dyn redis::ConnectionLike),
Cluster(&'a mut ClusterConnection),
}
impl redis::ConnectionLike for ConnectionRef<'_> {
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
match self {
Self::Node(connection) => connection.req_packed_command(cmd),
Self::Cluster(connection) => connection.req_packed_command(cmd),
}
}
fn req_packed_commands(
&mut self,
cmd: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
match self {
Self::Node(connection) => connection.req_packed_commands(cmd, offset, count),
Self::Cluster(connection) => connection.req_packed_commands(cmd, offset, count),
}
}
fn get_db(&self) -> i64 {
match self {
Self::Node(connection) => connection.get_db(),
Self::Cluster(connection) => redis::ConnectionLike::get_db(*connection),
}
}
fn supports_pipelining(&self) -> bool {
match self {
Self::Node(connection) => connection.supports_pipelining(),
Self::Cluster(connection) => redis::ConnectionLike::supports_pipelining(*connection),
}
}
fn check_connection(&mut self) -> bool {
match self {
Self::Node(connection) => connection.check_connection(),
Self::Cluster(connection) => connection.check_connection(),
}
}
fn is_open(&self) -> bool {
match self {
Self::Node(connection) => connection.is_open(),
Self::Cluster(connection) => redis::ConnectionLike::is_open(*connection),
}
}
}
impl ConnectionRef<'_> {
/// Runs `pipeline` and decodes its non-ignored replies as `T`. A cluster connection refuses
/// `Pipeline::query`, so there a transaction goes to its keys' slot as one MULTI/EXEC and
/// anything else is split per node by `ClusterPipeline`; either way the raw replies are
/// handed back to `pipeline` to decode.
pub(crate) fn query_pipeline<T: redis::FromRedisValue>(
&mut self,
pipeline: &redis::Pipeline,
) -> Result<T, Error> {
match self {
Self::Node(connection) => pipeline.query(*connection),
Self::Cluster(connection) if pipeline.is_transaction() => {
redis::ConnectionLike::req_packed_commands(
*connection,
&pipeline.get_packed_pipeline(),
pipeline.len() + 1,
1,
)
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
}
Self::Cluster(connection) => {
let mut cluster = ClusterPipeline::with_capacity(pipeline.len());
for command in pipeline.cmd_iter() {
cluster.add_command(command.clone());
}
cluster
.query(connection)
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
}
}
.map_err(Error::unavailable)
}
pub(crate) fn scan(
&mut self,
pattern: &str,
count: usize,
mut visit: impl FnMut(&mut Self, Vec<String>) -> Result<bool, Error>,
) -> Result<(), Error> {
let pages = match self {
Self::Node(connection) => {
let page = scan_command(0, pattern, count)
.query::<ScanPage>(*connection)
.map_err(Error::unavailable)?;
vec![(None, page)]
}
Self::Cluster(connection) => connection
.route_command(
&scan_command(0, pattern, count),
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllMasters,
Some(ResponsePolicy::Special),
)),
)
.map_err(Error::unavailable)
.and_then(primary_pages)?
.into_iter()
.map(|(node, page)| (Some(node), page))
.collect(),
};
for (node, (mut cursor, mut keys)) in pages {
loop {
if !visit(self, keys)? {
return Ok(());
}
if cursor == 0 {
break;
}
(cursor, keys) = self.scan_page(node.as_ref(), cursor, pattern, count)?;
}
}
Ok(())
}
pub(crate) fn ping(&mut self) -> Result<bool, redis::RedisError> {
let command = redis::cmd("PING");
match self {
Self::Node(connection) => command
.query::<String>(*connection)
.map(|response| response == "PONG"),
Self::Cluster(connection) => connection
.route_command(
&command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllNodes,
Some(ResponsePolicy::AllSucceeded),
)),
)
.map(|_| true),
}
}
pub(crate) fn node_text(&mut self, command: &redis::Cmd) -> Result<String, Error> {
match self {
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
Self::Cluster(connection) => {
let value = connection
.route_command(
command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllNodes,
Some(ResponsePolicy::Special),
)),
)
.map_err(Error::unavailable)?;
let redis::Value::Map(entries) = value else {
return Err(Error::unavailable(TopologyFault::FanOutReply));
};
let mut replies = entries
.into_iter()
.map(|(node, reply)| {
Ok((
redis::from_redis_value::<String>(node).map_err(Error::unavailable)?,
redis::from_redis_value::<String>(reply).map_err(Error::unavailable)?,
))
})
.collect::<Result<Vec<(String, String)>, Error>>()?;
replies.sort();
Ok(replies
.into_iter()
.map(|(_, reply)| reply)
.collect::<Vec<_>>()
.join("\n"))
}
}
}
pub(crate) fn flushall(&mut self) -> Result<(), Error> {
let command = redis::cmd("FLUSHALL");
match self {
Self::Node(connection) => command.query(*connection).map_err(Error::unavailable),
Self::Cluster(connection) => connection
.route_command(
&command,
RoutingInfo::MultiNode((
MultipleNodeRoutingInfo::AllMasters,
Some(ResponsePolicy::AllSucceeded),
)),
)
.map(|_| ())
.map_err(Error::unavailable),
}
}
fn scan_page(
&mut self,
node: Option<&NodeAddress>,
cursor: u64,
pattern: &str,
count: usize,
) -> Result<ScanPage, Error> {
let command = scan_command(cursor, pattern, count);
match (self, node) {
(Self::Node(connection), None) => {
command.query(*connection).map_err(Error::unavailable)
}
(Self::Cluster(connection), Some(node)) => connection
.route_command(
&command,
RoutingInfo::SingleNode(SingleNodeRoutingInfo::ByAddress {
host: node.host().to_string(),
port: node.port(),
}),
)
.map_err(Error::unavailable)
.and_then(|value| redis::from_redis_value(value).map_err(Error::unavailable)),
_ => Err(Error::unavailable(TopologyFault::ScanTarget)),
}
}
}
type ScanPage = (u64, Vec<String>);
fn primary_pages(value: redis::Value) -> Result<Vec<(NodeAddress, ScanPage)>, Error> {
let redis::Value::Map(entries) = value else {
return Err(Error::unavailable(TopologyFault::FanOutReply));
};
entries
.into_iter()
.map(|(node, page)| {
let node = redis::from_redis_value::<String>(node).map_err(Error::unavailable)?;
let node = NodeAddress::try_from(node.as_str()).map_err(Error::unavailable)?;
let page = redis::from_redis_value::<ScanPage>(page).map_err(Error::unavailable)?;
Ok((node, page))
})
.collect()
}
fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd {
let mut command = redis::cmd("SCAN");
command
.cursor_arg(cursor)
.arg("MATCH")
.arg(pattern)
.arg("COUNT")
.arg(count);
command
}
/// Hands already received pipeline replies to `Pipeline::query`, so it applies its own
/// ignore and error handling to replies a cluster pipeline gathered from several nodes.
struct Replies(Option<Vec<redis::Value>>);
impl redis::ConnectionLike for Replies {
fn req_packed_command(&mut self, _: &[u8]) -> redis::RedisResult<redis::Value> {
Err((redis::ErrorKind::Client, "replies hold a pipeline only").into())
}
fn req_packed_commands(
&mut self,
_: &[u8],
_: usize,
_: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
self.0
.take()
.ok_or_else(|| (redis::ErrorKind::Client, "replies were already read").into())
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}

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

@ -1,205 +0,0 @@
use std::time::Duration;
use litellm_cache::{
BoundedCounterCache, CacheCodec, CountReadCache, CounterCache, Error, ExactCacheContext,
IncrementOperation,
};
use crate::{
cache::{RedisCache, ttl_seconds},
connection::ConnectionRef,
store::mget,
};
const INCREMENT_SCRIPT: &str = concat!(
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
"if redis.call('TTL', KEYS[1]) == -1 then ",
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
);
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
"return count"
);
const SET_MAX_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
"redis.call('SET', KEYS[1], ARGV[1]); ",
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
"return ARGV[1]; end; return current"
);
impl<S, C> CounterCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(context.ttl);
self.execute(|connection| increment(connection, key, amount, ttl, false))
}
/// Python `_incrbyfloat_with_ttl`: without `refresh_ttl` the TTL is set only on a key that
/// has none, in one atomic script; with it, every increment re-arms the TTL.
async fn async_increment(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
refresh_ttl: bool,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(context.ttl);
self.run(move |connection| increment(connection, key, amount, ttl, refresh_ttl))
.await
}
async fn async_increment_pipeline(
&self,
operations: Vec<IncrementOperation>,
) -> Result<Vec<f64>, Error> {
if operations.is_empty() {
return Ok(Vec::new());
}
let mut pipeline = redis::pipe();
for operation in operations {
let key = self.namespaced_key(&operation.key);
pipeline.cmd("INCRBYFLOAT").arg(&key).arg(operation.amount);
if let Some(ttl) = operation.ttl {
pipeline
.cmd("EXPIRE")
.arg(key)
.arg(ttl_seconds(ttl))
.ignore();
}
}
self.run(move |connection| connection.query_pipeline(&pipeline))
.await
}
}
fn increment(
connection: &mut ConnectionRef<'_>,
key: String,
amount: f64,
ttl: u64,
refresh_ttl: bool,
) -> Result<f64, Error> {
if !refresh_ttl {
return redis::cmd("EVAL")
.arg(INCREMENT_SCRIPT)
.arg(1)
.arg(key)
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(Error::unavailable);
}
connection
.query_pipeline(
redis::pipe()
.cmd("INCRBYFLOAT")
.arg(&key)
.arg(amount)
.cmd("EXPIRE")
.arg(&key)
.arg(ttl)
.ignore(),
)
.map(|(value,)| value)
}
impl<S, C> CountReadCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, Error> {
let keys = self.namespaced_keys(keys);
self.execute(|connection| mget(connection, keys))?
.into_iter()
.map(count)
.collect()
}
async fn async_batch_get_counts(&self, keys: Vec<String>) -> Result<Vec<Option<i64>>, Error> {
let keys = self.namespaced_keys(&keys);
self.run(move |connection| mget(connection, keys))
.await?
.into_iter()
.map(count)
.collect()
}
}
fn count(value: redis::Value) -> Result<Option<i64>, Error> {
redis::from_redis_value(value).map_err(Error::invalid_entry)
}
impl<S, C> BoundedCounterCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn increment_with_floor(&self, key: &str, amount: i64, ttl: Duration) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = ttl_seconds(ttl);
self.execute(|connection| increment_with_floor(connection, key, amount, ttl))
}
async fn async_increment_with_floor(
&self,
key: &str,
amount: i64,
ttl: Duration,
) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = ttl_seconds(ttl);
self.run(move |connection| increment_with_floor(connection, key, amount, ttl))
.await
}
async fn async_set_max(
&self,
key: &str,
value: f64,
ttl: Option<Duration>,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(ttl);
self.run(move |connection| {
redis::cmd("EVAL")
.arg(SET_MAX_SCRIPT)
.arg(1)
.arg(key)
.arg(value)
.arg(ttl)
.query(connection)
.map_err(Error::unavailable)
})
.await
}
}
fn increment_with_floor(
connection: &mut ConnectionRef<'_>,
key: String,
amount: i64,
ttl: u64,
) -> Result<i64, Error> {
redis::cmd("EVAL")
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
.arg(1)
.arg(key)
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(Error::unavailable)
}

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

@ -1,63 +0,0 @@
use std::time::Duration;
use litellm_cache::{CacheCodec, Error, RefreshTtlCache, ScanCache, TtlCache};
use crate::cache::RedisCache;
impl<S, C> TtlCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
let key = self.namespaced_key(key);
let ttl = self
.run(move |connection| {
redis::cmd("TTL")
.arg(key)
.query::<i64>(connection)
.map_err(Error::unavailable)
})
.await?;
Ok(u64::try_from(ttl).ok().map(Duration::from_secs))
}
}
impl<S, C> RefreshTtlCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_refresh_ttl(&self, key: &str, ttl: Option<Duration>) -> Result<bool, Error> {
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(ttl);
self.run(move |connection| {
redis::cmd("EXPIRE")
.arg(key)
.arg(ttl)
.query(connection)
.map_err(Error::unavailable)
})
.await
}
}
impl<S, C> ScanCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
let pattern = format!("{}*", self.namespaced_key(pattern));
self.run(move |connection| {
let mut matches = Vec::new();
connection.scan(&pattern, count, |_, keys| {
matches.extend(keys);
Ok(matches.len() < count)
})?;
matches.truncate(count);
Ok(matches)
})
.await
}
}

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

@ -1,15 +0,0 @@
mod cache;
mod claim;
pub mod connection;
mod counter;
mod keys;
mod lifecycle;
mod queue;
mod script;
mod store;
mod topology;
pub use cache::RedisCache;
pub use queue::{RedisLpopOperation, RedisLpopResult, RedisRpushOperation};
pub use script::{RedisArg, RedisScript};
pub use topology::{RedisNode, RedisTopology};

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

@ -1,86 +0,0 @@
use litellm_cache::{
CacheCodec, CacheConnectionResult, CacheConnectionStatus, ClientInfoCache, ConnectionCache,
DisconnectCache, Error, PingCache,
};
use crate::{cache::RedisCache, topology::RedisTopology};
impl<S, C> PingCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn sync_ping(&self) -> Result<bool, Error> {
self.execute(|connection| connection.ping().map_err(Error::unavailable))
}
async fn ping(&self) -> Result<bool, Error> {
self.run(|connection| connection.ping().map_err(Error::unavailable))
.await
}
}
impl<S, C> ConnectionCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
/// Python `RedisCache.test_connection`, or `RedisClusterCache.test_connection` for a
/// cluster topology, which differs only in its messages.
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
let label = match self.topology {
RedisTopology::Standalone => "Redis",
RedisTopology::Cluster { .. } => "Redis Cluster",
};
let ping = self
.run(|connection| Ok(connection.ping().map_err(Error::unavailable)))
.await
.and_then(|ping| ping)
.map_err(|error| litellm_error::Report::new(&error).to_string());
Ok(match ping {
Ok(true) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: format!("{label} connection test successful"),
error: None,
},
Ok(false) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("{label} ping returned False"),
error: None,
},
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("{label} connection failed: {error}"),
error: Some(error),
},
})
}
}
impl<S, C> DisconnectCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn disconnect(&self) -> Result<(), Error> {
self.connections.disconnect();
Ok(())
}
}
impl<S, C> ClientInfoCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type ClientList = String;
type Info = String;
fn client_list(&self) -> Result<Self::ClientList, Error> {
self.execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST")))
}
fn info(&self) -> Result<Self::Info, Error> {
self.execute(|connection| connection.node_text(&redis::cmd("INFO")))
}
}

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

@ -1,226 +0,0 @@
use std::time::Duration;
use litellm_cache::{CacheCodec, Error, PopOperation, PushOperation, QueueCache, SetCache};
use crate::{cache::RedisCache, script::RedisArg};
pub type RedisRpushOperation = PushOperation<RedisArg>;
pub type RedisLpopOperation = PopOperation;
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RedisLpopResult {
Missing,
Value(Vec<u8>),
Values(Vec<Vec<u8>>),
}
impl<S, C> SetCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type SetValue = RedisArg;
type SetResult = usize;
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry(None));
}
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(ttl);
let mut pipeline = redis::pipe();
pipeline
.cmd("SADD")
.arg(&key)
.arg(values)
.cmd("EXPIRE")
.arg(&key)
.arg(ttl)
.ignore();
self.run(move |connection| connection.query_pipeline(&pipeline))
.await
.map(|(added,)| added)
}
}
impl<S, C> QueueCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type QueueValue = RedisArg;
type PopResult = RedisLpopResult;
async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry(None));
}
let key = self.namespaced_key(key);
self.run(move |connection| {
redis::cmd("RPUSH")
.arg(key)
.arg(values)
.query(connection)
.map_err(Error::unavailable)
})
.await
}
async fn async_rpush_and_trim(
&self,
key: &str,
values: Vec<Self::QueueValue>,
max_len: usize,
) -> Result<usize, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry(None));
}
let key = self.namespaced_key(key);
let start = i64::try_from(max_len).map_or(i64::MIN, |max_len| -max_len);
let mut pipeline = redis::pipe();
pipeline
.atomic()
.cmd("RPUSH")
.arg(&key)
.arg(values)
.cmd("LTRIM")
.arg(&key)
.arg(start)
.arg(-1)
.ignore();
self.run(move |connection| connection.query_pipeline(&pipeline))
.await
.map(|(length,)| length)
}
async fn async_rpush_pipeline(
&self,
operations: Vec<RedisRpushOperation>,
) -> Result<Vec<usize>, Error> {
if operations.is_empty() {
return Ok(Vec::new());
}
let mut pipeline = redis::pipe();
for operation in operations {
if operation.values.is_empty() {
return Err(Error::InvalidEntry(None));
}
pipeline
.cmd("RPUSH")
.arg(self.namespaced_key(&operation.key))
.arg(operation.values);
}
self.run(move |connection| connection.query_pipeline(&pipeline))
.await
}
async fn async_lpop(&self, key: &str, count: Option<usize>) -> Result<Self::PopResult, Error> {
if let Some(count) = count
&& self.major_version().await < 7
{
return self.lpop_one_at_a_time(key, count).await;
}
let command = lpop(self.namespaced_key(key), count);
let value = self
.run(move |connection| {
command
.query::<redis::Value>(connection)
.map_err(Error::unavailable)
})
.await?;
lpop_result(value, count.is_some())
}
async fn async_lpop_pipeline(
&self,
operations: Vec<RedisLpopOperation>,
) -> Result<Vec<Self::PopResult>, Error> {
if operations.is_empty() {
return Ok(Vec::new());
}
if operations.iter().any(|operation| operation.count.is_some())
&& self.major_version().await < 7
{
let mut results = Vec::with_capacity(operations.len());
for operation in &operations {
results.push(self.async_lpop(&operation.key, operation.count).await?);
}
return Ok(results);
}
let multiple = operations
.iter()
.map(|operation| operation.count.is_some())
.collect::<Vec<_>>();
let mut pipeline = redis::pipe();
for operation in operations {
pipeline.add_command(lpop(self.namespaced_key(&operation.key), operation.count));
}
self.run(move |connection| connection.query_pipeline::<Vec<redis::Value>>(&pipeline))
.await?
.into_iter()
.zip(multiple)
.map(|(value, multiple)| lpop_result(value, multiple))
.collect()
}
}
fn lpop(key: String, count: Option<usize>) -> redis::Cmd {
let mut command = redis::cmd("LPOP");
command.arg(key);
if let Some(count) = count {
command.arg(count);
}
command
}
fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
match value {
redis::Value::BulkString(bytes) => Ok(bytes),
redis::Value::SimpleString(text) => Ok(text.into_bytes()),
_ => Err(Error::InvalidEntry(None)),
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
/// `handle_lpop_count_for_older_redis_versions`: `count` single-`LPOP` pipelines, keeping
/// only the values actually popped.
async fn lpop_one_at_a_time(&self, key: &str, count: usize) -> Result<RedisLpopResult, Error> {
let key = self.namespaced_key(key);
let mut values = Vec::new();
for _ in 0..count {
let mut pipeline = redis::pipe();
pipeline.add_command(lpop(key.clone(), None));
let replies = self
.run(move |connection| connection.query_pipeline::<Vec<redis::Value>>(&pipeline))
.await?;
for reply in replies {
if reply != redis::Value::Nil {
values.push(redis_bytes(reply)?);
}
}
}
Ok(RedisLpopResult::Values(values))
}
}
fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, Error> {
match value {
redis::Value::Nil => Ok(RedisLpopResult::Missing),
redis::Value::Array(values) if multiple => values
.into_iter()
.map(redis_bytes)
.collect::<Result<Vec<_>, _>>()
.map(RedisLpopResult::Values),
value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
_ => Err(Error::InvalidEntry(None)),
}
}

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

@ -1,136 +0,0 @@
use std::sync::Arc;
use litellm_cache::{CacheCodec, CacheScript, Error, ScriptCache};
use crate::{
cache::{RedisCache, namespaced_key},
connection::Connections,
};
#[derive(Clone, Debug, PartialEq)]
pub enum RedisArg {
Bytes(Vec<u8>),
Integer(i64),
Float(f64),
}
impl From<&str> for RedisArg {
fn from(value: &str) -> Self {
Self::Bytes(value.as_bytes().to_vec())
}
}
impl From<String> for RedisArg {
fn from(value: String) -> Self {
Self::Bytes(value.into_bytes())
}
}
impl From<Vec<u8>> for RedisArg {
fn from(value: Vec<u8>) -> Self {
Self::Bytes(value)
}
}
impl From<i64> for RedisArg {
fn from(value: i64) -> Self {
Self::Integer(value)
}
}
impl From<f64> for RedisArg {
fn from(value: f64) -> Self {
Self::Float(value)
}
}
impl redis::ToRedisArgs for RedisArg {
fn write_redis_args<W>(&self, out: &mut W)
where
W: ?Sized + redis::RedisWrite,
{
match self {
Self::Bytes(value) => value.write_redis_args(out),
Self::Integer(value) => value.write_redis_args(out),
Self::Float(value) => value.write_redis_args(out),
}
}
}
pub struct RedisScript<C> {
connections: Arc<Connections<C>>,
namespace: Option<String>,
source: String,
}
impl<C> CacheScript for RedisScript<C>
where
C: redis::ConnectionLike + Send + 'static,
{
type Argument = RedisArg;
type Output = redis::Value;
async fn invoke(
&self,
keys: Vec<String>,
arguments: Vec<Self::Argument>,
) -> Result<Self::Output, Error> {
let keys = keys
.into_iter()
.map(|key| namespaced_key(self.namespace.as_deref(), &key))
.collect::<Vec<_>>();
let source = self.source.clone();
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
eval(connection, &source, keys, arguments)
})
.await
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub async fn async_eval(
&self,
script: String,
keys: Vec<String>,
arguments: Vec<RedisArg>,
) -> Result<redis::Value, Error> {
let keys = self.namespaced_keys(&keys);
self.run(move |connection| eval(connection, &script, keys, arguments))
.await
}
}
fn eval(
connection: &mut impl redis::ConnectionLike,
script: &str,
keys: Vec<String>,
arguments: Vec<RedisArg>,
) -> Result<redis::Value, Error> {
redis::cmd("EVAL")
.arg(script)
.arg(keys.len())
.arg(keys)
.arg(arguments)
.query(connection)
.map_err(Error::unavailable)
}
impl<S, C> ScriptCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Script = RedisScript<C>;
fn async_register_script(&self, source: String) -> Self::Script {
RedisScript {
connections: Arc::clone(&self.connections),
namespace: self.namespace.clone(),
source,
}
}
}

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

@ -1,232 +0,0 @@
use std::time::Duration;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheCodec, DeleteCache, Error,
ExactCacheContext, FlushAllCache, FlushCache, TtlPipelineCache,
};
use redis::Commands;
use crate::{cache::RedisCache, connection::ConnectionRef};
impl<S, C> BaseCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = S::Value;
type Context = ExactCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(self.default_ttl))
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &ExactCacheContext,
) -> Result<(), Error> {
let payload = self.codec.encode(&value)?;
let ttl = self.ttl_or_default(context.ttl);
let key = self.namespaced_key(key);
self.execute(|connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(Error::unavailable)
})
}
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = self.execute(|connection| {
connection
.get::<_, redis::Value>(key)
.map_err(Error::unavailable)
})?;
self.decode_response(value)
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: ExactCacheContext,
) -> Result<(), Error> {
let payload = self.codec.encode(&value)?;
let key = self.namespaced_key(key);
let ttl = self.ttl_or_default(context.ttl);
self.run(move |connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(Error::unavailable)
})
.await
}
async fn async_get_cache(
&self,
key: &str,
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = self
.run(move |connection| {
connection
.get::<_, redis::Value>(key)
.map_err(Error::unavailable)
})
.await?;
self.decode_response(value)
}
async fn async_set_cache_pipeline(
&self,
cache_list: Vec<(String, Self::Value)>,
context: ExactCacheContext,
) -> Result<(), Error> {
self.async_set_cache_pipeline_with_ttls(
cache_list
.into_iter()
.map(|(key, value)| (key, value, context.ttl))
.collect(),
)
.await
}
}
impl<S, C> TtlPipelineCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_set_cache_pipeline_with_ttls(
&self,
entries: Vec<(String, Self::Value, Option<Duration>)>,
) -> Result<(), Error> {
if entries.is_empty() {
return Ok(());
}
let mut pipeline = redis::pipe();
for (key, value, ttl) in entries {
pipeline
.cmd("SETEX")
.arg(self.namespaced_key(&key))
.arg(self.ttl_or_default(ttl))
.arg(self.codec.encode(&value)?)
.ignore();
}
self.run(move |connection| connection.query_pipeline(&pipeline))
.await
}
}
impl<S, C> BatchCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn batch_get_cache(
&self,
keys: &[String],
_: &ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = self.namespaced_keys(keys);
self.execute(|connection| mget(connection, keys))?
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
_: ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = self.namespaced_keys(&keys);
self.run(move |connection| mget(connection, keys))
.await?
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
}
pub(crate) fn mget(
connection: &mut ConnectionRef<'_>,
keys: Vec<String>,
) -> Result<Vec<redis::Value>, Error> {
redis::cmd("MGET")
.arg(keys)
.query(connection)
.map_err(Error::unavailable)
}
impl<S, C> DeleteCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
self.execute(|connection| connection.del::<_, ()>(key).map_err(Error::unavailable))
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
self.run(move |connection| connection.del::<_, ()>(key).map_err(Error::unavailable))
.await
}
}
impl<S, C> BulkDeleteCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
if keys.is_empty() {
return Ok(0);
}
let keys = self.namespaced_keys(&keys);
self.run(move |connection| connection.del(keys).map_err(Error::unavailable))
.await
}
}
impl<S, C> FlushCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
self.execute(|connection| flush_matching(connection, &pattern))
}
async fn async_flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
self.run(move |connection| flush_matching(connection, &pattern))
.await
}
}
impl<S, C> FlushAllCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn flushall(&self) -> Result<(), Error> {
self.execute(|connection| connection.flushall())
}
}
fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> {
connection.scan(pattern, 1000, |connection, keys| {
if !keys.is_empty() {
connection
.del::<_, usize>(keys)
.map_err(Error::unavailable)?;
}
Ok(true)
})
}

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

@ -1,14 +0,0 @@
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RedisNode {
pub host: String,
pub port: u16,
}
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub enum RedisTopology {
#[default]
Standalone,
Cluster {
startup_nodes: Vec<RedisNode>,
},
}

1164
litellm-rust/crates/cache-redis/tests/cache.rs
View file

File diff suppressed because it is too large Load diff

563
litellm-rust/crates/cache-redis/tests/cluster.rs
View file

@ -1,563 +0,0 @@
//! Tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a comma
//! separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them.
mod support;
use std::{collections::HashSet, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, BoundedCounterCache, BulkDeleteCache, CacheConnectionStatus,
CacheScript, ClaimCache, ClientInfoCache, ConnectionCache, CounterCache, DeleteCache,
DisconnectCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
PingCache, QueueCache, RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache,
TtlPipelineCache,
};
use litellm_cache_redis::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisNode, RedisRpushOperation,
RedisTopology,
};
use redis::cluster_routing::Slot;
use rstest::{fixture, rstest};
use serde_json::json;
use support::{JsonCache, cluster_cache, cluster_url};
type Counter = RedisCache<JsonCodec<f64>>;
#[fixture]
fn cache(#[default("cache")] label: &str) -> Option<JsonCache> {
cluster_cache(label, Duration::from_secs(120), JsonCodec::new())
}
#[fixture]
fn counter(#[default("counter")] label: &str) -> Option<Counter> {
cluster_cache(label, Duration::from_secs(60), JsonCodec::new())
}
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
fn multi_slot_keys(count: usize) -> Vec<String> {
let keys: Vec<String> = (0..count).map(|index| format!("key-{index}")).collect();
let slots: HashSet<Slot> = keys.iter().map(Slot::for_key).collect();
assert!(slots.len() > 1, "keys must span multiple slots");
keys
}
fn seconds(seconds: u64) -> Option<Duration> {
Some(Duration::from_secs(seconds))
}
#[rstest]
#[case::no_startup_nodes("redis://127.0.0.1:7000", Vec::new())]
#[case::unix_socket_url(
"redis+unix:///tmp/redis.sock",
vec![RedisNode { host: "127.0.0.1".into(), port: 7000 }]
)]
fn constructor_rejects_unusable_cluster_configs(
#[case] url: &str,
#[case] startup_nodes: Vec<RedisNode>,
) {
let error = JsonCache::connect(
url,
&RedisTopology::Cluster { startup_nodes },
None,
JsonCodec::new(),
)
.err();
assert!(matches!(error, Some(Error::Unavailable(_))));
}
#[rstest]
#[tokio::test]
async fn single_key_operations_round_trip_with_ttl_rounding(
#[with("single")] cache: Option<JsonCache>,
) {
let Some(cache) = cache else { return };
let context = ExactCacheContext {
ttl: Some(Duration::from_millis(1500)),
};
let keys = multi_slot_keys(12);
for (index, key) in keys.iter().enumerate() {
cache
.set_cache(key, json!({ "index": index }), &context)
.unwrap();
}
for (index, key) in keys.iter().enumerate() {
assert_eq!(
cache.get_cache(key, &context).unwrap(),
Some(json!({ "index": index }))
);
}
assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(2));
assert!(
cache
.async_refresh_ttl(&keys[0], seconds(40))
.await
.unwrap()
);
assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(40));
cache.delete_cache(&keys[0]).unwrap();
assert_eq!(cache.get_cache(&keys[0], &context).unwrap(), None);
assert!(!cache.async_refresh_ttl(&keys[0], None).await.unwrap());
assert!(cache.sync_ping().unwrap());
cache.async_flush_cache().await.unwrap();
}
#[rstest]
#[tokio::test]
async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries(
#[with("batch")] cache: Option<JsonCache>,
context: ExactCacheContext,
) {
let Some(cache) = cache else { return };
let keys = multi_slot_keys(40);
for (index, key) in keys.iter().enumerate() {
if index % 5 == 0 {
continue;
}
cache
.async_set_cache(key, json!(index), context.clone())
.await
.unwrap();
}
let mut raw = redis::cluster::ClusterClient::new(vec![cluster_url()])
.unwrap()
.get_connection()
.unwrap();
let malformed = format!("{}:{}", cache.namespace().unwrap(), keys[1]);
redis::cmd("SET")
.arg(&malformed)
.arg("not json")
.exec(&mut raw)
.unwrap();
let entries = cache
.async_batch_get_cache(keys.clone(), context.clone())
.await
.unwrap();
assert_eq!(entries.len(), keys.len());
for (index, entry) in entries.iter().enumerate() {
let expected = if index == 1 {
BatchEntry::Invalid
} else if index % 5 == 0 {
BatchEntry::Miss
} else {
BatchEntry::Hit(json!(index))
};
assert_eq!(*entry, expected, "entry {index}");
}
assert_eq!(cache.batch_get_cache(&keys, &context).unwrap(), entries);
cache.delete_cache_keys(keys.clone()).await.unwrap();
let entries = cache.async_batch_get_cache(keys, context).await.unwrap();
assert!(entries.iter().all(|entry| *entry == BatchEntry::Miss));
}
#[rstest]
#[tokio::test]
async fn pipelines_group_by_slot_and_return_results_in_submission_order(
#[with("pipeline")] cache: Option<JsonCache>,
counter: Option<Counter>,
context: ExactCacheContext,
) {
let (Some(cache), Some(counter)) = (cache, counter) else {
return;
};
let keys = multi_slot_keys(30);
cache
.async_set_cache_pipeline(
keys.iter()
.enumerate()
.map(|(index, key)| (key.clone(), json!(index)))
.collect(),
context.clone(),
)
.await
.unwrap();
let hits = cache
.async_batch_get_cache(keys.clone(), context.clone())
.await
.unwrap();
assert!(
hits.iter()
.enumerate()
.all(|(index, entry)| *entry == BatchEntry::Hit(json!(index)))
);
cache
.async_set_cache_pipeline_with_ttls(
keys.iter()
.enumerate()
.map(|(index, key)| (key.clone(), json!(index), seconds(index as u64 + 10)))
.collect(),
)
.await
.unwrap();
for (index, key) in keys.iter().enumerate() {
assert_eq!(
cache.async_get_ttl(key).await.unwrap(),
seconds(index as u64 + 10),
"{key}"
);
}
let queues: Vec<String> = keys.iter().map(|key| format!("queue:{key}")).collect();
let pushed = cache
.async_rpush_pipeline(
queues
.iter()
.enumerate()
.map(|(index, key)| RedisRpushOperation {
key: key.clone(),
values: (0..=index)
.map(|value| RedisArg::Integer(value as i64))
.collect(),
})
.collect(),
)
.await
.unwrap();
assert_eq!(pushed, (1..=keys.len()).collect::<Vec<_>>());
let popped = cache
.async_lpop_pipeline(
queues
.iter()
.enumerate()
.map(|(index, key)| RedisLpopOperation {
key: key.clone(),
count: (index % 2 == 0).then_some(2),
})
.collect(),
)
.await
.unwrap();
for (index, result) in popped.into_iter().enumerate() {
match result {
RedisLpopResult::Value(value) => {
assert_eq!(index % 2, 1, "queue {index}");
assert_eq!(value, b"0");
}
RedisLpopResult::Values(values) => {
assert_eq!(index % 2, 0, "queue {index}");
let expected: Vec<Vec<u8>> = (0..=index)
.take(2)
.map(|value| value.to_string().into_bytes())
.collect();
assert_eq!(values, expected);
}
other => panic!("queue {index}: {other:?}"),
}
}
let counters: Vec<String> = keys.iter().map(|key| format!("counter:{key}")).collect();
let totals = counter
.async_increment_pipeline(
counters
.iter()
.enumerate()
.map(|(index, key)| IncrementOperation {
key: key.clone(),
amount: index as f64 + 0.5,
ttl: (index % 3 == 0).then_some(Duration::from_secs(30)),
})
.collect(),
)
.await
.unwrap();
let expected: Vec<f64> = (0..keys.len()).map(|index| index as f64 + 0.5).collect();
assert_eq!(totals, expected);
assert_eq!(
counter.async_get_ttl(&counters[0]).await.unwrap(),
seconds(30)
);
assert_eq!(counter.async_get_ttl(&counters[1]).await.unwrap(), None);
counter.async_flush_cache().await.unwrap();
cache.async_flush_cache().await.unwrap();
}
#[rstest]
#[tokio::test]
async fn rpush_and_trim_is_one_transaction_on_the_key_slot(
#[with("trim")] cache: Option<JsonCache>,
) {
let Some(cache) = cache else { return };
let values = |values: &[&str]| values.iter().map(|value| RedisArg::from(*value)).collect();
assert_eq!(
cache
.async_rpush_and_trim("buf", values(&["a", "b"]), 3)
.await
.unwrap(),
2
);
assert_eq!(
cache
.async_rpush_and_trim("buf", values(&["c", "d"]), 3)
.await
.unwrap(),
4
);
assert_eq!(
cache.async_lpop("buf", Some(10)).await.unwrap(),
RedisLpopResult::Values(vec![b"b".to_vec(), b"c".to_vec(), b"d".to_vec()])
);
cache.async_flush_cache().await.unwrap();
}
#[rstest]
#[tokio::test]
async fn scan_and_scoped_flush_cover_every_primary(
#[with("flush")] cache: Option<JsonCache>,
#[from(cache)]
#[with("other")]
other: Option<JsonCache>,
context: ExactCacheContext,
) {
let (Some(cache), Some(other)) = (cache, other) else {
return;
};
let keys = multi_slot_keys(60);
for key in &keys {
cache
.async_set_cache(key, json!(true), context.clone())
.await
.unwrap();
other
.async_set_cache(key, json!(true), context.clone())
.await
.unwrap();
}
let mut scanned = cache.async_scan_iter("key-", 1000).await.unwrap();
scanned.sort();
let mut expected: Vec<String> = keys
.iter()
.map(|key| format!("{}:{key}", cache.namespace().unwrap()))
.collect();
expected.sort();
assert_eq!(scanned, expected);
assert_eq!(cache.async_scan_iter("key-", 7).await.unwrap().len(), 7);
cache.flush_cache().unwrap();
let flushed = cache
.async_batch_get_cache(keys.clone(), context.clone())
.await
.unwrap();
assert!(flushed.iter().all(|entry| *entry == BatchEntry::Miss));
let kept = other.async_batch_get_cache(keys, context).await.unwrap();
assert!(
kept.iter()
.all(|entry| *entry == BatchEntry::Hit(json!(true)))
);
other.async_flush_cache().await.unwrap();
}
fn ping_calls_per_node(startup: &redis::Client) -> Vec<(String, u64)> {
let mut connection = startup.get_connection().unwrap();
let nodes: String = redis::cmd("CLUSTER")
.arg("NODES")
.query(&mut connection)
.unwrap();
let mut counts: Vec<(String, u64)> = nodes
.lines()
.map(|line| {
let address = line.split_whitespace().nth(1).unwrap();
let address = address.split('@').next().unwrap();
let mut node = redis::Client::open(format!("redis://{address}"))
.unwrap()
.get_connection()
.unwrap();
let stats: String = redis::cmd("INFO")
.arg("commandstats")
.query(&mut node)
.unwrap();
let calls = stats
.lines()
.find_map(|stat| stat.strip_prefix("cmdstat_ping:calls="))
.and_then(|rest| rest.split(',').next())
.map_or(0, |calls| calls.parse().unwrap());
(address.to_string(), calls)
})
.collect();
counts.sort();
counts
}
#[rstest]
#[tokio::test]
async fn ping_reaches_every_node(#[with("ping")] cache: Option<JsonCache>) {
let Some(cache) = cache else { return };
let startup = redis::Client::open(cluster_url()).unwrap();
let before = ping_calls_per_node(&startup);
assert!(before.len() >= 2, "{before:?}");
assert!(cache.ping().await.unwrap());
let after = ping_calls_per_node(&startup);
for ((node, calls_before), (_, calls_after)) in before.iter().zip(&after) {
assert!(calls_after > calls_before, "{node} was not pinged");
}
assert!(cache.sync_ping().unwrap());
let result = cache.test_connection().await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert_eq!(result.message, "Redis Cluster connection test successful");
}
#[rstest]
#[tokio::test]
async fn disconnect_closes_idle_connections_and_reconnects_on_demand(
#[with("disconnect")] cache: Option<JsonCache>,
) {
let Some(cache) = cache else { return };
assert!(cache.ping().await.unwrap());
cache.disconnect().await.unwrap();
assert!(cache.ping().await.unwrap());
}
#[rstest]
#[case::keep_existing_ttl(false)]
#[case::refresh_ttl(true)]
#[tokio::test]
async fn increments_refresh_the_ttl_only_when_asked(
counter: Option<Counter>,
#[case] refresh_ttl: bool,
) {
let Some(counter) = counter else { return };
let context = ExactCacheContext { ttl: seconds(60) };
counter
.async_set_cache("spend", 0.0, ExactCacheContext { ttl: seconds(600) })
.await
.unwrap();
assert_eq!(
counter
.async_increment("spend", 1.5, context.clone(), refresh_ttl)
.await
.unwrap(),
1.5
);
assert_eq!(
counter
.async_increment("spend", 2.0, context, refresh_ttl)
.await
.unwrap(),
3.5
);
let ttl = counter.async_get_ttl("spend").await.unwrap().unwrap();
assert_eq!(ttl <= Duration::from_secs(60), refresh_ttl, "{ttl:?}");
counter.async_flush_cache().await.unwrap();
}
#[rstest]
#[tokio::test]
async fn counters_claims_scripts_and_sets_work_on_the_cluster(
counter: Option<Counter>,
#[with("claim")] cache: Option<JsonCache>,
context: ExactCacheContext,
) {
let (Some(counter), Some(cache)) = (counter, cache) else {
return;
};
assert_eq!(
counter
.increment_cache("spend", 1.5, context.clone())
.unwrap(),
1.5
);
assert_eq!(
counter
.async_increment("spend", 2.0, context.clone(), false)
.await
.unwrap(),
3.5
);
assert_eq!(
counter
.increment_with_floor("budget", -3, Duration::from_secs(30))
.unwrap(),
0
);
assert_eq!(
counter
.async_increment_with_floor("budget", 7, Duration::from_secs(30))
.await
.unwrap(),
7
);
assert_eq!(counter.async_set_max("peak", 4.0, None).await.unwrap(), 4.0);
assert_eq!(counter.async_set_max("peak", 2.0, None).await.unwrap(), 4.0);
counter.flush_cache().unwrap();
let owner = json!("owner-a");
let rival = json!("owner-b");
assert_eq!(
cache
.claim_cache("lock", owner.clone(), &[], context.clone())
.unwrap(),
owner
);
assert_eq!(
cache
.async_claim_cache("lock", rival.clone(), vec![owner.clone()], context.clone())
.await
.unwrap(),
owner
);
assert_eq!(
cache
.claim_cache("lock", rival.clone(), &[], context.clone())
.unwrap(),
owner
);
assert_eq!(
cache
.async_claim_cache("lock", rival.clone(), vec![rival.clone()], context.clone())
.await
.unwrap(),
rival
);
let script = cache
.async_register_script("return redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2])".into());
let reply = script
.invoke(
vec!["scripted".into()],
vec![RedisArg::Bytes(b"payload".to_vec()), RedisArg::Integer(5)],
)
.await
.unwrap();
assert_eq!(reply, redis::Value::Okay);
assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), seconds(5));
let evaluated = cache
.async_eval(
"return redis.call('GET', KEYS[1])".into(),
vec!["scripted".into()],
Vec::new(),
)
.await
.unwrap();
assert_eq!(evaluated, redis::Value::BulkString(b"payload".to_vec()));
assert_eq!(
cache
.async_set_cache_sadd(
"members",
vec![
RedisArg::Bytes(b"a".to_vec()),
RedisArg::Bytes(b"b".to_vec())
],
seconds(9),
)
.await
.unwrap(),
2
);
assert_eq!(cache.async_get_ttl("members").await.unwrap(), seconds(9));
let result = cache.test_connection().await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert!(cache.ping().await.unwrap());
let info = cache.info().unwrap();
assert!(info.matches("redis_version").count() > 1, "{info}");
assert!(cache.client_list().unwrap().contains("id="));
cache.async_flush_cache().await.unwrap();
assert_eq!(cache.async_get_ttl("members").await.unwrap(), None);
assert_eq!(cache.get_cache("lock", &context).unwrap(), None);
}

95
litellm-rust/crates/cache-redis/tests/contract.rs
View file

@ -1,95 +0,0 @@
//! The shared cache contracts, run against the in-process fake connection and, when
//! `LITELLM_TEST_REDIS_CLUSTER_NODES` is set, against a live Redis Cluster.
mod support;
use std::time::Duration;
use litellm_cache::{ExactCacheContext, JsonCodec};
use litellm_cache_testing as contract;
use rstest::rstest;
use serde_json::json;
use support::{JsonCache, cluster_cache, fake_cache};
const PREFIX: &str = "contract:";
#[derive(Clone, Copy, Debug)]
enum Contract {
HitAndMiss,
SyncAsyncEquivalence,
OverwriteReplaces,
PipelineWritesEveryEntry,
BatchPreservesOrder,
DeleteRemovesKey,
FlushClears,
CounterAccumulates,
}
#[derive(Clone, Copy, Debug)]
enum Server {
Fake,
Cluster,
}
async fn check<C>(contract: Contract, cache: &JsonCache<C>)
where
C: redis::ConnectionLike + Send + 'static,
{
let context = ExactCacheContext::default();
match contract {
Contract::HitAndMiss => {
contract::hit_and_miss(cache, context, PREFIX, json!({"answer": 42})).await
}
Contract::SyncAsyncEquivalence => {
contract::sync_async_equivalence(cache, context, PREFIX, json!("first"), json!([2]))
.await
}
Contract::OverwriteReplaces => {
contract::overwrite_replaces(cache, context, PREFIX, json!(1), json!({"b": 2})).await
}
Contract::PipelineWritesEveryEntry => {
contract::pipeline_writes_every_entry(
cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await
}
Contract::BatchPreservesOrder => {
contract::batch_preserves_order(cache, context, PREFIX, json!("first"), json!(2)).await
}
Contract::DeleteRemovesKey => {
contract::delete_removes_key(cache, context, PREFIX, json!("value")).await
}
Contract::FlushClears => {
contract::flush_clears(cache, context, PREFIX, json!("value")).await
}
Contract::CounterAccumulates => contract::counter_accumulates(cache, context, PREFIX).await,
}
}
#[rstest]
#[case::hit_and_miss(Contract::HitAndMiss)]
#[case::sync_async_equivalence(Contract::SyncAsyncEquivalence)]
#[case::overwrite_replaces(Contract::OverwriteReplaces)]
#[case::pipeline_writes_every_entry(Contract::PipelineWritesEveryEntry)]
#[case::batch_preserves_order(Contract::BatchPreservesOrder)]
#[case::delete_removes_key(Contract::DeleteRemovesKey)]
#[case::flush_clears(Contract::FlushClears)]
#[case::counter_accumulates(Contract::CounterAccumulates)]
#[tokio::test]
async fn redis_satisfies_the_cache_contract(
#[case] contract: Contract,
#[values(Server::Fake, Server::Cluster)] server: Server,
) {
match server {
Server::Fake => check(contract, &fake_cache("contract")).await,
Server::Cluster => {
let label = format!("{contract:?}");
if let Some(cache) = cluster_cache(&label, Duration::from_secs(120), JsonCodec::new()) {
check(contract, &cache).await;
}
}
}
}

231
litellm-rust/crates/cache-redis/tests/support/mod.rs
View file

@ -1,231 +0,0 @@
#![allow(dead_code)]
use std::{
collections::BTreeMap,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{CacheCodec, Error, JsonCodec};
use litellm_cache_redis::{RedisCache, RedisNode, RedisTopology};
/// Encodes a byte behind a tag, so a value written with another tag decodes as invalid.
pub struct TaggedByteCodec(pub u8);
impl CacheCodec for TaggedByteCodec {
type Value = u8;
fn encode(&self, value: &u8) -> Result<Vec<u8>, Error> {
if *value > 127 {
return Err(Error::InvalidEntry(None));
}
Ok(vec![self.0, *value])
}
fn decode(&self, bytes: &[u8]) -> Result<u8, Error> {
match bytes {
[tag, value] if *tag == self.0 => Ok(*value),
_ => Err(Error::InvalidEntry(None)),
}
}
}
/// A stateful in-process stand-in for a Redis server that understands the string commands the
/// shared contracts exercise, so they run without a live server. TTLs are accepted and ignored.
#[derive(Default)]
pub struct FakeRedis {
strings: BTreeMap<Vec<u8>, Vec<u8>>,
}
impl FakeRedis {
fn run(&mut self, command: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
let name = String::from_utf8_lossy(&command[0]).to_ascii_uppercase();
let args = &command[1..];
Ok(match name.as_str() {
"PING" => redis::Value::SimpleString("PONG".into()),
"SET" | "SETEX" => {
let value = if name == "SET" { &args[1] } else { &args[2] };
self.strings.insert(args[0].clone(), value.clone());
redis::Value::Okay
}
"GET" => self.get(&args[0]),
"MGET" => redis::Value::Array(args.iter().map(|key| self.get(key)).collect()),
"DEL" => {
let removed = args
.iter()
.filter(|key| self.strings.remove(*key).is_some())
.count();
redis::Value::Int(removed as i64)
}
"SCAN" => {
let pattern = &args[2];
let keys = self
.strings
.keys()
.filter(|key| glob(pattern, key))
.map(|key| redis::Value::BulkString(key.clone()))
.collect();
redis::Value::Array(vec![
redis::Value::BulkString(b"0".to_vec()),
redis::Value::Array(keys),
])
}
"EVAL" if args[0].windows(11).any(|window| window == b"INCRBYFLOAT") => {
self.increment_by_float(&args[2], &args[3])
}
"INCRBYFLOAT" => self.increment_by_float(&args[0], &args[1]),
_ => {
return Err(redis::RedisError::from((
redis::ErrorKind::Client,
"unsupported command",
name,
)));
}
})
}
fn get(&self, key: &[u8]) -> redis::Value {
self.strings.get(key).map_or(redis::Value::Nil, |value| {
redis::Value::BulkString(value.clone())
})
}
fn increment_by_float(&mut self, key: &[u8], amount: &[u8]) -> redis::Value {
let current = self
.strings
.get(key)
.map_or(0.0, |value| parse_float(value));
let total = format!("{}", current + parse_float(amount));
self.strings
.insert(key.to_vec(), total.clone().into_bytes());
redis::Value::BulkString(total.into_bytes())
}
}
fn parse_float(bytes: &[u8]) -> f64 {
std::str::from_utf8(bytes).unwrap().parse().unwrap()
}
/// Redis `MATCH` globbing for `*`, `?` and backslash escapes.
fn glob(pattern: &[u8], key: &[u8]) -> bool {
match pattern.split_first() {
None => key.is_empty(),
Some((b'*', rest)) => (0..=key.len()).any(|skip| glob(rest, &key[skip..])),
Some((b'?', rest)) => !key.is_empty() && glob(rest, &key[1..]),
Some((b'\\', [escaped, rest @ ..])) => {
key.first() == Some(escaped) && glob(rest, &key[1..])
}
Some((literal, rest)) => key.first() == Some(literal) && glob(rest, &key[1..]),
}
}
/// Splits RESP request bytes into the commands they carry.
fn commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
fn line<'a>(bytes: &mut &'a [u8]) -> &'a [u8] {
let end = bytes
.windows(2)
.position(|window| window == b"\r\n")
.unwrap();
let (line, rest) = bytes.split_at(end);
*bytes = &rest[2..];
line
}
fn length(line: &[u8]) -> usize {
std::str::from_utf8(&line[1..]).unwrap().parse().unwrap()
}
let mut commands = Vec::new();
while !bytes.is_empty() {
let count = length(line(&mut bytes));
let command = (0..count)
.map(|_| {
let size = length(line(&mut bytes));
let (argument, rest) = bytes.split_at(size);
bytes = &rest[2..];
argument.to_vec()
})
.collect();
commands.push(command);
}
commands
}
impl redis::ConnectionLike for FakeRedis {
fn req_packed_command(&mut self, cmd: &[u8]) -> redis::RedisResult<redis::Value> {
let command = commands(cmd).into_iter().next().unwrap();
self.run(command)
}
fn req_packed_commands(
&mut self,
cmd: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
let replies = commands(cmd)
.into_iter()
.map(|command| self.run(command))
.collect::<redis::RedisResult<Vec<_>>>()?;
Ok(replies.into_iter().skip(offset).take(count).collect())
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}
pub type JsonCache<C = redis::Connection> = RedisCache<JsonCodec<serde_json::Value>, C>;
pub fn fake_cache(namespace: &str) -> JsonCache<FakeRedis> {
RedisCache::with_connection(FakeRedis::default(), None, JsonCodec::new())
.with_namespace(Some(namespace.into()))
}
/// Startup nodes from `LITELLM_TEST_REDIS_CLUSTER_NODES` (`host:port,host:port`); tests that
/// need a live cluster skip when it is unset.
pub fn cluster_topology() -> Option<RedisTopology> {
let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?;
let startup_nodes = nodes
.split(',')
.map(|node| {
let (host, port) = node.trim().rsplit_once(':').expect("host:port");
RedisNode {
host: host.to_string(),
port: port.parse().expect("port"),
}
})
.collect();
Some(RedisTopology::Cluster { startup_nodes })
}
pub fn cluster_url() -> String {
std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL")
.unwrap_or_else(|_| "redis://127.0.0.1:7000".into())
}
pub fn unique_namespace(label: &str) -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
format!("cluster-test:{label}:{nanos}")
}
pub fn cluster_cache<S: CacheCodec>(
label: &str,
default_ttl: Duration,
codec: S,
) -> Option<RedisCache<S>> {
let topology = cluster_topology()?;
Some(
RedisCache::connect(&cluster_url(), &topology, Some(default_ttl), codec)
.expect("cluster connection")
.with_namespace(Some(unique_namespace(label))),
)
}

3
litellm-rust/crates/cache-response/Cargo.toml
View file

@ -16,9 +16,6 @@ sha2.workspace = true
litellm-cache-gcs.workspace = true
litellm-http = { workspace = true, features = ["test-support"] }
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
redis.workspace = true
redis-test.workspace = true
rstest.workspace = true
tokio.workspace = true
wiremock.workspace = true

72
litellm-rust/crates/cache-response/tests/connection.rs
View file

@ -2,79 +2,11 @@ mod support;
use std::{sync::Arc, time::Duration};
use litellm_cache::CacheConnectionStatus;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_response::{
CacheEntry, ConnectionProbe, ExactResponseCache, ResponseCache, ResponseCacheRequest,
};
use redis_test::MockCmd;
use litellm_cache_response::{CacheEntry, ExactResponseCache, ResponseCache, ResponseCacheRequest};
use rstest::rstest;
use serde_json::json;
use support::{keyed, memory, redis, request};
#[rstest]
#[case::reachable(
Ok("PONG"),
CacheConnectionStatus::Success,
"Redis connection test successful",
false
)]
#[case::unexpected_reply(
Ok("NOPE"),
CacheConnectionStatus::Failed,
"Redis ping returned False",
false
)]
#[case::connection_refused(
Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))),
CacheConnectionStatus::Failed,
"Redis connection failed:",
true
)]
#[tokio::test]
async fn connection_backends_are_reachable_as_a_probe(
#[case] reply: redis::RedisResult<&'static str>,
#[case] status: CacheConnectionStatus,
#[case] message: &str,
#[case] has_error: bool,
) {
let probe: Arc<dyn ConnectionProbe> =
Arc::new(redis(vec![MockCmd::new(redis::cmd("PING"), reply)], None));
let result = probe.test_connection().await.unwrap();
assert_eq!(result.status, status);
assert!(result.message.starts_with(message), "{}", result.message);
assert_eq!(result.error.is_some(), has_error);
}
#[rstest]
#[tokio::test]
async fn one_service_serves_both_the_exact_cache_and_its_probe(request: ResponseCacheRequest) {
let service = Arc::new(redis(
vec![
MockCmd::new(redis::cmd("PING"), Ok("PONG")),
MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(br#"{"timestamp":100.0,"response":{"ok":true}}"#.to_vec()),
),
],
Some("tenant"),
));
let probe: Arc<dyn ConnectionProbe> = service.clone();
let exact: Arc<dyn ExactResponseCache> = service;
assert_eq!(
probe.test_connection().await.unwrap().status,
CacheConnectionStatus::Success
);
assert_eq!(
exact
.async_lookup(&request, Duration::from_secs(100))
.await
.unwrap(),
Some(json!({"ok": true}))
);
}
use support::{keyed, memory, request};
/// The in-memory backend has no `test_connection`, as in Python, and still serves every response
/// operation.

97
litellm-rust/crates/cache-response/tests/response.rs
View file

@ -17,10 +17,9 @@ use litellm_cache_response::{
CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest,
WriteBuffer, cache_key,
};
use redis_test::MockCmd;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::{keyed, memory, redis, request};
use support::{keyed, memory, request};
type Memory = Arc<ResponseCache<InMemoryCache<CacheEntry>>>;
@ -331,79 +330,6 @@ async fn directives_skip_io_and_keep_reads_and_writes_independent(
);
}
#[rstest]
#[case::python_sync_literal(
br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.as_slice()
)]
#[case::python_async_json(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice())]
#[tokio::test]
async fn redis_consumer_reads_python_sync_and_async_envelopes(
request: ResponseCacheRequest,
#[case] stored: &[u8],
#[values(false, true)] asynchronous: bool,
) {
let cache = redis(
vec![MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(stored.to_vec()),
)],
Some("tenant"),
);
let now = Duration::from_secs(101);
let found = if asynchronous {
cache.async_lookup(&request, now).await.unwrap()
} else {
cache.lookup(&request, now).unwrap()
};
assert_eq!(found, Some(json!({"ok": true, "text": "cached"})));
}
#[rstest]
#[case::object(
json!({"ok": true, "text": "cached"}),
br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()
)]
#[case::array(json!([1, 2]), br#"{"timestamp":100.0,"response":"[1,2]"}"#.as_slice())]
#[tokio::test]
async fn redis_consumer_writes_python_compatible_json(
request: ResponseCacheRequest,
#[case] response: Value,
#[case] wire: &[u8],
) {
let cache = redis(
vec![MockCmd::new(
redis::cmd("SETEX").arg("tenant:key").arg(600).arg(wire),
Ok("OK"),
)],
Some("tenant"),
);
cache
.async_store(&request, response, Duration::from_secs(100))
.await
.unwrap();
}
#[rstest]
#[tokio::test]
async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis(
mut request: ResponseCacheRequest,
) {
let cache = redis(
vec![MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(b"invalid".to_vec()),
)],
None,
);
request.controls.no_cache = true;
assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
request.controls.no_cache = false;
assert_eq!(
cache.async_lookup(&request, Duration::ZERO).await.unwrap(),
None
);
}
#[rstest]
#[tokio::test]
async fn captured_service_keeps_the_selected_backend_for_background_writes(
@ -541,27 +467,6 @@ async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses(
assert_eq!(memory.lookup(&requests[2], now).unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn batch_lookup_with_no_readable_request_skips_the_backend(
#[values(false, true)] asynchronous: bool,
) {
let cache = redis(Vec::new(), None);
let mut request = keyed("key");
request.controls.no_cache = true;
let requests = [request.clone(), request];
let partial = if asynchronous {
cache
.async_lookup_batch(&requests, Duration::ZERO)
.await
.unwrap()
} else {
cache.lookup_batch(&requests, Duration::ZERO).unwrap()
};
assert_eq!(partial.values, vec![None, None]);
assert_eq!(partial.missing_indices, vec![0, 1]);
}
#[rstest]
#[tokio::test]
async fn deferred_entries_keep_the_time_they_were_produced(

17
litellm-rust/crates/cache-response/tests/support/mod.rs
View file

@ -1,15 +1,9 @@
use std::{sync::Arc, time::Duration};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::RedisCache;
use litellm_cache_response::{
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
};
use redis_test::{MockCmd, MockRedisConnection};
use litellm_cache_response::{CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheRequest};
use rstest::fixture;
pub type MockedRedis = RedisCache<ResponseCacheCodec, MockRedisConnection>;
#[fixture]
pub fn memory() -> Arc<ResponseCache<InMemoryCache<CacheEntry>>> {
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
@ -29,12 +23,3 @@ pub fn keyed(key: &str) -> ResponseCacheRequest {
..Default::default()
})
}
/// A Redis response cache that must receive exactly `commands`, in order.
pub fn redis(commands: Vec<MockCmd>, namespace: Option<&str>) -> ResponseCache<MockedRedis> {
let connection = MockRedisConnection::new(commands).assert_all_commands_consumed();
ResponseCache::new(Arc::new(
RedisCache::with_connection(connection, None, ResponseCacheCodec)
.with_namespace(namespace.map(str::to_owned)),
))
}

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

@ -1,22 +0,0 @@
[package]
name = "litellm-cache-valkey-semantic"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
[dependencies]
litellm-cache.workspace = true
litellm-cache-redis.workspace = true
redis = { workspace = true, features = ["tls-rustls"] }
sha2.workspace = true
thiserror.workspace = true
uuid = { version = "1", features = ["v4"] }
[dev-dependencies]
litellm-error.workspace = true
litellm-cache-testing.workspace = true
redis-test.workspace = true
rstest.workspace = true
serde_json.workspace = true
tokio.workspace = true

245
litellm-rust/crates/cache-valkey-semantic/src/cache.rs
View file

@ -1,245 +0,0 @@
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::{
BaseCache, CacheCodec, Error, SemanticCacheContext,
semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
};
use litellm_cache_redis::{RedisTopology, connection::Connections};
use crate::{
ValkeySemanticConfig,
index::IndexState,
search::{embedding_bytes, scope_tag, search_document, write_document},
};
/// `ValkeySemanticCache`: a semantic cache on valkey-search's TAG + VECTOR index. Values go
/// through the injected codec, so the response layer decides what a cached entry is.
pub struct ValkeySemanticCache<E, S, C = redis::Connection> {
connections: Arc<Connections<C>>,
embedder: E,
codec: S,
config: ValkeySemanticConfig,
index_dimension: Arc<Mutex<Option<usize>>>,
}
impl<E: Embedder, S: CacheCodec> ValkeySemanticCache<E, S> {
pub fn new(
url: &str,
embedder: E,
codec: S,
config: ValkeySemanticConfig,
) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
codec,
config,
index_dimension: Arc::new(Mutex::new(None)),
})
}
}
impl<E, S, C> ValkeySemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub fn with_connection(
connection: C,
embedder: E,
codec: S,
config: ValkeySemanticConfig,
) -> Self {
Self {
connections: Arc::new(Connections::fixed(connection)),
embedder,
codec,
config,
index_dimension: Arc::new(Mutex::new(None)),
}
}
pub fn similarity_threshold(&self) -> f64 {
self.config.similarity_threshold
}
pub fn index_name(&self) -> &str {
&self.config.index_name
}
fn index_state(&self) -> IndexState {
IndexState {
name: self.config.index_name.clone(),
prefix: format!("{}:", self.config.index_name),
dimension: Arc::clone(&self.index_dimension),
similarity_threshold: self.config.similarity_threshold,
}
}
fn decode(&self, lookup: SemanticLookup<Vec<u8>>) -> Result<SemanticLookup<S::Value>, Error> {
Ok(SemanticLookup {
value: lookup
.value
.map(|bytes| self.codec.decode(&bytes))
.transpose()?,
similarity: lookup.similarity,
})
}
}
impl<E, S, C> ValkeySemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec + Clone,
C: redis::ConnectionLike + Send + 'static,
{
/// The same index and connections behind a different embedder.
pub fn with_embedder<E2: Embedder>(&self, embedder: E2) -> ValkeySemanticCache<E2, S, C> {
ValkeySemanticCache {
connections: Arc::clone(&self.connections),
embedder,
codec: self.codec.clone(),
config: self.config.clone(),
index_dimension: Arc::clone(&self.index_dimension),
}
}
}
impl<E, S, C> BaseCache for ValkeySemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = S::Value;
type Context = SemanticCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<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 embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let response = self.codec.encode(&value)?;
let index = self.index_state();
self.connections.execute(|connection| {
write_document(
connection,
&index,
&scope_tag(key),
&prompt,
response,
embedding_bytes(&embedding),
self.get_ttl(context),
)
})
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.get_cache_with_similarity(key, context)
.map(|lookup| lookup.value)
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> Result<(), Error> {
let Some(prompt) = prompt_from_context(&context) else {
return Ok(());
};
let embedding = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let response = self.codec.encode(&value)?;
let index = self.index_state();
let scope = scope_tag(key);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
write_document(
connection,
&index,
&scope,
&prompt,
response,
embedding_bytes(&embedding),
context.ttl,
)
})
.await
}
async fn async_get_cache(
&self,
key: &str,
context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.async_get_cache_with_similarity(key, context)
.await
.map(|lookup| lookup.value)
}
}
/// Python stamps a similarity of `0.0` when there is no prompt or no document in the key's
/// scope, and the closest document's similarity even when it misses the threshold.
impl<E, S, C> SemanticCache for ValkeySemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let index = self.index_state();
let lookup = self.connections.execute(|connection| {
search_document(
connection,
&index,
&scope_tag(key),
embedding_bytes(&embedding),
)
})?;
self.decode(lookup)
}
async fn async_get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let embedding = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let index = self.index_state();
let scope = scope_tag(key);
let lookup = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
search_document(connection, &index, &scope, embedding_bytes(&embedding))
})
.await?;
self.decode(lookup)
}
}

8
litellm-rust/crates/cache-valkey-semantic/src/config.rs
View file

@ -1,8 +0,0 @@
/// `ValkeySemanticCache.DEFAULT_VALKEY_INDEX_NAME`.
pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index";
#[derive(Clone, Debug, PartialEq)]
pub struct ValkeySemanticConfig {
pub similarity_threshold: f64,
pub index_name: String,
}

117
litellm-rust/crates/cache-valkey-semantic/src/index.rs
View file

@ -1,117 +0,0 @@
use std::sync::{Arc, Mutex};
use litellm_cache::Error;
use litellm_cache_redis::connection::ConnectionRef;
use crate::search::value_text;
/// The valkey-search index one cache writes to, with the dimension it was last ensured for.
#[derive(Clone)]
pub(crate) struct IndexState {
pub(crate) name: String,
pub(crate) prefix: String,
pub(crate) dimension: Arc<Mutex<Option<usize>>>,
pub(crate) similarity_threshold: f64,
}
/// `_ensure_index_sync` / `_ensure_index_async`: create the TAG + HNSW index once per dimension,
/// and accept an existing index unless it reports a different dimension.
pub(crate) fn ensure_index(
connection: &mut ConnectionRef<'_>,
index: &IndexState,
dimension: usize,
) -> Result<(), Error> {
if index
.dimension
.lock()?
.is_some_and(|existing| existing == dimension)
{
return Ok(());
}
let create = redis::cmd("FT.CREATE")
.arg(&index.name)
.arg("ON")
.arg("HASH")
.arg("PREFIX")
.arg(1)
.arg(&index.prefix)
.arg("SCHEMA")
.arg("litellm_cache_key")
.arg("TAG")
.arg("embedding")
.arg("VECTOR")
.arg("HNSW")
.arg(6)
.arg("TYPE")
.arg("FLOAT32")
.arg("DIM")
.arg(dimension)
.arg("DISTANCE_METRIC")
.arg("COSINE")
.query::<String>(connection)
.map(|_| ());
if let Err(error) = create {
if !index_exists(&error) {
return Err(Error::unavailable(error));
}
let info = redis::cmd("FT.INFO")
.arg(&index.name)
.query::<redis::Value>(connection)
.map_err(Error::unavailable)?;
if let Some(existing) =
index_dimension_from_info(&info).filter(|existing| *existing != dimension)
{
return Err(Error::unavailable(DimensionMismatch {
existing,
dimension,
}));
}
}
*index.dimension.lock()? = Some(dimension);
Ok(())
}
fn index_exists(error: &redis::RedisError) -> bool {
error.code().is_some()
&& error
.detail()
.is_some_and(|detail| detail.to_ascii_lowercase().contains("already exists"))
}
#[derive(Debug, thiserror::Error)]
#[error("index exists with dimension {existing}, not {dimension}")]
struct DimensionMismatch {
existing: usize,
dimension: usize,
}
/// `_extract_index_dim`: flatten each attribute one level and read the value after
/// `dimensions`.
fn index_dimension_from_info(value: &redis::Value) -> Option<usize> {
let redis::Value::Array(values) = value else {
return None;
};
let attributes = values.windows(2).find_map(|pair| {
(value_text(&pair[0]).as_deref() == Some("attributes")).then_some(&pair[1])
})?;
let redis::Value::Array(fields) = attributes else {
return None;
};
fields.iter().find_map(|field| {
let redis::Value::Array(values) = field else {
return None;
};
let values = values
.iter()
.flat_map(|value| match value {
redis::Value::Array(values) => values.as_slice(),
_ => std::slice::from_ref(value),
})
.collect::<Vec<_>>();
values.windows(2).find_map(|pair| {
(value_text(pair[0]).as_deref() == Some("dimensions"))
.then(|| value_text(pair[1]).and_then(|value| value.parse().ok()))
.flatten()
})
})
}

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

@ -1,7 +0,0 @@
mod cache;
mod config;
mod index;
mod search;
pub use cache::ValkeySemanticCache;
pub use config::{DEFAULT_INDEX_NAME, ValkeySemanticConfig};

161
litellm-rust/crates/cache-valkey-semantic/src/search.rs
View file

@ -1,161 +0,0 @@
use std::time::Duration;
use litellm_cache::{Error, semantic::SemanticLookup};
use litellm_cache_redis::connection::ConnectionRef;
use sha2::{Digest, Sha256};
use uuid::Uuid;
use crate::index::{IndexState, ensure_index};
/// `_scope_tag`: valkey-search TAG fields cannot match arbitrary keys verbatim, so scopes are
/// the key's lowercase SHA-256.
pub(crate) fn scope_tag(key: &str) -> String {
let digest = Sha256::digest(key.as_bytes());
digest.iter().map(|byte| format!("{byte:02x}")).collect()
}
pub(crate) fn embedding_bytes(embedding: &[f32]) -> Vec<u8> {
embedding
.iter()
.flat_map(|value| value.to_le_bytes())
.collect()
}
/// `HSET` a fresh `<prefix><scope>:<uuid4>` document, then `EXPIRE` it when a TTL is set.
pub(crate) fn write_document(
connection: &mut ConnectionRef<'_>,
index: &IndexState,
scope: &str,
prompt: &str,
response: Vec<u8>,
vector: Vec<u8>,
ttl: Option<Duration>,
) -> Result<(), Error> {
ensure_index(connection, index, vector.len() / size_of::<f32>())?;
let document = format!("{}{scope}:{}", index.prefix, Uuid::new_v4());
let mut pipeline = redis::pipe();
pipeline
.cmd("HSET")
.arg(&document)
.arg("litellm_cache_key")
.arg(scope)
.arg("prompt")
.arg(prompt)
.arg("response")
.arg(response)
.arg("embedding")
.arg(vector)
.ignore();
if let Some(ttl) = ttl {
pipeline
.cmd("EXPIRE")
.arg(&document)
.arg(ttl.as_secs())
.ignore();
}
pipeline.query::<()>(connection).map_err(Error::unavailable)
}
/// The KNN-1 search within `scope`: the closest document's similarity, and its stored response
/// when that similarity reaches the threshold. No document reads as a similarity of `0.0`.
pub(crate) fn search_document(
connection: &mut ConnectionRef<'_>,
index: &IndexState,
scope: &str,
vector: Vec<u8>,
) -> Result<SemanticLookup<Vec<u8>>, Error> {
ensure_index(connection, index, vector.len() / size_of::<f32>())?;
let query =
format!("(@litellm_cache_key:{{{scope}}})=>[KNN 1 @embedding $vec AS vector_distance]");
let response = redis::cmd("FT.SEARCH")
.arg(&index.name)
.arg(query)
.arg("PARAMS")
.arg(2)
.arg("vec")
.arg(vector)
.arg("RETURN")
.arg(2)
.arg("response")
.arg("vector_distance")
.arg("DIALECT")
.arg(2)
.query::<redis::Value>(connection)
.map_err(Error::unavailable)?;
let Some(fields) = search_fields(response)? else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let field = |name: &str| {
fields
.iter()
.find_map(|(field, value)| (field == name).then(|| value.clone()))
.ok_or(Error::InvalidEntry(None))
};
let response = field("response")?;
let similarity = 1.0 - parse_f64(&field("vector_distance")?)?;
Ok(SemanticLookup {
value: (similarity >= index.similarity_threshold).then_some(response),
similarity: Some(similarity),
})
}
type SearchFields = Vec<(String, Vec<u8>)>;
fn search_fields(value: redis::Value) -> Result<Option<SearchFields>, Error> {
let redis::Value::Array(values) = value else {
return Err(Error::InvalidEntry(None));
};
let total = parse_i64(values.first().ok_or(Error::InvalidEntry(None))?)?;
if total <= 0 || values.len() < 3 {
return Ok(None);
}
let redis::Value::Array(fields) = &values[2] else {
return Err(Error::InvalidEntry(None));
};
let (pairs, remainder) = fields.as_chunks::<2>();
if !remainder.is_empty() {
return Err(Error::InvalidEntry(None));
}
pairs
.iter()
.map(|pair| {
Ok((
value_text(&pair[0]).ok_or(Error::InvalidEntry(None))?,
value_bytes(&pair[1])?,
))
})
.collect::<Result<Vec<_>, Error>>()
.map(Some)
}
fn parse_i64(value: &redis::Value) -> Result<i64, Error> {
value_text(value)
.ok_or(Error::InvalidEntry(None))?
.parse()
.map_err(Error::invalid_entry)
}
fn parse_f64(value: &[u8]) -> Result<f64, Error> {
std::str::from_utf8(value)
.map_err(Error::invalid_entry)?
.parse()
.map_err(Error::invalid_entry)
}
pub(crate) fn value_text(value: &redis::Value) -> Option<String> {
match value {
redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
redis::Value::SimpleString(value) => Some(value.clone()),
redis::Value::Int(value) => Some(value.to_string()),
_ => None,
}
}
fn value_bytes(value: &redis::Value) -> Result<Vec<u8>, Error> {
match value {
redis::Value::BulkString(bytes) => Ok(bytes.clone()),
redis::Value::SimpleString(value) => Ok(value.as_bytes().to_vec()),
redis::Value::Int(value) => Ok(value.to_string().into_bytes()),
_ => Err(Error::InvalidEntry(None)),
}
}

430
litellm-rust/crates/cache-valkey-semantic/tests/cache.rs
View file

@ -1,430 +0,0 @@
mod support;
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::{
BaseCache, Error, JsonCodec, SemanticCacheContext,
semantic::{PreparedEmbedding, SemanticCache, SemanticLookup},
};
use litellm_cache_valkey_semantic::{
DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig,
};
use redis_test::MockRedisConnection;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::{EmbedCalls, FakeEmbedder, RecordingConnection};
type Requests = Arc<Mutex<Vec<Vec<u8>>>>;
type RecordingCache = ValkeySemanticCache<FakeEmbedder, JsonCodec<Value>, RecordingConnection>;
const KEY_SCOPE: &str = "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683";
struct Recording {
cache: RecordingCache,
requests: Requests,
calls: EmbedCalls,
}
impl Recording {
fn text(&self) -> String {
self.requests
.lock()
.unwrap()
.iter()
.map(|request| String::from_utf8_lossy(request).into_owned())
.collect::<Vec<_>>()
.join("\n")
}
}
fn config() -> ValkeySemanticConfig {
ValkeySemanticConfig {
similarity_threshold: 0.8,
index_name: "test".into(),
}
}
fn recording(replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>) -> Recording {
let connection = RecordingConnection::new(replies);
let requests = connection.requests();
let embedder = FakeEmbedder::new(&[]);
let calls = Arc::clone(&embedder.calls);
Recording {
cache: ValkeySemanticCache::with_connection(
connection,
embedder,
JsonCodec::new(),
config(),
),
requests,
calls,
}
}
#[fixture]
fn entry() -> Value {
json!({"timestamp": 1.0, "response": {"answer": "ok"}})
}
#[fixture]
fn context() -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(json!([{"role": "user", "content": "hello"}])),
metadata: Some(json!({"source": "test"})),
..Default::default()
}
}
fn ok() -> redis::RedisResult<redis::Value> {
Ok(redis::Value::SimpleString("OK".into()))
}
fn already_exists() -> redis::RedisResult<redis::Value> {
Err(support::server_error("Index already exists"))
}
fn bulk(value: &[u8]) -> redis::Value {
redis::Value::BulkString(value.to_vec())
}
fn search_hit(response: &[u8], distance: &str) -> redis::Value {
redis::Value::Array(vec![
redis::Value::Int(1),
bulk(b"test:document"),
redis::Value::Array(vec![
bulk(b"response"),
bulk(response),
bulk(b"vector_distance"),
bulk(distance.as_bytes()),
]),
])
}
fn encoded(value: &Value) -> Vec<u8> {
serde_json::to_vec(value).unwrap()
}
/// `FT.INFO` with the vector field's dimension nested one level down, as valkey-search reports.
fn nested_dimension_info(dimension: i64) -> redis::Value {
redis::Value::Array(vec![
redis::Value::SimpleString("attributes".into()),
redis::Value::Array(vec![redis::Value::Array(vec![
redis::Value::SimpleString("embedding".into()),
redis::Value::Array(vec![
redis::Value::SimpleString("dimensions".into()),
redis::Value::Int(dimension),
]),
])]),
])
}
/// `FT.INFO` with `dimensions` as a sibling string of the identifier.
fn flat_dimension_info(dimension: &str) -> redis::Value {
redis::Value::Array(vec![
redis::Value::SimpleString("attributes".into()),
redis::Value::Array(vec![redis::Value::Array(vec![
redis::Value::SimpleString("identifier".into()),
redis::Value::SimpleString("embedding".into()),
redis::Value::Array(vec![
redis::Value::SimpleString("dimensions".into()),
redis::Value::SimpleString(dimension.into()),
]),
])]),
])
}
#[rstest]
#[case::string_content(json!([{"content": "hello"}]), None, Some("hello"))]
#[case::text_parts(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))]
#[case::non_object_parts_are_skipped(json!([{"content": ["raw", {"text": "hello"}]}]), None, Some("hello"))]
#[case::search_results(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))]
#[case::responses_string_input(json!([]), Some(json!(" hello ")), Some("hello"))]
#[case::responses_item_input(json!([]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))]
#[case::blank_input(json!([]), Some(json!(" ")), None)]
fn prompt_shapes_follow_redis_semantic_extraction(
#[case] messages: Value,
#[case] input: Option<Value>,
#[case] expected: Option<&str>,
) {
let recording = recording([ok(), Ok(redis::Value::Array(vec![redis::Value::Int(0)]))]);
let context = SemanticCacheContext {
messages: Some(messages),
input,
..Default::default()
};
assert_eq!(recording.cache.get_cache("key", &context).unwrap(), None);
let prompts = recording
.calls
.lock()
.unwrap()
.iter()
.map(|(prompt, _)| prompt.clone())
.collect::<Vec<_>>();
assert_eq!(
prompts,
expected.into_iter().map(str::to_owned).collect::<Vec<_>>()
);
assert_eq!(
recording.requests.lock().unwrap().is_empty(),
expected.is_none()
);
}
#[rstest]
#[case::key("key", KEY_SCOPE)]
#[case::empty_key("", "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855")]
fn documents_are_scoped_by_the_keys_sha256(
#[case] key: &str,
#[case] scope: &str,
entry: Value,
context: SemanticCacheContext,
) {
let recording = recording([ok()]);
recording.cache.set_cache(key, entry, &context).unwrap();
let text = recording.text();
assert!(text.contains(&format!("test:{scope}:")));
assert!(text.contains(&format!("litellm_cache_key\r\n$64\r\n{scope}\r\n")));
}
#[rstest]
#[case::no_ttl(None, None)]
#[case::whole_seconds(Some(Duration::from_secs(5)), Some("5"))]
#[case::fractional_seconds_truncate(Some(Duration::from_millis(1900)), Some("1"))]
fn set_writes_hset_and_expires_only_with_a_ttl(
#[case] ttl: Option<Duration>,
#[case] expire: Option<&str>,
entry: Value,
context: SemanticCacheContext,
) {
let recording = recording([ok()]);
recording
.cache
.set_cache(
"key",
entry,
&SemanticCacheContext {
ttl,
..context.clone()
},
)
.unwrap();
let text = recording.text();
assert!(text.contains("FT.CREATE"));
assert!(text.contains("HSET"));
let expire_seconds = text
.split_once("EXPIRE\r\n")
.and_then(|(_, rest)| rest.split("\r\n").nth(3));
assert_eq!(expire_seconds, expire);
assert_eq!(
*recording.calls.lock().unwrap(),
vec![("hello".to_owned(), context.metadata)]
);
}
#[rstest]
fn second_set_skips_create_after_dimension_is_cached(entry: Value, context: SemanticCacheContext) {
let recording = recording([ok()]);
recording
.cache
.set_cache("key", entry.clone(), &context)
.unwrap();
recording.cache.set_cache("key", entry, &context).unwrap();
let text = recording.text();
assert_eq!(text.matches("FT.CREATE").count(), 1);
assert_eq!(text.matches("HSET").count(), 2);
}
#[rstest]
#[case::nested_matching(nested_dimension_info(3), Ok(()))]
#[case::nested_mismatch(
nested_dimension_info(2),
Err("cache is unavailable: index exists with dimension 2, not 3")
)]
#[case::flat_matching(flat_dimension_info("3"), Ok(()))]
#[case::flat_mismatch(
flat_dimension_info("2"),
Err("cache is unavailable: index exists with dimension 2, not 3")
)]
#[case::unreported_dimension_is_accepted(redis::Value::Array(vec![]), Ok(()))]
fn existing_index_dimension_must_match_embedding(
#[case] info: redis::Value,
#[case] expected: Result<(), &str>,
entry: Value,
context: SemanticCacheContext,
) {
let recording = recording([already_exists(), Ok(info)]);
let result = recording.cache.set_cache("key", entry, &context);
assert_eq!(
result.map_err(|error| litellm_error::Report::new(&error).to_string()),
expected.map_err(str::to_owned)
);
}
#[rstest]
#[case::create_failure(Err(redis::RedisError::from((redis::ErrorKind::Io, "boom"))))]
fn index_creation_failures_are_unavailable(
#[case] reply: redis::RedisResult<redis::Value>,
entry: Value,
context: SemanticCacheContext,
) {
let recording = recording([reply]);
assert!(matches!(
recording.cache.set_cache("key", entry, &context),
Err(Error::Unavailable(Some(_)))
));
}
#[rstest]
#[case::within_threshold(search_hit(&encoded(&entry()), "0.1"), Ok(Some(entry())))]
#[case::at_threshold(search_hit(&encoded(&entry()), "0.2"), Ok(Some(entry())))]
#[case::beyond_threshold(search_hit(&encoded(&entry()), "0.5"), Ok(None))]
#[case::zero_documents(redis::Value::Array(vec![redis::Value::Int(0)]), Ok(None))]
#[case::missing_response(
redis::Value::Array(vec![
redis::Value::Int(1),
bulk(b"document"),
redis::Value::Array(vec![bulk(b"vector_distance"), bulk(b"0.1")]),
]),
Err(Error::InvalidEntry(None))
)]
#[case::unparsable_distance(search_hit(b"not-json", "abc"), Err(Error::InvalidEntry(None)))]
#[case::undecodable_response(search_hit(b"not-json", "0.1"), Err(Error::InvalidEntry(None)))]
fn get_applies_threshold_and_decodes_entry(
#[case] reply: redis::Value,
#[case] expected: Result<Option<Value>, Error>,
context: SemanticCacheContext,
) {
let recording = recording([ok(), Ok(reply)]);
let kind = |result: Result<Option<Value>, Error>| {
result.map_err(|error| std::mem::discriminant(&error))
};
assert_eq!(
kind(recording.cache.get_cache("key", &context)),
kind(expected)
);
}
#[rstest]
#[case::hit(context(), Some(search_hit(&encoded(&entry()), "0.1")), Some(entry()), Some(1.0 - 0.1))]
#[case::below_threshold(context(), Some(search_hit(&encoded(&entry()), "0.5")), None, Some(1.0 - 0.5))]
#[case::no_results(context(), Some(redis::Value::Array(vec![redis::Value::Int(0)])), None, Some(0.0))]
#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))]
#[tokio::test]
async fn lookup_reports_python_semantic_similarity(
#[case] context: SemanticCacheContext,
#[case] reply: Option<redis::Value>,
#[case] value: Option<Value>,
#[case] similarity: Option<f64>,
#[values(false, true)] use_async: bool,
) {
let searched = reply.is_some();
let recording = recording(reply.map_or_else(Vec::new, |reply| vec![ok(), Ok(reply)]));
let lookup = if use_async {
recording
.cache
.async_get_cache_with_similarity("key", &context)
.await
} else {
recording.cache.get_cache_with_similarity("key", &context)
};
assert_eq!(lookup, Ok(SemanticLookup { value, similarity }));
assert_eq!(recording.text().contains("FT.SEARCH"), searched);
}
#[rstest]
#[tokio::test]
async fn missing_prompt_does_not_touch_valkey(entry: Value) {
let cache = ValkeySemanticCache::with_connection(
MockRedisConnection::new([]).assert_all_commands_consumed(),
FakeEmbedder::new(&[]),
JsonCodec::new(),
config(),
);
let context = SemanticCacheContext::default();
cache.set_cache("key", entry.clone(), &context).unwrap();
assert_eq!(cache.get_cache("key", &context).unwrap(), None);
cache
.async_set_cache("key", entry, context.clone())
.await
.unwrap();
assert_eq!(cache.async_get_cache("key", &context).await.unwrap(), None);
assert_eq!(cache.get_ttl(&context), None);
}
#[rstest]
fn with_embedder_shares_index_state_and_connections(entry: Value, context: SemanticCacheContext) {
let recording = recording([ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]);
recording
.cache
.set_cache("key", entry.clone(), &context)
.unwrap();
let prepared = recording
.cache
.with_embedder(PreparedEmbedding(vec![0.1, 0.2, 0.3]));
assert_eq!(prepared.get_cache("key", &context).unwrap(), Some(entry));
assert_eq!(recording.text().matches("FT.CREATE").count(), 1);
}
#[rstest]
fn accessors_report_the_config() {
let recording = recording([]);
assert_eq!(recording.cache.index_name(), "test");
assert_eq!(recording.cache.similarity_threshold(), 0.8);
assert_eq!(DEFAULT_INDEX_NAME, "litellm_semantic_cache_index");
}
#[rstest]
#[tokio::test]
async fn async_set_and_get_use_shared_document_helpers(
entry: Value,
context: SemanticCacheContext,
) {
let recording = recording([ok(), ok(), ok(), Ok(search_hit(&encoded(&entry), "0.1"))]);
let context = SemanticCacheContext {
ttl: Some(Duration::from_millis(1900)),
..context
};
recording
.cache
.async_set_cache("key", entry.clone(), context.clone())
.await
.unwrap();
assert_eq!(
recording
.cache
.async_get_cache("key", &context)
.await
.unwrap(),
Some(entry)
);
let text = recording.text();
assert!(text.contains("FT.CREATE"));
assert!(text.contains("HSET"));
assert!(text.contains("EXPIRE"));
assert_eq!(
*recording.calls.lock().unwrap(),
vec![("hello".to_owned(), context.metadata.clone()); 2]
);
}

62
litellm-rust/crates/cache-valkey-semantic/tests/contract.rs
View file

@ -1,62 +0,0 @@
//! `overwrite_replaces` does not apply: like Python, every write is a new `<scope>:<uuid4>`
//! document, so a second write with the same prompt adds a tie instead of replacing the first.
mod support;
use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
use litellm_cache_testing as contract;
use litellm_cache_valkey_semantic::{
DEFAULT_INDEX_NAME, ValkeySemanticCache, ValkeySemanticConfig,
};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::FakeSearch;
type Cache = ValkeySemanticCache<PreparedEmbedding, JsonCodec<Value>, FakeSearch>;
const PREFIX: &str = "contract:";
#[fixture]
fn cache() -> Cache {
ValkeySemanticCache::with_connection(
FakeSearch::default(),
PreparedEmbedding(vec![0.6, 0.8]),
JsonCodec::new(),
ValkeySemanticConfig {
similarity_threshold: 0.9,
index_name: DEFAULT_INDEX_NAME.into(),
},
)
}
#[fixture]
fn context() -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(json!([{"role": "user", "content": "contract prompt"}])),
..Default::default()
}
}
#[rstest]
#[tokio::test]
async fn hit_and_miss(cache: Cache, context: SemanticCacheContext) {
contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await;
}
#[rstest]
#[tokio::test]
async fn sync_async_equivalence(cache: Cache, context: SemanticCacheContext) {
contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await;
}
#[rstest]
#[tokio::test]
async fn pipeline_writes_every_entry(cache: Cache, context: SemanticCacheContext) {
contract::pipeline_writes_every_entry(
&cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
}

352
litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs
View file

@ -1,352 +0,0 @@
#![allow(dead_code)]
use std::{
collections::{BTreeMap, HashMap, VecDeque},
sync::{Arc, Mutex},
};
use litellm_cache::{Error, semantic::Embedder};
use serde_json::Value;
pub type EmbedCalls = Arc<Mutex<Vec<(String, Option<Value>)>>>;
/// Embeds known prompts to fixed vectors, anything else to `[0.1, 0.2, 0.3]`, and records every
/// prompt with its metadata.
pub struct FakeEmbedder {
vectors: HashMap<String, Vec<f32>>,
pub calls: EmbedCalls,
}
impl FakeEmbedder {
pub fn new(vectors: &[(&str, &[f32])]) -> Self {
Self {
vectors: vectors
.iter()
.map(|(prompt, vector)| ((*prompt).to_owned(), vector.to_vec()))
.collect(),
calls: EmbedCalls::default(),
}
}
}
impl Embedder for FakeEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
self.calls
.lock()
.unwrap()
.push((prompt.to_owned(), metadata.cloned()));
Ok(self
.vectors
.get(prompt)
.cloned()
.unwrap_or_else(|| vec![0.1, 0.2, 0.3]))
}
async fn async_embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
self.embed(prompt, metadata)
}
}
struct FakeIndex {
prefix: Vec<u8>,
dims: usize,
vector_field: String,
}
#[derive(Default)]
struct SearchState {
indexes: HashMap<String, FakeIndex>,
hashes: BTreeMap<Vec<u8>, BTreeMap<String, Vec<u8>>>,
}
/// An in-memory valkey-search speaking the `FT.*`, `HSET` and `EXPIRE` subset the semantic cache
/// sends, with exact cosine KNN over the hashes under an index prefix.
#[derive(Clone, Default)]
pub struct FakeSearch {
state: Arc<Mutex<SearchState>>,
}
impl FakeSearch {
fn run(&self, args: Vec<Vec<u8>>) -> redis::RedisResult<redis::Value> {
let mut state = self.state.lock().unwrap();
let text = |index: usize| String::from_utf8_lossy(&args[index]).into_owned();
match text(0).to_uppercase().as_str() {
"FT.CREATE" => {
let name = text(1);
if state.indexes.contains_key(&name) {
return Err(server_error("Index already exists"));
}
let position = |token: &str| args.iter().position(|arg| arg == token.as_bytes());
let prefix = args[position("PREFIX").unwrap() + 2].clone();
let dims = text(position("DIM").unwrap() + 1).parse().unwrap();
let vector_field = text(position("VECTOR").unwrap() - 1);
state.indexes.insert(
name,
FakeIndex {
prefix,
dims,
vector_field,
},
);
Ok(redis::Value::Okay)
}
"FT.INFO" => {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| server_error("Unknown index name"))?;
Ok(index_info(index))
}
"FT.DROPINDEX" => {
state.indexes.remove(&text(1));
Ok(redis::Value::Okay)
}
"HSET" => {
let hash = state.hashes.entry(args[1].clone()).or_default();
for pair in args[2..].chunks(2) {
hash.insert(
String::from_utf8_lossy(&pair[0]).into_owned(),
pair[1].clone(),
);
}
Ok(redis::Value::Int(((args.len() - 2) / 2) as i64))
}
"EXPIRE" => Ok(redis::Value::Int(i64::from(
state.hashes.contains_key(&args[1]),
))),
"FT.SEARCH" => {
let index = state
.indexes
.get(&text(1))
.ok_or_else(|| server_error("no such index"))?;
let query = text(2);
let tag = query_tag(&query);
let params = args.iter().position(|arg| arg == b"PARAMS").unwrap();
let vector = floats(&args[params + 3]);
let best = state
.hashes
.iter()
.filter(|(key, _)| key.starts_with(&index.prefix))
.filter(|(_, fields)| {
fields.get("litellm_cache_key").map(Vec::as_slice) == Some(tag.as_bytes())
})
.filter_map(|(key, fields)| {
let stored = floats(fields.get(&index.vector_field)?);
(stored.len() == index.dims)
.then(|| (key, fields, 1.0 - cosine(&vector, &stored)))
})
.min_by(|left, right| left.2.total_cmp(&right.2));
let Some((key, fields, distance)) = best else {
return Ok(redis::Value::Array(vec![redis::Value::Int(0)]));
};
let mut reply = fields
.iter()
.filter(|(name, _)| **name != index.vector_field)
.flat_map(|(name, value)| [bulk(name.as_bytes()), bulk(value)])
.collect::<Vec<_>>();
reply.extend([
bulk(b"vector_distance"),
bulk(distance.to_string().as_bytes()),
]);
Ok(redis::Value::Array(vec![
redis::Value::Int(1),
bulk(key),
redis::Value::Array(reply),
]))
}
"PING" => Ok(redis::Value::SimpleString("PONG".into())),
_ => Err(server_error("unsupported command")),
}
}
}
impl redis::ConnectionLike for FakeSearch {
fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
let mut commands = parse_commands(command);
self.run(commands.remove(0))
}
fn req_packed_commands(
&mut self,
commands: &[u8],
offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
let replies = parse_commands(commands)
.into_iter()
.map(|args| self.run(args))
.collect::<redis::RedisResult<Vec<_>>>()?;
Ok(replies.into_iter().skip(offset).take(count).collect())
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}
pub fn server_error(reply: &str) -> redis::RedisError {
redis::parse_redis_value(format!("-{reply}\r\n").as_bytes())
.expect("error reply")
.extract_error()
.expect_err("server error")
}
fn bulk(bytes: &[u8]) -> redis::Value {
redis::Value::BulkString(bytes.to_vec())
}
fn index_info(index: &FakeIndex) -> redis::Value {
redis::Value::Array(vec![
bulk(b"index_name"),
bulk(b"fake"),
bulk(b"attributes"),
redis::Value::Array(vec![
redis::Value::Array(vec![
bulk(b"identifier"),
bulk(b"litellm_cache_key"),
bulk(b"type"),
bulk(b"TAG"),
]),
redis::Value::Array(vec![
bulk(b"identifier"),
bulk(index.vector_field.as_bytes()),
bulk(b"type"),
bulk(b"VECTOR"),
bulk(b"index"),
redis::Value::Array(vec![
bulk(b"dimensions"),
redis::Value::Int(index.dims as i64),
]),
]),
]),
])
}
/// The tag inside `@litellm_cache_key:{...}`, with query escapes removed.
fn query_tag(query: &str) -> String {
let start = query.find("@litellm_cache_key:{").unwrap() + "@litellm_cache_key:{".len();
let mut tag = String::new();
let mut characters = query[start..].chars();
while let Some(character) = characters.next() {
match character {
'\\' => tag.extend(characters.next()),
'}' => break,
character => tag.push(character),
}
}
tag
}
fn floats(bytes: &[u8]) -> Vec<f32> {
bytes
.as_chunks::<4>()
.0
.iter()
.map(|chunk| f32::from_le_bytes(*chunk))
.collect()
}
fn cosine(left: &[f32], right: &[f32]) -> f64 {
let dot = left
.iter()
.zip(right)
.map(|(left, right)| f64::from(*left) * f64::from(*right))
.sum::<f64>();
let norm = |vector: &[f32]| {
vector
.iter()
.map(|value| f64::from(*value).powi(2))
.sum::<f64>()
.sqrt()
};
dot / (norm(left) * norm(right))
}
/// Splits a packed RESP request into each command's arguments.
fn parse_commands(mut bytes: &[u8]) -> Vec<Vec<Vec<u8>>> {
let line = |bytes: &mut &[u8]| {
let end = bytes
.windows(2)
.position(|window| window == b"\r\n")
.unwrap();
let text = String::from_utf8(bytes[1..end].to_vec()).unwrap();
*bytes = &bytes[end + 2..];
text.parse::<usize>().unwrap()
};
let mut commands = Vec::new();
while !bytes.is_empty() {
let count = line(&mut bytes);
let mut args = Vec::with_capacity(count);
for _ in 0..count {
let length = line(&mut bytes);
args.push(bytes[..length].to_vec());
bytes = &bytes[length + 2..];
}
commands.push(args);
}
commands
}
/// Records every packed request and answers from a script, `OK` once the script runs out.
pub struct RecordingConnection {
requests: Arc<Mutex<Vec<Vec<u8>>>>,
replies: Mutex<VecDeque<redis::RedisResult<redis::Value>>>,
}
impl RecordingConnection {
pub fn new(replies: impl IntoIterator<Item = redis::RedisResult<redis::Value>>) -> Self {
Self {
requests: Arc::default(),
replies: Mutex::new(replies.into_iter().collect()),
}
}
pub fn requests(&self) -> Arc<Mutex<Vec<Vec<u8>>>> {
Arc::clone(&self.requests)
}
fn reply(&self) -> redis::RedisResult<redis::Value> {
self.replies
.lock()
.unwrap()
.pop_front()
.unwrap_or_else(|| Ok(redis::Value::SimpleString("OK".into())))
}
}
impl redis::ConnectionLike for RecordingConnection {
fn req_packed_command(&mut self, command: &[u8]) -> redis::RedisResult<redis::Value> {
self.requests.lock().unwrap().push(command.to_vec());
self.reply()
}
fn req_packed_commands(
&mut self,
command: &[u8],
_offset: usize,
count: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
self.requests.lock().unwrap().push(command.to_vec());
(0..count).map(|_| self.reply()).collect()
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}

2
litellm-rust/crates/cost/src/price.rs
View file

@ -151,6 +151,8 @@ pub enum Uplift {
Residency(Residency),
/// Vertex non-global endpoint, from `regional_endpoint_uplift_multiplier`
NonGlobalEndpoint,
/// OpenAI FedRAMP endpoint (gov.api.openai.com), from `fedramp_uplift_multiplier`
FedrampEndpoint,
/// Anthropic `speed: fast`, from `provider_specific_entry.fast`
Speed(Speed),
/// Anthropic and Bedrock geo, from `provider_specific_entry.<geo>`, lowercased as on main

10
litellm-rust/crates/model-catalog/src/lower.rs
View file

@ -104,6 +104,7 @@ impl From<&EntryWire> for Adjustments {
(Uplift::NonGlobalEndpoint, wire.regional_endpoint_uplift),
(Uplift::Residency(Residency::Eu), wire.residency_uplift_eu),
(Uplift::Residency(Residency::Us), wire.residency_uplift_us),
(Uplift::FedrampEndpoint, wire.fedramp_uplift),
]
.into_iter()
.filter_map(|(uplift, factor)| Some((uplift, factor?))),
@ -746,7 +747,8 @@ mod tests {
fn uplifts_are_adjustments_not_prices() {
let wire = wire(json!({
"regional_endpoint_uplift_multiplier": 1.1,
"regional_processing_uplift_multiplier_eu": "1.2"
"regional_processing_uplift_multiplier_eu": "1.2",
"fedramp_uplift_multiplier": 1.1
}));
let adjustments = Adjustments::from(&wire);
@ -764,6 +766,12 @@ mod tests {
Some(1.2)
);
assert_eq!(adjustments.get(&Uplift::Residency(Residency::Us)), None);
assert_eq!(
adjustments
.get(&Uplift::FedrampEndpoint)
.map(Multiplier::get),
Some(1.1)
);
}
#[rstest]

1
litellm-rust/crates/model-catalog/src/model_info.rs
View file

@ -96,6 +96,7 @@ pub struct ModelInfoFields {
pub output_cost_per_token_ultrafast: Option<f64>,
pub regional_processing_uplift_multiplier_eu: Option<f64>,
pub regional_processing_uplift_multiplier_us: Option<f64>,
pub fedramp_uplift_multiplier: Option<f64>,
pub regional_endpoint_uplift_multiplier: Option<f64>,
pub output_cost_per_audio_token: Option<f64>,
pub output_cost_per_character: Option<f64>,

2
litellm-rust/crates/model-catalog/src/schema.rs
View file

@ -29,6 +29,8 @@ struct PriceBlocks {
regional_processing_uplift_multiplier_eu: Option<Multiplier>,
/// Multiplier applied to all token costs for US data residency (e.g. 1.10 = +10%).
regional_processing_uplift_multiplier_us: Option<Multiplier>,
/// Multiplier applied to all token costs when served from the OpenAI FedRAMP endpoint (e.g. 1.10 = +10%).
fedramp_uplift_multiplier: Option<Multiplier>,
/// USD cost per web search query, keyed by search context size.
search_context_cost_per_query: Option<SearchContextCostPerQuery>,
/// Context-length or result-count tiered rates; each tier's costs apply within its range.

8
litellm-rust/crates/model-catalog/src/wire/attributes.rs
View file

@ -93,6 +93,10 @@ attributes! {
search_free_requests: SearchFreeRequests,
/// URL of the provider pricing/model page this entry was taken from.
source: String,
/// Model the provider recommends migrating to before `deprecation_date`, as a catalog key when one exists.
successor_model: String,
/// URL of the provider page or feed `deprecation_date` was read from.
deprecation_source: String,
/// Audio container formats the model can return.
supported_audio_formats: Vec<AudioFormat>,
/// OpenAI-style API routes this model can be called through, e.g. /v1/chat/completions.
@ -303,6 +307,10 @@ mod tests {
#[case::known_mode("mode", json!("chat"), true)]
#[case::unknown_mode("mode", json!("chatt"), false)]
#[case::negative_dbu("input_dbu_cost_per_token", json!(-1), false)]
#[case::successor_model("successor_model", json!("gpt-6-luna"), true)]
#[case::successor_model_not_a_name("successor_model", json!(["gpt-6-luna"]), false)]
#[case::deprecation_source("deprecation_source", json!("https://developers.openai.com/api/docs/deprecations"), true)]
#[case::deprecation_source_not_a_url_string("deprecation_source", json!(7), false)]
#[case::thinking_on_by_default("thinking_on_by_default", json!(true), true)]
#[case::thinking_on_by_default_not_a_flag("thinking_on_by_default", json!("yes"), false)]
#[case::reasoning_contexts("reasoning_contexts", json!(["reasoning", "non_reasoning"]), true)]

2
litellm-rust/crates/model-catalog/src/wire/entry.rs
View file

@ -24,6 +24,7 @@ pub struct EntryWire {
pub regional_endpoint_uplift: Option<Multiplier>,
pub residency_uplift_eu: Option<Multiplier>,
pub residency_uplift_us: Option<Multiplier>,
pub fedramp_uplift: Option<Multiplier>,
pub unknown: Fields,
pub problems: Vec<Error>,
}
@ -57,6 +58,7 @@ impl EntryWire {
"regional_processing_uplift_multiplier_us" => {
self.residency_uplift_us = reader.value(path(), value)
}
"fedramp_uplift_multiplier" => self.fedramp_uplift = reader.value(path(), value),
_ => self.read_open_key(reader, key, value),
}
}

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

@ -29,8 +29,6 @@ bytes.workspace = true
futures-util.workspace = true
litellm-cache.workspace = true
litellm-error.workspace = true
litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
litellm-cache-response.workspace = true
serde.workspace = true
litellm-auth.workspace = true
@ -56,7 +54,6 @@ pyo3.workspace = true
prost.workspace = true
pyo3-async-runtimes.workspace = true
reqwest.workspace = true
redis = { workspace = true, features = ["tls-rustls"] }
serde_json.workspace = true
strum.workspace = true
veil.workspace = true

6
litellm-rust/crates/python-bridge/src/cache/AGENTS.md vendored
View file

@ -1,7 +1,7 @@
# Cache boundary
This folder owns how Rust inference reaches the selected cache: global cache selection, route admission, delegation to a Python cache and the experimental V2 native handles. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
This folder owns how Rust inference reaches the selected Python cache. Cache algorithms, storage protocols and response-cache semantics belong to their cache crates
`mod.rs` exposes the cache boundary to routes and module registration. Adapter directories remain private. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache, and owns route admission and inference protocol composition for both adapters
`mod.rs` exposes the cache boundary to routes. `selection.rs` asks `litellm.rust_bridge.host.cache` for the live configured cache and owns route admission and inference protocol composition
`python/` delegates operations to the selected Python cache without discovering configuration. `native/` owns the experimental V2 native cache handles. Neither adapter depends on shared selection or the other adapter. Shared composition depends on the adapters, and routes use only the parent module's exports
`python/` delegates operations to the selected Python cache without discovering configuration. Routes use only the parent module's exports

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

@ -1,32 +1,15 @@
mod native;
mod python;
mod selection;
pub(crate) use native::NativeCacheHandle;
pub(crate) use python::{CacheCall, PythonCache};
pub(crate) use selection::{Cached, Selection, configure, configured_native};
pub(crate) use selection::{Cached, Selection, configure, ensure_cache_disabled};
use litellm_cache::Error;
use pyo3::{
exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError},
prelude::*,
};
use pyo3::{exceptions::PyRuntimeError, prelude::*};
/// A cache call the bridge rejects itself, before or after the cache ran.
#[derive(Debug, thiserror::Error)]
pub(crate) enum HandleError {
#[error("recreate the v2 cache after fork")]
Forked,
#[error("cache limits must be positive")]
Limits,
#[error("namespace and a positive cache limit are required")]
RedisLimits,
#[error("cache durations must be finite and positive")]
Duration(#[source] Option<std::time::TryFromFloatSecsError>),
#[error("unsupported v2 cache control: {0}")]
Control(String),
#[error("native inference needs a v2 cache, the catalog keeps Python caches on Python")]
NotNative,
#[error("this inference route does not support configured caches")]
Unsupported,
#[error("cache operation without configured cache")]
Unconfigured,
#[error("cache arguments unavailable")]
@ -39,66 +22,10 @@ impl From<HandleError> for PyErr {
fn from(error: HandleError) -> Self {
let message = litellm_error::Report::new(&error).to_string();
match error {
HandleError::Limits
| HandleError::RedisLimits
| HandleError::Duration(_)
| HandleError::Control(_) => PyValueError::new_err(message),
HandleError::Forked
| HandleError::NotNative
HandleError::Unsupported
| HandleError::Unconfigured
| HandleError::MissingArguments
| HandleError::NothingPending => PyRuntimeError::new_err(message),
}
}
}
fn cache_error(error: Error) -> PyErr {
let message = litellm_error::Report::new(&error).to_string();
match error {
Error::InvalidEntry(_) => PyValueError::new_err(message),
Error::UnsupportedOperation => PyNotImplementedError::new_err(message),
Error::Unavailable(_) | Error::UnscopedFlush | Error::MissingPrompt => {
PyRuntimeError::new_err(message)
}
}
}
#[cfg(test)]
mod tests {
use pyo3::{
exceptions::{PyRuntimeError, PyValueError},
prelude::*,
};
use rstest::rstest;
use super::{Error, cache_error};
fn refused() -> std::io::Error {
std::io::Error::other("connection refused")
}
#[rstest]
#[case::unavailable(
Error::unavailable(refused()),
false,
"cache is unavailable: connection refused"
)]
#[case::invalid_entry(
Error::invalid_entry(refused()),
true,
"invalid cache entry: connection refused"
)]
fn python_sees_the_class_and_the_whole_chain(
#[case] error: Error,
#[case] caller_fault: bool,
#[case] message: &str,
) {
Python::initialize();
Python::attach(|py| {
let failure = cache_error(error);
assert_eq!(failure.is_instance_of::<PyValueError>(py), caller_fault);
assert_eq!(failure.is_instance_of::<PyRuntimeError>(py), !caller_fault);
assert_eq!(failure.value(py).to_string(), message);
});
}
}

9
litellm-rust/crates/python-bridge/src/cache/native/AGENTS.md vendored
View file

@ -1,9 +0,0 @@
# Native cache bindings
This directory exposes the experimental V2 native cache handles (`NativeCacheHandle`) to Python. Cache algorithms and storage protocols remain in their cache crates
Accept the cache object selected by the parent module. Do not read global `litellm.cache`, decide route admission, or select the Python cache adapter here
Keep Python-facing cache classes and method signatures stable when reorganizing modules. Native internals stay private to this directory unless the shared cache boundary or Python module registration needs them
Verify changes with the V2 cache tests using a freshly built extension. Test cache behavior, not module paths or file structure

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

@ -1,3 +0,0 @@
pub(super) mod v2;
pub(crate) use v2::NativeCacheHandle;

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

@ -1,336 +0,0 @@
use crate::cache::{HandleError, cache_error};
use std::{
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{DeleteCache, DisconnectCache, PingCache};
use litellm_host_python::{from_py, release_gil, to_py};
use serde_json::Value;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
use litellm_cache_response::{
CacheEntry, CacheKeyInput, ExactResponseCache, ResponseCache, ResponseCacheCodec,
ResponseCacheConfig, ResponseCacheRequest, ResponseCacheService,
};
use pyo3::{prelude::*, types::PyDict};
#[pyclass(
frozen,
name = "NativeCacheHandle",
module = "litellm.rust_bridge._native"
)]
pub(crate) struct NativeCacheHandle {
service: Arc<dyn ResponseCacheService>,
backend: Arc<dyn ExactResponseCache>,
storage: Storage,
pid: u32,
}
#[derive(Clone)]
enum Storage {
Memory(Arc<InMemoryCache<CacheEntry>>),
Redis(Arc<RedisCache<ResponseCacheCodec>>),
}
impl NativeCacheHandle {
fn check_process(&self) -> PyResult<()> {
if self.pid != std::process::id() {
return Err(HandleError::Forked.into());
}
Ok(())
}
}
fn request(key: String, ttl: Option<f64>) -> PyResult<ResponseCacheRequest> {
let mut request: ResponseCacheRequest = ResponseCacheRequest::new(CacheKeyInput {
preset: Some(key),
..Default::default()
});
request.context.ttl = ttl.map(duration).transpose()?;
Ok(request)
}
#[pymethods]
impl NativeCacheHandle {
#[staticmethod]
#[pyo3(signature = (*, ttl=600.0, capacity=200, max_entry_bytes=4194304))]
fn memory(ttl: f64, capacity: usize, max_entry_bytes: usize) -> PyResult<Self> {
let ttl = duration(ttl)?;
if capacity == 0 || max_entry_bytes == 0 {
return Err(HandleError::Limits.into());
}
let storage = Arc::new(InMemoryCache::new(Some(capacity), Some(ttl)));
let backend = Arc::new(ResponseCache::new(storage.clone()).with_config(
ResponseCacheConfig {
namespace: "sdk".into(),
max_entry_bytes,
},
));
Ok(Self {
service: backend.clone(),
backend,
storage: Storage::Memory(storage),
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (url, *, namespace, ttl=600.0, max_entry_bytes=4194304))]
fn redis(
py: Python<'_>,
url: &str,
namespace: String,
ttl: f64,
max_entry_bytes: usize,
) -> PyResult<Self> {
let ttl = duration(ttl)?;
if namespace.is_empty() || max_entry_bytes == 0 {
return Err(HandleError::RedisLimits.into());
}
let storage = Arc::new(
release_gil(py, || {
RedisCache::connect(
url,
&RedisTopology::Standalone,
Some(ttl),
ResponseCacheCodec,
)
})
.map_err(cache_error)?
.with_namespace(Some(namespace.clone())),
);
let backend = Arc::new(ResponseCache::new(storage.clone()).with_config(
ResponseCacheConfig {
namespace,
max_entry_bytes,
},
));
Ok(Self {
service: backend.clone(),
backend,
storage: Storage::Redis(storage),
pid: std::process::id(),
})
}
fn get(&self, py: Python<'_>, key: String) -> PyResult<Py<PyAny>> {
self.check_process()?;
let request = request(key, None)?;
let value =
release_gil(py, || self.backend.lookup(&request, now())).map_err(cache_error)?;
to_py(py, &value)
}
#[pyo3(signature = (key, value, *, ttl=None))]
fn set(
&self,
py: Python<'_>,
key: String,
value: &Bound<'_, PyAny>,
ttl: Option<f64>,
) -> PyResult<()> {
self.check_process()?;
let request = request(key, ttl)?;
let value: Value = from_py(value)?;
release_gil(py, || self.backend.store(&request, value, now())).map_err(cache_error)
}
fn async_get<'py>(&self, py: Python<'py>, key: String) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let request = request(key, None)?;
let backend = self.backend.clone();
crate::execution::run_async(
py,
async move { backend.async_lookup(&request, now()).await },
cache_error,
)
}
#[pyo3(signature = (key, value, *, ttl=None))]
fn async_set<'py>(
&self,
py: Python<'py>,
key: String,
value: &Bound<'_, PyAny>,
ttl: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let request = request(key, ttl)?;
let value: Value = from_py(value)?;
let backend = self.backend.clone();
crate::execution::run_async(
py,
async move { backend.async_store(&request, value, now()).await },
cache_error,
)
}
#[pyo3(signature = (entries, *, ttl=None))]
fn async_set_many<'py>(
&self,
py: Python<'py>,
entries: &Bound<'_, PyAny>,
ttl: Option<f64>,
) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let entries: Vec<(String, Value)> = from_py(entries)?;
let entries = entries
.into_iter()
.map(|(key, value)| Ok((request(key, ttl)?, value)))
.collect::<PyResult<Vec<_>>>()?;
let backend = self.backend.clone();
crate::execution::run_async(
py,
async move { backend.async_store_batch(entries, now()).await },
cache_error,
)
}
fn flush(&self, py: Python<'_>) -> PyResult<Py<PyAny>> {
self.check_process()?;
let backend = self.backend.clone();
crate::execution::run_sync(py, async move { backend.async_flush().await }, cache_error)
}
fn async_flush<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let backend = self.backend.clone();
crate::execution::run_async(py, async move { backend.async_flush().await }, cache_error)
}
fn ping<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let storage = self.storage.clone();
crate::execution::run_async(
py,
async move {
match storage {
Storage::Memory(_) => Ok(true),
Storage::Redis(cache) => cache.ping().await,
}
},
cache_error,
)
}
fn disconnect<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let storage = self.storage.clone();
crate::execution::run_async(
py,
async move {
match storage {
Storage::Memory(cache) => cache.disconnect().await,
Storage::Redis(cache) => cache.disconnect().await,
}
},
cache_error,
)
}
fn delete<'py>(&self, py: Python<'py>, keys: Vec<String>) -> PyResult<Bound<'py, PyAny>> {
self.check_process()?;
let storage = self.storage.clone();
crate::execution::run_async(
py,
async move {
for key in keys {
match &storage {
Storage::Memory(cache) => cache.async_delete_cache(&key).await?,
Storage::Redis(cache) => cache.async_delete_cache(&key).await?,
}
}
Ok(())
},
cache_error,
)
}
}
fn duration(seconds: f64) -> Result<Duration, HandleError> {
match Duration::try_from_secs_f64(seconds) {
Ok(duration) if !duration.is_zero() => Ok(duration),
Ok(_) => Err(HandleError::Duration(None)),
Err(error) => Err(HandleError::Duration(Some(error))),
}
}
pub(in crate::cache) fn native_handle<'py>(
configured: &Bound<'py, PyAny>,
) -> PyResult<Option<Bound<'py, PyAny>>> {
Ok(configured
.getattr_opt("cache")?
.map(|backend| backend.getattr_opt("native_handle"))
.transpose()?
.flatten()
.filter(|handle| handle.is_instance_of::<NativeCacheHandle>()))
}
pub(in crate::cache) fn configured(
configured: &Bound<'_, PyAny>,
kwargs: &Bound<'_, PyDict>,
) -> PyResult<(
Option<Arc<dyn ResponseCacheService>>,
litellm_cache_response::CacheOptions,
)> {
let handle = native_handle(configured)?.ok_or(HandleError::NotNative)?;
let cache = handle.extract::<PyRef<'_, NativeCacheHandle>>()?;
cache.check_process()?;
let controls = kwargs.get_item("cache")?.filter(|value| !value.is_none());
let controls = controls
.as_ref()
.map(|value| value.cast::<PyDict>())
.transpose()?;
if let Some(controls) = controls {
for name in controls.keys() {
let name = name.extract::<String>()?;
if !matches!(
name.as_str(),
"no-cache" | "no-store" | "ttl" | "s-maxage" | "s-max-age" | "use-cache"
) {
return Err(HandleError::Control(name).into());
}
}
}
let boolean = |name: &str| -> PyResult<bool> {
controls
.map(|values| values.get_item(name))
.transpose()?
.flatten()
.map(|value| value.extract())
.transpose()
.map(|value| value.unwrap_or(false))
};
let seconds = |name: &str| -> PyResult<Option<Duration>> {
controls
.map(|values| values.get_item(name))
.transpose()?
.flatten()
.map(|value| Ok(duration(value.extract()?)?))
.transpose()
};
Ok((
Some(cache.service.clone()),
litellm_cache_response::CacheOptions {
policy: litellm_cache_response::CachePolicy {
caching: kwargs
.get_item("caching")?
.filter(|value| !value.is_none())
.map(|value| value.extract())
.transpose()?,
no_cache: boolean("no-cache")?,
no_store: boolean("no-store")?,
ttl: seconds("ttl")?,
max_age: seconds("s-max-age")?.or(seconds("s-maxage")?),
},
scope: litellm_cache_response::CacheScope::Shared,
},
))
}
fn now() -> Duration {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
}

45
litellm-rust/crates/python-bridge/src/cache/selection.rs vendored
View file

@ -1,14 +1,10 @@
use super::{native, python};
use litellm_cache_response::{
CacheOptions, CachePolicy, CacheScope, ResponseCacheService, ScopedCache,
};
use super::{HandleError, python};
use litellm_cache_response::{CacheOptions, CachePolicy, CacheScope, ScopedCache};
use litellm_host::{
machine::{HostServices, MachineFault},
protocol::Protocol,
};
use pyo3::{prelude::*, types::PyDict};
use std::sync::Arc;
const CACHE_HOST_MODULE: &str = "litellm.rust_bridge.host.cache";
pub(crate) struct Cached<P>(std::marker::PhantomData<P>);
@ -24,7 +20,6 @@ impl<P: Protocol> Protocol for Cached<P> {
enum Backend {
Disabled,
Native(Arc<dyn ResponseCacheService>),
Python { namespace: String },
}
@ -43,7 +38,6 @@ impl Selection {
{
let service = match self.backend {
Backend::Disabled => None,
Backend::Native(service) => Some(service),
Backend::Python { namespace } => Some(python::service(services, namespace)),
};
(
@ -76,37 +70,21 @@ fn selected_cache<'py>(
Ok(Some(configured))
}
pub(crate) fn configured_native(
pub(crate) fn ensure_cache_disabled(
py: Python<'_>,
kwargs: &Bound<'_, PyDict>,
call_type: &str,
) -> PyResult<(
Option<Arc<dyn ResponseCacheService>>,
litellm_cache_response::CacheOptions,
)> {
) -> PyResult<()> {
let Some(configured) = selected_cache(py, kwargs, call_type)? else {
return Ok((
None,
litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared),
));
return Ok(());
};
native_configuration(&configured, kwargs)
}
fn native_configuration(
configured: &Bound<'_, PyAny>,
kwargs: &Bound<'_, PyDict>,
) -> PyResult<(Option<Arc<dyn ResponseCacheService>>, CacheOptions)> {
if !configured
if configured
.call_method("should_use_cache", (), Some(kwargs))?
.extract::<bool>()?
{
return Ok((
None,
litellm_cache_response::CacheOptions::new(litellm_cache_response::CacheScope::Shared),
));
return Err(HandleError::Unsupported.into());
}
native::v2::configured(configured, kwargs)
Ok(())
}
pub(crate) fn configure(
@ -122,13 +100,6 @@ pub(crate) fn configure(
options: CacheOptions::new(CacheScope::Shared),
});
};
if native::v2::native_handle(&cache)?.is_some() {
let (native, options) = native_configuration(&cache, arguments)?;
return Ok(Selection {
backend: native.map_or(Backend::Disabled, Backend::Native),
options,
});
}
let enabled = cache
.call_method("should_use_cache", (), Some(arguments))?
.extract::<bool>()?;

5
litellm-rust/crates/python-bridge/src/lib.rs
View file

@ -65,10 +65,6 @@ mod _native {
fn init(module: &Bound<'_, PyModule>) -> PyResult<()> {
let py = module.py();
let dict = module.dict();
dict.set_item(
"NativeCacheHandle",
py.get_type::<crate::cache::NativeCacheHandle>(),
)?;
dict.set_item(
"_SecretManagerRuntime",
py.get_type::<crate::secrets::runtime::NativeSecretManager>(),
@ -92,7 +88,6 @@ mod tests {
Python::initialize();
Python::attach(|py| {
let mut expected = vec![
"NativeCacheHandle",
"RustBridgeDeclined",
"RustUpstreamError",
"ForkedAfterNativeRuntimeStarted",

6
litellm-rust/crates/python-bridge/src/routes/chat_completions/mod.rs
View file

@ -4,7 +4,7 @@ use std::sync::Arc;
use host::ChatCompletionsPythonHost;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, ScopedCache};
use litellm_cache_response::CachePolicy;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_inference_chat::{ChatCompletionsRoute, route::ChatCompletions};
@ -57,10 +57,6 @@ impl InferenceRoute for ChatCompletionsRoute {
Self::new(http, auth, secrets)
}
fn with_cache(self, cache: ScopedCache) -> Self {
self.with_cache(cache)
}
fn machine(
self,
call: <ChatCompletions as Protocol>::Request,

11
litellm-rust/crates/python-bridge/src/routes/inference.rs
View file

@ -1,7 +1,7 @@
use std::sync::Arc;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, CacheScope, ScopedCache};
use litellm_cache_response::CachePolicy;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_core_utils::get_llm_provider_logic::get_custom_llm_provider;
use litellm_host::{call::HostedMachine, protocol::Protocol};
@ -142,7 +142,6 @@ pub(super) trait InferenceRoute: Sized + 'static {
auth: Arc<AuthServices>,
secrets: Arc<dyn SecretSource>,
) -> Self;
fn with_cache(self, cache: ScopedCache) -> Self;
fn machine(
self,
call: <Self::Protocol as Protocol>::Request,
@ -176,12 +175,8 @@ where
crate::http::resources().auth.clone(),
crate::secrets::source(py)?,
);
let (cache, cache_options) = crate::cache::configured_native(py, arguments, call_type)?;
let route = match cache {
Some(cache) => route.with_cache(ScopedCache::new(cache, CacheScope::Shared)),
None => route,
};
Ok(route.machine(request, cache_options.policy))
crate::cache::ensure_cache_disabled(py, arguments, call_type)?;
Ok(route.machine(request, CachePolicy::default()))
},
host,
hooks,

6
litellm-rust/crates/python-bridge/src/routes/responses/mod.rs
View file

@ -5,7 +5,7 @@ use std::sync::Arc;
use host::ResponsesPythonHost;
use litellm_auth::AuthServices;
use litellm_cache_response::{CachePolicy, ScopedCache};
use litellm_cache_response::CachePolicy;
use litellm_callbacks_legacy_python::LoggingOperation;
use litellm_host::{call::HostedMachine, protocol::Protocol};
use litellm_inference_responses::{ResponsesRoute, route::Responses};
@ -49,10 +49,6 @@ impl InferenceRoute for ResponsesRoute {
Self::new(http, auth, secrets)
}
fn with_cache(self, cache: ScopedCache) -> Self {
self.with_cache(cache)
}
fn machine(
self,
call: <Responses as Protocol>::Request,

6
litellm/__init__.py
View file

@ -688,6 +688,7 @@ darkbloom_models: Set = set()
openinfer_models: Final[Set[str]] = set()
kyma_models: Final[set[str]] = set() # mutable-ok: populated when the shared model catalog loads
streamlake_models: Final[Set[str]] = set()
wally_models: Final[Set[str]] = set()
zerogpu_models: Final[Set[str]] = set()
v0_models: Set = set()
morph_models: Set = set()
@ -984,6 +985,8 @@ def _populate_provider_model_sets(model_cost_map: Mapping[str, Mapping[str, obje
kyma_models.add(key)
elif value.get("litellm_provider") == "streamlake":
streamlake_models.add(key)
elif value.get("litellm_provider") == "wally":
wally_models.add(key)
elif value.get("litellm_provider") == "zerogpu":
zerogpu_models.add(key)
elif value.get("litellm_provider") == "v0":
@ -1071,7 +1074,6 @@ _populate_provider_model_sets(model_cost)
# Azure returns gpt-35-turbo in their responses, we need to map this to azure/gpt-3.5-turbo for token counting
azure_llms = {
"gpt-35-turbo": "azure/gpt-35-turbo",
"gpt-35-turbo-16k": "azure/gpt-35-turbo-16k",
"gpt-35-turbo-instruct": "azure/gpt-35-turbo-instruct",
"azure/gpt-41": "gpt-4.1",
"azure/gpt-41-mini": "gpt-4.1-mini",
@ -1166,6 +1168,7 @@ model_list = list(
| openinfer_models
| kyma_models
| streamlake_models
| wally_models
| zerogpu_models
| v0_models
| morph_models
@ -1299,6 +1302,7 @@ def _build_models_by_provider() -> dict:
"openinfer": openinfer_models,
"kyma": kyma_models,
"streamlake": streamlake_models,
"wally": wally_models,
"zerogpu": zerogpu_models,
"v0": v0_models,
"morph": morph_models,

3
litellm/_v2/AGENTS.md
View file

@ -1,3 +0,0 @@
Everything here is experimental and should not be documented
Use this directory to explore alternative APIs where the Rust migration makes backward compatibility difficult. The gateway can use them for performance, but keep them behind the v2 flag to avoid breaking SDK users

3
litellm/_v2/__init__.py
View file

@ -1,3 +0,0 @@
from litellm._v2.cache import Cache
__all__ = ("Cache",)

11
litellm/_v2/cache/AGENTS.md vendored
View file

@ -1,11 +0,0 @@
# Python v2 cache
Keep `litellm._v2.cache.Cache` import-compatible when reorganizing this package. This package owns Python cache factories and the adapter between the existing `BaseCache` interface and `NativeCacheHandle`
Construct native handles at this boundary and inject the adapter through the existing cache facade's `_backend` parameter. Keep the facade's `type`, namespace, and TTL consistent with the configured native backend
Keep storage implementation in the Rust storage crates and response-cache policy in `litellm-cache-response` and core. Do not duplicate cache-key generation, freshness rules, response encoding, or inference orchestration here
Preserve synchronous and asynchronous cache operations, including TTL forwarding and lifecycle methods. Validate Python values before passing them to typed native interfaces. Keep native extension imports lazy so importing the package does not require loading the extension
Extend the existing v2 cache tests in `tests/test_litellm_rust/test_v2.py` for behavioral changes, following that directory's `AGENTS.md`. Test observable cache behavior rather than package layout or implementation structure

72
litellm/_v2/cache/__init__.py vendored
View file

@ -1,72 +0,0 @@
from __future__ import annotations
from collections.abc import Sequence
from typing import TYPE_CHECKING, Final
from pydantic import TypeAdapter
from litellm.caching.base_cache import BaseCache
from litellm.caching.caching import Cache as CacheFacade
from litellm.types.caching import LiteLLMCacheType
if TYPE_CHECKING:
from litellm.rust_bridge._native import NativeCacheHandle
_DURATION: Final[TypeAdapter[float | None]] = TypeAdapter(float | None)
class NativeBackend(BaseCache):
def __init__(self, handle: NativeCacheHandle) -> None:
self.native_handle = handle
def get_cache(self, key: str, **kwargs: object) -> object:
return self.native_handle.get(key)
async def async_get_cache(self, key: str, **kwargs: object) -> object:
return await self.native_handle.async_get(key)
def set_cache(self, key: str, value: object, **kwargs: object) -> None:
self.native_handle.set(key, value, ttl=_DURATION.validate_python(kwargs.get("ttl")))
async def async_set_cache(self, key: str, value: object, **kwargs: object) -> None:
await self.native_handle.async_set(key, value, ttl=_DURATION.validate_python(kwargs.get("ttl")))
async def async_set_cache_pipeline(self, cache_list: Sequence[tuple[str, object]], **kwargs: object) -> None:
await self.native_handle.async_set_many(cache_list, ttl=_DURATION.validate_python(kwargs.get("ttl")))
async def batch_cache_write(self, key: str, value: object, **kwargs: object) -> None:
await self.async_set_cache(key, value, **kwargs)
def flush_cache(self) -> None:
self.native_handle.flush()
async def async_flush_cache(self) -> None:
await self.native_handle.async_flush()
async def ping(self) -> bool:
return await self.native_handle.ping()
async def disconnect(self) -> None:
await self.native_handle.disconnect()
async def delete_cache_keys(self, keys: Sequence[str]) -> None:
await self.native_handle.delete(keys)
async def test_connection(self) -> dict[str, str]:
return {"status": "success" if await self.ping() else "failed"}
class Cache:
@staticmethod
def memory(*, ttl: float = 600, capacity: int = 200, max_entry_bytes: int = 4194304) -> CacheFacade:
from litellm.rust_bridge._native import NativeCacheHandle
handle: Final = NativeCacheHandle.memory(ttl=ttl, capacity=capacity, max_entry_bytes=max_entry_bytes)
return CacheFacade(type=LiteLLMCacheType.LOCAL, ttl=ttl, _backend=NativeBackend(handle))
@staticmethod
def redis(url: str, *, namespace: str, ttl: float = 600, max_entry_bytes: int = 4194304) -> CacheFacade:
from litellm.rust_bridge._native import NativeCacheHandle
handle: Final = NativeCacheHandle.redis(url, namespace=namespace, ttl=ttl, max_entry_bytes=max_entry_bytes)
return CacheFacade(type=LiteLLMCacheType.REDIS, namespace=namespace, ttl=ttl, _backend=NativeBackend(handle))

1
litellm/anthropic_beta_headers_manager.py
View file

@ -198,7 +198,6 @@ def _load_beta_headers_config() -> dict:
if _BETA_HEADERS_CONFIG is not None:
return _BETA_HEADERS_CONFIG
# Get the URL from environment or use default
from litellm import anthropic_beta_headers_url
_BETA_HEADERS_CONFIG = get_beta_headers_config(url=anthropic_beta_headers_url)

586
litellm/blog_posts.json
View file

@ -1,10 +1,588 @@
{
"posts": [
{
"title": "Realtime WebRTC HTTP Endpoints",
"description": "Use the LiteLLM proxy to route OpenAI-style WebRTC realtime via HTTP: client_secrets and SDP exchange.",
"date": "2026-03-12",
"url": "https://docs.litellm.ai/blog/realtime_webrtc_http_endpoints"
"title": "Release cadence changes and an LTS version pilot",
"description": "Starting the week of November 2, 2026, release candidates and stable releases ship on Tuesday mornings, and LiteLLM is piloting a long-term support line on 1.108 for one month.",
"date": "2026-11-02",
"url": "https://docs.litellm.ai/blog/release-cadence-and-lts-pilot",
"tags": [
"release",
"support"
]
},
{
"title": "How to build an agent with grounding on LiteLLM",
"description": "Build a grounded chat agent in four steps with LiteLLM's AI Gateway. Connect your evidence, call a grounding guardrail, and show its verdict in chat.",
"date": "2026-10-31",
"url": "https://docs.litellm.ai/blog/grounded-agent",
"tags": [
"engineering",
"guardrails",
"agents"
]
},
{
"title": "Introducing LiteLLM’s Unified Cost Accounting Engine",
"description": "Consistent cost reporting, independently scalable accounting, and spend analysis by team, user, model, or custom dimensions.",
"date": "2026-10-31",
"url": "https://docs.litellm.ai/blog/cost-infrastructure-makeover",
"tags": [
"engineering",
"reliability"
]
},
{
"title": "What we are changing after October 21",
"description": "Why the controls built after March did not stop the v1.107.0 release, and the identity, credential and detection changes we are making in response.",
"date": "2026-10-30",
"url": "https://docs.litellm.ai/blog/security-changes-after-oct-21-2026",
"tags": [
"security",
"incident-report"
]
},
{
"title": "Use Hex to analyze LiteLLM costs, then set budgets and routing",
"description": "Connect Hex to LiteLLM’s Postgres data, import a cost dashboard, then use the LiteLLM UI to set a budget or change model routing.",
"date": "2026-10-30",
"url": "https://docs.litellm.ai/blog/hex-litellm-roi",
"tags": [
"observability",
"cost tracking",
"ai-gateway"
]
},
{
"title": "Introducing the LiteLLM Private Registry",
"description": "LiteLLM is committing to a private registry to provide early access to fixes for the most severe security vulnerabilities, ahead of public release.",
"date": "2026-10-27",
"url": "https://docs.litellm.ai/blog/litellm-private-registry",
"tags": [
"security"
]
},
{
"title": "Morph open-weight models on LiteLLM",
"description": "Use Morph's Kimi, GLM, and DeepSeek models through LiteLLM with tool calling, structured output, streaming, and prompt-cache cost tracking.",
"date": "2026-10-27",
"url": "https://docs.litellm.ai/blog/morph-provider",
"tags": []
},
{
"title": "Abliterated models on LiteLLM",
"description": "Call abliterated models from Refuseless and Abliteration through LiteLLM, with gateway access controls, streaming, tool calls, and cached-token pricing.",
"date": "2026-10-26",
"url": "https://docs.litellm.ai/blog/abliterated-models",
"tags": [
"engineering"
]
},
{
"title": "GPT-6.1 Terra API Pricing and Day 0 Support on LiteLLM",
"description": "GPT-6.1 Terra costs $1 per 1M input tokens and $5 per 1M output tokens, half the price of GPT-6.1 Sol. LiteLLM supports it on day 0.",
"date": "2026-10-26",
"url": "https://docs.litellm.ai/blog/gpt-6-1-terra",
"tags": [
"openai",
"gpt-6.1",
"gpt-6.1-terra",
"pricing",
"day 0 support",
"card_title:GPT-6.1 Terra",
"card_summary:Half the price of GPT-6.1 Sol at $1 / $5 per 1M tokens, 50.3% on Terminal-Bench 4.0, 1.05M context."
]
},
{
"title": "Rust Migration: /v1/messages end-to-end in Rust by Dec 1",
"description": "By December 1st, the proxy will handle Anthropic messages requests 20x faster",
"date": "2026-10-26",
"url": "https://docs.litellm.ai/blog/rust-migration-rerouted",
"tags": [
"rust",
"rust-migration",
"ai-gateway",
"engineering"
]
},
{
"title": "OpenCode Go is now a LiteLLM provider",
"description": "Use your OpenCode Go subscription through LiteLLM. Connect once, route models to the right API, preserve sessions, and track estimated usage.",
"date": "2026-10-24",
"url": "https://docs.litellm.ai/blog/opencode-go",
"tags": [
"providers",
"tutorial"
]
},
{
"title": "The attack chain behind the v1.107.0 LiteLLM release",
"description": "How an internal key committed to a public repository in May led to three stolen GitHub tokens and the compromised litellm v1.107.0 release.",
"date": "2026-10-23",
"url": "https://docs.litellm.ai/blog/security-incident-oct-21-2026-attack-chain",
"tags": [
"security",
"incident-report"
]
},
{
"title": "Tokenize text across vLLM and Watsonx",
"description": "Choose token counting for prompt budgets or /tokenize to inspect tokens. Use one request format across vLLM and Watsonx through LiteLLM.",
"date": "2026-10-23",
"url": "https://docs.litellm.ai/blog/unified-tokenize",
"tags": [
"engineering"
]
},
{
"title": "October Townhall Updates: 426 Bug Fixes, 88.4% Coverage, and the Decisions API",
"description": "A recap of the October LiteLLM town hall: security changes that need action before you upgrade, 426 bug fixes, test coverage at 88.4%, the Decisions API, Lens, Moyai, and LiteLLM Relay.",
"date": "2026-10-23",
"url": "https://docs.litellm.ai/blog/october-townhall-updates",
"tags": [
"townhall",
"security",
"reliability",
"product"
]
},
{
"title": "Contained: Oct 21, 2026 LiteLLM release v1.107.0 Compromised",
"description": "LiteLLM release v1.107.0 was compromised by a malicious actor.",
"date": "2026-10-21",
"url": "https://docs.litellm.ai/blog/security-incident-oct-21-2026",
"tags": [
"security",
"incident-report"
]
},
{
"title": "JevRank: decision models through LiteLLM’s rerank API",
"description": "Preview JevRank for low-cost reranking with Perplexity and TypeSafe decision models. Explore SciFact results, reproducible data and LiteLLM setup.",
"date": "2026-10-20",
"url": "https://docs.litellm.ai/blog/jevrank",
"tags": [
"engineering",
"reranking",
"decisions"
]
},
{
"title": "Cache-Aware Routing for Self-Hosted Models",
"description": "Route each agent session to the same vLLM pod on Kubernetes to reuse cached prompt prefixes.",
"date": "2026-10-20",
"url": "https://docs.litellm.ai/blog/cache-aware-routing",
"tags": [
"routing",
"kubernetes",
"vllm",
"ai-gateway"
]
},
{
"title": "Day 0 Support: TypeSafe Jev 1.14: Adds images, box questions, and long context",
"description": "TypeSafe Jev 1.14 reads images, answers box questions with coordinates and takes long-context requests. LiteLLM serves all three on /v1/systemone and /v1/decisions behind TypeSafe's opt-in headers.",
"date": "2026-10-15",
"url": "https://docs.litellm.ai/blog/typesafe_jev_images",
"tags": [
"typesafe",
"jev",
"decisions",
"day 0 support",
"card_title:Jev 1.14",
"card_summary:Jev 1.14 reads PNG and JPEG images, answers box questions with coordinates, and takes long-context requests on LiteLLM."
]
},
{
"title": "Decision models on LiteLLM: /v1/decisions and /v1/systemone",
"description": "Decision models return probabilities instead of text, which makes routing, triage and guardrail checks cheap and deterministic. LiteLLM serves nine providers of them behind one request shape.",
"date": "2026-10-14",
"url": "https://docs.litellm.ai/blog/decision-models",
"tags": [
"openai",
"typesafe",
"decisions",
"day 0 support",
"card_title:Decision Models",
"card_summary:Typed yes/no, choice and score answers as probabilities, from 9 providers on /v1/decisions and /v1/systemone."
]
},
{
"title": "Workshop: Build a self-hosted cloud agent.",
"description": "Build a self-hosted cloud agent with the LiteLLM team. A two-week remote workshop, October 19–30, to build a working v0 and deploy an internal beta on your infrastructure.",
"date": "2026-10-13",
"url": "https://docs.litellm.ai/blog/cloud-agent-workshop",
"tags": [
"agents",
"infrastructure",
"announcement"
]
},
{
"title": "Gemini 4 Argon API Pricing and Day 0 Support on LiteLLM",
"description": "Gemini 4 Argon costs $2.00 per 1M input tokens and $10.00 per 1M output tokens until Nov 30. LiteLLM supports it on day 0 on AI Studio and Google Cloud.",
"date": "2026-10-13",
"url": "https://docs.litellm.ai/blog/gemini-4-argon",
"tags": [
"gemini",
"gemini 4",
"gemini-4-argon",
"pricing",
"day 0 support",
"card_title:Gemini 4 Argon",
"card_summary:Google's new flagship: 77.9% on DeepSWE v1.1, 1M context, thinking levels, $2 / $10 per 1M tokens until Nov 30."
]
},
{
"title": "October Townhall: Product + Roadmap Updates",
"description": "Join the LiteLLM October townhall on Friday, 23 October at 7:30 AM PT to learn about LiteLLM's product updates and roadmap.",
"date": "2026-10-10",
"url": "https://docs.litellm.ai/blog/october-townhall-announcement",
"tags": [
"announcement",
"townhall"
]
},
{
"title": "Introducing Microsoft 365 Copilot in LiteLLM",
"description": "Use Microsoft 365 Copilot from your applications through LiteLLM, with access to your Microsoft 365 work data.",
"date": "2026-10-09",
"url": "https://docs.litellm.ai/blog/microsoft-365-copilot",
"tags": [
"announcement",
"ai-gateway"
]
},
{
"title": "Introducing litellm-core: the first step toward a leaner SDK",
"description": "Ongoing work to reduce LiteLLM's dependencies and installation footprint, improve runtime loading, and make imports faster",
"date": "2026-10-08",
"url": "https://docs.litellm.ai/blog/introducing-litellm-core",
"tags": []
},
{
"title": "Claude Haiku 5.5 API Pricing and Day 0 Support on LiteLLM",
"description": "Claude Haiku 5.5 costs $0.10 per 1M input tokens and $0.50 per 1M output tokens. LiteLLM supports it on day 0 on Anthropic, AWS, Google Cloud and Azure.",
"date": "2026-10-07",
"url": "https://docs.litellm.ai/blog/claude-haiku-5-5",
"tags": [
"anthropic",
"claude",
"claude haiku",
"haiku 5.5",
"pricing",
"day 0 support",
"card_title:Claude Haiku 5.5",
"card_summary:90% cheaper than Haiku 4.5 on short prompts, 1M context, 72.4% on OSWorld 2.1, $0.10 / $0.50 per 1M tokens."
]
},
{
"title": "Open Sourcing Moyai: Self-Hosted Cloud Coding Agent",
"description": "Run cloud coding agents on your own infrastructure with Moyai. Choose your agent and models through LiteLLM, track costs, and follow the setup guide.",
"date": "2026-10-07",
"url": "https://docs.litellm.ai/blog/moyai-open-source",
"tags": [
"agents",
"open-source",
"infrastructure"
]
},
{
"title": "Day 0 Support: Nano Banana 2.1",
"description": "Day 0 support for Gemini Nano Banana 2.1 on LiteLLM, on Google AI Studio and Gemini Enterprise Agent Platform, with image output at half Nano Banana 2's price.",
"date": "2026-10-06",
"url": "https://docs.litellm.ai/blog/nano_banana_2_1",
"tags": [
"gemini",
"gemini enterprise agent platform",
"nano banana",
"image generation",
"day 0 support",
"card_title:Nano Banana 2.1",
"card_summary:Gemini image model with image output at half Nano Banana 2's price: $30 / MTok image out, $1.50 text in."
]
},
{
"title": "How LiteLLM Lens finds repeated failures across 1,000s of agent traces",
"description": "Inside Lens: parallel trace review, Python tools for large traces, and investigations that connect repeated failures to source evidence.",
"date": "2026-10-05",
"url": "https://docs.litellm.ai/blog/lens-failure-patterns",
"tags": [
"lens",
"agent-tracing",
"engineering"
]
},
{
"title": "How we cut time to first byte by 94% for long prompts",
"description": "A 440k-token benchmark went from 553 ms to 35 ms median time to first byte. We removed unnecessary prompt-cache routing work before the model call.",
"date": "2026-10-03",
"url": "https://docs.litellm.ai/blog/long-prompt-latency",
"tags": [
"performance",
"proxy",
"engineering",
"ai-gateway"
]
},
{
"title": "How we built our own internal Devin in 2 days",
"description": "How we built Moyai Devin with Render, Modal, Hermes, Temporal, and LiteLLM: durable sessions, parallel agents, Slack, and shared organization connections.",
"date": "2026-10-03",
"url": "https://docs.litellm.ai/blog/internal-devin-two-days",
"tags": [
"engineering",
"agents",
"infrastructure",
"slack"
]
},
{
"title": "Adding Self-hosted Auto Router Classifiers: Laya & Nimble",
"description": "Use Laya or Bespoke Nimble to classify Auto Router requests on your own infrastructure. Control where classification runs, which model you serve, and how you provision it.",
"date": "2026-10-02",
"url": "https://docs.litellm.ai/blog/laya-nimble-classifiers",
"tags": [
"auto-router",
"product",
"ai-gateway"
]
},
{
"title": "Launching LiteLLM Lens",
"description": "LiteLLM Lens turns the traces flowing through your gateway into findings your agents can act on. Built for agent swarms generating 200K+ traces.",
"date": "2026-10-01",
"url": "https://docs.litellm.ai/blog/litellm-lens-launch",
"tags": [
"lens",
"agent-tracing"
]
},
{
"title": "How we cut LiteLLM's Redis round trips per request by 64%",
"description": "A governed request to the LiteLLM AI Gateway waited on Redis 22 times. It now waits 8 times: one pipeline per Redis backend before the model call, one after.",
"date": "2026-10-01",
"url": "https://docs.litellm.ai/blog/redis-request-pipelines",
"tags": [
"performance",
"redis",
"proxy",
"engineering",
"ai-gateway"
]
},
{
"title": "How we made the LiteLLM Usage page 120x faster",
"description": "From six minutes to 3.2 seconds: how moving aggregation into Postgres made LiteLLM's Usage page roughly 120x faster in our benchmark.",
"date": "2026-09-30",
"url": "https://docs.litellm.ai/blog/usage-page-bounded-loading",
"tags": [
"performance",
"admin-ui",
"postgres",
"engineering",
"ai-gateway"
]
},
{
"title": "Day 0 Support: GPT-6.1 Sol",
"description": "Day 0 support for GPT-6.1 Sol on LiteLLM, with cached input at half GPT-6 Sol's price.",
"date": "2026-09-29",
"url": "https://docs.litellm.ai/blog/gpt_6_1_sol",
"tags": [
"openai",
"gpt-6",
"gpt-6.1-sol",
"completion",
"day 0 support",
"card_title:GPT-6.1 Sol",
"card_summary:GPT-6 Sol successor with cached input at half GPT-6 Sol's price, same request family and params."
]
},
{
"title": "Learn the LiteLLM gateway with a guided course",
"description": "Follow a request through the LiteLLM gateway, Router, and SDK. A guided course for teams deploying the gateway and contributors making code changes.",
"date": "2026-09-28",
"url": "https://docs.litellm.ai/blog/litellm-gateway-course",
"tags": [
"proxy",
"product"
]
},
{
"title": "Day 0 Support: Claude Sonnet 5.5",
"description": "Day 0 support for Claude Sonnet 5.5 on the LiteLLM AI Gateway. Use it across Anthropic, Bedrock, Gemini Enterprise Agent Platform, and Azure.",
"date": "2026-09-28",
"url": "https://docs.litellm.ai/blog/claude_sonnet_5_5",
"tags": [
"anthropic",
"claude",
"sonnet 5.5",
"day 0 support",
"card_title:Claude Sonnet 5.5",
"card_summary:Same $2 / $10 price as Sonnet 5, 30%+ faster, 70.6% on Terminal-Bench 4.0 (Sonnet 5: 10.3%), 1M context."
]
},
{
"title": "Introducing the LiteLLM ROI Calculator",
"description": "Compare each engineer's LiteLLM gateway spend with the estimated effort of their merged pull requests. Open source and self-hosted.",
"date": "2026-09-26",
"url": "https://docs.litellm.ai/blog/litellm-roi-calculator",
"tags": [
"cost",
"product",
"announcement"
]
},
{
"title": "Introducing LiteAgents",
"description": "Switch agent harnesses without rewriting your agent. Keep your tools and model configuration, use native harness controls, and add Temporal when you need durable runs.",
"date": "2026-09-25",
"url": "https://docs.litellm.ai/blog/liteagents-sdk",
"tags": [
"agents",
"sdk"
]
},
{
"title": "September Townhall Updates: 583 Bug Fixes, OCR on Rust, and 83.5% Coverage",
"description": "A recap of the September LiteLLM town hall: security and stability updates, OCR running on Rust by default, test coverage past the 80% target, and new Fusion and liteagents launches.",
"date": "2026-09-24",
"url": "https://docs.litellm.ai/blog/september-townhall-updates",
"tags": [
"townhall",
"reliability",
"product"
]
},
{
"title": "Day 0 support: Gemini 3.8 Flash TTS and Flash-Lite TTS",
"description": "Day 0 support for Gemini 3.8 Flash TTS and Flash-Lite TTS on LiteLLM's /v1/audio/speech, with launch pricing tracked through 2026.",
"date": "2026-09-23",
"url": "https://docs.litellm.ai/blog/gemini_3_8_flash_tts",
"tags": [
"gemini",
"gemini-3.8-flash-tts",
"text to speech",
"day 0 support",
"card_title:Gemini 3.8 Flash TTS and Flash-Lite TTS",
"card_summary:Gemini speech models on /v1/audio/speech, Flash and Flash-Lite, with launch pricing tracked through 2026."
]
},
{
"title": "Introducing LiteAdmin MCP",
"description": "Give your agent tools to create keys, add models, and manage budgets with LiteAdmin MCP. Use the same connector through LiteAdmin, our Slack admin agent.",
"date": "2026-09-23",
"url": "https://docs.litellm.ai/blog/litellm-admin-mcp",
"tags": [
"mcp",
"agents",
"ai-gateway"
]
},
{
"title": "Day 0 Support: GPT-6 Sol and GPT-6 Luna",
"description": "Day 0 support for GPT-6 Sol and GPT-6 Luna on LiteLLM, at half the price of their GPT-5.6 counterparts.",
"date": "2026-09-22",
"url": "https://docs.litellm.ai/blog/gpt_6_sol_luna",
"tags": [
"openai",
"gpt-6",
"gpt-6-sol",
"gpt-6-luna",
"completion",
"day 0 support",
"card_title:GPT-6 Sol and GPT-6 Luna",
"card_summary:Two GPT-6 variants at half the price of their GPT-5.6 counterparts, routed as the GPT-5 request family."
]
},
{
"title": "Day 0 Support: Claude Opus 5.5",
"description": "Day 0 support for Claude Opus 5.5 on the LiteLLM AI Gateway. Use it across Anthropic, Bedrock, Gemini Enterprise Agent Platform, and Azure.",
"date": "2026-09-22",
"url": "https://docs.litellm.ai/blog/claude_opus_5_5",
"tags": [
"anthropic",
"claude",
"opus 5.5",
"day 0 support",
"card_title:Claude Opus 5.5",
"card_summary:20% cheaper than Opus 5 at $4 / $20 per 1M, 30% faster output, Terminal-Bench 4.0 66.4%, thinking always on."
]
},
{
"title": "Getting Started with Fireworks AI on LiteLLM",
"description": "A beginner tutorial: run an open model on Fireworks AI behind a LiteLLM gateway, call it with the OpenAI SDK you already use, then add a second model and a fallback.",
"date": "2026-09-22",
"url": "https://docs.litellm.ai/blog/fireworks-getting-started",
"tags": [
"tutorial",
"providers",
"fireworks",
"getting-started"
]
},
{
"title": "Day 0 Support: Xiaomi MiMo V2.6",
"description": "Day 0 support for Xiaomi MiMo V2.6 Pro and Flash on LiteLLM, priced on the native route for the first time.",
"date": "2026-09-22",
"url": "https://docs.litellm.ai/blog/mimo_v2_6",
"tags": [
"xiaomi",
"mimo",
"mimo-v2.6",
"day 0 support",
"card_title:Xiaomi MiMo V2.6",
"card_summary:Xiaomi's Pro ($0.435 / $0.87 per 1M) and Flash ($0.14 / $0.28), same rates as V2.5, priced on the native route."
]
},
{
"title": "Claude Code server-side auto mode through LiteLLM",
"description": "Anthropic is moving Claude Code auto mode's safety classifier server-side. LiteLLM's native /v1/messages route now forwards the safeguards contract to the Anthropic API, Bedrock InvokeModel, Bedrock Mantle and Vertex AI. Here is the contract, what changed, which releases carry it, and how to verify it.",
"date": "2026-09-21",
"url": "https://docs.litellm.ai/blog/claude-code-server-side-auto-mode",
"tags": [
"announcement",
"claude-code",
"anthropic",
"ai-gateway"
]
},
{
"title": "Day 0 Support: Grok 4.7",
"description": "Day 0 support for Grok 4.7 on LiteLLM, at the same price as Grok 4.6.",
"date": "2026-09-21",
"url": "https://docs.litellm.ai/blog/grok_4_7",
"tags": [
"xai",
"grok",
"grok-4.7",
"day 0 support",
"card_title:Grok 4.7",
"card_summary:SpaceXAI's update to Grok 4.6 at the same price and limits, better at checking its own work and long context."
]
},
{
"title": "JEV Classifier: 5.43x as Fast as Haiku, 96% Lower Cost",
"description": "JEV classified requests 5.43x as fast as Haiku by median latency in our AI Gateway benchmark. Explore the setup, cost savings and methodology.",
"date": "2026-09-20",
"url": "https://docs.litellm.ai/blog/jev-auto-router-benchmark",
"tags": [
"routing",
"complexity-router",
"engineering",
"ai-gateway"
]
},
{
"title": "TypeSafe Jev on LiteLLM",
"description": "TypeSafe AI's Jev lands in LiteLLM v1.103.0-rc: call it through the proxy with logging and cost tracking.",
"date": "2026-09-20",
"url": "https://docs.litellm.ai/blog/typesafe_jev",
"tags": [
"typesafe",
"jev",
"product"
]
}
]
}

4
litellm/constants.py
View file

@ -135,6 +135,7 @@ ANTHROPIC_DEFAULT_FILE_LIST_LIMIT: Final = 20
DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = get_env_int("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10)
DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = get_env_int("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1)
GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS: Final = 10
DEFAULT_MAX_REQUEST_HEADER_SIZE: Final = 16 * 1024
budget_reservation_disabled_info_emitted = False
DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = get_env_int("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1)
DEFAULT_SQS_BATCH_SIZE: Final = get_env_int("DEFAULT_SQS_BATCH_SIZE", 512)
@ -1092,6 +1093,7 @@ openai_compatible_endpoints: Final[list[str]] = [
"https://api.scx.ai/v1",
"https://api.prisminference.com/v1",
"https://api.reka.ai/v1",
"https://inference.runanywhere.ai/v1",
"https://gigachat.devices.sberbank.ru/api/v1",
"https://api.openinfer.ai/v1",
"https://vanchin.streamlake.ai/api/gateway/v1/endpoints",
@ -1179,6 +1181,7 @@ openai_compatible_providers: Final[list[str]] = [
"sail",
"reka",
"streamlake",
"wally",
]
OPENAI_AUDIO_TRANSCRIPTION_PROVIDERS: Final = frozenset({"openai"} | frozenset(openai_compatible_providers))
@ -2067,6 +2070,7 @@ SPEND_LOG_KEY_METADATA_CACHE_TTL: Final = 600
SPEND_LOG_KEY_METADATA_MISS_CACHE_TTL: Final = 30
SPEND_LOG_KEY_METADATA_CACHE_MAX_ITEMS: Final = 10000
SPEND_LOG_KEY_METADATA_QUERY_TIMEOUT_MS: Final = 5000
SPEND_LOG_READ_STATEMENT_TIMEOUT_SECONDS: Final = 30
SPEND_LOG_KEY_METADATA_ROWS_PER_PROBE: Final = 100
# Short TTL for negative MCP access-group existence lookups. Keeps unauthenticated
# callers from forcing a DB query per request for unknown names, while bounding

84
litellm/cost_calculator.py
View file

@ -49,6 +49,8 @@ from litellm.litellm_core_utils.llm_cost_calc.utils import (
generic_cost_per_token,
get_batch_cost_rates,
get_billable_input_tokens,
get_cost_per_unit,
get_modality_token_cost_key,
get_regional_uplift_multiplier,
get_service_tier_cost_key,
get_token_type_cost_breakdown,
@ -442,6 +444,14 @@ def _per_second_pricing_cost(
model_info: Final = cached_get_model_info_helper(model=model, custom_llm_provider=custom_llm_provider)
except Exception: # noqa: BLE001 # the lookup raises plain Exception for an unmapped model
return None
return _per_second_pricing_cost_from_info(model_info, response_time_ms, audio_seconds)
def _per_second_pricing_cost_from_info(
model_info: ModelInfoBase,
response_time_ms: float | None,
audio_seconds: float = 0.0,
) -> tuple[float, float] | None:
if _has_token_or_tiered_pricing(model_info) or not _bills_wall_clock_seconds(model_info):
return None
cost_per_second: Final = model_info.get("cost_per_second")
@ -464,7 +474,7 @@ def _per_second_pricing_cost(
)
verbose_logger.debug(
"For model=%s - cost_per_second: %s; response time: %s",
model,
model_info.get("key"),
resolved_cost_per_second,
response_time_ms,
)
@ -487,6 +497,7 @@ def cost_per_token(
### CUSTOM PRICING ###
custom_cost_per_token: CostPerToken | None = None,
custom_cost_per_second: float | None = None,
requested_model_info_override: ModelInfo | None = None,
### NUMBER OF QUERIES ###
number_of_queries: int | None = None,
### USAGE OBJECT ###
@ -598,6 +609,20 @@ def cost_per_token(
if response_cost is not None:
return response_cost[0], response_cost[1]
if requested_model_info_override is not None:
per_second: Final = _per_second_pricing_cost_from_info(
requested_model_info_override, response_time_ms, audio_transcription_file_duration
)
if per_second is not None:
return per_second
return generic_cost_per_token(
model=model,
usage=usage_block,
custom_llm_provider=custom_llm_provider or "",
service_tier=service_tier,
model_info=requested_model_info_override,
)
model_cost_ref: Final = litellm.model_cost
# Only callers that explicitly pass `custom_llm_provider` get the
# dedup/prefix-join treatment. When provider is omitted, preserve legacy
@ -1102,20 +1127,19 @@ def _cost_map_model_name(model: str, custom_llm_provider: str | None, region_nam
return _strip_unregistered_leading_segments(f"{provider_prefix}/{model}", region_name)
def _rule_declares_rates(model: str, custom_llm_provider: str | None) -> bool:
rule_info: Final = unmapped_model_generalization(model, custom_llm_provider)
return rule_info is not None and _cost_map_entry_prices_anything(rule_info)
def declared_cost_map_pricing(key: str, model: str, custom_llm_provider: str | None) -> Mapping[str, object]:
entry: Final = _raw_cost_map_entry(key)
if entry is None:
return unmapped_model_generalization(model, custom_llm_provider) or MappingProxyType({})
fill_missing: Final = mapped_model_fill_missing(key, entry)
return entry if fill_missing is None else MappingProxyType({**fill_missing, **entry})
def _resolves_to_rates(model: str, custom_llm_provider: str | None) -> bool:
key: Final = cached_cost_map_key(model, custom_llm_provider)
if key is None:
return False
entry: Final = _raw_cost_map_entry(key)
if entry is None:
return _rule_declares_rates(model, custom_llm_provider)
fill_missing: Final = mapped_model_fill_missing(key, entry)
return _cost_map_entry_prices_anything(entry if fill_missing is None else {**fill_missing, **entry})
return key is not None and _cost_map_entry_prices_anything(
declared_cost_map_pricing(key, model, custom_llm_provider)
)
def _served_or_priced_requested_model(
@ -1748,6 +1772,7 @@ def completion_cost(
### CUSTOM PRICING ###
custom_cost_per_token: CostPerToken | None = None,
custom_cost_per_second: float | None = None,
requested_model_info_override: ModelInfo | None = None,
optional_params: dict | None = None,
custom_pricing: bool | None = None,
base_model: str | None = None,
@ -1801,6 +1826,7 @@ def completion_cost(
quality=quality,
n=n,
custom_cost_per_token=custom_cost_per_token,
requested_model_info_override=requested_model_info_override,
custom_cost_per_second=custom_cost_per_second,
optional_params=optional_params,
custom_pricing=custom_pricing,
@ -2275,6 +2301,7 @@ def completion_cost(
region_name=None if explicit_pricing else region_name,
custom_cost_per_second=custom_cost_per_second,
custom_cost_per_token=custom_cost_per_token,
requested_model_info_override=requested_model_info_override,
prompt_characters=prompt_characters,
completion_characters=completion_characters,
cache_creation_input_tokens=cache_creation_input_tokens,
@ -2420,18 +2447,20 @@ def completion_cost(
raise e
def get_response_cost_from_hidden_params(
hidden_params: dict | BaseModel,
) -> float | None:
if isinstance(hidden_params, BaseModel):
_hidden_params_dict = hidden_params.model_dump()
else:
_hidden_params_dict = hidden_params
_ADDITIONAL_HEADERS: Final = TypeAdapter(Mapping[str, object])
additional_headers: Final = _hidden_params_dict.get("additional_headers", {})
if additional_headers and "llm_provider-x-litellm-response-cost" in additional_headers:
return _billable_header_cost(additional_headers["llm_provider-x-litellm-response-cost"])
return None
def get_response_cost_from_hidden_params(
hidden_params: Mapping[str, object] | BaseModel,
) -> float | None:
params: Final[Mapping[str, object]] = (
hidden_params.model_dump() if isinstance(hidden_params, BaseModel) else hidden_params
)
additional_headers: Final = params.get("additional_headers")
if not isinstance(additional_headers, Mapping):
return None
headers: Final = _ADDITIONAL_HEADERS.validate_python(additional_headers)
return _billable_header_cost(headers.get("llm_provider-x-litellm-response-cost"))
def _billable_header_cost(value: object) -> float | None:
@ -2507,6 +2536,7 @@ def response_cost_calculator(
region_name: str | None = None,
custom_cost_per_token: CostPerToken | None = None,
custom_cost_per_second: float | None = None,
requested_model_info_override: ModelInfo | None = None,
) -> float:
"""
Returns
@ -2563,6 +2593,7 @@ def response_cost_calculator(
custom_pricing=custom_pricing,
base_model=base_model,
custom_cost_per_token=custom_cost_per_token,
requested_model_info_override=requested_model_info_override,
custom_cost_per_second=custom_cost_per_second,
prompt=prompt,
standard_built_in_tools_params=standard_built_in_tools_params,
@ -3157,8 +3188,12 @@ def _batch_rate(
"input_cost_per_video_token_batches",
],
fallback: float,
*,
prompt_tokens: int = 0,
) -> float:
rate: Final = model_info.get(key)
context_key: Final = get_modality_token_cost_key(model_info, key, prompt_tokens)
context_rate: Final = get_cost_per_unit(model_info, context_key, None)
rate: Final = model_info.get(key) if context_rate is None else context_rate
return fallback if rate is None else rate
@ -3308,7 +3343,8 @@ def _batch_prompt_cost(
)
return (
text_tokens * input_rate
+ audio_tokens * _batch_rate(model_info, "input_cost_per_audio_token_batches", input_rate)
+ audio_tokens
* _batch_rate(model_info, "input_cost_per_audio_token_batches", input_rate, prompt_tokens=usage.prompt_tokens)
+ image_tokens * _batch_rate(model_info, "input_cost_per_image_token_batches", input_rate)
+ video_tokens * _batch_rate(model_info, "input_cost_per_video_token_batches", input_rate)
+ (cached_tokens - cached_audio_tokens) * (cache_read_rate or 0.0)

3
litellm/exceptions.py
View file

@ -867,6 +867,7 @@ class BadGatewayError(openai.APIStatusError):
class InternalServerError(openai.InternalServerError):
response: httpx.Response # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX response
request: httpx.Request # pyright: ignore[reportIncompatibleVariableOverride] # LiteLLM constructs an HTTPX request
provider_status_code: int | None
def __init__(
self,
@ -878,8 +879,10 @@ class InternalServerError(openai.InternalServerError):
max_retries: int | None = None,
num_retries: int | None = None,
body: object | None = None,
provider_status_code: int | None = None,
):
self.status_code = 500
self.provider_status_code = provider_status_code
self.message = f"litellm.InternalServerError: {message}"
self.llm_provider = llm_provider
self.model = model

14
litellm/images/main.py
View file

@ -67,6 +67,10 @@ from litellm.utils import (
)
_PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints)
_OPTIONAL_OBJECT_DICT_ADAPTER: Final = TypeAdapter[dict[str, object] | None](dict[str, object] | None)
_OPTIONAL_STRING_DICT_ADAPTER: Final = TypeAdapter[dict[str, str] | None](dict[str, str] | None)
_OBJECT_DICT_ADAPTER: Final = TypeAdapter[dict[str, object]](dict[str, object])
_STRING_DICT_ADAPTER: Final = TypeAdapter[dict[str, str]](dict[str, str])
_PROVIDERS_RESOLVING_THEIR_OWN_IMAGE_BASE: Final = (
LlmProviders.DASHSCOPE,
LlmProviders.QWENCLOUD,
@ -93,8 +97,8 @@ def _chatgpt_image_params(kwargs: Mapping[str, object]) -> GenericLiteLLMParams:
from litellm.responses.utils import ResponsesAPIRequestUtils
headers: Final = ResponsesAPIRequestUtils.merge_client_forwarded_headers(
TypeAdapter[dict[str, object] | None](dict[str, object] | None).validate_python(kwargs.get("extra_headers")),
TypeAdapter[dict[str, str] | None](dict[str, str] | None).validate_python(kwargs.get("headers")),
_OPTIONAL_OBJECT_DICT_ADAPTER.validate_python(kwargs.get("extra_headers")),
_OPTIONAL_STRING_DICT_ADAPTER.validate_python(kwargs.get("headers")),
)
return GenericLiteLLMParams.model_validate({**kwargs, "extra_headers": headers})
@ -167,7 +171,7 @@ async def aimage_generation(*args, **kwargs) -> ImageResponse:
custom_llm_provider=provider_hints.get("custom_llm_provider"),
api_base=provider_hints.get("api_base"),
litellm_params=(
_chatgpt_image_params(TypeAdapter(dict[str, object]).validate_python(kwargs))
_chatgpt_image_params(_OBJECT_DICT_ADAPTER.validate_python(kwargs))
if provider_hints.get("custom_llm_provider") == "chatgpt"
or (isinstance(model, str) and model.startswith("chatgpt/"))
else None
@ -299,7 +303,7 @@ def image_generation(
headers.update(extra_headers)
model_response: ImageResponse = litellm.utils.ImageResponse()
chatgpt_params: Final = (
_chatgpt_image_params(TypeAdapter(dict[str, object]).validate_python({**kwargs, "api_key": api_key}))
_chatgpt_image_params(_OBJECT_DICT_ADAPTER.validate_python({**kwargs, "api_key": api_key}))
if custom_llm_provider == "chatgpt" or (model is not None and model.startswith("chatgpt/"))
else None
)
@ -582,7 +586,7 @@ def image_generation(
_chatgpt_image_headers(
model,
chatgpt_params.model_copy(update={"api_key": dynamic_api_key or api_key}),
TypeAdapter(dict[str, str]).validate_python(headers),
_STRING_DICT_ADAPTER.validate_python(headers),
)
)
if chatgpt_params is not None

128
litellm/integrations/code_interpreter_interception/handler.py
View file

@ -10,6 +10,7 @@ import json
import time
import uuid
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypedDict, runtime_checkable
from pydantic import ValidationError
@ -73,7 +74,7 @@ class CodeInterpreterLogOutput(TypedDict):
class CodeInterpreterCall(TypedDict):
id: str
type: Literal["code_interpreter_call"]
status: Literal["completed"]
status: Literal["completed", "failed"]
code: str
container_id: str | None
outputs: list[CodeInterpreterLogOutput]
@ -141,6 +142,39 @@ class _SupportsOutput(Protocol):
_CachedContainer: TypeAlias = tuple[ContainerHandle, SandboxToolParams | None, float, str | None]
@dataclass(frozen=True, slots=True)
class _ReadySandbox:
container: ContainerHandle
params: SandboxToolParams | None
@dataclass(frozen=True, slots=True)
class _SandboxUnavailable:
output: str
_Sandbox: TypeAlias = _ReadySandbox | _SandboxUnavailable
@dataclass(frozen=True, slots=True)
class _ToolRun:
output: str
status: Literal["completed", "failed"]
_SANDBOX_UNAVAILABLE_OUTPUT: Final = "The code did not run because the code execution sandbox is unavailable."
_SANDBOX_LOST_RESULT_OUTPUT: Final = (
"The code execution sandbox failed while running the code, so the code may have run but its result is lost."
)
def _sandbox_unavailable(step: str, error: Exception, output: str = _SANDBOX_UNAVAILABLE_OUTPUT) -> _SandboxUnavailable:
verbose_logger.warning(
"CodeInterpreterInterception: %s failed, answering the tool call with the failure: %s", step, error
)
return _SandboxUnavailable(output=output)
def _output_item_type(item: object) -> object:
if isinstance(item, dict):
item_mapping: Final[dict[str, object]] = item
@ -189,6 +223,14 @@ def _extract_identity(kwargs: Mapping[str, object]) -> str:
return identity if isinstance(identity, str) else ""
def _bridged_from_responses(kwargs: Mapping[str, object], call_type: CallTypes | None) -> bool:
route: Final = getattr(kwargs.get("litellm_logging_obj"), "call_type", None)
return call_type in (CallTypes.completion, CallTypes.acompletion) and route in (
CallTypes.responses,
CallTypes.aresponses,
)
def _resolve_sandbox_tool(sandbox_tool_name: str | None) -> SandboxToolParams | None:
if sandbox_tool_name is None:
return None
@ -261,7 +303,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
async def async_pre_call_deployment_hook(
self, kwargs: dict[str, object], call_type: CallTypes | None
) -> dict | None:
if not kwargs.get("_agentic_loop_depth"):
if not kwargs.get("_agentic_loop_depth") and not _bridged_from_responses(kwargs, call_type):
kwargs.pop(_INTERCEPTION_ACTIVE_KEY, None)
kwargs.pop(_SANDBOX_KEY, None)
self._strip_interception_metadata(kwargs)
@ -451,16 +493,16 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
sandbox_key: Final = self._extract_sandbox_key(kwargs)
is_session: Final = bool(kwargs.get(_SESSION_SCOPED_KEY))
identity: Final = _extract_identity(kwargs) if is_session else None
container, params = await self._get_or_create_container(cache_key=sandbox_key, identity=identity)
sandbox: Final = await self._acquire_sandbox(cache_key=sandbox_key, identity=identity)
try:
container_id: Final = self._container_id(container)
container_id: Final = self._sandbox_container_id(sandbox)
input_list: Final = self._normalize_messages(messages)
code_interpreter_calls: Final[list[CodeInterpreterCall]] = []
for tool_call in tool_calls:
arguments = tool_call.get("arguments", "")
code = self._parse_code(arguments)
stdout = await self._run_tool_call(container=container, params=params, arguments=arguments)
tool_run = await self._sandbox_tool_output(sandbox=sandbox, arguments=arguments)
input_list.append(
{
"type": "function_call",
@ -473,22 +515,23 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
{
"type": "function_call_output",
"call_id": tool_call.get("call_id"),
"output": stdout,
"output": tool_run.output,
}
)
code_interpreter_calls.append(
{
"id": f"ci_{uuid.uuid4().hex}",
"type": "code_interpreter_call",
"status": "completed",
"status": tool_run.status,
"code": code,
"container_id": container_id,
"outputs": ([{"type": "logs", "logs": stdout}] if stdout else []),
"outputs": ([{"type": "logs", "logs": tool_run.output}] if tool_run.output else []),
}
)
except Exception:
await self._delete_container_for_cache_key(sandbox_key)
raise
await self._evict_sandbox_if_a_run_failed(sandbox_key, sandbox, code_interpreter_calls)
optional_params: Final = anthropic_messages_optional_request_params
request_patch: Final = AgenticLoopRequestPatch(
@ -526,14 +569,13 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
sandbox_key: Final = self._extract_sandbox_key(kwargs)
is_session: Final = bool(kwargs.get(_SESSION_SCOPED_KEY))
identity: Final = _extract_identity(kwargs) if is_session else None
container, params = await self._get_or_create_container(cache_key=sandbox_key, identity=identity)
sandbox: Final = await self._acquire_sandbox(cache_key=sandbox_key, identity=identity)
try:
container_id: Final = self._container_id(container)
container_id: Final = self._sandbox_container_id(sandbox)
tool_results: Final = [
await self._build_chat_completion_tool_result(
container=container,
params=params,
sandbox=sandbox,
tool_call=tool_call,
container_id=container_id,
)
@ -544,6 +586,7 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
raise
tool_messages: Final = [result[0] for result in tool_results]
code_interpreter_calls: Final = [result[1] for result in tool_results]
await self._evict_sandbox_if_a_run_failed(sandbox_key, sandbox, code_interpreter_calls)
request_patch: Final = AgenticLoopRequestPatch(
model=model,
@ -569,9 +612,13 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
)
@staticmethod
def _container_id(container: ContainerHandle) -> str | None:
container_id: Final[object] = getattr(container, "id", None)
return container_id if isinstance(container_id, str) else None
def _sandbox_container_id(sandbox: _Sandbox) -> str | None:
match sandbox:
case _ReadySandbox(container=container):
container_id: Final[object] = getattr(container, "id", None)
return container_id if isinstance(container_id, str) else None
case _SandboxUnavailable():
return None
@staticmethod
def _agentic_tool_calls(tools: dict[str, object]) -> list[CodeExecutionToolCall]:
@ -588,28 +635,27 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
async def _build_chat_completion_tool_result(
self,
container: ContainerHandle,
params: SandboxToolParams | None,
sandbox: _Sandbox,
tool_call: CodeExecutionToolCall,
container_id: str | None,
) -> tuple[ChatCompletionToolMessage, CodeInterpreterCall]:
arguments: Final = tool_call.get("arguments", "")
code: Final = self._parse_code(arguments)
stdout: Final = await self._run_tool_call(container=container, params=params, arguments=arguments)
tool_run: Final = await self._sandbox_tool_output(sandbox=sandbox, arguments=arguments)
tool_call_id: Final = tool_call.get("id") or tool_call.get("call_id") or uuid.uuid4().hex
return (
{
"role": "tool",
"tool_call_id": tool_call_id,
"content": stdout,
"content": tool_run.output,
},
{
"id": f"ci_{uuid.uuid4().hex}",
"type": "code_interpreter_call",
"status": "completed",
"status": tool_run.status,
"code": code,
"container_id": container_id,
"outputs": [{"type": "logs", "logs": stdout}] if stdout else [],
"outputs": [{"type": "logs", "logs": tool_run.output}] if tool_run.output else [],
},
)
@ -694,19 +740,55 @@ class CodeInterpreterInterceptionLogger(CustomLogger):
except (json.JSONDecodeError, TypeError, AttributeError):
return ""
async def _run_tool_call(self, container: ContainerHandle, params: SandboxToolParams | None, arguments: str) -> str:
async def _evict_sandbox_if_a_run_failed(
self, sandbox_key: str | None, sandbox: _Sandbox, calls: Sequence[CodeInterpreterCall]
) -> None:
if not sandbox_key or not isinstance(sandbox, _ReadySandbox):
return
if not any(call["status"] == "failed" for call in calls):
return
cached: Final = self._container_cache.get(sandbox_key)
if cached is None or cached[0] is not sandbox.container:
return
self._container_cache.pop(sandbox_key)
await self._delete_container(container=cached[0], params=cached[1])
async def _sandbox_tool_output(self, sandbox: _Sandbox, arguments: str) -> _ToolRun:
match sandbox:
case _SandboxUnavailable(output=output):
return _ToolRun(output=output, status="failed")
case _ReadySandbox(container=container, params=params):
match await self._run_tool_call(container=container, params=params, arguments=arguments):
case _SandboxUnavailable(output=output):
return _ToolRun(output=output, status="failed")
case str() as output:
return _ToolRun(output=output, status="completed")
async def _run_tool_call(
self, container: ContainerHandle, params: SandboxToolParams | None, arguments: str
) -> str | _SandboxUnavailable:
try:
code: Final = json.loads(arguments).get("code", "") if arguments else ""
except (json.JSONDecodeError, TypeError):
return "[invalid tool arguments: could not parse code]"
result: Final = await self._run_code(container=container, params=params, code=code)
try:
result: Final = await self._run_code(container=container, params=params, code=code)
except Exception as e:
return _sandbox_unavailable("running the code", e, _SANDBOX_LOST_RESULT_OUTPUT)
if getattr(result, "error", None):
error: Final = result.error
message: Final = error.get("value") or error.get("name") if isinstance(error, dict) else str(error)
return f"[execution error] {message}"
return getattr(result, "stdout", "") or ""
async def _acquire_sandbox(self, cache_key: str | None, identity: str | None) -> _Sandbox:
try:
container, params = await self._get_or_create_container(cache_key=cache_key, identity=identity)
except Exception as e:
return _sandbox_unavailable("creating the sandbox", e)
return _ReadySandbox(container=container, params=params)
async def _get_or_create_container(
self,
cache_key: str | None,

1
litellm/integrations/otel/model/spans.py
View file

@ -341,6 +341,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset(
"LiteLLM_AccountingDelivery",
"LiteLLM_AccountingEvent",
"LiteLLM_ClaudeCodePluginTable",
"LiteLLM_ClaudeCodePluginUpload",
"LiteLLM_MemoryTable",
"LiteLLM_AdaptiveRouterState",
"LiteLLM_AdaptiveRouterSession",

36
litellm/litellm_core_utils/agentic_followup_failure.py Normal file
View file

@ -0,0 +1,36 @@
from dataclasses import dataclass
from typing import Final, TypeAlias
from litellm.types.utils import Usage
ServedRound: TypeAlias = tuple[Usage, float]
_MARK_ATTRIBUTE: Final = "_litellm_agentic_followup_mark"
@dataclass(frozen=True, slots=True)
class _FollowUpMark:
served_rounds: tuple[ServedRound, ...]
def _mark_of(error: BaseException) -> _FollowUpMark | None:
mark: Final[object] = getattr(error, _MARK_ATTRIBUTE, None)
return mark if isinstance(mark, _FollowUpMark) else None
def mark_agentic_followup_failure(error: BaseException, served_rounds: tuple[ServedRound, ...]) -> None:
setattr(error, _MARK_ATTRIBUTE, _FollowUpMark(served_rounds))
def is_agentic_followup_failure(error: BaseException) -> bool:
return _mark_of(error) is not None
def rounds_served_before(error: BaseException) -> tuple[ServedRound, ...]:
mark: Final = _mark_of(error)
return () if mark is None else mark.served_rounds
def carry_agentic_followup_mark(source: BaseException, mapped: BaseException) -> None:
if mapped is not source and is_agentic_followup_failure(source):
mark_agentic_followup_failure(mapped, rounds_served_before(source))

107
litellm/litellm_core_utils/agentic_loop_refusal.py Normal file
View file

@ -0,0 +1,107 @@
import asyncio
from collections.abc import Awaitable, Iterable, Mapping
from itertools import chain
from types import MappingProxyType
from typing import Final
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.prompt_templates.mid_conversation_system import message_field, parts_of
from litellm.types.llms.openai import ResponsesAPIResponse
from litellm.types.utils import (
ChatCompletionMessageCustomToolCall,
ChatCompletionMessageToolCall,
Choices,
ModelResponse,
)
def _strings(values: Iterable[object]) -> frozenset[str]:
return frozenset(value for value in values if isinstance(value, str))
def _claimed_calls(tool_calls: object) -> tuple[object, ...]:
wrapped: Final = message_field(tool_calls, "tool_calls")
return parts_of(tool_calls if wrapped is None else wrapped)
async def _calls_a_gate_claims(gate: Awaitable[tuple[bool, object]]) -> tuple[object, ...]:
try:
should_run, tool_calls = await gate
except Exception as e:
verbose_logger.debug("LiteLLM.AgenticLoopRefused: skipping a gate that failed with %s", type(e).__name__)
return ()
return _claimed_calls(tool_calls) if should_run else ()
async def every_claimed_tool_call(
refused_tool_calls: object, other_gates: Iterable[Awaitable[tuple[bool, object]]]
) -> Mapping[str, tuple[object, ...]]:
claims: Final = await asyncio.gather(*(_calls_a_gate_claims(gate) for gate in other_gates))
return MappingProxyType({"tool_calls": (*_claimed_calls(refused_tool_calls), *chain.from_iterable(claims))})
def _call_ids(call: object) -> frozenset[str]:
return _strings((message_field(call, "id"), message_field(call, "call_id")))
def refused_tool_call_identifiers(tool_calls: object) -> tuple[frozenset[str], frozenset[str]]:
calls: Final = _claimed_calls(tool_calls)
ids: Final = frozenset(chain.from_iterable(_call_ids(call) for call in calls))
names: Final = _strings(message_field(call, "name") for call in calls if not _call_ids(call))
return ids, names
def is_refused_call(call_id: object, name: object, refused_ids: frozenset[str], refused_names: frozenset[str]) -> bool:
return (isinstance(call_id, str) and call_id in refused_ids) or name in refused_names
def is_refused_item(
item: object, item_type: str, id_field: str, refused_ids: frozenset[str], refused_names: frozenset[str]
) -> bool:
if message_field(item, "type") != item_type:
return False
return is_refused_call(message_field(item, id_field), message_field(item, "name"), refused_ids, refused_names)
def finalize_refused_responses_turn(response: ResponsesAPIResponse, tool_calls: object) -> ResponsesAPIResponse:
refused_ids, refused_names = refused_tool_call_identifiers(tool_calls)
kept: Final = [
item
for item in parts_of(message_field(response, "output"))
if not is_refused_item(item, "function_call", "call_id", refused_ids, refused_names)
]
return response.model_copy(update={"output": kept})
def _is_refused_chat_tool_call(call: object, refused_ids: frozenset[str], refused_names: frozenset[str]) -> bool:
name: Final = message_field(message_field(call, "function"), "name")
return is_refused_call(message_field(call, "id"), name, refused_ids, refused_names)
def _at_position(
call: ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall, position: int
) -> ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall:
return call if message_field(call, "index") is None else call.model_copy(update={"index": position})
def _finalized_choice(choice: object, refused_ids: frozenset[str], refused_names: frozenset[str]) -> object:
if not isinstance(choice, Choices) or not choice.message.tool_calls:
return choice
surviving: Final = (
call for call in choice.message.tool_calls if not _is_refused_chat_tool_call(call, refused_ids, refused_names)
)
kept: Final = [_at_position(call, position) for position, call in enumerate(surviving)]
finish_reason: Final = "stop" if not kept and choice.finish_reason == "tool_calls" else choice.finish_reason
return choice.model_copy(
update={
"message": choice.message.model_copy(update={"tool_calls": kept or None}),
"finish_reason": finish_reason,
}
)
def finalize_refused_chat_turn(response: ModelResponse, tool_calls: object) -> ModelResponse:
refused_ids, refused_names = refused_tool_call_identifiers(tool_calls)
return response.model_copy(
update={"choices": [_finalized_choice(choice, refused_ids, refused_names) for choice in response.choices]}
)

2
litellm/litellm_core_utils/audio_utils/transcription_streaming.py
View file

@ -177,7 +177,7 @@ class _TranscriptionEventCollector:
def _billed_usage(response: TranscriptionResponse) -> Usage:
usage: Final[object] = getattr(response, "usage", None)
usage: Final = response.usage
token_usage: Final = (
TranscriptionUsageObjectTransformation.transform_transcription_usage_object(usage)
if isinstance(usage, TranscriptionUsageTokensObject)

371
litellm/litellm_core_utils/chat_completion_agentic_loop.py
View file

@ -1,29 +1,35 @@
# this is a patch to allow for agentic loops covering llm_http_handler.py and openai sdk based calling flows for the .completion() api
import json
from collections.abc import Mapping
from collections.abc import Awaitable, Iterator, Mapping, Sequence
from itertools import chain
from types import MappingProxyType
from typing import Final
from typing import Final, TypeVar
from pydantic import TypeAdapter, ValidationError
from litellm._internal_context import internal_sub_call
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.agentic_followup_failure import mark_agentic_followup_failure, rounds_served_before
from litellm.litellm_core_utils.agentic_followup_kwargs import build_agentic_followup_kwargs
from litellm.litellm_core_utils.agentic_loop_refusal import every_claimed_tool_call, finalize_refused_chat_turn
from litellm.litellm_core_utils.agentic_loop_settings import (
DEFAULT_MAX_AGENTIC_LOOPS,
validated_max_agentic_loops,
)
from litellm.litellm_core_utils.core_helpers import qualify_agentic_followup_model
from litellm.litellm_core_utils.hidden_params import get_hidden_params, served_from_cache
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import is_str_keyed_dict, object_items
from litellm.llms.base_llm.base_model_iterator import MockResponseIterator
from litellm.types.integrations.custom_logger import (
CHAT_COMPLETION_AGENTIC_SURFACE,
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
AgenticLoopLimitError,
AgenticLoopPlan,
AgenticLoopRequestPatch,
AgenticLoopSafetyError,
converted_stream_requested,
is_interception_internal_key,
stashed_stream_options,
@ -31,6 +37,8 @@ from litellm.types.integrations.custom_logger import (
from litellm.types.utils import ModelResponse
from litellm.utils import CustomStreamWrapper
_DispatchResult: Final = TypeVar("_DispatchResult")
_FOLLOWUP_INTERNAL_PARAMS: Final = frozenset(
(
"acompletion",
@ -44,22 +52,41 @@ _FOLLOWUP_INTERNAL_PARAMS: Final = frozenset(
)
async def bill_first_round_on_failure(
dispatch: Awaitable[_DispatchResult], *, logging_obj: object, response: object, stream: bool
) -> _DispatchResult:
try:
return await dispatch
except Exception as error:
if stream or not isinstance(logging_obj, LiteLLMLoggingObject):
mark_agentic_followup_failure(error, rounds_served_before(error))
raise
this_round: Final = logging_obj.served_round(response)
served_rounds: Final = (*rounds_served_before(error), *(() if this_round is None else (this_round,)))
mark_agentic_followup_failure(error, served_rounds)
logging_obj.record_served_rounds_for_failure(served_rounds)
raise
def _gate_overridden(callback: CustomLogger) -> bool:
base: Final = CustomLogger.async_should_run_agentic_loop
func: Final = type(callback).async_should_run_agentic_loop
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
return type(callback).async_should_run_agentic_loop is not CustomLogger.async_should_run_agentic_loop
def _build_plan_overridden(callback: CustomLogger) -> bool:
base: Final = CustomLogger.async_build_agentic_loop_plan
func: Final = type(callback).async_build_agentic_loop_plan
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
return type(callback).async_build_agentic_loop_plan is not CustomLogger.async_build_agentic_loop_plan
def _chat_build_plan_overridden(callback: CustomLogger) -> bool:
return (
type(callback).async_build_chat_completion_agentic_loop_plan
is not CustomLogger.async_build_chat_completion_agentic_loop_plan
)
def _post_hook_overridden(callback: CustomLogger) -> bool:
base: Final = CustomLogger.async_post_agentic_loop_response_hook
func: Final = type(callback).async_post_agentic_loop_response_hook
return getattr(func, "__func__", func) is not getattr(base, "__func__", base)
return (
type(callback).async_post_agentic_loop_response_hook is not CustomLogger.async_post_agentic_loop_response_hook
)
_STREAM_OPTIONS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
@ -76,14 +103,14 @@ def _coerce_int(value: object, default: int) -> int:
return int(value) if isinstance(value, (int, str)) else default
def _agentic_loop_settings(kwargs: dict[str, object]) -> tuple[int, int, list[str]]:
def _agentic_loop_settings(kwargs: Mapping[str, object]) -> tuple[int, int, tuple[str, ...]]:
depth: Final = _coerce_int(kwargs.get("_agentic_loop_depth"), 0)
configured: Final = validated_max_agentic_loops(
kwargs.get("max_agentic_loops"), field="litellm_params.max_agentic_loops"
)
max_loops: Final = DEFAULT_MAX_AGENTIC_LOOPS if configured is None else configured
raw_fingerprints: Final = kwargs.get("_agentic_loop_fingerprints")
fingerprints: Final = [str(fp) for fp in raw_fingerprints] if isinstance(raw_fingerprints, list) else []
fingerprints: Final = tuple(str(fp) for fp in raw_fingerprints) if isinstance(raw_fingerprints, list) else ()
return depth, max_loops, fingerprints
@ -96,16 +123,16 @@ def _fingerprint_tools(tool_calls: object) -> str:
def _check_agentic_loop_safety(
tool_calls: object,
fingerprints: list[str],
fingerprints: Sequence[str],
depth: int,
max_loops: int,
model: str,
) -> str:
fingerprint: Final = _fingerprint_tools(tool_calls)
if fingerprint in fingerprints:
raise ValueError("Agentic loop detected repeated tool-call fingerprint; aborting rerun")
raise AgenticLoopSafetyError("Agentic loop detected repeated tool-call fingerprint; aborting rerun")
if depth >= max_loops:
raise ValueError(f"Exceeded max_agentic_loops={max_loops} for model={model}")
raise AgenticLoopLimitError(f"Exceeded max_agentic_loops={max_loops} for model={model}")
return fingerprint
@ -134,6 +161,83 @@ def _converted_stream_replay(
)
def _end_refused_turn(
finalized: ModelResponse,
*,
kwargs: Mapping[str, object],
depth: int,
model: str,
custom_llm_provider: str,
logging_obj: object,
) -> ModelResponse | CustomStreamWrapper:
replay: Final = _converted_stream_replay(
finalized,
kwargs=kwargs,
depth=depth,
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
return finalized if replay is None else replay
def _request_tools(optional_params: Mapping[str, object]) -> list[dict[str, object]]:
return [tool for tool in object_items(optional_params.get("tools")) if is_str_keyed_dict(tool)]
def _custom_loggers(logging_obj: object) -> tuple[CustomLogger, ...]:
import litellm
dynamic: Final = getattr(logging_obj, "dynamic_success_callbacks", None) or []
return tuple(callback for callback in (*litellm.callbacks, *dynamic) if isinstance(callback, CustomLogger))
async def _refused_chat_turn(
callbacks: Sequence[CustomLogger],
refusing: CustomLogger,
refused_tool_calls: object,
*,
response: ModelResponse,
model: str,
messages: list[dict[str, object]],
tools: list[dict[str, object]] | None,
custom_llm_provider: str,
kwargs: Mapping[str, object],
) -> ModelResponse:
surface_kwargs: Final = {
**kwargs,
"_agentic_loop_api_surface": CHAT_COMPLETION_AGENTIC_SURFACE,
"custom_llm_provider": custom_llm_provider,
}
def other_gates() -> Iterator[Awaitable[tuple[bool, object]]]:
for callback in callbacks:
if callback is refusing:
continue
yield callback.async_should_run_chat_completion_agentic_loop(
response=response,
model=model,
messages=messages,
tools=tools,
stream=False,
custom_llm_provider=custom_llm_provider,
kwargs=dict(kwargs),
)
if _gate_overridden(callback):
yield callback.async_should_run_agentic_loop(
response=response,
model=model,
messages=messages,
tools=tools,
stream=False,
custom_llm_provider=custom_llm_provider,
kwargs=dict(surface_kwargs),
)
claimed: Final = await every_claimed_tool_call(refused_tool_calls, other_gates())
return finalize_refused_chat_turn(response, claimed)
def _with_agentic_loop_metadata(kwargs_for_followup: Mapping[str, object]) -> Mapping[str, object]:
metadata: Final = kwargs_for_followup.get("litellm_metadata")
return MappingProxyType(
@ -164,6 +268,12 @@ def _filter_followup_kwargs(source: dict[str, object]) -> dict[str, object]:
}
def _leave_pricing_to_the_request(response: object) -> None:
hidden_params: Final = get_hidden_params(response)
if hidden_params is not None and not served_from_cache(response):
hidden_params.pop("response_cost", None)
async def _execute_chat_completion_agentic_plan(
*,
plan: AgenticLoopPlan,
@ -175,7 +285,7 @@ async def _execute_chat_completion_agentic_plan(
custom_llm_provider: str,
depth: int,
max_loops: int,
fingerprints: list[str],
fingerprints: Sequence[str],
fingerprint: str,
) -> object:
import litellm
@ -226,6 +336,7 @@ async def _execute_chat_completion_agentic_plan(
model,
str(e),
)
_leave_pricing_to_the_request(response_followup)
replay: Final = _converted_stream_replay(
response_followup,
kwargs=kwargs,
@ -252,23 +363,18 @@ async def maybe_run_chat_completion_agentic_loop(
*,
response: ModelResponse,
model: str,
messages: list,
optional_params: dict,
messages: list[dict[str, object]],
optional_params: dict[str, object],
kwargs: dict[str, object],
logging_obj: object,
custom_llm_provider: str,
stream: bool,
) -> ModelResponse | CustomStreamWrapper | None:
import litellm
callbacks: Final = litellm.callbacks + (getattr(logging_obj, "dynamic_success_callbacks", None) or [])
callbacks: Final = _custom_loggers(logging_obj)
depth, max_loops, fingerprints = _agentic_loop_settings(kwargs)
tools: Final = optional_params.get("tools", [])
tools: Final = _request_tools(optional_params)
for callback in callbacks:
if not isinstance(callback, CustomLogger):
continue
if not _gate_overridden(callback):
continue
@ -297,29 +403,39 @@ async def maybe_run_chat_completion_agentic_loop(
if not should_run:
continue
fingerprint = _check_agentic_loop_safety(
tool_calls=tool_calls,
fingerprints=fingerprints,
depth=depth,
max_loops=max_loops,
model=model,
)
try:
if not _build_plan_overridden(callback):
return await callback.async_run_agentic_loop(
tools=tool_calls,
fingerprint = _check_agentic_loop_safety(
tool_calls=tool_calls,
fingerprints=fingerprints,
depth=depth,
max_loops=max_loops,
model=model,
)
except AgenticLoopSafetyError as e:
if stream:
raise
verbose_logger.warning("LiteLLM.AgenticLoopRefused: ending turn [rail=%s]", type(e).__name__)
return _end_refused_turn(
await _refused_chat_turn(
callbacks,
callback,
tool_calls,
response=response,
model=model,
messages=messages,
response=response,
anthropic_messages_provider_config=None,
anthropic_messages_optional_request_params=optional_params,
logging_obj=logging_obj,
stream=stream,
kwargs=hook_kwargs,
)
tools=tools,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
),
kwargs=kwargs,
depth=depth,
model=model,
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
plan = await callback.async_build_agentic_loop_plan(
if not _build_plan_overridden(callback):
return await callback.async_run_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
@ -331,31 +447,38 @@ async def maybe_run_chat_completion_agentic_loop(
kwargs=hook_kwargs,
)
if plan.response_override is not None:
return plan.response_override
if plan.terminate:
return response
if not plan.run_agentic_loop:
continue
plan = await callback.async_build_agentic_loop_plan(
tools=tool_calls,
model=model,
messages=messages,
response=response,
anthropic_messages_provider_config=None,
anthropic_messages_optional_request_params=optional_params,
logging_obj=logging_obj,
stream=stream,
kwargs=hook_kwargs,
)
return await _execute_chat_completion_agentic_plan(
plan=plan,
callback=callback,
model=model,
optional_params=optional_params,
kwargs=kwargs,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
depth=depth,
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
)
except Exception as e:
verbose_logger.exception(
"LiteLLM.AgenticHookError: Exception in chat completion agentic hooks: %s",
str(e),
)
if plan.response_override is not None:
return plan.response_override
if plan.terminate:
return response
if not plan.run_agentic_loop:
continue
return await _execute_chat_completion_agentic_plan(
plan=plan,
callback=callback,
model=model,
optional_params=optional_params,
kwargs=kwargs,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
depth=depth,
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
)
return _converted_stream_replay(
response,
@ -365,3 +488,113 @@ async def maybe_run_chat_completion_agentic_loop(
custom_llm_provider=custom_llm_provider,
logging_obj=logging_obj,
)
async def run_chat_completion_agentic_hooks(
*,
response: ModelResponse,
model: str,
messages: list[dict[str, object]],
optional_params: dict[str, object],
kwargs: dict[str, object],
logging_obj: LiteLLMLoggingObject,
custom_llm_provider: str,
stream: bool,
) -> object | None:
callbacks: Final = _custom_loggers(logging_obj)
depth, max_loops, fingerprints = _agentic_loop_settings(kwargs)
tools: Final = _request_tools(optional_params)
hook_kwargs: Final = {**kwargs, "custom_llm_provider": custom_llm_provider}
for callback in callbacks:
try:
should_run, tool_calls = await callback.async_should_run_chat_completion_agentic_loop(
response=response,
model=model,
messages=messages,
tools=tools,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
except Exception:
verbose_logger.exception(
"LiteLLM.AgenticHookError: Exception in async_should_run_chat_completion_agentic_loop"
)
continue
if not should_run:
continue
try:
fingerprint = _check_agentic_loop_safety(
tool_calls=tool_calls,
fingerprints=fingerprints,
depth=depth,
max_loops=max_loops,
model=model,
)
except AgenticLoopSafetyError as e:
if stream:
raise
verbose_logger.warning("LiteLLM.AgenticLoopRefused: ending turn [rail=%s]", type(e).__name__)
return await _refused_chat_turn(
callbacks,
callback,
tool_calls,
response=response,
model=model,
messages=messages,
tools=tools,
custom_llm_provider=custom_llm_provider,
kwargs=kwargs,
)
if not _chat_build_plan_overridden(callback):
return await callback.async_run_chat_completion_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
response=response,
optional_params=optional_params,
logging_obj=logging_obj,
stream=stream,
kwargs=hook_kwargs,
)
plan = await callback.async_build_chat_completion_agentic_loop_plan(
tools=tool_calls,
model=model,
messages=messages,
response=response,
optional_params=optional_params,
logging_obj=logging_obj,
stream=stream,
kwargs=hook_kwargs,
)
if plan.response_override is not None:
return plan.response_override
if plan.terminate:
verbose_logger.debug(
"Agentic chat loop terminated by callback=%s reason=%s", type(callback).__name__, plan.stop_reason
)
return response
if not plan.run_agentic_loop:
continue
return await _execute_chat_completion_agentic_plan(
plan=plan,
callback=callback,
model=model,
optional_params=optional_params,
kwargs=hook_kwargs,
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
depth=depth,
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
)
return None

10
litellm/litellm_core_utils/deployment_pricing.py
View file

@ -1,4 +1,4 @@
from collections.abc import Iterator, Mapping
from collections.abc import Iterable, Iterator, Mapping
from types import MappingProxyType
from typing import Final
@ -25,6 +25,14 @@ def deployment_pricing_models(deployment: object) -> tuple[str, ...]:
return pricing_models(info.get("base_model") or params.get("base_model"), params.get("model"))
def deployment_backend_model(deployment: object) -> str:
return next(iter(deployment_pricing_models(deployment)), "")
def deployment_backend_models(deployments: Iterable[object]) -> tuple[str, ...]:
return tuple(map(deployment_backend_model, deployments))
def _named_models(*candidates: object) -> Iterator[str]:
return (candidate for candidate in candidates if isinstance(candidate, str) and candidate)

4
litellm/litellm_core_utils/error_normalization.py
View file

@ -133,6 +133,8 @@ _STATUS_CODE_MAP: Final[Mapping[str, str]] = MappingProxyType(
}
)
_OVERLOADED_STATUS_CODE: Final = "529"
_INTERNAL_STATE_EXCEPTIONS: Final[tuple[type[BaseException], ...]] = (
TypeError,
KeyError,
@ -201,6 +203,8 @@ def normalize_error(exc: Exception | None, status_code: str, message: str) -> st
)
if by_message is not None:
return by_message
if status_code == _OVERLOADED_STATUS_CODE:
return PROVIDER_OVERLOADED
by_class: Final = _classify_by_class(exc)
if by_class is not None:
return by_class

10
litellm/litellm_core_utils/exception_mapping_utils.py
View file

@ -11,6 +11,7 @@ from pydantic import TypeAdapter, ValidationError
import litellm
from litellm._logging import _ENABLE_SECRET_REDACTION, redact_string, verbose_logger
from litellm.litellm_core_utils.agentic_followup_failure import carry_agentic_followup_mark
from litellm.litellm_core_utils.bug_report import (
bug_report_notice,
build_bug_report,
@ -605,6 +606,11 @@ def _map_openai_exception(
)
def _anthropic_overload_status(original_exception: _ProviderHTTPException) -> int | None:
upstream_status: Final = getattr(original_exception, "status_code", None)
return 529 if isinstance(upstream_status, int) and upstream_status >= 500 else None
def _map_anthropic_exception(
*,
model: str,
@ -630,6 +636,7 @@ def _map_anthropic_exception(
message=f"AnthropicError - {error_str}",
model=model,
llm_provider="anthropic",
provider_status_code=_anthropic_overload_status(original_exception),
)
if "Invalid API Key" in error_str:
raise AuthenticationError(
@ -694,6 +701,7 @@ def _map_anthropic_exception(
llm_provider="anthropic",
model=model,
response=getattr(original_exception, "response", None),
provider_status_code=original_exception.status_code,
)
elif original_exception.status_code == 502:
raise BadGatewayError(
@ -2895,6 +2903,7 @@ def exception_type(
request=httpx.Request(method="POST", url="https://api.openai.com/v1/"), # stub the request
)
except Exception as e:
carry_agentic_followup_mark(source=original_exception, mapped=e)
# LOGGING
exception_logging(
logger_fn=None,
@ -2920,6 +2929,7 @@ def exception_type(
model="",
)
setattr(raised_exc, "litellm_response_headers", litellm_response_headers)
carry_agentic_followup_mark(source=original_exception, mapped=raised_exc)
raise raised_exc

4
litellm/litellm_core_utils/get_litellm_params.py
View file

@ -33,6 +33,9 @@ AWS_CREDENTIAL_KWARGS_KEYS: Final = frozenset(
)
PROVIDER_AFFINITY_HEADER_KWARG_KEY: Final = "provider_affinity_header"
AGENTIC_LOOP_CONTROL_KWARGS_KEYS: Final = frozenset(
{"_agentic_loop_depth", "_agentic_loop_fingerprints", "max_agentic_loops"}
)
OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS: Final = frozenset(
{
"token_exchange_endpoint",
@ -94,6 +97,7 @@ OPTIONAL_KWARGS_KEYS: Final = (
| ANTHROPIC_WIF_KWARGS_KEYS
| OPENAI_WIF_KWARGS_KEYS
| OAUTH_TOKEN_EXCHANGE_KWARGS_KEYS
| AGENTIC_LOOP_CONTROL_KWARGS_KEYS
| frozenset(CustomPricingLiteLLMParams.model_fields)
)

5
litellm/litellm_core_utils/get_llm_provider_logic.py
View file

@ -13,6 +13,7 @@ from litellm.litellm_core_utils.fallback_generalizations import (
match_routing_generalization,
)
from litellm.llms.openai_like.json_loader import JSONProviderRegistry
from litellm.llms.opencode.common_utils import is_opencode_api_base
from litellm.secret_managers.main import get_secret, get_secret_str
from ..types.router import GenericLiteLLMParams, LiteLLM_Params
@ -1008,6 +1009,4 @@ def _get_openai_compatible_provider_info(
def get_provider_affinity_header(api_base: str | None) -> str | None:
from litellm.llms.opencode.transformation import OpenCodeProviderAdapter
return OpenCodeProviderAdapter().get_provider_affinity_header(api_base)
return "x-opencode-session" if is_opencode_api_base(api_base) else None

4
litellm/litellm_core_utils/get_model_cost_map.py
View file

@ -320,9 +320,7 @@ def _classify_fetch_content(content: bytes, url: str, etag: str | None) -> _Fetc
return ModelCostMapReloadUnavailable(reason=f"invalid JSON from {url}: {e}")
if not isinstance(parsed, dict):
return ModelCostMapReloadUnavailable(reason=f"expected a JSON object from {url}, got {type(parsed).__name__}")
model_cost_map: Final = cast( # cast-ok: JSON object keys are strings by definition
dict[str, object], parsed
)
model_cost_map: Final = _BUNDLED_CATALOG_ADAPTER.validate_python(parsed)
return ModelCostMapReloaded(model_cost_map=model_cost_map, revision=git_blob_id(content), etag=etag)

12
litellm/litellm_core_utils/get_provider_specific_headers.py
View file

@ -12,10 +12,18 @@ _PROVIDER_HEADERS: Final[TypeAdapter[ProviderSpecificHeader | Sequence[ProviderS
class ProviderSpecificHeaderUtils:
@staticmethod
def forwarded_authorization(provider_specific_header: object, custom_llm_provider: str | None) -> str | None:
headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
def validated_provider_specific_headers(
provider_specific_header: object, custom_llm_provider: str | None
) -> Mapping[str, str]:
return ProviderSpecificHeaderUtils.get_provider_specific_headers(
_PROVIDER_HEADERS.validate_python(provider_specific_header), custom_llm_provider
)
@staticmethod
def forwarded_authorization(provider_specific_header: object, custom_llm_provider: str | None) -> str | None:
headers: Final = ProviderSpecificHeaderUtils.validated_provider_specific_headers(
provider_specific_header, custom_llm_provider
)
return next((value for key, value in headers.items() if key.lower() == "authorization"), None)
@staticmethod

205
litellm/litellm_core_utils/litellm_logging.py
View file

@ -54,8 +54,11 @@ from litellm.constants import (
REDACTED_BY_LITELLM,
)
from litellm.cost_calculator import (
BaseTokenUsageProcessor,
RealtimeAPITokenUsageProcessor,
ResponsesWebSocketTokenUsageProcessor,
declared_cost_map_pricing,
get_response_cost_from_hidden_params,
get_usage_object,
pricing_entry_for_cost_calc,
select_model_name_for_cost_calc,
@ -74,6 +77,8 @@ from litellm.integrations.deepeval.deepeval import DeepEvalLogger
from litellm.integrations.langtrace import langtrace_trace_endpoint
from litellm.integrations.mlflow import MlflowLogger
from litellm.integrations.sqs import SQSLogger
from litellm.litellm_core_utils.add_retry_fallback_headers import get_hidden_params_dict
from litellm.litellm_core_utils.agentic_followup_failure import ServedRound
from litellm.litellm_core_utils.classifier_logging import (
classifier_audit_fields,
classifier_input_snapshot,
@ -215,7 +220,10 @@ from litellm.types.utils import (
batch_chat_usage,
)
from litellm.types.videos.main import VideoObject
from litellm.utils import get_base_model_from_metadata, print_verbose
from litellm.utils import (
get_base_model_from_metadata,
print_verbose,
)
from ..integrations.argilla import ArgillaLogger
from ..integrations.arize.arize_phoenix import ArizePhoenixLogger
@ -716,6 +724,7 @@ class Logging(LiteLLMLoggingBaseClass):
custom_pricing: bool = False
custom_cost_per_token: CostPerToken | None = None
custom_cost_per_second: float | None = None
custom_model_info: ModelInfo | None = None
replaced_global_callback_keys: frozenset[str] = frozenset()
stream_options = None
litellm_request_debug: bool = False
@ -937,6 +946,11 @@ class Logging(LiteLLMLoggingBaseClass):
self._sync_success_scheduled: bool = False
self.on_detached_stream_failure: Callable[[Exception], Awaitable[None]] | None = None
self._stream_failure_claimed: bool = False
self._stash_owed_on_success: bool = False
self._stash_billed_by_continuation: bool = False
self._failed_attempts_owed_on_success: tuple[ServedRound, ...] = ()
self._failed_attempt_ids_owed_on_success: tuple[str, ...] = ()
self.failed_attempts_billed_on_success: tuple[ServedRound, ...] = ()
self.shadow_eval_request_snapshot: GuardrailRequestSnapshot | None = None
def copy_for_new_call(self, litellm_call_id: str, start_time: datetime.datetime) -> "Logging":
@ -2221,6 +2235,7 @@ class Logging(LiteLLMLoggingBaseClass):
"custom_pricing": custom_pricing,
"custom_cost_per_token": self.custom_cost_per_token if custom_pricing else None,
"custom_cost_per_second": self.custom_cost_per_second if custom_pricing else None,
"requested_model_info_override": self.custom_model_info if custom_pricing else None,
"prompt": prompt,
"standard_built_in_tools_params": self.standard_built_in_tools_params,
"router_model_id": router_model_id,
@ -2312,12 +2327,19 @@ class Logging(LiteLLMLoggingBaseClass):
return 0.0
def uncached_price(self, result: object) -> UncachedPrice:
return self._price_without_recording(result, cache_hit=False, ignore_stamped_cost=True)
def _price_without_recording(
self, result: object, *, cache_hit: bool | None, ignore_stamped_cost: bool
) -> UncachedPrice:
previous_cost_breakdown: Final = self.cost_breakdown
previous_billed_token_rates: Final = self.billed_token_rates
debug_key: Final = "response_cost_failure_debug_information"
previous_failure_debug_info: Final = self.model_call_details.get(debug_key)
try:
cost: Final = self.response_cost_calculator(result=result, cache_hit=False, ignore_stamped_cost=True)
cost: Final = self.response_cost_calculator(
result=result, cache_hit=cache_hit, ignore_stamped_cost=ignore_stamped_cost
)
return UncachedPrice(cost=cost, cost_breakdown=self.cost_breakdown)
finally:
self.cost_breakdown = previous_cost_breakdown
@ -2417,6 +2439,10 @@ class Logging(LiteLLMLoggingBaseClass):
additional_headers=additional_headers,
)
def set_per_call_pricing(self, custom_model_info: ModelInfo | None) -> None:
"""Price this call from its own full price entry instead of ``litellm.model_cost``."""
self.custom_model_info = custom_model_info
def custom_pricing_for(self, result: object) -> bool:
if self.spilled_over_ptu_for(result):
return False
@ -2582,6 +2608,33 @@ class Logging(LiteLLMLoggingBaseClass):
**kwargs,
)
def _response_price_is_known(self, result: object, cost: float, payload: StandardLoggingPayload) -> bool:
if (
cost != 0
or self._provider_reported_cost == 0
or get_response_cost_from_hidden_params(get_hidden_params_dict(result)) == 0
):
return True
if self.model_call_details.get("cache_hit") is True or _served_from_cache_at_no_charge(result):
return True
params: Final = _LITELLM_PARAMS.validate_python(self._model_call_details_view.get("litellm_params") or {})
if is_unbilled_non_inference_call(
payload["call_type"], StandardLoggingPayloadSetup.merge_litellm_metadata(params), result
):
return True
if self.custom_cost_per_token is not None or self.custom_cost_per_second is not None:
return True
passthrough: Final = _LITELLM_PARAMS.validate_python(
self._model_call_details_view.get("passthrough_logging_payload") or {}
)
if passthrough.get("cost_per_request") is not None:
return True
model_info: Final = payload["model_map_information"]["model_map_value"]
if model_info is None:
return False
prices: Final = declared_cost_map_pricing(model_info["key"], model_info["key"], payload["custom_llm_provider"])
return any(value is not None and ("cost" in key or "pricing" in key) for key, value in prices.items())
def prepare_response_cost(self, response: object, now: datetime.datetime) -> PreparedCost | None:
if self.prepared_cost is not None:
return self.prepared_cost
@ -2634,6 +2687,7 @@ class Logging(LiteLLMLoggingBaseClass):
"cost_breakdown": copy.deepcopy(payload["cost_breakdown"]),
},
usage,
price_known=self._response_price_is_known(prepared[2], raw_cost, payload),
)
self.record_partial_usage_for_failure(
usage, raw_cost, observation=self.model_call_details["_accounting_usage_observed"]
@ -2737,27 +2791,43 @@ class Logging(LiteLLMLoggingBaseClass):
own usage while the request's failure log still bills every attempt."""
usage: Final = self._model_call_details_view.get("combined_usage_object")
response_cost: Final = self._model_call_details_view.get("response_cost")
if isinstance(usage, Usage) and type(self).usage_capture_callback is not None:
billed_by_continuation: Final = self._stash_billed_by_continuation
owed_on_success: Final = self._stash_owed_on_success
self._stash_billed_by_continuation = False
self._stash_owed_on_success = False
captured_attempt_ids: Final = (
self.capture_failed_attempt_usage(usage)
if isinstance(usage, Usage) and type(self).usage_capture_callback is not None
else ()
)
self.model_call_details.pop("_accounting_capture", None)
self.model_call_details.pop("_accounting_pricing_snapshot", None)
self.model_call_details.pop("_accounting_usage_observed", None)
self.model_call_details.pop("combined_usage_object", None)
self.model_call_details.pop("response_cost", None)
if isinstance(usage, Usage):
self.model_call_details["current_attempt_api_call_start_time"] = None
self.model_call_details["failed_attempts_usage"] = (
*self.model_call_details.get("failed_attempts_usage", ()),
(usage, float(response_cost) if isinstance(response_cost, (int, float)) else 0.0),
)
if not isinstance(usage, Usage):
return
attempt: Final = (usage, float(response_cost) if isinstance(response_cost, (int, float)) else 0.0)
self.model_call_details["current_attempt_api_call_start_time"] = None
self.model_call_details["failed_attempts_usage"] = (
*self.model_call_details.get("failed_attempts_usage", ()),
attempt,
)
if not owed_on_success or billed_by_continuation:
return
self._failed_attempts_owed_on_success = (*self._failed_attempts_owed_on_success, attempt)
self._failed_attempt_ids_owed_on_success = (*self._failed_attempt_ids_owed_on_success, *captured_attempt_ids)
def capture_failed_attempt_usage(self, usage: Usage) -> None:
def mark_stash_billed_by_continuation(self) -> None:
self._stash_billed_by_continuation = True
def capture_failed_attempt_usage(self, usage: Usage) -> tuple[str, ...]:
from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY
from litellm.proxy.spend_tracking.usage_capture import CapturedUsage
capture: Final = type(self).usage_capture_callback
if capture is None:
return
return ()
attempt: Final = self.attempt_scoped_copy()
attempt.prepared_cost = None
attempt.finalized_cost = None
@ -2792,11 +2862,13 @@ class Logging(LiteLLMLoggingBaseClass):
datetime.datetime.now(datetime.timezone.utc),
False,
)
if isinstance(captured, CapturedUsage):
self.model_call_details["_accounting_prior_captures"] = (
*self.model_call_details.get("_accounting_prior_captures", ()),
captured,
)
if not isinstance(captured, CapturedUsage):
return ()
self.model_call_details["_accounting_prior_captures"] = (
*self.model_call_details.get("_accounting_prior_captures", ()),
captured,
)
return tuple(write.event.event_id for write in captured.writes)
def record_continued_partial(self, usage: Usage, response_cost: float, response_id: str) -> None:
"""Bill what a cut stream already consumed on this continuation's success log, under the id the caller saw."""
@ -2804,20 +2876,17 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["continued_partial_cost"] = response_cost
self.model_call_details["continued_response_id"] = response_id
def _add_continued_partial_spend(self, response: object) -> None:
def _add_continued_partial_spend(self, response: object) -> Usage | None:
partial_usage: Final = self.model_call_details.get("continued_partial_usage")
partial_cost: Final = self.model_call_details.get("continued_partial_cost")
own_usage: Final = getattr(response, "usage", None)
own_cost: Final = self.model_call_details.get("response_cost")
if not (isinstance(partial_usage, Usage) and isinstance(partial_cost, float) and isinstance(own_usage, Usage)):
return
from litellm.cost_calculator import BaseTokenUsageProcessor
self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects(
usage_objects=[own_usage, partial_usage]
)
return own_usage if isinstance(own_usage, Usage) else None
combined_usage: Final = BaseTokenUsageProcessor.combine_usage_objects(usage_objects=[own_usage, partial_usage])
self.model_call_details["combined_usage_object"] = combined_usage
if not isinstance(own_cost, (int, float)):
return
return combined_usage
self.model_call_details["response_cost"] = own_cost + partial_cost
if self.cost_breakdown is not None:
self.cost_breakdown = self.cost_breakdown | CostBreakdown(
@ -2827,6 +2896,27 @@ class Logging(LiteLLMLoggingBaseClass):
"mid_stream_partial_cost": partial_cost,
},
)
return combined_usage
def _add_failed_attempts_spend(self, own_usage: Usage | None) -> None:
own_cost: Final = self.model_call_details.get("response_cost")
if own_usage is None or not isinstance(own_cost, (int, float)):
return
failed_attempts: Final = self._failed_attempts_owed_on_success
if not failed_attempts:
return
failed_attempts_cost: Final = sum(cost for _, cost in failed_attempts)
self.model_call_details.pop("failed_attempts_usage", None)
self.failed_attempts_billed_on_success = failed_attempts
self.model_call_details["failed_attempts_billed_on_success"] = failed_attempts
if self._failed_attempt_ids_owed_on_success:
self.model_call_details["_accounting_included_attempt_ids"] = self._failed_attempt_ids_owed_on_success
self.model_call_details["combined_usage_object"] = BaseTokenUsageProcessor.combine_usage_objects(
usage_objects=[own_usage, *(usage for usage, _ in failed_attempts)]
)
self.model_call_details["response_cost"] = own_cost + failed_attempts_cost
if self.cost_breakdown is not None:
self.cost_breakdown = _with_failed_attempts_cost(self.cost_breakdown, own_cost, failed_attempts_cost)
def claim_stream_failure_report(self) -> bool:
if self._stream_failure_claimed:
@ -2834,6 +2924,24 @@ class Logging(LiteLLMLoggingBaseClass):
self._stream_failure_claimed = True
return True
def served_round(self, response: object) -> ServedRound | None:
usage: Final = get_usage_object(completion_response=response)
if usage is None:
return None
price: Final = self._price_without_recording(response, cache_hit=None, ignore_stamped_cost=False)
return usage, price.cost or 0.0
def record_served_rounds_for_failure(self, served_rounds: tuple[ServedRound, ...]) -> None:
if not served_rounds:
return
self._stash_owed_on_success = True
self.record_partial_usage_for_failure(
served_rounds[0][0]
if len(served_rounds) == 1
else BaseTokenUsageProcessor.combine_usage_objects(usage_objects=[usage for usage, _ in served_rounds]),
sum(cost for _, cost in served_rounds),
)
def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None:
"""Bill a fully streamed response on the failure log when a post-call hook rejects it."""
usage: Final = getattr(assembled, "usage", None)
@ -3138,9 +3246,11 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["response_cost"] = self.finalized_cost.amount
elif self.model_call_details.get("cache_hit") is True:
self.model_call_details["response_cost"] = 0.0
self._add_failed_attempts_spend(get_usage_object(completion_response=logging_result))
elif "response_cost" in hidden_params:
self.model_call_details["response_cost"] = hidden_params["response_cost"]
self._record_zero_cost_diagnostic(logging_result, hidden_params["response_cost"])
self._add_failed_attempts_spend(get_usage_object(completion_response=logging_result))
elif (existing_cost := self.model_call_details.get("response_cost")) is not None and (
existing_cost != 0 or self.call_type == CallTypes.pass_through.value # pyright: ignore[reportUnknownMemberType] # Logging.call_type is untyped
):
@ -3150,6 +3260,7 @@ class Logging(LiteLLMLoggingBaseClass):
pass
else:
self.model_call_details["response_cost"] = self.response_cost_calculator(result=logging_result)
self._add_failed_attempts_spend(get_usage_object(completion_response=logging_result))
if not build_logging_payload:
return None
@ -3278,21 +3389,12 @@ class Logging(LiteLLMLoggingBaseClass):
set_hidden_params(served, {**(get_hidden_params(translated) or {}), **(get_hidden_params(response) or {})})
return served
def _with_continued_response_id(
def _with_response_id(
self,
response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse,
response_id: str | None,
) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse:
response_id: Final = self.model_call_details.get("continued_response_id")
if not isinstance(response_id, str) or getattr(response, "id", None) == response_id:
return response
return response.model_copy(update={"id": response_id})
def _with_client_facing_response_id(
self,
response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse,
) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse:
response_id: Final = self.client_facing_response_id
if response_id is None or getattr(response, "id", None) == response_id:
if response_id is None or response.id == response_id:
return response
return response.model_copy(update={"id": response_id})
@ -3300,8 +3402,13 @@ class Logging(LiteLLMLoggingBaseClass):
self,
response: ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse,
) -> ModelResponse | TextCompletionResponse | ResponsesAPIResponse | InteractionsAPIResponse:
return self._with_client_facing_response_id(
self._with_continued_response_id(self._with_client_facing_stream_model(response))
continued_response_id: Final = self.model_call_details.get("continued_response_id")
return self._with_response_id(
self._with_response_id(
self._with_client_facing_stream_model(response),
continued_response_id if isinstance(continued_response_id, str) else None,
),
self.client_facing_response_id,
)
def _success_handler_helper_fn(
@ -3713,7 +3820,7 @@ class Logging(LiteLLMLoggingBaseClass):
self.model_call_details["response_cost"] = self.response_cost_calculator(
result=complete_streaming_response if provider_transcription is None else provider_transcription
)
self._add_continued_partial_spend(complete_streaming_response)
self._add_failed_attempts_spend(self._add_continued_partial_spend(complete_streaming_response))
self._merge_hidden_params_from_response_into_metadata(complete_streaming_response)
logged_streaming_response: Final = self._client_facing_streaming_response(
self._as_served_chat_response(complete_streaming_response)
@ -4309,7 +4416,7 @@ class Logging(LiteLLMLoggingBaseClass):
self.model,
)
self.model_call_details["response_cost"] = None
self._add_continued_partial_spend(complete_streaming_response)
self._add_failed_attempts_spend(self._add_continued_partial_spend(complete_streaming_response))
self._merge_hidden_params_from_response_into_metadata(complete_streaming_response)
@ -6388,7 +6495,7 @@ def _init_custom_logger_compatible_class(
logging_integration: str,
internal_usage_cache: DualCache | None,
llm_router: object,
custom_logger_init_args: dict[str, object] | None = {},
custom_logger_init_args: dict[str, object] | None = None,
raise_import_errors: bool = False,
) -> CustomLogger | None:
"""
@ -7802,6 +7909,20 @@ def _cache_hit_cost_breakdown(served: CostBreakdown | None) -> CostBreakdown:
)
def _with_failed_attempts_cost(breakdown: CostBreakdown, own_cost: float, failed_attempts_cost: float) -> CostBreakdown:
return breakdown | CostBreakdown(
total_cost=own_cost + failed_attempts_cost,
additional_costs={**breakdown.get("additional_costs", {}), "failed_attempts_cost": failed_attempts_cost},
)
def _cached_cost_breakdown(served: CostBreakdown | None, failed_attempts: tuple[ServedRound, ...]) -> CostBreakdown:
cached: Final = _cache_hit_cost_breakdown(served)
if not failed_attempts:
return cached
return _with_failed_attempts_cost(cached, 0.0, sum(cost for _, cost in failed_attempts))
def _finalized_cost_of(logging_obj: object) -> FinalizedCost | None:
"""Duck-typed logging objects (and test doubles) carry no finalized cost."""
finalized: Final = getattr(logging_obj, "finalized_cost", None)
@ -7995,7 +8116,7 @@ def get_standard_logging_object_payload(
if finalized_cost is not None
else cost_breakdown_with_guardrail(
(
_cache_hit_cost_breakdown(uncached_price.cost_breakdown)
_cached_cost_breakdown(uncached_price.cost_breakdown, logging_obj.failed_attempts_billed_on_success)
if uncached_price is not None
else logging_obj.cost_breakdown
),

1
litellm/litellm_core_utils/llm_cost_calc/finalization.py
View file

@ -17,6 +17,7 @@ class PreparedCost:
amount: float
payload: StandardLoggingPayload
usage: Usage
price_known: bool = True
@dataclass(frozen=True, slots=True)

55
litellm/litellm_core_utils/llm_cost_calc/utils.py
View file

@ -1203,21 +1203,55 @@ def parse_completion_tokens_details(usage: Usage) -> CompletionTokensDetailsResu
)
def get_modality_token_cost_key(
model_info: ModelInfo,
cost_key: str,
prompt_tokens: int,
*,
threshold_is_inclusive: bool = False,
) -> str:
suffix: Final = next((suffix for suffix in _NON_STANDARD_THRESHOLD_SUFFIXES if cost_key.endswith(suffix)), "")
base_key: Final = cost_key.removesuffix(suffix)
threshold_key: Final = max(
(
key
for key in model_info
if key.startswith(f"{base_key}_above_")
and key.endswith(("_tokens", f"_tokens{suffix}"))
and model_info.get(key) is not None
and _prompt_exceeds_threshold(prompt_tokens, _parse_above_token_threshold(key), threshold_is_inclusive)
),
key=_parse_above_token_threshold,
default=None,
)
if threshold_key is None:
return cost_key
return f"{threshold_key.removesuffix(suffix)}{suffix}"
def _modality_token_cost(
model_info: ModelInfo,
modality_cost_key: str,
tokens: int,
prompt_base_cost: float,
service_tier: str | None,
*,
prompt_tokens: int = 0,
threshold_is_inclusive: bool = False,
) -> float:
"""
Audio, image and video prompt tokens bill at their modality rate when the model has one.
Without one they are ordinary prompt tokens, so they bill at `prompt_base_cost`, which
`_get_token_base_cost` has already resolved for the service tier and long-context tier.
"""
tiered_key: Final = get_service_tier_cost_key(modality_cost_key, service_tier)
cost_key: Final = tiered_key if model_info.get(tiered_key) is not None else modality_cost_key
if model_info.get(cost_key) is None:
cost_key: Final = get_modality_token_cost_key(
model_info,
get_service_tier_cost_key(modality_cost_key, service_tier),
prompt_tokens,
threshold_is_inclusive=threshold_is_inclusive,
)
standard_key: Final = cost_key.removesuffix(get_service_tier_cost_key("", service_tier))
if model_info.get(cost_key) is None and model_info.get(standard_key) is None:
return float(tokens) * prompt_base_cost
return calculate_cost_component(model_info, cost_key, tokens)
@ -1230,6 +1264,9 @@ def _calculate_input_cost(
cache_creation_cost: float,
cache_creation_cost_by_ttl: Mapping[str, float],
service_tier: str | None = None,
*,
prompt_tokens: int = 0,
threshold_is_inclusive: bool = False,
) -> float:
"""
Calculates the input cost for a given model, prompt tokens, and completion tokens.
@ -1260,6 +1297,8 @@ def _calculate_input_cost(
prompt_tokens_details["audio_tokens"],
prompt_base_cost,
service_tier,
prompt_tokens=prompt_tokens,
threshold_is_inclusive=threshold_is_inclusive,
)
### IMAGE TOKEN COST
@ -1618,6 +1657,8 @@ def generic_cost_per_token(
cache_creation_cost=base_costs.cache_creation_rate,
cache_creation_cost_by_ttl=base_costs.cache_creation_rates_by_ttl,
service_tier=service_tier,
prompt_tokens=usage.prompt_tokens,
threshold_is_inclusive=_uses_inclusive_token_thresholds(custom_llm_provider),
)
## CALCULATE OUTPUT COST
@ -1675,11 +1716,9 @@ def generic_cost_per_token(
## IMAGE COST
if not is_text_tokens_total and image_tokens and image_tokens > 0:
_output_cost_per_image_token = get_cost_per_unit(resolved_model_info, "output_cost_per_image_token", None)
_output_cost_per_image_token = (
_output_cost_per_image_token if _output_cost_per_image_token is not None else completion_base_cost
)
completion_cost += float(image_tokens) * _output_cost_per_image_token
image_cost_key: Final = get_service_tier_cost_key("output_cost_per_image_token", service_tier)
image_rate: Final = get_cost_per_unit(resolved_model_info, image_cost_key, None)
completion_cost += float(image_tokens) * (completion_base_cost if image_rate is None else image_rate)
## VIDEO COST
if not is_text_tokens_total and video_tokens and video_tokens > 0:

104
litellm/litellm_core_utils/per_call_pricing.py Normal file
View file

@ -0,0 +1,104 @@
"""The price entry for a single call that sets its own custom pricing, built without writing ``litellm.model_cost``."""
from collections.abc import Mapping
from typing import Final, Literal
from pydantic import TypeAdapter
import litellm
from litellm.litellm_core_utils.ptu_pricing import zeroed_ptu_pricing
from litellm.types.utils import CustomPricingLiteLLMParams, ModelInfo
_STR_MAPPING: Final = TypeAdapter(Mapping[str, object])
_MODEL_INFO: Final = TypeAdapter(ModelInfo)
def _router_deployment_model_info(kwargs: Mapping[str, object]) -> Mapping[str, object] | None:
"""A Router stamps the deployment's ``model_info`` (with its ``id``) into the call's metadata and registers that
deployment's full custom pricing under that id when the deployment is added."""
for metadata_key in ("litellm_metadata", "metadata"):
metadata = kwargs.get(metadata_key)
if not isinstance(metadata, Mapping):
continue
model_info = _STR_MAPPING.validate_python(metadata).get("model_info")
if isinstance(model_info, Mapping) and isinstance(model_info.get("id"), str):
return _STR_MAPPING.validate_python(model_info)
return None
def per_call_model_info(
model: str, custom_llm_provider: str, kwargs: Mapping[str, object], mode: Literal["chat", "embedding"]
) -> ModelInfo | None:
"""The full price entry for one call, or None when the cost map (or the Router's registered deployment entry)
already prices it. The cost map row for the model is overlaid with the registered deployment entry and then with
the custom pricing fields the request itself set, so a Router call pays request overrides over deployment prices.
The entry travels on the logging object and is never written into ``litellm.model_cost``, so the next call for the
same model name is priced from the cost map again."""
deployment: Final = _router_deployment_model_info(kwargs)
registered: Final = litellm.model_cost.get(str(deployment["id"])) if deployment is not None else None
custom: Final = {
field: kwargs[field]
for field in CustomPricingLiteLLMParams.model_fields
if kwargs.get(field) is not None and (registered is None or registered.get(field) != kwargs[field])
}
if not custom:
return None
unmapped: Final[Mapping[str, object]] = {
"key": model,
"max_tokens": None,
"max_input_tokens": None,
"max_output_tokens": None,
"input_cost_per_token": None,
"output_cost_per_token": None,
"litellm_provider": custom_llm_provider,
"mode": mode,
"supported_openai_params": None,
}
row: Final = _catalog_row(model, custom_llm_provider) or unmapped
catalog: Final = (
{**row, **{key: value for key, value in registered.items() if value is not None}} if registered else row
)
zeroed: Final = zeroed_ptu_pricing(deployment, custom) if deployment is not None else None
return _MODEL_INFO.validate_python(
{
**catalog,
**_token_rates_cleared_for_per_second_overrides(custom),
**custom,
**_tiers_with_flat_overrides(catalog, custom),
**(zeroed or {}),
"mode": catalog.get("mode") or mode,
}
)
_FLAT_TOKEN_RATE_FIELDS: Final = ("input_cost_per_token", "output_cost_per_token")
_PER_SECOND_RATE_FIELDS: Final = ("cost_per_second", "input_cost_per_second", "output_cost_per_second")
_TIERS: Final = TypeAdapter(tuple[Mapping[str, object], ...])
def _token_rates_cleared_for_per_second_overrides(custom: Mapping[str, object]) -> Mapping[str, object]:
"""Token rates in the entry make the cost calculator skip per-second pricing, so a caller who sets only a
per-second rate gets the catalog's token rates and tier table cleared for this call."""
if not any(key in custom for key in _PER_SECOND_RATE_FIELDS) or any(
key in custom for key in _FLAT_TOKEN_RATE_FIELDS
):
return {}
return {"input_cost_per_token": None, "output_cost_per_token": None, "tiered_pricing": None}
def _tiers_with_flat_overrides(catalog: Mapping[str, object], custom: Mapping[str, object]) -> Mapping[str, object]:
"""A ``tiered_pricing`` table wins over flat rates in the cost calculator, so a flat rate the caller set is written
into every tier of the table in force (the caller's own, else the inherited one) for that direction, and the other
direction keeps each tier's own rate."""
overrides: Final = {key: custom[key] for key in _FLAT_TOKEN_RATE_FIELDS if key in custom}
tiers: Final = custom.get("tiered_pricing", catalog.get("tiered_pricing"))
if not overrides or not isinstance(tiers, (list, tuple)) or not tiers:
return {}
return {"tiered_pricing": [{**tier, **overrides} for tier in _TIERS.validate_python(tiers)]}
def _catalog_row(model: str, custom_llm_provider: str) -> ModelInfo | None:
try:
return litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider)
except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model
return None

68
litellm/litellm_core_utils/prompt_templates/common_utils.py
View file

@ -399,6 +399,74 @@ def split_media_for_tpm(value: object) -> tuple[object, tuple[Mapping[str, objec
return projection.content, projection.media
_TPM_EMBEDDING_FILE_MIME_PREFIXES: Final = ("application/pdf", "video/")
def _tpm_embedding_file(data_uri: str) -> tuple[Mapping[str, object], ...]:
from litellm.litellm_core_utils.prompt_templates.factory import parse_mime_type
mime: Final = (parse_mime_type(data_uri[:256]) or "").strip().lower()
if not mime.startswith(_TPM_EMBEDDING_FILE_MIME_PREFIXES):
return ()
return ({"type": "file", "file": {"file_data": data_uri, "format": mime}},)
_BASE64_PAYLOAD: Final = re.compile(r"[A-Za-z0-9+/]*={0,2}")
def _is_base64_data_uri(value: str) -> bool:
header_end: Final = value.find(",", 0, 256)
if header_end < 0 or not value[:header_end].endswith(";base64"):
return False
payload_start: Final = header_end + 1
if (len(value) - payload_start) % 4 or _BASE64_PAYLOAD.fullmatch(value, payload_start) is None:
return False
last_quantum: Final = value[max(payload_start, len(value) - 4) :]
return base64.b64encode(base64.b64decode(last_quantum)).decode("ascii") == last_quantum
def _tpm_embedding_file_block(part: Mapping[str, object]) -> _TPMMediaProjection:
projection: Final = _tpm_file_media(part)
fields: Final = as_str_mapping(part.get("file"))
data: Final = fields.get("file_data") if fields is not None else None
if projection.media or fields is None or not isinstance(data, str) or not _is_base64_data_uri(data):
return projection
media: Final = _tpm_embedding_file(data)
return _TPMMediaProjection({**part, "file": _tpm_blank_fields(fields, ("file_data",))} if media else part, media)
def _tpm_embedding_leaf(value: object) -> _TPMMediaProjection:
part: Final = as_str_mapping(value)
if part is not None and part.get("type") == "file":
return _tpm_embedding_file_block(part)
if not isinstance(value, str) or not value.startswith("data:"):
return _tpm_media_leaf(value)
if not _is_base64_data_uri(value):
return _TPMMediaProjection(value)
media: Final = _tpm_file_media({"file_data": value}).media or _tpm_embedding_file(value)
return _TPMMediaProjection("" if media else value, media)
def _is_json_array(value: object) -> bool:
return isinstance(value, (list, tuple))
def _tpm_embedding_item(value: object) -> _TPMMediaProjection:
return _tpm_media_list(value, _tpm_embedding_leaf) if _is_json_array(value) else _tpm_embedding_leaf(value)
def split_embedding_media_for_tpm(value: object) -> tuple[object, tuple[Mapping[str, object], ...]]:
projection: Final = (
_tpm_media_list(value, _tpm_embedding_item) if _is_json_array(value) else _tpm_embedding_leaf(value)
)
return projection.content, projection.media
def embedding_input_is_all_media(value: object) -> bool:
items: Final = cast(Sequence[object], value) if _is_json_array(value) else (value,) # cast-ok: JSON input array
return all(_tpm_embedding_leaf(item).media for item in items)
def _content_parts_contain_image(parts: Sequence[object]) -> bool:
"""Depth-bounded frontier walk over nested content lists, iterative because the repo bans
recursion; an Anthropic tool_result nests its image parts exactly one level down."""

2
litellm/litellm_core_utils/streaming_handler.py
View file

@ -2826,8 +2826,6 @@ class CustomStreamWrapper:
return None if is_model_response_stream_empty(model_response=usage_free_chunk) else usage_free_chunk
def _accumulate_streamed_delta(self, delta: object) -> None:
# Shared by the sync, async, and non-aiohttp iteration sites so answer
# text and the disqualifying-content latch stay in step across all three.
get: Final = getattr(delta, "get", None)
content: Final = get("content") if callable(get) else None
self.response_uptil_now += content if isinstance(content, str) else ""

121
litellm/litellm_core_utils/token_counter.py
View file

@ -610,6 +610,127 @@ def estimate_image_tokens(
return max(_image_policy_tokens(family, dimensions, detail) for family in families)
_MEDIA_FILE_CHARS_PER_TOKEN: Final = 4
_EmbeddingInputT: Final = TypeVar("_EmbeddingInputT")
def _max_input_tokens(model: str) -> int | None:
try:
limit: Final = litellm.get_model_info(model).get("max_input_tokens")
except Exception: # noqa: BLE001 # get_model_info raises a bare Exception for an unmapped model
return None
return limit if isinstance(limit, int) and limit > 0 else None
def _media_file_tokens(data: str, models: Sequence[str]) -> int:
text_estimate: Final = len(data) // _MEDIA_FILE_CHARS_PER_TOKEN
ceilings: Final = tuple(map(_max_input_tokens, models))
known_ceilings: Final = tuple(ceiling for ceiling in ceilings if ceiling is not None)
if not known_ceilings or len(known_ceilings) != len(ceilings):
return text_estimate
return min(text_estimate, max(known_ceilings))
def _block_str_field(block: Mapping[str, object], container: str, field: str) -> str | None:
inner: Final = as_str_mapping(block.get(container))
value: Final = inner.get(field) if inner is not None else None
return value if isinstance(value, str) else None
def estimate_media_block_tokens(block: Mapping[str, object], models: Sequence[str], fallback_model: str | None) -> int:
from litellm.litellm_core_utils.audio_utils.utils import estimate_audio_tokens
match block.get("type"):
case "input_audio":
return estimate_audio_tokens(
_block_str_field(block, "input_audio", "data"), models, fallback_model=fallback_model
)
case "file":
return _media_file_tokens(_block_str_field(block, "file", "file_data") or "", models)
case _:
detail: Final = _block_str_field(block, "image_url", "detail")
return estimate_image_tokens(
models,
_block_str_field(block, "image_url", "url") or "",
"auto" if detail is None else detail,
fallback_model=fallback_model,
)
_PER_ITEM_DATA_URI_EMBEDDING_PROVIDERS: Final = frozenset({"vertex_ai", "gemini", "azure_ai"})
_DataUriMediaScope: TypeAlias = Literal["per_item", "whole_batch"]
def _is_pdf_block(block: Mapping[str, object]) -> bool:
return (_block_str_field(block, "file", "file_data") or "").startswith("data:application/pdf")
def _bedrock_data_uri_media_scope(
resolved_model: str, media: Sequence[Mapping[str, object]]
) -> _DataUriMediaScope | None:
from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM
from litellm.llms.bedrock.common_utils import BedrockModelInfo
from litellm.utils import supports_embedding_image_input
base_model: Final = BedrockModelInfo.get_base_model(resolved_model.removeprefix("async_invoke/"))
if not supports_embedding_image_input(f"bedrock/{base_model}"):
return None
match BaseAWSLLM.get_bedrock_embedding_provider(base_model):
case "cohere":
return "whole_batch"
case "twelvelabs":
return None
case "nova" if any(map(_is_pdf_block, media)):
return None
case _:
return "per_item"
def _data_uri_media_scope(model: str, media: Sequence[Mapping[str, object]]) -> _DataUriMediaScope | None:
from litellm.utils import supports_embedding_image_input
try:
resolved_model, provider, _, _ = litellm.get_llm_provider(model=model)
except litellm.exceptions.BadRequestError:
return None
if provider == "bedrock":
return _bedrock_data_uri_media_scope(resolved_model, media)
if not supports_embedding_image_input(model):
return None
if provider == "cohere":
return "whole_batch"
return "per_item" if provider in _PER_ITEM_DATA_URI_EMBEDDING_PROVIDERS else None
def split_embedding_input_tokens(
value: _EmbeddingInputT, media_models: Callable[[], Sequence[str]], fallback_model: str | None = None
) -> tuple[_EmbeddingInputT, int]:
from litellm.litellm_core_utils.prompt_templates.common_utils import (
embedding_input_is_all_media,
split_embedding_media_for_tpm,
)
text, media = split_embedding_media_for_tpm(value)
if not media:
return value, 0
models: Final = tuple(media_models())
scopes: Final = frozenset(_data_uri_media_scope(model, media) for model in models)
if not models or None in scopes or ("whole_batch" in scopes and not embedding_input_is_all_media(value)):
return value, 0
media_tokens: Final = sum(estimate_media_block_tokens(block, models, fallback_model) for block in media)
return cast(_EmbeddingInputT, text), media_tokens # cast-ok: the projection keeps the input's container shape
def embedding_aware_token_count(
messages: Sequence[Mapping[str, str]] | None,
text: str | list[str] | None, # mutable-ok: litellm.token_counter only takes text as list[str]
media_models: Callable[[], Sequence[str]],
) -> int:
embedding_text, media_tokens = split_embedding_input_tokens(text, media_models)
return litellm.token_counter(messages=messages, text=embedding_text) + media_tokens
TokenCounterFunction = Callable[[str], int]
"""
Type for a function that counts tokens in a string.

28
litellm/llms/anthropic/chat/transformation.py
View file

@ -28,7 +28,11 @@ from litellm.constants import (
RESPONSE_FORMAT_TOOL_NAME,
)
from litellm.litellm_core_utils.asyncify import is_event_loop_running
from litellm.litellm_core_utils.core_helpers import drops_unsupported_params, map_finish_reason
from litellm.litellm_core_utils.core_helpers import (
coerce_token_limit,
drops_unsupported_params,
map_finish_reason,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
DEFS_MAX_INLINED_BYTES,
invalid_type_message,
@ -581,12 +585,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
@classmethod
def get_config(cls, *, model: str | None = None):
def get_config(cls, *, model: str | None = None, configured_max_output_tokens: int | None = None):
config: Final = super().get_config()
# anthropic requires a default value for max_tokens
if config.get("max_tokens") is None:
config["max_tokens"] = cls.get_max_tokens_for_model(model)
config["max_tokens"] = (
configured_max_output_tokens
if configured_max_output_tokens is not None
else cls.get_max_tokens_for_model(model)
)
return config
@ -2539,7 +2547,19 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
optional_params["tools"] = tools
## Load Config
config: Final[Mapping[str, object]] = MappingProxyType(litellm.AnthropicConfig.get_config(model=model))
model_info: Final[object] = cast( # cast-ok: request parameters are a Mapping
Mapping[str, object], litellm_params
).get("model_info")
configured_max_output_tokens: Final = (
coerce_token_limit(
cast(Mapping[str, object], model_info).get("max_output_tokens") # cast-ok: runtime Mapping check
)
if isinstance(model_info, Mapping)
else None
)
config: Final[Mapping[str, object]] = MappingProxyType(
litellm.AnthropicConfig.get_config(model=model, configured_max_output_tokens=configured_max_output_tokens)
)
for k, v in config.items():
if (
k not in optional_params

9
litellm/llms/anthropic/messages/caller_model_checks.py
View file

@ -243,8 +243,7 @@ async def caller_within_model_budget(
return "allowed"
model_max_budget: Final = getattr(user_api_key_auth, "model_max_budget", None)
token: Final = getattr(user_api_key_auth, "token", None)
if isinstance(model_max_budget, dict) and model_max_budget and token is not None:
if isinstance(model_max_budget, dict) and model_max_budget and user_api_key_auth.token is not None:
try:
await model_max_budget_limiter.is_key_within_model_budget(
user_api_key_dict=user_api_key_auth,
@ -300,10 +299,8 @@ async def caller_within_model_budget(
)
return "denied_for_now"
end_user_model_max_budget: Final[dict[str, object] | None] = getattr(
user_api_key_auth, "end_user_model_max_budget", None
)
end_user_id: Final[str | None] = getattr(user_api_key_auth, "end_user_id", None)
end_user_model_max_budget: Final[dict[str, object] | None] = user_api_key_auth.end_user_model_max_budget
end_user_id: Final[str | None] = user_api_key_auth.end_user_id
if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None:
try:
await model_max_budget_limiter.is_end_user_within_model_budget(

23
litellm/llms/anthropic/messages/native/agentic_streaming_iterator.py
View file

@ -22,6 +22,7 @@ from typing import TYPE_CHECKING, Any, Final, cast
from litellm._logging import verbose_logger
from litellm.constants import STREAM_SSE_KEEPALIVE_PING_BYTES
from litellm.types.integrations.custom_logger import AgenticLoopLimitError, AgenticLoopSafetyError
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -35,10 +36,16 @@ SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES: Final = (
b'data: {"type": "error", "error": {"type": "api_error", "message": '
b'"Server-side tool retrieval failed, so this turn could not be completed. Please retry."}}\n\n'
)
AGENTIC_LOOP_LIMIT_ERROR_SSE_BYTES: Final = (
b"event: error\n"
b'data: {"type": "error", "error": {"type": "api_error", "message": '
b'"The agentic loop reached its limit before this turn finished. '
b'Raise max_agentic_loops to allow more loops."}}\n\n'
)
def is_server_fulfilled_tool_leak_error(chunk: object) -> bool:
return chunk == SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES
return chunk in (SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES, AGENTIC_LOOP_LIMIT_ERROR_SSE_BYTES)
async def _anext_or_none(iterator: AsyncIterator) -> bytes | None:
@ -216,6 +223,7 @@ class AgenticAnthropicStreamingIterator:
self._follow_up_chunk_task: asyncio.Task | None = None
self._replay_index = 0
self._error_emitted = False
self._loop_limit_reached = False
@property
def has_buffered_provider_output(self) -> bool:
@ -288,6 +296,8 @@ class AgenticAnthropicStreamingIterator:
"tool_use [model=%s]; emitting an SSE error instead of leaking the tool call to the client",
self._model,
)
if self._loop_limit_reached:
return AGENTIC_LOOP_LIMIT_ERROR_SSE_BYTES
return SERVER_FULFILLED_TOOL_LEAK_ERROR_SSE_BYTES
if self._replay_index < len(self._collected_bytes):
@ -393,14 +403,9 @@ class AgenticAnthropicStreamingIterator:
"AgenticStreamingIterator: Unexpected result type from hooks: %s",
type(result).__name__,
)
except Exception as e:
_call_id: Final = getattr(self._logging_obj, "litellm_call_id", "unknown")
verbose_logger.exception(
"AgenticStreamingIterator: Error in agentic hook processing [call_id=%s model=%s]: %s",
_call_id,
self._model,
str(e),
)
except AgenticLoopSafetyError as e:
self._loop_limit_reached = isinstance(e, AgenticLoopLimitError)
verbose_logger.warning("AgenticStreamingIterator: agentic loop refused [rail=%s]", type(e).__name__)
@staticmethod
def _rebuild_anthropic_response_from_sse(

1
litellm/llms/aws_polly/text_to_speech/transformation.py
View file

@ -182,7 +182,6 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM):
else:
mapped_params["output_format"] = self.DEFAULT_OUTPUT_FORMAT
# model arrives without the "aws_polly/" prefix, e.g. "standard"
engine: Final = self._extract_engine_from_model(model)
mapped_params["engine"] = engine

6
litellm/llms/azure_ai/rerank/transformation.py
View file

@ -111,10 +111,4 @@ class AzureAIRerankConfig(CohereRerankConfig):
return rerank_response
def _get_base_model(self, azure_model_group: str | None) -> str | None:
if azure_model_group is None:
return None
if azure_model_group == "offer-cohere-rerank-mul-paygo":
return "azure_ai/cohere-rerank-v3-multilingual"
if azure_model_group == "offer-cohere-rerank-eng-paygo":
return "azure_ai/cohere-rerank-v3-english"
return azure_model_group

12
litellm/llms/bedrock/chat/invoke_transformations/anthropic_claude_transformation.py
View file

@ -109,8 +109,16 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
self._read_configured_max_tokens: Final = read_configured_max_tokens
@classmethod
def get_config(cls, *, model: str | None = None) -> dict[str, object]: # mutable-ok: overrides BaseConfig's dict
return {key: value for key, value in super().get_config(model=model).items() if key != "max_tokens"}
def get_config(
cls, *, model: str | None = None, configured_max_output_tokens: int | None = None
) -> dict[str, object]: # mutable-ok: overrides BaseConfig's dict
return {
key: value
for key, value in super()
.get_config(model=model, configured_max_output_tokens=configured_max_output_tokens)
.items()
if key != "max_tokens"
}
@property
def custom_llm_provider(self) -> str | None:

4
litellm/llms/bedrock/count_tokens/kimi_k3.py
View file

@ -455,9 +455,7 @@ def _converse_request(
if not all(_text_only(message) for message in messages):
return None
try:
chat_tools: Final = tuple(
chat_tool for tool in tools or () for chat_tool in _as_converse_input(tool)
) # comprehension-ok: flattens each tool's translation
chat_tools: Final = tuple(chain.from_iterable(_as_converse_input(tool) for tool in tools or ()))
return _ConverseRequest.model_validate(
litellm.AmazonConverseConfig().transform_request(
model=model,

7
litellm/llms/chatgpt/authenticator.py
View file

@ -66,6 +66,13 @@ def get_cached_authenticator(auth_file: str) -> "Authenticator":
return Authenticator(auth_file=auth_file)
def resolve_chatgpt_authenticator(
litellm_params: "Mapping[str, object] | GenericLiteLLMParams | None", default: "Authenticator"
) -> "Authenticator":
auth_file: Final = get_chatgpt_auth_file(litellm_params)
return get_cached_authenticator(auth_file) if auth_file else default
def _token_refresh_timeout() -> httpx.Timeout:
configured: Final = get_configured_request_timeout()
seconds: Final = (

12
litellm/llms/chatgpt/chat/transformation.py
View file

@ -6,7 +6,7 @@ from litellm.llms.openai.openai import OpenAIConfig
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
from ..authenticator import Authenticator, get_cached_authenticator, get_chatgpt_auth_file
from ..authenticator import Authenticator, get_chatgpt_auth_file, resolve_chatgpt_authenticator
from ..common_utils import (
GetAccessTokenError,
chatgpt_client_api_base,
@ -30,12 +30,6 @@ class ChatGPTConfig(OpenAIConfig):
super().__init__()
self.authenticator = authenticator if authenticator is not None else Authenticator()
def _resolve_authenticator(
self, litellm_params: Mapping[str, object] | GenericLiteLLMParams | None
) -> Authenticator:
auth_file: Final = get_chatgpt_auth_file(litellm_params)
return get_cached_authenticator(auth_file) if auth_file else self.authenticator
@staticmethod
def _get_access_token_or_raise(authenticator: Authenticator, model: str, llm_provider: str) -> str:
try:
@ -63,7 +57,7 @@ class ChatGPTConfig(OpenAIConfig):
return chatgpt_client_api_base(api_base), client_credential.access_token, custom_llm_provider
if api_key:
return api_base or self.api_base_without_login(), api_key, custom_llm_provider
authenticator: Final = self._resolve_authenticator(litellm_params)
authenticator: Final = resolve_chatgpt_authenticator(litellm_params, self.authenticator)
dynamic_api_base: Final = authenticator.get_api_base()
dynamic_api_key: Final = self._get_access_token_or_raise(authenticator, model, custom_llm_provider)
return dynamic_api_base, dynamic_api_key, custom_llm_provider
@ -98,7 +92,7 @@ class ChatGPTConfig(OpenAIConfig):
),
**strip_chatgpt_client_credential_headers(headers),
)
authenticator: Final = self._resolve_authenticator(litellm_params)
authenticator: Final = resolve_chatgpt_authenticator(litellm_params, self.authenticator)
resolved_api_key: Final = (
self._get_access_token_or_raise(authenticator, model, "chatgpt")
if get_chatgpt_auth_file(litellm_params)

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