Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add Short/Long term memory agent support. Add streaming agent support. Add mixing transofmers support

Browse Source
This commit is contained in:
Christian Tzolov
2024-04-25 12:58:24 +02:00
parent d99629243c
commit cb0743f40c
42 changed files with 2516 additions and 94 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
.gitignore vendored
View File

@@ -29,9 +29,13 @@ out
vscode
settings.json
node
node_modules
package-lock.json
package.json
.vscode
.antlr
shell.log
shell.log
.profiler

8
.mvn/extensions.xml Normal file
View File

@@ -0,0 +1,8 @@
<?xml version="1.0" encoding="UTF-8"?>
<extensions>
<extension>
<groupId>fr.jcgay.maven</groupId>
<artifactId>maven-profiler</artifactId>
<version>3.2</version>
</extension>
</extensions>

35
models/spring-ai-openai/pom.xml
View File

@@ -1,5 +1,6 @@
<?xml version="1.0" encoding="UTF-8"?>
<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<project xmlns="http://maven.apache.org/POM/4.0.0"
xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
<modelVersion>4.0.0</modelVersion>
<parent>
<groupId>org.springframework.ai</groupId>
@@ -74,6 +75,38 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-qdrant</artifactId>
<version>${project.version}</version>
<exclusions>
<exclusion>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-openai</artifactId>
</exclusion>
</exclusions>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>qdrant</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>testcontainers</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
</project>

107
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/MessageChatHistoryChatAgentIT.java Normal file
View File

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

16
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/OpenAiDefaultChatAgentIT.java
View File

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

108
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/OpenAiMemoryChatAgentIT.java Normal file
View File

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

223
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/TextChatHistoryChatAgent3IT.java Normal file
View File

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

131
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/agent/TextChatHistoryChatAgentIT.java Normal file
View File

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

16
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/AgentResponse.java
View File

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

16
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgent.java
View File

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

29
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/ChatAgentListener.java
View File

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

38
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/DefaultChatAgent.java
View File

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

146
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/DefaultStreamingChatAgent.java Normal file
View File

@@ -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<PromptTransformer> retrievers;
private List<PromptTransformer> documentPostProcessors;
private List<PromptTransformer> augmentors;
private List<ChatAgentListener> chatAgentListeners;
public DefaultStreamingChatAgent(StreamingChatClient chatClient, List<PromptTransformer> retrievers,
List<PromptTransformer> documentPostProcessors, List<PromptTransformer> augmentors,
List<ChatAgentListener> 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<ChatResponse> 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<PromptTransformer> retrievers = new ArrayList<>();
private List<PromptTransformer> documentPostProcessors = new ArrayList<>();
private List<PromptTransformer> augmentors = new ArrayList<>();
private List<ChatAgentListener> chatAgentListeners = new ArrayList<>();
public DefaultChatAgentBuilder withChatClient(StreamingChatClient chatClient) {
this.chatClient = chatClient;
return this;
}
public DefaultChatAgentBuilder withRetrievers(List<PromptTransformer> retrievers) {
this.retrievers = retrievers;
return this;
}
public DefaultChatAgentBuilder withDocumentPostProcessors(List<PromptTransformer> documentPostProcessors) {
this.documentPostProcessors = documentPostProcessors;
return this;
}
public DefaultChatAgentBuilder withAugmentors(List<PromptTransformer> augmentors) {
this.augmentors = augmentors;
return this;
}
public DefaultChatAgentBuilder withChatAgentListeners(List<ChatAgentListener> chatAgentListeners) {
this.chatAgentListeners = chatAgentListeners;
return this;
}
public DefaultStreamingChatAgent build() {
return new DefaultStreamingChatAgent(chatClient, retrievers, documentPostProcessors, augmentors,
chatAgentListeners);
}
}
}

72
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/StreamingAgentResponse.java Normal file
View File

@@ -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> chatResponse;
public StreamingAgentResponse(PromptContext promptContext, Flux<ChatResponse> chatResponse) {
this.promptContext = promptContext;
this.chatResponse = chatResponse;
}
public PromptContext getPromptContext() {
return promptContext;
}
public Flux<ChatResponse> 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;
}
}

