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 131

Merge pull request #42319 from BerriAI/litellm_native_redis_semantic_cache

Browse source 
feat(cache): serve Redis Semantic caches natively in Rust
This commit is contained in:
yujonglee 2026-09-21 17:05:49 -07:00 • committed by GitHub
commit ef8103a5a8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 2944 additions and 75 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

16
litellm-rust/Cargo.lock generated
View file

@ -2741,6 +2741,21 @@ dependencies = [
"tokio",
]
[[package]]
name = "litellm-cache-redis-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-response",
"r2d2",
"redis",
"redis-test",
"serde_json",
"sha2 0.10.9",
"tokio",
]
[[package]]
name = "litellm-cache-response"
version = "0.1.0"
@ -2958,6 +2973,7 @@ dependencies = [
"litellm-cache-gcs",
"litellm-cache-memory",
"litellm-cache-redis",
"litellm-cache-redis-semantic",
"litellm-cache-response",
"litellm-cache-s3",
"litellm-cache-valkey-semantic",

1
litellm-rust/Cargo.toml
View file

@ -36,6 +36,7 @@ litellm-cache-redis = { path = "crates/cache-redis" }
litellm-cache-s3 = { path = "crates/cache-s3" }
litellm-cache-gcs = { path = "crates/cache-gcs" }
litellm-cache-disk = { path = "crates/cache-disk" }
litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" }
litellm-cache-response = { path = "crates/cache-response" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-token-counter-fast = { path = "crates/token-counter-fast" }

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

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

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

@ -0,0 +1,618 @@
use std::{
future::Future,
sync::{Arc, OnceLock},
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
SemanticCacheContext,
};
use litellm_cache_redis::{
RedisTopology,
connection::{ConnectionRef, Connections},
};
use litellm_cache_response::{CacheEntry, ResponseCacheCodec};
use serde_json::Value;
use sha2::{Digest, Sha256};
use crate::prompt::prompt_from_context;
const CACHE_KEY_FIELD: &str = "litellm_cache_key";
const VECTOR_FIELD: &str = "prompt_vector";
pub trait Embedder: Send + Sync + 'static {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error>;
fn async_embed(
&self,
prompt: &str,
metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
}
#[derive(Clone, Debug)]
pub struct RedisSemanticConfig {
pub index_name: String,
pub similarity_threshold: f32,
}
struct Inner {
index_name: String,
distance_threshold: f64,
resolved_index: OnceLock<String>,
codec: ResponseCacheCodec,
clock: fn() -> f64,
}
impl Inner {
fn new(config: RedisSemanticConfig) -> Self {
Self {
index_name: config.index_name,
distance_threshold: 1.0 - f64::from(config.similarity_threshold),
resolved_index: OnceLock::new(),
codec: ResponseCacheCodec,
clock: timestamp,
}
}
fn ensure_index(
&self,
connection: &mut ConnectionRef<'_>,
dims: usize,
) -> Result<String, Error> {
if let Some(name) = self.resolved_index.get() {
return Ok(name.clone());
}
let name = match index_compatible(connection, &self.index_name, dims)? {
Some(true) => self.index_name.clone(),
Some(false) => self.isolated_index(connection, dims)?,
None => match create_index(connection, &self.index_name, dims) {
Ok(()) => self.index_name.clone(),
Err(_) => match index_compatible(connection, &self.index_name, dims)? {
Some(true) => self.index_name.clone(),
Some(false) => self.isolated_index(connection, dims)?,
None => return Err(Error::Unavailable),
},
},
};
let _ = self.resolved_index.set(name.clone());
Ok(name)
}
fn isolated_index(
&self,
connection: &mut ConnectionRef<'_>,
dims: usize,
) -> Result<String, Error> {
let name = format!("{}_isolated", self.index_name);
match index_compatible(connection, &name, dims)? {
Some(true) => Ok(name),
Some(false) => {
redis::cmd("FT.DROPINDEX")
.arg(&name)
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
create_index(connection, &name, dims)?;
Ok(name)
}
None => {
create_index(connection, &name, dims)?;
Ok(name)
}
}
}
fn store(
&self,
connection: &mut ConnectionRef<'_>,
tag: &str,
value: &CacheEntry,
prompt: &str,
vector: &[f32],
ttl: Option<Duration>,
) -> Result<(), Error> {
let index = self.ensure_index(connection, vector.len())?;
let entry_id = entry_id(prompt, tag);
let hash_key = format!("{index}:{entry_id}");
let response = self.codec.encode(value)?;
redis::cmd("HSET")
.arg(&hash_key)
.arg("entry_id")
.arg(&entry_id)
.arg("prompt")
.arg(prompt)
.arg("response")
.arg(response)
.arg(VECTOR_FIELD)
.arg(vector_buffer(vector))
.arg("inserted_at")
.arg(format!("{}", (self.clock)()))
.arg("updated_at")
.arg(format!("{}", (self.clock)()))
.arg(CACHE_KEY_FIELD)
.arg(tag)
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
if let Some(ttl) = ttl {
redis::cmd("EXPIRE")
.arg(&hash_key)
.arg(ttl_seconds(ttl))
.query::<()>(connection)
.map_err(|_| Error::Unavailable)?;
}
Ok(())
}
fn lookup(
&self,
connection: &mut ConnectionRef<'_>,
tag: &str,
vector: &[f32],
) -> Result<Option<CacheEntry>, Error> {
let index = self.ensure_index(connection, vector.len())?;
let query = format!(
"(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]",
escape_tag(tag)
);
let result = redis::cmd("FT.SEARCH")
.arg(&index)
.arg(query)
.arg("RETURN")
.arg(8)
.arg("entry_id")
.arg("prompt")
.arg("response")
.arg("inserted_at")
.arg("updated_at")
.arg("metadata")
.arg(CACHE_KEY_FIELD)
.arg("vector_distance")
.arg("SORTBY")
.arg("vector_distance")
.arg("ASC")
.arg("DIALECT")
.arg(2)
.arg("LIMIT")
.arg(0)
.arg(1)
.arg("PARAMS")
.arg(2)
.arg("vector")
.arg(vector_buffer(vector))
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)?;
let Some(fields) = first_document(&result) else {
return Ok(None);
};
if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) {
return Ok(None);
}
if number_field(fields, "vector_distance")
.is_none_or(|distance| distance > self.distance_threshold)
{
return Ok(None);
}
let Some(response) = bytes_field(fields, "response") else {
return Ok(None);
};
self.codec.decode(&response).map(Some)
}
}
pub struct RedisSemanticCache<E: Embedder, C = redis::Connection> {
connections: Arc<Connections<C>>,
embedder: E,
inner: Arc<Inner>,
}
impl<E: Embedder> RedisSemanticCache<E> {
pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
inner: Arc::new(Inner::new(config)),
})
}
}
impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<E, C> {
pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self {
Self {
connections: Arc::new(Connections::fixed(connection)),
embedder,
inner: Arc::new(Inner::new(config)),
}
}
pub fn with_clock(self, clock: fn() -> f64) -> Self {
Self {
inner: Arc::new(Inner {
index_name: self.inner.index_name.clone(),
distance_threshold: self.inner.distance_threshold,
resolved_index: OnceLock::new(),
codec: self.inner.codec,
clock,
}),
..self
}
}
pub fn embedder(&self) -> &E {
&self.embedder
}
pub fn index_name(&self) -> &str {
&self.inner.index_name
}
pub fn similarity_threshold(&self) -> f32 {
(1.0 - self.inner.distance_threshold) as f32
}
fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str {
context.scope.as_deref().unwrap_or(key)
}
}
impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
for RedisSemanticCache<E, C>
{
type Value = CacheEntry;
type Context = SemanticCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &Self::Context,
) -> Result<(), Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(());
};
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context).to_string();
self.connections.execute(|connection| {
self.inner
.store(connection, &tag, &value, &prompt, &vector, context.ttl)
})
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(None);
};
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context).to_string();
self.connections
.execute(|connection| self.inner.lookup(connection, &tag, &vector))
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: Self::Context,
) -> Result<(), Error> {
let Some(prompt) = prompt_from_context(&context) else {
return Ok(());
};
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, &context).to_string();
let inner = Arc::clone(&self.inner);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
inner.store(connection, &tag, &value, &prompt, &vector, context.ttl)
})
.await
}
async fn async_get_cache(
&self,
key: &str,
context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(None);
};
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, context).to_string();
let inner = Arc::clone(&self.inner);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
inner.lookup(connection, &tag, &vector)
})
.await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
match Connections::run_blocking(Arc::clone(&self.connections), |connection| {
Ok(match redis::cmd("PING").query::<String>(connection) {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Redis cache connection test successful".into(),
error: None,
},
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Redis connection failed: {error}"),
error: Some(error.to_string()),
},
})
})
.await
{
Ok(result) => Ok(result),
Err(error) => Ok(CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Redis connection failed: {error}"),
error: Some(error.to_string()),
}),
}
}
}
fn timestamp() -> f64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|duration| duration.as_secs_f64())
.unwrap_or_default()
}
fn entry_id(prompt: &str, tag: &str) -> String {
let mut digest = Sha256::new();
digest.update(prompt.as_bytes());
digest.update(CACHE_KEY_FIELD.as_bytes());
digest.update(tag.as_bytes());
format!("{:x}", digest.finalize())
}
fn vector_buffer(vector: &[f32]) -> Vec<u8> {
vector
.iter()
.flat_map(|component| component.to_le_bytes())
.collect()
}
fn escape_tag(value: &str) -> String {
value
.chars()
.flat_map(|ch| {
if matches!(
ch,
',' | '.'
| '<'
| '>'
| '{'
| '}'
| '['
| ']'
| '\\'
| '"'
| '\''
| ':'
| ';'
| '!'
| '@'
| '#'
| '$'
| '%'
| '^'
| '&'
| '*'
| '('
| ')'
| '-'
| '+'
| '='
| '~'
| '|'
| '/'
| ' '
| '?'
) {
vec!['\\', ch]
} else {
vec![ch]
}
})
.collect()
}
fn create_index(connection: &mut ConnectionRef<'_>, name: &str, dims: usize) -> Result<(), Error> {
redis::cmd("FT.CREATE")
.arg(name)
.arg("ON")
.arg("HASH")
.arg("PREFIX")
.arg(1)
.arg(name)
.arg("SCORE")
.arg(1.0)
.arg("SCHEMA")
.arg("prompt")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("response")
.arg("TEXT")
.arg("WEIGHT")
.arg(1)
.arg("inserted_at")
.arg("NUMERIC")
.arg("updated_at")
.arg("NUMERIC")
.arg(VECTOR_FIELD)
.arg("VECTOR")
.arg("FLAT")
.arg(6)
.arg("TYPE")
.arg("FLOAT32")
.arg("DIM")
.arg(dims)
.arg("DISTANCE_METRIC")
.arg("COSINE")
.arg(CACHE_KEY_FIELD)
.arg("TAG")
.arg("SEPARATOR")
.arg(",")
.query::<()>(connection)
.map_err(|_| Error::Unavailable)
}
fn index_compatible(
connection: &mut ConnectionRef<'_>,
name: &str,
dims: usize,
) -> Result<Option<bool>, Error> {
let info = match redis::cmd("FT.INFO")
.arg(name)
.query::<redis::Value>(connection)
{
Ok(info) => info,
Err(error) if unknown_index(&error) => return Ok(None),
Err(_) => return Err(Error::Unavailable),
};
Ok(Some(schema_compatible(&info, dims)))
}
fn unknown_index(error: &redis::RedisError) -> bool {
let message = error.to_string().to_lowercase();
message.contains("unknown") && message.contains("index")
}
fn schema_compatible(info: &redis::Value, dims: usize) -> bool {
let redis::Value::Array(entries) = info else {
return false;
};
let attributes = entries
.as_chunks::<2>()
.0
.iter()
.find(|pair| string_value(&pair[0]).as_deref() == Some("attributes"))
.map(|pair| &pair[1]);
let Some(redis::Value::Array(attributes)) = attributes else {
return false;
};
let fields = attributes
.iter()
.map(|attribute| {
let redis::Value::Array(attribute) = attribute else {
return (None, None, None, None, None);
};
let mut name = None;
let mut field_type = None;
let mut dim = None;
let mut data_type = None;
let mut distance_metric = None;
for pair in attribute.as_chunks::<2>().0 {
match string_value(&pair[0]).as_deref() {
Some("identifier") => name = string_value(&pair[1]),
Some("type") => field_type = string_value(&pair[1]),
Some("dim") => dim = number_value(&pair[1]),
Some("data_type") => data_type = string_value(&pair[1]),
Some("distance_metric") => distance_metric = string_value(&pair[1]),
_ => {}
}
}
(name, field_type, dim, data_type, distance_metric)
})
.collect::<Vec<_>>();
let has_field = |name: &str, field_type: &str| {
fields
.iter()
.any(|(n, t, ..)| n.as_deref() == Some(name) && t.as_deref() == Some(field_type))
};
has_field("prompt", "TEXT")
&& has_field("response", "TEXT")
&& has_field("inserted_at", "NUMERIC")
&& has_field("updated_at", "NUMERIC")
&& has_field(CACHE_KEY_FIELD, "TAG")
&& fields.iter().any(|(n, t, d, data, metric)| {
n.as_deref() == Some(VECTOR_FIELD)
&& t.as_deref() == Some("VECTOR")
&& *d == Some(dims as f64)
&& data
.as_deref()
.is_some_and(|data| data.eq_ignore_ascii_case("float32"))
&& metric
.as_deref()
.is_some_and(|metric| metric.eq_ignore_ascii_case("cosine"))
})
}
fn string_value(value: &redis::Value) -> Option<String> {
match value {
redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
redis::Value::SimpleString(text) => Some(text.clone()),
redis::Value::VerbatimString { text, .. } => Some(text.clone()),
_ => None,
}
}
fn number_value(value: &redis::Value) -> Option<f64> {
match value {
redis::Value::Int(number) => Some(*number as f64),
redis::Value::Double(number) => Some(*number),
_ => string_value(value).and_then(|text| text.parse().ok()),
}
}
fn first_document(result: &redis::Value) -> Option<&[redis::Value]> {
let redis::Value::Array(items) = result else {
return None;
};
let [count, _document_id, fields, ..] = items.as_slice() else {
return None;
};
if !matches!(count, redis::Value::Int(count) if *count > 0) {
return None;
}
match fields {
redis::Value::Array(fields) => Some(fields.as_slice()),
_ => None,
}
}
fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> {
fields
.as_chunks::<2>()
.0
.iter()
.find(|pair| string_value(&pair[0]).as_deref() == Some(name))
.map(|pair| &pair[1])
}
fn string_field(fields: &[redis::Value], name: &str) -> Option<String> {
field_value(fields, name).and_then(string_value)
}
fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
field_value(fields, name).and_then(number_value)
}
fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
match field_value(fields, name)? {
redis::Value::BulkString(bytes) => Some(bytes.clone()),
redis::Value::SimpleString(text) => Some(text.clone().into_bytes()),
_ => None,
}
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}

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

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

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

