From dfb8bf6a44c63ca1cbcfb34bb284f3655f9e251c Mon Sep 17 00:00:00 2001 From: Mark Pollack Date: Thu, 18 Apr 2024 13:09:47 -0400 Subject: [PATCH] Add a new abstraction to simplify implementation of common ChatBot use cases * Add ChatBot and basic DefaultChatBot * Add streaming ChatBot support. * Add Evaluator interface and RelevancyEvaluator implementation * Add Content data type abstraction for Document and Message * Renaming and package refactoring * update .gitignore to allow node package name * Add List to node and move ai.transformer package to ai.prompt.transformer * Add Short/Long term memory support. * Add mixing transformers support Docs TBD --- .gitignore | 4 +- .mvn/extensions.xml | 8 + .../ai/huggingface/client/ClientIT.java | 4 +- models/spring-ai-openai/pom.xml | 35 ++- .../ChatMemoryLongTermSystemPromptIT.java | 131 +++++++++ .../ChatMemoryShortTermMessageListIT.java | 107 ++++++++ .../ChatMemoryShortTermSystemPromptIT.java | 108 ++++++++ .../LongShortTermChatMemoryWithRagIT.java | 252 ++++++++++++++++++ .../chat/chatbot/OpenAiDefaultChatBotIT.java | 160 +++++++++++ .../ai/chat/chatbot/ChatAgentListener.java | 36 +++ .../ai/chat/chatbot/ChatBot.java | 40 +++ .../ai/chat/chatbot/ChatBotResponse.java | 69 +++++ .../ai/chat/chatbot/DefaultChatBot.java | 137 ++++++++++ .../chat/chatbot/DefaultStreamingChatBot.java | 144 ++++++++++ .../ai/chat/chatbot/StreamingChatBot.java | 40 +++ .../chatbot/StreamingChatBotResponse.java | 72 +++++ .../ai/chat/history/ChatMemory.java | 39 +++ .../chat/history/ChatMemoryAgentListener.java | 60 +++++ .../ai/chat/history/ChatMemoryRetriever.java | 81 ++++++ .../ai/chat/history/InMemoryChatMemory.java | 50 ++++ .../LastMaxTokenSizeContentTransformer.java | 121 +++++++++ .../history/MessageChatMemoryAugmentor.java | 70 +++++ .../SystemPromptChatMemoryAugmentor.java | 110 ++++++++ .../VectorStoreChatMemoryAgentListener.java | 98 +++++++ .../VectorStoreChatMemoryRetriever.java | 88 ++++++ .../ai/chat/messages/AbstractMessage.java | 64 ++--- .../ai/chat/messages/AssistantMessage.java | 2 +- .../ai/chat/messages/FunctionMessage.java | 2 +- .../ai/chat/messages/Message.java | 11 +- .../ai/chat/messages/MessageAggregator.java | 74 +++++ .../ai/chat/messages/SystemMessage.java | 2 +- .../ai/chat/messages/UserMessage.java | 7 +- .../ai/chat/prompt/Prompt.java | 41 ++- .../prompt/transformer/PromptContext.java | 185 +++++++++++++ .../prompt/transformer/PromptTransformer.java | 39 +++ .../transformer/QuestionContextAugmentor.java | 91 +++++++ .../transformer/TransformerContentType.java | 34 +++ .../transformer/VectorStoreRetriever.java | 94 +++++++ .../springframework/ai/document/Document.java | 22 +- .../ai/evaluation/EvaluationRequest.java | 63 +++++ .../ai/evaluation/EvaluationResponse.java | 60 +++++ .../ai/evaluation/Evaluator.java | 8 + .../ai/evaluation/RelevancyEvaluator.java | 91 +++++++ .../org/springframework/ai/model/Content.java | 23 ++ .../tokenizer/JTokkitTokenCountEstimator.java | 85 ++++++ .../ai/tokenizer/TokenCountEstimator.java | 51 ++++ .../ai/chat/history/ChatMemoryTests.java | 147 ++++++++++ .../modules/ROOT/pages/api/chatclient.adoc | 17 +- .../antora/modules/ROOT/pages/api/prompt.adoc | 16 +- .../ai/evaluation/BaseMemoryTest.java | 108 ++++++++ .../ai/vectorstore/Neo4jVectorStore.java | 2 +- .../ai/vectorstore/PineconeVectorStoreIT.java | 2 +- vector-stores/spring-ai-qdrant/pom.xml | 50 ++-- .../qdrant/QdrantVectorStoreIT.java | 20 +- 54 files changed, 3376 insertions(+), 99 deletions(-) create mode 100644 .mvn/extensions.xml create mode 100644 models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryLongTermSystemPromptIT.java create mode 100644 models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryShortTermMessageListIT.java create mode 100644 models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryShortTermSystemPromptIT.java create mode 100644 models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/LongShortTermChatMemoryWithRagIT.java create mode 100644 models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/OpenAiDefaultChatBotIT.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatAgentListener.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatBot.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatBotResponse.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/DefaultChatBot.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/DefaultStreamingChatBot.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/StreamingChatBot.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/StreamingChatBotResponse.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemory.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemoryAgentListener.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemoryRetriever.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/history/InMemoryChatMemory.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/history/LastMaxTokenSizeContentTransformer.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/history/MessageChatMemoryAugmentor.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/history/SystemPromptChatMemoryAugmentor.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/history/VectorStoreChatMemoryAgentListener.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/history/VectorStoreChatMemoryRetriever.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/messages/MessageAggregator.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/PromptContext.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/PromptTransformer.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/QuestionContextAugmentor.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/TransformerContentType.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/VectorStoreRetriever.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/evaluation/EvaluationRequest.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/evaluation/EvaluationResponse.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/evaluation/Evaluator.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/evaluation/RelevancyEvaluator.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/model/Content.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/tokenizer/JTokkitTokenCountEstimator.java create mode 100644 spring-ai-core/src/main/java/org/springframework/ai/tokenizer/TokenCountEstimator.java create mode 100644 spring-ai-core/src/test/java/org/springframework/ai/chat/history/ChatMemoryTests.java create mode 100644 spring-ai-test/src/main/java/org/springframework/ai/evaluation/BaseMemoryTest.java 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> mapRef = new AtomicReference<>(); + + return fluxChatResponse.doOnSubscribe(subscription -> { + // logger.info("Aggregation Subscribe:" + subscription); + stringBufferRef.set(new StringBuilder()); + mapRef.set(new HashMap<>()); + }).doOnNext(chatResponse -> { + // logger.info("Aggregation Next:" + chatResponse); + if (chatResponse.getResult() != null) { + if (chatResponse.getResult().getOutput().getContent() != null) { + stringBufferRef.get().append(chatResponse.getResult().getOutput().getContent()); + } + if (chatResponse.getResult().getOutput().getMetadata() != null) { + mapRef.get().putAll(chatResponse.getResult().getOutput().getMetadata()); + } + } + }).doOnComplete(() -> { + // logger.debug("Aggregation Complete"); + onAggregationComplete + .accept(new ChatResponse(List.of(new Generation(stringBufferRef.get().toString(), mapRef.get())))); + stringBufferRef.set(new StringBuilder()); + mapRef.set(new HashMap<>()); + }).doOnError(e -> { + logger.error("Aggregation Error", e); + }); + } + +} \ No newline at end of file diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/SystemMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/SystemMessage.java index 0c0a20e04..8a9dc5eaa 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/SystemMessage.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/SystemMessage.java @@ -36,7 +36,7 @@ public class SystemMessage extends AbstractMessage { @Override public String toString() { - return "SystemMessage{" + "content='" + getContent() + '\'' + ", properties=" + properties + ", messageType=" + return "SystemMessage{" + "content='" + getContent() + '\'' + ", properties=" + metadata + ", messageType=" + messageType + '}'; } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/UserMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/UserMessage.java index 4c9229516..e792c985b 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/UserMessage.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/UserMessage.java @@ -17,6 +17,7 @@ package org.springframework.ai.chat.messages; import java.util.Arrays; import java.util.List; +import java.util.Map; import org.springframework.core.io.Resource; @@ -43,9 +44,13 @@ public class UserMessage extends AbstractMessage { this(textContent, Arrays.asList(media)); } + public UserMessage(String textContent, List mediaList, Map metadata) { + super(MessageType.USER, textContent, mediaList, metadata); + } + @Override public String toString() { - return "UserMessage{" + "content='" + getContent() + '\'' + ", properties=" + properties + ", messageType=" + return "UserMessage{" + "content='" + getContent() + '\'' + ", properties=" + metadata + ", messageType=" + messageType + '}'; } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/Prompt.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/Prompt.java index 86aab947f..1e4b97ed7 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/Prompt.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/Prompt.java @@ -15,15 +15,19 @@ */ package org.springframework.ai.chat.prompt; -import org.springframework.ai.model.ModelOptions; -import org.springframework.ai.model.ModelRequest; -import org.springframework.ai.chat.messages.Message; -import org.springframework.ai.chat.messages.UserMessage; - +import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Objects; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.FunctionMessage; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.SystemMessage; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.model.ModelOptions; +import org.springframework.ai.model.ModelRequest; + public class Prompt implements ModelRequest> { private final List messages; @@ -92,4 +96,31 @@ public class Prompt implements ModelRequest> { return Objects.hash(this.messages, this.modelOptions); } + public Prompt copy() { + return new Prompt(instructionsCopy(), this.modelOptions); + } + + private List instructionsCopy() { + List messagesCopy = new ArrayList<>(); + this.messages.forEach(message -> { + if (message instanceof UserMessage) { + messagesCopy.add(new UserMessage(message.getContent(), message.getMedia(), message.getMetadata())); + } + else if (message instanceof SystemMessage) { + messagesCopy.add(new SystemMessage(message.getContent())); + } + else if (message instanceof AssistantMessage) { + messagesCopy.add(new AssistantMessage(message.getContent(), message.getMetadata())); + } + else if (message instanceof FunctionMessage) { + messagesCopy.add(new FunctionMessage(message.getContent(), message.getMetadata())); + } + else { + throw new IllegalArgumentException("Unsupported message type: " + message.getClass().getName()); + } + }); + + return messagesCopy; + } + } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/PromptContext.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/PromptContext.java new file mode 100644 index 000000000..e069491d6 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/PromptContext.java @@ -0,0 +1,185 @@ +/* + * 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.prompt.transformer; + +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.model.Content; + +import java.util.*; + +/** + * The shared, at the moment, mutable, data structure that can be used to implement + * ChatBot functionality. + * + * @author Mark Pollack + * @author Christian Tzolov + * @since 1.0 M1 + */ +public class PromptContext { + + private Prompt prompt; // The most up-to-date prompt to use + + private List contents; // The most up-to-date data to use + + private List promptHistory; + + private String conversationId = "default"; + + private Map metadata = new HashMap<>(); + + public PromptContext(Prompt prompt) { + this(prompt, new ArrayList<>()); + } + + public PromptContext(Prompt prompt, String conversationId) { + this(prompt, new ArrayList<>()); + this.conversationId = conversationId; + } + + public PromptContext(Prompt prompt, List contents) { + this.prompt = prompt; + this.promptHistory = new ArrayList<>(); + this.promptHistory.add(prompt); + this.contents = contents; + } + + public Prompt getPrompt() { + return prompt; + } + + public void setPrompt(Prompt prompt) { + this.prompt = prompt; + } + + public void addData(Content datum) { + this.contents.add(datum); + } + + public List getContents() { + return contents; + } + + public void setContents(List contents) { + this.contents = contents; + } + + public void addPromptHistory(Prompt prompt) { + this.promptHistory.add(prompt); + } + + public List getPromptHistory() { + return promptHistory; + } + + public String getConversationId() { + return conversationId; + } + + public Map getMetadata() { + return metadata; + } + + public static Builder from(PromptContext promptContext) { + return PromptContext.builder() + .withContents( + new ArrayList<>(promptContext.getContents() != null ? promptContext.getContents() : List.of())) + .withPrompt(promptContext.getPrompt().copy()) // deep copy + .withMetadata(new HashMap<>(promptContext.getMetadata() != null ? promptContext.getMetadata() : Map.of())) + .withPromptHistory(new ArrayList<>( + promptContext.getPromptHistory() != null ? promptContext.getPromptHistory() : List.of())) + .withConversationId(promptContext.getConversationId()); + } + + public static Builder builder() { + return new Builder(); + } + + public static class Builder { + + private Prompt prompt; + + private List contents; + + private List promptHistory; + + private String conversationId; + + private Map metadata = new HashMap<>(); + + public Builder withPrompt(Prompt prompt) { + this.prompt = prompt; + return this; + } + + public Builder withContents(List contents) { + this.contents = new ArrayList<>(contents); + return this; + } + + public Builder withPromptHistory(List promptHistory) { + this.promptHistory = new ArrayList<>(promptHistory); + return this; + } + + public Builder addPromptHistory(Prompt prompt) { + this.promptHistory.add(prompt); + return this; + } + + public Builder withConversationId(String conversationId) { + this.conversationId = conversationId; + return this; + } + + public Builder withMetadata(Map metadata) { + this.metadata = new HashMap<>(metadata); + return this; + } + + public PromptContext build() { + PromptContext promptContext = new PromptContext(prompt, contents); + promptContext.promptHistory = promptHistory; + promptContext.conversationId = conversationId; + promptContext.metadata = metadata; + return promptContext; + } + + } + + @Override + public String toString() { + return "PromptContext{" + "prompt=" + prompt + ", contents=" + contents + ", promptHistory=" + promptHistory + + ", conversationId='" + conversationId + '\'' + ", metadata=" + metadata + '}'; + } + + @Override + public boolean equals(Object o) { + if (this == o) + return true; + if (!(o instanceof PromptContext that)) + return false; + return Objects.equals(prompt, that.prompt) && Objects.equals(contents, that.contents) + && Objects.equals(promptHistory, that.promptHistory) + && Objects.equals(conversationId, that.conversationId) && Objects.equals(metadata, that.metadata); + } + + @Override + public int hashCode() { + return Objects.hash(prompt, contents, promptHistory, conversationId, metadata); + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/PromptTransformer.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/PromptTransformer.java new file mode 100644 index 000000000..7d596a0e9 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/PromptTransformer.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.prompt.transformer; + +/** + * Responsible for transforming a Prompt. The PromptContext contains the necessary data to + * make the transformation + * + * Implementations may retrieve data and modify the Prompt object in the PromptContext as + * needed. + * + * @author Mark Pollack + * @since 1.0 M1 + */ +@FunctionalInterface +public interface PromptTransformer { + + /** + * Transforms the given PromptContext. + * @param context the PromptContext to transform + * @return the transformed PromptContext + */ + PromptContext transform(PromptContext context); + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/QuestionContextAugmentor.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/QuestionContextAugmentor.java new file mode 100644 index 000000000..7a99d9eb8 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/QuestionContextAugmentor.java @@ -0,0 +1,91 @@ +/* + * 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.prompt.transformer; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.MessageType; +import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.chat.prompt.PromptTemplate; +import org.springframework.ai.model.Content; + +/** + * Transforms the Prompt by taking to the current prompt in the Prompt Context and adding + * additional context to create a new prompt. The default user text contains the + * placeholder names "question" and "context". The "question" placeholder is filled using + * the value of the current UserMessage and the "context" placeholder is filled with + * Documents contained in the PromptContext's Nodes. + */ +public class QuestionContextAugmentor implements PromptTransformer { + + private static final String DEFAULT_USER_PROMPT_TEXT = """ + "Context information is below.\\n" + "---------------------\\n" + "{context}\\n" + "---------------------\\n" + "Given the context and provided history information and not prior knowledge, " + "reply to the user comment. If the answer is not in the context, inform " + "the user that you can't answer the question.\\n" + "User comment: {question}\\n" + "Answer: " + """; + + @Override + public PromptContext transform(PromptContext promptContext) { + String context = doCreateContext(promptContext.getContents()); + Map contextMap = doCreateContextMap(promptContext.getPrompt(), context); + Prompt prompt = doCreatePrompt(promptContext.getPrompt(), contextMap); + promptContext.setPrompt(prompt); + promptContext.addPromptHistory(prompt); // BUG? shouldn't this be original + // promptContext.getPrompt()? + // For now return the modified instance instead of a copy + return promptContext; + } + + protected String doCreateContext(List data) { + return data.stream() + .filter(content -> content.getMetadata().containsKey(TransformerContentType.EXTERNAL_KNOWLEDGE)) + .map(Content::getContent) + .collect(Collectors.joining(System.lineSeparator())); + } + + private Map doCreateContextMap(Prompt prompt, String context) { + String originalUserMessage = prompt.getInstructions() + .stream() + .filter(m -> m.getMessageType() == MessageType.USER) + .map(m -> m.getContent()) + .collect(Collectors.joining(System.lineSeparator())); + + return Map.of("context", context, "question", originalUserMessage); + } + + protected Prompt doCreatePrompt(Prompt originalPrompt, Map contextMap) { + PromptTemplate promptTemplate = new PromptTemplate(DEFAULT_USER_PROMPT_TEXT); + Message userMessageToAppend = promptTemplate.createMessage(contextMap); + List messageList = originalPrompt.getInstructions() + .stream() + .filter(m -> m.getMessageType() != MessageType.USER) + .collect(Collectors.toList()); + messageList.add(userMessageToAppend); + return new Prompt(messageList, (ChatOptions) originalPrompt.getOptions()); + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/TransformerContentType.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/TransformerContentType.java new file mode 100644 index 000000000..270e85943 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/TransformerContentType.java @@ -0,0 +1,34 @@ +/* + * 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.prompt.transformer; + +/** + * @author Christian Tzolov + */ +public class TransformerContentType { + + public static final String MEMORY = "MEMORY_TYPE"; + + public static final String LONG_TERM_MEMORY = "LONG_TERM_MEMORY_TYPE"; + + public static final String SHORT_TERM_MEMORY = "SHORT_TERM_MEMORY_TYPE"; + + public static final String CONVERSATION_ID = "conversationId"; + + public static final String EXTERNAL_KNOWLEDGE = "externalKnowledge"; + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/VectorStoreRetriever.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/VectorStoreRetriever.java new file mode 100644 index 000000000..92608b3d9 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/VectorStoreRetriever.java @@ -0,0 +1,94 @@ +/* + * 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.prompt.transformer; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.MessageType; +import org.springframework.ai.document.Document; +import org.springframework.ai.vectorstore.SearchRequest; +import org.springframework.ai.vectorstore.VectorStore; + +import java.util.List; +import java.util.Objects; +import java.util.stream.Collectors; + +/** + * Transforms the PromptContext by retrieving documents from a VectorStore + */ +public class VectorStoreRetriever implements PromptTransformer { + + private final Logger logger = LoggerFactory.getLogger(getClass()); + + private final VectorStore vectorStore; + + private final SearchRequest searchRequest; + + public VectorStoreRetriever(VectorStore vectorStore, SearchRequest searchRequest) { + this.vectorStore = vectorStore; + this.searchRequest = searchRequest; + } + + public VectorStore getVectorStore() { + return vectorStore; + } + + public SearchRequest getSearchRequest() { + return searchRequest; + } + + @Override + public PromptContext transform(PromptContext promptContext) { + List instructions = promptContext.getPrompt().getInstructions(); + String userMessage = instructions.stream() + .filter(m -> m.getMessageType() == MessageType.USER) + .map(m -> m.getContent()) + .collect(Collectors.joining(System.lineSeparator())); + + List documents = vectorStore.similaritySearch(searchRequest.withQuery(userMessage) + .withFilterExpression(TransformerContentType.EXTERNAL_KNOWLEDGE + "=='true'")); + + logger.info("Retrieved {} documents for user message {}", documents.size(), userMessage); + for (Document document : documents) { + var content = new Document(document.getContent(), document.getMetadata()); + // content.getMetadata().put(TransformerContentType.DOMAIN_DATA, true); + promptContext.addData(content); + } + return promptContext; + } + + @Override + public String toString() { + return "VectorStoreRetriever{" + "vectorStore=" + vectorStore + ", searchRequest=" + searchRequest + '}'; + } + + @Override + public boolean equals(Object o) { + if (this == o) + return true; + if (!(o instanceof VectorStoreRetriever that)) + return false; + return Objects.equals(vectorStore, that.vectorStore) && Objects.equals(searchRequest, that.searchRequest); + } + + @Override + public int hashCode() { + return Objects.hash(vectorStore, searchRequest); + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java b/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java index 7c54b4fec..30c4b479a 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/document/Document.java @@ -25,8 +25,10 @@ import com.fasterxml.jackson.annotation.JsonIgnore; import com.fasterxml.jackson.annotation.JsonIgnoreProperties; import com.fasterxml.jackson.annotation.JsonProperty; +import org.springframework.ai.chat.messages.Media; import org.springframework.ai.document.id.IdGenerator; import org.springframework.ai.document.id.RandomIdGenerator; +import org.springframework.ai.model.Content; import org.springframework.util.Assert; /** @@ -34,7 +36,7 @@ import org.springframework.util.Assert; * the document's unique ID and an optional embedding. */ @JsonIgnoreProperties({ "contentFormatter" }) -public class Document { +public class Document implements Content { public final static ContentFormatter DEFAULT_CONTENT_FORMATTER = DefaultContentFormatter.defaultConfig(); @@ -54,6 +56,8 @@ public class Document { */ private String content; + private List media; + /** * Embedding of the document. Note: ephemeral field. */ @@ -75,17 +79,26 @@ public class Document { this(content, metadata, new RandomIdGenerator()); } + public Document(String content, List media, Map metadata) { + this(new RandomIdGenerator().generateId(content, metadata), content, media, metadata); + } + public Document(String content, Map metadata, IdGenerator idGenerator) { this(idGenerator.generateId(content, metadata), content, metadata); } public Document(String id, String content, Map metadata) { + this(id, content, List.of(), metadata); + } + + public Document(String id, String content, List media, Map metadata) { Assert.hasText(id, "id must not be null"); Assert.hasText(content, "content must not be null"); Assert.notNull(metadata, "metadata must not be null"); this.id = id; this.content = content; + this.media = media; this.metadata = metadata; } @@ -93,10 +106,16 @@ public class Document { return id; } + @Override public String getContent() { return this.content; } + @Override + public List getMedia() { + return this.media; + } + @JsonIgnore public String getFormattedContent() { return this.getFormattedContent(MetadataMode.ALL); @@ -129,6 +148,7 @@ public class Document { this.contentFormatter = contentFormatter; } + @Override public Map getMetadata() { return this.metadata; } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/evaluation/EvaluationRequest.java b/spring-ai-core/src/main/java/org/springframework/ai/evaluation/EvaluationRequest.java new file mode 100644 index 000000000..657370e86 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/evaluation/EvaluationRequest.java @@ -0,0 +1,63 @@ +package org.springframework.ai.evaluation; + +import org.springframework.ai.chat.ChatResponse; +import org.springframework.ai.chat.chatbot.ChatBotResponse; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.model.Content; + +import java.util.List; +import java.util.Objects; + +public class EvaluationRequest { + + private final Prompt prompt; + + private final List dataList; + + private final ChatResponse chatResponse; + + public EvaluationRequest(ChatBotResponse chatBotResponse) { + this(chatBotResponse.getPromptContext().getPromptHistory().get(0), + chatBotResponse.getPromptContext().getContents(), chatBotResponse.getChatResponse()); + } + + public EvaluationRequest(Prompt prompt, List dataList, ChatResponse chatResponse) { + this.prompt = prompt; + this.dataList = dataList; + this.chatResponse = chatResponse; + } + + public Prompt getPrompt() { + return prompt; + } + + public List getDataList() { + return dataList; + } + + public ChatResponse getChatResponse() { + return chatResponse; + } + + @Override + public String toString() { + return "EvaluationRequest{" + "prompt=" + prompt + ", dataList=" + dataList + ", chatResponse=" + chatResponse + + '}'; + } + + @Override + public boolean equals(Object o) { + if (this == o) + return true; + if (!(o instanceof EvaluationRequest that)) + return false; + return Objects.equals(prompt, that.prompt) && Objects.equals(dataList, that.dataList) + && Objects.equals(chatResponse, that.chatResponse); + } + + @Override + public int hashCode() { + return Objects.hash(prompt, dataList, chatResponse); + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/evaluation/EvaluationResponse.java b/spring-ai-core/src/main/java/org/springframework/ai/evaluation/EvaluationResponse.java new file mode 100644 index 000000000..a22d738c6 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/evaluation/EvaluationResponse.java @@ -0,0 +1,60 @@ +package org.springframework.ai.evaluation; + +import java.util.Map; +import java.util.Objects; + +public class EvaluationResponse { + + private boolean pass; + + private float score; + + private String feedback; + + Map metadata; + + public EvaluationResponse(boolean pass, float score, String feedback, Map metadata) { + this.pass = pass; + this.score = score; + this.feedback = feedback; + this.metadata = metadata; + } + + public boolean isPass() { + return pass; + } + + public float getScore() { + return score; + } + + public String getFeedback() { + return feedback; + } + + public Map getMetadata() { + return metadata; + } + + @Override + public String toString() { + return "EvaluationResponse{" + "pass=" + pass + ", score=" + score + ", feedback='" + feedback + '\'' + + ", metadata=" + metadata + '}'; + } + + @Override + public boolean equals(Object o) { + if (this == o) + return true; + if (!(o instanceof EvaluationResponse that)) + return false; + return pass == that.pass && Float.compare(score, that.score) == 0 && Objects.equals(feedback, that.feedback) + && Objects.equals(metadata, that.metadata); + } + + @Override + public int hashCode() { + return Objects.hash(pass, score, feedback, metadata); + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/evaluation/Evaluator.java b/spring-ai-core/src/main/java/org/springframework/ai/evaluation/Evaluator.java new file mode 100644 index 000000000..7cfdbbf67 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/evaluation/Evaluator.java @@ -0,0 +1,8 @@ +package org.springframework.ai.evaluation; + +@FunctionalInterface +public interface Evaluator { + + EvaluationResponse evaluate(EvaluationRequest evaluationRequest); + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/evaluation/RelevancyEvaluator.java b/spring-ai-core/src/main/java/org/springframework/ai/evaluation/RelevancyEvaluator.java new file mode 100644 index 000000000..8afc9dcff --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/evaluation/RelevancyEvaluator.java @@ -0,0 +1,91 @@ +package org.springframework.ai.evaluation; + +import org.springframework.ai.chat.ChatClient; +import org.springframework.ai.chat.ChatResponse; +import org.springframework.ai.chat.messages.Message; +import org.springframework.ai.chat.messages.MessageType; +import org.springframework.ai.chat.prompt.ChatOptions; +import org.springframework.ai.chat.prompt.ChatOptionsBuilder; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.chat.prompt.PromptTemplate; +import org.springframework.ai.model.Content; + +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +public class RelevancyEvaluator implements Evaluator { + + private static final String DEFAULT_EVALUATION_PROMPT_TEXT = """ + Your task is to evaluate if the response for the query + is in line with the context information provided.\\n + You have two options to answer. Either YES/ NO.\\n + Answer - YES, if the response for the query + is in line with context information otherwise NO.\\n + Query: \\n {query}\\n + Response: \\n {response}\\n + Context: \\n {context}\\n + Answer: " + """; + + private final ChatOptions chatOptions; + + private ChatClient chatClient; + + public RelevancyEvaluator(ChatClient chatClient) { + this(chatClient, ChatOptionsBuilder.builder().build()); + } + + public RelevancyEvaluator(ChatClient chatClient, ChatOptions chatOptions) { + this.chatClient = chatClient; + this.chatOptions = chatOptions; + } + + @Override + public EvaluationResponse evaluate(EvaluationRequest evaluationRequest) { + var query = doGetUserQuestion(evaluationRequest); + var response = doGetResponse(evaluationRequest); + var context = doGetSupportingData(evaluationRequest); + + var promptTemplate = new PromptTemplate(DEFAULT_EVALUATION_PROMPT_TEXT); + Message message = promptTemplate + .createMessage(Map.of("query", query, "response", response, "context", context)); + + ChatResponse chatResponse = this.chatClient.call(new Prompt(message, this.chatOptions)); + + var evaluationResponse = chatResponse.getResult().getOutput().getContent(); + boolean passing = false; + float score = 0; + if (evaluationResponse.toLowerCase().contains("yes")) { + passing = true; + score = 1; + } + + return new EvaluationResponse(passing, score, "", Collections.emptyMap()); + } + + protected String doGetResponse(EvaluationRequest evaluationRequest) { + return evaluationRequest.getChatResponse().getResult().getOutput().getContent(); + } + + protected String doGetSupportingData(EvaluationRequest evaluationRequest) { + List data = evaluationRequest.getDataList(); + String supportingData = data.stream() + .filter(node -> node != null && node.getContent() instanceof String) + .map(node -> (Content) node) + .map(Content::getContent) + .collect(Collectors.joining(System.lineSeparator())); + return supportingData; + } + + protected String doGetUserQuestion(EvaluationRequest evaluationRequest) { + List instructions = evaluationRequest.getPrompt().getInstructions(); + String userMessage = instructions.stream() + .filter(m -> m.getMessageType() == MessageType.USER) + .map(m -> m.getContent()) + .collect(Collectors.joining(System.lineSeparator())); + return userMessage; + } + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/model/Content.java b/spring-ai-core/src/main/java/org/springframework/ai/model/Content.java new file mode 100644 index 000000000..7100f6797 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/model/Content.java @@ -0,0 +1,23 @@ +package org.springframework.ai.model; + +import org.springframework.ai.chat.messages.Media; + +import java.util.List; +import java.util.Map; + +/** + * A simple data structure that contains content and metadata. + * + * @param the type of content in the node + * @author Mark Pollack + * @since 1.0 M1 + */ +public interface Content { + + String getContent(); + + List getMedia(); + + Map getMetadata(); + +} diff --git a/spring-ai-core/src/main/java/org/springframework/ai/tokenizer/JTokkitTokenCountEstimator.java b/spring-ai-core/src/main/java/org/springframework/ai/tokenizer/JTokkitTokenCountEstimator.java new file mode 100644 index 000000000..39dc9781f --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/tokenizer/JTokkitTokenCountEstimator.java @@ -0,0 +1,85 @@ +/* + * 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.tokenizer; + +import com.knuddels.jtokkit.Encodings; +import com.knuddels.jtokkit.api.Encoding; +import com.knuddels.jtokkit.api.EncodingType; + +import org.springframework.ai.chat.messages.Media; +import org.springframework.ai.model.Content; +import org.springframework.util.CollectionUtils; + +/** + * @author Christian Tzolov + */ +public class JTokkitTokenCountEstimator implements TokenCountEstimator { + + private final Encoding estimator; + + public JTokkitTokenCountEstimator() { + this.estimator = Encodings.newLazyEncodingRegistry().getEncoding(EncodingType.CL100K_BASE); + } + + public JTokkitTokenCountEstimator(Encoding tokenEncoding) { + this.estimator = tokenEncoding; + } + + @Override + public int estimate(String text) { + if (text == null) { + return 0; + } + return this.estimator.countTokens(text); + } + + @Override + public int estimate(Content content) { + int tokenCount = 0; + + if (content.getContent() != null) { + tokenCount += this.estimate(content.getContent()); + } + + if (!CollectionUtils.isEmpty(content.getMedia())) { + + for (Media media : content.getMedia()) { + + tokenCount += this.estimate(media.getMimeType().toString()); + + if (media.getData() instanceof String textData) { + tokenCount += this.estimate(textData); + } + else if (media.getData() instanceof byte[] binaryData) { + tokenCount += binaryData.length; // This is likely incorrect. + } + } + } + + return tokenCount; + } + + @Override + public int estimate(Iterable contents) { + int totalSize = 0; + for (Content content : contents) { + totalSize += this.estimate(content); + } + return totalSize; + } + +} \ No newline at end of file diff --git a/spring-ai-core/src/main/java/org/springframework/ai/tokenizer/TokenCountEstimator.java b/spring-ai-core/src/main/java/org/springframework/ai/tokenizer/TokenCountEstimator.java new file mode 100644 index 000000000..0364005f2 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/tokenizer/TokenCountEstimator.java @@ -0,0 +1,51 @@ +/* + * 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.tokenizer; + +import org.springframework.ai.model.Content; + +/** + * Estimates the number of tokens in a given text or message. + * + * @author Christian Tzolov + * @since 1.0.0 + */ +public interface TokenCountEstimator { + + /** + * Estimates the number of tokens in the given text. + * @param text the text to estimate the number of tokens for. + * @return the estimated number of tokens. + */ + int estimate(String text); + + /** + * Estimates the number of tokens in the given message. + * @param content the content (Message or Document) to estimate the number of tokens + * for. + * @return the estimated number of tokens. + */ + int estimate(Content content); + + /** + * Estimates the number of tokens in the given messages. + * @param messages the messages to estimate the number of tokens for. + * @return the estimated number of tokens. + */ + int estimate(Iterable messages); + +} \ No newline at end of file diff --git a/spring-ai-core/src/test/java/org/springframework/ai/chat/history/ChatMemoryTests.java b/spring-ai-core/src/test/java/org/springframework/ai/chat/history/ChatMemoryTests.java new file mode 100644 index 000000000..38099408c --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/chat/history/ChatMemoryTests.java @@ -0,0 +1,147 @@ +/* + * 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.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Captor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import org.springframework.ai.chat.ChatClient; +import org.springframework.ai.chat.ChatResponse; +import org.springframework.ai.chat.Generation; +import org.springframework.ai.chat.StreamingChatClient; +import org.springframework.ai.chat.chatbot.ChatBotResponse; +import org.springframework.ai.chat.chatbot.DefaultChatBot; +import org.springframework.ai.chat.messages.Message; +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.model.Content; +import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.when; + +/** + * @author Christian Tzolov + */ +@ExtendWith(MockitoExtension.class) +public class ChatMemoryTests { + + @Mock + ChatClient chatClient; + + @Mock + StreamingChatClient streamingChatClient; + + @Captor + ArgumentCaptor promptCaptor; + + @Test + public void chatMemoryMessageListAugmentor() { + + ChatMemory chatHistory = new InMemoryChatMemory(); + + DefaultChatBot chatAgent = DefaultChatBot.builder(chatClient) + .withRetrievers(List.of(new ChatMemoryRetriever(chatHistory))) + .withContentPostProcessors( + List.of(new LastMaxTokenSizeContentTransformer(new JTokkitTokenCountEstimator(), 10))) + .withAugmentors(List.of(new MessageChatMemoryAugmentor())) + .withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory))) + .build(); + + chatClientUserMessages(chatAgent, chatHistory); + } + + @Test + public void chatMemorySystemPromptAugmentor() { + + ChatMemory chatHistory = new InMemoryChatMemory(); + + DefaultChatBot chatAgent = DefaultChatBot.builder(chatClient) + .withRetrievers(List.of(new ChatMemoryRetriever(chatHistory))) + .withContentPostProcessors( + List.of(new LastMaxTokenSizeContentTransformer(new JTokkitTokenCountEstimator(), 10))) + .withAugmentors(List.of(new SystemPromptChatMemoryAugmentor())) + .withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory))) + .build(); + + chatClientUserMessages(chatAgent, chatHistory); + } + + public void chatClientUserMessages(DefaultChatBot chatAgent, ChatMemory chatHistory) { + + when(chatClient.call(promptCaptor.capture())) + .thenReturn(new ChatResponse(List.of(new Generation("assistant:1")))) + .thenReturn(new ChatResponse(List.of(new Generation("assistant:2")))) + .thenReturn(new ChatResponse(List.of(new Generation("assistant:3")))); + + var promptContext = PromptContext.builder() + .withConversationId("test-session-id") + .withPrompt(new Prompt( + List.of(new UserMessage("user:1"), new UserMessage("user:2"), new UserMessage("user:3"), + new UserMessage("user:4"), new UserMessage("user:5")))) + .build(); + + ChatBotResponse response1 = chatAgent.call(promptContext); + + assertThat(response1.getChatResponse().getResult().getOutput().getContent()).isEqualTo("assistant:1"); + + List contents = response1.getPromptContext().getContents(); + assertThat(contents).hasSize(0); + + List history = chatHistory.get("test-session-id", 1000); + assertThat(history).hasSize(6); + + ChatBotResponse response2 = chatAgent.call(PromptContext.builder() + .withConversationId("test-session-id") + .withPrompt(new Prompt( + List.of(new UserMessage("user:6"), new UserMessage("user:7"), new UserMessage("user:8")))) + .build()); + + assertThat(response2.getChatResponse().getResult().getOutput().getContent()).isEqualTo("assistant:2"); + + history = chatHistory.get("test-session-id", 1000); + assertThat(history).hasSize(10); + + contents = response2.getPromptContext().getContents(); + assertThat(contents).hasSize(3); + assertThat(contents.get(0).getContent()).isEqualTo("user:4"); + assertThat(contents.get(1).getContent()).isEqualTo("user:5"); + assertThat(contents.get(2).getContent()).isEqualTo("assistant:1"); + + ChatBotResponse response3 = chatAgent.call(PromptContext.builder() + .withConversationId("test-session-id") + .withPrompt(new Prompt(List.of(new UserMessage("user:9")))).build()); + assertThat(response3.getChatResponse().getResult().getOutput().getContent()).isEqualTo("assistant:3"); + + history = chatHistory.get("test-session-id", 1000); + assertThat(history).hasSize(12); + + contents = response3.getPromptContext().getContents(); + assertThat(contents).hasSize(3); + assertThat(contents.get(0).getContent()).isEqualTo("user:7"); + assertThat(contents.get(1).getContent()).isEqualTo("user:8"); + assertThat(contents.get(2).getContent()).isEqualTo("assistant:2"); + } + +} diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc index 902d0d59b..42f7857a5 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc @@ -78,16 +78,29 @@ The `Message` interface encapsulates a textual message, a collection of attribut [source,java] ---- -public interface Message { +public interface Message extends Node { String getContent(); - Map getProperties(); + List getMedia(); MessageType getMessageType(); } ---- + +and the Node interface is + +```java + +public interface Node { + + T getContent(); + + Map getMetadata(); +} +``` + The `Message` interface has various implementations that correspond to the categories of messages that an AI model can process. Some models, like OpenAI's chat completion endpoint, distinguish between message categories based on conversational roles, effectively mapped by the `MessageType`. diff --git a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/prompt.adoc b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/prompt.adoc index 3f1052981..96fab52f4 100644 --- a/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/prompt.adoc +++ b/spring-ai-docs/src/main/antora/modules/ROOT/pages/api/prompt.adoc @@ -49,19 +49,29 @@ The `Message` interface encapsulates a textual message, a collection of attribut The interface is defined as follows: ```java -public interface Message { +public interface Message extends Node { String getContent(); List getMedia(); - Map getProperties(); - MessageType getMessageType(); } ``` +and the Node interface is + +```java + +public interface Node { + + T getContent(); + + Map getMetadata(); +} +``` + Various implementations of the `Message` interface correspond to different categories of messages that an AI model can process. Some models, like those from OpenAI, distinguish between message categories based on conversational roles. These roles are effectively mapped by the `MessageType`, as discussed below. diff --git a/spring-ai-test/src/main/java/org/springframework/ai/evaluation/BaseMemoryTest.java b/spring-ai-test/src/main/java/org/springframework/ai/evaluation/BaseMemoryTest.java new file mode 100644 index 000000000..133b2b002 --- /dev/null +++ b/spring-ai-test/src/main/java/org/springframework/ai/evaluation/BaseMemoryTest.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.evaluation; + +import java.util.stream.Collectors; + +import org.junit.jupiter.api.Test; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import org.springframework.ai.chat.chatbot.ChatBot; +import org.springframework.ai.chat.chatbot.StreamingChatBot; +import org.springframework.ai.chat.messages.UserMessage; +import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.chat.prompt.transformer.PromptContext; + +import static org.assertj.core.api.Assertions.assertThat; + +/** + * @author Christian Tzolov + */ +public class BaseMemoryTest { + + protected final Logger logger = LoggerFactory.getLogger(getClass()); + + protected RelevancyEvaluator relevancyEvaluator; + + protected ChatBot chatBot; + + protected StreamingChatBot streamingChatBot; + + public BaseMemoryTest(RelevancyEvaluator relevancyEvaluator, ChatBot chatBot, + StreamingChatBot streamingChatClient) { + this.relevancyEvaluator = relevancyEvaluator; + this.chatBot = chatBot; + this.streamingChatBot = streamingChatClient; + } + + @Test + void memoryChatAgent() { + + var prompt = new Prompt(new UserMessage("my name John Vincent Atanasoff")); + PromptContext promptContext = new PromptContext(prompt); + + var chatBotResponse1 = this.chatBot.call(promptContext); + + logger.info("Response1: " + chatBotResponse1.getChatResponse().getResult().getOutput().getContent()); + assertThat(chatBotResponse1.getChatResponse().getResult().getOutput().getContent()).contains("John"); + + var chatBotResponse2 = this.chatBot.call(new PromptContext(new Prompt(new String("What is my name?")))); + logger.info("Response2: " + chatBotResponse2.getChatResponse().getResult().getOutput().getContent()); + assertThat(chatBotResponse2.getChatResponse().getResult().getOutput().getContent()) + .contains("John Vincent Atanasoff"); + + EvaluationResponse evaluationResponse = this.relevancyEvaluator + .evaluate(new EvaluationRequest(chatBotResponse2)); + logger.info("" + evaluationResponse); + } + + @Test + void memoryStreamingChatBot() { + + var prompt = new Prompt(new UserMessage("my name John Vincent Atanasoff")); + PromptContext promptContext = new PromptContext(prompt); + + var fluxChatBotResponse1 = this.streamingChatBot.stream(promptContext); + + String chatBotResponse1 = fluxChatBotResponse1.getChatResponse() + .collectList() + .block() + .stream() + .filter(response -> response.getResult() != null) + .map(response -> response.getResult().getOutput().getContent()) + .collect(Collectors.joining()); + + logger.info("Response1: " + chatBotResponse1); + assertThat(chatBotResponse1).contains("John"); + + var fluxChatBotResponse2 = this.streamingChatBot + .stream(new PromptContext(new Prompt(new String("What is my name?")))); + + String chatBotResponse2 = fluxChatBotResponse2.getChatResponse() + .collectList() + .block() + .stream() + .filter(response -> response.getResult() != null) + .map(response -> response.getResult().getOutput().getContent()) + .collect(Collectors.joining()); + + logger.info("Response2: " + chatBotResponse2); + assertThat(chatBotResponse2).contains("John Vincent Atanasoff"); + } + +} diff --git a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java index 4337fa954..1c50ccb90 100644 --- a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java +++ b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java @@ -176,7 +176,7 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean { */ public Builder withLabel(String newLabel) { - Assert.hasText(newLabel, "Node label may not be null or blank"); + Assert.hasText(newLabel, "Content label may not be null or blank"); this.label = newLabel; return this; diff --git a/vector-stores/spring-ai-pinecone/src/test/java/org/springframework/ai/vectorstore/PineconeVectorStoreIT.java b/vector-stores/spring-ai-pinecone/src/test/java/org/springframework/ai/vectorstore/PineconeVectorStoreIT.java index d2ee47be2..9e8ef5f45 100644 --- a/vector-stores/spring-ai-pinecone/src/test/java/org/springframework/ai/vectorstore/PineconeVectorStoreIT.java +++ b/vector-stores/spring-ai-pinecone/src/test/java/org/springframework/ai/vectorstore/PineconeVectorStoreIT.java @@ -292,7 +292,7 @@ public class PineconeVectorStoreIT { } @Bean - public EmbeddingClient embeddingClient() { + public TransformersEmbeddingClient embeddingClient() { return new TransformersEmbeddingClient(); } diff --git a/vector-stores/spring-ai-qdrant/pom.xml b/vector-stores/spring-ai-qdrant/pom.xml index c6040decc..95c578d28 100644 --- a/vector-stores/spring-ai-qdrant/pom.xml +++ b/vector-stores/spring-ai-qdrant/pom.xml @@ -32,16 +32,10 @@ ${project.parent.version} - - org.springframework - spring-web - - - + + org.springframework + spring-web + io.qdrant @@ -55,19 +49,20 @@ ${protobuf-java.version} - - - org.springframework.ai - spring-ai-openai - ${project.parent.version} - test - + + + org.springframework.ai + spring-ai-azure-openai + ${project.parent.version} + test + true + - - org.springframework.boot - spring-boot-starter-test - test - + + org.springframework.boot + spring-boot-starter-test + test + org.testcontainers @@ -75,11 +70,10 @@ test - - org.testcontainers - junit-jupiter - test - + + org.testcontainers + junit-jupiter + test + - 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); } }