diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 66270501312..464edb3b104 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -650,6 +650,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" @@ -1439,7 +1482,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" dependencies = [ "libc", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -2241,7 +2284,7 @@ dependencies = [ "libc", "percent-encoding", "pin-project-lite", - "socket2 0.5.10", + "socket2 0.6.5", "tokio", "tower-service", "tracing", @@ -2729,6 +2772,26 @@ 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", + "tokio-stream", + "tonic", + "tonic-prost", + "uuid", +] + [[package]] name = "litellm-cache-redis" version = "0.1.0" @@ -2972,6 +3035,7 @@ dependencies = [ "litellm-cache-disk", "litellm-cache-gcs", "litellm-cache-memory", + "litellm-cache-qdrant-semantic", "litellm-cache-redis", "litellm-cache-redis-semantic", "litellm-cache-response", @@ -2987,7 +3051,9 @@ dependencies = [ "litellm-types", "pyo3", "pyo3-async-runtimes", + "qdrant-client", "redis", + "reqwest 0.12.28", "rstest", "serde", "serde_json", @@ -2995,6 +3061,7 @@ dependencies = [ "sha2 0.10.9", "tokio", "tokio-tungstenite", + "url", ] [[package]] @@ -3251,6 +3318,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 = "md-5" version = "0.11.0" @@ -3874,6 +3947,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" @@ -3903,7 +3997,7 @@ dependencies = [ "quinn-udp", "rustc-hash", "rustls 0.23.42", - "socket2 0.5.10", + "socket2 0.6.5", "thiserror 2.0.19", "tokio", "tracing", @@ -3942,9 +4036,9 @@ dependencies = [ "cfg_aliases", "libc", "once_cell", - "socket2 0.5.10", + "socket2 0.6.5", "tracing", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4430,7 +4524,7 @@ dependencies = [ "errno", "libc", "linux-raw-sys", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -4501,7 +4595,7 @@ dependencies = [ "security-framework", "security-framework-sys", "webpki-root-certs", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5073,10 +5167,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" dependencies = [ "fastrand", - "getrandom 0.3.4", + "getrandom 0.4.3", "once_cell", "rustix", - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] @@ -5365,8 +5459,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", @@ -5376,6 +5474,7 @@ dependencies = [ "percent-encoding", "pin-project", "rustls-native-certs", + "socket2 0.6.5", "sync_wrapper", "tokio", "tokio-rustls 0.26.4", @@ -5949,7 +6048,7 @@ version = "0.1.11" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" dependencies = [ - "windows-sys 0.52.0", + "windows-sys 0.61.2", ] [[package]] diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 4e370c13766..0e8941cb8e6 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -38,6 +38,7 @@ litellm-cache-gcs = { path = "crates/cache-gcs" } litellm-cache-disk = { path = "crates/cache-disk" } litellm-cache-redis-semantic = { path = "crates/cache-redis-semantic" } litellm-cache-response = { path = "crates/cache-response" } +litellm-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" } @@ -55,6 +56,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..09d6a9637f3 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml @@ -0,0 +1,24 @@ +[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" +tokio-stream = "0.1" 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..47b898d6f4e --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/src/embedder.rs @@ -0,0 +1,75 @@ +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, + timeout: Option, +} + +pub struct OpenAiEmbedderConfig { + pub api_base: String, + pub api_key: String, + pub model: String, + pub timeout: Option, +} + +impl OpenAiEmbedder { + 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, + } + } +} + +impl Embedder for OpenAiEmbedder { + fn model(&self) -> &str { + &self.model + } + + async fn embed(&self, input: &str) -> Result, Error> { + let request = 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", + })); + 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) + .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..fb165ed5a8e --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs @@ -0,0 +1,262 @@ +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 { + let Some(messages) = context.messages.as_ref().and_then(Value::as_array) else { + return Err(Error::MissingPrompt); + }; + if messages.is_empty() { + return Err(Error::MissingPrompt); + } + Ok(prompt_from_messages(messages)) + } + + 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, + )], + ) + .wait(true), + ) + .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); + }; + 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) + .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/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs new file mode 100644 index 00000000000..6b09448fde8 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs @@ -0,0 +1,166 @@ +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 { + 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)); + 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); + 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() + ); + stream.write_all(response.as_bytes()).await.unwrap(); + }); + Self { + address, + request, + 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( + 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(); + 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(reqwest::Client::new(), config(server.base_url(), None)); + assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable)); + + let server = TestHttpServer::response_after( + "200 OK", + r#"{"data":[{"embedding":[0.1,0.2]}]}"#, + Duration::from_millis(500), + ) + .await; + let embedder = OpenAiEmbedder::new( + reqwest::Client::new(), + config(server.base_url(), Some(Duration::from_millis(200))), + ); + assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable)); + + let server = TestHttpServer::response_after( + "200 OK", + r#"{"data":[{"embedding":[0.1,0.2]}]}"#, + Duration::from_millis(100), + ) + .await; + let embedder = OpenAiEmbedder::new( + reqwest::Client::new(), + config(server.base_url(), Some(Duration::from_secs(1))), + ); + assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]); +} + +#[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/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"}"# + ); +} 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..c7522c0b313 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -0,0 +1,422 @@ +#[path = "support/mod.rs"] +mod support; + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use litellm_cache::{BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext}; +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: Some(json!([{"role": "user", "content": prompt}])), + ..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")] +#[expect( + deprecated, + reason = "the test verifies Qdrant's legacy always_ram quantization contract" +)] +async fn connect_sets_collection_quantization_and_index() { + for (quantization, expected) in [ + (Quantization::Binary, 0), + (Quantization::Scalar, 1), + (Quantization::Product, 2), + ] { + let server = FakeQdrant::start(FakeState::default()).await; + let client = Qdrant::from_url(&server.url()).build().unwrap(); + QdrantSemanticCache::connect( + client, + FixedEmbedder::new([]), + ResponseCacheCodec, + config(quantization), + tokio::runtime::Handle::current(), + ) + .await + .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!( + server.state.lock().unwrap().upsert_waits, + vec![Some(true), Some(true), Some(true)] + ); + assert_eq!(cache.get_ttl(&context("one")), None); + assert_eq!( + cache.test_connection().await, + Err(Error::UnsupportedOperation) + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn response_payloads_decode_and_invalid_entries_fail() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await; + for (key, response) in [ + ("python", json!("{'timestamp': 1.0, 'response': {'a': 1}}")), + ("garbage", json!("not json")), + ("missing", json!("unused")), + ] { + let mut payload = serde_json::Map::new(); + payload.insert("litellm_cache_key".to_owned(), json!(key)); + if key != "missing" { + payload.insert("response".to_owned(), response); + } + server.insert_point(StoredPoint { + id: Some(PointId::from(key.len() as u64)), + vector: vec![1.0, 0.0], + payload: Payload::try_from(JsonValue::Object(payload)) + .unwrap() + .into(), + }); + } + assert_eq!( + cache + .async_get_cache("python", &context("hello")) + .await + .unwrap(), + Some(value(json!({"a": 1}))) + ); + assert_eq!( + cache.async_get_cache("garbage", &context("hello")).await, + Err(Error::InvalidEntry) + ); + assert_eq!( + cache.async_get_cache("missing", &context("hello")).await, + Err(Error::InvalidEntry) + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn response_cache_facade_turns_invalid_entry_into_miss() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await); + let request = ResponseCacheRequest::::new(CacheKeyInput { + preset: Some("key".to_owned()), + ..Default::default() + }) + .with_context(context("hello")); + let response = json!({"answer": 42}); + let facade = ResponseCache::new(cache.clone()); + facade + .async_store(&request, response.clone(), Duration::from_secs(1)) + .await + .unwrap(); + assert_eq!( + facade + .async_lookup(&request, Duration::from_secs(1)) + .await + .unwrap(), + Some(response) + ); + { + let mut state = server.state.lock().unwrap(); + state.points[0] + .payload + .insert("response".to_owned(), Value::from("not json")); + } + assert_eq!( + facade + .async_lookup(&request, Duration::from_secs(1)) + .await + .unwrap(), + None + ); + 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..9a556ae7df5 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs @@ -0,0 +1,342 @@ +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 upsert_waits: 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 request = request.into_inner(); + let mut state = self.state.lock().unwrap(); + state.upsert_waits.push(request.wait); + for point in request.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) +} diff --git a/litellm-rust/crates/cache-response/src/response.rs b/litellm-rust/crates/cache-response/src/response.rs index 2f949d511de..5088402f125 100644 --- a/litellm-rust/crates/cache-response/src/response.rs +++ b/litellm-rust/crates/cache-response/src/response.rs @@ -32,6 +32,17 @@ impl ResponseCacheRequest { } } +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, diff --git a/litellm-rust/crates/cache-response/tests/response.rs b/litellm-rust/crates/cache-response/tests/response.rs index e4f78dae8b2..dcfc0301148 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: Some(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/cache/src/error.rs b/litellm-rust/crates/cache/src/error.rs index 1a381d0afd8..79ef9cd18b1 100644 --- a/litellm-rust/crates/cache/src/error.rs +++ b/litellm-rust/crates/cache/src/error.rs @@ -8,4 +8,6 @@ pub enum Error { UnscopedFlush, #[error("operation is not supported by this cache")] UnsupportedOperation, + #[error("semantic cache requires request messages")] + MissingPrompt, } diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index bbe727f9345..6b4476d897c 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -29,6 +29,8 @@ litellm-cache-gcs.workspace = true litellm-cache-disk.workspace = true litellm-cache-redis-semantic.workspace = true litellm-cache-response.workspace = true +litellm-cache-qdrant-semantic.workspace = true +qdrant-client.workspace = true litellm-cache-valkey-semantic = { path = "../cache-valkey-semantic" } serde.workspace = true litellm-auth.workspace = true @@ -44,8 +46,10 @@ 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 redis = { version = "1.7.0", features = ["tls-rustls"] } serde_json.workspace = true +url.workspace = true tokio = { workspace = true, features = ["rt", "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 ef2a03ac9ba..d161b401928 100644 --- a/litellm-rust/crates/python-bridge/src/cache/config.rs +++ b/litellm-rust/crates/python-bridge/src/cache/config.rs @@ -2,6 +2,7 @@ use std::{path::PathBuf, time::Duration}; use litellm_auth_aws::AwsAuthConfig; use litellm_cache::CacheType; +use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, QdrantSemanticConfig, Quantization}; use litellm_cache_redis::{RedisNode, RedisTopology}; use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; use pyo3::{ @@ -125,6 +126,27 @@ pub(super) struct ValkeySemanticCacheConfig { pub(super) connection: RedisConnectionConfig, } +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), @@ -134,6 +156,7 @@ pub(super) enum CacheBackendConfig { Disk(DiskCacheConfig), AzureBlob(AzureBlobCacheConfig), RedisSemantic(Box), + QdrantSemantic(Box), } #[allow(dead_code, reason = "consumed by the cache activation follow-up")] @@ -153,6 +176,8 @@ pub(super) enum UnsupportedCacheConfig { S3Option, GcsBucket, DiskStore, + QdrantEndpoint, + SemanticEmbedding, } impl UnsupportedCacheConfig { @@ -168,6 +193,10 @@ impl UnsupportedCacheConfig { Self::S3Option => "native S3 configuration requires Python", Self::GcsBucket => "native GCS cache requires a configured bucket name", Self::DiskStore => "native disk cache requires the built-in diskcache store", + Self::QdrantEndpoint => { + "native Qdrant requires the default REST port so the gRPC port can be derived" + } + Self::SemanticEmbedding => "native semantic embedding requires Python", } } } @@ -238,6 +267,13 @@ 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::AzureBlob) => project_azure_blob(&backend).map(|backend| { CacheConfigProjection::Native(Box::new(Self { policy, @@ -250,7 +286,7 @@ impl NativeCacheConfig { backend: CacheBackendConfig::RedisSemantic(Box::new(backend)), })) }), - Some(CacheType::QdrantSemantic) | None => Ok(CacheConfigProjection::Unsupported( + None => Ok(CacheConfigProjection::Unsupported( UnsupportedCacheConfig::Backend, )), } @@ -265,7 +301,8 @@ impl NativeCacheConfig { CacheBackendConfig::Disk(_) | CacheBackendConfig::AzureBlob(_) | CacheBackendConfig::Gcs(_) - | CacheBackendConfig::RedisSemantic(_) => None, + | CacheBackendConfig::RedisSemantic(_) + | CacheBackendConfig::QdrantSemantic(_) => None, }; if !matches!(self.backend, CacheBackendConfig::ValkeySemantic(_)) && service.default_ttl() != default_ttl @@ -373,11 +410,35 @@ impl NativeCacheConfig { Some("facade and native backend index names must match") } CacheBackendConfig::RedisSemantic(config) - if service.similarity_threshold() != Some(config.similarity_threshold as f32) => + if service.similarity_threshold() != Some(config.similarity_threshold) => { Some("facade and native backend similarity thresholds must match") } CacheBackendConfig::RedisSemantic(_) => None, + 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, CacheBackendConfig::AzureBlob(config) => match service.azure_blob_identity() { None => Some("facade and native backend types must match"), Some((account_url, container)) @@ -391,6 +452,105 @@ impl NativeCacheConfig { } } +#[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() != Some(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); + + if optional_attribute(backend, "embedding_max_input_tokens")? + .is_some_and(|value| !value.is_none()) + { + 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((configured_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) = std::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 = std::env::var("OPENAI_BASE_URL") + .or_else(|_| std::env::var("OPENAI_API_BASE")) + .unwrap_or_else(|_| "https://api.openai.com/v1".to_owned()); + let timeout = optional_attribute(backend, "embedding_timeout")? + .map(|value| value.extract::>()) + .transpose()? + .flatten() + .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_azure_blob(backend: &Bound<'_, PyAny>) -> PyResult { let client = backend.getattr("container_client")?; diff --git a/litellm-rust/crates/python-bridge/src/cache/facade.rs b/litellm-rust/crates/python-bridge/src/cache/facade.rs index f1389b745e9..43b53584943 100644 --- a/litellm-rust/crates/python-bridge/src/cache/facade.rs +++ b/litellm-rust/crates/python-bridge/src/cache/facade.rs @@ -368,6 +368,11 @@ impl FacadeGuard { "RedisSemanticCache", "redis-semantic", ), + ("qdrant_semantic", _) => ( + "litellm.caching.qdrant_semantic_cache", + "QdrantSemanticCache", + "qdrant-semantic", + ), ("redis", true) => ( "litellm.caching.redis_cluster_cache", "RedisClusterCache", @@ -443,6 +448,10 @@ impl FacadeGuard { "embedding_model", "embedding_max_input_tokens", "embedding_timeout", + "qdrant_api_base", + "qdrant_api_key", + "collection_name", + "vector_size", "_index_name", "_redis_url", "similarity_threshold", diff --git a/litellm-rust/crates/python-bridge/src/cache/handle.rs b/litellm-rust/crates/python-bridge/src/cache/handle.rs index e6edf48d72f..51e1b02c405 100644 --- a/litellm-rust/crates/python-bridge/src/cache/handle.rs +++ b/litellm-rust/crates/python-bridge/src/cache/handle.rs @@ -1,18 +1,26 @@ use litellm_auth_aws::AwsAuthConfig; use litellm_cache_gcs::{DEFAULT_ENDPOINT, GcsConfig}; +use litellm_cache_qdrant_semantic::{OpenAiEmbedderConfig, Quantization}; use litellm_cache_redis::{RedisNode, RedisTopology}; use litellm_cache_redis_semantic::RedisSemanticConfig; use litellm_cache_s3::{S3CacheConfig, S3Endpoint}; use litellm_host_python::{release_gil, run_sync_value}; +use litellm_http::ClientVariant; use pyo3::{ PyTraverseError, PyVisit, exceptions::{PyRuntimeError, PyTypeError}, prelude::*, + types::PyDict, }; +use url::Url; use super::{ - cache_error, config::project_redis_semantic, embedder::PythonEmbedder, facade::FacadeGuard, - native::NativeResponseCache, request::duration, + cache_error, + config::{QdrantSemanticCacheConfig, project_redis_semantic}, + embedder::PythonEmbedder, + facade::FacadeGuard, + native::NativeResponseCache, + request::duration, }; #[pyclass(frozen, name = "_CacheTestHandle")] @@ -146,6 +154,105 @@ 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() != Some(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(|| { + std::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(|| { + std::env::var("OPENAI_BASE_URL") + .or_else(|_| std::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 http_config = crate::http::call_config(py, &PyDict::new(py), true)?; + let client = crate::http::pool() + .client(&http_config, ClientVariant::Provider) + .map_err(crate::http::client_error)?; + let service = run_sync_value(py, async move { + let handle = tokio::runtime::Handle::current(); + NativeResponseCache::qdrant_semantic(config, client, handle) + .await + .map_err(cache_error) + })?; + Ok(Self { + service, + guard: None, + pid: std::process::id(), + }) + } + #[staticmethod] #[pyo3(signature = (url, similarity_threshold, index_name, embedder))] fn valkey_semantic( diff --git a/litellm-rust/crates/python-bridge/src/cache/native.rs b/litellm-rust/crates/python-bridge/src/cache/native.rs index acc457d8c9c..7ca62f10ec8 100644 --- a/litellm-rust/crates/python-bridge/src/cache/native.rs +++ b/litellm-rust/crates/python-bridge/src/cache/native.rs @@ -7,6 +7,7 @@ use litellm_cache_azure_blob::AzureBlobCache; use litellm_cache_disk::DiskCache; use litellm_cache_gcs::{GcsCache, GcsConfig, StaticTokenSource}; use litellm_cache_memory::InMemoryCache; +use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, QdrantSemanticCache}; use litellm_cache_redis::{RedisCache, RedisTopology}; use litellm_cache_redis_semantic::{RedisSemanticCache, RedisSemanticConfig}; use litellm_cache_response::{ @@ -19,6 +20,7 @@ use pyo3::{PyTraverseError, PyVisit, prelude::*}; use serde_json::Value; use super::{ + config::QdrantSemanticCacheConfig, embedder::PythonEmbedder, request::NativeRequest, semantic::{SemanticBody, SemanticOperation, drive}, @@ -92,6 +94,7 @@ pub(super) enum NativeResponseCache { cache: Arc>>, embedder: PythonEmbedder, }, + QdrantSemantic(Arc>>), Disk(Arc>>), AzureBlob(Arc>>), } @@ -168,6 +171,31 @@ impl NativeResponseCache { }) } + pub async fn qdrant_semantic( + config: QdrantSemanticCacheConfig, + client: reqwest::Client, + runtime: tokio::runtime::Handle, + ) -> Result { + 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(client, config.embedding); + let cache = QdrantSemanticCache::connect( + qdrant, + embedder, + ResponseCacheCodec, + qdrant_config, + runtime, + ) + .await?; + Ok(Self::QdrantSemantic(Arc::new(ResponseCache::new( + Arc::new(cache), + )))) + } + pub fn disk(directory: &str) -> Result { let cache = DiskCache::open(directory, ResponseCacheCodec)?; Ok(Self::Disk(Arc::new(ResponseCache::new(Arc::new(cache))))) @@ -209,6 +237,7 @@ impl NativeResponseCache { | Self::S3(_) | Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } + | Self::QdrantSemantic(_) | Self::Disk(_) | Self::Gcs(_) => None, } @@ -223,7 +252,7 @@ impl NativeResponseCache { } } - pub(super) fn redis_semantic_request( + pub(super) fn semantic_request( request: &NativeRequest, ) -> ResponseCacheRequest { ResponseCacheRequest { @@ -289,6 +318,7 @@ impl NativeResponseCache { Self::Gcs(_) => "gcs", Self::ValkeySemantic { .. } => "valkey-semantic", Self::RedisSemantic { .. } => "redis_semantic", + Self::QdrantSemantic(_) => "qdrant_semantic", Self::Disk(_) => "disk", Self::AzureBlob(_) => "azure-blob", } @@ -302,6 +332,7 @@ impl NativeResponseCache { Self::Gcs(cache) => cache.default_ttl(), Self::ValkeySemantic { cache, .. } => cache.default_ttl(), Self::RedisSemantic { cache, .. } => cache.default_ttl(), + Self::QdrantSemantic(_) => None, Self::Disk(cache) => cache.default_ttl(), Self::AzureBlob(cache) => cache.default_ttl(), } @@ -341,6 +372,7 @@ impl NativeResponseCache { | Self::S3(_) | Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } + | Self::QdrantSemantic(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, @@ -354,6 +386,7 @@ impl NativeResponseCache { | Self::S3(_) | Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } + | Self::QdrantSemantic(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, @@ -368,6 +401,7 @@ impl NativeResponseCache { | Self::S3(_) | Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } + | Self::QdrantSemantic(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, @@ -381,6 +415,7 @@ impl NativeResponseCache { | Self::S3(_) | Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } + | Self::QdrantSemantic(_) | Self::Disk(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, @@ -395,6 +430,7 @@ impl NativeResponseCache { | Self::S3(_) | Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } + | Self::QdrantSemantic(_) | Self::AzureBlob(_) | Self::Gcs(_) => None, } @@ -421,9 +457,33 @@ impl NativeResponseCache { } } - pub fn similarity_threshold(&self) -> Option { + pub fn similarity_threshold(&self) -> Option { match self { - Self::RedisSemantic { cache, .. } => Some(cache.backend().similarity_threshold()), + Self::RedisSemantic { cache, .. } => { + Some(f64::from(cache.backend().similarity_threshold())) + } + Self::QdrantSemantic(cache) => Some(cache.backend().similarity_threshold()), + _ => None, + } + } + + pub fn collection_name(&self) -> Option<&str> { + match self { + Self::QdrantSemantic(cache) => Some(cache.backend().collection_name()), + _ => 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, } } @@ -451,9 +511,10 @@ impl NativeResponseCache { cache.lookup(&Self::semantic(request, scope), now) } Self::RedisSemantic { cache, .. } => { - cache.lookup(&Self::redis_semantic_request(request), now) + cache.lookup(&Self::semantic_request(request), now) } Self::Gcs(cache) => cache.lookup(&Self::exact(request), now), + Self::QdrantSemantic(cache) => cache.lookup(&Self::semantic_request(request), now), Self::Disk(cache) => cache.lookup(&Self::exact(request), now), Self::AzureBlob(cache) => cache.lookup(&Self::exact(request), now), } @@ -473,9 +534,12 @@ impl NativeResponseCache { cache.store(&Self::semantic(request, scope), response, now) } Self::RedisSemantic { cache, .. } => { - cache.store(&Self::redis_semantic_request(request), response, now) + cache.store(&Self::semantic_request(request), response, now) } Self::Gcs(cache) => cache.store(&Self::exact(request), response, now), + Self::QdrantSemantic(cache) => { + cache.store(&Self::semantic_request(request), response, now) + } Self::Disk(cache) => cache.store(&Self::exact(request), response, now), Self::AzureBlob(cache) => cache.store(&Self::exact(request), response, now), } @@ -498,7 +562,7 @@ impl NativeResponseCache { Self::S3(cache) => { cache.lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) } - Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { Err(Error::UnsupportedOperation) } Self::Gcs(cache) => { @@ -529,7 +593,12 @@ impl NativeResponseCache { } Self::RedisSemantic { cache, .. } => { cache - .async_lookup(&Self::redis_semantic_request(request), now) + .async_lookup(&Self::semantic_request(request), now) + .await + } + Self::QdrantSemantic(cache) => { + cache + .async_lookup(&Self::semantic_request(request), now) .await } Self::Gcs(cache) => cache.async_lookup(&Self::exact(request), now).await, @@ -573,6 +642,14 @@ impl NativeResponseCache { py, SemanticBody::new(self.clone(), SemanticOperation::Lookup(request)), ), + Self::QdrantSemantic(_) => { + let service = self.clone(); + litellm_host_python::run_async( + py, + async move { service.async_lookup(&request, super::request::now()).await }, + super::cache_error, + ) + } } } @@ -616,7 +693,12 @@ impl NativeResponseCache { } Self::RedisSemantic { cache, .. } => { cache - .async_store(&Self::redis_semantic_request(request), response, now) + .async_store(&Self::semantic_request(request), response, now) + .await + } + Self::QdrantSemantic(cache) => { + cache + .async_store(&Self::semantic_request(request), response, now) .await } Self::Gcs(cache) => { @@ -678,6 +760,18 @@ impl NativeResponseCache { py, SemanticBody::new(self.clone(), SemanticOperation::Store(request, response)), ), + Self::QdrantSemantic(_) => { + let service = self.clone(); + litellm_host_python::run_async( + py, + async move { + service + .async_store(&request, response, super::request::now()) + .await + }, + super::cache_error, + ) + } } } @@ -700,7 +794,7 @@ impl NativeResponseCache { .async_lookup_batch(&requests.iter().map(Self::exact).collect::>(), now) .await } - Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { Err(Error::UnsupportedOperation) } Self::Gcs(cache) => { @@ -755,7 +849,9 @@ impl NativeResponseCache { .collect(); cache.async_store_batch(entries, now).await } - Self::RedisSemantic { .. } => Err(Error::UnsupportedOperation), + Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { + Err(Error::UnsupportedOperation) + } Self::Gcs(cache) => { let entries = entries .into_iter() @@ -826,6 +922,7 @@ impl NativeResponseCache { py, SemanticBody::new(self.clone(), SemanticOperation::StoreBatch(entries.into())), ), + Self::QdrantSemantic(_) => Err(super::cache_error(Error::UnsupportedOperation)), } } @@ -839,7 +936,7 @@ impl NativeResponseCache { cache.async_flush().await } Self::S3(cache) => cache.async_flush().await, - Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } => { + Self::ValkeySemantic { .. } | Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { Err(Error::UnsupportedOperation) } Self::Gcs(cache) => cache.async_flush().await, @@ -854,7 +951,9 @@ impl NativeResponseCache { Self::Redis { cache, .. } => cache.test_connection().await, Self::S3(cache) => cache.test_connection().await, Self::ValkeySemantic { cache, .. } => cache.test_connection().await, - Self::RedisSemantic { .. } => Err(Error::UnsupportedOperation), + Self::RedisSemantic { .. } | Self::QdrantSemantic(_) => { + Err(Error::UnsupportedOperation) + } Self::Gcs(cache) => cache.test_connection().await, Self::Disk(cache) => cache.test_connection().await, Self::AzureBlob(cache) => cache.test_connection().await, diff --git a/litellm-rust/crates/python-bridge/src/cache/semantic.rs b/litellm-rust/crates/python-bridge/src/cache/semantic.rs index 7661598c5f9..934de01e721 100644 --- a/litellm-rust/crates/python-bridge/src/cache/semantic.rs +++ b/litellm-rust/crates/python-bridge/src/cache/semantic.rs @@ -101,7 +101,7 @@ impl ExecutionBody for SemanticBody { let (request, _) = self.pending.as_ref().ok_or_else(|| { PyRuntimeError::new_err("semantic execution has no pending operation") })?; - let semantic = NativeResponseCache::redis_semantic_request(request); + let semantic = NativeResponseCache::semantic_request(request); let Some(prompt) = prompt_from_context(&semantic.context) else { return self.backend_step(py, Err(Error::Unavailable)); }; diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 058cc8a1579..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,6 +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_FOR_INDEXING, json=data, ) @@ -422,6 +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_FOR_INDEXING, json=data, ) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index c0bb93daadc..31623a53dd0 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -167,6 +167,20 @@ class _CacheTestHandle: @staticmethod def disk(directory: str) -> _CacheTestHandle: ... @staticmethod + def qdrant_semantic( + url: str, + *, + collection_name: str, + similarity_threshold: float, + vector_size: int, + embedding_model: str = "text-embedding-3-small", + api_key: str | None = None, + embedding_api_key: str | None = None, + embedding_api_base: str | None = None, + embedding_timeout_seconds: float | None = None, + quantization: str = "binary", + ) -> _CacheTestHandle: ... + @staticmethod def azure_blob(account_url: str, container: str) -> _CacheTestHandle: ... @staticmethod def redis_semantic(backend: object) -> _CacheTestHandle: ... 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 45f7214ecbf..c4962b533c4 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -2,6 +2,7 @@ import asyncio import contextvars import gc import hashlib +import http.server import json import math import os @@ -59,6 +60,71 @@ def request(key: str = "key") -> dict[str, object]: return {"key": {"preset": key}} +def qdrant_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: Final = 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") @@ -1691,3 +1757,167 @@ def test_redis_semantic_handle_rejects_wrong_backends( ) with pytest.raises(TypeError, match="must be the native embedder"): _CacheTestHandle.redis_semantic(facade.cache)._bind_facade(replacement_facade) + + +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(qdrant_request("python-key", messages)) == {"id": "py"} + binding.store(qdrant_request("native-key", 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(qdrant_request("native-key", unrelated)) is None + assert facade.cache.get_cache("native-key", messages=unrelated) is None + assert binding.lookup(qdrant_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(qdrant_request("python-key", messages)) == {"id": "py"} + await binding.async_store(qdrant_request("native-key", 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( + 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(qdrant_request(key, messages)) is None + with pytest.raises(RuntimeError, match="operation is not supported"): + binding.lookup_batch([qdrant_request(key, messages)]) + with pytest.raises(RuntimeError, match="operation is not supported"): + await binding.async_flush() + with pytest.raises(RuntimeError, match="operation is not supported"): + 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(qdrant_request("persistent-key", messages, ttl_seconds=1.0), {"id": "persistent"}) + time.sleep(1.2) + assert binding.lookup(qdrant_request("persistent-key", 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( + 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.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}") + 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)