diff --git a/.gitignore b/.gitignore index 74151f12f..c2fb8fc27 100644 --- a/.gitignore +++ b/.gitignore @@ -29,9 +29,13 @@ out vscode settings.json +node +node_modules package-lock.json 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-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/agent/MessageChatHistoryChatAgentIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/MessageChatHistoryChatAgentIT.java new file mode 100644 index 000000000..9d116b43d --- /dev/null +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/MessageChatHistoryChatAgentIT.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.agent; + +import java.util.List; + +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; + +import org.springframework.ai.chat.agent.ChatAgent; +import org.springframework.ai.chat.agent.DefaultChatAgent; +import org.springframework.ai.chat.agent.DefaultStreamingChatAgent; +import org.springframework.ai.chat.agent.StreamingChatAgent; +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 = MessageChatHistoryChatAgentIT.Config.class) +@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") +public class MessageChatHistoryChatAgentIT extends BaseMemoryTest { + + @Autowired + public MessageChatHistoryChatAgentIT(RelevancyEvaluator relevancyEvaluator, ChatAgent chatAgent, + StreamingChatAgent streamingChatAgent) { + super(relevancyEvaluator, chatAgent, streamingChatAgent); + } + + @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 ChatAgent memoryChatAgent(OpenAiChatClient chatClient, ChatMemory chatHistory, + TokenCountEstimator tokenCountEstimator) { + + return DefaultChatAgent.builder(chatClient) + .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 StreamingChatAgent memoryStreamingChatAgent(OpenAiChatClient streamingChatClient, ChatMemory chatHistory, + TokenCountEstimator tokenCountEstimator) { + + return DefaultStreamingChatAgent.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/agent/OpenAiDefaultChatAgentIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/OpenAiDefaultChatAgentIT.java index e92d6530e..ce10a8167 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/OpenAiDefaultChatAgentIT.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/OpenAiDefaultChatAgentIT.java @@ -1,3 +1,19 @@ +/* + * 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.agent; import org.junit.jupiter.api.Test; diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/OpenAiMemoryChatAgentIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/OpenAiMemoryChatAgentIT.java new file mode 100644 index 000000000..05cb39288 --- /dev/null +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/OpenAiMemoryChatAgentIT.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.agent; + +import java.util.List; + +import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; + +import org.springframework.ai.chat.agent.ChatAgent; +import org.springframework.ai.chat.agent.DefaultChatAgent; +import org.springframework.ai.chat.agent.DefaultStreamingChatAgent; +import org.springframework.ai.chat.agent.StreamingChatAgent; +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 = OpenAiMemoryChatAgentIT.Config.class) +@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") +public class OpenAiMemoryChatAgentIT extends BaseMemoryTest { + + @Autowired + public OpenAiMemoryChatAgentIT(RelevancyEvaluator relevancyEvaluator, ChatAgent chatAgent, + StreamingChatAgent streamingChatAgent) { + super(relevancyEvaluator, chatAgent, streamingChatAgent); + } + + @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 ChatAgent memoryChatAgent(OpenAiChatClient chatClient, ChatMemory chatHistory, + TokenCountEstimator tokenCountEstimator) { + + return DefaultChatAgent.builder(chatClient) + .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 StreamingChatAgent memoryStreamingChatAgent(OpenAiChatClient streamingChatClient, ChatMemory chatHistory, + TokenCountEstimator tokenCountEstimator) { + + return DefaultStreamingChatAgent.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/agent/TextChatHistoryChatAgent3IT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/TextChatHistoryChatAgent3IT.java new file mode 100644 index 000000000..24b5662a9 --- /dev/null +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/TextChatHistoryChatAgent3IT.java @@ -0,0 +1,223 @@ +/* + * 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.agent; + +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.testcontainers.junit.jupiter.Container; +import org.testcontainers.junit.jupiter.Testcontainers; +import org.testcontainers.qdrant.QdrantContainer; + +import org.springframework.ai.chat.agent.ChatAgent; +import org.springframework.ai.chat.agent.DefaultChatAgent; +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.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; + +@Testcontainers +@SpringBootTest(classes = TextChatHistoryChatAgent3IT.Config.class) +@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") +public class TextChatHistoryChatAgent3IT { + + 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 + ChatAgent chatAgent; + + @Autowired + RelevancyEvaluator relevancyEvaluator; + + @Autowired + VectorStore vectorStore; + + @Value("classpath:/data/acme/bikes.json") + private Resource bikesResource; + + void loadData() { + JsonReader jsonReader = new JsonReader(bikesResource, "name", "price", "shortDescription", "description"); + var textSplitter = new TokenTextSplitter(); + vectorStore.accept(textSplitter.apply(jsonReader.get())); + } + + // @Autowired + // StreamingChatAgent streamingChatAgent; + + @Test + void memoryChatAgent() { + + loadData(); + + var prompt = new Prompt(new UserMessage("My name is Christian and I like mountain bikes.")); + PromptContext promptContext = new PromptContext(prompt); + + var agentResponse1 = this.chatAgent.call(promptContext); + + logger.info("Response1: " + agentResponse1.getChatResponse().getResult().getOutput().getContent()); + + var agentResponse2 = this.chatAgent.call( + new PromptContext(new Prompt(new String("What is my name and what bike model would suggest for me?")))); + logger.info("Response2: " + agentResponse2.getChatResponse().getResult().getOutput().getContent()); + + logger.info(agentResponse2.getPromptContext().getContents().toString()); + assertThat(agentResponse2.getChatResponse().getResult().getOutput().getContent()).contains("Christian", + "mountain bikes"); + + EvaluationResponse evaluationResponse = this.relevancyEvaluator.evaluate(new EvaluationRequest(agentResponse2)); + logger.info("" + evaluationResponse); + } + + @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 ChatAgent memoryChatAgent(OpenAiChatClient chatClient, VectorStore vectorStore, + TokenCountEstimator tokenCountEstimator, ChatMemory chatHistory) { + + return DefaultChatAgent.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, "")))) + .withDocumentPostProcessors(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.QA)))) + .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 StreamingChatAgent memoryStreamingChatAgent(OpenAiChatClient + // streamingChatClient, + // VectorStore vectorStore, TokenCountEstimator tokenCountEstimator, ChatHistory + // chatHistory) { + + // return DefaultStreamingChatAgent.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) { + return new RelevancyEvaluator(chatClient); + } + + } + +} diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/TextChatHistoryChatAgentIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/TextChatHistoryChatAgentIT.java new file mode 100644 index 000000000..2d6c24b49 --- /dev/null +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/TextChatHistoryChatAgentIT.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.agent; + +import java.util.List; + +import io.qdrant.client.QdrantClient; +import io.qdrant.client.QdrantGrpcClient; +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.chat.agent.ChatAgent; +import org.springframework.ai.chat.agent.DefaultChatAgent; +import org.springframework.ai.chat.agent.DefaultStreamingChatAgent; +import org.springframework.ai.chat.agent.StreamingChatAgent; +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 = TextChatHistoryChatAgentIT.Config.class) +@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+") +public class TextChatHistoryChatAgentIT 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 TextChatHistoryChatAgentIT(RelevancyEvaluator relevancyEvaluator, ChatAgent chatAgent, + StreamingChatAgent streamingChatAgent) { + super(relevancyEvaluator, chatAgent, streamingChatAgent); + } + + @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 ChatAgent memoryChatAgent(OpenAiChatClient chatClient, VectorStore vectorStore, + TokenCountEstimator tokenCountEstimator) { + + return DefaultChatAgent.builder(chatClient) + .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 StreamingChatAgent memoryStreamingChatAgent(OpenAiChatClient streamingChatClient, + VectorStore vectorStore, TokenCountEstimator tokenCountEstimator) { + + return DefaultStreamingChatAgent.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/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/AgentResponse.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/AgentResponse.java index 21fe6eccf..3fa89c5b3 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/AgentResponse.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/AgentResponse.java @@ -1,3 +1,19 @@ +/* + * 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.agent; import org.springframework.ai.chat.ChatResponse; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgent.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgent.java index de9a4e3fa..ce664d445 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgent.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgent.java @@ -1,3 +1,19 @@ +/* + * 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.agent; import org.springframework.ai.chat.prompt.transformer.PromptContext; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgentListener.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgentListener.java index 53822c59e..4675dc25b 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgentListener.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgentListener.java @@ -1,7 +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.agent; +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 ChatAgent execution. + * + * @author Mark Pollack + * @author Christian Tzolov + */ public interface ChatAgentListener { + default void onStart(PromptContext promptContext) { + + } + void onComplete(AgentResponse agentResponse); } diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/DefaultChatAgent.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/DefaultChatAgent.java index 9c7a40e27..0408156d8 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/DefaultChatAgent.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/DefaultChatAgent.java @@ -1,3 +1,18 @@ +/* + * 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.agent; import org.springframework.ai.chat.ChatClient; @@ -9,6 +24,10 @@ import java.util.ArrayList; import java.util.List; import java.util.Objects; +/** + * @author Mark Pollack + * @author Christian Tzolov + */ public class DefaultChatAgent implements ChatAgent { private ChatClient chatClient; @@ -39,27 +58,34 @@ public class DefaultChatAgent implements ChatAgent { @Override public AgentResponse call(PromptContext promptContext) { + PromptContext promptContextOnStart = PromptContext.from(promptContext).build(); + // Perform retrieval of documents and messages - for (PromptTransformer retriever : retrievers) { + for (PromptTransformer retriever : this.retrievers) { promptContext = retriever.transform(promptContext); } - // Perform post procesing of all retrieved documents and messages - for (PromptTransformer documentPostProcessor : documentPostProcessors) { + // Perform post processing of all retrieved documents and messages + for (PromptTransformer documentPostProcessor : this.documentPostProcessors) { promptContext = documentPostProcessor.transform(promptContext); } // Perform prompt augmentation - for (PromptTransformer augmentor : augmentors) { + for (PromptTransformer augmentor : this.augmentors) { promptContext = augmentor.transform(promptContext); } + // Invoke Listeners onStart + for (ChatAgentListener listener : this.chatAgentListeners) { + listener.onStart(promptContextOnStart); + } + // Perform generation - ChatResponse chatResponse = chatClient.call(promptContext.getPrompt()); + ChatResponse chatResponse = this.chatClient.call(promptContext.getPrompt()); // Invoke Listeners onComplete AgentResponse agentResponse = new AgentResponse(promptContext, chatResponse); - for (ChatAgentListener listener : chatAgentListeners) { + for (ChatAgentListener listener : this.chatAgentListeners) { listener.onComplete(agentResponse); } return agentResponse; diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/DefaultStreamingChatAgent.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/DefaultStreamingChatAgent.java new file mode 100644 index 000000000..cd7452354 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/DefaultStreamingChatAgent.java @@ -0,0 +1,146 @@ +/* + * 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.agent; + +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 DefaultStreamingChatAgent implements StreamingChatAgent { + + private StreamingChatClient streamingChatClient; + + private List retrievers; + + private List documentPostProcessors; + + private List augmentors; + + private List chatAgentListeners; + + public DefaultStreamingChatAgent(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 StreamingAgentResponse 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 AgentResponse(promptContext2, chatResponse)); + } + }); + + // Invoke Listeners onComplete + StreamingAgentResponse agentResponse = new StreamingAgentResponse(promptContext, fluxChatResponse); + + return agentResponse; + } + + 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 DefaultStreamingChatAgent build() { + return new DefaultStreamingChatAgent(chatClient, retrievers, documentPostProcessors, augmentors, + chatAgentListeners); + } + + } + +} \ No newline at end of file diff --git a/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/StreamingAgentResponse.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/StreamingAgentResponse.java new file mode 100644 index 000000000..b1bc6ced9 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/StreamingAgentResponse.java @@ -0,0 +1,72 @@ +package org.springframework.ai.chat.agent; + +import reactor.core.publisher.Flux; + +import org.springframework.ai.chat.ChatResponse; +import org.springframework.ai.chat.prompt.transformer.PromptContext; + +/** + * Encapsulates the response from the ChatAgent. Contains the most up-to-date + * PromptContext and the final ChatResponse + * + * @author Mark Pollack + * @since 1.0 M1 + */ +public class StreamingAgentResponse { + + private final PromptContext promptContext; + + private final Flux chatResponse; + + public StreamingAgentResponse(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 "AgentResponse{" + "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; + StreamingAgentResponse other = (StreamingAgentResponse) 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/agent/StreamingChatAgent.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/StreamingChatAgent.java new file mode 100644 index 000000000..c1cb9fb8f --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/agent/StreamingChatAgent.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.agent; + +import org.springframework.ai.chat.prompt.transformer.PromptContext; + +/** + * A ChatAgent encapsulates common AI workflows such as Retrieval Augmented Generation. + * + * @author Mark Pollack + * @author Christian Tzolov + * @since 1.0 M1 + */ +public interface StreamingChatAgent { + + /** + * Call the chat agent to execute a workflow + * @param promptContext A shared data structure that can be used in components that + * implement the workflow. Contains the initial Prompt and a conversation ID at the + * start of the workflow. + * @return the AgentResponse that contains the ChatResponse and the latest + * PromptContext + */ + StreamingAgentResponse stream(PromptContext promptContext); + +} 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..798b88ada --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemory.java @@ -0,0 +1,38 @@ +/* + * 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 messages) { + this.add(conversationId, List.of(messages)); + } + + void add(String conversationId, List messages); + + List get(String conversationId); + + 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..23d5289fb --- /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.agent.AgentResponse; +import org.springframework.ai.chat.agent.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(AgentResponse agentResponse) { + List assistantMessages = agentResponse.getChatResponse() + .getResults() + .stream() + .map(g -> (Message) g.getOutput()) + .toList(); + this.chatHistory.add(agentResponse.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..6671c41aa --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemoryRetriever.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.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.TransformerContentType; +import org.springframework.ai.chat.prompt.transformer.InnerContent; +import org.springframework.ai.chat.prompt.transformer.PromptContext; +import org.springframework.ai.chat.prompt.transformer.PromptTransformer; +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; + + public ChatMemoryRetriever(ChatMemory chatHistory) { + this(chatHistory, Map.of()); + } + + public ChatMemoryRetriever(ChatMemory chatHistory, Map additionalMetadata) { + this.chatHistory = chatHistory; + this.additionalMetadata = additionalMetadata; + } + + @Override + public PromptContext transform(PromptContext promptContext) { + + List messageHistory = this.chatHistory.get(promptContext.getConversationId()); + + List historyContent = (messageHistory != null) + ? messageHistory.stream().filter(m -> m.getMessageType() != MessageType.SYSTEM).map(m -> { + Content content = new InnerContent(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..3f17f0ca0 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/InMemoryChatMemory.java @@ -0,0 +1,49 @@ +/* + * 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) { + return this.conversationHistory.get(conversationId); + } + + @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..67b579330 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/LastMaxTokenSizeContentTransformer.java @@ -0,0 +1,114 @@ +/* + * 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 doGetDatum(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.doGetDatum(promptContext); + + // int totalSize = this.tokenCountEstimator.estimate(nonSystemChatMessages) - + // retrievalRequest.getTokenRunningTotal(); + int totalSize = this.doEstimateTokenCount(datum); + + if (totalSize <= this.maxTokenSize) { + return promptContext; + } + + List newSessionMessages = this.purgeExcess(datum, totalSize); + + return PromptContext.from(promptContext).withContents(newSessionMessages).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.tokenCountEstimator.estimate(oldDatum); + 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..f48461227 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/history/VectorStoreChatMemoryAgentListener.java @@ -0,0 +1,97 @@ +/* + * 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.agent.AgentResponse; +import org.springframework.ai.chat.agent.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(AgentResponse agentResponse) { + if (!CollectionUtils.isEmpty(agentResponse.getChatResponse().getResults())) { + List assistantMessages = agentResponse.getChatResponse() + .getResults() + .stream() + .map(g -> (org.springframework.ai.chat.messages.Message) g.getOutput()) + .toList(); + + List docs = toDocuments(assistantMessages, agentResponse.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 ff78138ca..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; @@ -36,6 +38,8 @@ import org.springframework.util.StreamUtils; */ public abstract class AbstractMessage implements Message { + public static final String MESSAGE_TYPE = "messageType"; + protected final MessageType messageType; protected final String textContent; @@ -48,7 +52,7 @@ public abstract class AbstractMessage implements Message { 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 metadata) { @@ -56,11 +60,12 @@ public abstract class AbstractMessage implements Message { this.messageType = messageType; this.textContent = content; this.mediaData = new ArrayList<>(); - this.metadata = metadata; + 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, @@ -73,7 +78,8 @@ public abstract class AbstractMessage implements Message { this.messageType = messageType; this.textContent = textContent; this.mediaData = new ArrayList<>(mediaData); - this.metadata = metadata; + this.metadata = new HashMap<>(metadata); + this.metadata.put(MESSAGE_TYPE, messageType); } protected AbstractMessage(MessageType messageType, Resource resource) { @@ -86,7 +92,8 @@ public abstract class AbstractMessage implements Message { Assert.notNull(resource, "Resource must not be null"); this.messageType = messageType; - this.metadata = metadata; + this.metadata = new HashMap<>(metadata); + this.metadata.put(MESSAGE_TYPE, messageType); this.mediaData = new ArrayList<>(); try (InputStream inputStream = resource.getInputStream()) { @@ -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 + ((metadata == null) ? 0 : metadata.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 (metadata == null) { - if (other.metadata != null) - return false; - } - else if (!metadata.equals(other.metadata)) - 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/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/UserMessage.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/messages/UserMessage.java index 6f101579e..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,6 +44,10 @@ 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=" + metadata + ", 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/InnerContent.java b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/InnerContent.java new file mode 100644 index 000000000..6f02f7785 --- /dev/null +++ b/spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/InnerContent.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.prompt.transformer; + +import java.util.List; +import java.util.Map; + +import org.springframework.ai.chat.messages.Media; +import org.springframework.ai.model.Content; + +/** + * @author Christian Tzolov + */ +public class InnerContent implements Content { + + private final String content; + + private final List media; + + private final Map metadata; + + public InnerContent(String content) { + this(content, Map.of()); + } + + public InnerContent(String content, Map metadata) { + this(content, List.of(), metadata); + } + + public InnerContent(String content, List media, Map metadata) { + this.content = content; + this.media = media; + this.metadata = metadata; + } + + @Override + public String getContent() { + return this.content; + } + + @Override + public List getMedia() { + return this.media; + } + + @Override + public Map getMetadata() { + return this.metadata; + } + + @Override + public String toString() { + return "InnerContent [content=" + content + ", media=" + media + ", metadata=" + metadata + "]"; + } + +} 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 index 92db874c3..8c70f8d32 100644 --- 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 @@ -1,3 +1,19 @@ +/* + * 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; @@ -10,6 +26,7 @@ import java.util.*; * ChatAgent functionality. * * @author Mark Pollack + * @author Christian Tzolov * @since 1.0 M1 */ public class PromptContext { @@ -20,7 +37,7 @@ public class PromptContext { private List promptHistory; - private String conversationId; + private String conversationId = "default"; private Map metadata = new HashMap<>(); @@ -52,11 +69,11 @@ public class PromptContext { this.contents.add(datum); } - public List getNodes() { + public List getContents() { return contents; } - public void setNodes(List contents) { + public void setContents(List contents) { this.contents = contents; } @@ -76,6 +93,73 @@ public class PromptContext { 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 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 index b6fd2a91e..7d596a0e9 100644 --- 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 @@ -1,3 +1,19 @@ +/* + * 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; /** 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 index c50bdb8fe..9219ad0dc 100644 --- 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 @@ -1,16 +1,31 @@ +/* + * 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.document.Document; -import org.springframework.ai.node.Content; - -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; +import org.springframework.ai.model.Content; /** * Transforms the Prompt by taking to the current prompt in the Prompt Context and adding @@ -26,28 +41,28 @@ public class QuestionContextAugmentor implements PromptTransformer { "---------------------\\n" "{context}\\n" "---------------------\\n" - "Given the context information and not prior knowledge, " - "answer the question. If the answer is not in the context, inform " + "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" - "Question: {question}\\n" + "User comment: {question}\\n" "Answer: " """; @Override public PromptContext transform(PromptContext promptContext) { - String context = doCreateContext(promptContext.getNodes()); + String context = doCreateContext(promptContext.getContents()); Map contextMap = doCreateContextMap(promptContext.getPrompt(), context); Prompt prompt = doCreatePrompt(promptContext.getPrompt(), contextMap); promptContext.setPrompt(prompt); - promptContext.addPromptHistory(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(node -> node instanceof Document) - .map(node -> (Document) node) + .filter(content -> content.getMetadata().containsKey(TransformerContentType.QA)) .map(Content::getContent) .collect(Collectors.joining(System.lineSeparator())); } 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..8b56cbc9b --- /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 QA = "QA"; + +} 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 index f59c858f7..cffca5f50 100644 --- 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 @@ -1,3 +1,19 @@ +/* + * 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.messages.Message; @@ -39,9 +55,21 @@ public class VectorStoreRetriever implements PromptTransformer { .filter(m -> m.getMessageType() == MessageType.USER) .map(m -> m.getContent()) .collect(Collectors.joining(System.lineSeparator())); + List documents = vectorStore.similaritySearch(searchRequest.withQuery(userMessage)); + for (Document document : documents) { - promptContext.addData(document); + if (!document.getMetadata().containsKey(TransformerContentType.MEMORY)) { // TODO: + // Bad + // coupling + // with + // other + // transformers + // types. + var content = new InnerContent(document.getContent(), document.getMetadata()); + content.getMetadata().put(TransformerContentType.QA, true); + promptContext.addData(content); + } } return promptContext; } 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 index 7bd9324d0..d81d9f184 100644 --- 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 @@ -17,7 +17,7 @@ public class EvaluationRequest { private final ChatResponse chatResponse; public EvaluationRequest(AgentResponse agentResponse) { - this(agentResponse.getPromptContext().getPromptHistory().get(0), agentResponse.getPromptContext().getNodes(), + this(agentResponse.getPromptContext().getPromptHistory().get(0), agentResponse.getPromptContext().getContents(), agentResponse.getChatResponse()); } 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/ChatHistoryTests.java b/spring-ai-core/src/test/java/org/springframework/ai/chat/history/ChatHistoryTests.java new file mode 100644 index 000000000..001c66134 --- /dev/null +++ b/spring-ai-core/src/test/java/org/springframework/ai/chat/history/ChatHistoryTests.java @@ -0,0 +1,152 @@ +/* + * 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.agent.AgentResponse; +import org.springframework.ai.chat.agent.DefaultChatAgent; +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 ChatHistoryTests { + + @Mock + ChatClient chatClient; + + @Mock + StreamingChatClient streamingChatClient; + + @Captor + ArgumentCaptor promptCaptor; + + @Test + public void chatAgentMessageHistory() { + + ChatMemory chatHistory = new InMemoryChatMemory(); + + DefaultChatAgent chatAgent = DefaultChatAgent.builder(chatClient) + .withRetrievers(List.of(new ChatMemoryRetriever(chatHistory))) + .withDocumentPostProcessors( + 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 chatAgentTextHistory() { + + ChatMemory chatHistory = new InMemoryChatMemory(); + + DefaultChatAgent chatAgent = DefaultChatAgent.builder(chatClient) + .withRetrievers(List.of(new ChatMemoryRetriever(chatHistory))) + .withDocumentPostProcessors( + 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(DefaultChatAgent 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(); + + AgentResponse response1 = chatAgent.call(promptContext); + + assertThat(response1.getChatResponse().getResult().getOutput().getContent()).isEqualTo("assistant:1"); + + // List messages2 = promptCaptor.getValue().getInstructions(); + // assertThat(messages2) + // .isEqualTo(List.of(new UserMessage("user:1"), new UserMessage("user:2"), new UserMessage("user:3"), + // new UserMessage("user:4"), new UserMessage("user:5"))); + + List contents = response1.getPromptContext().getContents(); + assertThat(contents).hasSize(0); + + List history = chatHistory.get("test-session-id"); + assertThat(history).hasSize(6); + + AgentResponse 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"); + 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"); + + AgentResponse 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"); + 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-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..7d3c2f0bb --- /dev/null +++ b/spring-ai-test/src/main/java/org/springframework/ai/evaluation/BaseMemoryTest.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.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.agent.ChatAgent; +import org.springframework.ai.chat.agent.StreamingChatAgent; +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 ChatAgent chatAgent; + + protected StreamingChatAgent streamingChatAgent; + + public BaseMemoryTest(RelevancyEvaluator relevancyEvaluator, ChatAgent chatAgent, + StreamingChatAgent streamingChatClient) { + this.relevancyEvaluator = relevancyEvaluator; + this.chatAgent = chatAgent; + this.streamingChatAgent = streamingChatClient; + } + + @Test + void memoryChatAgent() { + + var prompt = new Prompt(new UserMessage("my name John Vincent Atanasoff")); + PromptContext promptContext = new PromptContext(prompt); + + var agentResponse1 = this.chatAgent.call(promptContext); + + logger.info("Response1: " + agentResponse1.getChatResponse().getResult().getOutput().getContent()); + assertThat(agentResponse1.getChatResponse().getResult().getOutput().getContent()).contains("John"); + + var agentResponse2 = this.chatAgent.call(new PromptContext(new Prompt(new String("What is my name?")))); + logger.info("Response2: " + agentResponse2.getChatResponse().getResult().getOutput().getContent()); + assertThat(agentResponse2.getChatResponse().getResult().getOutput().getContent()) + .contains("John Vincent Atanasoff"); + + EvaluationResponse evaluationResponse = this.relevancyEvaluator.evaluate(new EvaluationRequest(agentResponse2)); + logger.info("" + evaluationResponse); + } + + @Test + void memoryStreamingChatAgent() { + + var prompt = new Prompt(new UserMessage("my name John Vincent Atanasoff")); + PromptContext promptContext = new PromptContext(prompt); + + var fluxAgentResponse1 = this.streamingChatAgent.stream(promptContext); + + String agentResponse1 = fluxAgentResponse1.getChatResponse() + .collectList() + .block() + .stream() + .filter(response -> response.getResult() != null) + .map(response -> response.getResult().getOutput().getContent()) + .collect(Collectors.joining()); + + logger.info("Response1: " + agentResponse1); + assertThat(agentResponse1).contains("John"); + + var fluxAgentResponse2 = this.streamingChatAgent + .stream(new PromptContext(new Prompt(new String("What is my name?")))); + + String agentResponse2 = fluxAgentResponse2.getChatResponse() + .collectList() + .block() + .stream() + .filter(response -> response.getResult() != null) + .map(response -> response.getResult().getOutput().getContent()) + .collect(Collectors.joining()); + + logger.info("Response2: " + agentResponse2); + assertThat(agentResponse2).contains("John Vincent Atanasoff"); + } + +} 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); } }