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

Allow PromptTemplate keys to be set multiple times

Browse Source 
Fixes #97
This commit is contained in:
Mark Pollack
2024-01-10 11:40:27 -05:00
parent 0b87df07e5
commit 5e9e4dd233
2 changed files with 106 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -79,6 +79,7 @@ public class PromptTemplate implements PromptTemplateActions, PromptTemplateMess
this.st = new ST(this.template, '{', '}');
for (Entry<String, Object> entry : model.entrySet()) {
add(entry.getKey(), entry.getValue());
dynamicModel.put(entry.getKey(), entry.getValue());
}
}
catch (Exception ex) {
@@ -98,6 +99,7 @@ public class PromptTemplate implements PromptTemplateActions, PromptTemplateMess
this.st = new ST(this.template, '{', '}');
for (Entry<String, Object> entry : model.entrySet()) {
add(entry.getKey(), entry.getValue());
dynamicModel.put(entry.getKey(), entry.getValue());
}
}
catch (Exception ex) {
@@ -130,6 +132,7 @@ public class PromptTemplate implements PromptTemplateActions, PromptTemplateMess
// Render Methods
@Override
public String render() {
validate(this.dynamicModel);
return st.render();
}
@@ -137,25 +140,33 @@ public class PromptTemplate implements PromptTemplateActions, PromptTemplateMess
public String render(Map<String, Object> model) {
validate(model);
for (Entry<String, Object> entry : model.entrySet()) {
if (st.getAttribute(entry.getKey()) == null) {
if (entry.getValue() instanceof Resource) {
st.add(entry.getKey(), renderResource((Resource) entry.getValue()));
}
else {
st.add(entry.getKey(), entry.getValue());
}
if (st.getAttribute(entry.getKey()) != null) {
st.remove(entry.getKey());
}
if (entry.getValue() instanceof Resource) {
st.add(entry.getKey(), renderResource((Resource) entry.getValue()));
}
else {
st.add(entry.getKey(), entry.getValue());
}
}
return st.render();
}
private String renderResource(Resource resource) {
try (InputStream inputStream = resource.getInputStream()) {
return StreamUtils.copyToString(inputStream, Charset.defaultCharset());
try {
return resource.getContentAsString(Charset.defaultCharset());
}
catch (IOException ex) {
throw new RuntimeException(ex);
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

84
spring-ai-core/src/test/java/org/springframework/ai/prompt/PromptTemplateTest.java Normal file
View File

@@ -0,0 +1,84 @@
package org.springframework.ai.prompt;
import org.junit.jupiter.api.Disabled;
import org.junit.jupiter.api.Test;
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.HashMap;
import java.util.Map;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
public class PromptTemplateTest {
@Test
public void testRender() {
// Create a map with string keys and object values to serve as a model for testing
Map<String, Object> model = new HashMap<>();
model.put("key1", "value1");
model.put("key2", true);
model.put("key3", 100);
// Create a simple template with placeholders for keys in the model
String template = "This is a {key1}, it is {key2}, and it costs {key3}";
PromptTemplate promptTemplate = new PromptTemplate(template, model);
// The expected result after rendering the template with the model
String expected = "This is a value1, it is true, and it costs 100";
String result = promptTemplate.render();
// Check that the rendered string matches the expected result
assertEquals(expected, result);
model.put("key3", 200);
expected = "This is a value1, it is true, and it costs 200";
result = promptTemplate.render(model);
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 model for testing
Map<String, Object> model = new HashMap<>();
model.put("key1", "value1");
model.put("key2", true);
// Create an input stream for the resource
InputStream inputStream = new ByteArrayInputStream("it costs 100".getBytes(Charset.defaultCharset()));
Resource resource = new InputStreamResource(inputStream);
model.put("key3", resource);
// Create a simple template with placeholders for keys in the model
String template = "{key1}, {key2}, {key3}";
PromptTemplate promptTemplate = new PromptTemplate(template, model);
// The expected result after rendering the template with the model
String expected = "value1, true, it costs 100";
String result = promptTemplate.render();
// Check that the rendered string matches the expected result
assertEquals(expected, result);
}
@Test
public void testRenderFailure() {
// Create a map with string keys and object values to serve as a model for testing
Map<String, Object> model = new HashMap<>();
model.put("key1", "value1");
// Create a simple template that includes a key not present in the model
String template = "This is a {key2}!";
PromptTemplate promptTemplate = new PromptTemplate(template, model);
// Rendering the template with a missing key should throw an exception
assertThrows(IllegalStateException.class, promptTemplate::render);
}
}