Simplify chat memory advisor hierarchy and remove deprecated API

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

@@ -20,24 +20,25 @@
 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 @@
  * @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,71 +78,84 @@ 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) {
 		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;
+	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 @@ private void after(ChatClientResponse chatClientResponse) {
 				.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 @@ else if (message instanceof AssistantMessage assistantMessage) {
 		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);
-		}
+		private Integer topK = DEFAULT_TOP_K;
 
-		public Builder systemTextAdvise(String systemTextAdvise) {
-			this.systemPromptTemplate = new PromptTemplate(systemTextAdvise);
-			return this;
+		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);
 		}
 
 	}

models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/client/OpenAiChatClientMemoryAdvisorReproIT.java

@@ -39,7 +39,7 @@ void messageChatMemoryAdvisor_withPromptMessages_throwsException() {
 		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();