restructuring

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;

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> memory = Optional.empty();
 
-	private boolean returnOnlyOutputs;
-
-	private List<String> inputKeys;
-
-	private List<String> 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<Memory> getMemory() {
+		return this.memory;
 	}
 
 	@Override
-	public List<String> getInputKeys() {
-		return this.inputKeys;
-	}
+	public abstract List<String> getInputKeys();
 
 	@Override
-	public List<String> getOutputKeys() {
-		return this.outputKeys;
-	}
+	public abstract List<String> getOutputKeys();
 
 	// TODO validation of input/outputs
 
+	@Override
+	public Map<String, Object> apply(Map<String, Object> inputMap) {
+		Map<String, Object> inputMapToUse = processBeforeApply(inputMap);
+		Map<String, Object> outputMap = doApply(inputMapToUse);
+		Map<String, Object> outputMapToUse = processAfterApply(inputMapToUse, outputMap);
+		return outputMapToUse;
+	}
+
+	protected Map<String, Object> processBeforeApply(Map<String, Object> inputMap) {
+		validateInputs(inputMap);
+		return inputMap;
+	}
+
+	protected abstract Map<String, Object> doApply(Map<String, Object> inputMap);
+
+	private Map<String, Object> processAfterApply(Map<String, Object> inputMap, Map<String, Object> outputMap) {
+		validateOutputs(outputMap);
+		Map<String, Object> combindedMap = new HashMap<>();
+		combindedMap.putAll(inputMap);
+		combindedMap.putAll(outputMap);
+		return combindedMap;
+	}
+
+	protected void validateOutputs(Map<String, Object> outputMap) {
+		Set<String> missingKeys = new HashSet<>(getOutputKeys());
+		missingKeys.removeAll(outputMap.keySet());
+		if (!missingKeys.isEmpty()) {
+			throw new IllegalArgumentException("Missing some output keys: " + missingKeys);
+		}
+	}
+
+	protected void validateInputs(Map<String, Object> inputMap) {
+		Set<String> missingKeys = new HashSet<>(getInputKeys());
+		missingKeys.removeAll(inputMap.keySet());
+		if (!missingKeys.isEmpty()) {
+			throw new IllegalArgumentException("Missing some input keys: " + missingKeys);
+		}
+	}
+
 }

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<ChainValues, ChainValues> {
+public interface Chain extends Function<Map<String, Object>, Map<String, Object>> {
 
 	List<String> getInputKeys();
 

spring-ai-core/src/main/java/org/springframework/ai/core/chain/ChainValues.java

@@ -1,17 +0,0 @@
-package org.springframework.ai.core.chain;
-
-import java.util.Map;
-
-public class ChainValues {
-
-	private final Map<String, Object> values;
-
-	public ChainValues(Map<String, Object> values) {
-		this.values = values;
-	}
-
-	public Map<String, Object> getValues() {
-		return values;
-	}
-
-}

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;

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;

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> {
+
+	T parse(List<Generation> output);
+
+}

spring-ai-core/src/main/java/org/springframework/ai/core/prompt/AbstractPromptTemplate.java

@@ -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;
-	}
-
-}

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<PromptTemplate> promptTemplates;
 
-	public ChatPromptTemplate(MessageType messageType, String template) {
-		super(template);
-		this.messageType = messageType;
+	public ChatPromptTemplate(List<PromptTemplate> 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<String, Object> model) {
+		StringBuilder sb = new StringBuilder();
+		for (PromptTemplate promptTemplate : promptTemplates) {
+			sb.append(promptTemplate.render(model));
+		}
+		return sb.toString();
+	}
+
+	@Override
+	public List<Message> createMessages() {
+		List<Message> messages = new ArrayList<>();
+		for (PromptTemplate promptTemplate : promptTemplates) {
+			messages.addAll(promptTemplate.createMessages());
+		}
+		return messages;
+	}
+
+	@Override
+	public List<Message> createMessages(Map<String, Object> model) {
+		List<Message> 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<Message> messages = createMessages();
+		return new Prompt(messages);
 	}
 
 	@Override
 	public Prompt create(Map<String, Object> model) {
-		return new Prompt(new ChatMessage(this.messageType, render(model)));
+		List<Message> messages = createMessages(model);
+		return new Prompt(messages);
 	}
 
 }

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 {

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<String, Object> model);
 
+	List<Message> createMessages();
+
+	List<Message> createMessages(Map<String, Object> model);
+
 	Prompt create();
 
 	Prompt create(Map<String, Object> model);

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<String, Object> 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<String, Object> 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<String, Object> 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<Message> createMessages() {
+		return List.of(new UserMessage(render()));
+	}
+
+	@Override
+	public List<Message> createMessages(Map<String, Object> model) {
+		return List.of(new UserMessage(render(model)));
+	}
+
 	@Override
 	public Prompt create() {
 		return new Prompt(render(new HashMap<>()));

spring-ai-core/src/main/java/org/springframework/ai/core/prompt/UserPromptTemplate.java

@@ -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<String, Object> model) {
-		return new Prompt(new UserMessage(render(model)));
-	}
-
-}

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<String, Object> 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<String> inputKeys = List.of("foo");
+
+		public FakeChain() {
+		}
+
+		public FakeChain(boolean beCorrect) {
+			this.beCorrect = beCorrect;
+		}
+
+		public FakeChain(List<String> inputKeys) {
+			this.inputKeys = inputKeys;
+		}
+
+		@Override
+		public List<String> getInputKeys() {
+			return this.inputKeys;
+		}
+
+		@Override
+		public List<String> getOutputKeys() {
+			return List.of("bar");
+		}
+
+		@Override
+		protected Map<String, Object> doApply(Map<String, Object> inputMap) {
+			if (beCorrect) {
+				return Map.of("bar", "baz");
+			}
+			else {
+				return Map.of("baz", "bar");
+			}
+		}
+
+	}
+
+	class FakeMemory implements Memory {
+
+		@Override
+		public List<String> getKeys() {
+			return List.of("baz");
+		}
+
+		@Override
+		public Map<String, Object> load(Map<String, Object> inputs) {
+			return Map.of("baz", "foo");
+		}
+
+		@Override
+		public void save(Map<String, Object> inputs, Map<String, Object> outputs) {
+
+		}
+
+	}
+
+}