Refactor Chat Memory Advisors to reduce code duplication and add extensive tests

Core Architecture Changes:
1. New Abstract Class: AbstractConversationHistoryAdvisor
   - Created a new abstract class that extends AbstractChatMemoryAdvisor<ChatMemory>
   - Provides common functionality for managing conversation history
   - Implements methods for retrieving messages and applying them to requests
   - Simplifies the implementation of concrete advisors

2. Refactored AbstractChatMemoryAdvisor
   - Removed the defaultChatMemoryRetrieveSize parameter from constructors
   - Enhanced the builder pattern with generic type parameters for better type safety
   - Added abstract before(ChatClientRequest, String) method for subclasses to implement
   - Improved logging for conversation ID handling

3. Refactored MessageChatMemoryAdvisor
   - Now extends AbstractConversationHistoryAdvisor instead of directly extending AbstractChatMemoryAdvisor
   - Simplified implementation by leveraging parent class methods
   - Fixed handling of multiple user messages in a single prompt
   - Updated the builder to properly extend AbstractChatMemoryAdvisor.AbstractBuilder<ChatMemory, Builder>

4. Refactored PromptChatMemoryAdvisor
   - Now extends AbstractConversationHistoryAdvisor instead of directly extending AbstractChatMemoryAdvisor
   - Fixed handling of multiple user messages by using getUserMessages() instead of getUserMessage()
   - Enhanced logging for better debugging
   - Updated the builder to properly extend AbstractChatMemoryAdvisor.AbstractBuilder<ChatMemory, Builder>

Test Changes:
1. New Test Classes
   - Added MessageChatMemoryAdvisorIT for integration testing of MessageChatMemoryAdvisor
   - Added PromptChatMemoryAdvisorIT for integration testing of PromptChatMemoryAdvisor
   - Both extend from AbstractChatMemoryAdvisorIT to share common test logic

2. Test Coverage
   - Added tests for handling multiple user messages in a single prompt
   - Added tests for custom conversation IDs
   - Added tests for maintaining separate conversations
   - Added tests for reactive mode operation

Key Improvements:
1. Code Duplication Reduction: Moved common functionality to the parent class
2. Bug Fix: Fixed a bug in PromptChatMemoryAdvisor where it was only storing the last user message
3. Enhanced Type Safety: Improved the builder pattern with proper generic type parameters
4. Better Logging: Added detailed logging for better debugging and traceability
5. Simplified API: Removed unnecessary parameters from constructors
6. Improved Test Coverage: Added comprehensive tests for various scenarios

Signed-off-by: Mark Pollack <[email protected]>

Author: markpollack

advisors/spring-ai-advisors-vector-store/src/main/java/org/springframework/ai/chat/client/advisor/vectorstore/ParameterizedChatMemory.java

@@ -0,0 +1,24 @@
+package org.springframework.ai.chat.client.advisor.vectorstore;
+
+import java.util.List;
+
+import org.springframework.ai.chat.memory.ChatMemory;
+import org.springframework.ai.chat.messages.Message;
+
+/**
+ * Interface for chat memories that support parameterized retrieval. Implementations can
+ * define their own parameter types.
+ *
+ * @param <P> The type of parameters used for retrieval
+ */
+public interface ParameterizedChatMemory<P> extends ChatMemory {
+
+	/**
+	 * Retrieve messages based on the provided parameters.
+	 * @param conversationId The conversation identifier
+	 * @param parameters The retrieval parameters
+	 * @return List of retrieved messages
+	 */
+	List<Message> retrieve(String conversationId, P parameters);
+
+}

advisors/spring-ai-advisors-vector-store/src/main/java/org/springframework/ai/chat/client/advisor/vectorstore/VectorSearchParameters.java

