diff --git a/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/llm/AzureOpenAiClient.java b/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/llm/AzureOpenAiClient.java index 9c38df877..e606370b3 100644 --- a/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/llm/AzureOpenAiClient.java +++ b/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/llm/AzureOpenAiClient.java @@ -22,7 +22,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.ai.core.llm.LLMResult; import org.springframework.ai.core.llm.LlmClient; -import org.springframework.ai.core.prompt.Generation; +import org.springframework.ai.core.llm.Generation; import org.springframework.ai.core.prompt.Prompt; import org.springframework.ai.core.prompt.messages.Message; import org.springframework.util.Assert; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/chain/AbstractChain.java b/spring-ai-core/src/main/java/org/springframework/ai/core/chain/AbstractChain.java index abf8a66c6..aecffb242 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/chain/AbstractChain.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/chain/AbstractChain.java @@ -2,43 +2,61 @@ package org.springframework.ai.core.chain; import org.springframework.ai.core.memory.Memory; -import java.util.List; +import java.util.*; public abstract class AbstractChain implements Chain { - private Memory memory; + private Optional memory = Optional.empty(); - private boolean returnOnlyOutputs; - - private List inputKeys; - - private List outputKeys; - - /** - * @return A string that uniquely identifies the type of chain - */ - protected abstract String getType(); - - protected abstract Memory getMemory(); - - public void setReturnOnlyOutputs(boolean returnOnlyOutputs) { - this.returnOnlyOutputs = returnOnlyOutputs; - } - - public boolean isReturnOnlyOutputs() { - return this.returnOnlyOutputs; + protected Optional getMemory() { + return this.memory; } @Override - public List getInputKeys() { - return this.inputKeys; - } + public abstract List getInputKeys(); @Override - public List getOutputKeys() { - return this.outputKeys; - } + public abstract List getOutputKeys(); // TODO validation of input/outputs + @Override + public Map apply(Map inputMap) { + Map inputMapToUse = processBeforeApply(inputMap); + Map outputMap = doApply(inputMapToUse); + Map outputMapToUse = processAfterApply(inputMapToUse, outputMap); + return outputMapToUse; + } + + protected Map processBeforeApply(Map inputMap) { + validateInputs(inputMap); + return inputMap; + } + + protected abstract Map doApply(Map inputMap); + + private Map processAfterApply(Map inputMap, Map outputMap) { + validateOutputs(outputMap); + Map combindedMap = new HashMap<>(); + combindedMap.putAll(inputMap); + combindedMap.putAll(outputMap); + return combindedMap; + } + + protected void validateOutputs(Map outputMap) { + Set missingKeys = new HashSet<>(getOutputKeys()); + missingKeys.removeAll(outputMap.keySet()); + if (!missingKeys.isEmpty()) { + throw new IllegalArgumentException("Missing some output keys: " + missingKeys); + } + } + + protected void validateInputs(Map inputMap) { + Set missingKeys = new HashSet<>(getInputKeys()); + missingKeys.removeAll(inputMap.keySet()); + if (!missingKeys.isEmpty()) { + throw new IllegalArgumentException("Missing some input keys: " + missingKeys); + } + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/chain/Chain.java b/spring-ai-core/src/main/java/org/springframework/ai/core/chain/Chain.java index 9e6626b43..ef59ebd30 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/chain/Chain.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/chain/Chain.java @@ -1,9 +1,10 @@ package org.springframework.ai.core.chain; import java.util.List; +import java.util.Map; import java.util.function.Function; -public interface Chain extends Function { +public interface Chain extends Function, Map> { List getInputKeys(); diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/chain/ChainValues.java b/spring-ai-core/src/main/java/org/springframework/ai/core/chain/ChainValues.java deleted file mode 100644 index e02023171..000000000 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/chain/ChainValues.java +++ /dev/null @@ -1,17 +0,0 @@ -package org.springframework.ai.core.chain; - -import java.util.Map; - -public class ChainValues { - - private final Map values; - - public ChainValues(Map values) { - this.values = values; - } - - public Map getValues() { - return values; - } - -} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/Generation.java b/spring-ai-core/src/main/java/org/springframework/ai/core/llm/Generation.java similarity index 95% rename from spring-ai-core/src/main/java/org/springframework/ai/core/prompt/Generation.java rename to spring-ai-core/src/main/java/org/springframework/ai/core/llm/Generation.java index d2777bf34..bee820312 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/Generation.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/llm/Generation.java @@ -14,7 +14,7 @@ * limitations under the License. */ -package org.springframework.ai.core.prompt; +package org.springframework.ai.core.llm; import java.util.HashMap; import java.util.Map; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/llm/LLMResult.java b/spring-ai-core/src/main/java/org/springframework/ai/core/llm/LLMResult.java index 32d91fd21..a980c3c49 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/llm/LLMResult.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/llm/LLMResult.java @@ -15,8 +15,6 @@ */ package org.springframework.ai.core.llm; -import org.springframework.ai.core.prompt.Generation; - import java.util.HashMap; import java.util.List; import java.util.Map; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/parser/OutputParser.java b/spring-ai-core/src/main/java/org/springframework/ai/core/parser/OutputParser.java new file mode 100644 index 000000000..7157850e2 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/parser/OutputParser.java @@ -0,0 +1,11 @@ +package org.springframework.ai.core.parser; + +import org.springframework.ai.core.llm.Generation; + +import java.util.List; + +public interface OutputParser { + + T parse(List output); + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/AbstractPromptTemplate.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/AbstractPromptTemplate.java deleted file mode 100644 index de1bbdf6b..000000000 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/AbstractPromptTemplate.java +++ /dev/null @@ -1,39 +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.prompt; - -public abstract class AbstractPromptTemplate implements PromptOperations { - - protected String template; - - protected TemplateFormat templateFormat = TemplateFormat.ST; - - public AbstractPromptTemplate(String template) { - this.template = template; - } - - @Override - public String getTemplate() { - return this.template; - } - - @Override - public TemplateFormat getTemplateFormat() { - return this.templateFormat; - } - -} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/ChatPromptTemplate.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/ChatPromptTemplate.java index 2130ad92d..a9cf46f06 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/ChatPromptTemplate.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/ChatPromptTemplate.java @@ -17,31 +17,72 @@ package org.springframework.ai.core.prompt; import org.springframework.ai.core.prompt.messages.ChatMessage; +import org.springframework.ai.core.prompt.messages.Message; import org.springframework.ai.core.prompt.messages.MessageType; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; import java.util.Map; /** * A PromptTemplate that lets you specify the role as a string should the current * implementations and their roles not suffice for your needs. */ -public class ChatPromptTemplate extends PromptTemplate { +public class ChatPromptTemplate implements PromptOperations { - private MessageType messageType; + private final List promptTemplates; - public ChatPromptTemplate(MessageType messageType, String template) { - super(template); - this.messageType = messageType; + public ChatPromptTemplate(List promptTemplates) { + this.promptTemplates = promptTemplates; + } + + @Override + public String render() { + StringBuilder sb = new StringBuilder(); + for (PromptTemplate promptTemplate : promptTemplates) { + sb.append(promptTemplate.render()); + } + return sb.toString(); + } + + @Override + public String render(Map model) { + StringBuilder sb = new StringBuilder(); + for (PromptTemplate promptTemplate : promptTemplates) { + sb.append(promptTemplate.render(model)); + } + return sb.toString(); + } + + @Override + public List createMessages() { + List messages = new ArrayList<>(); + for (PromptTemplate promptTemplate : promptTemplates) { + messages.addAll(promptTemplate.createMessages()); + } + return messages; + } + + @Override + public List createMessages(Map model) { + List messages = new ArrayList<>(); + for (PromptTemplate promptTemplate : promptTemplates) { + messages.addAll(promptTemplate.createMessages(model)); + } + return messages; } @Override public Prompt create() { - return new Prompt(new ChatMessage(this.messageType, render())); + List messages = createMessages(); + return new Prompt(messages); } @Override public Prompt create(Map model) { - return new Prompt(new ChatMessage(this.messageType, render(model))); + List messages = createMessages(model); + return new Prompt(messages); } } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/OutputParser.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/OutputParser.java index 2c1fc3412..c9df0b71a 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/OutputParser.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/OutputParser.java @@ -16,6 +16,8 @@ package org.springframework.ai.core.prompt; +import org.springframework.ai.core.llm.Generation; + import java.util.List; public interface OutputParser { diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/PromptOperations.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/PromptOperations.java index c49a10b4d..376167cac 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/PromptOperations.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/PromptOperations.java @@ -16,20 +16,21 @@ package org.springframework.ai.core.prompt; +import org.springframework.ai.core.prompt.messages.Message; + +import java.util.List; import java.util.Map; public interface PromptOperations { - String getTemplate(); - - TemplateFormat getTemplateFormat(); - - void add(String name, Object value); - String render(); String render(Map model); + List createMessages(); + + List createMessages(Map model); + Prompt create(); Prompt create(Map model); diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/PromptTemplate.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/PromptTemplate.java index 565a0dc73..eda3feb4a 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/PromptTemplate.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/PromptTemplate.java @@ -18,25 +18,30 @@ package org.springframework.ai.core.prompt; import org.antlr.runtime.Token; import org.antlr.runtime.TokenStream; +import org.springframework.ai.core.prompt.messages.Message; +import org.springframework.ai.core.prompt.messages.UserMessage; import org.stringtemplate.v4.ST; import org.stringtemplate.v4.compiler.STLexer; -import java.util.HashMap; -import java.util.HashSet; -import java.util.Map; +import java.util.*; import java.util.Map.Entry; -import java.util.Set; import java.util.stream.Collectors; import java.util.stream.IntStream; -public class PromptTemplate extends AbstractPromptTemplate { +public class PromptTemplate implements PromptOperations { private ST st; private Map dynamicModel = new HashMap<>(); + protected String template; + + protected TemplateFormat templateFormat = TemplateFormat.ST; + + private OutputParser outputParser; + public PromptTemplate(String template) { - super(template); + this.template = template; // If the template string is not valid, an exception will be thrown try { this.st = new ST(this.template, '{', '}'); @@ -46,12 +51,42 @@ public class PromptTemplate extends AbstractPromptTemplate { } } - @Override + public PromptTemplate(String template, Map model) { + this.template = template; + // If the template string is not valid, an exception will be thrown + try { + this.st = new ST(this.template, '{', '}'); + for (Entry entry : model.entrySet()) { + add(entry.getKey(), entry.getValue()); + } + } + catch (Exception ex) { + throw new IllegalArgumentException("The template string is not valid.", ex); + } + } + + public OutputParser getOutputParser() { + return outputParser; + } + + public void setOutputParser(OutputParser outputParser) { + Objects.requireNonNull(outputParser, "Output Parser can not be null"); + this.outputParser = outputParser; + } + public void add(String name, Object value) { this.st.add(name, value); this.dynamicModel.put(name, value); } + public String getTemplate() { + return this.template; + } + + public TemplateFormat getTemplateFormat() { + return this.templateFormat; + } + // Render Methods public String render() { return st.render(); @@ -68,6 +103,16 @@ public class PromptTemplate extends AbstractPromptTemplate { return st.render().trim(); } + @Override + public List createMessages() { + return List.of(new UserMessage(render())); + } + + @Override + public List createMessages(Map model) { + return List.of(new UserMessage(render(model))); + } + @Override public Prompt create() { return new Prompt(render(new HashMap<>())); diff --git a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/UserPromptTemplate.java b/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/UserPromptTemplate.java deleted file mode 100644 index 98d577478..000000000 --- a/spring-ai-core/src/main/java/org/springframework/ai/core/prompt/UserPromptTemplate.java +++ /dev/null @@ -1,40 +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.prompt; - -import org.springframework.ai.core.prompt.messages.SystemMessage; -import org.springframework.ai.core.prompt.messages.UserMessage; - -import java.util.Map; - -public class UserPromptTemplate extends PromptTemplate { - - public UserPromptTemplate(String template) { - super(template); - } - - @Override - public Prompt create() { - return new Prompt(new SystemMessage(render())); - } - - @Override - public Prompt create(Map model) { - return new Prompt(new UserMessage(render(model))); - } - -} diff --git a/spring-ai-core/src/test/java/org/springframework/ai/core/chain/ChainTests.java b/spring-ai-core/src/test/java/org/springframework/ai/core/chain/ChainTests.java new file mode 100644 index 000000000..b19a1ceba --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/core/chain/ChainTests.java @@ -0,0 +1,86 @@ +package org.springframework.ai.core.chain; + +import org.junit.jupiter.api.Test; +import org.springframework.ai.core.memory.Memory; + +import java.util.List; +import java.util.Map; + +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.assertj.core.api.AssertionsForInterfaceTypes.assertThat; + +class ChainTests { + + @Test + void badInputs() { + Chain chain = new FakeChain(); + assertThatThrownBy(() -> chain.apply(Map.of("foobar", "baz"))).isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Missing some input keys"); + } + + @Test + void correctInputs() { + Chain chain = new FakeChain(); + Map output = chain.apply(Map.of("foo", "bar")); + assertThat(output).containsEntry("foo", "bar").containsEntry("bar", "baz"); + } + + class FakeChain extends AbstractChain { + + private boolean beCorrect = true; + + private List inputKeys = List.of("foo"); + + public FakeChain() { + } + + public FakeChain(boolean beCorrect) { + this.beCorrect = beCorrect; + } + + public FakeChain(List inputKeys) { + this.inputKeys = inputKeys; + } + + @Override + public List getInputKeys() { + return this.inputKeys; + } + + @Override + public List getOutputKeys() { + return List.of("bar"); + } + + @Override + protected Map doApply(Map inputMap) { + if (beCorrect) { + return Map.of("bar", "baz"); + } + else { + return Map.of("baz", "bar"); + } + } + + } + + class FakeMemory implements Memory { + + @Override + public List getKeys() { + return List.of("baz"); + } + + @Override + public Map load(Map inputs) { + return Map.of("baz", "foo"); + } + + @Override + public void save(Map inputs, Map outputs) { + + } + + } + +}