test: add prompt_template_adapter

Author: udbhav4

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

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"