Fix PromptLab templating (mlflow#10341)

Signed-off-by: Daniel Lok <[email protected]>

mlflow/_promptlab.py

@@ -1,5 +1,5 @@
 import os
-from string import Template
+import re
 from typing import List
 
 import yaml
@@ -14,7 +14,7 @@ def __init__(self, prompt_template, prompt_parameters, model_parameters, model_r
         self.prompt_parameters = prompt_parameters
         self.model_parameters = model_parameters
         self.model_route = model_route
-        self.prompt_template = Template(prompt_template)
+        self.prompt_template = prompt_template
 
     def predict(self, inputs: pd.DataFrame) -> List[str]:
         from mlflow.gateway import query
@@ -24,7 +24,12 @@ def predict(self, inputs: pd.DataFrame) -> List[str]:
             prompt_parameters_as_dict = {
                 param.key: inputs[param.key][idx] for param in self.prompt_parameters
             }
-            prompt = self.prompt_template.substitute(prompt_parameters_as_dict)
+
+            # copy replacement logic from PromptEngineering.utils.ts for consistency
+            prompt = self.prompt_template
+            for key, value in prompt_parameters_as_dict.items():
+                prompt = re.sub(r"\{\{\s*" + key + r"\s*\}\}", value, prompt)
+
             model_parameters_as_dict = {param.key: param.value for param in self.model_parameters}
             result = query(
                 route=self.model_route, data={"prompt": prompt, **model_parameters_as_dict}

tests/promptlab/test_promptlab_model.py

@@ -0,0 +1,39 @@
+from unittest import mock
+
+import pandas as pd
+
+from mlflow._promptlab import _PromptlabModel
+from mlflow.entities.param import Param
+
+
+def test_promptlab_prompt_replacement():
+    data = pd.DataFrame(
+        data=[
+            {"thing": "books"},
+            {"thing": "coffee"},
+            {"thing": "nothing"},
+        ]
+    )
+
+    prompt_parameters = [Param(key="thing", value="books")]
+    model_parameters = [Param(key="temperature", value=0.5), Param(key="max_tokens", value=10)]
+    prompt_template = "Write me a story about {{ thing }}."
+    model_route = "completions"
+
+    model = _PromptlabModel(prompt_template, prompt_parameters, model_parameters, model_route)
+    with mock.patch("mlflow.gateway.query") as mock_query:
+        model.predict(data)
+
+        calls = [
+            mock.call(
+                route="completions",
+                data={
+                    "prompt": f"Write me a story about {thing}.",
+                    "temperature": 0.5,
+                    "max_tokens": 10,
+                },
+            )
+            for thing in data["thing"]
+        ]
+
+        mock_query.assert_has_calls(calls, any_order=True)