scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-25 01:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

Merge remote-tracking branch 'origin/main' into litellm_team_usage_export_per_user

Browse source
This commit is contained in:
yuneng 2026-09-21 23:36:49 +00:00
commit f9dabdbba2
29 changed files with 3644 additions and 602 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

19
litellm-rust/Cargo.lock generated
View file

@ -2772,6 +2772,22 @@ dependencies = [
"wiremock",
]
[[package]]
name = "litellm-cache-valkey-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-response",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"tokio",
"uuid",
]
[[package]]
name = "litellm-callbacks-legacy-python"
version = "0.1.0"
@ -2944,6 +2960,7 @@ dependencies = [
"litellm-cache-redis",
"litellm-cache-response",
"litellm-cache-s3",
"litellm-cache-valkey-semantic",
"litellm-callbacks-legacy-python",
"litellm-core",
"litellm-core-utils",
@ -2954,10 +2971,12 @@ dependencies = [
"litellm-types",
"pyo3",
"pyo3-async-runtimes",
"redis",
"rstest",
"serde",
"serde_json",
"serde_with",
"sha2 0.10.9",
"tokio",
"tokio-tungstenite",
]

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

@ -14,7 +14,7 @@ use crate::topology::RedisTopology;
mod connection;
mod operations;
pub(crate) use connection::ConnectionRef;
pub use connection::ConnectionRef;
use connection::{ClusterConnectionManager, ConnectionManager};
pub use operations::{
@ -40,7 +40,8 @@ const CLAIM_SCRIPT: &str = concat!(
);
const CLAIM_ATTEMPTS: usize = 8;
enum Connections<C> {
#[allow(private_interfaces)]
pub enum Connections<C> {
Pool(r2d2::Pool<ConnectionManager>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
@ -50,7 +51,7 @@ impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
fn execute<T>(
pub fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
@ -73,6 +74,29 @@ where
}
}
}
pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
}
pub fn fixed(connection: C) -> Self {
Self::Fixed(Mutex::new(connection))
}
pub fn open(url: &str, topology: &RedisTopology) -> Result<Self, Error> {
match topology {
RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)),
RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool(
ClusterConnectionManager::open(url, startup_nodes)?,
)?)),
}
}
}
pub struct RedisCache<S, C = redis::Connection> {
@ -94,12 +118,7 @@ impl<S: CacheCodec> RedisCache<S> {
default_ttl: Option<Duration>,
codec: S,
) -> Result<Self, Error> {
let connections = match topology {
RedisTopology::Standalone => Connections::Pool(pool(ConnectionManager::open(url)?)?),
RedisTopology::Cluster { startup_nodes } => {
Connections::Cluster(pool(ClusterConnectionManager::open(url, startup_nodes)?)?)
}
};
let connections = Connections::open(url, topology)?;
Ok(Self {
connections: Arc::new(connections),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
@ -127,7 +146,7 @@ where
{
pub fn with_connection(connection: C, default_ttl: Option<Duration>, codec: S) -> Self {
Self {
connections: Arc::new(Connections::Fixed(Mutex::new(connection))),
connections: Arc::new(Connections::fixed(connection)),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
codec,
namespace: None,
@ -203,16 +222,6 @@ where
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
async fn run_blocking<T, F>(connections: Arc<Connections<C>>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
}
}
fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
@ -271,7 +280,7 @@ where
let payload = self.codec.encode(&value)?;
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)
@ -285,7 +294,7 @@ where
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)
@ -311,7 +320,7 @@ where
if entries.is_empty() {
return Ok(());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = entries
.into_iter()
.map(|(key, payload)| {
@ -330,7 +339,7 @@ where
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
match Self::run_blocking(Arc::clone(&self.connections), |connection| {
match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
Ok(match connection.ping() {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
@ -391,7 +400,7 @@ where
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
@ -418,7 +427,7 @@ where
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
.await
@ -438,7 +447,7 @@ where
async fn async_flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
Self::flush_matching(connection, &pattern)
})
.await
@ -470,7 +479,7 @@ where
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
increment(connection, key, amount, ttl)
})
.await
@ -581,7 +590,7 @@ where
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
let codec = self.codec.clone();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
claim(connection, &codec, &key, candidate, &eligible, ttl)
})
.await

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

@ -117,7 +117,7 @@ impl r2d2::ManageConnection for ClusterConnectionManager {
}
}
pub(crate) enum ConnectionRef<'a> {
pub enum ConnectionRef<'a> {
Node(&'a mut dyn redis::ConnectionLike),
Cluster(&'a mut ClusterConnection),
}

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

@ -144,7 +144,7 @@ where
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del(keys).map_err(|_| Error::Unavailable)
})
.await
@ -172,7 +172,7 @@ where
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
@ -188,7 +188,7 @@ where
}
pub async fn ping(&self) -> Result<bool, Error> {
Self::run_blocking(Arc::clone(&self.connections), |connection| {
Connections::run_blocking(Arc::clone(&self.connections), |connection| {
connection.ping().map_err(|_| Error::Unavailable)
})
.await
@ -196,7 +196,7 @@ where
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<i64>, Error> {
let key = self.namespaced_key(key);
let ttl = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let ttl = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("TTL")
.arg(key)
.query::<i64>(connection)
@ -208,7 +208,7 @@ where
pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
let pattern = format!("{}*", self.namespaced_key(pattern));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut matches = Vec::new();
connection.scan(&pattern, count, |_, keys| {
matches.extend(keys);
@ -231,7 +231,7 @@ where
}
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut sadd = redis::cmd("SADD");
sadd.arg(&key).arg(values);
let mut expire = redis::cmd("EXPIRE");
@ -253,7 +253,7 @@ where
return Err(Error::InvalidEntry);
}
let key = self.namespaced_key(key);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("RPUSH")
.arg(key)
.arg(values)
@ -279,7 +279,7 @@ where
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = operations
.into_iter()
.map(|(key, values)| {
@ -304,7 +304,7 @@ where
) -> Result<RedisLpopResult, Error> {
let key = self.namespaced_key(key);
let multiple = count.is_some();
let value = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut command = redis::cmd("LPOP");
command.arg(key);
if let Some(count) = count {
@ -333,7 +333,7 @@ where
.iter()
.map(|(_, count)| count.is_some())
.collect::<Vec<_>>();
let values = Self::run_blocking(Arc::clone(&self.connections), move |connection| {
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = operations
.into_iter()
.map(|(key, count)| {
@ -365,7 +365,7 @@ where
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(script)
.arg(keys.len())
@ -426,7 +426,7 @@ where
if operations.is_empty() {
return Ok(Vec::new());
}
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut commands = Vec::with_capacity(operations.len() * 2);
let mut increments = Vec::with_capacity(operations.len());
for (key, amount, ttl) in operations {
@ -460,7 +460,7 @@ where
) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl);
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
increment_with_floor(connection, key, amount, ttl)
})
.await
@ -474,7 +474,7 @@ where
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
Self::run_blocking(Arc::clone(&self.connections), move |connection| {
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(SET_MAX_SCRIPT)
.arg(1)

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

@ -1,6 +1,10 @@
mod cache;
mod topology;
pub mod connection {
pub use crate::cache::{ConnectionRef, Connections};
}
pub use cache::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
};

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

@ -58,4 +58,4 @@ Verify typed values, TTL precedence, missing entries, serialization failures, na
Public SDK, Router, and proxy activation still need constructor parity, stream replay, embedding partial-batch integration, response reconstruction, callback scheduling, and failure-policy integration. This foundation does not switch those request paths
Redis cluster, disk, cloud stores, and semantic caching remain follow-ups. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees
Redis cluster, disk, and cloud stores remain follow-ups. Semantic backends plug in through `SemanticCacheContext`, which carries the prompt inputs and metadata alongside the cache TTL. The generic dual cache takes read, write, and remote-failure policies, runs its async operations through the async L2 methods, and provides L2-first counters and atomic affinity claims. Errors propagate by default, and `RemoteFailurePolicy::UseLocal` opts key-value operations and claims into the local tier when L2 is unavailable. Claims compare decoded values, so a pin written by Python still matches. Public Router integration remains follow-up work. Reservations and pubsub still need explicit capabilities owned by their consuming features. Adding a cache backend does not establish those guarantees

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

@ -1,21 +1,21 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, FlushCache,
};
use serde_json::Value;
use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key};
#[derive(Clone)]
pub struct ResponseCacheRequest {
pub struct ResponseCacheRequest<C: CacheContext = litellm_cache::ExactCacheContext> {
pub key: CacheKeyInput,
pub controls: CacheControls,
pub context: ExactCacheContext,
pub context: C,
pub max_age: Option<Duration>,
}
impl ResponseCacheRequest {
impl<C: CacheContext + Default> ResponseCacheRequest<C> {
pub fn new(key: CacheKeyInput) -> Self {
Self {
key,
@ -26,17 +26,24 @@ impl ResponseCacheRequest {
default_on: true,
..Default::default()
},
context: ExactCacheContext::default(),
context: C::default(),
max_age: None,
}
}
}
pub struct ResponseCache<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> {
pub struct ResponseCache<B: BaseCache<Value = CacheEntry>>
where
B::Context: Default + PartialEq,
{
backend: Arc<B>,
}
impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCache<B> {
impl<B> ResponseCache<B>
where
B: BaseCache<Value = CacheEntry>,
B::Context: Default + PartialEq,
{
pub fn new(backend: Arc<B>) -> Self {
Self { backend }
}
@ -45,8 +52,12 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
&self.backend
}
pub fn backend_arc(&self) -> &Arc<B> {
&self.backend
}
pub fn default_ttl(&self) -> Option<Duration> {
self.backend.get_ttl(&ExactCacheContext::default())
self.backend.get_ttl(&B::Context::default())
}
pub async fn async_flush(&self) -> Result<(), Error>
@ -62,7 +73,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn lookup(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
now: Duration,
) -> Result<Option<Value>, Error> {
if !request.controls.reads() {
@ -81,7 +92,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_lookup(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
now: Duration,
) -> Result<Option<Value>, Error> {
if !request.controls.reads() {
@ -101,7 +112,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[ResponseCacheRequest<B::Context>],
now: Duration,
) -> Result<PartialHits, Error>
where
@ -126,7 +137,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[ResponseCacheRequest<B::Context>],
now: Duration,
) -> Result<PartialHits, Error>
where
@ -153,7 +164,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub fn store(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
response: Value,
now: Duration,
) -> Result<(), Error> {
@ -172,7 +183,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_store(
&self,
request: &ResponseCacheRequest,
request: &ResponseCacheRequest<B::Context>,
response: Value,
now: Duration,
) -> Result<(), Error> {
@ -193,7 +204,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
pub async fn async_store_batch(
&self,
entries: Vec<(ResponseCacheRequest, Value)>,
entries: Vec<(ResponseCacheRequest<B::Context>, Value)>,
now: Duration,
) -> Result<(), Error> {
self.async_store_entries(
@ -209,7 +220,7 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
/// the freshness of its original response.
pub async fn async_store_entries(
&self,
entries: Vec<(ResponseCacheRequest, Value, Duration)>,
entries: Vec<(ResponseCacheRequest<B::Context>, Value, Duration)>,
) -> Result<(), Error> {
let writable = entries
.into_iter()
@ -249,8 +260,8 @@ impl<B: BaseCache<Value = CacheEntry, Context = ExactCacheContext>> ResponseCach
}
fn partial_hits(
requests: &[ResponseCacheRequest],
readable: Vec<(usize, &ResponseCacheRequest)>,
requests: &[ResponseCacheRequest<B::Context>],
readable: Vec<(usize, &ResponseCacheRequest<B::Context>)>,
entries: Vec<BatchEntry<CacheEntry>>,
now: Duration,
) -> Result<PartialHits, Error> {

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

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

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

File diff suppressed because it is too large Load diff

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

@ -32,6 +32,28 @@ impl CacheContext for ExactCacheContext {
}
}
#[derive(Clone, Debug, Default, PartialEq)]
pub struct SemanticCacheContext {
pub input: Option<serde_json::Value>,
pub messages: Option<serde_json::Value>,
pub metadata: Option<serde_json::Value>,
pub scope: Option<String>,
pub ttl: Option<Duration>,
}
impl CacheContext for SemanticCacheContext {
fn ttl(&self) -> Option<Duration> {
self.ttl
}
fn with_ttl(&self, ttl: Option<Duration>) -> Self {
Self {
ttl,
..self.clone()
}
}
}
#[derive(Clone, Copy, Debug, Deserialize, Serialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum CacheConnectionStatus {
@ -105,3 +127,31 @@ pub trait BaseCache: Send + Sync {
fn test_connection(&self) -> impl Future<Output = Result<CacheConnectionResult, Error>> + Send;
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use serde_json::json;
use super::{CacheContext, SemanticCacheContext};
#[test]
fn semantic_context_with_ttl_only_replaces_ttl() {
let context = SemanticCacheContext {
input: Some(json!({"input": "hello"})),
messages: Some(json!([{"role": "user", "content": "hello"}])),
metadata: Some(json!({"tenant": "team"})),
scope: Some("scope".into()),
ttl: Some(Duration::from_secs(10)),
};
let updated = context.with_ttl(Some(Duration::from_secs(20)));
assert_eq!(updated.ttl, Some(Duration::from_secs(20)));
assert_eq!(updated.input, context.input);
assert_eq!(updated.messages, context.messages);
assert_eq!(updated.metadata, context.metadata);
assert_eq!(updated.scope, context.scope);
}
}

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

@ -8,7 +8,7 @@ mod error;
pub use base_cache::{
BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext,
ExactCacheContext,
ExactCacheContext, SemanticCacheContext,
};
pub use cache_type::CacheType;
pub use caching::{Cache, CacheBackend, get_cache, set_cache};

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

@ -28,6 +28,7 @@ litellm-cache-s3.workspace = true
litellm-cache-gcs.workspace = true
litellm-cache-disk.workspace = true
litellm-cache-response.workspace = true
litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" }
serde.workspace = true
litellm-auth.workspace = true
litellm-auth-aws.workspace = true
@ -42,6 +43,7 @@ litellm-host-python.workspace = true
litellm-token-counter = { path = "../token-counter", default-features = false }
pyo3.workspace = true
pyo3-async-runtimes.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
serde_json.workspace = true
tokio = { workspace = true, features = ["sync"] }
@ -51,6 +53,7 @@ serde_with.workspace = true
criterion.workspace = true
futures-util.workspace = true
rstest.workspace = true
sha2.workspace = true
tokio-tungstenite.workspace = true
[[bench]]

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

@ -56,12 +56,7 @@ impl ResolvedCache {
CacheBinding::Disabled => ready_none(py)?,
CacheBinding::Native(service) => {
let request = request(input)?;
let service = service.clone();
run_async(
py,
async move { service.async_lookup(&request, now()).await },
cache_error,
)?
service.async_lookup_py(py, request)?
}
CacheBinding::PythonCallback(callback) => callback.async_lookup(py, kwargs)?,
};
@ -179,12 +174,7 @@ impl ResolvedCache {
CacheBinding::Native(service) => {
let request = self::request(request)?;
let response: Value = from_py(response)?;
let service = service.clone();
run_async(
py,
async move { service.async_store(&request, response, now()).await },
cache_error,
)
service.async_store_py(py, request, response)
}
CacheBinding::PythonCallback(callback) => {
callback.async_store(py, response, callback_kwargs)
@ -241,12 +231,7 @@ impl ResolvedCache {
));
}
let entries = requests.into_iter().zip(responses).collect();
let service = service.clone();
run_async(
py,
async move { service.async_store_batch(entries, now()).await },
cache_error,
)
service.async_store_batch_py(py, entries)
}
CacheBinding::PythonCallback(callback) => {
callback.async_store_batch(py, callback_result, callback_kwargs)

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

@ -88,11 +88,6 @@ pub(super) struct GcsCacheConfig {
pub(super) path_service_account: Option<String>,
}
pub(super) struct AzureBlobCacheConfig {
pub(super) account_url: String,
pub(super) container: String,
}
struct RedisClientProjection<'py> {
topology: RedisTopology,
host: String,
@ -104,11 +99,25 @@ struct RedisClientProjection<'py> {
const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31;
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
pub(super) struct ValkeySemanticCacheConfig {
pub(super) similarity_threshold: f64,
pub(super) index_name: String,
pub(super) embedding_model: String,
pub(super) connection: RedisConnectionConfig,
}
pub(super) struct AzureBlobCacheConfig {
pub(super) account_url: String,
pub(super) container: String,
}
pub(super) enum CacheBackendConfig {
Memory(MemoryCacheConfig),
Redis(Box<RedisCacheConfig>),
S3(Box<S3CacheConfig>),
Gcs(GcsCacheConfig),
ValkeySemantic(Box<ValkeySemanticCacheConfig>),
Disk(DiskCacheConfig),
AzureBlob(AzureBlobCacheConfig),
}
@ -201,6 +210,13 @@ impl NativeCacheConfig {
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::ValkeySemantic) => match project_valkey_semantic(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::ValkeySemantic(Box::new(backend)),
}))),
Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)),
},
Some(CacheType::Disk) => match project_disk(&backend)? {
Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self {
policy,
@ -214,12 +230,9 @@ impl NativeCacheConfig {
backend: CacheBackendConfig::AzureBlob(backend),
}))
}),
Some(
CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::QdrantSemantic,
)
| None => Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
)),
Some(CacheType::RedisSemantic | CacheType::QdrantSemantic) | None => Ok(
CacheConfigProjection::Unsupported(UnsupportedCacheConfig::Backend),
),
}
}
@ -228,11 +241,14 @@ impl NativeCacheConfig {
CacheBackendConfig::Memory(config) => Some(config.default_ttl),
CacheBackendConfig::Redis(config) => Some(config.default_ttl),
CacheBackendConfig::S3(_) => None,
CacheBackendConfig::ValkeySemantic(_) => Some(Duration::ZERO),
CacheBackendConfig::Disk(_)
| CacheBackendConfig::AzureBlob(_)
| CacheBackendConfig::Gcs(_) => None,
};
if service.default_ttl() != default_ttl {
if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_))
&& service.default_ttl() != default_ttl
{
return Some("facade and native backend default TTLs must match");
}
match &self.backend {
@ -306,6 +322,16 @@ impl NativeCacheConfig {
Some("facade and native backend credentials must match")
}
CacheBackendConfig::Gcs(_) => None,
CacheBackendConfig::ValkeySemantic(config) => {
if service.kind() != "valkey-semantic" {
return Some("facade and native backend types must match");
}
let Some((threshold, index_name)) = service.semantic_config() else {
return Some("facade and native backend types must match");
};
(threshold != config.similarity_threshold || index_name != config.index_name)
.then_some("facade and native semantic settings must match")
}
CacheBackendConfig::Disk(_) if service.kind() != "disk" => {
Some("facade and native backend types must match")
}
@ -588,7 +614,7 @@ fn project_standalone_client<'py>(
#[inline(never)]
fn project_cluster_client<'py>(
source: &Bound<'py, PyDict>,
source: &Bound<'_, PyDict>,
client: &Bound<'py, PyAny>,
) -> PyResult<Result<RedisClientProjection<'py>, UnsupportedCacheConfig>> {
let Some(startup_nodes) = startup_nodes(source)? else {
@ -676,6 +702,71 @@ fn port(value: i64) -> PyResult<u16> {
u16::try_from(value).map_err(|_| PyValueError::new_err("invalid Redis port"))
}
#[inline(never)]
fn project_valkey_semantic(
backend: &Bound<'_, PyAny>,
) -> PyResult<Result<ValkeySemanticCacheConfig, UnsupportedCacheConfig>> {
let client = backend.getattr("sync_client")?;
let pool = client.getattr("connection_pool")?;
let Ok((resolved, is_tls)) = project_connection_pool(&pool)? else {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
};
for key in ["credential_provider", "redis_connect_func"] {
if has_value(&resolved, key)? {
return Ok(Err(UnsupportedCacheConfig::RedisCredentials));
}
}
if is_tls {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let connection = RedisConnectionConfig {
host: required_string(&resolved, "host")?,
port: u16::try_from(required_i64(&resolved, "port")?)
.map_err(|_| PyValueError::new_err("invalid Redis port"))?,
database: optional_i64(&resolved, "db")?.unwrap_or(0),
username: optional_dict_string(&resolved, "username")?,
password: optional_dict_string(&resolved, "password")?,
protocol: RedisProtocol::Resp2,
pool_size: pool.getattr("max_connections")?.extract::<usize>()?,
read_timeout: None,
connect_timeout: None,
socket_keepalive: None,
health_check_interval: Duration::ZERO,
client_name: None,
tls: None,
};
if connection.host.is_empty() {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
Ok(Ok(ValkeySemanticCacheConfig {
similarity_threshold: backend.getattr("similarity_threshold")?.extract()?,
index_name: backend.getattr("index_name")?.extract()?,
embedding_model: backend.getattr("embedding_model")?.extract()?,
connection,
}))
}
#[inline(never)]
fn project_connection_pool<'py>(
pool: &Bound<'py, PyAny>,
) -> PyResult<Result<(Bound<'py, PyDict>, bool), UnsupportedCacheConfig>> {
if !instance_class_is(pool, "redis.connection", "ConnectionPool")? {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
}
let resolved = pool.getattr("connection_kwargs")?.cast_into::<PyDict>()?;
let connection_class = resolved
.get_item("connection_class")?
.unwrap_or(pool.getattr("connection_class")?);
let is_tls = if class_is(&connection_class, "redis.connection", "Connection")? {
false
} else if class_is(&connection_class, "redis.connection", "SSLConnection")? {
true
} else {
return Ok(Err(UnsupportedCacheConfig::RedisConnection));
};
Ok(Ok((resolved, is_tls)))
}
#[inline(never)]
fn project_tls(values: &Bound<'_, PyDict>) -> PyResult<RedisTlsConfig> {
Ok(RedisTlsConfig {
@ -869,16 +960,13 @@ mod tests {
use pyo3::{prelude::*, types::PyDict};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use litellm_host_python::run_sync_value;
use litellm_cache_redis::{RedisNode, RedisTopology};
use super::{
CacheBackendConfig, CacheConfigProjection, CachePolicy, CertificateRequirement,
DiskCacheConfig, GcsCacheConfig, NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
CacheBackendConfig, CacheConfigProjection, CertificateRequirement, GcsCacheConfig,
NativeCacheConfig, RedisProtocol, UnsupportedCacheConfig,
};
use crate::cache::native::NativeResponseCache;
use litellm_cache_redis::{RedisNode, RedisTopology};
fn cluster_facade<'py>(py: Python<'py>, startup_nodes: &str, hook: &str) -> Bound<'py, PyAny> {
facade(
@ -951,6 +1039,194 @@ mod tests {
});
}
#[test]
fn projects_resolved_redis_tls_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = SSLConnection\n\
pool.max_connections = 29\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6380, 'db': 4, 'username': 'user', 'password': 'secret', 'protocol': 3, 'socket_timeout': 7.5, 'socket_connect_timeout': 2, 'socket_keepalive': True, 'health_check_interval': 15, 'client_name': 'litellm', 'ssl_cert_reqs': 'optional', 'ssl_check_hostname': True, 'ssl_ca_certs': '/ca.pem', 'ssl_ca_data': 'CA DATA', 'ssl_certfile': '/client.pem', 'ssl_keyfile': '/client.key'}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\
facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("Redis cache should be supported");
};
let CacheBackendConfig::Redis(redis) = config.backend else {
panic!("expected Redis configuration");
};
assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777));
assert_eq!(redis.namespace.as_deref(), Some("team"));
assert_eq!(redis.flush_size, 31);
assert_eq!(redis.connection.host, "cache.internal");
assert_eq!(redis.connection.port, 6380);
assert_eq!(redis.connection.database, 4);
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
assert_eq!(redis.connection.pool_size, 29);
let tls = redis.connection.tls.unwrap();
assert_eq!(
tls.certificate_requirement,
CertificateRequirement::Optional
);
assert!(tls.check_hostname);
assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem"));
assert_eq!(tls.ca_data.as_deref(), Some("CA DATA"));
assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem"));
assert_eq!(tls.client_key.as_deref(), Some("/client.key"));
});
}
#[test]
fn projects_valkey_semantic_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = Connection\n\
pool.max_connections = 12\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'db': 2}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("Valkey semantic cache should be supported");
};
let CacheBackendConfig::ValkeySemantic(valkey) = config.backend else {
panic!("expected Valkey semantic configuration");
};
assert_eq!(valkey.similarity_threshold, 0.85);
assert_eq!(valkey.index_name, "semantic_idx");
assert_eq!(valkey.embedding_model, "text-embedding-3-small");
assert_eq!(valkey.connection.host, "cache.internal");
assert_eq!(valkey.connection.port, 6390);
assert_eq!(valkey.connection.database, 2);
assert_eq!(valkey.connection.pool_size, 12);
assert_eq!(valkey.connection.protocol, RedisProtocol::Resp2);
assert!(valkey.connection.tls.is_none());
});
}
#[test]
fn valkey_semantic_tls_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = SSLConnection\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("TLS Valkey semantic cache should stay on Python");
};
assert_eq!(
reason.message(),
"native Redis connection type is not implemented"
);
});
}
#[test]
fn valkey_semantic_dynamic_auth_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = Connection\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'credential_provider': object()}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\
facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("dynamic Valkey authentication must stay on Python");
};
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[test]
fn dynamic_redis_auth_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\
facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("dynamic authentication must stay on Python");
};
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[test]
fn projects_cluster_startup_nodes_as_redis_topology() {
Python::initialize();
Python::attach(|py| {
let facade = cluster_facade(
py,
"[{'host': 'node-a', 'port': 7000}, {'host': 'node-b', 'port': 7001}]",
"client.on_connect",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("cluster startup nodes should project natively");
};
let CacheBackendConfig::Redis(redis) = &config.backend else {
panic!("expected Redis configuration");
};
let expected = RedisTopology::Cluster {
startup_nodes: vec![
RedisNode {
host: "node-a".into(),
port: 7000,
},
RedisNode {
host: "node-b".into(),
port: 7001,
},
],
};
assert_eq!(redis.topology, expected);
assert_eq!(redis.connection.host, "node-a");
assert_eq!(redis.connection.port, 7000);
assert_eq!(redis.connection.password.as_deref(), Some("secret"));
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
assert_eq!(
redis
.connection
.tls
.as_ref()
.unwrap()
.certificate_requirement,
CertificateRequirement::None
);
});
}
#[test]
fn projects_gcs_configuration() {
Python::initialize();
@ -1016,375 +1292,6 @@ mod tests {
});
}
#[test]
fn projects_resolved_redis_tls_configuration() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"pool = ConnectionPool()\n\
pool.connection_class = SSLConnection\n\
pool.max_connections = 29\n\
pool.connection_kwargs = {'host': 'cache.internal', 'port': 6380, 'db': 4, 'username': 'user', 'password': 'secret', 'protocol': 3, 'socket_timeout': 7.5, 'socket_connect_timeout': 2, 'socket_keepalive': True, 'health_check_interval': 15, 'client_name': 'litellm', 'ssl_cert_reqs': 'optional', 'ssl_check_hostname': True, 'ssl_ca_certs': '/ca.pem', 'ssl_ca_data': 'CA DATA', 'ssl_certfile': '/client.pem', 'ssl_keyfile': '/client.key'}\n\
client = SimpleNamespace(connection_pool=pool)\n\
backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\
facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("Redis cache should be supported");
};
let CacheBackendConfig::Redis(redis) = config.backend else {
panic!("expected Redis configuration");
};
assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777));
assert_eq!(redis.namespace.as_deref(), Some("team"));
assert_eq!(redis.flush_size, 31);
assert_eq!(redis.connection.host, "cache.internal");
assert_eq!(redis.connection.port, 6380);
assert_eq!(redis.connection.database, 4);
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
assert_eq!(redis.connection.pool_size, 29);
let tls = redis.connection.tls.unwrap();
assert_eq!(
tls.certificate_requirement,
CertificateRequirement::Optional
);
assert!(tls.check_hostname);
assert_eq!(tls.ca_certificate.as_deref(), Some("/ca.pem"));
assert_eq!(tls.ca_data.as_deref(), Some("CA DATA"));
assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem"));
assert_eq!(tls.client_key.as_deref(), Some("/client.key"));
});
}
#[test]
fn dynamic_redis_auth_stays_on_python() {
Python::initialize();
Python::attach(|py| {
let facade = facade(
py,
"backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\
facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("dynamic authentication must stay on Python");
};
assert_eq!(reason.message(), "native Redis credentials require Python");
});
}
#[test]
fn projects_builtin_disk_configuration_and_rejects_custom_stores() {
Python::initialize();
Python::attach(|py| {
let root =
std::env::temp_dir().join(format!("litellm-disk-config-{}", std::process::id()));
let directory = root.to_string_lossy();
let disk_facade = facade(
py,
&format!(
"Cache = type('Cache', (), {{'__module__': 'diskcache.core'}})\n\
Disk = type('Disk', (), {{'__module__': 'diskcache.core'}})\n\
store = Cache()\n\
store._disk = Disk()\n\
store.directory = {directory:?}\n\
backend = SimpleNamespace(disk_cache=store)\n\
facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)"
),
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&disk_facade).unwrap()
else {
panic!("disk cache should be supported");
};
let CacheBackendConfig::Disk(disk) = config.backend else {
panic!("expected disk configuration");
};
assert_eq!(disk.directory, root);
let matching = NativeResponseCache::disk(&directory).unwrap();
assert_eq!(
(NativeCacheConfig {
policy: config.policy,
backend: CacheBackendConfig::Disk(disk),
})
.service_mismatch(&matching),
None
);
let other = NativeResponseCache::disk(&root.join("other").to_string_lossy()).unwrap();
let mismatch = NativeCacheConfig {
policy: CachePolicy {
mode: "default-on".into(),
ttl: None,
namespace: None,
supported_call_types: None,
redis_flush_size: None,
semantic_cache_scope: "key".into(),
},
backend: CacheBackendConfig::Disk(DiskCacheConfig {
directory: root.clone(),
}),
};
assert_eq!(
mismatch.service_mismatch(&other),
Some("facade and native backend directories must match")
);
let custom = facade(
py,
&format!(
"CustomCache = type('CustomCache', (), {{'__module__': 'mypkg'}})\n\
CustomDisk = type('CustomDisk', (), {{'__module__': 'mypkg'}})\n\
store = CustomCache()\n\
store._disk = CustomDisk()\n\
store.directory = {directory:?}\n\
backend = SimpleNamespace(disk_cache=store)\n\
facade = SimpleNamespace(type='disk', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)"
),
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&custom).unwrap()
else {
panic!("custom disk store must stay on Python");
};
assert_eq!(
reason.message(),
"native disk cache requires the built-in diskcache store"
);
});
}
fn s3_facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> {
let locals = PyDict::new(py);
py.run(
&CString::new(format!(
"from types import SimpleNamespace\n\
S3Client = type('S3', (), {{'__module__': 'botocore.client'}})\n\
client = S3Client()\n\
client.meta = SimpleNamespace(region_name='us-east-1', endpoint_url='https://example.test', config=SimpleNamespace(s3=None, proxies=None, client_cert=None, signature_version='s3v4'))\n\
client._endpoint = SimpleNamespace(http_session=SimpleNamespace(_verify=True))\n\
client._request_signer = SimpleNamespace(_credentials=SimpleNamespace(method='explicit', access_key='key', secret_key='secret', token='token'))\n\
backend = SimpleNamespace(bucket_name='bucket', key_prefix='team/', s3_client=client)\n\
facade = SimpleNamespace(type='s3', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)\n\
{body}"
))
.unwrap(),
None,
Some(&locals),
)
.unwrap();
locals.get_item("facade").unwrap().unwrap()
}
#[test]
fn projects_s3_configuration_with_explicit_credentials_and_custom_endpoint() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(py, "");
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert_eq!(s3.bucket, "bucket");
assert_eq!(s3.key_prefix, "team/");
assert_eq!(s3.region, "us-east-1");
assert_eq!(
s3.endpoint.map(|endpoint| endpoint.url).as_deref(),
Some("https://example.test")
);
assert_eq!(s3.auth.access_key_id.as_deref(), Some("key"));
assert_eq!(s3.auth.secret_access_key.as_deref(), Some("secret"));
assert_eq!(s3.auth.session_token.as_deref(), Some("token"));
assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1"));
});
}
#[test]
fn default_s3_endpoint_projects_no_custom_endpoint() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(
py,
"facade.cache.s3_client.meta.endpoint_url = 'https://s3.us-east-1.amazonaws.com'",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert!(s3.endpoint.is_none());
});
}
#[test]
fn non_sigv4_proxies_and_disabled_verification_stay_on_python() {
Python::initialize();
Python::attach(|py| {
for (body, message) in [
(
"facade.cache.s3_client.meta.config.signature_version = 's3'",
"native S3 configuration requires Python",
),
(
"facade.cache.s3_client.meta.config.proxies = {'https': 'proxy'}",
"native S3 configuration requires Python",
),
(
"facade.cache.s3_client._endpoint.http_session._verify = False",
"native S3 configuration requires Python",
),
(
"del facade.cache.s3_client._endpoint.http_session._verify",
"native S3 configuration requires Python",
),
] {
let facade = s3_facade(py, body);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("{body} must stay on Python");
};
assert_eq!(reason.message(), message);
}
let facade = s3_facade(py, "facade.cache.s3_client = SimpleNamespace()");
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("non-botocore client must stay on Python");
};
assert_eq!(reason.message(), "native S3 client type is not implemented");
let facade = s3_facade(
py,
"facade.cache.s3_client._request_signer._credentials = None",
);
let CacheConfigProjection::Unsupported(reason) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("missing credentials must stay on Python");
};
assert_eq!(reason.message(), "native S3 credentials require Python");
});
}
#[test]
fn non_explicit_s3_credentials_use_the_default_chain() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(
py,
"facade.cache.s3_client._request_signer._credentials = SimpleNamespace(method='sso', access_key=None, secret_key=None, token=None)",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("default-chain credentials should be supported");
};
let CacheBackendConfig::S3(s3) = config.backend else {
panic!("expected S3 configuration");
};
assert_eq!(s3.auth.access_key_id, None);
assert_eq!(s3.auth.secret_access_key, None);
assert_eq!(s3.auth.region_name.as_deref(), Some("us-east-1"));
});
}
fn s3_service(py: Python<'_>, region: &str, endpoint: Option<&str>) -> NativeResponseCache {
let config = S3CacheConfig {
bucket: "bucket".to_string(),
key_prefix: "team/".to_string(),
region: region.to_string(),
endpoint: endpoint.map(|url| S3Endpoint {
url: url.to_string(),
}),
auth: AwsAuthConfig::default(),
};
run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) }).unwrap()
}
#[test]
fn s3_binding_rejects_region_and_endpoint_mismatches() {
Python::initialize();
Python::attach(|py| {
let facade = s3_facade(py, "");
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("S3 cache should be supported");
};
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", Some("https://example.test"))),
None
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-west-2", Some("https://example.test"))),
Some("facade and native backend regions must match")
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", Some("https://other.test"))),
Some("facade and native backend endpoints must match")
);
assert_eq!(
config.service_mismatch(&s3_service(py, "us-east-1", None)),
Some("facade and native backend endpoints must match")
);
});
}
#[test]
fn projects_cluster_startup_nodes_as_redis_topology() {
Python::initialize();
Python::attach(|py| {
let facade = cluster_facade(
py,
"[{'host': 'node-a', 'port': 7000}, {'host': 'node-b', 'port': 7001}]",
"client.on_connect",
);
let CacheConfigProjection::Native(config) =
NativeCacheConfig::project(&facade).unwrap()
else {
panic!("cluster startup nodes should project natively");
};
let CacheBackendConfig::Redis(redis) = &config.backend else {
panic!("expected Redis configuration");
};
let expected = RedisTopology::Cluster {
startup_nodes: vec![
RedisNode {
host: "node-a".into(),
port: 7000,
},
RedisNode {
host: "node-b".into(),
port: 7001,
},
],
};
assert_eq!(redis.topology, expected);
assert_eq!(redis.connection.host, "node-a");
assert_eq!(redis.connection.port, 7000);
assert_eq!(redis.connection.password.as_deref(), Some("secret"));
assert_eq!(redis.connection.protocol, RedisProtocol::Resp3);
assert_eq!(
redis
.connection
.tls
.as_ref()
.unwrap()
.certificate_requirement,
CertificateRequirement::None
);
});
}
#[test]
fn malformed_startup_nodes_and_foreign_connect_hooks_stay_on_python() {
Python::initialize();

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

@ -0,0 +1,63 @@
use std::{future::Future, sync::Arc};
use litellm_cache::Error;
use litellm_cache_valkey_semantic::Embedder;
use litellm_host_python::to_py;
use pyo3::{PyTraverseError, PyVisit, prelude::*};
use serde_json::Value;
#[derive(Clone)]
pub(super) struct PythonEmbedder {
sync_embed: Arc<Py<PyAny>>,
async_embed_callable: Arc<Py<PyAny>>,
}
impl PythonEmbedder {
pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
sync_embed: Arc::new(backend.getattr("_get_embedding")?.unbind()),
async_embed_callable: Arc::new(backend.getattr("_get_async_embedding")?.unbind()),
})
}
pub(super) fn async_embed_awaitable<'py>(
&self,
py: Python<'py>,
prompt: &str,
metadata: &Option<Value>,
) -> PyResult<Bound<'py, PyAny>> {
let metadata = to_py(py, metadata)?;
self.async_embed_callable.bind(py).call1((prompt, metadata))
}
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&*self.sync_embed)?;
visit.call(&*self.async_embed_callable)
}
}
impl Embedder for PythonEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
let result = Python::attach(|py| -> PyResult<Vec<f64>> {
let metadata = to_py(py, &metadata)?;
self.sync_embed
.bind(py)
.call1((prompt, metadata))?
.extract()
})
.map_err(|_| Error::Unavailable)?;
Ok(result.into_iter().map(|value| value as f32).collect())
}
#[expect(
clippy::manual_async_fn,
reason = "the shared Embedder trait uses an impl Future return"
)]
fn async_embed(
&self,
_prompt: &str,
_metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send {
async { Err(Error::Unavailable) }
}
}

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

