From 8d41336a1e1bea925801eaff33b84010b6c186f7 Mon Sep 17 00:00:00 2001 From: Yujong Lee Date: Mon, 21 Sep 2026 20:33:00 +0000 Subject: [PATCH] test(cache-qdrant-semantic): cover backend contract against an in-process Qdrant Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm-rust/Cargo.lock | 1 + .../crates/cache-qdrant-semantic/Cargo.toml | 1 + .../cache-qdrant-semantic/src/semantic.rs | 6 +- .../cache-qdrant-semantic/tests/embedder.rs | 141 ++++++ .../cache-qdrant-semantic/tests/qdrant.rs | 418 ++++++++++++++++++ .../tests/support/mod.rs | 339 ++++++++++++++ 6 files changed, 903 insertions(+), 3 deletions(-) create mode 100644 litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs create mode 100644 litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs create mode 100644 litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index 8868b25bb18..fb75c2241b5 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -2537,6 +2537,7 @@ dependencies = [ "serde_json", "thiserror 2.0.19", "tokio", + "tokio-stream", "tonic", "tonic-prost", "uuid", diff --git a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml index 7e215cab837..09d6a9637f3 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml +++ b/litellm-rust/crates/cache-qdrant-semantic/Cargo.toml @@ -21,3 +21,4 @@ litellm-cache-response.workspace = true rstest.workspace = true tonic = "0.14" tonic-prost = "0.14" +tokio-stream = "0.1" diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs index 16f11286d3f..d761364f1ad 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs @@ -167,13 +167,13 @@ impl QdrantSemanticCache { let Some(point) = result.result.into_iter().next() else { return Ok(None); }; - if f64::from(point.score) < self.config.similarity_threshold { - return Ok(None); - } let payload: Map = Payload::from(point.payload).into(); if payload.get("litellm_cache_key").and_then(Value::as_str) != Some(key) { return Ok(None); } + if f64::from(point.score) < self.config.similarity_threshold { + return Ok(None); + } let response = payload .get("response") .and_then(Value::as_str) diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs new file mode 100644 index 00000000000..adce70654a8 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/embedder.rs @@ -0,0 +1,141 @@ +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; + +use litellm_cache::Error; +use litellm_cache_qdrant_semantic::{Embedder, OpenAiEmbedder, OpenAiEmbedderConfig}; +use serde_json::Value; +use tokio::{ + io::{AsyncReadExt, AsyncWriteExt}, + net::TcpListener, +}; + +struct TestHttpServer { + address: std::net::SocketAddr, + request: Arc>>>, + task: tokio::task::JoinHandle<()>, +} + +impl TestHttpServer { + async fn response(status: &str, body: &str) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let request = Arc::new(Mutex::new(None)); + let captured = request.clone(); + let status = status.to_owned(); + let body = body.to_owned(); + let task = tokio::spawn(async move { + let (mut stream, _) = listener.accept().await.unwrap(); + let request_bytes = read_request(&mut stream).await; + *captured.lock().unwrap() = Some(request_bytes); + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: application/json\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}", + body.len() + ); + stream.write_all(response.as_bytes()).await.unwrap(); + }); + Self { + address, + request, + task, + } + } + + async fn hanging() -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let task = tokio::spawn(async move { + let (_stream, _) = listener.accept().await.unwrap(); + std::future::pending::<()>().await; + }); + Self { + address, + request: Arc::new(Mutex::new(None)), + task, + } + } + + fn base_url(&self) -> String { + format!("http://{}", self.address) + } +} + +impl Drop for TestHttpServer { + fn drop(&mut self) { + self.task.abort(); + } +} + +async fn read_request(stream: &mut tokio::net::TcpStream) -> Vec { + let mut bytes = Vec::new(); + let header_end = loop { + let mut chunk = [0_u8; 1024]; + let count = stream.read(&mut chunk).await.unwrap(); + assert_ne!(count, 0); + bytes.extend_from_slice(&chunk[..count]); + if let Some(end) = bytes.windows(4).position(|window| window == b"\r\n\r\n") { + break end + 4; + } + }; + let headers = String::from_utf8_lossy(&bytes[..header_end]); + let content_length = headers + .lines() + .find_map(|line| { + line.split_once(':') + .filter(|(name, _)| name.eq_ignore_ascii_case("content-length")) + .map(|(_, value)| value.trim()) + }) + .unwrap() + .parse::() + .unwrap(); + while bytes.len() < header_end + content_length { + let mut chunk = [0_u8; 1024]; + let count = stream.read(&mut chunk).await.unwrap(); + assert_ne!(count, 0); + bytes.extend_from_slice(&chunk[..count]); + } + bytes +} + +fn config(base: String, timeout: Option) -> OpenAiEmbedderConfig { + OpenAiEmbedderConfig { + api_base: base, + api_key: "test-key".to_owned(), + model: "test-model".to_owned(), + timeout, + } +} + +#[tokio::test] +async fn posts_embeddings_request_and_parses_vector() { + let server = TestHttpServer::response("200 OK", r#"{"data":[{"embedding":[0.1,0.2]}]}"#).await; + let embedder = OpenAiEmbedder::new(config( + format!("{}/", server.base_url()), + Some(Duration::from_secs(1)), + )) + .unwrap(); + assert_eq!(embedder.embed("hello").await.unwrap(), vec![0.1, 0.2]); + let request = server.request.lock().unwrap().clone().unwrap(); + let request_text = String::from_utf8(request).unwrap(); + assert!(request_text.starts_with("POST /embeddings HTTP/1.1\r\n")); + assert!(request_text.contains("\r\nauthorization: Bearer test-key\r\n")); + let body = request_text.split("\r\n\r\n").nth(1).unwrap(); + let body: Value = serde_json::from_str(body).unwrap(); + assert_eq!(body["model"], "test-model"); + assert_eq!(body["input"], "hello"); + assert_eq!(body["encoding_format"], "float"); +} + +#[tokio::test] +async fn status_and_timeout_errors_are_unavailable() { + let server = TestHttpServer::response("500 Internal Server Error", "{}").await; + let embedder = + OpenAiEmbedder::new(config(server.base_url(), Some(Duration::from_secs(1)))).unwrap(); + assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable)); + + let server = TestHttpServer::hanging().await; + let embedder = + OpenAiEmbedder::new(config(server.base_url(), Some(Duration::from_millis(200)))).unwrap(); + assert_eq!(embedder.embed("hello").await, Err(Error::Unavailable)); +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs new file mode 100644 index 00000000000..d5ecaf7217d --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -0,0 +1,418 @@ +#[path = "support/mod.rs"] +mod support; + +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use litellm_cache::{ + BaseCache, CacheCodec, CacheContext, Error, SemanticCacheContext, SemanticCacheScope, +}; +use litellm_cache_qdrant_semantic::{ + Embedder, QdrantSemanticCache, QdrantSemanticConfig, Quantization, +}; +use litellm_cache_response::{ + CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, +}; +use qdrant_client::Payload; +use qdrant_client::{ + Qdrant, + qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams}, +}; +use serde_json::{Value as JsonValue, json}; + +use support::{FakeQdrant, FakeState, StoredPoint}; + +#[derive(Clone)] +struct FixedEmbedder { + vectors: Arc>>, +} + +impl FixedEmbedder { + fn new(vectors: impl IntoIterator)>) -> Self { + Self { + vectors: Arc::new( + vectors + .into_iter() + .map(|(prompt, vector)| (prompt.to_owned(), vector)) + .collect(), + ), + } + } +} + +impl Embedder for FixedEmbedder { + fn model(&self) -> &str { + "fixed" + } + + async fn embed(&self, input: &str) -> Result, Error> { + self.vectors.get(input).cloned().ok_or(Error::Unavailable) + } +} + +fn config(quantization: Quantization) -> QdrantSemanticConfig { + QdrantSemanticConfig { + collection_name: "semantic".to_owned(), + similarity_threshold: 0.9, + vector_size: 2, + quantization, + } +} + +fn context(prompt: &str) -> SemanticCacheContext { + SemanticCacheContext { + messages: vec![json!({"role": "user", "content": prompt})], + scope: SemanticCacheScope::default(), + ..Default::default() + } +} + +fn value(response: JsonValue) -> CacheEntry { + CacheEntry { + timestamp: Some(1.0), + response, + } +} + +async fn connect( + server: &FakeQdrant, + vectors: impl IntoIterator)>, +) -> QdrantSemanticCache { + let client = Qdrant::from_url(&server.url()).build().unwrap(); + QdrantSemanticCache::connect( + client, + FixedEmbedder::new(vectors), + ResponseCacheCodec, + config(Quantization::Binary), + tokio::runtime::Handle::current(), + ) + .await + .unwrap() +} + +#[tokio::test(flavor = "multi_thread")] +#[allow(deprecated)] +async fn connect_sets_collection_quantization_and_index() { + for (quantization, expected) in [ + (Quantization::Binary, 0), + (Quantization::Scalar, 1), + (Quantization::Product, 2), + ] { + let server = FakeQdrant::start(FakeState::default()).await; + let client = Qdrant::from_url(&server.url()).build().unwrap(); + QdrantSemanticCache::connect( + client, + FixedEmbedder::new([]), + ResponseCacheCodec, + config(quantization), + tokio::runtime::Handle::current(), + ) + .await + .unwrap(); + let state = server.state.lock().unwrap(); + let request = &state.created_collections[0]; + let Some(qdrant::vectors_config::Config::Params(VectorParams { size, distance, .. })) = + request + .vectors_config + .as_ref() + .and_then(|config| config.config.clone()) + else { + panic!("missing vector params"); + }; + assert_eq!(size, 2); + assert_eq!(distance, Distance::Cosine as i32); + let quantization_config = request + .quantization_config + .as_ref() + .unwrap() + .quantization + .unwrap(); + match (expected, quantization_config) { + (0, qdrant::quantization_config::Quantization::Binary(binary)) => { + assert_eq!(binary.always_ram, Some(false)); + } + (1, qdrant::quantization_config::Quantization::Scalar(scalar)) => { + assert_eq!(scalar.r#type, QuantizationType::Int8 as i32); + assert_eq!(scalar.quantile, Some(0.99)); + assert_eq!(scalar.always_ram, Some(false)); + } + (2, qdrant::quantization_config::Quantization::Product(product)) => { + assert_eq!(product.compression, CompressionRatio::X16 as i32); + assert_eq!(product.always_ram, Some(false)); + } + _ => panic!("unexpected quantization"), + } + assert!(state.index_creations >= 1); + assert_eq!(state.field_indexes[0].collection_name, "semantic"); + assert_eq!(state.field_indexes[0].field_name, "litellm_cache_key"); + assert_eq!( + state.field_indexes[0].field_type, + Some(qdrant::FieldType::Keyword as i32) + ); + server.stop(); + } +} + +#[tokio::test(flavor = "multi_thread")] +async fn existing_collection_skips_create_and_index_failure_is_non_fatal() { + let server = FakeQdrant::start(FakeState { + collections: ["semantic".to_owned()].into_iter().collect(), + fail_field_index: true, + ..Default::default() + }) + .await; + let _cache = connect(&server, [("hello", vec![1.0, 0.0])]).await; + let state = server.state.lock().unwrap(); + assert!(state.created_collections.is_empty()); + assert!(state.index_creations >= 1); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn async_and_sync_set_get_store_exact_payload() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await); + let ctx = context("hello"); + let entry = value(json!({"answer": 42})); + cache + .async_set_cache("key", entry.clone(), ctx.clone()) + .await + .unwrap(); + assert_eq!( + cache.async_get_cache("key", &ctx).await.unwrap().as_ref(), + Some(&entry) + ); + { + let state = server.state.lock().unwrap(); + let payload = &state.points[0].payload; + let mut payload_keys = payload.keys().cloned().collect::>(); + payload_keys.sort(); + assert_eq!(payload_keys, ["litellm_cache_key", "response", "text"]); + assert_eq!(payload["litellm_cache_key"], Value::from("key")); + assert_eq!( + payload["response"], + Value::from(String::from_utf8(ResponseCacheCodec.encode(&entry).unwrap()).unwrap()) + ); + } + let sync_entry = entry.clone(); + let sync_cache = cache.clone(); + let sync_ctx = ctx.clone(); + tokio::task::spawn_blocking(move || { + sync_cache + .set_cache("sync", sync_entry.clone(), &sync_ctx) + .unwrap(); + assert_eq!( + sync_cache.get_cache("sync", &sync_ctx).unwrap(), + Some(sync_entry) + ); + }) + .await + .unwrap(); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn misses_and_payload_validation_are_safe() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = connect( + &server, + [("hello", vec![1.0, 0.0]), ("near", vec![0.7, 0.71414286])], + ) + .await; + let entry = value(json!({"answer": 1})); + cache + .async_set_cache("key", entry, context("hello")) + .await + .unwrap(); + assert_eq!( + cache + .async_get_cache("other", &context("hello")) + .await + .unwrap(), + None + ); + assert_eq!( + cache + .async_get_cache("key", &context("near")) + .await + .unwrap(), + None + ); + server.insert_point(StoredPoint { + id: Some(PointId::from(99_u64)), + vector: vec![1.0, 0.0], + payload: Payload::try_from(json!({ + "litellm_cache_key": 99, + "response": "{}", + })) + .unwrap() + .into(), + }); + assert_eq!( + cache + .async_get_cache("99", &context("hello")) + .await + .unwrap(), + None + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn decoding_errors_missing_prompt_pipeline_and_ttl_behave_as_required() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = connect(&server, [("one", vec![1.0, 0.0]), ("two", vec![0.0, 1.0])]).await; + let empty = SemanticCacheContext::default(); + assert_eq!( + cache + .async_set_cache("key", value(json!({})), empty.clone()) + .await, + Err(Error::MissingPrompt) + ); + assert_eq!( + cache.async_get_cache("key", &empty).await, + Err(Error::MissingPrompt) + ); + assert_eq!( + cache.async_get_cache("key", &context("unknown")).await, + Err(Error::Unavailable) + ); + cache + .async_set_cache( + "ttl", + value(json!({"ttl": true})), + context("one").with_ttl(Some(Duration::from_secs(1))), + ) + .await + .unwrap(); + tokio::time::sleep(Duration::from_millis(1_100)).await; + assert!( + cache + .async_get_cache( + "ttl", + &context("one").with_ttl(Some(Duration::from_secs(1))), + ) + .await + .unwrap() + .is_some() + ); + cache + .async_set_cache_pipeline( + vec![ + ("one".to_owned(), value(json!({"n": 1}))), + ("two".to_owned(), value(json!({"n": 2}))), + ], + context("one"), + ) + .await + .unwrap(); + assert!( + cache + .async_get_cache("one", &context("one")) + .await + .unwrap() + .is_some() + ); + assert!( + cache + .async_get_cache("two", &context("one")) + .await + .unwrap() + .is_some() + ); + assert_eq!(cache.get_ttl(&context("one")), None); + assert_eq!( + cache.test_connection().await, + Err(Error::UnsupportedOperation) + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn response_payloads_decode_and_invalid_entries_fail() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await; + for (key, response) in [ + ("python", json!("{'timestamp': 1.0, 'response': {'a': 1}}")), + ("garbage", json!("not json")), + ("missing", json!("unused")), + ] { + let mut payload = serde_json::Map::new(); + payload.insert("litellm_cache_key".to_owned(), json!(key)); + if key != "missing" { + payload.insert("response".to_owned(), response); + } + server.insert_point(StoredPoint { + id: Some(PointId::from(key.len() as u64)), + vector: vec![1.0, 0.0], + payload: Payload::try_from(JsonValue::Object(payload)) + .unwrap() + .into(), + }); + } + assert_eq!( + cache + .async_get_cache("python", &context("hello")) + .await + .unwrap(), + Some(value(json!({"a": 1}))) + ); + assert_eq!( + cache.async_get_cache("garbage", &context("hello")).await, + Err(Error::InvalidEntry) + ); + assert_eq!( + cache.async_get_cache("missing", &context("hello")).await, + Err(Error::InvalidEntry) + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn response_cache_facade_turns_invalid_entry_into_miss() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = Arc::new(connect(&server, [("hello", vec![1.0, 0.0])]).await); + let request = ResponseCacheRequest::::new(CacheKeyInput { + preset: Some("key".to_owned()), + ..Default::default() + }) + .with_context(context("hello")); + let response = json!({"answer": 42}); + let facade = ResponseCache::new(cache.clone()); + facade + .async_store(&request, response.clone(), Duration::from_secs(1)) + .await + .unwrap(); + assert_eq!( + facade + .async_lookup(&request, Duration::from_secs(1)) + .await + .unwrap(), + Some(response) + ); + { + let mut state = server.state.lock().unwrap(); + state.points[0] + .payload + .insert("response".to_owned(), Value::from("not json")); + } + assert_eq!( + facade + .async_lookup(&request, Duration::from_secs(1)) + .await + .unwrap(), + None + ); + server.stop(); +} + +#[tokio::test(flavor = "multi_thread")] +async fn stopped_qdrant_server_maps_to_unavailable() { + let server = FakeQdrant::start(FakeState::default()).await; + let cache = connect(&server, [("hello", vec![1.0, 0.0])]).await; + server.stop(); + tokio::time::sleep(Duration::from_millis(50)).await; + assert_eq!( + cache.async_get_cache("key", &context("hello")).await, + Err(Error::Unavailable) + ); +} diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs new file mode 100644 index 00000000000..860213a1703 --- /dev/null +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs @@ -0,0 +1,339 @@ +use std::{ + collections::{HashMap, HashSet}, + net::SocketAddr, + sync::{Arc, Mutex}, +}; + +use qdrant_client::qdrant::collections_server::CollectionsServer; +use qdrant_client::qdrant::{ + self, CollectionExists, CollectionExistsRequest, CollectionExistsResponse, + CollectionOperationResponse, CreateCollection, CreateFieldIndexCollection, Filter, PointId, + PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors, + collections_server::Collections, + points_server::{Points, PointsServer}, +}; +use tokio::sync::oneshot; +use tokio_stream::wrappers::TcpListenerStream; +use tonic::{Request, Response, Status, transport::Server}; + +#[derive(Clone, Debug)] +pub struct StoredPoint { + pub id: Option, + pub vector: Vec, + pub payload: HashMap, +} + +#[derive(Default)] +pub struct FakeState { + pub collections: HashSet, + pub created_collections: Vec, + pub field_indexes: Vec, + pub points: Vec, + pub index_creations: usize, + pub fail_field_index: bool, +} + +#[derive(Clone)] +pub struct FakeQdrant { + pub state: Arc>, + pub address: SocketAddr, + shutdown: Arc>>>, +} + +impl FakeQdrant { + pub async fn start(state: FakeState) -> Self { + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let state = Arc::new(Mutex::new(state)); + let service = FakeService { + state: state.clone(), + }; + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + tokio::spawn(async move { + Server::builder() + .add_service(CollectionsServer::new(service.clone())) + .add_service(PointsServer::new(service)) + .serve_with_incoming_shutdown(TcpListenerStream::new(listener), async { + let _ = shutdown_rx.await; + }) + .await + .unwrap(); + }); + Self { + state, + address, + shutdown: Arc::new(Mutex::new(Some(shutdown_tx))), + } + } + + pub fn url(&self) -> String { + format!("http://{}", self.address) + } + + pub fn stop(&self) { + self.shutdown + .lock() + .unwrap() + .take() + .unwrap() + .send(()) + .unwrap(); + } + + pub fn insert_point(&self, point: StoredPoint) { + self.state.lock().unwrap().points.push(point); + } +} + +#[derive(Clone)] +struct FakeService { + state: Arc>, +} + +macro_rules! unimplemented_collections { + ($($name:ident, $request:ty, $response:ty);* $(;)?) => { + $( + fn $name<'life0, 'async_trait>( + &'life0 self, + _: Request<$request>, + ) -> std::pin::Pin< + Box< + dyn std::future::Future< + Output = Result, Status>, + > + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async { Err(Status::unimplemented(stringify!($name))) }) + } + )* + }; +} + +macro_rules! unimplemented_points { + ($($name:ident, $request:ty, $response:ty);* $(;)?) => { + $( + fn $name<'life0, 'async_trait>( + &'life0 self, + _: Request<$request>, + ) -> std::pin::Pin< + Box< + dyn std::future::Future< + Output = Result, Status>, + > + Send + + 'async_trait, + >, + > + where + 'life0: 'async_trait, + Self: 'async_trait, + { + Box::pin(async { Err(Status::unimplemented(stringify!($name))) }) + } + )* + }; +} + +#[tonic::async_trait] +impl Collections for FakeService { + async fn create( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + let mut state = self.state.lock().unwrap(); + state.collections.insert(request.collection_name.clone()); + state.created_collections.push(request); + Ok(Response::new(CollectionOperationResponse { + result: true, + ..Default::default() + })) + } + + async fn collection_exists( + &self, + request: Request, + ) -> Result, Status> { + let exists = self + .state + .lock() + .unwrap() + .collections + .contains(&request.into_inner().collection_name); + Ok(Response::new(CollectionExistsResponse { + result: Some(CollectionExists { exists }), + ..Default::default() + })) + } + + unimplemented_collections!( + get, qdrant::GetCollectionInfoRequest, qdrant::GetCollectionInfoResponse; + list, qdrant::ListCollectionsRequest, qdrant::ListCollectionsResponse; + update, qdrant::UpdateCollection, qdrant::CollectionOperationResponse; + delete, qdrant::DeleteCollection, qdrant::CollectionOperationResponse; + update_aliases, qdrant::ChangeAliases, qdrant::CollectionOperationResponse; + list_collection_aliases, qdrant::ListCollectionAliasesRequest, qdrant::ListAliasesResponse; + list_aliases, qdrant::ListAliasesRequest, qdrant::ListAliasesResponse; + collection_cluster_info, qdrant::CollectionClusterInfoRequest, qdrant::CollectionClusterInfoResponse; + update_collection_cluster_setup, qdrant::UpdateCollectionClusterSetupRequest, qdrant::UpdateCollectionClusterSetupResponse; + create_shard_key, qdrant::CreateShardKeyRequest, qdrant::CreateShardKeyResponse; + delete_shard_key, qdrant::DeleteShardKeyRequest, qdrant::DeleteShardKeyResponse; + list_shard_keys, qdrant::ListShardKeysRequest, qdrant::ListShardKeysResponse; + ); +} + +#[tonic::async_trait] +impl Points for FakeService { + async fn create_field_index( + &self, + request: Request, + ) -> Result, Status> { + let mut state = self.state.lock().unwrap(); + state.index_creations += 1; + state.field_indexes.push(request.into_inner()); + if state.fail_field_index { + return Err(Status::internal("field index failure")); + } + Ok(Response::new(PointsOperationResponse::default())) + } + + async fn upsert( + &self, + request: Request, + ) -> Result, Status> { + let mut state = self.state.lock().unwrap(); + for point in request.into_inner().points { + let stored = StoredPoint { + id: point.id.clone(), + vector: dense_vector(point.vectors)?, + payload: point.payload, + }; + if let Some(existing) = state + .points + .iter_mut() + .find(|existing| existing.id == stored.id) + { + *existing = stored; + } else { + state.points.push(stored); + } + } + Ok(Response::new(PointsOperationResponse::default())) + } + + async fn search( + &self, + request: Request, + ) -> Result, Status> { + let request = request.into_inner(); + let key_filter = keyword_filter(request.filter.as_ref()); + let state = self.state.lock().unwrap(); + let mut results = state + .points + .iter() + .filter(|point| { + key_filter.as_ref().is_none_or(|(field, expected)| { + point + .payload + .get(field) + .and_then(|value| { + let value: serde_json::Value = value.clone().into(); + value + .as_str() + .map(str::to_owned) + .or_else(|| value.as_i64().map(|value| value.to_string())) + }) + .is_some_and(|value| value == *expected) + }) + }) + .map(|point| ScoredPoint { + id: point.id.clone(), + payload: point.payload.clone(), + score: cosine(&request.vector, &point.vector), + ..Default::default() + }) + .collect::>(); + results.sort_by(|left, right| right.score.total_cmp(&left.score)); + results.truncate(request.limit as usize); + Ok(Response::new(SearchResponse { + result: results, + ..Default::default() + })) + } + + unimplemented_points!( + delete, qdrant::DeletePoints, qdrant::PointsOperationResponse; + get, qdrant::GetPoints, qdrant::GetResponse; + update_vectors, qdrant::UpdatePointVectors, qdrant::PointsOperationResponse; + delete_vectors, qdrant::DeletePointVectors, qdrant::PointsOperationResponse; + set_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse; + overwrite_payload, qdrant::SetPayloadPoints, qdrant::PointsOperationResponse; + delete_payload, qdrant::DeletePayloadPoints, qdrant::PointsOperationResponse; + clear_payload, qdrant::ClearPayloadPoints, qdrant::PointsOperationResponse; + delete_field_index, qdrant::DeleteFieldIndexCollection, qdrant::PointsOperationResponse; + create_vector_name, qdrant::CreateVectorNameRequest, qdrant::PointsOperationResponse; + delete_vector_name, qdrant::DeleteVectorNameRequest, qdrant::PointsOperationResponse; + search_batch, qdrant::SearchBatchPoints, qdrant::SearchBatchResponse; + search_groups, qdrant::SearchPointGroups, qdrant::SearchGroupsResponse; + scroll, qdrant::ScrollPoints, qdrant::ScrollResponse; + recommend, qdrant::RecommendPoints, qdrant::RecommendResponse; + recommend_batch, qdrant::RecommendBatchPoints, qdrant::RecommendBatchResponse; + recommend_groups, qdrant::RecommendPointGroups, qdrant::RecommendGroupsResponse; + discover, qdrant::DiscoverPoints, qdrant::DiscoverResponse; + discover_batch, qdrant::DiscoverBatchPoints, qdrant::DiscoverBatchResponse; + count, qdrant::CountPoints, qdrant::CountResponse; + update_batch, qdrant::UpdateBatchPoints, qdrant::UpdateBatchResponse; + query, qdrant::QueryPoints, qdrant::QueryResponse; + query_batch, qdrant::QueryBatchPoints, qdrant::QueryBatchResponse; + query_groups, qdrant::QueryPointGroups, qdrant::QueryGroupsResponse; + facet, qdrant::FacetCounts, qdrant::FacetResponse; + search_matrix_pairs, qdrant::SearchMatrixPoints, qdrant::SearchMatrixPairsResponse; + search_matrix_offsets, qdrant::SearchMatrixPoints, qdrant::SearchMatrixOffsetsResponse; + ); +} + +fn dense_vector(vectors: Option) -> Result, Status> { + let Some(Vectors { + vectors_options: + Some(qdrant::vectors::VectorsOptions::Vector(Vector { + vector: Some(qdrant::vector::Vector::Dense(qdrant::DenseVector { data })), + .. + })), + }) = vectors + else { + return Err(Status::invalid_argument("expected dense vector")); + }; + Ok(data) +} + +fn keyword_filter(filter: Option<&Filter>) -> Option<(String, String)> { + filter? + .must + .iter() + .find_map(|condition| match condition.condition_one_of.as_ref()? { + qdrant::condition::ConditionOneOf::Field(field) => { + let qdrant::r#match::MatchValue::Keyword(value) = + field.r#match.as_ref()?.match_value.as_ref()? + else { + return None; + }; + Some((field.key.clone(), value.clone())) + } + _ => None, + }) +} + +fn cosine(left: &[f32], right: &[f32]) -> f32 { + let dot = left + .iter() + .zip(right) + .map(|(left, right)| left * right) + .sum::(); + let left_norm = left.iter().map(|value| value * value).sum::().sqrt(); + let right_norm = right.iter().map(|value| value * value).sum::().sqrt(); + dot / (left_norm * right_norm) +}