diff --git a/haystack/components/builders/prompt_builder.py b/haystack/components/builders/prompt_builder.py index f6c51af34..f78610b06 100644 --- a/haystack/components/builders/prompt_builder.py +++ b/haystack/components/builders/prompt_builder.py @@ -1,8 +1,8 @@ -from typing import Any +from typing import Any, Dict from jinja2 import Template, meta -from haystack import component +from haystack import component, default_to_dict @component @@ -32,6 +32,9 @@ class PromptBuilder: for var in template_variables: component.set_input_type(self, var, Any, "") + def to_dict(self) -> Dict[str, Any]: + return default_to_dict(self, template=self._template_string) + @component.output_types(prompt=str) def run(self, **kwargs): """ diff --git a/test/components/builders/test_prompt_builder.py b/test/components/builders/test_prompt_builder.py index e08de568c..e1c86f7bd 100644 --- a/test/components/builders/test_prompt_builder.py +++ b/test/components/builders/test_prompt_builder.py @@ -8,6 +8,15 @@ def test_init(): assert builder._template_string == "This is a {{ variable }}" +def test_to_dict(): + builder = PromptBuilder(template="This is a {{ variable }}") + res = builder.to_dict() + assert res == { + "type": "haystack.components.builders.prompt_builder.PromptBuilder", + "init_parameters": {"template": "This is a {{ variable }}"}, + } + + def test_run(): builder = PromptBuilder(template="This is a {{ variable }}") res = builder.run(variable="test")