From d4b9d071ba80bb5f28aa44445980407d49d1b56c Mon Sep 17 00:00:00 2001 From: Mark Pollack Date: Tue, 15 Aug 2023 23:22:32 -0400 Subject: [PATCH] Add toString to Message classes, and allow Resource to be used in PromptTemplate --- .../org/springframework/ai/prompt/Prompt.java | 11 +++-- .../ai/prompt/PromptTemplate.java | 42 +++++++++++++++++-- .../ai/prompt/messages/AssistantMessage.java | 6 +++ .../ai/prompt/messages/FunctionMessage.java | 6 +++ .../ai/prompt/messages/SystemMessage.java | 6 +++ .../ai/prompt/messages/UserMessage.java | 6 +++ 6 files changed, 70 insertions(+), 7 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/prompt/Prompt.java b/spring-ai-core/src/main/java/org/springframework/ai/prompt/Prompt.java index 475d6e5e3..d157af908 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/prompt/Prompt.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/prompt/Prompt.java @@ -16,12 +16,12 @@ package org.springframework.ai.prompt; +import org.springframework.ai.prompt.messages.Message; +import org.springframework.ai.prompt.messages.UserMessage; + import java.util.Collections; import java.util.List; -import org.springframework.ai.prompt.messages.UserMessage; -import org.springframework.ai.prompt.messages.Message; - public class Prompt { private List messages; @@ -50,4 +50,9 @@ public class Prompt { return this.messages; } + @Override + public String toString() { + return "Prompt{" + "messages=" + messages + '}'; + } + } 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 2c8a7a31d..5bfa134f5 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 @@ -20,9 +20,14 @@ import org.antlr.runtime.Token; import org.antlr.runtime.TokenStream; import org.springframework.ai.prompt.messages.Message; import org.springframework.ai.prompt.messages.UserMessage; +import org.springframework.core.io.Resource; +import org.springframework.util.StreamUtils; import org.stringtemplate.v4.ST; import org.stringtemplate.v4.compiler.STLexer; +import java.io.IOException; +import java.io.InputStream; +import java.nio.charset.Charset; import java.util.*; import java.util.Map.Entry; import java.util.stream.Collectors; @@ -40,6 +45,21 @@ public class PromptTemplate implements PromptTemplateActions { private OutputParser outputParser; + public PromptTemplate(Resource resource) { + try (InputStream inputStream = resource.getInputStream()) { + this.template = StreamUtils.copyToString(inputStream, Charset.defaultCharset()); + } + catch (IOException ex) { + throw new RuntimeException("Failed to read resource", ex); + } + try { + this.st = new ST(this.template, '{', '}'); + } + catch (Exception ex) { + throw new IllegalArgumentException("The template string is not valid.", ex); + } + } + public PromptTemplate(String template) { this.template = template; // If the template string is not valid, an exception will be thrown @@ -95,12 +115,26 @@ public class PromptTemplate implements PromptTemplateActions { @Override public String render(Map model) { validate(model); - for (Entry stringObjectEntry : model.entrySet()) { - if (st.getAttribute(stringObjectEntry.getKey()) == null) { - st.add(stringObjectEntry.getKey(), stringObjectEntry.getValue()); + 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()); + } } } - return st.render().trim(); + return st.render(); + } + + private String renderResource(Resource resource) { + 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/main/java/org/springframework/ai/prompt/messages/AssistantMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/AssistantMessage.java index 26682ebe6..2e064bd65 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/AssistantMessage.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/AssistantMessage.java @@ -34,4 +34,10 @@ public class AssistantMessage extends AbstractMessage { super(MessageType.ASSISTANT, content, properties); } + @Override + public String toString() { + return "AssistantMessage{" + "content='" + content + '\'' + ", properties=" + properties + ", messageType=" + + messageType + '}'; + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/FunctionMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/FunctionMessage.java index 5ffc7701a..c00520dcd 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/FunctionMessage.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/FunctionMessage.java @@ -28,4 +28,10 @@ public class FunctionMessage extends AbstractMessage { super(MessageType.SYSTEM, content, properties); } + @Override + public String toString() { + return "FunctionMessage{" + "content='" + content + '\'' + ", properties=" + properties + ", messageType=" + + messageType + '}'; + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/SystemMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/SystemMessage.java index a3802f52a..a9e2b0e3c 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/SystemMessage.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/SystemMessage.java @@ -34,4 +34,10 @@ public class SystemMessage extends AbstractMessage { super(MessageType.SYSTEM, content, properties); } + @Override + public String toString() { + return "SystemMessage{" + "content='" + content + '\'' + ", properties=" + properties + ", messageType=" + + messageType + '}'; + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/UserMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/UserMessage.java index 45850f441..106f6ce66 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/UserMessage.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/prompt/messages/UserMessage.java @@ -33,4 +33,10 @@ public class UserMessage extends AbstractMessage { super(MessageType.USER, message, properties); } + @Override + public String toString() { + return "UserMessage{" + "content='" + content + '\'' + ", properties=" + properties + ", messageType=" + + messageType + '}'; + } + }