scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-11 08:38:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 359

Merge remote-tracking branch 'origin/main' into litellm_feat_955c0d
Some checks failed
Cost map pins recheck / rerun-cost-map-guard (push) Has been cancelled
Details
LiteLLM Rust DB / db-queries (push) Has been cancelled
Details
LiteLLM Rust / rust-lint (push) Has been cancelled
Details
LiteLLM Rust / rust-test (push) Has been cancelled
Details
LiteLLM Rust / rust-wheel (push) Has been cancelled
Details

Browse source 
# Conflicts:
#	litellm/proxy/deps/repositories.py
This commit is contained in:
Yuneng Jiang 2026-11-04 09:26:30 -08:00
commit d7e14d2144
No known key found for this signature in database
353 changed files with 18813 additions and 13641 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
}
}

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))
@ -2072,6 +2075,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

46
litellm/cost_calculator.py
View file

@ -49,6 +49,9 @@ 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_fedramp_uplift_multiplier,
get_modality_token_cost_key,
get_regional_uplift_multiplier,
get_service_tier_cost_key,
get_token_type_cost_breakdown,
@ -442,6 +445,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 +475,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 +498,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 +610,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
@ -1747,6 +1773,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,
@ -1800,6 +1827,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,
@ -2274,6 +2302,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,
@ -2508,6 +2537,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
@ -2564,6 +2594,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,
@ -3158,8 +3189,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
@ -3277,7 +3312,9 @@ def batch_cost_calculator(
text_tokens: Final = usage.completion_tokens - image_tokens
total_completion_cost = text_tokens * text_rate + image_tokens * image_rate
uplift: Final = get_regional_uplift_multiplier(model_info, data_residency)
uplift: Final = get_regional_uplift_multiplier(model_info, data_residency) * get_fedramp_uplift_multiplier(
model_info, data_residency
)
if uplift != 1.0:
total_prompt_cost *= uplift
total_completion_cost *= uplift
@ -3309,7 +3346,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

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

169
litellm/litellm_core_utils/litellm_logging.py
View file

@ -54,6 +54,7 @@ from litellm.constants import (
REDACTED_BY_LITELLM,
)
from litellm.cost_calculator import (
BaseTokenUsageProcessor,
RealtimeAPITokenUsageProcessor,
ResponsesWebSocketTokenUsageProcessor,
declared_cost_map_pricing,
@ -77,6 +78,7 @@ 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,
@ -722,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
@ -943,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":
@ -2227,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,
@ -2318,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
@ -2423,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
@ -2771,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
@ -2826,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."""
@ -2838,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(
@ -2861,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:
@ -2868,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)
@ -3172,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
):
@ -3184,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
@ -3312,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})
@ -3334,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(
@ -3747,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)
@ -4343,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)
@ -6422,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:
"""
@ -7836,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)
@ -8029,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
),

