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"}"# + ); +}