39
spring-ai-core/src/main/java/org/springframework/ai/chat/agent/StreamingChatAgent.java Normal file
View File

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

38
spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemory.java Normal file
View File

@@ -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<Message> messages);
List<Message> get(String conversationId);
void clear(String conversationId);
}

60
spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemoryAgentListener.java Normal file
View File

@@ -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<Message> assistantMessages = agentResponse.getChatResponse()
.getResults()
.stream()
.map(g -> (Message) g.getOutput())
.toList();
this.chatHistory.add(agentResponse.getPromptContext().getConversationId(), assistantMessages);
}
}

74
spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemoryRetriever.java Normal file
View File

@@ -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<String, Object> additionalMetadata;
public ChatMemoryRetriever(ChatMemory chatHistory) {
this(chatHistory, Map.of());
}
public ChatMemoryRetriever(ChatMemory chatHistory, Map<String, Object> additionalMetadata) {
this.chatHistory = chatHistory;
this.additionalMetadata = additionalMetadata;
}
@Override
public PromptContext transform(PromptContext promptContext) {
List<Message> messageHistory = this.chatHistory.get(promptContext.getConversationId());
List<Content> 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<Content> updatedContents = new ArrayList<>(
promptContext.getContents() != null ? promptContext.getContents() : List.of());
updatedContents.addAll(historyContent);
return PromptContext.from(promptContext).withContents(updatedContents).build();
}
}

49
spring-ai-core/src/main/java/org/springframework/ai/chat/history/InMemoryChatMemory.java Normal file
View File

@@ -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<String, List<Message>> conversationHistory = new ConcurrentHashMap<>();
@Override
public void add(String conversationId, List<Message> messages) {
this.conversationHistory.putIfAbsent(conversationId, new ArrayList<>());
this.conversationHistory.get(conversationId).addAll(messages);
}
@Override
public List<Message> get(String conversationId) {
return this.conversationHistory.get(conversationId);
}
@Override
public void clear(String conversationId) {
this.conversationHistory.remove(conversationId);
}
}

114
spring-ai-core/src/main/java/org/springframework/ai/chat/history/LastMaxTokenSizeContentTransformer.java Normal file
View File

@@ -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<String> filterTags;
public LastMaxTokenSizeContentTransformer(TokenCountEstimator tokenCountEstimator, int maxTokenSize) {
this(tokenCountEstimator, maxTokenSize, Set.of());
}
public LastMaxTokenSizeContentTransformer(TokenCountEstimator tokenCountEstimator, int maxTokenSize,
Set<String> filterTags) {
this.tokenCountEstimator = tokenCountEstimator;
this.maxTokenSize = maxTokenSize;
this.filterTags = filterTags;
}
protected List<Content> 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<Content> datum) {
return datum.stream().mapToInt(this::doEstimateTokenCount).sum();
}
@Override
public PromptContext transform(PromptContext promptContext) {
List<Content> datum = this.doGetDatum(promptContext);
// int totalSize = this.tokenCountEstimator.estimate(nonSystemChatMessages) -
// retrievalRequest.getTokenRunningTotal();
int totalSize = this.doEstimateTokenCount(datum);
if (totalSize <= this.maxTokenSize) {
return promptContext;
}
List<Content> newSessionMessages = this.purgeExcess(datum, totalSize);
return PromptContext.from(promptContext).withContents(newSessionMessages).build();
}
protected List<Content> purgeExcess(List<Content> datum, int totalSize) {
int index = 0;
List<Content> 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;
}
}

70
spring-ai-core/src/main/java/org/springframework/ai/chat/history/MessageChatMemoryAugmentor.java Normal file
View File

@@ -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<Message> 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();
}
}

110
spring-ai-core/src/main/java/org/springframework/ai/chat/history/SystemPromptChatMemoryAugmentor.java Normal file
View File

@@ -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<String> filterTags;
public SystemPromptChatMemoryAugmentor() {
this(DEFAULT_HISTORY_PROMPT, new HashSet<>());
}
public SystemPromptChatMemoryAugmentor(Set<String> filterTags) {
this(DEFAULT_HISTORY_PROMPT, filterTags);
}
public SystemPromptChatMemoryAugmentor(String historyPrompt, Set<String> 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<Message> systemMessages = (originalPrompt.getInstructions() != null) ? originalPrompt.getInstructions()
.stream()
.filter(m -> m.getMessageType() == MessageType.SYSTEM)
.toList() : List.of();
List<Message> 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<Message> 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();
}
}