96
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
@ -1402,6 +1441,29 @@ def get_regional_uplift_multiplier(model_info: ModelInfo, data_residency: str |
_get_regional_uplift_multiplier = get_regional_uplift_multiplier
def get_fedramp_uplift_multiplier(model_info: ModelInfo, data_residency: str | None) -> float:
"""
Resolve the per-model uplift multiplier for the OpenAI FedRAMP endpoint.
OpenAI applies a flat percentage uplift (e.g. +10%) on all token costs for
requests served from gov.api.openai.com. The multiplier is stored on the
model entry as ``fedramp_uplift_multiplier`` (e.g. 1.10).
Returns 1.0 (no uplift) unless ``data_residency`` is ``"fedramp"`` and the
model has a multiplier configured.
"""
if data_residency is None or data_residency.lower() != DataResidency.FEDRAMP.value:
return 1.0
multiplier: Final = model_info.get("fedramp_uplift_multiplier")
if multiplier is None:
return 1.0
try:
return float(multiplier)
except (TypeError, ValueError):
verbose_logger.exception("Invalid fedramp_uplift_multiplier for model; defaulting to 1.0")
return 1.0
def get_vertex_regional_endpoint_uplift(model_info: ModelInfo, vertex_location: str | None) -> float:
"""
Resolve the per-model uplift multiplier for Vertex AI non-global (regional and
@ -1511,8 +1573,8 @@ def generic_cost_per_token(
Input:
- model: str, the model name without provider prefix
- usage: LiteLLM Usage block, containing anthropic caching information
- data_residency: optional OpenAI data-residency region (e.g. "eu", "us"),
used to apply the per-model regional-processing uplift multiplier.
- data_residency: optional OpenAI data-residency region (e.g. "eu", "us") or
"fedramp", used to apply the per-model regional-processing or FedRAMP uplift multiplier.
- vertex_location: optional Vertex AI location the request was served from
(e.g. "us-east5", "global"), used to apply the per-model
regional-endpoint uplift multiplier when non-global.
@ -1618,6 +1680,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 +1739,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:
@ -1689,10 +1751,9 @@ def generic_cost_per_token(
)
completion_cost += float(video_tokens) * _output_cost_per_video_token
## REGIONAL DATA-RESIDENCY UPLIFT
# Applied as a flat multiplier across all token costs for the request
# when the upstream is a regionalized OpenAI host (eu./us.api.openai.com).
uplift: Final = get_regional_uplift_multiplier(resolved_model_info, data_residency)
uplift: Final = get_regional_uplift_multiplier(resolved_model_info, data_residency) * get_fedramp_uplift_multiplier(
resolved_model_info, data_residency
)
if uplift != 1.0:
prompt_cost *= uplift
completion_cost *= uplift
@ -1835,6 +1896,7 @@ def _cost_map_billed_rates(
)
multiplier: Final = (
get_regional_uplift_multiplier(model_info, data_residency)
* get_fedramp_uplift_multiplier(model_info, data_residency)
* get_vertex_regional_endpoint_uplift(model_info, vertex_location)
* get_provider_specific_geo_multiplier(model_info=model_info, usage=usage)
* get_provider_specific_speed_multiplier(model_info=model_info, usage=usage)
@ -2020,8 +2082,10 @@ def calculate_prompt_caching_savings(
{ttl: rate - prompt_base_cost for ttl, rate in base_costs.cache_creation_rates_by_ttl.items()}
),
)
uplift: Final = get_regional_uplift_multiplier(model_info, data_residency) * get_vertex_regional_endpoint_uplift(
model_info, vertex_location
uplift: Final = (
get_regional_uplift_multiplier(model_info, data_residency)
* get_fedramp_uplift_multiplier(model_info, data_residency)
* get_vertex_regional_endpoint_uplift(model_info, vertex_location)
)
return (read_discount - write_premium) * uplift

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

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

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

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

10
litellm/llms/chatgpt/common_utils.py
View file

@ -7,7 +7,7 @@ import hashlib
import json
import os
import platform
from collections.abc import Mapping, Sequence
from collections.abc import Mapping
from dataclasses import dataclass
from types import MappingProxyType
from typing import TYPE_CHECKING, Final, TypeAlias
@ -20,14 +20,10 @@ from pydantic import JsonValue, TypeAdapter
from litellm.litellm_core_utils.get_provider_specific_headers import ProviderSpecificHeaderUtils
from litellm.litellm_core_utils.provider_affinity import get_caller_litellm_session_id, get_requested_session_id
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.utils import ProviderSpecificHeader
if TYPE_CHECKING:
from litellm.types.router import GenericLiteLLMParams
_SCOPED_HEADERS: Final[TypeAdapter[ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None]] = TypeAdapter(
ProviderSpecificHeader | Sequence[ProviderSpecificHeader] | None
)
_STRING_HEADERS: Final = TypeAdapter(Mapping[str, str])
_PARAMS: Final = TypeAdapter(Mapping[str, object])
@ -357,8 +353,8 @@ def get_chatgpt_request_credential(
normalized: Final = _normalize_litellm_params(params)
scoped: Final = normalized.get("provider_specific_header")
extra: Final = normalized.get("extra_headers")
headers: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers(
provider_specific_header=_SCOPED_HEADERS.validate_python(scoped), custom_llm_provider="chatgpt"
headers: Final = ProviderSpecificHeaderUtils.validated_provider_specific_headers(
provider_specific_header=scoped, custom_llm_provider="chatgpt"
)
return get_chatgpt_client_credential(headers) or get_chatgpt_client_credential(
_STRING_HEADERS.validate_python(extra) if extra is not None else MappingProxyType({})

8
litellm/llms/chatgpt/responses/transformation.py
View file

@ -28,7 +28,7 @@ from litellm.types.llms.openai import (
from litellm.types.router import GenericLiteLLMParams
from litellm.types.utils import LlmProviders
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 (
CHATGPT_API_BASE,
GetAccessTokenError,
@ -64,10 +64,6 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
def custom_llm_provider(self) -> LlmProviders:
return LlmProviders.CHATGPT
def _resolve_authenticator(self, litellm_params: GenericLiteLLMParams | None) -> Authenticator:
auth_file: Final = get_chatgpt_auth_file(litellm_params)
return get_cached_authenticator(auth_file) if auth_file else self.authenticator
def validate_environment(
self,
headers: dict[str, str],
@ -81,7 +77,7 @@ class ChatGPTResponsesAPIConfig(OpenAIResponsesAPIConfig):
**get_chatgpt_default_headers(client_credential.access_token, client_credential.account_id, session_id),
**strip_chatgpt_client_credential_headers(headers),
}
authenticator: Final = self._resolve_authenticator(litellm_params)
authenticator: Final = resolve_chatgpt_authenticator(litellm_params, self.authenticator)
safe_headers: Final = strip_chatgpt_client_credential_headers(headers)
configured_api_key: Final = litellm_params.api_key if litellm_params is not None else None
if configured_api_key and not get_chatgpt_auth_file(litellm_params):

374
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -47,6 +47,12 @@ from litellm.anthropic_beta_headers_manager import (
from litellm.constants import MAX_FILE_LIST_LIMIT, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
from litellm.files.types import FileContentStreamingResult
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_responses_turn,
is_refused_item,
refused_tool_call_identifiers,
)
from litellm.litellm_core_utils.agentic_loop_settings import (
DEFAULT_MAX_AGENTIC_LOOPS,
validated_max_agentic_loops,
@ -56,11 +62,11 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import (
SUBTITLE_RESPONSE_FORMATS,
synthesize_subtitle_document,
)
from litellm.litellm_core_utils.core_helpers import (
as_str_mapping,
qualify_agentic_followup_model,
set_provider_response_headers_in_hidden_params,
from litellm.litellm_core_utils.chat_completion_agentic_loop import (
bill_first_round_on_failure,
run_chat_completion_agentic_hooks,
)
from litellm.litellm_core_utils.core_helpers import as_str_mapping, set_provider_response_headers_in_hidden_params
from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header
@ -150,6 +156,7 @@ from litellm.types.extraction import ExtractionResponse
from litellm.types.files import StreamingMediaUploadConfig, TwoStepFileUploadConfig
from litellm.types.integrations.custom_logger import (
NON_CODE_INTERPRETER_INTERCEPTION_INTERNAL_PREFIXES,
AgenticLoopLimitError,
AgenticLoopPlan,
AgenticLoopRequestPatch,
AgenticLoopSafetyError,
@ -6005,7 +6012,7 @@ class BaseLLMHTTPHandler:
@staticmethod
def _check_agentic_loop_safety(
tool_calls: object,
fingerprints: list[str],
fingerprints: Sequence[str],
depth: int,
max_loops: int,
model: str,
@ -6024,7 +6031,7 @@ class BaseLLMHTTPHandler:
if fingerprint in fingerprints:
raise AgenticLoopSafetyError("Agentic loop detected repeated tool-call fingerprint; aborting rerun")
if depth >= max_loops:
raise AgenticLoopSafetyError(f"Exceeded max_agentic_loops={max_loops} for model={model}")
raise AgenticLoopLimitError(f"Exceeded max_agentic_loops={max_loops} for model={model}")
return fingerprint
@staticmethod
@ -6034,26 +6041,6 @@ class BaseLLMHTTPHandler:
except Exception:
return str(tools)
@staticmethod
def _refused_agentic_tool_identifiers(tool_calls: object) -> tuple[frozenset[str], frozenset[str]]:
"""
Collect the ids and names of the tool calls a safety rail just refused.
Callbacks hand back either a bare list of tool calls or a dict wrapping
that list under ``tool_calls``, and both the anthropic and responses
shapes carry an ``id`` (or ``call_id``) plus a ``name``.
"""
calls: Final = tool_calls.get("tool_calls") if isinstance(tool_calls, dict) else tool_calls
if not isinstance(calls, list):
return frozenset(), frozenset()
dict_calls: Final = (call for call in calls if isinstance(call, dict))
fields: Final = tuple((call.get("id"), call.get("call_id"), call.get("name")) for call in dict_calls)
ids: Final = frozenset(
value for call_id, caller_id, _ in fields for value in (call_id, caller_id) if isinstance(value, str)
)
names: Final = frozenset(name for _, _, name in fields if isinstance(name, str))
return ids, names
@staticmethod
def _is_refused_tool_use_block(block: object, refused_ids: frozenset[str], refused_names: frozenset[str]) -> bool:
"""
@ -6064,22 +6051,13 @@ class BaseLLMHTTPHandler:
share a name with a refused one. The name is only consulted for tool call
shapes that arrive without an id.
"""
if not isinstance(block, dict) or block.get("type") != "tool_use":
return False
block_id: Final = block.get("id")
if isinstance(block_id, str) and refused_ids:
return block_id in refused_ids
return block.get("name") in refused_names
return is_refused_item(block, "tool_use", "id", refused_ids, refused_names)
@staticmethod
def _can_replace_turn_with_terminal_response(stream: bool, api_surface: str) -> bool:
"""
Whether a refused rerun can still be answered with a finalized turn.
Only the anthropic messages surface can. The responses surface carries a
pydantic model the finalizer does not rewrite, so it keeps raising, which
is what every surface did before this path learned to end the turn.
The messages and responses call sites pass ``stream=False``, because
interception converts an intercepted stream to non-streaming before the
loop runs and rebuilds the SSE stream from the finalized turn
@ -6088,7 +6066,7 @@ class BaseLLMHTTPHandler:
finalized turn would reach the client as a second message rather than
as a replacement.
"""
return not stream and api_surface == "anthropic_messages"
return not stream and api_surface in ("anthropic_messages", "responses")
@staticmethod
def _finalize_refused_agentic_response(response: object, tool_calls: object) -> object:
@ -6105,10 +6083,12 @@ class BaseLLMHTTPHandler:
one is still in the response the stop reason stays ``tool_use`` so the
client knows to answer it.
"""
if isinstance(response, ResponsesAPIResponse):
return finalize_refused_responses_turn(response, tool_calls)
if not isinstance(response, dict):
return response
refused_ids, refused_names = BaseLLMHTTPHandler._refused_agentic_tool_identifiers(tool_calls)
refused_ids, refused_names = refused_tool_call_identifiers(tool_calls)
finalized: Final = dict(response)
content: Final = finalized.get("content")
if isinstance(content, list):
@ -6135,7 +6115,7 @@ class BaseLLMHTTPHandler:
kwargs: dict,
depth: int,
max_loops: int,
fingerprints: list[str],
fingerprints: Sequence[str],
fingerprint: str,
stream: bool = False,
callback: Optional["CustomLogger"] = None,
@ -6176,7 +6156,7 @@ class BaseLLMHTTPHandler:
kwargs_for_followup.update(patch.kwargs)
kwargs_for_followup["_agentic_loop_depth"] = depth + 1
kwargs_for_followup["max_agentic_loops"] = max_loops
kwargs_for_followup["_agentic_loop_fingerprints"] = fingerprints + [fingerprint]
kwargs_for_followup["_agentic_loop_fingerprints"] = [*fingerprints, fingerprint]
response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate(
**{
@ -6215,7 +6195,7 @@ class BaseLLMHTTPHandler:
kwargs: dict,
depth: int,
max_loops: int,
fingerprints: list[str],
fingerprints: Sequence[str],
fingerprint: str,
callback: Optional["CustomLogger"] = None,
) -> ResponsesAPIResponse | BaseResponsesAPIStreamingIterator:
@ -6336,67 +6316,6 @@ class BaseLLMHTTPHandler:
custom_llm_provider=custom_llm_provider,
)
async def _execute_chat_completion_agentic_plan(
self,
plan: AgenticLoopPlan,
model: str,
messages: list[dict],
optional_params: dict,
kwargs: dict,
custom_llm_provider: str,
depth: int,
max_loops: int,
fingerprints: list[str],
fingerprint: str,
logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> ModelResponse | CustomStreamWrapper:
patch: Final = plan.request_patch or AgenticLoopRequestPatch()
if patch.messages is None:
raise ValueError("Agentic loop plan missing patched messages")
full_model_name: Final = qualify_agentic_followup_model(patch.model, model, custom_llm_provider)
optional_params_for_followup: Final = dict(optional_params)
optional_params_for_followup.update(patch.optional_params)
if patch.tools is not None:
optional_params_for_followup["tools"] = patch.tools
internal_params: Final = {
"_websearch_interception",
"acompletion",
"litellm_logging_obj",
"custom_llm_provider",
"model_alias_map",
"stream_response",
"custom_prompt_dict",
}
kwargs_for_followup: Final = build_agentic_followup_kwargs(
request_kwargs=MappingProxyType(
{
k: v
for k, v in kwargs.items()
if not k.startswith("_websearch_interception")
and not k.startswith("_compression_interception")
and k not in internal_params
}
),
patch_kwargs=patch.kwargs,
request_params=frozenset((*optional_params_for_followup, "model", "messages")),
depth=depth,
max_loops=max_loops,
fingerprints=fingerprints,
fingerprint=fingerprint,
)
with internal_sub_call(logging_obj.litellm_call_id if logging_obj is not None else None):
response: Final = await litellm.acompletion(
model=full_model_name,
messages=patch.messages,
**optional_params_for_followup,
**kwargs_for_followup,
)
return response
def _maybe_wrap_in_fake_stream(
self,
response: _ResponseT,
@ -6513,26 +6432,38 @@ class BaseLLMHTTPHandler:
except AgenticLoopSafetyError as e:
if not self._can_replace_turn_with_terminal_response(stream, api_surface):
raise
verbose_logger.warning(
"LiteLLM.AgenticLoopRefused: ending turn [call_id=%s model=%s]: %s",
logging_obj.litellm_call_id,
model,
str(e),
)
verbose_logger.warning("LiteLLM.AgenticLoopRefused: ending turn [rail=%s]", type(e).__name__)
return self._maybe_wrap_in_fake_stream(
self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls),
self._finalize_refused_agentic_response(
response=response,
tool_calls=await every_claimed_tool_call(
tool_calls,
(
other.async_should_run_agentic_loop(
response=response,
model=model,
messages=messages,
tools=tools,
stream=stream,
custom_llm_provider=custom_llm_provider,
kwargs=hook_kwargs,
)
for other in callbacks
if isinstance(other, CustomLogger) and other is not callback
),
),
),
logging_obj,
api_surface,
)
try:
kwargs_with_provider = hook_kwargs.copy()
kwargs_with_provider["custom_llm_provider"] = custom_llm_provider
build_plan_overridden = (
callback.__class__.async_build_agentic_loop_plan is not CustomLogger.async_build_agentic_loop_plan
)
if not build_plan_overridden:
agentic_result: object = await callback.async_run_agentic_loop(
kwargs_with_provider = {**hook_kwargs, "custom_llm_provider": custom_llm_provider}
build_plan_overridden = (
callback.__class__.async_build_agentic_loop_plan is not CustomLogger.async_build_agentic_loop_plan
)
if not build_plan_overridden:
agentic_result: object = await bill_first_round_on_failure(
callback.async_run_agentic_loop(
tools=tool_calls,
model=model,
messages=messages,
@ -6542,10 +6473,15 @@ class BaseLLMHTTPHandler:
logging_obj=logging_obj,
stream=stream,
kwargs=kwargs_with_provider,
)
return self._maybe_wrap_in_fake_stream(agentic_result, logging_obj, api_surface)
),
logging_obj=logging_obj,
response=response,
stream=stream,
)
return self._maybe_wrap_in_fake_stream(agentic_result, logging_obj, api_surface)
plan = await callback.async_build_agentic_loop_plan(
plan = await bill_first_round_on_failure(
callback.async_build_agentic_loop_plan(
tools=tool_calls,
model=model,
messages=messages,
@ -6555,30 +6491,35 @@ class BaseLLMHTTPHandler:
logging_obj=logging_obj,
stream=stream,
kwargs=kwargs_with_provider,
),
logging_obj=logging_obj,
response=response,
stream=stream,
)
if plan.response_override is not None:
return self._maybe_wrap_in_fake_stream(plan.response_override, logging_obj, api_surface)
if plan.terminate:
verbose_logger.debug(
"Agentic loop terminated by callback=%s reason=%s",
callback.__class__.__name__,
plan.stop_reason,
)
return self._maybe_wrap_in_fake_stream(
await callback.async_post_agentic_loop_response_hook(
response=self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls),
plan=plan,
kwargs=kwargs_with_provider,
),
logging_obj,
api_surface,
)
if not plan.run_agentic_loop:
continue
if plan.response_override is not None:
return self._maybe_wrap_in_fake_stream(plan.response_override, logging_obj, api_surface)
if plan.terminate:
verbose_logger.debug(
"Agentic loop terminated by callback=%s reason=%s",
callback.__class__.__name__,
plan.stop_reason,
)
return self._maybe_wrap_in_fake_stream(
await callback.async_post_agentic_loop_response_hook(
response=self._finalize_refused_agentic_response(response=response, tool_calls=tool_calls),
plan=plan,
kwargs=kwargs_with_provider,
),
logging_obj,
api_surface,
)
if not plan.run_agentic_loop:
continue
if api_surface == "responses":
return await self._execute_responses_agentic_plan(
if api_surface == "responses":
return await bill_first_round_on_failure(
self._execute_responses_agentic_plan(
plan=plan,
model=model,
response_api_optional_request_params=anthropic_messages_optional_request_params,
@ -6589,10 +6530,15 @@ class BaseLLMHTTPHandler:
fingerprints=fingerprints,
fingerprint=fingerprint,
callback=callback,
)
),
logging_obj=logging_obj,
response=response,
stream=stream,
)
return self._maybe_wrap_in_fake_stream(
await self._execute_anthropic_agentic_plan(
return self._maybe_wrap_in_fake_stream(
await bill_first_round_on_failure(
self._execute_anthropic_agentic_plan(
plan=plan,
model=model,
messages=messages,
@ -6606,17 +6552,13 @@ class BaseLLMHTTPHandler:
stream=stream,
callback=callback,
),
logging_obj,
api_surface,
)
except Exception as e:
_call_id = getattr(logging_obj, "litellm_call_id", "unknown")
verbose_logger.exception(
"LiteLLM.AgenticHookError: Exception in agentic completion hooks [call_id=%s model=%s]: %s",
_call_id,
model,
str(e),
)
logging_obj=logging_obj,
response=response,
stream=stream,
),
logging_obj,
api_surface,
)
# Check if we need to convert response to fake stream
# This happens when:
@ -6648,111 +6590,23 @@ class BaseLLMHTTPHandler:
Returns the response from agentic loop, or None if no hook runs.
"""
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
callbacks: Final = litellm.callbacks + (logging_obj.dynamic_success_callbacks or [])
tools: Final = optional_params.get("tools", [])
depth, max_loops, fingerprints = self._get_agentic_loop_settings(kwargs=kwargs)
for callback in callbacks:
if not isinstance(callback, CustomLogger):
continue
if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"):
continue
should_run: bool = False
tool_calls: object = None
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 as e:
verbose_logger.exception(
"LiteLLM.AgenticHookError: Exception in async_should_run_chat_completion_agentic_loop: %s",
str(e),
)
continue
if not should_run:
continue
# Safety guards must run OUTSIDE the callback try/except — they are
# bounded-loop / cycle-break rails that must propagate to the caller.
fingerprint = self._check_agentic_loop_safety(
tool_calls=tool_calls,
fingerprints=fingerprints,
depth=depth,
max_loops=max_loops,
agentic_response: Final = await bill_first_round_on_failure(
run_chat_completion_agentic_hooks(
response=response,
model=model,
)
try:
kwargs_with_provider = kwargs.copy() if kwargs else {}
kwargs_with_provider["custom_llm_provider"] = custom_llm_provider
build_plan_overridden = (
callback.__class__.async_build_chat_completion_agentic_loop_plan
is not CustomLogger.async_build_chat_completion_agentic_loop_plan
)
if not build_plan_overridden:
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=kwargs_with_provider,
)
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=kwargs_with_provider,
)
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",
callback.__class__.__name__,
plan.stop_reason,
)
return response
if not plan.run_agentic_loop:
continue
return await self._execute_chat_completion_agentic_plan(
plan=plan,
model=model,
messages=messages,
optional_params=optional_params,
kwargs=kwargs_with_provider,
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", e)
messages=messages,
optional_params=optional_params,
kwargs=kwargs or {},
logging_obj=logging_obj,
custom_llm_provider=custom_llm_provider,
stream=stream,
),
logging_obj=logging_obj,
response=response,
stream=stream,
)
if agentic_response is not None:
return agentic_response
# Check if we need to convert response to fake stream for chat completions
# This happens when:

2
litellm/llms/databricks/chat/transformation.py
View file

@ -395,7 +395,7 @@ class DatabricksConfig(DatabricksBase, OpenAILikeChatConfig, AnthropicConfig):
return "databricks"
@classmethod
def get_config(cls, *, model: str | None = None):
def get_config(cls, *, model: str | None = None, configured_max_output_tokens: int | None = None):
return super().get_config()
def get_json_schema_from_pydantic_object(

2
litellm/llms/databricks/messages/transformation.py
View file

@ -192,7 +192,7 @@ class DatabricksAnthropicMessagesConfig(DatabricksBase, AnthropicMessagesConfig)
headers=headers,
)
self._normalize_system_role_messages( # pyright: ignore[reportUnknownMemberType] # the shared helper takes a bare dict
anthropic_messages_request, model=model, drop_params=getattr(litellm_params, "drop_params", None)
anthropic_messages_request, model=model, drop_params=litellm_params.drop_params
)
self._remove_scope_from_cache_control(anthropic_messages_request) # pyright: ignore[reportUnknownMemberType] # the shared helper takes a bare dict
return anthropic_messages_request

2
litellm/llms/gemini/realtime/transformation.py
View file

@ -218,7 +218,7 @@ def _system_instruction_text(system_instruction: HttpxContentType | None) -> str
class GeminiRealtimeConfig(BaseRealtimeConfig):
_TOOL_CALL_ID_TO_NAME_MAX = 256 # LRU cap for call_id→name mapping
def __init__(self, drop_params: bool = False):
def __init__(self, drop_params: bool = False) -> None:
super().__init__()
self.drop_params = drop_params
self._tool_call_id_to_name: OrderedDict[str, str] = OrderedDict()

6
litellm/llms/openai/cost_calculation.py
View file

@ -31,9 +31,9 @@ def cost_per_token(
Input:
- model: str, the model name without provider prefix
- usage: LiteLLM Usage block, containing anthropic caching information
- data_residency: optional OpenAI data-residency region (e.g. "eu", "us"),
inferred from api_base. Applies the model's regional-processing
uplift multiplier when set.
- data_residency: optional OpenAI data-residency mode (e.g. "eu", "us", "fedramp"),
inferred from api_base, including gov.api.openai.com. Applies the
model's regional-processing uplift multiplier when set.
Returns:
Tuple[float, float] - prompt_cost_in_usd, completion_cost_in_usd

21
litellm/llms/openai/data_residency.py
View file

@ -7,9 +7,12 @@ enabled and rejects requests sent to the wrong host, so the api_base hostname
is the authoritative signal of which region a request was processed in.
"""
import os
from typing import Final
from urllib.parse import urlparse
import litellm
# Mapping of OpenAI regional hostnames to the corresponding data-residency
# value used by the cost calculator. See
# https://developers.openai.com/api/docs/pricing for the regional-processing
@ -17,6 +20,7 @@ from urllib.parse import urlparse
_OPENAI_REGIONAL_HOSTS: Final[dict[str, str]] = {
"eu.api.openai.com": "eu",
"us.api.openai.com": "us",
"gov.api.openai.com": "fedramp",
}
@ -25,13 +29,22 @@ def infer_openai_data_residency(custom_llm_provider: str | None, api_base: str |
Derive the OpenAI data-residency region from an api_base URL.
Returns ``"eu"`` for the EU regional host, ``"us"`` for the US regional
host, and ``None`` for the default global host, any non-OpenAI provider,
or any non-OpenAI URL.
host, ``"fedramp"`` for ``gov.api.openai.com``, and ``None`` for the
default global host, any non-OpenAI provider, or any non-OpenAI URL.
When ``api_base`` is unset, falls back to the same defaults the OpenAI
handler uses (``litellm.api_base``, ``OPENAI_BASE_URL``, ``OPENAI_API_BASE``)
so env-configured regional endpoints are billed with their uplift.
"""
if custom_llm_provider != "openai" or not api_base:
if custom_llm_provider != "openai":
return None
resolved_api_base: Final = (
api_base or litellm.api_base or os.getenv("OPENAI_BASE_URL") or os.getenv("OPENAI_API_BASE")
)
if not resolved_api_base:
return None
try:
host: Final = urlparse(api_base).hostname
host: Final = urlparse(resolved_api_base).hostname
except (TypeError, ValueError):
return None
if not host:

125
litellm/llms/openai/openai.py
View file

@ -30,6 +30,10 @@ from litellm._logging import verbose_logger
from litellm.constants import DEFAULT_MAX_RETRIES
from litellm.exceptions import UnsupportedParamsError
from litellm.files.types import FileContentStreamingResult
from litellm.litellm_core_utils.chat_completion_agentic_loop import (
bill_first_round_on_failure,
run_chat_completion_agentic_hooks,
)
from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing
@ -663,67 +667,23 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
Returns the response from agentic loop, or None if no hook runs.
"""
from litellm._logging import verbose_logger
from litellm.integrations.custom_logger import CustomLogger
callbacks: Final = litellm.callbacks + (logging_obj.dynamic_success_callbacks or [])
# Avoid logging full callback objects to prevent leaking sensitive data
verbose_logger.debug("LiteLLM.AgenticHooks: callbacks_count=%s", len(callbacks))
tools: Final = optional_params.get("tools", [])
# Avoid logging full tools payloads; they may contain sensitive parameters
verbose_logger.debug(
"LiteLLM.AgenticHooks: tools_count=%s",
len(tools) if isinstance(tools, list) else 1 if tools else 0,
if not isinstance(response, ModelResponse):
return None
return await bill_first_round_on_failure(
run_chat_completion_agentic_hooks(
response=response,
model=model,
messages=messages,
optional_params=optional_params,
kwargs=litellm_params or {},
logging_obj=logging_obj,
custom_llm_provider=litellm_params.get("custom_llm_provider", "openai"),
stream=stream,
),
logging_obj=logging_obj,
response=response,
stream=stream,
)
# Get custom_llm_provider from litellm_params
custom_llm_provider: Final = litellm_params.get("custom_llm_provider", "openai")
for callback in callbacks:
try:
if isinstance(callback, CustomLogger):
# Check if the callback has the chat completion agentic loop methods
if not hasattr(callback, "async_should_run_chat_completion_agentic_loop"):
continue
# First: Check if agentic loop should run (using chat completion method)
(
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=litellm_params,
)
if should_run:
# Second: Execute agentic loop
kwargs_with_provider = litellm_params.copy() if litellm_params else {}
kwargs_with_provider["custom_llm_provider"] = custom_llm_provider
# For OpenAI Chat Completions, use the chat completion agentic loop method
agentic_response: object = 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=kwargs_with_provider,
)
# First hook that runs agentic loop wins
return agentic_response
except Exception as e:
verbose_logger.exception(
"LiteLLM.AgenticHookError: Exception in agentic completion hooks for OpenAI: %s", e
)
return None
def mock_streaming(
self,
@ -1074,8 +1034,28 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
hidden_params={"headers": headers},
_response_headers=headers,
)
except openai.UnprocessableEntityError as e:
if litellm.drop_params is True or drop_params is True:
data = drop_params_from_unprocessable_entity_error(e, data, model=model, drop_params_gated=True)
else:
raise e
except Exception as e:
exception_response: object = getattr(e, "response", None)
status_code = getattr(e, "status_code", 500)
exception_body = getattr(e, "body", None)
error_headers = getattr(e, "headers", None)
if error_headers is None and exception_response:
error_headers = getattr(exception_response, "headers", None)
message = getattr(e, "message", str(e))
# Call agentic completion hooks (e.g., for websearch_interception)
raise OpenAIError(
status_code=status_code,
message=message,
response=upstream_error_response(e),
headers=error_headers,
body=exception_body,
)
else:
agentic_response = await self._call_agentic_completion_hooks_openai(
response=final_response_obj,
model=model,
@ -1098,29 +1078,6 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
)
return final_response_obj
except openai.UnprocessableEntityError as e:
## check if body contains unprocessable params - related issue https://github.com/BerriAI/litellm/issues/4800
if litellm.drop_params is True or drop_params is True:
data = drop_params_from_unprocessable_entity_error(e, data, model=model, drop_params_gated=True)
else:
raise e
# e.message
except Exception as e:
exception_response: object = getattr(e, "response", None)
status_code = getattr(e, "status_code", 500)
exception_body = getattr(e, "body", None)
error_headers = getattr(e, "headers", None)
if error_headers is None and exception_response:
error_headers = getattr(exception_response, "headers", None)
message = getattr(e, "message", str(e))
raise OpenAIError(
status_code=status_code,
message=message,
response=upstream_error_response(e),
headers=error_headers,
body=exception_body,
)
def streaming(
self,

9
litellm/llms/openai_like/providers.json
View file

@ -302,5 +302,14 @@
"preserve_cache_control": true
},
"supported_endpoints": ["/v1/chat/completions"]
},
"wally": {
"base_url": "https://inference.runanywhere.ai/v1",
"api_key_env": "WALLY_API_KEY",
"api_base_env": "WALLY_API_BASE",
"provider_info": {
"supports_function_calling": true,
"supports_reasoning": true
}
}
}

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