diff --git a/.gitignore b/.gitignore
index 43e4193b0..c2fb8fc27 100644
--- a/.gitignore
+++ b/.gitignore
@@ -36,4 +36,6 @@ package.json
.vscode
.antlr
-shell.log
\ No newline at end of file
+shell.log
+
+.profiler
\ No newline at end of file
diff --git a/.mvn/extensions.xml b/.mvn/extensions.xml
new file mode 100644
index 000000000..c7e6507ac
--- /dev/null
+++ b/.mvn/extensions.xml
@@ -0,0 +1,8 @@
+
+
+
+ fr.jcgay.maven
+ maven-profiler
+ 3.2
+
+
\ No newline at end of file
diff --git a/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/client/ClientIT.java b/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/client/ClientIT.java
index 2fa2e89bb..f21c7b9ab 100644
--- a/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/client/ClientIT.java
+++ b/models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/client/ClientIT.java
@@ -57,8 +57,8 @@ public class ClientIT {
}
```""";
assertThat(chatResponse.getResult().getOutput().getContent()).isEqualTo(expectedResponse);
- assertThat(chatResponse.getResult().getOutput().getProperties()).containsKey("generated_tokens");
- assertThat(chatResponse.getResult().getOutput().getProperties()).containsEntry("generated_tokens", 39);
+ assertThat(chatResponse.getResult().getOutput().getMetadata()).containsKey("generated_tokens");
+ assertThat(chatResponse.getResult().getOutput().getMetadata()).containsEntry("generated_tokens", 39);
}
diff --git a/models/spring-ai-openai/pom.xml b/models/spring-ai-openai/pom.xml
index cadbf64a5..715e3474f 100644
--- a/models/spring-ai-openai/pom.xml
+++ b/models/spring-ai-openai/pom.xml
@@ -1,5 +1,6 @@
-
+
4.0.0
org.springframework.ai
@@ -74,6 +75,38 @@
test
+
+ org.springframework.ai
+ spring-ai-qdrant
+ ${project.version}
+
+
+ org.springframework.ai
+ spring-ai-openai
+
+
+ test
+
+
+
+ org.testcontainers
+ qdrant
+ test
+
+
+
+ org.testcontainers
+ testcontainers
+ test
+
+
+
+ org.testcontainers
+ junit-jupiter
+ test
+
+
+
diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryLongTermSystemPromptIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryLongTermSystemPromptIT.java
new file mode 100644
index 000000000..4f33bd993
--- /dev/null
+++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryLongTermSystemPromptIT.java
@@ -0,0 +1,131 @@
+/*
+ * Copyright 2024 - 2024 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.openai.chat.chatbot;
+
+import java.util.List;
+
+import io.qdrant.client.QdrantClient;
+import io.qdrant.client.QdrantGrpcClient;
+import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
+import org.springframework.ai.chat.chatbot.ChatBot;
+import org.testcontainers.junit.jupiter.Container;
+import org.testcontainers.junit.jupiter.Testcontainers;
+import org.testcontainers.qdrant.QdrantContainer;
+
+import org.springframework.ai.chat.chatbot.DefaultChatBot;
+import org.springframework.ai.chat.chatbot.DefaultStreamingChatBot;
+import org.springframework.ai.chat.chatbot.StreamingChatBot;
+import org.springframework.ai.chat.history.VectorStoreChatMemoryAgentListener;
+import org.springframework.ai.chat.history.VectorStoreChatMemoryRetriever;
+import org.springframework.ai.chat.history.LastMaxTokenSizeContentTransformer;
+import org.springframework.ai.chat.history.SystemPromptChatMemoryAugmentor;
+import org.springframework.ai.embedding.EmbeddingClient;
+import org.springframework.ai.evaluation.BaseMemoryTest;
+import org.springframework.ai.evaluation.RelevancyEvaluator;
+import org.springframework.ai.openai.OpenAiChatClient;
+import org.springframework.ai.openai.OpenAiEmbeddingClient;
+import org.springframework.ai.openai.api.OpenAiApi;
+import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
+import org.springframework.ai.tokenizer.TokenCountEstimator;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.SpringBootConfiguration;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.context.annotation.Bean;
+
+@Testcontainers
+@SpringBootTest(classes = ChatMemoryLongTermSystemPromptIT.Config.class)
+@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
+public class ChatMemoryLongTermSystemPromptIT extends BaseMemoryTest {
+
+ private static final String COLLECTION_NAME = "test_collection";
+
+ private static final int QDRANT_GRPC_PORT = 6334;
+
+ @Container
+ static QdrantContainer qdrantContainer = new QdrantContainer("qdrant/qdrant:v1.7.4");
+
+ @Autowired
+ public ChatMemoryLongTermSystemPromptIT(RelevancyEvaluator relevancyEvaluator, ChatBot chatBot,
+ StreamingChatBot streamingChatBot) {
+ super(relevancyEvaluator, chatBot, streamingChatBot);
+ }
+
+ @SpringBootConfiguration
+ static class Config {
+
+ @Bean
+ public OpenAiApi chatCompletionApi() {
+ return new OpenAiApi(System.getenv("OPENAI_API_KEY"));
+ }
+
+ @Bean
+ public OpenAiChatClient openAiClient(OpenAiApi openAiApi) {
+ return new OpenAiChatClient(openAiApi);
+ }
+
+ @Bean
+ public EmbeddingClient embeddingClient(OpenAiApi openAiApi) {
+ return new OpenAiEmbeddingClient(openAiApi);
+ }
+
+ @Bean
+ public VectorStore qdrantVectorStore(EmbeddingClient embeddingClient) {
+ QdrantClient qdrantClient = new QdrantClient(QdrantGrpcClient
+ .newBuilder(qdrantContainer.getHost(), qdrantContainer.getMappedPort(QDRANT_GRPC_PORT), false)
+ .build());
+ return new QdrantVectorStore(qdrantClient, COLLECTION_NAME, embeddingClient);
+ }
+
+ @Bean
+ public TokenCountEstimator tokenCountEstimator() {
+ return new JTokkitTokenCountEstimator();
+ }
+
+ @Bean
+ public ChatBot memoryChatAgent(OpenAiChatClient chatClient, VectorStore vectorStore,
+ TokenCountEstimator tokenCountEstimator) {
+
+ return DefaultChatBot.builder(chatClient)
+ .withRetrievers(List.of(new VectorStoreChatMemoryRetriever(vectorStore, 10)))
+ .withContentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
+ .withAugmentors(List.of(new SystemPromptChatMemoryAugmentor()))
+ .withChatAgentListeners(List.of(new VectorStoreChatMemoryAgentListener(vectorStore)))
+ .build();
+ }
+
+ @Bean
+ public StreamingChatBot memoryStreamingChatAgent(OpenAiChatClient streamingChatClient, VectorStore vectorStore,
+ TokenCountEstimator tokenCountEstimator) {
+
+ return DefaultStreamingChatBot.builder(streamingChatClient)
+ .withRetrievers(List.of(new VectorStoreChatMemoryRetriever(vectorStore, 10)))
+ .withDocumentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
+ .withAugmentors(List.of(new SystemPromptChatMemoryAugmentor()))
+ .withChatAgentListeners(List.of(new VectorStoreChatMemoryAgentListener(vectorStore)))
+ .build();
+ }
+
+ @Bean
+ public RelevancyEvaluator relevancyEvaluator(OpenAiChatClient chatClient) {
+ return new RelevancyEvaluator(chatClient);
+ }
+
+ }
+
+}
diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryShortTermMessageListIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryShortTermMessageListIT.java
new file mode 100644
index 000000000..1c857ddbb
--- /dev/null
+++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryShortTermMessageListIT.java
@@ -0,0 +1,107 @@
+/*
+ * Copyright 2024 - 2024 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.openai.chat.chatbot;
+
+import java.util.List;
+
+import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
+
+import org.springframework.ai.chat.chatbot.ChatBot;
+import org.springframework.ai.chat.chatbot.DefaultChatBot;
+import org.springframework.ai.chat.chatbot.DefaultStreamingChatBot;
+import org.springframework.ai.chat.chatbot.StreamingChatBot;
+import org.springframework.ai.chat.history.ChatMemory;
+import org.springframework.ai.chat.history.ChatMemoryAgentListener;
+import org.springframework.ai.chat.history.ChatMemoryRetriever;
+import org.springframework.ai.chat.history.InMemoryChatMemory;
+import org.springframework.ai.chat.history.LastMaxTokenSizeContentTransformer;
+import org.springframework.ai.chat.history.MessageChatMemoryAugmentor;
+import org.springframework.ai.evaluation.BaseMemoryTest;
+import org.springframework.ai.evaluation.RelevancyEvaluator;
+import org.springframework.ai.openai.OpenAiChatClient;
+import org.springframework.ai.openai.api.OpenAiApi;
+import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
+import org.springframework.ai.tokenizer.TokenCountEstimator;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.SpringBootConfiguration;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.context.annotation.Bean;
+
+@SpringBootTest(classes = ChatMemoryShortTermMessageListIT.Config.class)
+@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
+public class ChatMemoryShortTermMessageListIT extends BaseMemoryTest {
+
+ @Autowired
+ public ChatMemoryShortTermMessageListIT(RelevancyEvaluator relevancyEvaluator, ChatBot chatBot,
+ StreamingChatBot streamingChatBot) {
+ super(relevancyEvaluator, chatBot, streamingChatBot);
+ }
+
+ @SpringBootConfiguration
+ static class Config {
+
+ @Bean
+ public OpenAiApi chatCompletionApi() {
+ return new OpenAiApi(System.getenv("OPENAI_API_KEY"));
+ }
+
+ @Bean
+ public OpenAiChatClient openAiClient(OpenAiApi openAiApi) {
+ return new OpenAiChatClient(openAiApi);
+ }
+
+ @Bean
+ public ChatMemory chatHistory() {
+ return new InMemoryChatMemory();
+ }
+
+ @Bean
+ public TokenCountEstimator tokenCountEstimator() {
+ return new JTokkitTokenCountEstimator();
+ }
+
+ @Bean
+ public ChatBot memoryChatAgent(OpenAiChatClient chatClient, ChatMemory chatHistory,
+ TokenCountEstimator tokenCountEstimator) {
+
+ return DefaultChatBot.builder(chatClient)
+ .withRetrievers(List.of(new ChatMemoryRetriever(chatHistory)))
+ .withContentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
+ .withAugmentors(List.of(new MessageChatMemoryAugmentor()))
+ .withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory)))
+ .build();
+ }
+
+ @Bean
+ public StreamingChatBot memoryStreamingChatAgent(OpenAiChatClient streamingChatClient, ChatMemory chatHistory,
+ TokenCountEstimator tokenCountEstimator) {
+
+ return DefaultStreamingChatBot.builder(streamingChatClient)
+ .withRetrievers(List.of(new ChatMemoryRetriever(chatHistory)))
+ .withDocumentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
+ .withAugmentors(List.of(new MessageChatMemoryAugmentor()))
+ .withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory)))
+ .build();
+ }
+
+ @Bean
+ public RelevancyEvaluator relevancyEvaluator(OpenAiChatClient chatClient) {
+ return new RelevancyEvaluator(chatClient);
+ }
+
+ }
+
+}
diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryShortTermSystemPromptIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryShortTermSystemPromptIT.java
new file mode 100644
index 000000000..ac5b1d386
--- /dev/null
+++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryShortTermSystemPromptIT.java
@@ -0,0 +1,108 @@
+/*
+ * Copyright 2024 - 2024 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.openai.chat.chatbot;
+
+import java.util.List;
+
+import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
+
+import org.springframework.ai.chat.chatbot.ChatBot;
+import org.springframework.ai.chat.chatbot.DefaultChatBot;
+import org.springframework.ai.chat.chatbot.DefaultStreamingChatBot;
+import org.springframework.ai.chat.chatbot.StreamingChatBot;
+import org.springframework.ai.chat.history.ChatMemory;
+import org.springframework.ai.chat.history.ChatMemoryAgentListener;
+import org.springframework.ai.chat.history.ChatMemoryRetriever;
+import org.springframework.ai.chat.history.InMemoryChatMemory;
+import org.springframework.ai.chat.history.LastMaxTokenSizeContentTransformer;
+import org.springframework.ai.chat.history.SystemPromptChatMemoryAugmentor;
+import org.springframework.ai.evaluation.BaseMemoryTest;
+import org.springframework.ai.evaluation.RelevancyEvaluator;
+import org.springframework.ai.openai.OpenAiChatClient;
+import org.springframework.ai.openai.api.OpenAiApi;
+import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
+import org.springframework.ai.tokenizer.TokenCountEstimator;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.boot.SpringBootConfiguration;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.context.annotation.Bean;
+
+@SpringBootTest(classes = ChatMemoryShortTermSystemPromptIT.Config.class)
+@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
+public class ChatMemoryShortTermSystemPromptIT extends BaseMemoryTest {
+
+ @Autowired
+ public ChatMemoryShortTermSystemPromptIT(RelevancyEvaluator relevancyEvaluator, ChatBot chatBot,
+ StreamingChatBot streamingChatBot) {
+ super(relevancyEvaluator, chatBot, streamingChatBot);
+ }
+
+ @SpringBootConfiguration
+ static class Config {
+
+ @Bean
+ public OpenAiApi chatCompletionApi() {
+ return new OpenAiApi(System.getenv("OPENAI_API_KEY"));
+ }
+
+ @Bean
+ public OpenAiChatClient openAiClient(OpenAiApi openAiApi) {
+ return new OpenAiChatClient(openAiApi);
+ }
+
+ @Bean
+ public ChatMemory chatHistory() {
+ return new InMemoryChatMemory();
+ }
+
+ @Bean
+ public TokenCountEstimator tokenCountEstimator() {
+ return new JTokkitTokenCountEstimator();
+ }
+
+ @Bean
+ public ChatBot memoryChatAgent(OpenAiChatClient chatClient, ChatMemory chatHistory,
+ TokenCountEstimator tokenCountEstimator) {
+
+ return DefaultChatBot.builder(chatClient)
+ .withRetrievers(List.of(new ChatMemoryRetriever(chatHistory)))
+ .withContentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
+ .withAugmentors(List.of(new SystemPromptChatMemoryAugmentor()))
+ .withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory)))
+ .build();
+ }
+
+ @Bean
+ public StreamingChatBot memoryStreamingChatAgent(OpenAiChatClient streamingChatClient, ChatMemory chatHistory,
+ TokenCountEstimator tokenCountEstimator) {
+
+ return DefaultStreamingChatBot.builder(streamingChatClient)
+ .withRetrievers(List.of(new ChatMemoryRetriever(chatHistory)))
+ .withDocumentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
+ .withAugmentors(List.of(new SystemPromptChatMemoryAugmentor()))
+ .withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory)))
+ .build();
+ }
+
+ @Bean
+ public RelevancyEvaluator relevancyEvaluator(OpenAiChatClient chatClient) {
+ return new RelevancyEvaluator(chatClient);
+ }
+
+ }
+
+}
diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/LongShortTermChatMemoryWithRagIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/LongShortTermChatMemoryWithRagIT.java
new file mode 100644
index 000000000..c513c8776
--- /dev/null
+++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/LongShortTermChatMemoryWithRagIT.java
@@ -0,0 +1,252 @@
+/*
+ * Copyright 2024 - 2024 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.openai.chat.chatbot;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+
+import io.qdrant.client.QdrantClient;
+import io.qdrant.client.QdrantGrpcClient;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.ai.chat.chatbot.ChatBot;
+import org.springframework.ai.openai.OpenAiChatOptions;
+import org.testcontainers.junit.jupiter.Container;
+import org.testcontainers.junit.jupiter.Testcontainers;
+import org.testcontainers.qdrant.QdrantContainer;
+
+import org.springframework.ai.chat.chatbot.DefaultChatBot;
+import org.springframework.ai.chat.history.ChatMemory;
+import org.springframework.ai.chat.history.ChatMemoryAgentListener;
+import org.springframework.ai.chat.history.ChatMemoryRetriever;
+import org.springframework.ai.chat.history.InMemoryChatMemory;
+import org.springframework.ai.chat.history.LastMaxTokenSizeContentTransformer;
+import org.springframework.ai.chat.history.SystemPromptChatMemoryAugmentor;
+import org.springframework.ai.chat.history.VectorStoreChatMemoryAgentListener;
+import org.springframework.ai.chat.history.VectorStoreChatMemoryRetriever;
+import org.springframework.ai.chat.messages.UserMessage;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+import org.springframework.ai.chat.prompt.transformer.QuestionContextAugmentor;
+import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
+import org.springframework.ai.chat.prompt.transformer.VectorStoreRetriever;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.document.DocumentTransformer;
+import org.springframework.ai.embedding.EmbeddingClient;
+import org.springframework.ai.evaluation.EvaluationRequest;
+import org.springframework.ai.evaluation.EvaluationResponse;
+import org.springframework.ai.evaluation.RelevancyEvaluator;
+import org.springframework.ai.openai.OpenAiChatClient;
+import org.springframework.ai.openai.OpenAiEmbeddingClient;
+import org.springframework.ai.openai.api.OpenAiApi;
+import org.springframework.ai.reader.JsonReader;
+import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
+import org.springframework.ai.tokenizer.TokenCountEstimator;
+import org.springframework.ai.transformer.splitter.TokenTextSplitter;
+import org.springframework.ai.vectorstore.SearchRequest;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.boot.SpringBootConfiguration;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.context.annotation.Bean;
+import org.springframework.core.io.Resource;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.springframework.ai.openai.api.OpenAiApi.ChatModel.GPT_4_TURBO_PREVIEW;
+
+@Testcontainers
+@SpringBootTest(classes = LongShortTermChatMemoryWithRagIT.Config.class)
+@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
+public class LongShortTermChatMemoryWithRagIT {
+
+ protected final Logger logger = LoggerFactory.getLogger(getClass());
+
+ private static final String COLLECTION_NAME = "test_collection";
+
+ private static final int QDRANT_GRPC_PORT = 6334;
+
+ @Container
+ static QdrantContainer qdrantContainer = new QdrantContainer("qdrant/qdrant:v1.7.4");
+
+ @Autowired
+ ChatBot chatBot;
+
+ @Autowired
+ RelevancyEvaluator relevancyEvaluator;
+
+ @Autowired
+ VectorStore vectorStore;
+
+ @Value("classpath:/data/acme/bikes.json")
+ private Resource bikesResource;
+
+ void loadData() {
+
+ var metadataEnricher = new DocumentTransformer() {
+
+ @Override
+ public List apply(List documents) {
+ documents.forEach(d -> {
+ Map metadata = d.getMetadata();
+ metadata.put(TransformerContentType.EXTERNAL_KNOWLEDGE, "true");
+ });
+
+ return documents;
+ }
+
+ };
+
+ JsonReader jsonReader = new JsonReader(bikesResource, "name", "price", "shortDescription", "description");
+ var textSplitter = new TokenTextSplitter();
+ vectorStore.accept(metadataEnricher.apply(textSplitter.apply(jsonReader.get())));
+ }
+
+ // @Autowired
+ // StreamingChatBot streamingChatBot;
+
+ @Test
+ void memoryChatBot() {
+
+ loadData();
+
+ var prompt = new Prompt(new UserMessage("My name is Christian and I like mountain bikes."));
+ PromptContext promptContext = new PromptContext(prompt);
+
+ var chatBotResponse1 = this.chatBot.call(promptContext);
+
+ logger.info("Response1: " + chatBotResponse1.getChatResponse().getResult().getOutput().getContent());
+
+ var chatBotResponse2 = this.chatBot.call(new PromptContext(
+ new Prompt(new String("What is my name and what bike model would you suggest for me?"))));
+ logger.info("Response2: " + chatBotResponse2.getChatResponse().getResult().getOutput().getContent());
+
+ // logger.info(chatBotResponse2.getPromptContext().getContents().toString());
+ assertThat(chatBotResponse2.getChatResponse().getResult().getOutput().getContent()).contains("Christian");
+
+ EvaluationResponse evaluationResponse = this.relevancyEvaluator
+ .evaluate(new EvaluationRequest(chatBotResponse2));
+
+ assertTrue(evaluationResponse.isPass(), "Response is not relevant to the question");
+
+ }
+
+ @SpringBootConfiguration
+ static class Config {
+
+ @Bean
+ public ChatMemory chatHistory() {
+ return new InMemoryChatMemory();
+ }
+
+ @Bean
+ public OpenAiApi chatCompletionApi() {
+ return new OpenAiApi(System.getenv("OPENAI_API_KEY"));
+ }
+
+ @Bean
+ public OpenAiChatClient openAiClient(OpenAiApi openAiApi) {
+ return new OpenAiChatClient(openAiApi);
+ }
+
+ @Bean
+ public OpenAiEmbeddingClient embeddingClient(OpenAiApi openAiApi) {
+ return new OpenAiEmbeddingClient(openAiApi);
+ }
+
+ @Bean
+ public VectorStore qdrantVectorStore(EmbeddingClient embeddingClient) {
+ QdrantClient qdrantClient = new QdrantClient(QdrantGrpcClient
+ .newBuilder(qdrantContainer.getHost(), qdrantContainer.getMappedPort(QDRANT_GRPC_PORT), false)
+ .build());
+ return new QdrantVectorStore(qdrantClient, COLLECTION_NAME, embeddingClient);
+ }
+
+ @Bean
+ public TokenCountEstimator tokenCountEstimator() {
+ return new JTokkitTokenCountEstimator();
+ }
+
+ @Bean
+ public ChatBot memoryChatAgent(OpenAiChatClient chatClient, VectorStore vectorStore,
+ TokenCountEstimator tokenCountEstimator, ChatMemory chatHistory) {
+
+ return DefaultChatBot.builder(chatClient)
+ .withRetrievers(List.of(new VectorStoreRetriever(vectorStore, SearchRequest.defaults()),
+ new ChatMemoryRetriever(chatHistory, Map.of(TransformerContentType.SHORT_TERM_MEMORY, "")),
+ new VectorStoreChatMemoryRetriever(vectorStore, 10,
+ Map.of(TransformerContentType.LONG_TERM_MEMORY, ""))))
+
+ .withContentPostProcessors(List.of(
+ new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000,
+ Set.of(TransformerContentType.SHORT_TERM_MEMORY)),
+ new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000,
+ Set.of(TransformerContentType.LONG_TERM_MEMORY)),
+ new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 2000,
+ Set.of(TransformerContentType.EXTERNAL_KNOWLEDGE))))
+ .withAugmentors(List.of(new QuestionContextAugmentor(),
+ new SystemPromptChatMemoryAugmentor(
+ """
+ Use the long term conversation history from the LONG TERM HISTORY section to provide accurate answers.
+
+ LONG TERM HISTORY:
+ {history}
+ """,
+ Set.of(TransformerContentType.LONG_TERM_MEMORY)),
+ new SystemPromptChatMemoryAugmentor(Set.of(TransformerContentType.SHORT_TERM_MEMORY))))
+
+ .withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory),
+ new VectorStoreChatMemoryAgentListener(vectorStore,
+ Map.of(TransformerContentType.LONG_TERM_MEMORY, ""))))
+ .build();
+ }
+
+ // @Bean
+ // public StreamingChatBot memoryStreamingChatAgent(OpenAiChatClient
+ // streamingChatClient,
+ // VectorStore vectorStore, TokenCountEstimator tokenCountEstimator, ChatHistory
+ // chatHistory) {
+
+ // return DefaultStreamingChatBot.builder(streamingChatClient)
+ // .withRetrievers(List.of(new ChatHistoryRetriever(chatHistory), new
+ // DocumentChatHistoryRetriever(vectorStore, 10)))
+ // .withDocumentPostProcessors(List.of(new
+ // LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
+ // .withAugmentors(List.of(new TextChatHistoryAugmenter()))
+ // .withChatAgentListeners(List.of(new ChatHistoryAgentListener(chatHistory), new
+ // DocumentChatHistoryAgentListener(vectorStore)))
+ // .build();
+ // }
+
+ @Bean
+ public RelevancyEvaluator relevancyEvaluator(OpenAiChatClient chatClient) {
+ // Use GPT 4 as a better model for determining relevancy. gpt 3.5 makes basic
+ // mistakes
+ OpenAiChatOptions openAiChatOptions = OpenAiChatOptions.builder()
+ .withModel(GPT_4_TURBO_PREVIEW.getValue())
+ .build();
+ return new RelevancyEvaluator(chatClient, openAiChatOptions);
+ }
+
+ }
+
+}
diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/OpenAiDefaultChatBotIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/OpenAiDefaultChatBotIT.java
new file mode 100644
index 000000000..4c5aa313f
--- /dev/null
+++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/OpenAiDefaultChatBotIT.java
@@ -0,0 +1,160 @@
+/*
+ * Copyright 2024 - 2024 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.openai.chat.chatbot;
+
+import java.util.List;
+
+import io.qdrant.client.QdrantClient;
+import io.qdrant.client.QdrantGrpcClient;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
+import org.springframework.ai.chat.chatbot.ChatBot;
+import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.openai.OpenAiChatOptions;
+import org.testcontainers.junit.jupiter.Container;
+import org.testcontainers.junit.jupiter.Testcontainers;
+import org.testcontainers.qdrant.QdrantContainer;
+
+import org.springframework.ai.chat.ChatClient;
+import org.springframework.ai.chat.chatbot.DefaultChatBot;
+import org.springframework.ai.chat.messages.UserMessage;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+import org.springframework.ai.chat.prompt.transformer.QuestionContextAugmentor;
+import org.springframework.ai.chat.prompt.transformer.VectorStoreRetriever;
+import org.springframework.ai.embedding.EmbeddingClient;
+import org.springframework.ai.evaluation.EvaluationRequest;
+import org.springframework.ai.evaluation.EvaluationResponse;
+import org.springframework.ai.evaluation.RelevancyEvaluator;
+import org.springframework.ai.openai.OpenAiChatClient;
+import org.springframework.ai.openai.OpenAiEmbeddingClient;
+import org.springframework.ai.openai.api.OpenAiApi;
+import org.springframework.ai.reader.JsonReader;
+import org.springframework.ai.transformer.splitter.TokenTextSplitter;
+import org.springframework.ai.vectorstore.SearchRequest;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.boot.SpringBootConfiguration;
+import org.springframework.boot.test.context.SpringBootTest;
+import org.springframework.context.annotation.Bean;
+import org.springframework.core.io.Resource;
+
+import static org.junit.jupiter.api.Assertions.assertTrue;
+import static org.springframework.ai.openai.api.OpenAiApi.ChatModel.GPT_4_TURBO_PREVIEW;
+
+@Testcontainers
+@SpringBootTest(classes = OpenAiDefaultChatBotIT.Config.class)
+@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
+public class OpenAiDefaultChatBotIT {
+
+ private static final String COLLECTION_NAME = "test_collection";
+
+ private static final int QDRANT_GRPC_PORT = 6334;
+
+ @Container
+ static QdrantContainer qdrantContainer = new QdrantContainer("qdrant/qdrant:v1.7.4");
+
+ private final ChatClient chatClient;
+
+ private final VectorStore vectorStore;
+
+ @Value("classpath:/data/acme/bikes.json")
+ private Resource bikesResource;
+
+ private ChatBot chatBot;
+
+ @Autowired
+ public OpenAiDefaultChatBotIT(ChatClient chatClient, ChatBot chatBot, VectorStore vectorStore) {
+ this.chatClient = chatClient;
+ this.chatBot = chatBot;
+ this.vectorStore = vectorStore;
+ }
+
+ @Test
+ void simpleChat() {
+ loadData();
+
+ var prompt = new Prompt(new UserMessage("What bike is good for city commuting?"));
+ var chatBotResponse = this.chatBot.call(new PromptContext(prompt));
+ String answer = chatBotResponse.getChatResponse().getResult().getOutput().getContent();
+ assertTrue(answer.contains("Celerity"), "Response does not include 'Celerity'");
+
+ // Use GPT 4 as a better model for determining relevancy. gpt 3.5 makes basic
+ // mistakes
+ OpenAiChatOptions openAiChatOptions = OpenAiChatOptions.builder()
+ .withModel(GPT_4_TURBO_PREVIEW.getValue())
+ .build();
+ var relevancyEvaluator = new RelevancyEvaluator(this.chatClient, openAiChatOptions);
+ EvaluationRequest evaluationRequest = new EvaluationRequest(chatBotResponse);
+ EvaluationResponse evaluationResponse = relevancyEvaluator.evaluate(evaluationRequest);
+ assertTrue(evaluationResponse.isPass(), "Response is not relevant to the question");
+
+ }
+
+ void loadData() {
+ JsonReader jsonReader = new JsonReader(bikesResource, "name", "price", "shortDescription", "description");
+ var textSplitter = new TokenTextSplitter();
+ List splitDocuments = textSplitter.apply(jsonReader.get());
+
+ for (Document splitDocument : splitDocuments) {
+ splitDocument.getMetadata().put(TransformerContentType.EXTERNAL_KNOWLEDGE, "true");
+ }
+
+ vectorStore.accept(splitDocuments);
+ }
+
+ @SpringBootConfiguration
+ static class Config {
+
+ @Bean
+ public OpenAiApi chatCompletionApi() {
+ return new OpenAiApi(System.getenv("OPENAI_API_KEY"));
+ }
+
+ @Bean
+ public ChatClient openAiClient(OpenAiApi openAiApi) {
+ return new OpenAiChatClient(openAiApi);
+ }
+
+ @Bean
+ public EmbeddingClient embeddingClient(OpenAiApi openAiApi) {
+ return new OpenAiEmbeddingClient(openAiApi);
+ }
+
+ @Bean
+ public VectorStore qdrantVectorStore(EmbeddingClient embeddingClient) {
+ QdrantClient qdrantClient = new QdrantClient(QdrantGrpcClient
+ .newBuilder(qdrantContainer.getHost(), qdrantContainer.getMappedPort(QDRANT_GRPC_PORT), false)
+ .build());
+ return new QdrantVectorStore(qdrantClient, COLLECTION_NAME, embeddingClient);
+ }
+
+ @Bean
+ public ChatBot chatBot(ChatClient chatClient, VectorStore vectorStore) {
+ return DefaultChatBot.builder(chatClient)
+ .withRetrievers(List.of(new VectorStoreRetriever(vectorStore, SearchRequest.defaults())))
+ .withAugmentors(List.of(new QuestionContextAugmentor()))
+ .build();
+
+ }
+
+ }
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatAgentListener.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatAgentListener.java
new file mode 100644
index 000000000..68699a06f
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatAgentListener.java
@@ -0,0 +1,36 @@
+/*
+ * Copyright 2024 - 2024 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.chat.chatbot;
+
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+
+/**
+ * The ChatAgentListener is a callback interface that can be implemented by classes that
+ * want to be notified of the completion of a ChatBot execution.
+ *
+ * @author Mark Pollack
+ * @author Christian Tzolov
+ */
+public interface ChatAgentListener {
+
+ default void onStart(PromptContext promptContext) {
+
+ }
+
+ void onComplete(ChatBotResponse chatBotResponse);
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatBot.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatBot.java
new file mode 100644
index 000000000..ff4eeb510
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatBot.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2024 - 2024 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.chat.chatbot;
+
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+
+/**
+ * A ChatBot encapsulates the logic to perform common AI use cases such as Retrieval
+ * Augmented Generation.
+ *
+ * @author Mark Pollack
+ * @since 1.0 M1
+ */
+public interface ChatBot {
+
+ /**
+ * Call the chatbot to execute AI actions
+ * @param promptContext A shared data structure used by the ChatBot to perform
+ * processing of the Prompt. It includes the intial Prompt and a conversation ID at
+ * the start of execution.
+ * @return the ChatBotResponse that contains the ChatResponse and the latest
+ * PromptContext
+ */
+ ChatBotResponse call(PromptContext promptContext);
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatBotResponse.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatBotResponse.java
new file mode 100644
index 000000000..344d8d08b
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatBotResponse.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright 2024 - 2024 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.chat.chatbot;
+
+import org.springframework.ai.chat.ChatResponse;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+
+import java.util.Objects;
+
+/**
+ * Encapsulates the response from the ChatBot. Contains the most up-to-date PromptContext
+ * and the final ChatResponse
+ *
+ * @author Mark Pollack
+ * @since 1.0 M1
+ */
+public class ChatBotResponse {
+
+ private final PromptContext promptContext;
+
+ private final ChatResponse chatResponse;
+
+ public ChatBotResponse(PromptContext promptContext, ChatResponse chatResponse) {
+ this.promptContext = promptContext;
+ this.chatResponse = chatResponse;
+ }
+
+ public PromptContext getPromptContext() {
+ return promptContext;
+ }
+
+ public ChatResponse getChatResponse() {
+ return chatResponse;
+ }
+
+ @Override
+ public String toString() {
+ return "ChatBotResponse{" + "promptContext=" + promptContext + ", chatResponse=" + chatResponse + '}';
+ }
+
+ @Override
+ public boolean equals(Object o) {
+ if (this == o)
+ return true;
+ if (!(o instanceof ChatBotResponse that))
+ return false;
+ return Objects.equals(promptContext, that.promptContext) && Objects.equals(chatResponse, that.chatResponse);
+ }
+
+ @Override
+ public int hashCode() {
+ return Objects.hash(promptContext, chatResponse);
+ }
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/DefaultChatBot.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/DefaultChatBot.java
new file mode 100644
index 000000000..3c8559865
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/DefaultChatBot.java
@@ -0,0 +1,137 @@
+/*
+ * Copyright 2024 - 2024 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.chat.chatbot;
+
+import org.springframework.ai.chat.ChatClient;
+import org.springframework.ai.chat.ChatResponse;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+import org.springframework.ai.chat.prompt.transformer.PromptTransformer;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+
+/**
+ * @author Mark Pollack
+ * @author Christian Tzolov
+ */
+public class DefaultChatBot implements ChatBot {
+
+ private ChatClient chatClient;
+
+ private List retrievers;
+
+ private List documentPostProcessors;
+
+ private List augmentors;
+
+ private List chatAgentListeners;
+
+ public DefaultChatBot(ChatClient chatClient, List retrievers,
+ List documentPostProcessors, List augmentors,
+ List chatAgentListeners) {
+ Objects.requireNonNull(chatClient, "chatClient must not be null");
+ this.chatClient = chatClient;
+ this.retrievers = retrievers;
+ this.documentPostProcessors = documentPostProcessors;
+ this.augmentors = augmentors;
+ this.chatAgentListeners = chatAgentListeners;
+ }
+
+ public static DefaultChatAgentBuilder builder(ChatClient chatClient) {
+ return new DefaultChatAgentBuilder().withChatClient(chatClient);
+ }
+
+ @Override
+ public ChatBotResponse call(PromptContext promptContext) {
+
+ PromptContext promptContextOnStart = PromptContext.from(promptContext).build();
+
+ // Perform retrieval of documents and messages
+ for (PromptTransformer retriever : this.retrievers) {
+ promptContext = retriever.transform(promptContext);
+ }
+
+ // Perform post processing of all retrieved documents and messages
+ for (PromptTransformer documentPostProcessor : this.documentPostProcessors) {
+ promptContext = documentPostProcessor.transform(promptContext);
+ }
+
+ // Perform prompt augmentation
+ for (PromptTransformer augmentor : this.augmentors) {
+ promptContext = augmentor.transform(promptContext);
+ }
+
+ // Invoke Listeners onStart
+ for (ChatAgentListener listener : this.chatAgentListeners) {
+ listener.onStart(promptContextOnStart);
+ }
+
+ // Perform generation
+ ChatResponse chatResponse = this.chatClient.call(promptContext.getPrompt());
+
+ // Invoke Listeners onComplete
+ ChatBotResponse chatBotResponse = new ChatBotResponse(promptContext, chatResponse);
+ for (ChatAgentListener listener : this.chatAgentListeners) {
+ listener.onComplete(chatBotResponse);
+ }
+ return chatBotResponse;
+ }
+
+ public static class DefaultChatAgentBuilder {
+
+ private ChatClient chatClient;
+
+ private List retrievers = new ArrayList<>();
+
+ private List documentPostProcessors = new ArrayList<>();
+
+ private List augmentors = new ArrayList<>();
+
+ private List chatAgentListeners = new ArrayList<>();
+
+ public DefaultChatAgentBuilder withChatClient(ChatClient chatClient) {
+ this.chatClient = chatClient;
+ return this;
+ }
+
+ public DefaultChatAgentBuilder withRetrievers(List retrievers) {
+ this.retrievers = retrievers;
+ return this;
+ }
+
+ public DefaultChatAgentBuilder withContentPostProcessors(List documentPostProcessors) {
+ this.documentPostProcessors = documentPostProcessors;
+ return this;
+ }
+
+ public DefaultChatAgentBuilder withAugmentors(List augmentors) {
+ this.augmentors = augmentors;
+ return this;
+ }
+
+ public DefaultChatAgentBuilder withChatAgentListeners(List chatAgentListeners) {
+ this.chatAgentListeners = chatAgentListeners;
+ return this;
+ }
+
+ public DefaultChatBot build() {
+ return new DefaultChatBot(chatClient, retrievers, documentPostProcessors, augmentors, chatAgentListeners);
+ }
+
+ }
+
+}
\ No newline at end of file
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/DefaultStreamingChatBot.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/DefaultStreamingChatBot.java
new file mode 100644
index 000000000..15ea38f49
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/DefaultStreamingChatBot.java
@@ -0,0 +1,144 @@
+/*
+ * Copyright 2024 - 2024 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.chat.chatbot;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Objects;
+
+import reactor.core.publisher.Flux;
+
+import org.springframework.ai.chat.ChatResponse;
+import org.springframework.ai.chat.StreamingChatClient;
+import org.springframework.ai.chat.messages.MessageAggregator;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+import org.springframework.ai.chat.prompt.transformer.PromptTransformer;
+
+/**
+ * @author Mark Pollack
+ * @author Christian Tzolov
+ */
+public class DefaultStreamingChatBot implements StreamingChatBot {
+
+ private StreamingChatClient streamingChatClient;
+
+ private List retrievers;
+
+ private List documentPostProcessors;
+
+ private List augmentors;
+
+ private List chatAgentListeners;
+
+ public DefaultStreamingChatBot(StreamingChatClient chatClient, List retrievers,
+ List documentPostProcessors, List augmentors,
+ List chatAgentListeners) {
+ Objects.requireNonNull(chatClient, "chatClient must not be null");
+ this.streamingChatClient = chatClient;
+ this.retrievers = retrievers;
+ this.documentPostProcessors = documentPostProcessors;
+ this.augmentors = augmentors;
+ this.chatAgentListeners = chatAgentListeners;
+ }
+
+ public static DefaultChatAgentBuilder builder(StreamingChatClient chatClient) {
+ return new DefaultChatAgentBuilder().withChatClient(chatClient);
+ }
+
+ @Override
+ public StreamingChatBotResponse stream(PromptContext promptContext) {
+
+ PromptContext promptContextOnStart = PromptContext.from(promptContext).build();
+
+ // Perform retrieval of documents and messages
+ for (PromptTransformer retriever : this.retrievers) {
+ promptContext = retriever.transform(promptContext);
+ }
+
+ // Perform post processing of all retrieved documents and messages
+ for (PromptTransformer documentPostProcessor : this.documentPostProcessors) {
+ promptContext = documentPostProcessor.transform(promptContext);
+ }
+
+ // Perform prompt augmentation
+ for (PromptTransformer augmentor : this.augmentors) {
+ promptContext = augmentor.transform(promptContext);
+ }
+
+ // Invoke Listeners onStart
+ for (ChatAgentListener listener : this.chatAgentListeners) {
+ listener.onStart(promptContextOnStart);
+ }
+
+ // Perform generation
+ final var promptContext2 = promptContext;
+
+ Flux fluxChatResponse = new MessageAggregator()
+ .aggregate(this.streamingChatClient.stream(promptContext.getPrompt()), chatResponse -> {
+ for (ChatAgentListener listener : this.chatAgentListeners) {
+ listener.onComplete(new ChatBotResponse(promptContext2, chatResponse));
+ }
+ });
+
+ // Invoke Listeners onComplete
+ return new StreamingChatBotResponse(promptContext, fluxChatResponse);
+ }
+
+ public static class DefaultChatAgentBuilder {
+
+ private StreamingChatClient chatClient;
+
+ private List retrievers = new ArrayList<>();
+
+ private List documentPostProcessors = new ArrayList<>();
+
+ private List augmentors = new ArrayList<>();
+
+ private List chatAgentListeners = new ArrayList<>();
+
+ public DefaultChatAgentBuilder withChatClient(StreamingChatClient chatClient) {
+ this.chatClient = chatClient;
+ return this;
+ }
+
+ public DefaultChatAgentBuilder withRetrievers(List retrievers) {
+ this.retrievers = retrievers;
+ return this;
+ }
+
+ public DefaultChatAgentBuilder withDocumentPostProcessors(List documentPostProcessors) {
+ this.documentPostProcessors = documentPostProcessors;
+ return this;
+ }
+
+ public DefaultChatAgentBuilder withAugmentors(List augmentors) {
+ this.augmentors = augmentors;
+ return this;
+ }
+
+ public DefaultChatAgentBuilder withChatAgentListeners(List chatAgentListeners) {
+ this.chatAgentListeners = chatAgentListeners;
+ return this;
+ }
+
+ public DefaultStreamingChatBot build() {
+ return new DefaultStreamingChatBot(chatClient, retrievers, documentPostProcessors, augmentors,
+ chatAgentListeners);
+ }
+
+ }
+
+}
\ No newline at end of file
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/StreamingChatBot.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/StreamingChatBot.java
new file mode 100644
index 000000000..d5d4843b1
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/StreamingChatBot.java
@@ -0,0 +1,40 @@
+/*
+ * Copyright 2024 - 2024 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.chat.chatbot;
+
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+
+/**
+ * A ChatBot encapsulates the logic to perform common AI use cases such as Retrieval
+ * Augmented Generation.
+ *
+ * @author Mark Pollack
+ * @author Christian Tzolov
+ * @since 1.0 M1
+ */
+public interface StreamingChatBot {
+
+ /**
+ * Call the chatbot to execute AI actions
+ * @param promptContext A shared data structure used by the ChatBot to perform
+ * processing of the Prompt. It includes the intial Prompt and a conversation ID at
+ * the start of execution.
+ * @return the StreamingChatBotResponse that contains the ChatResponse and the latest
+ * PromptContext
+ */
+ StreamingChatBotResponse stream(PromptContext promptContext);
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/StreamingChatBotResponse.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/StreamingChatBotResponse.java
new file mode 100644
index 000000000..5e699278c
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/StreamingChatBotResponse.java
@@ -0,0 +1,72 @@
+package org.springframework.ai.chat.chatbot;
+
+import reactor.core.publisher.Flux;
+
+import org.springframework.ai.chat.ChatResponse;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+
+/**
+ * Encapsulates the response from the ChatBot. Contains the most up-to-date PromptContext
+ * and the final ChatResponse
+ *
+ * @author Mark Pollack
+ * @since 1.0 M1
+ */
+public class StreamingChatBotResponse {
+
+ private final PromptContext promptContext;
+
+ private final Flux chatResponse;
+
+ public StreamingChatBotResponse(PromptContext promptContext, Flux chatResponse) {
+ this.promptContext = promptContext;
+ this.chatResponse = chatResponse;
+ }
+
+ public PromptContext getPromptContext() {
+ return promptContext;
+ }
+
+ public Flux getChatResponse() {
+ return chatResponse;
+ }
+
+ @Override
+ public String toString() {
+ return "ChatBotResponse{" + "promptContext=" + promptContext + ", chatResponse=" + chatResponse + '}';
+ }
+
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((promptContext == null) ? 0 : promptContext.hashCode());
+ result = prime * result + ((chatResponse == null) ? 0 : chatResponse.hashCode());
+ return result;
+ }
+
+ @Override
+ public boolean equals(Object obj) {
+ if (this == obj)
+ return true;
+ if (obj == null)
+ return false;
+ if (getClass() != obj.getClass())
+ return false;
+ StreamingChatBotResponse other = (StreamingChatBotResponse) obj;
+ if (promptContext == null) {
+ if (other.promptContext != null)
+ return false;
+ }
+ else if (!promptContext.equals(other.promptContext))
+ return false;
+ if (chatResponse == null) {
+ if (other.chatResponse != null)
+ return false;
+ }
+ else if (!chatResponse.equals(other.chatResponse))
+ return false;
+ return true;
+ }
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemory.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemory.java
new file mode 100644
index 000000000..26d4f1342
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemory.java
@@ -0,0 +1,39 @@
+/*
+ * Copyright 2024-2024 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.chat.history;
+
+import java.util.List;
+
+import org.springframework.ai.chat.messages.Message;
+
+/**
+ * @author Christian Tzolov
+ *
+ */
+public interface ChatMemory {
+
+ default void add(String conversationId, Message message) {
+ this.add(conversationId, List.of(message));
+ }
+
+ void add(String conversationId, List messages);
+
+ List get(String conversationId, int lastN);
+
+ void clear(String conversationId);
+
+}
\ No newline at end of file
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemoryAgentListener.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemoryAgentListener.java
new file mode 100644
index 000000000..6b92c365c
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemoryAgentListener.java
@@ -0,0 +1,60 @@
+/*
+ * Copyright 2024-2024 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.chat.history;
+
+import java.util.List;
+
+import org.springframework.ai.chat.chatbot.ChatBotResponse;
+import org.springframework.ai.chat.chatbot.ChatAgentListener;
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.messages.MessageType;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
+
+/**
+ * @author Christian Tzolov
+ */
+public class ChatMemoryAgentListener implements ChatAgentListener {
+
+ private final ChatMemory chatHistory;
+
+ public ChatMemoryAgentListener(ChatMemory chatHistory) {
+ this.chatHistory = chatHistory;
+ }
+
+ @Override
+ public void onStart(PromptContext promptContext) {
+ var messagesToAdd = promptContext.getPrompt()
+ .getInstructions()
+ .stream()
+ .filter(m -> !m.getMetadata().containsKey(TransformerContentType.MEMORY))
+ .filter(m -> (m.getMessageType() == MessageType.ASSISTANT || m.getMessageType() == MessageType.USER))
+ .toList();
+ this.chatHistory.add(promptContext.getConversationId(), messagesToAdd);
+ }
+
+ @Override
+ public void onComplete(ChatBotResponse chatBotResponse) {
+ List assistantMessages = chatBotResponse.getChatResponse()
+ .getResults()
+ .stream()
+ .map(g -> (Message) g.getOutput())
+ .toList();
+ this.chatHistory.add(chatBotResponse.getPromptContext().getConversationId(), assistantMessages);
+ }
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemoryRetriever.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemoryRetriever.java
new file mode 100644
index 000000000..31d083846
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemoryRetriever.java
@@ -0,0 +1,81 @@
+/*
+ * Copyright 2024-2024 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.chat.history;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.messages.MessageType;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+import org.springframework.ai.chat.prompt.transformer.PromptTransformer;
+import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.model.Content;
+
+/**
+ * @author Christian Tzolov
+ */
+public class ChatMemoryRetriever implements PromptTransformer {
+
+ private final ChatMemory chatHistory;
+
+ /**
+ * Additional metadata to be assigned to the retrieved history messages.
+ */
+ private final Map additionalMetadata;
+
+ private final int maxHistorySize;
+
+ public ChatMemoryRetriever(ChatMemory chatHistory) {
+ this(chatHistory, Map.of());
+ }
+
+ public ChatMemoryRetriever(ChatMemory chatHistory, Map additionalMetadata) {
+ this(chatHistory, 1000, additionalMetadata);
+ }
+
+ public ChatMemoryRetriever(ChatMemory chatHistory, int maxHistorySize, Map additionalMetadata) {
+ this.chatHistory = chatHistory;
+ this.additionalMetadata = additionalMetadata;
+ this.maxHistorySize = maxHistorySize;
+ }
+
+ @Override
+ public PromptContext transform(PromptContext promptContext) {
+
+ List messageHistory = this.chatHistory.get(promptContext.getConversationId(), maxHistorySize);
+
+ List historyContent = (messageHistory != null)
+ ? messageHistory.stream().filter(m -> m.getMessageType() != MessageType.SYSTEM).map(m -> {
+ Content content = new Document(m.getContent(), new ArrayList<>(m.getMedia()),
+ new HashMap<>(m.getMetadata()));
+ content.getMetadata().putAll(this.additionalMetadata);
+ content.getMetadata().put(TransformerContentType.MEMORY, true);
+ return content;
+ }).toList() : List.of();
+
+ List updatedContents = new ArrayList<>(
+ promptContext.getContents() != null ? promptContext.getContents() : List.of());
+ updatedContents.addAll(historyContent);
+
+ return PromptContext.from(promptContext).withContents(updatedContents).build();
+ }
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/history/InMemoryChatMemory.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/InMemoryChatMemory.java
new file mode 100644
index 000000000..80bf6671c
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/InMemoryChatMemory.java
@@ -0,0 +1,50 @@
+/*
+ * Copyright 2024-2024 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.chat.history;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import org.springframework.ai.chat.messages.Message;
+
+/**
+ * @author Christian Tzolov
+ */
+public class InMemoryChatMemory implements ChatMemory {
+
+ Map> conversationHistory = new ConcurrentHashMap<>();
+
+ @Override
+ public void add(String conversationId, List messages) {
+ this.conversationHistory.putIfAbsent(conversationId, new ArrayList<>());
+ this.conversationHistory.get(conversationId).addAll(messages);
+ }
+
+ @Override
+ public List get(String conversationId, int lastN) {
+ List all = this.conversationHistory.get(conversationId);
+ return all != null ? all.stream().skip(Math.max(0, all.size() - lastN)).toList() : List.of();
+ }
+
+ @Override
+ public void clear(String conversationId) {
+ this.conversationHistory.remove(conversationId);
+ }
+
+}
\ No newline at end of file
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/history/LastMaxTokenSizeContentTransformer.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/LastMaxTokenSizeContentTransformer.java
new file mode 100644
index 000000000..01b07ecd4
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/LastMaxTokenSizeContentTransformer.java
@@ -0,0 +1,121 @@
+/*
+ * Copyright 2024-2024 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.chat.history;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Set;
+
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+import org.springframework.ai.chat.prompt.transformer.PromptTransformer;
+import org.springframework.ai.model.Content;
+import org.springframework.ai.tokenizer.TokenCountEstimator;
+
+/**
+ * Returns a new list of content (e.g list of messages of list of documents) that is a
+ * subset of the input list of contents and complies with the max token size constraint.
+ *
+ * The token estimator is used to estimate the token count of the datum.
+ *
+ * @author Christian Tzolov
+ */
+public class LastMaxTokenSizeContentTransformer implements PromptTransformer {
+
+ protected final TokenCountEstimator tokenCountEstimator;
+
+ protected final int maxTokenSize;
+
+ /**
+ * Only Content entries with the following metadata tags will be included in the
+ * history.
+ */
+ private final Set filterTags;
+
+ public LastMaxTokenSizeContentTransformer(TokenCountEstimator tokenCountEstimator, int maxTokenSize) {
+ this(tokenCountEstimator, maxTokenSize, Set.of());
+ }
+
+ public LastMaxTokenSizeContentTransformer(TokenCountEstimator tokenCountEstimator, int maxTokenSize,
+ Set filterTags) {
+ this.tokenCountEstimator = tokenCountEstimator;
+ this.maxTokenSize = maxTokenSize;
+ this.filterTags = filterTags;
+ }
+
+ protected List doGetDatumToModify(PromptContext promptContext) {
+ return promptContext.getContents()
+ .stream()
+ .filter(content -> this.filterTags.stream().allMatch(tag -> content.getMetadata().containsKey(tag)))
+ .toList();
+ }
+
+ protected List doGetDatumNotToModify(PromptContext promptContext) {
+ return promptContext.getContents()
+ .stream()
+ .filter(content -> !this.filterTags.stream().allMatch(tag -> content.getMetadata().containsKey(tag)))
+ .toList();
+ }
+
+ protected int doEstimateTokenCount(Content datum) {
+ return this.tokenCountEstimator.estimate(datum);
+ }
+
+ protected int doEstimateTokenCount(List datum) {
+ return datum.stream().mapToInt(this::doEstimateTokenCount).sum();
+ }
+
+ @Override
+ public PromptContext transform(PromptContext promptContext) {
+
+ List datum = this.doGetDatumToModify(promptContext);
+
+ int totalSize = this.doEstimateTokenCount(datum);
+
+ if (totalSize <= this.maxTokenSize) {
+ return promptContext;
+ }
+
+ List purgedContent = this.purgeExcess(datum, totalSize);
+
+ var updatedContent = new ArrayList<>(doGetDatumNotToModify(promptContext));
+ updatedContent.addAll(purgedContent);
+
+ return PromptContext.from(promptContext).withContents(updatedContent).build();
+ }
+
+ protected List purgeExcess(List datum, int totalSize) {
+
+ int index = 0;
+ List newList = new ArrayList<>();
+
+ while (index < datum.size() && totalSize > this.maxTokenSize) {
+ Content oldDatum = datum.get(index++);
+ int oldMessageTokenSize = this.doEstimateTokenCount(oldDatum);
+ totalSize = totalSize - oldMessageTokenSize;
+ }
+
+ if (index >= datum.size()) {
+ return List.of();
+ }
+
+ // add the rest of the messages.
+ newList.addAll(datum.subList(index, datum.size()));
+
+ return newList;
+ }
+
+}
\ No newline at end of file
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/history/MessageChatMemoryAugmentor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/MessageChatMemoryAugmentor.java
new file mode 100644
index 000000000..dcf3d6daa
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/MessageChatMemoryAugmentor.java
@@ -0,0 +1,70 @@
+/*
+ * Copyright 2024-2024 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.chat.history;
+
+import java.util.ArrayList;
+import java.util.List;
+
+import org.springframework.ai.chat.messages.AbstractMessage;
+import org.springframework.ai.chat.messages.AssistantMessage;
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.messages.MessageType;
+import org.springframework.ai.chat.messages.UserMessage;
+import org.springframework.ai.chat.prompt.ChatOptions;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+import org.springframework.ai.chat.prompt.transformer.PromptTransformer;
+
+/**
+ * @author Christian Tzolov
+ */
+public class MessageChatMemoryAugmentor implements PromptTransformer {
+
+ @Override
+ public PromptContext transform(PromptContext promptContext) {
+
+ var originalPrompt = promptContext.getPrompt();
+
+ // Convert the retrieved contents into a list of messages.
+ List historyMessages = promptContext.getContents()
+ .stream()
+ .filter(content -> content.getMetadata().containsKey(TransformerContentType.MEMORY))
+ .map(content -> {
+ MessageType messageType = MessageType
+ .valueOf("" + content.getMetadata().get(AbstractMessage.MESSAGE_TYPE));
+ Message message = null;
+ if (messageType == MessageType.ASSISTANT) {
+ message = new AssistantMessage(content.getContent(), content.getMetadata());
+ }
+ else if (messageType == MessageType.USER) {
+ message = new UserMessage(content.getContent(), List.of(), content.getMetadata());
+ }
+ return message;
+ })
+ .filter(m -> m != null)
+ .toList();
+
+ var promptMessages = new ArrayList<>(historyMessages);
+ promptMessages.addAll(originalPrompt.getInstructions());
+
+ Prompt newPrompt = new Prompt(promptMessages, (ChatOptions) originalPrompt.getOptions());
+
+ return PromptContext.from(promptContext).withPrompt(newPrompt).addPromptHistory(originalPrompt).build();
+ }
+
+}
\ No newline at end of file
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/history/SystemPromptChatMemoryAugmentor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/SystemPromptChatMemoryAugmentor.java
new file mode 100644
index 000000000..b507faf16
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/SystemPromptChatMemoryAugmentor.java
@@ -0,0 +1,110 @@
+/*
+ * Copyright 2024-2024 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.chat.history;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+import org.springframework.ai.chat.messages.AbstractMessage;
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.messages.MessageType;
+import org.springframework.ai.chat.messages.SystemMessage;
+import org.springframework.ai.chat.prompt.ChatOptions;
+import org.springframework.ai.chat.prompt.Prompt;
+import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+import org.springframework.ai.chat.prompt.transformer.PromptTransformer;
+import org.springframework.util.Assert;
+
+/**
+ * @author Christian Tzolov
+ */
+public class SystemPromptChatMemoryAugmentor implements PromptTransformer {
+
+ public static final String DEFAULT_HISTORY_PROMPT = """
+ Use the conversation history from the HISTORY section to provide accurate answers.
+
+ HISTORY:
+ {history}
+ """;
+
+ private final String historyPrompt;
+
+ /**
+ * Only Content entries with the following metadata tags will be included in the
+ * history.
+ */
+ private final Set filterTags;
+
+ public SystemPromptChatMemoryAugmentor() {
+ this(DEFAULT_HISTORY_PROMPT, new HashSet<>());
+ }
+
+ public SystemPromptChatMemoryAugmentor(Set filterTags) {
+ this(DEFAULT_HISTORY_PROMPT, filterTags);
+ }
+
+ public SystemPromptChatMemoryAugmentor(String historyPrompt, Set metadataFilterTags) {
+ Assert.hasText(historyPrompt, "The historyPrompt must not be empty!");
+ Assert.notNull(metadataFilterTags, "The metadataFilterTags must not be null!");
+ this.historyPrompt = historyPrompt;
+ this.filterTags = new HashSet<>(metadataFilterTags);
+
+ // Always include the message history type tag.
+ this.filterTags.add(TransformerContentType.MEMORY);
+ }
+
+ @Override
+ public PromptContext transform(PromptContext promptContext) {
+
+ var originalPrompt = promptContext.getPrompt();
+
+ List systemMessages = (originalPrompt.getInstructions() != null) ? originalPrompt.getInstructions()
+ .stream()
+ .filter(m -> m.getMessageType() == MessageType.SYSTEM)
+ .toList() : List.of();
+
+ List nonSystemMessages = (originalPrompt.getInstructions() != null) ? originalPrompt.getInstructions()
+ .stream()
+ .filter(m -> m.getMessageType() != MessageType.SYSTEM)
+ .toList() : List.of();
+
+ SystemMessage originalSystemMessage = (!systemMessages.isEmpty()) ? (SystemMessage) systemMessages.get(0)
+ : new SystemMessage("");
+
+ String historyContext = promptContext.getContents()
+ .stream()
+ .filter(content -> this.filterTags.stream().allMatch(tag -> content.getMetadata().containsKey(tag)))
+ .map(content -> content.getMetadata().get(AbstractMessage.MESSAGE_TYPE) + ": " + content.getContent())
+ .collect(Collectors.joining(System.lineSeparator()));
+
+ SystemMessage newSystemMessage = new SystemMessage(originalSystemMessage.getContent() + System.lineSeparator()
+ + this.historyPrompt.replace("{history}", historyContext));
+
+ List newPromptMessages = new ArrayList<>();
+ newPromptMessages.add(newSystemMessage);
+ newPromptMessages.addAll(nonSystemMessages);
+
+ Prompt newPrompt = new Prompt(newPromptMessages, (ChatOptions) originalPrompt.getOptions());
+
+ return PromptContext.from(promptContext).withPrompt(newPrompt).addPromptHistory(originalPrompt).build();
+ }
+
+}
\ No newline at end of file
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/history/VectorStoreChatMemoryAgentListener.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/VectorStoreChatMemoryAgentListener.java
new file mode 100644
index 000000000..72b32c9c5
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/VectorStoreChatMemoryAgentListener.java
@@ -0,0 +1,98 @@
+/*
+ * Copyright 2024-2024 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.chat.history;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.springframework.ai.chat.chatbot.ChatBotResponse;
+import org.springframework.ai.chat.chatbot.ChatAgentListener;
+import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.messages.MessageType;
+import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.util.CollectionUtils;
+
+/**
+ * @author Christian Tzolov
+ */
+public class VectorStoreChatMemoryAgentListener implements ChatAgentListener {
+
+ private final VectorStore vectorStore;
+
+ private final Map additionalMetadata;
+
+ public VectorStoreChatMemoryAgentListener(VectorStore vectorStore) {
+ this(vectorStore, new HashMap<>());
+ }
+
+ public VectorStoreChatMemoryAgentListener(VectorStore vectorStore, Map additionalMetadata) {
+ this.vectorStore = vectorStore;
+ this.additionalMetadata = additionalMetadata;
+ }
+
+ @Override
+ public void onStart(PromptContext promptContext) {
+
+ if (!CollectionUtils.isEmpty(promptContext.getPrompt().getInstructions())) {
+ List docs = toDocuments(promptContext.getPrompt().getInstructions(),
+ promptContext.getConversationId());
+
+ this.vectorStore.add(docs);
+ }
+ }
+
+ @Override
+ public void onComplete(ChatBotResponse chatBotResponse) {
+ if (!CollectionUtils.isEmpty(chatBotResponse.getChatResponse().getResults())) {
+ List assistantMessages = chatBotResponse.getChatResponse()
+ .getResults()
+ .stream()
+ .map(g -> (org.springframework.ai.chat.messages.Message) g.getOutput())
+ .toList();
+
+ List docs = toDocuments(assistantMessages,
+ chatBotResponse.getPromptContext().getConversationId());
+
+ this.vectorStore.add(docs);
+ }
+ }
+
+ private List toDocuments(List messages, String conversationId) {
+
+ List docs = messages.stream()
+ .filter(m -> m.getMessageType() == MessageType.USER || m.getMessageType() == MessageType.ASSISTANT)
+ .map(message -> {
+ var metadata = new HashMap<>(message.getMetadata() != null ? message.getMetadata() : new HashMap<>());
+ metadata.putAll(this.additionalMetadata);
+ metadata.put(TransformerContentType.CONVERSATION_ID, conversationId);
+ metadata.put("messageType", message.getMessageType().name());
+ metadata.put(TransformerContentType.MEMORY, true);
+ metadata.put(TransformerContentType.LONG_TERM_MEMORY, true);
+ var doc = new Document(message.getContent(), metadata);
+ return doc;
+ })
+ .toList();
+
+ return docs;
+
+ }
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/history/VectorStoreChatMemoryRetriever.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/VectorStoreChatMemoryRetriever.java
new file mode 100644
index 000000000..aff050179
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/VectorStoreChatMemoryRetriever.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright 2024-2024 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.chat.history;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import org.springframework.ai.chat.messages.MessageType;
+import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
+import org.springframework.ai.chat.prompt.transformer.PromptContext;
+import org.springframework.ai.chat.prompt.transformer.PromptTransformer;
+import org.springframework.ai.document.Document;
+import org.springframework.ai.model.Content;
+import org.springframework.ai.vectorstore.SearchRequest;
+import org.springframework.ai.vectorstore.VectorStore;
+import org.springframework.util.CollectionUtils;
+
+/**
+ * @author Christian Tzolov
+ */
+public class VectorStoreChatMemoryRetriever implements PromptTransformer {
+
+ private final VectorStore vectorStore;
+
+ private final int topK;
+
+ /**
+ * Additional metadata to be assigned to the retrieved history messages.
+ */
+ private final Map additionalMetadata;
+
+ public VectorStoreChatMemoryRetriever(VectorStore vectorStore, int topK) {
+ this(vectorStore, topK, Map.of());
+ }
+
+ public VectorStoreChatMemoryRetriever(VectorStore vectorStore, int topK, Map additionalMetadata) {
+ this.vectorStore = vectorStore;
+ this.topK = topK;
+ this.additionalMetadata = additionalMetadata;
+ }
+
+ @Override
+ public PromptContext transform(PromptContext promptContext) {
+ List updatedContents = new ArrayList<>(
+ promptContext.getContents() != null ? promptContext.getContents() : List.of());
+
+ String query = promptContext.getPrompt()
+ .getInstructions()
+ .stream()
+ .filter(m -> m.getMessageType() == MessageType.USER)
+ .map(m -> m.getContent())
+ .collect(Collectors.joining());
+
+ var searchRequest = SearchRequest.query(query)
+ .withTopK(this.topK)
+ .withFilterExpression(
+ TransformerContentType.CONVERSATION_ID + "=='" + promptContext.getConversationId() + "'");
+
+ List documents = this.vectorStore.similaritySearch(searchRequest);
+
+ if (!CollectionUtils.isEmpty(documents)) {
+ documents.forEach(d -> {
+ d.getMetadata().putAll(this.additionalMetadata);
+ d.getMetadata().put(TransformerContentType.MEMORY, true);
+ });
+ updatedContents.addAll(documents);
+ }
+
+ return PromptContext.from(promptContext).withContents(updatedContents).build();
+ }
+
+}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AbstractMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AbstractMessage.java
index 77b54afea..3c0d7a855 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AbstractMessage.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AbstractMessage.java
@@ -20,8 +20,10 @@ import java.io.InputStream;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Collections;
+import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.Objects;
import org.springframework.core.io.Resource;
import org.springframework.util.Assert;
@@ -29,13 +31,15 @@ import org.springframework.util.StreamUtils;
/**
* The AbstractMessage class is an abstract implementation of the Message interface. It
- * provides a base implementation for message content, media attachments, properties, and
+ * provides a base implementation for message content, media attachments, metadata, and
* message type.
*
* @see Message
*/
public abstract class AbstractMessage implements Message {
+ public static final String MESSAGE_TYPE = "messageType";
+
protected final MessageType messageType;
protected final String textContent;
@@ -45,26 +49,27 @@ public abstract class AbstractMessage implements Message {
/**
* Additional options for the message to influence the response, not a generative map.
*/
- protected final Map properties;
+ protected final Map metadata;
protected AbstractMessage(MessageType messageType, String content) {
- this(messageType, content, Map.of());
+ this(messageType, content, Map.of(MESSAGE_TYPE, messageType));
}
- protected AbstractMessage(MessageType messageType, String content, Map messageProperties) {
+ protected AbstractMessage(MessageType messageType, String content, Map metadata) {
Assert.notNull(messageType, "Message type must not be null");
this.messageType = messageType;
this.textContent = content;
this.mediaData = new ArrayList<>();
- this.properties = messageProperties;
+ this.metadata = new HashMap<>(metadata);
+ this.metadata.put(MESSAGE_TYPE, messageType);
}
protected AbstractMessage(MessageType messageType, String textContent, List mediaData) {
- this(messageType, textContent, mediaData, Map.of());
+ this(messageType, textContent, mediaData, Map.of(MESSAGE_TYPE, messageType));
}
protected AbstractMessage(MessageType messageType, String textContent, List mediaData,
- Map messageProperties) {
+ Map metadata) {
Assert.notNull(messageType, "Message type must not be null");
Assert.notNull(textContent, "Content must not be null");
@@ -73,7 +78,8 @@ public abstract class AbstractMessage implements Message {
this.messageType = messageType;
this.textContent = textContent;
this.mediaData = new ArrayList<>(mediaData);
- this.properties = messageProperties;
+ this.metadata = new HashMap<>(metadata);
+ this.metadata.put(MESSAGE_TYPE, messageType);
}
protected AbstractMessage(MessageType messageType, Resource resource) {
@@ -81,12 +87,13 @@ public abstract class AbstractMessage implements Message {
}
@SuppressWarnings("null")
- protected AbstractMessage(MessageType messageType, Resource resource, Map messageProperties) {
+ protected AbstractMessage(MessageType messageType, Resource resource, Map metadata) {
Assert.notNull(messageType, "Message type must not be null");
Assert.notNull(resource, "Resource must not be null");
this.messageType = messageType;
- this.properties = messageProperties;
+ this.metadata = new HashMap<>(metadata);
+ this.metadata.put(MESSAGE_TYPE, messageType);
this.mediaData = new ArrayList<>();
try (InputStream inputStream = resource.getInputStream()) {
@@ -108,8 +115,8 @@ public abstract class AbstractMessage implements Message {
}
@Override
- public Map getProperties() {
- return this.properties;
+ public Map getMetadata() {
+ return this.metadata;
}
@Override
@@ -119,38 +126,21 @@ public abstract class AbstractMessage implements Message {
@Override
public int hashCode() {
- final int prime = 31;
- int result = 1;
- result = prime * result + ((mediaData == null) ? 0 : mediaData.hashCode());
- result = prime * result + ((properties == null) ? 0 : properties.hashCode());
- result = prime * result + ((messageType == null) ? 0 : messageType.hashCode());
- return result;
+ return Objects.hash(this.messageType, this.textContent, this.mediaData, this.metadata);
}
@Override
public boolean equals(Object obj) {
- if (this == obj)
+ if (this == obj) {
return true;
- if (obj == null)
- return false;
- if (getClass() != obj.getClass())
+ }
+ if (obj == null || getClass() != obj.getClass()) {
return false;
+ }
AbstractMessage other = (AbstractMessage) obj;
- if (mediaData == null) {
- if (other.mediaData != null)
- return false;
- }
- else if (!mediaData.equals(other.mediaData))
- return false;
- if (properties == null) {
- if (other.properties != null)
- return false;
- }
- else if (!properties.equals(other.properties))
- return false;
- if (messageType != other.messageType)
- return false;
- return true;
+ return Objects.equals(this.messageType, other.messageType)
+ && Objects.equals(this.textContent, other.textContent)
+ && Objects.equals(this.mediaData, other.mediaData) && Objects.equals(this.metadata, other.metadata);
}
}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AssistantMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AssistantMessage.java
index c5f6831ab..f8b890416 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AssistantMessage.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AssistantMessage.java
@@ -35,7 +35,7 @@ public class AssistantMessage extends AbstractMessage {
@Override
public String toString() {
- return "AssistantMessage{" + "content='" + getContent() + '\'' + ", properties=" + properties + ", messageType="
+ return "AssistantMessage{" + "content='" + getContent() + '\'' + ", properties=" + metadata + ", messageType="
+ messageType + '}';
}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/FunctionMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/FunctionMessage.java
index 06485ac57..a05ef5226 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/FunctionMessage.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/FunctionMessage.java
@@ -33,7 +33,7 @@ public class FunctionMessage extends AbstractMessage {
@Override
public String toString() {
- return "FunctionMessage{" + "content='" + getContent() + '\'' + ", properties=" + properties + ", messageType="
+ return "FunctionMessage{" + "content='" + getContent() + '\'' + ", properties=" + metadata + ", messageType="
+ messageType + '}';
}
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/Message.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/Message.java
index 77e3b5aba..0945ba8de 100644
--- a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/Message.java
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/Message.java
@@ -15,8 +15,7 @@
*/
package org.springframework.ai.chat.messages;
-import java.util.List;
-import java.util.Map;
+import org.springframework.ai.model.Content;
/**
* The Message interface represents a message that can be sent or received in a chat
@@ -26,13 +25,7 @@ import java.util.Map;
* @see Media
* @see MessageType
*/
-public interface Message {
-
- String getContent();
-
- List getMedia();
-
- Map getProperties();
+public interface Message extends Content {
MessageType getMessageType();
diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/MessageAggregator.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/MessageAggregator.java
new file mode 100644
index 000000000..07aed7a98
--- /dev/null
+++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/MessageAggregator.java
@@ -0,0 +1,74 @@
+/*
+ * Copyright 2024-2024 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.chat.messages;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import reactor.core.publisher.Flux;
+
+import org.springframework.ai.chat.ChatResponse;
+import org.springframework.ai.chat.Generation;
+
+/**
+ * Helper that for streaming chat responses, aggregate the chat response messages into a
+ * single AssistantMessage. Job is performed in parallel to the chat response processing.
+ *
+ * @author Christian Tzolov
+ * @since 1.0.0
+ */
+public class MessageAggregator {
+
+ private static final Logger logger = LoggerFactory.getLogger(MessageAggregator.class);
+
+ public Flux aggregate(Flux fluxChatResponse,
+ Consumer onAggregationComplete) {
+
+ AtomicReference stringBufferRef = new AtomicReference<>(new StringBuilder());
+ AtomicReference
diff --git a/vector-stores/spring-ai-qdrant/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java b/vector-stores/spring-ai-qdrant/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java
index bc0bfc3c7..b7a8eada1 100644
--- a/vector-stores/spring-ai-qdrant/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java
+++ b/vector-stores/spring-ai-qdrant/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java
@@ -21,20 +21,23 @@ import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ExecutionException;
+import com.azure.ai.openai.OpenAIClient;
+import com.azure.ai.openai.OpenAIClientBuilder;
+import com.azure.core.credential.AzureKeyCredential;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import io.qdrant.client.grpc.Collections.Distance;
import io.qdrant.client.grpc.Collections.VectorParams;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.qdrant.QdrantContainer;
+import org.springframework.ai.azure.openai.AzureOpenAiEmbeddingClient;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingClient;
-import org.springframework.ai.openai.OpenAiEmbeddingClient;
-import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.boot.SpringBootConfiguration;
@@ -48,6 +51,8 @@ import static org.assertj.core.api.Assertions.assertThat;
* @since 0.8.1
*/
@Testcontainers
+@EnabledIfEnvironmentVariable(named = "AZURE_OPENAI_API_KEY", matches = ".+")
+@EnabledIfEnvironmentVariable(named = "AZURE_OPENAI_ENDPOINT", matches = ".+")
public class QdrantVectorStoreIT {
private static final String COLLECTION_NAME = "test_collection";
@@ -250,8 +255,15 @@ public class QdrantVectorStoreIT {
}
@Bean
- public EmbeddingClient embeddingClient() {
- return new OpenAiEmbeddingClient(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
+ public OpenAIClient openAIClient() {
+ return new OpenAIClientBuilder().credential(new AzureKeyCredential(System.getenv("AZURE_OPENAI_API_KEY")))
+ .endpoint(System.getenv("AZURE_OPENAI_ENDPOINT"))
+ .buildClient();
+ }
+
+ @Bean
+ public AzureOpenAiEmbeddingClient azureEmbeddingClient(OpenAIClient openAIClient) {
+ return new AzureOpenAiEmbeddingClient(openAIClient);
}
}