@@ -0,0 +1,37 @@
+package org.springframework.ai.chat.client.advisor.vectorstore;
+
+import org.springframework.ai.vectorstore.SearchRequest;
+
+/**
+ * Parameters for vector store similarity search with default values.
+ */
+public record VectorSearchParameters(int topK, String filter) {
+
+	public static final int DEFAULT_TOP_K = 100;
+
+	public VectorSearchParameters() {
+		this(VectorStoreChatMemoryAdvisor.DEFAULT_CHAT_MEMORY_RESPONSE_SIZE, null);
+	}
+
+	public static VectorSearchParameters of(int topK) {
+		return new VectorSearchParameters(topK, null);
+	}
+
+	public static VectorSearchParameters forConversation(String conversationId) {
+		return new VectorSearchParameters(DEFAULT_TOP_K, "conversationId=='" + conversationId + "'");
+	}
+
+	public VectorSearchParameters withTopK(int topK) {
+		return new VectorSearchParameters(topK, this.filter);
+	}
+
+	/**
+	 * Create a new instance with a different filter.
+	 * @param filter the new filter expression
+	 * @return a new VectorSearchParameters instance with the updated filter
+	 */
+	public VectorSearchParameters withFilter(String filter) {
+		return new VectorSearchParameters(this.topK, filter);
+	}
+
+}

advisors/spring-ai-advisors-vector-store/src/main/java/org/springframework/ai/chat/client/advisor/vectorstore/VectorStoreChatMemoryAdvisor.java

@@ -24,9 +24,9 @@
 
 import reactor.core.publisher.Flux;
 
-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.AbstractChatMemoryAdvisor;
 import org.springframework.ai.chat.client.advisor.api.CallAdvisorChain;
 import org.springframework.ai.chat.client.advisor.api.StreamAdvisorChain;
 import org.springframework.ai.chat.messages.AssistantMessage;
