From 3c40268d8022cf718de8850097bbe1f324451fce Mon Sep 17 00:00:00 2001 From: Shane Witbeck Date: Fri, 26 Apr 2024 10:58:53 -0700 Subject: [PATCH] Fix PromptTemplate to handle Arrays/Lists Fixes #631 --- .../ai/chat/prompt/PromptTemplate.java | 62 +++++++++++++------ .../ai/prompt/PromptTemplateTest.java | 49 ++++++++++++--- 2 files changed, 83 insertions(+), 28 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/PromptTemplate.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/PromptTemplate.java index 6e6155fe6..a8915214a 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/PromptTemplate.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/PromptTemplate.java @@ -17,10 +17,10 @@ package org.springframework.ai.chat.prompt; import org.antlr.runtime.Token; import org.antlr.runtime.TokenStream; -import org.springframework.ai.parser.OutputParser; import org.springframework.ai.chat.messages.Media; import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.parser.OutputParser; import org.springframework.core.io.Resource; import org.springframework.util.StreamUtils; import org.stringtemplate.v4.ST; @@ -31,8 +31,6 @@ import java.io.InputStream; import java.nio.charset.Charset; import java.util.*; import java.util.Map.Entry; -import java.util.stream.Collectors; -import java.util.stream.IntStream; public class PromptTemplate implements PromptTemplateActions, PromptTemplateMessageActions { @@ -161,12 +159,6 @@ public class PromptTemplate implements PromptTemplateActions, PromptTemplateMess catch (IOException e) { throw new RuntimeException(e); } - // try (InputStream inputStream = resource.getInputStream()) { - // return StreamUtils.copyToString(inputStream, Charset.defaultCharset()); - // } - // catch (IOException ex) { - // throw new RuntimeException(ex); - // } } @Override @@ -196,22 +188,54 @@ public class PromptTemplate implements PromptTemplateActions, PromptTemplateMess public Set getInputVariables() { TokenStream tokens = this.st.impl.tokens; - return IntStream.range(0, tokens.range()) - .mapToObj(tokens::get) - .filter(token -> token.getType() == STLexer.ID) - .map(Token::getText) - .collect(Collectors.toSet()); + Set inputVariables = new HashSet<>(); + boolean isInsideList = false; + + for (int i = 0; i < tokens.size(); i++) { + Token token = tokens.get(i); + + if (token.getType() == STLexer.LDELIM && i + 1 < tokens.size() + && tokens.get(i + 1).getType() == STLexer.ID) { + if (i + 2 < tokens.size() && tokens.get(i + 2).getType() == STLexer.COLON) { + inputVariables.add(tokens.get(i + 1).getText()); + isInsideList = true; + } + } + else if (token.getType() == STLexer.RDELIM) { + isInsideList = false; + } + else if (!isInsideList && token.getType() == STLexer.ID) { + inputVariables.add(token.getText()); + } + } + + return inputVariables; } - protected void validate(Map model) { + private Set getModelKeys(Map model) { Set dynamicVariableNames = new HashSet<>(this.dynamicModel.keySet()); Set modelVariables = new HashSet<>(model.keySet()); modelVariables.addAll(dynamicVariableNames); - Set missingEntries = new HashSet<>(getInputVariables()); - missingEntries.removeAll(modelVariables); - if (!missingEntries.isEmpty()) { + return modelVariables; + } + + protected void validate(Map model) { + + Set templateTokens = getInputVariables(); + Set modelKeys = getModelKeys(model); + + // Check if model provides all keys required by the template + if (!modelKeys.containsAll(templateTokens)) { + templateTokens.removeAll(modelKeys); throw new IllegalStateException( - "All template variables were not replaced. Missing variable names are " + missingEntries); + "All template variables were not replaced. Missing variable names are " + templateTokens); + } + + // Check if the template references any keys not provided by the model + if (!templateTokens.containsAll(modelKeys)) { + modelKeys.removeAll(templateTokens); + throw new IllegalStateException( + "All model variables were not replaced. Missing variable names are " + modelKeys); } } diff --git a/spring-ai-core/src/test/java/org/springframework/ai/prompt/PromptTemplateTest.java b/spring-ai-core/src/test/java/org/springframework/ai/prompt/PromptTemplateTest.java index b60fb46a9..d2e74c2e3 100644 --- a/spring-ai-core/src/test/java/org/springframework/ai/prompt/PromptTemplateTest.java +++ b/spring-ai-core/src/test/java/org/springframework/ai/prompt/PromptTemplateTest.java @@ -17,6 +17,7 @@ package org.springframework.ai.prompt; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; +import org.springframework.ai.chat.messages.Message; import org.springframework.ai.chat.prompt.PromptTemplate; import org.springframework.core.io.InputStreamResource; import org.springframework.core.io.Resource; @@ -24,21 +25,36 @@ import org.springframework.core.io.Resource; import java.io.ByteArrayInputStream; import java.io.InputStream; import java.nio.charset.Charset; +import java.util.Arrays; import java.util.HashMap; +import java.util.List; import java.util.Map; +import static org.assertj.core.api.Assertions.assertThatExceptionOfType; import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; public class PromptTemplateTest { + @Test + public void testRenderWithList() { + String templateString = "The items are:\n{items:{item | - {item}\n}}"; + List itemList = Arrays.asList("apple", "banana", "cherry"); + PromptTemplate promptTemplate = new PromptTemplate(templateString); + Message message = promptTemplate.createMessage(Map.of("items", itemList)); + + String expected = "The items are:\n" + "- apple\n" + "- banana\n" + "- cherry\n"; + + assertEquals(expected, message.getContent()); + + PromptTemplate unfilledPromptTemplate = new PromptTemplate(templateString); + assertThatExceptionOfType(IllegalStateException.class).isThrownBy(unfilledPromptTemplate::render) + .withMessage("All template variables were not replaced. Missing variable names are [items]"); + } + @Test public void testRender() { - // Create a map with string keys and object values to serve as a generative for - // testing - Map model = new HashMap<>(); - model.put("key1", "value1"); - model.put("key2", true); + Map model = createTestMap(); model.put("key3", 100); // Create a simple template with placeholders for keys in the generative @@ -58,14 +74,29 @@ public class PromptTemplateTest { assertEquals(expected, result); } - @Disabled("Need to improve PromptTemplate to better handle Resource toString and tracking with 'dynamicModel' for underlying StringTemplate") @Test - public void testRenderResource() throws Exception { - // Create a map with string keys and object values to serve as a generative for - // testing + public void testRenderResource() { + Map model = createTestMap(); + InputStream inputStream = new ByteArrayInputStream( + "key1's value is {key1} and key2's value is {key2}".getBytes(Charset.defaultCharset())); + Resource resource = new InputStreamResource(inputStream); + PromptTemplate promptTemplate = new PromptTemplate(resource, model); + String expected = "key1's value is value1 and key2's value is true"; + String result = promptTemplate.render(); + assertEquals(expected, result); + } + + private static Map createTestMap() { Map model = new HashMap<>(); model.put("key1", "value1"); model.put("key2", true); + return model; + } + + @Disabled("Need to improve PromptTemplate to better handle Resource toString and tracking with 'dynamicModel' for underlying StringTemplate") + @Test + public void testRenderResourceAsValue() throws Exception { + Map model = createTestMap(); // Create an input stream for the resource InputStream inputStream = new ByteArrayInputStream("it costs 100".getBytes(Charset.defaultCharset()));