@ -0,0 +1,97 @@
use litellm_cache::SemanticCacheContext;
use serde_json::Value;
pub fn prompt_from_context(context: &SemanticCacheContext) -> Option<String> {
if let Some(messages) = context.messages.as_ref().and_then(Value::as_array)
&& !messages.is_empty()
{
return Some(messages_text(messages));
}
let input = context.input.as_ref()?;
let mut parts = Vec::new();
collect_input_text(input, &mut parts);
let prompt = parts.join("\n").trim().to_string();
(!prompt.is_empty()).then_some(prompt)
}
fn messages_text(messages: &[Value]) -> String {
let mut text = String::new();
for message in messages {
let Some(message) = message.as_object() else {
continue;
};
match message.get("content") {
Some(Value::String(content)) => text.push_str(content),
Some(Value::Array(parts)) => {
for part in parts {
if let Some(text_content) = part.get("text").and_then(Value::as_str) {
text.push_str(text_content);
}
}
}
_ => {}
}
text.push_str(&search_results_text(message.get("search_results")));
}
text
}
fn search_results_text(search_results: Option<&Value>) -> String {
let Some(Value::Array(results)) = search_results else {
return String::new();
};
let mut text = String::new();
for result in results {
let Some(result) = result.as_object() else {
continue;
};
for key in ["source", "title"] {
if let Some(value) = result.get(key).and_then(Value::as_str) {
text.push_str(value);
}
}
if let Some(Value::Array(content)) = result.get("content") {
for block in content {
if let Some(value) = block.get("text").and_then(Value::as_str) {
text.push_str(value);
}
}
}
if let Some(citations) = result.get("citations") {
text.push_str(&citations.to_string());
}
}
text
}
fn collect_input_text(value: &Value, parts: &mut Vec<String>) {
match value {
Value::String(text) => {
let trimmed = text.trim();
if !trimmed.is_empty() {
parts.push(trimmed.to_string());
}
}
Value::Array(items) => {
for item in items {
collect_input_text(item, parts);
}
}
Value::Object(map) => {
if let Some(content) = map.get("content").filter(|content| !content.is_null()) {
collect_input_text(content, parts);
return;
}
for key in ["text", "output", "input_text", "output_text"] {
if let Some(Value::String(text)) = map.get(key) {
let trimmed = text.trim();
if !trimmed.is_empty() {
parts.push(trimmed.to_string());
return;
}
}
}
}
_ => {}
}
}

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

File diff suppressed because it is too large Load diff

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

