use proper mustache library to render prompt templates

Author: realark

build.gradle

@@ -66,6 +66,8 @@ dependencies {
   implementation 'org.apache.commons:commons-lang3:3.14.0'
   implementation 'com.google.code.findbugs:jsr305:3.0.2' // for @Nullable annotations
 
+  implementation "com.github.spullara.mustache.java:compiler:0.9.14"
+
   testImplementation "org.slf4j:slf4j-simple:${slf4jVersion}"
   testImplementation "io.opentelemetry:opentelemetry-sdk-testing:${otelVersion}"
   testImplementation "org.junit.jupiter:junit-jupiter:${junitVersion}"

src/main/java/dev/braintrust/instrumentation/openai/BraintrustOpenAI.java

@@ -31,7 +31,7 @@ public static OpenAIClient wrapOpenAI(OpenTelemetry openTelemetry, OpenAIClient
 
     @SneakyThrows
     public static ChatCompletionCreateParams buildChatCompletionsPrompt(
-            BraintrustPrompt prompt, Map<String, String> parameters) {
+            BraintrustPrompt prompt, Map<String, Object> parameters) {
         var promptMap = new HashMap<>(prompt.getOptions());
         promptMap.put("messages", prompt.renderMessages(parameters));
         var promptJson = ObjectMappers.jsonMapper().writeValueAsString(promptMap);

src/main/java/dev/braintrust/prompt/BraintrustPrompt.java

@@ -1,18 +1,17 @@
 package dev.braintrust.prompt;
 
+import com.github.mustachejava.DefaultMustacheFactory;
+import com.github.mustachejava.Mustache;
+import com.github.mustachejava.MustacheException;
 import dev.braintrust.api.BraintrustApiClient;
+import java.io.StringReader;
+import java.io.StringWriter;
 import java.util.ArrayList;
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
-import java.util.Set;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
 
 public class BraintrustPrompt {
-    private static final Pattern MUSTACHE_PATTERN = Pattern.compile("\\{\\{([^}]+)\\}\\}");
-
     private final BraintrustApiClient.Prompt apiPrompt;
     private final Map<String, String> defaults;
 
@@ -25,7 +24,7 @@ public BraintrustPrompt(BraintrustApiClient.Prompt apiPrompt, Map<String, String
         this.defaults = defaults;
     }
 
-    public List<Map<String, Object>> renderMessages(Map<String, String> parameters) {
+    public List<Map<String, Object>> renderMessages(Map<String, Object> parameters) {
         // get promptData->prompt->messages
         Map<String, Object> promptData = (Map<String, Object>) apiPrompt.promptData().prompt();
         List<Map<String, Object>> messages = (List<Map<String, Object>>) promptData.get("messages");
@@ -34,28 +33,20 @@ public List<Map<String, Object>> renderMessages(Map<String, String> parameters)
             throw new RuntimeException("No messages found in prompt data");
         }
 
-        Set<String> usedParameters = new HashSet<>();
         List<Map<String, Object>> renderedMessages = new ArrayList<>();
 
         for (Map<String, Object> message : messages) {
             Map<String, Object> renderedMessage = new HashMap<>(message);
             String content = (String) message.get("content");
 
             if (content != null) {
-                String renderedContent = renderTemplate(content, parameters, usedParameters);
+                String renderedContent = renderTemplate(content, parameters);
                 renderedMessage.put("content", renderedContent);
             }
 
             renderedMessages.add(renderedMessage);
         }
 
-        // Check if all parameters were used
-        Set<String> unusedParameters = new HashSet<>(parameters.keySet());
-        unusedParameters.removeAll(usedParameters);
-        if (!unusedParameters.isEmpty()) {
-            throw new RuntimeException("Unused parameters: " + unusedParameters);
-        }
-
         return renderedMessages;
     }
 
@@ -92,24 +83,22 @@ public Map<String, Object> getOptions() {
         return result;
     }
 
-    private String renderTemplate(
-            String template, Map<String, String> parameters, Set<String> usedParameters) {
-        Matcher matcher = MUSTACHE_PATTERN.matcher(template);
-        StringBuffer result = new StringBuffer();
-
-        while (matcher.find()) {
-            String paramName = matcher.group(1);
-            String paramValue = parameters.get(paramName);
-
-            if (paramValue == null) {
-                throw new RuntimeException("Missing parameter: " + paramName);
+    private String renderTemplate(String template, Map<String, Object> parameters) {
+        try {
+            DefaultMustacheFactory factory = new DefaultMustacheFactory();
+            Mustache mustache = factory.compile(new StringReader(template), "template");
+            StringWriter writer = new StringWriter();
+            mustache.execute(writer, parameters);
+            writer.flush();
+            return writer.toString();
+        } catch (MustacheException e) {
+            // If the template is malformed, just return it as-is
+            return template;
+        } catch (Exception e) {
+            if (e instanceof RuntimeException) {
+                throw (RuntimeException) e;
             }
-
-            usedParameters.add(paramName);
-            matcher.appendReplacement(result, Matcher.quoteReplacement(paramValue));
+            throw new RuntimeException("Failed to render template", e);
         }
-
-        matcher.appendTail(result);
-        return result.toString();
     }
 }

src/test/java/dev/braintrust/instrumentation/openai/BraintrustOpenAITest.java

@@ -498,7 +498,7 @@ void testBuildChatCompletionsPrompt() {
 
         BraintrustPrompt prompt = new BraintrustPrompt(promptObject);
 
-        Map<String, String> parameters = Map.of("name", "Alice");
+        Map<String, Object> parameters = Map.of("name", "Alice");
         ChatCompletionCreateParams renderedParams =
                 BraintrustOpenAI.buildChatCompletionsPrompt(prompt, parameters);
 

src/test/java/dev/braintrust/prompt/BraintrustPromptLoaderTest.java

@@ -47,7 +47,7 @@ void testLoadPromptBySlug() {
         assertNotNull(prompt);
 
         // Test rendering
-        Map<String, String> parameters = Map.of("name", "Bob");
+        Map<String, Object> parameters = Map.of("name", "Bob");
         List<Map<String, Object>> renderedMessages = prompt.renderMessages(parameters);
 
         assertEquals(2, renderedMessages.size());