Adapt tests and serialization

Author: bogdankostic

integrations/falkordb/src/haystack_integrations/components/retrievers/falkordb/cypher_retriever.py

@@ -4,7 +4,7 @@
 
 from typing import Any
 
-from haystack import component
+from haystack import component, default_from_dict, default_to_dict
 from haystack.dataclasses import Document
 
 from haystack_integrations.document_stores.falkordb import FalkorDBDocumentStore
@@ -58,6 +58,30 @@ def __init__(
         self.document_store = document_store
         self.custom_cypher_query = custom_cypher_query
 
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serialise the retriever to a dictionary.
+
+        :returns: Dictionary representation of the retriever.
+        """
+        return default_to_dict(
+            self,
+            document_store=self.document_store.to_dict(),
+            custom_cypher_query=self.custom_cypher_query,
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "FalkorDBCypherRetriever":
+        """
+        Deserialise a `FalkorDBCypherRetriever` produced by `to_dict`.
+
+        :param data: Serialised retriever dictionary.
+        :returns: Reconstructed `FalkorDBCypherRetriever` instance.
+        """
+        init_params = data["init_parameters"]
+        init_params["document_store"] = FalkorDBDocumentStore.from_dict(init_params["document_store"])
+        return default_from_dict(cls, data)
+
     @component.output_types(documents=list[Document])
     def run(
         self,

integrations/falkordb/src/haystack_integrations/components/retrievers/falkordb/embedding_retriever.py

@@ -4,7 +4,7 @@
 
 from typing import Any
 
-from haystack import component
+from haystack import component, default_from_dict, default_to_dict
 from haystack.dataclasses import Document
 from haystack.document_stores.types import FilterPolicy, apply_filter_policy
 
@@ -63,6 +63,34 @@ def __init__(
         self.top_k = top_k
         self.filter_policy = FilterPolicy(filter_policy) if isinstance(filter_policy, str) else filter_policy
 
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serialise the retriever to a dictionary.
+
+        :returns: Dictionary representation of the retriever.
+        """
+        return default_to_dict(
+            self,
+            document_store=self.document_store.to_dict(),
+            filters=self.filters,
+            top_k=self.top_k,
+            filter_policy=self.filter_policy.value,
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "FalkorDBEmbeddingRetriever":
+        """
+        Deserialise a `FalkorDBEmbeddingRetriever` produced by `to_dict`.
+
+        :param data: Serialised retriever dictionary.
+        :returns: Reconstructed `FalkorDBEmbeddingRetriever` instance.
+        """
+        init_params = data["init_parameters"]
+        init_params["document_store"] = FalkorDBDocumentStore.from_dict(init_params["document_store"])
+        if "filter_policy" in init_params:
+            init_params["filter_policy"] = FilterPolicy(init_params["filter_policy"])
+        return default_from_dict(cls, data)
+
     @component.output_types(documents=list[Document])
     def run(
         self,

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

@@ -9,11 +9,12 @@
 from datetime import datetime
 from typing import Any, Literal
 
+from haystack import default_from_dict, default_to_dict
 from haystack.dataclasses import Document
 from haystack.document_stores.errors import DocumentStoreError, DuplicateDocumentError
 from haystack.document_stores.types import DocumentStore, DuplicatePolicy
 from haystack.errors import FilterError
-from haystack.utils import Secret
+from haystack.utils import Secret, deserialize_secrets_inplace
 from redis.exceptions import ResponseError
 
 import falkordb  # type: ignore[import-untyped,import-not-found]
@@ -130,6 +131,39 @@ def __init__(
         if verify_connectivity:
             self._ensure_connected()
 
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serialise the store to a dictionary suitable for `from_dict`.
+
+        :returns: Dictionary representation of the store.
+        """
+        return default_to_dict(
+            self,
+            host=self.host,
+            port=self.port,
+            graph_name=self.graph_name,
+            username=self.username,
+            password=self.password.to_dict() if self.password is not None else None,
+            node_label=self.node_label,
+            embedding_dim=self.embedding_dim,
+            embedding_field=self.embedding_field,
+            similarity=self.similarity,
+            write_batch_size=self.write_batch_size,
+            recreate_graph=self.recreate_graph,
+            verify_connectivity=self.verify_connectivity,
+        )
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> FalkorDBDocumentStore:
+        """
+        Deserialise a `FalkorDBDocumentStore` produced by `to_dict`.
+
+        :param data: Serialised store dictionary.
+        :returns: Reconstructed `FalkorDBDocumentStore` instance.
+        """
+        deserialize_secrets_inplace(data["init_parameters"], keys=["password"])
+        return default_from_dict(cls, data)
+
     # ------------------------------------------------------------------
     # Internal connection helpers
     # ------------------------------------------------------------------

integrations/falkordb/tests/test_document_store.py

@@ -6,7 +6,6 @@
 import os
 
 import pytest
-from haystack import default_from_dict, default_to_dict
 from haystack.dataclasses import Document
 from haystack.testing.document_store import DocumentStoreBaseTests
 
@@ -30,28 +29,14 @@ def test_to_dict_from_dict(self):
             verify_connectivity=False,
         )
 
-        data = default_to_dict(
-            store,
-            host=store.host,
-            port=store.port,
-            graph_name=store.graph_name,
-            username=store.username,
-            password=store.password,
-            node_label=store.node_label,
-            embedding_dim=store.embedding_dim,
-            embedding_field=store.embedding_field,
-            similarity=store.similarity,
-            write_batch_size=store.write_batch_size,
-            recreate_graph=store.recreate_graph,
-            verify_connectivity=store.verify_connectivity,
-        )
+        data = store.to_dict()
 
         assert data["init_parameters"]["host"] == "myhost"
         assert data["init_parameters"]["port"] == 1234
         assert data["init_parameters"]["embedding_dim"] == 512
         assert data["init_parameters"]["similarity"] == "euclidean"
 
-        restored = default_from_dict(FalkorDBDocumentStore, data)
+        restored = FalkorDBDocumentStore.from_dict(data)
         assert restored.host == "myhost"
         assert restored.port == 1234
         assert restored.graph_name == "test_graph"