From 1b3d6a55dd50b5bbc277fabac37895ff8da37f6b Mon Sep 17 00:00:00 2001 From: Mark Pollack Date: Fri, 28 Jul 2023 14:38:15 -0400 Subject: [PATCH] Improve Prompt package design --- .../core/prompts/AbstractPromptTemplate.java | 25 +-- .../{StringPromptValue.java => Prompt.java} | 26 ++- .../ai/core/prompts/PromptBuilder.java | 37 ++++ .../ai/core/prompts/PromptInput.java | 42 ---- ...mplateInput.java => PromptOperations.java} | 14 +- .../ai/core/prompts/PromptTemplate.java | 188 +++++++++++++----- .../ai/core/prompts/PromptValue.java | 29 --- .../ai/core/prompts/TemplateFormat.java | 2 +- .../ai/core/prompts/messages/AiMessage.java | 14 +- .../ai/core/prompts/messages/ChatMessage.java | 2 +- .../prompts/messages/FunctionMessage.java | 2 +- .../ai/core/prompts/messages/MessageType.java | 2 +- .../ai/core/prompts/PromptTests.java | 81 +++++++- 13 files changed, 304 insertions(+), 160 deletions(-) rename spring-ai-core/src/main/java/org/springframework/ai/core/prompts/{StringPromptValue.java => Prompt.java} (64%) create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptBuilder.java delete mode 100644 spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptInput.java rename spring-ai-core/src/main/java/org/springframework/ai/core/prompts/{PromptTemplateInput.java => PromptOperations.java} (79%) delete mode 100644 spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptValue.java diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/AbstractPromptTemplate.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/AbstractPromptTemplate.java index 754ee3776..d577af644 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/AbstractPromptTemplate.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/AbstractPromptTemplate.java @@ -16,29 +16,24 @@ package org.springframework.ai.core.prompts; -import java.util.Map; -import java.util.Optional; +public abstract class AbstractPromptTemplate implements PromptOperations { -public abstract class AbstractPromptTemplate implements PromptInput { + protected String template; - private Optional outputParser = Optional.empty(); + protected TemplateFormat templateFormat = TemplateFormat.ST; - public AbstractPromptTemplate() { - this.outputParser = Optional.empty(); - } - - public AbstractPromptTemplate(OutputParser outputParser) { - this.outputParser = Optional.of(outputParser); + public AbstractPromptTemplate(String template) { + this.template = template; } @Override - public Optional getOutputParser() { - return this.outputParser; + public String getTemplate() { + return this.template; } - public PromptValue formatAsPrompt(Map inputVariables) { - String formattedPrompt = formatAsString(inputVariables); - return new StringPromptValue(formattedPrompt); + @Override + public TemplateFormat getTemplateFormat() { + return this.templateFormat; } } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/StringPromptValue.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/Prompt.java similarity index 64% rename from spring-ai-core/src/main/java/org/springframework/ai/core/prompts/StringPromptValue.java rename to spring-ai-core/src/main/java/org/springframework/ai/core/prompts/Prompt.java index 56007804f..36add1862 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/StringPromptValue.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/Prompt.java @@ -22,22 +22,28 @@ import java.util.List; import org.springframework.ai.core.prompts.messages.HumanMessage; import org.springframework.ai.core.prompts.messages.Message; -public class StringPromptValue implements PromptValue { +public class Prompt { - private String value; + private List messages; - public StringPromptValue(String formattedPrompt) { - this.value = formattedPrompt; + public Prompt(String contents) { + this.messages = Collections.singletonList(new HumanMessage(contents)); } - @Override - public String toStringValue() { - return this.value; + public Prompt(List messages) { + this.messages = messages; } - @Override - public List toMessages() { - return Collections.singletonList(new HumanMessage(this.value)); + public String getContents() { + StringBuilder sb = new StringBuilder(); + for (Message message : messages) { + sb.append(message.getContent()); + } + return sb.toString(); + } + + public List getMessages() { + return this.messages; } } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptBuilder.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptBuilder.java new file mode 100644 index 000000000..1768da711 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptBuilder.java @@ -0,0 +1,37 @@ +package org.springframework.ai.core.prompts;/* + * Copyright 2023 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import java.util.Map; + +public interface PromptBuilder { + + PromptBuilder system(); + + PromptBuilder human(); + + PromptBuilder ai(boolean containsExample); + + PromptBuilder chat(String chatRole); + + PromptBuilder function(String functionName); + + PromptBuilder usingModel(Map model); + + PromptBuilder withProperties(Map properties); + + Prompt create(); + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptInput.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptInput.java deleted file mode 100644 index d3741c4a6..000000000 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptInput.java +++ /dev/null @@ -1,42 +0,0 @@ -/* - * Copyright 2023 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.springframework.ai.core.prompts; - -import java.util.Map; -import java.util.Optional; - -public interface PromptInput { - - // *Output* - - Optional getOutputParser(); - - - // *Input* - - // This is the handoff point. These methods provide the "input", then the - // "Template" is "rendered" and the output is then used to construct a "Message" that gets sent to the "LLM" model. - - // Maybe should be called String renderAsString() and renderAsPrompt - // View in spring mvc has render(Map model, HttpServletRequest request, HttpServletResponse response) method. - String formatAsString(Map inputVariables); - - PromptValue formatAsPrompt(Map inputVariables); - - // Leave out Partial Input Variables for now - -} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptTemplateInput.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptOperations.java similarity index 79% rename from spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptTemplateInput.java rename to spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptOperations.java index 92a12e16c..8b1c6b70c 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptTemplateInput.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptOperations.java @@ -16,12 +16,20 @@ package org.springframework.ai.core.prompts; -public interface PromptTemplateInput extends PromptInput { +import java.util.Map; + +public interface PromptOperations { String getTemplate(); TemplateFormat getTemplateFormat(); - // *Validation* - void validate(); + void add(String name, Object value); + + String render(); + + String render(Map model); + + PromptBuilder prompt(); + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptTemplate.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptTemplate.java index 17409ed26..75ad6f99b 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptTemplate.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptTemplate.java @@ -16,9 +16,11 @@ package org.springframework.ai.core.prompts; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; import java.util.Map; import java.util.Map.Entry; -import java.util.Optional; import java.util.Set; import java.util.stream.Collectors; import java.util.stream.IntStream; @@ -28,59 +30,55 @@ import org.antlr.runtime.TokenStream; import org.stringtemplate.v4.ST; import org.stringtemplate.v4.compiler.STLexer; -public class PromptTemplate extends AbstractPromptTemplate implements PromptTemplateInput { +import org.springframework.ai.core.prompts.messages.AiMessage; +import org.springframework.ai.core.prompts.messages.ChatMessage; +import org.springframework.ai.core.prompts.messages.FunctionMessage; +import org.springframework.ai.core.prompts.messages.HumanMessage; +import org.springframework.ai.core.prompts.messages.Message; +import org.springframework.ai.core.prompts.messages.MessageType; +import org.springframework.ai.core.prompts.messages.SystemMessage; - private String template; +public class PromptTemplate extends AbstractPromptTemplate { - private TemplateFormat templateFormat = TemplateFormat.FSTRING; + private ST st; + + private Map dynamicModel = new HashMap<>(); public PromptTemplate(String template) { - super(); - this.template = template; - } - - public PromptTemplate(String template, TemplateFormat templateFormat) { - super(); - this.template = template; - this.templateFormat = templateFormat; - } - - - public PromptTemplate(String template, OutputParser outputParser) { - super(outputParser); - this.template = template; - } - - public PromptTemplate(String template, OutputParser outputParser, TemplateFormat templateFormat) { - super(outputParser); - this.template = template; - this.templateFormat = templateFormat; + super(template); + // If the template string is not valid, an exception will be thrown + try { + this.st = new ST(this.template, '{', '}'); + } + catch (Exception ex) { + throw new IllegalArgumentException("The template string is not valid.", ex); + } } @Override - public String formatAsString(Map inputVariables) { - validate(); - // Only "F-String" for now - ST st = new ST(this.template, '{', '}'); - for (Entry stringObjectEntry : inputVariables.entrySet()) { - st.add(stringObjectEntry.getKey(), stringObjectEntry.getValue()); - } + public void add(String name, Object value) { + this.st.add(name, value); + this.dynamicModel.put(name, value); + } + + // Render Methods + public String render() { return st.render(); } @Override - public String getTemplate() { - return this.template; - } - - @Override - public TemplateFormat getTemplateFormat() { - return this.templateFormat; + public String render(Map model) { + validate(model); + for (Entry stringObjectEntry : model.entrySet()) { + if (st.getAttribute(stringObjectEntry.getKey()) == null) { + st.add(stringObjectEntry.getKey(), stringObjectEntry.getValue()); + } + } + return st.render().trim(); } protected Set getInputVariables() { - ST st = new ST(this.template, '{', '}'); - TokenStream tokens = st.impl.tokens; + TokenStream tokens = this.st.impl.tokens; return IntStream.range(0, tokens.range()) .mapToObj(tokens::get) .filter(token -> token.getType() == STLexer.ID) @@ -88,19 +86,105 @@ public class PromptTemplate extends AbstractPromptTemplate implements PromptTemp .collect(Collectors.toSet()); } - public void validate() { - try { - ST st = new ST(this.template, '{', '}'); - // TODO is doing this test even necessary, if it parsed correctly in the ctor, there should be no issues. - Set inputVariables = getInputVariables(); - for (String inputVariable : inputVariables) { - st.add(inputVariable, "foo"); - } - st.render(); - } - catch (Exception ex) { - throw new IllegalArgumentException("The template string is not valid.", ex); + protected void validate(Map model) { + Set dynamicVariableNames = new HashSet<>(this.dynamicModel.keySet()); + Set modelVariables = new HashSet<>(model.keySet()); + modelVariables.addAll(dynamicVariableNames); + Set missingEntries = new HashSet<>(getInputVariables()); + missingEntries.removeAll(modelVariables); + if (!missingEntries.isEmpty()) { + throw new IllegalStateException( + "All template variables were not replaced. Missing variable names are " + missingEntries); } } + @Override + public PromptBuilder prompt() { + return new PromptTemplatePromptBuilder(); + } + + public class PromptTemplatePromptBuilder implements PromptBuilder { + + private MessageType messageType = MessageType.HUMAN; + + private Map model = new HashMap<>(); + + private Map properties = new HashMap<>(); + + private boolean containsExample; + + private String chatRole; + + private String functionName; + + @Override + public PromptBuilder system() { + this.messageType = MessageType.SYSTEM; + return this; + } + + @Override + public PromptBuilder human() { + this.messageType = MessageType.HUMAN; + return this; + } + + @Override + public PromptBuilder ai(boolean containsExample) { + this.messageType = MessageType.AI; + this.containsExample = containsExample; + return this; + } + + @Override + public PromptBuilder chat(String chatRole) { + this.messageType = MessageType.CHAT; + this.chatRole = chatRole; + return this; + } + + @Override + public PromptBuilder function(String functionName) { + this.messageType = MessageType.FUNCTION; + this.functionName = functionName; + return this; + } + + @Override + public PromptBuilder usingModel(Map model) { + this.model = model; + return this; + } + + @Override + public PromptBuilder withProperties(Map properties) { + this.properties = properties; + return this; + } + + @Override + public Prompt create() { + + switch (messageType) { + case HUMAN: + return newPrompt(new HumanMessage(render(model), properties)); + case AI: + return newPrompt(new AiMessage(render(model), containsExample, properties)); + case CHAT: + return newPrompt(new ChatMessage(render(model), chatRole, properties)); + case SYSTEM: + return newPrompt(new SystemMessage(render(model), properties)); + case FUNCTION: + return newPrompt(new FunctionMessage(render(model), functionName, properties)); + default: + throw new IllegalArgumentException("Invalid MessageType: " + messageType); + } + } + + private Prompt newPrompt(Message message) { + return new Prompt(Collections.singletonList(message)); + } + + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptValue.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptValue.java deleted file mode 100644 index 8743818f7..000000000 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptValue.java +++ /dev/null @@ -1,29 +0,0 @@ -/* - * Copyright 2023 the original author or authors. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package org.springframework.ai.core.prompts; - -import java.util.List; - -import org.springframework.ai.core.prompts.messages.Message; - -public interface PromptValue { - - String toStringValue(); - - List toMessages(); - -} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/TemplateFormat.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/TemplateFormat.java index daa8a240d..59ef75cc2 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/TemplateFormat.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/TemplateFormat.java @@ -18,7 +18,7 @@ package org.springframework.ai.core.prompts; public enum TemplateFormat { - FSTRING("f-string"); + ST("ST"); private final String value; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/AiMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/AiMessage.java index f9328071e..19bac0137 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/AiMessage.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/AiMessage.java @@ -20,12 +20,24 @@ import java.util.Map; public class AiMessage extends AbstractMessage { + private boolean example = false; + public AiMessage(String content) { super(MessageType.AI, content); } - public AiMessage(String content, Map properties) { + public AiMessage(String content, boolean example) { + super(MessageType.AI, content); + this.example = example; + } + + public AiMessage(String content, boolean example, Map properties) { super(MessageType.AI, content, properties); + this.example = example; + } + + public boolean isExample() { + return example; } } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/ChatMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/ChatMessage.java index 7e4818bc3..d656b3714 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/ChatMessage.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/ChatMessage.java @@ -27,7 +27,7 @@ public class ChatMessage extends AbstractMessage { this.role = role; } - public ChatMessage(String content, Map properties, String role) { + public ChatMessage(String content, String role, Map properties) { super(MessageType.SYSTEM, content, properties); this.role = role; } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/FunctionMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/FunctionMessage.java index 0ca950204..9cd15d1f5 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/FunctionMessage.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/FunctionMessage.java @@ -27,7 +27,7 @@ public class FunctionMessage extends AbstractMessage { this.functionName = functionName; } - public FunctionMessage(String content, Map properties, String functionName) { + public FunctionMessage(String content, String functionName, Map properties) { super(MessageType.SYSTEM, content, properties); this.functionName = functionName; } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/MessageType.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/MessageType.java index 5a8d22457..ca64c084f 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/MessageType.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/MessageType.java @@ -21,7 +21,7 @@ public enum MessageType { AI("ai"), - GENERIC("generic"), + CHAT("chat"), SYSTEM("system"), diff --git a/spring-ai-core/src/test/java/org/springframework/ai/core/prompts/PromptTests.java b/spring-ai-core/src/test/java/org/springframework/ai/core/prompts/PromptTests.java index 77abe671e..69e1e0a27 100644 --- a/spring-ai-core/src/test/java/org/springframework/ai/core/prompts/PromptTests.java +++ b/spring-ai-core/src/test/java/org/springframework/ai/core/prompts/PromptTests.java @@ -16,17 +16,91 @@ package org.springframework.ai.core.prompts; -import java.util.ArrayList; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Set; import org.assertj.core.api.Assertions; import org.junit.jupiter.api.Test; +import org.springframework.ai.core.prompts.messages.Message; + import static org.assertj.core.api.Assertions.assertThat; class PromptTests { + @Test + void newApiPlaygroundTests() { + // Create a String, a PromptValue or Messages + String templateText = "Hello '{firstName}' '{lastName}' from Unix"; + PromptTemplate pt = new PromptTemplate(templateText); + + final Map model = new HashMap<>(); + model.put("firstName", "Nick"); + + // Try to render with missing value for template variable, expect exception + Assertions.assertThatThrownBy(() -> { + String promptString = pt.render(model); + }) + .isInstanceOf(IllegalStateException.class) + .hasMessage("All template variables were not replaced. Missing variable names are [lastName]"); + + pt.add("lastName", "Park"); // TODO investigate partial + String promptString = pt.render(model); + assertThat(promptString).isEqualTo("Hello 'Nick' 'Park' from Unix"); + + promptString = pt.render(model); // render again + assertThat(promptString).isEqualTo("Hello 'Nick' 'Park' from Unix"); + + // to have access to Messages + Prompt prompt = pt.prompt().system().usingModel(model).create(); + System.out.println(prompt.getContents()); + // -> Hello Nick Park + List messages = prompt.getMessages(); + + prompt = pt.prompt().ai(true).usingModel(model).create(); + System.out.println("Contents: " + prompt.getContents()); + System.out.println("Messages: " + prompt.getMessages()); + + prompt = pt.prompt().system().usingModel(model).create(); // Can use this for + // MessageType that + // don't take + // additional + // arguments. + + String systemTemplate = "You are a helpful assistant that translates {input_language} to {output_language}."; + // system_message_prompt = SystemMessagePromptTemplate.from_template(template) + + Map systemModel = new HashMap(); + systemModel.put("input_language", "English"); + systemModel.put("output_language", "French"); + + String humanTemplate = "{text}"; + Map humanModel = new HashMap(); + humanModel.put("text", "I love programming"); + // human_message_prompt = HumanMessagePromptTemplate.from_template(human_template) + + /* + * chat_prompt = ChatPromptTemplate.from_messages([system_message_prompt, + * human_message_prompt]) + * + * # get a chat completion from the formatted messages + * chat_prompt.format_prompt(input_language="English", output_language="French", + * text="I love programming.").to_messages() + */ + PromptTemplate promptTemplate = new PromptTemplate(systemTemplate); + Prompt systemPrompt = promptTemplate.prompt().system().usingModel(systemModel).create(); + + promptTemplate = new PromptTemplate(humanTemplate); + Prompt humanPrompt = promptTemplate.prompt().human().usingModel(humanModel).create(); + + // ChatPromptTemplate chatPromptTemplate = new ChatPromptTemplate(systemPrompt, + // humanPrompt); + // Prompt chatPrompt chatPromptTemplate.create(model); + + } + @Test void testSingleInputVariable() { String template = "This is a {foo} test"; @@ -58,11 +132,10 @@ class PromptTests { } @Test - void testBadTemplateString() { + void testBadFormatOfTemplateString() { String template = "This is a {foo test"; Assertions.assertThatThrownBy(() -> { - PromptTemplate promptTemplate = new PromptTemplate(template); - promptTemplate.validate(); + new PromptTemplate(template); }).isInstanceOf(IllegalArgumentException.class).hasMessage("The template string is not valid."); }