fix(cache-qdrant-semantic): wait for Qdrant upserts to be indexed

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm-rust/crates/cache-qdrant-semantic/src/semantic.rs

@@ -131,7 +131,7 @@ impl<E: Embedder, C: CacheCodec> QdrantSemanticCache<E, C> {
                     vector,
                     payload,
                 )],
-            ))
+            ).wait(true))
             .await
             .map_err(|_| Error::Unavailable)?;
         Ok(())

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<String, Value> {
+    value
+        .as_object()
+        .unwrap()
+        .clone()
+        .into_iter()
+        .map(|(key, value)| (key, json_to_qdrant(value)))
+        .collect()
+}
+
 #[derive(Clone)]
 struct FixedEmbedder {
     vectors: Arc<HashMap<String, Vec<f32>>>,
@@ -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!(

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<String, Value>,
 }
 
+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<String>,
     pub created_collections: Vec<CreateCollection>,
     pub field_indexes: Vec<CreateFieldIndexCollection>,
     pub points: Vec<StoredPoint>,
+    pub upsert_waits: Vec<Option<bool>>,
     pub index_creations: usize,
     pub fail_field_index: bool,
 }
@@ -205,8 +228,10 @@ impl Points for FakeService {
         &self,
         request: Request<qdrant::UpsertPoints>,
     ) -> Result<Response<PointsOperationResponse>, 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)

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,
         )
 

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():

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}")