diff --git a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs index d761364f1ad..8fa27d68d05 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs @@ -131,7 +131,7 @@ impl QdrantSemanticCache { vector, payload, )], - )) + ).wait(true)) .await .map_err(|_| Error::Unavailable)?; Ok(()) diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs index e8d8d5040f0..70cde4e0d5d 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/qdrant.rs @@ -12,15 +12,50 @@ use litellm_cache_qdrant_semantic::{ use litellm_cache_response::{ CacheEntry, CacheKeyInput, ResponseCache, ResponseCacheCodec, ResponseCacheRequest, }; -use qdrant_client::Payload; use qdrant_client::{ Qdrant, - qdrant::{self, CompressionRatio, Distance, PointId, QuantizationType, Value, VectorParams}, + qdrant::{ + self, CompressionRatio, Distance, PointId, QuantizationType, Struct, Value, VectorParams, + value::Kind, + }, }; use serde_json::{Value as JsonValue, json}; use support::{FakeQdrant, FakeState, StoredPoint}; +fn json_to_qdrant(value: JsonValue) -> Value { + let kind = match value { + JsonValue::Null => Kind::NullValue(0), + JsonValue::Bool(value) => Kind::BoolValue(value), + JsonValue::Number(value) => value + .as_i64() + .map(Kind::IntegerValue) + .or_else(|| value.as_f64().map(Kind::DoubleValue)) + .unwrap(), + JsonValue::String(value) => Kind::StringValue(value), + JsonValue::Array(values) => Kind::ListValue(qdrant::ListValue { + values: values.into_iter().map(json_to_qdrant).collect(), + }), + JsonValue::Object(values) => Kind::StructValue(Struct { + fields: values + .into_iter() + .map(|(key, value)| (key, json_to_qdrant(value))) + .collect(), + }), + }; + Value { kind: Some(kind) } +} + +fn payload_from_json(value: JsonValue) -> HashMap { + value + .as_object() + .unwrap() + .clone() + .into_iter() + .map(|(key, value)| (key, json_to_qdrant(value))) + .collect() +} + #[derive(Clone)] struct FixedEmbedder { vectors: Arc>>, @@ -243,12 +278,10 @@ async fn misses_and_payload_validation_are_safe() { server.insert_point(StoredPoint { id: Some(PointId::from(99_u64)), vector: vec![1.0, 0.0], - payload: Payload::try_from(json!({ + payload: payload_from_json(json!({ "litellm_cache_key": 99, "response": "{}", - })) - .unwrap() - .into(), + })), }); assert_eq!( cache @@ -322,6 +355,10 @@ async fn decoding_errors_missing_prompt_pipeline_and_ttl_behave_as_required() { .unwrap() .is_some() ); + assert_eq!( + server.state.lock().unwrap().upsert_waits, + vec![Some(true), Some(true), Some(true)] + ); assert_eq!(cache.get_ttl(&context("one")), None); assert_eq!( cache.test_connection().await, @@ -347,9 +384,7 @@ async fn response_payloads_decode_and_invalid_entries_fail() { server.insert_point(StoredPoint { id: Some(PointId::from(key.len() as u64)), vector: vec![1.0, 0.0], - payload: Payload::try_from(JsonValue::Object(payload)) - .unwrap() - .into(), + payload: payload.into_iter().map(|(key, value)| (key, json_to_qdrant(value))).collect(), }); } assert_eq!( diff --git a/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs index 860213a1703..2a4cd45e007 100644 --- a/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs +++ b/litellm-rust/crates/cache-qdrant-semantic/tests/support/mod.rs @@ -10,8 +10,10 @@ use qdrant_client::qdrant::{ CollectionOperationResponse, CreateCollection, CreateFieldIndexCollection, Filter, PointId, PointsOperationResponse, ScoredPoint, SearchPoints, SearchResponse, Value, Vector, Vectors, collections_server::Collections, + value::Kind, points_server::{Points, PointsServer}, }; +use serde_json::Value as JsonValue; use tokio::sync::oneshot; use tokio_stream::wrappers::TcpListenerStream; use tonic::{Request, Response, Status, transport::Server}; @@ -23,12 +25,33 @@ pub struct StoredPoint { pub payload: HashMap, } +fn qdrant_value_to_json(value: Value) -> JsonValue { + match value.kind { + Some(Kind::NullValue(_)) | None => JsonValue::Null, + Some(Kind::DoubleValue(value)) => serde_json::json!(value), + Some(Kind::IntegerValue(value)) => serde_json::json!(value), + Some(Kind::StringValue(value)) => JsonValue::String(value), + Some(Kind::BoolValue(value)) => JsonValue::Bool(value), + Some(Kind::StructValue(value)) => JsonValue::Object( + value + .fields + .into_iter() + .map(|(key, value)| (key, qdrant_value_to_json(value))) + .collect(), + ), + Some(Kind::ListValue(value)) => { + JsonValue::Array(value.values.into_iter().map(qdrant_value_to_json).collect()) + } + } +} + #[derive(Default)] pub struct FakeState { pub collections: HashSet, pub created_collections: Vec, pub field_indexes: Vec, pub points: Vec, + pub upsert_waits: Vec>, pub index_creations: usize, pub fail_field_index: bool, } @@ -205,8 +228,10 @@ impl Points for FakeService { &self, request: Request, ) -> Result, Status> { + let request = request.into_inner(); let mut state = self.state.lock().unwrap(); - for point in request.into_inner().points { + state.upsert_waits.push(request.wait); + for point in request.points { let stored = StoredPoint { id: point.id.clone(), vector: dense_vector(point.vectors)?, @@ -241,7 +266,7 @@ impl Points for FakeService { .payload .get(field) .and_then(|value| { - let value: serde_json::Value = value.clone().into(); + let value = qdrant_value_to_json(value.clone()); value .as_str() .map(str::to_owned) diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index 058cc8a1579..868f36f7f21 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -313,6 +313,7 @@ class QdrantSemanticCache(BaseCache): self.sync_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points", headers=self.headers, + params={"wait": "true"}, json=data, ) @@ -422,6 +423,7 @@ class QdrantSemanticCache(BaseCache): await self.async_client.put( url=f"{self.qdrant_api_base}/collections/{self.collection_name}/points", headers=self.headers, + params={"wait": "true"}, json=data, ) diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/test_litellm/caching/test_qdrant_semantic_cache.py index a0a9b71787c..ca7303e4c6d 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/test_litellm/caching/test_qdrant_semantic_cache.py @@ -578,6 +578,7 @@ def test_qdrant_semantic_cache_set_cache(): assert ( upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key" ) + assert qdrant_cache.sync_client.put.call_args.kwargs["params"] == {"wait": "true"} @pytest.mark.asyncio @@ -650,6 +651,7 @@ async def test_qdrant_semantic_cache_async_set_cache(): assert ( upsert_payload[QdrantSemanticCache.CACHE_KEY_FIELD_NAME] == "test_key" ) + assert qdrant_cache.async_client.put.call_args.kwargs["params"] == {"wait": "true"} def test_qdrant_semantic_cache_custom_vector_size(): diff --git a/tests/test_litellm_rust/test_cache.py b/tests/test_litellm_rust/test_cache.py index 02c159825bc..e43af903462 100644 --- a/tests/test_litellm_rust/test_cache.py +++ b/tests/test_litellm_rust/test_cache.py @@ -605,15 +605,7 @@ async def test_qdrant_semantic_async_parity( messages=messages, ) - async def lookup_after_commit() -> object: - for _ in range(20): - value: Final = await binding.async_lookup(semantic_request("python-key", messages)) - if value is not None: - return value - await asyncio.sleep(0.1) - return None - - assert await lookup_after_commit() == {"id": "py"} + assert await binding.async_lookup(semantic_request("python-key", messages)) == {"id": "py"} await binding.async_store(semantic_request("native-key", messages), {"id": "native"}) python_value: Final = await facade.cache.async_get_cache("native-key", messages=messages) assert isinstance(python_value, dict) @@ -705,6 +697,8 @@ def test_qdrant_semantic_mutation_and_projection_fallback( vector_size=8, ) handle._bind_facade(facade) + facade.cache.qdrant_api_key = "rotated" + assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" facade.cache.similarity_threshold = 0.5 assert _native._CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback" unsupported: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}")