diff --git a/haystack/preview/document_stores/mixins.py b/haystack/preview/document_stores/mixins.py index 9e4a28ede..853529cc2 100644 --- a/haystack/preview/document_stores/mixins.py +++ b/haystack/preview/document_stores/mixins.py @@ -13,6 +13,10 @@ class StoreAwareMixin: """ _store: Optional[Store] = None + # This is necessary to ease serialisation when converting a Component that uses + # a Store into a dictionary. + # This is only set when calling `Pipeline.add_component()`. + _store_name: str = "" supported_stores: List[Type[Store]] # type: ignore # (see https://github.com/python/mypy/issues/4717) @property diff --git a/haystack/preview/pipeline.py b/haystack/preview/pipeline.py index 9a844b2ec..4d7411c0a 100644 --- a/haystack/preview/pipeline.py +++ b/haystack/preview/pipeline.py @@ -99,6 +99,7 @@ class Pipeline(CanalsPipeline): raise ValueError("Reusing components with stores is not supported (yet). Create a separate instance.") instance.store = self._stores[store] + instance._store_name = store elif store: raise ValueError(f"Component '{name}' doesn't support stores.") diff --git a/test/preview/pipeline/test_pipeline.py b/test/preview/pipeline/test_pipeline.py index 0e7ce5892..506afa4f5 100644 --- a/test/preview/pipeline/test_pipeline.py +++ b/test/preview/pipeline/test_pipeline.py @@ -102,6 +102,7 @@ def test_add_component_store_aware_component_receives_one_docstore(): pipe.add_store(name="second_store", store=store_2) pipe.add_component("component", mock, store="first_store") assert mock.store == store_1 + assert mock._store_name == "first_store" assert pipe.run(data={"component": {"value": 1}}) == {"component": {"value": 1}} @@ -188,6 +189,7 @@ def test_add_component_store_aware_component_receives_correct_docstore_type(): pipe.add_component("component", mock, store="second_store") assert mock.store == store_2 + assert mock._store_name == "second_store" @pytest.mark.unit @@ -217,6 +219,7 @@ def test_add_component_store_aware_component_is_reused(): pipe.add_component("component2", mock, store="first_store") assert mock.store == store_2 + assert mock._store_name == "second_store" @pytest.mark.unit @@ -243,8 +246,9 @@ def test_add_component_store_aware_component_receives_subclass_of_correct_docsto pipe.add_component("component", mock, store="first_store") assert mock.store == store_1 + assert mock._store_name == "first_store" pipe.add_component("component2", mock2, store="second_store") - assert mock2.store == store_2 + assert mock2._store_name == "second_store" @pytest.mark.unit