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

Improve Prompt package design

Browse Source
This commit is contained in:
Mark Pollack
2023-07-28 14:38:15 -04:00
parent 62704e0bd5
commit 1b3d6a55dd
13 changed files with 304 additions and 160 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/AbstractPromptTemplate.java
View File

@@ -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> 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<OutputParser> getOutputParser() {
return this.outputParser;
public String getTemplate() {
return this.template;
}
public PromptValue formatAsPrompt(Map<String, Object> inputVariables) {
String formattedPrompt = formatAsString(inputVariables);
return new StringPromptValue(formattedPrompt);
@Override
public TemplateFormat getTemplateFormat() {
return this.templateFormat;
}
}

26
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/StringPromptValue.java → spring-ai-core/src/main/java/org/springframework/ai/core/prompts/Prompt.java
View File

@@ -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<Message> 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<Message> messages) {
this.messages = messages;
}
@Override
public List<Message> 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<Message> getMessages() {
return this.messages;
}
}

37
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptBuilder.java Normal file
View File

@@ -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<String, Object> model);
PromptBuilder withProperties(Map<String, Object> properties);
Prompt create();
}

42
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptInput.java
View File

@@ -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<OutputParser> 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<String,?> model, HttpServletRequest request, HttpServletResponse response) method.
String formatAsString(Map<String, Object> inputVariables);
PromptValue formatAsPrompt(Map<String, Object> inputVariables);
// Leave out Partial Input Variables for now
}

14
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptTemplateInput.java → spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptOperations.java
View File

@@ -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<String, Object> model);
PromptBuilder prompt();
}

188
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptTemplate.java
View File

@@ -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<String, Object> 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<String, Object> inputVariables) {
validate();
// Only "F-String" for now
ST st = new ST(this.template, '{', '}');
for (Entry<String, Object> 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<String, Object> model) {
validate(model);
for (Entry<String, Object> stringObjectEntry : model.entrySet()) {
if (st.getAttribute(stringObjectEntry.getKey()) == null) {
st.add(stringObjectEntry.getKey(), stringObjectEntry.getValue());
}
}
return st.render().trim();
}
protected Set<String> 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<String> 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<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()) {
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<String, Object> model = new HashMap<>();
private Map<String, Object> 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<String, Object> model) {
this.model = model;
return this;
}
@Override
public PromptBuilder withProperties(Map<String, Object> 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));
}
}
}

29
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/PromptValue.java
View File

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

2
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/TemplateFormat.java
View File

@@ -18,7 +18,7 @@ package org.springframework.ai.core.prompts;
public enum TemplateFormat {
FSTRING("f-string");
ST("ST");
private final String value;

14
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/AiMessage.java
View File

@@ -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<String, Object> properties) {
public AiMessage(String content, boolean example) {
super(MessageType.AI, content);
this.example = example;
}
public AiMessage(String content, boolean example, Map<String, Object> properties) {
super(MessageType.AI, content, properties);
this.example = example;
}
public boolean isExample() {
return example;
}
}

2
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/ChatMessage.java
View File

@@ -27,7 +27,7 @@ public class ChatMessage extends AbstractMessage {
this.role = role;
}
public ChatMessage(String content, Map<String, Object> properties, String role) {
public ChatMessage(String content, String role, Map<String, Object> properties) {
super(MessageType.SYSTEM, content, properties);
this.role = role;
}

2
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/FunctionMessage.java
View File

@@ -27,7 +27,7 @@ public class FunctionMessage extends AbstractMessage {
this.functionName = functionName;
}
public FunctionMessage(String content, Map<String, Object> properties, String functionName) {
public FunctionMessage(String content, String functionName, Map<String, Object> properties) {
super(MessageType.SYSTEM, content, properties);
this.functionName = functionName;
}

2
spring-ai-core/src/main/java/org/springframework/ai/core/prompts/messages/MessageType.java
View File

@@ -21,7 +21,7 @@ public enum MessageType {
AI("ai"),
GENERIC("generic"),
CHAT("chat"),
SYSTEM("system"),

81
spring-ai-core/src/test/java/org/springframework/ai/core/prompts/PromptTests.java
View File

@@ -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<String, Object> 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<Message> 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<String, Object> systemModel = new HashMap();
systemModel.put("input_language", "English");
systemModel.put("output_language", "French");
String humanTemplate = "{text}";
Map<String, Object> 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.");
}