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

Add a new abstraction to simplify implementation of common ChatBot use cases

Browse Source 
* Add ChatBot and basic DefaultChatBot
* Add streaming ChatBot support.
* Add Evaluator interface and RelevancyEvaluator implementation
* Add Content data type abstraction for Document and Message
* Renaming and package refactoring
* update .gitignore to allow node package name
* Add List<Media> to node and move ai.transformer package to ai.prompt.transformer
* Add Short/Long term memory support.
* Add mixing transformers support

Docs TBD
This commit is contained in:
Mark Pollack
2024-04-18 13:09:47 -04:00
parent 012a2ad74a
commit dfb8bf6a44
54 changed files with 3376 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
.gitignore vendored
View File

@@ -36,4 +36,6 @@ 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>

4
models/spring-ai-huggingface/src/test/java/org/springframework/ai/huggingface/client/ClientIT.java
View File

@@ -57,8 +57,8 @@ public class ClientIT {
}
```""";
assertThat(chatResponse.getResult().getOutput().getContent()).isEqualTo(expectedResponse);
assertThat(chatResponse.getResult().getOutput().getProperties()).containsKey("generated_tokens");
assertThat(chatResponse.getResult().getOutput().getProperties()).containsEntry("generated_tokens", 39);
assertThat(chatResponse.getResult().getOutput().getMetadata()).containsKey("generated_tokens");
assertThat(chatResponse.getResult().getOutput().getMetadata()).containsEntry("generated_tokens", 39);
}

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>

131
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryLongTermSystemPromptIT.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.chatbot;
import java.util.List;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.chat.chatbot.ChatBot;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.qdrant.QdrantContainer;
import org.springframework.ai.chat.chatbot.DefaultChatBot;
import org.springframework.ai.chat.chatbot.DefaultStreamingChatBot;
import org.springframework.ai.chat.chatbot.StreamingChatBot;
import org.springframework.ai.chat.history.VectorStoreChatMemoryAgentListener;
import org.springframework.ai.chat.history.VectorStoreChatMemoryRetriever;
import org.springframework.ai.chat.history.LastMaxTokenSizeContentTransformer;
import org.springframework.ai.chat.history.SystemPromptChatMemoryAugmentor;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.evaluation.BaseMemoryTest;
import org.springframework.ai.evaluation.RelevancyEvaluator;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.ai.openai.OpenAiEmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
import org.springframework.ai.tokenizer.TokenCountEstimator;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
@Testcontainers
@SpringBootTest(classes = ChatMemoryLongTermSystemPromptIT.Config.class)
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
public class ChatMemoryLongTermSystemPromptIT extends BaseMemoryTest {
private static final String COLLECTION_NAME = "test_collection";
private static final int QDRANT_GRPC_PORT = 6334;
@Container
static QdrantContainer qdrantContainer = new QdrantContainer("qdrant/qdrant:v1.7.4");
@Autowired
public ChatMemoryLongTermSystemPromptIT(RelevancyEvaluator relevancyEvaluator, ChatBot chatBot,
StreamingChatBot streamingChatBot) {
super(relevancyEvaluator, chatBot, streamingChatBot);
}
@SpringBootConfiguration
static class Config {
@Bean
public OpenAiApi chatCompletionApi() {
return new OpenAiApi(System.getenv("OPENAI_API_KEY"));
}
@Bean
public OpenAiChatClient openAiClient(OpenAiApi openAiApi) {
return new OpenAiChatClient(openAiApi);
}
@Bean
public EmbeddingClient embeddingClient(OpenAiApi openAiApi) {
return new OpenAiEmbeddingClient(openAiApi);
}
@Bean
public VectorStore qdrantVectorStore(EmbeddingClient embeddingClient) {
QdrantClient qdrantClient = new QdrantClient(QdrantGrpcClient
.newBuilder(qdrantContainer.getHost(), qdrantContainer.getMappedPort(QDRANT_GRPC_PORT), false)
.build());
return new QdrantVectorStore(qdrantClient, COLLECTION_NAME, embeddingClient);
}
@Bean
public TokenCountEstimator tokenCountEstimator() {
return new JTokkitTokenCountEstimator();
}
@Bean
public ChatBot memoryChatAgent(OpenAiChatClient chatClient, VectorStore vectorStore,
TokenCountEstimator tokenCountEstimator) {
return DefaultChatBot.builder(chatClient)
.withRetrievers(List.of(new VectorStoreChatMemoryRetriever(vectorStore, 10)))
.withContentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
.withAugmentors(List.of(new SystemPromptChatMemoryAugmentor()))
.withChatAgentListeners(List.of(new VectorStoreChatMemoryAgentListener(vectorStore)))
.build();
}
@Bean
public StreamingChatBot memoryStreamingChatAgent(OpenAiChatClient streamingChatClient, VectorStore vectorStore,
TokenCountEstimator tokenCountEstimator) {
return DefaultStreamingChatBot.builder(streamingChatClient)
.withRetrievers(List.of(new VectorStoreChatMemoryRetriever(vectorStore, 10)))
.withDocumentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
.withAugmentors(List.of(new SystemPromptChatMemoryAugmentor()))
.withChatAgentListeners(List.of(new VectorStoreChatMemoryAgentListener(vectorStore)))
.build();
}
@Bean
public RelevancyEvaluator relevancyEvaluator(OpenAiChatClient chatClient) {
return new RelevancyEvaluator(chatClient);
}
}
}

107
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryShortTermMessageListIT.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.chatbot;
import java.util.List;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.chat.chatbot.ChatBot;
import org.springframework.ai.chat.chatbot.DefaultChatBot;
import org.springframework.ai.chat.chatbot.DefaultStreamingChatBot;
import org.springframework.ai.chat.chatbot.StreamingChatBot;
import org.springframework.ai.chat.history.ChatMemory;
import org.springframework.ai.chat.history.ChatMemoryAgentListener;
import org.springframework.ai.chat.history.ChatMemoryRetriever;
import org.springframework.ai.chat.history.InMemoryChatMemory;
import org.springframework.ai.chat.history.LastMaxTokenSizeContentTransformer;
import org.springframework.ai.chat.history.MessageChatMemoryAugmentor;
import org.springframework.ai.evaluation.BaseMemoryTest;
import org.springframework.ai.evaluation.RelevancyEvaluator;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
import org.springframework.ai.tokenizer.TokenCountEstimator;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
@SpringBootTest(classes = ChatMemoryShortTermMessageListIT.Config.class)
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
public class ChatMemoryShortTermMessageListIT extends BaseMemoryTest {
@Autowired
public ChatMemoryShortTermMessageListIT(RelevancyEvaluator relevancyEvaluator, ChatBot chatBot,
StreamingChatBot streamingChatBot) {
super(relevancyEvaluator, chatBot, streamingChatBot);
}
@SpringBootConfiguration
static class Config {
@Bean
public OpenAiApi chatCompletionApi() {
return new OpenAiApi(System.getenv("OPENAI_API_KEY"));
}
@Bean
public OpenAiChatClient openAiClient(OpenAiApi openAiApi) {
return new OpenAiChatClient(openAiApi);
}
@Bean
public ChatMemory chatHistory() {
return new InMemoryChatMemory();
}
@Bean
public TokenCountEstimator tokenCountEstimator() {
return new JTokkitTokenCountEstimator();
}
@Bean
public ChatBot memoryChatAgent(OpenAiChatClient chatClient, ChatMemory chatHistory,
TokenCountEstimator tokenCountEstimator) {
return DefaultChatBot.builder(chatClient)
.withRetrievers(List.of(new ChatMemoryRetriever(chatHistory)))
.withContentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
.withAugmentors(List.of(new MessageChatMemoryAugmentor()))
.withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory)))
.build();
}
@Bean
public StreamingChatBot memoryStreamingChatAgent(OpenAiChatClient streamingChatClient, ChatMemory chatHistory,
TokenCountEstimator tokenCountEstimator) {
return DefaultStreamingChatBot.builder(streamingChatClient)
.withRetrievers(List.of(new ChatMemoryRetriever(chatHistory)))
.withDocumentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
.withAugmentors(List.of(new MessageChatMemoryAugmentor()))
.withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory)))
.build();
}
@Bean
public RelevancyEvaluator relevancyEvaluator(OpenAiChatClient chatClient) {
return new RelevancyEvaluator(chatClient);
}
}
}

108
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/ChatMemoryShortTermSystemPromptIT.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.chatbot;
import java.util.List;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.chat.chatbot.ChatBot;
import org.springframework.ai.chat.chatbot.DefaultChatBot;
import org.springframework.ai.chat.chatbot.DefaultStreamingChatBot;
import org.springframework.ai.chat.chatbot.StreamingChatBot;
import org.springframework.ai.chat.history.ChatMemory;
import org.springframework.ai.chat.history.ChatMemoryAgentListener;
import org.springframework.ai.chat.history.ChatMemoryRetriever;
import org.springframework.ai.chat.history.InMemoryChatMemory;
import org.springframework.ai.chat.history.LastMaxTokenSizeContentTransformer;
import org.springframework.ai.chat.history.SystemPromptChatMemoryAugmentor;
import org.springframework.ai.evaluation.BaseMemoryTest;
import org.springframework.ai.evaluation.RelevancyEvaluator;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
import org.springframework.ai.tokenizer.TokenCountEstimator;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
@SpringBootTest(classes = ChatMemoryShortTermSystemPromptIT.Config.class)
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
public class ChatMemoryShortTermSystemPromptIT extends BaseMemoryTest {
@Autowired
public ChatMemoryShortTermSystemPromptIT(RelevancyEvaluator relevancyEvaluator, ChatBot chatBot,
StreamingChatBot streamingChatBot) {
super(relevancyEvaluator, chatBot, streamingChatBot);
}
@SpringBootConfiguration
static class Config {
@Bean
public OpenAiApi chatCompletionApi() {
return new OpenAiApi(System.getenv("OPENAI_API_KEY"));
}
@Bean
public OpenAiChatClient openAiClient(OpenAiApi openAiApi) {
return new OpenAiChatClient(openAiApi);
}
@Bean
public ChatMemory chatHistory() {
return new InMemoryChatMemory();
}
@Bean
public TokenCountEstimator tokenCountEstimator() {
return new JTokkitTokenCountEstimator();
}
@Bean
public ChatBot memoryChatAgent(OpenAiChatClient chatClient, ChatMemory chatHistory,
TokenCountEstimator tokenCountEstimator) {
return DefaultChatBot.builder(chatClient)
.withRetrievers(List.of(new ChatMemoryRetriever(chatHistory)))
.withContentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
.withAugmentors(List.of(new SystemPromptChatMemoryAugmentor()))
.withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory)))
.build();
}
@Bean
public StreamingChatBot memoryStreamingChatAgent(OpenAiChatClient streamingChatClient, ChatMemory chatHistory,
TokenCountEstimator tokenCountEstimator) {
return DefaultStreamingChatBot.builder(streamingChatClient)
.withRetrievers(List.of(new ChatMemoryRetriever(chatHistory)))
.withDocumentPostProcessors(List.of(new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
.withAugmentors(List.of(new SystemPromptChatMemoryAugmentor()))
.withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory)))
.build();
}
@Bean
public RelevancyEvaluator relevancyEvaluator(OpenAiChatClient chatClient) {
return new RelevancyEvaluator(chatClient);
}
}
}

252
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/LongShortTermChatMemoryWithRagIT.java Normal file
View File

@@ -0,0 +1,252 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.openai.chat.chatbot;
import java.util.List;
import java.util.Map;
import java.util.Set;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.chatbot.ChatBot;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.qdrant.QdrantContainer;
import org.springframework.ai.chat.chatbot.DefaultChatBot;
import org.springframework.ai.chat.history.ChatMemory;
import org.springframework.ai.chat.history.ChatMemoryAgentListener;
import org.springframework.ai.chat.history.ChatMemoryRetriever;
import org.springframework.ai.chat.history.InMemoryChatMemory;
import org.springframework.ai.chat.history.LastMaxTokenSizeContentTransformer;
import org.springframework.ai.chat.history.SystemPromptChatMemoryAugmentor;
import org.springframework.ai.chat.history.VectorStoreChatMemoryAgentListener;
import org.springframework.ai.chat.history.VectorStoreChatMemoryRetriever;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
import org.springframework.ai.chat.prompt.transformer.QuestionContextAugmentor;
import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
import org.springframework.ai.chat.prompt.transformer.VectorStoreRetriever;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentTransformer;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.evaluation.EvaluationRequest;
import org.springframework.ai.evaluation.EvaluationResponse;
import org.springframework.ai.evaluation.RelevancyEvaluator;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.ai.openai.OpenAiEmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
import org.springframework.ai.tokenizer.TokenCountEstimator;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.Resource;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.springframework.ai.openai.api.OpenAiApi.ChatModel.GPT_4_TURBO_PREVIEW;
@Testcontainers
@SpringBootTest(classes = LongShortTermChatMemoryWithRagIT.Config.class)
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
public class LongShortTermChatMemoryWithRagIT {
protected final Logger logger = LoggerFactory.getLogger(getClass());
private static final String COLLECTION_NAME = "test_collection";
private static final int QDRANT_GRPC_PORT = 6334;
@Container
static QdrantContainer qdrantContainer = new QdrantContainer("qdrant/qdrant:v1.7.4");
@Autowired
ChatBot chatBot;
@Autowired
RelevancyEvaluator relevancyEvaluator;
@Autowired
VectorStore vectorStore;
@Value("classpath:/data/acme/bikes.json")
private Resource bikesResource;
void loadData() {
var metadataEnricher = new DocumentTransformer() {
@Override
public List<Document> apply(List<Document> documents) {
documents.forEach(d -> {
Map<String, Object> metadata = d.getMetadata();
metadata.put(TransformerContentType.EXTERNAL_KNOWLEDGE, "true");
});
return documents;
}
};
JsonReader jsonReader = new JsonReader(bikesResource, "name", "price", "shortDescription", "description");
var textSplitter = new TokenTextSplitter();
vectorStore.accept(metadataEnricher.apply(textSplitter.apply(jsonReader.get())));
}
// @Autowired
// StreamingChatBot streamingChatBot;
@Test
void memoryChatBot() {
loadData();
var prompt = new Prompt(new UserMessage("My name is Christian and I like mountain bikes."));
PromptContext promptContext = new PromptContext(prompt);
var chatBotResponse1 = this.chatBot.call(promptContext);
logger.info("Response1: " + chatBotResponse1.getChatResponse().getResult().getOutput().getContent());
var chatBotResponse2 = this.chatBot.call(new PromptContext(
new Prompt(new String("What is my name and what bike model would you suggest for me?"))));
logger.info("Response2: " + chatBotResponse2.getChatResponse().getResult().getOutput().getContent());
// logger.info(chatBotResponse2.getPromptContext().getContents().toString());
assertThat(chatBotResponse2.getChatResponse().getResult().getOutput().getContent()).contains("Christian");
EvaluationResponse evaluationResponse = this.relevancyEvaluator
.evaluate(new EvaluationRequest(chatBotResponse2));
assertTrue(evaluationResponse.isPass(), "Response is not relevant to the question");
}
@SpringBootConfiguration
static class Config {
@Bean
public ChatMemory chatHistory() {
return new InMemoryChatMemory();
}
@Bean
public OpenAiApi chatCompletionApi() {
return new OpenAiApi(System.getenv("OPENAI_API_KEY"));
}
@Bean
public OpenAiChatClient openAiClient(OpenAiApi openAiApi) {
return new OpenAiChatClient(openAiApi);
}
@Bean
public OpenAiEmbeddingClient embeddingClient(OpenAiApi openAiApi) {
return new OpenAiEmbeddingClient(openAiApi);
}
@Bean
public VectorStore qdrantVectorStore(EmbeddingClient embeddingClient) {
QdrantClient qdrantClient = new QdrantClient(QdrantGrpcClient
.newBuilder(qdrantContainer.getHost(), qdrantContainer.getMappedPort(QDRANT_GRPC_PORT), false)
.build());
return new QdrantVectorStore(qdrantClient, COLLECTION_NAME, embeddingClient);
}
@Bean
public TokenCountEstimator tokenCountEstimator() {
return new JTokkitTokenCountEstimator();
}
@Bean
public ChatBot memoryChatAgent(OpenAiChatClient chatClient, VectorStore vectorStore,
TokenCountEstimator tokenCountEstimator, ChatMemory chatHistory) {
return DefaultChatBot.builder(chatClient)
.withRetrievers(List.of(new VectorStoreRetriever(vectorStore, SearchRequest.defaults()),
new ChatMemoryRetriever(chatHistory, Map.of(TransformerContentType.SHORT_TERM_MEMORY, "")),
new VectorStoreChatMemoryRetriever(vectorStore, 10,
Map.of(TransformerContentType.LONG_TERM_MEMORY, ""))))
.withContentPostProcessors(List.of(
new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000,
Set.of(TransformerContentType.SHORT_TERM_MEMORY)),
new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000,
Set.of(TransformerContentType.LONG_TERM_MEMORY)),
new LastMaxTokenSizeContentTransformer(tokenCountEstimator, 2000,
Set.of(TransformerContentType.EXTERNAL_KNOWLEDGE))))
.withAugmentors(List.of(new QuestionContextAugmentor(),
new SystemPromptChatMemoryAugmentor(
"""
Use the long term conversation history from the LONG TERM HISTORY section to provide accurate answers.
LONG TERM HISTORY:
{history}
""",
Set.of(TransformerContentType.LONG_TERM_MEMORY)),
new SystemPromptChatMemoryAugmentor(Set.of(TransformerContentType.SHORT_TERM_MEMORY))))
.withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory),
new VectorStoreChatMemoryAgentListener(vectorStore,
Map.of(TransformerContentType.LONG_TERM_MEMORY, ""))))
.build();
}
// @Bean
// public StreamingChatBot memoryStreamingChatAgent(OpenAiChatClient
// streamingChatClient,
// VectorStore vectorStore, TokenCountEstimator tokenCountEstimator, ChatHistory
// chatHistory) {
// return DefaultStreamingChatBot.builder(streamingChatClient)
// .withRetrievers(List.of(new ChatHistoryRetriever(chatHistory), new
// DocumentChatHistoryRetriever(vectorStore, 10)))
// .withDocumentPostProcessors(List.of(new
// LastMaxTokenSizeContentTransformer(tokenCountEstimator, 1000)))
// .withAugmentors(List.of(new TextChatHistoryAugmenter()))
// .withChatAgentListeners(List.of(new ChatHistoryAgentListener(chatHistory), new
// DocumentChatHistoryAgentListener(vectorStore)))
// .build();
// }
@Bean
public RelevancyEvaluator relevancyEvaluator(OpenAiChatClient chatClient) {
// Use GPT 4 as a better model for determining relevancy. gpt 3.5 makes basic
// mistakes
OpenAiChatOptions openAiChatOptions = OpenAiChatOptions.builder()
.withModel(GPT_4_TURBO_PREVIEW.getValue())
.build();
return new RelevancyEvaluator(chatClient, openAiChatOptions);
}
}
}

160
models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/chatbot/OpenAiDefaultChatBotIT.java Normal file
View File

@@ -0,0 +1,160 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.openai.chat.chatbot;
import java.util.List;
import io.qdrant.client.QdrantClient;
import io.qdrant.client.QdrantGrpcClient;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
import org.springframework.ai.chat.chatbot.ChatBot;
import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
import org.springframework.ai.document.Document;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.qdrant.QdrantContainer;
import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.chatbot.DefaultChatBot;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
import org.springframework.ai.chat.prompt.transformer.QuestionContextAugmentor;
import org.springframework.ai.chat.prompt.transformer.VectorStoreRetriever;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.ai.evaluation.EvaluationRequest;
import org.springframework.ai.evaluation.EvaluationResponse;
import org.springframework.ai.evaluation.RelevancyEvaluator;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.ai.openai.OpenAiEmbeddingClient;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.reader.JsonReader;
import org.springframework.ai.transformer.splitter.TokenTextSplitter;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.qdrant.QdrantVectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.Resource;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.springframework.ai.openai.api.OpenAiApi.ChatModel.GPT_4_TURBO_PREVIEW;
@Testcontainers
@SpringBootTest(classes = OpenAiDefaultChatBotIT.Config.class)
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
public class OpenAiDefaultChatBotIT {
private static final String COLLECTION_NAME = "test_collection";
private static final int QDRANT_GRPC_PORT = 6334;
@Container
static QdrantContainer qdrantContainer = new QdrantContainer("qdrant/qdrant:v1.7.4");
private final ChatClient chatClient;
private final VectorStore vectorStore;
@Value("classpath:/data/acme/bikes.json")
private Resource bikesResource;
private ChatBot chatBot;
@Autowired
public OpenAiDefaultChatBotIT(ChatClient chatClient, ChatBot chatBot, VectorStore vectorStore) {
this.chatClient = chatClient;
this.chatBot = chatBot;
this.vectorStore = vectorStore;
}
@Test
void simpleChat() {
loadData();
var prompt = new Prompt(new UserMessage("What bike is good for city commuting?"));
var chatBotResponse = this.chatBot.call(new PromptContext(prompt));
String answer = chatBotResponse.getChatResponse().getResult().getOutput().getContent();
assertTrue(answer.contains("Celerity"), "Response does not include 'Celerity'");
// Use GPT 4 as a better model for determining relevancy. gpt 3.5 makes basic
// mistakes
OpenAiChatOptions openAiChatOptions = OpenAiChatOptions.builder()
.withModel(GPT_4_TURBO_PREVIEW.getValue())
.build();
var relevancyEvaluator = new RelevancyEvaluator(this.chatClient, openAiChatOptions);
EvaluationRequest evaluationRequest = new EvaluationRequest(chatBotResponse);
EvaluationResponse evaluationResponse = relevancyEvaluator.evaluate(evaluationRequest);
assertTrue(evaluationResponse.isPass(), "Response is not relevant to the question");
}
void loadData() {
JsonReader jsonReader = new JsonReader(bikesResource, "name", "price", "shortDescription", "description");
var textSplitter = new TokenTextSplitter();
List<Document> splitDocuments = textSplitter.apply(jsonReader.get());
for (Document splitDocument : splitDocuments) {
splitDocument.getMetadata().put(TransformerContentType.EXTERNAL_KNOWLEDGE, "true");
}
vectorStore.accept(splitDocuments);
}
@SpringBootConfiguration
static class Config {
@Bean
public OpenAiApi chatCompletionApi() {
return new OpenAiApi(System.getenv("OPENAI_API_KEY"));
}
@Bean
public ChatClient openAiClient(OpenAiApi openAiApi) {
return new OpenAiChatClient(openAiApi);
}
@Bean
public EmbeddingClient embeddingClient(OpenAiApi openAiApi) {
return new OpenAiEmbeddingClient(openAiApi);
}
@Bean
public VectorStore qdrantVectorStore(EmbeddingClient embeddingClient) {
QdrantClient qdrantClient = new QdrantClient(QdrantGrpcClient
.newBuilder(qdrantContainer.getHost(), qdrantContainer.getMappedPort(QDRANT_GRPC_PORT), false)
.build());
return new QdrantVectorStore(qdrantClient, COLLECTION_NAME, embeddingClient);
}
@Bean
public ChatBot chatBot(ChatClient chatClient, VectorStore vectorStore) {
return DefaultChatBot.builder(chatClient)
.withRetrievers(List.of(new VectorStoreRetriever(vectorStore, SearchRequest.defaults())))
.withAugmentors(List.of(new QuestionContextAugmentor()))
.build();
}
}
}

36
spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatAgentListener.java Normal file
View File

@@ -0,0 +1,36 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.chatbot;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
/**
* The ChatAgentListener is a callback interface that can be implemented by classes that
* want to be notified of the completion of a ChatBot execution.
*
* @author Mark Pollack
* @author Christian Tzolov
*/
public interface ChatAgentListener {
default void onStart(PromptContext promptContext) {
}
void onComplete(ChatBotResponse chatBotResponse);
}

40
spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatBot.java Normal file
View File

@@ -0,0 +1,40 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.chatbot;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
/**
* A ChatBot encapsulates the logic to perform common AI use cases such as Retrieval
* Augmented Generation.
*
* @author Mark Pollack
* @since 1.0 M1
*/
public interface ChatBot {
/**
* Call the chatbot to execute AI actions
* @param promptContext A shared data structure used by the ChatBot to perform
* processing of the Prompt. It includes the intial Prompt and a conversation ID at
* the start of execution.
* @return the ChatBotResponse that contains the ChatResponse and the latest
* PromptContext
*/
ChatBotResponse call(PromptContext promptContext);
}

69
spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/ChatBotResponse.java Normal file
View File

@@ -0,0 +1,69 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.chatbot;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
import java.util.Objects;
/**
* Encapsulates the response from the ChatBot. Contains the most up-to-date PromptContext
* and the final ChatResponse
*
* @author Mark Pollack
* @since 1.0 M1
*/
public class ChatBotResponse {
private final PromptContext promptContext;
private final ChatResponse chatResponse;
public ChatBotResponse(PromptContext promptContext, ChatResponse chatResponse) {
this.promptContext = promptContext;
this.chatResponse = chatResponse;
}
public PromptContext getPromptContext() {
return promptContext;
}
public ChatResponse getChatResponse() {
return chatResponse;
}
@Override
public String toString() {
return "ChatBotResponse{" + "promptContext=" + promptContext + ", chatResponse=" + chatResponse + '}';
}
@Override
public boolean equals(Object o) {
if (this == o)
return true;
if (!(o instanceof ChatBotResponse that))
return false;
return Objects.equals(promptContext, that.promptContext) && Objects.equals(chatResponse, that.chatResponse);
}
@Override
public int hashCode() {
return Objects.hash(promptContext, chatResponse);
}
}

137
spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/DefaultChatBot.java Normal file
View File

@@ -0,0 +1,137 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.chatbot;
import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
import org.springframework.ai.chat.prompt.transformer.PromptTransformer;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
/**
* @author Mark Pollack
* @author Christian Tzolov
*/
public class DefaultChatBot implements ChatBot {
private ChatClient chatClient;
private List<PromptTransformer> retrievers;
private List<PromptTransformer> documentPostProcessors;
private List<PromptTransformer> augmentors;
private List<ChatAgentListener> chatAgentListeners;
public DefaultChatBot(ChatClient chatClient, List<PromptTransformer> retrievers,
List<PromptTransformer> documentPostProcessors, List<PromptTransformer> augmentors,
List<ChatAgentListener> chatAgentListeners) {
Objects.requireNonNull(chatClient, "chatClient must not be null");
this.chatClient = chatClient;
this.retrievers = retrievers;
this.documentPostProcessors = documentPostProcessors;
this.augmentors = augmentors;
this.chatAgentListeners = chatAgentListeners;
}
public static DefaultChatAgentBuilder builder(ChatClient chatClient) {
return new DefaultChatAgentBuilder().withChatClient(chatClient);
}
@Override
public ChatBotResponse call(PromptContext promptContext) {
PromptContext promptContextOnStart = PromptContext.from(promptContext).build();
// Perform retrieval of documents and messages
for (PromptTransformer retriever : this.retrievers) {
promptContext = retriever.transform(promptContext);
}
// Perform post processing of all retrieved documents and messages
for (PromptTransformer documentPostProcessor : this.documentPostProcessors) {
promptContext = documentPostProcessor.transform(promptContext);
}
// Perform prompt augmentation
for (PromptTransformer augmentor : this.augmentors) {
promptContext = augmentor.transform(promptContext);
}
// Invoke Listeners onStart
for (ChatAgentListener listener : this.chatAgentListeners) {
listener.onStart(promptContextOnStart);
}
// Perform generation
ChatResponse chatResponse = this.chatClient.call(promptContext.getPrompt());
// Invoke Listeners onComplete
ChatBotResponse chatBotResponse = new ChatBotResponse(promptContext, chatResponse);
for (ChatAgentListener listener : this.chatAgentListeners) {
listener.onComplete(chatBotResponse);
}
return chatBotResponse;
}
public static class DefaultChatAgentBuilder {
private ChatClient chatClient;
private List<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(ChatClient chatClient) {
this.chatClient = chatClient;
return this;
}
public DefaultChatAgentBuilder withRetrievers(List<PromptTransformer> retrievers) {
this.retrievers = retrievers;
return this;
}
public DefaultChatAgentBuilder withContentPostProcessors(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 DefaultChatBot build() {
return new DefaultChatBot(chatClient, retrievers, documentPostProcessors, augmentors, chatAgentListeners);
}
}
}

144
spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/DefaultStreamingChatBot.java Normal file
View File

@@ -0,0 +1,144 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.chatbot;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import reactor.core.publisher.Flux;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.StreamingChatClient;
import org.springframework.ai.chat.messages.MessageAggregator;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
import org.springframework.ai.chat.prompt.transformer.PromptTransformer;
/**
* @author Mark Pollack
* @author Christian Tzolov
*/
public class DefaultStreamingChatBot implements StreamingChatBot {
private StreamingChatClient streamingChatClient;
private List<PromptTransformer> retrievers;
private List<PromptTransformer> documentPostProcessors;
private List<PromptTransformer> augmentors;
private List<ChatAgentListener> chatAgentListeners;
public DefaultStreamingChatBot(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 StreamingChatBotResponse stream(PromptContext promptContext) {
PromptContext promptContextOnStart = PromptContext.from(promptContext).build();
// Perform retrieval of documents and messages
for (PromptTransformer retriever : this.retrievers) {
promptContext = retriever.transform(promptContext);
}
// Perform post processing of all retrieved documents and messages
for (PromptTransformer documentPostProcessor : this.documentPostProcessors) {
promptContext = documentPostProcessor.transform(promptContext);
}
// Perform prompt augmentation
for (PromptTransformer augmentor : this.augmentors) {
promptContext = augmentor.transform(promptContext);
}
// Invoke Listeners onStart
for (ChatAgentListener listener : this.chatAgentListeners) {
listener.onStart(promptContextOnStart);
}
// Perform generation
final var promptContext2 = promptContext;
Flux<ChatResponse> fluxChatResponse = new MessageAggregator()
.aggregate(this.streamingChatClient.stream(promptContext.getPrompt()), chatResponse -> {
for (ChatAgentListener listener : this.chatAgentListeners) {
listener.onComplete(new ChatBotResponse(promptContext2, chatResponse));
}
});
// Invoke Listeners onComplete
return new StreamingChatBotResponse(promptContext, fluxChatResponse);
}
public static class DefaultChatAgentBuilder {
private StreamingChatClient chatClient;
private List<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 DefaultStreamingChatBot build() {
return new DefaultStreamingChatBot(chatClient, retrievers, documentPostProcessors, augmentors,
chatAgentListeners);
}
}
}

40
spring-ai-core/src/main/java/org/springframework/ai/chat/chatbot/StreamingChatBot.java Normal file
View File

@@ -0,0 +1,40 @@
/*
* Copyright 2024 - 2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.chatbot;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
/**
* A ChatBot encapsulates the logic to perform common AI use cases such as Retrieval
* Augmented Generation.
*
* @author Mark Pollack
* @author Christian Tzolov
* @since 1.0 M1
*/
public interface StreamingChatBot {
/**
* Call the chatbot to execute AI actions
* @param promptContext A shared data structure used by the ChatBot to perform
* processing of the Prompt. It includes the intial Prompt and a conversation ID at
* the start of execution.
* @return the StreamingChatBotResponse that contains the ChatResponse and the latest
* PromptContext
*/
StreamingChatBotResponse stream(PromptContext promptContext);
}

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

@@ -0,0 +1,72 @@
package org.springframework.ai.chat.chatbot;
import reactor.core.publisher.Flux;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
/**
* Encapsulates the response from the ChatBot. Contains the most up-to-date PromptContext
* and the final ChatResponse
*
* @author Mark Pollack
* @since 1.0 M1
*/
public class StreamingChatBotResponse {
private final PromptContext promptContext;
private final Flux<ChatResponse> chatResponse;
public StreamingChatBotResponse(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 "ChatBotResponse{" + "promptContext=" + promptContext + ", chatResponse=" + chatResponse + '}';
}
@Override
public int hashCode() {
final int prime = 31;
int result = 1;
result = prime * result + ((promptContext == null) ? 0 : promptContext.hashCode());
result = prime * result + ((chatResponse == null) ? 0 : chatResponse.hashCode());
return result;
}
@Override
public boolean equals(Object obj) {
if (this == obj)
return true;
if (obj == null)
return false;
if (getClass() != obj.getClass())
return false;
StreamingChatBotResponse other = (StreamingChatBotResponse) obj;
if (promptContext == null) {
if (other.promptContext != null)
return false;
}
else if (!promptContext.equals(other.promptContext))
return false;
if (chatResponse == null) {
if (other.chatResponse != null)
return false;
}
else if (!chatResponse.equals(other.chatResponse))
return false;
return true;
}
}

39
spring-ai-core/src/main/java/org/springframework/ai/chat/history/ChatMemory.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.history;
import java.util.List;
import org.springframework.ai.chat.messages.Message;
/**
* @author Christian Tzolov
*
*/
public interface ChatMemory {
default void add(String conversationId, Message message) {
this.add(conversationId, List.of(message));
}
void add(String conversationId, List<Message> messages);
List<Message> get(String conversationId, int lastN);
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.chatbot.ChatBotResponse;
import org.springframework.ai.chat.chatbot.ChatAgentListener;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
/**
* @author Christian Tzolov
*/
public class ChatMemoryAgentListener implements ChatAgentListener {
private final ChatMemory chatHistory;
public ChatMemoryAgentListener(ChatMemory chatHistory) {
this.chatHistory = chatHistory;
}
@Override
public void onStart(PromptContext promptContext) {
var messagesToAdd = promptContext.getPrompt()
.getInstructions()
.stream()
.filter(m -> !m.getMetadata().containsKey(TransformerContentType.MEMORY))
.filter(m -> (m.getMessageType() == MessageType.ASSISTANT || m.getMessageType() == MessageType.USER))
.toList();
this.chatHistory.add(promptContext.getConversationId(), messagesToAdd);
}
@Override
public void onComplete(ChatBotResponse chatBotResponse) {
List<Message> assistantMessages = chatBotResponse.getChatResponse()
.getResults()
.stream()
.map(g -> (Message) g.getOutput())
.toList();
this.chatHistory.add(chatBotResponse.getPromptContext().getConversationId(), assistantMessages);
}
}

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

@@ -0,0 +1,81 @@
/*
* Copyright 2024-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.history;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
import org.springframework.ai.chat.prompt.transformer.PromptTransformer;
import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
import org.springframework.ai.document.Document;
import org.springframework.ai.model.Content;
/**
* @author Christian Tzolov
*/
public class ChatMemoryRetriever implements PromptTransformer {
private final ChatMemory chatHistory;
/**
* Additional metadata to be assigned to the retrieved history messages.
*/
private final Map<String, Object> additionalMetadata;
private final int maxHistorySize;
public ChatMemoryRetriever(ChatMemory chatHistory) {
this(chatHistory, Map.of());
}
public ChatMemoryRetriever(ChatMemory chatHistory, Map<String, Object> additionalMetadata) {
this(chatHistory, 1000, additionalMetadata);
}
public ChatMemoryRetriever(ChatMemory chatHistory, int maxHistorySize, Map<String, Object> additionalMetadata) {
this.chatHistory = chatHistory;
this.additionalMetadata = additionalMetadata;
this.maxHistorySize = maxHistorySize;
}
@Override
public PromptContext transform(PromptContext promptContext) {
List<Message> messageHistory = this.chatHistory.get(promptContext.getConversationId(), maxHistorySize);
List<Content> historyContent = (messageHistory != null)
? messageHistory.stream().filter(m -> m.getMessageType() != MessageType.SYSTEM).map(m -> {
Content content = new Document(m.getContent(), new ArrayList<>(m.getMedia()),
new HashMap<>(m.getMetadata()));
content.getMetadata().putAll(this.additionalMetadata);
content.getMetadata().put(TransformerContentType.MEMORY, true);
return content;
}).toList() : List.of();
List<Content> updatedContents = new ArrayList<>(
promptContext.getContents() != null ? promptContext.getContents() : List.of());
updatedContents.addAll(historyContent);
return PromptContext.from(promptContext).withContents(updatedContents).build();
}
}

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

@@ -0,0 +1,50 @@
/*
* Copyright 2024-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.history;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import org.springframework.ai.chat.messages.Message;
/**
* @author Christian Tzolov
*/
public class InMemoryChatMemory implements ChatMemory {
Map<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, int lastN) {
List<Message> all = this.conversationHistory.get(conversationId);
return all != null ? all.stream().skip(Math.max(0, all.size() - lastN)).toList() : List.of();
}
@Override
public void clear(String conversationId) {
this.conversationHistory.remove(conversationId);
}
}

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

@@ -0,0 +1,121 @@
/*
* Copyright 2024-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.history;
import java.util.ArrayList;
import java.util.List;
import java.util.Set;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
import org.springframework.ai.chat.prompt.transformer.PromptTransformer;
import org.springframework.ai.model.Content;
import org.springframework.ai.tokenizer.TokenCountEstimator;
/**
* Returns a new list of content (e.g list of messages of list of documents) that is a
* subset of the input list of contents and complies with the max token size constraint.
*
* The token estimator is used to estimate the token count of the datum.
*
* @author Christian Tzolov
*/
public class LastMaxTokenSizeContentTransformer implements PromptTransformer {
protected final TokenCountEstimator tokenCountEstimator;
protected final int maxTokenSize;
/**
* Only Content entries with the following metadata tags will be included in the
* history.
*/
private final Set<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> doGetDatumToModify(PromptContext promptContext) {
return promptContext.getContents()
.stream()
.filter(content -> this.filterTags.stream().allMatch(tag -> content.getMetadata().containsKey(tag)))
.toList();
}
protected List<Content> doGetDatumNotToModify(PromptContext promptContext) {
return promptContext.getContents()
.stream()
.filter(content -> !this.filterTags.stream().allMatch(tag -> content.getMetadata().containsKey(tag)))
.toList();
}
protected int doEstimateTokenCount(Content datum) {
return this.tokenCountEstimator.estimate(datum);
}
protected int doEstimateTokenCount(List<Content> datum) {
return datum.stream().mapToInt(this::doEstimateTokenCount).sum();
}
@Override
public PromptContext transform(PromptContext promptContext) {
List<Content> datum = this.doGetDatumToModify(promptContext);
int totalSize = this.doEstimateTokenCount(datum);
if (totalSize <= this.maxTokenSize) {
return promptContext;
}
List<Content> purgedContent = this.purgeExcess(datum, totalSize);
var updatedContent = new ArrayList<>(doGetDatumNotToModify(promptContext));
updatedContent.addAll(purgedContent);
return PromptContext.from(promptContext).withContents(updatedContent).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.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();
}
}

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

@@ -0,0 +1,98 @@
/*
* Copyright 2024-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.history;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.springframework.ai.chat.chatbot.ChatBotResponse;
import org.springframework.ai.chat.chatbot.ChatAgentListener;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.prompt.transformer.TransformerContentType;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.util.CollectionUtils;
/**
* @author Christian Tzolov
*/
public class VectorStoreChatMemoryAgentListener implements ChatAgentListener {
private final VectorStore vectorStore;
private final Map<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(ChatBotResponse chatBotResponse) {
if (!CollectionUtils.isEmpty(chatBotResponse.getChatResponse().getResults())) {
List<Message> assistantMessages = chatBotResponse.getChatResponse()
.getResults()
.stream()
.map(g -> (org.springframework.ai.chat.messages.Message) g.getOutput())
.toList();
List<Document> docs = toDocuments(assistantMessages,
chatBotResponse.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();
}
}

64
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;
@@ -29,13 +31,15 @@ import org.springframework.util.StreamUtils;
/**
* The AbstractMessage class is an abstract implementation of the Message interface. It
* provides a base implementation for message content, media attachments, properties, and
* provides a base implementation for message content, media attachments, metadata, and
* message type.
*
* @see Message
*/
public abstract class AbstractMessage implements Message {
public static final String MESSAGE_TYPE = "messageType";
protected final MessageType messageType;
protected final String textContent;
@@ -45,26 +49,27 @@ public abstract class AbstractMessage implements Message {
/**
* Additional options for the message to influence the response, not a generative map.
*/
protected final Map<String, Object> properties;
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> messageProperties) {
protected AbstractMessage(MessageType messageType, String content, Map<String, Object> metadata) {
Assert.notNull(messageType, "Message type must not be null");
this.messageType = messageType;
this.textContent = content;
this.mediaData = new ArrayList<>();
this.properties = messageProperties;
this.metadata = new HashMap<>(metadata);
this.metadata.put(MESSAGE_TYPE, messageType);
}
protected AbstractMessage(MessageType messageType, String textContent, List<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,
Map<String, Object> messageProperties) {
Map<String, Object> metadata) {
Assert.notNull(messageType, "Message type must not be null");
Assert.notNull(textContent, "Content must not be null");
@@ -73,7 +78,8 @@ public abstract class AbstractMessage implements Message {
this.messageType = messageType;
this.textContent = textContent;
this.mediaData = new ArrayList<>(mediaData);
this.properties = messageProperties;
this.metadata = new HashMap<>(metadata);
this.metadata.put(MESSAGE_TYPE, messageType);
}
protected AbstractMessage(MessageType messageType, Resource resource) {
@@ -81,12 +87,13 @@ public abstract class AbstractMessage implements Message {
}
@SuppressWarnings("null")
protected AbstractMessage(MessageType messageType, Resource resource, Map<String, Object> messageProperties) {
protected AbstractMessage(MessageType messageType, Resource resource, Map<String, Object> metadata) {
Assert.notNull(messageType, "Message type must not be null");
Assert.notNull(resource, "Resource must not be null");
this.messageType = messageType;
this.properties = messageProperties;
this.metadata = new HashMap<>(metadata);
this.metadata.put(MESSAGE_TYPE, messageType);
this.mediaData = new ArrayList<>();
try (InputStream inputStream = resource.getInputStream()) {
@@ -108,8 +115,8 @@ public abstract class AbstractMessage implements Message {
}
@Override
public Map<String, Object> getProperties() {
return this.properties;
public Map<String, Object> getMetadata() {
return this.metadata;
}
@Override
@@ -119,38 +126,21 @@ public abstract class AbstractMessage implements Message {
@Override
public int hashCode() {
final int prime = 31;
int result = 1;
result = prime * result + ((mediaData == null) ? 0 : mediaData.hashCode());
result = prime * result + ((properties == null) ? 0 : properties.hashCode());
result = prime * result + ((messageType == null) ? 0 : messageType.hashCode());
return result;
return Objects.hash(this.messageType, this.textContent, this.mediaData, this.metadata);
}
@Override
public boolean equals(Object obj) {
if (this == obj)
if (this == obj) {
return true;
if (obj == null)
return false;
if (getClass() != obj.getClass())
}
if (obj == null || getClass() != obj.getClass()) {
return false;
}
AbstractMessage other = (AbstractMessage) obj;
if (mediaData == null) {
if (other.mediaData != null)
return false;
}
else if (!mediaData.equals(other.mediaData))
return false;
if (properties == null) {
if (other.properties != null)
return false;
}
else if (!properties.equals(other.properties))
return false;
if (messageType != other.messageType)
return false;
return true;
return Objects.equals(this.messageType, other.messageType)
&& Objects.equals(this.textContent, other.textContent)
&& Objects.equals(this.mediaData, other.mediaData) && Objects.equals(this.metadata, other.metadata);
}
}

2
spring-ai-core/src/main/java/org/springframework/ai/chat/messages/AssistantMessage.java
View File

@@ -35,7 +35,7 @@ public class AssistantMessage extends AbstractMessage {
@Override
public String toString() {
return "AssistantMessage{" + "content='" + getContent() + '\'' + ", properties=" + properties + ", messageType="
return "AssistantMessage{" + "content='" + getContent() + '\'' + ", properties=" + metadata + ", messageType="
+ messageType + '}';
}

2
spring-ai-core/src/main/java/org/springframework/ai/chat/messages/FunctionMessage.java
View File

@@ -33,7 +33,7 @@ public class FunctionMessage extends AbstractMessage {
@Override
public String toString() {
return "FunctionMessage{" + "content='" + getContent() + '\'' + ", properties=" + properties + ", messageType="
return "FunctionMessage{" + "content='" + getContent() + '\'' + ", properties=" + metadata + ", messageType="
+ messageType + '}';
}

11
spring-ai-core/src/main/java/org/springframework/ai/chat/messages/Message.java
View File

@@ -15,8 +15,7 @@
*/
package org.springframework.ai.chat.messages;
import java.util.List;
import java.util.Map;
import org.springframework.ai.model.Content;
/**
* The Message interface represents a message that can be sent or received in a chat
@@ -26,13 +25,7 @@ import java.util.Map;
* @see Media
* @see MessageType
*/
public interface Message {
String getContent();
List<Media> getMedia();
Map<String, Object> getProperties();
public interface Message extends Content {
MessageType getMessageType();

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

2
spring-ai-core/src/main/java/org/springframework/ai/chat/messages/SystemMessage.java
View File

@@ -36,7 +36,7 @@ public class SystemMessage extends AbstractMessage {
@Override
public String toString() {
return "SystemMessage{" + "content='" + getContent() + '\'' + ", properties=" + properties + ", messageType="
return "SystemMessage{" + "content='" + getContent() + '\'' + ", properties=" + metadata + ", messageType="
+ messageType + '}';
}

7
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,9 +44,13 @@ 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=" + properties + ", messageType="
return "UserMessage{" + "content='" + getContent() + '\'' + ", properties=" + metadata + ", messageType="
+ 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;
}
}

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

@@ -0,0 +1,185 @@
/*
* Copyright 2024-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.prompt.transformer;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.Content;
import java.util.*;
/**
* The shared, at the moment, mutable, data structure that can be used to implement
* ChatBot functionality.
*
* @author Mark Pollack
* @author Christian Tzolov
* @since 1.0 M1
*/
public class PromptContext {
private Prompt prompt; // The most up-to-date prompt to use
private List<Content> contents; // The most up-to-date data to use
private List<Prompt> promptHistory;
private String conversationId = "default";
private Map<String, Object> metadata = new HashMap<>();
public PromptContext(Prompt prompt) {
this(prompt, new ArrayList<>());
}
public PromptContext(Prompt prompt, String conversationId) {
this(prompt, new ArrayList<>());
this.conversationId = conversationId;
}
public PromptContext(Prompt prompt, List<Content> contents) {
this.prompt = prompt;
this.promptHistory = new ArrayList<>();
this.promptHistory.add(prompt);
this.contents = contents;
}
public Prompt getPrompt() {
return prompt;
}
public void setPrompt(Prompt prompt) {
this.prompt = prompt;
}
public void addData(Content datum) {
this.contents.add(datum);
}
public List<Content> getContents() {
return contents;
}
public void setContents(List<Content> contents) {
this.contents = contents;
}
public void addPromptHistory(Prompt prompt) {
this.promptHistory.add(prompt);
}
public List<Prompt> getPromptHistory() {
return promptHistory;
}
public String getConversationId() {
return conversationId;
}
public Map<String, Object> getMetadata() {
return metadata;
}
public static Builder from(PromptContext promptContext) {
return PromptContext.builder()
.withContents(
new ArrayList<>(promptContext.getContents() != null ? promptContext.getContents() : List.of()))
.withPrompt(promptContext.getPrompt().copy()) // deep copy
.withMetadata(new HashMap<>(promptContext.getMetadata() != null ? promptContext.getMetadata() : Map.of()))
.withPromptHistory(new ArrayList<>(
promptContext.getPromptHistory() != null ? promptContext.getPromptHistory() : List.of()))
.withConversationId(promptContext.getConversationId());
}
public static Builder builder() {
return new Builder();
}
public static class Builder {
private Prompt prompt;
private List<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
+ ", conversationId='" + conversationId + '\'' + ", metadata=" + metadata + '}';
}
@Override
public boolean equals(Object o) {
if (this == o)
return true;
if (!(o instanceof PromptContext that))
return false;
return Objects.equals(prompt, that.prompt) && Objects.equals(contents, that.contents)
&& Objects.equals(promptHistory, that.promptHistory)
&& Objects.equals(conversationId, that.conversationId) && Objects.equals(metadata, that.metadata);
}
@Override
public int hashCode() {
return Objects.hash(prompt, contents, promptHistory, conversationId, metadata);
}
}

39
spring-ai-core/src/main/java/org/springframework/ai/chat/prompt/transformer/PromptTransformer.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.prompt.transformer;
/**
* Responsible for transforming a Prompt. The PromptContext contains the necessary data to
* make the transformation
*
* Implementations may retrieve data and modify the Prompt object in the PromptContext as
* needed.
*
* @author Mark Pollack
* @since 1.0 M1
*/
@FunctionalInterface
public interface PromptTransformer {
/**
* Transforms the given PromptContext.
* @param context the PromptContext to transform
* @return the transformed PromptContext
*/
PromptContext transform(PromptContext context);
}

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

@@ -0,0 +1,91 @@
/*
* Copyright 2024-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.prompt.transformer;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.model.Content;
/**
* Transforms the Prompt by taking to the current prompt in the Prompt Context and adding
* additional context to create a new prompt. The default user text contains the
* placeholder names "question" and "context". The "question" placeholder is filled using
* the value of the current UserMessage and the "context" placeholder is filled with
* Documents contained in the PromptContext's Nodes.
*/
public class QuestionContextAugmentor implements PromptTransformer {
private static final String DEFAULT_USER_PROMPT_TEXT = """
"Context information is below.\\n"
"---------------------\\n"
"{context}\\n"
"---------------------\\n"
"Given the context and provided history information and not prior knowledge, "
"reply to the user comment. If the answer is not in the context, inform "
"the user that you can't answer the question.\\n"
"User comment: {question}\\n"
"Answer: "
""";
@Override
public PromptContext transform(PromptContext promptContext) {
String context = doCreateContext(promptContext.getContents());
Map<String, Object> contextMap = doCreateContextMap(promptContext.getPrompt(), context);
Prompt prompt = doCreatePrompt(promptContext.getPrompt(), contextMap);
promptContext.setPrompt(prompt);
promptContext.addPromptHistory(prompt); // BUG? shouldn't this be original
// promptContext.getPrompt()?
// For now return the modified instance instead of a copy
return promptContext;
}
protected String doCreateContext(List<Content> data) {
return data.stream()
.filter(content -> content.getMetadata().containsKey(TransformerContentType.EXTERNAL_KNOWLEDGE))
.map(Content::getContent)
.collect(Collectors.joining(System.lineSeparator()));
}
private Map<String, Object> doCreateContextMap(Prompt prompt, String context) {
String originalUserMessage = prompt.getInstructions()
.stream()
.filter(m -> m.getMessageType() == MessageType.USER)
.map(m -> m.getContent())
.collect(Collectors.joining(System.lineSeparator()));
return Map.of("context", context, "question", originalUserMessage);
}
protected Prompt doCreatePrompt(Prompt originalPrompt, Map<String, Object> contextMap) {
PromptTemplate promptTemplate = new PromptTemplate(DEFAULT_USER_PROMPT_TEXT);
Message userMessageToAppend = promptTemplate.createMessage(contextMap);
List<Message> messageList = originalPrompt.getInstructions()
.stream()
.filter(m -> m.getMessageType() != MessageType.USER)
.collect(Collectors.toList());
messageList.add(userMessageToAppend);
return new Prompt(messageList, (ChatOptions) originalPrompt.getOptions());
}
}

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 EXTERNAL_KNOWLEDGE = "externalKnowledge";
}

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

@@ -0,0 +1,94 @@
/*
* Copyright 2024-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.prompt.transformer;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;
/**
* Transforms the PromptContext by retrieving documents from a VectorStore
*/
public class VectorStoreRetriever implements PromptTransformer {
private final Logger logger = LoggerFactory.getLogger(getClass());
private final VectorStore vectorStore;
private final SearchRequest searchRequest;
public VectorStoreRetriever(VectorStore vectorStore, SearchRequest searchRequest) {
this.vectorStore = vectorStore;
this.searchRequest = searchRequest;
}
public VectorStore getVectorStore() {
return vectorStore;
}
public SearchRequest getSearchRequest() {
return searchRequest;
}
@Override
public PromptContext transform(PromptContext promptContext) {
List<Message> instructions = promptContext.getPrompt().getInstructions();
String userMessage = instructions.stream()
.filter(m -> m.getMessageType() == MessageType.USER)
.map(m -> m.getContent())
.collect(Collectors.joining(System.lineSeparator()));
List<Document> documents = vectorStore.similaritySearch(searchRequest.withQuery(userMessage)
.withFilterExpression(TransformerContentType.EXTERNAL_KNOWLEDGE + "=='true'"));
logger.info("Retrieved {} documents for user message {}", documents.size(), userMessage);
for (Document document : documents) {
var content = new Document(document.getContent(), document.getMetadata());
// content.getMetadata().put(TransformerContentType.DOMAIN_DATA, true);
promptContext.addData(content);
}
return promptContext;
}
@Override
public String toString() {
return "VectorStoreRetriever{" + "vectorStore=" + vectorStore + ", searchRequest=" + searchRequest + '}';
}
@Override
public boolean equals(Object o) {
if (this == o)
return true;
if (!(o instanceof VectorStoreRetriever that))
return false;
return Objects.equals(vectorStore, that.vectorStore) && Objects.equals(searchRequest, that.searchRequest);
}
@Override
public int hashCode() {
return Objects.hash(vectorStore, searchRequest);
}
}

22
spring-ai-core/src/main/java/org/springframework/ai/document/Document.java
View File

@@ -25,8 +25,10 @@ import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;
import org.springframework.ai.chat.messages.Media;
import org.springframework.ai.document.id.IdGenerator;
import org.springframework.ai.document.id.RandomIdGenerator;
import org.springframework.ai.model.Content;
import org.springframework.util.Assert;
/**
@@ -34,7 +36,7 @@ import org.springframework.util.Assert;
* the document's unique ID and an optional embedding.
*/
@JsonIgnoreProperties({ "contentFormatter" })
public class Document {
public class Document implements Content {
public final static ContentFormatter DEFAULT_CONTENT_FORMATTER = DefaultContentFormatter.defaultConfig();
@@ -54,6 +56,8 @@ public class Document {
*/
private String content;
private List<Media> media;
/**
* Embedding of the document. Note: ephemeral field.
*/
@@ -75,17 +79,26 @@ public class Document {
this(content, metadata, new RandomIdGenerator());
}
public Document(String content, List<Media> media, Map<String, Object> metadata) {
this(new RandomIdGenerator().generateId(content, metadata), content, media, metadata);
}
public Document(String content, Map<String, Object> metadata, IdGenerator idGenerator) {
this(idGenerator.generateId(content, metadata), content, metadata);
}
public Document(String id, String content, Map<String, Object> metadata) {
this(id, content, List.of(), metadata);
}
public Document(String id, String content, List<Media> media, Map<String, Object> metadata) {
Assert.hasText(id, "id must not be null");
Assert.hasText(content, "content must not be null");
Assert.notNull(metadata, "metadata must not be null");
this.id = id;
this.content = content;
this.media = media;
this.metadata = metadata;
}
@@ -93,10 +106,16 @@ public class Document {
return id;
}
@Override
public String getContent() {
return this.content;
}
@Override
public List<Media> getMedia() {
return this.media;
}
@JsonIgnore
public String getFormattedContent() {
return this.getFormattedContent(MetadataMode.ALL);
@@ -129,6 +148,7 @@ public class Document {
this.contentFormatter = contentFormatter;
}
@Override
public Map<String, Object> getMetadata() {
return this.metadata;
}

63
spring-ai-core/src/main/java/org/springframework/ai/evaluation/EvaluationRequest.java Normal file
View File

@@ -0,0 +1,63 @@
package org.springframework.ai.evaluation;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.chatbot.ChatBotResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.model.Content;
import java.util.List;
import java.util.Objects;
public class EvaluationRequest {
private final Prompt prompt;
private final List<Content> dataList;
private final ChatResponse chatResponse;
public EvaluationRequest(ChatBotResponse chatBotResponse) {
this(chatBotResponse.getPromptContext().getPromptHistory().get(0),
chatBotResponse.getPromptContext().getContents(), chatBotResponse.getChatResponse());
}
public EvaluationRequest(Prompt prompt, List<Content> dataList, ChatResponse chatResponse) {
this.prompt = prompt;
this.dataList = dataList;
this.chatResponse = chatResponse;
}
public Prompt getPrompt() {
return prompt;
}
public List<Content> getDataList() {
return dataList;
}
public ChatResponse getChatResponse() {
return chatResponse;
}
@Override
public String toString() {
return "EvaluationRequest{" + "prompt=" + prompt + ", dataList=" + dataList + ", chatResponse=" + chatResponse
+ '}';
}
@Override
public boolean equals(Object o) {
if (this == o)
return true;
if (!(o instanceof EvaluationRequest that))
return false;
return Objects.equals(prompt, that.prompt) && Objects.equals(dataList, that.dataList)
&& Objects.equals(chatResponse, that.chatResponse);
}
@Override
public int hashCode() {
return Objects.hash(prompt, dataList, chatResponse);
}
}

60
spring-ai-core/src/main/java/org/springframework/ai/evaluation/EvaluationResponse.java Normal file
View File

@@ -0,0 +1,60 @@
package org.springframework.ai.evaluation;
import java.util.Map;
import java.util.Objects;
public class EvaluationResponse {
private boolean pass;
private float score;
private String feedback;
Map<String, Object> metadata;
public EvaluationResponse(boolean pass, float score, String feedback, Map<String, Object> metadata) {
this.pass = pass;
this.score = score;
this.feedback = feedback;
this.metadata = metadata;
}
public boolean isPass() {
return pass;
}
public float getScore() {
return score;
}
public String getFeedback() {
return feedback;
}
public Map<String, Object> getMetadata() {
return metadata;
}
@Override
public String toString() {
return "EvaluationResponse{" + "pass=" + pass + ", score=" + score + ", feedback='" + feedback + '\''
+ ", metadata=" + metadata + '}';
}
@Override
public boolean equals(Object o) {
if (this == o)
return true;
if (!(o instanceof EvaluationResponse that))
return false;
return pass == that.pass && Float.compare(score, that.score) == 0 && Objects.equals(feedback, that.feedback)
&& Objects.equals(metadata, that.metadata);
}
@Override
public int hashCode() {
return Objects.hash(pass, score, feedback, metadata);
}
}

8
spring-ai-core/src/main/java/org/springframework/ai/evaluation/Evaluator.java Normal file
View File

@@ -0,0 +1,8 @@
package org.springframework.ai.evaluation;
@FunctionalInterface
public interface Evaluator {
EvaluationResponse evaluate(EvaluationRequest evaluationRequest);
}

91
spring-ai-core/src/main/java/org/springframework/ai/evaluation/RelevancyEvaluator.java Normal file
View File

@@ -0,0 +1,91 @@
package org.springframework.ai.evaluation;
import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.MessageType;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.ChatOptionsBuilder;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.model.Content;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
public class RelevancyEvaluator implements Evaluator {
private static final String DEFAULT_EVALUATION_PROMPT_TEXT = """
Your task is to evaluate if the response for the query
is in line with the context information provided.\\n
You have two options to answer. Either YES/ NO.\\n
Answer - YES, if the response for the query
is in line with context information otherwise NO.\\n
Query: \\n {query}\\n
Response: \\n {response}\\n
Context: \\n {context}\\n
Answer: "
""";
private final ChatOptions chatOptions;
private ChatClient chatClient;
public RelevancyEvaluator(ChatClient chatClient) {
this(chatClient, ChatOptionsBuilder.builder().build());
}
public RelevancyEvaluator(ChatClient chatClient, ChatOptions chatOptions) {
this.chatClient = chatClient;
this.chatOptions = chatOptions;
}
@Override
public EvaluationResponse evaluate(EvaluationRequest evaluationRequest) {
var query = doGetUserQuestion(evaluationRequest);
var response = doGetResponse(evaluationRequest);
var context = doGetSupportingData(evaluationRequest);
var promptTemplate = new PromptTemplate(DEFAULT_EVALUATION_PROMPT_TEXT);
Message message = promptTemplate
.createMessage(Map.of("query", query, "response", response, "context", context));
ChatResponse chatResponse = this.chatClient.call(new Prompt(message, this.chatOptions));
var evaluationResponse = chatResponse.getResult().getOutput().getContent();
boolean passing = false;
float score = 0;
if (evaluationResponse.toLowerCase().contains("yes")) {
passing = true;
score = 1;
}
return new EvaluationResponse(passing, score, "", Collections.emptyMap());
}
protected String doGetResponse(EvaluationRequest evaluationRequest) {
return evaluationRequest.getChatResponse().getResult().getOutput().getContent();
}
protected String doGetSupportingData(EvaluationRequest evaluationRequest) {
List<Content> data = evaluationRequest.getDataList();
String supportingData = data.stream()
.filter(node -> node != null && node.getContent() instanceof String)
.map(node -> (Content) node)
.map(Content::getContent)
.collect(Collectors.joining(System.lineSeparator()));
return supportingData;
}
protected String doGetUserQuestion(EvaluationRequest evaluationRequest) {
List<Message> instructions = evaluationRequest.getPrompt().getInstructions();
String userMessage = instructions.stream()
.filter(m -> m.getMessageType() == MessageType.USER)
.map(m -> m.getContent())
.collect(Collectors.joining(System.lineSeparator()));
return userMessage;
}
}

23
spring-ai-core/src/main/java/org/springframework/ai/model/Content.java Normal file
View File

@@ -0,0 +1,23 @@
package org.springframework.ai.model;
import org.springframework.ai.chat.messages.Media;
import java.util.List;
import java.util.Map;
/**
* A simple data structure that contains content and metadata.
*
* @param <T> the type of content in the node
* @author Mark Pollack
* @since 1.0 M1
*/
public interface Content {
String getContent();
List<Media> getMedia();
Map<String, Object> getMetadata();
}

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

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

@@ -0,0 +1,147 @@
/*
* Copyright 2024-2024 the original author or authors.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.chat.history;
import java.util.List;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.ArgumentCaptor;
import org.mockito.Captor;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
import org.springframework.ai.chat.ChatClient;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.Generation;
import org.springframework.ai.chat.StreamingChatClient;
import org.springframework.ai.chat.chatbot.ChatBotResponse;
import org.springframework.ai.chat.chatbot.DefaultChatBot;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
import org.springframework.ai.model.Content;
import org.springframework.ai.tokenizer.JTokkitTokenCountEstimator;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.when;
/**
* @author Christian Tzolov
*/
@ExtendWith(MockitoExtension.class)
public class ChatMemoryTests {
@Mock
ChatClient chatClient;
@Mock
StreamingChatClient streamingChatClient;
@Captor
ArgumentCaptor<Prompt> promptCaptor;
@Test
public void chatMemoryMessageListAugmentor() {
ChatMemory chatHistory = new InMemoryChatMemory();
DefaultChatBot chatAgent = DefaultChatBot.builder(chatClient)
.withRetrievers(List.of(new ChatMemoryRetriever(chatHistory)))
.withContentPostProcessors(
List.of(new LastMaxTokenSizeContentTransformer(new JTokkitTokenCountEstimator(), 10)))
.withAugmentors(List.of(new MessageChatMemoryAugmentor()))
.withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory)))
.build();
chatClientUserMessages(chatAgent, chatHistory);
}
@Test
public void chatMemorySystemPromptAugmentor() {
ChatMemory chatHistory = new InMemoryChatMemory();
DefaultChatBot chatAgent = DefaultChatBot.builder(chatClient)
.withRetrievers(List.of(new ChatMemoryRetriever(chatHistory)))
.withContentPostProcessors(
List.of(new LastMaxTokenSizeContentTransformer(new JTokkitTokenCountEstimator(), 10)))
.withAugmentors(List.of(new SystemPromptChatMemoryAugmentor()))
.withChatAgentListeners(List.of(new ChatMemoryAgentListener(chatHistory)))
.build();
chatClientUserMessages(chatAgent, chatHistory);
}
public void chatClientUserMessages(DefaultChatBot chatAgent, ChatMemory chatHistory) {
when(chatClient.call(promptCaptor.capture()))
.thenReturn(new ChatResponse(List.of(new Generation("assistant:1"))))
.thenReturn(new ChatResponse(List.of(new Generation("assistant:2"))))
.thenReturn(new ChatResponse(List.of(new Generation("assistant:3"))));
var promptContext = PromptContext.builder()
.withConversationId("test-session-id")
.withPrompt(new Prompt(
List.of(new UserMessage("user:1"), new UserMessage("user:2"), new UserMessage("user:3"),
new UserMessage("user:4"), new UserMessage("user:5"))))
.build();
ChatBotResponse response1 = chatAgent.call(promptContext);
assertThat(response1.getChatResponse().getResult().getOutput().getContent()).isEqualTo("assistant:1");
List<Content> contents = response1.getPromptContext().getContents();
assertThat(contents).hasSize(0);
List<Message> history = chatHistory.get("test-session-id", 1000);
assertThat(history).hasSize(6);
ChatBotResponse response2 = chatAgent.call(PromptContext.builder()
.withConversationId("test-session-id")
.withPrompt(new Prompt(
List.of(new UserMessage("user:6"), new UserMessage("user:7"), new UserMessage("user:8"))))
.build());
assertThat(response2.getChatResponse().getResult().getOutput().getContent()).isEqualTo("assistant:2");
history = chatHistory.get("test-session-id", 1000);
assertThat(history).hasSize(10);
contents = response2.getPromptContext().getContents();
assertThat(contents).hasSize(3);
assertThat(contents.get(0).getContent()).isEqualTo("user:4");
assertThat(contents.get(1).getContent()).isEqualTo("user:5");
assertThat(contents.get(2).getContent()).isEqualTo("assistant:1");
ChatBotResponse response3 = chatAgent.call(PromptContext.builder()
.withConversationId("test-session-id")
.withPrompt(new Prompt(List.of(new UserMessage("user:9")))).build());
assertThat(response3.getChatResponse().getResult().getOutput().getContent()).isEqualTo("assistant:3");
history = chatHistory.get("test-session-id", 1000);
assertThat(history).hasSize(12);
contents = response3.getPromptContext().getContents();
assertThat(contents).hasSize(3);
assertThat(contents.get(0).getContent()).isEqualTo("user:7");
assertThat(contents.get(1).getContent()).isEqualTo("user:8");
assertThat(contents.get(2).getContent()).isEqualTo("assistant:2");
}
}

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/chatclient.adoc
View File

@@ -78,16 +78,29 @@ The `Message` interface encapsulates a textual message, a collection of attribut
[source,java]
----
public interface Message {
public interface Message extends Node<String> {
String getContent();
Map<String, Object> getProperties();
List<Media> getMedia();
MessageType getMessageType();
}
----
and the Node interface is
```java
public interface Node<T> {
T getContent();
Map<String, Object> getMetadata();
}
```
The `Message` interface has various implementations that correspond to the categories of messages that an AI model can process.
Some models, like OpenAI's chat completion endpoint, distinguish between message categories based on conversational roles, effectively mapped by the `MessageType`.

16
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/prompt.adoc
View File

@@ -49,19 +49,29 @@ The `Message` interface encapsulates a textual message, a collection of attribut
The interface is defined as follows:
```java
public interface Message {
public interface Message extends Node<String> {
String getContent();
List<Media> getMedia();
Map<String, Object> getProperties();
MessageType getMessageType();
}
```
and the Node interface is
```java
public interface Node<T> {
T getContent();
Map<String, Object> getMetadata();
}
```
Various implementations of the `Message` interface correspond to different categories of messages that an AI model can process. Some models, like those from OpenAI, distinguish between message categories based on conversational roles. These roles are effectively mapped by the `MessageType`, as discussed below.

108
spring-ai-test/src/main/java/org/springframework/ai/evaluation/BaseMemoryTest.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.evaluation;
import java.util.stream.Collectors;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.chatbot.ChatBot;
import org.springframework.ai.chat.chatbot.StreamingChatBot;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.transformer.PromptContext;
import static org.assertj.core.api.Assertions.assertThat;
/**
* @author Christian Tzolov
*/
public class BaseMemoryTest {
protected final Logger logger = LoggerFactory.getLogger(getClass());
protected RelevancyEvaluator relevancyEvaluator;
protected ChatBot chatBot;
protected StreamingChatBot streamingChatBot;
public BaseMemoryTest(RelevancyEvaluator relevancyEvaluator, ChatBot chatBot,
StreamingChatBot streamingChatClient) {
this.relevancyEvaluator = relevancyEvaluator;
this.chatBot = chatBot;
this.streamingChatBot = streamingChatClient;
}
@Test
void memoryChatAgent() {
var prompt = new Prompt(new UserMessage("my name John Vincent Atanasoff"));
PromptContext promptContext = new PromptContext(prompt);
var chatBotResponse1 = this.chatBot.call(promptContext);
logger.info("Response1: " + chatBotResponse1.getChatResponse().getResult().getOutput().getContent());
assertThat(chatBotResponse1.getChatResponse().getResult().getOutput().getContent()).contains("John");
var chatBotResponse2 = this.chatBot.call(new PromptContext(new Prompt(new String("What is my name?"))));
logger.info("Response2: " + chatBotResponse2.getChatResponse().getResult().getOutput().getContent());
assertThat(chatBotResponse2.getChatResponse().getResult().getOutput().getContent())
.contains("John Vincent Atanasoff");
EvaluationResponse evaluationResponse = this.relevancyEvaluator
.evaluate(new EvaluationRequest(chatBotResponse2));
logger.info("" + evaluationResponse);
}
@Test
void memoryStreamingChatBot() {
var prompt = new Prompt(new UserMessage("my name John Vincent Atanasoff"));
PromptContext promptContext = new PromptContext(prompt);
var fluxChatBotResponse1 = this.streamingChatBot.stream(promptContext);
String chatBotResponse1 = fluxChatBotResponse1.getChatResponse()
.collectList()
.block()
.stream()
.filter(response -> response.getResult() != null)
.map(response -> response.getResult().getOutput().getContent())
.collect(Collectors.joining());
logger.info("Response1: " + chatBotResponse1);
assertThat(chatBotResponse1).contains("John");
var fluxChatBotResponse2 = this.streamingChatBot
.stream(new PromptContext(new Prompt(new String("What is my name?"))));
String chatBotResponse2 = fluxChatBotResponse2.getChatResponse()
.collectList()
.block()
.stream()
.filter(response -> response.getResult() != null)
.map(response -> response.getResult().getOutput().getContent())
.collect(Collectors.joining());
logger.info("Response2: " + chatBotResponse2);
assertThat(chatBotResponse2).contains("John Vincent Atanasoff");
}
}

2
vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/Neo4jVectorStore.java
View File

@@ -176,7 +176,7 @@ public class Neo4jVectorStore implements VectorStore, InitializingBean {
*/
public Builder withLabel(String newLabel) {
Assert.hasText(newLabel, "Node label may not be null or blank");
Assert.hasText(newLabel, "Content label may not be null or blank");
this.label = newLabel;
return this;

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