@ -12,7 +12,7 @@ use redis::{
use super::REDIS_TIMEOUT;
use crate::topology::RedisNode;
pub(super) struct PooledConnection<C> {
pub struct PooledConnection<C> {
pub(super) connection: C,
pub(super) failed: bool,
}
@ -20,7 +20,7 @@ pub(super) struct PooledConnection<C> {
/// Pools connections without a checkout PING, which would double every operation's round trips.
/// A timed-out command leaves its reply on the socket while redis still reports the connection
/// open, so any connection whose operation failed is discarded instead of being reused.
pub(super) struct ConnectionManager(redis::Client);
pub struct ConnectionManager(redis::Client);
impl ConnectionManager {
pub(super) fn open(url: &str) -> Result<Self, Error> {
@ -54,7 +54,7 @@ impl r2d2::ManageConnection for ConnectionManager {
}
}
pub(super) struct ClusterConnectionManager(ClusterClient);
pub struct ClusterConnectionManager(ClusterClient);
impl ClusterConnectionManager {
pub(super) fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {

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

@ -1,7 +1,8 @@
use std::{sync::Mutex, time::Duration};
use litellm_cache::{
BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, get_cache,
BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, SemanticCacheContext,
get_cache,
};
struct TestCache {
@ -126,6 +127,24 @@ fn associated_context_preserves_backend_specific_lookup_inputs() {
);
}
#[test]
fn semantic_context_with_ttl_preserves_lookup_inputs() {
let context = SemanticCacheContext {
input: Some(serde_json::json!("text")),
messages: Some(serde_json::json!([{"role": "user", "content": "hi"}])),
metadata: Some(serde_json::json!({"key": "value"})),
scope: Some("scope".into()),
ttl: None,
};
let updated = context.with_ttl(Some(Duration::from_secs(30)));
assert_eq!(updated.ttl(), Some(Duration::from_secs(30)));
assert_eq!(updated.input, context.input);
assert_eq!(updated.messages, context.messages);
assert_eq!(updated.metadata, context.metadata);
assert_eq!(updated.scope, context.scope);
assert_eq!(context.with_ttl(None).ttl(), None);
}
#[tokio::test]
async fn default_batch_operations_use_async_writes_and_stop_on_failure() {
let cache = TestCache {

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

@ -27,6 +27,7 @@ litellm-cache-redis.workspace = true
litellm-cache-s3.workspace = true
litellm-cache-gcs.workspace = true
litellm-cache-disk.workspace = true
litellm-cache-redis-semantic.workspace = true
litellm-cache-response.workspace = true
litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" }
serde.workspace = true
@ -45,7 +46,7 @@ pyo3.workspace = true
pyo3-async-runtimes.workspace = true
redis = { version = "1.7.0", features = ["tls-rustls"] }
serde_json.workspace = true
tokio = { workspace = true, features = ["sync"] }
tokio = { workspace = true, features = ["rt", "sync"] }
[dev-dependencies]
serde.workspace = true

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

@ -88,6 +88,24 @@ pub(super) struct GcsCacheConfig {
pub(super) path_service_account: Option<String>,
}
pub(super) struct AzureBlobCacheConfig {
pub(super) account_url: String,
pub(super) container: String,
}
#[allow(
dead_code,
reason = "embedding settings are projected so drift falls back to Python"
)]
pub(super) struct RedisSemanticCacheConfig {
pub(super) redis_url: String,
pub(super) index_name: String,
pub(super) similarity_threshold: f64,
pub(super) embedding_model: String,
pub(super) embedding_max_input_tokens: Option<u64>,
pub(super) embedding_timeout: Option<f64>,
}
struct RedisClientProjection<'py> {
topology: RedisTopology,
host: String,
@ -107,11 +125,6 @@ pub(super) struct ValkeySemanticCacheConfig {
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>),
@ -120,6 +133,7 @@ pub(super) enum CacheBackendConfig {
ValkeySemantic(Box<ValkeySemanticCacheConfig>),
Disk(DiskCacheConfig),
AzureBlob(AzureBlobCacheConfig),
RedisSemantic(Box<RedisSemanticCacheConfig>),
}
#[allow(dead_code, reason = "consumed by the cache activation follow-up")]
@ -230,9 +244,15 @@ impl NativeCacheConfig {
backend: CacheBackendConfig::AzureBlob(backend),
}))
}),
Some(CacheType::RedisSemantic | CacheType::QdrantSemantic) | None => Ok(
CacheConfigProjection::Unsupported(UnsupportedCacheConfig::Backend),
),
Some(CacheType::RedisSemantic) => project_redis_semantic(&backend).map(|backend| {
CacheConfigProjection::Native(Box::new(Self {
policy,
backend: CacheBackendConfig::RedisSemantic(Box::new(backend)),
}))
}),
Some(CacheType::QdrantSemantic) | None => Ok(CacheConfigProjection::Unsupported(
UnsupportedCacheConfig::Backend,
)),
}
}
@ -244,7 +264,8 @@ impl NativeCacheConfig {
CacheBackendConfig::ValkeySemantic(_) => Some(Duration::ZERO),
CacheBackendConfig::Disk(_)
| CacheBackendConfig::AzureBlob(_)
| CacheBackendConfig::Gcs(_) => None,
| CacheBackendConfig::Gcs(_)
| CacheBackendConfig::RedisSemantic(_) => None,
};
if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_))
&& service.default_ttl() != default_ttl
@ -343,6 +364,20 @@ impl NativeCacheConfig {
let facade = std::fs::canonicalize(&config.directory).ok();
(native != facade).then_some("facade and native backend directories must match")
}
CacheBackendConfig::RedisSemantic(_) if service.kind() != "redis_semantic" => {
Some("facade and native backend types must match")
}
CacheBackendConfig::RedisSemantic(config)
if service.index_name() != Some(config.index_name.as_str()) =>
{
Some("facade and native backend index names must match")
}
CacheBackendConfig::RedisSemantic(config)
if service.similarity_threshold() != Some(config.similarity_threshold as f32) =>
{
Some("facade and native backend similarity thresholds must match")
}
CacheBackendConfig::RedisSemantic(_) => None,
CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() {
None => Some("facade and native backend types must match"),
Some((account_url, container))
@ -371,6 +406,27 @@ fn project_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult<AzureBlobCacheConf
})
}
#[inline(never)]
pub(super) fn project_redis_semantic(
backend: &Bound<'_, PyAny>,
) -> PyResult<RedisSemanticCacheConfig> {
Ok(RedisSemanticCacheConfig {
redis_url: backend.getattr("_redis_url")?.extract::<String>()?,
index_name: backend
.getattr("_index_name")?
.extract::<Option<String>>()?
.unwrap_or_else(|| "litellm_semantic_cache_index".into()),
similarity_threshold: backend.getattr("similarity_threshold")?.extract::<f64>()?,
embedding_model: backend.getattr("embedding_model")?.extract::<String>()?,
embedding_max_input_tokens: backend
.getattr("embedding_max_input_tokens")?
.extract::<Option<u64>>()?,
embedding_timeout: backend
.getattr("embedding_timeout")?
.extract::<Option<f64>>()?,
})
}
#[inline(never)]
fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult<MemoryCacheConfig> {
let max_size_kib = backend.getattr("max_size_per_item")?.extract::<usize>()?;

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

@ -1,23 +1,44 @@
use std::{future::Future, sync::Arc};
use std::future::Future;
use litellm_cache::Error;
use litellm_cache_valkey_semantic::Embedder;
use litellm_host_python::to_py;
use pyo3::{PyTraverseError, PyVisit, prelude::*};
use pyo3::{PyTraverseError, PyVisit, prelude::*, types::PyDict};
use serde_json::Value;
#[derive(Clone)]
pub(super) struct PythonEmbedder {
sync_embed: Arc<Py<PyAny>>,
async_embed_callable: Arc<Py<PyAny>>,
tokio::task_local! {
static PREPARED_EMBEDDING: Result<Vec<f32>, Error>;
}
pub(super) fn with_prepared_embedding<F: Future>(
vector: Result<Vec<f32>, Error>,
future: F,
) -> impl Future<Output = F::Output> {
PREPARED_EMBEDDING.scope(vector, future)
}
pub(super) struct PythonEmbedder(Py<PyAny>);
impl Clone for PythonEmbedder {
fn clone(&self) -> Self {
Python::attach(|py| Self(self.0.clone_ref(py)))
}
}
impl PythonEmbedder {
pub(super) fn new(object: Py<PyAny>) -> Self {
Self(object)
}
pub(super) fn from_backend(backend: &Bound<'_, PyAny>) -> PyResult<Self> {
Ok(Self {
sync_embed: Arc::new(backend.getattr("_get_embedding")?.unbind()),
async_embed_callable: Arc::new(backend.getattr("_get_async_embedding")?.unbind()),
})
Ok(Self(backend.clone().unbind()))
}
pub(super) fn object(&self) -> &Py<PyAny> {
&self.0
}
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&self.0)
}
pub(super) fn async_embed_awaitable<'py>(
@ -27,37 +48,109 @@ impl PythonEmbedder {
metadata: &Option<Value>,
) -> PyResult<Bound<'py, PyAny>> {
let metadata = to_py(py, metadata)?;
self.async_embed_callable.bind(py).call1((prompt, metadata))
self.0
.bind(py)
.call_method1("_get_async_embedding", (prompt, metadata))
}
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
visit.call(&*self.sync_embed)?;
visit.call(&*self.async_embed_callable)
fn metadata_kwargs<'py>(
py: Python<'py>,
metadata: Option<&Value>,
) -> PyResult<Bound<'py, PyDict>> {
let kwargs = PyDict::new(py);
kwargs.set_item("metadata", to_py(py, &metadata)?)?;
Ok(kwargs)
}
pub(super) fn async_embedding_coroutine(
&self,
py: Python<'_>,
prompt: &str,
metadata: Option<&Value>,
) -> PyResult<Py<PyAny>> {
let kwargs = Self::metadata_kwargs(py, metadata)?;
self.0
.bind(py)
.call_method("_get_async_embedding", (prompt,), Some(&kwargs))
.map(Bound::unbind)
}
pub(super) fn extract(vector: Bound<'_, PyAny>) -> PyResult<Vec<f32>> {
Ok(vector
.extract::<Vec<f64>>()?
.into_iter()
.map(|value| value as f32)
.collect())
}
}
impl Embedder for PythonEmbedder {
impl litellm_cache_valkey_semantic::Embedder for PythonEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
let result = Python::attach(|py| -> PyResult<Vec<f64>> {
let metadata = to_py(py, &metadata)?;
self.sync_embed
self.0
.bind(py)
.call1((prompt, metadata))?
.call_method1("_get_embedding", (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) }
let seeded = PREPARED_EMBEDDING
.try_with(Clone::clone)
.unwrap_or(Err(Error::Unavailable));
std::future::ready(seeded)
}
}
impl litellm_cache_redis_semantic::Embedder for PythonEmbedder {
fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
Python::attach(|py| {
let kwargs = Self::metadata_kwargs(py, metadata)?;
Self::extract(self.0.bind(py).call_method(
"_get_embedding",
(prompt,),
Some(&kwargs),
)?)
})
.map_err(|_| Error::Unavailable)
}
fn async_embed(
&self,
_prompt: &str,
_metadata: Option<&Value>,
) -> impl Future<Output = Result<Vec<f32>, Error>> + Send {
let seeded = PREPARED_EMBEDDING
.try_with(Clone::clone)
.unwrap_or(Err(Error::Unavailable));
std::future::ready(seeded)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn async_embed_returns_the_seeded_vector_or_unavailable() {
Python::initialize();
let object = Python::attach(|py| py.None());
let embedder = PythonEmbedder::new(object);
let scoped_embedder = embedder.clone();
let scoped = with_prepared_embedding(Ok(vec![0.25]), async move {
litellm_cache_redis_semantic::Embedder::async_embed(&scoped_embedder, "prompt", None)
.await
});
assert_eq!(scoped.await, Ok(vec![0.25]));
let unscoped =
litellm_cache_redis_semantic::Embedder::async_embed(&embedder, "prompt", None).await;
assert_eq!(unscoped, Err(Error::Unavailable));
}
}

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

@ -363,6 +363,11 @@ impl FacadeGuard {
let (module, name, cache_kind) = match (kind, cluster) {
("memory", _) => ("litellm.caching.in_memory_cache", "InMemoryCache", "local"),
("redis", false) => ("litellm.caching.redis_cache", "RedisCache", "redis"),
("redis_semantic", _) => (
"litellm.caching.redis_semantic_cache",
"RedisSemanticCache",
"redis-semantic",
),
("redis", true) => (
"litellm.caching.redis_cluster_cache",
"RedisClusterCache",
@ -400,6 +405,15 @@ impl FacadeGuard {
if let Some(message) = config.service_mismatch(service) {
return Err(PyTypeError::new_err(message));
}
if kind == "redis_semantic"
&& service
.embedder_object()
.is_none_or(|embedder| !backend.is(embedder.bind(py)))
{
return Err(PyTypeError::new_err(
"facade backend must be the native embedder",
));
}
Ok(Self {
outer: ObjectGuard::capture(
py,
@ -425,6 +439,13 @@ impl FacadeGuard {
"redis_kwargs",
"redis_flush_size",
"similarity_threshold",
"distance_threshold",
"embedding_model",
"embedding_max_input_tokens",
"embedding_timeout",
"_index_name",
"_redis_url",
"similarity_threshold",
"embedding_model",
"index_name",
"embedding_max_input_tokens",

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

@ -1,13 +1,18 @@
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig};
use litellm_cache_redis::{RedisNode, RedisTopology};
use litellm_cache_redis_semantic::RedisSemanticConfig;
use litellm_cache_s3::{S3CacheConfig, S3Endpoint};
use litellm_host_python::{release_gil, run_sync_value};
use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*};
use pyo3::{
PyTraverseError, PyVisit,
exceptions::{PyRuntimeError, PyTypeError},
prelude::*,
};
use super::{
cache_error, embedder::PythonEmbedder, facade::FacadeGuard, native::NativeResponseCache,
request::duration,
cache_error, config::project_redis_semantic, embedder::PythonEmbedder, facade::FacadeGuard,
native::NativeResponseCache, request::duration,
};
#[pyclass(frozen, name = "_CacheTestHandle")]
@ -179,6 +184,36 @@ impl CacheTestHandle {
})
}
#[staticmethod]
fn redis_semantic(py: Python<'_>, backend: Bound<'_, PyAny>) -> PyResult<Self> {
let class = py
.import("litellm.caching.redis_semantic_cache")?
.getattr("RedisSemanticCache")?;
if !backend.get_type().is(&class) {
return Err(PyTypeError::new_err(
"native redis-semantic handles require the built-in RedisSemanticCache",
));
}
let config = project_redis_semantic(&backend)?;
let embedder = PythonEmbedder::new(backend.unbind());
let service = release_gil(py, move || {
NativeResponseCache::redis_semantic(
&config.redis_url,
embedder,
RedisSemanticConfig {
index_name: config.index_name,
similarity_threshold: config.similarity_threshold as f32,
},
)
})
.map_err(cache_error)?;
Ok(Self {
service,
guard: None,
pid: std::process::id(),
})
}
#[getter]
fn backend(&self) -> &'static str {
self.service.kind()
@ -210,6 +245,7 @@ impl CacheTestHandle {
}
fn __traverse__(&self, visit: PyVisit<'_>) -> Result<(), PyTraverseError> {
self.service.traverse(&visit)?;
if let Some(guard) = &self.guard {
guard.traverse(visit)?;
}

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

@ -8,6 +8,7 @@ mod handle;
mod native;
mod request;
mod resolver;
mod semantic;
mod semantic_step;
use litellm_cache::Error;

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

@ -8,18 +8,20 @@ use litellm_cache_disk::DiskCache;
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::{RedisCache, RedisTopology};
use litellm_cache_redis_semantic::{RedisSemanticCache, RedisSemanticConfig};
use litellm_cache_response::{
CacheEntry, CacheKeyField, PartialHits, ResponseCache, ResponseCacheCodec,
ResponseCacheRequest, WriteBuffer,
};
use litellm_cache_s3::{S3Cache, S3CacheConfig};
use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig};
use pyo3::prelude::*;
use pyo3::{PyTraverseError, PyVisit, prelude::*};
use serde_json::Value;
use super::{
embedder::PythonEmbedder,
request::NativeRequest,
semantic::{SemanticBody, SemanticOperation, drive},
semantic_step::{SemanticEmbedExecution, drive_semantic},
};
@ -86,6 +88,10 @@ pub(super) enum NativeResponseCache {
embedder: PythonEmbedder,
scope: String,
},
RedisSemantic {
cache: Arc<ResponseCache<RedisSemanticCache<PythonEmbedder>>>,
embedder: PythonEmbedder,
},
Disk(Arc<ResponseCache<DiskCache<ResponseCacheCodec>>>),
AzureBlob(Arc<ResponseCache<AzureBlobCache<ResponseCacheCodec>>>),
}
@ -150,6 +156,18 @@ impl NativeResponseCache {
})
}
pub fn redis_semantic(
url: &str,
embedder: PythonEmbedder,
config: RedisSemanticConfig,
) -> Result<Self, Error> {
let backend = RedisSemanticCache::new(url, embedder.clone(), config)?;
Ok(Self::RedisSemantic {
cache: Arc::new(ResponseCache::new(Arc::new(backend))),
embedder,
})
}
pub fn disk(directory: &str) -> Result<Self, Error> {
let cache = DiskCache::open(directory, ResponseCacheCodec)?;
Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache)))))
@ -190,6 +208,7 @@ impl NativeResponseCache {
| Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::RedisSemantic { .. }
| Self::Disk(_)
| Self::Gcs(_) => None,
}
@ -204,6 +223,23 @@ impl NativeResponseCache {
}
}
pub(super) fn redis_semantic_request(
request: &NativeRequest,
) -> ResponseCacheRequest<SemanticCacheContext> {
ResponseCacheRequest {
key: request.key.clone(),
controls: request.controls,
context: SemanticCacheContext {
input: request.input.clone(),
messages: request.messages.clone(),
metadata: request.metadata.clone(),
scope: request.scope.clone(),
ttl: request.ttl,
},
max_age: request.max_age,
}
}
fn semantic(
request: &NativeRequest,
scope: &str,
@ -252,6 +288,7 @@ impl NativeResponseCache {
Self::S3(_) => "s3",
Self::Gcs(_) => "gcs",
Self::ValkeySemantic { .. } => "valkey-semantic",
Self::RedisSemantic { .. } => "redis_semantic",
Self::Disk(_) => "disk",
Self::AzureBlob(_) => "azure-blob",
}
@ -264,6 +301,7 @@ impl NativeResponseCache {
Self::S3(cache) => cache.default_ttl(),
Self::Gcs(cache) => cache.default_ttl(),
Self::ValkeySemantic { cache, .. } => cache.default_ttl(),
Self::RedisSemantic { cache, .. } => cache.default_ttl(),
Self::Disk(cache) => cache.default_ttl(),
Self::AzureBlob(cache) => cache.default_ttl(),
}
@ -302,6 +340,7 @@ impl NativeResponseCache {
Self::Memory(_)
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::RedisSemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
@ -314,6 +353,7 @@ impl NativeResponseCache {
Self::Memory(_)
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::RedisSemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
@ -327,6 +367,7 @@ impl NativeResponseCache {
Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::RedisSemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
@ -339,6 +380,7 @@ impl NativeResponseCache {
Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::RedisSemantic { .. }
| Self::Disk(_)
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
@ -352,6 +394,7 @@ impl NativeResponseCache {
| Self::Redis { .. }
| Self::S3(_)
| Self::ValkeySemantic { .. }
| Self::RedisSemantic { .. }
| Self::AzureBlob(_)
| Self::Gcs(_) => None,
}
@ -363,6 +406,38 @@ impl NativeResponseCache {
cache.backend().similarity_threshold(),
cache.backend().index_name(),
)),
Self::RedisSemantic { cache, .. } => Some((
f64::from(cache.backend().similarity_threshold()),
cache.backend().index_name(),
)),
_ => None,
}
}
pub fn index_name(&self) -> Option<&str> {
match self {
Self::RedisSemantic { cache, .. } => Some(cache.backend().index_name()),
_ => None,
}
}
pub fn similarity_threshold(&self) -> Option<f32> {
match self {
Self::RedisSemantic { cache, .. } => Some(cache.backend().similarity_threshold()),
_ => None,
}
}
pub fn semantic_embedder(&self) -> Option<&PythonEmbedder> {
match self {
Self::RedisSemantic { embedder, .. } => Some(embedder),
_ => None,
}
}
pub fn embedder_object(&self) -> Option<&Py<PyAny>> {
match self {
Self::RedisSemantic { embedder, .. } => Some(embedder.object()),
_ => None,
}
}
@ -375,6 +450,9 @@ impl NativeResponseCache {
Self::ValkeySemantic { cache, scope, .. } => {
cache.lookup(&Self::semantic(request, scope), now)
}
Self::RedisSemantic { cache, .. } => {
cache.lookup(&Self::redis_semantic_request(request), now)
}
Self::Gcs(cache) => cache.lookup(&Self::exact(request), now),
Self::Disk(cache) => cache.lookup(&Self::exact(request), now),
Self::AzureBlob(cache) => cache.lookup(&Self::exact(request), now),
@ -394,6 +472,9 @@ impl NativeResponseCache {
Self::ValkeySemantic { cache, scope, .. } => {
cache.store(&Self::semantic(request, scope), response, now)
}
Self::RedisSemantic { cache, .. } => {
cache.store(&Self::redis_semantic_request(request), response, now)
}
Self::Gcs(cache) => cache.store(&Self::exact(request), response, now),
Self::Disk(cache) => cache.store(&Self::exact(request), response, now),
Self::AzureBlob(cache) => cache.store(&Self::exact(request), response, now),
@ -417,7 +498,9 @@ impl NativeResponseCache {
Self::S3(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => {
Err(Error::UnsupportedOperation)
}
Self::Gcs(cache) => {
cache.lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
}
@ -444,6 +527,11 @@ impl NativeResponseCache {
.async_lookup(&Self::semantic(request, scope), now)
.await
}
Self::RedisSemantic { cache, .. } => {
cache
.async_lookup(&Self::redis_semantic_request(request), now)
.await
}
Self::Gcs(cache) => cache.async_lookup(&Self::exact(request), now).await,
Self::Disk(cache) => cache.async_lookup(&Self::exact(request), now).await,
Self::AzureBlob(cache) => cache.async_lookup(&Self::exact(request), now).await,
@ -481,6 +569,10 @@ impl NativeResponseCache {
Self::semantic(&request, scope),
),
),
Self::RedisSemantic { .. } => drive(
py,
SemanticBody::new(self.clone(), SemanticOperation::Lookup(request)),
),
}
}
@ -522,6 +614,11 @@ impl NativeResponseCache {
.async_store(&Self::semantic(request, scope), response, now)
.await
}
Self::RedisSemantic { cache, .. } => {
cache
.async_store(&Self::redis_semantic_request(request), response, now)
.await
}
Self::Gcs(cache) => {
cache
.async_store(&Self::exact(request), response, now)
@ -577,6 +674,10 @@ impl NativeResponseCache {
response,
),
),
Self::RedisSemantic { .. } => drive(
py,
SemanticBody::new(self.clone(), SemanticOperation::Store(request, response)),
),
}
}
@ -599,7 +700,9 @@ impl NativeResponseCache {
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
.await
}
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => {
Err(Error::UnsupportedOperation)
}
Self::Gcs(cache) => {
cache
.async_lookup_batch(&requests.iter().map(Self::exact).collect::<Vec<_>>(), now)
@ -652,6 +755,7 @@ impl NativeResponseCache {
.collect();
cache.async_store_batch(entries, now).await
}
Self::RedisSemantic { .. } => Err(Error::UnsupportedOperation),
Self::Gcs(cache) => {
let entries = entries
.into_iter()
@ -718,6 +822,10 @@ impl NativeResponseCache {
),
)
}
Self::RedisSemantic { .. } => drive(
py,
SemanticBody::new(self.clone(), SemanticOperation::StoreBatch(entries.into())),
),
}
}
@ -731,7 +839,9 @@ impl NativeResponseCache {
cache.async_flush().await
}
Self::S3(cache) => cache.async_flush().await,
Self::ValkeySemantic { .. } => Err(Error::UnsupportedOperation),
Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => {
Err(Error::UnsupportedOperation)
}
Self::Gcs(cache) => cache.async_flush().await,
Self::Disk(cache) => cache.async_flush().await,
Self::AzureBlob(cache) => cache.async_flush().await,
@ -744,12 +854,22 @@ impl NativeResponseCache {
Self::Redis { cache, .. } => cache.test_connection().await,
Self::S3(cache) => cache.test_connection().await,
Self::ValkeySemantic { cache, .. } => cache.test_connection().await,
Self::RedisSemantic { .. } => Err(Error::UnsupportedOperation),
Self::Gcs(cache) => cache.test_connection().await,
Self::Disk(cache) => cache.test_connection().await,
Self::AzureBlob(cache) => cache.test_connection().await,
}
}
pub(super) fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
match self {
Self::ValkeySemantic { embedder, .. } => embedder.traverse(visit)?,
Self::RedisSemantic { embedder, .. } => embedder.traverse(visit)?,
_ => {}
}
Ok(())
}
pub fn gcs_backend(&self) -> Option<&GcsCache<ResponseCacheCodec>> {
match self {
Self::Gcs(cache) => Some(cache.backend()),
@ -777,6 +897,7 @@ mod tests {
metadata: Some(metadata),
litellm_metadata: None,
litellm_params: None,
scope: None,
}
}

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

@ -19,8 +19,10 @@ struct RequestInput {
metadata: Option<Value>,
litellm_metadata: Option<Value>,
litellm_params: Option<Value>,
scope: Option<String>,
}
#[derive(Clone)]
pub(super) struct NativeRequest {
pub(super) key: CacheKeyInput,
pub(super) controls: CacheControls,
@ -31,6 +33,7 @@ pub(super) struct NativeRequest {
pub(super) metadata: Option<Value>,
pub(super) litellm_metadata: Option<Value>,
pub(super) litellm_params: Option<Value>,
pub(super) scope: Option<String>,
}
pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult<NativeRequest> {
@ -52,6 +55,7 @@ fn request_input(input: RequestInput) -> PyResult<NativeRequest> {
metadata: input.metadata,
litellm_metadata: input.litellm_metadata,
litellm_params: input.litellm_params,
scope: input.scope,
})
}

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

@ -0,0 +1,175 @@
use std::collections::VecDeque;
use litellm_cache::Error;
use litellm_cache_redis_semantic::prompt_from_context;
use litellm_host_python::{Execution, ExecutionBody, ExecutionStep, run_async};
use pyo3::{
PyTraverseError, PyVisit,
exceptions::{PyException, PyRuntimeError},
prelude::*,
};
use serde_json::Value;
use super::{
cache_error,
embedder::{PythonEmbedder, with_prepared_embedding},
native::NativeResponseCache,
request::{NativeRequest, now},
};
pub(super) enum SemanticOperation {
Lookup(NativeRequest),
Store(NativeRequest, Value),
StoreBatch(VecDeque<(NativeRequest, Value)>),
}
enum Phase {
Start,
AwaitingEmbedding,
AwaitingBackend,
}
pub(super) struct SemanticBody {
service: NativeResponseCache,
operation: SemanticOperation,
pending: Option<(NativeRequest, Option<Value>)>,
phase: Phase,
}
impl SemanticBody {
pub(super) fn new(service: NativeResponseCache, operation: SemanticOperation) -> Self {
Self {
service,
operation,
pending: None,
phase: Phase::Start,
}
}
fn backend_step(
&mut self,
py: Python<'_>,
seed: Result<Vec<f32>, Error>,
) -> PyResult<ExecutionStep> {
self.phase = Phase::AwaitingBackend;
let (request, response) = self.pending.take().ok_or_else(|| {
PyRuntimeError::new_err("semantic execution resumed without a pending operation")
})?;
let service = self.service.clone();
let future = async move {
match response {
None => service.async_lookup(&request, now()).await,
Some(response) => service
.async_store(&request, response, now())
.await
.map(|_| None),
}
};
let awaitable = run_async(py, with_prepared_embedding(seed, future), cache_error)?;
Ok(ExecutionStep::Await(awaitable.unbind()))
}
}
impl ExecutionBody for SemanticBody {
fn resume(&mut self, mut result: Option<PyResult<Py<PyAny>>>) -> PyResult<ExecutionStep> {
Python::attach(|py| {
loop {
match self.phase {
Phase::Start => {
if result.is_some() {
return Err(PyRuntimeError::new_err(
"semantic execution received a result before starting",
));
}
if self.pending.is_none() {
match &mut self.operation {
SemanticOperation::Lookup(request) => {
self.pending = Some((request.clone(), None));
}
SemanticOperation::Store(request, response) => {
let response = std::mem::replace(response, Value::Null);
self.pending = Some((request.clone(), Some(response)));
}
SemanticOperation::StoreBatch(queue) => {
let Some((request, response)) = queue.pop_front() else {
return Ok(ExecutionStep::Return(py.None()));
};
self.pending = Some((request, Some(response)));
}
}
}
let (request, _) = self.pending.as_ref().ok_or_else(|| {
PyRuntimeError::new_err("semantic execution has no pending operation")
})?;
let semantic = NativeResponseCache::redis_semantic_request(request);
let Some(prompt) = prompt_from_context(&semantic.context) else {
return self.backend_step(py, Err(Error::Unavailable));
};
let embedder = self.service.semantic_embedder().ok_or_else(|| {
PyRuntimeError::new_err(
"semantic execution requires a redis-semantic backend",
)
})?;
let coroutine = embedder.async_embedding_coroutine(
py,
&prompt,
semantic.context.metadata.as_ref(),
)?;
self.phase = Phase::AwaitingEmbedding;
return Ok(ExecutionStep::Await(coroutine));
}
Phase::AwaitingEmbedding => {
let result = result.take().ok_or_else(|| {
PyRuntimeError::new_err(
"semantic execution expected an embedding result",
)
})?;
let seed = match result {
Ok(value) => PythonEmbedder::extract(value.into_bound(py))
.map_err(|_| Error::Unavailable),
Err(error) => {
if !error.is_instance_of::<PyException>(py) {
return Err(error);
}
Err(Error::Unavailable)
}
};
return self.backend_step(py, seed);
}
Phase::AwaitingBackend => {
let result = result.take().ok_or_else(|| {
PyRuntimeError::new_err("semantic execution expected a backend result")
})?;
let value = match result {
Ok(value) => value,
Err(error) => return Err(error),
};
let more = matches!(
&self.operation,
SemanticOperation::StoreBatch(queue) if !queue.is_empty()
);
if more {
self.phase = Phase::Start;
continue;
}
return Ok(ExecutionStep::Return(value));
}
}
}
})
}
fn traverse(&self, visit: &PyVisit<'_>) -> Result<(), PyTraverseError> {
if let Some(embedder) = self.service.semantic_embedder() {
embedder.traverse(visit)?;
}
Ok(())
}
}
pub(super) fn drive(py: Python<'_>, body: SemanticBody) -> PyResult<Bound<'_, PyAny>> {
let execution = Py::new(py, Execution::new(body))?;
py.import("litellm.rust_bridge.lifecycle")?
.getattr("drive")?
.call1((execution,))
}

2
litellm/rust_bridge/_native.pyi
View file

@ -169,6 +169,8 @@ class _CacheTestHandle:
@staticmethod
def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ...
@staticmethod
def redis_semantic(backend: object) -> _CacheTestHandle: ...
@staticmethod
def valkey_semantic(
url: str,
similarity_threshold: float,

641
tests/test_litellm_rust/test_cache.py
View file

@ -1,19 +1,23 @@
import asyncio
import contextvars
import gc
import hashlib
import json
import math
import os
import threading
import time
import uuid
import weakref
from collections.abc import Generator
from collections.abc import Callable, Generator
from contextlib import ExitStack
from datetime import datetime
from pathlib import Path
from types import SimpleNamespace
from typing import Final, Protocol, cast
from unittest.mock import Mock
from urllib.parse import urlparse
from uuid import uuid4
import boto3
import botocore.config
@ -30,13 +34,19 @@ from litellm.caching.disk_cache import DiskCache
from litellm.caching.gcs_cache import GCSCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cluster_cache import RedisClusterCache
from litellm.caching.redis_semantic_cache import RedisSemanticCache
from litellm.caching.s3_cache import S3Cache
from litellm.rust_bridge import _native
from litellm.types.caching import LiteLLMCacheType
from litellm.types.llms.custom_llm import CustomLLMItem
from litellm.types.utils import EmbeddingResponse
from tests.test_litellm_rust.support.fake_gcs import FakeGcs
from tests.test_litellm_rust.support.isolation import rebound
from tests.test_litellm_rust.support.s3_stub import S3Stub
_CacheTestHandle: Final = _native._CacheTestHandle # pyright: ignore[reportPrivateUsage] # test-only handle has no public module name
_CacheTestResolver: Final = _native._CacheTestResolver # pyright: ignore[reportPrivateUsage] # test-only resolver has no public module name
pytestmark: Final = pytest.mark.requires_rust_extension
@ -114,14 +124,14 @@ def test_existing_constructor_and_global_are_unchanged() -> None:
assert type(facade.cache) is InMemoryCache
assert "_native_cache_handle" not in vars(facade)
with rebound(litellm, "cache", facade):
resolver: Final = _native._CacheTestResolver(litellm)
resolver: Final = _CacheTestResolver(litellm)
assert resolver.resolve().kind == "python_callback"
resolver.resolve().store(None, {"answer": 7}, callback_kwargs={"cache_key": "key"})
assert cast(CacheLookup, facade).get_cache(cache_key="key") == {"answer": 7}
def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> None:
resolver: Final = _native._CacheTestResolver(litellm)
resolver: Final = _CacheTestResolver(litellm)
enable_cache(type=LiteLLMCacheType.LOCAL, ttl=30)
enabled: Final = litellm.cache
@ -144,13 +154,13 @@ def test_existing_global_lifecycle_remains_the_resolver_source_of_truth() -> Non
async def test_native_bindings_survive_replacement_and_capture_writes_before_dispatch() -> None:
namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.memory())
resolver: Final = _native._CacheTestResolver(namespace)
namespace: Final = SimpleNamespace(cache=_CacheTestHandle.memory())
resolver: Final = _CacheTestResolver(namespace)
selected: Final = resolver.resolve()
assert selected.kind == "native"
selected.store(request(), {"answer": 1})
assert await selected.async_lookup(request()) == {"answer": 1}
with rebound(namespace, "cache", _native._CacheTestHandle.memory()):
with rebound(namespace, "cache", _CacheTestHandle.memory()):
replacement: Final = resolver.resolve()
await selected.async_store(request(), {"answer": 2})
assert replacement.lookup(request()) is None
@ -183,7 +193,7 @@ async def test_python_callback_preserves_identity_caller_task_context_and_errors
raise failure
namespace: Final = SimpleNamespace(cache=CustomCache())
binding: Final = _native._CacheTestResolver(namespace).resolve()
binding: Final = _CacheTestResolver(namespace).resolve()
assert binding.kind == "python_callback"
assert await binding.async_lookup(None, callback_kwargs={"marker": sentinel}) is sentinel
assert context.get() == "callback"
@ -204,7 +214,7 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None:
finally:
finished.set()
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve()
binding: Final = _CacheTestResolver(SimpleNamespace(cache=CustomCache())).resolve()
async def lookup() -> object:
return await binding.async_lookup(None, callback_kwargs={})
@ -219,9 +229,9 @@ async def test_callback_cancellation_stays_in_the_callers_task() -> None:
def test_registered_facade_uses_native_and_instance_overrides_fall_back() -> None:
facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
handle: Final = _native._CacheTestHandle.memory()
handle: Final = _CacheTestHandle.memory()
handle._bind_facade(facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade))
native: Final = resolver.resolve()
assert native.kind == "native"
native.store(request(), {"source": "native"})
@ -252,12 +262,12 @@ def test_facade_subclasses_backend_replacement_and_configuration_changes_are_not
class CustomCache(Cache):
pass
handle: Final = _native._CacheTestHandle.memory()
handle: Final = _CacheTestHandle.memory()
with pytest.raises(TypeError):
handle._bind_facade(CustomCache(type=LiteLLMCacheType.LOCAL))
facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
handle._bind_facade(facade)
resolver: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade))
resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade))
with rebound(facade, "cache", InMemoryCache()):
assert resolver.resolve().kind == "python_callback"
with rebound(facade, "ttl", 12):
@ -282,7 +292,7 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None:
def cyclic_reference() -> weakref.ReferenceType[CustomCache]:
callback: Final = CustomCache()
namespace: Final = SimpleNamespace(cache=callback)
binding: Final = _native._CacheTestResolver(namespace).resolve()
binding: Final = _CacheTestResolver(namespace).resolve()
setattr(callback, "binding", binding)
return weakref.ref(callback)
@ -293,8 +303,8 @@ def test_resolver_and_callback_cycles_can_be_collected() -> None:
async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidden_prefix(redis_url: str) -> None:
client: Final = redis.Redis.from_url(redis_url)
namespace: Final = SimpleNamespace(cache=_native._CacheTestHandle.redis(redis_url, namespace="team"))
binding: Final = _native._CacheTestResolver(namespace).resolve()
namespace: Final = SimpleNamespace(cache=_CacheTestHandle.redis(redis_url, namespace="team"))
binding: Final = _CacheTestResolver(namespace).resolve()
response: Final = {"choices": [{"text": "cached"}], "usage": {"total_tokens": 3}, "flag": True, "empty": None}
envelope: Final = {"timestamp": time.time(), "response": json.dumps(response)}
client.set("team:sync", str(envelope))
@ -316,33 +326,33 @@ async def test_redis_reads_python_sync_and_async_entries_and_writes_without_hidd
def test_invalid_duration_and_request_shape_fail_before_storage() -> None:
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve()
binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve()
for seconds in (-1.0, float("nan"), float("inf")):
with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"):
binding.store({**request(), "ttl_seconds": seconds}, {"answer": 1})
assert binding.lookup(request()) is None
with pytest.raises(ValueError, match="cache durations must be finite and nonnegative"):
_native._CacheTestHandle.memory(ttl_seconds=-1)
_CacheTestHandle.memory(ttl_seconds=-1)
async def test_memory_size_policy_is_applied_by_the_native_host() -> None:
handle: Final = _native._CacheTestHandle.memory(capacity=2, max_entry_bytes=128)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
handle: Final = _CacheTestHandle.memory(capacity=2, max_entry_bytes=128)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=handle)).resolve()
small: Final = {"answer": "ok"}
binding.store(request("small"), small)
assert await binding.async_lookup(request("small")) == small
await binding.async_store(request("large"), {"answer": "x" * 256})
assert binding.lookup(request("large")) is None
assert binding.lookup(request("small")) == small
disabled: Final = _native._CacheTestResolver(
SimpleNamespace(cache=_native._CacheTestHandle.memory(capacity=0))
disabled: Final = _CacheTestResolver(
SimpleNamespace(cache=_CacheTestHandle.memory(capacity=0))
).resolve()
await disabled.async_store(request(), small)
assert await disabled.async_lookup(request()) is None
async def test_native_batch_lookup_and_store_report_partial_hits() -> None:
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=_native._CacheTestHandle.memory())).resolve()
binding: Final = _CacheTestResolver(SimpleNamespace(cache=_CacheTestHandle.memory())).resolve()
requests: Final = [request("hit"), request("miss"), request("disabled")]
requests[2]["controls"] = {
"supported_call_type": True,
@ -380,7 +390,7 @@ async def test_python_batch_callbacks_use_the_builtin_cache_api() -> None:
) -> object:
return result, kwargs
binding: Final = _native._CacheTestResolver(
binding: Final = _CacheTestResolver(
SimpleNamespace(cache=CustomCache(type=LiteLLMCacheType.LOCAL))
).resolve()
assert binding.kind == "python_callback"
@ -410,7 +420,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None:
cache: Final = Cache(type=LiteLLMCacheType.LOCAL)
cache.cache.set_cache("key", "value")
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=cache)).resolve()
binding: Final = _CacheTestResolver(SimpleNamespace(cache=cache)).resolve()
assert binding.kind == "python_callback"
setattr(cache.cache, "ping", ping)
@ -422,7 +432,7 @@ async def test_unmodified_builtin_cache_callbacks_can_ping_and_flush() -> None:
def test_facade_registration_rejects_mismatched_capacity() -> None:
facade: Final = Cache(type=LiteLLMCacheType.LOCAL)
with pytest.raises(TypeError, match="capacities must match"):
_native._CacheTestHandle.memory(capacity=7)._bind_facade(facade)
_CacheTestHandle.memory(capacity=7)._bind_facade(facade)
def test_azure_blob_facade_serves_natively_and_python_reads_the_same_blobs(azure_blob_facade: Cache) -> None:
@ -525,19 +535,19 @@ async def test_redis_facade_buffers_native_async_writes(redis_url: str) -> None:
redis_flush_size=2,
)
with pytest.raises(TypeError, match="default TTLs must match"):
_native._CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade)
_CacheTestHandle.redis(redis_url, ttl_seconds=61)._bind_facade(facade)
with pytest.raises(TypeError, match="namespaces must match"):
_native._CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade)
_native._CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade)
binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
_CacheTestHandle.redis(redis_url, namespace="other")._bind_facade(facade)
_CacheTestHandle.redis(redis_url, ttl_seconds=60)._bind_facade(facade)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
client: Final = redis.Redis.from_url(redis_url)
with rebound(facade.cache, "redis_kwargs", {**facade.cache.redis_kwargs, "ssl": True}):
assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback"
assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback"
pool: Final = facade.cache.redis_client.connection_pool
with rebound(pool, "connection_kwargs", {**pool.connection_kwargs, "db": 1}):
assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback"
assert _CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback"
await binding.async_store(request("first"), {"value": 1})
assert client.get("first") is None
@ -1112,3 +1122,572 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n
client.delete("unscoped")
client.close()
facade.cache.redis_client.close()
PARAPHRASE_MARKER: Final = " (paraphrase)"
SEMANTIC_EMBEDDING_MODEL: Final = "semantic-test/deterministic"
SEMANTIC_INDEX_PREFIX: Final = "litellm_test_semantic_"
SEMANTIC_CONTEXT: Final = contextvars.ContextVar("semantic_test_context", default="unset")
def _normalized(vector: list[float]) -> list[float]:
norm: Final = math.sqrt(sum(component * component for component in vector))
return [component / norm for component in vector]
def _base_embedding(prompt: str) -> list[float]:
digest: Final = hashlib.sha256(prompt.encode("utf-8")).digest()
return _normalized([float(digest[index] + 1) for index in range(8)])
def _semantic_embedding(prompt: str) -> list[float]:
if PARAPHRASE_MARKER not in prompt:
return _base_embedding(prompt)
base: Final = _base_embedding(prompt.replace(PARAPHRASE_MARKER, "").strip())
pivot: Final = min(range(8), key=lambda index: abs(base[index]))
direction: Final = _normalized(
[
(1.0 - base[pivot] * base[pivot]) if index == pivot else -base[index] * base[pivot]
for index in range(8)
]
)
# Rotating an orthogonal unit direction by 0.329 produces ~0.05 cosine distance
return _normalized([base[index] + 0.329 * direction[index] for index in range(8)])
class DeterministicEmbedding(litellm.CustomLLM):
def __init__(self) -> None:
self.calls: list[dict[str, object]] = []
self.async_calls: list[dict[str, object]] = []
self.entered = asyncio.Event()
self.gate: asyncio.Event | None = None
def _respond(
self,
model: str,
input: object,
model_response: EmbeddingResponse,
) -> EmbeddingResponse:
texts: Final = cast(list[object], input if isinstance(input, list) else [input])
self.calls.append({"model": model, "input": texts})
model_response.model = model
model_response.data = [
{"object": "embedding", "index": index, "embedding": _semantic_embedding(str(text))}
for index, text in enumerate(texts)
]
return model_response
def embedding(
self,
model: str,
input: list[object],
model_response: EmbeddingResponse,
print_verbose: Callable[..., object],
logging_obj: object,
optional_params: dict[str, object],
api_key: object = None,
api_base: object = None,
timeout: object = None,
litellm_params: object = None,
) -> EmbeddingResponse:
return self._respond(model, input, model_response)
async def aembedding(
self,
model: str,
input: list[object],
model_response: EmbeddingResponse,
print_verbose: Callable[..., object],
logging_obj: object,
optional_params: dict[str, object],
api_key: object = None,
api_base: object = None,
timeout: object = None,
litellm_params: object = None,
) -> EmbeddingResponse:
texts: Final = cast(list[object], input if isinstance(input, list) else [input])
self.async_calls.append(
{
"model": model,
"input": texts,
"task": asyncio.current_task(),
"context": SEMANTIC_CONTEXT.get(),
}
)
SEMANTIC_CONTEXT.set("written-in-aembedding")
self.entered.set()
if self.gate is not None:
await self.gate.wait()
return self._respond(model, input, model_response)
@pytest.fixture
def semantic_embedding() -> Generator[DeterministicEmbedding]:
handler: Final = DeterministicEmbedding()
with ExitStack() as stack:
stack.enter_context(
rebound(
litellm,
"custom_provider_map",
[
*litellm.custom_provider_map,
cast(
CustomLLMItem,
{"provider": "semantic-test", "custom_handler": handler},
),
],
)
)
stack.enter_context(
rebound(
litellm,
"_custom_providers", # pyright: ignore[reportPrivateUsage] # no public provider-registration hook
[*litellm._custom_providers, "semantic-test"], # pyright: ignore[reportPrivateUsage] # no public provider-registration hook
)
)
stack.enter_context(
rebound(litellm, "provider_list", [*litellm.provider_list, "semantic-test"])
)
yield handler
@pytest.fixture
def redis_stack() -> Generator[tuple[str, str]]:
url: Final = os.environ.get("LITELLM_REDIS_STACK_URL")
if url is None:
pytest.skip("LITELLM_REDIS_STACK_URL is not set")
index: Final = f"{SEMANTIC_INDEX_PREFIX}{uuid4().hex}"
yield url, index
client: Final = redis.Redis.from_url(url)
try:
client.execute_command("FT.DROPINDEX", index, "DD") # pyright: ignore[reportUnknownMemberType] # redis-py leaves execute_command partially unknown
except redis.RedisError:
pass
client.close()
def semantic_request(key: str, prompt: str, **extra: object) -> dict[str, object]:
return {
"key": {"preset": key},
"messages": [{"role": "user", "content": prompt}],
**extra,
}
def semantic_messages(prompt: str) -> list[dict[str, object]]:
return [{"role": "user", "content": prompt}]
def semantic_entry_id(prompt: str, tag: str) -> str:
return hashlib.sha256(f"{prompt}litellm_cache_key{tag}".encode()).hexdigest()
def semantic_facade(url: str, index: str, *, similarity_threshold: float = 0.8) -> Cache:
facade: Final = Cache(
type=LiteLLMCacheType.REDIS_SEMANTIC,
redis_url=url,
similarity_threshold=similarity_threshold,
redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL,
redis_semantic_cache_index_name=index,
)
_CacheTestHandle.redis_semantic(facade.cache)._bind_facade(facade)
return facade
def test_redis_semantic_constructor_identity_and_provenance(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
backend: Final = cast(RedisSemanticCache, facade.cache)
assert backend.__class__.__module__ == "litellm.caching.redis_semantic_cache"
assert type(backend) is RedisSemanticCache
assert backend._redis_url == url # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config
assert backend._index_name == index # pyright: ignore[reportPrivateUsage] # provenance check needs the projected config
assert backend.similarity_threshold == 0.8
assert backend.embedding_model == SEMANTIC_EMBEDDING_MODEL
handle: Final = cast(object, getattr(facade, "_native_cache_handle"))
assert isinstance(handle, _CacheTestHandle)
assert handle.backend == "redis_semantic"
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
assert binding.kind == "native"
def test_redis_semantic_native_and_python_sync_entries_share_one_layout(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
client: Final = redis.Redis.from_url(url)
response: Final = {"choices": [{"text": "paris"}], "usage": {"total_tokens": 2}}
binding.store(semantic_request("geo", "what is the capital of france"), response)
native_hash_key: Final = f"{index}:{semantic_entry_id('what is the capital of france', 'geo')}"
stored: Final = client.hgetall(native_hash_key)
assert set(stored) == {
b"entry_id",
b"prompt",
b"response",
b"prompt_vector",
b"inserted_at",
b"updated_at",
b"litellm_cache_key",
}, stored
assert stored[b"entry_id"].decode() == native_hash_key.split(":", 1)[1]
assert stored[b"prompt"] == b"what is the capital of france"
assert stored[b"litellm_cache_key"] == b"geo"
assert len(stored[b"prompt_vector"]) == 32
decoded: Final = cast(dict[str, object], json.loads(stored[b"response"]))
assert decoded["response"] == response
assert (
cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class
"geo", messages=semantic_messages("what is the capital of france")
)
== decoded
)
assert semantic_embedding.calls == [
{"model": "deterministic", "input": ["what is the capital of france"]},
{"model": "deterministic", "input": ["what is the capital of france"]},
{"model": "deterministic", "input": ["dimension test"]},
]
cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class
"math",
json.dumps({"timestamp": 1700000000.0, "response": {"answer": 42}}),
messages=semantic_messages("what is 6 times 7"),
)
python_hash_key: Final = f"{index}:{semantic_entry_id('what is 6 times 7', 'math')}"
assert json.loads(cast(bytes, client.hget(python_hash_key, "response"))) == {
"timestamp": 1700000000.0,
"response": {"answer": 42},
}
assert binding.lookup(semantic_request("math", "what is 6 times 7")) == {"answer": 42}
client.close()
async def test_redis_semantic_async_paths_and_store_batch_share_one_layout(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
client: Final = redis.Redis.from_url(url)
await binding.async_store(
semantic_request("async", "name a primary color"), {"answer": "blue"}
)
hash_key: Final = f"{index}:{semantic_entry_id('name a primary color', 'async')}"
decoded: Final = cast(dict[str, object], json.loads(cast(bytes, client.hget(hash_key, "response"))))
python_read: Final = await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class
"async", messages=semantic_messages("name a primary color")
)
assert python_read == decoded
await binding.async_store_batch(
[
semantic_request("batch-one", "first batch prompt"),
semantic_request("batch-two", "second batch prompt"),
],
[{"answer": 1}, {"answer": 2}],
)
expected: Final = {
key: json.loads(
cast(bytes, client.hget(f"{index}:{semantic_entry_id(prompt, key)}", "response"))
)
for key, prompt in (
("batch-one", "first batch prompt"),
("batch-two", "second batch prompt"),
)
}
for key, prompt in (
("batch-one", "first batch prompt"),
("batch-two", "second batch prompt"),
):
assert cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class
key, messages=semantic_messages(prompt)
) == expected[key], key
cast(RedisSemanticCache, facade.cache).set_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class
"async-python",
json.dumps({"timestamp": 1700000000.0, "response": {"answer": "python"}}),
messages=semantic_messages("python written prompt"),
)
assert await binding.async_lookup(
semantic_request("async-python", "python written prompt")
) == {"answer": "python"}
client.close()
async def test_native_semantic_async_embedding_runs_inline_in_the_callers_task(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
assert binding.kind == "native"
caller: Final = asyncio.current_task()
SEMANTIC_CONTEXT.set("caller-sentinel")
response: Final = {"choices": [{"text": "paris"}]}
await binding.async_store(
semantic_request("inline", "what is the capital of france"), response
)
assert (
await binding.async_lookup(
semantic_request("inline", f"what is the capital of france{PARAPHRASE_MARKER}")
)
== response
)
assert await binding.async_lookup(semantic_request("inline", "python written prompt")) is None
assert SEMANTIC_CONTEXT.get() == "written-in-aembedding"
assert semantic_embedding.async_calls == [
{
"model": "deterministic",
"input": ["what is the capital of france"],
"task": caller,
"context": "caller-sentinel",
},
{
"model": "deterministic",
"input": [f"what is the capital of france{PARAPHRASE_MARKER}"],
"task": caller,
"context": "written-in-aembedding",
},
{
"model": "deterministic",
"input": ["python written prompt"],
"task": caller,
"context": "written-in-aembedding",
},
], semantic_embedding.async_calls
async def test_native_semantic_cancellation_during_embedding_skips_the_backend(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
assert binding.kind == "native"
semantic_embedding.gate = asyncio.Event()
async def lookup() -> object:
return await binding.async_lookup(
semantic_request("cancel", "cancelled prompt")
)
task: Final = asyncio.create_task(lookup())
await semantic_embedding.entered.wait()
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
semantic_embedding.gate.set()
assert len(semantic_embedding.async_calls) == 1
assert (
await cast(RedisSemanticCache, facade.cache).async_get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class
"cancel", messages=semantic_messages("cancelled prompt")
)
is None
)
def test_redis_semantic_similarity_tag_and_threshold_boundaries(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
binding.store(semantic_request("sim", "tell me a joke"), {"answer": "haha"})
paraphrase: Final = f"tell me a joke{PARAPHRASE_MARKER}"
assert binding.lookup(semantic_request("sim", paraphrase)) == {"answer": "haha"}
assert binding.lookup(semantic_request("sim", "an unrelated question about spreadsheets")) is None
assert binding.lookup(semantic_request("other-key", "tell me a joke")) is None
strict: Final = semantic_facade(url, index, similarity_threshold=0.99)
strict_binding: Final = _CacheTestResolver(SimpleNamespace(cache=strict)).resolve()
assert strict_binding.lookup(semantic_request("sim", paraphrase)) is None
assert strict_binding.lookup(semantic_request("sim", "tell me a joke")) == {"answer": "haha"}
def test_redis_semantic_ttl_is_written_only_when_requested(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
client: Final = redis.Redis.from_url(url)
binding.store(
{**semantic_request("ttl", "ttl prompt"), "ttl_seconds": 12.0}, {"answer": 1}
)
expiring: Final = f"{index}:{semantic_entry_id('ttl prompt', 'ttl')}"
assert 0 < client.ttl(expiring) <= 12
binding.store(semantic_request("ttl-none", "untimed prompt"), {"answer": 2})
persistent: Final = f"{index}:{semantic_entry_id('untimed prompt', 'ttl-none')}"
assert client.ttl(persistent) == -1
binding.store(
{**semantic_request("ttl-fraction", "fractional prompt"), "ttl_seconds": 1.5},
{"answer": 3},
)
fractional: Final = f"{index}:{semantic_entry_id('fractional prompt', 'ttl-fraction')}"
assert client.ttl(fractional) == 2
client.close()
def test_redis_semantic_malformed_response_is_a_miss_for_both_readers(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
client: Final = redis.Redis.from_url(url)
binding.store(semantic_request("bad", "corrupt me"), {"answer": 1})
hash_key: Final = f"{index}:{semantic_entry_id('corrupt me', 'bad')}"
client.hset(hash_key, "response", b"{not json")
assert binding.lookup(semantic_request("bad", "corrupt me")) is None
assert (
cast(RedisSemanticCache, facade.cache).get_cache( # pyright: ignore[reportUnknownMemberType] # **kwargs stays unknown on the backend class
"bad", messages=semantic_messages("corrupt me")
)
is None
)
client.close()
async def test_redis_semantic_unsupported_operations_raise_not_implemented(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
with pytest.raises(NotImplementedError):
binding.lookup_batch([semantic_request("batch", "prompt one")])
with pytest.raises(NotImplementedError):
await binding.async_lookup_batch([semantic_request("batch", "prompt one")])
with pytest.raises(NotImplementedError):
await binding.async_flush()
with pytest.raises(NotImplementedError):
await binding.ping()
def test_redis_semantic_requests_without_prompt_are_noops(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
client: Final = redis.Redis.from_url(url)
binding.store(request("plain"), {"answer": 1})
assert binding.lookup(request("plain")) is None
assert semantic_embedding.calls == []
assert client.keys(f"{index}:*") == []
client.close()
def test_redis_semantic_scope_overrides_the_tag_and_isolates_entries(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
binding: Final = _CacheTestResolver(SimpleNamespace(cache=facade)).resolve()
client: Final = redis.Redis.from_url(url)
scoped: Final = {**semantic_request("scoped", "scoped prompt"), "scope": "team-a"}
binding.store(scoped, {"answer": "kept"})
hash_key: Final = f"{index}:{semantic_entry_id('scoped prompt', 'team-a')}"
assert client.hget(hash_key, "litellm_cache_key") == b"team-a"
assert binding.lookup(scoped) == {"answer": "kept"}
assert binding.lookup(semantic_request("scoped", "scoped prompt")) is None
assert binding.lookup({**scoped, "scope": "team-b"}) is None
client.close()
def test_redis_semantic_configuration_drift_falls_back_to_python(
redis_stack: tuple[str, str],
semantic_embedding: DeterministicEmbedding,
monkeypatch: pytest.MonkeyPatch,
) -> None:
url, index = redis_stack
facade: Final = semantic_facade(url, index)
resolver: Final = _CacheTestResolver(SimpleNamespace(cache=facade))
assert resolver.resolve().kind == "native"
with rebound(facade.cache, "similarity_threshold", 0.5):
assert resolver.resolve().kind == "python_callback"
with rebound(facade, "semantic_cache_scope", "end_user"):
assert resolver.resolve().kind == "python_callback"
with rebound(facade.cache, "embedding_model", "other-model"):
assert resolver.resolve().kind == "python_callback"
with rebound(facade.cache, "_index_name", "other-index"):
assert resolver.resolve().kind == "python_callback"
with rebound(facade.cache, "CACHE_KEY_FIELD_NAME", "other-field"):
assert resolver.resolve().kind == "python_callback"
def patched_embedding(self: object, prompt: str, metadata: object = None) -> list[float]:
return _semantic_embedding(prompt)
monkeypatch.setattr(RedisSemanticCache, "_get_embedding", patched_embedding)
assert resolver.resolve().kind == "python_callback"
def test_redis_semantic_handle_rejects_wrong_backends(
redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding
) -> None:
url, index = redis_stack
class CustomSemanticCache(RedisSemanticCache):
pass
with pytest.raises(TypeError, match="built-in RedisSemanticCache"):
_CacheTestHandle.redis_semantic(object())
with pytest.raises(TypeError, match="built-in RedisSemanticCache"):
_CacheTestHandle.redis_semantic(
CustomSemanticCache(
redis_url=url,
similarity_threshold=0.8,
embedding_model=SEMANTIC_EMBEDDING_MODEL,
index_name=f"{index}_subclass",
)
)
facade: Final = semantic_facade(url, index)
with pytest.raises(TypeError, match="backend types must match"):
_CacheTestHandle.redis(url)._bind_facade(facade)
subclassed_facade: Final = Cache(
type=LiteLLMCacheType.REDIS_SEMANTIC,
redis_url=url,
similarity_threshold=0.8,
redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL,
redis_semantic_cache_index_name=index,
)
subclassed_facade.cache = CustomSemanticCache( # pyright: ignore[reportAttributeAccessIssue] # facade backend slot is not declared
redis_url=url,
similarity_threshold=0.8,
embedding_model=SEMANTIC_EMBEDDING_MODEL,
index_name=index,
)
with pytest.raises(TypeError):
_CacheTestHandle.redis_semantic(
subclassed_facade.cache
)._bind_facade(subclassed_facade)
replacement_facade: Final = Cache(
type=LiteLLMCacheType.REDIS_SEMANTIC,
redis_url=url,
similarity_threshold=0.8,
redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL,
redis_semantic_cache_index_name=index,
)
with pytest.raises(TypeError, match="must be the native embedder"):
_CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade)