From 0c72a94a84d87616b5df7a9b34cc8ee99210324c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:21:24 +0000 Subject: [PATCH 01/17] feat(cache): add SemanticCacheContext and semantic error variants Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/cache/src/error.rs | 4 ++ litellm-rust/crates/cache/src/lib.rs | 2 + litellm-rust/crates/cache/src/semantic.rs | 72 +++++++++++++++++++++++ 3 files changed, 78 insertions(+) create mode 100644 litellm-rust/crates/cache/src/semantic.rs diff --git a/litellm-rust/crates/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index ff3ff6572d4..72fb338e7d8 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -6,4 +6,8 @@ pub enum Error { InvalidEntry, #[error("flushing Redis requires an explicit namespace")] UnscopedFlush, + #[error("cache backend does not support this operation")] + UnsupportedOperation, + #[error("semantic cache requires request messages")] + MissingPrompt, } diff --git a/litellm-rust/crates/cache/src/lib.rs b/litellm-rust/crates/cache/src/lib.rs index ce9f93b6dc4..b9c720fa3ee 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; +mod semantic; pub use base_cache::{ BaseCache, BatchEntry, CacheConnectionResult, CacheConnectionStatus, CacheContext, @@ -19,3 +20,4 @@ pub use capabilities::{ pub use codec::{CacheCodec, JsonCodec}; pub use dual::{DualCache, ReadPolicy, RemoteFailurePolicy, WritePolicy}; pub use error::Error; +pub use semantic::{SemanticCacheContext, SemanticCacheScope}; diff --git a/litellm-rust/crates/cache/src/semantic.rs b/litellm-rust/crates/cache/src/semantic.rs new file mode 100644 index 00000000000..61f9023fa4c --- /dev/null +++ b/litellm-rust/crates/cache/src/semantic.rs @@ -0,0 +1,72 @@ +use std::time::Duration; + +use serde::{Deserialize, Serialize}; +use serde_json::{Map, Value}; + +use crate::CacheContext; + +#[derive(Clone, Copy, Debug, Default, Deserialize, Serialize, PartialEq, Eq)] +#[serde(rename_all = "snake_case")] +pub enum SemanticCacheScope { + #[default] + Key, + EndUser, +} + +#[derive(Clone, Debug, Default, PartialEq)] +pub struct SemanticCacheContext { + pub input: Option, + pub messages: Vec, + pub metadata: Map, + pub scope: SemanticCacheScope, + pub ttl: Option, +} + +impl CacheContext for SemanticCacheContext { + fn ttl(&self) -> Option { + self.ttl + } + + fn with_ttl(&self, ttl: Option) -> Self { + Self { + ttl, + ..self.clone() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn with_ttl_keeps_request_fields() { + let context = SemanticCacheContext { + input: Some("query".to_owned()), + messages: vec![serde_json::json!({"role": "user", "content": "hi"})], + metadata: Map::from_iter([("user".to_owned(), Value::from("u1"))]), + scope: SemanticCacheScope::EndUser, + ttl: None, + }; + + let updated = context.with_ttl(Some(Duration::from_secs(5))); + + assert_eq!(updated.ttl(), Some(Duration::from_secs(5))); + assert_eq!(updated.input, context.input); + assert_eq!(updated.messages, context.messages); + assert_eq!(updated.metadata, context.metadata); + assert_eq!(updated.scope, SemanticCacheScope::EndUser); + } + + #[test] + fn scope_serializes_like_python_cache_scope() { + assert_eq!( + serde_json::to_value(SemanticCacheScope::EndUser).unwrap(), + Value::from("end_user") + ); + assert_eq!( + serde_json::from_value::(Value::from("key")).unwrap(), + SemanticCacheScope::Key + ); + } +} From fc3844e9913829cc99ee39b1270dbc7bcafde163 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:27:04 +0000 Subject: [PATCH 02/17] refactor(cache-response): generalize ResponseCache over the backend context Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/cache-response/src/response.rs | 57 ++++++++----- .../crates/cache-response/tests/response.rs | 83 ++++++++++++++++++- .../crates/python-bridge/src/cache/request.rs | 13 ++- 3 files changed, 128 insertions(+), 25 deletions(-) diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index e50e68cdabb..eedbf2caf1a 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -1,21 +1,22 @@ use std::{sync::Arc, time::Duration}; use litellm_cache::{ - BaseCache, BatchCache, BatchEntry, CacheConnectionResult, Error, ExactCacheContext, FlushCache, + BaseCache, BatchCache, BatchEntry, CacheConnectionResult, CacheContext, Error, + ExactCacheContext, FlushCache, }; use serde_json::Value; use crate::{CacheControls, CacheEntry, CacheKeyInput, PartialHits, cache_key}; #[derive(Clone)] -pub struct ResponseCacheRequest { +pub struct ResponseCacheRequest { pub key: CacheKeyInput, pub controls: CacheControls, - pub context: ExactCacheContext, + pub context: C, pub max_age: Option, } -impl ResponseCacheRequest { +impl ResponseCacheRequest { pub fn new(key: CacheKeyInput) -> Self { Self { key, @@ -26,17 +27,35 @@ impl ResponseCacheRequest { default_on: true, ..Default::default() }, - context: ExactCacheContext::default(), + context: C::default(), max_age: None, } } } -pub struct ResponseCache> { +impl ResponseCacheRequest { + pub fn with_context(self, context: D) -> ResponseCacheRequest { + ResponseCacheRequest { + key: self.key, + controls: self.controls, + context, + max_age: self.max_age, + } + } +} + +pub struct ResponseCache> +where + B::Context: Default + PartialEq, +{ backend: Arc, } -impl> ResponseCache { +impl ResponseCache +where + B: BaseCache, + B::Context: Default + PartialEq, +{ pub fn new(backend: Arc) -> Self { Self { backend } } @@ -46,7 +65,7 @@ impl> ResponseCach } pub fn default_ttl(&self) -> Option { - self.backend.get_ttl(&ExactCacheContext::default()) + self.backend.get_ttl(&B::Context::default()) } pub async fn async_flush(&self) -> Result<(), Error> @@ -62,7 +81,7 @@ impl> ResponseCach pub fn lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -81,7 +100,7 @@ impl> ResponseCach pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { if !request.controls.reads() { @@ -101,7 +120,7 @@ impl> ResponseCach pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -126,7 +145,7 @@ impl> ResponseCach pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result where @@ -153,7 +172,7 @@ impl> ResponseCach pub fn store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -172,7 +191,7 @@ impl> ResponseCach pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { @@ -193,7 +212,7 @@ impl> ResponseCach pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(ResponseCacheRequest, Value)>, now: Duration, ) -> Result<(), Error> { self.async_store_entries( @@ -209,7 +228,7 @@ impl> ResponseCach /// the freshness of its original response. pub async fn async_store_entries( &self, - entries: Vec<(ResponseCacheRequest, Value, Duration)>, + entries: Vec<(ResponseCacheRequest, Value, Duration)>, ) -> Result<(), Error> { let writable = entries .into_iter() @@ -248,9 +267,9 @@ impl> ResponseCach Ok(()) } - fn partial_hits( - requests: &[ResponseCacheRequest], - readable: Vec<(usize, &ResponseCacheRequest)>, + fn partial_hits( + requests: &[ResponseCacheRequest], + readable: Vec<(usize, &ResponseCacheRequest)>, entries: Vec>, now: Duration, ) -> Result { diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index e4f78dae8b2..56589063291 100644 --- a/litellm-rust/crates/cache-response/tests/response.rs +++ b/litellm-rust/crates/cache-response/tests/response.rs @@ -1,12 +1,15 @@ use std::{ sync::{ - Arc, + Arc, Mutex, atomic::{AtomicU64, Ordering}, }, time::Duration, }; -use litellm_cache::{BaseCache, CacheCodec, Error}; +use litellm_cache::{ + BaseCache, CacheCodec, CacheConnectionResult, CacheConnectionStatus, Error, + SemanticCacheContext, +}; use litellm_cache_memory::InMemoryCache; use litellm_cache_redis::RedisCache; use litellm_cache_response::{ @@ -30,6 +33,82 @@ fn request() -> ResponseCacheRequest { }) } +struct SemanticBackend { + entries: Mutex>, + contexts: Mutex>, +} + +impl BaseCache for SemanticBackend { + type Value = CacheEntry; + type Context = SemanticCacheContext; + + fn get_ttl(&self, _: &Self::Context) -> Option { + None + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + self.contexts.lock().unwrap().push(context.clone()); + self.entries.lock().unwrap().push((key.to_owned(), value)); + Ok(()) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + self.contexts.lock().unwrap().push(context.clone()); + Ok(self + .entries + .lock() + .unwrap() + .iter() + .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()), + }); + let cache = ResponseCache::new(backend.clone()); + let context = SemanticCacheContext { + messages: vec![json!({"role": "user", "content": "hello"})], + ..Default::default() + }; + let request = request().with_context(context.clone()); + let response = json!({"answer": 42}); + + cache + .store(&request, response.clone(), Duration::from_secs(100)) + .unwrap(); + + assert_eq!( + cache.lookup(&request, Duration::from_secs(100)).unwrap(), + Some(response) + ); + assert_eq!( + backend.contexts.lock().unwrap().as_slice(), + &[context.clone(), context] + ); +} + #[tokio::test] async fn sync_and_async_consumers_share_keys_ttls_and_freshness() { let clock = Arc::new(AtomicU64::new(100)); diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 0c5343a63d0..0067fc4392b 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -1,5 +1,6 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use litellm_cache::ExactCacheContext; use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; @@ -14,13 +15,15 @@ struct RequestInput { max_age_seconds: Option, } -pub(super) fn request(value: &Bound<'_, PyAny>) -> PyResult { +pub(super) fn request( + value: &Bound<'_, PyAny>, +) -> PyResult> { let input: RequestInput = from_py(value)?; request_input(input) } -fn request_input(input: RequestInput) -> PyResult { - let mut request = ResponseCacheRequest::new(input.key); +fn request_input(input: RequestInput) -> PyResult> { + let mut request: ResponseCacheRequest = ResponseCacheRequest::new(input.key); if let Some(controls) = input.controls { request.controls = controls; } @@ -29,7 +32,9 @@ fn request_input(input: RequestInput) -> PyResult { Ok(request) } -pub(super) fn requests(value: &Bound<'_, PyAny>) -> PyResult> { +pub(super) fn requests( + value: &Bound<'_, PyAny>, +) -> PyResult>> { from_py::>(value)? .into_iter() .map(request_input) From 80c0ceb5e61ca85f7e865f8901d38368e82b047d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:27:06 +0000 Subject: [PATCH 03/17] feat(cache-qdrant-semantic): add native Qdrant semantic cache backend Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 94 +++++++ litellm-rust/Cargo.toml | 3 + .../crates/cache-qdrant-semantic/Cargo.toml | 23 ++ .../cache-qdrant-semantic/src/embedder.rs | 73 +++++ .../crates/cache-qdrant-semantic/src/lib.rs | 7 + .../cache-qdrant-semantic/src/prompt.rs | 59 ++++ .../cache-qdrant-semantic/src/semantic.rs | 256 ++++++++++++++++++ .../cache-qdrant-semantic/tests/prompt.rs | 38 +++ 8 files changed, 553 insertions(+) create mode 100644 litellm-rust/crates/cache-qdrant-semantic/Cargo.toml create mode 100644 litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs create mode 100644 litellm-rust/crates/cache-qdrant-semantic/src/lib.rs create mode 100644 litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs create mode 100644 litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs create mode 100644 litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 85d1e1ca8ef..8868b25bb18 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -547,6 +547,49 @@ dependencies = [ "tracing", ] +[[package]] +name = "axum" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31b698c5f9a010f6573133b09e0de5408834d0c82f8d7475a89fc1867a71cd90" +dependencies = [ + "axum-core", + "bytes", + "futures-util", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "itoa", + "matchit", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "sync_wrapper", + "tower", + "tower-layer", + "tower-service", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http 1.4.2", + "http-body 1.1.0", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", +] + [[package]] name = "azure_core" version = "1.1.0" @@ -2480,6 +2523,25 @@ dependencies = [ "tokio", ] +[[package]] +name = "litellm-cache-qdrant-semantic" +version = "0.1.0" +dependencies = [ + "futures-util", + "litellm-cache", + "litellm-cache-response", + "qdrant-client", + "reqwest 0.12.28", + "rstest", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tonic", + "tonic-prost", + "uuid", +] + [[package]] name = "litellm-cache-redis" version = "0.1.0" @@ -2896,6 +2958,12 @@ version = "0.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fc04a4c58212d57930a24bf47d3fa87485264a3a054e9c10e042eb373573ad3c" +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + [[package]] name = "memchr" version = "2.8.3" @@ -3509,6 +3577,27 @@ dependencies = [ "serde", ] +[[package]] +name = "qdrant-client" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dddc19df129bad7346ebd027288621ab1ac7e52678371f906b9a8622d7aaf87e" +dependencies = [ + "anyhow", + "derive_builder", + "futures", + "parking_lot", + "prost", + "prost-types", + "semver", + "serde", + "serde_json", + "thiserror 2.0.19", + "tokio", + "tonic", + "tonic-prost", +] + [[package]] name = "quick-error" version = "1.2.3" @@ -4866,8 +4955,12 @@ version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" dependencies = [ + "async-trait", + "axum", "base64 0.22.1", "bytes", + "flate2", + "h2 0.4.15", "http 1.4.2", "http-body 1.1.0", "http-body-util", @@ -4877,6 +4970,7 @@ dependencies = [ "percent-encoding", "pin-project", "rustls-native-certs", + "socket2 0.6.5", "sync_wrapper", "tokio", "tokio-rustls 0.26.4", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index d35df1eafb8..fb9365012d1 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -31,6 +31,7 @@ litellm-cache = { path = "crates/cache" } litellm-cache-memory = { path = "crates/cache-memory" } litellm-cache-redis = { path = "crates/cache-redis" } litellm-cache-response = { path = "crates/cache-response" } +litellm-cache-qdrant-semantic = { path = "crates/cache-qdrant-semantic" } 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" } @@ -48,6 +49,8 @@ pyo3-async-runtimes = { version = "0.29.0", features = ["tokio-runtime"] } pythonize = "0.29.0" rand = "0.8" reqwest = { version = "0.12", default-features = false, features = ["json", "multipart", "rustls-tls", "http2", "stream"] } +qdrant-client = { version = "1.19.0", default-features = false } +uuid = { version = "1", features = ["v4"] } rstest = "0.26.1" rstest_reuse = "0.7.0" rustls = { version = "0.23", default-features = false, features = ["ring", "std", "tls12"] } diff --git a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml new file mode 100644 index 00000000000..7e215cab837 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "litellm-cache-qdrant-semantic" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +futures-util.workspace = true +litellm-cache.workspace = true +qdrant-client = { workspace = true, features = ["serde"] } +reqwest.workspace = true +serde.workspace = true +serde_json.workspace = true +thiserror.workspace = true +tokio.workspace = true +uuid.workspace = true + +[dev-dependencies] +litellm-cache-response.workspace = true +rstest.workspace = true +tonic = "0.14" +tonic-prost = "0.14" diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs new file mode 100644 index 00000000000..0dde0318448 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs @@ -0,0 +1,73 @@ +use std::time::Duration; + +use litellm_cache::Error; +use reqwest::Client; +use serde_json::Value; + +use crate::Embedder; + +pub struct OpenAiEmbedder { + client: Client, + api_base: String, + api_key: String, + model: String, +} + +pub struct OpenAiEmbedderConfig { + pub api_base: String, + pub api_key: String, + pub model: String, + pub timeout: Option, +} + +impl OpenAiEmbedder { + pub fn new(config: OpenAiEmbedderConfig) -> Result { + let mut builder = Client::builder(); + if let Some(timeout) = config.timeout { + builder = builder.timeout(timeout); + } + let client = builder.build().map_err(|_| Error::Unavailable)?; + Ok(Self { + client, + api_base: config.api_base.trim_end_matches('/').to_owned(), + api_key: config.api_key, + model: config.model, + }) + } +} + +impl Embedder for OpenAiEmbedder { + fn model(&self) -> &str { + &self.model + } + + async fn embed(&self, input: &str) -> Result, Error> { + let response = self + .client + .post(format!("{}/embeddings", self.api_base)) + .bearer_auth(&self.api_key) + .json(&serde_json::json!({ + "model": self.model, + "input": input, + "encoding_format": "float", + })) + .send() + .await + .map_err(|_| Error::Unavailable)? + .error_for_status() + .map_err(|_| Error::Unavailable)?; + let body: Value = response.json().await.map_err(|_| Error::Unavailable)?; + body.get("data") + .and_then(Value::as_array) + .and_then(|data| data.first()) + .and_then(|item| item.get("embedding")) + .and_then(Value::as_array) + .and_then(|embedding| { + embedding + .iter() + .map(|value| value.as_f64().map(|value| value as f32)) + .collect::>>() + }) + .ok_or(Error::Unavailable) + } +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/lib.rs b/litellm-rust/crates/cache-qdrant-semantic/src/lib.rs new file mode 100644 index 00000000000..0f346a9155b --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/src/lib.rs @@ -0,0 +1,7 @@ +mod embedder; +mod prompt; +mod semantic; + +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 new file mode 100644 index 00000000000..ef1a2306658 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/src/prompt.rs @@ -0,0 +1,59 @@ +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/src/semantic.rs b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs new file mode 100644 index 00000000000..16f11286d3f --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs @@ -0,0 +1,256 @@ +use std::future::Future; + +use futures_util::future::try_join_all; +use litellm_cache::{BaseCache, CacheCodec, CacheConnectionResult, Error, SemanticCacheContext}; +use qdrant_client::{ + Payload, Qdrant, + qdrant::{ + BinaryQuantizationBuilder, CompressionRatio, Condition, CreateCollectionBuilder, + CreateFieldIndexCollectionBuilder, Distance, FieldType, Filter, PointStruct, + ProductQuantizationBuilder, QuantizationSearchParamsBuilder, ScalarQuantizationBuilder, + SearchParamsBuilder, SearchPointsBuilder, UpsertPointsBuilder, VectorParamsBuilder, + }, +}; +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, +} + +pub struct QdrantSemanticCache { + client: Qdrant, + embedder: E, + codec: C, + config: QdrantSemanticConfig, + runtime: tokio::runtime::Handle, +} + +impl QdrantSemanticCache { + pub async fn connect( + client: Qdrant, + embedder: E, + codec: C, + config: QdrantSemanticConfig, + runtime: tokio::runtime::Handle, + ) -> Result { + let exists = client + .collection_exists(config.collection_name.clone()) + .await + .map_err(|_| Error::Unavailable)?; + if !exists { + client + .create_collection( + CreateCollectionBuilder::new(config.collection_name.clone()) + .vectors_config(VectorParamsBuilder::new( + config.vector_size, + Distance::Cosine, + )) + .quantization_config(quantization(&config.quantization)), + ) + .await + .map_err(|_| Error::Unavailable)?; + } + let _ = client + .create_field_index(CreateFieldIndexCollectionBuilder::new( + config.collection_name.clone(), + "litellm_cache_key".to_owned(), + FieldType::Keyword, + )) + .await; + Ok(Self { + client, + embedder, + codec, + config, + runtime, + }) + } + + pub fn collection_name(&self) -> &str { + &self.config.collection_name + } + + pub fn similarity_threshold(&self) -> f64 { + self.config.similarity_threshold + } + + pub fn vector_size(&self) -> u64 { + self.config.vector_size + } + + pub fn embedder(&self) -> &E { + &self.embedder + } + + fn prompt(context: &SemanticCacheContext) -> Result { + if context.messages.is_empty() { + return Err(Error::MissingPrompt); + } + Ok(prompt_from_messages(&context.messages)) + } + + async fn set( + &self, + key: &str, + value: C::Value, + context: &SemanticCacheContext, + ) -> Result<(), Error> { + let prompt = Self::prompt(context)?; + let vector = self.embedder.embed(&prompt).await?; + let response = + String::from_utf8(self.codec.encode(&value)?).map_err(|_| Error::InvalidEntry)?; + let payload = Payload::try_from(json!({ + "litellm_cache_key": key, + "text": prompt, + "response": response, + })) + .map_err(|_| Error::InvalidEntry)?; + self.client + .upsert_points(UpsertPointsBuilder::new( + self.collection_name(), + vec![PointStruct::new( + Uuid::new_v4().to_string(), + vector, + payload, + )], + )) + .await + .map_err(|_| Error::Unavailable)?; + Ok(()) + } + + async fn get( + &self, + key: &str, + context: &SemanticCacheContext, + ) -> Result, Error> { + let prompt = Self::prompt(context)?; + let vector = self.embedder.embed(&prompt).await?; + let result = self + .client + .search_points( + SearchPointsBuilder::new(self.collection_name(), vector, 1) + .with_payload(true) + .filter(Filter::must([Condition::matches( + "litellm_cache_key", + key.to_owned(), + )])) + .params( + SearchParamsBuilder::default().quantization( + QuantizationSearchParamsBuilder::default() + .ignore(false) + .rescore(true) + .oversampling(3.0), + ), + ), + ) + .await + .map_err(|_| Error::Unavailable)?; + let Some(point) = result.result.into_iter().next() else { + return Ok(None); + }; + if f64::from(point.score) < self.config.similarity_threshold { + return Ok(None); + } + let payload: Map = Payload::from(point.payload).into(); + if payload.get("litellm_cache_key").and_then(Value::as_str) != Some(key) { + return Ok(None); + } + let response = payload + .get("response") + .and_then(Value::as_str) + .ok_or(Error::InvalidEntry)?; + self.codec.decode(response.as_bytes()).map(Some) + } +} + +fn quantization(value: &Quantization) -> qdrant_client::qdrant::quantization_config::Quantization { + match value { + Quantization::Binary => BinaryQuantizationBuilder::new(false).into(), + Quantization::Scalar => ScalarQuantizationBuilder::default() + .quantile(0.99) + .always_ram(false) + .into(), + Quantization::Product => ProductQuantizationBuilder::new(CompressionRatio::X16.into()) + .always_ram(false) + .into(), + } +} + +impl BaseCache for QdrantSemanticCache { + type Value = C::Value; + type Context = SemanticCacheContext; + + fn get_ttl(&self, _: &Self::Context) -> Option { + None + } + + fn set_cache( + &self, + key: &str, + value: Self::Value, + context: &Self::Context, + ) -> Result<(), Error> { + self.runtime.block_on(self.set(key, value, context)) + } + + fn get_cache(&self, key: &str, context: &Self::Context) -> Result, Error> { + self.runtime.block_on(self.get(key, context)) + } + + async fn async_set_cache( + &self, + key: &str, + value: Self::Value, + context: Self::Context, + ) -> Result<(), Error> { + self.set(key, value, &context).await + } + + async fn async_get_cache( + &self, + key: &str, + context: &Self::Context, + ) -> Result, Error> { + self.get(key, context).await + } + + async fn async_set_cache_pipeline( + &self, + entries: Vec<(String, Self::Value)>, + context: Self::Context, + ) -> Result<(), Error> { + try_join_all(entries.into_iter().map(|(key, value)| { + let context = context.clone(); + async move { self.async_set_cache(&key, value, context).await } + })) + .await + .map(|_| ()) + } + + async fn disconnect(&self) -> Result<(), Error> { + Ok(()) + } + + async fn test_connection(&self) -> Result { + Err(Error::UnsupportedOperation) + } +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs new file mode 100644 index 00000000000..38cd9e2f908 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/prompt.rs @@ -0,0 +1,38 @@ +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"}"# + ); +} From 8d41336a1e1bea925801eaff33b84010b6c186f7 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:33:00 +0000 Subject: [PATCH 04/17] test(cache-qdrant-semantic): cover backend contract against an in-process Qdrant Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + .../crates/cache-qdrant-semantic/Cargo.toml | 1 + .../cache-qdrant-semantic/src/semantic.rs | 6 +- .../cache-qdrant-semantic/tests/embedder.rs | 141 ++++++ .../cache-qdrant-semantic/tests/qdrant.rs | 418 ++++++++++++++++++ .../tests/support/mod.rs | 339 ++++++++++++++ 6 files changed, 903 insertions(+), 3 deletions(-) create mode 100644 litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs create mode 100644 litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs create mode 100644 litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 8868b25bb18..fb75c2241b5 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2537,6 +2537,7 @@ dependencies = [ "serde_json", "thiserror 2.0.19", "tokio", + "tokio-stream", "tonic", "tonic-prost", "uuid", diff --git a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml index 7e215cab837..09d6a9637f3 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml @@ -21,3 +21,4 @@ litellm-cache-response.workspace = true rstest.workspace = true tonic = "0.14" tonic-prost = "0.14" +tokio-stream = "0.1" diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs index 16f11286d3f..d761364f1ad 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs @@ -167,13 +167,13 @@ impl QdrantSemanticCache { let Some(point) = result.result.into_iter().next() else { return Ok(None); }; - if f64::from(point.score) < self.config.similarity_threshold { - return Ok(None); - } 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 f64::from(point.score) < self.config.similarity_threshold { + return Ok(None); + } let response = payload .get("response") .and_then(Value::as_str) diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs new file mode 100644 index 00000000000..adce70654a8 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs @@ -0,0 +1,141 @@ +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::Error; +use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, OpenAiEmbedderConfig}; +use serde_json::Value; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, +}; + +struct TestHttpServer { + address: std::net::SocketAddr, + request: Arc>>>, + task: tokio::task::JoinHandle<()>, +} + +impl TestHttpServer { + async fn response(status: &str, body: &str) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let request = Arc::new(Mutex::new(None)); + let captured = request.clone(); + let status = status.to_owned(); + let body = body.to_owned(); + let task = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let request_bytes = read_request(&mut stream).await; + *captured.lock().unwrap() = Some(request_bytes); + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + stream.write_all(response.as_bytes()).await.unwrap(); + }); + Self { + address, + request, + task, + } + } + + async fn hanging() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let task = tokio::spawn(async move { + let (_stream, _) = listener.accept().await.unwrap(); + std::future::pending::<()>().await; + }); + Self { + address, + request: Arc::new(Mutex::new(None)), + task, + } + } + + fn base_url(&self) -> String { + format!("http://{}", self.address) + } +} + +impl Drop for TestHttpServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +async fn read_request(stream: &mut tokio::net::TcpStream) -> Vec { + let mut bytes = Vec::new(); + let header_end = loop { + let mut chunk = [0_u8; 1024]; + let count = stream.read(&mut chunk).await.unwrap(); + assert_ne!(count, 0); + bytes.extend_from_slice(&chunk[..count]); + if let Some(end) = bytes.windows(4).position(|window| window == b"\r\n\r\n") { + break end + 4; + } + }; + let headers = String::from_utf8_lossy(&bytes[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + line.split_once(':') + .filter(|(name, _)| name.eq_ignore_ascii_case("content-length")) + .map(|(_, value)| value.trim()) + }) + .unwrap() + .parse::() + .unwrap(); + while bytes.len() < header_end + content_length { + let mut chunk = [0_u8; 1024]; + let count = stream.read(&mut chunk).await.unwrap(); + assert_ne!(count, 0); + bytes.extend_from_slice(&chunk[..count]); + } + bytes +} + +fn config(base: String, timeout: Option) -> OpenAiEmbedderConfig { + OpenAiEmbedderConfig { + api_base: base, + api_key: "test-key".to_owned(), + model: "test-model".to_owned(), + timeout, + } +} + +#[tokio::test] +async fn posts_embeddings_request_and_parses_vector() { + let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await; + let embedder = OpenAiEmbedder::new(config( + format!("{}/", server.base_url()), + Some(Duration::from_secs(1)), + )) + .unwrap(); + assert_eq!(embedder.embed("hello").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")); + assert!(request_text.contains("\r\nauthorization: Bearer test-key\r\n")); + let body = request_text.split("\r\n\r\n").nth(1).unwrap(); + let body: Value = serde_json::from_str(body).unwrap(); + assert_eq!(body["model"], "test-model"); + assert_eq!(body["input"], "hello"); + assert_eq!(body["encoding_format"], "float"); +} + +#[tokio::test] +async fn status_and_timeout_errors_are_unavailable() { + let server = TestHttpServer::response("500 Internal Server Error", "{}").await; + let embedder = + OpenAiEmbedder::new(config(server.base_url(), Some(Duration::from_secs(1)))).unwrap(); + assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable)); + + let server = TestHttpServer::hanging().await; + let embedder = + OpenAiEmbedder::new(config(server.base_url(), Some(Duration::from_millis(200)))).unwrap(); + assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable)); +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs new file mode 100644 index 00000000000..d5ecaf7217d --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -0,0 +1,418 @@ +#[path = "support/mod.rs"] +mod support; + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use litellm_cache::{ + BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext, SemanticCacheScope, +}; +use litellm_cache_qdrant_semantic::{ + Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization, +}; +use litellm_cache_response::{ + CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, +}; +use qdrant_client::Payload; +use qdrant_client::{ + Qdrant, + qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams}, +}; +use serde_json::{Value as JsonValue, json}; + +use support::{FakeQdrant, FakeState, StoredPoint}; + +#[derive(Clone)] +struct FixedEmbedder { + vectors: Arc>>, +} + +impl FixedEmbedder { + fn new(vectors: impl IntoIterator)>) -> Self { + Self { + vectors: Arc::new( + vectors + .into_iter() + .map(|(prompt, vector)| (prompt.to_owned(), vector)) + .collect(), + ), + } + } +} + +impl Embedder for FixedEmbedder { + fn model(&self) -> &str { + "fixed" + } + + async fn embed(&self, input: &str) -> Result, Error> { + self.vectors.get(input).cloned().ok_or(Error::Unavailable) + } +} + +fn config(quantization: Quantization) -> QdrantSemanticConfig { + QdrantSemanticConfig { + collection_name: "semantic".to_owned(), + similarity_threshold: 0.9, + vector_size: 2, + quantization, + } +} + +fn context(prompt: &str) -> SemanticCacheContext { + SemanticCacheContext { + messages: vec![json!({"role": "user", "content": prompt})], + scope: SemanticCacheScope::default(), + ..Default::default() + } +} + +fn value(response: JsonValue) -> CacheEntry { + CacheEntry { + timestamp: Some(1.0), + response, + } +} + +async fn connect( + server: &FakeQdrant, + vectors: impl IntoIterator)>, +) -> QdrantSemanticCache { + let client = Qdrant::from_url(&server.url()).build().unwrap(); + QdrantSemanticCache::connect( + client, + FixedEmbedder::new(vectors), + ResponseCacheCodec, + config(Quantization::Binary), + tokio::runtime::Handle::current(), + ) + .await + .unwrap() +} + +#[tokio::test(flavor = "multi_thread")] +#[allow(deprecated)] +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 + .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"), + } + 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(); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn existing_collection_skips_create_and_index_failure_is_non_fatal() { + let server = FakeQdrant::start(FakeState { + collections: ["semantic".to_owned()].into_iter().collect(), + fail_field_index: true, + ..Default::default() + }) + .await; + let _cache = connect(&server, [("hello", vec![1.0, 0.0])]).await; + let state = server.state.lock().unwrap(); + assert!(state.created_collections.is_empty()); + assert!(state.index_creations >= 1); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn async_and_sync_set_get_store_exact_payload() { + 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})); + cache + .async_set_cache("key", entry.clone(), ctx.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache("key", &ctx).await.unwrap().as_ref(), + Some(&entry) + ); + { + let state = server.state.lock().unwrap(); + let payload = &state.points[0].payload; + let mut payload_keys = payload.keys().cloned().collect::>(); + 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()) + ); + } + let sync_entry = entry.clone(); + let sync_cache = cache.clone(); + let sync_ctx = ctx.clone(); + tokio::task::spawn_blocking(move || { + sync_cache + .set_cache("sync", sync_entry.clone(), &sync_ctx) + .unwrap(); + assert_eq!( + sync_cache.get_cache("sync", &sync_ctx).unwrap(), + Some(sync_entry) + ); + }) + .await + .unwrap(); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn misses_and_payload_validation_are_safe() { + 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(); + 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(), + None + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn decoding_errors_missing_prompt_pipeline_and_ttl_behave_as_required() { + 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) + ); + assert_eq!( + cache.async_get_cache("key", &empty).await, + Err(Error::MissingPrompt) + ); + assert_eq!( + cache.async_get_cache("key", &context("unknown")).await, + Err(Error::Unavailable) + ); + cache + .async_set_cache( + "ttl", + value(json!({"ttl": true})), + context("one").with_ttl(Some(Duration::from_secs(1))), + ) + .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() + ); + cache + .async_set_cache_pipeline( + vec![ + ("one".to_owned(), value(json!({"n": 1}))), + ("two".to_owned(), value(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() + ); + 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 + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn stopped_qdrant_server_maps_to_unavailable() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await; + server.stop(); + tokio::time::sleep(Duration::from_millis(50)).await; + assert_eq!( + cache.async_get_cache("key", &context("hello")).await, + Err(Error::Unavailable) + ); +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs new file mode 100644 index 00000000000..860213a1703 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs @@ -0,0 +1,339 @@ +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, + points_server::{Points, PointsServer}, +}; +use tokio::sync::oneshot; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::{Request, Response, Status, transport::Server}; + +#[derive(Clone, Debug)] +pub struct StoredPoint { + pub id: Option, + pub vector: Vec, + pub payload: HashMap, +} + +#[derive(Default)] +pub struct FakeState { + pub collections: HashSet, + pub created_collections: Vec, + pub field_indexes: Vec, + pub points: Vec, + pub index_creations: usize, + pub fail_field_index: bool, +} + +#[derive(Clone)] +pub struct FakeQdrant { + pub state: Arc>, + pub address: SocketAddr, + shutdown: Arc>>>, +} + +impl FakeQdrant { + pub async fn start(state: FakeState) -> Self { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let state = Arc::new(Mutex::new(state)); + let service = FakeService { + state: state.clone(), + }; + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + tokio::spawn(async move { + Server::builder() + .add_service(CollectionsServer::new(service.clone())) + .add_service(PointsServer::new(service)) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }) + .await + .unwrap(); + }); + Self { + state, + address, + shutdown: Arc::new(Mutex::new(Some(shutdown_tx))), + } + } + + pub fn url(&self) -> String { + format!("http://{}", self.address) + } + + pub fn stop(&self) { + self.shutdown + .lock() + .unwrap() + .take() + .unwrap() + .send(()) + .unwrap(); + } + + pub fn insert_point(&self, point: StoredPoint) { + self.state.lock().unwrap().points.push(point); + } +} + +#[derive(Clone)] +struct FakeService { + state: Arc>, +} + +macro_rules! unimplemented_collections { + ($($name:ident, $request:ty, $response:ty);* $(;)?) => { + $( + fn $name<'life0, 'async_trait>( + &'life0 self, + _: Request<$request>, + ) -> std::pin::Pin< + Box< + dyn std::future::Future< + Output = Result, Status>, + > + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async { Err(Status::unimplemented(stringify!($name))) }) + } + )* + }; +} + +macro_rules! unimplemented_points { + ($($name:ident, $request:ty, $response:ty);* $(;)?) => { + $( + fn $name<'life0, 'async_trait>( + &'life0 self, + _: Request<$request>, + ) -> std::pin::Pin< + Box< + dyn std::future::Future< + Output = Result, Status>, + > + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async { Err(Status::unimplemented(stringify!($name))) }) + } + )* + }; +} + +#[tonic::async_trait] +impl Collections for FakeService { + async fn create( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + let mut state = self.state.lock().unwrap(); + state.collections.insert(request.collection_name.clone()); + state.created_collections.push(request); + Ok(Response::new(CollectionOperationResponse { + result: true, + ..Default::default() + })) + } + + async fn collection_exists( + &self, + request: Request, + ) -> Result, Status> { + let exists = self + .state + .lock() + .unwrap() + .collections + .contains(&request.into_inner().collection_name); + Ok(Response::new(CollectionExistsResponse { + result: Some(CollectionExists { exists }), + ..Default::default() + })) + } + + unimplemented_collections!( + get, qdrant::GetCollectionInfoRequest, qdrant::GetCollectionInfoResponse; + list, qdrant::ListCollectionsRequest, qdrant::ListCollectionsResponse; + update, qdrant::UpdateCollection, qdrant::CollectionOperationResponse; + delete, qdrant::DeleteCollection, qdrant::CollectionOperationResponse; + update_aliases, qdrant::ChangeAliases, qdrant::CollectionOperationResponse; + list_collection_aliases, qdrant::ListCollectionAliasesRequest, qdrant::ListAliasesResponse; + list_aliases, qdrant::ListAliasesRequest, qdrant::ListAliasesResponse; + collection_cluster_info, qdrant::CollectionClusterInfoRequest, qdrant::CollectionClusterInfoResponse; + update_collection_cluster_setup, qdrant::UpdateCollectionClusterSetupRequest, qdrant::UpdateCollectionClusterSetupResponse; + create_shard_key, qdrant::CreateShardKeyRequest, qdrant::CreateShardKeyResponse; + delete_shard_key, qdrant::DeleteShardKeyRequest, qdrant::DeleteShardKeyResponse; + list_shard_keys, qdrant::ListShardKeysRequest, qdrant::ListShardKeysResponse; + ); +} + +#[tonic::async_trait] +impl Points for FakeService { + async fn create_field_index( + &self, + request: Request, + ) -> Result, Status> { + let mut state = self.state.lock().unwrap(); + state.index_creations += 1; + state.field_indexes.push(request.into_inner()); + if state.fail_field_index { + return Err(Status::internal("field index failure")); + } + Ok(Response::new(PointsOperationResponse::default())) + } + + async fn upsert( + &self, + request: Request, + ) -> Result, Status> { + let mut state = self.state.lock().unwrap(); + for point in request.into_inner().points { + let stored = StoredPoint { + id: point.id.clone(), + vector: dense_vector(point.vectors)?, + payload: point.payload, + }; + if let Some(existing) = state + .points + .iter_mut() + .find(|existing| existing.id == stored.id) + { + *existing = stored; + } else { + state.points.push(stored); + } + } + Ok(Response::new(PointsOperationResponse::default())) + } + + async fn search( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + let key_filter = keyword_filter(request.filter.as_ref()); + let state = self.state.lock().unwrap(); + let mut results = state + .points + .iter() + .filter(|point| { + key_filter.as_ref().is_none_or(|(field, expected)| { + point + .payload + .get(field) + .and_then(|value| { + let value: serde_json::Value = value.clone().into(); + value + .as_str() + .map(str::to_owned) + .or_else(|| value.as_i64().map(|value| value.to_string())) + }) + .is_some_and(|value| value == *expected) + }) + }) + .map(|point| ScoredPoint { + id: point.id.clone(), + payload: point.payload.clone(), + score: cosine(&request.vector, &point.vector), + ..Default::default() + }) + .collect::>(); + results.sort_by(|left, right| right.score.total_cmp(&left.score)); + results.truncate(request.limit as usize); + Ok(Response::new(SearchResponse { + result: results, + ..Default::default() + })) + } + + unimplemented_points!( + delete, qdrant::DeletePoints, qdrant::PointsOperationResponse; + get, qdrant::GetPoints, qdrant::GetResponse; + update_vectors, qdrant::UpdatePointVectors, qdrant::PointsOperationResponse; + delete_vectors, qdrant::DeletePointVectors, qdrant::PointsOperationResponse; + set_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse; + overwrite_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse; + delete_payload, qdrant::DeletePayloadPoints, qdrant::PointsOperationResponse; + clear_payload, qdrant::ClearPayloadPoints, qdrant::PointsOperationResponse; + delete_field_index, qdrant::DeleteFieldIndexCollection, qdrant::PointsOperationResponse; + create_vector_name, qdrant::CreateVectorNameRequest, qdrant::PointsOperationResponse; + delete_vector_name, qdrant::DeleteVectorNameRequest, qdrant::PointsOperationResponse; + search_batch, qdrant::SearchBatchPoints, qdrant::SearchBatchResponse; + search_groups, qdrant::SearchPointGroups, qdrant::SearchGroupsResponse; + scroll, qdrant::ScrollPoints, qdrant::ScrollResponse; + recommend, qdrant::RecommendPoints, qdrant::RecommendResponse; + recommend_batch, qdrant::RecommendBatchPoints, qdrant::RecommendBatchResponse; + recommend_groups, qdrant::RecommendPointGroups, qdrant::RecommendGroupsResponse; + discover, qdrant::DiscoverPoints, qdrant::DiscoverResponse; + discover_batch, qdrant::DiscoverBatchPoints, qdrant::DiscoverBatchResponse; + count, qdrant::CountPoints, qdrant::CountResponse; + update_batch, qdrant::UpdateBatchPoints, qdrant::UpdateBatchResponse; + query, qdrant::QueryPoints, qdrant::QueryResponse; + query_batch, qdrant::QueryBatchPoints, qdrant::QueryBatchResponse; + query_groups, qdrant::QueryPointGroups, qdrant::QueryGroupsResponse; + facet, qdrant::FacetCounts, qdrant::FacetResponse; + search_matrix_pairs, qdrant::SearchMatrixPoints, qdrant::SearchMatrixPairsResponse; + search_matrix_offsets, qdrant::SearchMatrixPoints, qdrant::SearchMatrixOffsetsResponse; + ); +} + +fn dense_vector(vectors: Option) -> Result, Status> { + let Some(Vectors { + vectors_options: + Some(qdrant::vectors::VectorsOptions::Vector(Vector { + vector: Some(qdrant::vector::Vector::Dense(qdrant::DenseVector { data })), + .. + })), + }) = vectors + else { + return Err(Status::invalid_argument("expected dense vector")); + }; + Ok(data) +} + +fn keyword_filter(filter: Option<&Filter>) -> Option<(String, String)> { + filter? + .must + .iter() + .find_map(|condition| match condition.condition_one_of.as_ref()? { + qdrant::condition::ConditionOneOf::Field(field) => { + let qdrant::r#match::MatchValue::Keyword(value) = + field.r#match.as_ref()?.match_value.as_ref()? + else { + return None; + }; + Some((field.key.clone(), value.clone())) + } + _ => None, + }) +} + +fn cosine(left: &[f32], right: &[f32]) -> f32 { + let dot = left + .iter() + .zip(right) + .map(|(left, right)| left * right) + .sum::(); + let left_norm = left.iter().map(|value| value * value).sum::().sqrt(); + let right_norm = right.iter().map(|value| value * value).sum::().sqrt(); + dot / (left_norm * right_norm) +} From 93e98365247ffca3f2cd554c1129f5f8039a1bfb Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:45:15 +0000 Subject: [PATCH 05/17] feat(python-bridge): serve QdrantSemanticCache natively Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 3 + .../cache-qdrant-semantic/tests/qdrant.rs | 5 +- litellm-rust/crates/python-bridge/Cargo.toml | 3 + .../crates/python-bridge/src/cache/config.rs | 408 +++++++++++++++++- .../crates/python-bridge/src/cache/facade.rs | 38 +- .../crates/python-bridge/src/cache/handle.rs | 111 ++++- .../crates/python-bridge/src/cache/native.rs | 149 ++++++- .../crates/python-bridge/src/cache/request.rs | 28 +- tests/test_litellm_rust/test_cache.py | 251 +++++++++++ 9 files changed, 936 insertions(+), 60 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index fb75c2241b5..238b869cb0b 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2736,6 +2736,7 @@ dependencies = [ "litellm-auth-gcp", "litellm-cache", "litellm-cache-memory", + "litellm-cache-qdrant-semantic", "litellm-cache-redis", "litellm-cache-response", "litellm-callbacks-legacy-python", @@ -2748,12 +2749,14 @@ dependencies = [ "litellm-types", "pyo3", "pyo3-async-runtimes", + "qdrant-client", "rstest", "serde", "serde_json", "serde_with", "tokio", "tokio-tungstenite", + "url", ] [[package]] diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs index d5ecaf7217d..e8d8d5040f0 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -90,7 +90,10 @@ async fn connect( } #[tokio::test(flavor = "multi_thread")] -#[allow(deprecated)] +#[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), diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 1bba83922f9..b407844020d 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -24,6 +24,8 @@ litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-cache-redis.workspace = true litellm-cache-response.workspace = true +litellm-cache-qdrant-semantic.workspace = true +qdrant-client.workspace = true serde.workspace = true litellm-auth.workspace = true litellm-callbacks-legacy-python.workspace = true @@ -38,6 +40,7 @@ litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true serde_json.workspace = true +url.workspace = true tokio = { workspace = true, features = ["sync"] } [dev-dependencies] diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 5fe36f6c1fa..b21e671e431 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -1,6 +1,7 @@ -use std::time::Duration; +use std::{env, time::Duration}; use litellm_cache::CacheType; +use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, QdrantSemanticConfig, Quantization}; use litellm_cache_redis::{RedisNode, RedisTopology}; use pyo3::{ exceptions::{PyTypeError, PyValueError}, @@ -86,9 +87,30 @@ struct RedisClientProjection<'py> { const REDIS_PY_DEFAULT_MAX_CONNECTIONS: usize = 1 << 31; +pub(super) struct QdrantSemanticCacheConfig { + pub(super) grpc_url: String, + pub(super) api_key: Option, + pub(super) collection_name: String, + pub(super) similarity_threshold: f64, + pub(super) vector_size: u64, + pub(super) embedding: OpenAiEmbedderConfig, + pub(super) quantization: Quantization, +} + +impl QdrantSemanticCacheConfig { + pub(super) fn to_qdrant_config(&self) -> QdrantSemanticConfig { + QdrantSemanticConfig { + collection_name: self.collection_name.clone(), + similarity_threshold: self.similarity_threshold, + vector_size: self.vector_size, + quantization: self.quantization.clone(), + } + } +} pub(super) enum CacheBackendConfig { Memory(MemoryCacheConfig), Redis(Box), + QdrantSemantic(Box), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -103,6 +125,8 @@ pub(super) enum UnsupportedCacheConfig { RedisCredentials, RedisConnection, RedisOption, + QdrantEndpoint, + SemanticEmbedding, } impl UnsupportedCacheConfig { @@ -113,6 +137,10 @@ impl UnsupportedCacheConfig { Self::RedisCredentials => "native Redis credentials require Python", Self::RedisConnection => "native Redis connection type is not implemented", Self::RedisOption => "native Redis configuration requires Python", + Self::QdrantEndpoint => { + "native Qdrant requires the default REST port so the gRPC port can be derived" + } + Self::SemanticEmbedding => "native semantic embedding requires Python", } } } @@ -155,12 +183,18 @@ impl NativeCacheConfig { }))), Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), }, + Some(CacheType::QdrantSemantic) => match project_qdrant_semantic(&backend)? { + Ok(backend) => Ok(CacheConfigProjection::Native(Box::new(Self { + policy, + backend: CacheBackendConfig::QdrantSemantic(Box::new(backend)), + }))), + Err(reason) => Ok(CacheConfigProjection::Unsupported(reason)), + }, Some( CacheType::RedisSemantic | CacheType::ValkeySemantic | CacheType::S3 | CacheType::Disk - | CacheType::QdrantSemantic | CacheType::AzureBlob | CacheType::Gcs, ) @@ -171,18 +205,15 @@ impl NativeCacheConfig { } pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { - if service.default_ttl() - != Some(match &self.backend { - CacheBackendConfig::Memory(config) => config.default_ttl, - CacheBackendConfig::Redis(config) => config.default_ttl, - }) - { - return Some("facade and native backend default TTLs must match"); - } match &self.backend { - CacheBackendConfig::Memory(config) if service.kind() != "memory" => { + CacheBackendConfig::Memory(_) if service.kind() != "memory" => { Some("facade and native backend types must match") } + CacheBackendConfig::Memory(config) + if service.default_ttl() != Some(config.default_ttl) => + { + Some("facade and native backend default TTLs must match") + } CacheBackendConfig::Memory(config) if service.capacity() != Some(config.capacity) => { Some("facade and native backend capacities must match") } @@ -192,7 +223,7 @@ impl NativeCacheConfig { Some("facade and native backend item limits must match") } CacheBackendConfig::Memory(_) => None, - CacheBackendConfig::Redis(_) if service.kind() != "redis" => { + CacheBackendConfig::Redis(config) if service.kind() != "redis" => { Some("facade and native backend types must match") } CacheBackendConfig::Redis(config) if service.topology() != Some(&config.topology) => { @@ -200,11 +231,134 @@ impl NativeCacheConfig { } CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) - .then_some("facade and native backend namespaces must match"), + .then_some("facade and native backend namespaces must match") + .or_else(|| { + (service.default_ttl() != Some(config.default_ttl)) + .then_some("facade and native backend default TTLs must match") + }), + CacheBackendConfig::QdrantSemantic(config) if service.kind() != "qdrant_semantic" => { + Some("facade and native backend types must match") + } + CacheBackendConfig::QdrantSemantic(config) + if service.collection_name() != Some(config.collection_name.as_str()) => + { + Some("facade and native backend collections must match") + } + CacheBackendConfig::QdrantSemantic(config) + if service.similarity_threshold() != Some(config.similarity_threshold) => + { + Some("facade and native backend similarity thresholds must match") + } + CacheBackendConfig::QdrantSemantic(config) + if service.vector_size() != Some(config.vector_size) => + { + Some("facade and native backend vector sizes must match") + } + CacheBackendConfig::QdrantSemantic(config) + if service.embedding_model() != Some(config.embedding.model.as_str()) => + { + Some("facade and native backend embedding models must match") + } + CacheBackendConfig::QdrantSemantic(_) => None, } } } +#[inline(never)] +fn project_qdrant_semantic( + backend: &Bound<'_, PyAny>, +) -> PyResult> { + let rest_url = backend.getattr("qdrant_api_base")?.extract::()?; + let parsed = match url::Url::parse(&rest_url) { + Ok(value) => value, + Err(_) => return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint)), + }; + if !matches!(parsed.scheme(), "http" | "https") + || !parsed.path().is_empty() && parsed.path() != "/" + || parsed.query().is_some() + || parsed.host_str().is_none() + || parsed.port().is_some_and(|port| port != 6333) + { + return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint)); + } + let mut grpc_url = parsed; + if grpc_url.set_port(Some(6334)).is_err() { + return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint)); + } + grpc_url.set_path(""); + grpc_url.set_query(None); + + let embedding_max_input_tokens = optional_attribute_i64(backend, "embedding_max_input_tokens")?; + if embedding_max_input_tokens.is_some() { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + } + let configured_model = backend.getattr("embedding_model")?.extract::()?; + let embedding_model = configured_model + .strip_prefix("openai/") + .unwrap_or(&configured_model) + .to_owned(); + if !embedding_model.starts_with("text-embedding-") { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + } + let proxy_server = py_sys_module(backend.py())?; + if let Some(proxy_server) = proxy_server { + let router = proxy_server.getattr("llm_router")?; + let model_list = proxy_server.getattr("llm_model_list")?; + let embedding_router = backend.py().import("litellm.caching._embedding_router")?; + if !embedding_router + .getattr("resolve_embedding_router")? + .call1((embedding_model.as_str(), router, model_list))? + .is_none() + { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + } + } + let litellm = backend.py().import("litellm")?; + for name in ["api_key", "openai_key", "api_base"] { + if !litellm.getattr(name)?.is_none() { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + } + } + let Ok(embedding_api_key) = env::var("OPENAI_API_KEY") else { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + }; + if embedding_api_key.is_empty() { + return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); + } + let embedding_api_base = env::var("OPENAI_BASE_URL") + .or_else(|_| env::var("OPENAI_API_BASE")) + .unwrap_or_else(|_| "https://api.openai.com/v1".to_owned()); + let timeout = optional_attribute_f64(backend, "embedding_timeout")? + .map(duration) + .transpose()?; + Ok(Ok(QdrantSemanticCacheConfig { + grpc_url: grpc_url.to_string().trim_end_matches('/').to_owned(), + api_key: optional_string(backend.getattr("qdrant_api_key")?)?, + collection_name: backend.getattr("collection_name")?.extract()?, + similarity_threshold: backend.getattr("similarity_threshold")?.extract()?, + vector_size: backend.getattr("vector_size")?.extract::()?, + embedding: OpenAiEmbedderConfig { + api_base: embedding_api_base, + api_key: embedding_api_key, + model: embedding_model, + timeout, + }, + quantization: Quantization::Binary, + })) +} + +fn py_sys_module(py: Python<'_>) -> PyResult>> { + match py + .import("sys")? + .getattr("modules")? + .get_item("litellm.proxy.proxy_server") + { + Ok(module) => Ok(Some(module)), + Err(error) if error.is_instance_of::(py) => Ok(None), + Err(error) => Err(error), + } +} + #[inline(never)] fn project_memory(backend: &Bound<'_, PyAny>) -> PyResult { let max_size_kib = backend.getattr("max_size_per_item")?.extract::()?; @@ -514,6 +668,28 @@ fn optional_attribute_string(value: &Bound<'_, PyAny>, name: &str) -> PyResult, name: &str) -> PyResult> { + match value.getattr(name) { + Ok(attribute) => attribute.extract::>(), + Err(error) if error.is_instance_of::(value.py()) => { + Ok(None) + } + Err(error) => Err(error), + } +} + +#[inline(never)] +fn optional_attribute_f64(value: &Bound<'_, PyAny>, name: &str) -> PyResult> { + match value.getattr(name) { + Ok(attribute) => attribute.extract::>(), + Err(error) if error.is_instance_of::(value.py()) => { + Ok(None) + } + Err(error) => Err(error), + } +} + #[inline(never)] fn optional_string(value: Bound<'_, PyAny>) -> PyResult> { Ok(value @@ -597,6 +773,11 @@ fn optional_dict_duration(values: &Bound<'_, PyDict>, key: &str) -> PyResult &'static Mutex<()> { + static LOCK: OnceLock> = OnceLock::new(); + LOCK.get_or_init(|| Mutex::new(())) + } + + fn qdrant_facade<'py>(py: Python<'py>, extra: &str) -> Bound<'py, PyAny> { + install_fake_litellm(py); + facade( + py, + &format!( + "backend = SimpleNamespace(qdrant_api_base='https://qdrant.example:6333', qdrant_api_key='qdrant-key', collection_name='cache', similarity_threshold=0.99, embedding_model='openai/text-embedding-3-small', vector_size=8, embedding_max_input_tokens=None, embedding_timeout=None)\n\ + facade = SimpleNamespace(type='qdrant-semantic', mode='default-on', ttl=None, namespace=None, supported_call_types=None, redis_flush_size=None, semantic_cache_scope='key', cache=backend)\n\ + {extra}" + ), + ) + } + + fn install_fake_litellm(py: Python<'_>) { + py.run( + c" +import sys +import types +litellm = types.ModuleType('litellm') +litellm.api_key = None +litellm.openai_key = None +litellm.api_base = None +litellm.__path__ = [] +caching = types.ModuleType('litellm.caching') +caching.__path__ = [] +embedding_router = types.ModuleType('litellm.caching._embedding_router') +embedding_router.resolve_embedding_router = lambda *_args: None +caching._embedding_router = embedding_router +litellm.caching = caching +sys.modules['litellm'] = litellm +sys.modules['litellm.caching'] = caching +sys.modules['litellm.caching._embedding_router'] = embedding_router +", + None, + None, + ) + .unwrap(); + } + + fn configure_embedding_environment<'py>( + py: Python<'py>, + key: Option<&str>, + ) -> PyResult> { + let environ = py.import("os")?.getattr("environ")?; + let prior = environ.call_method1("get", ("OPENAI_API_KEY",))?; + match key { + Some(key) => environ.set_item("OPENAI_API_KEY", key)?, + None => environ.del_item("OPENAI_API_KEY")?, + } + Ok(prior) + } + fn facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> { let locals = PyDict::new(py); py.run( @@ -740,7 +977,6 @@ mod tests { assert_eq!(reason.message(), "native Redis credentials require Python"); }); } - #[test] fn projects_cluster_startup_nodes_as_redis_topology() { Python::initialize(); @@ -823,4 +1059,146 @@ mod tests { } }); } + #[test] + fn projects_qdrant_configuration_from_python() { + let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap(); + let facade = qdrant_facade(py, ""); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("Qdrant cache should be supported"); + }; + let CacheBackendConfig::QdrantSemantic(config) = config.backend else { + panic!("expected Qdrant configuration"); + }; + assert_eq!(config.grpc_url, "https://qdrant.example:6334"); + assert_eq!(config.api_key.as_deref(), Some("qdrant-key")); + assert_eq!(config.collection_name, "cache"); + assert_eq!(config.vector_size, 8); + assert_eq!(config.embedding.api_key, "embedding-key"); + assert_eq!(config.embedding.model, "text-embedding-3-small"); + let environ = py.import("os").unwrap().getattr("environ").unwrap(); + if prior.is_none() { + environ.del_item("OPENAI_API_KEY").unwrap(); + } else { + environ.set_item("OPENAI_API_KEY", prior).unwrap(); + } + }); + } + + #[test] + fn qdrant_projection_rejects_non_default_port() { + let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap(); + let facade = qdrant_facade( + py, + "backend.qdrant_api_base = 'https://qdrant.example:6332'", + ); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("non-default Qdrant port should stay on Python"); + }; + assert!(matches!(reason, UnsupportedCacheConfig::QdrantEndpoint)); + let environ = py.import("os").unwrap().getattr("environ").unwrap(); + if prior.is_none() { + environ.del_item("OPENAI_API_KEY").unwrap(); + } else { + environ.set_item("OPENAI_API_KEY", prior).unwrap(); + } + }); + } + + #[test] + fn qdrant_projection_rejects_python_embedding_features() { + let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap(); + let cases = [ + ("backend.embedding_max_input_tokens = 100", "semantic"), + ("backend.embedding_model = 'cohere/embed'", "semantic"), + ]; + for (extra, _) in cases { + let facade = qdrant_facade(py, extra); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("unsupported embedding should stay on Python"); + }; + assert!(matches!(reason, UnsupportedCacheConfig::SemanticEmbedding)); + } + let environ = py.import("os").unwrap().getattr("environ").unwrap(); + if prior.is_none() { + environ.del_item("OPENAI_API_KEY").unwrap(); + } else { + environ.set_item("OPENAI_API_KEY", prior).unwrap(); + } + }); + } + + #[test] + fn qdrant_projection_rejects_configured_litellm_base_or_missing_key() { + let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap(); + let facade = qdrant_facade(py, ""); + let litellm = py.import("litellm").unwrap(); + litellm + .setattr("api_base", "https://proxy.example") + .unwrap(); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("configured LiteLLM base should stay on Python"); + }; + assert!(matches!(reason, UnsupportedCacheConfig::SemanticEmbedding)); + litellm.setattr("api_base", py.None()).unwrap(); + configure_embedding_environment(py, None).unwrap(); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("missing embedding key should stay on Python"); + }; + assert!(matches!(reason, UnsupportedCacheConfig::SemanticEmbedding)); + let environ = py.import("os").unwrap().getattr("environ").unwrap(); + if prior.is_none() { + environ.del_item("OPENAI_API_KEY").unwrap(); + } else { + environ.set_item("OPENAI_API_KEY", prior).unwrap(); + } + }); + } + + #[test] + fn qdrant_service_mismatch_reports_type_before_starting_qdrant() { + let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap(); + let facade = qdrant_facade(py, ""); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("Qdrant cache should be supported"); + }; + let service = NativeResponseCache::memory(1, Duration::from_secs(1), 1024); + assert_eq!( + config.service_mismatch(&service), + Some("facade and native backend types must match") + ); + let environ = py.import("os").unwrap().getattr("environ").unwrap(); + if prior.is_none() { + environ.del_item("OPENAI_API_KEY").unwrap(); + } else { + environ.set_item("OPENAI_API_KEY", prior).unwrap(); + } + }); + } } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index a9ce2ae7756..14e1bfe91ca 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -227,6 +227,11 @@ impl FacadeGuard { "RedisClusterCache", "redis", ), + "qdrant_semantic" => ( + "litellm.caching.qdrant_semantic_cache", + "QdrantSemanticCache", + "qdrant-semantic", + ), _ => unreachable!(), }; let backend = facade.getattr("cache")?; @@ -246,6 +251,26 @@ impl FacadeGuard { if let Some(message) = config.service_mismatch(service) { return Err(PyTypeError::new_err(message)); } + let backend_config_names = match kind { + "memory" | "redis" => &[ + "namespace", + "default_ttl", + "max_size_in_memory", + "max_size_per_item", + "redis_kwargs", + "redis_flush_size", + ][..], + "qdrant_semantic" => &[ + "qdrant_api_base", + "collection_name", + "similarity_threshold", + "embedding_model", + "vector_size", + "embedding_max_input_tokens", + "embedding_timeout", + ][..], + _ => unreachable!(), + }; Ok(Self { outer: ObjectGuard::capture( py, @@ -260,18 +285,7 @@ impl FacadeGuard { "semantic_cache_scope", ], )?, - backend: ObjectGuard::capture( - py, - &backend, - &[ - "namespace", - "default_ttl", - "max_size_in_memory", - "max_size_per_item", - "redis_kwargs", - "redis_flush_size", - ], - )?, + backend: ObjectGuard::capture(py, &backend, backend_config_names)?, redis_pool: match (kind, cluster) { ("redis", false) => Some(RedisPoolGuard::capture(&backend, STANDALONE_POOL)?), ("redis", true) => Some(RedisPoolGuard::capture(&backend, CLUSTER_POOL)?), diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 2ee2b0c2c8c..9709bd68791 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,8 +1,16 @@ -use litellm_cache_redis::{RedisNode, RedisTopology}; -use litellm_host_python::release_gil; -use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use std::env; -use super::{cache_error, facade::FacadeGuard, native::NativeResponseCache, request::duration}; +use litellm_cache_redis::{RedisNode, RedisTopology}; + +use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, Quantization}; +use litellm_host_python::{release_gil, run_sync_value}; +use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use url::Url; + +use super::{ + cache_error, config::QdrantSemanticCacheConfig, facade::FacadeGuard, + native::NativeResponseCache, request::duration, +}; #[pyclass(frozen, name = "_CacheTestHandle")] pub(crate) struct CacheTestHandle { @@ -64,6 +72,101 @@ impl CacheTestHandle { }) } + #[staticmethod] + #[pyo3(signature = (url, *, collection_name, similarity_threshold, vector_size, embedding_model="text-embedding-3-small", api_key=None, embedding_api_key=None, embedding_api_base=None, embedding_timeout_seconds=None, quantization="binary"))] + #[expect( + clippy::too_many_arguments, + reason = "the test handle exposes the complete Qdrant constructor" + )] + fn qdrant_semantic( + py: Python<'_>, + url: String, + collection_name: String, + similarity_threshold: f64, + vector_size: u64, + embedding_model: &str, + api_key: Option, + embedding_api_key: Option, + embedding_api_base: Option, + embedding_timeout_seconds: Option, + quantization: &str, + ) -> PyResult { + let parsed = Url::parse(&url).map_err(|_| { + pyo3::exceptions::PyValueError::new_err( + "native Qdrant requires the default REST port so the gRPC port can be derived", + ) + })?; + if !matches!(parsed.scheme(), "http" | "https") + || (!parsed.path().is_empty() && parsed.path() != "/") + || parsed.query().is_some() + || parsed.host_str().is_none() + || parsed.port().is_some_and(|port| port != 6333) + { + return Err(pyo3::exceptions::PyValueError::new_err( + "native Qdrant requires the default REST port so the gRPC port can be derived", + )); + } + let mut grpc_url = parsed; + grpc_url.set_port(Some(6334)).map_err(|_| { + pyo3::exceptions::PyValueError::new_err( + "native Qdrant requires the default REST port so the gRPC port can be derived", + ) + })?; + grpc_url.set_path(""); + grpc_url.set_query(None); + let embedding_api_key = embedding_api_key + .or_else(|| { + env::var("OPENAI_API_KEY") + .ok() + .filter(|value| !value.is_empty()) + }) + .ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err( + "native semantic embedding requires an OpenAI API key", + ) + })?; + let embedding_api_base = embedding_api_base.unwrap_or_else(|| { + env::var("OPENAI_BASE_URL") + .or_else(|_| env::var("OPENAI_API_BASE")) + .unwrap_or_else(|_| "https://api.openai.com/v1".to_owned()) + }); + let quantization = match quantization { + "binary" => Quantization::Binary, + "scalar" => Quantization::Scalar, + "product" => Quantization::Product, + _ => { + return Err(pyo3::exceptions::PyValueError::new_err( + "unsupported Qdrant quantization", + )); + } + }; + let config = QdrantSemanticCacheConfig { + grpc_url: grpc_url.to_string().trim_end_matches('/').to_owned(), + api_key, + collection_name, + similarity_threshold, + vector_size, + embedding: OpenAiEmbedderConfig { + api_base: embedding_api_base, + api_key: embedding_api_key, + model: embedding_model.to_owned(), + timeout: embedding_timeout_seconds.map(duration).transpose()?, + }, + quantization, + }; + let service = run_sync_value(py, async move { + let handle = tokio::runtime::Handle::current(); + NativeResponseCache::qdrant_semantic(config, handle) + .await + .map_err(cache_error) + })?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[getter] fn backend(&self) -> &'static str { self.service.kind() diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index b23038dee65..6e28b8414d3 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -1,13 +1,16 @@ use std::{sync::Arc, time::Duration}; -use litellm_cache::{CacheCodec, CacheConnectionResult, Error}; +use litellm_cache::{CacheCodec, CacheConnectionResult, Error, SemanticCacheContext}; use litellm_cache_memory::InMemoryCache; +use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, QdrantSemanticCache}; use litellm_cache_redis::{RedisCache, RedisTopology}; use litellm_cache_response::{ CacheEntry, PartialHits, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, WriteBuffer, }; use serde_json::Value; +use super::{config::QdrantSemanticCacheConfig, request::exact}; + #[derive(Clone)] pub(super) enum NativeResponseCache { Memory(Arc>>), @@ -15,6 +18,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, buffer: Option>, }, + QdrantSemantic(Arc>>), } impl NativeResponseCache { @@ -45,6 +49,30 @@ impl NativeResponseCache { buffer: None, }) } + + pub async fn qdrant_semantic( + config: QdrantSemanticCacheConfig, + runtime: tokio::runtime::Handle, + ) -> Result { + let client = qdrant_client::Qdrant::from_url(&config.grpc_url) + .skip_compatibility_check() + .api_key(config.api_key.as_deref()) + .build() + .map_err(|_| Error::Unavailable)?; + let qdrant_config = config.to_qdrant_config(); + let embedder = OpenAiEmbedder::new(config.embedding)?; + let cache = QdrantSemanticCache::connect( + client, + embedder, + ResponseCacheCodec, + qdrant_config, + runtime, + ) + .await?; + Ok(Self::QdrantSemantic(Arc::new(ResponseCache::new( + Arc::new(cache), + )))) + } } impl NativeResponseCache { @@ -52,6 +80,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => "memory", Self::Redis { .. } => "redis", + Self::QdrantSemantic(_) => "qdrant_semantic", } } @@ -59,6 +88,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.default_ttl(), Self::Redis { cache, .. } => cache.default_ttl(), + Self::QdrantSemantic(_) => None, } } @@ -66,6 +96,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => None, Self::Redis { cache, .. } => cache.backend().namespace(), + Self::QdrantSemantic(_) => None, } } @@ -80,6 +111,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => Some(cache.backend().max_size_in_memory()), Self::Redis { .. } => None, + Self::QdrantSemantic(_) => None, } } @@ -87,6 +119,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.backend().max_entry_bytes(), Self::Redis { .. } => None, + Self::QdrantSemantic(_) => None, } } @@ -100,89 +133,157 @@ impl NativeResponseCache { } } + pub fn collection_name(&self) -> Option<&str> { + match self { + Self::QdrantSemantic(cache) => Some(cache.backend().collection_name()), + _ => None, + } + } + + pub fn similarity_threshold(&self) -> Option { + match self { + Self::QdrantSemantic(cache) => Some(cache.backend().similarity_threshold()), + _ => None, + } + } + + pub fn vector_size(&self) -> Option { + match self { + Self::QdrantSemantic(cache) => Some(cache.backend().vector_size()), + _ => None, + } + } + + pub fn embedding_model(&self) -> Option<&str> { + match self { + Self::QdrantSemantic(cache) => Some(cache.backend().embedder().model()), + _ => None, + } + } + pub fn lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { match self { - Self::Memory(cache) => cache.lookup(request, now), - Self::Redis { cache, .. } => cache.lookup(request, now), + Self::Memory(cache) => cache.lookup(&exact(request), now), + Self::Redis { cache, .. } => cache.lookup(&exact(request), now), + Self::QdrantSemantic(cache) => cache.lookup(request, now), } } pub fn store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.store(request, response, now), - Self::Redis { cache, .. } => cache.store(request, response, now), + Self::Memory(cache) => cache.store(&exact(request), response, now), + Self::Redis { cache, .. } => cache.store(&exact(request), response, now), + Self::QdrantSemantic(cache) => cache.store(request, response, now), } } pub fn lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result { match self { - Self::Memory(cache) => cache.lookup_batch(requests, now), - Self::Redis { cache, .. } => cache.lookup_batch(requests, now), + Self::Memory(cache) => { + cache.lookup_batch(&requests.iter().map(exact).collect::>(), now) + } + Self::Redis { cache, .. } => { + cache.lookup_batch(&requests.iter().map(exact).collect::>(), now) + } + Self::QdrantSemantic(_) => Err(Error::UnsupportedOperation), } } pub async fn async_lookup( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, now: Duration, ) -> Result, Error> { match self { - Self::Memory(cache) => cache.async_lookup(request, now).await, - Self::Redis { cache, .. } => cache.async_lookup(request, now).await, + Self::Memory(cache) => cache.async_lookup(&exact(request), now).await, + Self::Redis { cache, .. } => cache.async_lookup(&exact(request), now).await, + Self::QdrantSemantic(cache) => cache.async_lookup(request, now).await, } } pub async fn async_store( &self, - request: &ResponseCacheRequest, + request: &ResponseCacheRequest, response: Value, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_store(request, response, now).await, + Self::Memory(cache) => cache.async_store(&exact(request), response, now).await, Self::Redis { cache, buffer: None, - } => cache.async_store(request, response, now).await, + } => cache.async_store(&exact(request), response, now).await, Self::Redis { cache, buffer: Some(buffer), - } => buffer.async_store(cache, request, response, now).await, + } => { + let request = exact(request); + buffer.async_store(cache, &request, response, now).await + } + Self::QdrantSemantic(cache) => cache.async_store(request, response, now).await, } } pub async fn async_lookup_batch( &self, - requests: &[ResponseCacheRequest], + requests: &[ResponseCacheRequest], now: Duration, ) -> Result { match self { - Self::Memory(cache) => cache.async_lookup_batch(requests, now).await, - Self::Redis { cache, .. } => cache.async_lookup_batch(requests, now).await, + Self::Memory(cache) => { + let requests = requests.iter().map(exact).collect::>(); + cache.async_lookup_batch(&requests, now).await + } + Self::Redis { cache, .. } => { + let requests = requests.iter().map(exact).collect::>(); + cache.async_lookup_batch(&requests, now).await + } + Self::QdrantSemantic(_) => Err(Error::UnsupportedOperation), } } pub async fn async_store_batch( &self, - entries: Vec<(ResponseCacheRequest, Value)>, + entries: Vec<(ResponseCacheRequest, Value)>, now: Duration, ) -> Result<(), Error> { match self { - Self::Memory(cache) => cache.async_store_batch(entries, now).await, - Self::Redis { cache, .. } => cache.async_store_batch(entries, now).await, + Self::Memory(cache) => { + cache + .async_store_batch( + entries + .into_iter() + .map(|(request, value)| (exact(&request), value)) + .collect(), + now, + ) + .await + } + Self::Redis { cache, .. } => { + cache + .async_store_batch( + entries + .into_iter() + .map(|(request, value)| (exact(&request), value)) + .collect(), + now, + ) + .await + } + Self::QdrantSemantic(cache) => cache.async_store_batch(entries, now).await, } } @@ -195,6 +296,7 @@ impl NativeResponseCache { } cache.async_flush().await } + Self::QdrantSemantic(_) => Err(Error::UnsupportedOperation), } } @@ -202,6 +304,7 @@ impl NativeResponseCache { match self { Self::Memory(cache) => cache.test_connection().await, Self::Redis { cache, .. } => cache.test_connection().await, + Self::QdrantSemantic(_) => Err(Error::UnsupportedOperation), } } } diff --git a/litellm-rust/crates/python-bridge/src/cache/request.rs b/litellm-rust/crates/python-bridge/src/cache/request.rs index 0067fc4392b..0f50f33e6b9 100644 --- a/litellm-rust/crates/python-bridge/src/cache/request.rs +++ b/litellm-rust/crates/python-bridge/src/cache/request.rs @@ -1,10 +1,11 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH}; -use litellm_cache::ExactCacheContext; +use litellm_cache::{ExactCacheContext, SemanticCacheContext, SemanticCacheScope}; use litellm_cache_response::{CacheControls, CacheKeyInput, ResponseCacheRequest}; use litellm_host_python::from_py; use pyo3::{exceptions::PyValueError, prelude::*}; use serde::Deserialize; +use serde_json::{Map, Value}; #[derive(Deserialize)] #[serde(deny_unknown_fields)] @@ -13,34 +14,51 @@ struct RequestInput { controls: Option, ttl_seconds: Option, max_age_seconds: Option, + messages: Option>, + input: Option, + metadata: Option>, + scope: Option, } pub(super) fn request( value: &Bound<'_, PyAny>, -) -> PyResult> { +) -> PyResult> { let input: RequestInput = from_py(value)?; request_input(input) } -fn request_input(input: RequestInput) -> PyResult> { - let mut request: ResponseCacheRequest = ResponseCacheRequest::new(input.key); +fn request_input(input: RequestInput) -> PyResult> { + let mut request: ResponseCacheRequest = + ResponseCacheRequest::new(input.key); if let Some(controls) = input.controls { request.controls = controls; } request.context.ttl = input.ttl_seconds.map(duration).transpose()?; + request.context.messages = input.messages.unwrap_or_default(); + request.context.input = input.input; + request.context.metadata = input.metadata.unwrap_or_default(); + request.context.scope = input.scope.unwrap_or_default(); request.max_age = input.max_age_seconds.map(duration).transpose()?; Ok(request) } pub(super) fn requests( value: &Bound<'_, PyAny>, -) -> PyResult>> { +) -> PyResult>> { from_py::>(value)? .into_iter() .map(request_input) .collect() } +pub(super) fn exact( + request: &ResponseCacheRequest, +) -> ResponseCacheRequest { + request.clone().with_context(ExactCacheContext { + ttl: request.context.ttl, + }) +} + pub(super) fn duration(seconds: f64) -> PyResult { Duration::try_from_secs_f64(seconds) .map_err(|_| PyValueError::new_err("cache durations must be finite and nonnegative")) diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 3903ef65daa..d33eaec0c27 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -1,7 +1,10 @@ import asyncio import contextvars import gc +import hashlib +import http.server import json +import math import os import threading import time @@ -10,6 +13,7 @@ from collections.abc import Generator from types import SimpleNamespace from typing import Final, Protocol, cast from urllib.parse import urlparse +from uuid import uuid4 import fakeredis import pytest @@ -34,6 +38,71 @@ def request(key: str = "key") -> dict[str, object]: return {"key": {"preset": key}} +def semantic_request( + key: str, + messages: list[dict[str, object]], + **kwargs: object, +) -> dict[str, object]: + return {**request(key), "messages": messages, **kwargs} + + +def embedding_vector(text: str) -> list[float]: + raw = hashlib.sha256(text.encode()).digest()[:8] + values: Final = [byte / 127.5 - 1 for byte in raw] + norm: Final = math.sqrt(sum(value * value for value in values)) + return [value / norm for value in values] + + +@pytest.fixture +def qdrant_url() -> str: + value: Final[str | None] = os.environ.get("QDRANT_URL") + if not value: + pytest.skip("QDRANT_URL is required for Qdrant semantic cache tests") + return value.rstrip("/") + + +@pytest.fixture +def fake_embedding_endpoint(monkeypatch: pytest.MonkeyPatch) -> Generator[str]: + class EmbeddingHandler(http.server.BaseHTTPRequestHandler): + def do_POST(self) -> None: + length: Final = int(self.headers["Content-Length"]) + body: Final = json.loads(self.rfile.read(length)) + text: Final = body["input"] + response: Final = { + "object": "list", + "data": [ + { + "object": "embedding", + "index": 0, + "embedding": embedding_vector(text), + } + ], + "model": body["model"], + "usage": {"prompt_tokens": 1, "total_tokens": 1}, + } + encoded: Final = json.dumps(response).encode() + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(encoded))) + self.end_headers() + self.wfile.write(encoded) + + def log_message(self, *_args: object) -> None: + return + + server: Final = http.server.ThreadingHTTPServer(("127.0.0.1", 0), EmbeddingHandler) + worker: Final = threading.Thread(target=server.serve_forever, daemon=True) + worker.start() + monkeypatch.setenv("OPENAI_API_BASE", f"http://127.0.0.1:{server.server_address[1]}") + monkeypatch.setenv("OPENAI_API_KEY", "sk-test") + try: + yield f"http://127.0.0.1:{server.server_address[1]}" + finally: + server.shutdown() + server.server_close() + worker.join(timeout=5) + + @pytest.fixture def redis_url() -> Generator[str]: server: Final = fakeredis.TcpFakeServer(("127.0.0.1", 0), server_type="redis") @@ -464,3 +533,185 @@ async def test_redis_cluster_facade_serves_multi_slot_batches_and_scoped_flush_n client.delete("unscoped") client.close() facade.cache.redis_client.close() + + +def qdrant_facade( + qdrant_url: str, + collection_name: str, +) -> Cache: + return Cache( + type=LiteLLMCacheType.QDRANT_SEMANTIC, + qdrant_api_base=qdrant_url, + qdrant_collection_name=collection_name, + similarity_threshold=0.99, + qdrant_semantic_cache_embedding_model="text-embedding-3-small", + qdrant_semantic_cache_vector_size=8, + ) + + +def test_qdrant_semantic_facade_binds_native_and_shares_entries( + qdrant_url: str, + fake_embedding_endpoint: str, +) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "shared prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + facade.cache.set_cache( + "python-key", + {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, + messages=messages, + ) + handle: Final = _native._CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + assert binding.kind == "native" + assert binding.lookup(semantic_request("python-key", messages)) == {"id": "py"} + binding.store(semantic_request("native-key", messages), {"id": "native"}) + assert facade.cache.get_cache("native-key", messages=messages) == {"id": "native"} + unrelated: Final = [{"role": "user", "content": "unrelated prompt"}] + assert binding.lookup(semantic_request("native-key", unrelated)) is None + assert facade.cache.get_cache("native-key", messages=unrelated) is None + assert binding.lookup(semantic_request("different-key", messages)) is None + assert facade.cache.get_cache("different-key", messages=messages) is None + + +async def test_qdrant_semantic_async_parity( + qdrant_url: str, + fake_embedding_endpoint: str, +) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "async prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = _native._CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + await facade.cache.async_set_cache( + "python-key", + {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, + messages=messages, + ) + assert await binding.async_lookup(semantic_request("python-key", messages)) == {"id": "py"} + await binding.async_store(semantic_request("native-key", messages), {"id": "native"}) + assert await facade.cache.async_get_cache("native-key", messages=messages) == {"id": "native"} + + +async def test_qdrant_semantic_malformed_entries_and_unsupported_operations( + qdrant_url: str, + fake_embedding_endpoint: str, +) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "malformed prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = _native._CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + key: Final = "malformed-key" + response: Final = { + "points": [ + { + "id": str(uuid4()), + "vector": embedding_vector("malformed prompt"), + "payload": { + "litellm_cache_key": key, + "text": "malformed prompt", + "response": "not json", + }, + } + ] + } + facade.cache.sync_client.put( + url=f"{qdrant_url}/collections/{collection}/points", + headers=facade.cache.headers, + json=response, + ) + assert binding.lookup(semantic_request(key, messages)) is None + with pytest.raises(RuntimeError, match="does not support"): + binding.lookup_batch([semantic_request(key, messages)]) + with pytest.raises(RuntimeError, match="does not support"): + await binding.async_flush() + with pytest.raises(RuntimeError, match="does not support"): + await binding.ping() + + +def test_qdrant_semantic_ignores_request_expiry( + qdrant_url: str, + fake_embedding_endpoint: str, +) -> None: + del fake_embedding_endpoint + messages: Final = [{"role": "user", "content": "persistent prompt"}] + collection: Final = f"cache_{uuid4().hex}" + facade: Final = qdrant_facade(qdrant_url, collection) + handle: Final = _native._CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + binding: Final = _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve() + binding.store( + semantic_request("persistent-key", messages, ttl_seconds=1.0), + {"id": "persistent"}, + ) + time.sleep(1.2) + assert binding.lookup(semantic_request("persistent-key", messages)) == {"id": "persistent"} + assert facade.cache.get_cache("persistent-key", messages=messages) == {"id": "persistent"} + + +def test_qdrant_semantic_mutation_and_projection_fallback( + 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) + handle: Final = _native._CacheTestHandle.qdrant_semantic( + qdrant_url, + collection_name=collection, + similarity_threshold=0.99, + vector_size=8, + ) + handle._bind_facade(facade) + facade.cache.similarity_threshold = 0.5 + assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" + unsupported: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") + unsupported.cache.embedding_max_input_tokens = 100 + with pytest.raises(TypeError, match="requires Python"): + handle._bind_facade(unsupported) + unsupported.cache.embedding_max_input_tokens = None + unsupported.cache.qdrant_api_base = "http://127.0.0.1:7777" + with pytest.raises(TypeError, match="gRPC"): + handle._bind_facade(unsupported) + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + class CustomQdrantSemanticCache(QdrantSemanticCache): + pass + + subclass_facade: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") + subclass_facade.cache = CustomQdrantSemanticCache( + qdrant_api_base=qdrant_url, + collection_name=subclass_facade.cache.collection_name, + similarity_threshold=0.99, + embedding_model="text-embedding-3-small", + vector_size=8, + ) + with pytest.raises(TypeError): + handle._bind_facade(subclass_facade) From fbaa53565746bc8cc3655af856ef82ec54097828 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:52:51 +0000 Subject: [PATCH 06/17] fix(python-bridge): ignore class data defaults in the facade guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/python-bridge/src/cache/facade.rs | 97 +++++++++++++++++-- tests/test_litellm_rust/test_cache.py | 23 ++++- 2 files changed, 106 insertions(+), 14 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 14e1bfe91ca..dbcbd9bda92 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -59,6 +59,34 @@ pub(super) struct FacadeGuard { } impl ObjectGuard { + fn class_behaviors(class: &Bound<'_, PyType>) -> PyResult)>> { + let py = class.py(); + let builtins = py.import("builtins")?; + let property_type = builtins.getattr("property")?; + let staticmethod_type = builtins.getattr("staticmethod")?; + let classmethod_type = builtins.getattr("classmethod")?; + class + .getattr("__dict__")? + .call_method0("items")? + .try_iter()? + .map(|item| { + let item = item?; + let (name, value): (String, Py) = item.extract()?; + let value_bound = value.bind(py); + let is_behavior = value_bound.is_callable() + || value_bound.is_instance(&property_type)? + || value_bound.is_instance(&staticmethod_type)? + || value_bound.is_instance(&classmethod_type)?; + Ok(is_behavior.then_some((name, value))) + }) + .filter_map(|result| match result { + Ok(Some(attribute)) => Some(Ok(attribute)), + Ok(None) => None, + Err(error) => Some(Err(error)), + }) + .collect() + } + fn capture( py: Python<'_>, object: &Bound<'_, PyAny>, @@ -71,12 +99,7 @@ impl ObjectGuard { .iter() .map(|class| { let class = class.cast_into::()?; - let attributes = class - .getattr("__dict__")? - .call_method0("items")? - .try_iter()? - .map(|item| item?.extract::<(String, Py)>()) - .collect::>>()?; + let attributes = Self::class_behaviors(&class)?; Ok(ClassGuard { class: class.unbind(), attributes, @@ -129,15 +152,21 @@ impl ObjectGuard { } let instance = object.getattr("__dict__")?.cast_into::()?; for (class, expected) in mro.iter().zip(&self.classes) { + let class = class.cast_into::()?; if !class.is(expected.class.bind(py)) { return Ok(false); } - let attributes = class.getattr("__dict__")?; - if attributes.len()? != expected.attributes.len() { + let attributes = Self::class_behaviors(&class)?; + if attributes.len() != expected.attributes.len() { return Ok(false); } - for (name, value) in &expected.attributes { - if instance.contains(name)? || !attributes.get_item(name)?.is(value.bind(py)) { + for ((name, value), (expected_name, expected_value)) in + attributes.iter().zip(&expected.attributes) + { + if name != expected_name + || instance.contains(name)? + || !value.bind(py).is(expected_value.bind(py)) + { return Ok(false); } } @@ -342,3 +371,51 @@ pub(super) fn resolve( } handle.service().map(Some) } + +#[cfg(test)] +mod tests { + use super::ObjectGuard; + use pyo3::{prelude::*, types::PyDict}; + + #[test] + fn class_data_shadowing_is_ignored_but_method_mutations_are_rejected() { + Python::initialize(); + Python::attach(|py| { + let namespace = PyDict::new(py); + py.run( + c"class Example:\n data = 1\n def method(self):\n return 1\nobject = Example()\nobject.data = 2", + None, + Some(&namespace), + ) + .unwrap(); + let object = namespace.get_item("object").unwrap().unwrap(); + let guard = ObjectGuard::capture(py, &object, &[]).unwrap(); + + assert!(guard.matches(py, &object).unwrap()); + + py.run(c"object.method = lambda: 2", None, Some(&namespace)) + .unwrap(); + assert!(!guard.matches(py, &object).unwrap()); + }); + } + + #[test] + fn class_method_replacement_is_rejected() { + Python::initialize(); + Python::attach(|py| { + let namespace = PyDict::new(py); + py.run( + c"class Example:\n def method(self):\n return 1\nobject = Example()", + None, + Some(&namespace), + ) + .unwrap(); + let object = namespace.get_item("object").unwrap().unwrap(); + let guard = ObjectGuard::capture(py, &object, &[]).unwrap(); + + py.run(c"Example.method = lambda self: 2", None, Some(&namespace)) + .unwrap(); + assert!(!guard.matches(py, &object).unwrap()); + }); + } +} diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index d33eaec0c27..02c159825bc 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -573,7 +573,9 @@ def test_qdrant_semantic_facade_binds_native_and_shares_entries( assert binding.kind == "native" assert binding.lookup(semantic_request("python-key", messages)) == {"id": "py"} binding.store(semantic_request("native-key", messages), {"id": "native"}) - assert facade.cache.get_cache("native-key", messages=messages) == {"id": "native"} + python_value: Final = facade.cache.get_cache("native-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "native"} unrelated: Final = [{"role": "user", "content": "unrelated prompt"}] assert binding.lookup(semantic_request("native-key", unrelated)) is None assert facade.cache.get_cache("native-key", messages=unrelated) is None @@ -602,9 +604,20 @@ async def test_qdrant_semantic_async_parity( {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, messages=messages, ) - assert await binding.async_lookup(semantic_request("python-key", messages)) == {"id": "py"} + + async def lookup_after_commit() -> object: + for _ in range(20): + value: Final = await binding.async_lookup(semantic_request("python-key", messages)) + if value is not None: + return value + await asyncio.sleep(0.1) + return None + + assert await lookup_after_commit() == {"id": "py"} await binding.async_store(semantic_request("native-key", messages), {"id": "native"}) - assert await facade.cache.async_get_cache("native-key", messages=messages) == {"id": "native"} + python_value: Final = await facade.cache.async_get_cache("native-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "native"} async def test_qdrant_semantic_malformed_entries_and_unsupported_operations( @@ -673,7 +686,9 @@ def test_qdrant_semantic_ignores_request_expiry( ) time.sleep(1.2) assert binding.lookup(semantic_request("persistent-key", messages)) == {"id": "persistent"} - assert facade.cache.get_cache("persistent-key", messages=messages) == {"id": "persistent"} + python_value: Final = facade.cache.get_cache("persistent-key", messages=messages) + assert isinstance(python_value, dict) + assert python_value["response"] == {"id": "persistent"} def test_qdrant_semantic_mutation_and_projection_fallback( From 877d5da419545db207cf4fa67ec569398c302993 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:34:44 +0000 Subject: [PATCH 07/17] fix(cache-qdrant-semantic): wait for Qdrant upserts to be indexed Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../cache-qdrant-semantic/src/semantic.rs | 2 +- .../cache-qdrant-semantic/tests/qdrant.rs | 53 +++++++++++++++---- .../tests/support/mod.rs | 29 +++++++++- litellm/caching/qdrant_semantic_cache.py | 2 + .../caching/test_qdrant_semantic_cache.py | 2 + tests/test_litellm_rust/test_cache.py | 12 ++--- 6 files changed, 79 insertions(+), 21 deletions(-) diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs index d761364f1ad..8fa27d68d05 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs @@ -131,7 +131,7 @@ impl QdrantSemanticCache { vector, payload, )], - )) + ).wait(true)) .await .map_err(|_| Error::Unavailable)?; Ok(()) diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs index e8d8d5040f0..70cde4e0d5d 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -12,15 +12,50 @@ use litellm_cache_qdrant_semantic::{ use litellm_cache_response::{ CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, }; -use qdrant_client::Payload; use qdrant_client::{ Qdrant, - qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams}, + qdrant::{ + self, CompressionRatio, Distance, PointId, QuantizationType, Struct, Value, VectorParams, + value::Kind, + }, }; use serde_json::{Value as JsonValue, json}; use support::{FakeQdrant, FakeState, StoredPoint}; +fn json_to_qdrant(value: JsonValue) -> Value { + let kind = match value { + JsonValue::Null => Kind::NullValue(0), + JsonValue::Bool(value) => Kind::BoolValue(value), + JsonValue::Number(value) => value + .as_i64() + .map(Kind::IntegerValue) + .or_else(|| value.as_f64().map(Kind::DoubleValue)) + .unwrap(), + JsonValue::String(value) => Kind::StringValue(value), + JsonValue::Array(values) => Kind::ListValue(qdrant::ListValue { + values: values.into_iter().map(json_to_qdrant).collect(), + }), + JsonValue::Object(values) => Kind::StructValue(Struct { + fields: values + .into_iter() + .map(|(key, value)| (key, json_to_qdrant(value))) + .collect(), + }), + }; + Value { kind: Some(kind) } +} + +fn payload_from_json(value: JsonValue) -> HashMap { + value + .as_object() + .unwrap() + .clone() + .into_iter() + .map(|(key, value)| (key, json_to_qdrant(value))) + .collect() +} + #[derive(Clone)] struct FixedEmbedder { vectors: Arc>>, @@ -243,12 +278,10 @@ async fn misses_and_payload_validation_are_safe() { server.insert_point(StoredPoint { id: Some(PointId::from(99_u64)), vector: vec![1.0, 0.0], - payload: Payload::try_from(json!({ + payload: payload_from_json(json!({ "litellm_cache_key": 99, "response": "{}", - })) - .unwrap() - .into(), + })), }); assert_eq!( cache @@ -322,6 +355,10 @@ async fn decoding_errors_missing_prompt_pipeline_and_ttl_behave_as_required() { .unwrap() .is_some() ); + 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, @@ -347,9 +384,7 @@ async fn response_payloads_decode_and_invalid_entries_fail() { 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(), + payload: payload.into_iter().map(|(key, value)| (key, json_to_qdrant(value))).collect(), }); } assert_eq!( 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 860213a1703..2a4cd45e007 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs @@ -10,8 +10,10 @@ use qdrant_client::qdrant::{ CollectionOperationResponse, CreateCollection, CreateFieldIndexCollection, Filter, PointId, PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors, collections_server::Collections, + value::Kind, points_server::{Points, PointsServer}, }; +use serde_json::Value as JsonValue; use tokio::sync::oneshot; use tokio_stream::wrappers::TcpListenerStream; use tonic::{Request, Response, Status, transport::Server}; @@ -23,12 +25,33 @@ pub struct StoredPoint { pub payload: HashMap, } +fn qdrant_value_to_json(value: Value) -> JsonValue { + match value.kind { + Some(Kind::NullValue(_)) | None => JsonValue::Null, + Some(Kind::DoubleValue(value)) => serde_json::json!(value), + Some(Kind::IntegerValue(value)) => serde_json::json!(value), + Some(Kind::StringValue(value)) => JsonValue::String(value), + Some(Kind::BoolValue(value)) => JsonValue::Bool(value), + Some(Kind::StructValue(value)) => JsonValue::Object( + value + .fields + .into_iter() + .map(|(key, value)| (key, qdrant_value_to_json(value))) + .collect(), + ), + Some(Kind::ListValue(value)) => { + JsonValue::Array(value.values.into_iter().map(qdrant_value_to_json).collect()) + } + } +} + #[derive(Default)] pub struct FakeState { pub collections: HashSet, pub created_collections: Vec, pub field_indexes: Vec, pub points: Vec, + pub upsert_waits: Vec>, pub index_creations: usize, pub fail_field_index: bool, } @@ -205,8 +228,10 @@ impl Points for FakeService { &self, request: Request, ) -> Result, Status> { + let request = request.into_inner(); let mut state = self.state.lock().unwrap(); - for point in request.into_inner().points { + state.upsert_waits.push(request.wait); + for point in request.points { let stored = StoredPoint { id: point.id.clone(), vector: dense_vector(point.vectors)?, @@ -241,7 +266,7 @@ impl Points for FakeService { .payload .get(field) .and_then(|value| { - let value: serde_json::Value = value.clone().into(); + let value = qdrant_value_to_json(value.clone()); value .as_str() .map(str::to_owned) diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 058cc8a1579..868f36f7f21 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -313,6 +313,7 @@ class QdrantSemanticCache(BaseCache): self.sync_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points", headers=self.headers, + params={"wait": "true"}, json=data, ) @@ -422,6 +423,7 @@ class QdrantSemanticCache(BaseCache): await self.async_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points", headers=self.headers, + params={"wait": "true"}, json=data, ) diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/test_litellm/caching/test_qdrant_semantic_cache.py index a0a9b71787c..ca7303e4c6d 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/test_litellm/caching/test_qdrant_semantic_cache.py @@ -578,6 +578,7 @@ def test_qdrant_semantic_cache_set_cache(): assert ( upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key" ) + assert qdrant_cache.sync_client.put.call_args.kwargs["params"] == {"wait": "true"} @pytest.mark.asyncio @@ -650,6 +651,7 @@ async def test_qdrant_semantic_cache_async_set_cache(): assert ( upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key" ) + assert qdrant_cache.async_client.put.call_args.kwargs["params"] == {"wait": "true"} def test_qdrant_semantic_cache_custom_vector_size(): diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 02c159825bc..e43af903462 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -605,15 +605,7 @@ async def test_qdrant_semantic_async_parity( messages=messages, ) - async def lookup_after_commit() -> object: - for _ in range(20): - value: Final = await binding.async_lookup(semantic_request("python-key", messages)) - if value is not None: - return value - await asyncio.sleep(0.1) - return None - - assert await lookup_after_commit() == {"id": "py"} + assert await binding.async_lookup(semantic_request("python-key", messages)) == {"id": "py"} await binding.async_store(semantic_request("native-key", messages), {"id": "native"}) python_value: Final = await facade.cache.async_get_cache("native-key", messages=messages) assert isinstance(python_value, dict) @@ -705,6 +697,8 @@ def test_qdrant_semantic_mutation_and_projection_fallback( vector_size=8, ) handle._bind_facade(facade) + facade.cache.qdrant_api_key = "rotated" + assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" facade.cache.similarity_threshold = 0.5 assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" unsupported: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") From 632c95f87eb5ce4715a4f8b58dea86b30011bbdc Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:34:47 +0000 Subject: [PATCH 08/17] fix(python-bridge): make the Qdrant config tests tolerate an unset OPENAI_API_KEY Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/python-bridge/src/cache/config.rs | 49 +++++++------------ 1 file changed, 18 insertions(+), 31 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index b21e671e431..d7041b9c7da 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -854,11 +854,23 @@ sys.modules['litellm.caching._embedding_router'] = embedding_router let prior = environ.call_method1("get", ("OPENAI_API_KEY",))?; match key { Some(key) => environ.set_item("OPENAI_API_KEY", key)?, - None => environ.del_item("OPENAI_API_KEY")?, + None => { + environ.call_method1("pop", ("OPENAI_API_KEY", py.None()))?; + } } Ok(prior) } + fn restore_embedding_environment(py: Python<'_>, prior: Bound<'_, PyAny>) -> PyResult<()> { + let environ = py.import("os")?.getattr("environ")?; + if prior.is_none() { + environ.call_method1("pop", ("OPENAI_API_KEY", py.None()))?; + } else { + environ.set_item("OPENAI_API_KEY", prior)?; + } + Ok(()) + } + fn facade<'py>(py: Python<'py>, body: &str) -> Bound<'py, PyAny> { let locals = PyDict::new(py); py.run( @@ -1080,12 +1092,7 @@ sys.modules['litellm.caching._embedding_router'] = embedding_router assert_eq!(config.vector_size, 8); assert_eq!(config.embedding.api_key, "embedding-key"); assert_eq!(config.embedding.model, "text-embedding-3-small"); - let environ = py.import("os").unwrap().getattr("environ").unwrap(); - if prior.is_none() { - environ.del_item("OPENAI_API_KEY").unwrap(); - } else { - environ.set_item("OPENAI_API_KEY", prior).unwrap(); - } + restore_embedding_environment(py, prior).unwrap(); }); } @@ -1105,12 +1112,7 @@ sys.modules['litellm.caching._embedding_router'] = embedding_router panic!("non-default Qdrant port should stay on Python"); }; assert!(matches!(reason, UnsupportedCacheConfig::QdrantEndpoint)); - let environ = py.import("os").unwrap().getattr("environ").unwrap(); - if prior.is_none() { - environ.del_item("OPENAI_API_KEY").unwrap(); - } else { - environ.set_item("OPENAI_API_KEY", prior).unwrap(); - } + restore_embedding_environment(py, prior).unwrap(); }); } @@ -1133,12 +1135,7 @@ sys.modules['litellm.caching._embedding_router'] = embedding_router }; assert!(matches!(reason, UnsupportedCacheConfig::SemanticEmbedding)); } - let environ = py.import("os").unwrap().getattr("environ").unwrap(); - if prior.is_none() { - environ.del_item("OPENAI_API_KEY").unwrap(); - } else { - environ.set_item("OPENAI_API_KEY", prior).unwrap(); - } + restore_embedding_environment(py, prior).unwrap(); }); } @@ -1167,12 +1164,7 @@ sys.modules['litellm.caching._embedding_router'] = embedding_router panic!("missing embedding key should stay on Python"); }; assert!(matches!(reason, UnsupportedCacheConfig::SemanticEmbedding)); - let environ = py.import("os").unwrap().getattr("environ").unwrap(); - if prior.is_none() { - environ.del_item("OPENAI_API_KEY").unwrap(); - } else { - environ.set_item("OPENAI_API_KEY", prior).unwrap(); - } + restore_embedding_environment(py, prior).unwrap(); }); } @@ -1193,12 +1185,7 @@ sys.modules['litellm.caching._embedding_router'] = embedding_router config.service_mismatch(&service), Some("facade and native backend types must match") ); - let environ = py.import("os").unwrap().getattr("environ").unwrap(); - if prior.is_none() { - environ.del_item("OPENAI_API_KEY").unwrap(); - } else { - environ.set_item("OPENAI_API_KEY", prior).unwrap(); - } + restore_embedding_environment(py, prior).unwrap(); }); } } From acbec828db40a3365d5657b59c0cdc4fef159b8d Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:34:49 +0000 Subject: [PATCH 09/17] fix(python-bridge): fall back to Python when qdrant_api_key changes after binding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/src/cache/facade.rs | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index dbcbd9bda92..6aa075600b8 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -290,8 +290,9 @@ impl FacadeGuard { "redis_flush_size", ][..], "qdrant_semantic" => &[ - "qdrant_api_base", - "collection_name", + "qdrant_api_base", + "qdrant_api_key", + "collection_name", "similarity_threshold", "embedding_model", "vector_size", From ff2ca804b5633ee1d47315d7e6e3fd5f8dbdea6c Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:36:19 +0000 Subject: [PATCH 10/17] fix(python-bridge): preserve Qdrant facade dispatch after rebase Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/src/cache/facade.rs | 2 +- litellm-rust/crates/python-bridge/src/cache/native.rs | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index 6aa075600b8..f307d109fb1 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -256,7 +256,7 @@ impl FacadeGuard { "RedisClusterCache", "redis", ), - "qdrant_semantic" => ( + ("qdrant_semantic", _) => ( "litellm.caching.qdrant_semantic_cache", "QdrantSemanticCache", "qdrant-semantic", diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 6e28b8414d3..b93f15e59f3 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -104,6 +104,7 @@ impl NativeResponseCache { match self { Self::Memory(_) => None, Self::Redis { cache, .. } => Some(cache.backend().topology()), + Self::QdrantSemantic(_) => None, } } From 6a06f69972e962beed5066619af1d8adabfcc666 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:37:11 +0000 Subject: [PATCH 11/17] fix(python-bridge): restore Qdrant dispatch after rebase Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../cache-qdrant-semantic/src/semantic.rs | 19 +++++++++++-------- .../cache-qdrant-semantic/tests/qdrant.rs | 5 ++++- .../tests/support/mod.rs | 2 +- .../crates/python-bridge/src/cache/facade.rs | 6 +++--- 4 files changed, 19 insertions(+), 13 deletions(-) diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs index 8fa27d68d05..cc82118d280 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs @@ -124,14 +124,17 @@ impl QdrantSemanticCache { })) .map_err(|_| Error::InvalidEntry)?; self.client - .upsert_points(UpsertPointsBuilder::new( - self.collection_name(), - vec![PointStruct::new( - Uuid::new_v4().to_string(), - vector, - payload, - )], - ).wait(true)) + .upsert_points( + UpsertPointsBuilder::new( + self.collection_name(), + vec![PointStruct::new( + Uuid::new_v4().to_string(), + vector, + payload, + )], + ) + .wait(true), + ) .await .map_err(|_| Error::Unavailable)?; Ok(()) diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs index 70cde4e0d5d..70ba666e782 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -384,7 +384,10 @@ async fn response_payloads_decode_and_invalid_entries_fail() { server.insert_point(StoredPoint { id: Some(PointId::from(key.len() as u64)), vector: vec![1.0, 0.0], - payload: payload.into_iter().map(|(key, value)| (key, json_to_qdrant(value))).collect(), + payload: payload + .into_iter() + .map(|(key, value)| (key, json_to_qdrant(value))) + .collect(), }); } assert_eq!( 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 2a4cd45e007..5311a2473d5 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs @@ -10,8 +10,8 @@ use qdrant_client::qdrant::{ CollectionOperationResponse, CreateCollection, CreateFieldIndexCollection, Filter, PointId, PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors, collections_server::Collections, - value::Kind, points_server::{Points, PointsServer}, + value::Kind, }; use serde_json::Value as JsonValue; use tokio::sync::oneshot; diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index f307d109fb1..db37491653c 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -290,9 +290,9 @@ impl FacadeGuard { "redis_flush_size", ][..], "qdrant_semantic" => &[ - "qdrant_api_base", - "qdrant_api_key", - "collection_name", + "qdrant_api_base", + "qdrant_api_key", + "collection_name", "similarity_threshold", "embedding_model", "vector_size", From b71e64446c3dda4efbcf0d48bed4ef952c8654ae Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:44:34 +0000 Subject: [PATCH 12/17] fix(cache-qdrant-semantic): annotate indexing wait payloads Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/qdrant_semantic_cache.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 868f36f7f21..64764ce402c 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -313,7 +313,7 @@ class QdrantSemanticCache(BaseCache): self.sync_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points", headers=self.headers, - params={"wait": "true"}, + params={"wait": "true"}, # mutable-ok: Qdrant requires an explicit indexing wait json=data, ) @@ -423,7 +423,7 @@ class QdrantSemanticCache(BaseCache): await self.async_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points", headers=self.headers, - params={"wait": "true"}, + params={"wait": "true"}, # mutable-ok: Qdrant requires an explicit indexing wait json=data, ) From a4fda8f0d80715a952b96f90e14a5c3ab1b36aeb Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:48:24 +0000 Subject: [PATCH 13/17] refactor(cache-qdrant-semantic): reuse immutable indexing params Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/caching/qdrant_semantic_cache.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 64764ce402c..b99023c07fd 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -12,6 +12,7 @@ import ast import asyncio import json import os +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Protocol, cast import litellm @@ -36,6 +37,8 @@ from ._embedding_router import ( ) from .base_cache import BaseCache +_WAIT_FOR_INDEXING: Final = MappingProxyType({"wait": "true"}) + if TYPE_CHECKING: from litellm.router import Router @@ -313,7 +316,7 @@ class QdrantSemanticCache(BaseCache): self.sync_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points", headers=self.headers, - params={"wait": "true"}, # mutable-ok: Qdrant requires an explicit indexing wait + params=_WAIT_FOR_INDEXING, json=data, ) @@ -423,7 +426,7 @@ class QdrantSemanticCache(BaseCache): await self.async_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points", headers=self.headers, - params={"wait": "true"}, # mutable-ok: Qdrant requires an explicit indexing wait + params=_WAIT_FOR_INDEXING, json=data, ) From b0c1b863820121b10c2b6bca99081d038472cce2 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 21:48:27 +0000 Subject: [PATCH 14/17] refactor(cache-qdrant-semantic): reuse qdrant-client serde payload conversion in tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../cache-qdrant-semantic/tests/qdrant.rs | 52 ++++--------------- .../tests/support/mod.rs | 24 +-------- 2 files changed, 10 insertions(+), 66 deletions(-) diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs index 70ba666e782..d806a9cb455 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -12,50 +12,15 @@ use litellm_cache_qdrant_semantic::{ use litellm_cache_response::{ CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, }; +use qdrant_client::Payload; use qdrant_client::{ Qdrant, - qdrant::{ - self, CompressionRatio, Distance, PointId, QuantizationType, Struct, Value, VectorParams, - value::Kind, - }, + qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams}, }; use serde_json::{Value as JsonValue, json}; use support::{FakeQdrant, FakeState, StoredPoint}; -fn json_to_qdrant(value: JsonValue) -> Value { - let kind = match value { - JsonValue::Null => Kind::NullValue(0), - JsonValue::Bool(value) => Kind::BoolValue(value), - JsonValue::Number(value) => value - .as_i64() - .map(Kind::IntegerValue) - .or_else(|| value.as_f64().map(Kind::DoubleValue)) - .unwrap(), - JsonValue::String(value) => Kind::StringValue(value), - JsonValue::Array(values) => Kind::ListValue(qdrant::ListValue { - values: values.into_iter().map(json_to_qdrant).collect(), - }), - JsonValue::Object(values) => Kind::StructValue(Struct { - fields: values - .into_iter() - .map(|(key, value)| (key, json_to_qdrant(value))) - .collect(), - }), - }; - Value { kind: Some(kind) } -} - -fn payload_from_json(value: JsonValue) -> HashMap { - value - .as_object() - .unwrap() - .clone() - .into_iter() - .map(|(key, value)| (key, json_to_qdrant(value))) - .collect() -} - #[derive(Clone)] struct FixedEmbedder { vectors: Arc>>, @@ -278,10 +243,12 @@ async fn misses_and_payload_validation_are_safe() { server.insert_point(StoredPoint { id: Some(PointId::from(99_u64)), vector: vec![1.0, 0.0], - payload: payload_from_json(json!({ + payload: Payload::try_from(json!({ "litellm_cache_key": 99, "response": "{}", - })), + })) + .unwrap() + .into(), }); assert_eq!( cache @@ -384,10 +351,9 @@ async fn response_payloads_decode_and_invalid_entries_fail() { server.insert_point(StoredPoint { id: Some(PointId::from(key.len() as u64)), vector: vec![1.0, 0.0], - payload: payload - .into_iter() - .map(|(key, value)| (key, json_to_qdrant(value))) - .collect(), + payload: Payload::try_from(JsonValue::Object(payload)) + .unwrap() + .into(), }); } assert_eq!( 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 5311a2473d5..9a556ae7df5 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs @@ -11,9 +11,7 @@ use qdrant_client::qdrant::{ PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors, collections_server::Collections, points_server::{Points, PointsServer}, - value::Kind, }; -use serde_json::Value as JsonValue; use tokio::sync::oneshot; use tokio_stream::wrappers::TcpListenerStream; use tonic::{Request, Response, Status, transport::Server}; @@ -25,26 +23,6 @@ pub struct StoredPoint { pub payload: HashMap, } -fn qdrant_value_to_json(value: Value) -> JsonValue { - match value.kind { - Some(Kind::NullValue(_)) | None => JsonValue::Null, - Some(Kind::DoubleValue(value)) => serde_json::json!(value), - Some(Kind::IntegerValue(value)) => serde_json::json!(value), - Some(Kind::StringValue(value)) => JsonValue::String(value), - Some(Kind::BoolValue(value)) => JsonValue::Bool(value), - Some(Kind::StructValue(value)) => JsonValue::Object( - value - .fields - .into_iter() - .map(|(key, value)| (key, qdrant_value_to_json(value))) - .collect(), - ), - Some(Kind::ListValue(value)) => { - JsonValue::Array(value.values.into_iter().map(qdrant_value_to_json).collect()) - } - } -} - #[derive(Default)] pub struct FakeState { pub collections: HashSet, @@ -266,7 +244,7 @@ impl Points for FakeService { .payload .get(field) .and_then(|value| { - let value = qdrant_value_to_json(value.clone()); + let value: serde_json::Value = value.clone().into(); value .as_str() .map(str::to_owned) From ae32609b540d89409abaa20af716620b186d5555 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:18:13 +0000 Subject: [PATCH 15/17] refactor(cache-qdrant-semantic): inject the shared LiteLLM HTTP client into the embedder Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + .../cache-qdrant-semantic/src/embedder.rs | 32 ++++---- .../cache-qdrant-semantic/tests/embedder.rs | 73 +++++++++++++------ litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/cache/config.rs | 19 ++--- .../crates/python-bridge/src/cache/handle.rs | 10 ++- .../crates/python-bridge/src/cache/native.rs | 7 +- 7 files changed, 90 insertions(+), 53 deletions(-) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index f4b94ae44e2..382009c95af 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2821,6 +2821,7 @@ dependencies = [ "pyo3", "pyo3-async-runtimes", "qdrant-client", + "reqwest 0.12.28", "rstest", "serde", "serde_json", diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs index 0dde0318448..47b898d6f4e 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs @@ -11,6 +11,7 @@ pub struct OpenAiEmbedder { api_base: String, api_key: String, model: String, + timeout: Option, } pub struct OpenAiEmbedderConfig { @@ -21,18 +22,14 @@ pub struct OpenAiEmbedderConfig { } impl OpenAiEmbedder { - pub fn new(config: OpenAiEmbedderConfig) -> Result { - let mut builder = Client::builder(); - if let Some(timeout) = config.timeout { - builder = builder.timeout(timeout); - } - let client = builder.build().map_err(|_| Error::Unavailable)?; - Ok(Self { + pub fn new(client: Client, config: OpenAiEmbedderConfig) -> Self { + Self { client, api_base: config.api_base.trim_end_matches('/').to_owned(), api_key: config.api_key, model: config.model, - }) + timeout: config.timeout, + } } } @@ -42,7 +39,7 @@ impl Embedder for OpenAiEmbedder { } async fn embed(&self, input: &str) -> Result, Error> { - let response = self + let request = self .client .post(format!("{}/embeddings", self.api_base)) .bearer_auth(&self.api_key) @@ -50,12 +47,17 @@ impl Embedder for OpenAiEmbedder { "model": self.model, "input": input, "encoding_format": "float", - })) - .send() - .await - .map_err(|_| Error::Unavailable)? - .error_for_status() - .map_err(|_| Error::Unavailable)?; + })); + let response = if let Some(timeout) = self.timeout { + request.timeout(timeout) + } else { + request + } + .send() + .await + .map_err(|_| Error::Unavailable)? + .error_for_status() + .map_err(|_| Error::Unavailable)?; let body: Value = response.json().await.map_err(|_| Error::Unavailable)?; body.get("data") .and_then(Value::as_array) diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs index adce70654a8..6b09448fde8 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs @@ -19,6 +19,10 @@ struct TestHttpServer { impl TestHttpServer { async fn response(status: &str, body: &str) -> Self { + Self::response_after(status, body, Duration::ZERO).await + } + + async fn response_after(status: &str, body: &str, delay: Duration) -> Self { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let address = listener.local_addr().unwrap(); let request = Arc::new(Mutex::new(None)); @@ -29,6 +33,7 @@ impl TestHttpServer { let (mut stream, _) = listener.accept().await.unwrap(); let request_bytes = read_request(&mut stream).await; *captured.lock().unwrap() = Some(request_bytes); + tokio::time::sleep(delay).await; let response = format!( "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", body.len() @@ -42,20 +47,6 @@ impl TestHttpServer { } } - async fn hanging() -> Self { - let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); - let address = listener.local_addr().unwrap(); - let task = tokio::spawn(async move { - let (_stream, _) = listener.accept().await.unwrap(); - std::future::pending::<()>().await; - }); - Self { - address, - request: Arc::new(Mutex::new(None)), - task, - } - } - fn base_url(&self) -> String { format!("http://{}", self.address) } @@ -110,11 +101,13 @@ fn config(base: String, timeout: Option) -> OpenAiEmbedderConfig { #[tokio::test] async fn posts_embeddings_request_and_parses_vector() { let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await; - let embedder = OpenAiEmbedder::new(config( - format!("{}/", server.base_url()), - Some(Duration::from_secs(1)), - )) - .unwrap(); + let embedder = OpenAiEmbedder::new( + reqwest::Client::new(), + config( + format!("{}/", server.base_url()), + Some(Duration::from_secs(1)), + ), + ); assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]); let request = server.request.lock().unwrap().clone().unwrap(); let request_text = String::from_utf8(request).unwrap(); @@ -130,12 +123,44 @@ async fn posts_embeddings_request_and_parses_vector() { #[tokio::test] async fn status_and_timeout_errors_are_unavailable() { let server = TestHttpServer::response("500 Internal Server Error", "{}").await; - let embedder = - OpenAiEmbedder::new(config(server.base_url(), Some(Duration::from_secs(1)))).unwrap(); + let embedder = OpenAiEmbedder::new(reqwest::Client::new(), config(server.base_url(), None)); assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable)); - let server = TestHttpServer::hanging().await; - let embedder = - OpenAiEmbedder::new(config(server.base_url(), Some(Duration::from_millis(200)))).unwrap(); + 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]); +} + +#[tokio::test] +async fn uses_the_injected_client() { + let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await; + let client = reqwest::Client::builder() + .user_agent("litellm-embedder-test") + .build() + .unwrap(); + let embedder = OpenAiEmbedder::new(client, config(server.base_url(), None)); + assert_eq!(embedder.embed("hello").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/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index 826092522ea..0b9a0148d07 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -40,6 +40,7 @@ litellm-host-python.workspace = true litellm-token-counter = { path = "../token-counter", default-features = false } pyo3.workspace = true pyo3-async-runtimes.workspace = true +reqwest.workspace = true serde_json.workspace = true url.workspace = true tokio = { workspace = true, features = ["sync"] } diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 7772716f128..57b0561b122 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -216,18 +216,15 @@ impl NativeCacheConfig { } pub(super) fn service_mismatch(&self, service: &NativeResponseCache) -> Option<&'static str> { - let default_ttl = match &self.backend { - CacheBackendConfig::Memory(config) => Some(config.default_ttl), - CacheBackendConfig::Redis(config) => Some(config.default_ttl), - CacheBackendConfig::AzureBlob(_) | CacheBackendConfig::QdrantSemantic(_) => None, - }; - if service.default_ttl() != default_ttl { - return Some("facade and native backend default TTLs must match"); - } match &self.backend { CacheBackendConfig::Memory(_) if service.kind() != "memory" => { Some("facade and native backend types must match") } + CacheBackendConfig::Memory(config) + if service.default_ttl() != Some(config.default_ttl) => + { + Some("facade and native backend default TTLs must match") + } CacheBackendConfig::Memory(config) if service.capacity() != Some(config.capacity) => { Some("facade and native backend capacities must match") } @@ -245,7 +242,11 @@ impl NativeCacheConfig { } CacheBackendConfig::Redis(config) => (service.namespace() != config.namespace.as_deref()) - .then_some("facade and native backend namespaces must match"), + .then_some("facade and native backend namespaces must match") + .or_else(|| { + (service.default_ttl() != Some(config.default_ttl)) + .then_some("facade and native backend default TTLs must match") + }), CacheBackendConfig::QdrantSemantic(config) if service.kind() != "qdrant_semantic" => { Some("facade and native backend types must match") } diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index 94a2cfec2bd..e71ad9f435d 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,16 +1,18 @@ use std::env; use litellm_cache_redis::{RedisNode, RedisTopology}; +use litellm_http::ClientVariant; use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, Quantization}; use litellm_host_python::{release_gil, run_sync_value}; -use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*}; +use pyo3::{PyTraverseError, PyVisit, exceptions::PyRuntimeError, prelude::*, types::PyDict}; use url::Url; use super::{ cache_error, config::QdrantSemanticCacheConfig, facade::FacadeGuard, native::NativeResponseCache, request::duration, }; +use crate::http; #[pyclass(frozen, name = "_CacheTestHandle")] pub(crate) struct CacheTestHandle { @@ -154,9 +156,13 @@ impl CacheTestHandle { }, quantization, }; + let http_config = http::call_config(py, &PyDict::new(py), true)?; + let client = http::pool() + .client(&http_config, ClientVariant::Provider) + .map_err(http::client_error)?; let service = run_sync_value(py, async move { let handle = tokio::runtime::Handle::current(); - NativeResponseCache::qdrant_semantic(config, handle) + NativeResponseCache::qdrant_semantic(config, client, handle) .await .map_err(cache_error) })?; diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index 0e14b9cb39c..e4810a303c8 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -54,17 +54,18 @@ impl NativeResponseCache { pub async fn qdrant_semantic( config: QdrantSemanticCacheConfig, + client: reqwest::Client, runtime: tokio::runtime::Handle, ) -> Result { - let client = qdrant_client::Qdrant::from_url(&config.grpc_url) + let qdrant = qdrant_client::Qdrant::from_url(&config.grpc_url) .skip_compatibility_check() .api_key(config.api_key.as_deref()) .build() .map_err(|_| Error::Unavailable)?; let qdrant_config = config.to_qdrant_config(); - let embedder = OpenAiEmbedder::new(config.embedding)?; + let embedder = OpenAiEmbedder::new(client, config.embedding); let cache = QdrantSemanticCache::connect( - client, + qdrant, embedder, ResponseCacheCodec, qdrant_config, From 1dfd54579bdbdc4a4a5ebc7534279f27e27ae8a3 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:18:59 +0000 Subject: [PATCH 16/17] fix(python-bridge): keep the Azure Blob default TTL mismatch check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/crates/python-bridge/src/cache/config.rs | 3 +++ 1 file changed, 3 insertions(+) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 57b0561b122..27ecf0f8463 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -278,6 +278,9 @@ impl NativeCacheConfig { { Some("facade and native backend containers must match") } + Some(_) if service.default_ttl().is_some() => { + Some("facade and native backend default TTLs must match") + } Some(_) => None, }, } From ed8d4441a5591d1fe5a25643f9c68479ee2ade06 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 22:44:47 +0000 Subject: [PATCH 17/17] fix(python-bridge): harden Qdrant facade projection guards Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../crates/python-bridge/src/cache/config.rs | 70 ++++++++++++++++--- .../crates/python-bridge/src/cache/facade.rs | 2 +- 2 files changed, 63 insertions(+), 9 deletions(-) diff --git a/litellm-rust/crates/python-bridge/src/cache/config.rs b/litellm-rust/crates/python-bridge/src/cache/config.rs index 27ecf0f8463..2aba7012b54 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -300,7 +300,7 @@ fn project_qdrant_semantic( || !parsed.path().is_empty() && parsed.path() != "/" || parsed.query().is_some() || parsed.host_str().is_none() - || parsed.port().is_some_and(|port| port != 6333) + || parsed.port() != Some(6333) { return Ok(Err(UnsupportedCacheConfig::QdrantEndpoint)); } @@ -330,7 +330,7 @@ fn project_qdrant_semantic( let embedding_router = backend.py().import("litellm.caching._embedding_router")?; if !embedding_router .getattr("resolve_embedding_router")? - .call1((embedding_model.as_str(), router, model_list))? + .call1((configured_model.as_str(), router, model_list))? .is_none() { return Ok(Err(UnsupportedCacheConfig::SemanticEmbedding)); @@ -1140,16 +1140,70 @@ sys.modules['litellm.caching._embedding_router'] = embedding_router Python::initialize(); Python::attach(|py| { let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap(); - let facade = qdrant_facade( - py, - "backend.qdrant_api_base = 'https://qdrant.example:6332'", - ); + for endpoint in [ + "https://qdrant.example:6332", + "https://qdrant.example", + "http://qdrant.example", + ] { + let facade = qdrant_facade(py, &format!("backend.qdrant_api_base = '{endpoint}'")); + let CacheConfigProjection::Unsupported(reason) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("unsupported Qdrant endpoint should stay on Python"); + }; + assert!(matches!(reason, UnsupportedCacheConfig::QdrantEndpoint)); + } + let facade = qdrant_facade(py, ""); + let CacheConfigProjection::Native(config) = + NativeCacheConfig::project(&facade).unwrap() + else { + panic!("default Qdrant endpoint should use native"); + }; + let CacheBackendConfig::QdrantSemantic(config) = config.backend else { + panic!("expected Qdrant configuration"); + }; + assert!(config.grpc_url.ends_with(":6334")); + restore_embedding_environment(py, prior).unwrap(); + }); + } + + #[test] + fn qdrant_projection_passes_configured_embedding_model_to_router() { + let _guard = env_lock().lock().unwrap_or_else(|error| error.into_inner()); + Python::initialize(); + Python::attach(|py| { + let prior = configure_embedding_environment(py, Some("embedding-key")).unwrap(); + let facade = qdrant_facade(py, ""); + py.run( + c" +import sys +import types +proxy_server = types.ModuleType('litellm.proxy.proxy_server') +proxy_server.llm_router = None +proxy_server.llm_model_list = None +sys.modules['litellm.proxy.proxy_server'] = proxy_server +embedding_router = sys.modules['litellm.caching._embedding_router'] +embedding_router.resolve_embedding_router = lambda model, *_args: object() if model == 'openai/text-embedding-3-small' else None +", + None, + None, + ) + .unwrap(); let CacheConfigProjection::Unsupported(reason) = NativeCacheConfig::project(&facade).unwrap() else { - panic!("non-default Qdrant port should stay on Python"); + panic!("router-backed embedding should stay on Python"); }; - assert!(matches!(reason, UnsupportedCacheConfig::QdrantEndpoint)); + assert!(matches!(reason, UnsupportedCacheConfig::SemanticEmbedding)); + py.run( + c" +import sys +sys.modules.pop('litellm.proxy.proxy_server', None) +", + None, + None, + ) + .unwrap(); restore_embedding_environment(py, prior).unwrap(); }); } diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index f24c1b2e86d..c8aafe41816 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -354,7 +354,7 @@ impl FacadeGuard { return Err(PyTypeError::new_err(message)); } let backend_config_names = match kind { - "memory" | "redis" => &[ + "memory" | "redis" | "azure-blob" => &[ "namespace", "default_ttl", "max_size_in_memory",