From 5e9e4dd2334f8187fe4801aea14b8be4235dc364 Mon Sep 17 00:00:00 2001 From: Mark Pollack Date: Wed, 10 Jan 2024 11:40:27 -0500 Subject: [PATCH] Allow PromptTemplate keys to be set multiple times Fixes #97 --- .../ai/prompt/PromptTemplate.java | 33 +++++--- .../ai/prompt/PromptTemplateTest.java | 84 +++++++++++++++++++ 2 files changed, 106 insertions(+), 11 deletions(-) create mode 100644 spring-ai-core/src/test/java/org/springframework/ai/prompt/PromptTemplateTest.java diff --git a/spring-ai-core/src/main/java/org/springframework/ai/prompt/PromptTemplate.java b/spring-ai-core/src/main/java/org/springframework/ai/prompt/PromptTemplate.java index 29cbab10a..6207aefea 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/prompt/PromptTemplate.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/prompt/PromptTemplate.java @@ -79,6 +79,7 @@ public class PromptTemplate implements PromptTemplateActions, PromptTemplateMess this.st = new ST(this.template, '{', '}'); for (Entry 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 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 model) { validate(model); for (Entry 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 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 new file mode 100644 index 000000000..74abea55e --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/prompt/PromptTemplateTest.java @@ -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 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 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 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); + } + +} \ No newline at end of file