@ -31,17 +31,15 @@ struct RedisPoolGuard {
connection_class: Py<PyAny>,
connection_kwargs: Py<PyAny>,
max_connections: Option<usize>,
client_name: &'static str,
attributes: RedisPoolAttributes,
}
struct S3ClientGuard {
reference: Py<PyAny>,
}
struct DiskStoreGuard {
reference: Py<PyAny>,
directory: String,
}
struct AzureBlobClientGuard {
sync_client: Py<PyAny>,
async_client: Py<PyAny>,
@ -49,12 +47,18 @@ struct AzureBlobClientGuard {
container_name: String,
}
struct S3ClientGuard {
reference: Py<PyAny>,
}
enum ConnectionGuard {
None,
RedisPool(RedisPoolGuard),
AzureBlob(AzureBlobClientGuard),
S3(S3ClientGuard),
}
#[derive(Clone, Copy)]
struct RedisPoolAttributes {
pool: &'static str,
connection_class: &'static str,
@ -72,6 +76,9 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes {
connection_class: "connection_pool_class",
max_connections: None,
};
const VALKEY_POOL: RedisPoolAttributes = STANDALONE_POOL;
pub(super) struct FacadeGuard {
outer: ObjectGuard,
backend: ObjectGuard,
@ -158,7 +165,9 @@ impl ObjectGuard {
return Ok(false);
}
for (name, value) in &expected.attributes {
if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) {
if (instance.contains(name)? && !self.config_names.contains(&name.as_str()))
|| !attributes.get_item(name)?.is(value.bind(py))
{
return Ok(false);
}
}
@ -179,8 +188,12 @@ impl ObjectGuard {
}
impl RedisPoolGuard {
fn capture(backend: &Bound<'_, PyAny>, attributes: RedisPoolAttributes) -> PyResult<Self> {
let pool = backend.getattr("redis_client")?.getattr(attributes.pool)?;
fn capture(
backend: &Bound<'_, PyAny>,
client_name: &'static str,
attributes: RedisPoolAttributes,
) -> PyResult<Self> {
let pool = backend.getattr(client_name)?.getattr(attributes.pool)?;
Ok(Self {
reference: pool.clone().unbind(),
connection_class: pool.getattr(attributes.connection_class)?.unbind(),
@ -188,31 +201,30 @@ impl RedisPoolGuard {
.getattr("connection_kwargs")?
.call_method0("copy")?
.unbind(),
max_connections: Self::max_connections(&pool, &attributes)?,
max_connections: attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()?,
client_name,
attributes,
})
}
fn max_connections(
pool: &Bound<'_, PyAny>,
attributes: &RedisPoolAttributes,
) -> PyResult<Option<usize>> {
attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
let pool = backend
.getattr("redis_client")?
.getattr(self.client_name)?
.getattr(self.attributes.pool)?;
Ok(self.reference.bind(py).is(&pool)
&& self
.connection_class
.bind(py)
.is(&pool.getattr(self.attributes.connection_class)?)
&& self.max_connections == Self::max_connections(&pool, &self.attributes)?
&& self.max_connections
== self
.attributes
.max_connections
.map(|name| pool.getattr(name)?.extract::<usize>())
.transpose()?
&& self
.connection_kwargs
.bind(py)
@ -226,22 +238,6 @@ impl RedisPoolGuard {
}
}
impl S3ClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
reference: backend.getattr("s3_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
Ok(self.reference.bind(py).is(&backend.getattr("s3_client")?))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)
}
}
impl DiskStoreGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let store = backend.getattr("disk_cache")?;
@ -261,6 +257,7 @@ impl DiskStoreGuard {
visit.call(&self.reference)
}
}
impl AzureBlobClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
let sync_client = backend.getattr("container_client")?;
@ -289,11 +286,41 @@ impl AzureBlobClientGuard {
}
}
impl S3ClientGuard {
fn capture(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
reference: backend.getattr("s3_client")?.unbind(),
})
}
fn matches(&self, py: Python<'_>, backend: &Bound<'_, PyAny>) -> PyResult<bool> {
Ok(self.reference.bind(py).is(&backend.getattr("s3_client")?))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.reference)
}
}
impl ConnectionGuard {
fn capture(kind: &str, cluster: bool, backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(match (kind, cluster) {
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(backend, STANDALONE_POOL)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(backend, CLUSTER_POOL)?),
("redis", false) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"redis_client",
STANDALONE_POOL,
)?),
("redis", true) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"redis_client",
CLUSTER_POOL,
)?),
("valkey-semantic", _) => Self::RedisPool(RedisPoolGuard::capture(
backend,
"sync_client",
VALKEY_POOL,
)?),
("disk", _) => Self::None,
("azure-blob", _) => Self::AzureBlob(AzureBlobClientGuard::capture(backend)?),
("s3", _) => Self::S3(S3ClientGuard::capture(backend)?),
_ => Self::None,
@ -318,6 +345,7 @@ impl ConnectionGuard {
}
}
}
impl FacadeGuard {
pub(super) fn capture(
py: Python<'_>,
@ -340,14 +368,19 @@ impl FacadeGuard {
"RedisClusterCache",
"redis",
),
("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"),
("gcs", _) => ("litellm.caching.gcs_cache", "GCSCache", "gcs"),
("valkey-semantic", false) => (
"litellm.caching.valkey_semantic_cache",
"ValkeySemanticCache",
"valkey-semantic",
),
("disk", _) => ("litellm.caching.disk_cache", "DiskCache", "disk"),
("azure-blob", _) => (
"litellm.caching.azure_blob_cache",
"AzureBlobCache",
"azure-blob",
),
("s3", _) => ("litellm.caching.s3_cache", "S3Cache", "s3"),
_ => unreachable!(),
};
let backend = facade.getattr("cache")?;
@ -391,6 +424,11 @@ impl FacadeGuard {
"max_size_per_item",
"redis_kwargs",
"redis_flush_size",
"similarity_threshold",
"embedding_model",
"index_name",
"embedding_max_input_tokens",
"embedding_timeout",
"bucket_name",
"key_prefix",
"path_service_account",

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

@ -1,12 +1,14 @@
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use litellm_host_python::{release_gil, run_sync_value};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration};
use super::{
cache_error, embedder::PythonEmbedder, facade::FacadeGuard, native::NativeResponseCache,
request::duration,
};
#[pyclass(frozen, name = "_CacheTestHandle")]
pub(crate) struct CacheTestHandle {
@ -139,6 +141,29 @@ impl CacheTestHandle {
})
}
#[staticmethod]
#[pyo3(signature = (url, similarity_threshold, index_name, embedder))]
fn valkey_semantic(
url: String,
similarity_threshold: f64,
index_name: String,
embedder: &Bound<'_, PyAny>,
) -> PyResult<Self> {
let python_embedder = PythonEmbedder::from_backend(embedder)?;
let service = NativeResponseCache::valkey_semantic(
&url,
similarity_threshold,
index_name,
python_embedder,
)
.map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[staticmethod]
#[pyo3(signature = (account_url, container))]
fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult<Self> {
@ -153,6 +178,7 @@ impl CacheTestHandle {
pid: std::process::id(),
})
}
#[getter]
fn backend(&self) -> &'static str {
self.service.kind()
@ -161,11 +187,17 @@ impl CacheTestHandle {
fn _bind_facade(&self, py: Python<'_>, facade: &Bound<'_, PyAny>) -> PyResult<()> {
let service = self.service()?;
let guard = FacadeGuard::capture(py, facade, &service)?;
let service = service.with_redis_flush_size(
facade
.getattr("redis_flush_size")?
.extract::<Option<usize>>()?,
);
let service = service
.with_scope(
facade
.getattr("semantic_cache_scope")?
.extract::<String>()?,
)
.with_redis_flush_size(
facade
.getattr("redis_flush_size")?
.extract::<Option<usize>>()?,
);
let handle = Py::new(
py,
Self {

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

@ -1,12 +1,14 @@
mod binding;
mod callback;
mod config;
mod embedder;
mod facade;
mod future;
mod handle;
mod native;
mod request;
mod resolver;
mod semantic_step;
use litellm_cache::Error;
use pyo3::{

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

@ -1,17 +1,77 @@
use std::{path::Path, sync::Arc, time::Duration};
use litellm_cache::{CacheCodec, CacheConnectionResult, Error};
use litellm_cache::{
CacheCodec, CacheConnectionResult, Error, ExactCacheContext, SemanticCacheContext,
};
use litellm_cache_azure_blob::AzureBlobCache;
use litellm_cache_disk::DiskCache;
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
use litellm_cache_response::{
CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer,
CacheEntry, CacheKeyField, PartialHits, ResponseCache, ResponseCacheCodec,
ResponseCacheRequest, WriteBuffer,
};
use litellm_cache_s3::{S3Cache, S3CacheConfig};
use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig};
use pyo3::prelude::*;
use serde_json::Value;
use super::{
embedder::PythonEmbedder,
request::NativeRequest,
semantic_step::{SemanticEmbedExecution, drive_semantic},
};
fn semantic_key(request: &NativeRequest, scope: &str) -> litellm_cache_response::CacheKeyInput {
let mut key = request.key.clone();
if key.preset.is_some() {
return key;
}
key.fields
.retain(|field| !matches!(field.name.as_str(), "messages" | "prompt" | "input"));
const TENANT: [&str; 3] = [
"user_api_key",
"user_api_key_team_id",
"user_api_key_org_id",
];
let end_user = (scope == "end_user").then_some("user_api_key_end_user_id");
for name in TENANT.into_iter().chain(end_user) {
let sources = [
request.metadata.as_ref(),
request.litellm_metadata.as_ref(),
request
.litellm_params
.as_ref()
.and_then(|params| params.get("metadata")),
request
.litellm_params
.as_ref()
.and_then(|params| params.get("litellm_metadata")),
];
let Some(value) = sources.into_iter().flatten().find_map(|source| {
source
.as_object()
.and_then(|values| values.get(name))
.filter(|value| !value.is_null())
}) else {
continue;
};
let value = match value {
Value::Null => continue,
Value::String(text) => text.clone(),
other => other.to_string(),
};
key.fields.push(CacheKeyField {
name: name.to_owned(),
value: Some(value),
api_parameter: true,
internal_parameter: false,
});
}
key
}
#[derive(Clone)]
pub(super) enum NativeResponseCache {
Memory(Arc<ResponseCache<InMemoryCache<CacheEntry>>>),
@ -21,6 +81,11 @@ pub(super) enum NativeResponseCache {
},
S3(Arc<ResponseCache<S3Cache<ResponseCacheCodec>>>),
Gcs(Arc<ResponseCache<GcsCache<ResponseCacheCodec>>>),
ValkeySemantic {
cache: Arc<ResponseCache<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>>,
embedder: PythonEmbedder,
scope: String,
},
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
}
@ -53,6 +118,7 @@ impl NativeResponseCache {
buffer: None,
})
}
pub async fn s3(config: S3CacheConfig) -> Self {
let runtime = tokio::runtime::Handle::current();
Self::S3(Arc::new(ResponseCache::new(Arc::new(S3Cache::new(
@ -61,6 +127,29 @@ impl NativeResponseCache {
runtime,
)))))
}
pub fn valkey_semantic(
url: &str,
similarity_threshold: f64,
index_name: String,
embedder: PythonEmbedder,
) -> Result<Self, Error> {
let backend = ValkeySemanticCache::new(
url,
embedder.clone(),
ResponseCacheCodec,
ValkeySemanticConfig {
similarity_threshold,
index_name,
},
)?;
Ok(Self::ValkeySemantic {
cache: Arc::new(ResponseCache::new(Arc::new(backend))),
embedder,
scope: String::from("key"),
})
}
pub fn disk(directory: &str) -> Result<Self, Error> {
let cache = DiskCache::open(directory, ResponseCacheCodec)?;
Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache)))))
@ -97,20 +186,72 @@ impl NativeResponseCache {
cache.backend().account_url(),
cache.backend().container_name(),
)),
Self::Memory(_) | Self::Redis { .. } | Self::S3(_) | Self::Disk(_) | Self::Gcs(_) => {
None
}
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::Disk(_)
| Self::Gcs(_) => None,
}
}
fn exact(request: &NativeRequest) -> ResponseCacheRequest<ExactCacheContext> {
ResponseCacheRequest {
key: request.key.clone(),
controls: request.controls,
context: ExactCacheContext { ttl: request.ttl },
max_age: request.max_age,
}
}
fn semantic(
request: &NativeRequest,
scope: &str,
) -> ResponseCacheRequest<SemanticCacheContext> {
ResponseCacheRequest {
key: semantic_key(request, scope),
controls: request.controls,
context: SemanticCacheContext {
input: request.input.clone(),
messages: request.messages.clone(),
metadata: request.metadata.clone(),
scope: Some(scope.to_owned()),
ttl: request.ttl,
},
max_age: request.max_age,
}
}
pub fn with_redis_flush_size(self, flush_size: Option<usize>) -> Self {
match self {
Self::Redis { cache, .. } => Self::Redis {
cache,
buffer: flush_size.map(|size| Arc::new(WriteBuffer::new(size))),
},
value => value,
}
}
pub fn with_scope(self, scope: String) -> Self {
match self {
Self::ValkeySemantic {
cache, embedder, ..
} => Self::ValkeySemantic {
cache,
embedder,
scope,
},
value => value,
}
}
}
impl NativeResponseCache {
pub fn kind(&self) -> &'static str {
match self {
Self::Memory(_) => "memory",
Self::Redis { .. } => "redis",
Self::S3(_) => "s3",
Self::Gcs(_) => "gcs",
Self::ValkeySemantic { .. } => "valkey-semantic",
Self::Disk(_) => "disk",
Self::AzureBlob(_) => "azure-blob",
}
@ -122,6 +263,7 @@ impl NativeResponseCache {
Self::Redis { cache, .. } => cache.default_ttl(),
Self::S3(cache) => cache.default_ttl(),
Self::Gcs(cache) => cache.default_ttl(),
Self::ValkeySemantic { cache, .. } => cache.default_ttl(),
Self::Disk(cache) => cache.default_ttl(),
Self::AzureBlob(cache) => cache.default_ttl(),
}
@ -157,17 +299,25 @@ impl NativeResponseCache {
pub fn namespace(&self) -> Option<&str> {
match self {
Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) => None,
Self::Memory(_)
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
Self::Redis { cache, .. } => cache.backend().namespace(),
Self::S3(_) | Self::Gcs(_) => None,
}
}
pub fn topology(&self) -> Option<&RedisTopology> {
match self {
Self::Memory(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None,
Self::Memory(_)
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
Self::Redis { cache, .. } => Some(cache.backend().topology()),
Self::S3(_) => None,
}
}
@ -176,6 +326,7 @@ impl NativeResponseCache {
Self::Memory(cache) => Some(cache.backend().max_size_in_memory()),
Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
@ -187,144 +338,386 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.backend().max_entry_bytes(),
Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
}
}
pub fn with_redis_flush_size(self, flush_size: Option<usize>) -> Self {
match self {
Self::Redis { cache, .. } => Self::Redis {
cache,
buffer: flush_size.map(|flush_size| Arc::new(WriteBuffer::new(flush_size))),
},
other => other,
}
}
pub fn directory(&self) -> Option<&Path> {
match self {
Self::Disk(cache) => Some(cache.backend().directory()),
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
}
}
pub fn lookup(
&self,
request: &ResponseCacheRequest,
now: Duration,
) -> Result<Option<Value>, Error> {
pub fn semantic_config(&self) -> Option<(f64, &str)> {
match self {
Self::Memory(cache) => cache.lookup(request, now),
Self::Redis { cache, .. } => cache.lookup(request, now),
Self::S3(cache) => cache.lookup(request, now),
Self::Gcs(cache) => cache.lookup(request, now),
Self::Disk(cache) => cache.lookup(request, now),
Self::AzureBlob(cache) => cache.lookup(request, now),
Self::ValkeySemantic { cache, .. } => Some((
cache.backend().similarity_threshold(),
cache.backend().index_name(),
)),
_ => None,
}
}
pub fn lookup(&self, request: &NativeRequest, now: Duration) -> Result<Option<Value>, Error> {
match self {
Self::Memory(cache) => cache.lookup(&Self::exact(request), now),
Self::Redis { cache, .. } => cache.lookup(&Self::exact(request), now),
Self::S3(cache) => cache.lookup(&Self::exact(request), now),
Self::ValkeySemantic { cache, scope, .. } => {
cache.lookup(&Self::semantic(request, scope), now)
}
Self::Gcs(cache) => cache.lookup(&Self::exact(request), now),
Self::Disk(cache) => cache.lookup(&Self::exact(request), now),
Self::AzureBlob(cache) => cache.lookup(&Self::exact(request), now),
}
}
pub fn store(
&self,
request: &ResponseCacheRequest,
request: &NativeRequest,
response: Value,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.store(request, response, now),
Self::Redis { cache, .. } => cache.store(request, response, now),
Self::S3(cache) => cache.store(request, response, now),
Self::Gcs(cache) => cache.store(request, response, now),
Self::Disk(cache) => cache.store(request, response, now),
Self::AzureBlob(cache) => cache.store(request, response, now),
Self::Memory(cache) => cache.store(&Self::exact(request), response, now),
Self::Redis { cache, .. } => cache.store(&Self::exact(request), response, now),
Self::S3(cache) => cache.store(&Self::exact(request), response, now),
Self::ValkeySemantic { cache, scope, .. } => {
cache.store(&Self::semantic(request, scope), response, now)
}
Self::Gcs(cache) => cache.store(&Self::exact(request), response, now),
Self::Disk(cache) => cache.store(&Self::exact(request), response, now),
Self::AzureBlob(cache) => cache.store(&Self::exact(request), response, now),
}
}
pub fn lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[NativeRequest],
now: Duration,
) -> Result<PartialHits, Error> {
match self {
Self::Memory(cache) => cache.lookup_batch(requests, now),
Self::Redis { cache, .. } => cache.lookup_batch(requests, now),
Self::S3(cache) => cache.lookup_batch(requests, now),
Self::Gcs(cache) => cache.lookup_batch(requests, now),
Self::Disk(cache) => cache.lookup_batch(requests, now),
Self::AzureBlob(cache) => cache.lookup_batch(requests, now),
Self::Memory(cache) => {
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
cache.lookup_batch(&requests, now)
}
Self::Redis { cache, .. } => {
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
cache.lookup_batch(&requests, now)
}
Self::S3(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
Self::Gcs(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
Self::Disk(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
Self::AzureBlob(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
}
}
pub async fn async_lookup(
&self,
request: &ResponseCacheRequest,
request: &NativeRequest,
now: Duration,
) -> Result<Option<Value>, Error> {
match self {
Self::Memory(cache) => cache.async_lookup(request, now).await,
Self::Redis { cache, .. } => cache.async_lookup(request, now).await,
Self::S3(cache) => cache.async_lookup(request, now).await,
Self::Gcs(cache) => cache.async_lookup(request, now).await,
Self::Disk(cache) => cache.async_lookup(request, now).await,
Self::AzureBlob(cache) => cache.async_lookup(request, now).await,
Self::Memory(cache) => cache.async_lookup(&Self::exact(request), now).await,
Self::Redis { cache, .. } => cache.async_lookup(&Self::exact(request), now).await,
Self::S3(cache) => cache.async_lookup(&Self::exact(request), now).await,
Self::ValkeySemantic { cache, scope, .. } => {
cache
.async_lookup(&Self::semantic(request, scope), now)
.await
}
Self::Gcs(cache) => cache.async_lookup(&Self::exact(request), now).await,
Self::Disk(cache) => cache.async_lookup(&Self::exact(request), now).await,
Self::AzureBlob(cache) => cache.async_lookup(&Self::exact(request), now).await,
}
}
pub(super) fn async_lookup_py<'py>(
&self,
py: Python<'py>,
request: NativeRequest,
) -> PyResult<Bound<'py, PyAny>> {
match self {
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => {
let service = self.clone();
litellm_host_python::run_async(
py,
async move { service.async_lookup(&request, super::request::now()).await },
super::cache_error,
)
}
Self::ValkeySemantic {
cache,
embedder,
scope,
} => drive_semantic(
py,
SemanticEmbedExecution::lookup(
Arc::clone(cache.backend_arc()),
embedder.clone(),
Self::semantic(&request, scope),
),
),
}
}
pub async fn async_store(
&self,
request: &ResponseCacheRequest,
request: &NativeRequest,
response: Value,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.async_store(request, response, now).await,
Self::Memory(cache) => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
Self::Redis {
cache,
buffer: None,
} => cache.async_store(request, response, now).await,
} => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
Self::Redis {
cache,
buffer: Some(buffer),
} => buffer.async_store(cache, request, response, now).await,
Self::S3(cache) => cache.async_store(request, response, now).await,
Self::Gcs(cache) => cache.async_store(request, response, now).await,
Self::Disk(cache) => cache.async_store(request, response, now).await,
Self::AzureBlob(cache) => cache.async_store(request, response, now).await,
} => {
buffer
.async_store(cache, &Self::exact(request), response, now)
.await
}
Self::S3(cache) => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
Self::ValkeySemantic { cache, scope, .. } => {
cache
.async_store(&Self::semantic(request, scope), response, now)
.await
}
Self::Gcs(cache) => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
Self::Disk(cache) => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
Self::AzureBlob(cache) => {
cache
.async_store(&Self::exact(request), response, now)
.await
}
}
}
pub(super) fn async_store_py<'py>(
&self,
py: Python<'py>,
request: NativeRequest,
response: Value,
) -> PyResult<Bound<'py, PyAny>> {
match self {
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => {
let service = self.clone();
litellm_host_python::run_async(
py,
async move {
service
.async_store(&request, response, super::request::now())
.await
},
super::cache_error,
)
}
Self::ValkeySemantic {
cache,
embedder,
scope,
} => drive_semantic(
py,
SemanticEmbedExecution::store(
Arc::clone(cache.backend_arc()),
embedder.clone(),
Self::semantic(&request, scope),
response,
),
),
}
}
pub async fn async_lookup_batch(
&self,
requests: &[ResponseCacheRequest],
requests: &[NativeRequest],
now: Duration,
) -> Result<PartialHits, Error> {
match self {
Self::Memory(cache) => cache.async_lookup_batch(requests, now).await,
Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await,
Self::S3(cache) => cache.async_lookup_batch(requests, now).await,
Self::Gcs(cache) => cache.async_lookup_batch(requests, now).await,
Self::Disk(cache) => cache.async_lookup_batch(requests, now).await,
Self::AzureBlob(cache) => cache.async_lookup_batch(requests, now).await,
Self::Memory(cache) => {
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
cache.async_lookup_batch(&requests, now).await
}
Self::Redis { cache, .. } => {
let requests = requests.iter().map(Self::exact).collect::<Vec<_>>();
cache.async_lookup_batch(&requests, now).await
}
Self::S3(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
Self::Gcs(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
Self::Disk(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
Self::AzureBlob(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
}
}
pub async fn async_store_batch(
&self,
entries: Vec<(ResponseCacheRequest, Value)>,
entries: Vec<(NativeRequest, Value)>,
now: Duration,
) -> Result<(), Error> {
match self {
Self::Memory(cache) => cache.async_store_batch(entries, now).await,
Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await,
Self::S3(cache) => cache.async_store_batch(entries, now).await,
Self::Gcs(cache) => cache.async_store_batch(entries, now).await,
Self::Disk(cache) => cache.async_store_batch(entries, now).await,
Self::AzureBlob(cache) => cache.async_store_batch(entries, now).await,
Self::Memory(cache) => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::Redis { cache, .. } => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::S3(cache) => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::ValkeySemantic { cache, scope, .. } => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::semantic(&request, scope), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::Gcs(cache) => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::Disk(cache) => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
Self::AzureBlob(cache) => {
let entries = entries
.into_iter()
.map(|(request, value)| (Self::exact(&request), value))
.collect();
cache.async_store_batch(entries, now).await
}
}
}
pub(super) fn async_store_batch_py<'py>(
&self,
py: Python<'py>,
entries: Vec<(NativeRequest, Value)>,
) -> PyResult<Bound<'py, PyAny>> {
match self {
Self::Memory(_)
| Self::Redis { .. }
| Self::S3(_)
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => {
let service = self.clone();
litellm_host_python::run_async(
py,
async move {
service
.async_store_batch(entries, super::request::now())
.await
},
super::cache_error,
)
}
Self::ValkeySemantic {
cache,
embedder,
scope,
} => {
let (requests, responses): (Vec<_>, Vec<_>) = entries
.into_iter()
.map(|(request, response)| (Self::semantic(&request, scope), response))
.unzip();
drive_semantic(
py,
SemanticEmbedExecution::store_batch(
Arc::clone(cache.backend_arc()),
embedder.clone(),
requests,
responses,
),
)
}
}
}
@ -338,6 +731,7 @@ impl NativeResponseCache {
cache.async_flush().await
}
Self::S3(cache) => cache.async_flush().await,
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
Self::Gcs(cache) => cache.async_flush().await,
Self::Disk(cache) => cache.async_flush().await,
Self::AzureBlob(cache) => cache.async_flush().await,
@ -349,6 +743,7 @@ impl NativeResponseCache {
Self::Memory(cache) => cache.test_connection().await,
Self::Redis { cache, .. } => cache.test_connection().await,
Self::S3(cache) => cache.test_connection().await,
Self::ValkeySemantic { cache, .. } => cache.test_connection().await,
Self::Gcs(cache) => cache.test_connection().await,
Self::Disk(cache) => cache.test_connection().await,
Self::AzureBlob(cache) => cache.test_connection().await,
@ -362,3 +757,79 @@ impl NativeResponseCache {
}
}
}
#[cfg(test)]
mod tests {
use litellm_cache_response::{CacheControls, CacheKeyInput, cache_key};
use serde_json::json;
use sha2::{Digest, Sha256};
use super::*;
fn native_request(key: CacheKeyInput, metadata: Value) -> NativeRequest {
NativeRequest {
key,
controls: CacheControls::default(),
ttl: None,
max_age: None,
messages: Some(json!([{"role": "user", "content": "prompt"}])),
input: None,
metadata: Some(metadata),
litellm_metadata: None,
litellm_params: None,
}
}
#[test]
fn semantic_key_matches_python_scope_material() {
let key = CacheKeyInput {
fields: vec![
CacheKeyField {
name: "model".to_owned(),
value: Some("gpt-4.1".to_owned()),
api_parameter: true,
internal_parameter: false,
},
CacheKeyField {
name: "messages".to_owned(),
value: Some("prompt".to_owned()),
api_parameter: true,
internal_parameter: false,
},
],
..Default::default()
};
let request = native_request(
key,
json!({"user_api_key": "k1", "user_api_key_team_id": null}),
);
let expected = format!("{:x}", Sha256::digest(b"model: gpt-4.1user_api_key: k1"));
assert_eq!(cache_key(&semantic_key(&request, "key")), expected);
let end_user_request = native_request(
request.key.clone(),
json!({"user_api_key": "k1", "user_api_key_end_user_id": "u1"}),
);
let expected = format!(
"{:x}",
Sha256::digest(b"model: gpt-4.1user_api_key: k1user_api_key_end_user_id: u1")
);
assert_eq!(
cache_key(&semantic_key(&end_user_request, "end_user")),
expected
);
let preset_request = native_request(
CacheKeyInput {
preset: Some("preset-key".to_owned()),
..Default::default()
},
json!({"user_api_key": "k1"}),
);
assert_eq!(
semantic_key(&preset_request, "end_user").preset.as_deref(),
Some("preset-key")
);
assert!(semantic_key(&preset_request, "end_user").fields.is_empty());
}
}

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

@ -1,9 +1,11 @@
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_cache::ExactCacheContext;
use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest};
use litellm_host_python::from_py;
use pyo3::{exceptions::PyValueError, prelude::*};
use serde::Deserialize;
use serde_json::Value;
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
@ -12,24 +14,48 @@ struct RequestInput {
controls: Option<CacheControls>,
ttl_seconds: Option<f64>,
max_age_seconds: Option<f64>,
messages: Option<Value>,
input: Option<Value>,
metadata: Option<Value>,
litellm_metadata: Option<Value>,
litellm_params: Option<Value>,
}
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<ResponseCacheRequest> {
pub(super) struct NativeRequest {
pub(super) key: CacheKeyInput,
pub(super) controls: CacheControls,
pub(super) ttl: Option<Duration>,
pub(super) max_age: Option<Duration>,
pub(super) messages: Option<Value>,
pub(super) input: Option<Value>,
pub(super) metadata: Option<Value>,
pub(super) litellm_metadata: Option<Value>,
pub(super) litellm_params: Option<Value>,
}
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<NativeRequest> {
let input: RequestInput = from_py(value)?;
request_input(input)
}
fn request_input(input: RequestInput) -> PyResult<ResponseCacheRequest> {
let mut request = ResponseCacheRequest::new(input.key);
if let Some(controls) = input.controls {
request.controls = controls;
}
request.context.ttl = input.ttl_seconds.map(duration).transpose()?;
request.max_age = input.max_age_seconds.map(duration).transpose()?;
Ok(request)
fn request_input(input: RequestInput) -> PyResult<NativeRequest> {
let controls = input.controls.unwrap_or_else(|| {
ResponseCacheRequest::<ExactCacheContext>::new(input.key.clone()).controls
});
Ok(NativeRequest {
key: input.key,
controls,
ttl: input.ttl_seconds.map(duration).transpose()?,
max_age: input.max_age_seconds.map(duration).transpose()?,
messages: input.messages,
input: input.input,
metadata: input.metadata,
litellm_metadata: input.litellm_metadata,
litellm_params: input.litellm_params,
})
}
pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult<Vec<ResponseCacheRequest>> {
pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult<Vec<NativeRequest>> {
from_py::<Vec<RequestInput>>(value)?
.into_iter()
.map(request_input)

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

@ -0,0 +1,249 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::SemanticCacheContext;
use litellm_cache_response::{ResponseCache, ResponseCacheCodec, ResponseCacheRequest};
use litellm_cache_valkey_semantic::{PreparedEmbedding, ValkeySemanticCache, prompt_from_context};
use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use serde_json::Value;
use super::{cache_error, embedder::PythonEmbedder};
pub(super) enum Op {
Lookup,
Store(Value),
StoreBatch(Vec<Value>),
}
#[derive(Clone, Copy)]
enum State {
Start,
AwaitingEmbedding,
AwaitingStorage,
Done,
}
pub(super) struct SemanticEmbedExecution {
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
op: Op,
now: Option<Duration>,
prepared: Vec<Option<Vec<f32>>>,
index: usize,
state: State,
}
impl SemanticEmbedExecution {
pub(super) fn lookup(
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
request: ResponseCacheRequest<SemanticCacheContext>,
) -> Self {
Self {
backend,
embedder,
requests: vec![request],
op: Op::Lookup,
now: None,
prepared: vec![None],
index: 0,
state: State::Start,
}
}
pub(super) fn store(
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
request: ResponseCacheRequest<SemanticCacheContext>,
response: Value,
) -> Self {
Self {
backend,
embedder,
requests: vec![request],
op: Op::Store(response),
now: None,
prepared: vec![None],
index: 0,
state: State::Start,
}
}
pub(super) fn store_batch(
backend: Arc<ValkeySemanticCache<PythonEmbedder, ResponseCacheCodec>>,
embedder: PythonEmbedder,
requests: Vec<ResponseCacheRequest<SemanticCacheContext>>,
responses: Vec<Value>,
) -> Self {
Self {
backend,
embedder,
prepared: vec![None; requests.len()],
requests,
op: Op::StoreBatch(responses),
now: None,
index: 0,
state: State::Start,
}
}
fn start(&mut self, py: Python<'_>) -> PyResult<ExecutionStep> {
if self.now.is_none() {
self.now = Some(super::request::now());
}
while self.index < self.requests.len() {
let request = &self.requests[self.index];
let enabled = match &self.op {
Op::Lookup => request.controls.reads(),
Op::Store(_) | Op::StoreBatch(_) => request.controls.writes(),
};
if !enabled {
self.index += 1;
continue;
}
let Some(prompt) = prompt_from_context(&request.context) else {
self.index += 1;
continue;
};
let metadata = request.context.metadata.clone();
let awaitable = self
.embedder
.async_embed_awaitable(py, &prompt, &metadata)?;
self.state = State::AwaitingEmbedding;
return Ok(ExecutionStep::Await(awaitable.unbind()));
}
self.state = State::AwaitingStorage;
self.storage_step(py)
}
fn storage_step(&self, py: Python<'_>) -> PyResult<ExecutionStep> {
let requests = self.requests.clone();
let prepared = self.prepared.clone();
let backend = Arc::clone(&self.backend);
let now = self
.now
.ok_or_else(|| PyRuntimeError::new_err("semantic cache timestamp is unavailable"))?;
let awaitable = match &self.op {
Op::Lookup => {
let Some(request) = requests.into_iter().next() else {
return Err(PyRuntimeError::new_err(
"semantic lookup requires one request",
));
};
match prepared.into_iter().next().flatten() {
Some(values) => {
let backend = backend.with_embedder(PreparedEmbedding(values));
let cache = Arc::new(ResponseCache::new(Arc::new(backend)));
run_async(
py,
async move { cache.async_lookup(&request, now).await },
cache_error,
)?
}
None => {
let cache = Arc::new(ResponseCache::new(backend));
run_async(
py,
async move { cache.async_lookup(&request, now).await },
cache_error,
)?
}
}
}
Op::Store(response) => {
let Some(request) = requests.into_iter().next() else {
return Err(PyRuntimeError::new_err(
"semantic store requires one request",
));
};
let response = response.clone();
match prepared.into_iter().next().flatten() {
Some(values) => {
let backend = backend.with_embedder(PreparedEmbedding(values));
let cache = Arc::new(ResponseCache::new(Arc::new(backend)));
run_async(
py,
async move { cache.async_store(&request, response, now).await },
cache_error,
)?
}
None => {
let cache = Arc::new(ResponseCache::new(backend));
run_async(
py,
async move { cache.async_store(&request, response, now).await },
cache_error,
)?
}
}
}
Op::StoreBatch(responses) => {
let responses = responses.clone();
run_async(
py,
async move {
for ((request, response), prepared) in
requests.into_iter().zip(responses).zip(prepared)
{
let Some(values) = prepared else {
continue;
};
let backend = backend.with_embedder(PreparedEmbedding(values));
let cache = ResponseCache::new(Arc::new(backend));
cache.async_store(&request, response, now).await?;
}
Ok(())
},
cache_error,
)?
}
};
Ok(ExecutionStep::Await(awaitable.unbind()))
}
fn resume_py(
&mut self,
py: Python<'_>,
result: Option<PyResult<Py<PyAny>>>,
) -> PyResult<ExecutionStep> {
match (self.state, result) {
(State::Start, None) => self.start(py),
(State::AwaitingEmbedding, Some(Ok(value))) => {
let values = value.bind(py).extract::<Vec<f64>>()?;
self.prepared[self.index] =
Some(values.into_iter().map(|value| value as f32).collect());
self.index += 1;
self.start(py)
}
(State::AwaitingStorage, Some(Ok(value))) => {
self.state = State::Done;
Ok(ExecutionStep::Return(value))
}
(_, Some(Err(error))) => Err(error),
_ => Err(PyRuntimeError::new_err(
"invalid semantic cache execution state",
)),
}
}
}
impl ExecutionBody for SemanticEmbedExecution {
fn resume(&mut self, result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
Python::attach(|py| self.resume_py(py, result))
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
self.embedder.traverse(visit)
}
}
pub(super) fn drive_semantic<'py>(
py: Python<'py>,
body: SemanticEmbedExecution,
) -> PyResult<Bound<'py, PyAny>> {
let execution = Py::new(py, Execution::new(body))?;
py.import("litellm.rust_bridge.lifecycle")?
.getattr("drive")?
.call1((execution,))
}

4
litellm/litellm_core_utils/logging_worker.py
View file

@ -484,9 +484,13 @@ class LoggingWorker:
so it correctly handles items that have been dequeued but whose
callback hasn't finished yet — ``queue.empty()`` would return True in
that window and cause us to skip the wait.
``start()`` runs first so a queue left behind by a previous event loop
is carried onto this one and drained here instead of joined forever.
"""
if self._queue is None:
return
self.start()
await self._queue.join()
async def clear_queue(self):

66
litellm/model_prices_and_context_window_backup.json
View file

@ -76879,6 +76879,26 @@
"supports_vision": true,
"supports_web_search": true
},
"moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"global.moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
@ -76978,5 +76998,51 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"xiaomi_mimo/mimo-v2.6-pro": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 4.35e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 8.7e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
},
"xiaomi_mimo/mimo-v2.6-flash": {
"cache_read_input_token_cost": 2.8e-09,
"input_cost_per_token": 1.4e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2.8e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
}
}

66
model_prices_and_context_window.json
View file

@ -76879,6 +76879,26 @@
"supports_vision": true,
"supports_web_search": true
},
"moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
"input_cost_per_token": 3e-06,
"litellm_provider": "bedrock_converse",
"max_input_tokens": 1000000,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 1.5e-05,
"source": "https://aws.amazon.com/bedrock/pricing/",
"supports_audio_input": false,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_vision": true
},
"global.moonshotai.kimi-k3": {
"cache_creation_input_token_cost": 3.75e-06,
"cache_read_input_token_cost": 3e-07,
@ -76978,5 +76998,51 @@
"supports_tool_choice": true,
"supports_vision": true,
"supports_web_search": false
},
"xiaomi_mimo/mimo-v2.6-pro": {
"cache_read_input_token_cost": 3.6e-09,
"input_cost_per_token": 4.35e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 8.7e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
},
"xiaomi_mimo/mimo-v2.6-flash": {
"cache_read_input_token_cost": 2.8e-09,
"input_cost_per_token": 1.4e-07,
"litellm_provider": "xiaomi_mimo",
"max_input_tokens": 1048576,
"max_output_tokens": 131072,
"max_tokens": 131072,
"mode": "chat",
"output_cost_per_token": 2.8e-07,
"source": "https://platform.xiaomimimo.com/static/docs/price/pay-as-you-go.md",
"supported_endpoints": [
"/v1/chat/completions"
],
"supports_audio_input": true,
"supports_function_calling": true,
"supports_prompt_caching": true,
"supports_reasoning": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_tool_choice": true,
"supports_video_input": true,
"supports_vision": true
}
}

3
tests/e2e/coverage_registry/llm_conversational.yaml
View file

@ -89,6 +89,9 @@
- {id: llm.chat_completions.together_ai.multi_turn.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: multi_turn, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together tool result round trip"}
- {id: llm.chat_completions.together_ai.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together cost header and spend row match the registry price"}
- {id: llm.chat_completions.together_ai.thinking.nonstream.effort_none_disables, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: thinking, streaming: nonstream, assertions: [effort_none_disables], source: "llm_translation/test_together_ai_e2e.py", rationale: "reasoning_effort=none maps to Together's reasoning disable toggle on hybrid models"}
- {id: llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: basic, streaming: nonstream, assertions: [cost_logged], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "Native MiMo v2.6 rows price the cost header and spend row from the cost map"}
- {id: llm.chat_completions.xiaomi_mimo.thinking.stream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: thinking, streaming: stream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo reasoning deltas stream as reasoning_content"}
- {id: llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: xiaomi_mimo, capability: tool_use, streaming: nonstream, assertions: [works], source: "llm_translation/test_xiaomi_mimo_e2e.py", rationale: "MiMo tool calls are not dropped"}
- {id: llm.chat_completions.together_ai.structured_output.nonstream.works, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: structured_output, streaming: nonstream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "response_format json_schema reaches Together and constrains the reply"}
- {id: llm.chat_completions.together_ai.prompt_cache_5m.nonstream.cost_logged, module: llm, tier: P1, subject_endpoint: chat_completions, route: together_ai, capability: prompt_cache_5m, streaming: nonstream, assertions: [cache_hit, cost_logged], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together prefix-cache reads bill at cache_read_input_token_cost, not full input price"}
- {id: llm.messages.together_ai.basic.stream.works, module: llm, tier: P1, subject_endpoint: messages, route: together_ai, capability: basic, streaming: stream, assertions: [works], source: "llm_translation/test_together_ai_e2e.py", rationale: "Together over /v1/messages streaming"}

1
tests/e2e/coverage_registry/schema.py
View file

@ -58,6 +58,7 @@ LlmRoute = Literal[
"openai",
"together_ai",
"vertex",
"xiaomi_mimo",
]
LlmCapability = Literal[

214
tests/e2e/llm_translation/test_xiaomi_mimo_e2e.py Normal file
View file

@ -0,0 +1,214 @@
"""Live e2e: Xiaomi MiMo v2.6 through the gateway on /chat/completions.
Both native ``xiaomi_mimo/`` v2.6 rows (pro and flash) are registered via
``/model/new`` and driven against Xiaomi's own endpoint. What the gateway owes
us is that the reasoning chain surfaces as ``reasoning_content``, tool calls
survive translation, and the cost header plus spend row follow the proxy's own
cost-map price for the row (read back from ``/model/info``, never pinned here).
Requires XIAOMI_MIMO_API_KEY on the proxy; no skip gate.
"""
from __future__ import annotations
from typing import Final
import pytest
from e2e_config import unique_marker
from e2e_http import StreamingResponse, require_successful_call, unwrap
from lifecycle import ResourceManager
from models import (
ChatBody,
ChatMessage,
ChatResponse,
ChatTool,
ChatToolFunction,
CostMapEntry,
LiteLLMParamsBody,
OutMessage,
SpendLogRow,
)
from passthrough_client import PassthroughClient
from pydantic import BaseModel
pytestmark = pytest.mark.e2e
BACKENDS: Final = ("xiaomi_mimo/mimo-v2.6-pro", "xiaomi_mimo/mimo-v2.6-flash")
ARITHMETIC_PROMPT = "What is 17 + 26? Answer with just the number."
WEATHER_PROMPT = "What is the weather in Paris? Use the tool."
COUNTING_PROMPT = "Count from 1 to 50, one number per line."
WEATHER_TOOL = ChatTool(
function=ChatToolFunction(
name="get_weather",
description="Get the current weather for a location.",
parameters={
"type": "object",
"properties": {"location": {"type": "string"}},
"required": ["location"],
},
)
)
class _WeatherArgs(BaseModel):
location: str
class _StreamDelta(BaseModel):
content: str | None = None
reasoning_content: str | None = None
class _StreamChoice(BaseModel):
delta: _StreamDelta | None = None
class _StreamChunk(BaseModel):
choices: list[_StreamChoice] = []
def _approx_equal(actual: float, expected: float) -> bool:
return abs(actual - expected) <= max(1e-9, abs(expected) * 1e-2)
@pytest.fixture(scope="module")
def registry(client: PassthroughClient) -> dict[str, CostMapEntry]:
return client.proxy.model_cost_map()
def _register(client: PassthroughClient, resources: ResourceManager, backend: str) -> tuple[str, str]:
model = f"e2e-xiaomi-{unique_marker()}"
model_id = client.proxy.create_model(
model, LiteLLMParamsBody(model=backend, api_key="os.environ/XIAOMI_MIMO_API_KEY")
)
resources.defer(lambda: client.proxy.delete_model(model_id))
return model, resources.key()
def _message(response: ChatResponse) -> OutMessage:
assert response.choices, f"Xiaomi returned no choices: {response}"
message = response.choices[0].message
assert message is not None, f"Xiaomi choice has no message: {response}"
return message
def _deltas(result: StreamingResponse) -> list[_StreamDelta]:
require_successful_call(result)
assert result.is_streaming, f"response was not streamed: {result.headers}"
assert not result.stream_error, f"stream errored: {result.stream_error}"
assert result.stream_done, f"stream never reached [DONE]: {result.stream_events[-3:]}"
return [
choice.delta
for event in result.stream_events
for choice in _StreamChunk.model_validate_json(event).choices
if choice.delta is not None
]
@pytest.mark.parametrize("backend", BACKENDS)
class TestXiaomiMimoChatCompletions:
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.basic.nonstream.cost_logged")
def test_cost_header_and_spend_row_match_the_registry_price(
self,
client: PassthroughClient,
resources: ResourceManager,
registry: dict[str, CostMapEntry],
backend: str,
) -> None:
price = registry.get(backend)
assert price is not None, f"{backend} has no row in the proxy's cost map, so native calls would bill $0"
assert price.litellm_provider == "xiaomi_mimo", f"{backend} is filed under the wrong provider: {price}"
assert price.input_cost_per_token and price.output_cost_per_token, f"{backend} carries no price: {price}"
model, key = _register(client, resources, backend)
result = client.proxy.transport.send(
"/chat/completions",
headers=client.proxy.transport.bearer(key),
json=ChatBody(
model=model,
messages=[ChatMessage(role="user", content=f"{ARITHMETIC_PROMPT} {unique_marker()}")],
max_tokens=1024,
),
)
require_successful_call(result)
response = ChatResponse.model_validate_json(result.body)
message = _message(response)
assert message.content and "43" in message.content, f"answer lost: {message}"
assert message.reasoning_content, f"{backend} reasons, but no reasoning_content came back: {message}"
usage = response.usage
assert usage is not None and usage.prompt_tokens and usage.completion_tokens, (
f"response carries no usage, so the cost cannot be real: {result.body[:300]}"
)
header_cost = result.response_cost
assert header_cost is not None and header_cost > 0, (
f"x-litellm-response-cost header missing or non-positive: {result.headers}"
)
cached = (usage.prompt_tokens_details.cached_tokens or 0) if usage.prompt_tokens_details else 0
expected = (
(usage.prompt_tokens - cached) * price.input_cost_per_token
+ cached * (price.cache_read_input_token_cost or 0.0)
+ usage.completion_tokens * price.output_cost_per_token
)
assert _approx_equal(header_cost, expected), (
f"header cost {header_cost} disagrees with the registry price for {backend} at {usage}: expected {expected}"
)
def _priced(rows: list[SpendLogRow]) -> bool:
return any(row.spend is not None and row.spend > 0 for row in rows)
rows = client.proxy.poll_logs_for_key(key, predicate=_priced)
priced = [row for row in rows if row.spend is not None and row.spend > 0]
assert priced, f"no priced spend row landed for key {key}; got {rows}"
row = priced[0]
assert row.custom_llm_provider == "xiaomi_mimo", f"spend row misattributed: {row}"
assert row.spend is not None and _approx_equal(row.spend, header_cost), (
f"logged spend {row.spend} disagrees with the x-litellm-response-cost header {header_cost}"
)
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.thinking.stream.works")
def test_reasoning_and_answer_stream_as_deltas(
self, client: PassthroughClient, resources: ResourceManager, backend: str
) -> None:
model, key = _register(client, resources, backend)
deltas = _deltas(
client.proxy.chat_stream(
key,
ChatBody(
model=model,
messages=[ChatMessage(role="user", content=COUNTING_PROMPT)],
max_tokens=2048,
stream=True,
),
)
)
reasoning = "".join(delta.reasoning_content or "" for delta in deltas)
content = "".join(delta.content or "" for delta in deltas)
assert reasoning, f"stream carried no reasoning_content deltas: {deltas[:5]}"
assert "50" in content, f"streamed answer lost: {content[:300]!r}"
@pytest.mark.covers("llm.chat_completions.xiaomi_mimo.tool_use.nonstream.works")
def test_tool_call_is_returned(self, client: PassthroughClient, resources: ResourceManager, backend: str) -> None:
model, key = _register(client, resources, backend)
message = _message(
unwrap(
client.proxy.chat(
key,
ChatBody(
model=model,
messages=[ChatMessage(role="user", content=WEATHER_PROMPT)],
tools=[WEATHER_TOOL],
max_tokens=1024,
),
)
)
)
assert message.tool_calls, f"{backend} dropped the tool call: {message}"
call = message.tool_calls[0]
assert call.id, f"tool call carries no id, so a tool result cannot answer it: {call}"
assert call.function.name == "get_weather", f"wrong tool called: {call}"
assert call.function.arguments, f"tool call carries no arguments: {call}"
args = _WeatherArgs.model_validate_json(call.function.arguments)
assert "paris" in args.location.lower(), f"tool arguments lost the location: {args}"

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

@ -180,6 +180,53 @@ class TestLoggingWorker:
assert sorted(fired) == ["first", "second"]
@pytest.mark.parametrize("stranded", ["still_queued", "dequeued_never_started"])
def test_flush_on_new_loop_drains_tasks_stranded_on_previous_loop(self, stranded):
"""
Regression: ``flush()`` from a new event loop used to ``join()`` the queue bound to the
previous loop, whose unfinished counter nothing on the new loop ever decrements. The first
such flush hung until pytest-timeout killed it and every later one raised
``RuntimeError: ... is bound to a different event loop`` from the queue's Event.
"""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)
callback = AsyncMock()
async def enqueue_on_first_loop():
if stranded == "still_queued":
worker._ensure_queue()
worker.enqueue(callback())
return
worker.ensure_initialized_and_enqueue(callback())
asyncio.run(enqueue_on_first_loop())
assert worker._queue is not None
expected_shape = (1, 0) if stranded == "still_queued" else (0, 1)
assert (worker._queue.qsize(), len(worker._unstarted_dequeued_tasks())) == expected_shape
assert callback.await_count == 0, "precondition: the callback never ran before the first loop closed"
async def flush_twice_on_second_loop():
await asyncio.wait_for(worker.flush(), timeout=5)
await asyncio.wait_for(worker.flush(), timeout=5)
asyncio.run(flush_twice_on_second_loop())
assert callback.await_count == 1
def test_flush_starts_a_worker_when_the_queue_has_none(self):
"""``flush()`` must drain a queue that exists on the current loop without a running worker."""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)
callback = AsyncMock()
async def enqueue_then_flush():
worker._ensure_queue()
worker.enqueue(callback())
assert worker._worker_task is None, "precondition: nothing is draining the queue yet"
await asyncio.wait_for(worker.flush(), timeout=3)
asyncio.run(enqueue_then_flush())
assert callback.await_count == 1
def test_flush_on_exit_swallows_cancellation_and_drains_remaining(self):
"""A callback raising CancelledError must not abort the atexit flush of later events."""
worker = LoggingWorker(timeout=1.0, max_queue_size=10)

599
tests/test_litellm_rust/test_valkey_semantic_cache_native.py Normal file
View file

@ -0,0 +1,599 @@
import asyncio
import contextvars
import hashlib
import os
import struct
import threading
import time
from collections.abc import Generator, Mapping
from types import SimpleNamespace
from typing import Final, cast
from uuid import uuid4
import pytest
import redis
from litellm.caching.caching import Cache
from litellm.caching.valkey_semantic_cache import ValkeySemanticCache
from litellm.rust_bridge import _native
from litellm.types.caching import LiteLLMCacheType
pytestmark: Final = pytest.mark.requires_rust_extension
embedding_context: Final = contextvars.ContextVar("embedding_context")
@pytest.fixture
def valkey_url() -> str:
url: Final = os.environ.get("LITELLM_TEST_VALKEY_URL")
if url is None:
pytest.skip("LITELLM_TEST_VALKEY_URL is not set")
return url
@pytest.fixture
def index_name(valkey_url: str) -> Generator[str]:
index: Final = f"litellm_test_{uuid4().hex}"
yield index
client: Final = redis.Redis.from_url(valkey_url)
try:
client.ft(index).dropindex(delete_documents=True)
except redis.ResponseError:
pass
finally:
client.close()
def _request(prompt: str = "semantic cache prompt") -> dict[str, object]:
return {
"key": {"preset": "key"},
"messages": [{"role": "user", "content": prompt}],
}
def _field_request(
prompt: str,
metadata: Mapping[str, object],
*,
namespace: str | None = None,
litellm_metadata: Mapping[str, object] | None = None,
litellm_params: Mapping[str, object] | None = None,
) -> dict[str, object]:
request: Final = {
"key": {
"fields": [
{
"name": "model",
"value": "gpt-4.1",
"api_parameter": True,
"internal_parameter": False,
},
{
"name": "messages",
"value": prompt,
"api_parameter": True,
"internal_parameter": False,
},
],
"namespace": namespace,
},
"messages": [{"role": "user", "content": prompt}],
"metadata": dict(metadata),
}
if litellm_metadata is not None:
request["litellm_metadata"] = dict(litellm_metadata)
if litellm_params is not None:
request["litellm_params"] = dict(litellm_params)
return request
def _facade(
url: str,
index_name: str,
embeddings: Mapping[str, list[float]],
*,
namespace: str | None = None,
) -> Cache:
facade: Final = Cache(
type=LiteLLMCacheType.VALKEY_SEMANTIC,
redis_url=url,
similarity_threshold=0.8,
valkey_semantic_cache_index_name=index_name,
namespace=namespace,
)
vectors: Final = embeddings
def embed(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]:
return vectors[prompt]
async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
return vectors[prompt]
facade.cache._get_embedding = embed
facade.cache._get_async_embedding = async_embedding
return facade
def _backend(
url: str,
index_name: str,
embeddings: Mapping[str, list[float]] | None = None,
) -> ValkeySemanticCache:
vectors: Final = embeddings or {"semantic cache prompt": [1.0, 0.0]}
backend: Final = ValkeySemanticCache(
redis_url=url,
similarity_threshold=0.8,
index_name=index_name,
)
def embed(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]:
return vectors[prompt]
async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
return vectors[prompt]
backend._get_embedding = embed
backend._get_async_embedding = async_embedding
return backend
def test_python_write_native_read(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
response: Final = {"answer": "python"}
backend.set_cache("key", response, messages=_request()["messages"])
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
assert binding.lookup(_request()) == response
def test_native_write_python_read(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
response: Final = {"answer": "native"}
binding.store({**_request(), "ttl_seconds": 2.0}, response)
cached: Final = cast(Mapping[str, object], backend.get_cache("key", messages=_request()["messages"]))
assert cached["response"] == response
async def test_async_lookup_and_store(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
request: Final = {**_request(), "ttl_seconds": 2.0}
await binding.async_store(request, {"answer": "async"})
assert await binding.async_lookup(request) == {"answer": "async"}
async def test_disabled_cache_controls_skip_async_embedding(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
calls: Final = []
async def fail_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
calls.append(prompt)
raise AssertionError("embedding must not run")
backend._get_async_embedding = fail_embedding
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
controls: Final = {
"supported_call_type": True,
"configured": True,
"native_backend": True,
"default_on": True,
"caching": True,
"no_cache": False,
"no_store": False,
"use_cache": True,
}
no_read_request: Final = {**_request(), "controls": {**controls, "no_cache": True}}
assert await binding.async_lookup(no_read_request) is None
no_write_request: Final = {**_request(), "controls": {**controls, "no_store": True}}
await binding.async_store(no_write_request, {"answer": "blocked"})
assert calls == []
client: Final = redis.Redis.from_url(valkey_url)
assert list(client.scan_iter(f"{index_name}:*")) == []
client.close()
async def test_async_embedding_runs_inline_in_caller_task(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
observed: dict[str, object] = {}
async def async_embedding(prompt: str, metadata: dict[str, object] | None = None) -> list[float]:
observed["context"] = embedding_context.get("missing")
observed["task"] = asyncio.current_task()
observed["thread"] = threading.get_ident()
embedding_context.set("embedder")
return [1.0, 0.0]
backend._get_async_embedding = async_embedding
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
request: Final = {**_request(), "ttl_seconds": 2.0}
caller_task: Final = asyncio.current_task()
caller_thread: Final = threading.get_ident()
token: Final = embedding_context.set("caller")
try:
await binding.async_store(request, {"answer": "inline"})
assert observed["context"] == "caller"
assert observed["task"] is caller_task
assert observed["thread"] == caller_thread
assert embedding_context.get() == "embedder"
assert await binding.async_lookup(request) == {"answer": "inline"}
finally:
embedding_context.reset(token)
def test_facade_activation_and_mutation_fallback(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = Cache(
type=LiteLLMCacheType.VALKEY_SEMANTIC,
redis_url=valkey_url,
similarity_threshold=0.8,
valkey_semantic_cache_index_name=index_name,
)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
handle._bind_facade(facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
assert resolver.resolve().kind == "native"
facade.cache.similarity_threshold = 0.7
assert resolver.resolve().kind == "python_callback"
def test_batch_lookup_is_unsupported(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
backend,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
with pytest.raises(NotImplementedError):
binding.lookup_batch([_request()])
def test_ttl_expiry(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
binding.store({**_request(), "ttl_seconds": 1.0}, {"answer": "expires"})
client: Final = redis.Redis.from_url(valkey_url)
documents: Final = list(client.scan_iter(f"{index_name}:*"))
assert len(documents) == 1
assert client.ttl(documents[0]) > 0
time.sleep(1.5)
assert binding.lookup(_request()) is None
def test_no_ttl_is_persistent_and_python_reads_native_value(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
response: Final = {"answer": "persistent"}
binding.store(_request(), response)
client: Final = redis.Redis.from_url(valkey_url)
documents: Final = list(client.scan_iter(f"{index_name}:*"))
assert len(documents) == 1
assert client.ttl(documents[0]) == -1
cached: Final = cast(Mapping[str, object], backend.get_cache("key", messages=_request()["messages"]))
assert cached["response"] == response
def test_below_threshold_misses_on_native_and_python(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(
valkey_url,
index_name,
{"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]},
)
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
binding.store(_request("prompt A"), {"answer": "A"})
assert binding.lookup(_request("prompt B")) is None
assert backend.get_cache("key", messages=_request("prompt B")["messages"]) is None
def test_malformed_entry_is_a_miss_on_native_and_python(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
client: Final = redis.Redis.from_url(valkey_url)
scope: Final = hashlib.sha256(b"key").hexdigest()
document: Final = f"{index_name}:{scope}:{uuid4().hex}"
client.hset(
document,
mapping={
"litellm_cache_key": scope,
"prompt": "semantic cache prompt",
"response": "not json",
"embedding": struct.pack("<2f", 1.0, 0.0),
},
)
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
assert binding.lookup(_request()) is None
assert backend.get_cache("key", messages=_request()["messages"]) is None
def test_mixed_content_parts_match_python_semantic_behavior(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
messages: Final = [{"role": "user", "content": ["raw", {"text": "hello"}]}]
backend.set_cache("key", {"answer": "mixed"}, messages=messages)
assert backend.get_cache("key", messages=messages) is None
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
request: Final = {**_request(), "messages": messages}
binding.store(request, {"answer": "mixed"})
assert binding.lookup(request) is None
client: Final = redis.Redis.from_url(valkey_url)
assert list(client.scan_iter(f"{index_name}:*")) == []
client.close()
async def test_async_store_batch_and_lookup(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(
valkey_url,
index_name,
{"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]},
)
sync_calls: Final = []
async_tasks: Final = []
def sync_embedding(prompt: str, metadata: Mapping[str, object] | None = None) -> list[float]:
sync_calls.append(prompt)
return {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}[prompt]
async def async_embedding(
prompt: str,
metadata: dict[str, object] | None = None,
) -> list[float]:
async_tasks.append(asyncio.current_task())
return {"prompt A": [1.0, 0.0], "prompt B": [0.0, 1.0]}[prompt]
backend._get_embedding = sync_embedding
backend._get_async_embedding = async_embedding
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
requests: Final = [_request("prompt A"), _request("prompt B")]
responses: Final = [{"answer": "A"}, {"answer": "B"}]
caller_task: Final = asyncio.current_task()
await binding.async_store_batch(requests, responses)
assert sync_calls == []
assert async_tasks
assert all(task is caller_task for task in async_tasks)
assert await binding.async_lookup(requests[0]) == responses[0]
assert await binding.async_lookup(requests[1]) == responses[1]
def test_subclass_backend_falls_back_to_python(
valkey_url: str,
index_name: str,
) -> None:
class Custom(ValkeySemanticCache):
pass
facade: Final = Cache(
type=LiteLLMCacheType.VALKEY_SEMANTIC,
redis_url=valkey_url,
similarity_threshold=0.8,
valkey_semantic_cache_index_name=index_name,
)
facade.cache = Custom(redis_url=valkey_url, similarity_threshold=0.8, index_name=index_name)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
assert resolver.resolve().kind == "python_callback"
def test_field_key_matches_python_semantic_scope(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]})
metadata: Final = {"user_api_key": "k1"}
expected: Final = facade.get_cache_key(
model="gpt-4.1",
messages=[{"role": "user", "content": "semantic cache prompt"}],
metadata=metadata,
)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
binding.store(_field_request("semantic cache prompt", metadata), {"answer": "scoped"})
client: Final = redis.Redis.from_url(valkey_url)
documents: Final = list(client.scan_iter(f"{index_name}:*"))
assert len(documents) == 1
document_parts: Final = documents[0].decode().split(":")
assert document_parts[1] == hashlib.sha256(expected.encode()).hexdigest()
client.close()
def test_field_key_reads_all_python_tenant_metadata_sources(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]})
params_metadata: Final = {"user_api_key_team_id": "team-from-params"}
expected: Final = facade.get_cache_key(
model="gpt-4.1",
messages=[{"role": "user", "content": "semantic cache prompt"}],
metadata={},
litellm_params={"metadata": params_metadata},
)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
binding.store(
_field_request(
"semantic cache prompt",
{},
litellm_params={"metadata": params_metadata},
),
{"answer": "params"},
)
client: Final = redis.Redis.from_url(valkey_url)
documents: Final = list(client.scan_iter(f"{index_name}:*"))
assert len(documents) == 1
document_parts: Final = documents[0].decode().split(":")
assert document_parts[1] == hashlib.sha256(expected.encode()).hexdigest()
client.close()
assert (
binding.lookup(
_field_request(
"semantic cache prompt",
{},
litellm_metadata={"user_api_key_team_id": "team-from-litellm"},
)
)
is None
)
def test_namespace_isolates_semantic_entries(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = _facade(
valkey_url,
index_name,
{"semantic cache prompt": [1.0, 0.0]},
namespace="team-a",
)
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
team_a: Final = _field_request("semantic cache prompt", {}, namespace="team-a")
team_b: Final = _field_request("semantic cache prompt", {}, namespace="team-b")
binding.store(team_a, {"answer": "team-a"})
assert binding.lookup(team_b) is None
assert binding.lookup(team_a) == {"answer": "team-a"}
cached: Final = cast(
Mapping[str, object],
facade.get_cache(
model="gpt-4.1",
messages=[{"role": "user", "content": "semantic cache prompt"}],
),
)
assert cached == {"answer": "team-a"}
def test_field_key_isolates_tenant_scope(
valkey_url: str,
index_name: str,
) -> None:
facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]})
handle: Final = _native._CacheTestHandle.valkey_semantic(
valkey_url,
0.8,
index_name,
facade.cache,
)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
binding.store(
_field_request("semantic cache prompt", {"user_api_key": "k1"}),
{"answer": "tenant one"},
)
assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k2"})) is None
assert binding.lookup(_field_request("semantic cache prompt", {"user_api_key": "k1"})) == {"answer": "tenant one"}
def test_tls_valkey_facade_falls_back_to_python(
index_name: str,
) -> None:
facade: Final = Cache(
type=LiteLLMCacheType.VALKEY_SEMANTIC,
redis_url="rediss://127.0.0.1:6390/0",
similarity_threshold=0.8,
valkey_semantic_cache_index_name=index_name,
)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
assert resolver.resolve().kind == "python_callback"
async def test_ping_maps_unsupported_native_operation_to_not_implemented(
valkey_url: str,
index_name: str,
) -> None:
backend: Final = _backend(valkey_url, index_name)
handle: Final = _native._CacheTestHandle.valkey_semantic(valkey_url, 0.8, index_name, backend)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
with pytest.raises(NotImplementedError):
await binding.ping()