feat: add ability to specify qdrant client kwargs (#4048)

Author: maxdswain

integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py

@@ -150,6 +150,7 @@ def __init__(
         write_batch_size: int = 100,
         scroll_size: int = 10_000,
         payload_fields_to_index: list[dict] | None = None,
+        client_kwargs: dict[str, Any] | None = None,
     ) -> None:
         """
         Initializes a QdrantDocumentStore.
@@ -236,6 +237,14 @@ def __init__(
             The scroll size for reading documents.
         :param payload_fields_to_index:
             List of payload fields to index.
+        :param client_kwargs:
+            Additional keyword arguments passed to both `QdrantClient` and `AsyncQdrantClient`.
+            These arguments take precedence over named client parameters in case of overlap.
+            Serialized as-is by `to_dict()`: callbacks and `httpx.Limits` may not serialize, and
+            authentication headers will appear in plain text in pipeline YAML.
+            The same arguments are used for sync and async clients. For example, an async
+            `auth_token_provider` is also passed to `QdrantClient`, where it may not work with
+            synchronous requests.
         """
 
         self._client: qdrant_client.QdrantClient | None = None
@@ -255,6 +264,7 @@ def __init__(
         self.path = path
         self.force_disable_check_same_thread = force_disable_check_same_thread
         self.metadata = metadata or {}
+        self.client_kwargs = dict(client_kwargs or {})
 
         # Store the Qdrant collection specific attributes
         self.shard_number = shard_number
@@ -2485,6 +2495,7 @@ def _prepare_client_params(self) -> dict[str, Any]:
             # when using a hosted Qdrant service, which means running to_dict() exposes the api key.
             "metadata": {**self.metadata},
             "force_disable_check_same_thread": self.force_disable_check_same_thread,
+            **{key: {**value} if isinstance(value, dict) else value for key, value in self.client_kwargs.items()},
         }
 
     def _prepare_collection_config(

integrations/qdrant/tests/test_dict_converters.py

@@ -43,6 +43,7 @@ def test_to_dict():
             "write_batch_size": 100,
             "scroll_size": 10000,
             "payload_fields_to_index": None,
+            "client_kwargs": {},
         },
     }
 
@@ -73,6 +74,7 @@ def test_from_dict():
                 "write_batch_size": 1000,
                 "scroll_size": 10000,
                 "payload_fields_to_index": None,
+                "client_kwargs": {"check_compatibility": False},
             },
         }
     )
@@ -102,5 +104,6 @@ def test_from_dict():
             document_store.scroll_size == 10000,
             document_store.api_key == Secret.from_env_var("ENV_VAR", strict=False),
             document_store.payload_fields_to_index is None,
+            document_store.client_kwargs == {"check_compatibility": False},
         ]
     )

integrations/qdrant/tests/test_document_store.py

@@ -49,6 +49,7 @@ def test_to_dict(self, monkeypatch):
             wait_result_from_api=True,
             use_sparse_embeddings=False,
             api_key=Secret.from_env_var("QDRANT_API_KEY"),
+            client_kwargs={"check_compatibility": False},
         )
         expected_dict = {
             "type": "haystack_integrations.document_stores.qdrant.document_store.QdrantDocumentStore",
@@ -91,6 +92,7 @@ def test_to_dict(self, monkeypatch):
                 "write_batch_size": 100,
                 "scroll_size": 10000,
                 "payload_fields_to_index": None,
+                "client_kwargs": {"check_compatibility": False},
             },
         }
         assert doc_store.to_dict() == expected_dict
@@ -347,6 +349,30 @@ def test_prepare_client_params_does_not_mutate_metadata(self):
         params["metadata"]["added"] = "x"
         assert metadata == {"key": "value"}
 
+    def test_prepare_client_params_does_not_mutate_client_kwargs(self):
+        client_kwargs = {"grpc_options": {"custom_option": "value"}}
+        document_store = QdrantDocumentStore(location=":memory:", client_kwargs=client_kwargs)
+        params = document_store._prepare_client_params()
+        params["grpc_options"]["grpc.primary_user_agent"] = "qdrant-client"
+
+        assert client_kwargs == {"grpc_options": {"custom_option": "value"}}
+        assert document_store.to_dict()["init_parameters"]["client_kwargs"] == client_kwargs
+
+    @pytest.mark.parametrize("port", [6333, 7333])
+    def test_client_kwargs_are_forwarded_with_client_kwargs_taking_precedence(self, port):
+        document_store = QdrantDocumentStore(
+            location=":memory:", port=port, client_kwargs={"check_compatibility": False, "port": 8333}
+        )
+
+        with (
+            patch("haystack_integrations.document_stores.qdrant.document_store.qdrant_client.QdrantClient") as client,
+            patch.object(document_store, "_set_up_collection"),
+        ):
+            document_store._initialize_client()
+
+        assert client.call_args.kwargs["check_compatibility"] is False
+        assert client.call_args.kwargs["port"] == 8333
+
     def test_get_batches_from_generator(self):
         batches = list(get_batches_from_generator([1, 2, 3, 4, 5], 2))
         assert batches == [(1, 2), (3, 4), (5,)]

integrations/qdrant/tests/test_document_store_async.py

@@ -35,6 +35,23 @@
 
 @pytest.mark.asyncio
 class TestQdrantDocumentStoreAsyncUnit:
+    @pytest.mark.parametrize("port", [6333, 7333])
+    async def test_client_kwargs_are_forwarded_async(self, port):
+        document_store = QdrantDocumentStore(
+            location=":memory:", port=port, client_kwargs={"check_compatibility": False, "port": 8333}
+        )
+
+        with (
+            patch(
+                "haystack_integrations.document_stores.qdrant.document_store.qdrant_client.AsyncQdrantClient"
+            ) as client,
+            patch.object(document_store, "_set_up_collection_async", new_callable=AsyncMock),
+        ):
+            await document_store._initialize_async_client()
+
+        assert client.call_args.kwargs["check_compatibility"] is False
+        assert client.call_args.kwargs["port"] == 8333
+
     async def test_query_hybrid_search_batch_failure_async(self):
         document_store = QdrantDocumentStore(location=":memory:", use_sparse_embeddings=True)
         await document_store._initialize_async_client()

integrations/qdrant/tests/test_embedding_retriever.py

@@ -83,6 +83,7 @@ def test_to_dict(self):
                         "write_batch_size": 100,
                         "scroll_size": 10000,
                         "payload_fields_to_index": None,
+                        "client_kwargs": {},
                     },
                 },
                 "filters": None,

integrations/qdrant/tests/test_hybrid_retriever.py

@@ -81,6 +81,7 @@ def test_to_dict(self):
                         "write_batch_size": 100,
                         "scroll_size": 10000,
                         "payload_fields_to_index": None,
+                        "client_kwargs": {},
                     },
                 },
                 "filters": None,

integrations/qdrant/tests/test_sparse_embedding_retriever.py

@@ -82,6 +82,7 @@ def test_to_dict(self):
                         "write_batch_size": 100,
                         "scroll_size": 10000,
                         "payload_fields_to_index": None,
+                        "client_kwargs": {},
                     },
                 },
                 "filters": None,