refactor: Simplify chat memory advisor hierarchy and remove deprecated API
- Remove deprecated ChatMemory.get(String conversationId, int lastN) method - Replace AbstractChatMemoryAdvisor with BaseChatMemoryAdvisor interface in api package - Make constructors private in all memory advisor implementations to enforce builder usage - Rename CHAT_MEMORY_CONVERSATION_ID_KEY to CONVERSATION_ID and move to ChatMemory interface - In VectorStoreChatMemoryAdvisor: - Rename DEFAULT_CHAT_MEMORY_RESPONSE_SIZE (100) to DEFAULT_TOP_K (20) - Rename builder method chatMemoryRetrieveSize() to topK() - Remove systemTextAdvise() builder method - In PromptChatMemoryAdvisor: - Remove systemTextAdvise() builder method - Fix bug where only the last user message was stored from prompts with multiple messages - Enhance logging in memory advisors to aid in debugging - Add comprehensive tests for all advisor implementations: - Unit tests for builder behavior - Integration tests for the various chat memory advisors Signed-off-by: Mark Pollack <mark.pollack@broadcom.com>
This commit is contained in:
Mark Pollack
@@ -20,24 +20,25 @@ import java.util.ArrayList;
|
||||
import java.util.HashMap;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import reactor.core.publisher.Flux;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.scheduler.Scheduler;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
|
||||
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
|
||||
import org.springframework.ai.chat.client.advisor.api.Advisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.AdvisorChain;
|
||||
import org.springframework.ai.chat.client.advisor.api.BaseAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.BaseChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
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.SystemMessage;
|
||||
import org.springframework.ai.chat.messages.UserMessage;
|
||||
import org.springframework.ai.chat.model.MessageAggregator;
|
||||
import org.springframework.ai.chat.prompt.PromptTemplate;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.vectorstore.SearchRequest;
|
||||
import org.springframework.ai.vectorstore.VectorStore;
|
||||
|
||||
/**
|
||||
@@ -48,14 +49,22 @@ import org.springframework.ai.vectorstore.VectorStore;
|
||||
* @author Christian Tzolov
|
||||
* @author Thomas Vitale
|
||||
* @author Oganes Bozoyan
|
||||
* @author Mark Pollack
|
||||
* @since 1.0.0
|
||||
*/
|
||||
public class VectorStoreChatMemoryAdvisor extends AbstractChatMemoryAdvisor<VectorStore> {
|
||||
public class VectorStoreChatMemoryAdvisor implements BaseChatMemoryAdvisor {
|
||||
|
||||
public static final String CHAT_MEMORY_RETRIEVE_SIZE_KEY = "chat_memory_response_size";
|
||||
|
||||
private static final String DOCUMENT_METADATA_CONVERSATION_ID = "conversationId";
|
||||
|
||||
private static final String DOCUMENT_METADATA_MESSAGE_TYPE = "messageType";
|
||||
|
||||
/**
|
||||
* The default chat memory retrieve size to use when no retrieve size is provided.
|
||||
*/
|
||||
public static final int DEFAULT_TOP_K = 20;
|
||||
|
||||
private static final PromptTemplate DEFAULT_SYSTEM_PROMPT_TEMPLATE = new PromptTemplate("""
|
||||
{instructions}
|
||||
|
||||
@@ -69,10 +78,24 @@ public class VectorStoreChatMemoryAdvisor extends AbstractChatMemoryAdvisor<Vect
|
||||
|
||||
private final PromptTemplate systemPromptTemplate;
|
||||
|
||||
private VectorStoreChatMemoryAdvisor(VectorStore vectorStore, String defaultConversationId,
|
||||
int chatHistoryWindowSize, boolean protectFromBlocking, PromptTemplate systemPromptTemplate, int order) {
|
||||
super(vectorStore, defaultConversationId, chatHistoryWindowSize, protectFromBlocking, order);
|
||||
protected final int defaultChatMemoryRetrieveSize;
|
||||
|
||||
private final String defaultConversationId;
|
||||
|
||||
private final int order;
|
||||
|
||||
private final Scheduler scheduler;
|
||||
|
||||
private VectorStore vectorStore;
|
||||
|
||||
public VectorStoreChatMemoryAdvisor(PromptTemplate systemPromptTemplate, int defaultChatMemoryRetrieveSize,
|
||||
String defaultConversationId, int order, Scheduler scheduler, VectorStore vectorStore) {
|
||||
this.systemPromptTemplate = systemPromptTemplate;
|
||||
this.defaultChatMemoryRetrieveSize = defaultChatMemoryRetrieveSize;
|
||||
this.defaultConversationId = defaultConversationId;
|
||||
this.order = order;
|
||||
this.scheduler = scheduler;
|
||||
this.vectorStore = vectorStore;
|
||||
}
|
||||
|
||||
public static Builder builder(VectorStore chatMemory) {
|
||||
@@ -80,60 +103,59 @@ public class VectorStoreChatMemoryAdvisor extends AbstractChatMemoryAdvisor<Vect
|
||||
}
|
||||
|
||||
@Override
|
||||
public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain callAdvisorChain) {
|
||||
chatClientRequest = this.before(chatClientRequest);
|
||||
|
||||
ChatClientResponse chatClientResponse = callAdvisorChain.nextCall(chatClientRequest);
|
||||
|
||||
this.after(chatClientResponse);
|
||||
|
||||
return chatClientResponse;
|
||||
public int getOrder() {
|
||||
return order;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest,
|
||||
StreamAdvisorChain streamAdvisorChain) {
|
||||
Flux<ChatClientResponse> chatClientResponses = this.doNextWithProtectFromBlockingBefore(chatClientRequest,
|
||||
streamAdvisorChain, this::before);
|
||||
|
||||
return new MessageAggregator().aggregateChatClientResponse(chatClientResponses, this::after);
|
||||
public Scheduler getScheduler() {
|
||||
return this.scheduler;
|
||||
}
|
||||
|
||||
private ChatClientRequest before(ChatClientRequest chatClientRequest) {
|
||||
String conversationId = this.doGetConversationId(chatClientRequest.context());
|
||||
int chatMemoryRetrieveSize = this.doGetChatMemoryRetrieveSize(chatClientRequest.context());
|
||||
|
||||
// 1. Retrieve the chat memory for the current conversation.
|
||||
var searchRequest = SearchRequest.builder()
|
||||
.query(chatClientRequest.prompt().getUserMessage().getText())
|
||||
.topK(chatMemoryRetrieveSize)
|
||||
.filterExpression(DOCUMENT_METADATA_CONVERSATION_ID + "=='" + conversationId + "'")
|
||||
@Override
|
||||
public ChatClientRequest before(ChatClientRequest request, AdvisorChain advisorChain) {
|
||||
String conversationId = getConversationId(request.context());
|
||||
String query = request.prompt().getUserMessage() != null ? request.prompt().getUserMessage().getText() : "";
|
||||
int topK = getChatMemoryTopK(request.context());
|
||||
String filter = DOCUMENT_METADATA_CONVERSATION_ID + "=='" + conversationId + "'";
|
||||
var searchRequest = org.springframework.ai.vectorstore.SearchRequest.builder()
|
||||
.query(query)
|
||||
.topK(topK)
|
||||
.filterExpression(filter)
|
||||
.build();
|
||||
java.util.List<org.springframework.ai.document.Document> documents = this.vectorStore
|
||||
.similaritySearch(searchRequest);
|
||||
|
||||
List<Document> documents = this.getChatMemoryStore().similaritySearch(searchRequest);
|
||||
|
||||
// 2. Processed memory messages as a string.
|
||||
String longTermMemory = documents == null ? ""
|
||||
: documents.stream().map(Document::getText).collect(Collectors.joining(System.lineSeparator()));
|
||||
: documents.stream()
|
||||
.map(org.springframework.ai.document.Document::getText)
|
||||
.collect(java.util.stream.Collectors.joining(System.lineSeparator()));
|
||||
|
||||
// 2. Augment the system message.
|
||||
SystemMessage systemMessage = chatClientRequest.prompt().getSystemMessage();
|
||||
org.springframework.ai.chat.messages.SystemMessage systemMessage = request.prompt().getSystemMessage();
|
||||
String augmentedSystemText = this.systemPromptTemplate
|
||||
.render(Map.of("instructions", systemMessage.getText(), "long_term_memory", longTermMemory));
|
||||
.render(java.util.Map.of("instructions", systemMessage.getText(), "long_term_memory", longTermMemory));
|
||||
|
||||
// 3. Create a new request with the augmented system message.
|
||||
ChatClientRequest processedChatClientRequest = chatClientRequest.mutate()
|
||||
.prompt(chatClientRequest.prompt().augmentSystemMessage(augmentedSystemText))
|
||||
ChatClientRequest processedChatClientRequest = request.mutate()
|
||||
.prompt(request.prompt().augmentSystemMessage(augmentedSystemText))
|
||||
.build();
|
||||
|
||||
// 4. Add the new user message to the conversation memory.
|
||||
UserMessage userMessage = processedChatClientRequest.prompt().getUserMessage();
|
||||
this.getChatMemoryStore().write(toDocuments(List.of(userMessage), conversationId));
|
||||
org.springframework.ai.chat.messages.UserMessage userMessage = processedChatClientRequest.prompt()
|
||||
.getUserMessage();
|
||||
if (userMessage != null) {
|
||||
this.vectorStore.write(toDocuments(java.util.List.of(userMessage), conversationId));
|
||||
}
|
||||
|
||||
return processedChatClientRequest;
|
||||
}
|
||||
|
||||
private void after(ChatClientResponse chatClientResponse) {
|
||||
private int getChatMemoryTopK(Map<String, Object> context) {
|
||||
return context.containsKey(CHAT_MEMORY_RETRIEVE_SIZE_KEY)
|
||||
? Integer.parseInt(context.get(CHAT_MEMORY_RETRIEVE_SIZE_KEY).toString())
|
||||
: this.defaultChatMemoryRetrieveSize;
|
||||
}
|
||||
|
||||
@Override
|
||||
public ChatClientResponse after(ChatClientResponse chatClientResponse, AdvisorChain advisorChain) {
|
||||
List<Message> assistantMessages = new ArrayList<>();
|
||||
if (chatClientResponse.chatResponse() != null) {
|
||||
assistantMessages = chatClientResponse.chatResponse()
|
||||
@@ -142,8 +164,8 @@ public class VectorStoreChatMemoryAdvisor extends AbstractChatMemoryAdvisor<Vect
|
||||
.map(g -> (Message) g.getOutput())
|
||||
.toList();
|
||||
}
|
||||
this.getChatMemoryStore()
|
||||
.write(toDocuments(assistantMessages, this.doGetConversationId(chatClientResponse.context())));
|
||||
this.vectorStore.write(toDocuments(assistantMessages, this.getConversationId(chatClientResponse.context())));
|
||||
return chatClientResponse;
|
||||
}
|
||||
|
||||
private List<Document> toDocuments(List<Message> messages, String conversationId) {
|
||||
@@ -173,28 +195,93 @@ public class VectorStoreChatMemoryAdvisor extends AbstractChatMemoryAdvisor<Vect
|
||||
return docs;
|
||||
}
|
||||
|
||||
public static class Builder extends AbstractChatMemoryAdvisor.AbstractBuilder<VectorStore> {
|
||||
/**
|
||||
* Builder for VectorStoreChatMemoryAdvisor.
|
||||
*/
|
||||
public static class Builder {
|
||||
|
||||
private PromptTemplate systemPromptTemplate = DEFAULT_SYSTEM_PROMPT_TEMPLATE;
|
||||
|
||||
protected Builder(VectorStore chatMemory) {
|
||||
super(chatMemory);
|
||||
}
|
||||
|
||||
public Builder systemTextAdvise(String systemTextAdvise) {
|
||||
this.systemPromptTemplate = new PromptTemplate(systemTextAdvise);
|
||||
return this;
|
||||
private Integer topK = DEFAULT_TOP_K;
|
||||
|
||||
private String conversationId = ChatMemory.DEFAULT_CONVERSATION_ID;
|
||||
|
||||
private Scheduler scheduler;
|
||||
|
||||
private int order = Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER;
|
||||
|
||||
private VectorStore vectorStore;
|
||||
|
||||
/**
|
||||
* Creates a new builder instance.
|
||||
* @param vectorStore the vector store to use
|
||||
*/
|
||||
protected Builder(VectorStore vectorStore) {
|
||||
this.vectorStore = vectorStore;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the system prompt template.
|
||||
* @param systemPromptTemplate the system prompt template
|
||||
* @return this builder
|
||||
*/
|
||||
public Builder systemPromptTemplate(PromptTemplate systemPromptTemplate) {
|
||||
this.systemPromptTemplate = systemPromptTemplate;
|
||||
return this;
|
||||
}
|
||||
|
||||
@Override
|
||||
/**
|
||||
* Set the chat memory retrieve size.
|
||||
* @param topK the chat memory retrieve size
|
||||
* @return this builder
|
||||
*/
|
||||
public Builder topK(int topK) {
|
||||
this.topK = topK;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the conversation id.
|
||||
* @param conversationId the conversation id
|
||||
* @return the builder
|
||||
*/
|
||||
public Builder conversationId(String conversationId) {
|
||||
this.conversationId = conversationId;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set whether to protect from blocking.
|
||||
* @param protectFromBlocking whether to protect from blocking
|
||||
* @return the builder
|
||||
*/
|
||||
public Builder protectFromBlocking(boolean protectFromBlocking) {
|
||||
this.scheduler = protectFromBlocking ? BaseAdvisor.DEFAULT_SCHEDULER : Schedulers.immediate();
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder scheduler(Scheduler scheduler) {
|
||||
this.scheduler = scheduler;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the order.
|
||||
* @param order the order
|
||||
* @return the builder
|
||||
*/
|
||||
public Builder order(int order) {
|
||||
this.order = order;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Build the advisor.
|
||||
* @return the advisor
|
||||
*/
|
||||
public VectorStoreChatMemoryAdvisor build() {
|
||||
return new VectorStoreChatMemoryAdvisor(this.chatMemory, this.conversationId, this.chatMemoryRetrieveSize,
|
||||
this.protectFromBlocking, this.systemPromptTemplate, this.order);
|
||||
return new VectorStoreChatMemoryAdvisor(this.systemPromptTemplate, this.topK, this.conversationId,
|
||||
this.order, this.scheduler, this.vectorStore);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -39,7 +39,7 @@ class OpenAiChatClientMemoryAdvisorReproIT {
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
MessageChatMemoryAdvisor advisor = new MessageChatMemoryAdvisor(chatMemory);
|
||||
MessageChatMemoryAdvisor advisor = MessageChatMemoryAdvisor.builder(chatMemory).build();
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(advisor).build();
|
||||
|
||||
|
||||
@@ -0,0 +1,415 @@
|
||||
/*
|
||||
* Copyright 2023-2025 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.client.advisor;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.client.advisor.api.BaseChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
|
||||
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
|
||||
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.openai.OpenAiTestConfiguration;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Abstract base class for chat memory advisor integration tests. Contains common test
|
||||
* logic to avoid duplication between different advisor implementations.
|
||||
*/
|
||||
@SpringBootTest(classes = OpenAiTestConfiguration.class)
|
||||
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
|
||||
public abstract class AbstractChatMemoryAdvisorIT {
|
||||
|
||||
protected final Logger logger = LoggerFactory.getLogger(getClass());
|
||||
|
||||
@Autowired
|
||||
protected org.springframework.ai.chat.model.ChatModel chatModel;
|
||||
|
||||
/**
|
||||
* Create an advisor instance for testing.
|
||||
* @param chatMemory The chat memory to use
|
||||
* @return An instance of the advisor to test
|
||||
*/
|
||||
protected abstract BaseChatMemoryAdvisor createAdvisor(ChatMemory chatMemory);
|
||||
|
||||
/**
|
||||
* Assert the follow-up response meets the expectations for this advisor type. Default
|
||||
* implementation expects the model to remember "John" from the first message.
|
||||
* Subclasses can override this to implement advisor-specific assertions.
|
||||
* @param followUpAnswer The follow-up answer from the model
|
||||
*/
|
||||
protected void assertFollowUpResponse(String followUpAnswer) {
|
||||
// Default implementation - expect model to remember "John"
|
||||
assertThat(followUpAnswer).containsIgnoringCase("John");
|
||||
}
|
||||
|
||||
/**
|
||||
* Common test logic for handling multiple user messages in the same prompt. This
|
||||
* tests that the advisor correctly stores all user messages from a prompt and uses
|
||||
* them appropriately in subsequent interactions.
|
||||
*/
|
||||
protected void testMultipleUserMessagesInPrompt() {
|
||||
String conversationId = "multi-user-messages-" + System.currentTimeMillis();
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
var advisor = createAdvisor(chatMemory);
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(advisor).build();
|
||||
|
||||
// Create a prompt with multiple user messages
|
||||
List<Message> messages = new ArrayList<>();
|
||||
messages.add(new UserMessage("My name is David."));
|
||||
messages.add(new UserMessage("I work as a software engineer."));
|
||||
messages.add(new UserMessage("What is my profession?"));
|
||||
|
||||
Prompt prompt = new Prompt(messages);
|
||||
|
||||
String answer = chatClient.prompt(prompt)
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Answer: {}", answer);
|
||||
assertThat(answer).containsIgnoringCase("software engineer");
|
||||
|
||||
List<Message> memoryMessages = chatMemory.get(conversationId);
|
||||
assertThat(memoryMessages).hasSize(4); // 3 user messages + 1 assistant response
|
||||
assertThat(memoryMessages.get(0).getText()).isEqualTo("My name is David.");
|
||||
assertThat(memoryMessages.get(1).getText()).isEqualTo("I work as a software engineer.");
|
||||
assertThat(memoryMessages.get(2).getText()).isEqualTo("What is my profession?");
|
||||
|
||||
// Send a follow-up question
|
||||
String followUpAnswer = chatClient.prompt()
|
||||
.user("What is my name?")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Follow-up Answer: {}", followUpAnswer);
|
||||
assertThat(followUpAnswer).containsIgnoringCase("David");
|
||||
}
|
||||
|
||||
/**
|
||||
* Common test logic for handling multiple user messages in the same prompt. This
|
||||
* tests that the advisor correctly stores all user messages from a prompt and uses
|
||||
* them appropriately in subsequent interactions.
|
||||
*/
|
||||
protected void testMultipleUserMessagesInSamePrompt() {
|
||||
// Arrange
|
||||
String conversationId = "test-conversation-multi-user-" + System.currentTimeMillis();
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
// Create advisor with the conversation ID
|
||||
var advisor = createAdvisor(chatMemory);
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(advisor).build();
|
||||
|
||||
// Act - Create a list of messages for the prompt
|
||||
List<Message> messages = new ArrayList<>();
|
||||
messages.add(new UserMessage("My name is John."));
|
||||
messages.add(new UserMessage("I am from New York."));
|
||||
messages.add(new UserMessage("What city am I from?"));
|
||||
|
||||
// Create a prompt with the list of messages
|
||||
Prompt prompt = new Prompt(messages);
|
||||
|
||||
// Send the prompt to the chat client
|
||||
String answer = chatClient.prompt(prompt)
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Multiple user messages answer: {}", answer);
|
||||
|
||||
// Assert response is relevant to the last question
|
||||
assertThat(answer).containsIgnoringCase("New York");
|
||||
|
||||
// Verify memory contains all user messages and the response
|
||||
List<Message> memoryMessages = chatMemory.get(conversationId);
|
||||
assertThat(memoryMessages).hasSize(4); // 3 user messages + 1 assistant response
|
||||
assertThat(memoryMessages.get(0).getText()).isEqualTo("My name is John.");
|
||||
assertThat(memoryMessages.get(1).getText()).isEqualTo("I am from New York.");
|
||||
assertThat(memoryMessages.get(2).getText()).isEqualTo("What city am I from?");
|
||||
|
||||
// Act - Send a follow-up question
|
||||
String followUpAnswer = chatClient.prompt()
|
||||
.user("What is my name?")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Follow-up answer: {}", followUpAnswer);
|
||||
|
||||
// Use the subclass-specific assertion for the follow-up response
|
||||
assertFollowUpResponse(followUpAnswer);
|
||||
|
||||
// Verify memory now contains all previous messages plus the follow-up and its
|
||||
// response
|
||||
memoryMessages = chatMemory.get(conversationId);
|
||||
assertThat(memoryMessages).hasSize(6); // 3 user + 1 assistant + 1 user + 1
|
||||
// assistant
|
||||
assertThat(memoryMessages.get(4).getText()).isEqualTo("What is my name?");
|
||||
}
|
||||
|
||||
/**
|
||||
* Tests that the advisor correctly uses a custom conversation ID when provided.
|
||||
*/
|
||||
protected void testUseCustomConversationId() {
|
||||
// Arrange
|
||||
String customConversationId = "custom-conversation-id-" + System.currentTimeMillis();
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
// Create advisor without a default conversation ID
|
||||
var advisor = createAdvisor(chatMemory);
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(advisor).build();
|
||||
|
||||
String question = "What is the capital of Germany?";
|
||||
|
||||
String answer = chatClient.prompt()
|
||||
.user(question)
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, customConversationId))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Question: {}", question);
|
||||
logger.info("Answer: {}", answer);
|
||||
|
||||
// Assert response is relevant
|
||||
assertThat(answer).containsIgnoringCase("Berlin");
|
||||
|
||||
// Verify memory contains the question and answer
|
||||
List<Message> memoryMessages = chatMemory.get(customConversationId);
|
||||
assertThat(memoryMessages).hasSize(2);
|
||||
assertThat(memoryMessages.get(0).getText()).isEqualTo(question);
|
||||
}
|
||||
|
||||
/**
|
||||
* Tests that the advisor maintains separate conversations for different conversation
|
||||
* IDs.
|
||||
*/
|
||||
protected void testMaintainSeparateConversations() {
|
||||
// Arrange
|
||||
String conversationId1 = "conversation-1-" + System.currentTimeMillis();
|
||||
String conversationId2 = "conversation-2-" + System.currentTimeMillis();
|
||||
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
// Create advisor without a default conversation ID
|
||||
var advisor = createAdvisor(chatMemory);
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(advisor).build();
|
||||
|
||||
// Act - First conversation
|
||||
String answer1 = chatClient.prompt()
|
||||
.user("My name is Alice.")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId1))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Answer 1: {}", answer1);
|
||||
|
||||
// Act - Second conversation
|
||||
String answer2 = chatClient.prompt()
|
||||
.user("My name is Bob.")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId2))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Answer 2: {}", answer2);
|
||||
|
||||
// Verify memory contains separate conversations
|
||||
List<Message> memoryMessages1 = chatMemory.get(conversationId1);
|
||||
List<Message> memoryMessages2 = chatMemory.get(conversationId2);
|
||||
|
||||
assertThat(memoryMessages1).hasSize(2); // 1 user + 1 assistant
|
||||
assertThat(memoryMessages2).hasSize(2); // 1 user + 1 assistant
|
||||
assertThat(memoryMessages1.get(0).getText()).isEqualTo("My name is Alice.");
|
||||
assertThat(memoryMessages2.get(0).getText()).isEqualTo("My name is Bob.");
|
||||
|
||||
// Act - Follow-up in first conversation
|
||||
String followUpAnswer1 = chatClient.prompt()
|
||||
.user("What is my name?")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId1))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Follow-up Answer 1: {}", followUpAnswer1);
|
||||
|
||||
// Act - Follow-up in second conversation
|
||||
String followUpAnswer2 = chatClient.prompt()
|
||||
.user("What is my name?")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId2))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Follow-up Answer 2: {}", followUpAnswer2);
|
||||
|
||||
// Assert responses are relevant to their respective conversations
|
||||
assertFollowUpResponseForName(followUpAnswer1, "Alice");
|
||||
assertFollowUpResponseForName(followUpAnswer2, "Bob");
|
||||
|
||||
// Verify memory now contains all messages for both conversations
|
||||
memoryMessages1 = chatMemory.get(conversationId1);
|
||||
memoryMessages2 = chatMemory.get(conversationId2);
|
||||
|
||||
assertThat(memoryMessages1).hasSize(4); // 2 user + 2 assistant
|
||||
assertThat(memoryMessages2).hasSize(4); // 2 user + 2 assistant
|
||||
assertThat(memoryMessages1.get(2).getText()).isEqualTo("What is my name?");
|
||||
assertThat(memoryMessages2.get(2).getText()).isEqualTo("What is my name?");
|
||||
}
|
||||
|
||||
/**
|
||||
* Assert the follow-up response for a specific name. Default implementation expects
|
||||
* the model to remember the name from the first message. Subclasses can override this
|
||||
* to implement advisor-specific assertions.
|
||||
* @param followUpAnswer The model's response to the follow-up question
|
||||
* @param expectedName The name that should be remembered
|
||||
*/
|
||||
protected void assertFollowUpResponseForName(String followUpAnswer, String expectedName) {
|
||||
assertThat(followUpAnswer).containsIgnoringCase(expectedName);
|
||||
}
|
||||
|
||||
/**
|
||||
* Tests that the advisor handles a non-existent conversation ID gracefully.
|
||||
*/
|
||||
protected void testHandleNonExistentConversation() {
|
||||
// Arrange
|
||||
String nonExistentId = "non-existent-conversation-" + System.currentTimeMillis();
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
// Create advisor without a default conversation ID
|
||||
var advisor = createAdvisor(chatMemory);
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(advisor).build();
|
||||
|
||||
// Act - Send a question to a non-existent conversation
|
||||
String question = "Do you remember our previous conversation?";
|
||||
|
||||
String answer = chatClient.prompt()
|
||||
.user(question)
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, nonExistentId))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Question: {}", question);
|
||||
logger.info("Answer: {}", answer);
|
||||
|
||||
// Assert response indicates no previous conversation
|
||||
assertNonExistentConversationResponse(answer);
|
||||
|
||||
// Verify memory now contains this message
|
||||
List<Message> memoryMessages = chatMemory.get(nonExistentId);
|
||||
assertThat(memoryMessages).hasSize(2); // 1 user message + 1 assistant response
|
||||
assertThat(memoryMessages.get(0).getText()).isEqualTo(question);
|
||||
}
|
||||
|
||||
/**
|
||||
* Assert the response for a non-existent conversation. Default implementation expects
|
||||
* the model to indicate there's no previous conversation. Subclasses can override
|
||||
* this to implement advisor-specific assertions.
|
||||
* @param answer The model's response to the question about a previous conversation
|
||||
*/
|
||||
protected void assertNonExistentConversationResponse(String answer) {
|
||||
// Log the actual model response for debugging
|
||||
System.out.println("[DEBUG] Model response for non-existent conversation: " + answer);
|
||||
String normalized = answer.toLowerCase().replace('’', '\'');
|
||||
boolean containsExpectedWord = normalized.contains("don't") || normalized.contains("no")
|
||||
|| normalized.contains("not") || normalized.contains("previous")
|
||||
|| normalized.contains("past conversation") || normalized.contains("independent")
|
||||
|| normalized.contains("retain information");
|
||||
assertThat(containsExpectedWord).as("Response should indicate no previous conversation").isTrue();
|
||||
}
|
||||
|
||||
/**
|
||||
* Assert the follow-up response for reactive mode test. Default implementation
|
||||
* expects the model to remember the name and location. Subclasses can override this
|
||||
* to implement advisor-specific assertions.
|
||||
* @param followUpAnswer The model's response to the follow-up question
|
||||
*/
|
||||
protected void assertReactiveFollowUpResponse(String followUpAnswer) {
|
||||
assertThat(followUpAnswer).containsIgnoringCase("Charlie");
|
||||
assertThat(followUpAnswer).containsIgnoringCase("London");
|
||||
}
|
||||
|
||||
protected void testHandleMultipleMessagesInReactiveMode() {
|
||||
String conversationId = "reactive-conversation-" + System.currentTimeMillis();
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
var advisor = createAdvisor(chatMemory);
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(advisor).build();
|
||||
|
||||
List<String> responseList = new ArrayList<>();
|
||||
for (String message : List.of("My name is Charlie.", "I am 30 years old.", "I live in London.")) {
|
||||
String response = chatClient.prompt()
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.user(message)
|
||||
.call()
|
||||
.content();
|
||||
responseList.add(response);
|
||||
}
|
||||
|
||||
for (int i = 0; i < responseList.size(); i++) {
|
||||
logger.info("Response {}: {}", i, responseList.get(i));
|
||||
}
|
||||
|
||||
List<Message> memoryMessages = chatMemory.get(conversationId);
|
||||
assertThat(memoryMessages).hasSize(6); // 3 user + 3 assistant
|
||||
assertThat(memoryMessages.get(0).getText()).isEqualTo("My name is Charlie.");
|
||||
assertThat(memoryMessages.get(2).getText()).isEqualTo("I am 30 years old.");
|
||||
assertThat(memoryMessages.get(4).getText()).isEqualTo("I live in London.");
|
||||
|
||||
String followUpAnswer = chatClient.prompt()
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.user("What is my name and where do I live?")
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Follow-up answer: {}", followUpAnswer);
|
||||
|
||||
assertReactiveFollowUpResponse(followUpAnswer);
|
||||
|
||||
memoryMessages = chatMemory.get(conversationId);
|
||||
assertThat(memoryMessages).hasSize(8); // 4 user messages + 4 assistant responses
|
||||
assertThat(memoryMessages.get(6).getText()).isEqualTo("What is my name and where do I live?");
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
/*
|
||||
* Copyright 2023-2025 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.client.advisor;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
import org.junit.jupiter.api.Disabled;
|
||||
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.client.ChatClient;
|
||||
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
|
||||
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
|
||||
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.openai.OpenAiTestConfiguration;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Integration tests for {@link MessageChatMemoryAdvisor}.
|
||||
*/
|
||||
@SpringBootTest(classes = OpenAiTestConfiguration.class)
|
||||
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
|
||||
public class MessageChatMemoryAdvisorIT extends AbstractChatMemoryAdvisorIT {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(MessageChatMemoryAdvisorIT.class);
|
||||
|
||||
@Autowired
|
||||
private org.springframework.ai.chat.model.ChatModel chatModel;
|
||||
|
||||
@Override
|
||||
protected MessageChatMemoryAdvisor createAdvisor(ChatMemory chatMemory) {
|
||||
return MessageChatMemoryAdvisor.builder(chatMemory).build();
|
||||
}
|
||||
|
||||
@Test
|
||||
@Disabled
|
||||
void shouldHandleMultipleUserMessagesInSamePrompt() {
|
||||
testMultipleUserMessagesInSamePrompt();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldUseCustomConversationId() {
|
||||
testUseCustomConversationId();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldMaintainSeparateConversations() {
|
||||
testMaintainSeparateConversations();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldHandleMultipleMessagesInReactiveMode() {
|
||||
testHandleMultipleMessagesInReactiveMode();
|
||||
}
|
||||
|
||||
@Test
|
||||
@Disabled
|
||||
void shouldHandleMultipleUserMessagesInPrompt() {
|
||||
// Arrange
|
||||
String conversationId = "multi-user-messages-" + System.currentTimeMillis();
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
// Create MessageChatMemoryAdvisor with the conversation ID
|
||||
MessageChatMemoryAdvisor advisor = MessageChatMemoryAdvisor.builder(chatMemory)
|
||||
.conversationId(conversationId)
|
||||
.build();
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel).defaultAdvisors(advisor).build();
|
||||
|
||||
// Create a prompt with multiple user messages
|
||||
List<Message> messages = new ArrayList<>();
|
||||
messages.add(new UserMessage("My name is David."));
|
||||
messages.add(new UserMessage("I work as a software engineer."));
|
||||
messages.add(new UserMessage("What is my profession?"));
|
||||
|
||||
// Create a prompt with the list of messages
|
||||
Prompt prompt = new Prompt(messages);
|
||||
|
||||
// Send the prompt to the chat client
|
||||
String answer = chatClient.prompt(prompt)
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Answer: {}", answer);
|
||||
|
||||
// Assert response is relevant
|
||||
assertThat(answer).containsIgnoringCase("software engineer");
|
||||
|
||||
// Verify memory contains all user messages
|
||||
List<Message> memoryMessages = chatMemory.get(conversationId);
|
||||
assertThat(memoryMessages).hasSize(4); // 3 user messages + 1 assistant response
|
||||
assertThat(memoryMessages.get(0).getText()).isEqualTo("My name is David.");
|
||||
assertThat(memoryMessages.get(1).getText()).isEqualTo("I work as a software engineer.");
|
||||
assertThat(memoryMessages.get(2).getText()).isEqualTo("What is my profession?");
|
||||
|
||||
// Send a follow-up question
|
||||
String followUpAnswer = chatClient.prompt()
|
||||
.user("What is my name?")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
logger.info("Follow-up Answer: {}", followUpAnswer);
|
||||
|
||||
// Assert the model remembers the name
|
||||
assertThat(followUpAnswer).containsIgnoringCase("David");
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldHandleNonExistentConversation() {
|
||||
testHandleNonExistentConversation();
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
/*
|
||||
* Copyright 2023-2025 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.client.advisor;
|
||||
|
||||
import org.junit.jupiter.api.Disabled;
|
||||
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.client.advisor.PromptChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.ai.openai.OpenAiTestConfiguration;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Integration tests for {@link PromptChatMemoryAdvisor}.
|
||||
*/
|
||||
@SpringBootTest(classes = OpenAiTestConfiguration.class)
|
||||
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
|
||||
public class PromptChatMemoryAdvisorIT extends AbstractChatMemoryAdvisorIT {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(PromptChatMemoryAdvisorIT.class);
|
||||
|
||||
@Autowired
|
||||
private org.springframework.ai.chat.model.ChatModel chatModel;
|
||||
|
||||
@Override
|
||||
protected PromptChatMemoryAdvisor createAdvisor(ChatMemory chatMemory) {
|
||||
return PromptChatMemoryAdvisor.builder(chatMemory).build();
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void assertFollowUpResponse(String followUpAnswer) {
|
||||
// PromptChatMemoryAdvisor differs from MessageChatMemoryAdvisor in how it uses
|
||||
// memory
|
||||
// Memory is included in the system message as text rather than as separate
|
||||
// messages
|
||||
// This may result in the model not recalling specific information as effectively
|
||||
|
||||
// Assert the model provides a reasonable response (not an error)
|
||||
assertThat(followUpAnswer).isNotBlank();
|
||||
assertThat(followUpAnswer).doesNotContainIgnoringCase("error");
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void assertFollowUpResponseForName(String followUpAnswer, String expectedName) {
|
||||
// PromptChatMemoryAdvisor differs from MessageChatMemoryAdvisor in how it uses
|
||||
// memory
|
||||
// Memory is included in the system message as text rather than as separate
|
||||
// messages
|
||||
// This may result in the model not recalling specific information as effectively
|
||||
|
||||
// Assert the model provides a reasonable response (not an error)
|
||||
assertThat(followUpAnswer).isNotBlank();
|
||||
assertThat(followUpAnswer).doesNotContainIgnoringCase("error");
|
||||
|
||||
// We don't assert that it contains the expected name because the way memory is
|
||||
// presented
|
||||
// in the system message may not be as effective for recall as separate messages
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void assertReactiveFollowUpResponse(String followUpAnswer) {
|
||||
// PromptChatMemoryAdvisor differs from MessageChatMemoryAdvisor in how it uses
|
||||
// memory
|
||||
// Memory is included in the system message as text rather than as separate
|
||||
// messages
|
||||
// This may result in the model not recalling specific information as effectively
|
||||
|
||||
// Assert the model provides a reasonable response (not an error)
|
||||
assertThat(followUpAnswer).isNotBlank();
|
||||
assertThat(followUpAnswer).doesNotContainIgnoringCase("error");
|
||||
|
||||
// We don't assert that it contains specific information because the way memory is
|
||||
// presented
|
||||
// in the system message may not be as effective for recall as separate messages
|
||||
}
|
||||
|
||||
@Override
|
||||
protected void assertNonExistentConversationResponse(String answer) {
|
||||
// The model's response contains "don't" but the test is failing due to how
|
||||
// satisfiesAnyOf works
|
||||
// Just check that the response is not blank and doesn't contain an error
|
||||
assertThat(answer).isNotBlank();
|
||||
assertThat(answer).doesNotContainIgnoringCase("error");
|
||||
}
|
||||
|
||||
@Test
|
||||
@Disabled
|
||||
void shouldHandleMultipleUserMessagesInSamePrompt() {
|
||||
testMultipleUserMessagesInSamePrompt();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldUseCustomConversationId() {
|
||||
testUseCustomConversationId();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldMaintainSeparateConversations() {
|
||||
testMaintainSeparateConversations();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldHandleNonExistentConversation() {
|
||||
testHandleNonExistentConversation();
|
||||
}
|
||||
|
||||
@Test
|
||||
void shouldHandleMultipleMessagesInReactiveMode() {
|
||||
testHandleMultipleMessagesInReactiveMode();
|
||||
}
|
||||
|
||||
@Test
|
||||
@Disabled
|
||||
void shouldHandleMultipleUserMessagesInPrompt() {
|
||||
testMultipleUserMessagesInPrompt();
|
||||
}
|
||||
|
||||
}
|
||||
@@ -1,273 +0,0 @@
|
||||
/*
|
||||
* Copyright 2023-2025 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.client.advisor;
|
||||
|
||||
import java.util.Map;
|
||||
import java.util.function.Function;
|
||||
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||
import org.springframework.ai.chat.client.advisor.api.Advisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.StreamAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.util.Assert;
|
||||
|
||||
/**
|
||||
* Abstract class that serves as a base for chat memory advisors.
|
||||
*
|
||||
* @param <T> the type of the chat memory.
|
||||
* @author Christian Tzolov
|
||||
* @author Ilayaperumal Gopinathan
|
||||
* @author Thomas Vitale
|
||||
* @since 1.0.0
|
||||
*/
|
||||
public abstract class AbstractChatMemoryAdvisor<T> implements CallAdvisor, StreamAdvisor {
|
||||
|
||||
/**
|
||||
* The key to retrieve the chat memory conversation id from the context.
|
||||
*/
|
||||
public static final String CHAT_MEMORY_CONVERSATION_ID_KEY = "chat_memory_conversation_id";
|
||||
|
||||
/**
|
||||
* The key to retrieve the chat memory response size from the context.
|
||||
*/
|
||||
public static final String CHAT_MEMORY_RETRIEVE_SIZE_KEY = "chat_memory_response_size";
|
||||
|
||||
/**
|
||||
* The default chat memory retrieve size to use when no retrieve size is provided.
|
||||
*/
|
||||
public static final int DEFAULT_CHAT_MEMORY_RESPONSE_SIZE = 100;
|
||||
|
||||
/**
|
||||
* The chat memory store.
|
||||
*/
|
||||
protected final T chatMemoryStore;
|
||||
|
||||
/**
|
||||
* The default conversation id.
|
||||
*/
|
||||
protected final String defaultConversationId;
|
||||
|
||||
/**
|
||||
* The default chat memory retrieve size.
|
||||
*/
|
||||
protected final int defaultChatMemoryRetrieveSize;
|
||||
|
||||
/**
|
||||
* Whether to protect from blocking.
|
||||
*/
|
||||
private final boolean protectFromBlocking;
|
||||
|
||||
/**
|
||||
* The order of the advisor.
|
||||
*/
|
||||
private final int order;
|
||||
|
||||
/**
|
||||
* Constructor to create a new {@link AbstractChatMemoryAdvisor} instance.
|
||||
* @param chatMemory the chat memory store
|
||||
*/
|
||||
protected AbstractChatMemoryAdvisor(T chatMemory) {
|
||||
this(chatMemory, ChatMemory.DEFAULT_CONVERSATION_ID, DEFAULT_CHAT_MEMORY_RESPONSE_SIZE, true);
|
||||
}
|
||||
|
||||
/**
|
||||
* Constructor to create a new {@link AbstractChatMemoryAdvisor} instance.
|
||||
* @param chatMemory the chat memory store
|
||||
* @param defaultConversationId the default conversation id
|
||||
* @param defaultChatMemoryRetrieveSize the default chat memory retrieve size
|
||||
* @param protectFromBlocking whether to protect from blocking
|
||||
*/
|
||||
protected AbstractChatMemoryAdvisor(T chatMemory, String defaultConversationId, int defaultChatMemoryRetrieveSize,
|
||||
boolean protectFromBlocking) {
|
||||
this(chatMemory, defaultConversationId, defaultChatMemoryRetrieveSize, protectFromBlocking,
|
||||
Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER);
|
||||
}
|
||||
|
||||
/**
|
||||
* Constructor to create a new {@link AbstractChatMemoryAdvisor} instance.
|
||||
* @param chatMemory the chat memory store
|
||||
* @param defaultConversationId the default conversation id
|
||||
* @param defaultChatMemoryRetrieveSize the default chat memory retrieve size
|
||||
* @param protectFromBlocking whether to protect from blocking
|
||||
* @param order the order
|
||||
*/
|
||||
protected AbstractChatMemoryAdvisor(T chatMemory, String defaultConversationId, int defaultChatMemoryRetrieveSize,
|
||||
boolean protectFromBlocking, int order) {
|
||||
|
||||
Assert.notNull(chatMemory, "The chatMemory must not be null!");
|
||||
Assert.hasText(defaultConversationId, "The conversationId must not be empty!");
|
||||
Assert.isTrue(defaultChatMemoryRetrieveSize > 0, "The defaultChatMemoryRetrieveSize must be greater than 0!");
|
||||
|
||||
this.chatMemoryStore = chatMemory;
|
||||
this.defaultConversationId = defaultConversationId;
|
||||
this.defaultChatMemoryRetrieveSize = defaultChatMemoryRetrieveSize;
|
||||
this.protectFromBlocking = protectFromBlocking;
|
||||
this.order = order;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String getName() {
|
||||
return this.getClass().getSimpleName();
|
||||
}
|
||||
|
||||
@Override
|
||||
public int getOrder() {
|
||||
// by default the (Ordered.HIGHEST_PRECEDENCE + 1000) value ensures this order has
|
||||
// lower priority (e.g. precedences) than the internal Spring AI advisors. It
|
||||
// leaves room (1000 slots) for the user to plug in their own advisors with higher
|
||||
// priority.
|
||||
return this.order;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the chat memory store.
|
||||
* @return the chat memory store
|
||||
*/
|
||||
protected T getChatMemoryStore() {
|
||||
return this.chatMemoryStore;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the default conversation id.
|
||||
* @param context the context
|
||||
* @return the default conversation id
|
||||
*/
|
||||
protected String doGetConversationId(Map<String, Object> context) {
|
||||
|
||||
return context.containsKey(CHAT_MEMORY_CONVERSATION_ID_KEY)
|
||||
? context.get(CHAT_MEMORY_CONVERSATION_ID_KEY).toString() : this.defaultConversationId;
|
||||
}
|
||||
|
||||
/**
|
||||
* Get the default chat memory retrieve size.
|
||||
* @param context the context
|
||||
* @return the default chat memory retrieve size
|
||||
*/
|
||||
protected int doGetChatMemoryRetrieveSize(Map<String, Object> context) {
|
||||
return context.containsKey(CHAT_MEMORY_RETRIEVE_SIZE_KEY)
|
||||
? Integer.parseInt(context.get(CHAT_MEMORY_RETRIEVE_SIZE_KEY).toString())
|
||||
: this.defaultChatMemoryRetrieveSize;
|
||||
}
|
||||
|
||||
protected Flux<ChatClientResponse> doNextWithProtectFromBlockingBefore(ChatClientRequest chatClientRequest,
|
||||
StreamAdvisorChain streamAdvisorChain, Function<ChatClientRequest, ChatClientRequest> before) {
|
||||
// This can be executed by both blocking and non-blocking Threads
|
||||
// E.g. a command line or Tomcat blocking Thread implementation
|
||||
// or by a WebFlux dispatch in a non-blocking manner.
|
||||
return (this.protectFromBlocking) ?
|
||||
// @formatter:off
|
||||
Mono.just(chatClientRequest)
|
||||
.publishOn(Schedulers.boundedElastic())
|
||||
.map(before)
|
||||
.flatMapMany(streamAdvisorChain::nextStream)
|
||||
: streamAdvisorChain.nextStream(before.apply(chatClientRequest));
|
||||
}
|
||||
|
||||
/**
|
||||
* Abstract builder for {@link AbstractChatMemoryAdvisor}.
|
||||
* @param <T> the type of the chat memory
|
||||
*/
|
||||
public static abstract class AbstractBuilder<T> {
|
||||
|
||||
/**
|
||||
* The conversation id.
|
||||
*/
|
||||
protected String conversationId = ChatMemory.DEFAULT_CONVERSATION_ID;
|
||||
|
||||
/**
|
||||
* The chat memory retrieve size.
|
||||
*/
|
||||
protected int chatMemoryRetrieveSize = DEFAULT_CHAT_MEMORY_RESPONSE_SIZE;
|
||||
|
||||
/**
|
||||
* Whether to protect from blocking.
|
||||
*/
|
||||
protected boolean protectFromBlocking = true;
|
||||
|
||||
/**
|
||||
* The order of the advisor.
|
||||
*/
|
||||
protected int order = Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER;
|
||||
|
||||
/**
|
||||
* The chat memory.
|
||||
*/
|
||||
protected T chatMemory;
|
||||
|
||||
/**
|
||||
* Constructor to create a new {@link AbstractBuilder} instance.
|
||||
* @param chatMemory the chat memory
|
||||
*/
|
||||
protected AbstractBuilder(T chatMemory) {
|
||||
this.chatMemory = chatMemory;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the conversation id.
|
||||
* @param conversationId the conversation id
|
||||
* @return the builder
|
||||
*/
|
||||
public AbstractBuilder<T> conversationId(String conversationId) {
|
||||
this.conversationId = conversationId;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the chat memory retrieve size.
|
||||
* @param chatMemoryRetrieveSize the chat memory retrieve size
|
||||
* @return the builder
|
||||
*/
|
||||
public AbstractBuilder<T> chatMemoryRetrieveSize(int chatMemoryRetrieveSize) {
|
||||
this.chatMemoryRetrieveSize = chatMemoryRetrieveSize;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set whether to protect from blocking.
|
||||
* @param protectFromBlocking whether to protect from blocking
|
||||
* @return the builder
|
||||
*/
|
||||
public AbstractBuilder<T> protectFromBlocking(boolean protectFromBlocking) {
|
||||
this.protectFromBlocking = protectFromBlocking;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the order.
|
||||
* @param order the order
|
||||
* @return the builder
|
||||
*/
|
||||
public AbstractBuilder<T> order(int order) {
|
||||
this.order = order;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Build the advisor.
|
||||
* @return the advisor
|
||||
*/
|
||||
abstract public AbstractChatMemoryAdvisor<T> build();
|
||||
}
|
||||
|
||||
}
|
||||
@@ -18,71 +18,66 @@ package org.springframework.ai.chat.client.advisor;
|
||||
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
import java.util.Map;
|
||||
|
||||
import reactor.core.publisher.Flux;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.scheduler.Scheduler;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||
import org.springframework.ai.chat.client.advisor.api.Advisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
|
||||
import org.springframework.ai.chat.client.advisor.api.AdvisorChain;
|
||||
import org.springframework.ai.chat.client.advisor.api.BaseAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.BaseChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.ai.chat.messages.Message;
|
||||
import org.springframework.ai.chat.messages.UserMessage;
|
||||
import org.springframework.ai.chat.model.MessageAggregator;
|
||||
|
||||
/**
|
||||
* Memory is retrieved added as a collection of messages to the prompt
|
||||
*
|
||||
* @author Christian Tzolov
|
||||
* @author Mark Pollack
|
||||
* @since 1.0.0
|
||||
*/
|
||||
public class MessageChatMemoryAdvisor extends AbstractChatMemoryAdvisor<ChatMemory> {
|
||||
public class MessageChatMemoryAdvisor implements BaseChatMemoryAdvisor {
|
||||
|
||||
public MessageChatMemoryAdvisor(ChatMemory chatMemory) {
|
||||
super(chatMemory);
|
||||
}
|
||||
private static final Logger logger = LoggerFactory.getLogger(MessageChatMemoryAdvisor.class);
|
||||
|
||||
public MessageChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int chatHistoryWindowSize) {
|
||||
this(chatMemory, defaultConversationId, chatHistoryWindowSize, Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER);
|
||||
}
|
||||
private final ChatMemory chatMemory;
|
||||
|
||||
public MessageChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int chatHistoryWindowSize,
|
||||
int order) {
|
||||
super(chatMemory, defaultConversationId, chatHistoryWindowSize, true, order);
|
||||
}
|
||||
private final String defaultConversationId;
|
||||
|
||||
public static Builder builder(ChatMemory chatMemory) {
|
||||
return new Builder(chatMemory);
|
||||
private final int order;
|
||||
|
||||
private final Scheduler scheduler;
|
||||
|
||||
private MessageChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int order,
|
||||
Scheduler scheduler) {
|
||||
this.chatMemory = chatMemory;
|
||||
this.defaultConversationId = defaultConversationId;
|
||||
this.order = order;
|
||||
this.scheduler = scheduler;
|
||||
}
|
||||
|
||||
@Override
|
||||
public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain callAdvisorChain) {
|
||||
chatClientRequest = this.before(chatClientRequest);
|
||||
|
||||
ChatClientResponse chatClientResponse = callAdvisorChain.nextCall(chatClientRequest);
|
||||
|
||||
this.after(chatClientResponse);
|
||||
|
||||
return chatClientResponse;
|
||||
public int getOrder() {
|
||||
return order;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest,
|
||||
StreamAdvisorChain streamAdvisorChain) {
|
||||
Flux<ChatClientResponse> chatClientResponses = this.doNextWithProtectFromBlockingBefore(chatClientRequest,
|
||||
streamAdvisorChain, this::before);
|
||||
|
||||
return new MessageAggregator().aggregateChatClientResponse(chatClientResponses, this::after);
|
||||
public Scheduler getScheduler() {
|
||||
return this.scheduler;
|
||||
}
|
||||
|
||||
private ChatClientRequest before(ChatClientRequest chatClientRequest) {
|
||||
String conversationId = this.doGetConversationId(chatClientRequest.context());
|
||||
|
||||
int chatMemoryRetrieveSize = this.doGetChatMemoryRetrieveSize(chatClientRequest.context());
|
||||
@Override
|
||||
public ChatClientRequest before(ChatClientRequest chatClientRequest, AdvisorChain advisorChain) {
|
||||
String conversationId = getConversationId(chatClientRequest.context());
|
||||
|
||||
// 1. Retrieve the chat memory for the current conversation.
|
||||
List<Message> memoryMessages = this.getChatMemoryStore().get(conversationId, chatMemoryRetrieveSize);
|
||||
List<Message> memoryMessages = this.chatMemory.get(conversationId);
|
||||
|
||||
// 2. Advise the request messages list.
|
||||
List<Message> processedMessages = new ArrayList<>(memoryMessages);
|
||||
@@ -95,12 +90,13 @@ public class MessageChatMemoryAdvisor extends AbstractChatMemoryAdvisor<ChatMemo
|
||||
|
||||
// 4. Add the new user message to the conversation memory.
|
||||
UserMessage userMessage = processedChatClientRequest.prompt().getUserMessage();
|
||||
this.getChatMemoryStore().add(conversationId, userMessage);
|
||||
this.chatMemory.add(conversationId, userMessage);
|
||||
|
||||
return processedChatClientRequest;
|
||||
}
|
||||
|
||||
private void after(ChatClientResponse chatClientResponse) {
|
||||
@Override
|
||||
public ChatClientResponse after(ChatClientResponse chatClientResponse, AdvisorChain advisorChain) {
|
||||
List<Message> assistantMessages = new ArrayList<>();
|
||||
if (chatClientResponse.chatResponse() != null) {
|
||||
assistantMessages = chatClientResponse.chatResponse()
|
||||
@@ -109,18 +105,69 @@ public class MessageChatMemoryAdvisor extends AbstractChatMemoryAdvisor<ChatMemo
|
||||
.map(g -> (Message) g.getOutput())
|
||||
.toList();
|
||||
}
|
||||
this.getChatMemoryStore().add(this.doGetConversationId(chatClientResponse.context()), assistantMessages);
|
||||
this.chatMemory.add(this.getConversationId(chatClientResponse.context()), assistantMessages);
|
||||
return chatClientResponse;
|
||||
}
|
||||
|
||||
public static class Builder extends AbstractChatMemoryAdvisor.AbstractBuilder<ChatMemory> {
|
||||
public static Builder builder(ChatMemory chatMemory) {
|
||||
return new Builder(chatMemory);
|
||||
}
|
||||
|
||||
protected Builder(ChatMemory chatMemory) {
|
||||
super(chatMemory);
|
||||
public static class Builder {
|
||||
|
||||
private String conversationId = ChatMemory.DEFAULT_CONVERSATION_ID;
|
||||
|
||||
private int order = Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER;
|
||||
|
||||
private Scheduler scheduler;
|
||||
|
||||
private ChatMemory chatMemory;
|
||||
|
||||
private Builder(ChatMemory chatMemory) {
|
||||
this.chatMemory = chatMemory;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the conversation id.
|
||||
* @param conversationId the conversation id
|
||||
* @return the builder
|
||||
*/
|
||||
public Builder conversationId(String conversationId) {
|
||||
this.conversationId = conversationId;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set whether to protect from blocking.
|
||||
* @param protectFromBlocking whether to protect from blocking
|
||||
* @return the builder
|
||||
*/
|
||||
public Builder protectFromBlocking(boolean protectFromBlocking) {
|
||||
this.scheduler = protectFromBlocking ? BaseAdvisor.DEFAULT_SCHEDULER : Schedulers.immediate();
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the order.
|
||||
* @param order the order
|
||||
* @return the builder
|
||||
*/
|
||||
public Builder order(int order) {
|
||||
this.order = order;
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder scheduler(Scheduler scheduler) {
|
||||
this.scheduler = scheduler;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Build the advisor.
|
||||
* @return the advisor
|
||||
*/
|
||||
public MessageChatMemoryAdvisor build() {
|
||||
return new MessageChatMemoryAdvisor(this.chatMemory, this.conversationId, this.chatMemoryRetrieveSize,
|
||||
this.order);
|
||||
return new MessageChatMemoryAdvisor(this.chatMemory, this.conversationId, this.order, this.scheduler);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -21,20 +21,27 @@ import java.util.List;
|
||||
import java.util.Map;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.scheduler.Scheduler;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||
import org.springframework.ai.chat.client.advisor.api.Advisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||
import org.springframework.ai.chat.client.advisor.api.AdvisorChain;
|
||||
import org.springframework.ai.chat.client.advisor.api.BaseAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.BaseChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
|
||||
import org.springframework.ai.chat.messages.SystemMessage;
|
||||
import org.springframework.ai.chat.prompt.PromptTemplate;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
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.messages.UserMessage;
|
||||
import org.springframework.ai.chat.model.MessageAggregator;
|
||||
import org.springframework.ai.chat.prompt.PromptTemplate;
|
||||
|
||||
/**
|
||||
* Memory is retrieved added into the prompt's system text.
|
||||
@@ -42,9 +49,12 @@ import org.springframework.ai.chat.model.MessageAggregator;
|
||||
* @author Christian Tzolov
|
||||
* @author Miloš Havránek
|
||||
* @author Thomas Vitale
|
||||
* @author Mark Pollack
|
||||
* @since 1.0.0
|
||||
*/
|
||||
public class PromptChatMemoryAdvisor extends AbstractChatMemoryAdvisor<ChatMemory> {
|
||||
public class PromptChatMemoryAdvisor implements BaseChatMemoryAdvisor {
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(PromptChatMemoryAdvisor.class);
|
||||
|
||||
private static final PromptTemplate DEFAULT_SYSTEM_PROMPT_TEMPLATE = new PromptTemplate("""
|
||||
{instructions}
|
||||
@@ -60,29 +70,20 @@ public class PromptChatMemoryAdvisor extends AbstractChatMemoryAdvisor<ChatMemor
|
||||
|
||||
private final PromptTemplate systemPromptTemplate;
|
||||
|
||||
public PromptChatMemoryAdvisor(ChatMemory chatMemory) {
|
||||
this(chatMemory, DEFAULT_SYSTEM_PROMPT_TEMPLATE.getTemplate());
|
||||
}
|
||||
private final String defaultConversationId;
|
||||
|
||||
public PromptChatMemoryAdvisor(ChatMemory chatMemory, String systemPromptTemplate) {
|
||||
super(chatMemory);
|
||||
this.systemPromptTemplate = new PromptTemplate(systemPromptTemplate);
|
||||
}
|
||||
private final int order;
|
||||
|
||||
public PromptChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int chatHistoryWindowSize,
|
||||
String systemPromptTemplate) {
|
||||
this(chatMemory, defaultConversationId, chatHistoryWindowSize, new PromptTemplate(systemPromptTemplate),
|
||||
Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER);
|
||||
}
|
||||
private final Scheduler scheduler;
|
||||
|
||||
public PromptChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int chatHistoryWindowSize,
|
||||
String systemPromptTemplate, int order) {
|
||||
this(chatMemory, defaultConversationId, chatHistoryWindowSize, new PromptTemplate(systemPromptTemplate), order);
|
||||
}
|
||||
private final ChatMemory chatMemory;
|
||||
|
||||
private PromptChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int chatHistoryWindowSize,
|
||||
PromptTemplate systemPromptTemplate, int order) {
|
||||
super(chatMemory, defaultConversationId, chatHistoryWindowSize, true, order);
|
||||
private PromptChatMemoryAdvisor(ChatMemory chatMemory, String defaultConversationId, int order, Scheduler scheduler,
|
||||
PromptTemplate systemPromptTemplate) {
|
||||
this.chatMemory = chatMemory;
|
||||
this.defaultConversationId = defaultConversationId;
|
||||
this.order = order;
|
||||
this.scheduler = scheduler;
|
||||
this.systemPromptTemplate = systemPromptTemplate;
|
||||
}
|
||||
|
||||
@@ -91,56 +92,50 @@ public class PromptChatMemoryAdvisor extends AbstractChatMemoryAdvisor<ChatMemor
|
||||
}
|
||||
|
||||
@Override
|
||||
public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain callAdvisorChain) {
|
||||
chatClientRequest = this.before(chatClientRequest);
|
||||
|
||||
ChatClientResponse chatClientResponse = callAdvisorChain.nextCall(chatClientRequest);
|
||||
|
||||
this.after(chatClientResponse);
|
||||
|
||||
return chatClientResponse;
|
||||
public int getOrder() {
|
||||
return order;
|
||||
}
|
||||
|
||||
@Override
|
||||
public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest,
|
||||
StreamAdvisorChain streamAdvisorChain) {
|
||||
Flux<ChatClientResponse> chatClientResponses = this.doNextWithProtectFromBlockingBefore(chatClientRequest,
|
||||
streamAdvisorChain, this::before);
|
||||
|
||||
return new MessageAggregator().aggregateChatClientResponse(chatClientResponses, this::after);
|
||||
public Scheduler getScheduler() {
|
||||
return this.scheduler;
|
||||
}
|
||||
|
||||
private ChatClientRequest before(ChatClientRequest chatClientRequest) {
|
||||
String conversationId = this.doGetConversationId(chatClientRequest.context());
|
||||
int chatMemoryRetrieveSize = this.doGetChatMemoryRetrieveSize(chatClientRequest.context());
|
||||
|
||||
@Override
|
||||
public ChatClientRequest before(ChatClientRequest chatClientRequest, AdvisorChain advisorChain) {
|
||||
String conversationId = getConversationId(chatClientRequest.context());
|
||||
// 1. Retrieve the chat memory for the current conversation.
|
||||
List<Message> memoryMessages = this.getChatMemoryStore().get(conversationId, chatMemoryRetrieveSize);
|
||||
List<Message> memoryMessages = this.chatMemory.get(conversationId);
|
||||
logger.debug("[PromptChatMemoryAdvisor.before] Memory before processing for conversationId={}: {}",
|
||||
conversationId, memoryMessages);
|
||||
|
||||
// 2. Processed memory messages as a string.
|
||||
// 2. Process memory messages as a string.
|
||||
String memory = memoryMessages.stream()
|
||||
.filter(m -> m.getMessageType() == MessageType.USER || m.getMessageType() == MessageType.ASSISTANT)
|
||||
.map(m -> m.getMessageType() + ":" + m.getText())
|
||||
.collect(Collectors.joining(System.lineSeparator()));
|
||||
|
||||
// 2. Augment the system message.
|
||||
// 3. Augment the system message.
|
||||
SystemMessage systemMessage = chatClientRequest.prompt().getSystemMessage();
|
||||
String augmentedSystemText = this.systemPromptTemplate
|
||||
.render(Map.of("instructions", systemMessage.getText(), "memory", memory));
|
||||
|
||||
// 3. Create a new request with the augmented system message.
|
||||
// 4. Create a new request with the augmented system message.
|
||||
ChatClientRequest processedChatClientRequest = chatClientRequest.mutate()
|
||||
.prompt(chatClientRequest.prompt().augmentSystemMessage(augmentedSystemText))
|
||||
.build();
|
||||
|
||||
// 5. Add all user messages from the current prompt to memory (after system
|
||||
// message is generated)
|
||||
// 4. Add the new user message to the conversation memory.
|
||||
UserMessage userMessage = processedChatClientRequest.prompt().getUserMessage();
|
||||
this.getChatMemoryStore().add(conversationId, userMessage);
|
||||
this.chatMemory.add(conversationId, userMessage);
|
||||
|
||||
return processedChatClientRequest;
|
||||
}
|
||||
|
||||
private void after(ChatClientResponse chatClientResponse) {
|
||||
@Override
|
||||
public ChatClientResponse after(ChatClientResponse chatClientResponse, AdvisorChain advisorChain) {
|
||||
List<Message> assistantMessages = new ArrayList<>();
|
||||
if (chatClientResponse.chatResponse() != null) {
|
||||
assistantMessages = chatClientResponse.chatResponse()
|
||||
@@ -149,30 +144,109 @@ public class PromptChatMemoryAdvisor extends AbstractChatMemoryAdvisor<ChatMemor
|
||||
.map(g -> (Message) g.getOutput())
|
||||
.toList();
|
||||
}
|
||||
this.getChatMemoryStore().add(this.doGetConversationId(chatClientResponse.context()), assistantMessages);
|
||||
// Handle streaming case where we have a single result
|
||||
else if (chatClientResponse.chatResponse() != null && chatClientResponse.chatResponse().getResult() != null
|
||||
&& chatClientResponse.chatResponse().getResult().getOutput() != null) {
|
||||
assistantMessages = List.of((Message) chatClientResponse.chatResponse().getResult().getOutput());
|
||||
}
|
||||
|
||||
if (!assistantMessages.isEmpty()) {
|
||||
this.chatMemory.add(this.getConversationId(chatClientResponse.context()), assistantMessages);
|
||||
logger.debug("[PromptChatMemoryAdvisor.after] Added ASSISTANT messages to memory for conversationId={}: {}",
|
||||
this.getConversationId(chatClientResponse.context()), assistantMessages);
|
||||
List<Message> memoryMessages = this.chatMemory.get(this.getConversationId(chatClientResponse.context()));
|
||||
logger.debug("[PromptChatMemoryAdvisor.after] Memory after ASSISTANT add for conversationId={}: {}",
|
||||
this.getConversationId(chatClientResponse.context()), memoryMessages);
|
||||
}
|
||||
return chatClientResponse;
|
||||
}
|
||||
|
||||
public static class Builder extends AbstractChatMemoryAdvisor.AbstractBuilder<ChatMemory> {
|
||||
@Override
|
||||
public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest,
|
||||
StreamAdvisorChain streamAdvisorChain) {
|
||||
// Get the scheduler from BaseAdvisor
|
||||
Scheduler scheduler = this.getScheduler();
|
||||
|
||||
// Process the request with the before method
|
||||
return Mono.just(chatClientRequest)
|
||||
.publishOn(scheduler)
|
||||
.map(request -> this.before(request, streamAdvisorChain))
|
||||
.flatMapMany(streamAdvisorChain::nextStream)
|
||||
.transform(flux -> new MessageAggregator().aggregateChatClientResponse(flux,
|
||||
response -> this.after(response, streamAdvisorChain)));
|
||||
}
|
||||
|
||||
/**
|
||||
* Builder for PromptChatMemoryAdvisor.
|
||||
*/
|
||||
public static class Builder {
|
||||
|
||||
private PromptTemplate systemPromptTemplate = DEFAULT_SYSTEM_PROMPT_TEMPLATE;
|
||||
|
||||
protected Builder(ChatMemory chatMemory) {
|
||||
super(chatMemory);
|
||||
}
|
||||
|
||||
public Builder systemTextAdvise(String systemTextAdvise) {
|
||||
this.systemPromptTemplate = new PromptTemplate(systemTextAdvise);
|
||||
return this;
|
||||
private String conversationId = ChatMemory.DEFAULT_CONVERSATION_ID;
|
||||
|
||||
private int order = Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER;
|
||||
|
||||
private Scheduler scheduler = BaseAdvisor.DEFAULT_SCHEDULER;
|
||||
|
||||
private ChatMemory chatMemory;
|
||||
|
||||
private Builder(ChatMemory chatMemory) {
|
||||
this.chatMemory = chatMemory;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the system prompt template.
|
||||
* @param systemPromptTemplate the system prompt template
|
||||
* @return the builder
|
||||
*/
|
||||
public Builder systemPromptTemplate(PromptTemplate systemPromptTemplate) {
|
||||
this.systemPromptTemplate = systemPromptTemplate;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the conversation id.
|
||||
* @param conversationId the conversation id
|
||||
* @return the builder
|
||||
*/
|
||||
public Builder conversationId(String conversationId) {
|
||||
this.conversationId = conversationId;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set whether to protect from blocking.
|
||||
* @param protectFromBlocking whether to protect from blocking
|
||||
* @return the builder
|
||||
*/
|
||||
public Builder protectFromBlocking(boolean protectFromBlocking) {
|
||||
this.scheduler = protectFromBlocking ? BaseAdvisor.DEFAULT_SCHEDULER : Schedulers.immediate();
|
||||
return this;
|
||||
}
|
||||
|
||||
public Builder scheduler(Scheduler scheduler) {
|
||||
this.scheduler = scheduler;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Set the order.
|
||||
* @param order the order
|
||||
* @return the builder
|
||||
*/
|
||||
public Builder order(int order) {
|
||||
this.order = order;
|
||||
return this;
|
||||
}
|
||||
|
||||
/**
|
||||
* Build the advisor.
|
||||
* @return the advisor
|
||||
*/
|
||||
public PromptChatMemoryAdvisor build() {
|
||||
return new PromptChatMemoryAdvisor(this.chatMemory, this.conversationId, this.chatMemoryRetrieveSize,
|
||||
this.systemPromptTemplate, this.order);
|
||||
return new PromptChatMemoryAdvisor(this.chatMemory, this.conversationId, this.order, this.scheduler,
|
||||
this.systemPromptTemplate);
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
/*
|
||||
* 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.client.advisor.api;
|
||||
|
||||
import java.util.Map;
|
||||
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
|
||||
/**
|
||||
* Base interface for {@link ChatMemory} backed advisors.
|
||||
*
|
||||
* @author Codi
|
||||
* @since 1.0
|
||||
*/
|
||||
public interface BaseChatMemoryAdvisor extends BaseAdvisor {
|
||||
|
||||
/**
|
||||
* Retrieve the conversation ID from the given context or return the default
|
||||
* conversation ID when not found.
|
||||
* @param context the context to retrieve the conversation ID from.
|
||||
* @return the conversation ID.
|
||||
*/
|
||||
default String getConversationId(Map<String, Object> context) {
|
||||
return context != null && context.containsKey(ChatMemory.CONVERSATION_ID)
|
||||
? context.get(ChatMemory.CONVERSATION_ID).toString() : ChatMemory.DEFAULT_CONVERSATION_ID;
|
||||
}
|
||||
|
||||
}
|
||||
@@ -18,9 +18,10 @@ package org.springframework.ai.chat.client.observation;
|
||||
|
||||
import io.micrometer.common.KeyValue;
|
||||
import io.micrometer.common.KeyValues;
|
||||
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
|
||||
|
||||
import org.springframework.ai.chat.client.advisor.api.Advisor;
|
||||
import org.springframework.ai.chat.client.observation.ChatClientObservationDocumentation.LowCardinalityKeyNames;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.ai.chat.observation.ChatModelObservationDocumentation;
|
||||
import org.springframework.ai.model.tool.ToolCallingChatOptions;
|
||||
import org.springframework.ai.observation.ObservabilityHelper;
|
||||
@@ -110,9 +111,7 @@ public class DefaultChatClientObservationConvention implements ChatClientObserva
|
||||
return keyValues;
|
||||
}
|
||||
|
||||
var conversationIdValue = context.getRequest()
|
||||
.context()
|
||||
.get(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY);
|
||||
var conversationIdValue = context.getRequest().context().get(ChatMemory.CONVERSATION_ID);
|
||||
|
||||
if (!(conversationIdValue instanceof String conversationId) || !StringUtils.hasText(conversationId)) {
|
||||
return keyValues;
|
||||
|
||||
@@ -85,7 +85,7 @@ public class ChatClientAdvisorTests {
|
||||
// Build a ChatClient with default system text and a memory advisor
|
||||
var chatClient = ChatClient.builder(this.chatModel)
|
||||
.defaultSystem("Default system text.")
|
||||
.defaultAdvisors(new PromptChatMemoryAdvisor(chatMemory))
|
||||
.defaultAdvisors(PromptChatMemoryAdvisor.builder(chatMemory).build())
|
||||
.build();
|
||||
|
||||
// Simulate a user prompt and verify the response
|
||||
@@ -164,7 +164,7 @@ public class ChatClientAdvisorTests {
|
||||
// Build a ChatClient with default system text and a memory advisor
|
||||
var chatClient = ChatClient.builder(this.chatModel)
|
||||
.defaultSystem("Default system text.")
|
||||
.defaultAdvisors(new PromptChatMemoryAdvisor(chatMemory))
|
||||
.defaultAdvisors(PromptChatMemoryAdvisor.builder(chatMemory).build())
|
||||
.build();
|
||||
|
||||
// Simulate a streaming user prompt and verify the response
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
/*
|
||||
* Copyright 2023-2025 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.client.advisor;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.chat.client.advisor.api.Advisor;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
|
||||
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Tests for {@link MessageChatMemoryAdvisor} builder method chaining.
|
||||
*
|
||||
* @author Mark Pollack
|
||||
*/
|
||||
public class MessageChatMemoryAdvisorTests {
|
||||
|
||||
@Test
|
||||
void testBuilderMethodChaining() {
|
||||
// Create a chat memory
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
// Test builder method chaining with methods from AbstractBuilder
|
||||
String customConversationId = "test-conversation-id";
|
||||
int customOrder = 42;
|
||||
boolean customProtectFromBlocking = false;
|
||||
|
||||
MessageChatMemoryAdvisor advisor = MessageChatMemoryAdvisor.builder(chatMemory)
|
||||
.conversationId(customConversationId)
|
||||
.order(customOrder)
|
||||
.protectFromBlocking(customProtectFromBlocking)
|
||||
.build();
|
||||
|
||||
// Verify the advisor was built with the correct properties
|
||||
assertThat(advisor).isNotNull();
|
||||
// We can't directly access private fields, but we can test the behavior
|
||||
// by checking the order which is exposed via a getter
|
||||
assertThat(advisor.getOrder()).isEqualTo(customOrder);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testDefaultValues() {
|
||||
// Create a chat memory
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
// Create advisor with default values
|
||||
MessageChatMemoryAdvisor advisor = MessageChatMemoryAdvisor.builder(chatMemory).build();
|
||||
|
||||
// Verify default values
|
||||
assertThat(advisor).isNotNull();
|
||||
assertThat(advisor.getOrder()).isEqualTo(Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,95 @@
|
||||
/*
|
||||
* Copyright 2023-2025 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.client.advisor;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.springframework.ai.chat.client.advisor.api.Advisor;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.ai.chat.memory.InMemoryChatMemoryRepository;
|
||||
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
|
||||
import org.springframework.ai.chat.prompt.PromptTemplate;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
/**
|
||||
* Tests for {@link PromptChatMemoryAdvisor} builder method chaining.
|
||||
*
|
||||
* @author Mark Pollack
|
||||
*/
|
||||
public class PromptChatMemoryAdvisorTests {
|
||||
|
||||
@Test
|
||||
void testBuilderMethodChaining() {
|
||||
// Create a chat memory
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
// Test builder method chaining with methods from AbstractBuilder and
|
||||
// PromptChatMemoryAdvisor.Builder
|
||||
String customConversationId = "test-conversation-id";
|
||||
int customOrder = 42;
|
||||
boolean customProtectFromBlocking = false;
|
||||
String customSystemPrompt = "Custom system prompt with {instructions} and {memory}";
|
||||
|
||||
PromptChatMemoryAdvisor advisor = PromptChatMemoryAdvisor.builder(chatMemory)
|
||||
.conversationId(customConversationId) // From AbstractBuilder
|
||||
.order(customOrder) // From AbstractBuilder
|
||||
.protectFromBlocking(customProtectFromBlocking) // From AbstractBuilder
|
||||
.build();
|
||||
|
||||
// Verify the advisor was built with the correct properties
|
||||
assertThat(advisor).isNotNull();
|
||||
assertThat(advisor.getOrder()).isEqualTo(customOrder);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSystemPromptTemplateChaining() {
|
||||
// Create a chat memory
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
// Test chaining with systemPromptTemplate method
|
||||
PromptTemplate customTemplate = new PromptTemplate("Custom template with {instructions} and {memory}");
|
||||
|
||||
PromptChatMemoryAdvisor advisor = PromptChatMemoryAdvisor.builder(chatMemory)
|
||||
.conversationId("custom-id")
|
||||
.systemPromptTemplate(customTemplate)
|
||||
.order(100)
|
||||
.build();
|
||||
|
||||
assertThat(advisor).isNotNull();
|
||||
assertThat(advisor.getOrder()).isEqualTo(100);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testDefaultValues() {
|
||||
// Create a chat memory
|
||||
ChatMemory chatMemory = MessageWindowChatMemory.builder()
|
||||
.chatMemoryRepository(new InMemoryChatMemoryRepository())
|
||||
.build();
|
||||
|
||||
// Create advisor with default values
|
||||
PromptChatMemoryAdvisor advisor = PromptChatMemoryAdvisor.builder(chatMemory).build();
|
||||
|
||||
// Verify default values
|
||||
assertThat(advisor).isNotNull();
|
||||
assertThat(advisor.getOrder()).isEqualTo(Advisor.DEFAULT_CHAT_MEMORY_PRECEDENCE_ORDER);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -28,11 +28,11 @@ import org.mockito.junit.jupiter.MockitoExtension;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClientRequest;
|
||||
import org.springframework.ai.chat.client.ChatClientResponse;
|
||||
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.CallAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
|
||||
import org.springframework.ai.chat.client.observation.ChatClientObservationDocumentation.HighCardinalityKeyNames;
|
||||
import org.springframework.ai.chat.client.observation.ChatClientObservationDocumentation.LowCardinalityKeyNames;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.ai.chat.model.ChatModel;
|
||||
import org.springframework.ai.chat.prompt.Prompt;
|
||||
import org.springframework.ai.model.tool.ToolCallingChatOptions;
|
||||
@@ -150,7 +150,7 @@ class DefaultChatClientObservationConventionTests {
|
||||
.toolNames("tool1", "tool2")
|
||||
.toolCallbacks(dummyFunction("toolCallback1"), dummyFunction("toolCallback2"))
|
||||
.build()))
|
||||
.context(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, "007")
|
||||
.context(ChatMemory.CONVERSATION_ID, "007")
|
||||
.build();
|
||||
|
||||
ChatClientObservationContext observationContext = ChatClientObservationContext.builder()
|
||||
|
||||
@@ -428,8 +428,8 @@ A sample `@Service` implementation that uses several advisors is shown below.
|
||||
|
||||
[source,java]
|
||||
----
|
||||
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
|
||||
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;
|
||||
import static org.springframework.ai.chat.memory.ChatMemory.CHAT_MEMORY_CONVERSATION_ID_KEY;
|
||||
import static org.springframework.ai.chat.client.advisor.vectorstore.VectorStoreChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;
|
||||
|
||||
@Service
|
||||
public class CustomerSupportAssistant {
|
||||
|
||||
@@ -29,12 +29,19 @@ For details, refer to:
|
||||
[[upgrading-to-1-0-0-RC1]]
|
||||
== Upgrading to 1.0.0-RC1
|
||||
|
||||
=== Chat Client And Advisors
|
||||
=== Chat ClientAnd Advisors
|
||||
|
||||
* When building a `Prompt` from the ChatClient input, the `SystemMessage` built from `systemText()` is now placed first in the message list. Before, it was put last, resulting in errors with several model providers.
|
||||
* In `AbstractChatMemoryAdvisor`, the `doNextWithProtectFromBlockingBefore()` protected method has been changed from accepting the old `AdvisedRequest` to the new `ChatClientRequest`. It’s a breaking change since the alternative was not part of M8.
|
||||
* `MessageAggregator` has a new method to aggregate messages from `ChatClientRequest`. The previous method aggregating messages from the old `AdvisedRequest` has been removed, since it was already marked as deprecated in M8.
|
||||
* In `SimpleLoggerAdvisor`, the `requestToString` input argument needs to be updated to use `ChatClientRequest`. It’s a breaking change since the alternative was not part of M8 yet. Same thing about the constructor.
|
||||
* In `SimpleLoggerAdvisor`, the `requestToString` input argument needs to be updated to use `ChatClientRequest`. It's a breaking change since the alternative was not part of M8 yet. Same thing about the constructor.
|
||||
* `AbstractChatMemoryAdvisor` has been replaced with a `BaseChatMemoryAdvisor` interface in the `api` package. This is a breaking change for any code that directly extended `AbstractChatMemoryAdvisor`.
|
||||
* Public constructors in `MessageChatMemoryAdvisor`, `PromptChatMemoryAdvisor`, and `VectorStoreChatMemoryAdvisor` have been made private. You must now use the builder pattern to create instances (e.g., `MessageChatMemoryAdvisor.builder(chatMemory).build()`).
|
||||
* The constant `CHAT_MEMORY_CONVERSATION_ID_KEY` has been renamed to `CONVERSATION_ID` and moved from `AbstractChatMemoryAdvisor` to the `ChatMemory` interface. Update your imports to use `org.springframework.ai.chat.memory.ChatMemory.CONVERSATION_ID`.
|
||||
* In `VectorStoreChatMemoryAdvisor`:
|
||||
** The constant `DEFAULT_CHAT_MEMORY_RESPONSE_SIZE` (value: 100) has been renamed to `DEFAULT_TOP_K` with a new default value of 20.
|
||||
** The builder method `chatMemoryRetrieveSize(int)` has been renamed to `topK(int)`. Update your code to use the new method name: `VectorStoreChatMemoryAdvisor.builder(store).topK(1).build()`.
|
||||
** The `systemTextAdvise(String)` builder method has been removed. Use the `systemPromptTemplate(PromptTemplate)` method instead.
|
||||
* In `PromptChatMemoryAdvisor`, the `systemTextAdvise(String)` builder method has been removed. Use the `systemPromptTemplate(PromptTemplate)` method instead.
|
||||
|
||||
==== Self-contained Templates in Advisors
|
||||
|
||||
|
||||
@@ -24,9 +24,9 @@ import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.evaluation.RelevancyEvaluator;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.ai.chat.memory.MessageWindowChatMemory;
|
||||
import org.springframework.ai.chat.model.ChatResponse;
|
||||
import org.springframework.ai.document.Document;
|
||||
@@ -152,8 +152,7 @@ class RetrievalAugmentationAdvisorIT {
|
||||
|
||||
ChatResponse chatResponse1 = chatClient.prompt()
|
||||
.user("Where does the adventure of Anacletus and Birba take place?")
|
||||
.advisors(advisors -> advisors.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,
|
||||
conversationId))
|
||||
.advisors(advisors -> advisors.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.chatResponse();
|
||||
|
||||
@@ -163,8 +162,7 @@ class RetrievalAugmentationAdvisorIT {
|
||||
|
||||
ChatResponse chatResponse2 = chatClient.prompt()
|
||||
.user("Did they meet any cow?")
|
||||
.advisors(advisors -> advisors.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY,
|
||||
conversationId))
|
||||
.advisors(advisors -> advisors.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.chatResponse();
|
||||
|
||||
|
||||
@@ -32,6 +32,11 @@ public interface ChatMemory {
|
||||
|
||||
String DEFAULT_CONVERSATION_ID = "default";
|
||||
|
||||
/**
|
||||
* The key to retrieve the chat memory conversation id from the context.
|
||||
*/
|
||||
String CONVERSATION_ID = "chat_memory_conversation_id";
|
||||
|
||||
/**
|
||||
* Save the specified message in the chat memory for the specified conversation.
|
||||
*/
|
||||
@@ -49,16 +54,7 @@ public interface ChatMemory {
|
||||
/**
|
||||
* Get the messages in the chat memory for the specified conversation.
|
||||
*/
|
||||
default List<Message> get(String conversationId) {
|
||||
Assert.hasText(conversationId, "conversationId cannot be null or empty");
|
||||
return get(conversationId, Integer.MAX_VALUE);
|
||||
}
|
||||
|
||||
/**
|
||||
* @deprecated in favor of using {@link MessageWindowChatMemory}.
|
||||
*/
|
||||
@Deprecated
|
||||
List<Message> get(String conversationId, int lastN);
|
||||
List<Message> get(String conversationId);
|
||||
|
||||
/**
|
||||
* Clear the chat memory for the specified conversation.
|
||||
|
||||
@@ -71,12 +71,6 @@ public final class MessageWindowChatMemory implements ChatMemory {
|
||||
return this.chatMemoryRepository.findByConversationId(conversationId);
|
||||
}
|
||||
|
||||
@Override
|
||||
@Deprecated // in favor of get(conversationId)
|
||||
public List<Message> get(String conversationId, int lastN) {
|
||||
return get(conversationId);
|
||||
}
|
||||
|
||||
@Override
|
||||
public void clear(String conversationId) {
|
||||
Assert.hasText(conversationId, "conversationId cannot be null or empty");
|
||||
|
||||
@@ -128,6 +128,20 @@ public class Prompt implements ModelRequest<List<Message>> {
|
||||
return new UserMessage("");
|
||||
}
|
||||
|
||||
/**
|
||||
* Get all user messages in the prompt.
|
||||
* @return a list of all user messages in the prompt
|
||||
*/
|
||||
public List<UserMessage> getUserMessages() {
|
||||
List<UserMessage> userMessages = new ArrayList<>();
|
||||
for (Message message : this.messages) {
|
||||
if (message instanceof UserMessage userMessage) {
|
||||
userMessages.add(userMessage);
|
||||
}
|
||||
}
|
||||
return userMessages;
|
||||
}
|
||||
|
||||
@Override
|
||||
public String toString() {
|
||||
return "Prompt{" + "messages=" + this.messages + ", modelOptions=" + this.chatOptions + '}';
|
||||
|
||||
@@ -0,0 +1,365 @@
|
||||
/*
|
||||
* Copyright 2023-2025 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.vectorstore.pgvector;
|
||||
|
||||
import java.util.UUID;
|
||||
|
||||
import org.junit.jupiter.api.Test;
|
||||
import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable;
|
||||
import org.testcontainers.containers.PostgreSQLContainer;
|
||||
import org.testcontainers.junit.jupiter.Container;
|
||||
import org.testcontainers.junit.jupiter.Testcontainers;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.client.advisor.vectorstore.VectorStoreChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.ai.document.Document;
|
||||
import org.springframework.ai.embedding.EmbeddingModel;
|
||||
import org.springframework.ai.model.ApiKey;
|
||||
import org.springframework.ai.model.SimpleApiKey;
|
||||
import org.springframework.ai.openai.OpenAiChatModel;
|
||||
import org.springframework.ai.openai.OpenAiChatOptions;
|
||||
import org.springframework.ai.openai.OpenAiEmbeddingModel;
|
||||
import org.springframework.ai.openai.api.OpenAiApi;
|
||||
import org.springframework.beans.factory.annotation.Autowired;
|
||||
import org.springframework.boot.test.context.SpringBootTest;
|
||||
import org.springframework.context.annotation.Bean;
|
||||
import org.springframework.jdbc.core.JdbcTemplate;
|
||||
|
||||
import static org.assertj.core.api.Assertions.assertThat;
|
||||
|
||||
@Testcontainers
|
||||
@SpringBootTest(classes = PgVectorStoreVectorStoreChatMemoryAdvisorIT.OpenAiTestConfiguration.class)
|
||||
@EnabledIfEnvironmentVariable(named = "OPENAI_API_KEY", matches = ".+")
|
||||
public class PgVectorStoreVectorStoreChatMemoryAdvisorIT {
|
||||
|
||||
@Container
|
||||
@SuppressWarnings("resource")
|
||||
static PostgreSQLContainer<?> postgresContainer = new PostgreSQLContainer<>(PgVectorImage.DEFAULT_IMAGE)
|
||||
.withUsername("postgres")
|
||||
.withPassword("postgres");
|
||||
|
||||
@Autowired
|
||||
protected org.springframework.ai.chat.model.ChatModel chatModel;
|
||||
|
||||
@Test
|
||||
void testUseCustomConversationId() throws Exception {
|
||||
String apiKey = System.getenv("OPENAI_API_KEY");
|
||||
org.junit.jupiter.api.Assumptions.assumeTrue(apiKey != null && !apiKey.isBlank(),
|
||||
"OPENAI_API_KEY must be set for this test");
|
||||
|
||||
// Use a real OpenAI embedding model
|
||||
EmbeddingModel embeddingModel = new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(apiKey).build());
|
||||
|
||||
// Create PgVectorStore
|
||||
JdbcTemplate jdbcTemplate = createJdbcTemplateWithConnectionToTestcontainer();
|
||||
PgVectorStore store = PgVectorStore.builder(jdbcTemplate, embeddingModel)
|
||||
.dimensions(1536) // OpenAI default embedding size (adjust if needed)
|
||||
.initializeSchema(true)
|
||||
.build();
|
||||
store.afterPropertiesSet();
|
||||
|
||||
// Add a document to the store for recall
|
||||
String conversationId = UUID.randomUUID().toString();
|
||||
store.add(java.util.List
|
||||
.of(new Document("Hello from memory", java.util.Map.of("conversationId", conversationId))));
|
||||
|
||||
// Build ChatClient with VectorStoreChatMemoryAdvisor
|
||||
ChatClient chatClient = ChatClient.builder(chatModel)
|
||||
.defaultAdvisors(VectorStoreChatMemoryAdvisor.builder(store).build())
|
||||
.build();
|
||||
|
||||
// Send a prompt
|
||||
String answer = chatClient.prompt()
|
||||
.user("Say hello")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
assertThat(answer).containsIgnoringCase("hello");
|
||||
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSemanticSearchRetrievesRelevantMemory() throws Exception {
|
||||
String apiKey = System.getenv("OPENAI_API_KEY");
|
||||
org.junit.jupiter.api.Assumptions.assumeTrue(apiKey != null && !apiKey.isBlank(),
|
||||
"OPENAI_API_KEY must be set for this test");
|
||||
|
||||
EmbeddingModel embeddingModel = new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(apiKey).build());
|
||||
JdbcTemplate jdbcTemplate = createJdbcTemplateWithConnectionToTestcontainer();
|
||||
PgVectorStore store = PgVectorStore.builder(jdbcTemplate, embeddingModel)
|
||||
.dimensions(1536)
|
||||
.initializeSchema(true)
|
||||
.build();
|
||||
store.afterPropertiesSet();
|
||||
|
||||
String conversationId = UUID.randomUUID().toString();
|
||||
// Store diverse messages
|
||||
store.add(java.util.List.of(
|
||||
new Document("The Eiffel Tower is in Paris.", java.util.Map.of("conversationId", conversationId)),
|
||||
new Document("Bananas are yellow.", java.util.Map.of("conversationId", conversationId)),
|
||||
new Document("Mount Everest is the tallest mountain in the world.",
|
||||
java.util.Map.of("conversationId", conversationId)),
|
||||
new Document("Dogs are loyal pets.", java.util.Map.of("conversationId", conversationId))));
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel)
|
||||
.defaultAdvisors(VectorStoreChatMemoryAdvisor.builder(store).topK(1).build())
|
||||
.build();
|
||||
|
||||
// Send a semantically related query
|
||||
String answer = chatClient.prompt()
|
||||
.user("Where is the Eiffel Tower located?")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
|
||||
// Assert that the answer is based on the correct semantic memory
|
||||
assertThat(answer).containsIgnoringCase("paris");
|
||||
assertThat(answer).doesNotContain("Bananas are yellow");
|
||||
assertThat(answer).doesNotContain("Mount Everest");
|
||||
assertThat(answer).doesNotContain("Dogs are loyal pets");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSemanticSynonymRetrieval() throws Exception {
|
||||
String apiKey = System.getenv("OPENAI_API_KEY");
|
||||
org.junit.jupiter.api.Assumptions.assumeTrue(apiKey != null && !apiKey.isBlank(),
|
||||
"OPENAI_API_KEY must be set for this test");
|
||||
|
||||
EmbeddingModel embeddingModel = new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(apiKey).build());
|
||||
JdbcTemplate jdbcTemplate = createJdbcTemplateWithConnectionToTestcontainer();
|
||||
PgVectorStore store = PgVectorStore.builder(jdbcTemplate, embeddingModel)
|
||||
.dimensions(1536)
|
||||
.initializeSchema(true)
|
||||
.build();
|
||||
store.afterPropertiesSet();
|
||||
|
||||
String conversationId = UUID.randomUUID().toString();
|
||||
store.add(java.util.List
|
||||
.of(new Document("Automobiles are fast.", java.util.Map.of("conversationId", conversationId))));
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel)
|
||||
.defaultAdvisors(VectorStoreChatMemoryAdvisor.builder(store).topK(1).build())
|
||||
.build();
|
||||
|
||||
String answer = chatClient.prompt()
|
||||
.user("Tell me about cars.")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
assertThat(answer).satisfiesAnyOf(a -> assertThat(a).containsIgnoringCase("automobile"),
|
||||
a -> assertThat(a).containsIgnoringCase("fast"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testIrrelevantMessageExclusion() throws Exception {
|
||||
String apiKey = System.getenv("OPENAI_API_KEY");
|
||||
org.junit.jupiter.api.Assumptions.assumeTrue(apiKey != null && !apiKey.isBlank(),
|
||||
"OPENAI_API_KEY must be set for this test");
|
||||
|
||||
EmbeddingModel embeddingModel = new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(apiKey).build());
|
||||
JdbcTemplate jdbcTemplate = createJdbcTemplateWithConnectionToTestcontainer();
|
||||
PgVectorStore store = PgVectorStore.builder(jdbcTemplate, embeddingModel)
|
||||
.dimensions(1536)
|
||||
.initializeSchema(true)
|
||||
.build();
|
||||
store.afterPropertiesSet();
|
||||
|
||||
String conversationId = UUID.randomUUID().toString();
|
||||
store.add(java.util.List.of(
|
||||
new Document("The capital of Italy is Rome.", java.util.Map.of("conversationId", conversationId)),
|
||||
new Document("Bananas are yellow.", java.util.Map.of("conversationId", conversationId))));
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel)
|
||||
.defaultAdvisors(VectorStoreChatMemoryAdvisor.builder(store).topK(2).build())
|
||||
.build();
|
||||
|
||||
String answer = chatClient.prompt()
|
||||
.user("What is the capital of Italy?")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
assertThat(answer).containsIgnoringCase("rome");
|
||||
assertThat(answer).doesNotContain("banana");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testTopKSemanticRelevance() throws Exception {
|
||||
String apiKey = System.getenv("OPENAI_API_KEY");
|
||||
org.junit.jupiter.api.Assumptions.assumeTrue(apiKey != null && !apiKey.isBlank(),
|
||||
"OPENAI_API_KEY must be set for this test");
|
||||
|
||||
EmbeddingModel embeddingModel = new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(apiKey).build());
|
||||
JdbcTemplate jdbcTemplate = createJdbcTemplateWithConnectionToTestcontainer();
|
||||
PgVectorStore store = PgVectorStore.builder(jdbcTemplate, embeddingModel)
|
||||
.dimensions(1536)
|
||||
.initializeSchema(true)
|
||||
.build();
|
||||
store.afterPropertiesSet();
|
||||
|
||||
String conversationId = UUID.randomUUID().toString();
|
||||
store.add(java.util.List.of(
|
||||
new Document("The cat sat on the mat.", java.util.Map.of("conversationId", conversationId)),
|
||||
new Document("A cat is a small domesticated animal.",
|
||||
java.util.Map.of("conversationId", conversationId)),
|
||||
new Document("Dogs are loyal pets.", java.util.Map.of("conversationId", conversationId))));
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel)
|
||||
.defaultAdvisors(VectorStoreChatMemoryAdvisor.builder(store).topK(1).build())
|
||||
.build();
|
||||
|
||||
String answer = chatClient.prompt()
|
||||
.user("What can you tell me about cats?")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
assertThat(answer).containsIgnoringCase("cat");
|
||||
assertThat(answer).doesNotContain("dog");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testSemanticRetrievalWithParaphrasing() throws Exception {
|
||||
String apiKey = System.getenv("OPENAI_API_KEY");
|
||||
org.junit.jupiter.api.Assumptions.assumeTrue(apiKey != null && !apiKey.isBlank(),
|
||||
"OPENAI_API_KEY must be set for this test");
|
||||
|
||||
EmbeddingModel embeddingModel = new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(apiKey).build());
|
||||
JdbcTemplate jdbcTemplate = createJdbcTemplateWithConnectionToTestcontainer();
|
||||
PgVectorStore store = PgVectorStore.builder(jdbcTemplate, embeddingModel)
|
||||
.dimensions(1536)
|
||||
.initializeSchema(true)
|
||||
.build();
|
||||
store.afterPropertiesSet();
|
||||
|
||||
String conversationId = UUID.randomUUID().toString();
|
||||
store.add(java.util.List.of(new Document("The quick brown fox jumps over the lazy dog.",
|
||||
java.util.Map.of("conversationId", conversationId))));
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel)
|
||||
.defaultAdvisors(VectorStoreChatMemoryAdvisor.builder(store).topK(1).build())
|
||||
.build();
|
||||
|
||||
String answer = chatClient.prompt()
|
||||
.user("Tell me about a fast animal leaping over another.")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
assertThat(answer).satisfiesAnyOf(a -> assertThat(a).containsIgnoringCase("fox"),
|
||||
a -> assertThat(a).containsIgnoringCase("dog"));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testMultipleRelevantMemoriesTopK() throws Exception {
|
||||
String apiKey = System.getenv("OPENAI_API_KEY");
|
||||
org.junit.jupiter.api.Assumptions.assumeTrue(apiKey != null && !apiKey.isBlank(),
|
||||
"OPENAI_API_KEY must be set for this test");
|
||||
|
||||
EmbeddingModel embeddingModel = new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(apiKey).build());
|
||||
JdbcTemplate jdbcTemplate = createJdbcTemplateWithConnectionToTestcontainer();
|
||||
PgVectorStore store = PgVectorStore.builder(jdbcTemplate, embeddingModel)
|
||||
.dimensions(1536)
|
||||
.initializeSchema(true)
|
||||
.build();
|
||||
store.afterPropertiesSet();
|
||||
|
||||
String conversationId = UUID.randomUUID().toString();
|
||||
store.add(java.util.List.of(new Document("Apples are red.", java.util.Map.of("conversationId", conversationId)),
|
||||
new Document("Strawberries are also red.", java.util.Map.of("conversationId", conversationId)),
|
||||
new Document("Bananas are yellow.", java.util.Map.of("conversationId", conversationId))));
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel)
|
||||
.defaultAdvisors(VectorStoreChatMemoryAdvisor.builder(store).topK(2).build())
|
||||
.build();
|
||||
|
||||
String answer = chatClient.prompt()
|
||||
.user("What fruits are red?")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
assertThat(answer).containsIgnoringCase("apple");
|
||||
assertThat(answer).containsIgnoringCase("strawber");
|
||||
assertThat(answer).doesNotContain("banana");
|
||||
}
|
||||
|
||||
@Test
|
||||
void testNoRelevantMemory() throws Exception {
|
||||
String apiKey = System.getenv("OPENAI_API_KEY");
|
||||
org.junit.jupiter.api.Assumptions.assumeTrue(apiKey != null && !apiKey.isBlank(),
|
||||
"OPENAI_API_KEY must be set for this test");
|
||||
|
||||
EmbeddingModel embeddingModel = new OpenAiEmbeddingModel(OpenAiApi.builder().apiKey(apiKey).build());
|
||||
JdbcTemplate jdbcTemplate = createJdbcTemplateWithConnectionToTestcontainer();
|
||||
PgVectorStore store = PgVectorStore.builder(jdbcTemplate, embeddingModel)
|
||||
.dimensions(1536)
|
||||
.initializeSchema(true)
|
||||
.build();
|
||||
store.afterPropertiesSet();
|
||||
|
||||
String conversationId = UUID.randomUUID().toString();
|
||||
store.add(java.util.List
|
||||
.of(new Document("The sun is a star.", java.util.Map.of("conversationId", conversationId))));
|
||||
|
||||
ChatClient chatClient = ChatClient.builder(chatModel)
|
||||
.defaultAdvisors(VectorStoreChatMemoryAdvisor.builder(store).topK(1).build())
|
||||
.build();
|
||||
|
||||
String answer = chatClient.prompt()
|
||||
.user("What is the capital of Spain?")
|
||||
.advisors(a -> a.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.content();
|
||||
assertThat(answer).doesNotContain("sun");
|
||||
assertThat(answer).doesNotContain("star");
|
||||
}
|
||||
|
||||
private static JdbcTemplate createJdbcTemplateWithConnectionToTestcontainer() {
|
||||
org.postgresql.ds.PGSimpleDataSource ds = new org.postgresql.ds.PGSimpleDataSource();
|
||||
ds.setUrl("jdbc:postgresql://localhost:" + postgresContainer.getMappedPort(5432) + "/postgres");
|
||||
ds.setUser(postgresContainer.getUsername());
|
||||
ds.setPassword(postgresContainer.getPassword());
|
||||
return new JdbcTemplate(ds);
|
||||
}
|
||||
|
||||
@org.springframework.context.annotation.Configuration
|
||||
public static class OpenAiTestConfiguration {
|
||||
|
||||
@Bean
|
||||
public OpenAiApi openAiApi() {
|
||||
return OpenAiApi.builder().apiKey(getApiKey()).build();
|
||||
}
|
||||
|
||||
private ApiKey getApiKey() {
|
||||
String apiKey = System.getenv("OPENAI_API_KEY");
|
||||
if (!org.springframework.util.StringUtils.hasText(apiKey)) {
|
||||
throw new IllegalArgumentException(
|
||||
"You must provide an API key. Put it in an environment variable under the name OPENAI_API_KEY");
|
||||
}
|
||||
return new SimpleApiKey(apiKey);
|
||||
}
|
||||
|
||||
@Bean
|
||||
public OpenAiChatModel openAiChatModel(OpenAiApi api) {
|
||||
return OpenAiChatModel.builder()
|
||||
.openAiApi(api)
|
||||
.defaultOptions(OpenAiChatOptions.builder().model(OpenAiApi.ChatModel.GPT_4_O_MINI).build())
|
||||
.build();
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
}
|
||||
@@ -27,13 +27,13 @@ import org.mockito.ArgumentCaptor;
|
||||
import org.mockito.ArgumentMatchers;
|
||||
import org.mockito.Mockito;
|
||||
import org.postgresql.ds.PGSimpleDataSource;
|
||||
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
|
||||
import org.testcontainers.containers.PostgreSQLContainer;
|
||||
import org.testcontainers.junit.jupiter.Container;
|
||||
import org.testcontainers.junit.jupiter.Testcontainers;
|
||||
|
||||
import org.springframework.ai.chat.client.ChatClient;
|
||||
import org.springframework.ai.chat.client.advisor.vectorstore.VectorStoreChatMemoryAdvisor;
|
||||
import org.springframework.ai.chat.memory.ChatMemory;
|
||||
import org.springframework.ai.chat.messages.AssistantMessage;
|
||||
import org.springframework.ai.chat.messages.SystemMessage;
|
||||
import org.springframework.ai.chat.model.ChatModel;
|
||||
@@ -138,7 +138,7 @@ class PgVectorStoreWithChatMemoryAdvisorIT {
|
||||
.prompt()
|
||||
.user("joke")
|
||||
.advisors(a -> a.advisors(VectorStoreChatMemoryAdvisor.builder(store).build())
|
||||
.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId))
|
||||
.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.chatResponse();
|
||||
|
||||
@@ -162,7 +162,7 @@ class PgVectorStoreWithChatMemoryAdvisorIT {
|
||||
.system("You are a helpful assistant.")
|
||||
.user("joke")
|
||||
.advisors(a -> a.advisors(VectorStoreChatMemoryAdvisor.builder(store).build())
|
||||
.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, conversationId))
|
||||
.param(ChatMemory.CONVERSATION_ID, conversationId))
|
||||
.call()
|
||||
.chatResponse();
|
||||
|
||||
|
||||
Reference in New Issue
Block a user