Fix PromptTemplate to handle Arrays/Lists

Fixes #631

Author: Shane Witbeck

spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/PromptTemplate.java

@@ -17,10 +17,10 @@
 
 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.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 @@ private String renderResource(Resource resource) {
 		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 Prompt create(Map<String, Object> model) {
 
 	public Set<String> 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<String> 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<String, Object> model) {
+	private Set<String> getModelKeys(Map<String, Object> model) {
 		Set<String> dynamicVariableNames = new HashSet<>(this.dynamicModel.keySet());
 		Set<String> modelVariables = new HashSet<>(model.keySet());
 		modelVariables.addAll(dynamicVariableNames);
-		Set<String> missingEntries = new HashSet<>(getInputVariables());
-		missingEntries.removeAll(modelVariables);
-		if (!missingEntries.isEmpty()) {
+		return modelVariables;
+	}
+
+	protected void validate(Map<String, Object> model) {
+
+		Set<String> templateTokens = getInputVariables();
+		Set<String> 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 " + templateTokens);
+		}
+
+		// Check if the template references any keys not provided by the model
+		if (!templateTokens.containsAll(modelKeys)) {
+			modelKeys.removeAll(templateTokens);
 			throw new IllegalStateException(
-					"All template variables were not replaced. Missing variable names are " + missingEntries);
+					"All model variables were not replaced. Missing variable names are " + modelKeys);
 		}
 	}
 

spring-ai-core/src/test/java/org/springframework/ai/prompt/PromptTemplateTest.java

@@ -17,28 +17,44 @@
 
 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;
 
 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<String> 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<String, Object> model = new HashMap<>();
-		model.put("key1", "value1");
-		model.put("key2", true);
+		Map<String, Object> model = createTestMap();
 		model.put("key3", 100);
 
 		// Create a simple template with placeholders for keys in the generative
@@ -58,14 +74,29 @@ public void testRender() {
 		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<String, Object> 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<String, Object> createTestMap() {
 		Map<String, Object> 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<String, Object> model = createTestMap();
 
 		// Create an input stream for the resource
 		InputStream inputStream = new ByteArrayInputStream("it costs 100".getBytes(Charset.defaultCharset()));