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

refactor(rust): align the cache crates with Python and wire every native backend (#42530)

Browse source 
* refactor(rust): align the cache crates with Python and activate every backend

The cache port had drifted: lifecycle and Redis-only operations sat on
`BaseCache`, counters were pinned to `f64`, each semantic backend defined its
own embedder and prompt handling, and only the in-memory backend could be
selected natively.

- Split `disconnect` and `test_connection` out of `BaseCache` into optional
  capabilities, implemented only where the Python class defines them, and give
  every Redis-only operation its own capability trait.
- Decouple counters from the stored value type, so one backend can serve both
  responses and counters as Python's `RedisCache` does.
- Share one `Embedder` and prompt contract in `litellm_cache::semantic`, and
  make the Redis and Valkey semantic backends generic over their codec.
- Port the Python operations that were missing: `async_refresh_ttl`,
  `async_rpush_and_trim`, `async_set_cache_pipeline_with_ttls`, the DualCache
  pipeline, sadd, bulk delete and TTL reads, and the semantic-similarity
  write-back.
- Take the HTTP client from the host pool in the GCS, S3 and Azure backends.
- Activate all nine backends through the Rust catalog, whose rules all stay
  `PYTHON_ONLY`, and route the `Cache` facade's storage calls to the native
  runtime when one is selected.
- Give every crate the same layout, move all tests to `tests/` on rstest, and
  add the shared `litellm-cache-testing` contract suite.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* fix: freeze native cache request kwargs and batch entries for type discipline

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix: declare semantic lookup methods in the native stub

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* refactor(rust): align the cache crates with Python and activate every backend

The cache port had drifted: lifecycle and Redis-only operations sat on
`BaseCache`, counters were pinned to `f64`, each semantic backend defined its
own embedder and prompt handling, and only the in-memory backend could be
selected natively.

- Split `disconnect` and `test_connection` out of `BaseCache` into optional
  capabilities, implemented only where the Python class defines them, and give
  every Redis-only operation its own capability trait.
- Decouple counters from the stored value type, so one backend can serve both
  responses and counters as Python's `RedisCache` does.
- Share one `Embedder` and prompt contract in `litellm_cache::semantic`, and
  make the Redis and Valkey semantic backends generic over their codec.
- Port the Python operations that were missing: `async_refresh_ttl`,
  `async_rpush_and_trim`, `async_set_cache_pipeline_with_ttls`, the DualCache
  pipeline, sadd, bulk delete and TTL reads, and the semantic-similarity
  write-back.
- Take the HTTP client from the host pool in the GCS, S3 and Azure backends.
- Activate all nine backends through the Rust catalog, whose rules all stay
  `PYTHON_ONLY`, and route the `Cache` facade's storage calls to the native
  runtime when one is selected.
- Give every crate the same layout, move all tests to `tests/` on rstest, and
  add the shared `litellm-cache-testing` contract suite.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

* fix: freeze native cache request kwargs and batch entries for type discipline

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix: declare semantic lookup methods in the native stub

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(rust): opt the native Messages and tokenizer suites into Rust explicitly

#42517 made the Messages, token counter and tokenizer routes Python-only, so
tests/test_litellm_rust silently exercised the Python path or failed outright.
Each suite now prepends a RUST_OPT_IN rule for its route, keeping native
coverage without changing the shipped default.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(rust): pop one at a time in the Redis 6 lpop pipeline and drop explanatory comments

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Yujong Lee <yujong@berri.ai>
Co-authored-by: Claude Opus 5 <noreply@anthropic.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-22 13:13:02 -07:00 • committed by GitHub
parent 153e5ed185
commit 4677f1028e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
122 changed files with 12665 additions and 6944 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm-rust/Cargo.lock generated
View file

@ -2727,9 +2727,13 @@ dependencies = [
"litellm-auth-types",
"litellm-cache",
"litellm-cache-response",
"litellm-cache-testing",
"reqwest 0.12.28",
"rstest",
"serde_json",
"tokio",
"url",
"wiremock",
]
[[package]]
@ -2737,6 +2741,7 @@ name = "litellm-cache-disk"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-testing",
"py_literal",
"rand 0.8.7",
"rstest",
@ -2755,8 +2760,10 @@ dependencies = [
"litellm-auth-gcp",
"litellm-auth-types",
"litellm-cache",
"litellm-cache-testing",
"percent-encoding",
"reqwest 0.12.28",
"rstest",
"serde_json",
"tokio",
"wiremock",
@ -2767,8 +2774,8 @@ name = "litellm-cache-memory"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-testing",
"rstest",
"serde_json",
"tokio",
]
@ -2776,9 +2783,10 @@ dependencies = [
name = "litellm-cache-qdrant-semantic"
version = "0.1.0"
dependencies = [
"futures-executor",
"futures-util",
"litellm-cache",
"litellm-cache-response",
"litellm-cache-testing",
"qdrant-client",
"reqwest 0.12.28",
"rstest",
@ -2797,9 +2805,11 @@ name = "litellm-cache-redis"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-testing",
"r2d2",
"redis",
"redis-test",
"rstest",
"serde_json",
"tokio",
]
@ -2810,10 +2820,10 @@ version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-response",
"r2d2",
"litellm-cache-testing",
"redis",
"redis-test",
"rstest",
"serde_json",
"sha2 0.10.9",
"tokio",
@ -2829,6 +2839,7 @@ dependencies = [
"py_literal",
"redis",
"redis-test",
"rstest",
"serde",
"serde_json",
"sha2 0.10.9",
@ -2841,22 +2852,35 @@ version = "0.1.0"
dependencies = [
"aws-credential-types",
"aws-sdk-s3",
"aws-smithy-runtime-api",
"aws-smithy-types",
"aws-types",
"futures-util",
"http 1.4.2",
"litellm-auth-aws",
"litellm-cache",
"litellm-cache-testing",
"reqwest 0.12.28",
"rstest",
"serde_json",
"tokio",
"wiremock",
]
[[package]]
name = "litellm-cache-testing"
version = "0.1.0"
dependencies = [
"litellm-cache",
]
[[package]]
name = "litellm-cache-valkey-semantic"
version = "0.1.0"
dependencies = [
"litellm-cache",
"litellm-cache-redis",
"litellm-cache-response",
"litellm-cache-testing",
"redis",
"redis-test",
"rstest",

1
litellm-rust/Cargo.toml
View file

@ -39,6 +39,7 @@ litellm-cache-disk = { path = "crates/cache-disk" }
litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" }
litellm-cache-response = { path = "crates/cache-response" }
litellm-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" }
litellm-cache-testing = { path = "crates/cache-testing" }
litellm-token-counter = { path = "crates/token-counter" }
litellm-token-counter-fast = { path = "crates/token-counter-fast" }
litellm-token-counter-huggingface = { path = "crates/token-counter-huggingface" }

5
litellm-rust/crates/cache-azure-blob/Cargo.toml
View file

@ -14,9 +14,14 @@ async-trait = "0.1"
azure_core = "1.1.0"
azure_storage_blob = "1.1.0"
futures-util.workspace = true
reqwest.workspace = true
tokio.workspace = true
url.workspace = true
[dev-dependencies]
litellm-cache-response.workspace = true
litellm-cache-testing.workspace = true
rstest.workspace = true
serde_json.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
wiremock = "0.6.5"

48
litellm-rust/crates/cache-azure-blob/src/cache.rs
View file

@ -3,7 +3,7 @@ use std::{sync::Arc, time::Duration};
use azure_core::{
credentials::TokenCredential,
error::ErrorKind,
http::{ClientOptions, RequestContent},
http::{ClientOptions, RequestContent, Transport},
};
use azure_storage_blob::{
BlobContainerClient, BlobContainerClientOptions,
@ -11,13 +11,12 @@ use azure_storage_blob::{
};
use futures_util::{TryStreamExt, future::try_join_all};
use litellm_cache::{
BaseCache, BatchCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
ExactCacheContext, FlushCache,
BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache,
};
use tokio::runtime::Handle;
use url::Url;
use crate::credential::AzureBlobCredential;
use crate::{credential::AzureBlobCredential, transport::ReqwestTransport};
pub struct AzureBlobCache<C> {
container: BlobContainerClient,
@ -28,9 +27,11 @@ pub struct AzureBlobCache<C> {
}
impl<C: CacheCodec> AzureBlobCache<C> {
/// `http` is the host's pooled client; the SDK sends every request through it.
pub async fn connect(
account_url: &str,
container: &str,
http: reqwest::Client,
codec: C,
runtime: Handle,
) -> Result<Self, Error> {
@ -38,7 +39,10 @@ impl<C: CacheCodec> AzureBlobCache<C> {
account_url,
container,
Some(Arc::new(AzureBlobCredential::default())),
ClientOptions::default(),
ClientOptions {
transport: Some(Transport::new(Arc::new(ReqwestTransport(http)))),
..ClientOptions::default()
},
codec,
runtime,
)
@ -152,7 +156,11 @@ impl<C: CacheCodec> AzureBlobCache<C> {
}
fn block_on<T>(&self, future: impl Future<Output = T>) -> T {
self.runtime.block_on(future)
if Handle::try_current().is_ok() {
tokio::task::block_in_place(|| self.runtime.block_on(future))
} else {
self.runtime.block_on(future)
}
}
}
@ -217,25 +225,6 @@ impl<C: CacheCodec> BaseCache for AzureBlobCache<C> {
.await
.map(drop)
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Ok(match self.container.get_properties(None).await {
Ok(_) => CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Azure Blob cache connection test successful".into(),
error: None,
},
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Azure Blob connection failed: {error}"),
error: Some(error.to_string()),
},
})
}
}
impl<C: CacheCodec> BatchCache for AzureBlobCache<C> {}
@ -250,5 +239,10 @@ impl<C: CacheCodec> FlushCache for AzureBlobCache<C> {
}
}
#[cfg(test)]
mod tests;
impl<C: CacheCodec> DisconnectCache for AzureBlobCache<C> {
/// Python closes its two SDK clients; the Rust clients hold no connection of their own
/// (the pooled transport belongs to the host), so there is nothing to release.
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
}

746
litellm-rust/crates/cache-azure-blob/src/cache/tests.rs vendored
View file

@ -1,746 +0,0 @@
use std::{
collections::BTreeMap,
sync::{Arc, Mutex},
time::Duration,
};
use azure_core::http::{
AsyncRawResponse, Body, ClientOptions, HttpClient, Method, Request, StatusCode, Transport,
headers::{HeaderName, Headers},
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionStatus, Error, ExactCacheContext, FlushCache,
};
use litellm_cache_response::{
CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
ResponseCacheRequest, cache_key,
};
use serde_json::json;
use tokio::runtime::Runtime;
use super::AzureBlobCache;
const ACCOUNT_URL: &str = "https://example.blob.core.windows.net";
const CONTAINER: &str = "litellm-cache";
const IF_NONE_MATCH: HeaderName = HeaderName::from_static("if-none-match");
const ERROR_CODE: HeaderName = HeaderName::from_static("x-ms-error-code");
#[derive(Clone, Debug, PartialEq, Eq)]
struct RecordedRequest {
method: Method,
path: String,
query: String,
if_none_match: Option<String>,
}
#[derive(Default)]
struct FakeState {
container_exists: bool,
blobs: BTreeMap<String, Vec<u8>>,
requests: Vec<RecordedRequest>,
failing: bool,
precondition_conflicts: bool,
}
#[derive(Clone, Default)]
struct FakeBlobService {
state: Arc<Mutex<FakeState>>,
}
impl std::fmt::Debug for FakeBlobService {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("FakeBlobService")
}
}
impl FakeBlobService {
fn with_existing_container() -> Self {
let service = Self::default();
service.state.lock().unwrap().container_exists = true;
service
}
fn blob(&self, name: &str) -> Option<Vec<u8>> {
self.state.lock().unwrap().blobs.get(name).cloned()
}
fn blob_names(&self) -> Vec<String> {
self.state.lock().unwrap().blobs.keys().cloned().collect()
}
fn seed_blob(&self, name: &str, bytes: &[u8]) {
self.state
.lock()
.unwrap()
.blobs
.insert(name.to_string(), bytes.to_vec());
}
fn set_failing(&self, failing: bool) {
self.state.lock().unwrap().failing = failing;
}
fn set_precondition_conflicts(&self, enabled: bool) {
self.state.lock().unwrap().precondition_conflicts = enabled;
}
fn requests(&self) -> Vec<RecordedRequest> {
self.state.lock().unwrap().requests.clone()
}
fn container_exists(&self) -> bool {
self.state.lock().unwrap().container_exists
}
fn respond(status: StatusCode, error_code: Option<&str>, body: Vec<u8>) -> AsyncRawResponse {
let mut headers = Headers::new();
if let Some(code) = error_code {
headers.insert(ERROR_CODE, code.to_string());
}
AsyncRawResponse::from_bytes(status, headers, body)
}
fn list_body(state: &FakeState) -> Vec<u8> {
let mut xml = String::from(
r#"<?xml version="1.0" encoding="utf-8"?><EnumerationResults ServiceEndpoint="https://example.blob.core.windows.net/" ContainerName="litellm-cache"><Blobs>"#,
);
for name in state.blobs.keys() {
xml.push_str(&format!(
"<Blob><Name>{name}</Name><Properties><BlobType>BlockBlob</BlobType></Properties></Blob>"
));
}
xml.push_str("</Blobs><NextMarker /></EnumerationResults>");
xml.into_bytes()
}
}
#[async_trait::async_trait]
impl HttpClient for FakeBlobService {
async fn execute_request(&self, request: &Request) -> azure_core::Result<AsyncRawResponse> {
let mut state = self.state.lock().unwrap();
let path = request.url().path().to_string();
let query = request.url().query().unwrap_or_default().to_string();
let if_none_match = request
.headers()
.get_optional_str(&IF_NONE_MATCH)
.map(str::to_owned);
state.requests.push(RecordedRequest {
method: request.method(),
path: path.clone(),
query: query.clone(),
if_none_match: if_none_match.clone(),
});
if state.failing {
return Ok(Self::respond(
StatusCode::Forbidden,
Some("AuthorizationFailure"),
Vec::new(),
));
}
let container_path = format!("/{CONTAINER}");
let blob_name = path
.strip_prefix(&format!("{container_path}/"))
.map(str::to_owned);
let is_container = path == container_path && query.contains("restype=container");
let response = match (request.method(), is_container, blob_name) {
(Method::Put, true, None) if state.container_exists => Self::respond(
StatusCode::Conflict,
Some("ContainerAlreadyExists"),
Vec::new(),
),
(Method::Put, true, None) => {
state.container_exists = true;
Self::respond(StatusCode::Created, None, Vec::new())
}
(Method::Get, true, None) if query.contains("comp=list") => {
Self::respond(StatusCode::Ok, None, Self::list_body(&state))
}
(Method::Get, true, None) if state.container_exists => {
Self::respond(StatusCode::Ok, None, Vec::new())
}
(Method::Get, true, None) => {
Self::respond(StatusCode::NotFound, Some("ContainerNotFound"), Vec::new())
}
(Method::Put, false, Some(name)) => {
if if_none_match.as_deref() == Some("*") && state.blobs.contains_key(&name) {
if state.precondition_conflicts {
Self::respond(
StatusCode::PreconditionFailed,
Some("ConditionNotMet"),
Vec::new(),
)
} else {
Self::respond(StatusCode::Conflict, Some("BlobAlreadyExists"), Vec::new())
}
} else {
let bytes = match request.body() {
Body::Bytes(bytes) => bytes.to_vec(),
Body::SeekableStream(_) => panic!("unexpected streaming upload"),
};
state.blobs.insert(name, bytes);
Self::respond(StatusCode::Created, None, Vec::new())
}
}
(Method::Get, false, Some(name)) => match state.blobs.get(&name) {
Some(bytes) => Self::respond(StatusCode::Ok, None, bytes.clone()),
None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
},
(Method::Delete, false, Some(name)) => match state.blobs.remove(&name) {
Some(_) => Self::respond(StatusCode::Accepted, None, Vec::new()),
None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
},
(method, _, _) => panic!("unexpected request {method:?} {path}?{query}"),
};
Ok(response)
}
}
struct Fixture {
runtime: Runtime,
service: FakeBlobService,
cache: Arc<AzureBlobCache<ResponseCacheCodec>>,
}
impl Fixture {
fn new(service: FakeBlobService) -> Self {
let runtime = Runtime::new().unwrap();
let cache = runtime
.block_on(Self::connect(&service, runtime.handle().clone()))
.unwrap();
Self {
runtime,
service,
cache: Arc::new(cache),
}
}
async fn connect(
service: &FakeBlobService,
handle: tokio::runtime::Handle,
) -> Result<AzureBlobCache<ResponseCacheCodec>, Error> {
AzureBlobCache::connect_with_options(
ACCOUNT_URL,
CONTAINER,
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(service.clone()))),
..ClientOptions::default()
},
ResponseCacheCodec,
handle,
)
.await
}
fn response_cache(&self) -> ResponseCache<AzureBlobCache<ResponseCacheCodec>> {
ResponseCache::new(self.cache.clone())
}
fn stored_json(&self, key: &str) -> serde_json::Value {
serde_json::from_slice(&self.service.blob(key).expect("blob should exist")).unwrap()
}
}
fn request(model: &str) -> ResponseCacheRequest {
ResponseCacheRequest::new(CacheKeyInput {
fields: vec![CacheKeyField {
name: "model".into(),
value: Some(model.into()),
api_parameter: true,
internal_parameter: false,
}],
preset: None,
namespace: None,
include_provider_parameters: false,
})
}
fn now() -> Duration {
Duration::from_secs(1_700_000_000)
}
fn entry(value: serde_json::Value) -> CacheEntry {
CacheEntry {
timestamp: Some(1_700_000_000.5),
response: value,
}
}
fn no_ttl() -> ExactCacheContext {
ExactCacheContext::default()
}
fn with_ttl(seconds: u64) -> ExactCacheContext {
ExactCacheContext {
ttl: Some(Duration::from_secs(seconds)),
}
}
#[test]
fn connect_creates_the_container_once() {
let fixture = Fixture::new(FakeBlobService::default());
assert!(fixture.service.container_exists());
assert_eq!(
fixture.service.requests(),
vec![RecordedRequest {
method: Method::Put,
path: format!("/{CONTAINER}"),
query: "restype=container".into(),
if_none_match: None,
}]
);
assert_eq!(fixture.cache.account_url(), ACCOUNT_URL);
assert_eq!(fixture.cache.container_name(), CONTAINER);
}
#[test]
fn connect_accepts_an_existing_container() {
let fixture = Fixture::new(FakeBlobService::with_existing_container());
assert!(fixture.service.container_exists());
assert_eq!(fixture.service.requests().len(), 1);
}
#[test]
fn connect_accepts_account_urls_with_trailing_slash() {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
let cache = runtime
.block_on(AzureBlobCache::connect_with_options(
"https://example.blob.core.windows.net/",
CONTAINER,
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(service.clone()))),
..ClientOptions::default()
},
ResponseCacheCodec,
runtime.handle().clone(),
))
.unwrap();
assert_eq!(service.requests()[0].path, format!("/{CONTAINER}"));
assert_eq!(cache.account_url(), "https://example.blob.core.windows.net");
}
#[test]
fn connect_keeps_account_url_query_parameters_on_the_container_path() {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
runtime
.block_on(AzureBlobCache::connect_with_options(
"https://example.blob.core.windows.net/?sv=2024-01-01&sig=abc",
CONTAINER,
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(service.clone()))),
..ClientOptions::default()
},
ResponseCacheCodec,
runtime.handle().clone(),
))
.unwrap();
let create = &service.requests()[0];
assert_eq!(create.path, format!("/{CONTAINER}"));
assert!(create.query.contains("sig=abc"));
}
#[test]
fn connect_surfaces_service_failures() {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
service.set_failing(true);
let result = runtime.block_on(Fixture::connect(&service, runtime.handle().clone()));
assert!(matches!(result, Err(Error::Unavailable)));
}
#[test]
fn sync_set_and_get_round_trip_python_json_shape() {
let fixture = Fixture::new(FakeBlobService::default());
let value = entry(json!({"choices": [{"message": {"content": "héllo 🌍"}}]}));
fixture
.cache
.set_cache("key-1", value.clone(), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key-1"),
json!({
"timestamp": 1_700_000_000.5,
"response": {"choices": [{"message": {"content": "héllo 🌍"}}]}
})
);
assert_eq!(
fixture.cache.get_cache("key-1", &no_ttl()).unwrap(),
Some(value)
);
}
#[test]
fn sync_set_does_not_overwrite_an_existing_blob() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("key", entry(json!({"v": "first"})), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("key", entry(json!({"v": "second"})), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "first"})
);
let uploads: Vec<_> = fixture
.service
.requests()
.into_iter()
.filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
.collect();
assert_eq!(uploads.len(), 2);
assert!(
uploads
.iter()
.all(|request| request.if_none_match.as_deref() == Some("*"))
);
}
#[test]
fn sync_set_treats_a_precondition_conflict_as_an_existing_blob() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.set_precondition_conflicts(true);
fixture
.cache
.set_cache("key", entry(json!({"v": "first"})), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("key", entry(json!({"v": "second"})), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "first"})
);
}
#[test]
fn async_set_overwrites_an_existing_blob() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.runtime.block_on(async {
fixture
.cache
.async_set_cache("key", entry(json!({"v": "first"})), no_ttl())
.await
.unwrap();
fixture
.cache
.async_set_cache("key", entry(json!({"v": "second"})), no_ttl())
.await
.unwrap();
assert_eq!(
fixture
.cache
.async_get_cache("key", &no_ttl())
.await
.unwrap(),
Some(entry(json!({"v": "second"})))
);
});
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "second"})
);
assert!(
fixture
.service
.requests()
.iter()
.filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
.all(|request| request.if_none_match.is_none())
);
}
#[test]
fn missing_blobs_are_misses() {
let fixture = Fixture::new(FakeBlobService::default());
assert_eq!(fixture.cache.get_cache("absent", &no_ttl()).unwrap(), None);
assert_eq!(
fixture
.runtime
.block_on(fixture.cache.async_get_cache("absent", &no_ttl()))
.unwrap(),
None
);
}
#[test]
fn ttl_is_ignored_and_entries_never_expire() {
let fixture = Fixture::new(FakeBlobService::default());
assert_eq!(fixture.cache.get_ttl(&with_ttl(1)), None);
assert_eq!(fixture.cache.get_ttl(&no_ttl()), None);
fixture
.cache
.set_cache("key", entry(json!("value")), &with_ttl(1))
.unwrap();
std::thread::sleep(Duration::from_millis(1100));
assert_eq!(
fixture.cache.get_cache("key", &with_ttl(1)).unwrap(),
Some(entry(json!("value")))
);
assert!(
fixture
.service
.requests()
.iter()
.all(|request| !request.query.contains("expiry"))
);
}
#[test]
fn malformed_blobs_are_invalid_entries_and_response_cache_misses() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.seed_blob("broken-json", b"{not json");
fixture
.service
.seed_blob("broken-utf8", &[0xff, 0xfe, 0x22]);
fixture
.service
.seed_blob("wrong-shape", br#"{"timestamp": "yesterday"}"#);
for key in ["broken-json", "broken-utf8", "wrong-shape"] {
assert!(matches!(
fixture.cache.get_cache(key, &no_ttl()),
Err(Error::InvalidEntry)
));
}
let response_cache = fixture.response_cache();
let broken = request("broken");
fixture
.service
.seed_blob(&cache_key(&broken.key), b"{not json");
assert_eq!(response_cache.lookup(&broken, now()).unwrap(), None);
assert_eq!(
fixture
.runtime
.block_on(response_cache.async_lookup(&broken, now()))
.unwrap(),
None
);
}
#[test]
fn batch_get_preserves_order_and_marks_misses_and_invalid_entries() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("a", entry(json!("A")), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("c", entry(json!("C")), &no_ttl())
.unwrap();
fixture.service.seed_blob("bad", b"nope");
let keys = ["c", "missing", "a", "bad"].map(String::from);
let sync = fixture.cache.batch_get_cache(&keys, &no_ttl()).unwrap();
assert_eq!(
sync,
vec![
BatchEntry::Hit(entry(json!("C"))),
BatchEntry::Miss,
BatchEntry::Hit(entry(json!("A"))),
BatchEntry::Invalid,
]
);
let asynchronous = fixture
.runtime
.block_on(fixture.cache.async_batch_get_cache(keys.to_vec(), no_ttl()))
.unwrap();
assert_eq!(asynchronous, sync);
let response_cache = fixture.response_cache();
let requests = [request("hit"), request("missing"), request("bad")];
response_cache
.store(&requests[0], json!("HIT"), now())
.unwrap();
fixture
.service
.seed_blob(&cache_key(&requests[2].key), b"nope");
let hits = response_cache.lookup_batch(&requests, now()).unwrap();
assert_eq!(hits.values, vec![Some(json!("HIT")), None, None]);
assert_eq!(hits.missing_indices, vec![1, 2]);
let async_hits = fixture
.runtime
.block_on(response_cache.async_lookup_batch(&requests, now()))
.unwrap();
assert_eq!(async_hits.values, hits.values);
}
#[test]
fn async_pipeline_writes_every_entry_with_overwrite() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.seed_blob("k2", b"stale");
fixture
.runtime
.block_on(fixture.cache.async_set_cache_pipeline(
vec![
("k1".into(), entry(json!({"n": 1}))),
("k2".into(), entry(json!({"n": 2}))),
("k3".into(), entry(json!({"n": 3}))),
],
with_ttl(30),
))
.unwrap();
assert_eq!(fixture.service.blob_names(), ["k1", "k2", "k3"]);
assert_eq!(fixture.stored_json("k2")["response"], json!({"n": 2}));
}
#[test]
fn flush_deletes_every_blob_in_the_container() {
let fixture = Fixture::new(FakeBlobService::default());
for key in ["x", "y", "z"] {
fixture
.cache
.set_cache(key, entry(json!(key)), &no_ttl())
.unwrap();
}
fixture.cache.flush_cache().unwrap();
assert!(fixture.service.blob_names().is_empty());
assert!(fixture.service.container_exists());
fixture
.cache
.set_cache("again", entry(json!(1)), &no_ttl())
.unwrap();
fixture
.runtime
.block_on(fixture.cache.async_flush_cache())
.unwrap();
assert!(fixture.service.blob_names().is_empty());
}
#[test]
fn service_failures_map_to_unavailable() {
let fixture = Fixture::new(FakeBlobService::default());
fixture.service.set_failing(true);
assert!(matches!(
fixture.cache.get_cache("key", &no_ttl()),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.cache.flush_cache(),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.runtime.block_on(
fixture
.cache
.async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl())
),
Err(Error::Unavailable)
));
}
#[test]
fn test_connection_reports_container_reachability() {
let fixture = Fixture::new(FakeBlobService::default());
let ok = fixture
.runtime
.block_on(fixture.cache.test_connection())
.unwrap();
assert_eq!(ok.status, CacheConnectionStatus::Success);
assert!(ok.error.is_none());
fixture.service.set_failing(true);
let failed = fixture
.runtime
.block_on(fixture.cache.test_connection())
.unwrap();
assert_eq!(failed.status, CacheConnectionStatus::Failed);
assert!(failed.error.is_some());
}
#[test]
fn disconnect_is_idempotent_and_keeps_data() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("key", entry(json!(1)), &no_ttl())
.unwrap();
fixture.runtime.block_on(async {
fixture.cache.disconnect().await.unwrap();
fixture.cache.disconnect().await.unwrap();
});
assert_eq!(
fixture.cache.get_cache("key", &no_ttl()).unwrap(),
Some(entry(json!(1)))
);
}
#[test]
fn response_cache_stores_and_reads_through_the_backend() {
let fixture = Fixture::new(FakeBlobService::default());
let response_cache = fixture.response_cache();
let mut request = request("gpt");
request.context = with_ttl(60);
let response = json!({"id": "chatcmpl-1"});
response_cache
.store(&request, response.clone(), now())
.unwrap();
assert_eq!(
fixture.stored_json(&cache_key(&request.key)),
json!({"timestamp": 1_700_000_000.0, "response": {"id": "chatcmpl-1"}})
);
assert_eq!(
response_cache
.lookup(&request, now() + Duration::from_secs(3600))
.unwrap(),
Some(response.clone())
);
assert_eq!(
fixture
.runtime
.block_on(response_cache.async_lookup(&request, now() + Duration::from_secs(3600)))
.unwrap(),
Some(response.clone())
);
fixture.runtime.block_on(async {
response_cache
.async_store(&request, json!("replaced"), now())
.await
.unwrap();
assert_eq!(
response_cache.async_lookup(&request, now()).await.unwrap(),
Some(json!("replaced"))
);
response_cache.async_flush().await.unwrap();
assert_eq!(
response_cache.async_lookup(&request, now()).await.unwrap(),
None
);
});
}
#[test]
fn non_object_responses_are_written_serialized_like_python() {
let fixture = Fixture::new(FakeBlobService::default());
fixture
.cache
.set_cache("s", entry(json!("plain")), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("s"),
json!({"timestamp": 1_700_000_000.5, "response": "\"plain\""})
);
assert_eq!(
fixture.cache.get_cache("s", &no_ttl()).unwrap(),
Some(entry(json!("plain")))
);
}

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

@ -1,5 +1,7 @@
mod cache;
mod credential;
mod transport;
pub use cache::AzureBlobCache;
pub use credential::AzureBlobCredential;
pub use transport::ReqwestTransport;

0
litellm-rust/crates/cache-azure-blob/src/tests.rs
View file

49
litellm-rust/crates/cache-azure-blob/src/transport.rs Normal file
View file

@ -0,0 +1,49 @@
use azure_core::{
error::ErrorKind,
http::{
AsyncRawResponse, Body, HttpClient, Request,
headers::{HeaderName, HeaderValue, Headers},
},
};
use futures_util::TryStreamExt;
#[derive(Debug)]
pub struct ReqwestTransport(pub reqwest::Client);
#[async_trait::async_trait]
impl HttpClient for ReqwestTransport {
async fn execute_request(&self, request: &Request) -> azure_core::Result<AsyncRawResponse> {
let method = reqwest::Method::from_bytes(request.method().as_ref().as_bytes())
.map_err(|error| azure_core::Error::new(ErrorKind::Other, error))?;
let mut outgoing = self.0.request(method, request.url().as_str());
for (name, value) in request.headers().iter() {
outgoing = outgoing.header(name.as_str(), value.as_str());
}
let outgoing = match request.body().clone() {
Body::Bytes(bytes) => outgoing.body(bytes),
Body::SeekableStream(stream) => outgoing.body(reqwest::Body::wrap_stream(stream)),
};
let response = outgoing.send().await.map_err(|error| {
let kind = if error.is_connect() {
ErrorKind::Connection
} else {
ErrorKind::Io
};
azure_core::Error::new(kind, error)
})?;
let status = response.status().as_u16().into();
let mut headers = Headers::new();
for (name, value) in response.headers() {
if let Ok(value) = value.to_str() {
headers.insert(
HeaderName::from(name.as_str().to_owned()),
HeaderValue::from(value.to_owned()),
);
}
}
let body = response
.bytes_stream()
.map_err(|error| azure_core::Error::new(ErrorKind::Io, error));
Ok(AsyncRawResponse::new(status, headers, Box::pin(body)))
}
}

494
litellm-rust/crates/cache-azure-blob/tests/cache.rs Normal file
View file

@ -0,0 +1,494 @@
mod support;
use std::{sync::Arc, time::Duration};
use azure_core::http::Method;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, DisconnectCache, Error, ExactCacheContext, FlushCache,
};
use litellm_cache_azure_blob::AzureBlobCache;
use litellm_cache_response::{
CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
ResponseCacheRequest, cache_key,
};
use rstest::{fixture, rstest};
use serde_json::json;
use support::{ACCOUNT_URL, CONTAINER, FakeBlobService, RecordedRequest};
use tokio::runtime::Runtime;
type Fixture = support::Fixture<ResponseCacheCodec>;
#[fixture]
fn fixture() -> Fixture {
Fixture::new(FakeBlobService::default(), ResponseCacheCodec)
}
fn response_cache(fixture: &Fixture) -> ResponseCache<AzureBlobCache<ResponseCacheCodec>> {
ResponseCache::new(fixture.cache.clone())
}
fn request(model: &str) -> ResponseCacheRequest {
ResponseCacheRequest::new(CacheKeyInput {
fields: vec![CacheKeyField {
name: "model".into(),
value: Some(model.into()),
api_parameter: true,
internal_parameter: false,
}],
preset: None,
namespace: None,
include_provider_parameters: false,
})
}
fn now() -> Duration {
Duration::from_secs(1_700_000_000)
}
fn entry(value: serde_json::Value) -> CacheEntry {
CacheEntry {
timestamp: Some(1_700_000_000.5),
response: value,
}
}
fn no_ttl() -> ExactCacheContext {
ExactCacheContext::default()
}
fn with_ttl(seconds: u64) -> ExactCacheContext {
ExactCacheContext {
ttl: Some(Duration::from_secs(seconds)),
}
}
fn connect_to(account_url: &str) -> (FakeBlobService, AzureBlobCache<ResponseCacheCodec>) {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
let cache = runtime
.block_on(support::connect(
&service,
account_url,
ResponseCacheCodec,
runtime.handle().clone(),
))
.unwrap();
(service, cache)
}
#[rstest]
fn connect_creates_the_container_once(fixture: Fixture) {
assert!(fixture.service.container_exists());
assert_eq!(
fixture.service.requests(),
vec![RecordedRequest {
method: Method::Put,
path: format!("/{CONTAINER}"),
query: "restype=container".into(),
if_none_match: None,
}]
);
assert_eq!(fixture.cache.account_url(), ACCOUNT_URL);
assert_eq!(fixture.cache.container_name(), CONTAINER);
}
#[rstest]
fn connect_accepts_an_existing_container() {
let fixture = Fixture::new(
FakeBlobService::with_existing_container(),
ResponseCacheCodec,
);
assert!(fixture.service.container_exists());
assert_eq!(fixture.service.requests().len(), 1);
}
#[rstest]
fn connect_accepts_account_urls_with_trailing_slash() {
let (service, cache) = connect_to("https://example.blob.core.windows.net/");
assert_eq!(service.requests()[0].path, format!("/{CONTAINER}"));
assert_eq!(cache.account_url(), "https://example.blob.core.windows.net");
}
#[rstest]
fn connect_keeps_account_url_query_parameters_on_the_container_path() {
let (service, _) = connect_to("https://example.blob.core.windows.net/?sv=2024-01-01&sig=abc");
let create = &service.requests()[0];
assert_eq!(create.path, format!("/{CONTAINER}"));
assert!(create.query.contains("sig=abc"));
}
#[rstest]
fn connect_surfaces_service_failures() {
let runtime = Runtime::new().unwrap();
let service = FakeBlobService::default();
service.set_failing(true);
let result = runtime.block_on(support::connect(
&service,
ACCOUNT_URL,
ResponseCacheCodec,
runtime.handle().clone(),
));
assert!(matches!(result, Err(Error::Unavailable)));
}
#[rstest]
fn sync_set_and_get_round_trip_python_json_shape(fixture: Fixture) {
let value = entry(json!({"choices": [{"message": {"content": "héllo 🌍"}}]}));
fixture
.cache
.set_cache("key-1", value.clone(), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key-1"),
json!({
"timestamp": 1_700_000_000.5,
"response": {"choices": [{"message": {"content": "héllo 🌍"}}]}
})
);
assert_eq!(
fixture.cache.get_cache("key-1", &no_ttl()).unwrap(),
Some(value)
);
}
#[rstest]
#[case::blob_already_exists(false)]
#[case::precondition_conflict(true)]
fn sync_set_does_not_overwrite_an_existing_blob(fixture: Fixture, #[case] precondition: bool) {
fixture.service.set_precondition_conflicts(precondition);
fixture
.cache
.set_cache("key", entry(json!({"v": "first"})), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("key", entry(json!({"v": "second"})), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "first"})
);
let uploads: Vec<_> = fixture
.service
.requests()
.into_iter()
.filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
.collect();
assert_eq!(uploads.len(), 2);
assert!(
uploads
.iter()
.all(|request| request.if_none_match.as_deref() == Some("*"))
);
}
#[rstest]
fn async_set_overwrites_an_existing_blob(fixture: Fixture) {
fixture.runtime.block_on(async {
fixture
.cache
.async_set_cache("key", entry(json!({"v": "first"})), no_ttl())
.await
.unwrap();
fixture
.cache
.async_set_cache("key", entry(json!({"v": "second"})), no_ttl())
.await
.unwrap();
assert_eq!(
fixture
.cache
.async_get_cache("key", &no_ttl())
.await
.unwrap(),
Some(entry(json!({"v": "second"})))
);
});
assert_eq!(
fixture.stored_json("key")["response"],
json!({"v": "second"})
);
assert!(
fixture
.service
.requests()
.iter()
.filter(|request| request.method == Method::Put && request.path.ends_with("/key"))
.all(|request| request.if_none_match.is_none())
);
}
#[rstest]
fn missing_blobs_are_misses(fixture: Fixture) {
assert_eq!(fixture.cache.get_cache("absent", &no_ttl()).unwrap(), None);
assert_eq!(
fixture
.runtime
.block_on(fixture.cache.async_get_cache("absent", &no_ttl()))
.unwrap(),
None
);
}
#[rstest]
fn ttl_is_ignored_and_entries_never_expire(fixture: Fixture) {
assert_eq!(fixture.cache.get_ttl(&with_ttl(1)), None);
assert_eq!(fixture.cache.get_ttl(&no_ttl()), None);
fixture
.cache
.set_cache("key", entry(json!("value")), &with_ttl(1))
.unwrap();
std::thread::sleep(Duration::from_millis(1100));
assert_eq!(
fixture.cache.get_cache("key", &with_ttl(1)).unwrap(),
Some(entry(json!("value")))
);
assert!(
fixture
.service
.requests()
.iter()
.all(|request| !request.query.contains("expiry"))
);
}
#[rstest]
#[case::broken_json("broken-json", b"{not json".as_slice())]
#[case::broken_utf8("broken-utf8", &[0xff, 0xfe, 0x22])]
#[case::wrong_shape("wrong-shape", br#"{"timestamp": "yesterday"}"#.as_slice())]
fn malformed_blobs_are_invalid_entries(fixture: Fixture, #[case] key: &str, #[case] bytes: &[u8]) {
fixture.service.seed_blob(key, bytes);
assert!(matches!(
fixture.cache.get_cache(key, &no_ttl()),
Err(Error::InvalidEntry)
));
}
#[rstest]
fn malformed_blobs_are_response_cache_misses(fixture: Fixture) {
let response_cache = response_cache(&fixture);
let broken = request("broken");
fixture
.service
.seed_blob(&cache_key(&broken.key), b"{not json");
assert_eq!(response_cache.lookup(&broken, now()).unwrap(), None);
assert_eq!(
fixture
.runtime
.block_on(response_cache.async_lookup(&broken, now()))
.unwrap(),
None
);
}
#[rstest]
fn batch_get_preserves_order_and_marks_misses_and_invalid_entries(fixture: Fixture) {
fixture
.cache
.set_cache("a", entry(json!("A")), &no_ttl())
.unwrap();
fixture
.cache
.set_cache("c", entry(json!("C")), &no_ttl())
.unwrap();
fixture.service.seed_blob("bad", b"nope");
let keys = ["c", "missing", "a", "bad"].map(String::from);
let sync = fixture.cache.batch_get_cache(&keys, &no_ttl()).unwrap();
assert_eq!(
sync,
vec![
BatchEntry::Hit(entry(json!("C"))),
BatchEntry::Miss,
BatchEntry::Hit(entry(json!("A"))),
BatchEntry::Invalid,
]
);
let asynchronous = fixture
.runtime
.block_on(fixture.cache.async_batch_get_cache(keys.to_vec(), no_ttl()))
.unwrap();
assert_eq!(asynchronous, sync);
let response_cache = response_cache(&fixture);
let requests = [request("hit"), request("missing"), request("bad")];
response_cache
.store(&requests[0], json!("HIT"), now())
.unwrap();
fixture
.service
.seed_blob(&cache_key(&requests[2].key), b"nope");
let hits = response_cache.lookup_batch(&requests, now()).unwrap();
assert_eq!(hits.values, vec![Some(json!("HIT")), None, None]);
assert_eq!(hits.missing_indices, vec![1, 2]);
let async_hits = fixture
.runtime
.block_on(response_cache.async_lookup_batch(&requests, now()))
.unwrap();
assert_eq!(async_hits.values, hits.values);
}
#[rstest]
fn async_pipeline_writes_every_entry_with_overwrite(fixture: Fixture) {
fixture.service.seed_blob("k2", b"stale");
fixture
.runtime
.block_on(fixture.cache.async_set_cache_pipeline(
vec![
("k1".into(), entry(json!({"n": 1}))),
("k2".into(), entry(json!({"n": 2}))),
("k3".into(), entry(json!({"n": 3}))),
],
with_ttl(30),
))
.unwrap();
assert_eq!(fixture.service.blob_names(), ["k1", "k2", "k3"]);
assert_eq!(fixture.stored_json("k2")["response"], json!({"n": 2}));
}
#[rstest]
fn flush_deletes_every_blob_in_the_container(fixture: Fixture) {
for key in ["x", "y", "z"] {
fixture
.cache
.set_cache(key, entry(json!(key)), &no_ttl())
.unwrap();
}
fixture.cache.flush_cache().unwrap();
assert!(fixture.service.blob_names().is_empty());
assert!(fixture.service.container_exists());
fixture
.cache
.set_cache("again", entry(json!(1)), &no_ttl())
.unwrap();
fixture
.runtime
.block_on(fixture.cache.async_flush_cache())
.unwrap();
assert!(fixture.service.blob_names().is_empty());
}
#[rstest]
fn service_failures_map_to_unavailable(fixture: Fixture) {
fixture.service.set_failing(true);
assert!(matches!(
fixture.cache.get_cache("key", &no_ttl()),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.cache.set_cache("key", entry(json!(1)), &no_ttl()),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.cache.flush_cache(),
Err(Error::Unavailable)
));
assert!(matches!(
fixture.runtime.block_on(
fixture
.cache
.async_set_cache_pipeline(vec![("k".into(), entry(json!(1)))], no_ttl())
),
Err(Error::Unavailable)
));
}
#[rstest]
fn disconnect_is_idempotent_and_keeps_data(fixture: Fixture) {
fixture
.cache
.set_cache("key", entry(json!(1)), &no_ttl())
.unwrap();
fixture.runtime.block_on(async {
fixture.cache.disconnect().await.unwrap();
fixture.cache.disconnect().await.unwrap();
});
assert_eq!(
fixture.cache.get_cache("key", &no_ttl()).unwrap(),
Some(entry(json!(1)))
);
}
#[rstest]
fn response_cache_stores_and_reads_through_the_backend(fixture: Fixture) {
let response_cache = response_cache(&fixture);
let mut request = request("gpt");
request.context = with_ttl(60);
let response = json!({"id": "chatcmpl-1"});
response_cache
.store(&request, response.clone(), now())
.unwrap();
assert_eq!(
fixture.stored_json(&cache_key(&request.key)),
json!({"timestamp": 1_700_000_000.0, "response": {"id": "chatcmpl-1"}})
);
assert_eq!(
response_cache
.lookup(&request, now() + Duration::from_secs(3600))
.unwrap(),
Some(response.clone())
);
assert_eq!(
fixture
.runtime
.block_on(response_cache.async_lookup(&request, now() + Duration::from_secs(3600)))
.unwrap(),
Some(response.clone())
);
fixture.runtime.block_on(async {
response_cache
.async_store(&request, json!("replaced"), now())
.await
.unwrap();
assert_eq!(
response_cache.async_lookup(&request, now()).await.unwrap(),
Some(json!("replaced"))
);
response_cache.async_flush().await.unwrap();
assert_eq!(
response_cache.async_lookup(&request, now()).await.unwrap(),
None
);
});
}
#[rstest]
fn non_object_responses_are_written_serialized_like_python(fixture: Fixture) {
fixture
.cache
.set_cache("s", entry(json!("plain")), &no_ttl())
.unwrap();
assert_eq!(
fixture.stored_json("s"),
json!({"timestamp": 1_700_000_000.5, "response": "\"plain\""})
);
assert_eq!(
fixture.cache.get_cache("s", &no_ttl()).unwrap(),
Some(entry(json!("plain")))
);
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn sync_methods_block_inside_a_multi_thread_runtime() {
let service = FakeBlobService::default();
let cache = support::connect(
&service,
ACCOUNT_URL,
ResponseCacheCodec,
tokio::runtime::Handle::current(),
)
.await
.map(Arc::new)
.unwrap();
cache.set_cache("key", entry(json!(1)), &no_ttl()).unwrap();
assert_eq!(
cache.get_cache("key", &no_ttl()).unwrap(),
Some(entry(json!(1)))
);
}

81
litellm-rust/crates/cache-azure-blob/tests/contract.rs Normal file
View file

@ -0,0 +1,81 @@
mod support;
use litellm_cache::{ExactCacheContext, JsonCodec};
use litellm_cache_azure_blob::AzureBlobCache;
use litellm_cache_testing as contract;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::{ACCOUNT_URL, FakeBlobService};
use tokio::runtime::Handle;
#[fixture]
async fn azure() -> AzureBlobCache<JsonCodec<Value>> {
support::connect(
&FakeBlobService::default(),
ACCOUNT_URL,
JsonCodec::new(),
Handle::current(),
)
.await
.unwrap()
}
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
const PREFIX: &str = "contract:";
// `overwrite_replaces` does not apply: sync `set_cache` never overwrites a blob, as in Python.
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn hit_and_miss(
#[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
context: ExactCacheContext,
) {
contract::hit_and_miss(&azure, context, PREFIX, json!({"answer": 42})).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn sync_async_equivalence(
#[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
context: ExactCacheContext,
) {
contract::sync_async_equivalence(&azure, context, PREFIX, json!("first"), json!([2])).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn pipeline_writes_every_entry(
#[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
context: ExactCacheContext,
) {
contract::pipeline_writes_every_entry(
&azure,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn batch_preserves_order(
#[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
context: ExactCacheContext,
) {
contract::batch_preserves_order(&azure, context, PREFIX, json!("first"), json!(2)).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn flush_clears(
#[future(awt)] azure: AzureBlobCache<JsonCodec<Value>>,
context: ExactCacheContext,
) {
contract::flush_clears(&azure, context, PREFIX, json!("value")).await;
}

239
litellm-rust/crates/cache-azure-blob/tests/support/mod.rs Normal file
View file

@ -0,0 +1,239 @@
#![allow(dead_code)]
use std::{
collections::BTreeMap,
sync::{Arc, Mutex},
};
use azure_core::http::{
AsyncRawResponse, Body, ClientOptions, HttpClient, Method, Request, StatusCode, Transport,
headers::{HeaderName, Headers},
};
use litellm_cache::{CacheCodec, Error};
use litellm_cache_azure_blob::AzureBlobCache;
use tokio::runtime::{Handle, Runtime};
pub const ACCOUNT_URL: &str = "https://example.blob.core.windows.net";
pub const CONTAINER: &str = "litellm-cache";
const IF_NONE_MATCH: HeaderName = HeaderName::from_static("if-none-match");
const ERROR_CODE: HeaderName = HeaderName::from_static("x-ms-error-code");
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RecordedRequest {
pub method: Method,
pub path: String,
pub query: String,
pub if_none_match: Option<String>,
}
#[derive(Default)]
struct FakeState {
container_exists: bool,
blobs: BTreeMap<String, Vec<u8>>,
requests: Vec<RecordedRequest>,
failing: bool,
precondition_conflicts: bool,
}
#[derive(Clone, Default)]
pub struct FakeBlobService {
state: Arc<Mutex<FakeState>>,
}
impl std::fmt::Debug for FakeBlobService {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("FakeBlobService")
}
}
impl FakeBlobService {
pub fn with_existing_container() -> Self {
let service = Self::default();
service.state.lock().unwrap().container_exists = true;
service
}
pub fn blob(&self, name: &str) -> Option<Vec<u8>> {
self.state.lock().unwrap().blobs.get(name).cloned()
}
pub fn blob_names(&self) -> Vec<String> {
self.state.lock().unwrap().blobs.keys().cloned().collect()
}
pub fn seed_blob(&self, name: &str, bytes: &[u8]) {
self.state
.lock()
.unwrap()
.blobs
.insert(name.to_string(), bytes.to_vec());
}
pub fn set_failing(&self, failing: bool) {
self.state.lock().unwrap().failing = failing;
}
pub fn set_precondition_conflicts(&self, enabled: bool) {
self.state.lock().unwrap().precondition_conflicts = enabled;
}
pub fn requests(&self) -> Vec<RecordedRequest> {
self.state.lock().unwrap().requests.clone()
}
pub fn container_exists(&self) -> bool {
self.state.lock().unwrap().container_exists
}
fn respond(status: StatusCode, error_code: Option<&str>, body: Vec<u8>) -> AsyncRawResponse {
let mut headers = Headers::new();
if let Some(code) = error_code {
headers.insert(ERROR_CODE, code.to_string());
}
AsyncRawResponse::from_bytes(status, headers, body)
}
fn list_body(state: &FakeState) -> Vec<u8> {
let mut xml = String::from(
r#"<?xml version="1.0" encoding="utf-8"?><EnumerationResults ServiceEndpoint="https://example.blob.core.windows.net/" ContainerName="litellm-cache"><Blobs>"#,
);
for name in state.blobs.keys() {
xml.push_str(&format!(
"<Blob><Name>{name}</Name><Properties><BlobType>BlockBlob</BlobType></Properties></Blob>"
));
}
xml.push_str("</Blobs><NextMarker /></EnumerationResults>");
xml.into_bytes()
}
}
#[async_trait::async_trait]
impl HttpClient for FakeBlobService {
async fn execute_request(&self, request: &Request) -> azure_core::Result<AsyncRawResponse> {
let mut state = self.state.lock().unwrap();
let path = request.url().path().to_string();
let query = request.url().query().unwrap_or_default().to_string();
let if_none_match = request
.headers()
.get_optional_str(&IF_NONE_MATCH)
.map(str::to_owned);
state.requests.push(RecordedRequest {
method: request.method(),
path: path.clone(),
query: query.clone(),
if_none_match: if_none_match.clone(),
});
if state.failing {
return Ok(Self::respond(
StatusCode::Forbidden,
Some("AuthorizationFailure"),
Vec::new(),
));
}
let container_path = format!("/{CONTAINER}");
let blob_name = path
.strip_prefix(&format!("{container_path}/"))
.map(str::to_owned);
let is_container = path == container_path && query.contains("restype=container");
let response = match (request.method(), is_container, blob_name) {
(Method::Put, true, None) if state.container_exists => Self::respond(
StatusCode::Conflict,
Some("ContainerAlreadyExists"),
Vec::new(),
),
(Method::Put, true, None) => {
state.container_exists = true;
Self::respond(StatusCode::Created, None, Vec::new())
}
(Method::Get, true, None) if query.contains("comp=list") => {
Self::respond(StatusCode::Ok, None, Self::list_body(&state))
}
(Method::Get, true, None) if state.container_exists => {
Self::respond(StatusCode::Ok, None, Vec::new())
}
(Method::Get, true, None) => {
Self::respond(StatusCode::NotFound, Some("ContainerNotFound"), Vec::new())
}
(Method::Put, false, Some(name)) => {
if if_none_match.as_deref() == Some("*") && state.blobs.contains_key(&name) {
if state.precondition_conflicts {
Self::respond(
StatusCode::PreconditionFailed,
Some("ConditionNotMet"),
Vec::new(),
)
} else {
Self::respond(StatusCode::Conflict, Some("BlobAlreadyExists"), Vec::new())
}
} else {
let bytes = match request.body() {
Body::Bytes(bytes) => bytes.to_vec(),
Body::SeekableStream(_) => panic!("unexpected streaming upload"),
};
state.blobs.insert(name, bytes);
Self::respond(StatusCode::Created, None, Vec::new())
}
}
(Method::Get, false, Some(name)) => match state.blobs.get(&name) {
Some(bytes) => Self::respond(StatusCode::Ok, None, bytes.clone()),
None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
},
(Method::Delete, false, Some(name)) => match state.blobs.remove(&name) {
Some(_) => Self::respond(StatusCode::Accepted, None, Vec::new()),
None => Self::respond(StatusCode::NotFound, Some("BlobNotFound"), Vec::new()),
},
(method, _, _) => panic!("unexpected request {method:?} {path}?{query}"),
};
Ok(response)
}
}
pub async fn connect<C: CacheCodec>(
service: &FakeBlobService,
account_url: &str,
codec: C,
handle: Handle,
) -> Result<AzureBlobCache<C>, Error> {
AzureBlobCache::connect_with_options(
account_url,
CONTAINER,
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(service.clone()))),
..ClientOptions::default()
},
codec,
handle,
)
.await
}
/// A cache on its own fake service and runtime, so sync methods run outside any runtime.
pub struct Fixture<C> {
pub runtime: Runtime,
pub service: FakeBlobService,
pub cache: Arc<AzureBlobCache<C>>,
}
impl<C: CacheCodec> Fixture<C> {
pub fn new(service: FakeBlobService, codec: C) -> Self {
let runtime = Runtime::new().unwrap();
let cache = runtime
.block_on(connect(
&service,
ACCOUNT_URL,
codec,
runtime.handle().clone(),
))
.unwrap();
Self {
runtime,
service,
cache: Arc::new(cache),
}
}
pub fn stored_json(&self, key: &str) -> serde_json::Value {
serde_json::from_slice(&self.service.blob(key).expect("blob should exist")).unwrap()
}
}

90
litellm-rust/crates/cache-azure-blob/tests/transport.rs Normal file
View file

@ -0,0 +1,90 @@
use std::sync::Arc;
use azure_core::http::{ClientOptions, Transport};
use litellm_cache::{BaseCache, ExactCacheContext, JsonCodec};
use litellm_cache_azure_blob::{AzureBlobCache, ReqwestTransport};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use tokio::runtime::Handle;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_json, header, method, path, query_param},
};
#[fixture]
async fn server() -> MockServer {
let server = MockServer::start().await;
Mock::given(method("PUT"))
.and(path("/litellm-cache"))
.and(query_param("restype", "container"))
.respond_with(ResponseTemplate::new(201))
.expect(1)
.mount(&server)
.await;
server
}
async fn connect(server: &MockServer) -> AzureBlobCache<JsonCodec<Value>> {
AzureBlobCache::connect_with_options(
&server.uri(),
"litellm-cache",
None,
ClientOptions {
transport: Some(Transport::new(Arc::new(ReqwestTransport(
reqwest::Client::new(),
)))),
..ClientOptions::default()
},
JsonCodec::new(),
Handle::current(),
)
.await
.unwrap()
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn uploads_go_through_the_host_client(#[future(awt)] server: MockServer) {
Mock::given(method("PUT"))
.and(path("/litellm-cache/key"))
.and(header("if-none-match", "*"))
.and(body_json(json!({"answer": 1})))
.respond_with(ResponseTemplate::new(201))
.expect(1)
.mount(&server)
.await;
connect(&server)
.await
.set_cache("key", json!({"answer": 1}), &ExactCacheContext::default())
.unwrap();
}
#[rstest]
#[case::hit(
ResponseTemplate::new(200).set_body_json(json!({"answer": 2})),
Some(json!({"answer": 2}))
)]
#[case::blob_not_found(
ResponseTemplate::new(404).insert_header("x-ms-error-code", "BlobNotFound"),
None
)]
#[tokio::test(flavor = "multi_thread")]
async fn downloads_map_the_host_client_response(
#[future(awt)] server: MockServer,
#[case] response: ResponseTemplate,
#[case] expected: Option<Value>,
) {
Mock::given(method("GET"))
.and(path("/litellm-cache/key"))
.respond_with(response)
.mount(&server)
.await;
assert_eq!(
connect(&server)
.await
.async_get_cache("key", &ExactCacheContext::default())
.await
.unwrap(),
expected
);
}

1
litellm-rust/crates/cache-disk/Cargo.toml
View file

@ -15,5 +15,6 @@ serde_json.workspace = true
tokio.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
rstest.workspace = true
tempfile = "3.27.0"

38
litellm-rust/crates/cache-disk/src/cache.rs
View file

@ -5,8 +5,8 @@ use std::{
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, DisconnectCache,
Error, ExactCacheContext, FlushCache,
};
use crate::{DiskStore, DiskcacheSqliteStore, PythonDiskCacheAdapter, StoredValue, ValueAdapter};
@ -150,29 +150,6 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BaseCache for DiskCache<S, D,
})
.await
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
let result = Self::run_blocking(Arc::clone(&self.store), |store| {
store.probe().map(|_| CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "Disk cache connection test successful".into(),
error: None,
})
})
.await;
Ok(match result {
Ok(result) => result,
Err(error) => CacheConnectionResult {
status: CacheConnectionStatus::Failed,
message: format!("Disk cache connection failed: {error}"),
error: Some(error.to_string()),
},
})
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> BatchCache for DiskCache<S, D, A> {
@ -241,9 +218,13 @@ impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> FlushCache for DiskCache<S, D
}
}
impl<S: CacheCodec<Value = f64>, D: DiskStore, A: ValueAdapter> CounterCache
for DiskCache<S, D, A>
{
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> DisconnectCache for DiskCache<S, D, A> {
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
}
impl<S: CacheCodec, D: DiskStore, A: ValueAdapter> CounterCache for DiskCache<S, D, A> {
fn increment_cache(
&self,
key: &str,
@ -264,6 +245,7 @@ impl<S: CacheCodec<Value = f64>, D: DiskStore, A: ValueAdapter> CounterCache
key: &str,
amount: f64,
context: ExactCacheContext,
_refresh_ttl: bool,
) -> Result<f64, Error> {
let key = key.to_string();
let adapter = Arc::clone(&self.adapter);

12
litellm-rust/crates/cache-disk/src/sqlite.rs
View file

@ -544,18 +544,6 @@ impl DiskStore for DiskcacheSqliteStore {
}
}
}
fn probe(&self) -> Result<(), Error> {
let connection = self.connection.lock().map_err(|_| Error::Unavailable)?;
connection
.query_row(
"SELECT value FROM Settings WHERE key = 'count'",
[],
|row| row.get::<_, i64>(0),
)
.map(|_| ())
.map_err(|_| Error::Unavailable)
}
}
fn default_settings() -> HashMap<String, Value> {

1
litellm-rust/crates/cache-disk/src/store.rs
View file

@ -29,5 +29,4 @@ pub trait DiskStore: Send + Sync + 'static {
now: f64,
apply: &mut dyn FnMut(Option<StoredValue>) -> Result<(StoredValue, Option<f64>), Error>,
) -> Result<(), Error>;
fn probe(&self) -> Result<(), Error>;
}

103
litellm-rust/crates/cache-disk/tests/cache.rs
View file

@ -7,8 +7,8 @@ use std::{
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, ExactCacheContext,
FlushCache, JsonCodec,
BaseCache, BatchCache, BatchEntry, CacheCodec, CounterCache, DeleteCache, DisconnectCache,
ExactCacheContext, FlushCache, JsonCodec,
};
use litellm_cache_disk::{DiskCache, DiskStore, DiskcacheSqliteStore, StoredValue, ValueAdapter};
use rstest::{fixture, rstest};
@ -395,7 +395,7 @@ fn delete_flush_and_spilled_file_replacement_clean_up_storage(sandbox: Sandbox)
#[rstest]
#[tokio::test]
async fn async_operations_connection_and_delete_match_sync_operations(sandbox: Sandbox) {
async fn async_operations_disconnect_and_delete_match_sync_operations(sandbox: Sandbox) {
let cache = sandbox.cache::<Value>();
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
@ -424,8 +424,101 @@ async fn async_operations_connection_and_delete_match_sync_operations(sandbox: S
);
cache.async_delete_cache("a").await.unwrap();
cache.async_flush_cache().await.unwrap();
cache.disconnect().await.unwrap();
}
#[derive(Clone, Copy, Debug)]
enum Increment {
Sync,
Async { refresh_ttl: bool },
}
impl Increment {
async fn apply(
self,
cache: &DiskCache<JsonCodec<Value>>,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> f64 {
match self {
Self::Sync => cache.increment_cache(key, amount, context).unwrap(),
Self::Async { refresh_ttl } => cache
.async_increment(key, amount, context, refresh_ttl)
.await
.unwrap(),
}
}
}
#[rstest]
#[case::sync_missing(Increment::Sync, None, 3.0, 3.0)]
#[case::sync_existing_int(Increment::Sync, Some(json!(7)), 5.0, 12.0)]
#[case::sync_non_int(Increment::Sync, Some(json!("not-a-number")), 4.0, 4.0)]
#[case::async_missing(Increment::Async { refresh_ttl: false }, None, 2.0, 2.0)]
#[case::async_existing_int(Increment::Async { refresh_ttl: false }, Some(json!(10)), 5.0, 15.0)]
#[case::async_non_int(Increment::Async { refresh_ttl: false }, Some(json!("corrupt")), 9.0, 9.0)]
#[case::async_refresh_ttl_is_ignored(Increment::Async { refresh_ttl: true }, Some(json!(1)), 1.0, 2.0)]
#[tokio::test]
async fn increments_read_back_through_get_cache(
sandbox: Sandbox,
#[case] increment: Increment,
#[case] initial: Option<Value>,
#[case] amount: f64,
#[case] expected: f64,
) {
let cache = sandbox.cache::<Value>();
let context = ExactCacheContext::default();
if let Some(initial) = initial {
cache
.async_set_cache("counter", initial, context.clone())
.await
.unwrap();
}
assert_eq!(
cache.test_connection().await.unwrap().status,
litellm_cache::CacheConnectionStatus::Success
increment
.apply(&cache, "counter", amount, context.clone())
.await,
expected
);
assert_eq!(
cache.get_cache("counter", &context).unwrap(),
Some(json!(expected as i64))
);
}
#[rstest]
#[case::without_refresh(false)]
#[case::with_refresh(true)]
#[tokio::test]
async fn async_increment_rewrites_ttl_on_every_write(sandbox: Sandbox, #[case] refresh_ttl: bool) {
let cache = sandbox.cache::<Value>();
let expiry = || {
sandbox
.db()
.query_row(
"SELECT expire_time IS NOT NULL FROM Cache WHERE key = 'counter'",
[],
|row| row.get::<_, bool>(0),
)
.unwrap()
};
let ttl = ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
};
cache
.async_increment("counter", 1.0, ttl.clone(), refresh_ttl)
.await
.unwrap();
assert!(expiry());
cache
.async_increment("counter", 1.0, ExactCacheContext::default(), refresh_ttl)
.await
.unwrap();
assert!(!expiry());
cache
.async_increment("counter", 1.0, ttl, refresh_ttl)
.await
.unwrap();
assert!(expiry());
}

82
litellm-rust/crates/cache-disk/tests/contract.rs Normal file
View file

@ -0,0 +1,82 @@
use litellm_cache::{ExactCacheContext, JsonCodec};
use litellm_cache_disk::DiskCache;
use litellm_cache_testing as contract;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use tempfile::TempDir;
struct Disk {
cache: DiskCache<JsonCodec<Value>>,
_directory: TempDir,
}
#[fixture]
fn disk() -> Disk {
let directory = tempfile::tempdir().unwrap();
Disk {
cache: DiskCache::open(directory.path(), JsonCodec::new()).unwrap(),
_directory: directory,
}
}
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
const PREFIX: &str = "contract:";
#[rstest]
#[tokio::test]
async fn hit_and_miss(disk: Disk, context: ExactCacheContext) {
contract::hit_and_miss(&disk.cache, context, PREFIX, json!({"answer": 42})).await;
}
#[rstest]
#[tokio::test]
async fn sync_async_equivalence(disk: Disk, context: ExactCacheContext) {
contract::sync_async_equivalence(&disk.cache, context, PREFIX, json!("first"), json!([2]))
.await;
}
#[rstest]
#[tokio::test]
async fn overwrite_replaces(disk: Disk, context: ExactCacheContext) {
contract::overwrite_replaces(&disk.cache, context, PREFIX, json!(1), json!({"b": 2})).await;
}
#[rstest]
#[tokio::test]
async fn pipeline_writes_every_entry(disk: Disk, context: ExactCacheContext) {
contract::pipeline_writes_every_entry(
&disk.cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
}
#[rstest]
#[tokio::test]
async fn batch_preserves_order(disk: Disk, context: ExactCacheContext) {
contract::batch_preserves_order(&disk.cache, context, PREFIX, json!("first"), json!(2)).await;
}
#[rstest]
#[tokio::test]
async fn delete_removes_key(disk: Disk, context: ExactCacheContext) {
contract::delete_removes_key(&disk.cache, context, PREFIX, json!("value")).await;
}
#[rstest]
#[tokio::test]
async fn flush_clears(disk: Disk, context: ExactCacheContext) {
contract::flush_clears(&disk.cache, context, PREFIX, json!("value")).await;
}
#[rstest]
#[tokio::test]
async fn counter_accumulates(disk: Disk, context: ExactCacheContext) {
contract::counter_accumulates(&disk.cache, context, PREFIX).await;
}

2
litellm-rust/crates/cache-gcs/Cargo.toml
View file

@ -15,6 +15,8 @@ reqwest.workspace = true
tokio.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
rstest.workspace = true
serde_json.workspace = true
tokio.workspace = true
wiremock = "0.6.5"

36
litellm-rust/crates/cache-gcs/src/cache.rs
View file

@ -2,7 +2,7 @@ use std::{future::Future, sync::Arc, time::Duration};
use futures_util::future::try_join_all;
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, Error, ExactCacheContext,
BaseCache, BatchCache, BatchEntry, CacheCodec, DisconnectCache, Error, ExactCacheContext,
FlushCache,
};
use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, percent_encode};
@ -53,25 +53,25 @@ pub struct GcsCache<S: CacheCodec> {
}
impl<S: CacheCodec> GcsCache<S> {
pub fn new(config: GcsConfig, codec: S) -> Result<Self, Error> {
pub fn new(config: GcsConfig, client: Client, codec: S) -> Self {
let token = Arc::new(GcpTokenSource::new(config.path_service_account.clone()));
Self::with_token_source(config, codec, token)
Self::with_token_source(config, client, codec, token)
}
pub fn with_token_source(
config: GcsConfig,
client: Client,
codec: S,
token: Arc<dyn TokenSource>,
) -> Result<Self, Error> {
let client = Client::builder().build().map_err(|_| Error::Unavailable)?;
) -> Self {
let key_prefix = key_prefix(config.gcs_path.as_deref());
Ok(Self {
Self {
config,
key_prefix,
client,
token,
codec,
})
}
}
pub fn bucket_name(&self) -> &str {
@ -154,26 +154,26 @@ impl<S: CacheCodec> GcsCache<S> {
F: Future<Output = Result<T, Error>> + Send,
T: Send,
{
let run = || {
let run = |future: F| {
tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.map_err(|_| Error::Unavailable)
.and_then(|runtime| runtime.block_on(future))
};
if let Ok(handle) = tokio::runtime::Handle::try_current() {
if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread {
return tokio::task::block_in_place(run);
match tokio::runtime::Handle::try_current() {
Ok(handle) if handle.runtime_flavor() == tokio::runtime::RuntimeFlavor::MultiThread => {
tokio::task::block_in_place(|| handle.block_on(future))
}
return std::thread::scope(|scope| {
Ok(_) => std::thread::scope(|scope| {
scope
.spawn(run)
.spawn(|| run(future))
.join()
.map_err(|_| Error::Unavailable)
.and_then(|result| result)
});
}),
Err(_) => run(future),
}
run()
}
}
@ -222,14 +222,12 @@ impl<S: CacheCodec> BaseCache for GcsCache<S> {
.await
.map(|_| ())
}
}
impl<S: CacheCodec> DisconnectCache for GcsCache<S> {
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
}
}
impl<S: CacheCodec> BatchCache for GcsCache<S> {

351
litellm-rust/crates/cache-gcs/tests/cache.rs
View file

@ -1,37 +1,32 @@
use std::{sync::Arc, time::Duration};
mod support;
use std::{future::Future, pin::Pin, sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheContext, Error, ExactCacheContext, FlushCache,
JsonCodec,
BaseCache, BatchCache, BatchEntry, CacheContext, DisconnectCache, Error, ExactCacheContext,
FlushCache,
};
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource, key_prefix};
use serde_json::json;
use litellm_cache_gcs::{GcsCache, GcsConfig, TokenSource, key_prefix};
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::FakeBucket;
use wiremock::{
Mock, MockServer, ResponseTemplate,
matchers::{body_bytes, header, method, path, query_param},
};
fn config(server: &MockServer, gcs_path: Option<&str>) -> GcsConfig {
GcsConfig {
bucket_name: "bucket".into(),
gcs_path: gcs_path.map(str::to_string),
path_service_account: None,
endpoint: server.uri(),
}
#[fixture]
async fn server() -> MockServer {
MockServer::start().await
}
fn cache(server: &MockServer, gcs_path: Option<&str>) -> GcsCache<JsonCodec<serde_json::Value>> {
GcsCache::with_token_source(
config(server, gcs_path),
JsonCodec::new(),
Arc::new(StaticTokenSource("tok".into())),
)
.unwrap()
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
#[rstest]
#[tokio::test]
async fn set_writes_encoded_object_and_headers() {
let server = MockServer::start().await;
async fn set_writes_encoded_object_and_headers(#[future(awt)] server: MockServer) {
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.and(query_param("uploadType", "media"))
@ -42,12 +37,8 @@ async fn set_writes_encoded_object_and_headers() {
.expect(1)
.mount(&server)
.await;
cache(&server, Some("cache/"))
.set_cache(
"team:a b/c",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
support::cache(&server, Some("cache/"))
.set_cache("team:a b/c", json!({"value": "entry"}), &context())
.unwrap();
let requests = server.received_requests().await.unwrap();
assert_eq!(requests.len(), 1);
@ -57,103 +48,108 @@ async fn set_writes_encoded_object_and_headers() {
);
}
#[rstest]
#[case::hit(
"hit",
ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})),
Ok(Some(json!({"value": "entry"})))
)]
#[case::missing("missing", ResponseTemplate::new(404), Ok(None))]
#[case::server_error("server-error", ResponseTemplate::new(500), Err(Error::Unavailable))]
#[case::invalid(
"invalid",
ResponseTemplate::new(200).set_body_string("not json"),
Err(Error::InvalidEntry)
)]
#[tokio::test]
async fn get_maps_statuses_and_decode_failures() {
let server = MockServer::start().await;
async fn get_maps_statuses_and_decode_failures(
#[future(awt)] server: MockServer,
#[case] key: &str,
#[case] response: ResponseTemplate,
#[case] expected: Result<Option<Value>, Error>,
) {
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/hit"))
.and(path(format!("/storage/v1/b/bucket/o/{key}")))
.and(query_param("alt", "media"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.respond_with(response)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/missing"))
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/server-error"))
.respond_with(ResponseTemplate::new(500))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/invalid"))
.respond_with(ResponseTemplate::new(200).set_body_string("not json"))
.mount(&server)
.await;
let cache = cache(&server, None);
assert_eq!(
cache
.get_cache("hit", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
assert_eq!(
cache
.get_cache("missing", &ExactCacheContext::default())
.unwrap(),
None
);
assert_eq!(
cache
.get_cache("server-error", &ExactCacheContext::default())
.unwrap_err(),
Error::Unavailable
);
assert_eq!(
cache
.get_cache("invalid", &ExactCacheContext::default())
.unwrap_err(),
Error::InvalidEntry
);
let cache = support::cache(&server, None);
assert_eq!(cache.get_cache(key, &context()), expected);
assert_eq!(cache.async_get_cache(key, &context()).await, expected);
}
#[test]
fn key_prefix_normalizes_paths() {
assert_eq!(key_prefix(None), "");
assert_eq!(key_prefix(Some("a/b/")), "a/b/");
assert_eq!(key_prefix(Some("a/b")), "a/b/");
assert_eq!(key_prefix(Some("")), "");
#[rstest]
#[case::none(None, "")]
#[case::trailing_slash(Some("a/b/"), "a/b/")]
#[case::no_trailing_slash(Some("a/b"), "a/b/")]
#[case::empty(Some(""), "")]
fn key_prefix_normalizes_paths(#[case] gcs_path: Option<&str>, #[case] expected: &str) {
assert_eq!(key_prefix(gcs_path), expected);
}
#[rstest]
#[tokio::test]
async fn object_names_use_python_quote_encoding() {
let server = MockServer::start().await;
async fn cache_exposes_its_configuration(#[future(awt)] server: MockServer) {
let cache = GcsCache::new(
GcsConfig {
path_service_account: Some("/secrets/sa.json".into()),
..support::config(&server, Some("folder"))
},
reqwest::Client::new(),
litellm_cache::JsonCodec::<Value>::new(),
);
assert_eq!(cache.bucket_name(), "bucket");
assert_eq!(cache.key_prefix(), "folder/");
assert_eq!(cache.path_service_account(), Some("/secrets/sa.json"));
assert_eq!(cache.object_name("k"), "folder/k");
}
#[rstest]
#[case::punctuation("a~b-c_d.e/f g%h", "uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h")]
#[case::utf8("ключ", "uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87")]
#[tokio::test]
async fn object_names_use_python_quote_encoding(
#[future(awt)] server: MockServer,
#[case] key: &str,
#[case] query: &str,
) {
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.and(query_param("uploadType", "media"))
.respond_with(ResponseTemplate::new(200))
.expect(1)
.mount(&server)
.await;
support::cache(&server, Some("p/"))
.async_set_cache(key, json!({"value": key}), context())
.await
.unwrap();
let requests = server.received_requests().await.unwrap();
assert_eq!(requests[0].url.query(), Some(query));
}
#[rstest]
#[tokio::test]
async fn object_names_are_encoded_in_the_download_path(#[future(awt)] server: MockServer) {
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/p%2Fa%3Ab%20c"))
.and(query_param("alt", "media"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!(1)))
.expect(2)
.mount(&server)
.await;
let cache = cache(&server, Some("p/"));
cache
.set_cache(
"a~b-c_d.e/f g%h",
json!({"value": "punctuation"}),
&ExactCacheContext::default(),
)
.unwrap();
cache
.set_cache(
"ключ",
json!({"value": "utf8"}),
&ExactCacheContext::default(),
)
.unwrap();
let requests = server.received_requests().await.unwrap();
let queries: Vec<_> = requests
.iter()
.filter_map(|request| request.url.query())
.collect();
assert!(queries.contains(&"uploadType=media&name=p%2Fa~b-c_d.e%2Ff%20g%25h"));
assert!(queries.contains(&"uploadType=media&name=p%2F%D0%BA%D0%BB%D1%8E%D1%87"));
let cache = support::cache(&server, Some("p"));
assert_eq!(cache.get_cache("a:b c", &context()), Ok(Some(json!(1))));
assert_eq!(
cache.async_get_cache("a:b c", &context()).await,
Ok(Some(json!(1)))
);
}
#[rstest]
#[tokio::test]
async fn ignores_ttl_and_writes_pipeline_concurrently() {
let server = MockServer::start().await;
async fn ignores_ttl_and_writes_pipeline_concurrently(#[future(awt)] server: MockServer) {
for key in ["one", "two", "three"] {
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
@ -164,12 +160,10 @@ async fn ignores_ttl_and_writes_pipeline_concurrently() {
.mount(&server)
.await;
}
let cache = cache(&server, None);
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
assert_eq!(
cache.get_ttl(&ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5)))),
None
);
let cache = support::cache(&server, None);
let with_ttl = context().with_ttl(Some(Duration::from_secs(5)));
assert_eq!(cache.get_ttl(&context()), None);
assert_eq!(cache.get_ttl(&with_ttl), None);
cache
.async_set_cache_pipeline(
vec![
@ -177,15 +171,15 @@ async fn ignores_ttl_and_writes_pipeline_concurrently() {
("two".into(), json!({"key": "two"})),
("three".into(), json!({"key": "three"})),
],
ExactCacheContext::default().with_ttl(Some(Duration::from_secs(5))),
with_ttl,
)
.await
.unwrap();
}
#[rstest]
#[tokio::test]
async fn async_batch_get_preserves_hits_misses_and_invalid_entries() {
let server = MockServer::start().await;
async fn batch_get_preserves_hits_misses_and_invalid_entries(#[future(awt)] server: MockServer) {
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/hit"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
@ -201,122 +195,83 @@ async fn async_batch_get_preserves_hits_misses_and_invalid_entries() {
.respond_with(ResponseTemplate::new(200).set_body_string("not json"))
.mount(&server)
.await;
let cache = support::cache(&server, None);
let keys = vec!["hit".to_string(), "missing".into(), "invalid".into()];
let expected = vec![
BatchEntry::Hit(json!({"value": "entry"})),
BatchEntry::Miss,
BatchEntry::Invalid,
];
assert_eq!(cache.batch_get_cache(&keys, &context()).unwrap(), expected);
assert_eq!(
cache(&server, None)
.async_batch_get_cache(
vec!["hit".into(), "missing".into(), "invalid".into()],
ExactCacheContext::default(),
)
.await
.unwrap(),
vec![
BatchEntry::Hit(json!({"value": "entry"})),
BatchEntry::Miss,
BatchEntry::Invalid,
]
cache.async_batch_get_cache(keys, context()).await.unwrap(),
expected
);
}
#[rstest]
#[tokio::test]
async fn lifecycle_operations_are_noops_and_connection_test_is_unsupported() {
let server = MockServer::start().await;
let cache = cache(&server, None);
async fn flush_and_disconnect_are_noops_like_python(#[future(awt)] server: MockServer) {
let cache = support::cache(&server, None);
assert_eq!(cache.flush_cache(), Ok(()));
assert_eq!(cache.async_flush_cache().await, Ok(()));
assert_eq!(cache.disconnect().await, Ok(()));
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
);
assert!(server.received_requests().await.unwrap().is_empty());
}
#[test]
fn round_trip(cache: &support::JsonGcsCache) -> Result<Option<Value>, Error> {
cache.set_cache("key", json!({"value": "entry"}), &context())?;
cache.get_cache("key", &context())
}
#[rstest]
fn sync_operations_work_without_an_active_runtime() {
let runtime = tokio::runtime::Builder::new_multi_thread()
.enable_all()
.build()
.unwrap();
let server = runtime.block_on(MockServer::start());
runtime.block_on(
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.respond_with(ResponseTemplate::new(200))
.mount(&server),
);
runtime.block_on(
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server),
);
let cache = cache(&server, None);
cache
.set_cache(
"key",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
.unwrap();
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
let server = runtime.block_on(FakeBucket::serve());
let cache = support::cache(&server, None);
assert_eq!(round_trip(&cache), Ok(Some(json!({"value": "entry"}))));
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn sync_operations_work_inside_a_multi_thread_runtime() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path("/upload/storage/v1/b/bucket/o"))
.respond_with(ResponseTemplate::new(200))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/storage/v1/b/bucket/o/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"value": "entry"})))
.mount(&server)
.await;
let cache = cache(&server, None);
cache
.set_cache(
"key",
json!({"value": "entry"}),
&ExactCacheContext::default(),
)
.unwrap();
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(json!({"value": "entry"}))
);
let server = FakeBucket::serve().await;
let cache = support::cache(&server, None);
assert_eq!(round_trip(&cache), Ok(Some(json!({"value": "entry"}))));
}
#[rstest]
#[tokio::test]
async fn sync_operations_work_inside_a_current_thread_runtime() {
let server = FakeBucket::serve().await;
let cache = support::cache(&server, None);
assert_eq!(round_trip(&cache), Ok(Some(json!({"value": "entry"}))));
}
struct FailingTokenSource;
impl TokenSource for FailingTokenSource {
fn bearer_token(
&self,
) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<String, Error>> + Send + '_>>
{
fn bearer_token(&self) -> Pin<Box<dyn Future<Output = Result<String, Error>> + Send + '_>> {
Box::pin(async { Err(Error::Unavailable) })
}
}
#[rstest]
#[tokio::test]
async fn token_source_failure_skips_http() {
let server = MockServer::start().await;
let cache = GcsCache::with_token_source(
config(&server, None),
JsonCodec::<serde_json::Value>::new(),
Arc::new(FailingTokenSource),
)
.unwrap();
async fn token_source_failure_skips_http(#[future(awt)] server: MockServer) {
let cache = support::cache_with_token(&server, None, Arc::new(FailingTokenSource));
assert_eq!(
cache.get_cache("key", &context()).unwrap_err(),
Error::Unavailable
);
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.async_set_cache("key", json!(1), context())
.await
.unwrap_err(),
Error::Unavailable
);

65
litellm-rust/crates/cache-gcs/tests/contract.rs Normal file
View file

@ -0,0 +1,65 @@
mod support;
use litellm_cache::ExactCacheContext;
use litellm_cache_testing as contract;
use rstest::{fixture, rstest};
use serde_json::json;
use support::{FakeBucket, JsonGcsCache};
use wiremock::MockServer;
struct Gcs {
cache: JsonGcsCache,
_server: MockServer,
}
#[fixture]
async fn gcs() -> Gcs {
let server = FakeBucket::serve().await;
Gcs {
cache: support::cache(&server, Some("contract")),
_server: server,
}
}
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
const PREFIX: &str = "contract:";
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn hit_and_miss(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
contract::hit_and_miss(&gcs.cache, context, PREFIX, json!({"answer": 42})).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn sync_async_equivalence(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
contract::sync_async_equivalence(&gcs.cache, context, PREFIX, json!("first"), json!([2])).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn overwrite_replaces(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
contract::overwrite_replaces(&gcs.cache, context, PREFIX, json!(1), json!({"b": 2})).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn pipeline_writes_every_entry(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
contract::pipeline_writes_every_entry(
&gcs.cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn batch_preserves_order(#[future(awt)] gcs: Gcs, context: ExactCacheContext) {
contract::batch_preserves_order(&gcs.cache, context, PREFIX, json!("first"), json!(2)).await;
}

87
litellm-rust/crates/cache-gcs/tests/support/mod.rs Normal file
View file

@ -0,0 +1,87 @@
#![allow(dead_code)]
use std::{
collections::HashMap,
sync::{Arc, Mutex},
};
use litellm_cache::JsonCodec;
use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource, TokenSource};
use percent_encoding::percent_decode_str;
use serde_json::Value;
use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate, http::Method, matchers::any};
pub type JsonGcsCache = GcsCache<JsonCodec<Value>>;
pub fn config(server: &MockServer, gcs_path: Option<&str>) -> GcsConfig {
GcsConfig {
bucket_name: "bucket".into(),
gcs_path: gcs_path.map(str::to_string),
path_service_account: None,
endpoint: server.uri(),
}
}
pub fn cache_with_token(
server: &MockServer,
gcs_path: Option<&str>,
token: Arc<dyn TokenSource>,
) -> JsonGcsCache {
GcsCache::with_token_source(
config(server, gcs_path),
reqwest::Client::new(),
JsonCodec::new(),
token,
)
}
pub fn cache(server: &MockServer, gcs_path: Option<&str>) -> JsonGcsCache {
cache_with_token(server, gcs_path, Arc::new(StaticTokenSource("tok".into())))
}
/// An in-memory bucket speaking the JSON API's media upload and `alt=media` download.
#[derive(Clone, Default)]
pub struct FakeBucket {
objects: Arc<Mutex<HashMap<String, Vec<u8>>>>,
}
impl FakeBucket {
pub async fn serve() -> MockServer {
let server = MockServer::start().await;
Mock::given(any())
.respond_with(Self::default())
.mount(&server)
.await;
server
}
}
impl Respond for FakeBucket {
fn respond(&self, request: &Request) -> ResponseTemplate {
let mut objects = self.objects.lock().unwrap();
match request.method {
Method::POST => {
let name = request
.url
.query_pairs()
.find_map(|(key, value)| (key == "name").then(|| value.into_owned()))
.expect("uploads carry the object name");
objects.insert(name, request.body.clone());
ResponseTemplate::new(200)
}
Method::GET => {
let encoded = request
.url
.path()
.strip_prefix("/storage/v1/b/bucket/o/")
.expect("downloads address an object");
let name = percent_decode_str(encoded).decode_utf8().unwrap();
match objects.get(name.as_ref()) {
Some(body) => ResponseTemplate::new(200).set_body_bytes(body.clone()),
None => ResponseTemplate::new(404),
}
}
_ => ResponseTemplate::new(405),
}
}
}

2
litellm-rust/crates/cache-memory/Cargo.toml
View file

@ -9,6 +9,6 @@ repository.workspace = true
litellm-cache.workspace = true
[dev-dependencies]
serde_json.workspace = true
litellm-cache-testing.workspace = true
rstest.workspace = true
tokio.workspace = true

197
litellm-rust/crates/cache-memory/src/cache.rs
View file

@ -7,8 +7,8 @@ use std::{
};
use litellm_cache::{
BaseCache, BatchCache, CacheConnectionResult, CacheConnectionStatus, ClaimCache, CounterCache,
DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache,
BaseCache, BatchCache, ClaimCache, CounterCache, DeleteCache, DisconnectCache, Error,
ExactCacheContext, FlushCache, SetCache, TtlCache,
};
const DEFAULT_MAX_SIZE_IN_MEMORY: usize = 200;
@ -75,7 +75,9 @@ impl<V: Clone> InMemoryCache<V> {
expiration_heap: BinaryHeap::new(),
}),
max_size_in_memory: max_size_in_memory.unwrap_or(DEFAULT_MAX_SIZE_IN_MEMORY),
default_ttl: default_ttl.unwrap_or(DEFAULT_TTL),
default_ttl: default_ttl
.filter(|ttl| !ttl.is_zero())
.unwrap_or(DEFAULT_TTL),
max_entry_bytes,
measure_value,
now: Arc::new(now),
@ -91,21 +93,9 @@ impl<V: Clone> InMemoryCache<V> {
if self.max_size_in_memory == 0 {
return Ok(CacheWrite::Disabled);
}
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
&& measure(&value)? > limit
{
return Ok(CacheWrite::TooLarge);
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let key = key.into();
Self::evict(&mut state, self.max_size_in_memory, now, &key);
let expiration = state.expirations.get(&key).copied();
if expiration.is_none_or(|expiration| expiration < now) {
Self::set_expiration(&mut state, &key, now + ttl.unwrap_or(self.default_ttl));
}
state.values.insert(key, value);
Ok(CacheWrite::Stored)
self.store(&mut state, key.into(), value, ttl, now)
}
pub fn get_cache(&self, key: &str) -> Result<Option<V>, Error> {
@ -121,6 +111,70 @@ impl<V: Clone> InMemoryCache<V> {
Ok(state.values.get(key).cloned())
}
/// `check_value_size`: whether `value` fits `max_entry_bytes`. Always `true` without a
/// limit and a measure, since typed values have no generic size.
pub fn check_value_size(&self, value: &V) -> Result<bool, Error> {
match (self.max_entry_bytes, &self.measure_value) {
(Some(limit), Some(measure)) => Ok(measure(value)? <= limit),
_ => Ok(true),
}
}
/// `evict_cache`: drops expired entries, then the earliest-expiring ones until a new key
/// fits.
pub fn evict_cache(&self) -> Result<(), Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, None);
Ok(())
}
/// `evict_element_if_expired`: `true` when `key` had expired and was removed.
pub fn evict_element_if_expired(&self, key: &str) -> Result<bool, Error> {
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
let expired = state
.expirations
.get(key)
.is_some_and(|expiration| *expiration < now);
if expired {
Self::remove(&mut state, key);
}
Ok(expired)
}
/// `allow_ttl_override`: a write may set the TTL when the key has none or it has passed.
pub fn allow_ttl_override(&self, key: &str) -> Result<bool, Error> {
let now = (self.now)();
Ok(self
.expires_at(key)?
.is_none_or(|expiration| expiration < now))
}
/// The number of stored entries, expired ones included until they are evicted.
pub fn len(&self) -> Result<usize, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.values
.len())
}
pub fn is_empty(&self) -> Result<bool, Error> {
Ok(self.len()? == 0)
}
/// Entries in the expiration heap, stale ones included; bounded by eviction.
pub fn expiration_heap_len(&self) -> Result<usize, Error> {
Ok(self
.state
.lock()
.map_err(|_| Error::Unavailable)?
.expiration_heap
.len())
}
pub fn max_size_in_memory(&self) -> usize {
self.max_size_in_memory
}
@ -172,7 +226,9 @@ impl<V: Clone> InMemoryCache<V> {
Ok(())
}
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration, key: &str) {
/// Writing an existing `key` never evicts another entry, unlike Python, which pops the
/// earliest-expiring entry whenever the cache is full.
fn evict(state: &mut CacheState<V>, capacity: usize, now: Duration, key: Option<&str>) {
while let Some(Reverse((expiration, key))) = state.expiration_heap.peek().cloned() {
if state.expirations.get(&key).copied() != Some(expiration) {
state.expiration_heap.pop();
@ -183,7 +239,7 @@ impl<V: Clone> InMemoryCache<V> {
break;
}
}
if state.values.contains_key(key) {
if key.is_some_and(|key| state.values.contains_key(key)) {
return;
}
while state.values.len() >= capacity {
@ -209,6 +265,40 @@ impl<V: Clone> InMemoryCache<V> {
state.values.remove(key);
state.expirations.remove(key);
}
/// `get_cache` under the held lock: an expired entry is removed and reads as missing.
fn live(state: &mut CacheState<V>, key: &str, now: Duration) -> Option<V> {
if state
.expirations
.get(key)
.is_some_and(|expiration| *expiration < now)
{
Self::remove(state, key);
}
state.values.get(key).cloned()
}
/// Python `set_cache` under the held lock: evict first (even when `key` already exists),
/// then skip oversized values, then write, keeping a live key's expiry.
fn store(
&self,
state: &mut CacheState<V>,
key: String,
value: V,
ttl: Option<Duration>,
now: Duration,
) -> Result<CacheWrite, Error> {
Self::evict(state, self.max_size_in_memory, now, None);
if !self.check_value_size(&value)? {
return Ok(CacheWrite::TooLarge);
}
let expiration = state.expirations.get(&key).copied();
if expiration.is_none_or(|expiration| expiration < now) {
Self::set_expiration(state, &key, now + ttl.unwrap_or(self.default_ttl));
}
state.values.insert(key, value);
Ok(CacheWrite::Stored)
}
}
impl<V> ClaimCache for InMemoryCache<V>
@ -227,7 +317,7 @@ where
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, key);
Self::evict(&mut state, self.max_size_in_memory, now, Some(key));
let existing = state
.values
.get(key)
@ -262,38 +352,12 @@ impl CounterCache for InMemoryCache<f64> {
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, key);
let value = state.values.get(key).copied().unwrap_or_default() + amount;
if !state.expirations.contains_key(key) {
Self::set_expiration(
&mut state,
key,
now + self.get_ttl(&context).unwrap_or(self.default_ttl),
);
}
state.values.insert(key.into(), value);
let value = Self::live(&mut state, key, now).unwrap_or_default() + amount;
self.store(&mut state, key.into(), value, self.get_ttl(&context), now)?;
Ok(value)
}
}
impl InMemoryCache<f64> {
pub async fn async_increment_pipeline(
&self,
operations: Vec<IncrementOperation>,
) -> Result<Vec<f64>, Error> {
operations
.into_iter()
.map(|operation| {
self.increment_cache(
&operation.key,
operation.amount,
ExactCacheContext { ttl: operation.ttl },
)
})
.collect()
}
}
impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
type Value = V;
type Context = ExactCacheContext;
@ -315,18 +379,12 @@ impl<V: Clone + Send + Sync + 'static> BaseCache for InMemoryCache<V> {
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
self.get_cache(key)
}
}
impl<V: Clone + Send + Sync + 'static> DisconnectCache for InMemoryCache<V> {
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "In-memory cache connection test successful".into(),
error: None,
})
}
}
impl<V: Clone + Send + Sync + 'static> BatchCache for InMemoryCache<V> {}
@ -367,34 +425,9 @@ where
}
let now = (self.now)();
let mut state = self.state.lock().map_err(|_| Error::Unavailable)?;
Self::evict(&mut state, self.max_size_in_memory, now, key);
let mut stored = state.values.get(key).cloned().unwrap_or_default();
let mut stored = Self::live(&mut state, key, now).unwrap_or_default();
stored.extend(values.iter().cloned());
if let (Some(limit), Some(measure)) = (self.max_entry_bytes, &self.measure_value)
&& measure(&stored)? > limit
{
return Ok(values);
}
if !state.expirations.contains_key(key) {
Self::set_expiration(&mut state, key, now + ttl.unwrap_or(self.default_ttl));
}
state.values.insert(key.into(), stored);
self.store(&mut state, key.into(), stored, ttl, now)?;
Ok(values)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn repeated_increments_keep_one_heap_entry_per_expiration() {
let cache = InMemoryCache::<f64>::new(Some(4), None);
for _ in 0..100 {
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
}
assert_eq!(cache.state.lock().unwrap().expiration_heap.len(), 1);
}
}

687
litellm-rust/crates/cache-memory/tests/cache.rs
View file

@ -8,111 +8,296 @@ use std::{
};
use litellm_cache::{
BaseCache, CacheBackend, CacheConnectionStatus, ClaimCache, CounterCache, DeleteCache, Error,
ExactCacheContext, IncrementOperation, SetCache, get_cache, set_cache,
BaseCache, BatchCache, BatchEntry, CacheBackend, ClaimCache, CounterCache, DeleteCache,
DisconnectCache, Error, ExactCacheContext, FlushCache, IncrementOperation, SetCache, TtlCache,
get_cache, set_cache,
};
use litellm_cache_memory::{CacheWrite, InMemoryCache};
use rstest::{fixture, rstest};
type Clock = Arc<AtomicU64>;
#[fixture]
fn clock() -> Arc<AtomicU64> {
fn clock() -> Clock {
Arc::new(AtomicU64::new(100))
}
fn cache(clock: Arc<AtomicU64>, capacity: usize) -> InMemoryCache<String> {
fn cache_with<V: Clone>(clock: &Clock, capacity: usize) -> InMemoryCache<V> {
let clock = clock.clone();
InMemoryCache::with_clock(Some(capacity), Some(Duration::from_secs(60)), move || {
Duration::from_secs(clock.load(Ordering::SeqCst))
Duration::from_millis(clock.load(Ordering::SeqCst) * 1000)
})
}
fn cache(clock: &Clock, capacity: usize) -> InMemoryCache<String> {
cache_with(clock, capacity)
}
fn at(clock: &Clock, seconds: u64) {
clock.store(seconds, Ordering::SeqCst);
}
fn secs(seconds: u64) -> Option<Duration> {
Some(Duration::from_secs(seconds))
}
fn ttl(seconds: u64) -> ExactCacheContext {
ExactCacheContext { ttl: secs(seconds) }
}
fn measured(capacity: usize) -> InMemoryCache<String> {
InMemoryCache::with_clock_and_size_measurement(
Some(capacity),
secs(60),
Some(4),
Some(Arc::new(|value: &String| {
if value.is_empty() {
return Err(Error::InvalidEntry);
}
Ok(value.len())
})),
|| Duration::from_secs(100),
)
}
#[rstest]
fn default_explicit_and_override_ttls_follow_python_rules(clock: Arc<AtomicU64>) {
let cache = cache(clock.clone(), 4);
fn default_explicit_and_override_ttls_follow_python_rules(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("key", "first".into(), None).unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(160))
);
cache
.set_cache("key", "second".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(160))
);
clock.store(160, Ordering::SeqCst);
assert_eq!(cache.expires_at("key").unwrap(), secs(160));
cache.set_cache("key", "second".into(), secs(10)).unwrap();
assert_eq!(cache.expires_at("key").unwrap(), secs(160));
at(&clock, 160);
assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
clock.store(161, Ordering::SeqCst);
at(&clock, 161);
assert_eq!(cache.get_cache("key").unwrap(), None);
cache
.set_cache("key", "third".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(171))
);
assert_eq!(cache.expires_at("key").unwrap(), None);
cache.set_cache("key", "third".into(), secs(10)).unwrap();
assert_eq!(cache.expires_at("key").unwrap(), secs(171));
}
#[rstest]
fn write_at_expiry_boundary_refreshes_ttl(clock: Arc<AtomicU64>) {
let cache = cache(clock.clone(), 4);
cache
.set_cache("key", "first".into(), Some(Duration::from_secs(10)))
.unwrap();
clock.store(110, Ordering::SeqCst);
cache
.set_cache("key", "second".into(), Some(Duration::from_secs(10)))
.unwrap();
assert_eq!(
cache.expires_at("key").unwrap(),
Some(Duration::from_secs(120))
);
clock.store(115, Ordering::SeqCst);
#[case::unset(None, secs(600))]
#[case::zero_falls_back_like_python_or(Some(Duration::ZERO), secs(600))]
#[case::explicit(secs(5), secs(5))]
fn default_ttl_falls_back_to_ten_minutes(
#[case] default_ttl: Option<Duration>,
#[case] expected: Option<Duration>,
) {
let cache = InMemoryCache::<String>::with_clock(None, default_ttl, || Duration::ZERO);
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), expected);
cache.set_cache("key", "value".into(), None).unwrap();
assert_eq!(cache.expires_at("key").unwrap(), expected);
assert_eq!(cache.max_size_in_memory(), 200);
}
#[rstest]
fn write_at_expiry_boundary_refreshes_ttl(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("key", "first".into(), secs(10)).unwrap();
at(&clock, 110);
cache.set_cache("key", "second".into(), secs(10)).unwrap();
assert_eq!(cache.expires_at("key").unwrap(), secs(120));
at(&clock, 115);
assert_eq!(cache.get_cache("key").unwrap(), Some("second".into()));
}
#[rstest]
fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Arc<AtomicU64>) {
let cache = cache(clock, 2);
cache
.set_cache("early", "a".into(), Some(Duration::from_secs(10)))
.unwrap();
cache
.set_cache("late", "b".into(), Some(Duration::from_secs(20)))
.unwrap();
fn expired_key_without_a_read_allows_a_ttl_override(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("key", "first".into(), secs(1)).unwrap();
assert_eq!(cache.allow_ttl_override("key"), Ok(false));
at(&clock, 102);
assert_eq!(cache.allow_ttl_override("key"), Ok(true));
cache.set_cache("key", "second".into(), secs(1)).unwrap();
assert_eq!(cache.expires_at("key").unwrap(), secs(103));
assert_eq!(cache.allow_ttl_override("missing"), Ok(true));
}
#[rstest]
fn capacity_evicts_earliest_and_ignores_stale_heap_entries(clock: Clock) {
let cache = cache(&clock, 2);
cache.set_cache("early", "a".into(), secs(10)).unwrap();
cache.set_cache("late", "b".into(), secs(20)).unwrap();
cache.delete_cache("early").unwrap();
cache
.set_cache("new", "c".into(), Some(Duration::from_secs(30)))
.unwrap();
cache.set_cache("new", "c".into(), secs(30)).unwrap();
assert_eq!(cache.get_cache("late").unwrap(), Some("b".into()));
cache
.set_cache("last", "d".into(), Some(Duration::from_secs(40)))
.unwrap();
cache.set_cache("last", "d".into(), secs(40)).unwrap();
assert_eq!(cache.get_cache("late").unwrap(), None);
}
#[test]
fn disabled_size_limited_and_validated_writes_are_observable() {
let cache = |capacity| {
InMemoryCache::with_clock_and_size_measurement(
Some(capacity),
Some(Duration::from_secs(60)),
Some(4),
Some(Arc::new(|value: &String| {
if value.is_empty() {
return Err(Error::InvalidEntry);
}
Ok(value.len())
})),
|| Duration::from_secs(100),
)
};
let disabled = cache(0);
#[rstest]
fn max_size_is_respected_when_every_item_has_a_long_ttl(clock: Clock) {
let cache = cache(&clock, 3);
for index in 0..3 {
at(&clock, 100 + index);
cache
.set_cache(
format!("key_{index}"),
format!("value_{index}"),
secs(86_400),
)
.unwrap();
}
assert_eq!(cache.len(), Ok(3));
cache
.set_cache("key_3", "value_3".into(), secs(86_400))
.unwrap();
assert_eq!(cache.len(), Ok(3));
assert_eq!(cache.get_cache("key_0").unwrap(), None);
assert_eq!(cache.expires_at("key_0").unwrap(), None);
for key in ["key_1", "key_2", "key_3"] {
assert!(cache.get_cache(key).unwrap().is_some(), "{key}");
}
}
#[rstest]
fn expired_items_are_evicted_before_live_ones(clock: Clock) {
let cache = cache(&clock, 3);
cache.set_cache("expired_1", "1".into(), secs(1)).unwrap();
cache.set_cache("expired_2", "2".into(), secs(1)).unwrap();
cache
.set_cache("long_lived", "3".into(), secs(86_400))
.unwrap();
assert_eq!(cache.len(), Ok(3));
at(&clock, 102);
cache
.set_cache("new_item", "4".into(), secs(86_400))
.unwrap();
assert_eq!(cache.len(), Ok(2));
assert_eq!(cache.get_cache("long_lived").unwrap(), Some("3".into()));
assert_eq!(cache.get_cache("new_item").unwrap(), Some("4".into()));
for key in ["expired_1", "expired_2"] {
assert_eq!(cache.expires_at(key).unwrap(), None, "{key}");
}
}
#[rstest]
fn injected_clock_controls_expiry_and_eviction(clock: Clock) {
let cache = cache(&clock, 2);
at(&clock, 0);
cache
.set_cache("first", "original".into(), secs(10))
.unwrap();
at(&clock, 9);
cache.set_cache("second", "survivor".into(), None).unwrap();
assert_eq!(cache.get_cache("first").unwrap(), Some("original".into()));
at(&clock, 11);
assert_eq!(cache.get_cache("first").unwrap(), None);
cache
.set_cache("third", "replacement".into(), None)
.unwrap();
assert_eq!(cache.get_cache("second").unwrap(), Some("survivor".into()));
at(&clock, 70);
cache.set_cache("fourth", "new".into(), None).unwrap();
assert_eq!(cache.get_cache("second").unwrap(), None);
assert_eq!(
disabled.set_cache("a", "x".into(), None).unwrap(),
cache.get_cache("third").unwrap(),
Some("replacement".into())
);
assert_eq!(cache.get_cache("fourth").unwrap(), Some("new".into()));
}
#[rstest]
fn rewriting_one_key_keeps_one_heap_entry(clock: Clock) {
let cache = cache(&clock, 10);
for index in 0..1_000 {
cache
.set_cache("hot_key", format!("value_{index}"), secs(60))
.unwrap();
}
assert_eq!(cache.expiration_heap_len(), Ok(1));
}
#[rstest]
fn repeated_increments_keep_one_heap_entry_per_expiration() {
let cache = InMemoryCache::<f64>::new(Some(4), None);
for _ in 0..100 {
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap();
}
assert_eq!(cache.expiration_heap_len(), Ok(1));
}
#[rstest]
fn reinserting_expired_keys_below_capacity_prunes_the_heap(clock: Clock) {
let cache = cache(&clock, 200);
for cycle in 0..3 {
for index in 0..5 {
cache
.set_cache(format!("key_{index}"), format!("value_{cycle}"), secs(1))
.unwrap();
}
at(&clock, 100 + 2 * (cycle + 1));
}
for index in 0..5 {
cache
.set_cache(format!("key_{index}"), "final".into(), secs(1))
.unwrap();
}
assert_eq!(cache.len(), Ok(5));
assert_eq!(cache.expiration_heap_len(), Ok(5));
}
#[rstest]
fn evict_cache_drops_expired_entries_then_makes_room(clock: Clock) {
let cache = cache(&clock, 2);
assert_eq!(cache.is_empty(), Ok(true));
cache.set_cache("short", "a".into(), secs(1)).unwrap();
cache.set_cache("long", "b".into(), secs(50)).unwrap();
at(&clock, 102);
cache.evict_cache().unwrap();
assert_eq!(cache.len(), Ok(1));
assert_eq!(cache.expires_at("short").unwrap(), None);
cache.set_cache("longer", "c".into(), secs(90)).unwrap();
cache.evict_cache().unwrap();
assert_eq!(cache.len(), Ok(1));
assert_eq!(cache.get_cache("long").unwrap(), None);
assert_eq!(cache.get_cache("longer").unwrap(), Some("c".into()));
}
#[rstest]
fn evict_element_if_expired_reports_removal(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("key", "value".into(), secs(10)).unwrap();
assert_eq!(cache.evict_element_if_expired("key"), Ok(false));
assert_eq!(cache.evict_element_if_expired("missing"), Ok(false));
at(&clock, 110);
assert_eq!(cache.evict_element_if_expired("key"), Ok(false));
at(&clock, 111);
assert_eq!(cache.evict_element_if_expired("key"), Ok(true));
assert_eq!(cache.len(), Ok(0));
assert_eq!(cache.expires_at("key").unwrap(), None);
}
#[rstest]
#[case::fits("ok", Ok(true))]
#[case::at_limit("four", Ok(true))]
#[case::too_large("oversized", Ok(false))]
#[case::measure_error("", Err(Error::InvalidEntry))]
fn check_value_size_applies_the_entry_limit(
#[case] value: &str,
#[case] expected: Result<bool, Error>,
) {
assert_eq!(measured(2).check_value_size(&value.to_string()), expected);
}
#[rstest]
fn values_are_unbounded_without_a_measure() {
let cache = InMemoryCache::<String>::default();
assert_eq!(cache.max_entry_bytes(), None);
assert_eq!(cache.check_value_size(&"x".repeat(1 << 20)), Ok(true));
}
#[rstest]
fn disabled_size_limited_and_validated_writes_are_observable() {
assert_eq!(
measured(0).set_cache("a", "x".into(), None).unwrap(),
CacheWrite::Disabled
);
let cache = cache(2);
let cache = measured(2);
assert_eq!(cache.max_entry_bytes(), Some(4));
assert_eq!(
cache.set_cache("large", "oversized".into(), None).unwrap(),
CacheWrite::TooLarge
@ -132,30 +317,21 @@ fn disabled_size_limited_and_validated_writes_are_observable() {
assert_eq!(cache.get_cache("small").unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn connection_test_matches_python_result_contract() {
async fn disconnect_is_a_no_op_that_keeps_entries() {
let cache = InMemoryCache::<String>::default();
let result = BaseCache::test_connection(&cache).await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert_eq!(result.message, "In-memory cache connection test successful");
assert_eq!(result.error, None);
assert_eq!(
serde_json::to_value(result).unwrap(),
serde_json::json!({
"status": "success",
"message": "In-memory cache connection test successful"
})
);
cache.set_cache("key", "value".into(), None).unwrap();
cache.disconnect().await.unwrap();
assert_eq!(cache.get_cache("key").unwrap(), Some("value".into()));
}
#[rstest]
#[tokio::test]
async fn generic_consumers_share_typed_values_and_honor_expiration() {
let clock = clock();
let cache: CacheBackend<InMemoryCache<String>> = Arc::new(cache(clock.clone(), 4));
async fn generic_consumers_share_typed_values_and_honor_expiration(clock: Clock) {
let cache: CacheBackend<InMemoryCache<String>> = Arc::new(self::cache(&clock, 4));
let reader = Arc::clone(&cache);
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(5)),
};
let context = ttl(5);
set_cache(cache.as_ref(), "sync", "first".into(), &context).unwrap();
assert_eq!(
get_cache(reader.as_ref(), "sync", &context).unwrap(),
@ -181,7 +357,7 @@ async fn generic_consumers_share_typed_values_and_honor_expiration() {
reader.async_get_cache("async", &context).await.unwrap(),
None
);
clock.store(106, Ordering::SeqCst);
at(&clock, 106);
assert_eq!(get_cache(reader.as_ref(), "sync", &context).unwrap(), None);
assert_eq!(
reader.async_get_cache("batch", &context).await.unwrap(),
@ -189,34 +365,95 @@ async fn generic_consumers_share_typed_values_and_honor_expiration() {
);
}
#[test]
fn claims_are_atomic_and_refresh_eligible_winners() {
let clock = clock();
let cache = InMemoryCache::with_clock(Some(4), Some(Duration::from_secs(60)), {
let clock = clock.clone();
move || Duration::from_secs(clock.load(Ordering::SeqCst))
});
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(10)),
};
#[rstest]
#[case::context_ttl(ttl(5), secs(105))]
#[case::default_ttl(ExactCacheContext::default(), secs(160))]
#[tokio::test]
async fn pipeline_writes_use_the_context_ttl_or_the_default(
clock: Clock,
#[case] context: ExactCacheContext,
#[case] expected: Option<Duration>,
) {
let cache = cache(&clock, 4);
cache
.async_set_cache_pipeline(
vec![("a".into(), "1".into()), ("b".into(), "2".into())],
context,
)
.await
.unwrap();
assert_eq!(cache.expires_at("a").unwrap(), expected);
assert_eq!(cache.expires_at("b").unwrap(), expected);
}
#[rstest]
#[tokio::test]
async fn batch_reads_return_one_entry_per_key_and_drop_expired_ones(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("short", "a".into(), secs(1)).unwrap();
cache.set_cache("long", "b".into(), secs(50)).unwrap();
let keys = vec!["short".to_string(), "missing".into(), "long".into()];
assert_eq!(
cache
.batch_get_cache(&keys, &ExactCacheContext::default())
.unwrap(),
[
BatchEntry::Hit("a".to_string()),
BatchEntry::Miss,
BatchEntry::Hit("b".into()),
]
);
at(&clock, 102);
assert_eq!(
cache
.async_batch_get_cache(keys, ExactCacheContext::default())
.await
.unwrap(),
[
BatchEntry::Miss,
BatchEntry::Miss,
BatchEntry::Hit("b".into())
]
);
}
#[rstest]
#[tokio::test]
async fn flush_clears_values_and_expirations(clock: Clock) {
let cache = cache(&clock, 4);
cache.set_cache("a", "1".into(), None).unwrap();
cache.set_cache("b", "2".into(), None).unwrap();
cache.flush_cache().unwrap();
assert_eq!(cache.len(), Ok(0));
assert_eq!(cache.expiration_heap_len(), Ok(0));
cache.set_cache("c", "3".into(), None).unwrap();
FlushCache::async_flush_cache(&cache).await.unwrap();
assert_eq!(cache.is_empty(), Ok(true));
assert_eq!(
cache.async_get_oldest_n_keys(5).await.unwrap(),
Vec::<String>::new()
);
}
#[rstest]
fn claims_are_atomic_and_refresh_eligible_winners(clock: Clock) {
let cache = cache(&clock, 4);
let context = ttl(10);
assert_eq!(
cache
.claim_cache("affinity", "first".to_string(), &[], context.clone())
.unwrap(),
"first"
);
clock.store(103, Ordering::SeqCst);
at(&clock, 103);
assert_eq!(
cache
.claim_cache("affinity", "second".to_string(), &[], context.clone())
.unwrap(),
"first"
);
assert_eq!(
cache.expires_at("affinity").unwrap(),
Some(Duration::from_secs(110))
);
clock.store(105, Ordering::SeqCst);
assert_eq!(cache.expires_at("affinity").unwrap(), secs(110));
at(&clock, 105);
assert_eq!(
cache
.claim_cache(
@ -228,13 +465,10 @@ fn claims_are_atomic_and_refresh_eligible_winners() {
.unwrap(),
"first"
);
assert_eq!(
cache.expires_at("affinity").unwrap(),
Some(Duration::from_secs(115))
);
assert_eq!(cache.expires_at("affinity").unwrap(), secs(115));
}
#[test]
#[rstest]
fn counters_increment_under_one_lock() {
let cache = InMemoryCache::<f64>::default();
assert_eq!(
@ -250,44 +484,107 @@ fn counters_increment_under_one_lock() {
}
#[rstest]
fn rewriting_an_existing_key_at_capacity_keeps_other_entries(clock: Arc<AtomicU64>) {
let cache = cache(clock, 2);
cache
.set_cache("hot", "1".into(), Some(Duration::from_secs(10)))
.unwrap();
cache
.set_cache("cold", "2".into(), Some(Duration::from_secs(20)))
.unwrap();
fn concurrent_increments_are_atomic() {
let cache = Arc::new(InMemoryCache::<f64>::default());
cache.set_cache("counter", 1000.0, None).unwrap();
let threads = (0..8)
.map(|_| {
let cache = cache.clone();
std::thread::spawn(move || {
cache
.increment_cache("counter", 1.0, ExactCacheContext::default())
.unwrap()
})
})
.collect::<Vec<_>>();
for thread in threads {
thread.join().unwrap();
}
assert_eq!(cache.get_cache("counter").unwrap(), Some(1008.0));
}
#[rstest]
#[case::window_semantics(false)]
#[case::refresh_ttl_is_ignored(true)]
#[tokio::test]
async fn async_increment_delegates_to_the_locked_sync_path(
clock: Clock,
#[case] refresh_ttl: bool,
) {
let cache = cache_with::<f64>(&clock, 4);
assert_eq!(
cache
.async_increment("counter", 2.0, ttl(10), refresh_ttl)
.await,
Ok(2.0)
);
at(&clock, 105);
assert_eq!(
cache
.async_increment("counter", 3.0, ttl(10), refresh_ttl)
.await,
Ok(5.0)
);
assert_eq!(cache.get_cache("counter").unwrap(), Some(5.0));
assert_eq!(cache.expires_at("counter").unwrap(), secs(110));
}
#[rstest]
fn expired_counters_restart_from_zero_with_a_new_ttl(clock: Clock) {
let cache = cache_with::<f64>(&clock, 4);
cache.increment_cache("counter", 2.0, ttl(10)).unwrap();
at(&clock, 111);
assert_eq!(cache.increment_cache("counter", 1.0, ttl(10)), Ok(1.0));
assert_eq!(cache.expires_at("counter").unwrap(), secs(121));
}
/// Python `InMemoryCache.set_cache` runs `evict_cache()` before every insert, and step 2 evicts
/// the earliest expiry while `len(cache_dict) >= max_size_in_memory`, even when the key being
/// written already exists.
#[rstest]
fn overwriting_an_existing_key_at_capacity_evicts_the_earliest_expiry_like_python(clock: Clock) {
let cache = cache(&clock, 2);
cache.set_cache("hot", "1".into(), secs(10)).unwrap();
cache.set_cache("cold", "2".into(), secs(20)).unwrap();
cache.set_cache("cold", "3".into(), None).unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
assert_eq!(cache.get_cache("hot").unwrap(), None);
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
}
/// `claim_cache` has no Python counterpart; it never evicts another entry for a key it holds.
#[rstest]
fn claiming_an_existing_key_at_capacity_keeps_other_entries(clock: Clock) {
let cache = cache(&clock, 2);
cache.set_cache("hot", "1".into(), secs(10)).unwrap();
cache.set_cache("cold", "2".into(), secs(20)).unwrap();
cache
.claim_cache("cold", "4".into(), &[], ExactCacheContext::default())
.unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), Some("1".into()));
cache.set_cache("new", "5".into(), None).unwrap();
assert_eq!(cache.get_cache("hot").unwrap(), None);
assert_eq!(cache.get_cache("cold").unwrap(), Some("3".into()));
assert_eq!(cache.get_cache("new").unwrap(), Some("5".into()));
assert_eq!(cache.get_cache("cold").unwrap(), Some("2".into()));
}
#[test]
fn incrementing_an_existing_counter_at_capacity_keeps_every_counter() {
let cache = InMemoryCache::<f64>::new(Some(2), None);
/// Python `increment_cache` is `get_cache` then `set_cache`, so at capacity the write evicts
/// the earliest expiry first: equal expiries tie-break on the key, and the value read before
/// eviction is the one written back.
#[rstest]
fn incrementing_at_capacity_evicts_the_earliest_expiry_like_python(clock: Clock) {
let cache = cache_with::<f64>(&clock, 2);
for key in ["a", "b", "a", "b"] {
cache
.increment_cache(key, 1.0, ExactCacheContext::default())
.unwrap();
}
assert_eq!(cache.get_cache("a").unwrap(), Some(2.0));
assert_eq!(cache.get_cache("a").unwrap(), None);
assert_eq!(cache.get_cache("b").unwrap(), Some(2.0));
}
#[test]
fn disabled_cache_does_not_retain_claims_or_counters() {
#[rstest]
#[tokio::test]
async fn disabled_cache_does_not_retain_claims_counters_or_sets() {
let claims = InMemoryCache::<String>::new(Some(0), None);
assert_eq!(
claims
@ -305,62 +602,128 @@ fn disabled_cache_does_not_retain_claims_or_counters() {
2.0
);
assert_eq!(counters.get_cache("key").unwrap(), None);
let sets = InMemoryCache::<HashSet<String>>::new(Some(0), None);
assert_eq!(
sets.async_set_cache_sadd("key", vec!["a".into()], None)
.await
.unwrap(),
["a"]
);
assert_eq!(sets.get_cache("key").unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn ttl_and_oldest_key_operations_use_the_stored_expirations() {
let clock = Arc::new(AtomicU64::new(100));
let cache = cache(clock, 3);
cache
.set_cache("later", "2".into(), Some(Duration::from_secs(20)))
.unwrap();
cache
.set_cache("first", "1".into(), Some(Duration::from_secs(10)))
.unwrap();
async fn ttl_and_oldest_key_operations_use_the_stored_expirations(clock: Clock) {
let cache = cache(&clock, 3);
cache.set_cache("later", "2".into(), secs(20)).unwrap();
cache.set_cache("first", "1".into(), secs(10)).unwrap();
cache.set_cache("latest", "3".into(), secs(30)).unwrap();
assert_eq!(cache.async_get_ttl("first").await.unwrap(), secs(110));
assert_eq!(
cache.async_get_ttl("first").await.unwrap(),
Some(Duration::from_secs(110))
TtlCache::async_get_ttl(&cache, "later").await.unwrap(),
secs(120)
);
assert_eq!(cache.async_get_oldest_n_keys(1).await.unwrap(), ["first"]);
assert_eq!(
cache.async_get_oldest_n_keys(10).await.unwrap(),
["first", "later", "latest"]
);
assert_eq!(
cache.async_get_oldest_n_keys(0).await.unwrap(),
Vec::<String>::new()
);
assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn increment_pipeline_preserves_operation_order() {
let cache = InMemoryCache::<f64>::new(Some(3), None);
async fn increment_pipeline_preserves_operation_order(clock: Clock) {
let cache = cache_with::<f64>(&clock, 3);
let operation = |key: &str, amount, ttl| IncrementOperation {
key: key.into(),
amount,
ttl: secs(ttl),
};
assert_eq!(
cache
.async_increment_pipeline(vec![
IncrementOperation {
key: "a".into(),
amount: 1.0,
ttl: Some(Duration::from_secs(10)),
},
IncrementOperation {
key: "a".into(),
amount: 2.0,
ttl: Some(Duration::from_secs(20)),
},
operation("a", 1.0, 10),
operation("b", 5.0, 30),
operation("a", 2.0, 20),
])
.await
.unwrap(),
[1.0, 3.0]
[1.0, 5.0, 3.0]
);
assert_eq!(cache.get_cache("a").unwrap(), Some(3.0));
assert_eq!(cache.expires_at("a").unwrap(), secs(110));
assert_eq!(cache.expires_at("b").unwrap(), secs(130));
assert_eq!(
cache.async_increment_pipeline(Vec::new()).await.unwrap(),
Vec::<f64>::new()
);
}
#[rstest]
#[tokio::test]
async fn set_capability_preserves_python_result_and_deduplicates_storage() {
let cache = InMemoryCache::<HashSet<String>>::new(None, None);
let inserted = vec!["a".into(), "a".into(), "b".into()];
async fn set_capability_preserves_python_result_and_deduplicates_storage(clock: Clock) {
let cache = cache_with::<HashSet<String>>(&clock, 4);
let inserted = vec!["a".to_string(), "a".into(), "b".into()];
assert_eq!(
cache
.async_set_cache_sadd("members", inserted.clone(), None)
.async_set_cache_sadd("members", inserted.clone(), secs(10))
.await
.unwrap(),
inserted
);
assert_eq!(
cache
.async_set_cache_sadd("members", vec!["c".into()], secs(99))
.await
.unwrap(),
["c"]
);
assert_eq!(
cache.get_cache("members").unwrap(),
Some(HashSet::from(["a".into(), "b".into(), "c".into()]))
);
assert_eq!(cache.expires_at("members").unwrap(), secs(110));
at(&clock, 111);
cache
.async_set_cache_sadd("members", vec!["d".into()], None)
.await
.unwrap();
assert_eq!(
cache.get_cache("members").unwrap(),
Some(HashSet::from(["d".into()]))
);
assert_eq!(cache.expires_at("members").unwrap(), secs(171));
}
#[rstest]
#[tokio::test]
async fn oversized_set_additions_are_not_stored() {
let cache = InMemoryCache::<HashSet<String>>::with_clock_and_size_measurement(
Some(4),
None,
Some(2),
Some(Arc::new(|value: &HashSet<String>| Ok(value.len()))),
|| Duration::ZERO,
);
cache
.async_set_cache_sadd("members", vec!["a".into(), "b".into()], None)
.await
.unwrap();
assert_eq!(
cache
.async_set_cache_sadd("members", vec!["c".into()], None)
.await
.unwrap(),
["c"]
);
assert_eq!(
cache.get_cache("members").unwrap(),
Some(HashSet::from(["a".into(), "b".into()]))

98
litellm-rust/crates/cache-memory/tests/contract.rs Normal file
View file

@ -0,0 +1,98 @@
use std::time::Duration;
use litellm_cache::ExactCacheContext;
use litellm_cache_memory::InMemoryCache;
use litellm_cache_testing as contract;
use rstest::{fixture, rstest};
#[fixture]
fn strings() -> InMemoryCache<String> {
InMemoryCache::new(Some(16), None)
}
#[fixture]
fn counters() -> InMemoryCache<f64> {
InMemoryCache::new(Some(16), None)
}
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext {
ttl: Some(Duration::from_secs(60)),
}
}
#[rstest]
#[tokio::test]
async fn hit_and_miss(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::hit_and_miss(&strings, context, "memory:", "value".into()).await;
}
#[rstest]
#[tokio::test]
async fn sync_async_equivalence(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::sync_async_equivalence(
&strings,
context,
"memory:",
"first".into(),
"second".into(),
)
.await;
}
#[rstest]
#[tokio::test]
async fn overwrite_replaces(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::overwrite_replaces(
&strings,
context,
"memory:",
"first".into(),
"second".into(),
)
.await;
}
#[rstest]
#[tokio::test]
async fn pipeline_writes_every_entry(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::pipeline_writes_every_entry(
&strings,
context,
"memory:",
vec!["a".into(), "b".into(), "c".into()],
)
.await;
}
#[rstest]
#[tokio::test]
async fn batch_preserves_order(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::batch_preserves_order(
&strings,
context,
"memory:",
"first".into(),
"second".into(),
)
.await;
}
#[rstest]
#[tokio::test]
async fn delete_removes_key(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::delete_removes_key(&strings, context, "memory:", "value".into()).await;
}
#[rstest]
#[tokio::test]
async fn flush_clears(strings: InMemoryCache<String>, context: ExactCacheContext) {
contract::flush_clears(&strings, context, "memory:", "value".into()).await;
}
#[rstest]
#[tokio::test]
async fn counter_accumulates(counters: InMemoryCache<f64>, context: ExactCacheContext) {
contract::counter_accumulates(&counters, context, "memory:").await;
}

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

@ -17,7 +17,8 @@ tokio.workspace = true
uuid.workspace = true
[dev-dependencies]
litellm-cache-response.workspace = true
futures-executor = "0.3"
litellm-cache-testing.workspace = true
rstest.workspace = true
tonic = "0.14"
tonic-prost = "0.14"

105
litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs → litellm-rust/crates/cache-qdrant-semantic/src/cache.rs
View file

@ -1,7 +1,8 @@
use std::future::Future;
use futures_util::future::try_join_all;
use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext};
use litellm_cache::{
BaseCache, CacheCodec, Error, SemanticCacheContext,
semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_messages},
};
use qdrant_client::{
Payload, Qdrant,
qdrant::{
@ -14,26 +15,7 @@ use qdrant_client::{
use serde_json::{Map, Value, json};
use uuid::Uuid;
use crate::prompt_from_messages;
pub trait Embedder: Send + Sync + 'static {
fn model(&self) -> &str;
fn embed(&self, input: &str) -> impl Future<Output = Result<Vec<f32>, Error>> + Send;
}
#[derive(Clone, Debug, PartialEq)]
pub enum Quantization {
Binary,
Scalar,
Product,
}
pub struct QdrantSemanticConfig {
pub collection_name: String,
pub similarity_threshold: f64,
pub vector_size: u64,
pub quantization: Quantization,
}
use crate::{QdrantSemanticConfig, Quantization};
pub struct QdrantSemanticCache<E: Embedder, C: CacheCodec> {
client: Qdrant,
@ -100,14 +82,9 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
&self.embedder
}
/// Python reads `kwargs["messages"]` unguarded, so a request without messages fails.
fn prompt(context: &SemanticCacheContext) -> Result<String, Error> {
let Some(messages) = context.messages.as_ref().and_then(Value::as_array) else {
return Err(Error::MissingPrompt);
};
if messages.is_empty() {
return Err(Error::MissingPrompt);
}
Ok(prompt_from_messages(messages))
prompt_from_messages(context).ok_or(Error::MissingPrompt)
}
async fn set(
@ -117,7 +94,10 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
context: &SemanticCacheContext,
) -> Result<(), Error> {
let prompt = Self::prompt(context)?;
let vector = self.embedder.embed(&prompt).await?;
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let response =
String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?;
let payload = Payload::try_from(json!({
@ -147,9 +127,12 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
&self,
key: &str,
context: &SemanticCacheContext,
) -> Result<Option<C::Value>, Error> {
) -> Result<SemanticLookup<C::Value>, Error> {
let prompt = Self::prompt(context)?;
let vector = self.embedder.embed(&prompt).await?;
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let result = self
.client
.search_points(
@ -171,20 +154,27 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
.await
.map_err(|_| Error::Unavailable)?;
let Some(point) = result.result.into_iter().next() else {
return Ok(None);
return Ok(SemanticLookup::miss(Some(0.0)));
};
let payload: Map<String, Value> = Payload::from(point.payload).into();
if payload.get("litellm_cache_key").and_then(Value::as_str) != Some(key) {
return Ok(None);
if !payload
.get("litellm_cache_key")
.is_some_and(|cached| python_str(cached).as_deref() == Some(key))
{
return Ok(SemanticLookup::miss(Some(0.0)));
}
if f64::from(point.score) < self.config.similarity_threshold {
return Ok(None);
let similarity = f64::from(point.score);
if similarity < self.config.similarity_threshold {
return Ok(SemanticLookup::miss(Some(similarity)));
}
let response = payload
.get("response")
.and_then(Value::as_str)
.ok_or(Error::InvalidEntry)?;
self.codec.decode(response.as_bytes()).map(Some)
Ok(SemanticLookup {
value: Some(self.codec.decode(response.as_bytes())?),
similarity: Some(similarity),
})
}
}
@ -219,7 +209,8 @@ impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.runtime.block_on(self.get(key, context))
self.get_cache_with_similarity(key, context)
.map(|lookup| lookup.value)
}
async fn async_set_cache(
@ -236,7 +227,7 @@ impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
key: &str,
context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.get(key, context).await
self.get(key, context).await.map(|lookup| lookup.value)
}
async fn async_set_cache_pipeline(
@ -251,12 +242,36 @@ impl<E: Embedder, C: CacheCodec> BaseCache for QdrantSemanticCache<E, C> {
.await
.map(|_| ())
}
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
/// Python stamps the top point's score, even below the threshold, and `0.0` when there is no
/// point or it belongs to another key. A request without messages fails before any search.
impl<E: Embedder, C: CacheCodec> SemanticCache for QdrantSemanticCache<E, C> {
fn get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
self.runtime.block_on(self.get(key, context))
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
async fn async_get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
self.get(key, context).await
}
}
/// `str(value)` for the scalar payload values `_payload_matches_cache_key` compares; `None` for
/// null (a pre-isolation point without a key) and for containers, which never equal a key.
fn python_str(value: &Value) -> Option<String> {
match value {
Value::String(text) => Some(text.clone()),
Value::Number(number) => Some(number.to_string()),
Value::Bool(true) => Some("True".into()),
Value::Bool(false) => Some("False".into()),
Value::Null | Value::Array(_) | Value::Object(_) => None,
}
}

13
litellm-rust/crates/cache-qdrant-semantic/src/config.rs Normal file
View file

@ -0,0 +1,13 @@
#[derive(Clone, Debug, PartialEq)]
pub enum Quantization {
Binary,
Scalar,
Product,
}
pub struct QdrantSemanticConfig {
pub collection_name: String,
pub similarity_threshold: f64,
pub vector_size: u64,
pub quantization: Quantization,
}

14
litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs
View file

@ -1,11 +1,9 @@
use std::time::Duration;
use litellm_cache::Error;
use litellm_cache::{Error, semantic::Embedder};
use reqwest::Client;
use serde_json::Value;
use crate::Embedder;
pub struct OpenAiEmbedder {
client: Client,
api_base: String,
@ -31,14 +29,16 @@ impl OpenAiEmbedder {
timeout: config.timeout,
}
}
}
impl Embedder for OpenAiEmbedder {
fn model(&self) -> &str {
pub fn model(&self) -> &str {
&self.model
}
}
async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
/// An OpenAI-compatible `/embeddings` call. It has no router to route on, so `metadata` is
/// unused, and it only embeds asynchronously: sync cache calls block on the cache's runtime.
impl Embedder for OpenAiEmbedder {
async fn async_embed(&self, input: &str, _metadata: Option<&Value>) -> Result<Vec<f32>, Error> {
let request = self
.client
.post(format!("{}/embeddings", self.api_base))

8
litellm-rust/crates/cache-qdrant-semantic/src/lib.rs
View file

@ -1,7 +1,7 @@
mod cache;
mod config;
mod embedder;
mod prompt;
mod semantic;
pub use cache::QdrantSemanticCache;
pub use config::{QdrantSemanticConfig, Quantization};
pub use embedder::{OpenAiEmbedder, OpenAiEmbedderConfig};
pub use prompt::prompt_from_messages;
pub use semantic::{Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization};

59
litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs
View file

@ -1,59 +0,0 @@
use serde_json::Value;
fn search_results_text(search_results: Option<&Value>) -> String {
let Some(Value::Array(results)) = search_results else {
return String::new();
};
results
.iter()
.filter_map(Value::as_object)
.flat_map(|result| {
let source = result
.get("source")
.and_then(Value::as_str)
.map(str::to_owned);
let title = result
.get("title")
.and_then(Value::as_str)
.map(str::to_owned);
let content = result
.get("content")
.and_then(Value::as_array)
.into_iter()
.flatten()
.filter_map(Value::as_object)
.filter_map(|block| block.get("text").and_then(Value::as_str).map(str::to_owned));
let citations = result
.get("citations")
.filter(|value| !value.is_null())
.map(|value| serde_json::to_string(value).unwrap_or_default());
source
.into_iter()
.chain(title)
.chain(content)
.chain(citations)
})
.collect()
}
pub fn prompt_from_messages(messages: &[Value]) -> String {
messages
.iter()
.filter_map(Value::as_object)
.map(|message| {
let content = match message.get("content") {
Some(Value::String(content)) => content.clone(),
Some(Value::Array(parts)) => parts
.iter()
.filter_map(Value::as_object)
.filter_map(|part| part.get("text").and_then(Value::as_str))
.collect(),
_ => String::new(),
};
format!(
"{content}{}",
search_results_text(message.get("search_results"))
)
})
.collect()
}

91
litellm-rust/crates/cache-qdrant-semantic/tests/contract.rs Normal file
View file

@ -0,0 +1,91 @@
//! `overwrite_replaces` does not apply: like Python, every write upserts a new `uuid4` point,
//! so a second write with the same prompt adds a tie instead of replacing the first.
mod support;
use std::future::Future;
use litellm_cache::{JsonCodec, SemanticCacheContext, semantic::PreparedEmbedding};
use litellm_cache_qdrant_semantic::{QdrantSemanticCache, QdrantSemanticConfig, Quantization};
use litellm_cache_testing as contract;
use qdrant_client::Qdrant;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use support::{FakeQdrant, FakeState};
type Cache = QdrantSemanticCache<PreparedEmbedding, JsonCodec<Value>>;
const PREFIX: &str = "contract:";
#[fixture]
fn context() -> SemanticCacheContext {
SemanticCacheContext {
messages: Some(json!([{"role": "user", "content": "contract prompt"}])),
..Default::default()
}
}
/// Runs a contract against a fresh fake Qdrant. The sync cache methods block on the runtime, so
/// the contract is polled on a blocking thread outside the runtime's own executor.
async fn run<F, Fut>(check: F)
where
F: FnOnce(Cache) -> Fut + Send + 'static,
Fut: Future<Output = ()>,
{
let server = FakeQdrant::start(FakeState::default()).await;
let runtime = tokio::runtime::Handle::current();
let cache = QdrantSemanticCache::connect(
Qdrant::from_url(&server.url()).build().unwrap(),
PreparedEmbedding(vec![0.6, 0.8]),
JsonCodec::new(),
QdrantSemanticConfig {
collection_name: "contract".to_owned(),
similarity_threshold: 0.9,
vector_size: 2,
quantization: Quantization::Binary,
},
runtime.clone(),
)
.await
.unwrap();
tokio::task::spawn_blocking(move || {
let _guard = runtime.enter();
futures_executor::block_on(check(cache));
})
.await
.unwrap();
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn hit_and_miss(context: SemanticCacheContext) {
run(|cache| async move {
contract::hit_and_miss(&cache, context, PREFIX, json!({"answer": 42})).await;
})
.await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn sync_async_equivalence(context: SemanticCacheContext) {
run(|cache| async move {
contract::sync_async_equivalence(&cache, context, PREFIX, json!("first"), json!([2])).await;
})
.await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn pipeline_writes_every_entry(context: SemanticCacheContext) {
run(|cache| async move {
contract::pipeline_writes_every_entry(
&cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
})
.await;
}

91
litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs
View file

@ -3,9 +3,10 @@ use std::{
time::Duration,
};
use litellm_cache::Error;
use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, OpenAiEmbedderConfig};
use serde_json::Value;
use litellm_cache::{Error, semantic::Embedder};
use litellm_cache_qdrant_semantic::{OpenAiEmbedder, OpenAiEmbedderConfig};
use rstest::rstest;
use serde_json::{Value, json};
use tokio::{
io::{AsyncReadExt, AsyncWriteExt},
net::TcpListener,
@ -98,6 +99,7 @@ fn config(base: String, timeout: Option<Duration>) -> OpenAiEmbedderConfig {
}
}
#[rstest]
#[tokio::test]
async fn posts_embeddings_request_and_parses_vector() {
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
@ -108,7 +110,14 @@ async fn posts_embeddings_request_and_parses_vector() {
Some(Duration::from_secs(1)),
),
);
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
assert_eq!(embedder.model(), "test-model");
assert_eq!(
embedder
.async_embed("hello", Some(&json!({"ignored": true})))
.await
.unwrap(),
vec![0.1, 0.2]
);
let request = server.request.lock().unwrap().clone().unwrap();
let request_text = String::from_utf8(request).unwrap();
assert!(request_text.starts_with("POST /embeddings HTTP/1.1\r\n"));
@ -120,37 +129,50 @@ async fn posts_embeddings_request_and_parses_vector() {
assert_eq!(body["encoding_format"], "float");
}
#[rstest]
#[case::error_status("500 Internal Server Error", "{}", 0, None, Err(Error::Unavailable))]
#[case::timed_out(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
500,
Some(Duration::from_millis(200)),
Err(Error::Unavailable)
)]
#[case::within_timeout(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
100,
Some(Duration::from_secs(1)),
Ok(vec![0.1, 0.2])
)]
#[case::missing_embedding("200 OK", r#"{"data":[]}"#, 0, None, Err(Error::Unavailable))]
#[tokio::test]
async fn status_and_timeout_errors_are_unavailable() {
let server = TestHttpServer::response("500 Internal Server Error", "{}").await;
let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), None));
assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
let server = TestHttpServer::response_after(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
Duration::from_millis(500),
)
.await;
let embedder = OpenAiEmbedder::new(
reqwest::Client::new(),
config(server.base_url(), Some(Duration::from_millis(200))),
);
assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable));
let server = TestHttpServer::response_after(
"200 OK",
r#"{"data":[{"embedding":[0.1,0.2]}]}"#,
Duration::from_millis(100),
)
.await;
let embedder = OpenAiEmbedder::new(
reqwest::Client::new(),
config(server.base_url(), Some(Duration::from_secs(1))),
);
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
async fn status_timeout_and_body_errors_are_unavailable(
#[case] status: &str,
#[case] body: &str,
#[case] delay_ms: u64,
#[case] timeout: Option<Duration>,
#[case] expected: Result<Vec<f32>, Error>,
) {
let server =
TestHttpServer::response_after(status, body, Duration::from_millis(delay_ms)).await;
let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), timeout));
assert_eq!(embedder.async_embed("hello", None).await, expected);
}
#[rstest]
fn sync_embedding_is_unsupported() {
let embedder = OpenAiEmbedder::new(
reqwest::Client::new(),
config("http://127.0.0.1:9".to_owned(), None),
);
assert_eq!(
embedder.embed("hello", None),
Err(Error::UnsupportedOperation)
);
}
#[rstest]
#[tokio::test]
async fn uses_the_injected_client() {
let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await;
@ -159,7 +181,10 @@ async fn uses_the_injected_client() {
.build()
.unwrap();
let embedder = OpenAiEmbedder::new(client, config(server.base_url(), None));
assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]);
assert_eq!(
embedder.async_embed("hello", None).await.unwrap(),
vec![0.1, 0.2]
);
let request = server.request.lock().unwrap().clone().unwrap();
let request_text = String::from_utf8(request).unwrap();
assert!(request_text.contains("\r\nuser-agent: litellm-embedder-test\r\n"));

38
litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs
View file

@ -1,38 +0,0 @@
use litellm_cache_qdrant_semantic::prompt_from_messages;
use serde_json::json;
#[test]
fn prompt_matches_python_message_content_rules() {
let messages = vec![
json!({"role": "user", "content": "hello"}),
json!({
"role": "user",
"content": [
{"type": "text", "text": "world"},
{"type": "image_url", "image_url": {"url": "ignored"}},
{"type": "text", "text": "!"},
],
}),
];
assert_eq!(prompt_from_messages(&messages), "helloworld!");
}
#[test]
fn prompt_includes_search_result_text_and_compact_citations() {
let messages = vec![json!({
"role": "tool",
"content": null,
"search_results": [{
"source": "source",
"title": "title",
"content": [{"text": "body"}],
"citations": {"page": 1, "section": "intro"},
}],
})];
assert_eq!(
prompt_from_messages(&messages),
r#"sourcetitlebody{"page":1,"section":"intro"}"#
);
}

595
litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs
View file

@ -1,27 +1,32 @@
#[path = "support/mod.rs"]
mod support;
use std::{collections::HashMap, sync::Arc, time::Duration};
use std::{
collections::HashMap,
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::{BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext};
use litellm_cache_qdrant_semantic::{
Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization,
use litellm_cache::{
BaseCache, CacheContext, Error, JsonCodec, SemanticCacheContext,
semantic::{Embedder, SemanticCache, SemanticLookup},
};
use litellm_cache_response::{
CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest,
};
use qdrant_client::Payload;
use litellm_cache_qdrant_semantic::{QdrantSemanticCache, QdrantSemanticConfig, Quantization};
use qdrant_client::{
Qdrant,
Payload, Qdrant,
qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams},
};
use rstest::{fixture, rstest};
use serde_json::{Value as JsonValue, json};
use support::{FakeQdrant, FakeState, StoredPoint};
type Calls = Arc<Mutex<Vec<(String, Option<JsonValue>)>>>;
type Cache = QdrantSemanticCache<FixedEmbedder, JsonCodec<JsonValue>>;
/// Embeds known prompts, fails on anything else, and records every call.
#[derive(Clone)]
struct FixedEmbedder {
vectors: Arc<HashMap<String, Vec<f32>>>,
calls: Calls,
}
impl FixedEmbedder {
@ -33,16 +38,21 @@ impl FixedEmbedder {
.map(|(prompt, vector)| (prompt.to_owned(), vector))
.collect(),
),
calls: Calls::default(),
}
}
}
impl Embedder for FixedEmbedder {
fn model(&self) -> &str {
"fixed"
}
async fn embed(&self, input: &str) -> Result<Vec<f32>, Error> {
async fn async_embed(
&self,
input: &str,
metadata: Option<&JsonValue>,
) -> Result<Vec<f32>, Error> {
self.calls
.lock()
.unwrap()
.push((input.to_owned(), metadata.cloned()));
self.vectors.get(input).cloned().ok_or(Error::Unavailable)
}
}
@ -63,22 +73,20 @@ fn context(prompt: &str) -> SemanticCacheContext {
}
}
fn value(response: JsonValue) -> CacheEntry {
CacheEntry {
timestamp: Some(1.0),
response,
}
#[fixture]
fn entry() -> JsonValue {
json!({"timestamp": 1.0, "response": {"answer": 42}})
}
async fn connect(
server: &FakeQdrant,
vectors: impl IntoIterator<Item = (&'static str, Vec<f32>)>,
) -> QdrantSemanticCache<FixedEmbedder, ResponseCacheCodec> {
) -> Cache {
let client = Qdrant::from_url(&server.url()).build().unwrap();
QdrantSemanticCache::connect(
client,
FixedEmbedder::new(vectors),
ResponseCacheCodec,
JsonCodec::new(),
config(Quantization::Binary),
tokio::runtime::Handle::current(),
)
@ -86,72 +94,70 @@ async fn connect(
.unwrap()
}
#[rstest]
#[case::binary(Quantization::Binary)]
#[case::scalar(Quantization::Scalar)]
#[case::product(Quantization::Product)]
#[tokio::test(flavor = "multi_thread")]
#[expect(
deprecated,
reason = "the test verifies Qdrant's legacy always_ram quantization contract"
)]
async fn connect_sets_collection_quantization_and_index() {
for (quantization, expected) in [
(Quantization::Binary, 0),
(Quantization::Scalar, 1),
(Quantization::Product, 2),
] {
let server = FakeQdrant::start(FakeState::default()).await;
let client = Qdrant::from_url(&server.url()).build().unwrap();
QdrantSemanticCache::connect(
client,
FixedEmbedder::new([]),
ResponseCacheCodec,
config(quantization),
tokio::runtime::Handle::current(),
)
.await
async fn connect_sets_collection_quantization_and_index(#[case] quantization: Quantization) {
let server = FakeQdrant::start(FakeState::default()).await;
let client = Qdrant::from_url(&server.url()).build().unwrap();
QdrantSemanticCache::connect(
client,
FixedEmbedder::new([]),
JsonCodec::<JsonValue>::new(),
config(quantization.clone()),
tokio::runtime::Handle::current(),
)
.await
.unwrap();
let state = server.state.lock().unwrap();
let request = &state.created_collections[0];
let Some(qdrant::vectors_config::Config::Params(VectorParams { size, distance, .. })) = request
.vectors_config
.as_ref()
.and_then(|config| config.config.clone())
else {
panic!("missing vector params");
};
assert_eq!(size, 2);
assert_eq!(distance, Distance::Cosine as i32);
let quantization_config = request
.quantization_config
.as_ref()
.unwrap()
.quantization
.unwrap();
let state = server.state.lock().unwrap();
let request = &state.created_collections[0];
let Some(qdrant::vectors_config::Config::Params(VectorParams { size, distance, .. })) =
request
.vectors_config
.as_ref()
.and_then(|config| config.config.clone())
else {
panic!("missing vector params");
};
assert_eq!(size, 2);
assert_eq!(distance, Distance::Cosine as i32);
let quantization_config = request
.quantization_config
.as_ref()
.unwrap()
.quantization
.unwrap();
match (expected, quantization_config) {
(0, qdrant::quantization_config::Quantization::Binary(binary)) => {
assert_eq!(binary.always_ram, Some(false));
}
(1, qdrant::quantization_config::Quantization::Scalar(scalar)) => {
assert_eq!(scalar.r#type, QuantizationType::Int8 as i32);
assert_eq!(scalar.quantile, Some(0.99));
assert_eq!(scalar.always_ram, Some(false));
}
(2, qdrant::quantization_config::Quantization::Product(product)) => {
assert_eq!(product.compression, CompressionRatio::X16 as i32);
assert_eq!(product.always_ram, Some(false));
}
_ => panic!("unexpected quantization"),
#[expect(
deprecated,
reason = "the test verifies Qdrant's legacy always_ram quantization contract"
)]
match (quantization, quantization_config) {
(Quantization::Binary, qdrant::quantization_config::Quantization::Binary(binary)) => {
assert_eq!(binary.always_ram, Some(false));
}
assert!(state.index_creations >= 1);
assert_eq!(state.field_indexes[0].collection_name, "semantic");
assert_eq!(state.field_indexes[0].field_name, "litellm_cache_key");
assert_eq!(
state.field_indexes[0].field_type,
Some(qdrant::FieldType::Keyword as i32)
);
server.stop();
(Quantization::Scalar, qdrant::quantization_config::Quantization::Scalar(scalar)) => {
assert_eq!(scalar.r#type, QuantizationType::Int8 as i32);
assert_eq!(scalar.quantile, Some(0.99));
assert_eq!(scalar.always_ram, Some(false));
}
(Quantization::Product, qdrant::quantization_config::Quantization::Product(product)) => {
assert_eq!(product.compression, CompressionRatio::X16 as i32);
assert_eq!(product.always_ram, Some(false));
}
_ => panic!("unexpected quantization"),
}
assert!(state.index_creations >= 1);
assert_eq!(state.field_indexes[0].collection_name, "semantic");
assert_eq!(state.field_indexes[0].field_name, "litellm_cache_key");
assert_eq!(
state.field_indexes[0].field_type,
Some(qdrant::FieldType::Keyword as i32)
);
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn existing_collection_skips_create_and_index_failure_is_non_fatal() {
let server = FakeQdrant::start(FakeState {
@ -160,19 +166,25 @@ async fn existing_collection_skips_create_and_index_failure_is_non_fatal() {
..Default::default()
})
.await;
let _cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
assert_eq!(cache.collection_name(), "semantic");
assert_eq!(cache.similarity_threshold(), 0.9);
assert_eq!(cache.vector_size(), 2);
let state = server.state.lock().unwrap();
assert!(state.created_collections.is_empty());
assert!(state.index_creations >= 1);
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn async_and_sync_set_get_store_exact_payload() {
async fn async_and_sync_set_get_store_exact_payload(entry: JsonValue) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
let ctx = context("hello");
let entry = value(json!({"answer": 42}));
let ctx = SemanticCacheContext {
metadata: Some(json!({"tenant": "team"})),
..context("hello")
};
cache
.async_set_cache("key", entry.clone(), ctx.clone())
.await
@ -188,10 +200,8 @@ async fn async_and_sync_set_get_store_exact_payload() {
payload_keys.sort();
assert_eq!(payload_keys, ["litellm_cache_key", "response", "text"]);
assert_eq!(payload["litellm_cache_key"], Value::from("key"));
assert_eq!(
payload["response"],
Value::from(String::from_utf8(ResponseCacheCodec.encode(&entry).unwrap()).unwrap())
);
assert_eq!(payload["text"], Value::from("hello"));
assert_eq!(payload["response"], Value::from(entry.to_string()));
}
let sync_entry = entry.clone();
let sync_cache = cache.clone();
@ -207,208 +217,276 @@ async fn async_and_sync_set_get_store_exact_payload() {
})
.await
.unwrap();
assert_eq!(
*cache.embedder().calls.lock().unwrap(),
vec![("hello".to_owned(), ctx.metadata.clone()); 4]
);
server.stop();
}
#[rstest]
#[case::content_parts_skip_images(
json!([
{"role": "user", "content": "hello"},
{
"role": "user",
"content": [
{"type": "text", "text": "world"},
{"type": "image_url", "image_url": {"url": "ignored"}},
{"type": "text", "text": "!"},
],
},
]),
"helloworld!"
)]
#[case::search_results_and_compact_citations(
json!([{
"role": "tool",
"content": null,
"search_results": [{
"source": "source",
"title": "title",
"content": [{"text": "body"}],
"citations": {"page": 1, "section": "intro"},
}],
}]),
r#"sourcetitlebody{"page":1,"section":"intro"}"#
)]
#[tokio::test(flavor = "multi_thread")]
async fn misses_and_payload_validation_are_safe() {
async fn prompt_matches_python_message_rules(
#[case] messages: JsonValue,
#[case] prompt: &'static str,
entry: JsonValue,
) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [(prompt, vec![1.0, 0.0])]).await;
let context = SemanticCacheContext {
messages: Some(messages),
..Default::default()
};
cache.async_set_cache("key", entry, context).await.unwrap();
assert_eq!(cache.embedder().calls.lock().unwrap()[0].0, prompt);
assert_eq!(
server.state.lock().unwrap().points[0].payload["text"],
Value::from(prompt)
);
server.stop();
}
#[rstest]
#[case::no_messages(SemanticCacheContext::default())]
#[case::empty_messages(SemanticCacheContext { messages: Some(json!([])), ..Default::default() })]
#[case::responses_input_is_not_read(SemanticCacheContext { input: Some(json!("hello")), ..Default::default() })]
#[tokio::test(flavor = "multi_thread")]
async fn requests_without_messages_are_missing_a_prompt(
#[case] context: SemanticCacheContext,
entry: JsonValue,
) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
assert_eq!(
cache.async_set_cache("key", entry, context.clone()).await,
Err(Error::MissingPrompt)
);
assert_eq!(
cache.async_get_cache("key", &context).await,
Err(Error::MissingPrompt)
);
assert!(cache.embedder().calls.lock().unwrap().is_empty());
server.stop();
}
#[rstest]
#[case::other_key("other", "hello", None)]
#[case::below_similarity_threshold("key", "near", None)]
#[tokio::test(flavor = "multi_thread")]
async fn misses_and_payload_validation_are_safe(
#[case] key: &str,
#[case] prompt: &str,
#[case] numeric_key_point: Option<u64>,
entry: JsonValue,
) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(
&server,
[("hello", vec![1.0, 0.0]), ("near", vec![0.7, 0.71414286])],
)
.await;
let entry = value(json!({"answer": 1}));
cache
.async_set_cache("key", entry, context("hello"))
.await
.unwrap();
if let Some(id) = numeric_key_point {
server.insert_point(StoredPoint {
id: Some(PointId::from(id)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(json!({
"litellm_cache_key": id,
"response": "{}",
}))
.unwrap()
.into(),
});
}
assert_eq!(
cache
.async_get_cache("other", &context("hello"))
.await
.unwrap(),
None
);
assert_eq!(
cache
.async_get_cache("key", &context("near"))
.await
.unwrap(),
None
);
server.insert_point(StoredPoint {
id: Some(PointId::from(99_u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(json!({
"litellm_cache_key": 99,
"response": "{}",
}))
.unwrap()
.into(),
});
assert_eq!(
cache
.async_get_cache("99", &context("hello"))
.await
.unwrap(),
cache.async_get_cache(key, &context(prompt)).await.unwrap(),
None
);
server.stop();
}
#[rstest]
#[case::hit("key", context("hello"), Ok((true, Some(1.0))))]
#[case::below_similarity_threshold("key", context("near"), Ok((false, Some(0.7))))]
#[case::no_results("other", context("hello"), Ok((false, Some(0.0))))]
#[case::no_prompt("key", SemanticCacheContext::default(), Err(Error::MissingPrompt))]
#[tokio::test(flavor = "multi_thread")]
async fn decoding_errors_missing_prompt_pipeline_and_ttl_behave_as_required() {
async fn lookup_reports_python_semantic_similarity(
#[case] key: &'static str,
#[case] context: SemanticCacheContext,
#[case] expected: Result<(bool, Option<f64>), Error>,
#[values(false, true)] use_async: bool,
entry: JsonValue,
) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("one", vec![1.0, 0.0]), ("two", vec![0.0, 1.0])]).await;
let empty = SemanticCacheContext::default();
assert_eq!(
cache
.async_set_cache("key", value(json!({})), empty.clone())
.await,
Err(Error::MissingPrompt)
let cache = Arc::new(
connect(
&server,
[("hello", vec![1.0, 0.0]), ("near", vec![0.7, 0.71414286])],
)
.await,
);
cache
.async_set_cache("key", entry.clone(), self::context("hello"))
.await
.unwrap();
server.insert_point(StoredPoint {
id: Some(PointId::from(99_u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(json!({"litellm_cache_key": 99, "response": "{}"}))
.unwrap()
.into(),
});
let lookup = if use_async {
cache.async_get_cache_with_similarity(key, &context).await
} else {
let cache = Arc::clone(&cache);
tokio::task::spawn_blocking(move || cache.get_cache_with_similarity(key, &context))
.await
.unwrap()
};
match (lookup, expected) {
(Ok(SemanticLookup { value, similarity }), Ok((hit, expected))) => {
assert_eq!(value, hit.then_some(entry));
assert_eq!(similarity.is_some(), expected.is_some());
if let (Some(similarity), Some(expected)) = (similarity, expected) {
assert!((similarity - expected).abs() < 1e-6, "{similarity}");
}
}
(lookup, expected) => assert_eq!(lookup.map(|_| ()), expected.map(|_| ())),
}
server.stop();
}
#[rstest]
#[case::codec_decodes_the_payload(Some(json!("{\"a\":1}")), Ok(Some(json!({"a": 1}))))]
#[case::undecodable_response(Some(json!("not json")), Err(Error::InvalidEntry))]
#[case::non_string_response(Some(json!(1)), Err(Error::InvalidEntry))]
#[case::missing_response(None, Err(Error::InvalidEntry))]
#[tokio::test(flavor = "multi_thread")]
async fn stored_responses_go_through_the_codec(
#[case] response: Option<JsonValue>,
#[case] expected: Result<Option<JsonValue>, Error>,
) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
let mut payload = serde_json::Map::new();
payload.insert("litellm_cache_key".to_owned(), json!("key"));
if let Some(response) = response {
payload.insert("response".to_owned(), response);
}
server.insert_point(StoredPoint {
id: Some(PointId::from(1_u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(JsonValue::Object(payload))
.unwrap()
.into(),
});
assert_eq!(
cache.async_get_cache("key", &empty).await,
Err(Error::MissingPrompt)
cache.async_get_cache("key", &context("hello")).await,
expected
);
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn embedding_failures_propagate() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, []).await;
assert_eq!(
cache.async_get_cache("key", &context("unknown")).await,
Err(Error::Unavailable)
);
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn ttl_is_ignored_and_entries_do_not_expire(entry: JsonValue) {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("one", vec![1.0, 0.0])]).await;
let ctx = context("one").with_ttl(Some(Duration::from_secs(1)));
assert_eq!(cache.get_ttl(&ctx), None);
cache
.async_set_cache(
"ttl",
value(json!({"ttl": true})),
context("one").with_ttl(Some(Duration::from_secs(1))),
)
.async_set_cache("ttl", entry, ctx.clone())
.await
.unwrap();
tokio::time::sleep(Duration::from_millis(1_100)).await;
assert!(
cache
.async_get_cache(
"ttl",
&context("one").with_ttl(Some(Duration::from_secs(1))),
)
.await
.unwrap()
.is_some()
);
assert!(cache.async_get_cache("ttl", &ctx).await.unwrap().is_some());
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn pipeline_upserts_each_entry_and_waits_for_indexing() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("one", vec![1.0, 0.0])]).await;
cache
.async_set_cache_pipeline(
vec![
("one".to_owned(), value(json!({"n": 1}))),
("two".to_owned(), value(json!({"n": 2}))),
("one".to_owned(), json!({"n": 1})),
("two".to_owned(), json!({"n": 2})),
],
context("one"),
)
.await
.unwrap();
assert!(
cache
.async_get_cache("one", &context("one"))
.await
.unwrap()
.is_some()
);
assert!(
cache
.async_get_cache("two", &context("one"))
.await
.unwrap()
.is_some()
);
for (key, value) in [("one", json!({"n": 1})), ("two", json!({"n": 2}))] {
assert_eq!(
cache.async_get_cache(key, &context("one")).await.unwrap(),
Some(value)
);
}
assert_eq!(
server.state.lock().unwrap().upsert_waits,
vec![Some(true), Some(true), Some(true)]
);
assert_eq!(cache.get_ttl(&context("one")), None);
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
);
server.stop();
}
#[tokio::test(flavor = "multi_thread")]
async fn response_payloads_decode_and_invalid_entries_fail() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
for (key, response) in [
("python", json!("{'timestamp': 1.0, 'response': {'a': 1}}")),
("garbage", json!("not json")),
("missing", json!("unused")),
] {
let mut payload = serde_json::Map::new();
payload.insert("litellm_cache_key".to_owned(), json!(key));
if key != "missing" {
payload.insert("response".to_owned(), response);
}
server.insert_point(StoredPoint {
id: Some(PointId::from(key.len() as u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(JsonValue::Object(payload))
.unwrap()
.into(),
});
}
assert_eq!(
cache
.async_get_cache("python", &context("hello"))
.await
.unwrap(),
Some(value(json!({"a": 1})))
);
assert_eq!(
cache.async_get_cache("garbage", &context("hello")).await,
Err(Error::InvalidEntry)
);
assert_eq!(
cache.async_get_cache("missing", &context("hello")).await,
Err(Error::InvalidEntry)
);
server.stop();
}
#[tokio::test(flavor = "multi_thread")]
async fn response_cache_facade_turns_invalid_entry_into_miss() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await);
let request = ResponseCacheRequest::<SemanticCacheContext>::new(CacheKeyInput {
preset: Some("key".to_owned()),
..Default::default()
})
.with_context(context("hello"));
let response = json!({"answer": 42});
let facade = ResponseCache::new(cache.clone());
facade
.async_store(&request, response.clone(), Duration::from_secs(1))
.await
.unwrap();
assert_eq!(
facade
.async_lookup(&request, Duration::from_secs(1))
.await
.unwrap(),
Some(response)
);
{
let mut state = server.state.lock().unwrap();
state.points[0]
.payload
.insert("response".to_owned(), Value::from("not json"));
}
assert_eq!(
facade
.async_lookup(&request, Duration::from_secs(1))
.await
.unwrap(),
None
vec![Some(true), Some(true)]
);
server.stop();
}
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn stopped_qdrant_server_maps_to_unavailable() {
let server = FakeQdrant::start(FakeState::default()).await;
@ -420,3 +498,28 @@ async fn stopped_qdrant_server_maps_to_unavailable() {
Err(Error::Unavailable)
);
}
/// `_payload_matches_cache_key` compares `str(cached_key) == str(key)`, so a point whose stored
/// key is the number 99 answers a lookup for `"99"`.
#[rstest]
#[tokio::test(flavor = "multi_thread")]
async fn numeric_stored_cache_keys_match_like_python_str() {
let server = FakeQdrant::start(FakeState::default()).await;
let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await;
server.insert_point(StoredPoint {
id: Some(PointId::from(99_u64)),
vector: vec![1.0, 0.0],
payload: Payload::try_from(json!({"litellm_cache_key": 99, "response": "{}"}))
.unwrap()
.into(),
});
let lookup = cache
.async_get_cache_with_similarity("99", &context("hello"))
.await
.unwrap();
assert_eq!(lookup.value, Some(json!({})));
assert!((lookup.similarity.unwrap() - 1.0).abs() < 1e-6);
server.stop();
}

5
litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs
View file

@ -1,15 +1,16 @@
#![allow(dead_code)]
use std::{
collections::{HashMap, HashSet},
net::SocketAddr,
sync::{Arc, Mutex},
};
use qdrant_client::qdrant::collections_server::CollectionsServer;
use qdrant_client::qdrant::{
self, CollectionExists, CollectionExistsRequest, CollectionExistsResponse,
CollectionOperationResponse, CreateCollection, CreateFieldIndexCollection, Filter, PointId,
PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors,
collections_server::Collections,
collections_server::{Collections, CollectionsServer},
points_server::{Points, PointsServer},
};
use tokio::sync::oneshot;

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

@ -8,14 +8,12 @@ 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]
litellm-cache-testing.workspace = true
redis-test = "1.0.4"
rstest.workspace = true
serde_json.workspace = true
tokio.workspace = true

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

@ -1,105 +1,36 @@
use std::{
future::Future,
sync::{Arc, OnceLock},
sync::Arc,
time::{Duration, SystemTime, UNIX_EPOCH},
};
use litellm_cache::{
BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
SemanticCacheContext,
BaseCache, CacheCodec, Error, SemanticCacheContext,
semantic::{Embedder, SemanticCache, SemanticLookup, prompt_from_context},
};
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,
}
use crate::{
RedisSemanticConfig,
index::{CACHE_KEY_FIELD, Index, VECTOR_FIELD},
reply::{bytes_field, first_document, number_field, string_field},
};
struct Inner {
index_name: String,
index: Index,
distance_threshold: f64,
resolved_index: OnceLock<String>,
codec: ResponseCacheCodec,
clock: fn() -> f64,
}
impl Inner {
fn new(config: RedisSemanticConfig) -> Self {
fn new(config: RedisSemanticConfig, clock: fn() -> f64) -> Self {
Self {
index_name: config.index_name,
index: Index::new(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)
}
clock,
}
}
@ -107,15 +38,14 @@ impl Inner {
&self,
connection: &mut ConnectionRef<'_>,
tag: &str,
value: &CacheEntry,
response: Vec<u8>,
prompt: &str,
vector: &[f32],
ttl: Option<Duration>,
) -> Result<(), Error> {
let index = self.ensure_index(connection, vector.len())?;
let index = self.index.ensure(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")
@ -149,8 +79,8 @@ impl Inner {
connection: &mut ConnectionRef<'_>,
tag: &str,
vector: &[f32],
) -> Result<Option<CacheEntry>, Error> {
let index = self.ensure_index(connection, vector.len())?;
) -> Result<SemanticLookup<Vec<u8>>, Error> {
let index = self.index.ensure(connection, vector.len())?;
let query = format!(
"(@{CACHE_KEY_FIELD}:{{{}}})=>[KNN 1 @{VECTOR_FIELD} $vector AS vector_distance]",
escape_tag(tag)
@ -183,57 +113,80 @@ impl Inner {
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)?;
let Some(fields) = first_document(&result) else {
return Ok(None);
return Ok(SemanticLookup::miss(Some(0.0)));
};
if string_field(fields, CACHE_KEY_FIELD).as_deref() != Some(tag) {
return Ok(None);
return Ok(SemanticLookup::miss(Some(0.0)));
}
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);
// redisvl's range query only returns entries within the distance threshold, so a
// farther hit reads as no result.
let Some(distance) = number_field(fields, "vector_distance")
.filter(|distance| *distance <= self.distance_threshold)
else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
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)),
let Some(response) = bytes_field(fields, "response") else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
Ok(SemanticLookup {
value: Some(response),
similarity: Some(1.0 - distance),
})
}
}
impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<E, C> {
pub fn with_connection(connection: C, embedder: E, config: RedisSemanticConfig) -> Self {
/// `RedisSemanticCache`: a redisvl-compatible semantic index on Redis Stack. Values go through
/// the injected codec, so the response layer decides what a cached entry is.
pub struct RedisSemanticCache<E, S, C = redis::Connection> {
connections: Arc<Connections<C>>,
embedder: E,
codec: S,
inner: Arc<Inner>,
}
impl<E: Embedder, S: CacheCodec> RedisSemanticCache<E, S> {
pub fn new(
url: &str,
embedder: E,
codec: S,
config: RedisSemanticConfig,
) -> Result<Self, Error> {
Ok(Self {
connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?),
embedder,
codec,
inner: Arc::new(Inner::new(config, timestamp)),
})
}
}
impl<E, S, C> RedisSemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub fn with_connection(
connection: C,
embedder: E,
codec: S,
config: RedisSemanticConfig,
) -> Self {
Self {
connections: Arc::new(Connections::fixed(connection)),
embedder,
inner: Arc::new(Inner::new(config)),
codec,
inner: Arc::new(Inner::new(config, timestamp)),
}
}
pub fn with_clock(self, clock: fn() -> f64) -> Self {
let config = RedisSemanticConfig {
index_name: self.index_name().to_owned(),
similarity_threshold: self.similarity_threshold(),
};
Self {
inner: Arc::new(Inner {
index_name: self.inner.index_name.clone(),
distance_threshold: self.inner.distance_threshold,
resolved_index: OnceLock::new(),
codec: self.inner.codec,
clock,
}),
inner: Arc::new(Inner::new(config, clock)),
..self
}
}
@ -243,7 +196,7 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<
}
pub fn index_name(&self) -> &str {
&self.inner.index_name
self.inner.index.name()
}
pub fn similarity_threshold(&self) -> f32 {
@ -253,12 +206,25 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> RedisSemanticCache<
fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str {
context.scope.as_deref().unwrap_or(key)
}
fn decode(&self, lookup: SemanticLookup<Vec<u8>>) -> Result<SemanticLookup<S::Value>, Error> {
Ok(SemanticLookup {
value: lookup
.value
.map(|bytes| self.codec.decode(&bytes))
.transpose()?,
similarity: lookup.similarity,
})
}
}
impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
for RedisSemanticCache<E, C>
impl<E, S, C> BaseCache for RedisSemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = CacheEntry;
type Value = S::Value;
type Context = SemanticCacheContext;
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
@ -274,22 +240,18 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
let Some(prompt) = prompt_from_context(context) else {
return Ok(());
};
let response = self.codec.encode(&value)?;
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context).to_string();
let tag = Self::tag(key, context);
self.connections.execute(|connection| {
self.inner
.store(connection, &tag, &value, &prompt, &vector, context.ttl)
.store(connection, tag, response, &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))
self.get_cache_with_similarity(key, context)
.map(|lookup| lookup.value)
}
async fn async_set_cache(
@ -301,14 +263,15 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
let Some(prompt) = prompt_from_context(&context) else {
return Ok(());
};
let response = self.codec.encode(&value)?;
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, &context).to_string();
let tag = Self::tag(key, &context).to_owned();
let inner = Arc::clone(&self.inner);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
inner.store(connection, &tag, &value, &prompt, &vector, context.ttl)
inner.store(connection, &tag, response, &prompt, &vector, context.ttl)
})
.await
}
@ -318,49 +281,54 @@ impl<E: Embedder, C: redis::ConnectionLike + Send + 'static> BaseCache
key: &str,
context: &Self::Context,
) -> Result<Option<Self::Value>, Error> {
self.async_get_cache_with_similarity(key, context)
.await
.map(|lookup| lookup.value)
}
}
/// Python stamps a similarity of `0.0` when there is no prompt or no hit in the key's scope.
impl<E, S, C> SemanticCache for RedisSemanticCache<E, S, C>
where
E: Embedder,
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(None);
return Ok(SemanticLookup::miss(Some(0.0)));
};
let vector = self.embedder.embed(&prompt, context.metadata.as_ref())?;
let tag = Self::tag(key, context);
let lookup = self
.connections
.execute(|connection| self.inner.lookup(connection, tag, &vector))?;
self.decode(lookup)
}
async fn async_get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
let Some(prompt) = prompt_from_context(context) else {
return Ok(SemanticLookup::miss(Some(0.0)));
};
let vector = self
.embedder
.async_embed(&prompt, context.metadata.as_ref())
.await?;
let tag = Self::tag(key, context).to_string();
let tag = Self::tag(key, context).to_owned();
let inner = Arc::clone(&self.inner);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let lookup = 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()),
}),
}
.await?;
self.decode(lookup)
}
}
@ -387,228 +355,46 @@ fn vector_buffer(vector: &[f32]) -> Vec<u8> {
}
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,
let mut escaped = String::with_capacity(value.len());
for ch in value.chars() {
if matches!(
ch,
',' | '.'
| '<'
| '>'
| '{'
| '}'
| '['
| ']'
| '\\'
| '"'
| '\''
| ':'
| ';'
| '!'
| '@'
| '#'
| '$'
| '%'
| '^'
| '&'
| '*'
| '('
| ')'
| '-'
| '+'
| '='
| '~'
| '|'
| '/'
| ' '
| '?'
) {
escaped.push('\\');
}
escaped.push(ch);
}
escaped
}
fn ttl_seconds(ttl: Duration) -> u64 {

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

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

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

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

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

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

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

@ -1,97 +0,0 @@
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;
}
}
}
}
_ => {}
}
}

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

@ -0,0 +1,57 @@
pub(crate) fn string_value(value: &redis::Value) -> Option<String> {
match value {
redis::Value::BulkString(bytes) => String::from_utf8(bytes.clone()).ok(),
redis::Value::SimpleString(text) => Some(text.clone()),
redis::Value::VerbatimString { text, .. } => Some(text.clone()),
_ => None,
}
}
pub(crate) fn number_value(value: &redis::Value) -> Option<f64> {
match value {
redis::Value::Int(number) => Some(*number as f64),
redis::Value::Double(number) => Some(*number),
_ => string_value(value).and_then(|text| text.parse().ok()),
}
}
pub(crate) fn first_document(result: &redis::Value) -> Option<&[redis::Value]> {
let redis::Value::Array(items) = result else {
return None;
};
let [count, _document_id, fields, ..] = items.as_slice() else {
return None;
};
if !matches!(count, redis::Value::Int(count) if *count > 0) {
return None;
}
match fields {
redis::Value::Array(fields) => Some(fields.as_slice()),
_ => None,
}
}
fn field_value<'a>(fields: &'a [redis::Value], name: &str) -> Option<&'a redis::Value> {
fields
.as_chunks::<2>()
.0
.iter()
.find(|pair| string_value(&pair[0]).as_deref() == Some(name))
.map(|pair| &pair[1])
}
pub(crate) fn string_field(fields: &[redis::Value], name: &str) -> Option<String> {
field_value(fields, name).and_then(string_value)
}
pub(crate) fn number_field(fields: &[redis::Value], name: &str) -> Option<f64> {
field_value(fields, name).and_then(number_value)
}
pub(crate) fn bytes_field(fields: &[redis::Value], name: &str) -> Option<Vec<u8>> {
match field_value(fields, name)? {
redis::Value::BulkString(bytes) => Some(bytes.clone()),
redis::Value::SimpleString(text) => Some(text.clone().into_bytes()),
_ => None,
}
}

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

File diff suppressed because it is too large Load diff

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

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

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

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

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

@ -12,5 +12,7 @@ r2d2 = "0.8.10"
tokio.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
redis-test = "1.0.4"
rstest.workspace = true
serde_json.workspace = true

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

@ -1,110 +1,25 @@
use std::{
sync::{Arc, Mutex},
sync::{Arc, OnceLock},
time::Duration,
};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionResult, CacheConnectionStatus,
ClaimCache, CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache,
};
use redis::Commands;
use litellm_cache::{BatchEntry, CacheCodec, Error};
use crate::topology::RedisTopology;
mod connection;
mod operations;
pub use connection::ConnectionRef;
use connection::{ClusterConnectionManager, ConnectionManager};
pub use operations::{
RedisArg, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, RedisScript,
use crate::{
connection::{ConnectionRef, Connections},
topology::RedisTopology,
};
const DEFAULT_TTL: Duration = Duration::from_secs(600);
const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
const REDIS_POOL_SIZE: u32 = 16;
const INCREMENT_SCRIPT: &str = concat!(
"local value = redis.call('INCRBYFLOAT', KEYS[1], ARGV[1]); ",
"if redis.call('TTL', KEYS[1]) == -1 then ",
"redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return value"
);
const CLAIM_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if ARGV[1] == '' then if current ~= false and current ~= '' then return 0; end; ",
"elseif current ~= ARGV[1] then return 0; end; ",
"if ARGV[3] ~= '' then redis.call('SET', KEYS[1], ARGV[3], 'EX', ARGV[2]); ",
"elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1"
);
const CLAIM_ATTEMPTS: usize = 8;
#[allow(private_interfaces)]
pub enum Connections<C> {
Pool(r2d2::Pool<ConnectionManager>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
}
impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
pub fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
match self {
Self::Pool(pool) => {
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let result = operation(&mut ConnectionRef::Node(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
result
}
Self::Cluster(pool) => {
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
result
}
Self::Fixed(connection) => {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
operation(&mut ConnectionRef::Node(&mut *connection))
}
}
}
pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
}
pub fn fixed(connection: C) -> Self {
Self::Fixed(Mutex::new(connection))
}
pub fn open(url: &str, topology: &RedisTopology) -> Result<Self, Error> {
match topology {
RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)),
RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool(
ClusterConnectionManager::open(url, startup_nodes)?,
)?)),
}
}
}
pub struct RedisCache<S, C = redis::Connection> {
connections: Arc<Connections<C>>,
default_ttl: Duration,
codec: S,
namespace: Option<String>,
topology: RedisTopology,
pub(crate) connections: Arc<Connections<C>>,
pub(crate) default_ttl: Duration,
pub(crate) codec: S,
pub(crate) namespace: Option<String>,
pub(crate) topology: RedisTopology,
/// The server's major version, read from `INFO` once, like Python's `redis_version`.
pub(crate) major_version: Arc<OnceLock<u32>>,
}
impl<S: CacheCodec> RedisCache<S> {
@ -125,20 +40,11 @@ impl<S: CacheCodec> RedisCache<S> {
codec,
namespace: None,
topology: topology.clone(),
major_version: Arc::default(),
})
}
}
fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
r2d2::Pool::builder()
.max_size(REDIS_POOL_SIZE)
.min_idle(Some(0))
.connection_timeout(REDIS_TIMEOUT)
.test_on_check_out(false)
.build(manager)
.map_err(|_| Error::Unavailable)
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
@ -151,6 +57,7 @@ where
codec,
namespace: None,
topology: RedisTopology::Standalone,
major_version: Arc::default(),
}
}
@ -169,11 +76,52 @@ where
&self.topology
}
fn namespaced_key(&self, key: &str) -> String {
pub(crate) fn namespaced_key(&self, key: &str) -> String {
namespaced_key(self.namespace.as_deref(), key)
}
fn namespaced_pattern(&self) -> Result<String, Error> {
pub(crate) fn namespaced_keys(&self, keys: &[String]) -> Vec<String> {
keys.iter().map(|key| self.namespaced_key(key)).collect()
}
/// Whole seconds for `ttl`, falling back to the default TTL like Python's `get_ttl`.
pub(crate) fn ttl_or_default(&self, ttl: Option<Duration>) -> u64 {
ttl_seconds(ttl.unwrap_or(self.default_ttl))
}
pub(crate) fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
self.connections.execute(operation)
}
/// `_parse_redis_major_version`: the major version from `INFO`, or
/// `DEFAULT_REDIS_MAJOR_VERSION` when `INFO` fails or its version does not parse. The first
/// answer is kept, as Python reads `redis_version` once at construction.
pub(crate) async fn major_version(&self) -> u32 {
if let Some(version) = self.major_version.get() {
return *version;
}
let info = self
.run(|connection| connection.node_text(&redis::cmd("INFO")))
.await;
let version = info
.ok()
.and_then(|info| parse_major_version(&info))
.unwrap_or_else(default_major_version);
*self.major_version.get_or_init(|| version)
}
pub(crate) async fn run<T, F>(&self, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
Connections::run_blocking(Arc::clone(&self.connections), operation).await
}
pub(crate) fn namespaced_pattern(&self) -> Result<String, Error> {
let namespace = self.namespace.as_ref().ok_or(Error::UnscopedFlush)?;
let escaped: String = namespace
.chars()
@ -188,18 +136,7 @@ where
Ok(format!("{escaped}:*"))
}
fn flush_matching(connection: &mut ConnectionRef<'_>, pattern: &str) -> Result<(), Error> {
connection.scan(pattern, 1000, |connection, keys| {
if !keys.is_empty() {
connection
.del::<_, usize>(keys)
.map_err(|_| Error::Unavailable)?;
}
Ok(true)
})
}
fn decode_response(&self, value: redis::Value) -> Result<Option<S::Value>, Error> {
pub(crate) fn decode_response(&self, value: redis::Value) -> Result<Option<S::Value>, Error> {
match value {
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => self.codec.decode(&bytes).map(Some),
@ -208,7 +145,10 @@ where
}
}
fn decode_batch_response(&self, value: redis::Value) -> Result<BatchEntry<S::Value>, Error> {
pub(crate) fn decode_batch_response(
&self,
value: redis::Value,
) -> Result<BatchEntry<S::Value>, Error> {
match self.decode_response(value) {
Ok(Some(value)) => Ok(BatchEntry::Hit(value)),
Ok(None) => Ok(BatchEntry::Miss),
@ -216,15 +156,9 @@ where
Err(error) => Err(error),
}
}
fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
}
fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
pub(crate) fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
match namespace {
Some(namespace) if !key.starts_with(&format!("{namespace}:")) => {
format!("{namespace}:{key}")
@ -233,469 +167,26 @@ fn namespaced_key(namespace: Option<&str>, key: &str) -> String {
}
}
impl<S, C> BaseCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Value = S::Value;
type Context = ExactCacheContext;
pub(crate) fn ttl_seconds(ttl: Duration) -> u64 {
ttl.as_secs()
.saturating_add(u64::from(ttl.subsec_nanos() > 0))
.max(1)
}
fn get_ttl(&self, context: &Self::Context) -> Option<Duration> {
context.ttl.or(Some(self.default_ttl))
}
fn set_cache(
&self,
key: &str,
value: Self::Value,
context: &ExactCacheContext,
) -> Result<(), Error> {
let payload = self.codec.encode(&value)?;
let ttl = Self::ttl_seconds(self.get_ttl(context).unwrap_or(self.default_ttl));
let key = self.namespaced_key(key);
self.connections.execute(|connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)
})
}
fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = self.connections.execute(|connection| {
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)
})?;
self.decode_response(value)
}
async fn async_set_cache(
&self,
key: &str,
value: Self::Value,
context: ExactCacheContext,
) -> Result<(), Error> {
let payload = self.codec.encode(&value)?;
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.set_ex::<_, _, ()>(key, payload, ttl)
.map_err(|_| Error::Unavailable)
})
.await
}
async fn async_get_cache(
&self,
key: &str,
_: &ExactCacheContext,
) -> Result<Option<Self::Value>, Error> {
let key = self.namespaced_key(key);
let value = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)
})
.await?;
self.decode_response(value)
}
async fn async_set_cache_pipeline(
&self,
cache_list: Vec<(String, Self::Value)>,
context: ExactCacheContext,
) -> Result<(), Error> {
let entries = cache_list
.into_iter()
.map(|(key, value)| {
self.codec
.encode(&value)
.map(|payload| (self.namespaced_key(&key), payload))
})
.collect::<Result<Vec<_>, _>>()?;
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
if entries.is_empty() {
return Ok(());
}
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = entries
.into_iter()
.map(|(key, payload)| {
let mut command = redis::cmd("SETEX");
command.arg(key).arg(ttl).arg(payload);
command
})
.collect();
connection.pipeline(commands).map(drop)
})
.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 connection.ping() {
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 parse_major_version(info: &str) -> Option<u32> {
let version = info
.lines()
.find_map(|line| line.trim().strip_prefix("redis_version:"))?
.trim();
match version.split_once('.') {
Some((major, _)) => major.parse().ok(),
None => version.parse::<f64>().ok().map(|major| major as u32),
}
}
impl<S, C> BatchCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn batch_get_cache(
&self,
keys: &[String],
_: &ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = self.connections.execute(|connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})?;
values
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
async fn async_batch_get_cache(
&self,
keys: Vec<String>,
_: ExactCacheContext,
) -> Result<Vec<BatchEntry<Self::Value>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
values
.into_iter()
.map(|value| self.decode_batch_response(value))
.collect()
}
}
impl<S, C> DeleteCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
self.connections
.execute(|connection| connection.del::<_, ()>(key).map_err(|_| Error::Unavailable))
}
async fn async_delete_cache(&self, key: &str) -> Result<(), Error> {
let key = self.namespaced_key(key);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del::<_, ()>(key).map_err(|_| Error::Unavailable)
})
.await
}
}
impl<S, C> FlushCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
fn flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
self.connections
.execute(|connection| Self::flush_matching(connection, &pattern))
}
async fn async_flush_cache(&self) -> Result<(), Error> {
let pattern = self.namespaced_pattern()?;
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
Self::flush_matching(connection, &pattern)
})
.await
}
}
impl<S, C> CounterCache for RedisCache<S, C>
where
S: CacheCodec<Value = f64>,
C: redis::ConnectionLike + Send + 'static,
{
fn increment_cache(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
self.connections
.execute(|connection| increment(connection, key, amount, ttl))
}
async fn async_increment(
&self,
key: &str,
amount: f64,
context: ExactCacheContext,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
increment(connection, key, amount, ttl)
})
.await
}
}
fn increment(
connection: &mut ConnectionRef<'_>,
key: String,
amount: f64,
ttl: u64,
) -> Result<f64, Error> {
redis::cmd("EVAL")
.arg(INCREMENT_SCRIPT)
.arg(1)
.arg(key)
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
}
fn stored_bytes(value: redis::Value) -> Result<Option<Vec<u8>>, Error> {
match value {
redis::Value::Nil => Ok(None),
redis::Value::BulkString(bytes) => Ok(Some(bytes)),
redis::Value::SimpleString(text) => Ok(Some(text.into_bytes())),
_ => Err(Error::InvalidEntry),
}
}
/// Eligibility is decided on decoded values, so a pin written by another encoder (Python's
/// `json.dumps` spacing or key order) still matches. The write is a compare-and-set on the
/// bytes that decision was made on, retried when another claimant wins the race.
fn claim<S: CacheCodec>(
connection: &mut ConnectionRef<'_>,
codec: &S,
key: &str,
candidate: S::Value,
eligible: &[S::Value],
ttl: u64,
) -> Result<S::Value, Error>
where
S::Value: PartialEq,
{
let payload = codec.encode(&candidate)?;
if payload.is_empty() {
return Err(Error::InvalidEntry);
}
for _ in 0..CLAIM_ATTEMPTS {
let current = stored_bytes(
connection
.get::<_, redis::Value>(key)
.map_err(|_| Error::Unavailable)?,
)?
.filter(|bytes| !bytes.is_empty());
let existing = current
.as_deref()
.and_then(|bytes| codec.decode(bytes).ok())
.filter(|existing| eligible.is_empty() || eligible.contains(existing));
let refresh = existing
.as_ref()
.is_some_and(|existing| !eligible.is_empty() || *existing == candidate);
let write: &[u8] = if existing.is_some() { b"" } else { &payload };
let applied = redis::cmd("EVAL")
.arg(CLAIM_SCRIPT)
.arg(1)
.arg(key)
.arg(current.as_deref().unwrap_or_default())
.arg(ttl)
.arg(write)
.arg(u8::from(refresh))
.query::<bool>(connection)
.map_err(|_| Error::Unavailable)?;
if applied {
return Ok(existing.unwrap_or(candidate));
}
}
Err(Error::Unavailable)
}
impl<S, C> ClaimCache for RedisCache<S, C>
where
S: CacheCodec + Clone + 'static,
S::Value: PartialEq,
C: redis::ConnectionLike + Send + 'static,
{
fn claim_cache(
&self,
key: &str,
candidate: S::Value,
eligible: &[S::Value],
context: ExactCacheContext,
) -> Result<S::Value, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
self.connections
.execute(|connection| claim(connection, &self.codec, &key, candidate, eligible, ttl))
}
async fn async_claim_cache(
&self,
key: &str,
candidate: S::Value,
eligible: Vec<S::Value>,
context: ExactCacheContext,
) -> Result<S::Value, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(self.get_ttl(&context).unwrap_or(self.default_ttl));
let codec = self.codec.clone();
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
claim(connection, &codec, &key, candidate, &eligible, ttl)
})
.await
}
}
#[cfg(test)]
mod tests {
use std::time::Duration;
use litellm_cache::{
BaseCache, CacheCodec, DeleteCache, ExactCacheContext, FlushCache, JsonCodec,
};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
use super::RedisCache;
fn entry() -> serde_json::Value {
json!({"deployment": "model-a", "cooldown_seconds": 30})
}
#[test]
fn ttl_seconds_rounds_up_and_keeps_expiration_positive() {
assert_eq!(
RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::ZERO),
1
);
assert_eq!(
RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::from_millis(1500)),
2
);
assert_eq!(
RedisCache::<JsonCodec<serde_json::Value>>::ttl_seconds(Duration::from_secs(15)),
15
);
}
#[test]
fn redis_commands_round_trip_entries_and_delete_only_namespaced_keys() {
let value = entry();
let payload = JsonCodec::<serde_json::Value>::new()
.encode(&value)
.unwrap();
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SETEX")
.arg("litellm-cache:key")
.arg(600)
.arg(payload.clone()),
Ok("OK"),
),
MockCmd::new(redis::cmd("GET").arg("litellm-cache:key"), Ok(payload)),
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
.with_namespace(Some("litellm-cache".into()));
cache
.set_cache("key", value.clone(), &ExactCacheContext::default())
.unwrap();
assert_eq!(
cache
.get_cache("key", &ExactCacheContext::default())
.unwrap(),
Some(value)
);
cache.delete_cache("key").unwrap();
}
#[test]
fn flush_scans_and_deletes_only_cache_keys() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("SCAN")
.cursor_arg(0)
.arg("MATCH")
.arg("litellm-cache:*")
.arg("COUNT")
.arg(1000),
Ok(redis_test::redis_value!(["0", ["litellm-cache:key"]])),
),
MockCmd::new(redis::cmd("DEL").arg("litellm-cache:key"), Ok(1u32)),
])
.assert_all_commands_consumed();
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
.with_namespace(Some("litellm-cache".into()));
cache.flush_cache().unwrap();
}
#[tokio::test]
async fn test_connection_runs_ping_off_executor() {
let connection = MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Ok("PONG"))])
.assert_all_commands_consumed();
let cache =
RedisCache::with_connection(connection, None, JsonCodec::<serde_json::Value>::new())
.with_namespace(Some("litellm-cache".into()));
assert_eq!(
cache.test_connection().await.unwrap().status,
litellm_cache::CacheConnectionStatus::Success
);
}
fn default_major_version() -> u32 {
std::env::var("DEFAULT_REDIS_MAJOR_VERSION")
.ok()
.and_then(|value| value.parse().ok())
.unwrap_or(7)
}

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

@ -1,632 +0,0 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{
CacheCodec, CacheScript, ClientInfoCache, Error, IncrementOperation, QueueCache, ScanCache,
ScriptCache, SetCache, TtlCache,
};
use redis::Commands;
use super::{ConnectionRef, Connections, RedisCache, namespaced_key};
const INCREMENT_WITH_FLOOR_SCRIPT: &str = concat!(
"local count = redis.call('INCRBY', KEYS[1], ARGV[1]); ",
"if count < 0 then count = redis.call('INCRBY', KEYS[1], -count); end; ",
"if redis.call('TTL', KEYS[1]) < 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
"return count"
);
const SET_MAX_SCRIPT: &str = concat!(
"local current = redis.call('GET', KEYS[1]); ",
"if current == false or tonumber(current) < tonumber(ARGV[1]) then ",
"redis.call('SET', KEYS[1], ARGV[1]); ",
"if tonumber(ARGV[2]) > 0 then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; ",
"return ARGV[1]; end; return current"
);
#[derive(Clone, Debug, PartialEq)]
pub enum RedisArg {
Bytes(Vec<u8>),
Integer(i64),
Float(f64),
}
impl From<&str> for RedisArg {
fn from(value: &str) -> Self {
Self::Bytes(value.as_bytes().to_vec())
}
}
impl From<String> for RedisArg {
fn from(value: String) -> Self {
Self::Bytes(value.into_bytes())
}
}
impl From<Vec<u8>> for RedisArg {
fn from(value: Vec<u8>) -> Self {
Self::Bytes(value)
}
}
impl From<i64> for RedisArg {
fn from(value: i64) -> Self {
Self::Integer(value)
}
}
impl From<f64> for RedisArg {
fn from(value: f64) -> Self {
Self::Float(value)
}
}
impl redis::ToRedisArgs for RedisArg {
fn write_redis_args<W>(&self, out: &mut W)
where
W: ?Sized + redis::RedisWrite,
{
match self {
Self::Bytes(value) => value.write_redis_args(out),
Self::Integer(value) => value.write_redis_args(out),
Self::Float(value) => value.write_redis_args(out),
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub struct RedisRpushOperation {
pub key: String,
pub values: Vec<RedisArg>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RedisLpopOperation {
pub key: String,
pub count: Option<usize>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum RedisLpopResult {
Missing,
Value(Vec<u8>),
Values(Vec<Vec<u8>>),
}
pub struct RedisScript<C> {
connections: Arc<Connections<C>>,
namespace: Option<String>,
source: String,
}
impl<C> CacheScript for RedisScript<C>
where
C: redis::ConnectionLike + Send + 'static,
{
type Argument = RedisArg;
type Output = redis::Value;
async fn invoke(
&self,
keys: Vec<String>,
arguments: Vec<Self::Argument>,
) -> Result<Self::Output, Error> {
let keys = keys
.into_iter()
.map(|key| namespaced_key(self.namespace.as_deref(), &key))
.collect::<Vec<_>>();
let connections = Arc::clone(&self.connections);
let source = self.source.clone();
tokio::task::spawn_blocking(move || {
connections.execute(|connection| {
redis::cmd("EVAL")
.arg(source)
.arg(keys.len())
.arg(keys)
.arg(arguments)
.query(connection)
.map_err(|_| Error::Unavailable)
})
})
.await
.map_err(|_| Error::Unavailable)?
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
pub async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
if keys.is_empty() {
return Ok(0);
}
let keys = keys
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
connection.del(keys).map_err(|_| Error::Unavailable)
})
.await
}
pub fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = self.connections.execute(|connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})?;
values.into_iter().map(count).collect()
}
pub async fn async_batch_get_counts(
&self,
keys: Vec<String>,
) -> Result<Vec<Option<i64>>, Error> {
let keys = keys
.iter()
.map(|key| self.namespaced_key(key))
.collect::<Vec<_>>();
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("MGET")
.arg(keys)
.query::<Vec<redis::Value>>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
values.into_iter().map(count).collect()
}
pub fn sync_ping(&self) -> Result<bool, Error> {
self.connections
.execute(|connection| connection.ping().map_err(|_| Error::Unavailable))
}
pub async fn ping(&self) -> Result<bool, Error> {
Connections::run_blocking(Arc::clone(&self.connections), |connection| {
connection.ping().map_err(|_| Error::Unavailable)
})
.await
}
pub async fn async_get_ttl(&self, key: &str) -> Result<Option<i64>, Error> {
let key = self.namespaced_key(key);
let ttl = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("TTL")
.arg(key)
.query::<i64>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
Ok((ttl >= 0).then_some(ttl))
}
pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
let pattern = format!("{}*", self.namespaced_key(pattern));
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let mut matches = Vec::new();
connection.scan(&pattern, count, |_, keys| {
matches.extend(keys);
Ok(matches.len() < count)
})?;
matches.truncate(count);
Ok(matches)
})
.await
}
pub async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<RedisArg>,
ttl: Option<Duration>,
) -> Result<usize, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry);
}
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
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");
expire.arg(&key).arg(ttl);
let replies = connection.pipeline(vec![sadd, expire])?;
replies
.into_iter()
.next()
.map(redis::from_redis_value::<usize>)
.transpose()
.map_err(|_| Error::Unavailable)?
.ok_or(Error::Unavailable)
})
.await
}
pub async fn async_rpush(&self, key: &str, values: Vec<RedisArg>) -> Result<usize, Error> {
if values.is_empty() {
return Err(Error::InvalidEntry);
}
let key = self.namespaced_key(key);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("RPUSH")
.arg(key)
.arg(values)
.query(connection)
.map_err(|_| Error::Unavailable)
})
.await
}
pub async fn async_rpush_pipeline(
&self,
operations: Vec<RedisRpushOperation>,
) -> Result<Vec<usize>, Error> {
let operations = operations
.into_iter()
.map(|operation| {
if operation.values.is_empty() {
return Err(Error::InvalidEntry);
}
Ok((self.namespaced_key(&operation.key), operation.values))
})
.collect::<Result<Vec<_>, _>>()?;
if operations.is_empty() {
return Ok(Vec::new());
}
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = operations
.into_iter()
.map(|(key, values)| {
let mut command = redis::cmd("RPUSH");
command.arg(key).arg(values);
command
})
.collect();
connection
.pipeline(commands)?
.into_iter()
.map(|value| redis::from_redis_value(value).map_err(|_| Error::Unavailable))
.collect()
})
.await
}
pub async fn async_lpop(
&self,
key: &str,
count: Option<usize>,
) -> Result<RedisLpopResult, Error> {
let key = self.namespaced_key(key);
let multiple = count.is_some();
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 {
command.arg(count);
}
command
.query::<redis::Value>(connection)
.map_err(|_| Error::Unavailable)
})
.await?;
lpop_result(value, multiple)
}
pub async fn async_lpop_pipeline(
&self,
operations: Vec<RedisLpopOperation>,
) -> Result<Vec<RedisLpopResult>, Error> {
let operations = operations
.into_iter()
.map(|operation| (self.namespaced_key(&operation.key), operation.count))
.collect::<Vec<_>>();
if operations.is_empty() {
return Ok(Vec::new());
}
let multiple = operations
.iter()
.map(|(_, count)| count.is_some())
.collect::<Vec<_>>();
let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
let commands = operations
.into_iter()
.map(|(key, count)| {
let mut command = redis::cmd("LPOP");
command.arg(key);
if let Some(count) = count {
command.arg(count);
}
command
})
.collect();
connection.pipeline(commands)
})
.await?;
values
.into_iter()
.zip(multiple)
.map(|(value, multiple)| lpop_result(value, multiple))
.collect()
}
pub async fn async_eval(
&self,
script: String,
keys: Vec<String>,
arguments: Vec<RedisArg>,
) -> Result<redis::Value, Error> {
let keys = keys
.into_iter()
.map(|key| self.namespaced_key(&key))
.collect::<Vec<_>>();
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(script)
.arg(keys.len())
.arg(keys)
.arg(arguments)
.query(connection)
.map_err(|_| Error::Unavailable)
})
.await
}
pub fn client_list(&self) -> Result<String, Error> {
self.connections
.execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST")))
}
pub fn info(&self) -> Result<String, Error> {
self.connections
.execute(|connection| connection.node_text(&redis::cmd("INFO")))
}
pub fn flushall(&self) -> Result<(), Error> {
self.connections.execute(|connection| connection.flushall())
}
}
impl<S, C> RedisCache<S, C>
where
S: CacheCodec<Value = f64>,
C: redis::ConnectionLike + Send + 'static,
{
pub fn increment_with_floor(
&self,
key: &str,
amount: i64,
ttl: Duration,
) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl);
self.connections
.execute(|connection| increment_with_floor(connection, key, amount, ttl))
}
pub async fn async_increment_pipeline(
&self,
operations: Vec<IncrementOperation>,
) -> Result<Vec<f64>, Error> {
let operations = operations
.into_iter()
.map(|operation| {
(
self.namespaced_key(&operation.key),
operation.amount,
operation.ttl.map(Self::ttl_seconds),
)
})
.collect::<Vec<_>>();
if operations.is_empty() {
return Ok(Vec::new());
}
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 {
let mut increment = redis::cmd("INCRBYFLOAT");
increment.arg(&key).arg(amount);
increments.push(commands.len());
commands.push(increment);
if let Some(ttl) = ttl {
let mut expire = redis::cmd("EXPIRE");
expire.arg(key).arg(ttl);
commands.push(expire);
}
}
let mut replies = connection.pipeline(commands)?;
increments
.into_iter()
.map(|index| {
redis::from_redis_value(std::mem::take(&mut replies[index]))
.map_err(|_| Error::Unavailable)
})
.collect()
})
.await
}
pub async fn async_increment_with_floor(
&self,
key: &str,
amount: i64,
ttl: Duration,
) -> Result<i64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl);
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
increment_with_floor(connection, key, amount, ttl)
})
.await
}
pub async fn async_set_max(
&self,
key: &str,
value: f64,
ttl: Option<Duration>,
) -> Result<f64, Error> {
let key = self.namespaced_key(key);
let ttl = Self::ttl_seconds(ttl.unwrap_or(self.default_ttl));
Connections::run_blocking(Arc::clone(&self.connections), move |connection| {
redis::cmd("EVAL")
.arg(SET_MAX_SCRIPT)
.arg(1)
.arg(key)
.arg(value)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
})
.await
}
}
fn redis_bytes(value: redis::Value) -> Result<Vec<u8>, Error> {
match value {
redis::Value::BulkString(bytes) => Ok(bytes),
redis::Value::SimpleString(text) => Ok(text.into_bytes()),
_ => Err(Error::InvalidEntry),
}
}
fn lpop_result(value: redis::Value, multiple: bool) -> Result<RedisLpopResult, Error> {
match value {
redis::Value::Nil => Ok(RedisLpopResult::Missing),
redis::Value::Array(values) if multiple => values
.into_iter()
.map(redis_bytes)
.collect::<Result<Vec<_>, _>>()
.map(RedisLpopResult::Values),
value if !multiple => redis_bytes(value).map(RedisLpopResult::Value),
_ => Err(Error::InvalidEntry),
}
}
fn count(value: redis::Value) -> Result<Option<i64>, Error> {
match value {
redis::Value::Nil => Ok(None),
redis::Value::Int(value) => Ok(Some(value)),
redis::Value::BulkString(value) => std::str::from_utf8(&value)
.ok()
.and_then(|value| value.parse().ok())
.map(Some)
.ok_or(Error::InvalidEntry),
redis::Value::SimpleString(value) => {
value.parse().map(Some).map_err(|_| Error::InvalidEntry)
}
_ => Err(Error::InvalidEntry),
}
}
fn increment_with_floor(
connection: &mut ConnectionRef<'_>,
key: String,
amount: i64,
ttl: u64,
) -> Result<i64, Error> {
redis::cmd("EVAL")
.arg(INCREMENT_WITH_FLOOR_SCRIPT)
.arg(1)
.arg(key)
.arg(amount)
.arg(ttl)
.query(connection)
.map_err(|_| Error::Unavailable)
}
impl<S, C> TtlCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
RedisCache::async_get_ttl(self, key)
.await
.map(|ttl| ttl.map(|seconds| Duration::from_secs(seconds as u64)))
}
}
impl<S, C> ScanCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result<Vec<String>, Error> {
RedisCache::async_scan_iter(self, pattern, count).await
}
}
impl<S, C> ClientInfoCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type ClientList = String;
type Info = String;
fn client_list(&self) -> Result<Self::ClientList, Error> {
RedisCache::client_list(self)
}
fn info(&self) -> Result<Self::Info, Error> {
RedisCache::info(self)
}
}
impl<S, C> SetCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type SetValue = RedisArg;
type SetResult = usize;
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<Self::SetValue>,
ttl: Option<Duration>,
) -> Result<Self::SetResult, Error> {
RedisCache::async_set_cache_sadd(self, key, values, ttl).await
}
}
impl<S, C> QueueCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type QueueValue = RedisArg;
type PopResult = RedisLpopResult;
async fn async_rpush(&self, key: &str, values: Vec<Self::QueueValue>) -> Result<usize, Error> {
RedisCache::async_rpush(self, key, values).await
}
async fn async_lpop(&self, key: &str, count: Option<usize>) -> Result<Self::PopResult, Error> {
RedisCache::async_lpop(self, key, count).await
}
}
impl<S, C> ScriptCache for RedisCache<S, C>
where
S: CacheCodec,
C: redis::ConnectionLike + Send + 'static,
{
type Script = RedisScript<C>;
fn async_register_script(&self, source: String) -> Self::Script {
RedisScript {
connections: Arc::clone(&self.connections),
namespace: self.namespace.clone(),
source,
}
}
}

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

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

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

@ -1,29 +1,126 @@
use std::collections::HashMap;
use std::{
sync::{Arc, Mutex},
time::Duration,
};
use litellm_cache::Error;
use redis::{
ConnectionAddr, ConnectionInfo, ConnectionLike, IntoConnectionInfo,
cluster::{ClusterClient, ClusterClientBuilder, ClusterConnection, NodeAddress},
ConnectionAddr, ConnectionInfo, IntoConnectionInfo,
cluster::{
ClusterClient, ClusterClientBuilder, ClusterConnection, ClusterPipeline, NodeAddress,
},
cluster_routing::{
MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo, Slot,
MultipleNodeRoutingInfo, ResponsePolicy, RoutingInfo, SingleNodeRoutingInfo,
},
};
use super::REDIS_TIMEOUT;
use crate::topology::RedisNode;
use crate::topology::{RedisNode, RedisTopology};
pub struct PooledConnection<C> {
pub(super) connection: C,
pub(super) failed: bool,
pub(crate) const REDIS_TIMEOUT: Duration = Duration::from_secs(5);
const REDIS_POOL_SIZE: u32 = 16;
#[allow(private_interfaces)]
pub enum Connections<C> {
Pool(r2d2::Pool<ConnectionManager>),
Cluster(r2d2::Pool<ClusterConnectionManager>),
Fixed(Mutex<C>),
}
impl<C> Connections<C>
where
C: redis::ConnectionLike + Send + 'static,
{
pub fn execute<T>(
&self,
operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error>,
) -> Result<T, Error> {
match self {
Self::Pool(pool) => {
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let result = operation(&mut ConnectionRef::Node(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
result
}
Self::Cluster(pool) => {
let mut pooled = pool.get().map_err(|_| Error::Unavailable)?;
let result = operation(&mut ConnectionRef::Cluster(&mut pooled.connection));
pooled.failed = matches!(result, Err(Error::Unavailable));
result
}
Self::Fixed(connection) => {
let mut connection = connection.lock().map_err(|_| Error::Unavailable)?;
operation(&mut ConnectionRef::Node(&mut *connection))
}
}
}
pub async fn run_blocking<T, F>(connections: Arc<Self>, operation: F) -> Result<T, Error>
where
T: Send + 'static,
F: FnOnce(&mut ConnectionRef<'_>) -> Result<T, Error> + Send + 'static,
{
tokio::task::spawn_blocking(move || connections.execute(operation))
.await
.map_err(|_| Error::Unavailable)?
}
pub fn fixed(connection: C) -> Self {
Self::Fixed(Mutex::new(connection))
}
pub fn open(url: &str, topology: &RedisTopology) -> Result<Self, Error> {
match topology {
RedisTopology::Standalone => Ok(Self::Pool(pool(ConnectionManager::open(url)?)?)),
RedisTopology::Cluster { startup_nodes } => Ok(Self::Cluster(pool(
ClusterConnectionManager::open(url, startup_nodes)?,
)?)),
}
}
/// Closes every idle pooled connection; the next operation opens a fresh one. Connections
/// checked out right now return to the pool, and a caller-owned connection stays open.
pub fn disconnect(&self) {
match self {
Self::Pool(pool) => close_idle(pool),
Self::Cluster(pool) => close_idle(pool),
Self::Fixed(_) => {}
}
}
}
fn pool<M: r2d2::ManageConnection>(manager: M) -> Result<r2d2::Pool<M>, Error> {
r2d2::Pool::builder()
.max_size(REDIS_POOL_SIZE)
.min_idle(Some(0))
.connection_timeout(REDIS_TIMEOUT)
.test_on_check_out(false)
.build(manager)
.map_err(|_| Error::Unavailable)
}
fn close_idle<M, T>(pool: &r2d2::Pool<M>)
where
M: r2d2::ManageConnection<Connection = PooledConnection<T>>,
{
let mut idle = Vec::new();
while let Some(mut connection) = pool.try_get() {
connection.failed = true;
idle.push(connection);
}
}
pub(crate) struct PooledConnection<C> {
connection: C,
failed: bool,
}
/// Pools connections without a checkout PING, which would double every operation's round trips.
/// A timed-out command leaves its reply on the socket while redis still reports the connection
/// open, so any connection whose operation failed is discarded instead of being reused.
pub struct ConnectionManager(redis::Client);
pub(crate) struct ConnectionManager(redis::Client);
impl ConnectionManager {
pub(super) fn open(url: &str) -> Result<Self, Error> {
fn open(url: &str) -> Result<Self, Error> {
redis::Client::open(url)
.map(Self)
.map_err(|_| Error::Unavailable)
@ -54,10 +151,10 @@ impl r2d2::ManageConnection for ConnectionManager {
}
}
pub struct ClusterConnectionManager(ClusterClient);
pub(crate) struct ClusterConnectionManager(ClusterClient);
impl ClusterConnectionManager {
pub(super) fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
fn open(url: &str, startup_nodes: &[RedisNode]) -> Result<Self, Error> {
if startup_nodes.is_empty() {
return Err(Error::Unavailable);
}
@ -172,43 +269,36 @@ impl redis::ConnectionLike for ConnectionRef<'_> {
}
impl ConnectionRef<'_> {
pub(crate) fn pipeline(
/// Runs `pipeline` and decodes its non-ignored replies as `T`. A cluster connection refuses
/// `Pipeline::query`, so there a transaction goes to its keys' slot as one MULTI/EXEC and
/// anything else is split per node by `ClusterPipeline`; either way the raw replies are
/// handed back to `pipeline` to decode.
pub(crate) fn query_pipeline<T: redis::FromRedisValue>(
&mut self,
commands: Vec<redis::Cmd>,
) -> Result<Vec<redis::Value>, Error> {
pipeline: &redis::Pipeline,
) -> Result<T, Error> {
match self {
Self::Node(connection) => {
let mut pipeline = redis::pipe();
for command in &commands {
pipeline.add_command(command.clone());
}
pipeline
.query::<Vec<redis::Value>>(*connection)
.map_err(|_| Error::Unavailable)
Self::Node(connection) => pipeline.query(*connection),
Self::Cluster(connection) if pipeline.is_transaction() => {
redis::ConnectionLike::req_packed_commands(
*connection,
&pipeline.get_packed_pipeline(),
pipeline.len() + 1,
1,
)
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
}
Self::Cluster(connection) => {
let mut replies: Vec<Option<redis::Value>> = vec![None; commands.len()];
for indices in slot_groups(&commands).into_values() {
let mut pipeline = redis::pipe();
for index in &indices {
pipeline.add_command(commands[*index].clone());
}
let values = connection
.req_packed_commands(&pipeline.get_packed_pipeline(), 0, indices.len())
.map_err(|_| Error::Unavailable)?;
if values.len() != indices.len() {
return Err(Error::Unavailable);
}
for (index, value) in indices.into_iter().zip(values) {
replies[index] = Some(value);
}
let mut cluster = ClusterPipeline::with_capacity(pipeline.len());
for command in pipeline.cmd_iter() {
cluster.add_command(command.clone());
}
replies
.into_iter()
.collect::<Option<Vec<_>>>()
.ok_or(Error::Unavailable)
cluster
.query(connection)
.and_then(|replies| pipeline.query(&mut Replies(Some(replies))))
}
}
.map_err(|_| Error::Unavailable)
}
pub(crate) fn scan(
@ -379,14 +469,35 @@ fn scan_command(cursor: u64, pattern: &str, count: usize) -> redis::Cmd {
command
}
fn slot_groups(commands: &[redis::Cmd]) -> HashMap<Slot, Vec<usize>> {
let mut groups: HashMap<Slot, Vec<usize>> = HashMap::new();
for (index, command) in commands.iter().enumerate() {
let key = match command.args_iter().nth(1) {
Some(redis::Arg::Simple(key)) => key,
_ => b"",
};
groups.entry(Slot::for_key(key)).or_default().push(index);
/// Hands already received pipeline replies to `Pipeline::query`, so it applies its own
/// ignore and error handling to replies a cluster pipeline gathered from several nodes.
struct Replies(Option<Vec<redis::Value>>);
impl redis::ConnectionLike for Replies {
fn req_packed_command(&mut self, _: &[u8]) -> redis::RedisResult<redis::Value> {
Err((redis::ErrorKind::Client, "replies hold a pipeline only").into())
}
fn req_packed_commands(
&mut self,
_: &[u8],
_: usize,
_: usize,
) -> redis::RedisResult<Vec<redis::Value>> {
self.0
.take()
.ok_or_else(|| (redis::ErrorKind::Client, "replies were already read").into())
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
groups
}

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

File diff suppressed because it is too large Load diff

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

@ -1,100 +1,67 @@
//! Contract tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a
//! comma separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them.
//! Tests against a real Redis Cluster. Set `LITELLM_TEST_REDIS_CLUSTER_NODES` to a comma
//! separated `host:port` list (for example `127.0.0.1:7000,127.0.0.1:7001`) to run them.
use std::time::{Duration, SystemTime, UNIX_EPOCH};
mod support;
use std::{collections::HashSet, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionStatus, CacheScript, ClaimCache,
CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
ScriptCache,
BaseCache, BatchCache, BatchEntry, BoundedCounterCache, BulkDeleteCache, CacheConnectionStatus,
CacheScript, ClaimCache, ClientInfoCache, ConnectionCache, CounterCache, DeleteCache,
DisconnectCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec,
PingCache, QueueCache, RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache,
TtlPipelineCache,
};
use litellm_cache_redis::{
RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisNode, RedisRpushOperation,
RedisTopology,
};
use redis::cluster_routing::Slot;
use rstest::{fixture, rstest};
use serde_json::json;
use support::{JsonCache, cluster_cache, cluster_url};
type Cache = RedisCache<JsonCodec<serde_json::Value>>;
type Counter = RedisCache<JsonCodec<f64>>;
fn topology() -> Option<RedisTopology> {
let nodes = std::env::var("LITELLM_TEST_REDIS_CLUSTER_NODES").ok()?;
let startup_nodes = nodes
.split(',')
.map(|node| {
let (host, port) = node.trim().rsplit_once(':').expect("host:port");
RedisNode {
host: host.to_string(),
port: port.parse().expect("port"),
}
})
.collect();
Some(RedisTopology::Cluster { startup_nodes })
#[fixture]
fn cache(#[default("cache")] label: &str) -> Option<JsonCache> {
cluster_cache(label, Duration::from_secs(120), JsonCodec::new())
}
fn namespace(label: &str) -> String {
let nanos = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
format!("cluster-test:{label}:{nanos}")
#[fixture]
fn counter(#[default("counter")] label: &str) -> Option<Counter> {
cluster_cache(label, Duration::from_secs(60), JsonCodec::new())
}
fn cluster_url() -> String {
std::env::var("LITELLM_TEST_REDIS_CLUSTER_URL")
.unwrap_or_else(|_| "redis://127.0.0.1:7000".into())
}
fn cluster_cache(label: &str) -> Option<Cache> {
let topology = topology()?;
Some(
Cache::connect(
&cluster_url(),
&topology,
Some(Duration::from_secs(120)),
JsonCodec::new(),
)
.expect("cluster connection")
.with_namespace(Some(namespace(label))),
)
}
fn counter_cache(label: &str) -> Option<RedisCache<JsonCodec<f64>>> {
let topology = topology()?;
Some(
RedisCache::connect(
&cluster_url(),
&topology,
Some(Duration::from_secs(60)),
JsonCodec::new(),
)
.expect("cluster connection")
.with_namespace(Some(namespace(label))),
)
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
fn multi_slot_keys(count: usize) -> Vec<String> {
let keys: Vec<String> = (0..count).map(|index| format!("key-{index}")).collect();
let slots: std::collections::HashSet<Slot> = keys.iter().map(Slot::for_key).collect();
let slots: HashSet<Slot> = keys.iter().map(Slot::for_key).collect();
assert!(slots.len() > 1, "keys must span multiple slots");
keys
}
macro_rules! cluster_or_skip {
($label:expr) => {
match cluster_cache($label) {
Some(cache) => cache,
None => return,
}
};
fn seconds(seconds: u64) -> Option<Duration> {
Some(Duration::from_secs(seconds))
}
#[test]
fn constructor_rejects_clusters_without_startup_nodes() {
let error = Cache::connect(
"redis://127.0.0.1:7000",
&RedisTopology::Cluster {
startup_nodes: Vec::new(),
},
#[rstest]
#[case::no_startup_nodes("redis://127.0.0.1:7000", Vec::new())]
#[case::unix_socket_url(
"redis+unix:///tmp/redis.sock",
vec![RedisNode { host: "127.0.0.1".into(), port: 7000 }]
)]
fn constructor_rejects_unusable_cluster_configs(
#[case] url: &str,
#[case] startup_nodes: Vec<RedisNode>,
) {
let error = JsonCache::connect(
url,
&RedisTopology::Cluster { startup_nodes },
None,
JsonCodec::new(),
)
@ -102,60 +69,56 @@ fn constructor_rejects_clusters_without_startup_nodes() {
assert!(matches!(error, Some(Error::Unavailable)));
}
#[test]
fn constructor_rejects_unix_socket_urls_for_clusters() {
let error = Cache::connect(
"redis+unix:///tmp/redis.sock",
&RedisTopology::Cluster {
startup_nodes: vec![RedisNode {
host: "127.0.0.1".into(),
port: 7000,
}],
},
None,
JsonCodec::new(),
)
.err();
assert!(matches!(error, Some(Error::Unavailable)));
}
#[test]
fn single_key_operations_round_trip_with_ttl_rounding() {
let cache = cluster_or_skip!("single");
#[rstest]
#[tokio::test]
async fn single_key_operations_round_trip_with_ttl_rounding(
#[with("single")] cache: Option<JsonCache>,
) {
let Some(cache) = cache else { return };
let context = ExactCacheContext {
ttl: Some(Duration::from_millis(1500)),
};
let keys = multi_slot_keys(12);
for (index, key) in keys.iter().enumerate() {
cache
.set_cache(key, serde_json::json!({ "index": index }), &context)
.set_cache(key, json!({ "index": index }), &context)
.unwrap();
}
for (index, key) in keys.iter().enumerate() {
assert_eq!(
cache.get_cache(key, &context).unwrap(),
Some(serde_json::json!({ "index": index }))
Some(json!({ "index": index }))
);
}
let runtime = tokio::runtime::Runtime::new().unwrap();
let ttl = runtime.block_on(cache.async_get_ttl(&keys[0])).unwrap();
assert_eq!(ttl, Some(2));
assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(2));
assert!(
cache
.async_refresh_ttl(&keys[0], seconds(40))
.await
.unwrap()
);
assert_eq!(cache.async_get_ttl(&keys[0]).await.unwrap(), seconds(40));
cache.delete_cache(&keys[0]).unwrap();
assert_eq!(cache.get_cache(&keys[0], &context).unwrap(), None);
assert!(!cache.async_refresh_ttl(&keys[0], None).await.unwrap());
assert!(cache.sync_ping().unwrap());
cache.async_flush_cache().await.unwrap();
}
#[rstest]
#[tokio::test]
async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries() {
let cache = cluster_or_skip!("batch");
let context = ExactCacheContext::default();
async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries(
#[with("batch")] cache: Option<JsonCache>,
context: ExactCacheContext,
) {
let Some(cache) = cache else { return };
let keys = multi_slot_keys(40);
for (index, key) in keys.iter().enumerate() {
if index % 5 == 0 {
continue;
}
cache
.async_set_cache(key, serde_json::json!(index), context.clone())
.async_set_cache(key, json!(index), context.clone())
.await
.unwrap();
}
@ -181,41 +144,65 @@ async fn batch_reads_span_slots_and_preserve_order_with_malformed_entries() {
} else if index % 5 == 0 {
BatchEntry::Miss
} else {
BatchEntry::Hit(serde_json::json!(index))
BatchEntry::Hit(json!(index))
};
assert_eq!(*entry, expected, "entry {index}");
}
let sync_entries = cache.batch_get_cache(&keys, &context).unwrap();
assert_eq!(sync_entries, entries);
assert_eq!(cache.batch_get_cache(&keys, &context).unwrap(), entries);
cache.delete_cache_keys(keys.clone()).await.unwrap();
let entries = cache.async_batch_get_cache(keys, context).await.unwrap();
assert!(entries.iter().all(|entry| *entry == BatchEntry::Miss));
}
#[rstest]
#[tokio::test]
async fn pipelines_group_by_slot_and_return_results_in_submission_order() {
let cache = cluster_or_skip!("pipeline");
async fn pipelines_group_by_slot_and_return_results_in_submission_order(
#[with("pipeline")] cache: Option<JsonCache>,
counter: Option<Counter>,
context: ExactCacheContext,
) {
let (Some(cache), Some(counter)) = (cache, counter) else {
return;
};
let keys = multi_slot_keys(30);
let entries = keys
.iter()
.enumerate()
.map(|(index, key)| (key.clone(), serde_json::json!(index)))
.collect();
cache
.async_set_cache_pipeline(entries, ExactCacheContext::default())
.async_set_cache_pipeline(
keys.iter()
.enumerate()
.map(|(index, key)| (key.clone(), json!(index)))
.collect(),
context.clone(),
)
.await
.unwrap();
let hits = cache
.async_batch_get_cache(keys.clone(), ExactCacheContext::default())
.async_batch_get_cache(keys.clone(), context.clone())
.await
.unwrap();
assert!(
hits.iter()
.enumerate()
.all(|(index, entry)| *entry == BatchEntry::Hit(serde_json::json!(index)))
.all(|(index, entry)| *entry == BatchEntry::Hit(json!(index)))
);
cache
.async_set_cache_pipeline_with_ttls(
keys.iter()
.enumerate()
.map(|(index, key)| (key.clone(), json!(index), seconds(index as u64 + 10)))
.collect(),
)
.await
.unwrap();
for (index, key) in keys.iter().enumerate() {
assert_eq!(
cache.async_get_ttl(key).await.unwrap(),
seconds(index as u64 + 10),
"{key}"
);
}
let queues: Vec<String> = keys.iter().map(|key| format!("queue:{key}")).collect();
let pushed = cache
.async_rpush_pipeline(
@ -265,9 +252,6 @@ async fn pipelines_group_by_slot_and_return_results_in_submission_order() {
}
let counters: Vec<String> = keys.iter().map(|key| format!("counter:{key}")).collect();
let Some(counter) = counter_cache("counter") else {
return;
};
let totals = counter
.async_increment_pipeline(
counters
@ -284,25 +268,63 @@ async fn pipelines_group_by_slot_and_return_results_in_submission_order() {
.unwrap();
let expected: Vec<f64> = (0..keys.len()).map(|index| index as f64 + 0.5).collect();
assert_eq!(totals, expected);
assert_eq!(counter.async_get_ttl(&counters[0]).await.unwrap(), Some(30));
assert_eq!(
counter.async_get_ttl(&counters[0]).await.unwrap(),
seconds(30)
);
assert_eq!(counter.async_get_ttl(&counters[1]).await.unwrap(), None);
counter.async_flush_cache().await.unwrap();
cache.async_flush_cache().await.unwrap();
}
#[rstest]
#[tokio::test]
async fn scan_and_scoped_flush_cover_every_primary() {
let cache = cluster_or_skip!("flush");
let other = cluster_or_skip!("other");
let context = ExactCacheContext::default();
async fn rpush_and_trim_is_one_transaction_on_the_key_slot(
#[with("trim")] cache: Option<JsonCache>,
) {
let Some(cache) = cache else { return };
let values = |values: &[&str]| values.iter().map(|value| RedisArg::from(*value)).collect();
assert_eq!(
cache
.async_rpush_and_trim("buf", values(&["a", "b"]), 3)
.await
.unwrap(),
2
);
assert_eq!(
cache
.async_rpush_and_trim("buf", values(&["c", "d"]), 3)
.await
.unwrap(),
4
);
assert_eq!(
cache.async_lpop("buf", Some(10)).await.unwrap(),
RedisLpopResult::Values(vec![b"b".to_vec(), b"c".to_vec(), b"d".to_vec()])
);
cache.async_flush_cache().await.unwrap();
}
#[rstest]
#[tokio::test]
async fn scan_and_scoped_flush_cover_every_primary(
#[with("flush")] cache: Option<JsonCache>,
#[from(cache)]
#[with("other")]
other: Option<JsonCache>,
context: ExactCacheContext,
) {
let (Some(cache), Some(other)) = (cache, other) else {
return;
};
let keys = multi_slot_keys(60);
for key in &keys {
cache
.async_set_cache(key, serde_json::json!(true), context.clone())
.async_set_cache(key, json!(true), context.clone())
.await
.unwrap();
other
.async_set_cache(key, serde_json::json!(true), context.clone())
.async_set_cache(key, json!(true), context.clone())
.await
.unwrap();
}
@ -325,7 +347,7 @@ async fn scan_and_scoped_flush_cover_every_primary() {
let kept = other.async_batch_get_cache(keys, context).await.unwrap();
assert!(
kept.iter()
.all(|entry| *entry == BatchEntry::Hit(serde_json::json!(true)))
.all(|entry| *entry == BatchEntry::Hit(json!(true)))
);
other.async_flush_cache().await.unwrap();
}
@ -361,9 +383,10 @@ fn ping_calls_per_node(startup: &redis::Client) -> Vec<(String, u64)> {
counts
}
#[rstest]
#[tokio::test]
async fn ping_reaches_every_node() {
let cache = cluster_or_skip!("ping");
async fn ping_reaches_every_node(#[with("ping")] cache: Option<JsonCache>) {
let Some(cache) = cache else { return };
let startup = redis::Client::open(cluster_url()).unwrap();
let before = ping_calls_per_node(&startup);
assert!(before.len() >= 2, "{before:?}");
@ -375,14 +398,63 @@ async fn ping_reaches_every_node() {
assert!(cache.sync_ping().unwrap());
let result = cache.test_connection().await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);
assert_eq!(result.message, "Redis Cluster connection test successful");
}
#[rstest]
#[tokio::test]
async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
let Some(counter) = counter_cache("counter") else {
async fn disconnect_closes_idle_connections_and_reconnects_on_demand(
#[with("disconnect")] cache: Option<JsonCache>,
) {
let Some(cache) = cache else { return };
assert!(cache.ping().await.unwrap());
cache.disconnect().await.unwrap();
assert!(cache.ping().await.unwrap());
}
#[rstest]
#[case::keep_existing_ttl(false)]
#[case::refresh_ttl(true)]
#[tokio::test]
async fn increments_refresh_the_ttl_only_when_asked(
counter: Option<Counter>,
#[case] refresh_ttl: bool,
) {
let Some(counter) = counter else { return };
let context = ExactCacheContext { ttl: seconds(60) };
counter
.async_set_cache("spend", 0.0, ExactCacheContext { ttl: seconds(600) })
.await
.unwrap();
assert_eq!(
counter
.async_increment("spend", 1.5, context.clone(), refresh_ttl)
.await
.unwrap(),
1.5
);
assert_eq!(
counter
.async_increment("spend", 2.0, context, refresh_ttl)
.await
.unwrap(),
3.5
);
let ttl = counter.async_get_ttl("spend").await.unwrap().unwrap();
assert_eq!(ttl <= Duration::from_secs(60), refresh_ttl, "{ttl:?}");
counter.async_flush_cache().await.unwrap();
}
#[rstest]
#[tokio::test]
async fn counters_claims_scripts_and_sets_work_on_the_cluster(
counter: Option<Counter>,
#[with("claim")] cache: Option<JsonCache>,
context: ExactCacheContext,
) {
let (Some(counter), Some(cache)) = (counter, cache) else {
return;
};
let context = ExactCacheContext::default();
assert_eq!(
counter
.increment_cache("spend", 1.5, context.clone())
@ -391,7 +463,7 @@ async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
);
assert_eq!(
counter
.async_increment("spend", 2.0, context.clone())
.async_increment("spend", 2.0, context.clone(), false)
.await
.unwrap(),
3.5
@ -413,9 +485,8 @@ async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
assert_eq!(counter.async_set_max("peak", 2.0, None).await.unwrap(), 4.0);
counter.flush_cache().unwrap();
let cache = cluster_or_skip!("claim");
let owner = serde_json::json!("owner-a");
let rival = serde_json::json!("owner-b");
let owner = json!("owner-a");
let rival = json!("owner-b");
assert_eq!(
cache
.claim_cache("lock", owner.clone(), &[], context.clone())
@ -453,8 +524,8 @@ async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
.await
.unwrap();
assert_eq!(reply, redis::Value::Okay);
assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), Some(5));
let evaluated: redis::Value = cache
assert_eq!(cache.async_get_ttl("scripted").await.unwrap(), seconds(5));
let evaluated = cache
.async_eval(
"return redis.call('GET', KEYS[1])".into(),
vec!["scripted".into()],
@ -472,13 +543,13 @@ async fn counters_claims_scripts_and_sets_work_on_the_cluster() {
RedisArg::Bytes(b"a".to_vec()),
RedisArg::Bytes(b"b".to_vec())
],
Some(Duration::from_secs(9)),
seconds(9),
)
.await
.unwrap(),
2
);
assert_eq!(cache.async_get_ttl("members").await.unwrap(), Some(9));
assert_eq!(cache.async_get_ttl("members").await.unwrap(), seconds(9));
let result = cache.test_connection().await.unwrap();
assert_eq!(result.status, CacheConnectionStatus::Success);

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

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

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

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

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

@ -17,4 +17,5 @@ litellm-cache-memory.workspace = true
litellm-cache-redis.workspace = true
redis = "1.7.0"
redis-test = "1.0.4"
rstest.workspace = true
tokio.workspace = true

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

@ -1,14 +1,16 @@
# Response cache foundation
# Response cache
`ResponseCache<B>` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache<Value = CacheEntry>`
## Ownership
`litellm-cache` defines typed storage and codec traits. Memory and Redis implement those traits without depending on response policy. Other consumers can store their own value types using the same backend implementations
`litellm-cache` defines typed storage, codec, and capability traits. `BaseCache` is only get, set, TTL, and pipeline writes. Everything else is an optional capability a backend implements only where its Python class defines the method: `DisconnectCache`, `ConnectionCache` (`test_connection`), `PingCache`, `BatchCache`, `DeleteCache`, `FlushCache`, counters, queues, TTL, scan, and scripts. Memory, Redis, disk, S3, GCS, and Azure Blob implement those traits without depending on response policy, so other consumers can store their own value types in the same backends
Semantic backends (Redis, Valkey, Qdrant) are generic over their embedder and codec, and share one prompt and embedding contract from `litellm_cache::semantic`. They take a `SemanticCacheContext`, so `ResponseCache` drives them the same way it drives exact backends
`litellm-cache-response` owns response keys, controls, entries, the Python-compatible response codec, and `WriteBuffer`, the backend-neutral deferred-write policy. It has no runtime dependency on a specific cache backend or Python
The Python bridge constructs backends and selects them through its private `NativeResponseCache` enum, which only dispatches. Generic Rust callers inject their backend directly. A native gateway can construct the same generic response service in its own host
`ExactResponseCache` is the object-safe view of a `ResponseCache` over an exact backend. `ConnectionProbe` is the object-safe `test_connection`, implemented only when the backend implements `ConnectionCache`, so a host holds one next to its `ExactResponseCache` and reports the operation as unsupported otherwise, as Python's `BaseCache` does. Lookup, store, batch, and flush never require it
## Native Rust use
@ -28,36 +30,22 @@ cache.store(&request, json!({"answer": 7}), now)?;
assert_eq!(cache.async_lookup(&request, now).await?, Some(json!({"answer": 7})));
```
For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved. Sync operations check out independent connections from a bounded pool, while async callers, including counters and claims, move that blocking work off the executor. The pool skips the checkout PING and instead discards any connection whose command failed
For Redis, inject `RedisCache::new(url, ttl, ResponseCacheCodec)` instead. Namespaces are optional and existing namespace prefixes are preserved
Callers supply Unix time for response freshness. Backend TTL uses its own clock. A read can reject an entry through `max_age` even while the backend still retains it
## Python integration boundary
## Python integration
The extension keeps a private test harness for memory and Redis single and batch response lookup and storage. Batch lookup returns ordered values plus missing indices for embedding partial-hit wiring. No bridge-only cache type is part of the public API
The bridge activates backends through the Rust catalog in `litellm/rust_bridge/catalog.py`. Every cache rule ships as `PYTHON_ONLY`, so SDK, Router, and proxy calls stay on Python and construct no native cache resources until a rule is changed
The bridge also exposes a production-shaped response cache runtime selected through the Rust catalog. Its shipped rule set is empty, so current SDK, Router, and proxy calls stay on Python and do not construct native cache resources. Tests can inject a rule and build the native memory runtime from an ordinary Python `Cache` configuration without changing the legacy cache classes
When a rule selects a backend, the Python `Cache` facade builds the native runtime from its own configuration and routes its storage calls (sync and async lookup and store, and pipelined batch store) to it. Stream replay, embedding partial-hit merging, response reconstruction, and callbacks stay in Python on top of that native store. The Python backend object remains for its direct API
Object responses are written as they are, and every other response shape is written as a serialized string, which is the pair of shapes Python reads. A string on the wire is therefore always a serialized response, so string-valued responses round trip. Typed backends such as memory never pass through the codec
The resolver reads the namespace's `cache` attribute each time it resolves. A captured binding retains the selected service for its operation, including background writes. `None` disables caching. Custom Python cache objects keep their original methods, arguments, returned awaitables, exceptions, and caller-task execution
Python callbacks use the built-in `Cache` API, so a `Cache` subclass works unchanged. A batch lookup takes one original kwargs mapping per request and returns the list of `get_cache` or gathered `async_get_cache` results, while native bindings return `{values, missing_indices}`. A batch store hands the caller's original result to `async_add_cache_pipeline`. `ping` calls `ping`, and a flush goes to the facade's backend
The private facade test harness checks object identity, method overrides, effective TTL, Redis namespace, memory capacity, and later configuration changes before selecting native execution. Its snapshot includes Redis connection settings, so a later `redis_kwargs` change, including an SSL option, selects Python callback execution. Buffered async writes honor `redis_flush_size`. Public activation must construct the shared native service from the initial Python Redis settings, including `litellm.default_redis_ttl` and SSL options. A buffered entry keeps the time it was produced, and a failed flush drops its batch instead of growing the buffer during an outage. The harness does not migrate entries or replace Python methods. Until activation configures one shared service, the Python facade and native test service can hold separate data. Existing public cache constructors remain on Python
Native cache handles must be recreated after fork. The bridge releases the GIL around native operations, and Redis runs blocking connection operations off the async executor. Native errors propagate to the host, which owns the existing fail-open and logging policy
The Redis backend also provides the primitives needed to preserve its direct Python surface later: TLS URLs, ping, bulk delete, counter batches, TTL, scan, set membership, raw queue push and pop, queue and counter pipelines, counter floor and maximum operations, script evaluation, client information, namespaced flush, and full flush. These are backend operations only and are not exported to Python by this PR. Memory provides TTL, oldest-key, and counter-pipeline operations
Native cache handles must be recreated after fork. Native errors propagate to the host, which owns the existing fail-open and logging policy
## Adding another backend
Implement `BaseCache` for the backend with its associated value type, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache<B>` then works without another response implementation. Add a concrete bridge enum variant and constructor only when exposing that backend to Python
Implement `BaseCache` for the backend with its associated value type and the capability traits its Python class supports, and accept a `CacheCodec` when wire serialization is needed. `ResponseCache<B>` then works without another response implementation
Verify typed values, TTL precedence, missing entries, serialization failures, namespaces, batch ordering, and sync/async behavior. Run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before enabling a public facade
## Follow-up scope
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, 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
Run the `litellm-cache-testing` contract checks the backend's capabilities allow, and run response fixtures with `ResponseCacheCodec`, including both Python envelope encodings, before adding a catalog rule

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

@ -1,7 +1,8 @@
use std::{future::Future, pin::Pin, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
BaseCache, BatchCache, CacheConnectionResult, ConnectionCache, Error, ExactCacheContext,
FlushCache,
};
use serde_json::Value;
@ -61,10 +62,25 @@ pub trait ExactResponseCache: Send + Sync {
) -> BoxFuture<'a, Result<(), Error>>;
fn async_flush<'a>(&'a self) -> BoxFuture<'a, Result<(), Error>>;
}
/// Object-safe `test_connection` for the exact backends whose Python class defines it. Hosts hold
/// one next to their `ExactResponseCache` when the backend has it, and report the operation as
/// unsupported otherwise, as Python's `BaseCache.test_connection` does.
pub trait ConnectionProbe: Send + Sync {
fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result<CacheConnectionResult, Error>>;
}
impl<B> ConnectionProbe for ResponseCache<B>
where
B: ConnectionCache<Value = CacheEntry>,
B::Context: Default + PartialEq,
{
fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result<CacheConnectionResult, Error>> {
Box::pin(ResponseCache::test_connection(self))
}
}
impl<B> ExactResponseCache for ResponseCache<B>
where
B: BaseCache<Value = CacheEntry, Context = ExactCacheContext> + BatchCache + FlushCache,
@ -141,8 +157,4 @@ where
fn async_flush<'a>(&'a self) -> BoxFuture<'a, Result<(), Error>> {
Box::pin(ResponseCache::async_flush(self))
}
fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result<CacheConnectionResult, Error>> {
Box::pin(ResponseCache::test_connection(self))
}
}

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

@ -12,5 +12,5 @@ pub use caching::{
};
pub use codec::ResponseCacheCodec;
pub use embedding::PartialHits;
pub use exact::ExactResponseCache;
pub use exact::{ConnectionProbe, ExactResponseCache};
pub use response::{ResponseCache, ResponseCacheRequest};

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

@ -1,7 +1,9 @@
use std::{sync::Arc, time::Duration};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, FlushCache,
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, ConnectionCache, Error,
FlushCache,
semantic::{SemanticCache, SemanticLookup},
};
use serde_json::Value;
@ -78,7 +80,10 @@ where
self.backend.async_flush_cache().await
}
pub async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
pub async fn test_connection(&self) -> Result<CacheConnectionResult, Error>
where
B: ConnectionCache,
{
self.backend.test_connection().await
}
@ -121,6 +126,43 @@ where
Ok(Self::fresh_or_miss(entry, now, request.max_age))
}
/// `lookup` plus the similarity the semantic backend reports. Freshness applies to the
/// value only: Python stamps the similarity before its max-age check.
pub fn lookup_semantic(
&self,
request: &ResponseCacheRequest<B::Context>,
now: Duration,
) -> Result<SemanticLookup<Value>, Error>
where
B: SemanticCache,
{
if !request.controls.reads() {
return Ok(SemanticLookup::miss(None));
}
let lookup = self
.backend
.get_cache_with_similarity(&cache_key(&request.key), &request.context);
Self::fresh_semantic(lookup, now, request.max_age)
}
pub async fn async_lookup_semantic(
&self,
request: &ResponseCacheRequest<B::Context>,
now: Duration,
) -> Result<SemanticLookup<Value>, Error>
where
B: SemanticCache,
{
if !request.controls.reads() {
return Ok(SemanticLookup::miss(None));
}
let lookup = self
.backend
.async_get_cache_with_similarity(&cache_key(&request.key), &request.context)
.await;
Self::fresh_semantic(lookup, now, request.max_age)
}
pub fn lookup_batch(
&self,
requests: &[ResponseCacheRequest<B::Context>],
@ -290,6 +332,21 @@ where
Ok(PartialHits::new(values))
}
fn fresh_semantic(
lookup: Result<SemanticLookup<CacheEntry>, Error>,
now: Duration,
max_age: Option<Duration>,
) -> Result<SemanticLookup<Value>, Error> {
match lookup {
Ok(lookup) => Ok(SemanticLookup {
value: Self::fresh_or_miss(lookup.value, now, max_age),
similarity: lookup.similarity,
}),
Err(Error::InvalidEntry) => Ok(SemanticLookup::miss(None)),
Err(error) => Err(error),
}
}
fn fresh_or_miss(
entry: Option<CacheEntry>,
now: Duration,

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

@ -1,90 +1,174 @@
use litellm_cache_response::{
CacheControls, CacheKeyContext, CacheKeyField, CacheKeyInput, cache_key, get_cache_key,
should_use_cache,
};
use rstest::rstest;
use sha2::{Digest, Sha256};
#[test]
fn keys_match_python_order_groups_files_presets_and_namespaces() {
let mut input = CacheKeyInput {
fields: vec![
CacheKeyField {
name: "model".into(),
value: Some("deployment".into()),
api_parameter: true,
internal_parameter: false,
},
CacheKeyField {
name: "file".into(),
value: None,
api_parameter: true,
internal_parameter: false,
},
],
namespace: Some("team".into()),
..Default::default()
};
fn field(name: &str, value: Option<&str>) -> CacheKeyField {
CacheKeyField {
name: name.into(),
value: value.map(str::to_owned),
api_parameter: true,
internal_parameter: false,
}
}
fn hash(preimage: &[u8]) -> String {
format!("{:x}", Sha256::digest(preimage))
}
#[rstest]
#[case::caching_group_and_checksum(
CacheKeyContext {
model_group: Some("group".into()),
caching_groups: vec![(vec!["group".into()], "['group']".into())],
file_checksum: Some("checksum".into()),
..Default::default()
}
.apply(&mut input);
assert_eq!(
cache_key(&input),
format!(
"team:{:x}",
Sha256::digest(b"model: ['group']file: checksum")
)
);
input.preset = Some("preset".into());
},
Some("team"),
"team:",
b"model: ['group']file: checksum".as_slice(),
)]
#[case::model_group_outside_caching_groups(
CacheKeyContext {
model_group: Some("group".into()),
caching_groups: vec![(vec!["other".into()], "['other']".into())],
file_object_name: Some("object".into()),
..Default::default()
},
None,
"",
b"model: groupfile: object".as_slice(),
)]
#[case::metadata_file_name_before_parameters(
CacheKeyContext {
metadata_file_name: Some("metadata".into()),
parameters_file_name: Some("parameters".into()),
..Default::default()
},
Some(""),
"",
b"model: deploymentfile: metadata".as_slice(),
)]
#[case::parameters_file_name_last(
CacheKeyContext {
parameters_file_name: Some("parameters".into()),
..Default::default()
},
None,
"",
b"model: deploymentfile: parameters".as_slice(),
)]
#[case::no_context_keeps_the_request_model(
CacheKeyContext::default(),
Some("team"),
"team:",
b"model: deployment".as_slice(),
)]
fn keys_match_python_order_groups_files_and_namespaces(
#[case] context: CacheKeyContext,
#[case] namespace: Option<&str>,
#[case] prefix: &str,
#[case] preimage: &[u8],
) {
let mut input = CacheKeyInput {
fields: vec![field("model", Some("deployment")), field("file", None)],
namespace: namespace.map(str::to_owned),
..Default::default()
};
context.apply(&mut input);
let expected = format!("{prefix}{}", hash(preimage));
assert_eq!(cache_key(&input), expected);
assert_eq!(get_cache_key(&input), expected);
}
#[rstest]
#[case::api_parameter(true, false, false, true)]
#[case::provider_parameter_when_included(false, false, true, true)]
#[case::provider_parameter_when_excluded(false, false, false, false)]
#[case::internal_parameter_never(false, true, true, false)]
fn keys_hash_api_and_opted_in_provider_parameters(
#[case] api_parameter: bool,
#[case] internal_parameter: bool,
#[case] include_provider_parameters: bool,
#[case] hashed: bool,
) {
let input = CacheKeyInput {
fields: vec![
field("model", Some("a")),
CacheKeyField {
name: "extra".into(),
value: Some("x".into()),
api_parameter,
internal_parameter,
},
],
include_provider_parameters,
..Default::default()
};
let preimage: &[u8] = if hashed {
b"model: aextra: x"
} else {
b"model: a"
};
assert_eq!(cache_key(&input), hash(preimage));
}
#[rstest]
#[case::without_namespace(None)]
#[case::with_namespace(Some("team"))]
fn preset_keys_are_used_verbatim(#[case] namespace: Option<&str>) {
let input = CacheKeyInput {
fields: vec![field("model", Some("a"))],
preset: Some("preset".into()),
namespace: namespace.map(str::to_owned),
..Default::default()
};
assert_eq!(cache_key(&input), "preset");
assert_eq!(get_cache_key(&input), "preset");
}
#[test]
fn cache_controls_honor_default_modes_and_directives() {
let enabled = CacheControls {
supported_call_type: true,
configured: true,
default_on: true,
..Default::default()
};
assert!(enabled.reads());
assert!(enabled.writes());
assert!(
!CacheControls {
default_on: false,
..enabled
}
.reads()
);
assert!(
CacheControls {
default_on: false,
use_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_cache: true,
..enabled
}
.reads()
);
assert!(
!CacheControls {
no_store: true,
..enabled
}
.writes()
);
assert!(
!CacheControls {
caching: Some(false),
..enabled
}
.writes()
);
const ENABLED: CacheControls = CacheControls {
supported_call_type: true,
configured: true,
native_backend: false,
default_on: true,
caching: None,
no_cache: false,
no_store: false,
use_cache: false,
};
#[rstest]
#[case::enabled(ENABLED, true, true)]
#[case::default_off(CacheControls { default_on: false, ..ENABLED }, false, false)]
#[case::default_off_with_use_cache(
CacheControls { default_on: false, use_cache: true, ..ENABLED },
true,
true
)]
#[case::no_cache(CacheControls { no_cache: true, ..ENABLED }, false, true)]
#[case::no_store(CacheControls { no_store: true, ..ENABLED }, true, false)]
#[case::no_cache_and_no_store(
CacheControls { no_cache: true, no_store: true, ..ENABLED },
false,
false
)]
#[case::caching_disabled(CacheControls { caching: Some(false), ..ENABLED }, false, false)]
#[case::caching_enabled(CacheControls { caching: Some(true), ..ENABLED }, true, true)]
#[case::unsupported_call_type(
CacheControls { supported_call_type: false, ..ENABLED },
false,
false
)]
#[case::unconfigured(CacheControls { configured: false, ..ENABLED }, false, false)]
fn cache_controls_honor_default_modes_and_directives(
#[case] controls: CacheControls,
#[case] reads: bool,
#[case] writes: bool,
) {
assert_eq!(controls.reads(), reads);
assert_eq!(controls.writes(), writes);
assert_eq!(should_use_cache(controls), reads || writes);
}

110
litellm-rust/crates/cache-response/tests/codec.rs Normal file
View file

@ -0,0 +1,110 @@
use litellm_cache::{CacheCodec, Error};
use litellm_cache_response::{CacheEntry, ResponseCacheCodec};
use rstest::rstest;
use serde_json::{Value, json};
fn entry(response: Value) -> CacheEntry {
CacheEntry {
timestamp: Some(100.0),
response,
}
}
#[rstest]
#[case::python_literal_object(
br#"{'timestamp': 100.0, 'response': {'text': 'hello \\ world', 'flag': True, 'empty': None, 'list': [1, 2.5]}}"#.as_slice(),
json!({"text": "hello \\ world", "flag": true, "empty": null, "list": [1, 2.5]}),
)]
#[case::python_sync_string_response(
br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.as_slice(),
json!({"ok": true, "text": "cached"}),
)]
#[case::python_literal_string_response(
br#"{'timestamp': 100.0, 'response': "{'ok': True, 'items': (1, 2)}"}"#.as_slice(),
json!({"ok": true, "items": [1, 2]}),
)]
#[case::json_object_response(
br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice(),
json!({"ok": true, "text": "cached"}),
)]
#[case::json_string_response(
br#"{"timestamp": 100.0, "response": "[1,2]"}"#.as_slice(),
json!([1, 2]),
)]
fn decode_reads_every_python_envelope(#[case] bytes: &[u8], #[case] response: Value) {
assert_eq!(ResponseCacheCodec.decode(bytes).unwrap(), entry(response));
}
#[rstest]
#[case::code_is_not_executed(b"__import__('os').system('false')".to_vec())]
#[case::non_numeric_timestamp(b"{'timestamp': 'invalid', 'response': {}}".to_vec())]
#[case::infinite_timestamp(b"{'timestamp': 1e9999, 'response': {}}".to_vec())]
#[case::missing_response(br#"{"timestamp": 100.0}"#.to_vec())]
#[case::unserialized_string_response(br#"{"timestamp": 100.0, "response": "not serialized"}"#.to_vec())]
#[case::non_utf8(vec![0xff, 0xfe])]
#[case::deep_nesting(format!("{}None{}", "[".repeat(1000), "]".repeat(1000)).into_bytes())]
fn decode_rejects_invalid_entries(#[case] bytes: Vec<u8>) {
assert_eq!(
ResponseCacheCodec.decode(&bytes).unwrap_err(),
Error::InvalidEntry
);
}
#[rstest]
#[case::nan(f64::NAN)]
#[case::infinity(f64::INFINITY)]
fn encode_rejects_non_finite_timestamps(#[case] timestamp: f64) {
assert_eq!(
ResponseCacheCodec
.encode(&CacheEntry {
timestamp: Some(timestamp),
response: json!({}),
})
.unwrap_err(),
Error::InvalidEntry
);
}
#[rstest]
#[case::object(json!({"choices": [{"text": "cached"}]}), json!({"choices": [{"text": "cached"}]}))]
#[case::array(json!([1, 2]), json!("[1,2]"))]
#[case::number(json!(7), json!("7"))]
#[case::null(json!(null), json!("null"))]
#[case::string(json!("hello world"), json!("\"hello world\""))]
#[case::numeric_string(json!("123"), json!("\"123\""))]
#[case::null_string(json!("null"), json!("\"null\""))]
fn encode_writes_python_readable_envelopes_that_round_trip(
#[case] response: Value,
#[case] wire_response: Value,
) {
let wire = ResponseCacheCodec.encode(&entry(response.clone())).unwrap();
assert_eq!(
serde_json::from_slice::<Value>(&wire).unwrap(),
json!({"timestamp": 100.0, "response": wire_response})
);
assert_eq!(ResponseCacheCodec.decode(&wire).unwrap(), entry(response));
}
#[rstest]
fn object_entries_preserve_the_existing_json_representation() {
let entry = CacheEntry {
timestamp: Some(123.0),
response: json!({"choices": [{"text": "cached"}]}),
};
let bytes = ResponseCacheCodec.encode(&entry).unwrap();
assert_eq!(bytes, serde_json::to_vec(&entry).unwrap());
assert_eq!(ResponseCacheCodec.decode(&bytes).unwrap(), entry);
}
#[rstest]
#[case::json(br#"{"choices": [{"text": "legacy"}]}"#.as_slice())]
#[case::python_literal(br#"{'choices': [{'text': 'legacy'}]}"#.as_slice())]
fn values_without_timestamps_decode_as_bare_responses(#[case] bytes: &[u8]) {
assert_eq!(
ResponseCacheCodec.decode(bytes).unwrap(),
CacheEntry {
timestamp: None,
response: json!({"choices": [{"text": "legacy"}]}),
}
);
}

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

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

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

@ -1,3 +1,5 @@
mod support;
use std::{
sync::{
Arc, Mutex,
@ -7,32 +9,22 @@ use std::{
};
use litellm_cache::{
BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error,
SemanticCacheContext,
BaseCache, Error, SemanticCacheContext,
semantic::{SemanticCache, SemanticLookup},
};
use litellm_cache_memory::InMemoryCache;
use litellm_cache_redis::RedisCache;
use litellm_cache_response::{
CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheCodec,
ResponseCacheRequest, WriteBuffer,
CacheControls, CacheEntry, CacheKeyField, CacheKeyInput, ResponseCache, ResponseCacheRequest,
WriteBuffer, cache_key,
};
use redis_test::{MockCmd, MockRedisConnection};
use serde_json::json;
use redis_test::MockCmd;
use rstest::rstest;
use serde_json::{Value, json};
use support::{keyed, memory, redis, request};
fn memory() -> Arc<ResponseCache<InMemoryCache<CacheEntry>>> {
Arc::new(ResponseCache::new(Arc::new(InMemoryCache::new(
Some(8),
Some(Duration::from_secs(600)),
))))
}
fn request() -> ResponseCacheRequest {
ResponseCacheRequest::new(CacheKeyInput {
preset: Some("tenant:key".into()),
..Default::default()
})
}
type Memory = Arc<ResponseCache<InMemoryCache<CacheEntry>>>;
#[derive(Default)]
struct SemanticBackend {
entries: Mutex<Vec<(String, CacheEntry)>>,
contexts: Mutex<Vec<SemanticCacheContext>>,
@ -67,50 +59,133 @@ impl BaseCache for SemanticBackend {
.find(|(entry_key, _)| entry_key == key)
.map(|(_, entry)| entry.clone()))
}
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Ok(CacheConnectionResult {
status: CacheConnectionStatus::Success,
message: "ok".into(),
error: None,
})
}
}
#[test]
fn semantic_context_reaches_backend_for_store_and_lookup() {
let backend = Arc::new(SemanticBackend {
entries: Mutex::new(Vec::new()),
contexts: Mutex::new(Vec::new()),
});
#[rstest]
#[tokio::test]
async fn semantic_context_reaches_backend_for_store_and_lookup(
request: ResponseCacheRequest,
#[values(false, true)] asynchronous: bool,
) {
let backend = Arc::new(SemanticBackend::default());
let cache = ResponseCache::new(backend.clone());
let context = SemanticCacheContext {
messages: Some(json!([{"role": "user", "content": "hello"}])),
..Default::default()
};
let request = request().with_context(context.clone());
let request = request.with_context(context.clone());
let response = json!({"answer": 42});
let now = Duration::from_secs(100);
cache
.store(&request, response.clone(), Duration::from_secs(100))
.unwrap();
let hit = if asynchronous {
cache
.async_store(&request, response.clone(), now)
.await
.unwrap();
cache.async_lookup(&request, now).await.unwrap()
} else {
cache.store(&request, response.clone(), now).unwrap();
cache.lookup(&request, now).unwrap()
};
assert_eq!(
cache.lookup(&request, Duration::from_secs(100)).unwrap(),
Some(response)
);
assert_eq!(hit, Some(response));
assert_eq!(
backend.contexts.lock().unwrap().as_slice(),
&[context.clone(), context]
);
}
/// A semantic backend that answers every read with one fixed lookup.
struct ScoredBackend(Result<SemanticLookup<CacheEntry>, Error>);
impl BaseCache for ScoredBackend {
type Value = CacheEntry;
type Context = SemanticCacheContext;
fn get_ttl(&self, _: &Self::Context) -> Option<Duration> {
None
}
fn set_cache(&self, _: &str, _: Self::Value, _: &Self::Context) -> Result<(), Error> {
Ok(())
}
fn get_cache(&self, key: &str, context: &Self::Context) -> Result<Option<Self::Value>, Error> {
self.get_cache_with_similarity(key, context)
.map(|lookup| lookup.value)
}
}
impl SemanticCache for ScoredBackend {
fn get_cache_with_similarity(
&self,
_: &str,
_: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
self.0.clone()
}
async fn async_get_cache_with_similarity(
&self,
key: &str,
context: &Self::Context,
) -> Result<SemanticLookup<Self::Value>, Error> {
self.get_cache_with_similarity(key, context)
}
}
fn scored(timestamp: f64, similarity: f64) -> Result<SemanticLookup<CacheEntry>, Error> {
Ok(SemanticLookup {
value: Some(CacheEntry {
timestamp: Some(timestamp),
response: json!({"answer": 42}),
}),
similarity: Some(similarity),
})
}
#[rstest]
#[case::fresh_hit(scored(95.0, 0.95), true, Ok(SemanticLookup { value: Some(json!({"answer": 42})), similarity: Some(0.95) }))]
#[case::stale_hit_keeps_the_similarity(
scored(50.0, 0.95),
true,
Ok(SemanticLookup::miss(Some(0.95)))
)]
#[case::miss_keeps_the_similarity(
Ok(SemanticLookup::miss(Some(0.4))),
true,
Ok(SemanticLookup::miss(Some(0.4)))
)]
#[case::no_search(Ok(SemanticLookup::miss(None)), true, Ok(SemanticLookup::miss(None)))]
#[case::disabled_reads_skip_the_backend(scored(95.0, 0.95), false, Ok(SemanticLookup::miss(None)))]
#[case::invalid_entry_is_a_miss(Err(Error::InvalidEntry), true, Ok(SemanticLookup::miss(None)))]
#[case::backend_errors_propagate(Err(Error::Unavailable), true, Err(Error::Unavailable))]
#[tokio::test]
async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
async fn semantic_lookup_applies_freshness_to_the_value_only(
#[case] backend: Result<SemanticLookup<CacheEntry>, Error>,
#[case] reads: bool,
#[case] expected: Result<SemanticLookup<Value>, Error>,
#[values(false, true)] asynchronous: bool,
request: ResponseCacheRequest,
) {
let cache = ResponseCache::new(Arc::new(ScoredBackend(backend)));
let mut request = request.with_context(SemanticCacheContext::default());
request.max_age = Some(Duration::from_secs(10));
request.controls.no_cache = !reads;
let now = Duration::from_secs(100);
let lookup = if asynchronous {
cache.async_lookup_semantic(&request, now).await
} else {
cache.lookup_semantic(&request, now)
};
assert_eq!(lookup, expected);
}
#[rstest]
#[tokio::test]
async fn sync_and_async_consumers_share_keys_ttls_and_freshness(mut request: ResponseCacheRequest) {
let clock = Arc::new(AtomicU64::new(100));
let backend = Arc::new(InMemoryCache::with_clock(
Some(8),
@ -121,7 +196,6 @@ async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
},
));
let cache = ResponseCache::new(backend.clone());
let mut request = request();
request.context.ttl = Some(Duration::from_secs(10));
request.max_age = Some(Duration::from_secs(5));
cache
@ -172,80 +246,164 @@ async fn sync_and_async_consumers_share_keys_ttls_and_freshness() {
);
}
#[derive(Clone, Copy, Debug)]
enum Directive {
Plain,
NoCache,
NoStore,
DefaultOff,
UseCache,
CachingOff,
Unsupported,
}
impl Directive {
fn apply(self, controls: &mut CacheControls) {
match self {
Self::Plain => {}
Self::NoCache => controls.no_cache = true,
Self::NoStore => controls.no_store = true,
Self::DefaultOff => controls.default_on = false,
Self::UseCache => {
controls.default_on = false;
controls.use_cache = true;
}
Self::CachingOff => controls.caching = Some(false),
Self::Unsupported => controls.supported_call_type = false,
}
}
}
#[rstest]
#[case::plain(Directive::Plain, Directive::Plain, true)]
#[case::no_store_skips_the_write(Directive::NoStore, Directive::Plain, false)]
#[case::no_store_keeps_reads(Directive::Plain, Directive::NoStore, true)]
#[case::no_cache_keeps_writes(Directive::NoCache, Directive::Plain, true)]
#[case::no_cache_skips_the_read(Directive::Plain, Directive::NoCache, false)]
#[case::default_off_skips_the_write(Directive::DefaultOff, Directive::Plain, false)]
#[case::default_off_skips_the_read(Directive::Plain, Directive::DefaultOff, false)]
#[case::use_cache_opts_in_under_default_off(Directive::UseCache, Directive::UseCache, true)]
#[case::caching_off_skips_the_write(Directive::CachingOff, Directive::Plain, false)]
#[case::caching_off_skips_the_read(Directive::Plain, Directive::CachingOff, false)]
#[case::unsupported_call_type_skips_the_write(Directive::Unsupported, Directive::Plain, false)]
#[case::unsupported_call_type_skips_the_read(Directive::Plain, Directive::Unsupported, false)]
#[tokio::test]
async fn directives_skip_io_and_keep_reads_and_writes_independent() {
let cache = memory();
let mut request = request();
async fn directives_skip_io_and_keep_reads_and_writes_independent(
memory: Memory,
request: ResponseCacheRequest,
#[case] write: Directive,
#[case] read: Directive,
#[case] hit: bool,
#[values(false, true)] asynchronous: bool,
) {
let now = Duration::from_secs(100);
request.controls.no_store = true;
cache
.async_store(&request, json!({"v": 1}), now)
.await
.unwrap();
assert_eq!(cache.lookup(&request, now).unwrap(), None);
request.controls.no_store = false;
request.controls.no_cache = true;
cache.store(&request, json!({"v": 2}), now).unwrap();
assert_eq!(cache.async_lookup(&request, now).await.unwrap(), None);
request.controls.no_cache = false;
assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
request.controls.default_on = false;
cache.store(&request, json!({"v": 3}), now).unwrap();
assert_eq!(cache.lookup(&request, now).unwrap(), None);
request.controls.use_cache = true;
assert_eq!(cache.lookup(&request, now).unwrap(), Some(json!({"v": 2})));
request.controls.supported_call_type = false;
assert_eq!(cache.lookup(&request, now).unwrap(), None);
}
let mut writer = request.clone();
write.apply(&mut writer.controls);
let mut reader = request;
read.apply(&mut reader.controls);
#[tokio::test]
async fn redis_consumer_reads_python_sync_and_async_envelopes_and_writes_compatible_json() {
let connection = MockRedisConnection::new([
MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.to_vec()),
),
MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.to_vec()),
),
MockCmd::new(
redis::cmd("SETEX")
.arg("tenant:key")
.arg(600)
.arg(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()),
Ok("OK"),
),
])
.assert_all_commands_consumed();
let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec)
.with_namespace(Some("tenant".into()));
let cache = ResponseCache::new(Arc::new(backend));
let request = request();
let expected = json!({"ok": true, "text": "cached"});
assert_eq!(
cache.lookup(&request, Duration::from_secs(101)).unwrap(),
Some(expected.clone())
);
assert_eq!(
cache
.async_lookup(&request, Duration::from_secs(101))
if asynchronous {
memory
.async_store(&writer, json!({"v": 1}), now)
.await
.unwrap(),
Some(expected.clone())
.unwrap();
} else {
memory.store(&writer, json!({"v": 1}), now).unwrap();
}
let found = if asynchronous {
memory.async_lookup(&reader, now).await.unwrap()
} else {
memory.lookup(&reader, now).unwrap()
};
assert_eq!(
found,
hit.then(|| json!({"v": 1})),
"{write:?} then {read:?}"
);
}
#[rstest]
#[case::python_sync_literal(
br#"{'timestamp': 100.0, 'response': '{"ok": true, "text": "cached"}'}"#.as_slice()
)]
#[case::python_async_json(br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice())]
#[tokio::test]
async fn redis_consumer_reads_python_sync_and_async_envelopes(
request: ResponseCacheRequest,
#[case] stored: &[u8],
#[values(false, true)] asynchronous: bool,
) {
let cache = redis(
vec![MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(stored.to_vec()),
)],
Some("tenant"),
);
let now = Duration::from_secs(101);
let found = if asynchronous {
cache.async_lookup(&request, now).await.unwrap()
} else {
cache.lookup(&request, now).unwrap()
};
assert_eq!(found, Some(json!({"ok": true, "text": "cached"})));
}
#[rstest]
#[case::object(
json!({"ok": true, "text": "cached"}),
br#"{"timestamp":100.0,"response":{"ok":true,"text":"cached"}}"#.as_slice()
)]
#[case::array(json!([1, 2]), br#"{"timestamp":100.0,"response":"[1,2]"}"#.as_slice())]
#[tokio::test]
async fn redis_consumer_writes_python_compatible_json(
request: ResponseCacheRequest,
#[case] response: Value,
#[case] wire: &[u8],
) {
let cache = redis(
vec![MockCmd::new(
redis::cmd("SETEX").arg("tenant:key").arg(600).arg(wire),
Ok("OK"),
)],
Some("tenant"),
);
cache
.async_store(&request, expected, Duration::from_secs(100))
.async_store(&request, response, Duration::from_secs(100))
.await
.unwrap();
}
#[rstest]
#[tokio::test]
async fn captured_service_keeps_the_selected_backend_for_background_writes() {
let original = memory();
async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis(
mut request: ResponseCacheRequest,
) {
let cache = redis(
vec![MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(b"invalid".to_vec()),
)],
None,
);
request.controls.no_cache = true;
assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
request.controls.no_cache = false;
assert_eq!(
cache.async_lookup(&request, Duration::ZERO).await.unwrap(),
None
);
}
#[rstest]
#[tokio::test]
async fn captured_service_keeps_the_selected_backend_for_background_writes(
#[from(memory)] original: Memory,
#[from(memory)] replacement: Memory,
request: ResponseCacheRequest,
) {
let captured = original.clone();
let replacement = memory();
let request = request();
let writer = tokio::spawn({
let request = request.clone();
async move {
@ -271,9 +429,13 @@ async fn captured_service_keeps_the_selected_backend_for_background_writes() {
);
}
#[test]
fn generated_keys_preserve_namespace_and_explicit_keys() {
let cache = memory();
#[rstest]
#[case::with_namespace(Some("tenant"))]
#[case::without_namespace(None)]
fn generated_keys_preserve_namespace_and_explicit_keys(
memory: Memory,
#[case] namespace: Option<&str>,
) {
let key = CacheKeyInput {
fields: vec![CacheKeyField {
name: "model".into(),
@ -281,201 +443,125 @@ fn generated_keys_preserve_namespace_and_explicit_keys() {
api_parameter: true,
internal_parameter: false,
}],
namespace: Some("tenant".into()),
namespace: namespace.map(str::to_owned),
..Default::default()
};
let generated = ResponseCacheRequest::new(key.clone());
let explicit = ResponseCacheRequest::new(CacheKeyInput {
preset: Some(litellm_cache_response::cache_key(&key)),
..Default::default()
});
cache
let explicit = keyed(&cache_key(&key));
memory
.store(&generated, json!({"value": 7}), Duration::from_secs(100))
.unwrap();
assert_eq!(
cache.lookup(&explicit, Duration::from_secs(100)).unwrap(),
memory.lookup(&explicit, Duration::from_secs(100)).unwrap(),
Some(json!({"value":7}))
);
}
#[test]
fn response_codec_accepts_python_literals_without_executing_code() {
let bytes = br#"{'timestamp': 100.0, 'response': {'text': 'hello \\ world', 'flag': True, 'empty': None, 'list': [1, 2.5]}}"#;
let entry = ResponseCacheCodec.decode(bytes).unwrap();
assert_eq!(
entry.response,
json!({"text": "hello \\ world", "flag": true, "empty": null, "list": [1, 2.5]})
);
for bytes in [
b"__import__('os').system('false')".as_slice(),
b"{'timestamp': 'invalid', 'response': {}}",
b"{'timestamp': 1e9999, 'response': {}}",
] {
assert_eq!(
ResponseCacheCodec.decode(bytes).unwrap_err(),
Error::InvalidEntry
);
}
let deep = format!("{}None{}", "[".repeat(1000), "]".repeat(1000));
assert_eq!(
ResponseCacheCodec.decode(deep.as_bytes()).unwrap_err(),
Error::InvalidEntry
);
assert_eq!(
ResponseCacheCodec
.encode(&CacheEntry {
timestamp: Some(f64::NAN),
response: json!({})
})
.unwrap_err(),
Error::InvalidEntry
);
}
#[tokio::test]
async fn invalid_entries_are_misses_and_disabled_reads_do_not_touch_redis() {
let connection = MockRedisConnection::new([MockCmd::new(
redis::cmd("GET").arg("tenant:key"),
Ok(b"invalid".to_vec()),
)])
.assert_all_commands_consumed();
let backend = RedisCache::with_connection(connection, None, ResponseCacheCodec);
let cache = ResponseCache::new(Arc::new(backend));
let mut request = request();
request.controls.no_cache = true;
assert_eq!(cache.lookup(&request, Duration::ZERO).unwrap(), None);
request.controls.no_cache = false;
assert_eq!(
cache.async_lookup(&request, Duration::ZERO).await.unwrap(),
None
);
}
#[test]
fn string_responses_round_trip_through_typed_and_wire_backends() {
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
#[rstest]
#[case::text(json!("hello world"))]
#[case::numeric_text(json!("123"))]
#[case::null_text(json!("null"))]
#[case::array(json!([1, 2]))]
fn non_object_responses_round_trip_through_a_typed_backend(
memory: Memory,
request: ResponseCacheRequest,
#[case] response: Value,
) {
let now = Duration::from_secs(100);
for response in [json!("hello world"), json!("123"), json!("null")] {
cache.store(&request(), response.clone(), now).unwrap();
assert_eq!(
cache.lookup(&request(), now).unwrap(),
Some(response.clone())
);
let wire = ResponseCacheCodec
.encode(&CacheEntry {
timestamp: Some(100.0),
response: response.clone(),
})
.unwrap();
assert_eq!(ResponseCacheCodec.decode(&wire).unwrap().response, response);
}
memory.store(&request, response.clone(), now).unwrap();
assert_eq!(memory.lookup(&request, now).unwrap(), Some(response));
}
#[test]
fn non_object_responses_are_written_as_python_readable_serialized_strings() {
let wire = ResponseCacheCodec
.encode(&CacheEntry {
timestamp: Some(100.0),
response: json!([1, 2]),
})
.unwrap();
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&wire).unwrap(),
json!({"timestamp": 100.0, "response": "[1,2]"})
);
assert_eq!(
ResponseCacheCodec.decode(&wire).unwrap().response,
json!([1, 2])
);
assert_eq!(
ResponseCacheCodec.decode(br#"{"timestamp": 100.0, "response": "not serialized"}"#),
Err(Error::InvalidEntry)
);
}
#[test]
fn response_entries_preserve_the_existing_json_representation() {
let codec = ResponseCacheCodec;
let entry = CacheEntry {
timestamp: Some(123.0),
response: json!({"choices": [{"text": "cached"}]}),
};
let bytes = codec.encode(&entry).unwrap();
assert_eq!(bytes, serde_json::to_vec(&entry).unwrap());
assert_eq!(codec.decode(&bytes).unwrap(), entry);
}
#[test]
fn response_codec_preserves_values_without_timestamps() {
let codec = ResponseCacheCodec;
let raw = json!({"choices": [{"text": "legacy"}]});
let entry = codec.decode(&serde_json::to_vec(&raw).unwrap()).unwrap();
assert_eq!(entry.timestamp, None);
assert_eq!(entry.response, raw);
#[rstest]
fn entries_without_timestamps_are_always_fresh(request: ResponseCacheRequest) {
let backend = Arc::new(InMemoryCache::default());
BaseCache::set_cache(backend.as_ref(), "tenant:key", entry, &Default::default()).unwrap();
BaseCache::set_cache(
backend.as_ref(),
"tenant:key",
CacheEntry {
timestamp: None,
response: json!({"choices": [{"text": "legacy"}]}),
},
&Default::default(),
)
.unwrap();
let cache = ResponseCache::new(backend);
let mut request = request;
request.max_age = Some(Duration::from_secs(1));
assert_eq!(
cache.lookup(&request(), Duration::from_secs(100)).unwrap(),
cache.lookup(&request, Duration::from_secs(100)).unwrap(),
Some(json!({"choices": [{"text": "legacy"}]}))
);
}
#[rstest]
#[tokio::test]
async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses() {
let cache = memory();
let requests = ["hit", "miss", "disabled"].map(|key| {
ResponseCacheRequest::new(CacheKeyInput {
preset: Some(key.into()),
..Default::default()
})
});
cache
.store(&requests[0], json!({"value": 1}), Duration::from_secs(100))
async fn batch_lookup_reports_partial_hits_and_batch_store_populates_misses(
memory: Memory,
#[values(false, true)] asynchronous: bool,
) {
let now = Duration::from_secs(100);
let mut requests = ["hit", "miss", "disabled"].map(keyed).to_vec();
memory
.store(&requests[0], json!({"value": 1}), now)
.unwrap();
let mut requests = requests.to_vec();
requests[2].controls.caching = Some(false);
let partial = cache
.async_lookup_batch(&requests, Duration::from_secs(100))
.await
.unwrap();
let partial = if asynchronous {
memory.async_lookup_batch(&requests, now).await.unwrap()
} else {
memory.lookup_batch(&requests, now).unwrap()
};
assert_eq!(partial.values, vec![Some(json!({"value": 1})), None, None]);
assert_eq!(partial.missing_indices, vec![1, 2]);
cache
memory
.async_store_batch(
vec![
(requests[1].clone(), json!({"value": 2})),
(requests[2].clone(), json!({"value": 3})),
],
Duration::from_secs(100),
now,
)
.await
.unwrap();
assert_eq!(
cache
.lookup(&requests[1], Duration::from_secs(100))
.unwrap(),
memory.lookup(&requests[1], now).unwrap(),
Some(json!({"value": 2}))
);
requests[2].controls.caching = None;
assert_eq!(
cache
.lookup(&requests[2], Duration::from_secs(100))
.unwrap(),
None
);
assert_eq!(memory.lookup(&requests[2], now).unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn deferred_entries_keep_the_time_they_were_produced() {
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
let mut request = request();
async fn batch_lookup_with_no_readable_request_skips_the_backend(
#[values(false, true)] asynchronous: bool,
) {
let cache = redis(Vec::new(), None);
let mut request = keyed("key");
request.controls.no_cache = true;
let requests = [request.clone(), request];
let partial = if asynchronous {
cache
.async_lookup_batch(&requests, Duration::ZERO)
.await
.unwrap()
} else {
cache.lookup_batch(&requests, Duration::ZERO).unwrap()
};
assert_eq!(partial.values, vec![None, None]);
assert_eq!(partial.missing_indices, vec![0, 1]);
}
#[rstest]
#[tokio::test]
async fn deferred_entries_keep_the_time_they_were_produced(
memory: Memory,
mut request: ResponseCacheRequest,
) {
request.max_age = Some(Duration::from_secs(10));
cache
memory
.async_store_entries(vec![(
request.clone(),
json!({"answer": 7}),
@ -485,27 +571,29 @@ async fn deferred_entries_keep_the_time_they_were_produced() {
.unwrap();
assert_eq!(
cache.lookup(&request, Duration::from_secs(110)).unwrap(),
memory.lookup(&request, Duration::from_secs(110)).unwrap(),
Some(json!({"answer": 7}))
);
assert_eq!(
cache.lookup(&request, Duration::from_secs(111)).unwrap(),
memory.lookup(&request, Duration::from_secs(111)).unwrap(),
None
);
}
#[rstest]
#[tokio::test]
async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() {
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time(
memory: Memory,
request: ResponseCacheRequest,
) {
let buffer = WriteBuffer::new(2);
let mut first = request();
let mut first = request;
first.max_age = Some(Duration::from_secs(10));
let mut second = request();
second.key.preset = Some("tenant:other".into());
let second = keyed("tenant:other");
buffer
.async_store(
&cache,
memory.as_ref(),
&first,
json!({"answer": 7}),
Duration::from_secs(100),
@ -513,13 +601,13 @@ async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() {
.await
.unwrap();
assert_eq!(
cache.lookup(&first, Duration::from_secs(100)).unwrap(),
memory.lookup(&first, Duration::from_secs(100)).unwrap(),
None
);
buffer
.async_store(
&cache,
memory.as_ref(),
&second,
json!({"answer": 8}),
Duration::from_secs(200),
@ -527,37 +615,36 @@ async fn write_buffer_flushes_at_its_size_and_keeps_each_produced_time() {
.await
.unwrap();
assert_eq!(
cache.lookup(&first, Duration::from_secs(110)).unwrap(),
memory.lookup(&first, Duration::from_secs(110)).unwrap(),
Some(json!({"answer": 7}))
);
assert_eq!(
cache.lookup(&first, Duration::from_secs(111)).unwrap(),
memory.lookup(&first, Duration::from_secs(111)).unwrap(),
None
);
assert_eq!(
cache.lookup(&second, Duration::from_secs(200)).unwrap(),
memory.lookup(&second, Duration::from_secs(200)).unwrap(),
Some(json!({"answer": 8}))
);
}
#[rstest]
#[tokio::test]
async fn write_buffer_clear_drops_pending_entries() {
let cache = ResponseCache::new(Arc::new(InMemoryCache::default()));
async fn write_buffer_clear_drops_pending_entries(memory: Memory, request: ResponseCacheRequest) {
let buffer = WriteBuffer::new(2);
let mut other = request();
other.key.preset = Some("tenant:other".into());
let other = keyed("tenant:other");
let now = Duration::from_secs(100);
buffer
.async_store(&cache, &request(), json!({"answer": 7}), now)
.async_store(memory.as_ref(), &request, json!({"answer": 7}), now)
.await
.unwrap();
buffer.clear().unwrap();
buffer
.async_store(&cache, &other, json!({"answer": 8}), now)
.async_store(memory.as_ref(), &other, json!({"answer": 8}), now)
.await
.unwrap();
assert_eq!(cache.lookup(&request(), now).unwrap(), None);
assert_eq!(cache.lookup(&other, now).unwrap(), None);
assert_eq!(memory.lookup(&request, now).unwrap(), None);
assert_eq!(memory.lookup(&other, now).unwrap(), None);
}

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

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

8
litellm-rust/crates/cache-s3/Cargo.toml
View file

@ -10,11 +10,17 @@ litellm-cache.workspace = true
litellm-auth-aws.workspace = true
aws-sdk-s3 = { version = "1.146.1", default-features = false, features = ["rustls", "rt-tokio"] }
aws-credential-types = "1.3.0"
aws-smithy-types = "1.6.0"
aws-smithy-runtime-api = { version = "1.16.2", features = ["client", "http-1x"] }
aws-smithy-types = { version = "1.6.0", features = ["http-body-1-x"] }
aws-types = "1.6.0"
futures-util.workspace = true
http.workspace = true
reqwest.workspace = true
tokio.workspace = true
[dev-dependencies]
litellm-cache-testing.workspace = true
rstest.workspace = true
wiremock = "0.6.5"
serde_json.workspace = true
tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }

63
litellm-rust/crates/cache-s3/src/auth.rs
View file

@ -5,22 +5,22 @@ use aws_credential_types::{
use litellm_auth_aws::{AwsAuthConfig, resolve_credentials};
#[derive(Clone)]
pub(crate) struct Credentials {
pub struct S3Credentials {
config: AwsAuthConfig,
env: fn(&str) -> Option<String>,
}
impl Credentials {
pub(crate) fn new(config: AwsAuthConfig) -> Self {
impl S3Credentials {
pub fn new(config: AwsAuthConfig) -> Self {
Self::with_env(config, |name| std::env::var(name).ok())
}
pub(crate) fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
pub fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option<String>) -> Self {
Self { config, env }
}
}
impl ProvideCredentials for Credentials {
impl ProvideCredentials for S3Credentials {
fn provide_credentials<'a>(&'a self) -> future::ProvideCredentials<'a>
where
Self: 'a,
@ -45,57 +45,8 @@ impl ProvideCredentials for Credentials {
}
}
impl std::fmt::Debug for Credentials {
impl std::fmt::Debug for S3Credentials {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Credentials").finish_non_exhaustive()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn explicit_keys_ignore_an_ambient_session_token() {
let provider = Credentials::with_env(
AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
|name| (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()),
);
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.access_key_id(), "key");
assert_eq!(credentials.secret_access_key(), "secret");
assert_eq!(credentials.session_token(), None);
}
#[tokio::test]
async fn explicit_keys_keep_their_session_token() {
let provider = Credentials::new(AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
session_token: Some("t".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
});
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.session_token(), Some("t"));
}
#[tokio::test]
async fn environment_keys_resolve_with_their_session_token() {
let provider = Credentials::with_env(AwsAuthConfig::default(), |name| match name {
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
_ => None,
});
let credentials = provider.provide_credentials().await.unwrap();
assert_eq!(credentials.access_key_id(), "env-key");
assert_eq!(credentials.secret_access_key(), "env-secret");
assert_eq!(credentials.session_token(), Some("env-token"));
f.debug_struct("S3Credentials").finish_non_exhaustive()
}
}

31
litellm-rust/crates/cache-s3/src/cache.rs
View file

@ -10,13 +10,14 @@ use aws_sdk_s3::{
primitives::ByteStream,
};
use aws_smithy_types::{DateTime, date_time::Format};
use futures_util::future::try_join_all;
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::{
BaseCache, BatchCache, CacheCodec, CacheConnectionResult, Error, ExactCacheContext, FlushCache,
BaseCache, BatchCache, CacheCodec, DisconnectCache, Error, ExactCacheContext, FlushCache,
};
use tokio::runtime::Handle;
use crate::auth::Credentials;
use crate::{auth::S3Credentials, transport::ReqwestHttpClient};
pub struct S3Endpoint {
pub url: String,
@ -41,12 +42,13 @@ pub struct S3Cache<C: CacheCodec> {
}
impl<C: CacheCodec> S3Cache<C> {
pub fn new(config: S3CacheConfig, codec: C, runtime: Handle) -> Self {
pub fn new(config: S3CacheConfig, http: reqwest::Client, codec: C, runtime: Handle) -> Self {
let endpoint_url: Option<String> = config.endpoint.map(|endpoint| endpoint.url);
let base = aws_sdk_s3::Config::builder()
.behavior_version(BehaviorVersion::latest())
.region(Region::new(config.region.clone()))
.credentials_provider(Credentials::new(config.auth))
.http_client(ReqwestHttpClient(http))
.credentials_provider(S3Credentials::new(config.auth))
.request_checksum_calculation(RequestChecksumCalculation::WhenRequired)
.response_checksum_validation(ResponseChecksumValidation::WhenRequired);
let builder = match &endpoint_url {
@ -202,13 +204,26 @@ impl<C: CacheCodec> BaseCache for S3Cache<C> {
self.get(key).await
}
async fn async_set_cache_pipeline(
&self,
entries: Vec<(String, Self::Value)>,
context: Self::Context,
) -> Result<(), Error> {
let context = &context;
try_join_all(
entries
.into_iter()
.map(|(key, value)| async move { self.put(&key, value, context).await }),
)
.await
.map(drop)
}
}
impl<C: CacheCodec> DisconnectCache for S3Cache<C> {
async fn disconnect(&self) -> Result<(), Error> {
Ok(())
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
Err(Error::UnsupportedOperation)
}
}
impl<C: CacheCodec> BatchCache for S3Cache<C> {}

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

@ -1,4 +1,6 @@
mod auth;
mod cache;
mod transport;
pub use auth::S3Credentials;
pub use cache::{S3Cache, S3CacheConfig, S3Endpoint};

49
litellm-rust/crates/cache-s3/src/transport.rs Normal file
View file

@ -0,0 +1,49 @@
use aws_smithy_runtime_api::client::{
http::{
HttpClient, HttpConnector, HttpConnectorFuture, HttpConnectorSettings, SharedHttpConnector,
},
orchestrator::HttpRequest,
result::ConnectorError,
runtime_components::RuntimeComponents,
};
use aws_smithy_types::body::SdkBody;
#[derive(Clone, Debug)]
pub(crate) struct ReqwestHttpClient(pub(crate) reqwest::Client);
impl HttpClient for ReqwestHttpClient {
fn http_connector(
&self,
_: &HttpConnectorSettings,
_: &RuntimeComponents,
) -> SharedHttpConnector {
SharedHttpConnector::new(self.clone())
}
}
impl HttpConnector for ReqwestHttpClient {
fn call(&self, request: HttpRequest) -> HttpConnectorFuture {
let client = self.0.clone();
HttpConnectorFuture::new(async move {
let request = request
.try_into_http1x()
.map_err(|error| ConnectorError::other(error.into(), None))?
.map(reqwest::Body::wrap);
let request = reqwest::Request::try_from(request)
.map_err(|error| ConnectorError::other(error.into(), None))?;
let response = client.execute(request).await.map_err(|error| {
if error.is_timeout() {
ConnectorError::timeout(error.into())
} else {
ConnectorError::io(error.into())
}
})?;
let response = http::Response::from(response).map(SdkBody::from_body_1_x);
response
.try_into()
.map_err(|error: aws_smithy_runtime_api::http::HttpError| {
ConnectorError::other(error.into(), None)
})
})
}
}

57
litellm-rust/crates/cache-s3/tests/auth.rs Normal file
View file

@ -0,0 +1,57 @@
use aws_credential_types::provider::ProvideCredentials;
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache_s3::S3Credentials;
use rstest::rstest;
fn explicit(session_token: Option<&str>) -> AwsAuthConfig {
AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
session_token: session_token.map(str::to_string),
region_name: Some("us-east-1".to_string()),
..Default::default()
}
}
fn ambient_token(name: &str) -> Option<String> {
(name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string())
}
fn environment_keys(name: &str) -> Option<String> {
match name {
"AWS_ACCESS_KEY_ID" => Some("env-key".to_string()),
"AWS_SECRET_ACCESS_KEY" => Some("env-secret".to_string()),
"AWS_SESSION_TOKEN" => Some("env-token".to_string()),
_ => None,
}
}
#[rstest]
#[case::explicit_keys_ignore_an_ambient_session_token(
explicit(None), ambient_token, ("key", "secret", None)
)]
#[case::explicit_keys_keep_their_session_token(
explicit(Some("t")), ambient_token, ("key", "secret", Some("t"))
)]
#[case::environment_keys_resolve_with_their_session_token(
AwsAuthConfig::default(), environment_keys, ("env-key", "env-secret", Some("env-token"))
)]
#[tokio::test]
async fn credentials_resolve(
#[case] config: AwsAuthConfig,
#[case] env: fn(&str) -> Option<String>,
#[case] expected: (&str, &str, Option<&str>),
) {
let credentials = S3Credentials::with_env(config, env)
.provide_credentials()
.await
.unwrap();
assert_eq!(
(
credentials.access_key_id(),
credentials.secret_access_key(),
credentials.session_token(),
),
expected
);
}

303
litellm-rust/crates/cache-s3/tests/cache.rs
View file

@ -1,41 +1,23 @@
mod support;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use litellm_auth_aws::AwsAuthConfig;
use aws_smithy_types::{DateTime, date_time::Format};
use litellm_cache::{
BaseCache, BatchCache, BatchEntry, Error, ExactCacheContext, FlushCache, JsonCodec,
BaseCache, BatchCache, BatchEntry, DisconnectCache, Error, ExactCacheContext, FlushCache,
};
use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
use litellm_cache_s3::S3CacheConfig;
use rstest::{fixture, rstest};
use serde_json::{Value, json};
use tokio::runtime::Handle;
use support::FakeBucket;
use wiremock::{
Mock, MockServer, ResponseTemplate,
http::HeaderMap,
matchers::{method, path},
};
fn config(endpoint: String) -> S3CacheConfig {
S3CacheConfig {
bucket: "cache-bucket".to_string(),
key_prefix: "team/".to_string(),
region: "us-east-1".to_string(),
endpoint: Some(S3Endpoint { url: endpoint }),
auth: AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
}
}
fn cache(endpoint: &str) -> S3Cache<JsonCodec<Value>> {
S3Cache::new(
config(endpoint.to_string()),
JsonCodec::<Value>::new(),
Handle::current(),
)
}
async fn mock_server() -> MockServer {
#[fixture]
async fn server() -> MockServer {
let server = MockServer::start().await;
Mock::given(method("PUT"))
.respond_with(ResponseTemplate::new(200).insert_header("etag", "\"etag\""))
@ -44,23 +26,25 @@ async fn mock_server() -> MockServer {
server
}
fn http_date_from(headers: &wiremock::http::HeaderMap, name: &str) -> Option<SystemTime> {
use aws_smithy_types::{DateTime, date_time::Format};
fn http_date_from(headers: &HeaderMap, name: &str) -> Option<SystemTime> {
headers
.get(name)
.and_then(|value| DateTime::from_str(value.to_str().ok()?, Format::HttpDate).ok())
.map(|date| UNIX_EPOCH + Duration::new(date.secs() as u64, date.subsec_nanos()))
}
fn ttl(seconds: u64) -> ExactCacheContext {
ExactCacheContext {
ttl: Some(Duration::from_secs(seconds)),
}
}
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn set_writes_python_metadata_with_and_without_ttl() {
let server = mock_server().await;
let cache = cache(&server.uri());
let context = ExactCacheContext {
ttl: Some(Duration::from_secs(90)),
};
async fn set_writes_python_metadata_with_and_without_ttl(#[future(awt)] server: MockServer) {
let cache = support::cache(&server.uri());
cache
.set_cache("alpha:beta", json!({"answer": 1}), &context)
.set_cache("alpha:beta", json!({"answer": 1}), &ttl(90))
.unwrap();
cache
.set_cache("plain", json!({"answer": 2}), &ExactCacheContext::default())
@ -110,61 +94,84 @@ async fn set_writes_python_metadata_with_and_without_ttl() {
);
}
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_hit_miss_expired_and_invalid_entries() {
let server = mock_server().await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/hit"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 3})))
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/missing"))
.respond_with(
ResponseTemplate::new(404).set_body_string("<Error><Code>NoSuchKey</Code></Error>"),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/denied"))
.respond_with(
ResponseTemplate::new(403).set_body_string("<Error><Code>AccessDenied</Code></Error>"),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/expired"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("expires", "Thu, 01 Jan 1970 00:00:00 GMT")
.set_body_json(json!({"answer": 4})),
)
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/malformed"))
.respond_with(ResponseTemplate::new(200).set_body_string("not a cache entry"))
.mount(&server)
.await;
let cache = cache(&server.uri());
let context = ExactCacheContext::default();
async fn async_set_signs_with_the_configured_keys(#[future(awt)] server: MockServer) {
support::cache(&server.uri())
.async_set_cache("key", json!({"answer": 1}), ttl(3600))
.await
.unwrap();
assert_eq!(
cache.get_cache("hit", &context).unwrap(),
Some(json!({"answer": 3}))
let requests = server.received_requests().await.unwrap();
let request = &requests[0];
assert_eq!(request.url.path(), "/cache-bucket/team/key");
assert!(
request.headers["authorization"]
.to_str()
.unwrap()
.contains("Credential=key/")
);
assert_eq!(cache.get_cache("missing", &context).unwrap(), None);
assert_eq!(cache.get_cache("denied", &context).unwrap(), None);
assert_eq!(cache.get_cache("expired", &context).unwrap(), None);
assert!(request.headers.get("x-amz-security-token").is_none());
assert_eq!(
cache.get_cache("malformed", &context),
Err(Error::InvalidEntry)
request.headers["cache-control"].to_str().unwrap(),
"immutable, max-age=3600, s-maxage=3600"
);
}
#[rstest]
#[case::hit("hit", ResponseTemplate::new(200).set_body_json(json!({"answer": 3})), Ok(Some(json!({"answer": 3}))))]
#[case::no_such_key(
"missing",
ResponseTemplate::new(404).set_body_string("<Error><Code>NoSuchKey</Code></Error>"),
Ok(None)
)]
#[case::access_denied(
"denied",
ResponseTemplate::new(403).set_body_string("<Error><Code>AccessDenied</Code></Error>"),
Ok(None)
)]
#[case::expired(
"expired",
ResponseTemplate::new(200)
.insert_header("expires", "Thu, 01 Jan 1970 00:00:00 GMT")
.set_body_json(json!({"answer": 4})),
Ok(None)
)]
#[case::not_yet_expired(
"fresh",
ResponseTemplate::new(200)
.insert_header("expires", "Fri, 01 Jan 2100 00:00:00 GMT")
.set_body_json(json!({"answer": 5})),
Ok(Some(json!({"answer": 5})))
)]
#[case::malformed(
"malformed",
ResponseTemplate::new(200).set_body_string("not a cache entry"),
Err(Error::InvalidEntry)
)]
#[case::server_error("broken", ResponseTemplate::new(500), Err(Error::Unavailable))]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn batch_get_preserves_order_with_hits_misses_and_invalid() {
let server = mock_server().await;
async fn get_maps_s3_responses(
#[future(awt)] server: MockServer,
#[case] key: &str,
#[case] response: ResponseTemplate,
#[case] expected: Result<Option<Value>, Error>,
) {
Mock::given(method("GET"))
.and(path(format!("/cache-bucket/team/{key}")))
.respond_with(response)
.mount(&server)
.await;
let cache = support::cache(&server.uri());
let context = ExactCacheContext::default();
assert_eq!(cache.get_cache(key, &context), expected);
assert_eq!(cache.async_get_cache(key, &context).await, expected);
}
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn batch_get_preserves_order_with_hits_misses_and_invalid(#[future(awt)] server: MockServer) {
for (key, status, body) in [
("first", 200, "{\"answer\": 1}"),
("invalid", 200, "garbage"),
@ -180,91 +187,117 @@ async fn batch_get_preserves_order_with_hits_misses_and_invalid() {
.respond_with(ResponseTemplate::new(404))
.mount(&server)
.await;
let cache = cache(&server.uri());
let cache = support::cache(&server.uri());
let context = ExactCacheContext::default();
let keys = vec![
"first".to_string(),
"miss".to_string(),
"invalid".to_string(),
];
let expected = vec![
BatchEntry::Hit(json!({"answer": 1})),
BatchEntry::Miss,
BatchEntry::Invalid,
];
let entries = cache.batch_get_cache(&keys, &context).unwrap();
assert_eq!(cache.batch_get_cache(&keys, &context).unwrap(), expected);
assert_eq!(
entries,
vec![
BatchEntry::Hit(json!({"answer": 1})),
BatchEntry::Miss,
BatchEntry::Invalid,
]
cache.async_batch_get_cache(keys, context).await.unwrap(),
expected
);
}
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn unsupported_and_noop_capabilities_match_python() {
let server = mock_server().await;
let cache = cache(&server.uri());
async fn pipeline_writes_every_entry_with_the_shared_ttl() {
let server = FakeBucket::serve().await;
let cache = support::cache(&server.uri());
cache
.async_set_cache_pipeline(
vec![
("one".into(), json!({"n": 1})),
("two".into(), json!({"n": 2})),
],
ttl(30),
)
.await
.unwrap();
let requests = server.received_requests().await.unwrap();
assert_eq!(requests.len(), 2);
assert!(requests.iter().all(|request| {
request.headers["cache-control"].to_str().unwrap() == "immutable, max-age=30, s-maxage=30"
}));
assert_eq!(
cache.test_connection().await,
Err(Error::UnsupportedOperation)
cache
.get_cache("two", &ExactCacheContext::default())
.unwrap(),
Some(json!({"n": 2}))
);
}
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn flush_and_disconnect_are_noops_like_python(#[future(awt)] server: MockServer) {
let cache = support::cache(&server.uri());
cache.flush_cache().unwrap();
cache.async_flush_cache().await.unwrap();
cache.disconnect().await.unwrap();
assert_eq!(cache.get_ttl(&ExactCacheContext::default()), None);
assert_eq!(
cache.get_ttl(&ExactCacheContext {
ttl: Some(Duration::from_secs(45)),
}),
Some(Duration::from_secs(45))
);
assert!(server.received_requests().await.unwrap().is_empty());
}
#[test]
fn key_conversion_prefixes_and_splits_colons() {
#[rstest]
#[case::without_ttl(ExactCacheContext::default(), None)]
#[case::with_ttl(ttl(45), Some(Duration::from_secs(45)))]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn get_ttl_reports_the_request_ttl(
#[case] context: ExactCacheContext,
#[case] expected: Option<Duration>,
) {
assert_eq!(
support::cache("http://localhost").get_ttl(&context),
expected
);
}
#[rstest]
#[case::prefixed("team/", "a:b:c", "team/a/b/c")]
#[case::prefixed_plain("team/", "plain", "team/plain")]
#[case::unprefixed("", "a:b", "a/b")]
fn key_conversion_prefixes_and_splits_colons(
#[case] key_prefix: &str,
#[case] key: &str,
#[case] expected: &str,
) {
let runtime = tokio::runtime::Builder::new_multi_thread()
.worker_threads(1)
.enable_all()
.build()
.unwrap();
let _guard = runtime.enter();
let cache = S3Cache::new(
let cache = support::cache_with(
S3CacheConfig {
key_prefix: "team/".to_string(),
..config("http://localhost".to_string())
key_prefix: key_prefix.to_string(),
..support::config("http://localhost")
},
JsonCodec::<Value>::new(),
runtime.handle().clone(),
);
assert_eq!(cache.bucket(), "cache-bucket");
assert_eq!(cache.key_prefix(), "team/");
assert_eq!(cache.to_s3_key("a:b:c"), "team/a/b/c");
assert_eq!(cache.to_s3_key("plain"), "team/plain");
let unprefixed = S3Cache::new(
S3CacheConfig {
key_prefix: String::new(),
..config("http://localhost".to_string())
},
JsonCodec::<Value>::new(),
runtime.handle().clone(),
);
assert_eq!(unprefixed.to_s3_key("a:b"), "a/b");
assert_eq!(cache.key_prefix(), key_prefix);
assert_eq!(cache.region(), "us-east-1");
assert_eq!(cache.endpoint(), Some("http://localhost"));
assert_eq!(cache.to_s3_key(key), expected);
}
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sync_methods_block_inside_and_outside_the_runtime() {
let server = mock_server().await;
Mock::given(method("GET"))
.and(path("/cache-bucket/team/key"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"answer": 9})))
.mount(&server)
.await;
async fn sync_methods_block_outside_the_runtime() {
let server = FakeBucket::serve().await;
let uri = server.uri();
let cache = tokio::task::spawn_blocking(move || {
let cache = cache(&uri);
let handle = tokio::runtime::Handle::current();
let cached = tokio::task::spawn_blocking(move || {
let cache = support::cache_with(support::config(&uri), handle);
let context = ExactCacheContext::default();
cache
.set_cache("key", json!({"answer": 9}), &context)
@ -274,5 +307,5 @@ async fn sync_methods_block_inside_and_outside_the_runtime() {
.await
.unwrap();
assert_eq!(cache, Some(json!({"answer": 9})));
assert_eq!(cached, Some(json!({"answer": 9})));
}

65
litellm-rust/crates/cache-s3/tests/contract.rs Normal file
View file

@ -0,0 +1,65 @@
mod support;
use litellm_cache::ExactCacheContext;
use litellm_cache_testing as contract;
use rstest::{fixture, rstest};
use serde_json::json;
use support::{FakeBucket, JsonS3Cache};
use wiremock::MockServer;
struct S3 {
cache: JsonS3Cache,
_server: MockServer,
}
#[fixture]
async fn s3() -> S3 {
let server = FakeBucket::serve().await;
S3 {
cache: support::cache(&server.uri()),
_server: server,
}
}
#[fixture]
fn context() -> ExactCacheContext {
ExactCacheContext::default()
}
const PREFIX: &str = "contract:";
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn hit_and_miss(#[future(awt)] s3: S3, context: ExactCacheContext) {
contract::hit_and_miss(&s3.cache, context, PREFIX, json!({"answer": 42})).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn sync_async_equivalence(#[future(awt)] s3: S3, context: ExactCacheContext) {
contract::sync_async_equivalence(&s3.cache, context, PREFIX, json!("first"), json!([2])).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn overwrite_replaces(#[future(awt)] s3: S3, context: ExactCacheContext) {
contract::overwrite_replaces(&s3.cache, context, PREFIX, json!(1), json!({"b": 2})).await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn pipeline_writes_every_entry(#[future(awt)] s3: S3, context: ExactCacheContext) {
contract::pipeline_writes_every_entry(
&s3.cache,
context,
PREFIX,
vec![json!("a"), json!(2), json!({"c": true})],
)
.await;
}
#[rstest]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn batch_preserves_order(#[future(awt)] s3: S3, context: ExactCacheContext) {
contract::batch_preserves_order(&s3.cache, context, PREFIX, json!("first"), json!(2)).await;
}

77
litellm-rust/crates/cache-s3/tests/support/mod.rs Normal file
View file

@ -0,0 +1,77 @@
#![allow(dead_code)]
use std::{
collections::HashMap,
sync::{Arc, Mutex},
};
use litellm_auth_aws::AwsAuthConfig;
use litellm_cache::JsonCodec;
use litellm_cache_s3::{S3Cache, S3CacheConfig, S3Endpoint};
use serde_json::Value;
use tokio::runtime::Handle;
use wiremock::{Mock, MockServer, Request, Respond, ResponseTemplate, http::Method, matchers::any};
pub type JsonS3Cache = S3Cache<JsonCodec<Value>>;
pub fn config(endpoint: &str) -> S3CacheConfig {
S3CacheConfig {
bucket: "cache-bucket".to_string(),
key_prefix: "team/".to_string(),
region: "us-east-1".to_string(),
endpoint: Some(S3Endpoint {
url: endpoint.to_string(),
}),
auth: AwsAuthConfig {
access_key_id: Some("key".to_string()),
secret_access_key: Some("secret".to_string()),
region_name: Some("us-east-1".to_string()),
..Default::default()
},
}
}
pub fn cache_with(config: S3CacheConfig, runtime: Handle) -> JsonS3Cache {
S3Cache::new(config, reqwest::Client::new(), JsonCodec::new(), runtime)
}
pub fn cache(endpoint: &str) -> JsonS3Cache {
cache_with(config(endpoint), Handle::current())
}
/// An in-memory bucket: PUT stores the body under the request path, GET serves it or answers
/// `NoSuchKey`.
#[derive(Clone, Default)]
pub struct FakeBucket {
objects: Arc<Mutex<HashMap<String, Vec<u8>>>>,
}
impl FakeBucket {
pub async fn serve() -> MockServer {
let server = MockServer::start().await;
Mock::given(any())
.respond_with(Self::default())
.mount(&server)
.await;
server
}
}
impl Respond for FakeBucket {
fn respond(&self, request: &Request) -> ResponseTemplate {
let path = request.url.path().to_string();
let mut objects = self.objects.lock().unwrap();
match request.method {
Method::PUT => {
objects.insert(path, request.body.clone());
ResponseTemplate::new(200).insert_header("etag", "\"etag\"")
}
Method::GET => match objects.get(&path) {
Some(body) => ResponseTemplate::new(200).set_body_bytes(body.clone()),
None => ResponseTemplate::new(404)
.set_body_string("<Error><Code>NoSuchKey</Code></Error>"),
},
_ => ResponseTemplate::new(405),
}
}
}

10
litellm-rust/crates/cache-testing/Cargo.toml Normal file
View file

@ -0,0 +1,10 @@
[package]
name = "litellm-cache-testing"
version = "0.1.0"
edition.workspace = true
license.workspace = true
repository.workspace = true
publish = false
[dependencies]
litellm-cache.workspace = true

211
litellm-rust/crates/cache-testing/src/lib.rs Normal file
View file

@ -0,0 +1,211 @@
//! Backend-neutral contract checks every cache backend runs from its own `rstest` suite.
//!
//! Each check takes the cache under test, the context to call it with, a key `prefix` that
//! keeps runs apart on shared servers, and distinct sample values. A check panics with the
//! violated invariant, so a backend test is one `#[rstest]` case per contract.
use std::fmt::Debug;
use litellm_cache::{BaseCache, BatchCache, BatchEntry, CounterCache, DeleteCache, FlushCache};
fn key(prefix: &str, name: &str) -> String {
format!("{prefix}{name}")
}
/// A missing key reads as `None`, and a written key reads back through sync and async gets.
pub async fn hit_and_miss<B>(cache: &B, context: B::Context, prefix: &str, value: B::Value)
where
B: BaseCache,
B::Value: Debug + PartialEq,
{
let key = key(prefix, "hit-and-miss");
assert_eq!(
cache.get_cache(&key, &context).unwrap(),
None,
"unwritten key must miss"
);
assert_eq!(
cache.async_get_cache(&key, &context).await.unwrap(),
None,
"unwritten key must miss asynchronously"
);
cache.set_cache(&key, value.clone(), &context).unwrap();
assert_eq!(
cache.get_cache(&key, &context).unwrap(),
Some(value.clone())
);
assert_eq!(
cache.async_get_cache(&key, &context).await.unwrap(),
Some(value)
);
}
/// Sync and async writes land in the same store: each is visible to the other read path.
pub async fn sync_async_equivalence<B>(
cache: &B,
context: B::Context,
prefix: &str,
first: B::Value,
second: B::Value,
) where
B: BaseCache,
B::Value: Debug + PartialEq,
{
let async_written = key(prefix, "async-written");
let sync_written = key(prefix, "sync-written");
cache
.async_set_cache(&async_written, first.clone(), context.clone())
.await
.unwrap();
assert_eq!(
cache.get_cache(&async_written, &context).unwrap(),
Some(first)
);
cache
.set_cache(&sync_written, second.clone(), &context)
.unwrap();
assert_eq!(
cache
.async_get_cache(&sync_written, &context)
.await
.unwrap(),
Some(second)
);
}
/// A second write to a key replaces the first.
pub async fn overwrite_replaces<B>(
cache: &B,
context: B::Context,
prefix: &str,
first: B::Value,
second: B::Value,
) where
B: BaseCache,
B::Value: Debug + PartialEq,
{
let key = key(prefix, "overwrite");
cache.set_cache(&key, first, &context).unwrap();
cache.set_cache(&key, second.clone(), &context).unwrap();
assert_eq!(cache.get_cache(&key, &context).unwrap(), Some(second));
}
/// `async_set_cache_pipeline` writes every entry, and an empty pipeline succeeds.
pub async fn pipeline_writes_every_entry<B>(
cache: &B,
context: B::Context,
prefix: &str,
values: Vec<B::Value>,
) where
B: BaseCache,
B::Value: Debug + PartialEq,
{
cache
.async_set_cache_pipeline(Vec::new(), context.clone())
.await
.unwrap();
let entries = values
.iter()
.enumerate()
.map(|(index, value)| (key(prefix, &format!("pipeline-{index}")), value.clone()))
.collect::<Vec<_>>();
cache
.async_set_cache_pipeline(entries.clone(), context.clone())
.await
.unwrap();
for (key, value) in entries {
assert_eq!(
cache.get_cache(&key, &context).unwrap(),
Some(value),
"{key}"
);
}
}
/// Batch reads answer in request order, with a `Miss` in place of each absent key.
pub async fn batch_preserves_order<B>(
cache: &B,
context: B::Context,
prefix: &str,
first: B::Value,
second: B::Value,
) where
B: BatchCache,
B::Value: Debug + PartialEq,
{
let keys = vec![
key(prefix, "batch-first"),
key(prefix, "batch-missing"),
key(prefix, "batch-second"),
];
cache.set_cache(&keys[0], first.clone(), &context).unwrap();
cache.set_cache(&keys[2], second.clone(), &context).unwrap();
let expected = vec![
BatchEntry::Hit(first),
BatchEntry::Miss,
BatchEntry::Hit(second),
];
assert_eq!(cache.batch_get_cache(&keys, &context).unwrap(), expected);
assert_eq!(
cache.async_batch_get_cache(keys, context).await.unwrap(),
expected
);
}
/// Sync and async deletes remove only the named key, and deleting a missing key succeeds.
pub async fn delete_removes_key<B>(cache: &B, context: B::Context, prefix: &str, value: B::Value)
where
B: DeleteCache,
B::Value: Debug + PartialEq,
{
let sync_deleted = key(prefix, "delete-sync");
let async_deleted = key(prefix, "delete-async");
let kept = key(prefix, "delete-kept");
for key in [&sync_deleted, &async_deleted, &kept] {
cache.set_cache(key, value.clone(), &context).unwrap();
}
cache.delete_cache(&sync_deleted).unwrap();
cache.async_delete_cache(&async_deleted).await.unwrap();
cache
.delete_cache(&key(prefix, "delete-never-written"))
.unwrap();
assert_eq!(cache.get_cache(&sync_deleted, &context).unwrap(), None);
assert_eq!(cache.get_cache(&async_deleted, &context).unwrap(), None);
assert_eq!(cache.get_cache(&kept, &context).unwrap(), Some(value));
}
/// `flush_cache` and `async_flush_cache` each leave the cache empty.
pub async fn flush_clears<B>(cache: &B, context: B::Context, prefix: &str, value: B::Value)
where
B: FlushCache,
B::Value: Debug + PartialEq,
{
let key = key(prefix, "flush");
cache.set_cache(&key, value.clone(), &context).unwrap();
cache.flush_cache().unwrap();
assert_eq!(cache.get_cache(&key, &context).unwrap(), None);
cache.set_cache(&key, value, &context).unwrap();
cache.async_flush_cache().await.unwrap();
assert_eq!(cache.get_cache(&key, &context).unwrap(), None);
}
/// Sync and async increments accumulate on one counter, starting from zero. Whole-number
/// steps, since Python's disk cache restarts any counter whose stored value is not an `int`.
pub async fn counter_accumulates<B>(cache: &B, context: B::Context, prefix: &str)
where
B: CounterCache,
{
let key = key(prefix, "counter");
assert_eq!(
cache.increment_cache(&key, 1.0, context.clone()).unwrap(),
1.0
);
assert_eq!(
cache
.async_increment(&key, 2.0, context.clone(), false)
.await
.unwrap(),
3.0
);
assert_eq!(cache.increment_cache(&key, -1.0, context).unwrap(), 2.0);
}

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

@ -8,13 +8,13 @@ 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]
litellm-cache-testing.workspace = true
redis-test = "1.0.4"
rstest.workspace = true
serde_json.workspace = true
tokio.workspace = true

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

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

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

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

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

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

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

File diff suppressed because it is too large Load diff

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

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

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

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

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

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

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

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

2
litellm-rust/crates/cache/Cargo.toml vendored
View file

@ -7,7 +7,7 @@ repository.workspace = true
[dependencies]
serde.workspace = true
serde_json.workspace = true
serde_json = { workspace = true, features = ["preserve_order"] }
thiserror.workspace = true
[dev-dependencies]

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

@ -122,36 +122,4 @@ pub trait BaseCache: Send + Sync {
) -> impl Future<Output = Result<(), Error>> + Send {
self.async_set_cache(key, value, context)
}
fn disconnect(&self) -> impl Future<Output = Result<(), Error>> + Send;
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);
}
}

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

@ -55,31 +55,3 @@ impl CacheType {
.find(|cache_type| cache_type.as_python_name() == value)
}
}
#[cfg(test)]
mod tests {
use super::CacheType;
#[test]
fn every_python_cache_type_has_one_round_trip_identity() {
let names = CacheType::ALL.map(CacheType::as_python_name);
assert_eq!(
names,
[
"local",
"redis",
"redis-semantic",
"valkey-semantic",
"s3",
"disk",
"qdrant-semantic",
"azure-blob",
"gcs",
]
);
assert_eq!(
names.map(CacheType::from_python_name),
CacheType::ALL.map(Some)
);
}
}

140
litellm-rust/crates/cache/src/capabilities.rs vendored
View file

@ -1,6 +1,6 @@
use std::{future::Future, time::Duration};
use crate::{BaseCache, BatchEntry, Error};
use crate::{BaseCache, BatchEntry, CacheConnectionResult, CacheContext, Error};
#[derive(Clone, Debug, PartialEq)]
pub struct IncrementOperation {
@ -9,6 +9,35 @@ pub struct IncrementOperation {
pub ttl: Option<Duration>,
}
#[derive(Clone, Debug, PartialEq)]
pub struct PushOperation<V> {
pub key: String,
pub values: Vec<V>,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PopOperation {
pub key: String,
pub count: Option<usize>,
}
/// `disconnect`, for backends whose Python class releases connections or clients.
pub trait DisconnectCache: BaseCache {
fn disconnect(&self) -> impl Future<Output = Result<(), Error>> + Send;
}
/// `test_connection`, for backends whose Python class overrides the base `NotImplementedError`.
pub trait ConnectionCache: BaseCache {
fn test_connection(&self) -> impl Future<Output = Result<CacheConnectionResult, Error>> + Send;
}
/// `sync_ping` and `ping`.
pub trait PingCache: BaseCache {
fn sync_ping(&self) -> Result<bool, Error>;
fn ping(&self) -> impl Future<Output = Result<bool, Error>> + Send;
}
pub trait BatchCache: BaseCache {
fn batch_get_cache(
&self,
@ -45,6 +74,14 @@ pub trait BatchCache: BaseCache {
}
}
/// `async_set_cache_pipeline_with_ttls`: one pipeline where every entry carries its own TTL.
pub trait TtlPipelineCache: BaseCache {
fn async_set_cache_pipeline_with_ttls(
&self,
entries: Vec<(String, Self::Value, Option<Duration>)>,
) -> impl Future<Output = Result<(), Error>> + Send;
}
pub trait DeleteCache: BaseCache {
fn delete_cache(&self, key: &str) -> Result<(), Error>;
@ -53,6 +90,14 @@ pub trait DeleteCache: BaseCache {
}
}
/// `delete_cache_keys`: one round trip that reports how many keys existed.
pub trait BulkDeleteCache: DeleteCache {
fn delete_cache_keys(
&self,
keys: Vec<String>,
) -> impl Future<Output = Result<usize, Error>> + Send;
}
pub trait FlushCache: BaseCache {
fn flush_cache(&self) -> Result<(), Error>;
@ -61,18 +106,79 @@ pub trait FlushCache: BaseCache {
}
}
pub trait CounterCache: BaseCache<Value = f64> {
/// `flushall`: drops every key on the server, ignoring any namespace.
pub trait FlushAllCache: FlushCache {
fn flushall(&self) -> Result<(), Error>;
}
/// Numeric counters. Counters are independent of `Value`, so a response-valued backend can
/// expose them, the way one Python `RedisCache` serves both.
pub trait CounterCache: BaseCache {
fn increment_cache(&self, key: &str, amount: f64, context: Self::Context)
-> Result<f64, Error>;
/// `refresh_ttl` re-arms the TTL on every write instead of only when the key is new;
/// backends without expiring counters ignore it, as Python's `**kwargs` does.
fn async_increment(
&self,
key: &str,
amount: f64,
context: Self::Context,
_refresh_ttl: bool,
) -> impl Future<Output = Result<f64, Error>> + Send {
async move { self.increment_cache(key, amount, context) }
}
/// `async_increment_pipeline`, one result per operation in order. The default increments
/// one key at a time, as the in-memory cache does.
fn async_increment_pipeline(
&self,
operations: Vec<IncrementOperation>,
) -> impl Future<Output = Result<Vec<f64>, Error>> + Send
where
Self::Context: Default,
{
async move {
let mut results = Vec::with_capacity(operations.len());
for operation in operations {
let context = Self::Context::default().with_ttl(operation.ttl);
results.push(
self.async_increment(&operation.key, operation.amount, context, false)
.await?,
);
}
Ok(results)
}
}
}
/// `batch_get_counts` and `async_batch_get_counts`: counter values read in one round trip.
pub trait CountReadCache: CounterCache {
fn batch_get_counts(&self, keys: &[String]) -> Result<Vec<Option<i64>>, Error>;
fn async_batch_get_counts(
&self,
keys: Vec<String>,
) -> impl Future<Output = Result<Vec<Option<i64>>, Error>> + Send;
}
/// `increment_with_floor`, `async_increment_with_floor`, and `async_set_max`.
pub trait BoundedCounterCache: CounterCache {
fn increment_with_floor(&self, key: &str, amount: i64, ttl: Duration) -> Result<i64, Error>;
fn async_increment_with_floor(
&self,
key: &str,
amount: i64,
ttl: Duration,
) -> impl Future<Output = Result<i64, Error>> + Send;
fn async_set_max(
&self,
key: &str,
value: f64,
ttl: Option<Duration>,
) -> impl Future<Output = Result<f64, Error>> + Send;
}
pub trait ClaimCache: BaseCache
@ -105,6 +211,17 @@ pub trait TtlCache: BaseCache {
) -> impl Future<Output = Result<Option<Duration>, Error>> + Send;
}
pub trait RefreshTtlCache: TtlCache {
/// `async_refresh_ttl`: re-arms an existing key without touching its value. `ttl` falls
/// back to the backend default, and the result is `false` when the key is absent or
/// neither TTL is set.
fn async_refresh_ttl(
&self,
key: &str,
ttl: Option<Duration>,
) -> impl Future<Output = Result<bool, Error>> + Send;
}
pub trait SetCache: BaseCache {
type SetValue: Clone + Send + Sync + 'static;
type SetResult: Send + Sync + 'static;
@ -127,11 +244,30 @@ pub trait QueueCache: BaseCache {
values: Vec<Self::QueueValue>,
) -> impl Future<Output = Result<usize, Error>> + Send;
/// `async_rpush_and_trim`: pushes, then keeps only the newest `max_len` entries, atomically.
/// Returns the list length right after the push, before the trim.
fn async_rpush_and_trim(
&self,
key: &str,
values: Vec<Self::QueueValue>,
max_len: usize,
) -> impl Future<Output = Result<usize, Error>> + Send;
fn async_rpush_pipeline(
&self,
operations: Vec<PushOperation<Self::QueueValue>>,
) -> impl Future<Output = Result<Vec<usize>, Error>> + Send;
fn async_lpop(
&self,
key: &str,
count: Option<usize>,
) -> impl Future<Output = Result<Self::PopResult, Error>> + Send;
fn async_lpop_pipeline(
&self,
operations: Vec<PopOperation>,
) -> impl Future<Output = Result<Vec<Self::PopResult>, Error>> + Send;
}
pub trait ScanCache: BaseCache {

151
litellm-rust/crates/cache/src/dual.rs vendored
View file

@ -1,8 +1,8 @@
use std::{sync::Arc, time::Duration};
use crate::{
BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, ClaimCache,
CounterCache, DeleteCache, Error, FlushCache,
BaseCache, BatchCache, BatchEntry, BulkDeleteCache, CacheContext, ClaimCache, CounterCache,
DeleteCache, Error, FlushCache, IncrementOperation, SetCache, TtlCache,
};
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
@ -33,8 +33,12 @@ pub struct DualCache<L1, L2> {
write_policy: WritePolicy,
remote_failure_policy: RemoteFailurePolicy,
promotion_ttl: Option<Duration>,
delete_batch_size: usize,
}
/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`.
pub const DEFAULT_DELETE_BATCH_SIZE: usize = 1000;
impl<L1, L2> DualCache<L1, L2> {
pub fn new(l1: Arc<L1>, l2: Arc<L2>) -> Self {
Self {
@ -44,6 +48,14 @@ impl<L1, L2> DualCache<L1, L2> {
write_policy: WritePolicy::default(),
remote_failure_policy: RemoteFailurePolicy::default(),
promotion_ttl: None,
delete_batch_size: DEFAULT_DELETE_BATCH_SIZE,
}
}
pub fn with_delete_batch_size(self, delete_batch_size: usize) -> Self {
Self {
delete_batch_size,
..self
}
}
@ -217,15 +229,6 @@ where
}
self.l1.async_set_cache_pipeline(entries, context).await
}
async fn disconnect(&self) -> Result<(), Error> {
self.l2.disconnect().await?;
self.l1.disconnect().await
}
async fn test_connection(&self) -> Result<CacheConnectionResult, Error> {
self.l2.test_connection().await
}
}
impl<V, C, L1, L2> BatchCache for DualCache<L1, L2>
@ -320,26 +323,146 @@ where
}
}
impl<C, L1, L2> DualCache<L1, L2>
where
C: CacheContext,
L1: BaseCache<Value = f64, Context = C>,
{
/// Python's `local_only=True` increment: the local tier alone, read then written back.
fn increment_local(&self, key: &str, amount: f64, context: &C) -> Result<f64, Error> {
let value = self.l1.get_cache(key, context)?.unwrap_or(0.0) + amount;
self.l1.set_cache(key, value, context)?;
Ok(value)
}
}
impl<C, L1, L2> CounterCache for DualCache<L1, L2>
where
C: CacheContext,
L1: BaseCache<Value = f64, Context = C>,
L2: CounterCache<Context = C>,
L2: CounterCache<Value = f64, Context = C>,
{
fn increment_cache(&self, key: &str, amount: f64, context: C) -> Result<f64, Error> {
if !self.writes_remote() {
return self.increment_local(key, amount, &context);
}
let value = self.l2.increment_cache(key, amount, context.clone())?;
self.l1.set_cache(key, value, &context)?;
Ok(value)
}
async fn async_increment(&self, key: &str, amount: f64, context: C) -> Result<f64, Error> {
async fn async_increment(
&self,
key: &str,
amount: f64,
context: C,
refresh_ttl: bool,
) -> Result<f64, Error> {
if !self.writes_remote() {
return self.increment_local(key, amount, &context);
}
let value = self
.l2
.async_increment(key, amount, context.clone())
.async_increment(key, amount, context.clone(), refresh_ttl)
.await?;
self.l1.async_set_cache(key, value, context).await?;
Ok(value)
}
/// `async_increment_cache_pipeline`, L2-first like single increments: the local tier takes
/// each remote result.
async fn async_increment_pipeline(
&self,
operations: Vec<IncrementOperation>,
) -> Result<Vec<f64>, Error>
where
C: Default,
{
if !self.writes_remote() {
return operations
.iter()
.map(|operation| {
let context = C::default().with_ttl(operation.ttl);
self.increment_local(&operation.key, operation.amount, &context)
})
.collect();
}
let values = self.l2.async_increment_pipeline(operations.clone()).await?;
if values.len() != operations.len() {
return Err(Error::Unavailable);
}
for (operation, value) in operations.iter().zip(&values) {
self.l1
.async_set_cache(&operation.key, *value, C::default().with_ttl(operation.ttl))
.await?;
}
Ok(values)
}
}
/// `async_set_cache_sadd`: local set first, then the remote one unless writes stay local.
impl<V, C, S, L1, L2> SetCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
S: Clone + Send + Sync + 'static,
L1: SetCache<Value = V, Context = C, SetValue = S>,
L2: SetCache<Value = V, Context = C, SetValue = S>,
{
type SetValue = S;
type SetResult = ();
async fn async_set_cache_sadd(
&self,
key: &str,
values: Vec<S>,
ttl: Option<Duration>,
) -> Result<(), Error> {
self.l1
.async_set_cache_sadd(key, values.clone(), ttl)
.await?;
if self.writes_remote() {
self.l2.async_set_cache_sadd(key, values, ttl).await?;
}
Ok(())
}
}
/// `async_delete_cache_keys`: every key leaves the local tier, then the remote tier in chunks
/// of `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`, since Redis takes a chunk as one command.
impl<V, C, L1, L2> BulkDeleteCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: DeleteCache<Value = V, Context = C>,
L2: BulkDeleteCache<Value = V, Context = C>,
{
async fn delete_cache_keys(&self, keys: Vec<String>) -> Result<usize, Error> {
for key in &keys {
self.l1.delete_cache(key)?;
}
let mut deleted = 0;
for chunk in keys.chunks(self.delete_batch_size.max(1)) {
deleted += self.l2.delete_cache_keys(chunk.to_vec()).await?;
}
Ok(deleted)
}
}
/// `async_get_ttl`: the local TTL, or the remote one when the local tier has none.
impl<V, C, L1, L2> TtlCache for DualCache<L1, L2>
where
V: Clone + Send + Sync + 'static,
C: CacheContext,
L1: TtlCache<Value = V, Context = C>,
L2: TtlCache<Value = V, Context = C>,
{
async fn async_get_ttl(&self, key: &str) -> Result<Option<Duration>, Error> {
match self.l1.async_get_ttl(key).await? {
Some(ttl) => Ok(Some(ttl)),
None => self.l2.async_get_ttl(key).await,
}
}
}
impl<V, C, L1, L2> ClaimCache for DualCache<L1, L2>

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

@ -5,6 +5,7 @@ mod capabilities;
mod codec;
mod dual;
mod error;
pub mod semantic;
pub use base_cache::{
BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext,
@ -13,9 +14,13 @@ pub use base_cache::{
pub use cache_type::CacheType;
pub use caching::{Cache, CacheBackend, get_cache, set_cache};
pub use capabilities::{
BatchCache, CacheScript, ClaimCache, ClientInfoCache, CounterCache, DeleteCache, FlushCache,
IncrementOperation, QueueCache, ScanCache, ScriptCache, SetCache, TtlCache,
BatchCache, BoundedCounterCache, BulkDeleteCache, CacheScript, ClaimCache, ClientInfoCache,
ConnectionCache, CountReadCache, CounterCache, DeleteCache, DisconnectCache, FlushAllCache,
FlushCache, IncrementOperation, PingCache, PopOperation, PushOperation, QueueCache,
RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache, TtlPipelineCache,
};
pub use codec::{CacheCodec, JsonCodec};
pub use dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy};
pub use dual::{
DEFAULT_DELETE_BATCH_SIZE, DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy,
};
pub use error::Error;

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