From 45a48346fa328c17d0a510abb9de6c6401769daa Mon Sep 17 00:00:00 2001
From: Udbhav Narayan Sharma <udbhav.ds@gmail.com>
Date: Mon, 13 Apr 2026 16:59:53 +0200
Subject: [PATCH 1/2] added PromptTemplateAdapter

---
 .../dspy/prompt_template_adapter.py           | 31 +++++++++++++++++++
 1 file changed, 31 insertions(+)
 create mode 100644 integrations/dspy/src/haystack_integrations/components/generators/dspy/prompt_template_adapter.py

diff --git a/integrations/dspy/src/haystack_integrations/components/generators/dspy/prompt_template_adapter.py b/integrations/dspy/src/haystack_integrations/components/generators/dspy/prompt_template_adapter.py
new file mode 100644
index 0000000000..43052d1b84
--- /dev/null
+++ b/integrations/dspy/src/haystack_integrations/components/generators/dspy/prompt_template_adapter.py
@@ -0,0 +1,31 @@
+class PromptTemplateAdapter:
+    """ Makes DSPy base prompts better with custom model templates, for an enhanced text output. """
+
+    templates = {
+        "mistral": "[INST] {prompt} [/INST]",
+        "llama":   "<s>[INST] {prompt} [/INST]",
+        "chatml":  "<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant",
+        "default": "{prompt}",
+    }
+    # Templates currently made for only 3 models as samples. This can be extended.
+    # (function written below
+
+
+    def __init__(self, model_family: str = "default"):
+        if model_family not in self.templates:
+            self.template = self.templates["default"]
+        else:
+            self.template = self.templates[model_family]
+        
+        self.model_family = model_family
+
+
+    def wrap(self, prompt: str) -> str:
+        return self.template.format(prompt=prompt)
+
+
+    # Makes addition of models and the templates possible.
+    @classmethod
+    def register_template(cls, name: str, template: str) -> None:
+        """Allow users to add custom model families."""
+        cls.templates[name] = template
\ No newline at end of file

From e13d813e3172b2b49fb195ef7c48581c76dcd318 Mon Sep 17 00:00:00 2001
From: Udbhav Narayan Sharma <udbhav.ds@gmail.com>
Date: Mon, 13 Apr 2026 17:29:59 +0200
Subject: [PATCH 2/2] test: add prompt_template_adapter

---
 .../dspy/prompt_template_adapter.py           | 23 +++++++-----------
 .../tests/test_prompt_template_adapter.py     | 24 +++++++++++++++++++
 2 files changed, 33 insertions(+), 14 deletions(-)
 create mode 100644 integrations/dspy/tests/test_prompt_template_adapter.py

diff --git a/integrations/dspy/src/haystack_integrations/components/generators/dspy/prompt_template_adapter.py b/integrations/dspy/src/haystack_integrations/components/generators/dspy/prompt_template_adapter.py
index 43052d1b84..672040c9b4 100644
--- a/integrations/dspy/src/haystack_integrations/components/generators/dspy/prompt_template_adapter.py
+++ b/integrations/dspy/src/haystack_integrations/components/generators/dspy/prompt_template_adapter.py
@@ -1,31 +1,26 @@
+import typing
+
+
 class PromptTemplateAdapter:
-    """ Makes DSPy base prompts better with custom model templates, for an enhanced text output. """
+    """Makes DSPy base prompts better with custom model templates, for an enhanced text output."""
 
-    templates = {
+    templates: typing.ClassVar = {
         "mistral": "[INST] {prompt} [/INST]",
-        "llama":   "<s>[INST] {prompt} [/INST]",
-        "chatml":  "<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant",
+        "llama": "<s>[INST] {prompt} [/INST]",
+        "chatml": "<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant",
         "default": "{prompt}",
     }
     # Templates currently made for only 3 models as samples. This can be extended.
     # (function written below
 
-
     def __init__(self, model_family: str = "default"):
         if model_family not in self.templates:
             self.template = self.templates["default"]
         else:
             self.template = self.templates[model_family]
-        
-        self.model_family = model_family
 
+        self.model_family = model_family
 
     def wrap(self, prompt: str) -> str:
+        """Wqrap the prompt according to the model family template."""
         return self.template.format(prompt=prompt)
-
-
-    # Makes addition of models and the templates possible.
-    @classmethod
-    def register_template(cls, name: str, template: str) -> None:
-        """Allow users to add custom model families."""
-        cls.templates[name] = template
\ No newline at end of file
diff --git a/integrations/dspy/tests/test_prompt_template_adapter.py b/integrations/dspy/tests/test_prompt_template_adapter.py
new file mode 100644
index 0000000000..ee829d3a6a
--- /dev/null
+++ b/integrations/dspy/tests/test_prompt_template_adapter.py
@@ -0,0 +1,24 @@
+import sys
+
+sys.path.insert(0, "src/haystack_integrations/components/generators/dspy")
+from prompt_template_adapter import PromptTemplateAdapter
+
+
+def test_mistral_wraps_correctly():
+    adapter = PromptTemplateAdapter(model_family="mistral")
+    assert adapter.wrap("Hello") == "[INST] Hello [/INST]"
+
+
+def test_llama_wraps_correctly():
+    adapter = PromptTemplateAdapter(model_family="llama")
+    assert adapter.wrap("Hello") == "<s>[INST] Hello [/INST]"
+
+
+def test_default_is_passthrough():
+    adapter = PromptTemplateAdapter()
+    assert adapter.wrap("Hello") == "Hello"
+
+
+def test_unknown_model_falls_back_to_default():
+    adapter = PromptTemplateAdapter(model_family="unknown_model")
+    assert adapter.wrap("Hello") == "Hello"