restore to_dict method (#7261)

This commit is contained in:
Massimiliano Pippi 2024-02-29 14:30:06 +01:00 committed by GitHub
parent 14191deefa
commit 25a1a97be0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 14 additions and 2 deletions

View File

@ -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):
"""

View File

@ -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")