From 4677f1028e6cea75638ec41d954f5ca4a59d73db Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:13:02 -0700 Subject: [PATCH] refactor(rust): align the cache crates with Python and wire every native backend (#42530) * 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 * 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 * 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 * 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 Co-authored-by: Claude Opus 5 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 34 +- litellm-rust/Cargo.toml | 1 + .../crates/cache-azure-blob/Cargo.toml | 5 + .../crates/cache-azure-blob/src/cache.rs | 48 +- .../cache-azure-blob/src/cache/tests.rs | 746 ----------- .../crates/cache-azure-blob/src/lib.rs | 2 + .../crates/cache-azure-blob/src/tests.rs | 0 .../crates/cache-azure-blob/src/transport.rs | 49 + .../crates/cache-azure-blob/tests/cache.rs | 494 +++++++ .../crates/cache-azure-blob/tests/contract.rs | 81 ++ .../cache-azure-blob/tests/support/mod.rs | 239 ++++ .../cache-azure-blob/tests/transport.rs | 90 ++ litellm-rust/crates/cache-disk/Cargo.toml | 1 + litellm-rust/crates/cache-disk/src/cache.rs | 38 +- litellm-rust/crates/cache-disk/src/sqlite.rs | 12 - litellm-rust/crates/cache-disk/src/store.rs | 1 - litellm-rust/crates/cache-disk/tests/cache.rs | 103 +- .../crates/cache-disk/tests/contract.rs | 82 ++ litellm-rust/crates/cache-gcs/Cargo.toml | 2 + litellm-rust/crates/cache-gcs/src/cache.rs | 36 +- litellm-rust/crates/cache-gcs/tests/cache.rs | 351 +++-- .../crates/cache-gcs/tests/contract.rs | 65 + .../crates/cache-gcs/tests/support/mod.rs | 87 ++ litellm-rust/crates/cache-memory/Cargo.toml | 2 +- litellm-rust/crates/cache-memory/src/cache.rs | 197 +-- .../crates/cache-memory/tests/cache.rs | 687 +++++++--- .../crates/cache-memory/tests/contract.rs | 98 ++ .../crates/cache-qdrant-semantic/Cargo.toml | 3 +- .../src/{semantic.rs => cache.rs} | 105 +- .../cache-qdrant-semantic/src/config.rs | 13 + .../cache-qdrant-semantic/src/embedder.rs | 14 +- .../crates/cache-qdrant-semantic/src/lib.rs | 8 +- .../cache-qdrant-semantic/src/prompt.rs | 59 - .../cache-qdrant-semantic/tests/contract.rs | 91 ++ .../cache-qdrant-semantic/tests/embedder.rs | 91 +- .../cache-qdrant-semantic/tests/prompt.rs | 38 - .../cache-qdrant-semantic/tests/qdrant.rs | 595 +++++---- .../tests/support/mod.rs | 5 +- .../crates/cache-redis-semantic/Cargo.toml | 6 +- .../crates/cache-redis-semantic/src/cache.rs | 568 +++----- .../crates/cache-redis-semantic/src/config.rs | 8 + .../crates/cache-redis-semantic/src/index.rs | 205 +++ .../crates/cache-redis-semantic/src/lib.rs | 8 +- .../crates/cache-redis-semantic/src/prompt.rs | 97 -- .../crates/cache-redis-semantic/src/reply.rs | 57 + .../cache-redis-semantic/tests/cache.rs | 1181 +++++++---------- .../cache-redis-semantic/tests/contract.rs | 63 + .../cache-redis-semantic/tests/support/mod.rs | 299 +++++ litellm-rust/crates/cache-redis/Cargo.toml | 2 + litellm-rust/crates/cache-redis/src/cache.rs | 671 ++-------- .../cache-redis/src/cache/operations.rs | 632 --------- litellm-rust/crates/cache-redis/src/claim.rs | 105 ++ .../cache-redis/src/{cache => }/connection.rs | 215 ++- .../crates/cache-redis/src/counter.rs | 205 +++ litellm-rust/crates/cache-redis/src/keys.rs | 63 + litellm-rust/crates/cache-redis/src/lib.rs | 18 +- .../crates/cache-redis/src/lifecycle.rs | 85 ++ litellm-rust/crates/cache-redis/src/queue.rs | 226 ++++ litellm-rust/crates/cache-redis/src/script.rs | 136 ++ litellm-rust/crates/cache-redis/src/store.rs | 232 ++++ .../crates/cache-redis/tests/cache.rs | 1099 ++++++++++----- .../crates/cache-redis/tests/cluster.rs | 347 +++-- .../crates/cache-redis/tests/contract.rs | 95 ++ .../crates/cache-redis/tests/support/mod.rs | 231 ++++ litellm-rust/crates/cache-response/Cargo.toml | 1 + litellm-rust/crates/cache-response/README.md | 36 +- .../crates/cache-response/src/exact.rs | 22 +- litellm-rust/crates/cache-response/src/lib.rs | 2 +- .../crates/cache-response/src/response.rs | 61 +- .../crates/cache-response/tests/caching.rs | 236 ++-- .../crates/cache-response/tests/codec.rs | 110 ++ .../crates/cache-response/tests/connection.rs | 123 ++ .../crates/cache-response/tests/response.rs | 659 +++++---- .../cache-response/tests/support/mod.rs | 40 + litellm-rust/crates/cache-s3/Cargo.toml | 8 +- litellm-rust/crates/cache-s3/src/auth.rs | 63 +- litellm-rust/crates/cache-s3/src/cache.rs | 31 +- litellm-rust/crates/cache-s3/src/lib.rs | 2 + litellm-rust/crates/cache-s3/src/transport.rs | 49 + litellm-rust/crates/cache-s3/tests/auth.rs | 57 + litellm-rust/crates/cache-s3/tests/cache.rs | 303 +++-- .../crates/cache-s3/tests/contract.rs | 65 + .../crates/cache-s3/tests/support/mod.rs | 77 ++ litellm-rust/crates/cache-testing/Cargo.toml | 10 + litellm-rust/crates/cache-testing/src/lib.rs | 211 +++ .../crates/cache-valkey-semantic/Cargo.toml | 6 +- .../crates/cache-valkey-semantic/src/cache.rs | 245 ++++ .../cache-valkey-semantic/src/config.rs | 8 + .../crates/cache-valkey-semantic/src/index.rs | 100 ++ .../crates/cache-valkey-semantic/src/lib.rs | 1158 +--------------- .../cache-valkey-semantic/src/search.rs | 163 +++ .../cache-valkey-semantic/tests/cache.rs | 417 ++++++ .../cache-valkey-semantic/tests/contract.rs | 62 + .../tests/support/mod.rs | 349 +++++ litellm-rust/crates/cache/Cargo.toml | 2 +- litellm-rust/crates/cache/src/base_cache.rs | 32 - litellm-rust/crates/cache/src/cache_type.rs | 28 - litellm-rust/crates/cache/src/capabilities.rs | 140 +- litellm-rust/crates/cache/src/dual.rs | 151 ++- litellm-rust/crates/cache/src/lib.rs | 11 +- litellm-rust/crates/cache/src/semantic.rs | 274 ++++ litellm-rust/crates/cache/tests/cache_type.rs | 56 + litellm-rust/crates/cache/tests/caching.rs | 248 +++- litellm-rust/crates/cache/tests/codec.rs | 15 +- litellm-rust/crates/cache/tests/dual.rs | 521 +++++++- litellm-rust/crates/cache/tests/semantic.rs | 270 ++++ .../python-bridge/src/cache/activation.rs | 111 ++ .../crates/python-bridge/src/cache/binding.rs | 60 +- .../crates/python-bridge/src/cache/config.rs | 756 +++++++---- .../python-bridge/src/cache/embedder.rs | 23 +- .../crates/python-bridge/src/cache/facade.rs | 8 +- .../crates/python-bridge/src/cache/handle.rs | 19 +- .../crates/python-bridge/src/cache/mod.rs | 11 + .../crates/python-bridge/src/cache/native.rs | 183 ++- .../python-bridge/src/cache/semantic.rs | 34 +- litellm/caching/caching.py | 89 +- litellm/rust_bridge/_native.pyi | 2 + litellm/rust_bridge/response_cache.py | 11 + .../messages/test_callbacks.py | 10 + tests/test_litellm_rust/test_cache.py | 410 +++++- tests/test_litellm_rust/test_fork_guard.py | 9 +- .../test_valkey_semantic_cache_native.py | 24 +- 122 files changed, 12665 insertions(+), 6944 deletions(-) delete mode 100644 litellm-rust/crates/cache-azure-blob/src/cache/tests.rs delete mode 100644 litellm-rust/crates/cache-azure-blob/src/tests.rs create mode 100644 litellm-rust/crates/cache-azure-blob/src/transport.rs create mode 100644 litellm-rust/crates/cache-azure-blob/tests/cache.rs create mode 100644 litellm-rust/crates/cache-azure-blob/tests/contract.rs create mode 100644 litellm-rust/crates/cache-azure-blob/tests/support/mod.rs create mode 100644 litellm-rust/crates/cache-azure-blob/tests/transport.rs create mode 100644 litellm-rust/crates/cache-disk/tests/contract.rs create mode 100644 litellm-rust/crates/cache-gcs/tests/contract.rs create mode 100644 litellm-rust/crates/cache-gcs/tests/support/mod.rs create mode 100644 litellm-rust/crates/cache-memory/tests/contract.rs rename litellm-rust/crates/cache-qdrant-semantic/src/{semantic.rs => cache.rs} (72%) create mode 100644 litellm-rust/crates/cache-qdrant-semantic/src/config.rs delete mode 100644 litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs create mode 100644 litellm-rust/crates/cache-qdrant-semantic/tests/contract.rs delete mode 100644 litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/src/config.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/src/index.rs delete mode 100644 litellm-rust/crates/cache-redis-semantic/src/prompt.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/src/reply.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/tests/contract.rs create mode 100644 litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs delete mode 100644 litellm-rust/crates/cache-redis/src/cache/operations.rs create mode 100644 litellm-rust/crates/cache-redis/src/claim.rs rename litellm-rust/crates/cache-redis/src/{cache => }/connection.rs (65%) create mode 100644 litellm-rust/crates/cache-redis/src/counter.rs create mode 100644 litellm-rust/crates/cache-redis/src/keys.rs create mode 100644 litellm-rust/crates/cache-redis/src/lifecycle.rs create mode 100644 litellm-rust/crates/cache-redis/src/queue.rs create mode 100644 litellm-rust/crates/cache-redis/src/script.rs create mode 100644 litellm-rust/crates/cache-redis/src/store.rs create mode 100644 litellm-rust/crates/cache-redis/tests/contract.rs create mode 100644 litellm-rust/crates/cache-redis/tests/support/mod.rs create mode 100644 litellm-rust/crates/cache-response/tests/codec.rs create mode 100644 litellm-rust/crates/cache-response/tests/connection.rs create mode 100644 litellm-rust/crates/cache-response/tests/support/mod.rs create mode 100644 litellm-rust/crates/cache-s3/src/transport.rs create mode 100644 litellm-rust/crates/cache-s3/tests/auth.rs create mode 100644 litellm-rust/crates/cache-s3/tests/contract.rs create mode 100644 litellm-rust/crates/cache-s3/tests/support/mod.rs create mode 100644 litellm-rust/crates/cache-testing/Cargo.toml create mode 100644 litellm-rust/crates/cache-testing/src/lib.rs create mode 100644 litellm-rust/crates/cache-valkey-semantic/src/cache.rs create mode 100644 litellm-rust/crates/cache-valkey-semantic/src/config.rs create mode 100644 litellm-rust/crates/cache-valkey-semantic/src/index.rs create mode 100644 litellm-rust/crates/cache-valkey-semantic/src/search.rs create mode 100644 litellm-rust/crates/cache-valkey-semantic/tests/cache.rs create mode 100644 litellm-rust/crates/cache-valkey-semantic/tests/contract.rs create mode 100644 litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs create mode 100644 litellm-rust/crates/cache/src/semantic.rs create mode 100644 litellm-rust/crates/cache/tests/cache_type.rs create mode 100644 litellm-rust/crates/cache/tests/semantic.rs create mode 100644 litellm-rust/crates/python-bridge/src/cache/activation.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index ea63746d56d..d9425fc6bd7 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -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", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 813d0713128..65813a35214 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -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" } diff --git a/litellm-rust/crates/cache-azure-blob/Cargo.toml b/litellm-rust/crates/cache-azure-blob/Cargo.toml index 55abaff1975..baa1b0f5482 100644 --- a/litellm-rust/crates/cache-azure-blob/Cargo.toml +++ b/litellm-rust/crates/cache-azure-blob/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/cache-azure-blob/src/cache.rs b/litellm-rust/crates/cache-azure-blob/src/cache.rs index 6a872a0d6e6..489b08d485e 100644 --- a/litellm-rust/crates/cache-azure-blob/src/cache.rs +++ b/litellm-rust/crates/cache-azure-blob/src/cache.rs @@ -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 { container: BlobContainerClient, @@ -28,9 +27,11 @@ pub struct AzureBlobCache { } impl AzureBlobCache { + /// `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 { @@ -38,7 +39,10 @@ impl AzureBlobCache { 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 AzureBlobCache { } fn block_on(&self, future: impl Future) -> 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 BaseCache for AzureBlobCache { .await .map(drop) } - - async fn disconnect(&self) -> Result<(), Error> { - Ok(()) - } - - async fn test_connection(&self) -> Result { - 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 BatchCache for AzureBlobCache {} @@ -250,5 +239,10 @@ impl FlushCache for AzureBlobCache { } } -#[cfg(test)] -mod tests; +impl DisconnectCache for AzureBlobCache { + /// 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(()) + } +} diff --git a/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs b/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs deleted file mode 100644 index f8736ab069b..00000000000 --- a/litellm-rust/crates/cache-azure-blob/src/cache/tests.rs +++ /dev/null @@ -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, -} - -#[derive(Default)] -struct FakeState { - container_exists: bool, - blobs: BTreeMap>, - requests: Vec, - failing: bool, - precondition_conflicts: bool, -} - -#[derive(Clone, Default)] -struct FakeBlobService { - state: Arc>, -} - -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> { - self.state.lock().unwrap().blobs.get(name).cloned() - } - - fn blob_names(&self) -> Vec { - 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 { - 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) -> 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 { - let mut xml = String::from( - r#""#, - ); - for name in state.blobs.keys() { - xml.push_str(&format!( - "{name}BlockBlob" - )); - } - xml.push_str(""); - xml.into_bytes() - } -} - -#[async_trait::async_trait] -impl HttpClient for FakeBlobService { - async fn execute_request(&self, request: &Request) -> azure_core::Result { - 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>, -} - -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, 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> { - 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"))) - ); -} diff --git a/litellm-rust/crates/cache-azure-blob/src/lib.rs b/litellm-rust/crates/cache-azure-blob/src/lib.rs index 5ae752c111d..6bcfe130576 100644 --- a/litellm-rust/crates/cache-azure-blob/src/lib.rs +++ b/litellm-rust/crates/cache-azure-blob/src/lib.rs @@ -1,5 +1,7 @@ mod cache; mod credential; +mod transport; pub use cache::AzureBlobCache; pub use credential::AzureBlobCredential; +pub use transport::ReqwestTransport; diff --git a/litellm-rust/crates/cache-azure-blob/src/tests.rs b/litellm-rust/crates/cache-azure-blob/src/tests.rs deleted file mode 100644 index e69de29bb2d..00000000000 diff --git a/litellm-rust/crates/cache-azure-blob/src/transport.rs b/litellm-rust/crates/cache-azure-blob/src/transport.rs new file mode 100644 index 00000000000..ed038b8d69d --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/src/transport.rs @@ -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 { + 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))) + } +} diff --git a/litellm-rust/crates/cache-azure-blob/tests/cache.rs b/litellm-rust/crates/cache-azure-blob/tests/cache.rs new file mode 100644 index 00000000000..b9f50b824b3 --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/tests/cache.rs @@ -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; + +#[fixture] +fn fixture() -> Fixture { + Fixture::new(FakeBlobService::default(), ResponseCacheCodec) +} + +fn response_cache(fixture: &Fixture) -> ResponseCache> { + 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) { + 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))) + ); +} diff --git a/litellm-rust/crates/cache-azure-blob/tests/contract.rs b/litellm-rust/crates/cache-azure-blob/tests/contract.rs new file mode 100644 index 00000000000..585ca26bbd7 --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/tests/contract.rs @@ -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> { + 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>, + 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>, + 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>, + 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>, + 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>, + context: ExactCacheContext, +) { + contract::flush_clears(&azure, context, PREFIX, json!("value")).await; +} diff --git a/litellm-rust/crates/cache-azure-blob/tests/support/mod.rs b/litellm-rust/crates/cache-azure-blob/tests/support/mod.rs new file mode 100644 index 00000000000..f908151c22e --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/tests/support/mod.rs @@ -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, +} + +#[derive(Default)] +struct FakeState { + container_exists: bool, + blobs: BTreeMap>, + requests: Vec, + failing: bool, + precondition_conflicts: bool, +} + +#[derive(Clone, Default)] +pub struct FakeBlobService { + state: Arc>, +} + +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> { + self.state.lock().unwrap().blobs.get(name).cloned() + } + + pub fn blob_names(&self) -> Vec { + 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 { + 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) -> 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 { + let mut xml = String::from( + r#""#, + ); + for name in state.blobs.keys() { + xml.push_str(&format!( + "{name}BlockBlob" + )); + } + xml.push_str(""); + xml.into_bytes() + } +} + +#[async_trait::async_trait] +impl HttpClient for FakeBlobService { + async fn execute_request(&self, request: &Request) -> azure_core::Result { + 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( + service: &FakeBlobService, + account_url: &str, + codec: C, + handle: Handle, +) -> Result, 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 { + pub runtime: Runtime, + pub service: FakeBlobService, + pub cache: Arc>, +} + +impl Fixture { + 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() + } +} diff --git a/litellm-rust/crates/cache-azure-blob/tests/transport.rs b/litellm-rust/crates/cache-azure-blob/tests/transport.rs new file mode 100644 index 00000000000..cd1e10aa3d8 --- /dev/null +++ b/litellm-rust/crates/cache-azure-blob/tests/transport.rs @@ -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> { + 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, +) { + 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 + ); +} diff --git a/litellm-rust/crates/cache-disk/Cargo.toml b/litellm-rust/crates/cache-disk/Cargo.toml index b96994b3b55..5cb75f7f129 100644 --- a/litellm-rust/crates/cache-disk/Cargo.toml +++ b/litellm-rust/crates/cache-disk/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/cache-disk/src/cache.rs b/litellm-rust/crates/cache-disk/src/cache.rs index 8e1223309b4..b8a2c6ed1e4 100644 --- a/litellm-rust/crates/cache-disk/src/cache.rs +++ b/litellm-rust/crates/cache-disk/src/cache.rs @@ -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 BaseCache for DiskCache Result<(), Error> { - Ok(()) - } - - async fn test_connection(&self) -> Result { - 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 BatchCache for DiskCache { @@ -241,9 +218,13 @@ impl FlushCache for DiskCache, D: DiskStore, A: ValueAdapter> CounterCache - for DiskCache -{ +impl DisconnectCache for DiskCache { + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } +} + +impl CounterCache for DiskCache { fn increment_cache( &self, key: &str, @@ -264,6 +245,7 @@ impl, D: DiskStore, A: ValueAdapter> CounterCache key: &str, amount: f64, context: ExactCacheContext, + _refresh_ttl: bool, ) -> Result { let key = key.to_string(); let adapter = Arc::clone(&self.adapter); diff --git a/litellm-rust/crates/cache-disk/src/sqlite.rs b/litellm-rust/crates/cache-disk/src/sqlite.rs index 9a36f8af6ad..24c2be2e1c5 100644 --- a/litellm-rust/crates/cache-disk/src/sqlite.rs +++ b/litellm-rust/crates/cache-disk/src/sqlite.rs @@ -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 { diff --git a/litellm-rust/crates/cache-disk/src/store.rs b/litellm-rust/crates/cache-disk/src/store.rs index ed167317cf0..b5c12003cee 100644 --- a/litellm-rust/crates/cache-disk/src/store.rs +++ b/litellm-rust/crates/cache-disk/src/store.rs @@ -29,5 +29,4 @@ pub trait DiskStore: Send + Sync + 'static { now: f64, apply: &mut dyn FnMut(Option) -> Result<(StoredValue, Option), Error>, ) -> Result<(), Error>; - fn probe(&self) -> Result<(), Error>; } diff --git a/litellm-rust/crates/cache-disk/tests/cache.rs b/litellm-rust/crates/cache-disk/tests/cache.rs index dd1f2b1f04e..8c817a8a829 100644 --- a/litellm-rust/crates/cache-disk/tests/cache.rs +++ b/litellm-rust/crates/cache-disk/tests/cache.rs @@ -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::(); 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>, + 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, + #[case] amount: f64, + #[case] expected: f64, +) { + let cache = sandbox.cache::(); + 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::(); + 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()); +} diff --git a/litellm-rust/crates/cache-disk/tests/contract.rs b/litellm-rust/crates/cache-disk/tests/contract.rs new file mode 100644 index 00000000000..257252cb90c --- /dev/null +++ b/litellm-rust/crates/cache-disk/tests/contract.rs @@ -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>, + _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; +} diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml index 4ec60bcfa3b..da0acf554f9 100644 --- a/litellm-rust/crates/cache-gcs/Cargo.toml +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/cache-gcs/src/cache.rs b/litellm-rust/crates/cache-gcs/src/cache.rs index 65282ac99d5..a8a7fbc9a7b 100644 --- a/litellm-rust/crates/cache-gcs/src/cache.rs +++ b/litellm-rust/crates/cache-gcs/src/cache.rs @@ -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 { } impl GcsCache { - pub fn new(config: GcsConfig, codec: S) -> Result { + 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, - ) -> Result { - 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 GcsCache { F: Future> + 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 BaseCache for GcsCache { .await .map(|_| ()) } +} +impl DisconnectCache for GcsCache { async fn disconnect(&self) -> Result<(), Error> { Ok(()) } - - async fn test_connection(&self) -> Result { - Err(Error::UnsupportedOperation) - } } impl BatchCache for GcsCache { diff --git a/litellm-rust/crates/cache-gcs/tests/cache.rs b/litellm-rust/crates/cache-gcs/tests/cache.rs index 45eecf01cec..cdce6a00bdd 100644 --- a/litellm-rust/crates/cache-gcs/tests/cache.rs +++ b/litellm-rust/crates/cache-gcs/tests/cache.rs @@ -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> { - 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, 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::::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, 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> + Send + '_>> - { + fn bearer_token(&self) -> Pin> + 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::::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 ); diff --git a/litellm-rust/crates/cache-gcs/tests/contract.rs b/litellm-rust/crates/cache-gcs/tests/contract.rs new file mode 100644 index 00000000000..5841d0db642 --- /dev/null +++ b/litellm-rust/crates/cache-gcs/tests/contract.rs @@ -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; +} diff --git a/litellm-rust/crates/cache-gcs/tests/support/mod.rs b/litellm-rust/crates/cache-gcs/tests/support/mod.rs new file mode 100644 index 00000000000..6097f0ee1bd --- /dev/null +++ b/litellm-rust/crates/cache-gcs/tests/support/mod.rs @@ -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>; + +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, +) -> 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>>>, +} + +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), + } + } +} diff --git a/litellm-rust/crates/cache-memory/Cargo.toml b/litellm-rust/crates/cache-memory/Cargo.toml index 86ab01564c8..88124f5401e 100644 --- a/litellm-rust/crates/cache-memory/Cargo.toml +++ b/litellm-rust/crates/cache-memory/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/cache-memory/src/cache.rs b/litellm-rust/crates/cache-memory/src/cache.rs index 85850c1d925..74476559cbe 100644 --- a/litellm-rust/crates/cache-memory/src/cache.rs +++ b/litellm-rust/crates/cache-memory/src/cache.rs @@ -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 InMemoryCache { 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 InMemoryCache { 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, Error> { @@ -121,6 +111,70 @@ impl InMemoryCache { 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 { + 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 { + 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 { + 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 { + Ok(self + .state + .lock() + .map_err(|_| Error::Unavailable)? + .values + .len()) + } + + pub fn is_empty(&self) -> Result { + Ok(self.len()? == 0) + } + + /// Entries in the expiration heap, stale ones included; bounded by eviction. + pub fn expiration_heap_len(&self) -> Result { + 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 InMemoryCache { Ok(()) } - fn evict(state: &mut CacheState, 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, 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 InMemoryCache { 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 InMemoryCache { 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, key: &str, now: Duration) -> Option { + 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, + key: String, + value: V, + ttl: Option, + now: Duration, + ) -> Result { + 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 ClaimCache for InMemoryCache @@ -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 { } 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 { - pub async fn async_increment_pipeline( - &self, - operations: Vec, - ) -> Result, Error> { - operations - .into_iter() - .map(|operation| { - self.increment_cache( - &operation.key, - operation.amount, - ExactCacheContext { ttl: operation.ttl }, - ) - }) - .collect() - } -} - impl BaseCache for InMemoryCache { type Value = V; type Context = ExactCacheContext; @@ -315,18 +379,12 @@ impl BaseCache for InMemoryCache { fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, Error> { self.get_cache(key) } +} +impl DisconnectCache for InMemoryCache { async fn disconnect(&self) -> Result<(), Error> { Ok(()) } - - async fn test_connection(&self) -> Result { - Ok(CacheConnectionResult { - status: CacheConnectionStatus::Success, - message: "In-memory cache connection test successful".into(), - error: None, - }) - } } impl BatchCache for InMemoryCache {} @@ -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::::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); - } -} diff --git a/litellm-rust/crates/cache-memory/tests/cache.rs b/litellm-rust/crates/cache-memory/tests/cache.rs index 0df0319b990..ce7d6ac8a98 100644 --- a/litellm-rust/crates/cache-memory/tests/cache.rs +++ b/litellm-rust/crates/cache-memory/tests/cache.rs @@ -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; + #[fixture] -fn clock() -> Arc { +fn clock() -> Clock { Arc::new(AtomicU64::new(100)) } -fn cache(clock: Arc, capacity: usize) -> InMemoryCache { +fn cache_with(clock: &Clock, capacity: usize) -> InMemoryCache { + 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 { + cache_with(clock, capacity) +} + +fn at(clock: &Clock, seconds: u64) { + clock.store(seconds, Ordering::SeqCst); +} + +fn secs(seconds: u64) -> Option { + Some(Duration::from_secs(seconds)) +} + +fn ttl(seconds: u64) -> ExactCacheContext { + ExactCacheContext { ttl: secs(seconds) } +} + +fn measured(capacity: usize) -> InMemoryCache { + 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) { - 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) { - 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, + #[case] expected: Option, +) { + let cache = InMemoryCache::::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) { - 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::::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, +) { + assert_eq!(measured(2).check_value_size(&value.to_string()), expected); +} + +#[rstest] +fn values_are_unbounded_without_a_measure() { + let cache = InMemoryCache::::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::::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> = Arc::new(cache(clock.clone(), 4)); +async fn generic_consumers_share_typed_values_and_honor_expiration(clock: Clock) { + let cache: CacheBackend> = 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, +) { + 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::::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::::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) { - 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::::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::>(); + 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::(&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::(&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::::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::(&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::::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::>::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::::new() + ); assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None); } +#[rstest] #[tokio::test] -async fn increment_pipeline_preserves_operation_order() { - let cache = InMemoryCache::::new(Some(3), None); +async fn increment_pipeline_preserves_operation_order(clock: Clock) { + let cache = cache_with::(&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::::new() + ); } +#[rstest] #[tokio::test] -async fn set_capability_preserves_python_result_and_deduplicates_storage() { - let cache = InMemoryCache::>::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::>(&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::>::with_clock_and_size_measurement( + Some(4), + None, + Some(2), + Some(Arc::new(|value: &HashSet| 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()])) diff --git a/litellm-rust/crates/cache-memory/tests/contract.rs b/litellm-rust/crates/cache-memory/tests/contract.rs new file mode 100644 index 00000000000..860de1ed798 --- /dev/null +++ b/litellm-rust/crates/cache-memory/tests/contract.rs @@ -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 { + InMemoryCache::new(Some(16), None) +} + +#[fixture] +fn counters() -> InMemoryCache { + 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, context: ExactCacheContext) { + contract::hit_and_miss(&strings, context, "memory:", "value".into()).await; +} + +#[rstest] +#[tokio::test] +async fn sync_async_equivalence(strings: InMemoryCache, context: ExactCacheContext) { + contract::sync_async_equivalence( + &strings, + context, + "memory:", + "first".into(), + "second".into(), + ) + .await; +} + +#[rstest] +#[tokio::test] +async fn overwrite_replaces(strings: InMemoryCache, context: ExactCacheContext) { + contract::overwrite_replaces( + &strings, + context, + "memory:", + "first".into(), + "second".into(), + ) + .await; +} + +#[rstest] +#[tokio::test] +async fn pipeline_writes_every_entry(strings: InMemoryCache, 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, context: ExactCacheContext) { + contract::batch_preserves_order( + &strings, + context, + "memory:", + "first".into(), + "second".into(), + ) + .await; +} + +#[rstest] +#[tokio::test] +async fn delete_removes_key(strings: InMemoryCache, context: ExactCacheContext) { + contract::delete_removes_key(&strings, context, "memory:", "value".into()).await; +} + +#[rstest] +#[tokio::test] +async fn flush_clears(strings: InMemoryCache, context: ExactCacheContext) { + contract::flush_clears(&strings, context, "memory:", "value".into()).await; +} + +#[rstest] +#[tokio::test] +async fn counter_accumulates(counters: InMemoryCache, context: ExactCacheContext) { + contract::counter_accumulates(&counters, context, "memory:").await; +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml index 09d6a9637f3..950c2db7491 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml @@ -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" diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs b/litellm-rust/crates/cache-qdrant-semantic/src/cache.rs similarity index 72% rename from litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs rename to litellm-rust/crates/cache-qdrant-semantic/src/cache.rs index fb165ed5a8e..a140e0af174 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/cache.rs @@ -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, 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 { client: Qdrant, @@ -100,14 +82,9 @@ impl QdrantSemanticCache { &self.embedder } + /// Python reads `kwargs["messages"]` unguarded, so a request without messages fails. fn prompt(context: &SemanticCacheContext) -> Result { - 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 QdrantSemanticCache { 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 QdrantSemanticCache { &self, key: &str, context: &SemanticCacheContext, - ) -> Result, Error> { + ) -> 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 result = self .client .search_points( @@ -171,20 +154,27 @@ impl QdrantSemanticCache { .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 = 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 BaseCache for QdrantSemanticCache { } fn get_cache(&self, key: &str, context: &Self::Context) -> Result, 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 BaseCache for QdrantSemanticCache { key: &str, context: &Self::Context, ) -> Result, 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 BaseCache for QdrantSemanticCache { .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 SemanticCache for QdrantSemanticCache { + fn get_cache_with_similarity( + &self, + key: &str, + context: &Self::Context, + ) -> Result, Error> { + self.runtime.block_on(self.get(key, context)) } - async fn test_connection(&self) -> Result { - Err(Error::UnsupportedOperation) + async fn async_get_cache_with_similarity( + &self, + key: &str, + context: &Self::Context, + ) -> Result, 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 { + 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, } } diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/config.rs b/litellm-rust/crates/cache-qdrant-semantic/src/config.rs new file mode 100644 index 00000000000..2654c155f23 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/src/config.rs @@ -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, +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs index 47b898d6f4e..340393600f2 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs @@ -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, 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, Error> { let request = self .client .post(format!("{}/embeddings", self.api_base)) diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/lib.rs b/litellm-rust/crates/cache-qdrant-semantic/src/lib.rs index 0f346a9155b..3017bc695e2 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/lib.rs @@ -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}; diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs b/litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs deleted file mode 100644 index ef1a2306658..00000000000 --- a/litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs +++ /dev/null @@ -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() -} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/contract.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/contract.rs new file mode 100644 index 00000000000..896b799decf --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/contract.rs @@ -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>; + +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(check: F) +where + F: FnOnce(Cache) -> Fut + Send + 'static, + Fut: Future, +{ + 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; +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs index 6b09448fde8..de0fab0a66f 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs @@ -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) -> 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, + #[case] expected: Result, 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")); diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs deleted file mode 100644 index 38cd9e2f908..00000000000 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs +++ /dev/null @@ -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"}"# - ); -} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs index c7522c0b313..fe25c503a36 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -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)>>>; +type Cache = QdrantSemanticCache>; + +/// Embeds known prompts, fails on anything else, and records every call. #[derive(Clone)] struct FixedEmbedder { vectors: Arc>>, + 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, Error> { + async fn async_embed( + &self, + input: &str, + metadata: Option<&JsonValue>, + ) -> Result, 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)>, -) -> QdrantSemanticCache { +) -> 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::::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, + 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), 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, + #[case] expected: Result, 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::::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(); +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs index 9a556ae7df5..695fceeac44 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs @@ -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; diff --git a/litellm-rust/crates/cache-redis-semantic/Cargo.toml b/litellm-rust/crates/cache-redis-semantic/Cargo.toml index 9a8755a189e..fe5a317cec0 100644 --- a/litellm-rust/crates/cache-redis-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-redis-semantic/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/cache-redis-semantic/src/cache.rs b/litellm-rust/crates/cache-redis-semantic/src/cache.rs index e0ac31f3630..02feb3cbdf4 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/cache.rs @@ -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, Error>; - - fn async_embed( - &self, - prompt: &str, - metadata: Option<&Value>, - ) -> impl Future, 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, - 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 { - 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 { - 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, prompt: &str, vector: &[f32], ttl: Option, ) -> 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, Error> { - let index = self.ensure_index(connection, vector.len())?; + ) -> Result>, 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::(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 { - connections: Arc>, - embedder: E, - inner: Arc, -} - -impl RedisSemanticCache { - pub fn new(url: &str, embedder: E, config: RedisSemanticConfig) -> Result { - 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 RedisSemanticCache { - 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 { + connections: Arc>, + embedder: E, + codec: S, + inner: Arc, +} + +impl RedisSemanticCache { + pub fn new( + url: &str, + embedder: E, + codec: S, + config: RedisSemanticConfig, + ) -> Result { + Ok(Self { + connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), + embedder, + codec, + inner: Arc::new(Inner::new(config, timestamp)), + }) + } +} + +impl RedisSemanticCache +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 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 RedisSemanticCache< fn tag<'a>(key: &'a str, context: &'a SemanticCacheContext) -> &'a str { context.scope.as_deref().unwrap_or(key) } + + fn decode(&self, lookup: SemanticLookup>) -> Result, Error> { + Ok(SemanticLookup { + value: lookup + .value + .map(|bytes| self.codec.decode(&bytes)) + .transpose()?, + similarity: lookup.similarity, + }) + } } -impl BaseCache - for RedisSemanticCache +impl BaseCache for RedisSemanticCache +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 { @@ -274,22 +240,18 @@ impl 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, 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 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 BaseCache key: &str, context: &Self::Context, ) -> Result, 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 SemanticCache for RedisSemanticCache +where + E: Embedder, + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn get_cache_with_similarity( + &self, + key: &str, + context: &Self::Context, + ) -> Result, 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, 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 { - match Connections::run_blocking(Arc::clone(&self.connections), |connection| { - Ok(match redis::cmd("PING").query::(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 { } 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, Error> { - let info = match redis::cmd("FT.INFO") - .arg(name) - .query::(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::>(); - 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 { - 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 { - 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 { - field_value(fields, name).and_then(string_value) -} - -fn number_field(fields: &[redis::Value], name: &str) -> Option { - field_value(fields, name).and_then(number_value) -} - -fn bytes_field(fields: &[redis::Value], name: &str) -> Option> { - 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 { diff --git a/litellm-rust/crates/cache-redis-semantic/src/config.rs b/litellm-rust/crates/cache-redis-semantic/src/config.rs new file mode 100644 index 00000000000..6b810628be7 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/config.rs @@ -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, +} diff --git a/litellm-rust/crates/cache-redis-semantic/src/index.rs b/litellm-rust/crates/cache-redis-semantic/src/index.rs new file mode 100644 index 00000000000..c141e47003a --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/index.rs @@ -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 `_isolated`, recreated when that one is stale too. +pub(crate) struct Index { + name: String, + resolved: OnceLock, +} + +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 { + 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 { + 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, Error> { + let info = match redis::cmd("FT.INFO") + .arg(name) + .query::(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, + field_type: Option, + dim: Option, + data_type: Option, + distance_metric: Option, +} + +fn attribute(value: &redis::Value) -> Option { + 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::>(); + 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")) + }) +} diff --git a/litellm-rust/crates/cache-redis-semantic/src/lib.rs b/litellm-rust/crates/cache-redis-semantic/src/lib.rs index 51d0b4ba5f3..251323df6cf 100644 --- a/litellm-rust/crates/cache-redis-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-redis-semantic/src/lib.rs @@ -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}; diff --git a/litellm-rust/crates/cache-redis-semantic/src/prompt.rs b/litellm-rust/crates/cache-redis-semantic/src/prompt.rs deleted file mode 100644 index b9c38e98d77..00000000000 --- a/litellm-rust/crates/cache-redis-semantic/src/prompt.rs +++ /dev/null @@ -1,97 +0,0 @@ -use litellm_cache::SemanticCacheContext; -use serde_json::Value; - -pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { - 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) { - 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; - } - } - } - } - _ => {} - } -} diff --git a/litellm-rust/crates/cache-redis-semantic/src/reply.rs b/litellm-rust/crates/cache-redis-semantic/src/reply.rs new file mode 100644 index 00000000000..24cbd573cde --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/src/reply.rs @@ -0,0 +1,57 @@ +pub(crate) fn string_value(value: &redis::Value) -> Option { + 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 { + 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 { + field_value(fields, name).and_then(string_value) +} + +pub(crate) fn number_field(fields: &[redis::Value], name: &str) -> Option { + field_value(fields, name).and_then(number_value) +} + +pub(crate) fn bytes_field(fields: &[redis::Value], name: &str) -> Option> { + match field_value(fields, name)? { + redis::Value::BulkString(bytes) => Some(bytes.clone()), + redis::Value::SimpleString(text) => Some(text.clone().into_bytes()), + _ => None, + } +} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs index 233b87ec52f..fa4ac00e767 100644 --- a/litellm-rust/crates/cache-redis-semantic/tests/cache.rs +++ b/litellm-rust/crates/cache-redis-semantic/tests/cache.rs @@ -1,55 +1,26 @@ -use std::{ - collections::HashMap, - sync::{Arc, Mutex}, - time::Duration, -}; +mod support; -use litellm_cache::{BaseCache, CacheCodec, Error, SemanticCacheContext}; -use litellm_cache_redis_semantic::{Embedder, RedisSemanticCache, RedisSemanticConfig}; -use litellm_cache_response::{CacheEntry, ResponseCacheCodec}; +use std::time::Duration; + +use litellm_cache::{ + BaseCache, Error, JsonCodec, SemanticCacheContext, + semantic::{SemanticCache, SemanticLookup}, +}; +use litellm_cache_redis_semantic::{DEFAULT_INDEX_NAME, RedisSemanticCache, RedisSemanticConfig}; use redis_test::{MockCmd, MockRedisConnection}; +use rstest::{fixture, rstest}; use serde_json::{Value, json}; use sha2::{Digest, Sha256}; +use support::FakeEmbedder; -const INDEX: &str = "litellm_semantic_cache_index"; +const INDEX: &str = DEFAULT_INDEX_NAME; +const PROMPT: &str = "hello prompt"; +const CLOCK: fn() -> f64 = || 1700000000.5; +const VECTOR: [f32; 3] = [0.1, 0.2, 0.3]; -struct FakeEmbedder { - vectors: HashMap>, - calls: Arc>>, -} - -impl FakeEmbedder { - fn new(vectors: &[(&str, &[f32])]) -> (Self, Arc>>) { - let calls = Arc::new(Mutex::new(Vec::new())); - ( - Self { - vectors: vectors - .iter() - .map(|(prompt, vector)| (prompt.to_string(), vector.to_vec())) - .collect(), - calls: Arc::clone(&calls), - }, - calls, - ) - } -} - -impl Embedder for FakeEmbedder { - fn embed(&self, prompt: &str, _: Option<&Value>) -> Result, Error> { - self.calls.lock().unwrap().push(prompt.to_string()); - - 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, Error> { - self.embed(prompt, metadata) - } -} +type MockCache = RedisSemanticCache, MockRedisConnection>; +#[fixture] fn config() -> RedisSemanticConfig { RedisSemanticConfig { index_name: INDEX.into(), @@ -57,6 +28,16 @@ fn config() -> RedisSemanticConfig { } } +#[fixture] +fn entry() -> Value { + json!({"timestamp": 1.0, "response": {"answer": "yes"}}) +} + +#[fixture] +fn context() -> SemanticCacheContext { + messages_context(vec![json!({"role": "user", "content": PROMPT})]) +} + fn messages_context(messages: Vec) -> SemanticCacheContext { SemanticCacheContext { messages: Some(Value::Array(messages)), @@ -64,15 +45,18 @@ fn messages_context(messages: Vec) -> SemanticCacheContext { } } -fn entry() -> CacheEntry { - CacheEntry { - timestamp: Some(1.0), - response: json!({"answer": "yes"}), - } +fn cache(commands: Vec, embedder: FakeEmbedder) -> MockCache { + RedisSemanticCache::with_connection( + MockRedisConnection::new(commands).assert_all_commands_consumed(), + embedder, + JsonCodec::new(), + config(), + ) + .with_clock(CLOCK) } -fn encoded(entry: &CacheEntry) -> Vec { - ResponseCacheCodec.encode(entry).unwrap() +fn encoded(value: &Value) -> Vec { + serde_json::to_vec(value).unwrap() } fn vector_bytes(vector: &[f32]) -> Vec { @@ -98,6 +82,17 @@ fn unknown_index_error() -> redis::RedisError { redis::RedisError::from((redis::ErrorKind::Extension, "Unknown index name")) } +fn info_missing(index: &str) -> MockCmd { + MockCmd::new( + redis::cmd("FT.INFO").arg(index), + Err::(unknown_index_error()), + ) +} + +fn info(index: &str, value: redis::Value) -> MockCmd { + MockCmd::new(redis::cmd("FT.INFO").arg(index), Ok(value)) +} + fn attribute(name: &str, field_type: &str, extra: Vec) -> redis::Value { let mut parts = vec![ s("identifier"), @@ -137,10 +132,6 @@ fn vector_attribute_with(dims: i64, data_type: &str, distance_metric: &str) -> r ) } -fn vector_attribute(dims: i64) -> redis::Value { - vector_attribute_with(dims, "FLOAT32", "COSINE") -} - fn info_with_vector(vector: redis::Value) -> redis::Value { index_info(vec![ attribute("prompt", "TEXT", vec![]), @@ -153,7 +144,7 @@ fn info_with_vector(vector: redis::Value) -> redis::Value { } fn compatible_info(dims: i64) -> redis::Value { - info_with_vector(vector_attribute(dims)) + info_with_vector(vector_attribute_with(dims, "FLOAT32", "COSINE")) } fn unscoped_info(dims: i64) -> redis::Value { @@ -162,10 +153,14 @@ fn unscoped_info(dims: i64) -> redis::Value { attribute("response", "TEXT", vec![]), attribute("inserted_at", "NUMERIC", vec![]), attribute("updated_at", "NUMERIC", vec![]), - vector_attribute(dims), + vector_attribute_with(dims, "FLOAT32", "COSINE"), ]) } +fn create_index(name: &str, dims: usize) -> MockCmd { + MockCmd::new(create_index_command(name, dims), Ok("OK")) +} + fn create_index_command(name: &str, dims: usize) -> redis::Cmd { let mut command = redis::cmd("FT.CREATE"); command @@ -207,7 +202,34 @@ fn create_index_command(name: &str, dims: usize) -> redis::Cmd { command } -fn search_command(index: &str, tag: &str, vector: &[f32]) -> redis::Cmd { +fn hset(index: &str, prompt: &str, tag: &str, vector: &[f32], value: &Value) -> MockCmd { + MockCmd::new( + redis::cmd("HSET") + .arg(format!("{index}:{}", entry_id(prompt, tag))) + .arg("entry_id") + .arg(entry_id(prompt, tag)) + .arg("prompt") + .arg(prompt) + .arg("response") + .arg(encoded(value)) + .arg("prompt_vector") + .arg(vector_bytes(vector)) + .arg("inserted_at") + .arg("1700000000.5") + .arg("updated_at") + .arg("1700000000.5") + .arg("litellm_cache_key") + .arg(tag), + Ok(7), + ) +} + +fn search( + index: &str, + tag: &str, + vector: &[f32], + reply: redis::RedisResult, +) -> MockCmd { let mut command = redis::cmd("FT.SEARCH"); command .arg(index) @@ -236,33 +258,29 @@ fn search_command(index: &str, tag: &str, vector: &[f32]) -> redis::Cmd { .arg(2) .arg("vector") .arg(vector_bytes(vector)); - command + MockCmd::new(command, reply) } -fn hit_fields(tag: &str, distance: &str, response: Vec) -> redis::Value { - redis::Value::Array(vec![ - s("entry_id"), - s("stored-id"), - s("prompt"), - s("hello prompt"), - s("response"), - redis::Value::BulkString(response), - s("inserted_at"), - s("1700000000.5"), - s("updated_at"), - s("1700000000.5"), - s("litellm_cache_key"), - s(tag), - s("vector_distance"), - s(distance), - ]) -} - -fn search_result(fields: redis::Value) -> redis::Value { +fn hit(tag: &str, distance: &str, response: Vec) -> redis::Value { redis::Value::Array(vec![ redis::Value::Int(1), s("litellm_semantic_cache_index:stored-id"), - fields, + redis::Value::Array(vec![ + s("entry_id"), + s("stored-id"), + s("prompt"), + s(PROMPT), + s("response"), + redis::Value::BulkString(response), + s("inserted_at"), + s("1700000000.5"), + s("updated_at"), + s("1700000000.5"), + s("litellm_cache_key"), + s(tag), + s("vector_distance"), + s(distance), + ]), ]) } @@ -270,685 +288,452 @@ fn empty_result() -> redis::Value { redis::Value::Array(vec![redis::Value::Int(0)]) } -#[test] -fn store_creates_index_and_writes_hash_with_expire() { - let vector = vec![0.1f32, 0.2, 0.3]; - let prompt = "hello prompt"; - let tag = "key1"; - let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); - let value = entry(); - let connection = MockRedisConnection::new([ - MockCmd::new( - redis::cmd("FT.INFO").arg(INDEX), - Err::(unknown_index_error()), - ), - MockCmd::new(create_index_command(INDEX, 3), Ok("OK")), - MockCmd::new( - redis::cmd("HSET") - .arg(&hash_key) - .arg("entry_id") - .arg(entry_id(prompt, tag)) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(encoded(&value)) - .arg("prompt_vector") - .arg(vector_bytes(&vector)) - .arg("inserted_at") - .arg("1700000000.5") - .arg("updated_at") - .arg("1700000000.5") - .arg("litellm_cache_key") - .arg(tag), - Ok(7), - ), - MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(5), Ok(1)), - ]) - .assert_all_commands_consumed(); - let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()) - .with_clock(|| 1700000000.5); - - let context = SemanticCacheContext { - ttl: Some(Duration::from_secs(5)), - ..messages_context(vec![json!({"role": "user", "content": prompt})]) +#[rstest] +#[case::creates_index_and_expires(false, Some(Duration::from_secs(5)), Some(5))] +#[case::existing_index_without_ttl(true, None, None)] +#[case::fractional_ttl_rounds_up(true, Some(Duration::from_millis(1500)), Some(2))] +fn store_writes_the_redisvl_hash( + #[case] index_exists: bool, + #[case] ttl: Option, + #[case] expire: Option, + entry: Value, + context: SemanticCacheContext, +) { + let hash_key = format!("{INDEX}:{}", entry_id(PROMPT, "key1")); + let mut commands = if index_exists { + vec![info(INDEX, compatible_info(3))] + } else { + vec![info_missing(INDEX), create_index(INDEX, 3)] }; - cache.set_cache(tag, value, &context).unwrap(); -} - -#[test] -fn store_without_ttl_skips_expire() { - let prompt = "hello prompt"; - let connection = MockRedisConnection::new([ - MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), - MockCmd::new( - redis::cmd("HSET") - .arg(format!("{INDEX}:{}", entry_id(prompt, "key1"))) - .arg("entry_id") - .arg(entry_id(prompt, "key1")) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(encoded(&entry())) - .arg("prompt_vector") - .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) - .arg("inserted_at") - .arg("1700000000.5") - .arg("updated_at") - .arg("1700000000.5") - .arg("litellm_cache_key") - .arg("key1"), - Ok(7), - ), - ]) - .assert_all_commands_consumed(); - let (embedder, _) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()) - .with_clock(|| 1700000000.5); + commands.push(hset(INDEX, PROMPT, "key1", &VECTOR, &entry)); + commands.extend( + expire.map(|seconds| MockCmd::new(redis::cmd("EXPIRE").arg(&hash_key).arg(seconds), Ok(1))), + ); + let cache = cache(commands, FakeEmbedder::new(&[])); cache - .set_cache( - "key1", - entry(), - &messages_context(vec![json!({"role": "user", "content": prompt})]), - ) + .set_cache("key1", entry, &SemanticCacheContext { ttl, ..context }) .unwrap(); } -#[test] -fn lookup_returns_hit_below_distance_threshold() { - let vector = vec![0.1f32, 0.2, 0.3]; - let value = entry(); - let connection = MockRedisConnection::new([ - MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), - MockCmd::new( - search_command(INDEX, "key1", &vector), - Ok(search_result(hit_fields("key1", "0.05", encoded(&value)))), - ), - ]) - .assert_all_commands_consumed(); - let (embedder, _) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()); - - let hit = cache - .get_cache( - "key1", - &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]), - ) - .unwrap(); - assert_eq!(hit, Some(value)); -} - -#[test] -fn lookup_misses_above_distance_threshold_and_on_tag_mismatch() { - let vector = vec![0.1f32, 0.2, 0.3]; - let connection = MockRedisConnection::new([ - MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), - MockCmd::new( - search_command(INDEX, "key1", &vector), - Ok(search_result(hit_fields("key1", "0.5", encoded(&entry())))), - ), - MockCmd::new( - search_command(INDEX, "key1", &vector), - Ok(search_result(hit_fields( - "other", - "0.05", - encoded(&entry()), - ))), - ), - ]) - .assert_all_commands_consumed(); - let (embedder, _) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()); - let context = messages_context(vec![json!({"role": "user", "content": "hello prompt"})]); - - assert_eq!(cache.get_cache("key1", &context).unwrap(), None); - assert_eq!(cache.get_cache("key1", &context).unwrap(), None); -} - -#[test] -fn lookup_returns_invalid_entry_on_malformed_response() { - let vector = vec![0.1f32, 0.2, 0.3]; - let connection = MockRedisConnection::new([ - MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), - MockCmd::new( - search_command(INDEX, "key1", &vector), - Ok(search_result(hit_fields( +#[rstest] +#[case::below_distance_threshold("key1", "0.05", None, Ok(Some(entry())))] +#[case::above_distance_threshold("key1", "0.5", None, Ok(None))] +#[case::other_cache_key("other", "0.05", None, Ok(None))] +#[case::malformed_response("key1", "0.05", Some(b"not json!".as_slice()), Err(Error::InvalidEntry))] +fn lookup_applies_threshold_scope_and_codec( + #[case] stored_tag: &str, + #[case] distance: &str, + #[case] response: Option<&[u8]>, + #[case] expected: Result, Error>, + entry: Value, + context: SemanticCacheContext, +) { + let response = response.map_or_else(|| encoded(&entry), <[u8]>::to_vec); + let cache = cache( + vec![ + info(INDEX, compatible_info(3)), + search( + INDEX, "key1", - "0.05", - b"not json!".to_vec(), - ))), - ), - ]) - .assert_all_commands_consumed(); - let (embedder, _) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()); - - assert_eq!( - cache - .get_cache( - "key1", - &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) - ) - .unwrap_err(), - Error::InvalidEntry + &VECTOR, + Ok(hit(stored_tag, distance, response)), + ), + ], + FakeEmbedder::new(&[]), ); + + assert_eq!(cache.get_cache("key1", &context), expected); } -#[test] -fn missing_prompt_is_noop_and_never_embeds() { - let connection = MockRedisConnection::new(Vec::::new()).assert_all_commands_consumed(); - let (embedder, calls) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()); +#[rstest] +#[case::hit(context(), Some(hit("key1", "0.05", encoded(&entry()))), Some(entry()), Some(1.0 - 0.05))] +#[case::beyond_distance_threshold(context(), Some(hit("key1", "0.5", encoded(&entry()))), None, Some(0.0))] +#[case::no_results(context(), Some(empty_result()), None, Some(0.0))] +#[case::other_cache_key(context(), Some(hit("other", "0.05", encoded(&entry()))), None, Some(0.0))] +#[case::no_prompt(SemanticCacheContext::default(), None, None, Some(0.0))] +#[tokio::test] +async fn lookup_reports_python_semantic_similarity( + #[case] context: SemanticCacheContext, + #[case] reply: Option, + #[case] value: Option, + #[case] similarity: Option, + #[values(false, true)] use_async: bool, +) { + let commands = reply.map_or_else(Vec::new, |reply| { + vec![ + info(INDEX, compatible_info(3)), + search(INDEX, "key1", &VECTOR, Ok(reply)), + ] + }); + let cache = cache(commands, FakeEmbedder::new(&[])); + let lookup = if use_async { + cache + .async_get_cache_with_similarity("key1", &context) + .await + } else { + cache.get_cache_with_similarity("key1", &context) + }; + + assert_eq!(lookup, Ok(SemanticLookup { value, similarity })); +} + +#[rstest] +#[tokio::test] +async fn missing_prompt_is_a_noop_that_never_embeds(entry: Value) { + let embedder = FakeEmbedder::new(&[]); + let calls = embedder.calls.clone(); + let cache = cache(Vec::new(), embedder); let context = SemanticCacheContext::default(); - cache.set_cache("key1", entry(), &context).unwrap(); + + cache.set_cache("key1", entry.clone(), &context).unwrap(); assert_eq!(cache.get_cache("key1", &context).unwrap(), None); + cache + .async_set_cache("key1", entry, context.clone()) + .await + .unwrap(); + assert_eq!(cache.async_get_cache("key1", &context).await.unwrap(), None); assert!(calls.lock().unwrap().is_empty()); } -#[test] -fn scope_overrides_key_as_filter_tag() { - let vector = vec![0.1f32, 0.2, 0.3]; - let prompt = "hello prompt"; - let value = entry(); - let connection = MockRedisConnection::new([ - MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), - MockCmd::new( - redis::cmd("HSET") - .arg(format!("{INDEX}:{}", entry_id(prompt, "scope-a"))) - .arg("entry_id") - .arg(entry_id(prompt, "scope-a")) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(encoded(&value)) - .arg("prompt_vector") - .arg(vector_bytes(&vector)) - .arg("inserted_at") - .arg("1700000000.5") - .arg("updated_at") - .arg("1700000000.5") - .arg("litellm_cache_key") - .arg("scope-a"), - Ok(7), - ), - MockCmd::new( - search_command(INDEX, "scope\\-a", &vector), - Ok(search_result(hit_fields( - "scope-a", - "0.05", - encoded(&value), - ))), - ), - ]) - .assert_all_commands_consumed(); - let (embedder, _) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()) - .with_clock(|| 1700000000.5); +#[rstest] +fn scope_overrides_key_as_filter_tag(entry: Value, context: SemanticCacheContext) { + let cache = cache( + vec![ + info(INDEX, compatible_info(3)), + hset(INDEX, PROMPT, "scope-a", &VECTOR, &entry), + search( + INDEX, + "scope\\-a", + &VECTOR, + Ok(hit("scope-a", "0.05", encoded(&entry))), + ), + ], + FakeEmbedder::new(&[]), + ); let context = SemanticCacheContext { scope: Some("scope-a".into()), - ..messages_context(vec![json!({"role": "user", "content": prompt})]) + ..context }; - cache.set_cache("key1", value.clone(), &context).unwrap(); - assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(value)); + cache.set_cache("key1", entry.clone(), &context).unwrap(); + assert_eq!(cache.get_cache("key1", &context).unwrap(), Some(entry)); } -#[test] -fn incompatible_schema_falls_back_to_isolated_index() { - let prompt = "hello prompt"; - let tag = "key1"; +#[rstest] +#[case::unscoped_schema(unscoped_info(3))] +#[case::wrong_distance_metric(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2")))] +#[case::wrong_data_type(info_with_vector(vector_attribute_with(3, "FLOAT64", "COSINE")))] +fn incompatible_schema_falls_back_to_isolated_index( + #[case] base_info: redis::Value, + entry: Value, + context: SemanticCacheContext, +) { let isolated = format!("{INDEX}_isolated"); - let connection = MockRedisConnection::new([ - MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(unscoped_info(3))), - MockCmd::new( - redis::cmd("FT.INFO").arg(&isolated), - Err::(unknown_index_error()), - ), - MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), - MockCmd::new( - redis::cmd("HSET") - .arg(format!("{isolated}:{}", entry_id(prompt, tag))) - .arg("entry_id") - .arg(entry_id(prompt, tag)) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(encoded(&entry())) - .arg("prompt_vector") - .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) - .arg("inserted_at") - .arg("1700000000.5") - .arg("updated_at") - .arg("1700000000.5") - .arg("litellm_cache_key") - .arg(tag), - Ok(7), - ), - ]) - .assert_all_commands_consumed(); - let (embedder, _) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()) - .with_clock(|| 1700000000.5); - - cache - .set_cache( - tag, - entry(), - &messages_context(vec![json!({"role": "user", "content": prompt})]), - ) - .unwrap(); -} - -#[test] -fn create_index_race_rechecks_schema_and_stores() { - let prompt = "hello prompt"; - let tag = "key1"; - let connection = MockRedisConnection::new([ - MockCmd::new( - redis::cmd("FT.INFO").arg(INDEX), - Err::(unknown_index_error()), - ), - MockCmd::new( - create_index_command(INDEX, 3), - Err::<&str, _>(redis::RedisError::from(( - redis::ErrorKind::Extension, - "Index already exists", - ))), - ), - MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), - MockCmd::new( - redis::cmd("HSET") - .arg(format!("{INDEX}:{}", entry_id(prompt, tag))) - .arg("entry_id") - .arg(entry_id(prompt, tag)) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(encoded(&entry())) - .arg("prompt_vector") - .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) - .arg("inserted_at") - .arg("1700000000.5") - .arg("updated_at") - .arg("1700000000.5") - .arg("litellm_cache_key") - .arg(tag), - Ok(7), - ), - ]) - .assert_all_commands_consumed(); - let (embedder, _) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()) - .with_clock(|| 1700000000.5); - - cache - .set_cache( - tag, - entry(), - &messages_context(vec![json!({"role": "user", "content": prompt})]), - ) - .unwrap(); -} - -#[test] -fn wrong_distance_metric_falls_back_to_isolated_index() { - let prompt = "hello prompt"; - let tag = "key1"; - let isolated = format!("{INDEX}_isolated"); - let connection = MockRedisConnection::new([ - MockCmd::new( - redis::cmd("FT.INFO").arg(INDEX), - Ok(info_with_vector(vector_attribute_with(3, "FLOAT32", "L2"))), - ), - MockCmd::new( - redis::cmd("FT.INFO").arg(&isolated), - Err::(unknown_index_error()), - ), - MockCmd::new(create_index_command(&isolated, 3), Ok("OK")), - MockCmd::new( - redis::cmd("HSET") - .arg(format!("{isolated}:{}", entry_id(prompt, tag))) - .arg("entry_id") - .arg(entry_id(prompt, tag)) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(encoded(&entry())) - .arg("prompt_vector") - .arg(vector_bytes(&[0.1f32, 0.2, 0.3])) - .arg("inserted_at") - .arg("1700000000.5") - .arg("updated_at") - .arg("1700000000.5") - .arg("litellm_cache_key") - .arg(tag), - Ok(7), - ), - ]) - .assert_all_commands_consumed(); - let (embedder, _) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()) - .with_clock(|| 1700000000.5); - - cache - .set_cache( - tag, - entry(), - &messages_context(vec![json!({"role": "user", "content": prompt})]), - ) - .unwrap(); -} - -#[test] -fn tag_special_characters_are_escaped_in_search_filter() { - let vector = vec![0.1f32, 0.2, 0.3]; - let tag = "a:b, c|d"; - let connection = MockRedisConnection::new([ - MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), - MockCmd::new( - search_command(INDEX, "a\\:b\\,\\ c\\|d", &vector), - Ok(empty_result()), - ), - ]) - .assert_all_commands_consumed(); - let (embedder, _) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()); - - assert_eq!( - cache - .get_cache( - tag, - &messages_context(vec![json!({"role": "user", "content": "hello prompt"})]) - ) - .unwrap(), - None - ); -} - -#[test] -fn prompt_extraction_matches_python_message_and_input_shapes() { - let vector = vec![0.1f32, 0.2, 0.3]; - let lookups = 5; - let mut commands = vec![MockCmd::new( - redis::cmd("FT.INFO").arg(INDEX), - Ok(compatible_info(3)), - )]; - for _ in 0..lookups { - commands.push(MockCmd::new( - search_command(INDEX, "key1", &vector), - Ok(empty_result()), - )); - } - let connection = MockRedisConnection::new(commands).assert_all_commands_consumed(); - let (embedder, calls) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()); - - cache - .get_cache( - "key1", - &messages_context(vec![ - json!({"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]}), - json!({"role": "assistant", "content": "reply"}), - ]), - ) - .unwrap(); - cache - .get_cache( - "key1", - &SemanticCacheContext { - input: Some(json!(" plain input ")), - ..Default::default() - }, - ) - .unwrap(); - cache - .get_cache( - "key1", - &SemanticCacheContext { - input: Some( - json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"]), - ), - ..Default::default() - }, - ) - .unwrap(); - cache - .get_cache( - "key1", - &SemanticCacheContext { - input: Some(json!({"output_text": " result text "})), - ..Default::default() - }, - ) - .unwrap(); - cache - .get_cache( - "key1", - &messages_context(vec![json!({ - "role": "user", - "content": "question", - "search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}], - })]), - ) - .unwrap(); - - assert_eq!( - *calls.lock().unwrap(), + let cache = cache( vec![ - "firstsecondreply", - "plain input", - "nested\ntail", - "result text", - "questionsrctfound{\"a\":1}", - ] + info(INDEX, base_info), + info_missing(&isolated), + create_index(&isolated, 3), + hset(&isolated, PROMPT, "key1", &VECTOR, &entry), + ], + FakeEmbedder::new(&[]), ); + + cache.set_cache("key1", entry, &context).unwrap(); } -#[test] -fn ttl_passes_through_context_only() { - let (embedder, _) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection( - MockRedisConnection::new(Vec::::new()), - embedder, - config(), +#[rstest] +fn create_index_race_rechecks_schema_and_stores(entry: Value, context: SemanticCacheContext) { + let cache = cache( + vec![ + info_missing(INDEX), + MockCmd::new( + create_index_command(INDEX, 3), + Err::<&str, _>(redis::RedisError::from(( + redis::ErrorKind::Extension, + "Index already exists", + ))), + ), + info(INDEX, compatible_info(3)), + hset(INDEX, PROMPT, "key1", &VECTOR, &entry), + ], + FakeEmbedder::new(&[]), ); - assert_eq!(cache.get_ttl(&SemanticCacheContext::default()), None); + + cache.set_cache("key1", entry, &context).unwrap(); +} + +#[rstest] +#[case::punctuation_and_spaces("a:b, c|d", "a\\:b\\,\\ c\\|d")] +#[case::braces_and_dots("{x}.y", "\\{x\\}\\.y")] +#[case::plain("key1", "key1")] +fn tag_special_characters_are_escaped_in_search_filter( + #[case] tag: &str, + #[case] escaped: &str, + context: SemanticCacheContext, +) { + let cache = cache( + vec![ + info(INDEX, compatible_info(3)), + search(INDEX, escaped, &VECTOR, Ok(empty_result())), + ], + FakeEmbedder::new(&[]), + ); + + assert_eq!(cache.get_cache(tag, &context).unwrap(), None); +} + +#[rstest] +#[case::content_parts( + SemanticCacheContext { + messages: Some(json!([ + {"role": "user", "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}]}, + {"role": "assistant", "content": "reply"}, + ])), + ..Default::default() + }, + "firstsecondreply" +)] +#[case::responses_string_input( + SemanticCacheContext { input: Some(json!(" plain input ")), ..Default::default() }, + "plain input" +)] +#[case::responses_nested_input( + SemanticCacheContext { + input: Some(json!([{"content": [{"type": "input_text", "text": "nested"}]}, "tail"])), + ..Default::default() + }, + "nested\ntail" +)] +#[case::responses_output_text( + SemanticCacheContext { input: Some(json!({"output_text": " result text "})), ..Default::default() }, + "result text" +)] +#[case::search_results( + messages_context(vec![json!({ + "role": "user", + "content": "question", + "search_results": [{"source": "src", "title": "t", "content": [{"text": "found"}], "citations": {"a": 1}}], + })]), + "questionsrctfound{\"a\":1}" +)] +#[case::empty_messages_fall_back_to_input( + SemanticCacheContext { messages: Some(json!([])), input: Some(json!("fallback")), ..Default::default() }, + "fallback" +)] +fn prompt_extraction_matches_python_message_and_input_shapes( + #[case] context: SemanticCacheContext, + #[case] prompt: &str, +) { + let embedder = FakeEmbedder::new(&[]); + let calls = embedder.calls.clone(); + let cache = cache( + vec![ + info(INDEX, compatible_info(3)), + search(INDEX, "key1", &VECTOR, Ok(empty_result())), + ], + embedder, + ); + + cache.get_cache("key1", &context).unwrap(); + + assert_eq!(*calls.lock().unwrap(), vec![(prompt.to_owned(), None)]); +} + +#[rstest] +#[case(None)] +#[case(Some(Duration::from_secs(9)))] +fn ttl_passes_through_context_only(#[case] ttl: Option) { + let cache = cache(Vec::new(), FakeEmbedder::new(&[])); + assert_eq!( cache.get_ttl(&SemanticCacheContext { - ttl: Some(Duration::from_secs(9)), + ttl, ..Default::default() }), - Some(Duration::from_secs(9)) + ttl ); } +#[rstest] #[tokio::test] -async fn async_paths_embed_then_run_blocking_redis_work() { - let vector = vec![0.1f32, 0.2, 0.3]; - let prompt = "hello prompt"; - let tag = "key1"; - let hash_key = format!("{INDEX}:{}", entry_id(prompt, tag)); - let value = entry(); - let connection = MockRedisConnection::new([ - MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(3))), - MockCmd::new( - redis::cmd("HSET") - .arg(&hash_key) - .arg("entry_id") - .arg(entry_id(prompt, tag)) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(encoded(&value)) - .arg("prompt_vector") - .arg(vector_bytes(&vector)) - .arg("inserted_at") - .arg("1700000000.5") - .arg("updated_at") - .arg("1700000000.5") - .arg("litellm_cache_key") - .arg(tag), - Ok(7), - ), - MockCmd::new( - search_command(INDEX, tag, &vector), - Ok(search_result(hit_fields(tag, "0.05", encoded(&value)))), - ), - ]) - .assert_all_commands_consumed(); - let (embedder, _) = FakeEmbedder::new(&[]); - let cache = RedisSemanticCache::with_connection(connection, embedder, config()) - .with_clock(|| 1700000000.5); - let context = messages_context(vec![json!({"role": "user", "content": prompt})]); +async fn async_paths_embed_with_metadata_then_run_blocking_redis_work( + entry: Value, + context: SemanticCacheContext, +) { + let embedder = FakeEmbedder::new(&[]); + let calls = embedder.calls.clone(); + let cache = cache( + vec![ + info(INDEX, compatible_info(3)), + hset(INDEX, PROMPT, "key1", &VECTOR, &entry), + search( + INDEX, + "key1", + &VECTOR, + Ok(hit("key1", "0.05", encoded(&entry))), + ), + ], + embedder, + ); + let context = SemanticCacheContext { + metadata: Some(json!({"tenant": "team"})), + ..context + }; cache - .async_set_cache(tag, value.clone(), context.clone()) + .async_set_cache("key1", entry.clone(), context.clone()) .await .unwrap(); assert_eq!( - cache.async_get_cache(tag, &context).await.unwrap(), - Some(value) + cache.async_get_cache("key1", &context).await.unwrap(), + Some(entry) + ); + assert_eq!( + *calls.lock().unwrap(), + vec![(PROMPT.to_owned(), context.metadata.clone()); 2] ); } -#[test] -fn shared_base_index_across_dimensions_replaces_the_isolated_index() { +#[rstest] +fn accessors_report_the_config(config: RedisSemanticConfig) { + let cache = cache(Vec::new(), FakeEmbedder::new(&[])); + + assert_eq!(cache.index_name(), config.index_name); + assert!((cache.similarity_threshold() - config.similarity_threshold).abs() < 1e-6); +} + +#[rstest] +fn shared_base_index_across_dimensions_replaces_the_isolated_index(entry: Value) { // Pins parity with Python's `_isolated` + overwrite=True flow. let prompt = "shared prompt"; - let tag = "key1"; let isolated = format!("{INDEX}_isolated"); - let value = entry(); let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); - let store_hash = |index: &str, vector: &[f32]| { - MockCmd::new( - redis::cmd("HSET") - .arg(format!("{index}:{}", entry_id(prompt, tag))) - .arg("entry_id") - .arg(entry_id(prompt, tag)) - .arg("prompt") - .arg(prompt) - .arg("response") - .arg(encoded(&value)) - .arg("prompt_vector") - .arg(vector_bytes(vector)) - .arg("inserted_at") - .arg("1700000000.5") - .arg("updated_at") - .arg("1700000000.5") - .arg("litellm_cache_key") - .arg(tag), - Ok(7), - ) - }; let vector_a = vec![0.1f32; 8]; - let connection_a = MockRedisConnection::new([ - MockCmd::new( - redis::cmd("FT.INFO").arg(INDEX), - Err::(unknown_index_error()), - ), - MockCmd::new(create_index_command(INDEX, 8), Ok("OK")), - store_hash(INDEX, &vector_a), - ]) - .assert_all_commands_consumed(); - let (embedder_a, _) = FakeEmbedder::new(&[(prompt, &vector_a)]); - let worker_a = RedisSemanticCache::with_connection(connection_a, embedder_a, config()) - .with_clock(|| 1700000000.5); - worker_a.set_cache(tag, value.clone(), &context()).unwrap(); + let worker_a = cache( + vec![ + info_missing(INDEX), + create_index(INDEX, 8), + hset(INDEX, prompt, "key1", &vector_a, &entry), + ], + FakeEmbedder::new(&[(prompt, &vector_a)]), + ); + worker_a + .set_cache("key1", entry.clone(), &context()) + .unwrap(); let vector_b = vec![0.2f32; 4]; - let connection_b = MockRedisConnection::new([ - MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(8))), - MockCmd::new( - redis::cmd("FT.INFO").arg(&isolated), - Err::(unknown_index_error()), - ), - MockCmd::new(create_index_command(&isolated, 4), Ok("OK")), - store_hash(&isolated, &vector_b), - MockCmd::new( - search_command(&isolated, tag, &vector_b), - Ok(search_result(hit_fields(tag, "0.0", encoded(&value)))), - ), - MockCmd::new( - search_command(&isolated, tag, &vector_b), - Err::(redis::RedisError::from(( - redis::ErrorKind::Extension, - "Vector dimension mismatch", - ))), - ), - ]) - .assert_all_commands_consumed(); - let (embedder_b, _) = FakeEmbedder::new(&[(prompt, &vector_b)]); - let worker_b = RedisSemanticCache::with_connection(connection_b, embedder_b, config()) - .with_clock(|| 1700000000.5); - worker_b.set_cache(tag, value.clone(), &context()).unwrap(); + let worker_b = cache( + vec![ + info(INDEX, compatible_info(8)), + info_missing(&isolated), + create_index(&isolated, 4), + hset(&isolated, prompt, "key1", &vector_b, &entry), + search( + &isolated, + "key1", + &vector_b, + Ok(hit("key1", "0.0", encoded(&entry))), + ), + search( + &isolated, + "key1", + &vector_b, + Err(redis::RedisError::from(( + redis::ErrorKind::Extension, + "Vector dimension mismatch", + ))), + ), + ], + FakeEmbedder::new(&[(prompt, &vector_b)]), + ); + worker_b + .set_cache("key1", entry.clone(), &context()) + .unwrap(); assert_eq!( - worker_b.get_cache(tag, &context()).unwrap(), - Some(value.clone()) + worker_b.get_cache("key1", &context()).unwrap(), + Some(entry.clone()) ); let vector_c = vec![0.3f32; 16]; - let connection_c = MockRedisConnection::new([ - MockCmd::new(redis::cmd("FT.INFO").arg(INDEX), Ok(compatible_info(8))), - MockCmd::new(redis::cmd("FT.INFO").arg(&isolated), Ok(compatible_info(4))), - MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")), - MockCmd::new(create_index_command(&isolated, 16), Ok("OK")), - store_hash(&isolated, &vector_c), - ]) - .assert_all_commands_consumed(); - let (embedder_c, _) = FakeEmbedder::new(&[(prompt, &vector_c)]); - let worker_c = RedisSemanticCache::with_connection(connection_c, embedder_c, config()) - .with_clock(|| 1700000000.5); - worker_c.set_cache(tag, value.clone(), &context()).unwrap(); + let worker_c = cache( + vec![ + info(INDEX, compatible_info(8)), + info(&isolated, compatible_info(4)), + MockCmd::new(redis::cmd("FT.DROPINDEX").arg(&isolated), Ok("OK")), + create_index(&isolated, 16), + hset(&isolated, prompt, "key1", &vector_c, &entry), + ], + FakeEmbedder::new(&[(prompt, &vector_c)]), + ); + worker_c + .set_cache("key1", entry.clone(), &context()) + .unwrap(); assert_eq!( - worker_b.get_cache(tag, &context()).unwrap_err(), + worker_b.get_cache("key1", &context()).unwrap_err(), Error::Unavailable ); } -#[test] -fn live_shared_index_is_replaced_across_dimensions() { - let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { +#[fixture] +fn redis_stack_url() -> Option { + std::env::var("LITELLM_REDIS_STACK_URL").ok() +} + +fn live_cache( + url: &str, + index_name: &str, + prompt: &str, + vector: Vec, +) -> RedisSemanticCache> { + RedisSemanticCache::new( + url, + FakeEmbedder::new(&[(prompt, vector.as_slice())]), + JsonCodec::::new(), + RedisSemanticConfig { + index_name: index_name.to_owned(), + similarity_threshold: 0.9, + }, + ) + .unwrap() +} + +#[rstest] +fn live_shared_index_is_replaced_across_dimensions(redis_stack_url: Option, entry: Value) { + let Some(url) = redis_stack_url else { return; }; // Pins parity with Python's `_isolated` + overwrite=True flow. let base = format!("rust_semantic_shared_{}", std::process::id()); let isolated = format!("{base}_isolated"); let prompt = "shared live prompt"; - let tag = "key1"; let context = || messages_context(vec![json!({"role": "user", "content": prompt})]); - let value = entry(); - let worker = |vector: Vec| { - let (embedder, _) = FakeEmbedder::new(&[(prompt, vector.as_slice())]); - RedisSemanticCache::new( - &url, - embedder, - RedisSemanticConfig { - index_name: base.clone(), - similarity_threshold: 0.9, - }, - ) - .unwrap() - }; - let worker_a = worker(vec![0.1f32; 8]); - worker_a.set_cache(tag, value.clone(), &context()).unwrap(); + let worker_a = live_cache(&url, &base, prompt, vec![0.1f32; 8]); + worker_a + .set_cache("key1", entry.clone(), &context()) + .unwrap(); - let worker_b = worker(vec![0.2f32; 4]); - worker_b.set_cache(tag, value.clone(), &context()).unwrap(); + let worker_b = live_cache(&url, &base, prompt, vec![0.2f32; 4]); + worker_b + .set_cache("key1", entry.clone(), &context()) + .unwrap(); assert_eq!( - worker_b.get_cache(tag, &context()).unwrap(), - Some(value.clone()) + worker_b.get_cache("key1", &context()).unwrap(), + Some(entry.clone()) ); - let worker_c = worker(vec![0.3f32; 16]); - worker_c.set_cache(tag, value.clone(), &context()).unwrap(); + let worker_c = live_cache(&url, &base, prompt, vec![0.3f32; 16]); + worker_c + .set_cache("key1", entry.clone(), &context()) + .unwrap(); assert_eq!( - worker_b.get_cache(tag, &context()).unwrap_err(), + worker_b.get_cache("key1", &context()).unwrap_err(), Error::Unavailable ); @@ -961,39 +746,29 @@ fn live_shared_index_is_replaced_across_dimensions() { } } -#[test] -fn live_store_lookup_and_ttl_against_redis_stack() { - let Ok(url) = std::env::var("LITELLM_REDIS_STACK_URL") else { +#[rstest] +fn live_store_lookup_and_ttl_against_redis_stack(redis_stack_url: Option, entry: Value) { + let Some(url) = redis_stack_url else { return; }; - let vector = vec![0.1f32, 0.2, 0.3, 0.4]; let prompt = "rust semantic cache live prompt"; - let tag = "live-key"; let index_name = format!("rust_semantic_test_{}", std::process::id()); - let (embedder, _) = FakeEmbedder::new(&[(prompt, &vector)]); - let cache = RedisSemanticCache::new( - &url, - embedder, - RedisSemanticConfig { - index_name: index_name.clone(), - similarity_threshold: 0.9, - }, - ) - .unwrap(); + let cache = live_cache(&url, &index_name, prompt, vec![0.1, 0.2, 0.3, 0.4]); let context = SemanticCacheContext { ttl: Some(Duration::from_secs(120)), ..messages_context(vec![json!({"role": "user", "content": prompt})]) }; - let value = entry(); - cache.set_cache(tag, value.clone(), &context).unwrap(); - assert_eq!(cache.get_cache(tag, &context).unwrap(), Some(value)); + cache + .set_cache("live-key", entry.clone(), &context) + .unwrap(); + assert_eq!(cache.get_cache("live-key", &context).unwrap(), Some(entry)); assert_eq!(cache.get_cache("other-key", &context).unwrap(), None); let mut connection = redis::Client::open(url).unwrap().get_connection().unwrap(); let ttl: i64 = redis::Commands::ttl( &mut connection, - format!("{index_name}:{}", entry_id(prompt, tag)), + format!("{index_name}:{}", entry_id(prompt, "live-key")), ) .unwrap(); assert!( diff --git a/litellm-rust/crates/cache-redis-semantic/tests/contract.rs b/litellm-rust/crates/cache-redis-semantic/tests/contract.rs new file mode 100644 index 00000000000..fb3cda9616e --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/tests/contract.rs @@ -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, 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; +} diff --git a/litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs new file mode 100644 index 00000000000..0f5c2985c36 --- /dev/null +++ b/litellm-rust/crates/cache-redis-semantic/tests/support/mod.rs @@ -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)>>>; + +/// 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>, + 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, 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, Error> { + self.embed(prompt, metadata) + } +} + +struct FakeIndex { + prefix: Vec, + dims: usize, + vector_field: String, +} + +#[derive(Default)] +struct SearchState { + indexes: HashMap, + hashes: BTreeMap, BTreeMap>>, +} + +/// 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>, +} + +impl FakeSearch { + fn run(&self, args: Vec>) -> redis::RedisResult { + 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::>(); + 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 { + 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> { + let replies = parse_commands(commands) + .into_iter() + .map(|args| self.run(args)) + .collect::>>()?; + 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 { + 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::(); + let norm = |vector: &[f32]| { + vector + .iter() + .map(|value| f64::from(*value).powi(2)) + .sum::() + .sqrt() + }; + dot / (norm(left) * norm(right)) +} + +/// Splits a packed RESP request into each command's arguments. +fn parse_commands(mut bytes: &[u8]) -> Vec>> { + 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::().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 +} diff --git a/litellm-rust/crates/cache-redis/Cargo.toml b/litellm-rust/crates/cache-redis/Cargo.toml index ea937098698..a234286a338 100644 --- a/litellm-rust/crates/cache-redis/Cargo.toml +++ b/litellm-rust/crates/cache-redis/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/cache-redis/src/cache.rs b/litellm-rust/crates/cache-redis/src/cache.rs index 24399c9b2f9..ebcf0b6916b 100644 --- a/litellm-rust/crates/cache-redis/src/cache.rs +++ b/litellm-rust/crates/cache-redis/src/cache.rs @@ -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 { - Pool(r2d2::Pool), - Cluster(r2d2::Pool), - Fixed(Mutex), -} - -impl Connections -where - C: redis::ConnectionLike + Send + 'static, -{ - pub fn execute( - &self, - operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, - ) -> Result { - 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(connections: Arc, operation: F) -> Result - where - T: Send + 'static, - F: FnOnce(&mut ConnectionRef<'_>) -> Result + 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 { - 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 { - connections: Arc>, - default_ttl: Duration, - codec: S, - namespace: Option, - topology: RedisTopology, + pub(crate) connections: Arc>, + pub(crate) default_ttl: Duration, + pub(crate) codec: S, + pub(crate) namespace: Option, + pub(crate) topology: RedisTopology, + /// The server's major version, read from `INFO` once, like Python's `redis_version`. + pub(crate) major_version: Arc>, } impl RedisCache { @@ -125,20 +40,11 @@ impl RedisCache { codec, namespace: None, topology: topology.clone(), + major_version: Arc::default(), }) } } -fn pool(manager: M) -> Result, 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 RedisCache 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 { + pub(crate) fn namespaced_keys(&self, keys: &[String]) -> Vec { + 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) -> u64 { + ttl_seconds(ttl.unwrap_or(self.default_ttl)) + } + + pub(crate) fn execute( + &self, + operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, + ) -> Result { + 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(&self, operation: F) -> Result + where + T: Send + 'static, + F: FnOnce(&mut ConnectionRef<'_>) -> Result + Send + 'static, + { + Connections::run_blocking(Arc::clone(&self.connections), operation).await + } + + pub(crate) fn namespaced_pattern(&self) -> Result { 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, Error> { + pub(crate) fn decode_response(&self, value: redis::Value) -> Result, 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, Error> { + pub(crate) fn decode_batch_response( + &self, + value: redis::Value, + ) -> Result, 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 BaseCache for RedisCache -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 { - 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, 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, 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::, _>>()?; - 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 { - 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 { + 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::().ok().map(|major| major as u32), } } -impl BatchCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn batch_get_cache( - &self, - keys: &[String], - _: &ExactCacheContext, - ) -> Result>, Error> { - let keys = keys - .iter() - .map(|key| self.namespaced_key(key)) - .collect::>(); - let values = self.connections.execute(|connection| { - redis::cmd("MGET") - .arg(keys) - .query::>(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, - _: ExactCacheContext, - ) -> Result>, Error> { - let keys = keys - .iter() - .map(|key| self.namespaced_key(key)) - .collect::>(); - let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - redis::cmd("MGET") - .arg(keys) - .query::>(connection) - .map_err(|_| Error::Unavailable) - }) - .await?; - values - .into_iter() - .map(|value| self.decode_batch_response(value)) - .collect() - } -} - -impl DeleteCache for RedisCache -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 FlushCache for RedisCache -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 CounterCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - fn increment_cache( - &self, - key: &str, - amount: f64, - context: ExactCacheContext, - ) -> Result { - 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 { - 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 { - 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>, 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( - connection: &mut ConnectionRef<'_>, - codec: &S, - key: &str, - candidate: S::Value, - eligible: &[S::Value], - ttl: u64, -) -> Result -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::(connection) - .map_err(|_| Error::Unavailable)?; - if applied { - return Ok(existing.unwrap_or(candidate)); - } - } - Err(Error::Unavailable) -} - -impl ClaimCache for RedisCache -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 { - 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, - context: ExactCacheContext, - ) -> Result { - 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::>::ttl_seconds(Duration::ZERO), - 1 - ); - assert_eq!( - RedisCache::>::ttl_seconds(Duration::from_millis(1500)), - 2 - ); - assert_eq!( - RedisCache::>::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::::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::::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::::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::::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) } diff --git a/litellm-rust/crates/cache-redis/src/cache/operations.rs b/litellm-rust/crates/cache-redis/src/cache/operations.rs deleted file mode 100644 index 4345ee879b3..00000000000 --- a/litellm-rust/crates/cache-redis/src/cache/operations.rs +++ /dev/null @@ -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), - Integer(i64), - Float(f64), -} - -impl From<&str> for RedisArg { - fn from(value: &str) -> Self { - Self::Bytes(value.as_bytes().to_vec()) - } -} - -impl From for RedisArg { - fn from(value: String) -> Self { - Self::Bytes(value.into_bytes()) - } -} - -impl From> for RedisArg { - fn from(value: Vec) -> Self { - Self::Bytes(value) - } -} - -impl From for RedisArg { - fn from(value: i64) -> Self { - Self::Integer(value) - } -} - -impl From for RedisArg { - fn from(value: f64) -> Self { - Self::Float(value) - } -} - -impl redis::ToRedisArgs for RedisArg { - fn write_redis_args(&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, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub struct RedisLpopOperation { - pub key: String, - pub count: Option, -} - -#[derive(Clone, Debug, PartialEq, Eq)] -pub enum RedisLpopResult { - Missing, - Value(Vec), - Values(Vec>), -} - -pub struct RedisScript { - connections: Arc>, - namespace: Option, - source: String, -} - -impl CacheScript for RedisScript -where - C: redis::ConnectionLike + Send + 'static, -{ - type Argument = RedisArg; - type Output = redis::Value; - - async fn invoke( - &self, - keys: Vec, - arguments: Vec, - ) -> Result { - let keys = keys - .into_iter() - .map(|key| namespaced_key(self.namespace.as_deref(), &key)) - .collect::>(); - 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 RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - pub async fn delete_cache_keys(&self, keys: Vec) -> Result { - if keys.is_empty() { - return Ok(0); - } - let keys = keys - .into_iter() - .map(|key| self.namespaced_key(&key)) - .collect::>(); - 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>, Error> { - let keys = keys - .iter() - .map(|key| self.namespaced_key(key)) - .collect::>(); - let values = self.connections.execute(|connection| { - redis::cmd("MGET") - .arg(keys) - .query::>(connection) - .map_err(|_| Error::Unavailable) - })?; - values.into_iter().map(count).collect() - } - - pub async fn async_batch_get_counts( - &self, - keys: Vec, - ) -> Result>, Error> { - let keys = keys - .iter() - .map(|key| self.namespaced_key(key)) - .collect::>(); - let values = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - redis::cmd("MGET") - .arg(keys) - .query::>(connection) - .map_err(|_| Error::Unavailable) - }) - .await?; - values.into_iter().map(count).collect() - } - - pub fn sync_ping(&self) -> Result { - self.connections - .execute(|connection| connection.ping().map_err(|_| Error::Unavailable)) - } - - pub async fn ping(&self) -> Result { - 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, Error> { - let key = self.namespaced_key(key); - let ttl = Connections::run_blocking(Arc::clone(&self.connections), move |connection| { - redis::cmd("TTL") - .arg(key) - .query::(connection) - .map_err(|_| Error::Unavailable) - }) - .await?; - Ok((ttl >= 0).then_some(ttl)) - } - - pub async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, 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, - ttl: Option, - ) -> Result { - 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::) - .transpose() - .map_err(|_| Error::Unavailable)? - .ok_or(Error::Unavailable) - }) - .await - } - - pub async fn async_rpush(&self, key: &str, values: Vec) -> Result { - 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, - ) -> Result, 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::, _>>()?; - 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, - ) -> Result { - 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::(connection) - .map_err(|_| Error::Unavailable) - }) - .await?; - lpop_result(value, multiple) - } - - pub async fn async_lpop_pipeline( - &self, - operations: Vec, - ) -> Result, Error> { - let operations = operations - .into_iter() - .map(|operation| (self.namespaced_key(&operation.key), operation.count)) - .collect::>(); - if operations.is_empty() { - return Ok(Vec::new()); - } - let multiple = operations - .iter() - .map(|(_, count)| count.is_some()) - .collect::>(); - 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, - arguments: Vec, - ) -> Result { - let keys = keys - .into_iter() - .map(|key| self.namespaced_key(&key)) - .collect::>(); - 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 { - self.connections - .execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST"))) - } - - pub fn info(&self) -> Result { - self.connections - .execute(|connection| connection.node_text(&redis::cmd("INFO"))) - } - - pub fn flushall(&self) -> Result<(), Error> { - self.connections.execute(|connection| connection.flushall()) - } -} - -impl RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - pub fn increment_with_floor( - &self, - key: &str, - amount: i64, - ttl: Duration, - ) -> Result { - 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, - ) -> Result, Error> { - let operations = operations - .into_iter() - .map(|operation| { - ( - self.namespaced_key(&operation.key), - operation.amount, - operation.ttl.map(Self::ttl_seconds), - ) - }) - .collect::>(); - 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 { - 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, - ) -> Result { - 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, 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 { - match value { - redis::Value::Nil => Ok(RedisLpopResult::Missing), - redis::Value::Array(values) if multiple => values - .into_iter() - .map(redis_bytes) - .collect::, _>>() - .map(RedisLpopResult::Values), - value if !multiple => redis_bytes(value).map(RedisLpopResult::Value), - _ => Err(Error::InvalidEntry), - } -} - -fn count(value: redis::Value) -> Result, 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 { - redis::cmd("EVAL") - .arg(INCREMENT_WITH_FLOOR_SCRIPT) - .arg(1) - .arg(key) - .arg(amount) - .arg(ttl) - .query(connection) - .map_err(|_| Error::Unavailable) -} - -impl TtlCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn async_get_ttl(&self, key: &str) -> Result, Error> { - RedisCache::async_get_ttl(self, key) - .await - .map(|ttl| ttl.map(|seconds| Duration::from_secs(seconds as u64))) - } -} - -impl ScanCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, Error> { - RedisCache::async_scan_iter(self, pattern, count).await - } -} - -impl ClientInfoCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type ClientList = String; - type Info = String; - - fn client_list(&self) -> Result { - RedisCache::client_list(self) - } - - fn info(&self) -> Result { - RedisCache::info(self) - } -} - -impl SetCache for RedisCache -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, - ttl: Option, - ) -> Result { - RedisCache::async_set_cache_sadd(self, key, values, ttl).await - } -} - -impl QueueCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type QueueValue = RedisArg; - type PopResult = RedisLpopResult; - - async fn async_rpush(&self, key: &str, values: Vec) -> Result { - RedisCache::async_rpush(self, key, values).await - } - - async fn async_lpop(&self, key: &str, count: Option) -> Result { - RedisCache::async_lpop(self, key, count).await - } -} - -impl ScriptCache for RedisCache -where - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type Script = RedisScript; - - fn async_register_script(&self, source: String) -> Self::Script { - RedisScript { - connections: Arc::clone(&self.connections), - namespace: self.namespace.clone(), - source, - } - } -} diff --git a/litellm-rust/crates/cache-redis/src/claim.rs b/litellm-rust/crates/cache-redis/src/claim.rs new file mode 100644 index 00000000000..e4cbe8c8d5c --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/claim.rs @@ -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>, 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( + connection: &mut ConnectionRef<'_>, + codec: &S, + key: &str, + candidate: S::Value, + eligible: &[S::Value], + ttl: u64, +) -> Result +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::(connection) + .map_err(|_| Error::Unavailable)?; + if applied { + return Ok(existing.unwrap_or(candidate)); + } + } + Err(Error::Unavailable) +} + +impl ClaimCache for RedisCache +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 { + 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, + context: ExactCacheContext, + ) -> Result { + 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 + } +} diff --git a/litellm-rust/crates/cache-redis/src/cache/connection.rs b/litellm-rust/crates/cache-redis/src/connection.rs similarity index 65% rename from litellm-rust/crates/cache-redis/src/cache/connection.rs rename to litellm-rust/crates/cache-redis/src/connection.rs index 013bf055f89..2f58e2a9b80 100644 --- a/litellm-rust/crates/cache-redis/src/cache/connection.rs +++ b/litellm-rust/crates/cache-redis/src/connection.rs @@ -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 { - 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 { + Pool(r2d2::Pool), + Cluster(r2d2::Pool), + Fixed(Mutex), +} + +impl Connections +where + C: redis::ConnectionLike + Send + 'static, +{ + pub fn execute( + &self, + operation: impl FnOnce(&mut ConnectionRef<'_>) -> Result, + ) -> Result { + 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(connections: Arc, operation: F) -> Result + where + T: Send + 'static, + F: FnOnce(&mut ConnectionRef<'_>) -> Result + 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 { + 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(manager: M) -> Result, 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(pool: &r2d2::Pool) +where + M: r2d2::ManageConnection>, +{ + let mut idle = Vec::new(); + while let Some(mut connection) = pool.try_get() { + connection.failed = true; + idle.push(connection); + } +} + +pub(crate) struct PooledConnection { + 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 { + fn open(url: &str) -> Result { 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 { + fn open(url: &str, startup_nodes: &[RedisNode]) -> Result { 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( &mut self, - commands: Vec, - ) -> Result, Error> { + pipeline: &redis::Pipeline, + ) -> Result { match self { - Self::Node(connection) => { - let mut pipeline = redis::pipe(); - for command in &commands { - pipeline.add_command(command.clone()); - } - pipeline - .query::>(*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> = 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::>>() - .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> { - let mut groups: HashMap> = 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>); + +impl redis::ConnectionLike for Replies { + fn req_packed_command(&mut self, _: &[u8]) -> redis::RedisResult { + Err((redis::ErrorKind::Client, "replies hold a pipeline only").into()) + } + + fn req_packed_commands( + &mut self, + _: &[u8], + _: usize, + _: usize, + ) -> redis::RedisResult> { + 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 } diff --git a/litellm-rust/crates/cache-redis/src/counter.rs b/litellm-rust/crates/cache-redis/src/counter.rs new file mode 100644 index 00000000000..adcd597b2f4 --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/counter.rs @@ -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 CounterCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn increment_cache( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + ) -> Result { + 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 { + 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, + ) -> Result, 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 { + 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 CountReadCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn batch_get_counts(&self, keys: &[String]) -> Result>, 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) -> Result>, 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, Error> { + redis::from_redis_value(value).map_err(|_| Error::InvalidEntry) +} + +impl BoundedCounterCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn increment_with_floor(&self, key: &str, amount: i64, ttl: Duration) -> Result { + 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 { + 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, + ) -> Result { + 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 { + redis::cmd("EVAL") + .arg(INCREMENT_WITH_FLOOR_SCRIPT) + .arg(1) + .arg(key) + .arg(amount) + .arg(ttl) + .query(connection) + .map_err(|_| Error::Unavailable) +} diff --git a/litellm-rust/crates/cache-redis/src/keys.rs b/litellm-rust/crates/cache-redis/src/keys.rs new file mode 100644 index 00000000000..b9d7bb20ff3 --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/keys.rs @@ -0,0 +1,63 @@ +use std::time::Duration; + +use litellm_cache::{CacheCodec, Error, RefreshTtlCache, ScanCache, TtlCache}; + +use crate::cache::RedisCache; + +impl TtlCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + async fn async_get_ttl(&self, key: &str) -> Result, Error> { + let key = self.namespaced_key(key); + let ttl = self + .run(move |connection| { + redis::cmd("TTL") + .arg(key) + .query::(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + Ok(u64::try_from(ttl).ok().map(Duration::from_secs)) + } +} + +impl RefreshTtlCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + async fn async_refresh_ttl(&self, key: &str, ttl: Option) -> Result { + 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 ScanCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + async fn async_scan_iter(&self, pattern: &str, count: usize) -> Result, 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 + } +} diff --git a/litellm-rust/crates/cache-redis/src/lib.rs b/litellm-rust/crates/cache-redis/src/lib.rs index efb0db931ac..037e39e5d40 100644 --- a/litellm-rust/crates/cache-redis/src/lib.rs +++ b/litellm-rust/crates/cache-redis/src/lib.rs @@ -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}; diff --git a/litellm-rust/crates/cache-redis/src/lifecycle.rs b/litellm-rust/crates/cache-redis/src/lifecycle.rs new file mode 100644 index 00000000000..6ab8c31509e --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/lifecycle.rs @@ -0,0 +1,85 @@ +use litellm_cache::{ + CacheCodec, CacheConnectionResult, CacheConnectionStatus, ClientInfoCache, ConnectionCache, + DisconnectCache, Error, PingCache, +}; + +use crate::{cache::RedisCache, topology::RedisTopology}; + +impl PingCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn sync_ping(&self) -> Result { + self.execute(|connection| connection.ping().map_err(|_| Error::Unavailable)) + } + + async fn ping(&self) -> Result { + self.run(|connection| connection.ping().map_err(|_| Error::Unavailable)) + .await + } +} + +impl ConnectionCache for RedisCache +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 { + 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 DisconnectCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + async fn disconnect(&self) -> Result<(), Error> { + self.connections.disconnect(); + Ok(()) + } +} + +impl ClientInfoCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type ClientList = String; + type Info = String; + + fn client_list(&self) -> Result { + self.execute(|connection| connection.node_text(redis::cmd("CLIENT").arg("LIST"))) + } + + fn info(&self) -> Result { + self.execute(|connection| connection.node_text(&redis::cmd("INFO"))) + } +} diff --git a/litellm-rust/crates/cache-redis/src/queue.rs b/litellm-rust/crates/cache-redis/src/queue.rs new file mode 100644 index 00000000000..623bce23640 --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/queue.rs @@ -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; +pub type RedisLpopOperation = PopOperation; + +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum RedisLpopResult { + Missing, + Value(Vec), + Values(Vec>), +} + +impl SetCache for RedisCache +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, + ttl: Option, + ) -> Result { + 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 QueueCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type QueueValue = RedisArg; + type PopResult = RedisLpopResult; + + async fn async_rpush(&self, key: &str, values: Vec) -> Result { + 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, + max_len: usize, + ) -> Result { + 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, + ) -> Result, 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) -> Result { + 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::(connection) + .map_err(|_| Error::Unavailable) + }) + .await?; + lpop_result(value, count.is_some()) + } + + async fn async_lpop_pipeline( + &self, + operations: Vec, + ) -> Result, 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::>(); + 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::>(&pipeline)) + .await? + .into_iter() + .zip(multiple) + .map(|(value, multiple)| lpop_result(value, multiple)) + .collect() + } +} + +fn lpop(key: String, count: Option) -> 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, Error> { + match value { + redis::Value::BulkString(bytes) => Ok(bytes), + redis::Value::SimpleString(text) => Ok(text.into_bytes()), + _ => Err(Error::InvalidEntry), + } +} + +impl RedisCache +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 { + 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::>(&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 { + match value { + redis::Value::Nil => Ok(RedisLpopResult::Missing), + redis::Value::Array(values) if multiple => values + .into_iter() + .map(redis_bytes) + .collect::, _>>() + .map(RedisLpopResult::Values), + value if !multiple => redis_bytes(value).map(RedisLpopResult::Value), + _ => Err(Error::InvalidEntry), + } +} diff --git a/litellm-rust/crates/cache-redis/src/script.rs b/litellm-rust/crates/cache-redis/src/script.rs new file mode 100644 index 00000000000..e8a16bcc2b5 --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/script.rs @@ -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), + Integer(i64), + Float(f64), +} + +impl From<&str> for RedisArg { + fn from(value: &str) -> Self { + Self::Bytes(value.as_bytes().to_vec()) + } +} + +impl From for RedisArg { + fn from(value: String) -> Self { + Self::Bytes(value.into_bytes()) + } +} + +impl From> for RedisArg { + fn from(value: Vec) -> Self { + Self::Bytes(value) + } +} + +impl From for RedisArg { + fn from(value: i64) -> Self { + Self::Integer(value) + } +} + +impl From for RedisArg { + fn from(value: f64) -> Self { + Self::Float(value) + } +} + +impl redis::ToRedisArgs for RedisArg { + fn write_redis_args(&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 { + connections: Arc>, + namespace: Option, + source: String, +} + +impl CacheScript for RedisScript +where + C: redis::ConnectionLike + Send + 'static, +{ + type Argument = RedisArg; + type Output = redis::Value; + + async fn invoke( + &self, + keys: Vec, + arguments: Vec, + ) -> Result { + let keys = keys + .into_iter() + .map(|key| namespaced_key(self.namespace.as_deref(), &key)) + .collect::>(); + let source = self.source.clone(); + Connections::run_blocking(Arc::clone(&self.connections), move |connection| { + eval(connection, &source, keys, arguments) + }) + .await + } +} + +impl RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + pub async fn async_eval( + &self, + script: String, + keys: Vec, + arguments: Vec, + ) -> Result { + 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, + arguments: Vec, +) -> Result { + redis::cmd("EVAL") + .arg(script) + .arg(keys.len()) + .arg(keys) + .arg(arguments) + .query(connection) + .map_err(|_| Error::Unavailable) +} + +impl ScriptCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type Script = RedisScript; + + fn async_register_script(&self, source: String) -> Self::Script { + RedisScript { + connections: Arc::clone(&self.connections), + namespace: self.namespace.clone(), + source, + } + } +} diff --git a/litellm-rust/crates/cache-redis/src/store.rs b/litellm-rust/crates/cache-redis/src/store.rs new file mode 100644 index 00000000000..111d481ad02 --- /dev/null +++ b/litellm-rust/crates/cache-redis/src/store.rs @@ -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 BaseCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + type Value = S::Value; + type Context = ExactCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + 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, 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, 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 TtlPipelineCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + async fn async_set_cache_pipeline_with_ttls( + &self, + entries: Vec<(String, Self::Value, Option)>, + ) -> 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 BatchCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn batch_get_cache( + &self, + keys: &[String], + _: &ExactCacheContext, + ) -> Result>, 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, + _: ExactCacheContext, + ) -> Result>, 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, +) -> Result, Error> { + redis::cmd("MGET") + .arg(keys) + .query(connection) + .map_err(|_| Error::Unavailable) +} + +impl DeleteCache for RedisCache +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 BulkDeleteCache for RedisCache +where + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + async fn delete_cache_keys(&self, keys: Vec) -> Result { + 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 FlushCache for RedisCache +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 FlushAllCache for RedisCache +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) + }) +} diff --git a/litellm-rust/crates/cache-redis/tests/cache.rs b/litellm-rust/crates/cache-redis/tests/cache.rs index 337f27984f8..70baeffd572 100644 --- a/litellm-rust/crates/cache-redis/tests/cache.rs +++ b/litellm-rust/crates/cache-redis/tests/cache.rs @@ -1,43 +1,84 @@ +mod support; + use std::time::Duration; use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, CacheCodec, CacheConnectionStatus, CacheScript, ClaimCache, - CounterCache, DeleteCache, Error, ExactCacheContext, FlushCache, IncrementOperation, JsonCodec, - ScriptCache, get_cache, set_cache, + BaseCache, BatchCache, BatchEntry, BoundedCounterCache, BulkDeleteCache, CacheCodec, + CacheConnectionStatus, CacheScript, ClaimCache, ClientInfoCache, ConnectionCache, + CountReadCache, CounterCache, DeleteCache, DisconnectCache, Error, ExactCacheContext, + FlushAllCache, FlushCache, IncrementOperation, JsonCodec, PingCache, QueueCache, + RefreshTtlCache, ScanCache, ScriptCache, SetCache, TtlCache, TtlPipelineCache, get_cache, + set_cache, }; use litellm_cache_redis::{ RedisArg, RedisCache, RedisLpopOperation, RedisLpopResult, RedisRpushOperation, }; use redis_test::{MockCmd, MockRedisConnection}; +use rstest::{fixture, rstest}; +use serde_json::json; +use support::TaggedByteCodec; -struct TaggedByteCodec(u8); +type Mocked = RedisCache; -impl CacheCodec for TaggedByteCodec { - type Value = u8; - - fn encode(&self, value: &u8) -> Result, Error> { - if *value > 127 { - return Err(Error::InvalidEntry); - } - Ok(vec![self.0, *value]) - } - - fn decode(&self, bytes: &[u8]) -> Result { - match bytes { - [tag, value] if *tag == self.0 => Ok(*value), - _ => Err(Error::InvalidEntry), - } - } +fn mock(commands: Vec) -> MockRedisConnection { + MockRedisConnection::new(commands).assert_all_commands_consumed() } -#[test] +fn tagged(commands: Vec) -> Mocked { + RedisCache::with_connection(mock(commands), None, TaggedByteCodec(42)) +} + +fn json_cache(commands: Vec) -> Mocked> { + RedisCache::with_connection(mock(commands), None, JsonCodec::new()) +} + +fn team(commands: Vec) -> Mocked> { + json_cache(commands).with_namespace(Some("team".into())) +} + +#[fixture] +fn context() -> ExactCacheContext { + ExactCacheContext::default() +} + +fn scan(pattern: &str, cursor: u64, count: usize, reply: redis::Value) -> MockCmd { + MockCmd::new( + redis::cmd("SCAN") + .cursor_arg(cursor) + .arg("MATCH") + .arg(pattern) + .arg("COUNT") + .arg(count), + Ok(reply), + ) +} + +#[rstest] fn constructor_rejects_invalid_urls() { assert!(RedisCache::new("not a redis url", None, JsonCodec::::new()).is_err()); } -#[test] +#[rstest] +#[case::zero_rounds_up_to_one(Some(Duration::ZERO), 1)] +#[case::fractions_round_up(Some(Duration::from_millis(1500)), 2)] +#[case::whole_seconds_are_kept(Some(Duration::from_secs(15)), 15)] +#[case::missing_ttl_uses_default(None, 600)] +fn writes_round_ttls_up_to_positive_seconds(#[case] ttl: Option, #[case] seconds: u64) { + let cache = tagged(vec![MockCmd::new( + redis::cmd("SETEX") + .arg("key") + .arg(seconds) + .arg([42u8, 7].as_slice()), + Ok("OK"), + )]); + cache + .set_cache("key", 7, &ExactCacheContext { ttl }) + .unwrap(); +} + +#[rstest] fn generic_helpers_use_the_injected_codec_and_ttl() { - let connection = MockRedisConnection::new([ + let cache = tagged(vec![ MockCmd::new( redis::cmd("SETEX") .arg("counter") @@ -46,9 +87,7 @@ fn generic_helpers_use_the_injected_codec_and_ttl() { Ok("OK"), ), MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])), - ]) - .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); + ]); let context = ExactCacheContext { ttl: Some(Duration::from_millis(1500)), }; @@ -56,34 +95,56 @@ fn generic_helpers_use_the_injected_codec_and_ttl() { assert_eq!(get_cache(&cache, "counter", &context).unwrap(), Some(7)); } -#[tokio::test] -async fn async_operations_preserve_codec_ttl_and_missing_values() { - let connection = MockRedisConnection::new([ +#[rstest] +fn commands_round_trip_entries_and_delete_only_namespaced_keys(context: ExactCacheContext) { + let value = json!({"deployment": "model-a", "cooldown_seconds": 30}); + let payload = JsonCodec::::new() + .encode(&value) + .unwrap(); + let cache = json_cache(vec![ MockCmd::new( redis::cmd("SETEX") - .arg("counter") - .arg(9) - .arg([42u8, 7].as_slice()), + .arg("litellm-cache:key") + .arg(600) + .arg(payload.clone()), Ok("OK"), ), - MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])), - MockCmd::new( - redis::cmd("SETEX") - .arg("batch") - .arg(2) - .arg([42u8, 8].as_slice()), - Ok("OK"), - ), - MockCmd::new(redis::cmd("DEL").arg("counter"), Ok(1u32)), - MockCmd::new(redis::cmd("GET").arg("counter"), Ok(redis::Value::Nil)), + 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(); + .with_namespace(Some("litellm-cache".into())); + + cache.set_cache("key", value.clone(), &context).unwrap(); + assert_eq!(cache.get_cache("key", &context).unwrap(), Some(value)); + cache.delete_cache("key").unwrap(); +} + +#[rstest] +#[tokio::test] +async fn async_operations_preserve_codec_ttl_and_missing_values(context: ExactCacheContext) { let cache = RedisCache::with_connection( - connection, + mock(vec![ + MockCmd::new( + redis::cmd("SETEX") + .arg("counter") + .arg(9) + .arg([42u8, 7].as_slice()), + Ok("OK"), + ), + MockCmd::new(redis::cmd("GET").arg("counter"), Ok(vec![42u8, 7])), + MockCmd::new( + redis::cmd("SETEX") + .arg("batch") + .arg(2) + .arg([42u8, 8].as_slice()), + Ok("OK"), + ), + MockCmd::new(redis::cmd("DEL").arg("counter"), Ok(1u32)), + MockCmd::new(redis::cmd("GET").arg("counter"), Ok(redis::Value::Nil)), + ]), Some(Duration::from_secs(9)), TaggedByteCodec(42), ); - let context = ExactCacheContext::default(); cache .batch_cache_write("counter", 7, context.clone()) .await @@ -108,15 +169,13 @@ async fn async_operations_preserve_codec_ttl_and_missing_values() { ); } +#[rstest] #[tokio::test] -async fn codec_errors_propagate_without_writing_partial_batches() { - let connection = MockRedisConnection::new([ +async fn codec_errors_propagate_without_writing_partial_batches(context: ExactCacheContext) { + let cache = tagged(vec![ MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])), MockCmd::new(redis::cmd("GET").arg("invalid"), Ok(vec![99u8, 7])), - ]) - .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); - let context = ExactCacheContext::default(); + ]); assert_eq!( cache.set_cache("invalid", 255, &context), Err(Error::InvalidEntry) @@ -134,6 +193,15 @@ async fn codec_errors_propagate_without_writing_partial_batches() { .await, Err(Error::InvalidEntry) ); + assert_eq!( + cache + .async_set_cache_pipeline_with_ttls(vec![ + ("valid".into(), 7, None), + ("invalid".into(), 255, None), + ]) + .await, + Err(Error::InvalidEntry) + ); assert_eq!( cache.get_cache("invalid", &context), Err(Error::InvalidEntry) @@ -144,221 +212,269 @@ async fn codec_errors_propagate_without_writing_partial_batches() { ); } -#[test] -fn namespaces_are_optional_and_existing_prefixes_are_not_duplicated() { - let connection = MockRedisConnection::new([ - MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)), - MockCmd::new(redis::cmd("GET").arg("team:key"), Ok(redis::Value::Nil)), - ]) - .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("team".into())); - assert_eq!( - cache - .get_cache("key", &ExactCacheContext::default()) - .unwrap(), - None - ); - assert_eq!( - cache - .get_cache("team:key", &ExactCacheContext::default()) - .unwrap(), - None - ); +#[rstest] +#[case::bare_key("key")] +#[case::already_prefixed("team:key")] +fn namespaces_are_added_once(#[case] key: &str, context: ExactCacheContext) { + let cache = team(vec![MockCmd::new( + redis::cmd("GET").arg("team:key"), + Ok(redis::Value::Nil), + )]); + assert_eq!(cache.get_cache(key, &context).unwrap(), None); } -#[test] -fn flush_requires_a_namespace_and_escapes_glob_metacharacters() { - let unscoped = RedisCache::with_connection( - MockRedisConnection::new([]).assert_all_commands_consumed(), - None, - JsonCodec::::new(), - ); +#[rstest] +#[case::empty(Some(String::new()))] +#[case::missing(None)] +fn empty_namespaces_leave_keys_unprefixed( + #[case] namespace: Option, + context: ExactCacheContext, +) { + let cache = json_cache(vec![MockCmd::new( + redis::cmd("GET").arg("key"), + Ok(redis::Value::Nil), + )]) + .with_namespace(namespace); + assert_eq!(cache.namespace(), None); + assert_eq!(cache.get_cache("key", &context).unwrap(), None); +} + +#[rstest] +#[tokio::test] +async fn flush_requires_a_namespace() { + let unscoped = json_cache(Vec::new()); assert_eq!(unscoped.flush_cache(), Err(Error::UnscopedFlush)); - let connection = MockRedisConnection::new([ - MockCmd::new( - redis::cmd("SCAN") - .cursor_arg(0) - .arg("MATCH") - .arg("team\\*:*") - .arg("COUNT") - .arg(1000), - Ok(redis_test::redis_value!(["0", ["team*:key"]])), - ), - MockCmd::new(redis::cmd("DEL").arg("team*:key"), Ok(1u32)), + assert_eq!( + unscoped.async_flush_cache().await, + Err(Error::UnscopedFlush) + ); +} + +#[rstest] +#[case::plain_namespace("litellm-cache", "litellm-cache:*", "litellm-cache:key")] +#[case::glob_metacharacters_are_escaped("team*", "team\\*:*", "team*:key")] +fn flush_scans_and_deletes_only_namespaced_keys( + #[case] namespace: &str, + #[case] pattern: &str, + #[case] key: &str, +) { + let cache = json_cache(vec![ + scan(pattern, 0, 1000, redis_test::redis_value!(["0", [key]])), + MockCmd::new(redis::cmd("DEL").arg(key), Ok(1u32)), ]) - .assert_all_commands_consumed(); - let scoped = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("team*".into())); - scoped.flush_cache().unwrap(); + .with_namespace(Some(namespace.into())); + cache.flush_cache().unwrap(); } +#[rstest] #[tokio::test] -async fn connection_failures_use_the_python_result_contract() { - let error = redis::RedisError::from((redis::ErrorKind::Io, "connection refused")); - let connection = - MockRedisConnection::new([MockCmd::new(redis::cmd("PING"), Err::(error))]) - .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); - - let result = cache.test_connection().await.unwrap(); - assert_eq!(result.status, CacheConnectionStatus::Failed); - assert!(result.message.starts_with("Redis connection failed:")); - assert!(result.error.is_some()); +async fn async_flush_deletes_each_scan_page_separately() { + let cache = team(vec![ + scan( + "team:*", + 0, + 1000, + redis_test::redis_value!(["7", ["team:a", "team:b"]]), + ), + MockCmd::new(redis::cmd("DEL").arg("team:a").arg("team:b"), Ok(2u32)), + scan( + "team:*", + 7, + 1000, + redis_test::redis_value!(["0", ["team:c"]]), + ), + MockCmd::new(redis::cmd("DEL").arg("team:c"), Ok(1u32)), + ]); + cache.async_flush_cache().await.unwrap(); } +#[rstest] +fn flushall_ignores_the_namespace() { + team(vec![MockCmd::new(redis::cmd("FLUSHALL"), Ok("OK"))]) + .flushall() + .unwrap(); +} + +#[rstest] #[tokio::test] -async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries() { - let connection = MockRedisConnection::new([MockCmd::new( +async fn batch_reads_keep_order_and_treat_invalid_values_as_invalid_entries( + context: ExactCacheContext, +) { + let cache = tagged(vec![MockCmd::new( redis::cmd("MGET").arg("hit").arg("miss").arg("invalid"), Ok(vec![ redis::Value::BulkString(vec![42, 7]), redis::Value::Nil, redis::Value::BulkString(vec![99, 7]), ]), - )]) - .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, TaggedByteCodec(42)); + )]); assert_eq!( cache - .async_batch_get_cache( - vec!["hit".into(), "miss".into(), "invalid".into()], - ExactCacheContext::default(), - ) + .async_batch_get_cache(vec!["hit".into(), "miss".into(), "invalid".into()], context) .await .unwrap(), vec![BatchEntry::Hit(7), BatchEntry::Miss, BatchEntry::Invalid] ); } +#[rstest] #[tokio::test] -async fn async_flush_deletes_each_scan_page_separately() { - let connection = MockRedisConnection::new([ - MockCmd::new( - redis::cmd("SCAN") - .cursor_arg(0) - .arg("MATCH") - .arg("team:*") - .arg("COUNT") - .arg(1000), - Ok(redis_test::redis_value!(["7", ["team:a", "team:b"]])), - ), - MockCmd::new(redis::cmd("DEL").arg("team:a").arg("team:b"), Ok(2u32)), - MockCmd::new( - redis::cmd("SCAN") - .cursor_arg(7) - .arg("MATCH") - .arg("team:*") - .arg("COUNT") - .arg(1000), - Ok(redis_test::redis_value!(["0", ["team:c"]])), - ), - MockCmd::new(redis::cmd("DEL").arg("team:c"), Ok(1u32)), - ]) - .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("team".into())); +async fn ttl_pipeline_keeps_each_entry_ttl_and_defaults_missing_ones() { + let mut pipeline = redis::pipe(); + pipeline + .cmd("SETEX") + .arg("ns:team_id:t1") + .arg(60u64) + .arg(r#"{"team_id":"t1"}"#) + .cmd("SETEX") + .arg("ns:u1") + .arg(7u64) + .arg(r#"{"user_id":"u1"}"#) + .cmd("SETEX") + .arg("ns:org_id:o1") + .arg(300u64) + .arg(r#"{"a":1}"#); + let cache = RedisCache::with_connection( + mock(vec![MockCmd::with_values( + pipeline, + Ok(vec!["OK", "OK", "OK"]), + )]), + Some(Duration::from_secs(300)), + JsonCodec::::new(), + ) + .with_namespace(Some("ns".into())); - cache.async_flush_cache().await.unwrap(); + cache + .async_set_cache_pipeline_with_ttls(vec![ + ( + "team_id:t1".into(), + json!({"team_id": "t1"}), + Some(Duration::from_secs(60)), + ), + ( + "u1".into(), + json!({"user_id": "u1"}), + Some(Duration::from_secs(7)), + ), + ("org_id:o1".into(), json!({"a": 1}), None), + ]) + .await + .unwrap(); } +#[rstest] #[tokio::test] -async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() { - let mut sadd_pipeline = redis::pipe(); - sadd_pipeline - .cmd("SADD") - .arg("team:members") - .arg("a") - .arg("b") - .cmd("EXPIRE") - .arg("team:members") - .arg(600u64) - .ignore(); - let connection = MockRedisConnection::new([ +async fn empty_pipelines_skip_the_round_trip(context: ExactCacheContext) { + let cache = json_cache(Vec::new()); + cache + .async_set_cache_pipeline(Vec::new(), context) + .await + .unwrap(); + cache + .async_set_cache_pipeline_with_ttls(Vec::new()) + .await + .unwrap(); + assert_eq!(cache.delete_cache_keys(Vec::new()).await.unwrap(), 0); + assert_eq!( + cache.async_rpush_pipeline(Vec::new()).await.unwrap(), + Vec::::new() + ); + assert_eq!( + cache.async_lpop_pipeline(Vec::new()).await.unwrap(), + Vec::::new() + ); + assert_eq!( + cache.async_increment_pipeline(Vec::new()).await.unwrap(), + Vec::::new() + ); +} + +#[rstest] +#[tokio::test] +async fn count_reads_parse_integers_and_keep_missing_counters() { + let mget = || { MockCmd::new( redis::cmd("MGET").arg("team:count").arg("team:missing"), Ok(redis_test::redis_value!(["7", nil])), - ), - MockCmd::new( - redis::cmd("MGET").arg("team:count").arg("team:missing"), - Ok(redis_test::redis_value!(["7", nil])), - ), + ) + }; + let cache = team(vec![mget(), mget()]); + let keys = vec!["count".to_string(), "missing".to_string()]; + + assert_eq!(cache.batch_get_counts(&keys).unwrap(), [Some(7), None]); + assert_eq!( + cache.async_batch_get_counts(keys).await.unwrap(), + [Some(7), None] + ); +} + +#[rstest] +#[tokio::test] +async fn pings_run_on_both_paths() { + let cache = team(vec![ MockCmd::new(redis::cmd("PING"), Ok("PONG")), MockCmd::new(redis::cmd("PING"), Ok("PONG")), - MockCmd::new(redis::cmd("TTL").arg("team:missing"), Ok(-2i64)), - MockCmd::new( - redis::cmd("SCAN") - .cursor_arg(0) - .arg("MATCH") - .arg("team:job-*") - .arg("COUNT") - .arg(25), - Ok(redis_test::redis_value!(["4", ["team:job-a"]])), + ]); + assert!(cache.sync_ping().unwrap()); + assert!(cache.ping().await.unwrap()); +} + +#[rstest] +#[case::remaining(12, Some(Duration::from_secs(12)))] +#[case::no_expiry(-1, None)] +#[case::missing(-2, None)] +#[tokio::test] +async fn ttl_reads_hide_negative_replies(#[case] reply: i64, #[case] ttl: Option) { + let cache = team(vec![MockCmd::new( + redis::cmd("TTL").arg("team:key"), + Ok(reply), + )]); + assert_eq!(cache.async_get_ttl("key").await.unwrap(), ttl); +} + +#[rstest] +#[case::explicit_ttl(Some(Duration::from_secs(30)), 30, 1, true)] +#[case::default_ttl(None, 600, 1, true)] +#[case::missing_key(Some(Duration::from_secs(30)), 30, 0, false)] +#[tokio::test] +async fn refresh_ttl_expires_existing_keys_only( + #[case] ttl: Option, + #[case] seconds: u64, + #[case] reply: i64, + #[case] refreshed: bool, +) { + let cache = team(vec![MockCmd::new( + redis::cmd("EXPIRE").arg("team:key").arg(seconds), + Ok(reply), + )]); + assert_eq!( + cache.async_refresh_ttl("key", ttl).await.unwrap(), + refreshed + ); +} + +#[rstest] +#[tokio::test] +async fn scan_stops_at_count_and_bulk_delete_reports_existing_keys() { + let cache = team(vec![ + scan( + "team:job-*", + 0, + 25, + redis_test::redis_value!(["4", ["team:job-a"]]), ), - MockCmd::new( - redis::cmd("SCAN") - .cursor_arg(4) - .arg("MATCH") - .arg("team:job-*") - .arg("COUNT") - .arg(25), - Ok(redis_test::redis_value!(["0", ["team:job-b"]])), + scan( + "team:job-*", + 4, + 25, + redis_test::redis_value!(["0", ["team:job-b"]]), ), MockCmd::new( redis::cmd("DEL").arg("team:job-a").arg("team:job-b"), Ok(2u32), ), - MockCmd::with_values( - sadd_pipeline, - Ok(vec![redis::Value::Int(2), redis::Value::Int(1)]), - ), - MockCmd::new( - redis::cmd("RPUSH").arg("team:queue").arg("a").arg("b"), - Ok(2u32), - ), - MockCmd::new( - redis::cmd("LPOP").arg("team:queue").arg(2usize), - Ok(redis_test::redis_value!(["a", "b"])), - ), - MockCmd::new( - redis::cmd("EVAL") - .arg("return KEYS[1]") - .arg(1usize) - .arg("team:key"), - Ok("team:key"), - ), - MockCmd::new( - redis::cmd("EVAL") - .arg("return KEYS[1]") - .arg(1usize) - .arg("team:key"), - Ok("team:key"), - ), - MockCmd::new(redis::cmd("CLIENT").arg("LIST"), Ok("id=1")), - MockCmd::new(redis::cmd("INFO"), Ok("redis_version:7")), - MockCmd::new(redis::cmd("FLUSHALL"), Ok("OK")), - ]) - .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("team".into())); - - assert_eq!( - cache - .batch_get_counts(&["count".into(), "missing".into()]) - .unwrap(), - [Some(7), None] - ); - assert_eq!( - cache - .async_batch_get_counts(vec!["count".into(), "missing".into()]) - .await - .unwrap(), - [Some(7), None] - ); - assert!(cache.sync_ping().unwrap()); - assert!(cache.ping().await.unwrap()); - assert_eq!(cache.async_get_ttl("missing").await.unwrap(), None); + ]); assert_eq!( cache.async_scan_iter("job-", 25).await.unwrap(), ["team:job-a", "team:job-b"] @@ -370,6 +486,25 @@ async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() { .unwrap(), 2 ); +} + +#[rstest] +#[tokio::test] +async fn sets_add_members_and_arm_the_default_ttl() { + let mut pipeline = redis::pipe(); + pipeline + .cmd("SADD") + .arg("team:members") + .arg("a") + .arg("b") + .cmd("EXPIRE") + .arg("team:members") + .arg(600u64) + .ignore(); + let cache = team(vec![MockCmd::with_values( + pipeline, + Ok(vec![redis::Value::Int(2), redis::Value::Int(1)]), + )]); assert_eq!( cache .async_set_cache_sadd("members", vec!["a".into(), "b".into()], None) @@ -377,6 +512,32 @@ async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() { .unwrap(), 2 ); + assert_eq!( + cache + .async_set_cache_sadd("members", Vec::new(), None) + .await, + Err(Error::InvalidEntry) + ); +} + +#[rstest] +#[tokio::test] +async fn queues_push_and_pop_namespaced_lists() { + let cache = team(vec![ + MockCmd::new( + redis::cmd("RPUSH").arg("team:queue").arg("a").arg("b"), + Ok(2u32), + ), + MockCmd::new( + redis::cmd("INFO"), + Ok("# Server\r\nredis_version:7.2.4\r\n"), + ), + MockCmd::new( + redis::cmd("LPOP").arg("team:queue").arg(2usize), + Ok(redis_test::redis_value!(["a", "b"])), + ), + MockCmd::new(redis::cmd("LPOP").arg("team:queue"), Ok("c")), + ]); assert_eq!( cache .async_rpush("queue", vec!["a".into(), "b".into()]) @@ -384,10 +545,159 @@ async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() { .unwrap(), 2 ); + assert_eq!( + cache.async_rpush("queue", Vec::new()).await, + Err(Error::InvalidEntry) + ); assert_eq!( cache.async_lpop("queue", Some(2)).await.unwrap(), RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()]) ); + assert_eq!( + cache.async_lpop("queue", None).await.unwrap(), + RedisLpopResult::Value(b"c".to_vec()) + ); +} + +/// Python `RedisCache.async_lpop` checks `redis_version` from `INFO` and, below major version 7, +/// pops a counted batch as `count` single-command `LPOP` pipelines, dropping `None` replies. +#[rstest] +#[tokio::test] +async fn counted_lpop_falls_back_to_single_pops_below_redis_7() { + let single_pop = || { + let mut pipeline = redis::pipe(); + pipeline.cmd("LPOP").arg("team:queue"); + pipeline + }; + let cache = team(vec![ + MockCmd::new( + redis::cmd("INFO"), + Ok("# Server\r\nredis_version:6.2.14\r\n"), + ), + MockCmd::with_values(single_pop(), Ok(vec![redis_test::redis_value!("a")])), + MockCmd::with_values(single_pop(), Ok(vec![redis_test::redis_value!("b")])), + MockCmd::with_values(single_pop(), Ok(vec![redis::Value::Nil])), + ]); + + assert_eq!( + cache.async_lpop("queue", Some(3)).await.unwrap(), + RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()]) + ); +} + +/// Python keeps `redis_version = "Unknown"` when `INFO` fails and then assumes +/// `DEFAULT_REDIS_MAJOR_VERSION` (7), so a counted pop is one `LPOP key count`. The version is read +/// once: the second pop sends no second `INFO`. +#[rstest] +#[tokio::test] +async fn counted_lpop_assumes_redis_7_when_info_fails() { + let cache = team(vec![ + MockCmd::new( + redis::cmd("INFO"), + Err::(redis::RedisError::from((redis::ErrorKind::Io, "down"))), + ), + MockCmd::new( + redis::cmd("LPOP").arg("team:queue").arg(2usize), + Ok(redis_test::redis_value!(["a", "b"])), + ), + MockCmd::new( + redis::cmd("LPOP").arg("team:queue").arg(1usize), + Ok(redis_test::redis_value!(["c"])), + ), + ]); + + assert_eq!( + cache.async_lpop("queue", Some(2)).await.unwrap(), + RedisLpopResult::Values(vec![b"a".to_vec(), b"b".to_vec()]) + ); + assert_eq!( + cache.async_lpop("queue", Some(1)).await.unwrap(), + RedisLpopResult::Values(vec![b"c".to_vec()]) + ); +} + +fn push_and_trim(start: i64) -> redis::Pipeline { + let mut pipeline = redis::pipe(); + pipeline + .atomic() + .cmd("RPUSH") + .arg("ns:buf") + .arg("c") + .arg("d") + .cmd("LTRIM") + .arg("ns:buf") + .arg(start) + .arg(-1); + pipeline +} + +#[rstest] +#[case::keeps_newest_entries(3, -3)] +#[case::zero_keeps_everything(0, 0)] +#[tokio::test] +async fn rpush_and_trim_runs_push_and_trim_in_one_transaction( + #[case] max_len: usize, + #[case] start: i64, +) { + let cache = json_cache(vec![MockCmd::with_values( + push_and_trim(start), + Ok(vec![redis::Value::Array(vec![ + redis::Value::Int(4), + redis::Value::Okay, + ])]), + )]) + .with_namespace(Some("ns".into())); + + assert_eq!( + cache + .async_rpush_and_trim("buf", vec!["c".into(), "d".into()], max_len) + .await + .unwrap(), + 4 + ); +} + +#[rstest] +#[tokio::test] +async fn rpush_and_trim_raises_when_a_queued_command_fails() { + let wrong_type = redis::parse_redis_value( + b"-WRONGTYPE Operation against a key holding the wrong kind of value\r\n", + ) + .unwrap(); + let cache = json_cache(vec![MockCmd::with_values( + push_and_trim(-3), + Ok(vec![redis::Value::Array(vec![ + wrong_type, + redis::Value::Okay, + ])]), + )]) + .with_namespace(Some("ns".into())); + + assert_eq!( + cache + .async_rpush_and_trim("buf", vec!["c".into(), "d".into()], 3) + .await, + Err(Error::Unavailable) + ); + assert_eq!( + cache.async_rpush_and_trim("buf", Vec::new(), 3).await, + Err(Error::InvalidEntry) + ); +} + +#[rstest] +#[tokio::test] +async fn scripts_and_eval_namespace_their_keys() { + let eval = || { + MockCmd::new( + redis::cmd("EVAL") + .arg("return KEYS[1]") + .arg(1usize) + .arg("team:key"), + Ok("team:key"), + ) + }; + let cache = team(vec![eval(), eval()]); assert_eq!( cache .async_eval("return KEYS[1]".into(), vec!["key".into()], Vec::new()) @@ -403,13 +713,21 @@ async fn direct_redis_operations_preserve_namespace_values_and_missing_ttls() { .unwrap(), redis::Value::BulkString(b"team:key".to_vec()) ); - assert_eq!(cache.client_list().unwrap(), "id=1"); - assert_eq!(cache.info().unwrap(), "redis_version:7"); - cache.flushall().unwrap(); } +#[rstest] +fn client_list_and_info_return_server_text() { + let cache = team(vec![ + MockCmd::new(redis::cmd("CLIENT").arg("LIST"), Ok("id=1")), + MockCmd::new(redis::cmd("INFO"), Ok("redis_version:7")), + ]); + assert_eq!(cache.client_list().unwrap(), "id=1"); + assert_eq!(cache.info().unwrap(), "redis_version:7"); +} + +#[rstest] #[tokio::test] -async fn direct_redis_pipelines_preserve_operation_order() { +async fn pipelines_preserve_operation_order() { let mut rpush_pipeline = redis::pipe(); rpush_pipeline .cmd("RPUSH") @@ -425,19 +743,20 @@ async fn direct_redis_pipelines_preserve_operation_order() { .arg(2usize) .cmd("LPOP") .arg("team:b"); - let connection = MockRedisConnection::new([ + let queue = team(vec![ MockCmd::with_values( rpush_pipeline, Ok(vec![redis::Value::Int(1), redis::Value::Int(2)]), ), + MockCmd::new( + redis::cmd("INFO"), + Ok("# Server\r\nredis_version:7.2.4\r\n"), + ), MockCmd::with_values( lpop_pipeline, Ok(vec![redis_test::redis_value!(["one"]), redis::Value::Nil]), ), - ]) - .assert_all_commands_consumed(); - let queue = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("team".into())); + ]); assert_eq!( queue @@ -474,7 +793,59 @@ async fn direct_redis_pipelines_preserve_operation_order() { RedisLpopResult::Missing, ] ); +} +/// Below major version 7 a counted `LPOP` is unsupported, so a pipeline that mixes counted and +/// plain pops runs each operation through `async_lpop`: `count` single-`LPOP` pipelines for the +/// counted ones, a bare `LPOP` for the rest. No `LPOP key count` reaches the connection. +#[rstest] +#[tokio::test] +async fn lpop_pipeline_pops_one_at_a_time_below_redis_7() { + let single_pop = |key: &str| { + let mut pipeline = redis::pipe(); + pipeline.cmd("LPOP").arg(key); + pipeline + }; + let cache = team(vec![ + MockCmd::new( + redis::cmd("INFO"), + Ok("# Server\r\nredis_version:6.2.14\r\n"), + ), + MockCmd::with_values( + single_pop("team:a"), + Ok(vec![redis_test::redis_value!("one")]), + ), + MockCmd::with_values( + single_pop("team:a"), + Ok(vec![redis_test::redis_value!("two")]), + ), + MockCmd::new(redis::cmd("LPOP").arg("team:b"), Ok(redis::Value::Nil)), + ]); + + assert_eq!( + cache + .async_lpop_pipeline(vec![ + RedisLpopOperation { + key: "a".into(), + count: Some(2), + }, + RedisLpopOperation { + key: "b".into(), + count: None, + }, + ]) + .await + .unwrap(), + [ + RedisLpopResult::Values(vec![b"one".to_vec(), b"two".to_vec()]), + RedisLpopResult::Missing, + ] + ); +} + +#[rstest] +#[tokio::test] +async fn increment_pipeline_expires_only_operations_with_a_ttl() { let mut increment_pipeline = redis::pipe(); increment_pipeline .cmd("INCRBYFLOAT") @@ -487,17 +858,14 @@ async fn direct_redis_pipelines_preserve_operation_order() { .cmd("INCRBYFLOAT") .arg("team:counter") .arg(2.0f64); - let connection = MockRedisConnection::new([MockCmd::with_values( + let counters = team(vec![MockCmd::with_values( increment_pipeline, Ok(vec![ redis::Value::BulkString(b"1.5".to_vec()), redis::Value::Int(1), redis::Value::BulkString(b"3.5".to_vec()), ]), - )]) - .assert_all_commands_consumed(); - let counters = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("team".into())); + )]); assert_eq!( counters .async_increment_pipeline(vec![ @@ -518,6 +886,11 @@ async fn direct_redis_pipelines_preserve_operation_order() { ); } +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; ", @@ -532,20 +905,93 @@ const SET_MAX_SCRIPT: &str = concat!( "return ARGV[1]; end; return current" ); +fn increment_script(amount: f64) -> MockCmd { + MockCmd::new( + redis::cmd("EVAL") + .arg(INCREMENT_SCRIPT) + .arg(1) + .arg("counter") + .arg(amount) + .arg(600), + Ok("4.5"), + ) +} + +#[rstest] +#[tokio::test] +async fn increments_keep_an_existing_ttl_in_one_atomic_script(context: ExactCacheContext) { + let cache = RedisCache::with_connection( + mock(vec![increment_script(2.5), increment_script(2.5)]), + None, + JsonCodec::::new(), + ); + + assert_eq!( + cache + .increment_cache("counter", 2.5, context.clone()) + .unwrap(), + 4.5 + ); + assert_eq!( + cache + .async_increment("counter", 2.5, context, false) + .await + .unwrap(), + 4.5 + ); +} + +#[rstest] +#[case::explicit_ttl(Some(Duration::from_secs(60)), 60u64)] +#[case::default_ttl(None, 600u64)] +#[tokio::test] +async fn refresh_ttl_increments_rearm_the_ttl_in_the_same_round_trip( + #[case] ttl: Option, + #[case] seconds: u64, +) { + let mut pipeline = redis::pipe(); + pipeline + .cmd("INCRBYFLOAT") + .arg("ns:spend:key:k") + .arg(1.5f64) + .cmd("EXPIRE") + .arg("ns:spend:key:k") + .arg(seconds); + let cache = json_cache(vec![MockCmd::with_values( + pipeline, + Ok(vec![ + redis::Value::BulkString(b"1.5".to_vec()), + redis::Value::Int(1), + ]), + )]) + .with_namespace(Some("ns".into())); + + assert_eq!( + cache + .async_increment("spend:key:k", 1.5, ExactCacheContext { ttl }, true) + .await + .unwrap(), + 1.5 + ); +} + +#[rstest] #[tokio::test] async fn counter_repairs_are_atomic_and_use_default_ttl() { let floor = || { - redis::cmd("EVAL") - .arg(INCREMENT_WITH_FLOOR_SCRIPT) - .arg(1) - .arg("team:counter") - .arg(-2i64) - .arg(30u64) - .clone() + MockCmd::new( + redis::cmd("EVAL") + .arg(INCREMENT_WITH_FLOOR_SCRIPT) + .arg(1) + .arg("team:counter") + .arg(-2i64) + .arg(30u64), + Ok(0i64), + ) }; - let connection = MockRedisConnection::new([ - MockCmd::new(floor(), Ok(0i64)), - MockCmd::new(floor(), Ok(0i64)), + let cache = team(vec![ + floor(), + floor(), MockCmd::new( redis::cmd("EVAL") .arg(SET_MAX_SCRIPT) @@ -555,10 +1001,7 @@ async fn counter_repairs_are_atomic_and_use_default_ttl() { .arg(600u64), Ok("4.5"), ), - ]) - .assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()) - .with_namespace(Some("team".into())); + ]); assert_eq!( cache @@ -587,38 +1030,37 @@ const CLAIM_SCRIPT: &str = concat!( "elseif ARGV[4] == '1' then redis.call('EXPIRE', KEYS[1], ARGV[2]); end; return 1" ); -fn claim_eval(expected: &str, write: &str, refresh: bool) -> redis::Cmd { - let mut cmd = redis::cmd("EVAL"); - cmd.arg(CLAIM_SCRIPT) - .arg(1) - .arg("pin") - .arg(expected) - .arg(600) - .arg(write) - .arg(u8::from(refresh)); - cmd +fn claim_eval(expected: &str, write: &str, refresh: bool, applied: i64) -> MockCmd { + MockCmd::new( + redis::cmd("EVAL") + .arg(CLAIM_SCRIPT) + .arg(1) + .arg("pin") + .arg(expected) + .arg(600) + .arg(write) + .arg(u8::from(refresh)), + Ok(applied), + ) } +#[rstest] #[tokio::test] -async fn claims_match_eligible_values_written_by_another_encoder() { +async fn claims_match_eligible_values_written_by_another_encoder(context: ExactCacheContext) { let python_payload = r#"{"model_id": "a", "deployment": "east"}"#; - let stored = serde_json::json!({"deployment": "east", "model_id": "a"}); - let candidate = serde_json::json!({"model_id": "b"}); - let connection = MockRedisConnection::new([ + let stored = json!({"deployment": "east", "model_id": "a"}); + let cache = json_cache(vec![ MockCmd::new(redis::cmd("GET").arg("pin"), Ok(python_payload)), - MockCmd::new(claim_eval(python_payload, "", true), Ok(1)), - ]) - .assert_all_commands_consumed(); - let cache = - RedisCache::with_connection(connection, None, JsonCodec::::new()); + claim_eval(python_payload, "", true, 1), + ]); assert_eq!( cache .async_claim_cache( "pin", - candidate, + json!({"model_id": "b"}), vec![stored.clone()], - ExactCacheContext::default() + context ) .await .unwrap(), @@ -626,78 +1068,91 @@ async fn claims_match_eligible_values_written_by_another_encoder() { ); } -#[test] -fn claims_retry_when_the_key_changes_and_replace_ineligible_winners() { - let candidate = serde_json::json!({"model_id": "b"}); +#[rstest] +fn claims_retry_when_the_key_changes_and_replace_ineligible_winners(context: ExactCacheContext) { + let candidate = json!({"model_id": "b"}); let payload = r#"{"model_id":"b"}"#; - let connection = MockRedisConnection::new([ + let cache = json_cache(vec![ MockCmd::new(redis::cmd("GET").arg("pin"), Ok(redis::Value::Nil)), - MockCmd::new(claim_eval("", payload, false), Ok(0)), + claim_eval("", payload, false, 0), MockCmd::new(redis::cmd("GET").arg("pin"), Ok(r#"{"model_id":"gone"}"#)), - MockCmd::new(claim_eval(r#"{"model_id":"gone"}"#, payload, false), Ok(1)), - ]) - .assert_all_commands_consumed(); - let cache = - RedisCache::with_connection(connection, None, JsonCodec::::new()); + claim_eval(r#"{"model_id":"gone"}"#, payload, false, 1), + ]); assert_eq!( cache .claim_cache( "pin", candidate.clone(), - &[serde_json::json!({"model_id": "a"})], - ExactCacheContext::default() + &[json!({"model_id": "a"})], + context ) .unwrap(), candidate ); } -#[test] -fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl() { +#[rstest] +fn claims_without_eligible_values_keep_the_winner_without_refreshing_its_ttl( + context: ExactCacheContext, +) { let stored = r#"{"model_id": "a"}"#; - let connection = MockRedisConnection::new([ + let cache = json_cache(vec![ MockCmd::new(redis::cmd("GET").arg("pin"), Ok(stored)), - MockCmd::new(claim_eval(stored, "", false), Ok(1)), - ]) - .assert_all_commands_consumed(); - let cache = - RedisCache::with_connection(connection, None, JsonCodec::::new()); + claim_eval(stored, "", false, 1), + ]); assert_eq!( cache - .claim_cache( - "pin", - serde_json::json!({"model_id": "b"}), - &[], - ExactCacheContext::default() - ) + .claim_cache("pin", json!({"model_id": "b"}), &[], context) .unwrap(), - serde_json::json!({"model_id": "a"}) + json!({"model_id": "a"}) ); } +#[rstest] #[tokio::test] -async fn async_increment_runs_the_atomic_script() { - let mut eval = redis::cmd("EVAL"); - eval.arg(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" - )) - .arg(1) - .arg("counter") - .arg(2.5f64) - .arg(600); - let connection = - MockRedisConnection::new([MockCmd::new(eval, Ok("4.5"))]).assert_all_commands_consumed(); - let cache = RedisCache::with_connection(connection, None, JsonCodec::::new()); +async fn test_connection_reports_success_with_the_python_message() { + let cache = team(vec![MockCmd::new(redis::cmd("PING"), Ok("PONG"))]); - assert_eq!( - cache - .async_increment("counter", 2.5, ExactCacheContext::default()) - .await - .unwrap(), - 4.5 - ); + let result = cache.test_connection().await.unwrap(); + assert_eq!(result.status, CacheConnectionStatus::Success); + assert_eq!(result.message, "Redis connection test successful"); + assert_eq!(result.error, None); +} + +#[rstest] +#[case::unexpected_reply(Ok("NOPE"), "Redis ping returned False", false)] +#[case::connection_refused( + Err(redis::RedisError::from((redis::ErrorKind::Io, "connection refused"))), + "Redis connection failed:", + true +)] +#[tokio::test] +async fn test_connection_failures_use_the_python_result_contract( + #[case] reply: redis::RedisResult<&'static str>, + #[case] message: &str, + #[case] has_error: bool, +) { + let cache = json_cache(vec![MockCmd::new(redis::cmd("PING"), reply)]); + + let result = cache.test_connection().await.unwrap(); + assert_eq!(result.status, CacheConnectionStatus::Failed); + assert!(result.message.starts_with(message), "{}", result.message); + assert_eq!(result.error.is_some(), has_error); +} + +#[rstest] +#[tokio::test] +async fn disconnect_keeps_a_caller_owned_connection_usable() { + let cache = team(vec![MockCmd::new(redis::cmd("PING"), Ok("PONG"))]); + cache.disconnect().await.unwrap(); + assert!(cache.ping().await.unwrap()); +} + +#[rstest] +#[tokio::test] +async fn disconnect_drains_an_idle_pool_without_connecting() { + let cache = RedisCache::new("redis://127.0.0.1:1", None, JsonCodec::::new()).unwrap(); + cache.disconnect().await.unwrap(); } diff --git a/litellm-rust/crates/cache-redis/tests/cluster.rs b/litellm-rust/crates/cache-redis/tests/cluster.rs index 2c3fc818b66..c1a9a70a5d2 100644 --- a/litellm-rust/crates/cache-redis/tests/cluster.rs +++ b/litellm-rust/crates/cache-redis/tests/cluster.rs @@ -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>; +type Counter = RedisCache>; -fn topology() -> Option { - 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 { + 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 { + 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 { - 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>> { - 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 { let keys: Vec = (0..count).map(|index| format!("key-{index}")).collect(); - let slots: std::collections::HashSet = keys.iter().map(Slot::for_key).collect(); + let slots: HashSet = 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 { + 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, +) { + 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, +) { + 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, + 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, + counter: Option, + 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 = 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 = 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 = (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, +) { + 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, + #[from(cache)] + #[with("other")] + other: Option, + 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) { + 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, +) { + 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, + #[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, + #[with("claim")] cache: Option, + 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); diff --git a/litellm-rust/crates/cache-redis/tests/contract.rs b/litellm-rust/crates/cache-redis/tests/contract.rs new file mode 100644 index 00000000000..85d043b90b2 --- /dev/null +++ b/litellm-rust/crates/cache-redis/tests/contract.rs @@ -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(contract: Contract, cache: &JsonCache) +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; + } + } + } +} diff --git a/litellm-rust/crates/cache-redis/tests/support/mod.rs b/litellm-rust/crates/cache-redis/tests/support/mod.rs new file mode 100644 index 00000000000..bff1e618582 --- /dev/null +++ b/litellm-rust/crates/cache-redis/tests/support/mod.rs @@ -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, Error> { + if *value > 127 { + return Err(Error::InvalidEntry); + } + Ok(vec![self.0, *value]) + } + + fn decode(&self, bytes: &[u8]) -> Result { + 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>, +} + +impl FakeRedis { + fn run(&mut self, command: Vec>) -> redis::RedisResult { + 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>> { + 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 { + 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> { + let replies = commands(cmd) + .into_iter() + .map(|command| self.run(command)) + .collect::>>()?; + 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 = RedisCache, C>; + +pub fn fake_cache(namespace: &str) -> JsonCache { + 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 { + 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( + label: &str, + default_ttl: Duration, + codec: S, +) -> Option> { + let topology = cluster_topology()?; + Some( + RedisCache::connect(&cluster_url(), &topology, Some(default_ttl), codec) + .expect("cluster connection") + .with_namespace(Some(unique_namespace(label))), + ) +} diff --git a/litellm-rust/crates/cache-response/Cargo.toml b/litellm-rust/crates/cache-response/Cargo.toml index 04affb9872d..42a1afb2ba0 100644 --- a/litellm-rust/crates/cache-response/Cargo.toml +++ b/litellm-rust/crates/cache-response/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/cache-response/README.md b/litellm-rust/crates/cache-response/README.md index d048afb69f8..dbad474c9e7 100644 --- a/litellm-rust/crates/cache-response/README.md +++ b/litellm-rust/crates/cache-response/README.md @@ -1,14 +1,16 @@ -# Response cache foundation +# Response cache `ResponseCache` adds request keys, independent read/write controls, response envelopes, and freshness checks to any `B: BaseCache` ## 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` 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` 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 diff --git a/litellm-rust/crates/cache-response/src/exact.rs b/litellm-rust/crates/cache-response/src/exact.rs index f5e86b2598c..16b79e4b11b 100644 --- a/litellm-rust/crates/cache-response/src/exact.rs +++ b/litellm-rust/crates/cache-response/src/exact.rs @@ -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>; } +impl ConnectionProbe for ResponseCache +where + B: ConnectionCache, + B::Context: Default + PartialEq, +{ + fn test_connection<'a>(&'a self) -> BoxFuture<'a, Result> { + Box::pin(ResponseCache::test_connection(self)) + } +} + impl ExactResponseCache for ResponseCache where B: BaseCache + 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> { - Box::pin(ResponseCache::test_connection(self)) - } } diff --git a/litellm-rust/crates/cache-response/src/lib.rs b/litellm-rust/crates/cache-response/src/lib.rs index ab9867ac8db..a6a4bb3eb64 100644 --- a/litellm-rust/crates/cache-response/src/lib.rs +++ b/litellm-rust/crates/cache-response/src/lib.rs @@ -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}; diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index 5088402f125..e761c7157db 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -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 { + pub async fn test_connection(&self) -> Result + 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, + now: Duration, + ) -> Result, 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, + now: Duration, + ) -> Result, 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], @@ -290,6 +332,21 @@ where Ok(PartialHits::new(values)) } + fn fresh_semantic( + lookup: Result, Error>, + now: Duration, + max_age: Option, + ) -> Result, 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, now: Duration, diff --git a/litellm-rust/crates/cache-response/tests/caching.rs b/litellm-rust/crates/cache-response/tests/caching.rs index 0e8ce9b3b1d..93c2eb3d16d 100644 --- a/litellm-rust/crates/cache-response/tests/caching.rs +++ b/litellm-rust/crates/cache-response/tests/caching.rs @@ -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); } diff --git a/litellm-rust/crates/cache-response/tests/codec.rs b/litellm-rust/crates/cache-response/tests/codec.rs new file mode 100644 index 00000000000..4fb5a5094cb --- /dev/null +++ b/litellm-rust/crates/cache-response/tests/codec.rs @@ -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) { + 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::(&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"}]}), + } + ); +} diff --git a/litellm-rust/crates/cache-response/tests/connection.rs b/litellm-rust/crates/cache-response/tests/connection.rs new file mode 100644 index 00000000000..bc24cf56846 --- /dev/null +++ b/litellm-rust/crates/cache-response/tests/connection.rs @@ -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 = + 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 = service.clone(); + let exact: Arc = 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>>, + request: ResponseCacheRequest, +) { + let cache: Arc = 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]); +} diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index dcfc0301148..ec5e16f1367 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -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>> { - 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>>; +#[derive(Default)] struct SemanticBackend { entries: Mutex>, contexts: Mutex>, @@ -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 { - 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, Error>); + +impl BaseCache for ScoredBackend { + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, _: &Self::Context) -> Option { + None + } + + fn set_cache(&self, _: &str, _: Self::Value, _: &Self::Context) -> Result<(), Error> { + Ok(()) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, 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, Error> { + self.0.clone() + } + + async fn async_get_cache_with_similarity( + &self, + key: &str, + context: &Self::Context, + ) -> Result, Error> { + self.get_cache_with_similarity(key, context) + } +} + +fn scored(timestamp: f64, similarity: f64) -> Result, 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, Error>, + #[case] reads: bool, + #[case] expected: Result, 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::(&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); } diff --git a/litellm-rust/crates/cache-response/tests/support/mod.rs b/litellm-rust/crates/cache-response/tests/support/mod.rs new file mode 100644 index 00000000000..b992a937b70 --- /dev/null +++ b/litellm-rust/crates/cache-response/tests/support/mod.rs @@ -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; + +#[fixture] +pub fn memory() -> Arc>> { + 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, namespace: Option<&str>) -> ResponseCache { + 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)), + )) +} diff --git a/litellm-rust/crates/cache-s3/Cargo.toml b/litellm-rust/crates/cache-s3/Cargo.toml index cdc17e732cb..c8150180e7c 100644 --- a/litellm-rust/crates/cache-s3/Cargo.toml +++ b/litellm-rust/crates/cache-s3/Cargo.toml @@ -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"] } diff --git a/litellm-rust/crates/cache-s3/src/auth.rs b/litellm-rust/crates/cache-s3/src/auth.rs index b7ca722cea3..fdf71fc011b 100644 --- a/litellm-rust/crates/cache-s3/src/auth.rs +++ b/litellm-rust/crates/cache-s3/src/auth.rs @@ -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, } -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) -> Self { + pub fn with_env(config: AwsAuthConfig, env: fn(&str) -> Option) -> 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() } } diff --git a/litellm-rust/crates/cache-s3/src/cache.rs b/litellm-rust/crates/cache-s3/src/cache.rs index 9c791f42c3b..91cd5e8ef54 100644 --- a/litellm-rust/crates/cache-s3/src/cache.rs +++ b/litellm-rust/crates/cache-s3/src/cache.rs @@ -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 { } impl S3Cache { - 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 = 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 BaseCache for S3Cache { 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 DisconnectCache for S3Cache { async fn disconnect(&self) -> Result<(), Error> { Ok(()) } - - async fn test_connection(&self) -> Result { - Err(Error::UnsupportedOperation) - } } impl BatchCache for S3Cache {} diff --git a/litellm-rust/crates/cache-s3/src/lib.rs b/litellm-rust/crates/cache-s3/src/lib.rs index f6126dfa908..9be0210aad3 100644 --- a/litellm-rust/crates/cache-s3/src/lib.rs +++ b/litellm-rust/crates/cache-s3/src/lib.rs @@ -1,4 +1,6 @@ mod auth; mod cache; +mod transport; +pub use auth::S3Credentials; pub use cache::{S3Cache, S3CacheConfig, S3Endpoint}; diff --git a/litellm-rust/crates/cache-s3/src/transport.rs b/litellm-rust/crates/cache-s3/src/transport.rs new file mode 100644 index 00000000000..3e5ce578c31 --- /dev/null +++ b/litellm-rust/crates/cache-s3/src/transport.rs @@ -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) + }) + }) + } +} diff --git a/litellm-rust/crates/cache-s3/tests/auth.rs b/litellm-rust/crates/cache-s3/tests/auth.rs new file mode 100644 index 00000000000..59c5aa2e091 --- /dev/null +++ b/litellm-rust/crates/cache-s3/tests/auth.rs @@ -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 { + (name == "AWS_SESSION_TOKEN").then(|| "ambient".to_string()) +} + +fn environment_keys(name: &str) -> Option { + 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, + #[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 + ); +} diff --git a/litellm-rust/crates/cache-s3/tests/cache.rs b/litellm-rust/crates/cache-s3/tests/cache.rs index 9a71656286b..53a51f8a289 100644 --- a/litellm-rust/crates/cache-s3/tests/cache.rs +++ b/litellm-rust/crates/cache-s3/tests/cache.rs @@ -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> { - S3Cache::new( - config(endpoint.to_string()), - JsonCodec::::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 { - use aws_smithy_types::{DateTime, date_time::Format}; +fn http_date_from(headers: &HeaderMap, name: &str) -> Option { 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("NoSuchKey"), - ) - .mount(&server) - .await; - Mock::given(method("GET")) - .and(path("/cache-bucket/team/denied")) - .respond_with( - ResponseTemplate::new(403).set_body_string("AccessDenied"), - ) - .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("NoSuchKey"), + Ok(None) +)] +#[case::access_denied( + "denied", + ResponseTemplate::new(403).set_body_string("AccessDenied"), + 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, 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, +) { + 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::::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::::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}))); } diff --git a/litellm-rust/crates/cache-s3/tests/contract.rs b/litellm-rust/crates/cache-s3/tests/contract.rs new file mode 100644 index 00000000000..9855e2868ff --- /dev/null +++ b/litellm-rust/crates/cache-s3/tests/contract.rs @@ -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; +} diff --git a/litellm-rust/crates/cache-s3/tests/support/mod.rs b/litellm-rust/crates/cache-s3/tests/support/mod.rs new file mode 100644 index 00000000000..046b042c66a --- /dev/null +++ b/litellm-rust/crates/cache-s3/tests/support/mod.rs @@ -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>; + +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>>>, +} + +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("NoSuchKey"), + }, + _ => ResponseTemplate::new(405), + } + } +} diff --git a/litellm-rust/crates/cache-testing/Cargo.toml b/litellm-rust/crates/cache-testing/Cargo.toml new file mode 100644 index 00000000000..674472df50d --- /dev/null +++ b/litellm-rust/crates/cache-testing/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/cache-testing/src/lib.rs b/litellm-rust/crates/cache-testing/src/lib.rs new file mode 100644 index 00000000000..c227c1e12dd --- /dev/null +++ b/litellm-rust/crates/cache-testing/src/lib.rs @@ -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(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( + 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( + 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( + cache: &B, + context: B::Context, + prefix: &str, + values: Vec, +) 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::>(); + 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( + 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(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(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(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); +} diff --git a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml index f98bb5a5fa8..da9bb60c853 100644 --- a/litellm-rust/crates/cache-valkey-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-valkey-semantic/Cargo.toml @@ -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 diff --git a/litellm-rust/crates/cache-valkey-semantic/src/cache.rs b/litellm-rust/crates/cache-valkey-semantic/src/cache.rs new file mode 100644 index 00000000000..947c0edc479 --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/src/cache.rs @@ -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 { + connections: Arc>, + embedder: E, + codec: S, + config: ValkeySemanticConfig, + index_dimension: Arc>>, +} + +impl ValkeySemanticCache { + pub fn new( + url: &str, + embedder: E, + codec: S, + config: ValkeySemanticConfig, + ) -> Result { + Ok(Self { + connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), + embedder, + codec, + config, + index_dimension: Arc::new(Mutex::new(None)), + }) + } +} + +impl ValkeySemanticCache +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>) -> Result, Error> { + Ok(SemanticLookup { + value: lookup + .value + .map(|bytes| self.codec.decode(&bytes)) + .transpose()?, + similarity: lookup.similarity, + }) + } +} + +impl ValkeySemanticCache +where + E: Embedder, + S: CacheCodec + Clone, + C: redis::ConnectionLike + Send + 'static, +{ + /// The same index and connections behind a different embedder. + pub fn with_embedder(&self, embedder: E2) -> ValkeySemanticCache { + ValkeySemanticCache { + connections: Arc::clone(&self.connections), + embedder, + codec: self.codec.clone(), + config: self.config.clone(), + index_dimension: Arc::clone(&self.index_dimension), + } + } +} + +impl BaseCache for ValkeySemanticCache +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 { + 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, 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, 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 SemanticCache for ValkeySemanticCache +where + E: Embedder, + S: CacheCodec, + C: redis::ConnectionLike + Send + 'static, +{ + fn get_cache_with_similarity( + &self, + key: &str, + context: &Self::Context, + ) -> Result, 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, 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) + } +} diff --git a/litellm-rust/crates/cache-valkey-semantic/src/config.rs b/litellm-rust/crates/cache-valkey-semantic/src/config.rs new file mode 100644 index 00000000000..c417652c27b --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/src/config.rs @@ -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, +} diff --git a/litellm-rust/crates/cache-valkey-semantic/src/index.rs b/litellm-rust/crates/cache-valkey-semantic/src/index.rs new file mode 100644 index 00000000000..a0b7be4ca7b --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/src/index.rs @@ -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>>, + 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::(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::(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 { + 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::>(); + 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() + }) + }) +} diff --git a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs index 6062ccc842c..f2f62bc95bd 100644 --- a/litellm-rust/crates/cache-valkey-semantic/src/lib.rs +++ b/litellm-rust/crates/cache-valkey-semantic/src/lib.rs @@ -1,1153 +1,7 @@ -use std::{ - future::Future, - sync::{Arc, Mutex}, - time::Duration, -}; +mod cache; +mod config; +mod index; +mod search; -use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext}; -use litellm_cache_redis::{ - RedisTopology, - connection::{ConnectionRef, Connections}, -}; -use litellm_cache_response::CacheEntry; -use serde_json::Value; -use sha2::{Digest, Sha256}; -use uuid::Uuid; - -pub trait Embedder: Send + Sync + 'static { - fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error>; - - fn async_embed( - &self, - prompt: &str, - metadata: Option<&Value>, - ) -> impl Future, Error>> + Send; -} - -pub struct PreparedEmbedding(pub Vec); - -impl Embedder for PreparedEmbedding { - fn embed(&self, _prompt: &str, _metadata: Option<&Value>) -> Result, Error> { - Ok(self.0.clone()) - } - - async fn async_embed( - &self, - _prompt: &str, - _metadata: Option<&Value>, - ) -> Result, Error> { - Ok(self.0.clone()) - } -} - -#[derive(Clone, Debug, PartialEq)] -pub struct ValkeySemanticConfig { - pub similarity_threshold: f64, - pub index_name: String, -} - -pub const DEFAULT_INDEX_NAME: &str = "litellm_semantic_cache_index"; - -#[derive(Clone)] -struct IndexState { - name: String, - prefix: String, - dimension: Arc>>, - similarity_threshold: f64, -} - -pub struct ValkeySemanticCache< - E: Embedder, - S: CacheCodec, - C = redis::Connection, -> { - connections: Arc>, - embedder: E, - codec: S, - config: ValkeySemanticConfig, - index_dimension: Arc>>, -} - -impl ValkeySemanticCache -where - E: Embedder, - S: CacheCodec, -{ - pub fn new( - url: &str, - embedder: E, - codec: S, - config: ValkeySemanticConfig, - ) -> Result { - Ok(Self { - connections: Arc::new(Connections::open(url, &RedisTopology::Standalone)?), - embedder, - codec, - config, - index_dimension: Arc::new(Mutex::new(None)), - }) - } -} - -impl ValkeySemanticCache -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, - } - } -} - -impl ValkeySemanticCache -where - E: Embedder, - S: CacheCodec + Clone, - C: redis::ConnectionLike + Send + 'static, -{ - pub fn with_embedder(&self, embedder: E2) -> ValkeySemanticCache { - ValkeySemanticCache { - connections: Arc::clone(&self.connections), - embedder, - codec: self.codec.clone(), - config: self.config.clone(), - index_dimension: Arc::clone(&self.index_dimension), - } - } -} - -impl BaseCache for ValkeySemanticCache -where - E: Embedder, - S: CacheCodec, - C: redis::ConnectionLike + Send + 'static, -{ - type Value = CacheEntry; - type Context = SemanticCacheContext; - - fn get_ttl(&self, context: &Self::Context) -> Option { - 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 scope = scope_tag(key); - let response = self.codec.encode(&value)?; - let vector = embedding_bytes(&embedding); - let index = self.index_state(); - self.connections.execute(|connection| { - write_document( - connection, - &index, - &scope, - &prompt, - response, - vector, - self.get_ttl(context), - ) - }) - } - - fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { - let Some(prompt) = prompt_from_context(context) else { - return Ok(None); - }; - let embedding = self.embedder.embed(&prompt, context.metadata.as_ref())?; - let scope = scope_tag(key); - let vector = embedding_bytes(&embedding); - let index = self.index_state(); - let response = self.connections.execute(|connection| { - search_document(connection, &index, &scope, vector, embedding.len()) - })?; - let Some(response) = response else { - return Ok(None); - }; - self.codec.decode(&response).map(Some) - } - - fn async_set_cache( - &self, - key: &str, - value: Self::Value, - context: Self::Context, - ) -> impl Future> + Send { - let key = key.to_owned(); - let prompt = prompt_from_context(&context); - let metadata = context.metadata.clone(); - async move { - let Some(prompt) = prompt else { - return Ok(()); - }; - let embedding = self - .embedder - .async_embed(&prompt, metadata.as_ref()) - .await?; - let connections = Arc::clone(&self.connections); - let index = self.index_state(); - let response = self.codec.encode(&value)?; - let vector = embedding_bytes(&embedding); - let scope = scope_tag(&key); - let ttl = context.ttl; - Connections::run_blocking(connections, move |connection| { - write_document(connection, &index, &scope, &prompt, response, vector, ttl) - }) - .await - } - } - - fn async_get_cache( - &self, - key: &str, - context: &Self::Context, - ) -> impl Future, Error>> + Send { - let key = key.to_owned(); - let prompt = prompt_from_context(context); - let metadata = context.metadata.clone(); - async move { - let Some(prompt) = prompt else { - return Ok(None); - }; - let embedding = self - .embedder - .async_embed(&prompt, metadata.as_ref()) - .await?; - let connections = Arc::clone(&self.connections); - let index = self.index_state(); - Connections::run_blocking(connections, move |connection| { - let scope = scope_tag(&key); - let vector = embedding_bytes(&embedding); - search_document(connection, &index, &scope, vector, embedding.len()) - }) - .await - .and_then(|response| response.map(|bytes| self.codec.decode(&bytes)).transpose()) - } - } - - async fn disconnect(&self) -> Result<(), Error> { - Ok(()) - } - - async fn test_connection(&self) -> Result { - Err(Error::UnsupportedOperation) - } -} - -pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { - if let Some(Value::Array(messages)) = context.messages.as_ref() - && !messages.is_empty() - { - return messages - .iter() - .filter_map(Value::as_object) - .map(message_text) - .collect(); - } - let input = context.input.as_ref()?; - let mut parts = Vec::new(); - collect_input_text(input, &mut parts); - let prompt = parts.join("\n").trim().to_owned(); - (!prompt.is_empty()).then_some(prompt) -} - -fn message_text(message: &serde_json::Map) -> Option { - let content = match message.get("content") { - Some(Value::String(value)) => value.clone(), - Some(Value::Array(parts)) => { - let mut content = String::new(); - for part in parts { - let part = part.as_object()?; - if let Some(text) = part.get("text").and_then(Value::as_str) { - content.push_str(text); - } - } - content - } - _ => String::new(), - }; - Some(format!( - "{content}{}", - search_results_text(message.get("search_results")) - )) -} - -fn search_results_text(value: Option<&Value>) -> String { - let Some(Value::Array(results)) = value else { - return String::new(); - }; - results - .iter() - .filter_map(Value::as_object) - .map(|result| { - let source = result.get("source").and_then(Value::as_str).unwrap_or(""); - let title = result.get("title").and_then(Value::as_str).unwrap_or(""); - let content = result - .get("content") - .and_then(Value::as_array) - .map(|blocks| { - blocks - .iter() - .filter_map(Value::as_object) - .filter_map(|block| block.get("text").and_then(Value::as_str)) - .collect::() - }) - .unwrap_or_default(); - let citations = result - .get("citations") - .filter(|value| !value.is_null()) - .and_then(|value| serde_json::to_string(value).ok()) - .unwrap_or_default(); - format!("{source}{title}{content}{citations}") - }) - .collect() -} - -fn collect_input_text(value: &Value, parts: &mut Vec) { - match value { - Value::String(value) => { - let value = value.trim(); - if !value.is_empty() { - parts.push(value.to_owned()); - } - } - Value::Array(values) => values - .iter() - .for_each(|value| collect_input_text(value, parts)), - Value::Object(object) => { - if let Some(content) = object.get("content").filter(|value| !value.is_null()) { - collect_input_text(content, parts); - return; - } - for key in ["text", "output", "input_text", "output_text"] { - if let Some(Value::String(value)) = object.get(key) { - let value = value.trim(); - if !value.is_empty() { - parts.push(value.to_owned()); - return; - } - } - } - } - _ => {} - } -} - -fn scope_tag(key: &str) -> String { - let digest = Sha256::digest(key.as_bytes()); - digest.iter().map(|byte| format!("{byte:02x}")).collect() -} - -fn embedding_bytes(embedding: &[f32]) -> Vec { - embedding - .iter() - .flat_map(|value| value.to_le_bytes()) - .collect() -} - -fn write_document( - connection: &mut ConnectionRef<'_>, - index: &IndexState, - scope: &str, - prompt: &str, - response: Vec, - vector: Vec, - ttl: Option, -) -> Result<(), Error> { - let dimension = vector.len() / std::mem::size_of::(); - ensure_index( - connection, - &index.name, - &index.prefix, - &index.dimension, - dimension, - )?; - 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) -} - -fn search_document( - connection: &mut ConnectionRef<'_>, - index: &IndexState, - scope: &str, - vector: Vec, - dimension: usize, -) -> Result>, Error> { - ensure_index( - connection, - &index.name, - &index.prefix, - &index.dimension, - dimension, - )?; - 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::(connection) - .map_err(|_| Error::Unavailable)?; - let Some(fields) = search_fields(response)? else { - return Ok(None); - }; - let response = fields - .iter() - .find_map(|(name, value)| (name == "response").then(|| value.clone())) - .ok_or(Error::InvalidEntry)?; - let distance = fields - .iter() - .find_map(|(name, value)| (name == "vector_distance").then(|| value.clone())) - .ok_or(Error::InvalidEntry)?; - let distance = parse_f64(&distance)?; - if 1.0 - distance < index.similarity_threshold { - return Ok(None); - } - Ok(Some(response)) -} - -fn ensure_index( - connection: &mut ConnectionRef<'_>, - index_name: &str, - prefix: &str, - index_dimension: &Mutex>, - 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(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::(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::(connection) - .map_err(|_| Error::Unavailable)?; - let existing = index_dimension_from_info(&info).ok_or(Error::Unavailable)?; - if existing != dimension { - return Err(Error::Unavailable); - } - } - *index_dimension.lock().map_err(|_| Error::Unavailable)? = Some(dimension); - Ok(()) -} - -fn index_dimension_from_info(value: &redis::Value) -> Option { - 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 flattened = values.iter().flat_map(|value| match value { - redis::Value::Array(values) => values.as_slice(), - _ => std::slice::from_ref(value), - }); - let values = flattened.collect::>(); - values.windows(2).find_map(|pair| { - if value_text(pair[0]).as_deref() == Some("dimensions") { - return value_text(pair[1]).and_then(|value| value.parse().ok()); - } - None - }) - }) -} - -type SearchFields = Vec<(String, Vec)>; - -fn search_fields(value: redis::Value) -> Result, 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); - } - let pairs = pairs - .iter() - .map(|pair| { - Ok(( - value_text(&pair[0]).ok_or(Error::InvalidEntry)?, - value_bytes(&pair[1])?, - )) - }) - .collect::, Error>>()?; - Ok(Some(pairs)) -} - -fn parse_i64(value: &redis::Value) -> Result { - value_text(value) - .ok_or(Error::InvalidEntry)? - .parse() - .map_err(|_| Error::InvalidEntry) -} - -fn parse_f64(value: &[u8]) -> Result { - std::str::from_utf8(value) - .map_err(|_| Error::InvalidEntry)? - .parse() - .map_err(|_| Error::InvalidEntry) -} - -fn value_text(value: &redis::Value) -> Option { - 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, 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), - } -} - -#[cfg(test)] -mod tests { - use std::{ - collections::VecDeque, - sync::{Arc, Mutex}, - time::Duration, - }; - - use litellm_cache::{BaseCache, CacheCodec}; - use litellm_cache_response::{ - CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, - }; - use redis_test::MockRedisConnection; - use rstest::rstest; - use serde_json::{Value, json}; - - use super::{ - Embedder, PreparedEmbedding, ValkeySemanticCache, ValkeySemanticConfig, - index_dimension_from_info, prompt_from_context, scope_tag, - }; - - #[derive(Clone)] - struct FixedEmbedder { - vector: Vec, - calls: EmbedderCalls, - } - - type EmbedderCalls = Arc)>>>; - type RecordingCache = - ValkeySemanticCache; - type RecordingSetup = (RecordingCache, Arc>>>, EmbedderCalls); - - impl Embedder for FixedEmbedder { - fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, super::Error> { - self.calls - .lock() - .unwrap() - .push((prompt.into(), metadata.cloned())); - Ok(self.vector.clone()) - } - - async fn async_embed( - &self, - prompt: &str, - metadata: Option<&Value>, - ) -> Result, super::Error> { - self.embed(prompt, metadata) - } - } - - struct RecordingConnection { - requests: Arc>>>, - replies: Mutex>>, - } - - impl RecordingConnection { - fn new(replies: impl IntoIterator>) -> Self { - Self { - requests: Arc::default(), - replies: Mutex::new(replies.into_iter().collect()), - } - } - - fn requests(&self) -> Arc>>> { - Arc::clone(&self.requests) - } - - fn reply(&self) -> redis::RedisResult { - 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 { - self.requests.lock().unwrap().push(command.to_vec()); - self.reply() - } - - fn req_packed_commands( - &mut self, - command: &[u8], - _offset: usize, - count: usize, - ) -> redis::RedisResult> { - 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 - } - } - - fn context( - messages: Option, - input: Option, - ) -> litellm_cache::SemanticCacheContext { - litellm_cache::SemanticCacheContext { - messages, - input, - ..Default::default() - } - } - - #[rstest] - #[case(json!([{"content": "hello"}]), None, Some("hello"))] - #[case(json!([{"content": [{"text": "hello"}, {"text": " world"}]}]), None, Some("hello world"))] - #[case(json!([{"content": ["raw", {"text": "hello"}]}]), None, None)] - #[case(json!([{"search_results": [{"source": "s", "title": "t", "content": [{"text": "c"}], "citations": ["x"]}]}]), None, Some(r#"stc["x"]"#))] - #[case(Value::Array(vec![]), Some(json!(" hello ")), Some("hello"))] - #[case(Value::Array(vec![]), Some(json!([{"content": "first"}, {"text": "second"}])), Some("first\nsecond"))] - #[case(Value::Array(vec![]), Some(json!(" ")), None)] - fn prompt_shapes( - #[case] messages: Value, - #[case] input: Option, - #[case] expected: Option<&str>, - ) { - assert_eq!( - prompt_from_context(&context(Some(messages), input)), - expected.map(str::to_owned) - ); - } - - #[test] - fn scope_tags_are_lowercase_sha256() { - assert_eq!( - scope_tag("key"), - "2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683" - ); - } - - #[test] - fn existing_index_dimension_is_read_from_attributes() { - let info = 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("2".into()), - ]), - ])]), - ]); - assert_eq!(index_dimension_from_info(&info), Some(2)); - } - - #[tokio::test] - async fn unsupported_connection_test_is_reported() { - let cache = ValkeySemanticCache::with_connection( - MockRedisConnection::new([]).assert_all_commands_consumed(), - FixedEmbedder { - vector: vec![1.0, 0.0], - calls: Arc::default(), - }, - ResponseCacheCodec, - ValkeySemanticConfig { - similarity_threshold: 0.8, - index_name: "test".into(), - }, - ); - assert_eq!( - cache.test_connection().await, - Err(super::Error::UnsupportedOperation) - ); - } - - #[tokio::test] - async fn prepared_embedding_returns_its_vector_for_any_prompt() { - let embedding = PreparedEmbedding(vec![1.0, 2.0]); - assert_eq!( - embedding - .async_embed("different prompt", None) - .await - .unwrap(), - vec![1.0, 2.0] - ); - } - - #[test] - fn with_embedder_shares_index_state_and_connections() { - let entry = CacheEntry { - timestamp: Some(1.0), - response: json!({"answer": "ok"}), - }; - let encoded = ResponseCacheCodec.encode(&entry).unwrap(); - let cache = ValkeySemanticCache::with_connection( - RecordingConnection::new([ok(), ok(), Ok(search_hit(encoded, "0.1"))]), - FixedEmbedder { - vector: vec![1.0, 0.0], - calls: Arc::default(), - }, - ResponseCacheCodec, - ValkeySemanticConfig { - similarity_threshold: 0.8, - index_name: "test".into(), - }, - ); - cache - .set_cache("key", entry.clone(), &semantic_context(None)) - .unwrap(); - let prepared = cache.with_embedder(PreparedEmbedding(vec![1.0, 0.0])); - assert_eq!( - prepared.get_cache("key", &semantic_context(None)).unwrap(), - Some(entry) - ); - } - - #[test] - fn missing_prompt_does_not_touch_redis() { - let cache = ValkeySemanticCache::with_connection( - MockRedisConnection::new([]).assert_all_commands_consumed(), - FixedEmbedder { - vector: vec![1.0, 0.0], - calls: Arc::default(), - }, - ResponseCacheCodec, - ValkeySemanticConfig { - similarity_threshold: 0.8, - index_name: "test".into(), - }, - ); - assert_eq!(cache.get_cache("key", &context(None, None)).unwrap(), None); - assert_eq!(cache.get_ttl(&context(None, None)), None); - } - - fn semantic_context(ttl: Option) -> litellm_cache::SemanticCacheContext { - litellm_cache::SemanticCacheContext { - messages: Some(json!([{"role": "user", "content": "hello"}])), - metadata: Some(json!({"source": "test"})), - ttl, - ..Default::default() - } - } - - fn cache_with_recording( - replies: impl IntoIterator>, - vector: Vec, - threshold: f64, - ) -> RecordingSetup { - let connection = RecordingConnection::new(replies); - let requests = connection.requests(); - let calls: EmbedderCalls = Arc::default(); - let cache = ValkeySemanticCache::with_connection( - connection, - FixedEmbedder { - vector, - calls: Arc::clone(&calls), - }, - ResponseCacheCodec, - ValkeySemanticConfig { - similarity_threshold: threshold, - index_name: "test".into(), - }, - ); - (cache, requests, calls) - } - - fn ok() -> redis::RedisResult { - Ok(redis::Value::SimpleString("OK".into())) - } - - fn already_exists() -> redis::RedisResult { - Err(redis::RedisError::from(( - redis::ErrorKind::Io, - "already exists", - ))) - } - - fn info_dimension(dimension: usize) -> 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 as i64), - ]), - ])]), - ]) - } - - fn search_hit(response: Vec, distance: &str) -> redis::Value { - redis::Value::Array(vec![ - redis::Value::Int(1), - redis::Value::BulkString(b"test:document".to_vec()), - redis::Value::Array(vec![ - redis::Value::BulkString(b"response".to_vec()), - redis::Value::BulkString(response), - redis::Value::BulkString(b"vector_distance".to_vec()), - redis::Value::BulkString(distance.as_bytes().to_vec()), - ]), - ]) - } - - fn requests_text(requests: &Arc>>>) -> String { - requests - .lock() - .unwrap() - .iter() - .map(|request| String::from_utf8_lossy(request)) - .collect::>() - .join("\n") - } - - #[test] - fn set_without_ttl_writes_hset_without_expire() { - let (cache, requests, calls) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); - cache - .set_cache( - "key", - CacheEntry { - timestamp: None, - response: json!({"answer": "ok"}), - }, - &semantic_context(None), - ) - .unwrap(); - let text = requests_text(&requests); - assert!(text.contains("FT.CREATE")); - assert!(text.contains("HSET")); - assert!( - text.contains("test:2c70e12b7a0646f92279f427c7b38e7334d8e5389cff167a1dc30e73f826b683:") - ); - assert!(!text.contains("EXPIRE")); - assert_eq!( - *calls.lock().unwrap(), - vec![("hello".into(), Some(json!({"source": "test"})))] - ); - } - - #[test] - fn set_with_ttl_truncates_expire_seconds() { - let (cache, requests, _) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); - cache - .set_cache( - "key", - CacheEntry { - timestamp: None, - response: json!({"answer": "ok"}), - }, - &semantic_context(Some(Duration::from_millis(1900))), - ) - .unwrap(); - let text = requests_text(&requests); - assert!(text.contains("EXPIRE")); - assert!(text.contains("\r\n$1\r\n1\r\n")); - } - - #[test] - fn second_set_skips_create_after_dimension_is_cached() { - let (cache, requests, _) = cache_with_recording([ok()], vec![1.0, 0.0], 0.8); - let context = semantic_context(None); - let entry = CacheEntry { - timestamp: None, - response: json!({"answer": "ok"}), - }; - cache.set_cache("key", entry.clone(), &context).unwrap(); - cache.set_cache("key", entry, &context).unwrap(); - let text = requests_text(&requests); - assert_eq!(text.matches("FT.CREATE").count(), 1); - assert_eq!(text.matches("HSET").count(), 2); - } - - #[test] - fn existing_index_dimension_must_match_embedding() { - let (cache, _, _) = cache_with_recording( - [already_exists(), Ok(info_dimension(2))], - vec![1.0, 0.0], - 0.8, - ); - cache - .set_cache( - "key", - CacheEntry { - timestamp: None, - response: json!({"answer": "ok"}), - }, - &semantic_context(None), - ) - .unwrap(); - - let (cache, _, _) = cache_with_recording( - [already_exists(), Ok(info_dimension(3))], - vec![1.0, 0.0], - 0.8, - ); - assert_eq!( - cache.set_cache( - "key", - CacheEntry { - timestamp: None, - response: json!({"answer": "ok"}), - }, - &semantic_context(None), - ), - Err(super::Error::Unavailable) - ); - } - - #[test] - fn get_applies_threshold_and_decodes_entry() { - let entry = CacheEntry { - timestamp: Some(1.0), - response: json!({"answer": "ok"}), - }; - let encoded = ResponseCacheCodec.encode(&entry).unwrap(); - let (cache, _, _) = cache_with_recording( - [ok(), Ok(search_hit(encoded.clone(), "0.1"))], - vec![1.0, 0.0], - 0.8, - ); - assert_eq!( - cache.get_cache("key", &semantic_context(None)).unwrap(), - Some(entry) - ); - - let (cache, _, _) = - cache_with_recording([ok(), Ok(search_hit(encoded, "0.5"))], vec![1.0, 0.0], 0.8); - assert_eq!( - cache.get_cache("key", &semantic_context(None)).unwrap(), - None - ); - } - - #[test] - fn get_zero_docs_is_a_miss() { - let (cache, _, _) = cache_with_recording( - [ok(), Ok(redis::Value::Array(vec![redis::Value::Int(0)]))], - vec![1.0, 0.0], - 0.8, - ); - assert_eq!( - cache.get_cache("key", &semantic_context(None)).unwrap(), - None - ); - } - - #[rstest] - #[case(redis::Value::Array(vec![ - redis::Value::Int(1), - redis::Value::BulkString(b"document".to_vec()), - redis::Value::Array(vec![ - redis::Value::BulkString(b"vector_distance".to_vec()), - redis::Value::BulkString(b"0.1".to_vec()), - ]), - ]))] - #[case(redis::Value::Array(vec![ - redis::Value::Int(1), - redis::Value::BulkString(b"document".to_vec()), - redis::Value::Array(vec![ - redis::Value::BulkString(b"response".to_vec()), - redis::Value::BulkString(b"not-json".to_vec()), - redis::Value::BulkString(b"vector_distance".to_vec()), - redis::Value::BulkString(b"abc".to_vec()), - ]), - ]))] - fn malformed_entries_are_invalid(#[case] search: redis::Value) { - let (cache, _, _) = cache_with_recording([ok(), Ok(search)], vec![1.0, 0.0], 0.8); - assert_eq!( - cache.get_cache("key", &semantic_context(None)), - Err(super::Error::InvalidEntry) - ); - } - - #[test] - fn response_cache_turns_invalid_entries_into_misses() { - let (cache, _, _) = cache_with_recording( - [ - ok(), - Ok(redis::Value::Array(vec![ - redis::Value::Int(1), - redis::Value::BulkString(b"document".to_vec()), - redis::Value::Array(vec![ - redis::Value::BulkString(b"response".to_vec()), - redis::Value::BulkString(b"not-json".to_vec()), - redis::Value::BulkString(b"vector_distance".to_vec()), - redis::Value::BulkString(b"0.1".to_vec()), - ]), - ])), - ], - vec![1.0, 0.0], - 0.8, - ); - let service = ResponseCache::new(Arc::new(cache)); - let request = ResponseCacheRequest { - key: CacheKeyInput { - preset: Some("key".into()), - ..Default::default() - }, - context: semantic_context(None), - ..ResponseCacheRequest::new(CacheKeyInput::default()) - }; - assert_eq!(service.lookup(&request, Duration::ZERO).unwrap(), None); - } - - #[tokio::test] - async fn async_set_and_get_use_shared_document_helpers() { - let entry = CacheEntry { - timestamp: Some(1.0), - response: json!({"answer": "ok"}), - }; - let encoded = ResponseCacheCodec.encode(&entry).unwrap(); - let (cache, requests, calls) = cache_with_recording( - [ok(), ok(), ok(), Ok(search_hit(encoded, "0.1"))], - vec![1.0, 0.0], - 0.8, - ); - let context = semantic_context(Some(Duration::from_millis(1900))); - cache - .async_set_cache("key", entry.clone(), context.clone()) - .await - .unwrap(); - assert_eq!( - cache.async_get_cache("key", &context).await.unwrap(), - Some(entry) - ); - let text = requests_text(&requests); - assert!(text.contains("FT.CREATE")); - assert!(text.contains("HSET")); - assert!(text.contains("EXPIRE")); - assert_eq!(calls.lock().unwrap().len(), 2); - } -} +pub use cache::ValkeySemanticCache; +pub use config::{DEFAULT_INDEX_NAME, ValkeySemanticConfig}; diff --git a/litellm-rust/crates/cache-valkey-semantic/src/search.rs b/litellm-rust/crates/cache-valkey-semantic/src/search.rs new file mode 100644 index 00000000000..3b4fae84317 --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/src/search.rs @@ -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 { + embedding + .iter() + .flat_map(|value| value.to_le_bytes()) + .collect() +} + +/// `HSET` a fresh `:` 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, + vector: Vec, + ttl: Option, +) -> Result<(), Error> { + ensure_index(connection, index, vector.len() / size_of::())?; + 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, +) -> Result>, Error> { + ensure_index(connection, index, vector.len() / size_of::())?; + 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::(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)>; + +fn search_fields(value: redis::Value) -> Result, 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::, Error>>() + .map(Some) +} + +fn parse_i64(value: &redis::Value) -> Result { + value_text(value) + .ok_or(Error::InvalidEntry)? + .parse() + .map_err(|_| Error::InvalidEntry) +} + +fn parse_f64(value: &[u8]) -> Result { + std::str::from_utf8(value) + .map_err(|_| Error::InvalidEntry)? + .parse() + .map_err(|_| Error::InvalidEntry) +} + +pub(crate) fn value_text(value: &redis::Value) -> Option { + 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, 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), + } +} diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs b/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs new file mode 100644 index 00000000000..1c693d0875d --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/tests/cache.rs @@ -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>>>; +type RecordingCache = ValkeySemanticCache, 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::>() + .join("\n") + } +} + +fn config() -> ValkeySemanticConfig { + ValkeySemanticConfig { + similarity_threshold: 0.8, + index_name: "test".into(), + } +} + +fn recording(replies: impl IntoIterator>) -> 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 { + Ok(redis::Value::SimpleString("OK".into())) +} + +fn already_exists() -> redis::RedisResult { + 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 { + 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, + #[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::>(); + assert_eq!( + prompts, + expected.into_iter().map(str::to_owned).collect::>() + ); + 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, + #[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, + 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, 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, + #[case] value: Option, + #[case] similarity: Option, + #[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] + ); +} diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/contract.rs b/litellm-rust/crates/cache-valkey-semantic/tests/contract.rs new file mode 100644 index 00000000000..71081fca92b --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/tests/contract.rs @@ -0,0 +1,62 @@ +//! `overwrite_replaces` does not apply: like Python, every write is a new `:` +//! 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, 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; +} diff --git a/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs new file mode 100644 index 00000000000..2c6ff06b134 --- /dev/null +++ b/litellm-rust/crates/cache-valkey-semantic/tests/support/mod.rs @@ -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)>>>; + +/// 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>, + 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, 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, Error> { + self.embed(prompt, metadata) + } +} + +struct FakeIndex { + prefix: Vec, + dims: usize, + vector_field: String, +} + +#[derive(Default)] +struct SearchState { + indexes: HashMap, + hashes: BTreeMap, BTreeMap>>, +} + +/// 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>, +} + +impl FakeSearch { + fn run(&self, args: Vec>) -> redis::RedisResult { + 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::>(); + 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 { + 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> { + let replies = parse_commands(commands) + .into_iter() + .map(|args| self.run(args)) + .collect::>>()?; + 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 { + 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::(); + let norm = |vector: &[f32]| { + vector + .iter() + .map(|value| f64::from(*value).powi(2)) + .sum::() + .sqrt() + }; + dot / (norm(left) * norm(right)) +} + +/// Splits a packed RESP request into each command's arguments. +fn parse_commands(mut bytes: &[u8]) -> Vec>> { + 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::().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>>>, + replies: Mutex>>, +} + +impl RecordingConnection { + pub fn new(replies: impl IntoIterator>) -> Self { + Self { + requests: Arc::default(), + replies: Mutex::new(replies.into_iter().collect()), + } + } + + pub fn requests(&self) -> Arc>>> { + Arc::clone(&self.requests) + } + + fn reply(&self) -> redis::RedisResult { + 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 { + self.requests.lock().unwrap().push(command.to_vec()); + self.reply() + } + + fn req_packed_commands( + &mut self, + command: &[u8], + _offset: usize, + count: usize, + ) -> redis::RedisResult> { + 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 + } +} diff --git a/litellm-rust/crates/cache/Cargo.toml b/litellm-rust/crates/cache/Cargo.toml index 0c504ab727a..f18dbd9cb26 100644 --- a/litellm-rust/crates/cache/Cargo.toml +++ b/litellm-rust/crates/cache/Cargo.toml @@ -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] diff --git a/litellm-rust/crates/cache/src/base_cache.rs b/litellm-rust/crates/cache/src/base_cache.rs index 5c10e7fd5c3..506d944b0e8 100644 --- a/litellm-rust/crates/cache/src/base_cache.rs +++ b/litellm-rust/crates/cache/src/base_cache.rs @@ -122,36 +122,4 @@ pub trait BaseCache: Send + Sync { ) -> impl Future> + Send { self.async_set_cache(key, value, context) } - - fn disconnect(&self) -> impl Future> + Send; - - fn test_connection(&self) -> impl Future> + 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); - } } diff --git a/litellm-rust/crates/cache/src/cache_type.rs b/litellm-rust/crates/cache/src/cache_type.rs index f0a97c04fd5..22d8c8c7cb5 100644 --- a/litellm-rust/crates/cache/src/cache_type.rs +++ b/litellm-rust/crates/cache/src/cache_type.rs @@ -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) - ); - } -} diff --git a/litellm-rust/crates/cache/src/capabilities.rs b/litellm-rust/crates/cache/src/capabilities.rs index f7307e5c7bd..ac9d8fc8764 100644 --- a/litellm-rust/crates/cache/src/capabilities.rs +++ b/litellm-rust/crates/cache/src/capabilities.rs @@ -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, } +#[derive(Clone, Debug, PartialEq)] +pub struct PushOperation { + pub key: String, + pub values: Vec, +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct PopOperation { + pub key: String, + pub count: Option, +} + +/// `disconnect`, for backends whose Python class releases connections or clients. +pub trait DisconnectCache: BaseCache { + fn disconnect(&self) -> impl Future> + Send; +} + +/// `test_connection`, for backends whose Python class overrides the base `NotImplementedError`. +pub trait ConnectionCache: BaseCache { + fn test_connection(&self) -> impl Future> + Send; +} + +/// `sync_ping` and `ping`. +pub trait PingCache: BaseCache { + fn sync_ping(&self) -> Result; + + fn ping(&self) -> impl Future> + 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)>, + ) -> impl Future> + 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, + ) -> impl Future> + Send; +} + pub trait FlushCache: BaseCache { fn flush_cache(&self) -> Result<(), Error>; @@ -61,18 +106,79 @@ pub trait FlushCache: BaseCache { } } -pub trait CounterCache: BaseCache { +/// `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; + /// `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> + 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, + ) -> impl Future, 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>, Error>; + + fn async_batch_get_counts( + &self, + keys: Vec, + ) -> impl Future>, 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; + + fn async_increment_with_floor( + &self, + key: &str, + amount: i64, + ttl: Duration, + ) -> impl Future> + Send; + + fn async_set_max( + &self, + key: &str, + value: f64, + ttl: Option, + ) -> impl Future> + Send; } pub trait ClaimCache: BaseCache @@ -105,6 +211,17 @@ pub trait TtlCache: BaseCache { ) -> impl Future, 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, + ) -> impl Future> + 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, ) -> impl Future> + 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, + max_len: usize, + ) -> impl Future> + Send; + + fn async_rpush_pipeline( + &self, + operations: Vec>, + ) -> impl Future, Error>> + Send; + fn async_lpop( &self, key: &str, count: Option, ) -> impl Future> + Send; + + fn async_lpop_pipeline( + &self, + operations: Vec, + ) -> impl Future, Error>> + Send; } pub trait ScanCache: BaseCache { diff --git a/litellm-rust/crates/cache/src/dual.rs b/litellm-rust/crates/cache/src/dual.rs index d68d4b2b69f..c4bc953239d 100644 --- a/litellm-rust/crates/cache/src/dual.rs +++ b/litellm-rust/crates/cache/src/dual.rs @@ -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 { write_policy: WritePolicy, remote_failure_policy: RemoteFailurePolicy, promotion_ttl: Option, + delete_batch_size: usize, } +/// `DEFAULT_MAX_REDIS_BATCH_CACHE_SIZE`. +pub const DEFAULT_DELETE_BATCH_SIZE: usize = 1000; + impl DualCache { pub fn new(l1: Arc, l2: Arc) -> Self { Self { @@ -44,6 +48,14 @@ impl DualCache { 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 { - self.l2.test_connection().await - } } impl BatchCache for DualCache @@ -320,26 +323,146 @@ where } } +impl DualCache +where + C: CacheContext, + L1: BaseCache, +{ + /// 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 { + let value = self.l1.get_cache(key, context)?.unwrap_or(0.0) + amount; + self.l1.set_cache(key, value, context)?; + Ok(value) + } +} + impl CounterCache for DualCache where C: CacheContext, L1: BaseCache, - L2: CounterCache, + L2: CounterCache, { fn increment_cache(&self, key: &str, amount: f64, context: C) -> Result { + 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 { + async fn async_increment( + &self, + key: &str, + amount: f64, + context: C, + refresh_ttl: bool, + ) -> Result { + 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, + ) -> Result, 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 SetCache for DualCache +where + V: Clone + Send + Sync + 'static, + C: CacheContext, + S: Clone + Send + Sync + 'static, + L1: SetCache, + L2: SetCache, +{ + type SetValue = S; + type SetResult = (); + + async fn async_set_cache_sadd( + &self, + key: &str, + values: Vec, + ttl: Option, + ) -> 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 BulkDeleteCache for DualCache +where + V: Clone + Send + Sync + 'static, + C: CacheContext, + L1: DeleteCache, + L2: BulkDeleteCache, +{ + async fn delete_cache_keys(&self, keys: Vec) -> Result { + 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 TtlCache for DualCache +where + V: Clone + Send + Sync + 'static, + C: CacheContext, + L1: TtlCache, + L2: TtlCache, +{ + async fn async_get_ttl(&self, key: &str) -> Result, Error> { + match self.l1.async_get_ttl(key).await? { + Some(ttl) => Ok(Some(ttl)), + None => self.l2.async_get_ttl(key).await, + } + } } impl ClaimCache for DualCache diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index 8364c635e3a..55d8a5fa28b 100644 --- a/litellm-rust/crates/cache/src/lib.rs +++ b/litellm-rust/crates/cache/src/lib.rs @@ -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; diff --git a/litellm-rust/crates/cache/src/semantic.rs b/litellm-rust/crates/cache/src/semantic.rs new file mode 100644 index 00000000000..c88706a213b --- /dev/null +++ b/litellm-rust/crates/cache/src/semantic.rs @@ -0,0 +1,274 @@ +//! The embedding and prompt contract every semantic backend shares. +//! +//! Python's semantic caches all read their prompt through `get_str_from_messages`, and +//! `RedisSemanticCache._get_prompt_from_kwargs` (inherited by Valkey) adds Responses API +//! `input`. Qdrant reads messages only. Each backend picks one of the two extractors here. + +use std::{future::Future, io}; + +use serde::Serialize; +use serde_json::{ + Value, + ser::{CharEscape, Formatter, Serializer}, +}; + +use crate::{BaseCache, Error, SemanticCacheContext}; + +/// Turns a prompt into the vector a semantic backend stores and searches with. +/// +/// `metadata` is the request metadata, which a host embedder may route on. Hosts that can only +/// embed asynchronously keep the default `embed`; backends that serve sync calls through a +/// runtime then block on `async_embed` instead. +pub trait Embedder: Send + Sync + 'static { + fn embed(&self, _prompt: &str, _metadata: Option<&Value>) -> Result, Error> { + Err(Error::UnsupportedOperation) + } + + fn async_embed( + &self, + prompt: &str, + metadata: Option<&Value>, + ) -> impl Future, Error>> + Send; +} + +/// One semantic read: the cached value, if any, and the similarity Python's backend writes to +/// `metadata["semantic-similarity"]`. `similarity` is `None` when the backend reports none. +#[derive(Clone, Debug, PartialEq)] +pub struct SemanticLookup { + pub value: Option, + pub similarity: Option, +} + +impl SemanticLookup { + /// A read that found nothing, with the similarity Python records for it. + pub fn miss(similarity: Option) -> Self { + Self { + value: None, + similarity, + } + } +} + +/// A semantic backend's read that also reports the similarity of the closest cached prompt, +/// the value Python stamps onto the request metadata as `semantic-similarity`. +pub trait SemanticCache: BaseCache { + fn get_cache_with_similarity( + &self, + key: &str, + context: &Self::Context, + ) -> Result, Error>; + + fn async_get_cache_with_similarity( + &self, + key: &str, + context: &Self::Context, + ) -> impl Future, Error>> + Send; +} + +/// An embedding computed ahead of time, for callers that already hold the vector. +#[derive(Clone, Debug, PartialEq)] +pub struct PreparedEmbedding(pub Vec); + +impl Embedder for PreparedEmbedding { + fn embed(&self, _prompt: &str, _metadata: Option<&Value>) -> Result, Error> { + Ok(self.0.clone()) + } + + async fn async_embed( + &self, + _prompt: &str, + _metadata: Option<&Value>, + ) -> Result, Error> { + Ok(self.0.clone()) + } +} + +/// `get_str_from_messages`: every message's text content followed by its search results. +pub fn str_from_messages(messages: &[Value]) -> String { + let mut text = String::new(); + for message in messages.iter().filter_map(Value::as_object) { + match message.get("content") { + Some(Value::String(content)) => text.push_str(content), + Some(Value::Array(parts)) => { + for part in parts { + if let Some(part_text) = part.get("text").and_then(Value::as_str) { + text.push_str(part_text); + } + } + } + _ => {} + } + push_search_results_text(&mut text, message.get("search_results")); + } + text +} + +/// The messages prompt Qdrant embeds: `None` when the request carries no messages. +pub fn prompt_from_messages(context: &SemanticCacheContext) -> Option { + let messages = context.messages.as_ref()?.as_array()?; + (!messages.is_empty()).then(|| str_from_messages(messages)) +} + +/// `RedisSemanticCache._get_prompt_from_kwargs`: chat messages first, then the text parts of a +/// Responses API `input`. `None` when neither yields a prompt. +pub fn prompt_from_context(context: &SemanticCacheContext) -> Option { + if let Some(messages) = context.messages.as_ref().and_then(Value::as_array) + && !messages.is_empty() + { + return Some(str_from_messages(messages)); + } + let input = context.input.as_ref()?; + let mut parts = Vec::new(); + collect_input_text(input, &mut parts); + let prompt = python_strip(&parts.join("\n")).to_owned(); + (!prompt.is_empty()).then_some(prompt) +} + +/// `extract_search_results_text`. +fn push_search_results_text(text: &mut String, search_results: Option<&Value>) { + let Some(Value::Array(results)) = search_results else { + return; + }; + for result in results.iter().filter_map(Value::as_object) { + 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.iter().filter_map(Value::as_object) { + if let Some(value) = block.get("text").and_then(Value::as_str) { + text.push_str(value); + } + } + } + if let Some(citations) = result.get("citations").filter(|value| !value.is_null()) { + text.push_str(&compact_json(citations)); + } + } +} + +fn collect_input_text(value: &Value, parts: &mut Vec) { + match value { + Value::String(text) => { + push_trimmed(text, parts); + } + 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) + && push_trimmed(text, parts) + { + return; + } + } + } + _ => {} + } +} + +/// Pushes `text` stripped as Python's `str.strip` does, reporting whether anything was left. +fn push_trimmed(text: &str, parts: &mut Vec) -> bool { + let trimmed = python_strip(text); + if trimmed.is_empty() { + return false; + } + parts.push(trimmed.to_owned()); + true +} + +/// `str.strip()`: Python's whitespace also covers the ASCII information separators. +fn python_strip(text: &str) -> &str { + text.trim_matches(|character: char| { + character.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&character) + }) +} + +/// `json.dumps(value, separators=(",", ":"))`: compact, key insertion order, `ensure_ascii`. +fn compact_json(value: &Value) -> String { + let mut output = Vec::new(); + // Serializing a `Value` into memory cannot fail. + let _ = value.serialize(&mut Serializer::with_formatter(&mut output, AsciiFormatter)); + String::from_utf8(output).unwrap_or_default() +} + +struct AsciiFormatter; + +impl Formatter for AsciiFormatter { + fn write_string_fragment(&mut self, writer: &mut W, fragment: &str) -> io::Result<()> + where + W: ?Sized + io::Write, + { + let mut start = 0; + for (index, character) in fragment.char_indices() { + if character.is_ascii() && character != '\u{7f}' { + continue; + } + writer.write_all(&fragment.as_bytes()[start..index])?; + let mut units = [0; 2]; + for unit in character.encode_utf16(&mut units) { + write!(writer, "\\u{unit:04x}")?; + } + start = index + character.len_utf8(); + } + writer.write_all(&fragment.as_bytes()[start..]) + } + + fn write_f64(&mut self, writer: &mut W, value: f64) -> io::Result<()> + where + W: ?Sized + io::Write, + { + writer.write_all(python_float_repr(value).as_bytes()) + } + + fn write_char_escape(&mut self, writer: &mut W, escape: CharEscape) -> io::Result<()> + where + W: ?Sized + io::Write, + { + match escape { + CharEscape::AsciiControl(byte) => write!(writer, "\\u{byte:04x}"), + escape => serde_json::ser::CompactFormatter.write_char_escape(writer, escape), + } + } +} + +/// `repr(float)`: the shortest round-trip digits, positional between `1e-4` and `1e16`, and +/// otherwise scientific with a signed exponent of at least two digits. +fn python_float_repr(value: f64) -> String { + // `{:e}` yields the shortest round-trip digits, e.g. `1.5e-7`. + let scientific = format!("{value:e}"); + let (mantissa, exponent) = scientific.split_once('e').unwrap_or((&scientific, "0")); + let exponent: i32 = exponent.parse().unwrap_or(0); + let (sign, mantissa) = mantissa + .strip_prefix('-') + .map_or(("", mantissa), |rest| ("-", rest)); + let digits = mantissa.replace('.', ""); + if !(-4..16).contains(&exponent) { + let fraction = &digits[1..]; + let mantissa = if fraction.is_empty() { + digits[..1].to_owned() + } else { + format!("{}.{fraction}", &digits[..1]) + }; + let exponent_sign = if exponent < 0 { '-' } else { '+' }; + return format!("{sign}{mantissa}e{exponent_sign}{:02}", exponent.abs()); + } + let point = exponent + 1; + let positional = if point <= 0 { + format!("0.{}{digits}", "0".repeat(point.unsigned_abs() as usize)) + } else if point as usize >= digits.len() { + format!("{digits}{}.0", "0".repeat(point as usize - digits.len())) + } else { + let (whole, fraction) = digits.split_at(point as usize); + format!("{whole}.{fraction}") + }; + format!("{sign}{positional}") +} diff --git a/litellm-rust/crates/cache/tests/cache_type.rs b/litellm-rust/crates/cache/tests/cache_type.rs new file mode 100644 index 00000000000..24aaba8e5fd --- /dev/null +++ b/litellm-rust/crates/cache/tests/cache_type.rs @@ -0,0 +1,56 @@ +use litellm_cache::CacheType; +use rstest::rstest; + +#[rstest] +#[case(CacheType::Local, "local")] +#[case(CacheType::Redis, "redis")] +#[case(CacheType::RedisSemantic, "redis-semantic")] +#[case(CacheType::ValkeySemantic, "valkey-semantic")] +#[case(CacheType::S3, "s3")] +#[case(CacheType::Disk, "disk")] +#[case(CacheType::QdrantSemantic, "qdrant-semantic")] +#[case(CacheType::AzureBlob, "azure-blob")] +#[case(CacheType::Gcs, "gcs")] +fn every_python_cache_type_has_one_round_trip_identity( + #[case] cache_type: CacheType, + #[case] name: &str, +) { + assert_eq!(cache_type.as_python_name(), name); + assert_eq!(CacheType::from_python_name(name), Some(cache_type)); + assert_eq!( + serde_json::to_value(cache_type).unwrap(), + serde_json::Value::from(name) + ); + assert_eq!( + CacheType::ALL + .iter() + .filter(|candidate| candidate.as_python_name() == name) + .count(), + 1 + ); +} + +#[rstest] +fn python_cache_types_are_listed_in_python_order() { + assert_eq!( + CacheType::ALL.map(CacheType::as_python_name), + [ + "local", + "redis", + "redis-semantic", + "valkey-semantic", + "s3", + "disk", + "qdrant-semantic", + "azure-blob", + "gcs", + ] + ); +} + +#[rstest] +#[case::unknown("memcached")] +#[case::case_sensitive("Redis")] +fn unknown_python_names_have_no_cache_type(#[case] name: &str) { + assert_eq!(CacheType::from_python_name(name), None); +} diff --git a/litellm-rust/crates/cache/tests/caching.rs b/litellm-rust/crates/cache/tests/caching.rs index 36307ac9b33..baf968f4659 100644 --- a/litellm-rust/crates/cache/tests/caching.rs +++ b/litellm-rust/crates/cache/tests/caching.rs @@ -1,15 +1,25 @@ use std::{sync::Mutex, time::Duration}; use litellm_cache::{ - BaseCache, CacheConnectionResult, CacheContext, Error, ExactCacheContext, SemanticCacheContext, - get_cache, + BaseCache, CacheContext, CounterCache, Error, ExactCacheContext, IncrementOperation, + SemanticCacheContext, get_cache, }; +use rstest::{fixture, rstest}; +use serde_json::json; struct TestCache { default_ttl: Duration, writes: Mutex>, } +#[fixture] +fn cache() -> TestCache { + TestCache { + default_ttl: Duration::from_secs(60), + writes: Mutex::default(), + } +} + #[derive(Clone)] struct SemanticContext { ttl: Option, @@ -46,14 +56,6 @@ impl BaseCache for SemanticCache { fn get_cache(&self, _: &str, context: &Self::Context) -> Result, Error> { Ok((context.query == "matching prompt").then(|| "semantic hit".into())) } - - async fn disconnect(&self) -> Result<(), Error> { - Ok(()) - } - - async fn test_connection(&self) -> Result { - unreachable!() - } } impl BaseCache for TestCache { @@ -87,70 +89,124 @@ impl BaseCache for TestCache { fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { Ok(None) } +} - async fn disconnect(&self) -> Result<(), Error> { +/// Records every `async_increment` so the default pipeline's calls are observable. +#[derive(Default)] +struct RecordingCounter { + increments: Mutex>, + total: Mutex, +} + +impl BaseCache for RecordingCounter { + type Value = f64; + type Context = ExactCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache(&self, _: &str, _: f64, _: &ExactCacheContext) -> Result<(), Error> { Ok(()) } - async fn test_connection(&self) -> Result { - unreachable!() + fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { + Ok(None) } } -#[test] -fn ttl_uses_default_and_allows_per_call_override() { - let cache = TestCache { - default_ttl: Duration::from_secs(60), - writes: Mutex::default(), - }; - assert_eq!( - cache.get_ttl(&ExactCacheContext::default()), - Some(Duration::from_secs(60)) - ); +impl CounterCache for RecordingCounter { + fn increment_cache(&self, key: &str, amount: f64, _: ExactCacheContext) -> Result { + if key == "unavailable" { + return Err(Error::Unavailable); + } + let mut total = self.total.lock().unwrap(); + *total += amount; + Ok(*total) + } + + async fn async_increment( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + refresh_ttl: bool, + ) -> Result { + self.increments + .lock() + .unwrap() + .push((key.into(), amount, context.clone(), refresh_ttl)); + self.increment_cache(key, amount, context) + } +} + +fn operation(key: &str, amount: f64, ttl: Option) -> IncrementOperation { + IncrementOperation { + key: key.into(), + amount, + ttl: ttl.map(Duration::from_secs), + } +} + +#[rstest] +#[case::default_ttl(None, Some(60))] +#[case::per_call_override(Some(5), Some(5))] +fn ttl_uses_default_and_allows_per_call_override( + cache: TestCache, + #[case] ttl: Option, + #[case] expected: Option, +) { assert_eq!( cache.get_ttl(&ExactCacheContext { - ttl: Some(Duration::from_secs(5)), + ttl: ttl.map(Duration::from_secs), }), - Some(Duration::from_secs(5)) + expected.map(Duration::from_secs) ); } -#[test] -fn associated_context_preserves_backend_specific_lookup_inputs() { +#[rstest] +#[case::matching("matching prompt", Some("semantic hit"))] +#[case::other("other prompt", None)] +fn associated_context_preserves_backend_specific_lookup_inputs( + #[case] query: &str, + #[case] expected: Option<&str>, +) { let context = SemanticContext { ttl: None, - query: "matching prompt".into(), + query: query.into(), }; assert_eq!( get_cache(&SemanticCache, "shared-key", &context).unwrap(), - Some("semantic hit".into()) + expected.map(String::from) ); } -#[test] -fn semantic_context_with_ttl_preserves_lookup_inputs() { +#[rstest] +#[case::set(None, Some(30))] +#[case::replaced(Some(10), Some(20))] +#[case::cleared(Some(10), None)] +fn semantic_context_with_ttl_only_replaces_ttl( + #[case] initial: Option, + #[case] updated: Option, +) { let context = SemanticCacheContext { - input: Some(serde_json::json!("text")), - messages: Some(serde_json::json!([{"role": "user", "content": "hi"}])), - metadata: Some(serde_json::json!({"key": "value"})), + input: Some(json!({"input": "hello"})), + messages: Some(json!([{"role": "user", "content": "hello"}])), + metadata: Some(json!({"tenant": "team"})), scope: Some("scope".into()), - ttl: None, + ttl: initial.map(Duration::from_secs), }; - let updated = context.with_ttl(Some(Duration::from_secs(30))); - assert_eq!(updated.ttl(), Some(Duration::from_secs(30))); - assert_eq!(updated.input, context.input); - assert_eq!(updated.messages, context.messages); - assert_eq!(updated.metadata, context.metadata); - assert_eq!(updated.scope, context.scope); - assert_eq!(context.with_ttl(None).ttl(), None); + let result = context.with_ttl(updated.map(Duration::from_secs)); + assert_eq!(result.ttl(), updated.map(Duration::from_secs)); + assert_eq!(result.input, context.input); + assert_eq!(result.messages, context.messages); + assert_eq!(result.metadata, context.metadata); + assert_eq!(result.scope, context.scope); } +#[rstest] #[tokio::test] -async fn default_batch_operations_use_async_writes_and_stop_on_failure() { - let cache = TestCache { - default_ttl: Duration::from_secs(60), - writes: Mutex::default(), - }; +async fn default_batch_operations_use_async_writes_and_stop_on_failure(cache: TestCache) { let entry = String::from("cached"); let context = ExactCacheContext { ttl: Some(Duration::from_secs(5)), @@ -180,3 +236,101 @@ async fn default_batch_operations_use_async_writes_and_stop_on_failure() { ] ); } + +#[rstest] +#[tokio::test] +async fn default_async_increment_delegates_to_the_sync_increment() { + struct SyncOnly; + + impl BaseCache for SyncOnly { + type Value = f64; + type Context = ExactCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache(&self, _: &str, _: f64, _: &ExactCacheContext) -> Result<(), Error> { + Ok(()) + } + + fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { + Ok(None) + } + } + + impl CounterCache for SyncOnly { + fn increment_cache( + &self, + _: &str, + amount: f64, + _: ExactCacheContext, + ) -> Result { + Ok(amount * 10.0) + } + } + + for refresh_ttl in [false, true] { + assert_eq!( + SyncOnly + .async_increment("key", 2.0, ExactCacheContext::default(), refresh_ttl) + .await, + Ok(20.0) + ); + } +} + +#[rstest] +#[case::empty(Vec::new(), Vec::new())] +#[case::one(vec![operation("a", 1.0, Some(10))], vec![1.0])] +#[case::in_order( + vec![operation("a", 1.0, Some(10)), operation("b", 2.5, None), operation("a", -0.5, Some(20))], + vec![1.0, 3.5, 3.0], +)] +#[tokio::test] +async fn default_increment_pipeline_increments_each_operation_in_order( + #[case] operations: Vec, + #[case] expected: Vec, +) { + let cache = RecordingCounter::default(); + assert_eq!( + cache.async_increment_pipeline(operations.clone()).await, + Ok(expected) + ); + assert_eq!( + *cache.increments.lock().unwrap(), + operations + .into_iter() + .map(|operation| ( + operation.key, + operation.amount, + ExactCacheContext { ttl: operation.ttl }, + false, + )) + .collect::>() + ); +} + +#[rstest] +#[tokio::test] +async fn default_increment_pipeline_stops_at_the_first_failure() { + let cache = RecordingCounter::default(); + assert_eq!( + cache + .async_increment_pipeline(vec![ + operation("a", 1.0, None), + operation("unavailable", 1.0, None), + operation("skipped", 1.0, None), + ]) + .await, + Err(Error::Unavailable) + ); + let keys = cache + .increments + .lock() + .unwrap() + .iter() + .map(|(key, ..)| key.clone()) + .collect::>(); + assert_eq!(keys, ["a", "unavailable"]); +} diff --git a/litellm-rust/crates/cache/tests/codec.rs b/litellm-rust/crates/cache/tests/codec.rs index e24545caad6..5320545fd49 100644 --- a/litellm-rust/crates/cache/tests/codec.rs +++ b/litellm-rust/crates/cache/tests/codec.rs @@ -1,6 +1,7 @@ use std::collections::BTreeMap; use litellm_cache::{CacheCodec, Error, JsonCodec}; +use rstest::rstest; use serde::{Deserialize, Serialize}; use serde_json::json; @@ -10,7 +11,7 @@ struct RoutingState { cooldown_seconds: u64, } -#[test] +#[rstest] fn json_codec_round_trips_typed_domain_values() { let codec = JsonCodec::::new(); let value = RoutingState { @@ -25,15 +26,15 @@ fn json_codec_round_trips_typed_domain_values() { ); } -#[test] -fn json_codec_rejects_malformed_and_wrongly_typed_entries() { +#[rstest] +#[case::malformed(b"not json")] +#[case::wrongly_typed(br#"{"deployment":12}"#)] +fn json_codec_rejects_malformed_and_wrongly_typed_entries(#[case] bytes: &[u8]) { let codec = JsonCodec::::new(); - for bytes in [b"not json".as_slice(), br#"{"deployment":12}"#.as_slice()] { - assert_eq!(codec.decode(bytes).unwrap_err(), Error::InvalidEntry); - } + assert_eq!(codec.decode(bytes).unwrap_err(), Error::InvalidEntry); } -#[test] +#[rstest] fn json_codec_propagates_encoding_errors() { let codec = JsonCodec::>::new(); let value = BTreeMap::from([((1, 2), "invalid JSON object key".into())]); diff --git a/litellm-rust/crates/cache/tests/dual.rs b/litellm-rust/crates/cache/tests/dual.rs index e7e8927f8d0..87dd76388ef 100644 --- a/litellm-rust/crates/cache/tests/dual.rs +++ b/litellm-rust/crates/cache/tests/dual.rs @@ -1,12 +1,15 @@ use std::{ + collections::HashMap, sync::{Arc, Mutex}, time::Duration, }; use litellm_cache::{ - BaseCache, BatchCache, CacheConnectionResult, ClaimCache, CounterCache, DeleteCache, DualCache, - Error, ExactCacheContext, FlushCache, ReadPolicy, RemoteFailurePolicy, WritePolicy, + BaseCache, BatchCache, BatchEntry, BulkDeleteCache, ClaimCache, CounterCache, DeleteCache, + DualCache, Error, ExactCacheContext, FlushCache, IncrementOperation, ReadPolicy, + RemoteFailurePolicy, SetCache, TtlCache, WritePolicy, }; +use rstest::{fixture, rstest}; struct TestCache { value: Mutex>, @@ -41,14 +44,6 @@ where fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { Ok(self.value.lock().unwrap().clone()) } - - async fn disconnect(&self) -> Result<(), Error> { - Ok(()) - } - - async fn test_connection(&self) -> Result { - unreachable!() - } } impl BatchCache for TestCache where V: Clone + Send + Sync + 'static {} @@ -111,7 +106,7 @@ where } } -#[test] +#[rstest] fn failed_l2_increment_leaves_l1_unchanged() { let l1 = Arc::new(TestCache::new(Some(10.0), false)); let cache = DualCache::new(l1.clone(), Arc::new(TestCache::new(Some(20.0), true))); @@ -127,7 +122,7 @@ fn failed_l2_increment_leaves_l1_unchanged() { ); } -#[test] +#[rstest] fn claim_uses_l1_fallback_without_overwriting_an_eligible_winner() { let l1 = Arc::new(TestCache::new(Some("first".to_string()), false)); let cache = DualCache::new(l1, Arc::new(TestCache::new(None, true))) @@ -193,14 +188,6 @@ impl BaseCache for SyncPanics { } Ok(()) } - - async fn disconnect(&self) -> Result<(), Error> { - Ok(()) - } - - async fn test_connection(&self) -> Result { - unreachable!() - } } impl BatchCache for SyncPanics { @@ -208,11 +195,11 @@ impl BatchCache for SyncPanics { &self, keys: Vec, context: ExactCacheContext, - ) -> Result>, Error> { + ) -> Result>, Error> { assert_eq!(keys, ["missing"]); Ok(vec![match self.0.get_cache("missing", &context)? { - Some(value) => litellm_cache::BatchEntry::Hit(value), - None => litellm_cache::BatchEntry::Miss, + Some(value) => BatchEntry::Hit(value), + None => BatchEntry::Miss, }]) } } @@ -233,6 +220,7 @@ impl FlushCache for SyncPanics { } } +#[rstest] #[tokio::test] async fn async_operations_use_the_async_l2_methods() { let l1 = Arc::new(TestCache::new(None, false)); @@ -260,7 +248,7 @@ async fn async_operations_use_the_async_l2_methods() { .async_batch_get_cache(vec!["missing".into()], context.clone()) .await .unwrap(), - [litellm_cache::BatchEntry::Hit("remote".to_string())] + [BatchEntry::Hit("remote".to_string())] ); cache .async_set_cache("missing", "written".into(), context.clone()) @@ -294,14 +282,6 @@ impl BaseCache for Unavailable { fn get_cache(&self, _: &str, _: &ExactCacheContext) -> Result, Error> { Err(Error::Unavailable) } - - async fn disconnect(&self) -> Result<(), Error> { - Ok(()) - } - - async fn test_connection(&self) -> Result { - unreachable!() - } } impl BatchCache for Unavailable {} @@ -330,7 +310,7 @@ impl ClaimCache for Unavailable { } } -#[test] +#[rstest] fn remote_failure_policy_selects_propagation_or_the_local_tier() { let context = ExactCacheContext::default(); let strict = DualCache::new(Arc::new(TestCache::new(None, false)), Arc::new(Unavailable)); @@ -353,7 +333,7 @@ fn remote_failure_policy_selects_propagation_or_the_local_tier() { assert_eq!(l1.get_cache("key", &context), Ok(None)); } -#[test] +#[rstest] fn claim_fallback_does_not_hide_non_availability_errors() { let cache = DualCache::new( Arc::new(TestCache::new(Some("first".to_string()), false)), @@ -371,7 +351,7 @@ fn claim_fallback_does_not_hide_non_availability_errors() { ); } -#[test] +#[rstest] fn local_only_policies_never_touch_l2() { let l2 = Arc::new(TestCache::new(Some("remote".to_string()), false)); let cache = DualCache::new(Arc::new(TestCache::new(None, false)), l2.clone()) @@ -383,3 +363,474 @@ fn local_only_policies_never_touch_l2() { cache.set_cache("key", "local".into(), &context).unwrap(); assert_eq!(l2.get_cache("key", &context), Ok(Some("remote".into()))); } + +type Log = Arc>>; +type StoredSet = (Vec, Option); + +/// A keyed tier that logs every call, so tests can assert which tier ran and in what order. +struct Tier { + name: &'static str, + log: Log, + fail: bool, + counters: Mutex>, + sets: Mutex>, + ttls: HashMap, +} + +impl Tier { + fn new(name: &'static str, log: &Log) -> Self { + Self { + name, + log: log.clone(), + fail: false, + counters: Mutex::default(), + sets: Mutex::default(), + ttls: HashMap::new(), + } + } + + fn failing(self) -> Self { + Self { fail: true, ..self } + } + + fn with_counter(self, key: &str, value: f64) -> Self { + self.counters + .lock() + .unwrap() + .insert(key.into(), (value, ExactCacheContext::default())); + self + } + + fn with_ttl(mut self, key: &str, seconds: u64) -> Self { + self.ttls.insert(key.into(), Duration::from_secs(seconds)); + self + } + + fn record(&self, event: String) { + self.log + .lock() + .unwrap() + .push(format!("{} {event}", self.name)); + } + + fn counter(&self, key: &str) -> Option<(f64, ExactCacheContext)> { + self.counters.lock().unwrap().get(key).cloned() + } + + fn check(&self) -> Result<(), Error> { + if self.fail { + return Err(Error::Unavailable); + } + Ok(()) + } +} + +impl BaseCache for Tier { + type Value = f64; + type Context = ExactCacheContext; + + fn get_ttl(&self, context: &Self::Context) -> Option { + context.ttl + } + + fn set_cache(&self, key: &str, value: f64, context: &ExactCacheContext) -> Result<(), Error> { + self.check()?; + self.record(format!("set {key}={value}")); + self.counters + .lock() + .unwrap() + .insert(key.into(), (value, context.clone())); + Ok(()) + } + + fn get_cache(&self, key: &str, _: &ExactCacheContext) -> Result, Error> { + Ok(self.counter(key).map(|(value, _)| value)) + } +} + +impl CounterCache for Tier { + fn increment_cache( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + ) -> Result { + self.check()?; + let mut counters = self.counters.lock().unwrap(); + let value = counters.get(key).map_or(0.0, |(value, _)| *value) + amount; + counters.insert(key.into(), (value, context)); + Ok(value) + } + + async fn async_increment( + &self, + key: &str, + amount: f64, + context: ExactCacheContext, + refresh_ttl: bool, + ) -> Result { + self.record(format!( + "increment {key}+{amount} refresh_ttl={refresh_ttl}" + )); + self.increment_cache(key, amount, context) + } + + async fn async_increment_pipeline( + &self, + operations: Vec, + ) -> Result, Error> { + let keys = operations + .iter() + .map(|operation| operation.key.as_str()) + .collect::>(); + self.record(format!("pipeline {}", keys.join(","))); + operations + .into_iter() + .map(|operation| { + self.increment_cache( + &operation.key, + operation.amount, + ExactCacheContext { ttl: operation.ttl }, + ) + }) + .collect() + } +} + +impl DeleteCache for Tier { + fn delete_cache(&self, key: &str) -> Result<(), Error> { + self.check()?; + self.record(format!("delete {key}")); + self.counters.lock().unwrap().remove(key); + Ok(()) + } +} + +impl BulkDeleteCache for Tier { + async fn delete_cache_keys(&self, keys: Vec) -> Result { + self.check()?; + self.record(format!("delete_keys {}", keys.join(","))); + let mut counters = self.counters.lock().unwrap(); + Ok(keys + .iter() + .filter(|key| counters.remove(key.as_str()).is_some()) + .count()) + } +} + +impl SetCache for Tier { + type SetValue = String; + type SetResult = (); + + async fn async_set_cache_sadd( + &self, + key: &str, + values: Vec, + ttl: Option, + ) -> Result<(), Error> { + self.check()?; + self.record(format!("sadd {key} {}", values.join(","))); + self.sets.lock().unwrap().insert(key.into(), (values, ttl)); + Ok(()) + } +} + +impl TtlCache for Tier { + async fn async_get_ttl(&self, key: &str) -> Result, Error> { + self.check()?; + self.record(format!("ttl {key}")); + Ok(self.ttls.get(key).copied()) + } +} + +#[fixture] +fn log() -> Log { + Log::default() +} + +fn events(log: &Log) -> Vec { + log.lock().unwrap().clone() +} + +fn seconds(ttl: u64) -> ExactCacheContext { + ExactCacheContext { + ttl: Some(Duration::from_secs(ttl)), + } +} + +#[rstest] +#[case::window_semantics(false)] +#[case::refresh_on_every_write(true)] +#[tokio::test] +async fn async_increment_passes_refresh_ttl_to_l2_and_stores_its_result_locally( + log: Log, + #[case] refresh_ttl: bool, +) { + let l1 = Arc::new(Tier::new("l1", &log).with_counter("counter", 1.0)); + let l2 = Arc::new(Tier::new("l2", &log).with_counter("counter", 10.0)); + let cache = DualCache::new(l1.clone(), l2.clone()); + + assert_eq!( + cache + .async_increment("counter", 2.0, seconds(30), refresh_ttl) + .await, + Ok(12.0) + ); + assert_eq!( + events(&log), + [ + format!("l2 increment counter+2 refresh_ttl={refresh_ttl}"), + "l1 set counter=12".into(), + ] + ); + assert_eq!(l1.counter("counter"), Some((12.0, seconds(30)))); + assert_eq!(l2.counter("counter"), Some((12.0, seconds(30)))); +} + +#[rstest] +#[tokio::test] +async fn failed_async_l2_increment_leaves_l1_unchanged(log: Log) { + let l1 = Arc::new(Tier::new("l1", &log).with_counter("counter", 1.0)); + let cache = DualCache::new(l1.clone(), Arc::new(Tier::new("l2", &log).failing())) + .with_remote_failure_policy(RemoteFailurePolicy::UseLocal); + + assert_eq!( + cache + .async_increment("counter", 2.0, seconds(30), true) + .await, + Err(Error::Unavailable) + ); + assert_eq!( + l1.counter("counter"), + Some((1.0, ExactCacheContext::default())) + ); +} + +#[rstest] +#[case::empty(Vec::new(), Vec::new())] +#[case::one_key(vec![("a", 1.0, Some(10))], vec![6.0])] +#[case::repeated_and_mixed_ttls( + vec![("a", 1.0, Some(10)), ("b", 2.0, None), ("a", 3.0, Some(20))], + vec![6.0, 2.0, 9.0], +)] +#[tokio::test] +async fn async_increment_pipeline_runs_l2_first_and_l1_takes_each_remote_result( + log: Log, + #[case] operations: Vec<(&str, f64, Option)>, + #[case] expected: Vec, +) { + let operations = operations + .into_iter() + .map(|(key, amount, ttl)| IncrementOperation { + key: key.into(), + amount, + ttl: ttl.map(Duration::from_secs), + }) + .collect::>(); + let l1 = Arc::new(Tier::new("l1", &log).with_counter("a", 100.0)); + let l2 = Arc::new(Tier::new("l2", &log).with_counter("a", 5.0)); + let cache = DualCache::new(l1.clone(), l2); + + assert_eq!( + cache.async_increment_pipeline(operations.clone()).await, + Ok(expected.clone()) + ); + let keys = operations + .iter() + .map(|operation| operation.key.as_str()) + .collect::>(); + let mut expected_events = vec![format!("l2 pipeline {}", keys.join(","))]; + expected_events.extend( + operations + .iter() + .zip(&expected) + .map(|(operation, value)| format!("l1 set {}={value}", operation.key)), + ); + assert_eq!(events(&log), expected_events); + if let Some((operation, value)) = operations.iter().zip(&expected).next_back() { + assert_eq!( + l1.counter(&operation.key), + Some((*value, ExactCacheContext { ttl: operation.ttl })) + ); + } +} + +#[rstest] +#[tokio::test] +async fn failed_l2_increment_pipeline_leaves_l1_unchanged(log: Log) { + let l1 = Arc::new(Tier::new("l1", &log).with_counter("a", 1.0)); + let cache = DualCache::new(l1.clone(), Arc::new(Tier::new("l2", &log).failing())); + + assert_eq!( + cache + .async_increment_pipeline(vec![IncrementOperation { + key: "a".into(), + amount: 1.0, + ttl: None, + }]) + .await, + Err(Error::Unavailable) + ); + assert_eq!(events(&log), ["l2 pipeline a"]); + assert_eq!(l1.counter("a"), Some((1.0, ExactCacheContext::default()))); +} + +#[rstest] +#[case::both_tiers(WritePolicy::Both, &["l1 sadd members a,b", "l2 sadd members a,b"])] +#[case::local_only(WritePolicy::LocalOnly, &["l1 sadd members a,b"])] +#[tokio::test] +async fn set_add_writes_locally_then_remotely_unless_local_only( + log: Log, + #[case] write_policy: WritePolicy, + #[case] expected: &[&str], +) { + let l1 = Arc::new(Tier::new("l1", &log)); + let l2 = Arc::new(Tier::new("l2", &log)); + let cache = DualCache::new(l1.clone(), l2.clone()).with_write_policy(write_policy); + let ttl = Some(Duration::from_secs(45)); + + cache + .async_set_cache_sadd("members", vec!["a".into(), "b".into()], ttl) + .await + .unwrap(); + assert_eq!(events(&log), expected); + let stored = Some((vec!["a".to_string(), "b".to_string()], ttl)); + assert_eq!(l1.sets.lock().unwrap().get("members").cloned(), stored); + assert_eq!( + l2.sets.lock().unwrap().get("members").cloned(), + stored.filter(|_| write_policy == WritePolicy::Both) + ); +} + +#[rstest] +#[tokio::test] +async fn failed_local_set_add_never_reaches_l2(log: Log) { + let cache = DualCache::new( + Arc::new(Tier::new("l1", &log).failing()), + Arc::new(Tier::new("l2", &log)), + ); + assert_eq!( + cache + .async_set_cache_sadd("members", vec!["a".into()], None) + .await, + Err(Error::Unavailable) + ); + assert!(events(&log).is_empty()); +} + +#[rstest] +#[case::empty(None, &[], &[])] +#[case::default_batch_size(None, &["a", "b", "c"], &["a,b,c"])] +#[case::chunked(Some(2), &["a", "b", "c", "d", "e"], &["a,b", "c,d", "e"])] +#[case::exact_chunks(Some(2), &["a", "b", "c", "d"], &["a,b", "c,d"])] +#[case::zero_means_one_per_chunk(Some(0), &["a", "b"], &["a", "b"])] +#[tokio::test] +async fn bulk_delete_removes_every_key_locally_then_remotely_in_chunks( + log: Log, + #[case] batch_size: Option, + #[case] keys: &[&str], + #[case] chunks: &[&str], +) { + let l1 = Tier::new("l1", &log); + let l2 = Tier::new("l2", &log); + for key in keys.iter().step_by(2) { + l2.counters + .lock() + .unwrap() + .insert((*key).into(), (1.0, ExactCacheContext::default())); + } + let l1 = Arc::new(l1); + let mut cache = DualCache::new(l1, Arc::new(l2)); + if let Some(batch_size) = batch_size { + cache = cache.with_delete_batch_size(batch_size); + } + + assert_eq!( + cache + .delete_cache_keys(keys.iter().map(|key| (*key).into()).collect()) + .await, + Ok(keys.len().div_ceil(2)) + ); + let expected = keys + .iter() + .map(|key| format!("l1 delete {key}")) + .chain(chunks.iter().map(|chunk| format!("l2 delete_keys {chunk}"))) + .collect::>(); + assert_eq!(events(&log), expected); +} + +#[rstest] +#[tokio::test] +async fn bulk_delete_stops_before_l2_when_the_local_delete_fails(log: Log) { + let cache = DualCache::new( + Arc::new(Tier::new("l1", &log).failing()), + Arc::new(Tier::new("l2", &log)), + ); + assert_eq!( + cache.delete_cache_keys(vec!["a".into()]).await, + Err(Error::Unavailable) + ); + assert!(events(&log).is_empty()); +} + +#[rstest] +#[case::local_hit("both", Some(10), &["l1 ttl both"])] +#[case::remote_fallback("remote", Some(20), &["l1 ttl remote", "l2 ttl remote"])] +#[case::missing_everywhere("missing", None, &["l1 ttl missing", "l2 ttl missing"])] +#[tokio::test] +async fn ttl_reads_local_then_remote( + log: Log, + #[case] key: &str, + #[case] expected: Option, + #[case] expected_events: &[&str], +) { + let cache = DualCache::new( + Arc::new(Tier::new("l1", &log).with_ttl("both", 10)), + Arc::new( + Tier::new("l2", &log) + .with_ttl("both", 99) + .with_ttl("remote", 20), + ), + ); + assert_eq!( + cache.async_get_ttl(key).await, + Ok(expected.map(Duration::from_secs)) + ); + assert_eq!(events(&log), expected_events); +} + +/// Python `local_only=True`: the increment and the pipeline stay on the local tier. +#[rstest] +#[tokio::test] +async fn local_only_writes_increment_the_local_tier_alone(log: Log) { + let l1 = Arc::new(Tier::new("l1", &log).with_counter("a", 1.0)); + let l2 = Arc::new(Tier::new("l2", &log).with_counter("a", 10.0)); + let cache = DualCache::new(l1.clone(), l2.clone()).with_write_policy(WritePolicy::LocalOnly); + + assert_eq!(cache.increment_cache("a", 2.0, seconds(30)), Ok(3.0)); + assert_eq!( + cache.async_increment("a", 1.0, seconds(30), true).await, + Ok(4.0) + ); + let operations = vec![ + IncrementOperation { + key: "a".into(), + amount: 1.0, + ttl: Some(Duration::from_secs(10)), + }, + IncrementOperation { + key: "b".into(), + amount: 2.0, + ttl: None, + }, + ]; + assert_eq!( + cache.async_increment_pipeline(operations).await, + Ok(vec![5.0, 2.0]) + ); + assert_eq!( + events(&log), + ["l1 set a=3", "l1 set a=4", "l1 set a=5", "l1 set b=2"] + ); + assert_eq!(l2.counter("a"), Some((10.0, ExactCacheContext::default()))); +} diff --git a/litellm-rust/crates/cache/tests/semantic.rs b/litellm-rust/crates/cache/tests/semantic.rs new file mode 100644 index 00000000000..97a552a8010 --- /dev/null +++ b/litellm-rust/crates/cache/tests/semantic.rs @@ -0,0 +1,270 @@ +use litellm_cache::{ + Error, SemanticCacheContext, + semantic::{ + Embedder, PreparedEmbedding, prompt_from_context, prompt_from_messages, str_from_messages, + }, +}; +use rstest::rstest; +use serde_json::{Value, json}; + +fn context(messages: Option, input: Option) -> SemanticCacheContext { + SemanticCacheContext { + messages, + input, + ..SemanticCacheContext::default() + } +} + +#[rstest] +#[case::empty(json!([]), "")] +#[case::string_content(json!([{"role": "user", "content": "hello"}]), "hello")] +#[case::concatenates_messages( + json!([{"role": "system", "content": "be brief. "}, {"role": "user", "content": "hello"}]), + "be brief. hello", +)] +#[case::text_parts( + json!([{"role": "user", "content": [ + {"type": "text", "text": "What is "}, + {"type": "image_url", "image_url": {"url": "https://example.com/a.png"}}, + {"type": "text", "text": "this?"}, + ]}]), + "What is this?", +)] +#[case::missing_null_and_empty_content( + json!([{"role": "assistant"}, {"role": "assistant", "content": null}, {"role": "user", "content": ""}]), + "", +)] +#[case::search_results_hidden_behind_small_content( + json!([{"role": "tool", "content": "small", "search_results": [ + {"source": "s", "title": "t", "content": [{"text": "hidden payload"}]}, + ]}]), + "smallsthidden payload", +)] +#[case::title_only_search_result( + json!([{"role": "tool", "content": "small", "search_results": [ + {"source": "s", "title": "long title", "content": []}, + ]}]), + "smallslong title", +)] +#[case::search_results_without_content( + json!([{"role": "tool", "search_results": [{"source": "s", "title": "t"}]}]), + "st", +)] +#[case::search_result_fields_in_python_order( + json!([{"role": "tool", "content": "c", "search_results": [ + {"citations": {"enabled": true}, "content": [{"text": "body"}], "title": "t", "source": "s"}, + {"source": "s2"}, + ]}]), + r#"cstbody{"enabled":true}s2"#, +)] +#[case::null_citations_skipped( + json!([{"role": "tool", "content": "c", "search_results": [ + {"source": "s", "citations": null}, + ]}]), + "cs", +)] +#[case::non_string_and_non_object_entries_skipped( + json!([{"role": "tool", "content": "c", "search_results": [ + "junk", + {"source": 1, "title": null, "content": ["junk", {"text": 3}, {"text": "kept"}]}, + ]}]), + "ckept", +)] +#[case::non_list_search_results_skipped( + json!([{"role": "tool", "content": "c", "search_results": {"source": "s"}}]), + "c", +)] +#[case::citations_compact_in_insertion_order( + json!([{"role": "tool", "search_results": [ + {"citations": {"z": 1, "a": [1.5, true, null], "m": {"k": "v"}}}, + ]}]), + r#"{"z":1,"a":[1.5,true,null],"m":{"k":"v"}}"#, +)] +#[case::citations_ensure_ascii( + json!([{"role": "tool", "search_results": [{"citations": ["caf\u{e9}", "\u{4e2d}"]}]}]), + r#"["caf\u00e9","\u4e2d"]"#, +)] +#[case::citations_astral_chars_as_surrogate_pairs( + json!([{"role": "tool", "search_results": [{"citations": "\u{1f600}"}]}]), + r#""\ud83d\ude00""#, +)] +#[case::citations_escapes( + json!([{"role": "tool", "search_results": [{"citations": "q\"\\\n\t\u{1}/"}]}]), + r#""q\"\\\n\t\u0001/""#, +)] +#[case::citations_large_float_exponent( + json!([{"role": "tool", "search_results": [{"citations": [1e20, 1.0]}]}]), + "[1e+20,1.0]", +)] +#[case::citations_scalars( + json!([{"role": "tool", "search_results": [{"citations": false}, {"citations": 3}]}]), + "false3", +)] +fn str_from_messages_matches_python(#[case] messages: Value, #[case] expected: &str) { + assert_eq!(str_from_messages(messages.as_array().unwrap()), expected); +} + +#[rstest] +#[case::no_messages(None, None)] +#[case::empty_messages(Some(json!([])), None)] +#[case::messages_not_a_list(Some(json!("hello")), None)] +#[case::messages(Some(json!([{"content": "hello"}])), Some("hello"))] +#[case::messages_without_text(Some(json!([{"content": null}])), Some(""))] +fn prompt_from_messages_reads_messages_only( + #[case] messages: Option, + #[case] expected: Option<&str>, +) { + let context = context(messages, Some(json!("responses prompt"))); + assert_eq!(prompt_from_messages(&context).as_deref(), expected); +} + +#[rstest] +#[case::prefers_messages( + Some(json!([{"content": "message prompt"}])), + Some(json!("responses prompt")), + Some("message prompt"), +)] +#[case::empty_messages_fall_back_to_input( + Some(json!([])), + Some(json!("responses prompt")), + Some("responses prompt"), +)] +#[case::messages_without_text_keep_an_empty_prompt( + Some(json!([{"content": null}])), + Some(json!("x")), + Some(""), +)] +#[case::nothing(None, None, None)] +#[case::null_input(None, Some(Value::Null), None)] +#[case::blank_string(None, Some(json!(" ")), None)] +#[case::trimmed_string( + None, + Some(json!(" What is the capital of France?\n")), + Some("What is the capital of France?"), +)] +#[case::image_only( + None, + Some(json!([{"type": "input_image", "image_url": "https://example.com"}])), + None, +)] +#[case::structured_input( + None, + Some(json!([{"role": "user", "content": [ + {"type": "input_text", "text": "What is the capital of France?"}, + {"type": "input_text", "text": "Answer briefly."}, + {"type": "input_image", "image_url": "https://example.com/paris.png"}, + ]}])), + Some("What is the capital of France?\nAnswer briefly."), +)] +#[case::model_objects_after_dump( + None, + Some(json!([ + {"content": [{"text": "model dump prompt"}]}, + {"content": [{"output_text": "dict prompt"}]}, + {"content": [{"input_text": "inline prompt"}]}, + {"content": [{"type": "input_image", "image_url": "https://example.com"}]}, + ])), + Some("model dump prompt\ndict prompt\ninline prompt"), +)] +#[case::object_content( + None, + Some(json!({"content": [{"text": "object content prompt"}]})), + Some("object content prompt"), +)] +#[case::string_content(None, Some(json!({"content": " inline "})), Some("inline"))] +#[case::null_content_uses_text_keys( + None, + Some(json!({"content": null, "output": "tool output"})), + Some("tool output"), +)] +#[case::content_wins_over_text(None, Some(json!({"content": [], "text": "ignored"})), None)] +#[case::text_key_precedence( + None, + Some(json!({"output_text": "d", "input_text": "c", "output": "b", "text": "a"})), + Some("a"), +)] +#[case::input_text_key(None, Some(json!({"input_text": "only input"})), Some("only input"))] +#[case::output_text_key(None, Some(json!({"output_text": "only output"})), Some("only output"))] +#[case::non_string_text_keys_skipped( + None, + Some(json!({"text": 1, "output": "fallback"})), + Some("fallback"), +)] +#[case::nested_lists(None, Some(json!([["a", [" b "]], "", "c"])), Some("a\nb\nc"))] +#[case::scalars_ignored(None, Some(json!([1, true, null, "kept"])), Some("kept"))] +fn prompt_from_context_matches_python( + #[case] messages: Option, + #[case] input: Option, + #[case] expected: Option<&str>, +) { + assert_eq!( + prompt_from_context(&context(messages, input)).as_deref(), + expected + ); +} + +/// Python `test_redis_semantic_cache_prompt_extraction_skips_blank_dict_text_keys`: a blank +/// text key falls through to the next one. +#[rstest] +#[case::blank_text_falls_through( + json!({"text": " ", "input_text": "fallback prompt"}), + "fallback prompt", +)] +fn prompt_from_context_skips_blank_text_keys(#[case] input: Value, #[case] expected: &str) { + assert_eq!( + prompt_from_context(&context(None, Some(input))).as_deref(), + Some(expected) + ); +} + +/// Where `json.dumps(..., separators=(",", ":"))` and Python `str.strip` differ from +/// `semantic.rs`: ensure_ascii escapes DEL, small floats keep Python's two-digit exponent, and +/// strip also removes the ASCII information separators. +#[rstest] +#[case::del_is_escaped(json!([{"search_results": [{"citations": "\u{7f}"}]}]), None, r#""\u007f""#)] +#[case::small_float_exponent(json!([{"search_results": [{"citations": 1.5e-7}]}]), None, "1.5e-07")] +#[case::float_at_positional_floor(json!([{"search_results": [{"citations": 1e-4}]}]), None, "0.0001")] +#[case::float_at_scientific_ceiling(json!([{"search_results": [{"citations": 1e16}]}]), None, "1e+16")] +#[case::large_float(json!([{"search_results": [{"citations": [1.25e20, -2.5, 3.0]}]}]), None, "[1.25e+20,-2.5,3.0]")] +#[case::strip_information_separators(json!([]), Some(json!("\u{1c}a\u{1f}")), "a")] +fn python_serialization_edge_cases( + #[case] messages: Value, + #[case] input: Option, + #[case] expected: &str, +) { + let actual = match input { + Some(input) => prompt_from_context(&context(None, Some(input))).unwrap_or_default(), + None => str_from_messages(messages.as_array().unwrap()), + }; + assert_eq!(actual, expected); +} + +#[rstest] +#[tokio::test] +async fn prepared_embedding_returns_its_vector_for_any_prompt() { + let embedding = PreparedEmbedding(vec![0.1, 0.2, 0.3]); + let metadata = json!({"tenant": "team"}); + assert_eq!(embedding.embed("a", None), Ok(vec![0.1, 0.2, 0.3])); + assert_eq!( + embedding.async_embed("b", Some(&metadata)).await, + Ok(vec![0.1, 0.2, 0.3]) + ); +} + +#[rstest] +#[tokio::test] +async fn embedders_default_to_async_only() { + struct AsyncOnly; + + impl Embedder for AsyncOnly { + async fn async_embed(&self, prompt: &str, _: Option<&Value>) -> Result, Error> { + Ok(vec![prompt.len() as f32]) + } + } + + assert_eq!( + AsyncOnly.embed("abc", None), + Err(Error::UnsupportedOperation) + ); + assert_eq!(AsyncOnly.async_embed("abc", None).await, Ok(vec![3.0])); +} diff --git a/litellm-rust/crates/python-bridge/src/cache/activation.rs b/litellm-rust/crates/python-bridge/src/cache/activation.rs new file mode 100644 index 00000000000..8d2c340dec8 --- /dev/null +++ b/litellm-rust/crates/python-bridge/src/cache/activation.rs @@ -0,0 +1,111 @@ +use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; +use litellm_cache_redis_semantic::RedisSemanticConfig; +use litellm_host_python::{release_gil, run_sync_value}; +use litellm_http::ClientVariant; +use pyo3::prelude::*; + +use super::{ + cache_error, + config::{CacheBackendConfig, NativeCacheConfig, UnsupportedCacheConfig}, + embedder::PythonEmbedder, + host_client, + native::NativeResponseCache, +}; +use crate::errors::RustBridgeDeclined; + +fn declined(reason: UnsupportedCacheConfig) -> PyErr { + RustBridgeDeclined::new_err(reason.message()) +} + +/// Builds the native backend a `Cache` facade's projected configuration describes. `backend` is +/// the facade's `.cache` object, which owns embedding for the Python-embedded semantic caches. +pub(super) fn activate( + py: Python<'_>, + backend: &Bound<'_, PyAny>, + config: NativeCacheConfig, +) -> PyResult { + let policy = config.policy; + let service = match config.backend { + CacheBackendConfig::Memory(memory) => { + NativeResponseCache::memory(memory.capacity, memory.default_ttl, memory.max_entry_bytes) + } + CacheBackendConfig::Redis(redis) => { + let url = redis.connection.native_url().map_err(declined)?; + let flush_size = policy.redis_flush_size.map(|_| redis.flush_size); + release_gil(py, move || { + NativeResponseCache::redis( + &url, + &redis.topology, + Some(redis.default_ttl), + redis.namespace, + ) + }) + .map_err(cache_error)? + .with_redis_flush_size(flush_size) + } + CacheBackendConfig::S3(s3) => { + let http = host_client(py, ClientVariant::NoRedirect)?; + run_sync_value( + py, + async move { Ok(NativeResponseCache::s3(*s3, http).await) }, + )? + } + CacheBackendConfig::Gcs(gcs) => NativeResponseCache::gcs( + GcsConfig { + bucket_name: gcs.bucket_name, + gcs_path: Some(gcs.key_prefix), + path_service_account: gcs.path_service_account, + endpoint: DEFAULT_ENDPOINT.to_owned(), + }, + host_client(py, ClientVariant::NoRedirect)?, + None, + ), + CacheBackendConfig::Disk(disk) => { + release_gil(py, move || NativeResponseCache::disk(&disk.directory)) + .map_err(cache_error)? + } + CacheBackendConfig::AzureBlob(azure) => { + let http = host_client(py, ClientVariant::NoRedirect)?; + run_sync_value(py, async move { + NativeResponseCache::azure_blob(&azure.account_url, &azure.container, http) + .await + .map_err(cache_error) + })? + } + CacheBackendConfig::RedisSemantic(semantic) => { + let url = semantic.native_url().map_err(declined)?.to_owned(); + let embedder = PythonEmbedder::new(backend.clone().unbind()); + let semantic_config = RedisSemanticConfig { + index_name: semantic.index_name, + similarity_threshold: semantic.similarity_threshold as f32, + }; + release_gil(py, move || { + NativeResponseCache::redis_semantic(&url, embedder, semantic_config) + }) + .map_err(cache_error)? + } + CacheBackendConfig::ValkeySemantic(valkey) => { + let url = valkey.connection.native_url().map_err(declined)?; + let embedder = PythonEmbedder::new(backend.clone().unbind()); + release_gil(py, move || { + NativeResponseCache::valkey_semantic( + &url, + valkey.similarity_threshold, + valkey.index_name, + embedder, + ) + }) + .map_err(cache_error)? + } + CacheBackendConfig::QdrantSemantic(qdrant) => { + let client = host_client(py, ClientVariant::Provider)?; + run_sync_value(py, async move { + let runtime = tokio::runtime::Handle::current(); + NativeResponseCache::qdrant_semantic(*qdrant, client, runtime) + .await + .map_err(cache_error) + })? + } + }; + Ok(service.with_scope(policy.semantic_cache_scope)) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/binding.rs b/litellm-rust/crates/python-bridge/src/cache/binding.rs index 273d3f9ca4e..56f1bff2853 100644 --- a/litellm-rust/crates/python-bridge/src/cache/binding.rs +++ b/litellm-rust/crates/python-bridge/src/cache/binding.rs @@ -9,11 +9,12 @@ use pyo3::{ use serde_json::Value; use super::{ + activation::activate, cache_error, callback::PythonCallback, - config::{CacheBackendConfig, CacheConfigProjection, NativeCacheConfig}, + config::{CacheConfigProjection, NativeCacheConfig}, future::{ready_none, ready_value}, - native::NativeResponseCache, + native::{NativeResponseCache, SemanticReply}, request::{now, request, requests}, }; use crate::errors::RustBridgeDeclined; @@ -76,23 +77,9 @@ impl ResolvedCache { return Err(RustBridgeDeclined::new_err(reason.message())); } }; - let service = match config.backend { - CacheBackendConfig::Memory(memory) => NativeResponseCache::memory( - memory.capacity, - memory.default_ttl, - memory.max_entry_bytes, - ), - _ => { - return Err(RustBridgeDeclined::new_err( - "native response cache activation is not implemented for this backend", - )); - } - }; - Ok(Self::new(CacheBinding::Native( - service - .with_scope(config.policy.semantic_cache_scope) - .with_redis_flush_size(config.policy.redis_flush_size), - ))) + let backend = cache.getattr("cache")?; + let service = activate(cache.py(), &backend, config)?; + Ok(Self::new(CacheBinding::Native(service))) } #[getter] @@ -127,6 +114,41 @@ impl ResolvedCache { } } + /// `(response, similarity)`: the similarity is `None` when the backend reports none. + fn lookup_semantic(&self, py: Python<'_>, request: &Bound<'_, PyAny>) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Native(service) => { + let request = self::request(request)?; + let service = service.clone(); + let lookup = release_gil(py, move || service.lookup_semantic(&request, now())) + .map_err(cache_error)?; + to_py(py, &SemanticReply::from(lookup)) + } + CacheBinding::Disabled => to_py(py, &SemanticReply(None, None)), + CacheBinding::PythonCallback(_) => Err(PyRuntimeError::new_err( + "semantic lookups require a native cache binding", + )), + } + } + + fn async_lookup_semantic<'py>( + &self, + py: Python<'py>, + request: &Bound<'py, PyAny>, + ) -> PyResult> { + self.check_process()?; + match &self.binding { + CacheBinding::Native(service) => { + service.async_lookup_semantic_py(py, self::request(request)?) + } + CacheBinding::Disabled => ready_value(py, &SemanticReply(None, None)), + CacheBinding::PythonCallback(_) => Err(PyRuntimeError::new_err( + "semantic lookups require a native cache binding", + )), + } + } + #[pyo3(signature = (request, response, *, callback_kwargs=None))] fn store( &self, diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index b6e08102e18..6e25f07efa1 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -13,12 +13,7 @@ use pyo3::{ use super::{identity::BackendIdentity, native::NativeResponseCache, request::duration}; -#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct CachePolicy { - pub(super) mode: String, - pub(super) ttl: Option, - pub(super) namespace: Option, - pub(super) supported_call_types: Option>, pub(super) redis_flush_size: Option, pub(super) semantic_cache_scope: String, } @@ -46,7 +41,6 @@ pub(super) enum CertificateRequirement { Required, } -#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct RedisTlsConfig { pub(super) certificate_requirement: CertificateRequirement, pub(super) check_hostname: bool, @@ -56,7 +50,6 @@ pub(super) struct RedisTlsConfig { pub(super) client_key: Option, } -#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct RedisConnectionConfig { pub(super) host: String, pub(super) port: u16, @@ -73,7 +66,6 @@ pub(super) struct RedisConnectionConfig { pub(super) tls: Option, } -#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct RedisCacheConfig { pub(super) default_ttl: Duration, pub(super) namespace: Option, @@ -94,17 +86,10 @@ pub(super) struct AzureBlobCacheConfig { pub(super) container: String, } -#[allow( - dead_code, - reason = "embedding settings are projected so drift falls back to Python" -)] pub(super) struct RedisSemanticCacheConfig { pub(super) redis_url: String, pub(super) index_name: String, pub(super) similarity_threshold: f64, - pub(super) embedding_model: String, - pub(super) embedding_max_input_tokens: Option, - pub(super) embedding_timeout: Option, } struct RedisClientProjection<'py> { @@ -118,11 +103,13 @@ struct RedisClientProjection<'py> { const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31; -#[allow(dead_code, reason = "consumed by the cache activation follow-up")] +/// The read and write timeout every native Redis connection uses, which is also `RedisCache`'s +/// default `socket_timeout`. +const NATIVE_REDIS_SOCKET_TIMEOUT: Duration = Duration::from_secs(5); + pub(super) struct ValkeySemanticCacheConfig { pub(super) similarity_threshold: f64, pub(super) index_name: String, - pub(super) embedding_model: String, pub(super) connection: RedisConnectionConfig, } @@ -147,6 +134,93 @@ impl QdrantSemanticCacheConfig { } } +impl RedisTlsConfig { + /// Whether redis-rs with rustls behaves like this redis-py `SSLConnection`: it verifies the + /// certificate chain against the system roots and always checks the hostname. + fn native(&self) -> Result<(), UnsupportedCacheConfig> { + if self.ca_certificate.is_some() + || self.ca_data.is_some() + || self.client_certificate.is_some() + || self.client_key.is_some() + { + return Err(UnsupportedCacheConfig::RedisTlsCertificates); + } + if self.certificate_requirement == CertificateRequirement::None || !self.check_hostname { + return Err(UnsupportedCacheConfig::RedisTlsVerification); + } + Ok(()) + } +} + +impl RedisConnectionConfig { + /// The redis-rs URL for this connection, or the first setting the native client cannot + /// honor. The native pool and socket timeouts are fixed, so only redis-py's unbounded pool, + /// its unset timeouts and `RedisCache`'s five-second `socket_timeout` map onto them. + pub(super) fn native_url(&self) -> Result { + if self.pool_size != REDIS_PY_DEFAULT_MAX_CONNECTIONS { + return Err(UnsupportedCacheConfig::RedisPoolSize); + } + if self + .read_timeout + .is_some_and(|timeout| timeout != NATIVE_REDIS_SOCKET_TIMEOUT) + || self.connect_timeout.is_some() + { + return Err(UnsupportedCacheConfig::RedisTimeout); + } + if self.socket_keepalive == Some(true) { + return Err(UnsupportedCacheConfig::RedisKeepalive); + } + if !self.health_check_interval.is_zero() { + return Err(UnsupportedCacheConfig::RedisHealthCheck); + } + if self.client_name.is_some() { + return Err(UnsupportedCacheConfig::RedisClientName); + } + let scheme = match &self.tls { + None => "redis", + Some(tls) => { + tls.native()?; + "rediss" + } + }; + let host = if self.host.contains(':') { + format!("[{}]", self.host) + } else { + self.host.clone() + }; + let mut url = url::Url::parse(&format!( + "{scheme}://{host}:{}/{}", + self.port, self.database + )) + .map_err(|_| UnsupportedCacheConfig::RedisConnection)?; + if let Some(username) = &self.username { + url.set_username(username) + .map_err(|()| UnsupportedCacheConfig::RedisConnection)?; + } + if let Some(password) = &self.password { + url.set_password(Some(password)) + .map_err(|()| UnsupportedCacheConfig::RedisConnection)?; + } + if self.protocol == RedisProtocol::Resp3 { + url.set_query(Some("protocol=resp3")); + } + Ok(url.into()) + } +} + +impl RedisSemanticCacheConfig { + /// redisvl hands `redis_url` to redis-py, which reads TLS and socket options from the URL; + /// redis-rs ignores those, so only a plain URL keeps its meaning. + pub(super) fn native_url(&self) -> Result<&str, UnsupportedCacheConfig> { + let url = url::Url::parse(&self.redis_url) + .map_err(|_| UnsupportedCacheConfig::RedisSemanticUrl)?; + if !matches!(url.scheme(), "redis" | "unix") || url.query().is_some() { + return Err(UnsupportedCacheConfig::RedisSemanticUrl); + } + Ok(&self.redis_url) + } +} + pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), @@ -159,7 +233,6 @@ pub(super) enum CacheBackendConfig { QdrantSemantic(Box), } -#[allow(dead_code, reason = "consumed by the cache activation follow-up")] pub(super) struct NativeCacheConfig { pub(super) policy: CachePolicy, pub(super) backend: CacheBackendConfig, @@ -178,6 +251,15 @@ pub(super) enum UnsupportedCacheConfig { DiskStore, QdrantEndpoint, SemanticEmbedding, + RedisPoolSize, + RedisTimeout, + RedisKeepalive, + RedisHealthCheck, + RedisClientName, + RedisTlsCertificates, + RedisTlsVerification, + RedisSemanticUrl, + ValkeyTls, } impl UnsupportedCacheConfig { @@ -197,6 +279,28 @@ impl UnsupportedCacheConfig { "native Qdrant requires the default REST port so the gRPC port can be derived" } Self::SemanticEmbedding => "native semantic embedding requires Python", + Self::RedisPoolSize => { + "native Redis uses a fixed connection pool; max_connections requires Python" + } + Self::RedisTimeout => { + "native Redis uses fixed socket timeouts; socket_timeout and \ + socket_connect_timeout require Python" + } + Self::RedisKeepalive => "native Redis does not support socket_keepalive", + Self::RedisHealthCheck => "native Redis does not support health_check_interval", + Self::RedisClientName => "native Redis does not support client_name", + Self::RedisTlsCertificates => { + "native Redis TLS does not support ssl_ca_certs, ssl_ca_data, ssl_certfile or \ + ssl_keyfile" + } + Self::RedisTlsVerification => { + "native Redis TLS always verifies the certificate and hostname; \ + ssl_cert_reqs=none and ssl_check_hostname=false require Python" + } + Self::RedisSemanticUrl => { + "native Redis semantic cache does not support TLS or query options in redis_url" + } + Self::ValkeyTls => "native Valkey semantic cache does not support TLS connections", } } } @@ -211,12 +315,6 @@ impl NativeCacheConfig { pub(super) fn project(facade: &Bound<'_, PyAny>) -> PyResult { let backend_name = facade.getattr("type")?.extract::()?; let policy = CachePolicy { - mode: facade.getattr("mode")?.extract::()?, - ttl: optional_duration(facade.getattr("ttl")?)?, - namespace: optional_string(facade.getattr("namespace")?)?, - supported_call_types: facade - .getattr("supported_call_types")? - .extract::>>()?, redis_flush_size: facade .getattr("redis_flush_size")? .extract::>()?, @@ -475,13 +573,6 @@ pub(super) fn project_redis_semantic( .extract::>()? .unwrap_or_else(|| "litellm_semantic_cache_index".into()), similarity_threshold: backend.getattr("similarity_threshold")?.extract::()?, - embedding_model: backend.getattr("embedding_model")?.extract::()?, - embedding_max_input_tokens: backend - .getattr("embedding_max_input_tokens")? - .extract::>()?, - embedding_timeout: backend - .getattr("embedding_timeout")? - .extract::>()?, }) } @@ -592,37 +683,48 @@ fn project_redis( if has_value(&resolved, "credential_provider")? { return Ok(Err(UnsupportedCacheConfig::RedisCredentials)); } - - let protocol = match optional_i64(&resolved, "protocol")?.unwrap_or(2) { - 2 => RedisProtocol::Resp2, - 3 => RedisProtocol::Resp3, - _ => return Err(PyValueError::new_err("unsupported Redis protocol version")), - }; - let health_check_interval = - duration(optional_f64(&resolved, "health_check_interval")?.unwrap_or(0.0))?; Ok(Ok(RedisCacheConfig { default_ttl: duration(backend.getattr("default_ttl")?.extract::()?)?, namespace: optional_attribute_string(backend, "namespace")?, flush_size: backend.getattr("redis_flush_size")?.extract::()?, topology, - connection: RedisConnectionConfig { - host, - port, - database: optional_i64(&resolved, "db")?.unwrap_or(0), - username: optional_dict_string(&resolved, "username")?, - password: optional_dict_string(&resolved, "password")?, - protocol, - pool_size, - read_timeout: optional_dict_duration(&resolved, "socket_timeout")?, - connect_timeout: optional_dict_duration(&resolved, "socket_connect_timeout")?, - socket_keepalive: optional_bool(&resolved, "socket_keepalive")?, - health_check_interval, - client_name: optional_dict_string(&resolved, "client_name")?, - tls, - }, + connection: resolved_connection(&resolved, host, port, pool_size, tls)?, })) } +/// The connection settings redis-py resolved for one client's pool. +#[inline(never)] +fn resolved_connection( + resolved: &Bound<'_, PyDict>, + host: String, + port: u16, + pool_size: usize, + tls: Option, +) -> PyResult { + let protocol = match optional_i64(resolved, "protocol")?.unwrap_or(2) { + 2 => RedisProtocol::Resp2, + 3 => RedisProtocol::Resp3, + _ => return Err(PyValueError::new_err("unsupported Redis protocol version")), + }; + Ok(RedisConnectionConfig { + host, + port, + database: optional_i64(resolved, "db")?.unwrap_or(0), + username: optional_dict_string(resolved, "username")?, + password: optional_dict_string(resolved, "password")?, + protocol, + pool_size, + read_timeout: optional_dict_duration(resolved, "socket_timeout")?, + connect_timeout: optional_dict_duration(resolved, "socket_connect_timeout")?, + socket_keepalive: optional_bool(resolved, "socket_keepalive")?, + health_check_interval: duration( + optional_f64(resolved, "health_check_interval")?.unwrap_or(0.0), + )?, + client_name: optional_dict_string(resolved, "client_name")?, + tls, + }) +} + #[inline(never)] fn project_s3( backend: &Bound<'_, PyAny>, @@ -831,31 +933,22 @@ fn project_valkey_semantic( } } if is_tls { + return Ok(Err(UnsupportedCacheConfig::ValkeyTls)); + } + let host = required_string(&resolved, "host")?; + if host.is_empty() { return Ok(Err(UnsupportedCacheConfig::RedisConnection)); } - let connection = RedisConnectionConfig { - host: required_string(&resolved, "host")?, - port: u16::try_from(required_i64(&resolved, "port")?) - .map_err(|_| PyValueError::new_err("invalid Redis port"))?, - database: optional_i64(&resolved, "db")?.unwrap_or(0), - username: optional_dict_string(&resolved, "username")?, - password: optional_dict_string(&resolved, "password")?, - protocol: RedisProtocol::Resp2, - pool_size: pool.getattr("max_connections")?.extract::()?, - read_timeout: None, - connect_timeout: None, - socket_keepalive: None, - health_check_interval: Duration::ZERO, - client_name: None, - tls: None, - }; - if connection.host.is_empty() { - return Ok(Err(UnsupportedCacheConfig::RedisConnection)); - } + let connection = resolved_connection( + &resolved, + host, + port(required_i64(&resolved, "port")?)?, + pool.getattr("max_connections")?.extract::()?, + None, + )?; Ok(Ok(ValkeySemanticCacheConfig { similarity_threshold: backend.getattr("similarity_threshold")?.extract()?, index_name: backend.getattr("index_name")?.extract()?, - embedding_model: backend.getattr("embedding_model")?.extract()?, connection, })) } @@ -946,11 +1039,6 @@ fn class_is(value: &Bound<'_, PyAny>, module: &str, name: &str) -> PyResult) -> PyResult> { - value.extract::>()?.map(duration).transpose() -} - #[inline(never)] fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { match value.getattr(name) { @@ -1070,16 +1158,17 @@ fn optional_dict_duration(values: &Bound<'_, PyDict>, key: &str) -> PyResult) -> NativeCacheConfig { + match NativeCacheConfig::project(facade).unwrap() { + CacheConfigProjection::Native(config) => *config, + CacheConfigProjection::Unsupported(reason) => panic!("{}", reason.message()), + } + } + + fn unsupported(facade: &Bound<'_, PyAny>) -> UnsupportedCacheConfig { + match NativeCacheConfig::project(facade).unwrap() { + CacheConfigProjection::Native(_) => panic!("configuration must stay on Python"), + CacheConfigProjection::Unsupported(reason) => reason, + } + } + + #[fixture] + fn interpreter() { Python::initialize(); + } + + #[fixture] + fn connection() -> RedisConnectionConfig { + RedisConnectionConfig { + host: "cache.internal".into(), + port: 6380, + database: 4, + username: None, + password: None, + protocol: RedisProtocol::Resp2, + pool_size: REDIS_PY_DEFAULT_MAX_CONNECTIONS, + read_timeout: None, + connect_timeout: None, + socket_keepalive: None, + health_check_interval: Duration::ZERO, + client_name: None, + tls: None, + } + } + + fn verified_tls() -> RedisTlsConfig { + RedisTlsConfig { + certificate_requirement: CertificateRequirement::Required, + check_hostname: true, + ca_certificate: None, + ca_data: None, + client_certificate: None, + client_key: None, + } + } + + #[rstest] + fn projects_effective_memory_configuration(_interpreter: ()) { Python::attach(|py| { let facade = facade( py, "backend = SimpleNamespace(default_ttl=913, max_size_in_memory=37, max_size_per_item=8)\n\ facade = SimpleNamespace(type='local', mode='default-on', ttl=11.5, namespace=None, supported_call_types=['completion'], redis_flush_size=None, semantic_cache_scope='key', cache=backend)", ); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("memory cache should be supported"); - }; - assert_eq!( - config.policy.ttl.unwrap(), - std::time::Duration::from_secs_f64(11.5) - ); + let config = native(&facade); + assert_eq!(config.policy.semantic_cache_scope, "key"); + assert_eq!(config.policy.redis_flush_size, None); let CacheBackendConfig::Memory(memory) = config.backend else { panic!("expected memory configuration"); }; - assert_eq!(memory.default_ttl, std::time::Duration::from_secs(913)); + assert_eq!(memory.default_ttl, Duration::from_secs(913)); assert_eq!(memory.capacity, 37); assert_eq!(memory.max_entry_bytes, 8192); - let matching = - NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8192); - let mismatched = - NativeResponseCache::memory(37, std::time::Duration::from_secs(913), 8191); + let matching = NativeResponseCache::memory(37, Duration::from_secs(913), 8192); + let mismatched = NativeResponseCache::memory(37, Duration::from_secs(913), 8191); let matching_config = NativeCacheConfig { policy: config.policy, backend: CacheBackendConfig::Memory(memory), @@ -1154,9 +1283,8 @@ mod tests { }); } - #[test] - fn redis_semantic_service_mismatch_accepts_backend_precision_threshold() { - Python::initialize(); + #[rstest] + fn redis_semantic_service_mismatch_accepts_backend_precision_threshold(_interpreter: ()) { Python::attach(|py| { let facade = facade( py, @@ -1165,12 +1293,7 @@ mod tests { ); let backend = facade.getattr("cache").unwrap(); let embedder = PythonEmbedder::new(backend.clone().unbind()); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("Redis semantic cache should be supported"); - }; - let CacheBackendConfig::RedisSemantic(config) = config.backend else { + let CacheBackendConfig::RedisSemantic(config) = native(&facade).backend else { panic!("expected Redis semantic configuration"); }; let service = NativeResponseCache::redis_semantic( @@ -1184,10 +1307,6 @@ mod tests { .unwrap(); let matching_config = NativeCacheConfig { policy: CachePolicy { - mode: "default-on".into(), - ttl: None, - namespace: None, - supported_call_types: None, redis_flush_size: None, semantic_cache_scope: "key".into(), }, @@ -1197,9 +1316,8 @@ mod tests { }); } - #[test] - fn projects_resolved_redis_tls_configuration() { - Python::initialize(); + #[rstest] + fn projects_resolved_redis_tls_configuration(_interpreter: ()) { Python::attach(|py| { let facade = facade( py, @@ -1211,15 +1329,12 @@ mod tests { backend = SimpleNamespace(default_ttl=777, namespace='team', redis_flush_size=31, redis_kwargs={}, redis_client=client)\n\ facade = SimpleNamespace(type='redis', mode='default-off', ttl=None, namespace='team', supported_call_types=None, redis_flush_size=31, semantic_cache_scope='key', cache=backend)", ); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("Redis cache should be supported"); - }; + let config = native(&facade); + assert_eq!(config.policy.redis_flush_size, Some(31)); let CacheBackendConfig::Redis(redis) = config.backend else { panic!("expected Redis configuration"); }; - assert_eq!(redis.default_ttl, std::time::Duration::from_secs(777)); + assert_eq!(redis.default_ttl, Duration::from_secs(777)); assert_eq!(redis.namespace.as_deref(), Some("team")); assert_eq!(redis.flush_size, 31); assert_eq!(redis.connection.host, "cache.internal"); @@ -1227,7 +1342,21 @@ mod tests { assert_eq!(redis.connection.database, 4); assert_eq!(redis.connection.protocol, RedisProtocol::Resp3); assert_eq!(redis.connection.pool_size, 29); - let tls = redis.connection.tls.unwrap(); + assert_eq!( + redis.connection.read_timeout, + Some(Duration::from_secs_f64(7.5)) + ); + assert_eq!( + redis.connection.connect_timeout, + Some(Duration::from_secs(2)) + ); + assert_eq!(redis.connection.socket_keepalive, Some(true)); + assert_eq!( + redis.connection.health_check_interval, + Duration::from_secs(15) + ); + assert_eq!(redis.connection.client_name.as_deref(), Some("litellm")); + let tls = redis.connection.tls.as_ref().unwrap(); assert_eq!( tls.certificate_requirement, CertificateRequirement::Optional @@ -1237,123 +1366,89 @@ mod tests { assert_eq!(tls.ca_data.as_deref(), Some("CA DATA")); assert_eq!(tls.client_certificate.as_deref(), Some("/client.pem")); assert_eq!(tls.client_key.as_deref(), Some("/client.key")); + assert!(matches!( + redis.connection.native_url(), + Err(UnsupportedCacheConfig::RedisPoolSize) + )); }); } - #[test] - fn projects_valkey_semantic_configuration() { - Python::initialize(); + #[rstest] + fn projects_valkey_semantic_configuration(_interpreter: ()) { Python::attach(|py| { let facade = facade( py, "pool = ConnectionPool()\n\ pool.connection_class = Connection\n\ pool.max_connections = 12\n\ - pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'db': 2}\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'db': 2, 'socket_timeout': 3}\n\ client = SimpleNamespace(connection_pool=pool)\n\ backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", ); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("Valkey semantic cache should be supported"); - }; - let CacheBackendConfig::ValkeySemantic(valkey) = config.backend else { + let CacheBackendConfig::ValkeySemantic(valkey) = native(&facade).backend else { panic!("expected Valkey semantic configuration"); }; assert_eq!(valkey.similarity_threshold, 0.85); assert_eq!(valkey.index_name, "semantic_idx"); - assert_eq!(valkey.embedding_model, "text-embedding-3-small"); assert_eq!(valkey.connection.host, "cache.internal"); assert_eq!(valkey.connection.port, 6390); assert_eq!(valkey.connection.database, 2); assert_eq!(valkey.connection.pool_size, 12); assert_eq!(valkey.connection.protocol, RedisProtocol::Resp2); + assert_eq!(valkey.connection.read_timeout, Some(Duration::from_secs(3))); assert!(valkey.connection.tls.is_none()); }); } - #[test] - fn valkey_semantic_tls_stays_on_python() { - Python::initialize(); + #[rstest] + #[case::valkey_tls( + "pool = ConnectionPool()\n\ + pool.connection_class = SSLConnection\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ + facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + "native Valkey semantic cache does not support TLS connections" + )] + #[case::valkey_dynamic_auth( + "pool = ConnectionPool()\n\ + pool.connection_class = Connection\n\ + pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'credential_provider': object()}\n\ + client = SimpleNamespace(connection_pool=pool)\n\ + backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ + facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + "native Redis credentials require Python" + )] + #[case::redis_dynamic_auth( + "backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\ + facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + "native Redis credentials require Python" + )] + #[case::gcs_without_bucket( + "backend = SimpleNamespace(bucket_name=None, key_prefix='', path_service_account=None)\n\ + facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", + "native GCS cache requires a configured bucket name" + )] + fn configurations_that_stay_on_python( + _interpreter: (), + #[case] body: &str, + #[case] message: &str, + ) { Python::attach(|py| { - let facade = facade( - py, - "pool = ConnectionPool()\n\ - pool.connection_class = SSLConnection\n\ - pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390}\n\ - client = SimpleNamespace(connection_pool=pool)\n\ - backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ - facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", - ); - let CacheConfigProjection::Unsupported(reason) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("TLS Valkey semantic cache should stay on Python"); - }; - assert_eq!( - reason.message(), - "native Redis connection type is not implemented" - ); + assert_eq!(unsupported(&facade(py, body)).message(), message); }); } - #[test] - fn valkey_semantic_dynamic_auth_stays_on_python() { - Python::initialize(); - Python::attach(|py| { - let facade = facade( - py, - "pool = ConnectionPool()\n\ - pool.connection_class = Connection\n\ - pool.connection_kwargs = {'host': 'cache.internal', 'port': 6390, 'credential_provider': object()}\n\ - client = SimpleNamespace(connection_pool=pool)\n\ - backend = SimpleNamespace(similarity_threshold=0.85, index_name='semantic_idx', embedding_model='text-embedding-3-small', sync_client=client)\n\ - facade = SimpleNamespace(type='valkey-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", - ); - let CacheConfigProjection::Unsupported(reason) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("dynamic Valkey authentication must stay on Python"); - }; - assert_eq!(reason.message(), "native Redis credentials require Python"); - }); - } - - #[test] - fn dynamic_redis_auth_stays_on_python() { - Python::initialize(); - Python::attach(|py| { - let facade = facade( - py, - "backend = SimpleNamespace(redis_kwargs={'credential_provider': object()})\n\ - facade = SimpleNamespace(type='redis', mode='default-on', ttl=None, namespace=None, supported_call_types=[], redis_flush_size=None, semantic_cache_scope='key', cache=backend)", - ); - let CacheConfigProjection::Unsupported(reason) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("dynamic authentication must stay on Python"); - }; - assert_eq!(reason.message(), "native Redis credentials require Python"); - }); - } - - #[test] - fn projects_cluster_startup_nodes_as_redis_topology() { - Python::initialize(); + #[rstest] + fn projects_cluster_startup_nodes_as_redis_topology(_interpreter: ()) { Python::attach(|py| { let facade = cluster_facade( py, "[{'host': 'node-a', 'port': 7000}, {'host': 'node-b', 'port': 7001}]", "client.on_connect", ); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("cluster startup nodes should project natively"); - }; - let CacheBackendConfig::Redis(redis) = &config.backend else { + let CacheBackendConfig::Redis(redis) = native(&facade).backend else { panic!("expected Redis configuration"); }; let expected = RedisTopology::Cluster { @@ -1382,23 +1477,22 @@ mod tests { .certificate_requirement, CertificateRequirement::None ); + assert!(matches!( + redis.connection.native_url(), + Err(UnsupportedCacheConfig::RedisTlsVerification) + )); }); } - #[test] - fn projects_gcs_configuration() { - Python::initialize(); + #[rstest] + fn projects_gcs_configuration(_interpreter: ()) { Python::attach(|py| { let facade = facade( py, "backend = SimpleNamespace(bucket_name='bucket', key_prefix='cache/', path_service_account='credentials.json')\n\ facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", ); - let CacheConfigProjection::Native(config) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("GCS cache should be supported"); - }; + let config = native(&facade); let CacheBackendConfig::Gcs(gcs) = config.backend else { panic!("expected GCS configuration"); }; @@ -1417,9 +1511,9 @@ mod tests { path_service_account: Some("credentials.json".into()), endpoint: litellm_cache_gcs::DEFAULT_ENDPOINT.into(), }, + reqwest::Client::new(), Some("token".into()), - ) - .unwrap(); + ); let matching_config = NativeCacheConfig { policy: config.policy, backend: CacheBackendConfig::Gcs(gcs), @@ -1428,62 +1522,194 @@ mod tests { }); } - #[test] - fn rejects_gcs_without_a_bucket_name() { - Python::initialize(); + #[rstest] + #[case::extra_node_field( + "[{'host': 'node-a', 'port': 7000, 'server_type': 'primary'}]", + "client.on_connect", + "native Redis topology is not implemented" + )] + #[case::non_numeric_port( + "[{'host': 'node-a', 'port': 'seven'}]", + "client.on_connect", + "native Redis topology is not implemented" + )] + #[case::empty("[]", "client.on_connect", "native Redis topology is not implemented")] + #[case::foreign_hook( + "[{'host': 'node-a', 'port': 7000}]", + "lambda connection: None", + "native Redis credentials require Python" + )] + fn malformed_startup_nodes_and_foreign_connect_hooks_stay_on_python( + _interpreter: (), + #[case] startup_nodes: &str, + #[case] hook: &str, + #[case] message: &str, + ) { Python::attach(|py| { - let facade = facade( - py, - "backend = SimpleNamespace(bucket_name=None, key_prefix='', path_service_account=None)\n\ - facade = SimpleNamespace(type='gcs', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)", - ); - let CacheConfigProjection::Unsupported(reason) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("GCS cache without a bucket should be unsupported"); - }; - assert!(matches!(&reason, UnsupportedCacheConfig::GcsBucket)); - assert_eq!( - reason.message(), - "native GCS cache requires a configured bucket name" - ); + let facade = cluster_facade(py, startup_nodes, hook); + assert_eq!(unsupported(&facade).message(), message); }); } - #[test] - fn malformed_startup_nodes_and_foreign_connect_hooks_stay_on_python() { - Python::initialize(); - Python::attach(|py| { - for (startup_nodes, hook, message) in [ - ( - "[{'host': 'node-a', 'port': 7000, 'server_type': 'primary'}]", - "client.on_connect", - "native Redis topology is not implemented", - ), - ( - "[{'host': 'node-a', 'port': 'seven'}]", - "client.on_connect", - "native Redis topology is not implemented", - ), - ( - "[]", - "client.on_connect", - "native Redis topology is not implemented", - ), - ( - "[{'host': 'node-a', 'port': 7000}]", - "lambda connection: None", - "native Redis credentials require Python", - ), - ] { - let facade = cluster_facade(py, startup_nodes, hook); - let CacheConfigProjection::Unsupported(reason) = - NativeCacheConfig::project(&facade).unwrap() - else { - panic!("{startup_nodes} with {hook} must stay on Python"); - }; - assert_eq!(reason.message(), message, "{startup_nodes} with {hook}"); + #[rstest] + #[case::plain(|_: &mut RedisConnectionConfig| {}, "redis://cache.internal:6380/4")] + #[case::credentials( + |connection: &mut RedisConnectionConfig| { + connection.username = Some("user".into()); + connection.password = Some("p@ss:word".into()); + }, + "redis://user:p%40ss%3Aword@cache.internal:6380/4" + )] + #[case::password_only( + |connection: &mut RedisConnectionConfig| connection.password = Some("secret".into()), + "redis://:secret@cache.internal:6380/4" + )] + #[case::resp3( + |connection: &mut RedisConnectionConfig| connection.protocol = RedisProtocol::Resp3, + "redis://cache.internal:6380/4?protocol=resp3" + )] + #[case::ipv6( + |connection: &mut RedisConnectionConfig| connection.host = "::1".into(), + "redis://[::1]:6380/4" + )] + #[case::verified_tls( + |connection: &mut RedisConnectionConfig| connection.tls = Some(verified_tls()), + "rediss://cache.internal:6380/4" + )] + #[case::optional_certificate( + |connection: &mut RedisConnectionConfig| { + connection.tls = Some(RedisTlsConfig { + certificate_requirement: CertificateRequirement::Optional, + ..verified_tls() + }); + }, + "rediss://cache.internal:6380/4" + )] + #[case::keepalive_off( + |connection: &mut RedisConnectionConfig| connection.socket_keepalive = Some(false), + "redis://cache.internal:6380/4" + )] + #[case::redis_cache_socket_timeout( + |connection: &mut RedisConnectionConfig| { + connection.read_timeout = Some(Duration::from_secs(5)); + }, + "redis://cache.internal:6380/4" + )] + fn native_url_encodes_the_resolved_connection( + mut connection: RedisConnectionConfig, + #[case] configure: fn(&mut RedisConnectionConfig), + #[case] expected: &str, + ) { + configure(&mut connection); + assert_eq!(connection.native_url().ok().as_deref(), Some(expected)); + } + + #[rstest] + #[case::pool_size( + |connection: &mut RedisConnectionConfig| connection.pool_size = 50, + "native Redis uses a fixed connection pool; max_connections requires Python" + )] + #[case::socket_timeout( + |connection: &mut RedisConnectionConfig| { + connection.read_timeout = Some(Duration::from_millis(100)); + }, + "native Redis uses fixed socket timeouts; socket_timeout and socket_connect_timeout require Python" + )] + #[case::connect_timeout( + |connection: &mut RedisConnectionConfig| { + connection.connect_timeout = Some(Duration::from_secs(1)); + }, + "native Redis uses fixed socket timeouts; socket_timeout and socket_connect_timeout require Python" + )] + #[case::keepalive( + |connection: &mut RedisConnectionConfig| connection.socket_keepalive = Some(true), + "native Redis does not support socket_keepalive" + )] + #[case::health_check( + |connection: &mut RedisConnectionConfig| { + connection.health_check_interval = Duration::from_secs(25); + }, + "native Redis does not support health_check_interval" + )] + #[case::client_name( + |connection: &mut RedisConnectionConfig| connection.client_name = Some("litellm".into()), + "native Redis does not support client_name" + )] + #[case::custom_ca( + |connection: &mut RedisConnectionConfig| { + connection.tls = Some(RedisTlsConfig { + ca_certificate: Some("/ca.pem".into()), + ..verified_tls() + }); + }, + "native Redis TLS does not support ssl_ca_certs, ssl_ca_data, ssl_certfile or ssl_keyfile" + )] + #[case::client_certificate( + |connection: &mut RedisConnectionConfig| { + connection.tls = Some(RedisTlsConfig { + client_certificate: Some("/client.pem".into()), + client_key: Some("/client.key".into()), + ..verified_tls() + }); + }, + "native Redis TLS does not support ssl_ca_certs, ssl_ca_data, ssl_certfile or ssl_keyfile" + )] + #[case::unverified( + |connection: &mut RedisConnectionConfig| { + connection.tls = Some(RedisTlsConfig { + certificate_requirement: CertificateRequirement::None, + check_hostname: false, + ..verified_tls() + }); + }, + "native Redis TLS always verifies the certificate and hostname; ssl_cert_reqs=none and ssl_check_hostname=false require Python" + )] + #[case::hostname_unchecked( + |connection: &mut RedisConnectionConfig| { + connection.tls = Some(RedisTlsConfig { + check_hostname: false, + ..verified_tls() + }); + }, + "native Redis TLS always verifies the certificate and hostname; ssl_cert_reqs=none and ssl_check_hostname=false require Python" + )] + fn native_url_declines_settings_the_native_client_cannot_honor( + mut connection: RedisConnectionConfig, + #[case] configure: fn(&mut RedisConnectionConfig), + #[case] message: &str, + ) { + configure(&mut connection); + let Err(reason) = connection.native_url() else { + panic!("{message}"); + }; + assert_eq!(reason.message(), message); + } + + #[rstest] + #[case::plain("redis://:secret@127.0.0.1:6379", true)] + #[case::database("redis://127.0.0.1:6379/2", true)] + #[case::unix("unix:///tmp/redis.sock", true)] + #[case::tls("rediss://cache.internal:6380", false)] + #[case::query_options("redis://127.0.0.1:6379?socket_timeout=1", false)] + #[case::malformed("not a url", false)] + fn redis_semantic_native_url_accepts_only_plain_urls(#[case] url: &str, #[case] native: bool) { + let config = RedisSemanticCacheConfig { + redis_url: url.into(), + index_name: "idx".into(), + similarity_threshold: 0.8, + }; + match config.native_url() { + Ok(value) => { + assert!(native, "{url} must decline"); + assert_eq!(value, url); } - }); + Err(reason) => { + assert!(!native, "{url} must be native"); + assert_eq!( + reason.message(), + "native Redis semantic cache does not support TLS or query options in redis_url" + ); + } + } } } diff --git a/litellm-rust/crates/python-bridge/src/cache/embedder.rs b/litellm-rust/crates/python-bridge/src/cache/embedder.rs index ffd72e33e1b..7eadd9bc4b0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/embedder.rs +++ b/litellm-rust/crates/python-bridge/src/cache/embedder.rs @@ -90,21 +90,7 @@ impl PythonEmbedder { } } -impl litellm_cache_valkey_semantic::Embedder for PythonEmbedder { - fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { - self.embed_sync(prompt, metadata) - } - - fn async_embed( - &self, - _prompt: &str, - _metadata: Option<&Value>, - ) -> impl Future, Error>> + Send { - std::future::ready(Self::seeded_embedding()) - } -} - -impl litellm_cache_redis_semantic::Embedder for PythonEmbedder { +impl litellm_cache::semantic::Embedder for PythonEmbedder { fn embed(&self, prompt: &str, metadata: Option<&Value>) -> Result, Error> { self.embed_sync(prompt, metadata) } @@ -129,15 +115,14 @@ mod tests { let embedder = PythonEmbedder::new(object); let scoped_embedder = embedder.clone(); let scoped = with_prepared_embedding(Ok(vec![0.25]), async move { - litellm_cache_redis_semantic::Embedder::async_embed(&scoped_embedder, "prompt", None) - .await + litellm_cache::semantic::Embedder::async_embed(&scoped_embedder, "prompt", None).await }); assert_eq!(scoped.await, Ok(vec![0.25])); let unscoped = - litellm_cache_redis_semantic::Embedder::async_embed(&embedder, "prompt", None).await; + litellm_cache::semantic::Embedder::async_embed(&embedder, "prompt", None).await; assert_eq!(unscoped, Err(Error::Unavailable)); let valkey = with_prepared_embedding(Ok(vec![0.5]), async move { - litellm_cache_valkey_semantic::Embedder::async_embed(&embedder, "prompt", None).await + litellm_cache::semantic::Embedder::async_embed(&embedder, "prompt", None).await }); assert_eq!(valkey.await, Ok(vec![0.5])); } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 17fa278ae5e..88fde2f6de0 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -80,6 +80,10 @@ const CLUSTER_POOL: RedisPoolAttributes = RedisPoolAttributes { const VALKEY_POOL: RedisPoolAttributes = STANDALONE_POOL; +/// Class-level defaults an instance overwrites with its own state rather than behavior: +/// `Cache._native_cache` holds the runtime `Cache.__init__` resolved. +const INSTANCE_STATE: &[&str] = &["_native_cache"]; + pub(super) struct FacadeGuard { outer: ObjectGuard, backend: ObjectGuard, @@ -166,7 +170,9 @@ impl ObjectGuard { return Ok(false); } for (name, value) in &expected.attributes { - if (instance.contains(name)? && !self.config_names.contains(&name.as_str())) + if (instance.contains(name)? + && !self.config_names.contains(&name.as_str()) + && !INSTANCE_STATE.contains(&name.as_str())) || !attributes.get_item(name)?.is(value.bind(py)) { return Ok(false); diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 61993f42279..dff8a771a23 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -10,7 +10,6 @@ use pyo3::{ PyTraverseError, PyVisit, exceptions::{PyRuntimeError, PyTypeError}, prelude::*, - types::PyDict, }; use url::Url; @@ -19,6 +18,7 @@ use super::{ config::{QdrantSemanticCacheConfig, project_redis_semantic}, embedder::PythonEmbedder, facade::FacadeGuard, + host_client, native::NativeResponseCache, request::duration, }; @@ -109,7 +109,10 @@ impl CacheTestHandle { ..Default::default() }, }; - let service = run_sync_value(py, async move { Ok(NativeResponseCache::s3(config).await) })?; + let http = host_client(py, ClientVariant::NoRedirect)?; + let service = run_sync_value(py, async move { + Ok(NativeResponseCache::s3(config, http).await) + })?; Ok(Self { service, guard: None, @@ -133,8 +136,8 @@ impl CacheTestHandle { path_service_account, endpoint: endpoint.unwrap_or_else(|| DEFAULT_ENDPOINT.to_string()), }; - let service = release_gil(py, move || NativeResponseCache::gcs(config, token)) - .map_err(cache_error)?; + let client = host_client(py, ClientVariant::NoRedirect)?; + let service = NativeResponseCache::gcs(config, client, token); Ok(Self { service, guard: None, @@ -236,10 +239,7 @@ impl CacheTestHandle { }, quantization, }; - let http_config = crate::http::call_config(py, &PyDict::new(py), true)?; - let client = crate::http::pool() - .client(&http_config, ClientVariant::Provider) - .map_err(crate::http::client_error)?; + let client = host_client(py, ClientVariant::Provider)?; let service = run_sync_value(py, async move { let handle = tokio::runtime::Handle::current(); NativeResponseCache::qdrant_semantic(config, client, handle) @@ -279,8 +279,9 @@ impl CacheTestHandle { #[staticmethod] #[pyo3(signature = (account_url, container))] fn azure_blob(py: Python<'_>, account_url: String, container: String) -> PyResult { + let http = host_client(py, ClientVariant::NoRedirect)?; let service = run_sync_value(py, async move { - NativeResponseCache::azure_blob(&account_url, &container) + NativeResponseCache::azure_blob(&account_url, &container, http) .await .map_err(cache_error) })?; diff --git a/litellm-rust/crates/python-bridge/src/cache/mod.rs b/litellm-rust/crates/python-bridge/src/cache/mod.rs index 28dd6c3e798..0cfd4ac8138 100644 --- a/litellm-rust/crates/python-bridge/src/cache/mod.rs +++ b/litellm-rust/crates/python-bridge/src/cache/mod.rs @@ -1,3 +1,4 @@ +mod activation; mod binding; mod callback; mod config; @@ -12,9 +13,11 @@ mod resolver; mod semantic; use litellm_cache::Error; +use litellm_http::ClientVariant; use pyo3::{ exceptions::{PyNotImplementedError, PyRuntimeError, PyValueError}, prelude::*, + types::PyDict, }; pub(crate) use self::{ @@ -28,3 +31,11 @@ fn cache_error(error: Error) -> PyErr { _ => PyRuntimeError::new_err(error.to_string()), } } + +/// The host's pooled HTTP client, configured from the proxy's HTTP settings. +fn host_client(py: Python<'_>, variant: ClientVariant) -> PyResult { + let http_config = crate::http::call_config(py, &PyDict::new(py), true)?; + crate::http::pool() + .client(&http_config, variant) + .map_err(crate::http::client_error) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 254b9cdea4d..3c260e70843 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,15 +1,16 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; +use litellm_cache::{CacheCodec, CacheConnectionResult, Error, semantic::SemanticLookup}; use litellm_cache_azure_blob::AzureBlobCache; use litellm_cache_disk::DiskCache; use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource}; use litellm_cache_memory::InMemoryCache; -use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, QdrantSemanticCache}; +use litellm_cache_qdrant_semantic::{OpenAiEmbedder, QdrantSemanticCache}; use litellm_cache_redis::{RedisCache, RedisTopology}; use litellm_cache_redis_semantic::{RedisSemanticCache, RedisSemanticConfig}; use litellm_cache_response::{ - ExactResponseCache, PartialHits, ResponseCache, ResponseCacheCodec, WriteBuffer, + ConnectionProbe, ExactResponseCache, PartialHits, ResponseCache, ResponseCacheCodec, + WriteBuffer, }; use litellm_cache_s3::{S3Cache, S3CacheConfig}; use litellm_cache_valkey_semantic::{ValkeySemanticCache, ValkeySemanticConfig}; @@ -33,20 +34,11 @@ pub(super) struct EmbeddingInput { /// An exact-match backend behind one pointer, with the identity its facade must reproduce. pub(super) struct ExactService { cache: Arc, + probe: Option>, buffer: Option, identity: BackendIdentity, } -impl ExactService { - fn new(cache: Arc, identity: BackendIdentity) -> Arc { - Arc::new(Self { - cache, - buffer: None, - identity, - }) - } -} - #[derive(Clone)] pub(super) enum NativeResponseCache { Exact(Arc), @@ -56,7 +48,7 @@ pub(super) enum NativeResponseCache { scope: String, }, RedisSemantic { - cache: Arc>>, + cache: Arc>>, embedder: PythonEmbedder, }, QdrantSemantic(Arc>>), @@ -94,12 +86,15 @@ impl NativeResponseCache { namespace: backend.namespace().map(str::to_owned), default_ttl: None, }; - Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) + Ok(Self::exact_probed( + ResponseCache::new(Arc::new(backend)), + identity, + )) } - pub async fn s3(config: S3CacheConfig) -> Self { + pub async fn s3(config: S3CacheConfig, http: reqwest::Client) -> Self { let runtime = tokio::runtime::Handle::current(); - let backend = S3Cache::new(config, ResponseCacheCodec, runtime); + let backend = S3Cache::new(config, http, ResponseCacheCodec, runtime); let identity = BackendIdentity::S3 { bucket: backend.bucket().to_owned(), key_prefix: backend.key_prefix().to_owned(), @@ -109,7 +104,7 @@ impl NativeResponseCache { Self::exact(ResponseCache::new(Arc::new(backend)), identity) } - pub fn disk(directory: &str) -> Result { + pub fn disk(directory: impl AsRef) -> Result { let backend = DiskCache::open(directory, ResponseCacheCodec)?; let identity = BackendIdentity::Disk { directory: backend.directory().to_path_buf(), @@ -117,27 +112,33 @@ impl NativeResponseCache { Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) } - pub fn gcs(config: GcsConfig, token: Option) -> Result { + pub fn gcs(config: GcsConfig, client: reqwest::Client, token: Option) -> Self { let backend = match token { Some(token) => GcsCache::with_token_source( config, + client, ResponseCacheCodec, Arc::new(StaticTokenSource(token)), - )?, - None => GcsCache::new(config, ResponseCacheCodec)?, + ), + None => GcsCache::new(config, client, ResponseCacheCodec), }; let identity = BackendIdentity::Gcs { bucket_name: backend.bucket_name().to_owned(), key_prefix: backend.key_prefix().to_owned(), path_service_account: backend.path_service_account().map(str::to_owned), }; - Ok(Self::exact(ResponseCache::new(Arc::new(backend)), identity)) + Self::exact(ResponseCache::new(Arc::new(backend)), identity) } - pub async fn azure_blob(account_url: &str, container: &str) -> Result { + pub async fn azure_blob( + account_url: &str, + container: &str, + http: reqwest::Client, + ) -> Result { let backend = AzureBlobCache::connect( account_url, container, + http, ResponseCacheCodec, tokio::runtime::Handle::current(), ) @@ -156,7 +157,25 @@ impl NativeResponseCache { B: litellm_cache::BaseCache, B::Context: Default + PartialEq, { - let cache: Arc = Arc::new(cache); + Self::exact_service(Arc::new(cache), None, identity) + } + + /// Wraps an exact backend whose Python class defines `test_connection`. + fn exact_probed(cache: ResponseCache, identity: BackendIdentity) -> Self + where + ResponseCache: ExactResponseCache + ConnectionProbe + 'static, + B: litellm_cache::BaseCache, + B::Context: Default + PartialEq, + { + let cache = Arc::new(cache); + Self::exact_service(cache.clone(), Some(cache), identity) + } + + fn exact_service( + cache: Arc, + probe: Option>, + identity: BackendIdentity, + ) -> Self { let default_ttl = cache.default_ttl(); let identity = match identity { BackendIdentity::Memory { @@ -179,7 +198,12 @@ impl NativeResponseCache { }, other => other, }; - Self::Exact(ExactService::new(cache, identity)) + Self::Exact(Arc::new(ExactService { + cache, + probe, + buffer: None, + identity, + })) } pub fn valkey_semantic( @@ -209,7 +233,7 @@ impl NativeResponseCache { embedder: PythonEmbedder, config: RedisSemanticConfig, ) -> Result { - let backend = RedisSemanticCache::new(url, embedder.clone(), config)?; + let backend = RedisSemanticCache::new(url, embedder.clone(), ResponseCacheCodec, config)?; Ok(Self::RedisSemantic { cache: Arc::new(ResponseCache::new(Arc::new(backend))), embedder, @@ -270,6 +294,7 @@ impl NativeResponseCache { Self::Exact(service) if matches!(service.identity, BackendIdentity::Redis { .. }) => { Self::Exact(Arc::new(ExactService { cache: Arc::clone(&service.cache), + probe: service.probe.clone(), buffer: flush_size.map(WriteBuffer::new), identity: service.identity.clone(), })) @@ -305,7 +330,7 @@ impl NativeResponseCache { Self::RedisSemantic { .. } => request.semantic().context, Self::Exact(_) | Self::QdrantSemantic(_) => return None, }; - let prompt = litellm_cache_redis_semantic::prompt_from_context(&context)?; + let prompt = litellm_cache::semantic::prompt_from_context(&context)?; Some(EmbeddingInput { prompt, metadata: context.metadata, @@ -344,6 +369,28 @@ impl NativeResponseCache { } } + /// `lookup` plus the similarity Python's semantic backend writes to the request metadata. + /// Exact backends report none. + pub fn lookup_semantic( + &self, + request: &NativeRequest, + now: Duration, + ) -> Result, Error> { + match self { + Self::Exact(service) => service + .cache + .lookup(&request.exact(), now) + .map(exact_lookup), + Self::ValkeySemantic { cache, scope, .. } => { + redis_family(cache.lookup_semantic(&request.scoped_semantic(scope), now)) + } + Self::RedisSemantic { cache, .. } => { + redis_family(cache.lookup_semantic(&request.semantic(), now)) + } + Self::QdrantSemantic(cache) => cache.lookup_semantic(&request.semantic(), now), + } + } + pub fn store( &self, request: &NativeRequest, @@ -390,6 +437,56 @@ impl NativeResponseCache { } } + pub async fn async_lookup_semantic( + &self, + request: &NativeRequest, + now: Duration, + ) -> Result, Error> { + match self { + Self::Exact(service) => service + .cache + .async_lookup(&request.exact(), now) + .await + .map(exact_lookup), + Self::ValkeySemantic { cache, scope, .. } => redis_family( + cache + .async_lookup_semantic(&request.scoped_semantic(scope), now) + .await, + ), + Self::RedisSemantic { cache, .. } => { + redis_family(cache.async_lookup_semantic(&request.semantic(), now).await) + } + Self::QdrantSemantic(cache) => { + cache.async_lookup_semantic(&request.semantic(), now).await + } + } + } + + pub(super) fn async_lookup_semantic_py<'py>( + &self, + py: Python<'py>, + request: NativeRequest, + ) -> PyResult> { + match self { + Self::Exact(_) | Self::QdrantSemantic(_) => { + let service = self.clone(); + litellm_host_python::run_async( + py, + async move { + service + .async_lookup_semantic(&request, now()) + .await + .map(SemanticReply::from) + }, + super::cache_error, + ) + } + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + self.python_semantic(py, SemanticOperation::LookupSemantic(request)) + } + } + } + pub(super) fn async_lookup_py<'py>( &self, py: Python<'py>, @@ -550,9 +647,11 @@ impl NativeResponseCache { pub async fn test_connection(&self) -> Result { match self { - Self::Exact(service) => service.cache.test_connection().await, - Self::ValkeySemantic { cache, .. } => cache.test_connection().await, - Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { + Self::Exact(service) => match &service.probe { + Some(probe) => probe.test_connection().await, + None => Err(Error::UnsupportedOperation), + }, + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { Err(Error::UnsupportedOperation) } } @@ -571,3 +670,27 @@ impl NativeResponseCache { fn exact_requests(requests: &[NativeRequest]) -> Vec { requests.iter().map(NativeRequest::exact).collect() } + +/// What `lookup_semantic` hands Python: the response and the similarity to stamp, if any. +#[derive(serde::Serialize)] +pub(super) struct SemanticReply(pub(super) Option, pub(super) Option); + +impl From> for SemanticReply { + fn from(lookup: SemanticLookup) -> Self { + Self(lookup.value, lookup.similarity) + } +} + +fn exact_lookup(value: Option) -> SemanticLookup { + SemanticLookup { + value, + similarity: None, + } +} + +/// Python's Redis and Valkey semantic caches catch every lookup failure and stamp `0.0`. +fn redis_family( + lookup: Result, Error>, +) -> Result, Error> { + Ok(lookup.unwrap_or_else(|_| SemanticLookup::miss(Some(0.0)))) +} diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/semantic.rs index 9f4d18d45cd..bea58b88885 100644 --- a/litellm-rust/crates/python-bridge/src/cache/semantic.rs +++ b/litellm-rust/crates/python-bridge/src/cache/semantic.rs @@ -12,12 +12,14 @@ use serde_json::Value; use super::{ cache_error, embedder::{PythonEmbedder, with_prepared_embedding}, - native::NativeResponseCache, + native::{NativeResponseCache, SemanticReply}, request::{NativeRequest, now}, }; pub(super) enum SemanticOperation { Lookup(NativeRequest), + /// A lookup that also reports the similarity, as `SemanticReply`. + LookupSemantic(NativeRequest), Store(NativeRequest, Value), StoreBatch(VecDeque<(NativeRequest, Value)>), } @@ -71,7 +73,9 @@ impl SemanticExecution { /// Takes the next entry of the operation; `None` once a batch is exhausted. fn next_pending(&mut self) -> Option<(NativeRequest, Option)> { match &mut self.operation { - SemanticOperation::Lookup(request) => Some((request.clone(), None)), + SemanticOperation::Lookup(request) | SemanticOperation::LookupSemantic(request) => { + Some((request.clone(), None)) + } SemanticOperation::Store(request, response) => { Some((request.clone(), Some(std::mem::take(response)))) } @@ -109,7 +113,7 @@ impl SemanticExecution { Ok(vector) => { PythonEmbedder::extract(vector.into_bound(py)).map_err(|_| Error::Unavailable) } - Err(error) => match self.failure { + Err(error) => match self.embedding_failure() { EmbeddingFailure::Propagate => return Err(error), EmbeddingFailure::Unavailable if error.is_instance_of::(py) => { Err(Error::Unavailable) @@ -120,6 +124,14 @@ impl SemanticExecution { self.backend_step(py, seed) } + /// Python's semantic lookups catch embedding errors and stamp a similarity of `0.0`. + fn embedding_failure(&self) -> EmbeddingFailure { + match self.operation { + SemanticOperation::LookupSemantic(_) => EmbeddingFailure::Unavailable, + _ => self.failure, + } + } + fn backend_step( &mut self, py: Python<'_>, @@ -131,13 +143,18 @@ impl SemanticExecution { })?; let service = self.service.clone(); let now = self.now; + let with_similarity = matches!(self.operation, SemanticOperation::LookupSemantic(_)); let future = async move { match response { - None => service.async_lookup(&request, now).await, + None if with_similarity => service + .async_lookup_semantic(&request, now) + .await + .map(|lookup| Reply::Semantic(lookup.into())), + None => service.async_lookup(&request, now).await.map(Reply::Plain), Some(response) => service .async_store(&request, response, now) .await - .map(|_| None), + .map(|_| Reply::Plain(None)), } }; let awaitable = run_async(py, with_prepared_embedding(seed, future), cache_error)?; @@ -171,6 +188,13 @@ impl SemanticExecution { } } +#[derive(serde::Serialize)] +#[serde(untagged)] +enum Reply { + Plain(Option), + Semantic(SemanticReply), +} + impl ExecutionBody for SemanticExecution { fn resume(&mut self, result: Option>>) -> PyResult { Python::attach(|py| self.resume_py(py, result)) diff --git a/litellm/caching/caching.py b/litellm/caching/caching.py index 6a98b104221..6f99eb616b0 100644 --- a/litellm/caching/caching.py +++ b/litellm/caching/caching.py @@ -15,7 +15,8 @@ import time import traceback from collections.abc import Mapping from enum import Enum -from typing import Any, Final +from types import MappingProxyType +from typing import TYPE_CHECKING, Any, Final from pydantic import BaseModel @@ -38,6 +39,25 @@ from .redis_cluster_cache import RedisClusterCache from .redis_semantic_cache import RedisSemanticCache from .s3_cache import S3Cache +if TYPE_CHECKING: + from litellm.rust_bridge.response_cache import NativeCacheRequest, ResponseCacheRuntime + + +def _native_response(result: object) -> object: + """The value Python's own reader would return for `result` once it is cached. + + Python stores a model as its JSON text and `json.loads` a string response on read, so the + native store receives the decoded value and writes the envelope shape Python reads. + """ + if isinstance(result, BaseModel): + return json.loads(result.model_dump_json()) + if isinstance(result, str): + try: + return json.loads(result) + except ValueError: + return result + return result + def print_verbose(print_statement): try: @@ -55,6 +75,8 @@ class CacheMode(str, Enum): #### LiteLLM.Completion / Embedding Cache #### class Cache: + _native_cache: "ResponseCacheRuntime | None" = None + def __init__( self, type: LiteLLMCacheType | None = LiteLLMCacheType.LOCAL, @@ -292,6 +314,12 @@ class Cache: if self.namespace is not None and isinstance(self.cache, RedisCache): self.cache.namespace = self.namespace + from litellm.rust_bridge.response_cache import resolve_response_cache + + # The Rust catalog picks the store per backend. When it selects Rust, the storage calls + # below go to the native runtime and the Python backend stays only for its direct API. + self._native_cache = resolve_response_cache(self) + # Params whose values carry prompt content. Excluded from semantic-cache # scope keys so differently worded prompts share a bucket and match via # vector similarity rather than being split into per-wording buckets. @@ -570,6 +598,13 @@ class Cache: if "semantic-similarity" in cache_lookup_metadata: original_metadata["semantic-similarity"] = cache_lookup_metadata["semantic-similarity"] + @staticmethod + def _stamp_semantic_similarity(kwargs: Mapping[str, object], similarity: float | None) -> None: + """Write a native semantic lookup's similarity where the Python backends put it.""" + metadata: Final = kwargs.get("metadata") + if similarity is not None and isinstance(metadata, dict): + metadata["semantic-similarity"] = similarity + def get_cache(self, dynamic_cache_object: BaseCache | None = None, **kwargs): """ Retrieves the cached result for the given arguments. @@ -588,6 +623,15 @@ class Cache: cache_key = kwargs["cache_key"] else: cache_key = self.get_cache_key(**kwargs) + if cache_key is not None and self._native_cache is not None: + request = self._native_cache.request(self, MappingProxyType({**kwargs, "cache_key": cache_key})) + if request is None: + return None + if not self._is_semantic_cache(): + return self._native_cache.lookup(request) + response, similarity = self._native_cache.lookup_semantic(request) + self._stamp_semantic_similarity(kwargs, similarity) + return response if cache_key is not None: cache_control_args: Final[DynamicCacheControl] = kwargs.get("cache", {}) max_age = cache_control_args.get("s-maxage") or cache_control_args.get("s-max-age") or float("inf") @@ -620,6 +664,15 @@ class Cache: cache_key = kwargs["cache_key"] else: cache_key = self.get_cache_key(**kwargs) + if cache_key is not None and self._native_cache is not None: + request = self._native_cache.request(self, MappingProxyType({**kwargs, "cache_key": cache_key})) + if request is None: + return None + if not self._is_semantic_cache(): + return await self._native_cache.async_lookup(request) + response, similarity = await self._native_cache.async_lookup_semantic(request) + self._stamp_semantic_similarity(kwargs, similarity) + return response if cache_key is not None: cache_control_args: Final = kwargs.get("cache", {}) max_age: Final = cache_control_args.get("s-max-age", cache_control_args.get("s-maxage", float("inf"))) @@ -676,6 +729,11 @@ class Cache: try: if self.should_use_cache(**kwargs) is not True: return + if self._native_cache is not None: + request = self._native_request(kwargs) + if request is not None: + self._native_cache.store(request, _native_response(result)) + return cache_key, cached_data, kwargs = self._add_cache_logic(result=result, **kwargs) self.cache.set_cache(cache_key, cached_data, **kwargs) except Exception as e: @@ -695,6 +753,11 @@ class Cache: try: if self.should_use_cache(**kwargs) is not True: return + if self._native_cache is not None: + request = self._native_request(kwargs) + if request is not None: + await self._native_cache.async_store(request, _native_response(result)) + return if self.type == "redis" and self.redis_flush_size is not None: # high traffic - fill in results in memory and then flush await self.batch_cache_write(result, **kwargs) @@ -879,13 +942,35 @@ class Cache: cache_key, cached_data, kwargs = self.add_embedding_response_to_cache(result, kwargs["input"], kwargs) cache_list.append((cache_key, cached_data)) - if dynamic_cache_object is not None: + if self._native_cache is not None: + entries: Final = tuple( + (request, cached_data["response"]) + for cache_key, cached_data in cache_list + if (request := self._native_request(MappingProxyType({**kwargs, "cache_key": cache_key}))) + is not None + ) + await self._native_cache.async_store_batch( + tuple(request for request, _ in entries), + tuple(response for _, response in entries), + ) + elif dynamic_cache_object is not None: await dynamic_cache_object.async_set_cache_pipeline(cache_list=cache_list, **kwargs) else: await self.cache.async_set_cache_pipeline(cache_list=cache_list, **kwargs) except Exception as e: self._log_add_cache_failure(e) + def _native_request(self, kwargs: Mapping[str, object]) -> "NativeCacheRequest | None": + if self._native_cache is None: + return None + cache_key: Final = kwargs.get("cache_key") + return self._native_cache.request( + self, + kwargs + if isinstance(cache_key, str) + else MappingProxyType({**kwargs, "cache_key": self.get_cache_key(**kwargs)}), + ) + def should_use_cache(self, **kwargs): """ Returns true if we should use the cache for LLM API calls diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index 61e597bf674..f794dc4a99e 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -105,6 +105,7 @@ class _ResponseCacheRuntime: *, callback_kwargs: Mapping[str, object] | Sequence[object] | None = None, ) -> object: ... + def lookup_semantic(self, request: object) -> tuple[object, float | None]: ... def store( self, request: object, @@ -124,6 +125,7 @@ class _ResponseCacheRuntime: *, callback_kwargs: Mapping[str, object] | None = None, ) -> Future[object]: ... + def async_lookup_semantic(self, request: object) -> Future[tuple[object, float | None]]: ... def async_store( self, request: object, diff --git a/litellm/rust_bridge/response_cache.py b/litellm/rust_bridge/response_cache.py index 82d27fce27b..a6fc121a3b5 100644 --- a/litellm/rust_bridge/response_cache.py +++ b/litellm/rust_bridge/response_cache.py @@ -46,9 +46,11 @@ class NativeResponseCacheRuntime(Protocol): def kind(self) -> str: ... def lookup(self, request: NativeCacheRequest) -> object: ... + def lookup_semantic(self, request: NativeCacheRequest) -> tuple[object, float | None]: ... def store(self, request: NativeCacheRequest, response: object) -> None: ... def lookup_batch(self, requests: Sequence[NativeCacheRequest]) -> object: ... def async_lookup(self, request: NativeCacheRequest) -> Awaitable[object]: ... + def async_lookup_semantic(self, request: NativeCacheRequest) -> Awaitable[tuple[object, float | None]]: ... def async_store(self, request: NativeCacheRequest, response: object) -> Awaitable[None]: ... def async_lookup_batch(self, requests: Sequence[NativeCacheRequest]) -> Awaitable[object]: ... def async_store_batch( @@ -108,6 +110,11 @@ class ResponseCacheRuntime: def lookup(self, request: NativeCacheRequest) -> object: return self.native.lookup(request) + def lookup_semantic(self, request: NativeCacheRequest) -> tuple[object, float | None]: + """The cached response and the similarity a semantic backend reports, if any.""" + response, similarity = self.native.lookup_semantic(request) + return response, similarity + def store(self, request: NativeCacheRequest, response: object) -> None: self.native.store(request, response) @@ -117,6 +124,10 @@ class ResponseCacheRuntime: async def async_lookup(self, request: NativeCacheRequest) -> object: return await self.native.async_lookup(request) + async def async_lookup_semantic(self, request: NativeCacheRequest) -> tuple[object, float | None]: + response, similarity = await self.native.async_lookup_semantic(request) + return response, similarity + async def async_store(self, request: NativeCacheRequest, response: object) -> None: await self.native.async_store(request, response) diff --git a/tests/test_litellm_rust/messages/test_callbacks.py b/tests/test_litellm_rust/messages/test_callbacks.py index b55bc47d640..d5dc5437398 100644 --- a/tests/test_litellm_rust/messages/test_callbacks.py +++ b/tests/test_litellm_rust/messages/test_callbacks.py @@ -5,7 +5,11 @@ import pytest import litellm from litellm.integrations.custom_logger import CustomLogger +from litellm.rust_bridge import catalog +from litellm.rust_bridge.catalog import Route, RouteRule +from litellm.rust_bridge.configuration import Rollout from tests.test_litellm_rust.support.callback_recorder import RecordingLogger, drain_logging +from tests.test_litellm_rust.support.isolation import rebound from tests.test_litellm_rust.support.recording_server import RecordingServer, ResponseSpec from tests.test_litellm_rust.support.requests import ( MESSAGES, @@ -20,6 +24,12 @@ pytestmark = pytest.mark.requires_rust_extension STREAM: Final = ResponseSpec(body=None, events=MESSAGES_EVENTS) +@pytest.fixture(autouse=True) +def opt_messages_into_rust() -> Iterator[None]: + with rebound(catalog, "RULES", (RouteRule(Route.MESSAGES, Rollout.RUST_OPT_IN), *catalog.RULES)): + yield + + @pytest.fixture def messages_server(recording_server: RecordingServer) -> RecordingServer: recording_server.default_response = ResponseSpec(body=MESSAGES_RESPONSE) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 0f389edaa27..e2e2f9f1819 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -15,7 +15,7 @@ from contextlib import ExitStack from datetime import datetime from pathlib import Path from types import SimpleNamespace -from typing import Final, Protocol, cast +from typing import Final, Protocol, TypeAlias, cast from unittest.mock import Mock from urllib.parse import urlparse from uuid import uuid4 @@ -37,10 +37,10 @@ from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.caching.redis_semantic_cache import RedisSemanticCache from litellm.caching.s3_cache import S3Cache -from litellm.rust_bridge import _native +from litellm.rust_bridge import _native, catalog from litellm.rust_bridge.catalog import CacheRule, Route, RouteRule, SecretManagerRule from litellm.rust_bridge.configuration import Rollout -from litellm.rust_bridge.response_cache import ResponseCacheRuntime, resolve_response_cache +from litellm.rust_bridge.response_cache import NativeResponseCacheRuntime, ResponseCacheRuntime, resolve_response_cache from litellm.types.caching import LiteLLMCacheType from litellm.types.llms.custom_llm import CustomLLMItem from litellm.types.utils import EmbeddingResponse @@ -1843,9 +1843,7 @@ async def test_qdrant_semantic_async_parity(qdrant_url: str, fake_embedding_endp assert python_value["response"] == {"id": "native"} -async def test_qdrant_semantic_async_store_batch_shares_entries( - qdrant_url: str, fake_embedding_endpoint: str -) -> None: +async def test_qdrant_semantic_async_store_batch_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: del fake_embedding_endpoint collection: Final = f"cache_{uuid4().hex}" facade: Final = qdrant_facade(qdrant_url, collection) @@ -1865,14 +1863,12 @@ async def test_qdrant_semantic_async_store_batch_shares_entries( assert binding.lookup(entries[0]) == {"id": "one"} assert binding.lookup(entries[1]) == {"id": "two"} - assert ( - (await facade.cache.async_get_cache("batch-one", messages=entries[0]["messages"]))["response"] - == {"id": "one"} - ) - assert ( - (await facade.cache.async_get_cache("batch-two", messages=entries[1]["messages"]))["response"] - == {"id": "two"} - ) + assert (await facade.cache.async_get_cache("batch-one", messages=entries[0]["messages"]))["response"] == { + "id": "one" + } + assert (await facade.cache.async_get_cache("batch-two", messages=entries[1]["messages"]))["response"] == { + "id": "two" + } async def test_qdrant_semantic_malformed_entries_and_unsupported_operations( @@ -1962,3 +1958,389 @@ def test_qdrant_semantic_mutation_and_projection_fallback(qdrant_url: str, fake_ unsupported.cache.qdrant_api_base = "http://127.0.0.1:7777" with pytest.raises(TypeError, match="gRPC"): handle._bind_facade(unsupported) + + +CacheFactory: TypeAlias = Callable[[], Cache] + + +def require_rust(monkeypatch: pytest.MonkeyPatch, backend: LiteLLMCacheType) -> None: + monkeypatch.setattr(catalog, "RULES", (CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({backend})),)) + + +def native_runtime(facade: Cache) -> ResponseCacheRuntime: + runtime: Final = facade._native_cache # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor + assert isinstance(runtime, ResponseCacheRuntime) + assert runtime.kind == "native" + return runtime + + +@pytest.fixture +def cache_factory(request: pytest.FixtureRequest, tmp_path: Path) -> CacheFactory: + backend: Final = cast(LiteLLMCacheType, request.param) + match backend: + case LiteLLMCacheType.LOCAL: + return lambda: Cache(type=backend) + case LiteLLMCacheType.DISK: + return lambda: Cache(type=backend, disk_cache_dir=str(tmp_path)) + case LiteLLMCacheType.REDIS: + parsed: Final = urlparse(cast(str, request.getfixturevalue("redis_url"))) + return lambda: Cache(type=backend, host=parsed.hostname, port=str(parsed.port)) + case LiteLLMCacheType.S3: + stub: Final = cast(S3Stub, request.getfixturevalue("s3_stub")) + return lambda: Cache( + type=backend, + s3_bucket_name="cache-bucket", + s3_region_name="us-east-1", + s3_endpoint_url=stub.url, + s3_aws_access_key_id="key", + s3_aws_secret_access_key="secret", + s3_path="team", + ) + case LiteLLMCacheType.GCS: + return lambda: Cache(type=backend, gcs_bucket_name="bucket", gcs_path="cache/") + case LiteLLMCacheType.REDIS_SEMANTIC: + return lambda: Cache( + type=backend, + redis_url="redis://127.0.0.1:6379", + similarity_threshold=0.8, + redis_semantic_cache_embedding_model="text-embedding-3-small", + ) + case LiteLLMCacheType.VALKEY_SEMANTIC: + return lambda: Cache(type=backend, redis_url="redis://127.0.0.1:6390/0", similarity_threshold=0.8) + case _: + raise AssertionError(f"no local factory for {backend}") + + +ROUND_TRIP_BACKENDS: Final = ( + LiteLLMCacheType.LOCAL, + LiteLLMCacheType.DISK, + LiteLLMCacheType.REDIS, + LiteLLMCacheType.S3, +) +SHARED_STORE_BACKENDS: Final = (LiteLLMCacheType.DISK, LiteLLMCacheType.REDIS, LiteLLMCacheType.S3) + + +def completion_kwargs(label: str) -> dict[str, object]: + return {"model": "gpt-4o", "messages": [{"role": "user", "content": f"{label} {uuid4().hex}"}]} + + +@pytest.mark.parametrize("backend", list(LiteLLMCacheType)) +def test_shipped_rules_keep_every_backend_on_python(backend: LiteLLMCacheType) -> None: + assert resolve_response_cache(cast(Cache, SimpleNamespace(type=backend))) is None + + +@pytest.mark.parametrize( + "cache_factory", + [ + LiteLLMCacheType.LOCAL, + LiteLLMCacheType.DISK, + LiteLLMCacheType.REDIS, + LiteLLMCacheType.S3, + LiteLLMCacheType.GCS, + LiteLLMCacheType.REDIS_SEMANTIC, + LiteLLMCacheType.VALKEY_SEMANTIC, + ], + indirect=True, +) +def test_shipped_rules_construct_python_backed_facades(cache_factory: CacheFactory) -> None: + assert cache_factory()._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor + + +@pytest.mark.parametrize( + "cache_factory", + [ + LiteLLMCacheType.LOCAL, + LiteLLMCacheType.DISK, + LiteLLMCacheType.REDIS, + LiteLLMCacheType.S3, + LiteLLMCacheType.GCS, + LiteLLMCacheType.REDIS_SEMANTIC, + LiteLLMCacheType.VALKEY_SEMANTIC, + ], + indirect=True, +) +def test_rust_required_rule_activates_the_native_backend( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + native_runtime(cache_factory()) + + +@pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) +async def test_facade_storage_calls_round_trip_through_the_native_backend( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + facade: Final = cache_factory() + native_runtime(facade) + + sync_kwargs: Final = completion_kwargs("sync") + facade.add_cache({"answer": 1}, **sync_kwargs) + assert facade.get_cache(**sync_kwargs) == {"answer": 1} + + async_kwargs: Final = completion_kwargs("async") + await facade.async_add_cache({"answer": 2}, **async_kwargs) + assert await facade.async_get_cache(**async_kwargs) == {"answer": 2} + assert facade.get_cache(**completion_kwargs("absent")) is None + + +async def test_memory_facade_writes_bypass_the_python_backend(monkeypatch: pytest.MonkeyPatch) -> None: + require_rust(monkeypatch, LiteLLMCacheType.LOCAL) + facade: Final = Cache(type=LiteLLMCacheType.LOCAL) + native_runtime(facade) + kwargs: Final = completion_kwargs("memory") + facade.add_cache({"answer": 1}, **kwargs) + assert facade.cache.get_cache(facade.get_cache_key(**kwargs)) is None + assert facade.get_cache(**kwargs) == {"answer": 1} + + +@pytest.mark.parametrize("cache_factory", SHARED_STORE_BACKENDS, indirect=True) +async def test_native_and_python_facades_share_one_wire_format( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + python_facade: Final = cache_factory() + assert python_facade._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + native_facade: Final = cache_factory() + native_runtime(native_facade) + + native_written: Final = completion_kwargs("native") + native_facade.add_cache({"writer": "native"}, **native_written) + assert python_facade.get_cache(**native_written) == {"writer": "native"} + + python_written: Final = completion_kwargs("python") + python_facade.add_cache({"writer": "python"}, **python_written) + assert native_facade.get_cache(**python_written) == {"writer": "python"} + + async_native: Final = completion_kwargs("async-native") + await native_facade.async_add_cache({"writer": "async-native"}, **async_native) + assert await python_facade.async_get_cache(**async_native) == {"writer": "async-native"} + + async_python: Final = completion_kwargs("async-python") + await python_facade.async_add_cache({"writer": "async-python"}, **async_python) + assert await native_facade.async_get_cache(**async_python) == {"writer": "async-python"} + + +@pytest.mark.parametrize("cache_factory", ROUND_TRIP_BACKENDS, indirect=True) +async def test_embedding_pipeline_stores_one_native_entry_per_input( + cache_factory: CacheFactory, monkeypatch: pytest.MonkeyPatch, request: pytest.FixtureRequest +) -> None: + require_rust(monkeypatch, cast(LiteLLMCacheType, request.node.callspec.params["cache_factory"])) + facade: Final = cache_factory() + native_runtime(facade) + inputs: Final = [f"alpha {uuid4().hex}", f"beta {uuid4().hex}"] + result: Final = EmbeddingResponse( + model="text-embedding-3-small", + data=[ + {"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}, + {"object": "embedding", "index": 1, "embedding": [0.3, 0.4]}, + ], + ) + await facade.async_add_cache_pipeline(result, model="text-embedding-3-small", input=inputs) + + keys: Final = [facade.get_cache_key(model="text-embedding-3-small", input=text) for text in inputs] + assert len(set(keys)) == len(inputs) + for text, expected in zip(inputs, ([0.1, 0.2], [0.3, 0.4]), strict=True): + cached = await facade.async_get_cache(model="text-embedding-3-small", input=text) + assert isinstance(cached, dict) + assert cached["embedding"] == expected + assert await facade.async_get_cache(model="text-embedding-3-small", input=inputs) is None + + +def redis_facade(redis_url: str, **settings: object) -> Cache: + parsed: Final = urlparse(redis_url) + return Cache(type=LiteLLMCacheType.REDIS, host=parsed.hostname, port=str(parsed.port), **settings) + + +@pytest.mark.parametrize( + ("settings", "message"), + [ + pytest.param({"max_connections": 10}, "max_connections requires Python", id="pool-size"), + pytest.param({"socket_timeout": 1.0}, "socket_timeout and socket_connect_timeout", id="socket-timeout"), + pytest.param( + {"socket_connect_timeout": 1.0}, "socket_timeout and socket_connect_timeout", id="connect-timeout" + ), + pytest.param({"socket_keepalive": True}, "does not support socket_keepalive", id="keepalive"), + pytest.param({"health_check_interval": 5}, "does not support health_check_interval", id="health-check"), + pytest.param({"client_name": "litellm"}, "does not support client_name", id="client-name"), + pytest.param({"ssl": True}, "ssl_check_hostname=false require Python", id="tls-default-hostname-check"), + pytest.param({"ssl": True, "ssl_cert_reqs": "none"}, "ssl_cert_reqs=none", id="tls-without-verification"), + pytest.param( + {"ssl": True, "ssl_check_hostname": True, "ssl_ca_certs": "/ca.pem"}, + "does not support ssl_ca_certs", + id="tls-custom-ca", + ), + pytest.param( + {"ssl": True, "ssl_check_hostname": True, "ssl_certfile": "/client.pem", "ssl_keyfile": "/client.key"}, + "does not support ssl_ca_certs, ssl_ca_data, ssl_certfile or ssl_keyfile", + id="tls-client-certificate", + ), + ], +) +def test_redis_settings_the_native_client_cannot_honor_decline( + redis_url: str, monkeypatch: pytest.MonkeyPatch, settings: dict[str, object], message: str +) -> None: + require_rust(monkeypatch, LiteLLMCacheType.REDIS) + with pytest.raises(RuntimeError, match=f"declined the cache: native Redis.*{message}"): + redis_facade(redis_url, **settings) + + +def test_redis_verified_tls_activates_natively(redis_url: str, monkeypatch: pytest.MonkeyPatch) -> None: + require_rust(monkeypatch, LiteLLMCacheType.REDIS) + native_runtime(redis_facade(redis_url, ssl=True, ssl_check_hostname=True)) + + +async def test_redis_flush_size_buffers_native_facade_writes(redis_url: str, monkeypatch: pytest.MonkeyPatch) -> None: + require_rust(monkeypatch, LiteLLMCacheType.REDIS) + facade: Final = redis_facade(redis_url, redis_flush_size=2, namespace="team") + native_runtime(facade) + client: Final = redis.Redis.from_url(redis_url) + first: Final = completion_kwargs("first") + await facade.async_add_cache({"value": 1}, **first) + first_key: Final = facade.get_cache_key(**first) + assert first_key.startswith("team:") + assert client.get(first_key) is None + second: Final = completion_kwargs("second") + await facade.async_add_cache({"value": 2}, **second) + assert client.get(first_key) is not None + assert client.get(facade.get_cache_key(**second)) is not None + client.close() + + +@pytest.mark.parametrize( + ("backend", "settings", "message"), + [ + pytest.param( + LiteLLMCacheType.VALKEY_SEMANTIC, + {"redis_url": "rediss://127.0.0.1:6390/0", "similarity_threshold": 0.8}, + "native Valkey semantic cache does not support TLS connections", + id="valkey-tls", + ), + pytest.param( + LiteLLMCacheType.VALKEY_SEMANTIC, + {"redis_url": "redis://127.0.0.1:6390/0?socket_timeout=1", "similarity_threshold": 0.8}, + "native Redis uses fixed socket timeouts; socket_timeout and socket_connect_timeout require Python", + id="valkey-socket-timeout", + ), + pytest.param( + LiteLLMCacheType.REDIS_SEMANTIC, + {"redis_url": "rediss://127.0.0.1:6380", "similarity_threshold": 0.8}, + "native Redis semantic cache does not support TLS or query options in redis_url", + id="redis-semantic-tls", + ), + pytest.param( + LiteLLMCacheType.REDIS_SEMANTIC, + {"redis_url": "redis://127.0.0.1:6379?socket_timeout=1", "similarity_threshold": 0.8}, + "native Redis semantic cache does not support TLS or query options in redis_url", + id="redis-semantic-query", + ), + ], +) +def test_semantic_settings_the_native_client_cannot_honor_decline( + monkeypatch: pytest.MonkeyPatch, backend: LiteLLMCacheType, settings: dict[str, object], message: str +) -> None: + require_rust(monkeypatch, backend) + with pytest.raises(RuntimeError, match=f"declined the cache: {message}"): + Cache(type=backend, **settings) + + +def test_rust_with_fallback_keeps_python_when_the_native_client_declines( + redis_url: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + catalog, + "RULES", + (CacheRule(Rollout.RUST_OPT_OUT, backends=frozenset({LiteLLMCacheType.REDIS})),), + ) + assert redis_facade(redis_url, socket_timeout=1.0)._native_cache is None # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor + + +def test_qdrant_semantic_rust_required_rule_activates_natively( + qdrant_url: str, fake_embedding_endpoint: str, monkeypatch: pytest.MonkeyPatch +) -> None: + del fake_embedding_endpoint + require_rust(monkeypatch, LiteLLMCacheType.QDRANT_SEMANTIC) + facade: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") + native_runtime(facade) + kwargs: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "qdrant activation"}]} + facade.add_cache({"answer": "qdrant"}, **kwargs) + assert facade.get_cache(**kwargs) == {"answer": "qdrant"} + + +async def test_redis_semantic_rust_required_rule_activates_natively( + redis_stack: tuple[str, str], semantic_embedding: DeterministicEmbedding, monkeypatch: pytest.MonkeyPatch +) -> None: + del semantic_embedding + url, index = redis_stack + require_rust(monkeypatch, LiteLLMCacheType.REDIS_SEMANTIC) + facade: Final = Cache( + type=LiteLLMCacheType.REDIS_SEMANTIC, + redis_url=url, + similarity_threshold=0.8, + redis_semantic_cache_embedding_model=SEMANTIC_EMBEDDING_MODEL, + redis_semantic_cache_index_name=index, + ) + native_runtime(facade) + kwargs: Final = {"model": "gpt-4o", "messages": semantic_messages("name a primary color")} + await facade.async_add_cache({"answer": "blue"}, **kwargs) + assert await facade.async_get_cache(**kwargs) == {"answer": "blue"} + + +async def test_azure_blob_rust_required_rule_activates_natively(monkeypatch: pytest.MonkeyPatch) -> None: + account_url: Final = os.environ.get("AZURE_BLOB_CACHE_ACCOUNT_URL") + if account_url is None: + pytest.skip( + "live Azure Blob parity needs AZURE_BLOB_CACHE_ACCOUNT_URL plus DefaultAzureCredential inputs in the environment" + ) + require_rust(monkeypatch, LiteLLMCacheType.AZURE_BLOB) + facade: Final = Cache( + type=LiteLLMCacheType.AZURE_BLOB, + azure_account_url=account_url, + azure_blob_container=f"litellm-parity-{uuid.uuid4().hex[:12]}", + ) + backend: Final = facade.cache + assert isinstance(backend, AzureBlobCache) + try: + native_runtime(facade) + kwargs: Final = completion_kwargs("azure") + await facade.async_add_cache({"answer": "azure"}, **kwargs) + assert await facade.async_get_cache(**kwargs) == {"answer": "azure"} + assert backend.get_cache(facade.get_cache_key(**kwargs))["response"] == {"answer": "azure"} + finally: + backend.container_client.delete_container() + await backend.disconnect() + + +class _SemanticHit: + """A native semantic runtime that answers every lookup with one cached response.""" + + kind: Final = "native" + + def lookup_semantic(self, request: object) -> tuple[object, float | None]: + return {"answer": 42}, 0.97 + + async def async_lookup_semantic(self, request: object) -> tuple[object, float | None]: + return {"answer": 42}, 0.97 + + +@pytest.mark.parametrize("semantic_type", [LiteLLMCacheType.QDRANT_SEMANTIC, LiteLLMCacheType.REDIS_SEMANTIC]) +@pytest.mark.parametrize("use_async", [False, True], ids=["sync", "async"]) +def test_native_semantic_hit_stamps_similarity_on_request_metadata( + semantic_type: LiteLLMCacheType, use_async: bool +) -> None: + """Python semantic backends write `metadata["semantic-similarity"]` on every lookup, and the + facade copies it to the caller's metadata; the native path must report it the same way.""" + facade: Final = Cache() + facade.type = semantic_type + facade._native_cache = ResponseCacheRuntime(cast(NativeResponseCacheRuntime, _SemanticHit())) # pyright: ignore[reportPrivateUsage] # the native path under test has no public setter + metadata: Final[dict[str, object]] = {} + kwargs: Final = { + "cache_key": "semantic-key", + "messages": [{"role": "user", "content": "hello"}], + "metadata": metadata, + } + + result: Final = asyncio.run(facade.async_get_cache(**kwargs)) if use_async else facade.get_cache(**kwargs) + + assert result == {"answer": 42} + assert metadata["semantic-similarity"] == 0.97 diff --git a/tests/test_litellm_rust/test_fork_guard.py b/tests/test_litellm_rust/test_fork_guard.py index 2a8fb6f9fca..02c454e0460 100644 --- a/tests/test_litellm_rust/test_fork_guard.py +++ b/tests/test_litellm_rust/test_fork_guard.py @@ -155,10 +155,17 @@ import asyncio import os import litellm from litellm.proxy.spend_tracking.input_tokens import count_input_tokens -from litellm.rust_bridge import _native +from litellm.rust_bridge import _native, catalog +from litellm.rust_bridge.catalog import Route, RouteRule +from litellm.rust_bridge.configuration import Rollout from litellm.litellm_core_utils.tokenizer import HuggingFaceTokenizer from litellm.utils import claude_json_str +catalog.RULES = ( + RouteRule(Route.TOKENIZER, Rollout.RUST_OPT_IN), + RouteRule(Route.TOKEN_COUNTER, Rollout.RUST_OPT_IN), + *catalog.RULES, +) litellm.anthropic_models = {*litellm.anthropic_models, "tokenizer-fork-fixture"} _native.reserve_process_for_forking() for create in ( diff --git a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py index c87a9f86a80..046f3a70ae8 100644 --- a/tests/test_litellm_rust/test_valkey_semantic_cache_native.py +++ b/tests/test_litellm_rust/test_valkey_semantic_cache_native.py @@ -15,7 +15,10 @@ import redis from litellm.caching.caching import Cache from litellm.caching.valkey_semantic_cache import ValkeySemanticCache -from litellm.rust_bridge import _native +from litellm.rust_bridge import _native, catalog +from litellm.rust_bridge.catalog import CacheRule +from litellm.rust_bridge.configuration import Rollout +from litellm.rust_bridge.response_cache import ResponseCacheRuntime from litellm.types.caching import LiteLLMCacheType pytestmark: Final = pytest.mark.requires_rust_extension @@ -597,3 +600,22 @@ async def test_ping_maps_unsupported_native_operation_to_not_implemented( binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=handle)).resolve() with pytest.raises(NotImplementedError): await binding.ping() + + +async def test_rust_required_rule_activates_the_facade_natively( + valkey_url: str, + index_name: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + catalog, + "RULES", + (CacheRule(Rollout.RUST_REQUIRED, backends=frozenset({LiteLLMCacheType.VALKEY_SEMANTIC})),), + ) + facade: Final = _facade(valkey_url, index_name, {"semantic cache prompt": [1.0, 0.0]}) + runtime: Final = facade._native_cache # pyright: ignore[reportPrivateUsage] # the activation under test has no public accessor + assert isinstance(runtime, ResponseCacheRuntime) + assert runtime.kind == "native" + kwargs: Final = {"model": "gpt-4o", "messages": _request()["messages"]} + await facade.async_add_cache({"answer": "valkey"}, **kwargs) + assert await facade.async_get_cache(**kwargs) == {"answer": "valkey"}