Repository: instructlab/sdg

Issue:
ConditionalLLMBlock not working with Jinja templates
We had a user discover that during the Jinja template transition we missed updating `ConditionalLLMBlock` to use Jinja templates. So, we need to fix that as well as plug the hole in our test coverage to ensure ConditionalLLMBlock works as well. What happens is we have Jinja Template objects in the code, but are still treating them as the old style python template strings instead of calling the proper Jinja template methods to render the templates.

The commit at https://github.com/instructlab/sdg/commit/c7066c3596aac81c0ad221b981588b211821c401 appears to fix this in manual testing, but we need to add a test case that fails and then passes with that change to ensure this is fixed.

Discussion:

Patch:
diff --git a/src/instructlab/sdg/blocks/llmblock.py b/src/instructlab/sdg/blocks/llmblock.py
index b940954..3350a77 100644
--- a/src/instructlab/sdg/blocks/llmblock.py
+++ b/src/instructlab/sdg/blocks/llmblock.py
@@ -323,14 +323,15 @@ class ConditionalLLMBlock(LLMBlock):
                 )
 
     def _format_prompt(self, sample: Dict) -> str:
+        # If prompt_template is a dict, use the selector column to select the correct template
         if isinstance(self.prompt_template, dict):
             return (
                 self.prompt_template[sample[self.selector_column_name]]
-                .format(**sample)
+                .render(sample)
                 .strip()
             )
-
-        return self.prompt_template.format(**sample).strip()
+        # Otherwise, use the sample to render the prompt without any selection
+        return self.prompt_template.render(sample).strip()
 
     def _validate(self, prompt_template: str, input_dict: Dict[str, Any]) -> bool:
         if isinstance(prompt_template, dict):


Test patch:
diff --git a/tests/test_llmblock.py b/tests/test_llmblock.py
index 2205ab5..6505861 100644
--- a/tests/test_llmblock.py
+++ b/tests/test_llmblock.py
@@ -9,6 +9,7 @@ import unittest
 # Third Party
 from datasets import Dataset, Features, Value
 from httpx import URL
+from jinja2 import StrictUndefined, Template, UndefinedError
 from openai import InternalServerError, NotFoundError
 import pytest
 
@@ -308,6 +309,80 @@ class TestConditionalLLMBlock(unittest.TestCase):
             block.prompt_template, {"selector": "_A_", "var1": "foo", "var2": "bar"}
         )
 
+    def test_format_prompt_with_jinja_templates(self, mock_load_config):
+        mock_load_config.return_value = {
+            "system": "{{var1}} {{var2}}",
+            "introduction": "introduction",
+            "principles": "principles",
+            "examples": "examples",
+            "generation": "generation",
+        }
+        # Create an instance of the block
+        block = ConditionalLLMBlock(
+            ctx=self.mock_ctx,
+            pipe=self.mock_pipe,
+            block_name="gen_knowledge",
+            config_paths=[["/foo/bar", "_A_"]],
+            output_cols=[],
+            selector_column_name="selector",
+        )
+
+        # Mock prompt_template as a Jinja template
+        block.prompt_template = {
+            "case1": Template("Hello, {{ name }}!", undefined=StrictUndefined),
+            "case2": Template("Goodbye, {{ name }}.", undefined=StrictUndefined),
+        }
+
+        # Sample input
+        sample1 = {"selector": "case1", "name": "Alice"}
+        sample2 = {"selector": "case2", "name": "Bob"}
+
+        # Test case 1
+        result1 = block._format_prompt(sample1)
+        self.assertEqual(result1, "Hello, Alice!")
+
+        # Test case 2
+        result2 = block._format_prompt(sample2)
+        self.assertEqual(result2, "Goodbye, Bob.")
+
+        # Verify error with missing keys in sample
+        with self.assertRaises(UndefinedError):
+            block._format_prompt({"selector": "case1"})  # 'name' is missing
+
+    def test_format_prompt_without_selector(self, mock_load_config):
+        mock_load_config.return_value = {
+            "system": "{{var1}} {{var2}}",
+            "introduction": "introduction",
+            "principles": "principles",
+            "examples": "examples",
+            "generation": "generation",
+        }
+        # Create an instance of the block
+        block = ConditionalLLMBlock(
+            ctx=self.mock_ctx,
+            pipe=self.mock_pipe,
+            block_name="gen_knowledge",
+            config_paths=[["/foo/bar", "_A_"]],
+            output_cols=[],
+            selector_column_name=None,  # No selector
+        )
+
+        # Mock prompt_template as a Jinja template
+        block.prompt_template = Template(
+            "Welcome, {{ user }}!", undefined=StrictUndefined
+        )
+
+        # Sample input
+        sample = {"user": "Connor"}
+
+        # Test
+        result = block._format_prompt(sample)
+        self.assertEqual(result, "Welcome, Connor!")
+
+        # Verify error with missing keys in sample
+        with self.assertRaises(UndefinedError):
+            block._format_prompt({})  # 'user' is missing
+
 
 @patch("src.instructlab.sdg.blocks.block.Block._load_config")
 class TestLLMLogProbBlock(unittest.TestCase):