97
spring-ai-core/src/main/java/org/springframework/ai/chat/history/VectorStoreChatMemoryAgentListener.java Normal file
View File

@@ -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<String, Object> additionalMetadata;
public VectorStoreChatMemoryAgentListener(VectorStore vectorStore) {
this(vectorStore, new HashMap<>());
}
public VectorStoreChatMemoryAgentListener(VectorStore vectorStore, Map<String, Object> additionalMetadata) {
this.vectorStore = vectorStore;
this.additionalMetadata = additionalMetadata;
}
@Override
public void onStart(PromptContext promptContext) {
if (!CollectionUtils.isEmpty(promptContext.getPrompt().getInstructions())) {
List<Document> docs = toDocuments(promptContext.getPrompt().getInstructions(),
promptContext.getConversationId());
this.vectorStore.add(docs);
}
}
@Override
public void onComplete(AgentResponse agentResponse) {
if (!CollectionUtils.isEmpty(agentResponse.getChatResponse().getResults())) {
List<Message> assistantMessages = agentResponse.getChatResponse()
.getResults()
.stream()
.map(g -> (org.springframework.ai.chat.messages.Message) g.getOutput())
.toList();
List<Document> docs = toDocuments(assistantMessages, agentResponse.getPromptContext().getConversationId());
this.vectorStore.add(docs);
}
}
private List<Document> toDocuments(List<Message> messages, String conversationId) {
List<Document> 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;
}
}

88
spring-ai-core/src/main/java/org/springframework/ai/chat/history/VectorStoreChatMemoryRetriever.java Normal file
View File

@@ -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<String, Object> additionalMetadata;
public VectorStoreChatMemoryRetriever(VectorStore vectorStore, int topK) {
this(vectorStore, topK, Map.of());
}
public VectorStoreChatMemoryRetriever(VectorStore vectorStore, int topK, Map<String, Object> additionalMetadata) {
this.vectorStore = vectorStore;
this.topK = topK;
this.additionalMetadata = additionalMetadata;
}
@Override
public PromptContext transform(PromptContext promptContext) {
List<Content> 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<Document> 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();
}
}

50
spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AbstractMessage.java
View File

@@ -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<String, Object> 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<String, Object> 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<Media> mediaData) {
this(messageType, textContent, mediaData, Map.of());
this(messageType, textContent, mediaData, Map.of(MESSAGE_TYPE, messageType));
}
protected AbstractMessage(MessageType messageType, String textContent, List<Media> 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);
}
}

74
spring-ai-core/src/main/java/org/springframework/ai/chat/messages/MessageAggregator.java Normal file
View File

@@ -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<ChatResponse> aggregate(Flux<ChatResponse> fluxChatResponse,
Consumer<ChatResponse> onAggregationComplete) {
AtomicReference<StringBuilder> stringBufferRef = new AtomicReference<>(new StringBuilder());
AtomicReference<Map<String, Object>> 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);
});
}
}

5
spring-ai-core/src/main/java/org/springframework/ai/chat/messages/UserMessage.java
View File

@@ -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<Media> mediaList, Map<String, Object> metadata) {
super(MessageType.USER, textContent, mediaList, metadata);
}
@Override
public String toString() {
return "UserMessage{" + "content='" + getContent() + '\'' + ", properties=" + metadata + ", messageType="

41
spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/Prompt.java
View File

@@ -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<List<Message>> {
private final List<Message> messages;
@@ -92,4 +96,31 @@ public class Prompt implements ModelRequest<List<Message>> {
return Objects.hash(this.messages, this.modelOptions);
}
public Prompt copy() {
return new Prompt(instructionsCopy(), this.modelOptions);
}
private List<Message> instructionsCopy() {
List<Message> 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;
}
}

70
spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/InnerContent.java Normal file
View File

@@ -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> media;
private final Map<String, Object> metadata;
public InnerContent(String content) {
this(content, Map.of());
}
public InnerContent(String content, Map<String, Object> metadata) {
this(content, List.of(), metadata);
}
public InnerContent(String content, List<Media> media, Map<String, Object> metadata) {
this.content = content;
this.media = media;
this.metadata = metadata;
}
@Override
public String getContent() {
return this.content;
}
@Override
public List<Media> getMedia() {
return this.media;
}
@Override
public Map<String, Object> getMetadata() {
return this.metadata;
}
@Override
public String toString() {
return "InnerContent [content=" + content + ", media=" + media + ", metadata=" + metadata + "]";
}
}