@@ -52,10 +52,17 @@
  */
 public class VectorStoreChatMemoryAdvisor extends AbstractChatMemoryAdvisor<VectorStore> {
 
+	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_CHAT_MEMORY_RESPONSE_SIZE = 100;
+
 	private static final PromptTemplate DEFAULT_SYSTEM_PROMPT_TEMPLATE = new PromptTemplate("""
 			{instructions}
 
@@ -69,71 +76,62 @@ 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;
+
+	public VectorStoreChatMemoryAdvisor(VectorStore chatMemory, String defaultConversationId,
+			int defaultChatMemoryRetrieveSize, boolean protectFromBlocking, PromptTemplate systemPromptTemplate,
+			int order) {
+		super(chatMemory, defaultConversationId, protectFromBlocking, order);
 		this.systemPromptTemplate = systemPromptTemplate;
+		this.defaultChatMemoryRetrieveSize = defaultChatMemoryRetrieveSize;
 	}
 
 	public static Builder builder(VectorStore chatMemory) {
 		return new Builder(chatMemory);
 	}
 
-	@Override
-	public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain callAdvisorChain) {
-		chatClientRequest = this.before(chatClientRequest);
-
-		ChatClientResponse chatClientResponse = callAdvisorChain.nextCall(chatClientRequest);
-
-		this.after(chatClientResponse);
-
-		return chatClientResponse;
+	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;
 	}
 
 	@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);
-	}
-
-	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 + "'")
+	protected ChatClientRequest before(ChatClientRequest request, String conversationId) {
+		String query = request.prompt().getUserMessage() != null ? request.prompt().getUserMessage().getText() : "";
+		int topK = doGetChatMemoryRetrieveSize(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.getChatMemoryStore()
+			.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.getChatMemoryStore().write(toDocuments(java.util.List.of(userMessage), conversationId));
+		}
 
 		return processedChatClientRequest;
 	}
 
-	private void after(ChatClientResponse chatClientResponse) {
+	protected void after(ChatClientResponse chatClientResponse) {
 		List<Message> assistantMessages = new ArrayList<>();
 		if (chatClientResponse.chatResponse() != null) {
 			assistantMessages = chatClientResponse.chatResponse()
@@ -146,6 +144,24 @@ private void after(ChatClientResponse chatClientResponse) {
 			.write(toDocuments(assistantMessages, this.doGetConversationId(chatClientResponse.context())));
 	}
 
+	protected ChatClientRequest applyMessagesToRequest(ChatClientRequest request, List<Message> memoryMessages) {
+		if (memoryMessages == null || memoryMessages.isEmpty()) {
+			return request;
+		}
+		// Convert memory messages to a string for the system prompt
+		String longTermMemory = memoryMessages.stream()
+			.filter(m -> m.getMessageType() == MessageType.USER || m.getMessageType() == MessageType.ASSISTANT)
+			.map(m -> m.getMessageType() + ":" + m.getText())
+			.collect(Collectors.joining(System.lineSeparator()));
+
+		SystemMessage systemMessage = request.prompt().getSystemMessage();
+		String augmentedSystemText = this.systemPromptTemplate
+			.render(Map.of("instructions", systemMessage.getText(), "long_term_memory", longTermMemory));
+
+		// Create a new request with the augmented system message
+		return request.mutate().prompt(request.prompt().augmentSystemMessage(augmentedSystemText)).build();
+	}
+
 	private List<Document> toDocuments(List<Message> messages, String conversationId) {
 		List<Document> docs = messages.stream()
 			.filter(m -> m.getMessageType() == MessageType.USER || m.getMessageType() == MessageType.ASSISTANT)
@@ -173,28 +189,71 @@ else if (message instanceof AssistantMessage assistantMessage) {
 		return docs;
 	}
 
-	public static class Builder extends AbstractChatMemoryAdvisor.AbstractBuilder<VectorStore> {
+	/**
+	 * Builder for VectorStoreChatMemoryAdvisor.
+	 */
+	public static class Builder extends AbstractChatMemoryAdvisor.AbstractBuilder<VectorStore, Builder> {
 
 		private PromptTemplate systemPromptTemplate = DEFAULT_SYSTEM_PROMPT_TEMPLATE;
 
-		protected Builder(VectorStore chatMemory) {
-			super(chatMemory);
+		private Integer defaultChatMemoryRetrieveSize = null;
+
+		/**
+		 * Creates a new builder instance.
+		 * @param vectorStore the vector store to use
+		 */
+		protected Builder(VectorStore vectorStore) {
+			super(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;
 		}
 
+		/**
+		 * Set the system prompt template using a text template.
+		 * @param systemTextAdvise the system prompt text template
+		 * @return this builder
+		 */
 		public Builder systemTextAdvise(String systemTextAdvise) {
 			this.systemPromptTemplate = new PromptTemplate(systemTextAdvise);
 			return this;
 		}
 
-		public Builder systemPromptTemplate(PromptTemplate systemPromptTemplate) {
-			this.systemPromptTemplate = systemPromptTemplate;
+		/**
+		 * Set the default chat memory retrieve size.
+		 * @param defaultChatMemoryRetrieveSize the default chat memory retrieve size
+		 * @return this builder
+		 */
+		public Builder defaultChatMemoryRetrieveSize(int defaultChatMemoryRetrieveSize) {
+			this.defaultChatMemoryRetrieveSize = defaultChatMemoryRetrieveSize;
+			return this;
+		}
+
+		@Override
+		protected Builder self() {
 			return this;
 		}
 
 		@Override
 		public VectorStoreChatMemoryAdvisor build() {
-			return new VectorStoreChatMemoryAdvisor(this.chatMemory, this.conversationId, this.chatMemoryRetrieveSize,
-					this.protectFromBlocking, this.systemPromptTemplate, this.order);
+			if (defaultChatMemoryRetrieveSize == null) {
+				// Default to legacy mode for backward compatibility
+				return new VectorStoreChatMemoryAdvisor(this.chatMemory, this.conversationId,
+						DEFAULT_CHAT_MEMORY_RESPONSE_SIZE, this.protectFromBlocking, this.systemPromptTemplate,
+						this.order);
+			}
+			else {
+				return new VectorStoreChatMemoryAdvisor(this.chatMemory, this.conversationId,
+						this.defaultChatMemoryRetrieveSize, this.protectFromBlocking, this.systemPromptTemplate,
+						this.order);
+			}
 		}
 
 	}