From 2acc41ea8503ceb47fe2d45f8cd637f267367c74 Mon Sep 17 00:00:00 2001
From: Silvano Cerza <3314350+silvanocerza@users.noreply.github.com>
Date: Tue, 5 Sep 2023 12:22:21 +0200
Subject: [PATCH] Add `PromptBuilder` (#5713)

* Add PromptBuilder

* Update release note

* Add test
---
 .../components/generators/prompt_builder.py   | 45 +++++++++++++++++++
 pyproject.toml                                |  1 +
 .../prompt-builder-d22954ef9c4a2a7b.yaml      |  3 ++
 .../generators/test_prompt_builder.py         | 44 ++++++++++++++++++
 4 files changed, 93 insertions(+)
 create mode 100644 haystack/preview/components/generators/prompt_builder.py
 create mode 100644 releasenotes/notes/prompt-builder-d22954ef9c4a2a7b.yaml
 create mode 100644 test/preview/components/generators/test_prompt_builder.py

diff --git a/haystack/preview/components/generators/prompt_builder.py b/haystack/preview/components/generators/prompt_builder.py
new file mode 100644
index 0000000000..26b0ebe078
--- /dev/null
+++ b/haystack/preview/components/generators/prompt_builder.py
@@ -0,0 +1,45 @@
+from typing import Dict, Any
+
+from jinja2 import Template, meta
+
+from haystack.preview import component
+from haystack.preview import default_to_dict, default_from_dict
+
+
+@component
+class PromptBuilder:
+    """
+    PromptBuilder is a component that renders a prompt from a template string using Jinja2 engine.
+    The template variables found in the template string are used as input types for the component and are all required.
+
+    Usage:
+    ```python
+    template = "Translate the following context to {{ target_language }}. Context: {{ snippet }}; Translation:"
+    builder = PromptBuilder(template=template)
+    builder.run(target_language="spanish", snippet="I can't speak spanish.")
+    ```
+    """
+
+    def __init__(self, template: str):
+        """
+        Initialize the component with a template string.
+
+        :param template: Jinja2 template string, e.g. "Summarize this document: {documents}\nSummary:"
+        :type template: str
+        """
+        self._template_string = template
+        self.template = Template(template)
+        ast = self.template.environment.parse(template)
+        template_variables = meta.find_undeclared_variables(ast)
+        component.set_input_types(self, **{var: Any for var in template_variables})
+
+    def to_dict(self) -> Dict[str, Any]:
+        return default_to_dict(self, template=self._template_string)
+
+    @classmethod
+    def from_dict(cls, data) -> "PromptBuilder":
+        return default_from_dict(cls, data)
+
+    @component.output_types(prompt=str)
+    def run(self, **kwargs):
+        return {"prompt": self.template.render(kwargs)}
diff --git a/pyproject.toml b/pyproject.toml
index 2ae5e91e3b..04a24a6fdd 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -80,6 +80,7 @@ dependencies = [
 
   # Preview
   "canals==0.8.0",
+  "Jinja2",
 
   # Agent events
   "events",
diff --git a/releasenotes/notes/prompt-builder-d22954ef9c4a2a7b.yaml b/releasenotes/notes/prompt-builder-d22954ef9c4a2a7b.yaml
new file mode 100644
index 0000000000..2e9561c44e
--- /dev/null
+++ b/releasenotes/notes/prompt-builder-d22954ef9c4a2a7b.yaml
@@ -0,0 +1,3 @@
+---
+preview:
+  - Add PromptBuilder component to render prompts from template strings
diff --git a/test/preview/components/generators/test_prompt_builder.py b/test/preview/components/generators/test_prompt_builder.py
new file mode 100644
index 0000000000..5915fa730d
--- /dev/null
+++ b/test/preview/components/generators/test_prompt_builder.py
@@ -0,0 +1,44 @@
+import pytest
+
+from haystack.preview.components.generators.prompt_builder import PromptBuilder
+
+
+@pytest.mark.unit
+def test_init():
+    builder = PromptBuilder(template="This is a {{ variable }}")
+    assert builder._template_string == "This is a {{ variable }}"
+
+
+@pytest.mark.unit
+def test_to_dict():
+    builder = PromptBuilder(template="This is a {{ variable }}")
+    res = builder.to_dict()
+    assert res == {"type": "PromptBuilder", "init_parameters": {"template": "This is a {{ variable }}"}}
+
+
+@pytest.mark.unit
+def test_from_dict():
+    data = {"type": "PromptBuilder", "init_parameters": {"template": "This is a {{ variable }}"}}
+    builder = PromptBuilder.from_dict(data)
+    builder._template_string == "This is a {{ variable }}"
+
+
+@pytest.mark.unit
+def test_run():
+    builder = PromptBuilder(template="This is a {{ variable }}")
+    res = builder.run(variable="test")
+    assert res == {"prompt": "This is a test"}
+
+
+@pytest.mark.unit
+def test_run_without_input():
+    builder = PromptBuilder(template="This is a template without input")
+    res = builder.run()
+    assert res == {"prompt": "This is a template without input"}
+
+
+@pytest.mark.unit
+def test_run_with_missing_input():
+    builder = PromptBuilder(template="This is a {{ variable }}")
+    res = builder.run()
+    assert res == {"prompt": "This is a "}