90
spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/PromptContext.java
View File

@@ -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<Prompt> promptHistory;
private String conversationId;
private String conversationId = "default";
private Map<String, Object> metadata = new HashMap<>();
@@ -52,11 +69,11 @@ public class PromptContext {
this.contents.add(datum);
}
public List<Content> getNodes() {
public List<Content> getContents() {
return contents;
}
public void setNodes(List<Content> contents) {
public void setContents(List<Content> 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<Content> contents;
private List<Prompt> promptHistory;
private String conversationId;
private Map<String, Object> metadata = new HashMap<>();
public Builder withPrompt(Prompt prompt) {
this.prompt = prompt;
return this;
}
public Builder withContents(List<Content> contents) {
this.contents = new ArrayList<>(contents);
return this;
}
public Builder withPromptHistory(List<Prompt> 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<String, Object> 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

16
spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/PromptTransformer.java
View File

@@ -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;
/**

41
spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/QuestionContextAugmentor.java
View File

@@ -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<String, Object> 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<Content> 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()));
}

34
spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/TransformerContentType.java Normal file
View File

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

30
spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/VectorStoreRetriever.java
View File

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

2
spring-ai-core/src/main/java/org/springframework/ai/evaluation/EvaluationRequest.java
View File

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

85
spring-ai-core/src/main/java/org/springframework/ai/tokenizer/JTokkitTokenCountEstimator.java Normal file
View File

@@ -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<Content> contents) {
int totalSize = 0;
for (Content content : contents) {
totalSize += this.estimate(content);
}
return totalSize;
}
}

51
spring-ai-core/src/main/java/org/springframework/ai/tokenizer/TokenCountEstimator.java Normal file
View File

@@ -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<Content> messages);
}

152
spring-ai-core/src/test/java/org/springframework/ai/chat/history/ChatHistoryTests.java Normal file
View File

@@ -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<Prompt> 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<Message> 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<Content> contents = response1.getPromptContext().getContents();
assertThat(contents).hasSize(0);
List<Message> 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");
}
}

107
spring-ai-test/src/main/java/org/springframework/ai/evaluation/BaseMemoryTest.java Normal file
View File

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

2
vector-stores/spring-ai-pinecone/src/test/java/org/springframework/ai/vectorstore/PineconeVectorStoreIT.java
View File

@@ -292,7 +292,7 @@ public class PineconeVectorStoreIT {
}
@Bean
public EmbeddingClient embeddingClient() {
public TransformersEmbeddingClient embeddingClient() {
return new TransformersEmbeddingClient();
}

50
vector-stores/spring-ai-qdrant/pom.xml
View File

@@ -32,16 +32,10 @@
<version>${project.parent.version}</version>
</dependency>
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-web</artifactId>
</dependency>
<!-- <dependency>
<groupId>io.grpc</groupId>
<artifactId>grpc-protobuf</artifactId>
<version>1.59.0</version>
</dependency> -->
<dependency>
<groupId>org.springframework</groupId>
<artifactId>spring-web</artifactId>
</dependency>
<dependency>
<groupId>io.qdrant</groupId>
@@ -55,19 +49,20 @@
<version>${protobuf-java.version}</version>
</dependency>
<!-- TESTING -->
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-openai</artifactId>
<version>${project.parent.version}</version>
<scope>test</scope>
</dependency>
<!-- TESTING -->
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-azure-openai</artifactId>
<version>${project.parent.version}</version>
<scope>test</scope>
<optional>true</optional>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-test</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-test</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
@@ -75,11 +70,10 @@
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<scope>test</scope>
</dependency>
</dependencies>
</project>

20
vector-stores/spring-ai-qdrant/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java
View File

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