Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Fix PromptTemplate to handle Arrays/Lists

Browse Source 
Fixes #631
This commit is contained in:
Shane Witbeck
2024-04-26 10:58:53 -07:00
committed by Mark Pollack
parent 7252ba19f9
commit 3c40268d80
2 changed files with 83 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -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<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 " + 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);
}
}

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

@@ -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<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 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<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()));