From e7e3ece3c8ede9b8d73c7bc2bbbec9eac851bda7 Mon Sep 17 00:00:00 2001 From: Madeesh Kannan Date: Wed, 14 Feb 2024 12:02:24 +0100 Subject: [PATCH] refactor!: Qdrant - update secret management (#405) * refactor!: Qdrant - update secret management * Linting --- integrations/qdrant/pyproject.toml | 2 +- .../document_stores/qdrant/document_store.py | 8 ++++++-- integrations/qdrant/tests/test_dict_converters.py | 3 +++ 3 files changed, 10 insertions(+), 3 deletions(-) diff --git a/integrations/qdrant/pyproject.toml b/integrations/qdrant/pyproject.toml index 58a3534c4..1db58ea0d 100644 --- a/integrations/qdrant/pyproject.toml +++ b/integrations/qdrant/pyproject.toml @@ -25,7 +25,7 @@ classifiers = [ "Programming Language :: Python :: Implementation :: CPython", "Programming Language :: Python :: Implementation :: PyPy", ] -dependencies = ["haystack-ai", "qdrant-client"] +dependencies = ["haystack-ai>=2.0.0b6", "qdrant-client"] [project.urls] Source = "https://github.com/deepset-ai/haystack-core-integrations" diff --git a/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py b/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py index 50dd0220c..4a47bf59e 100644 --- a/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py +++ b/integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py @@ -10,6 +10,7 @@ from haystack.dataclasses import Document from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError from haystack.document_stores.types import DuplicatePolicy +from haystack.utils import Secret, deserialize_secrets_inplace from haystack.utils.filters import convert from qdrant_client import grpc from qdrant_client.http import models as rest @@ -55,7 +56,7 @@ def __init__( grpc_port: int = 6334, prefer_grpc: bool = False, # noqa: FBT001, FBT002 https: Optional[bool] = None, - api_key: Optional[str] = None, + api_key: Optional[Secret] = None, prefix: Optional[str] = None, timeout: Optional[float] = None, host: Optional[str] = None, @@ -94,7 +95,7 @@ def __init__( grpc_port=grpc_port, prefer_grpc=prefer_grpc, https=https, - api_key=api_key, + api_key=api_key.resolve_value() if api_key else None, prefix=prefix, timeout=timeout, host=host, @@ -115,6 +116,7 @@ def __init__( self.host = host self.path = path self.metadata = metadata + self.api_key = api_key # Store the Qdrant collection specific attributes self.shard_number = shard_number @@ -232,6 +234,7 @@ def delete_documents(self, ids: List[str]): @classmethod def from_dict(cls, data: Dict[str, Any]) -> "QdrantDocumentStore": + deserialize_secrets_inplace(data["init_parameters"], keys=["api_key"]) return default_from_dict(cls, data) def to_dict(self) -> Dict[str, Any]: @@ -239,6 +242,7 @@ def to_dict(self) -> Dict[str, Any]: # All the __init__ params must be set as attributes # Set as init_parms without default values init_params = {k: getattr(self, k) for k in params} + init_params["api_key"] = self.api_key.to_dict() if self.api_key else None return default_to_dict( self, **init_params, diff --git a/integrations/qdrant/tests/test_dict_converters.py b/integrations/qdrant/tests/test_dict_converters.py index 1c9eb36e2..18940fbbf 100644 --- a/integrations/qdrant/tests/test_dict_converters.py +++ b/integrations/qdrant/tests/test_dict_converters.py @@ -1,3 +1,4 @@ +from haystack.utils import Secret from haystack_integrations.document_stores.qdrant import QdrantDocumentStore @@ -52,6 +53,7 @@ def test_from_dict(): { "type": "haystack_integrations.document_stores.qdrant.document_store.QdrantDocumentStore", "init_parameters": { + "api_key": {"env_vars": ["ENV_VAR"], "strict": False, "type": "env_var"}, "location": ":memory:", "index": "test", "embedding_dim": 768, @@ -98,5 +100,6 @@ def test_from_dict(): document_store.metadata == {}, document_store.write_batch_size == 1000, document_store.scroll_size == 10000, + document_store.api_key == Secret.from_env_var("ENV_VAR", strict=False), ] )