feat: add incremental refresh support to ChatMemoryRepository

Author: juntae6942

spring-ai-model/src/main/java/org/springframework/ai/chat/memory/ChatMemoryRepository.java

@@ -40,4 +40,15 @@ public interface ChatMemoryRepository {
 
 	void deleteByConversationId(String conversationId);
 
+	/**
+	 * Atomically removes the messages in {@code deletes} and adds the messages in {@code adds}
+	 * for the given conversation ID. This provides a more efficient way to update
+	 * the memory than reading the entire history and overwriting it.
+	 *
+	 * @param conversationId The ID of the conversation to update.
+	 * @param deletes A list of messages to be removed from the memory.
+	 * @param adds A list of new messages to be added to the memory.
+	 */
+	void refresh(String conversationId, List<Message> deletes, List<Message> adds);
+
 }

spring-ai-model/src/main/java/org/springframework/ai/chat/memory/InMemoryChatMemoryRepository.java

@@ -60,4 +60,17 @@ public void deleteByConversationId(String conversationId) {
 		this.chatMemoryStore.remove(conversationId);
 	}
 
+	@Override
+	public void refresh(String conversationId, List<Message> deletes, List<Message> adds) {
+		this.chatMemoryStore.compute(conversationId, (key, currentMessages) -> {
+			if (currentMessages == null) {
+				return new ArrayList<>(adds);
+			}
+			List<Message> updatedMessages = new ArrayList<>(currentMessages);
+			updatedMessages.removeAll(deletes);
+			updatedMessages.addAll(adds);
+			return updatedMessages;
+		});
+	}
+
 }

spring-ai-model/src/main/java/org/springframework/ai/chat/memory/MessageWindowChatMemory.java

@@ -18,6 +18,7 @@
 
 import java.util.ArrayList;
 import java.util.HashSet;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Set;
 
@@ -61,8 +62,10 @@ public void add(String conversationId, List<Message> messages) {
 		Assert.noNullElements(messages, "messages cannot contain null elements");
 
 		List<Message> memoryMessages = this.chatMemoryRepository.findByConversationId(conversationId);
-		List<Message> processedMessages = process(memoryMessages, messages);
-		this.chatMemoryRepository.saveAll(conversationId, processedMessages);
+		MessageChanges changes = process(memoryMessages, messages);
+		if (!changes.toDelete.isEmpty() || !changes.toAdd.isEmpty()) {
+			this.chatMemoryRepository.refresh(conversationId, changes.toDelete, changes.toAdd);
+		}
 	}
 
 	@Override
@@ -77,38 +80,58 @@ public void clear(String conversationId) {
 		this.chatMemoryRepository.deleteByConversationId(conversationId);
 	}
 
-	private List<Message> process(List<Message> memoryMessages, List<Message> newMessages) {
-		List<Message> processedMessages = new ArrayList<>();
+	private MessageChanges process(List<Message> memoryMessages, List<Message> newMessages) {
+		Set<Message> originalMessageSet = new LinkedHashSet<>(memoryMessages);
+		List<Message> uniqueNewMessages = newMessages.stream()
+				.filter(msg -> !originalMessageSet.contains(msg))
+				.toList();
+		boolean hasNewSystemMessage = uniqueNewMessages.stream().anyMatch(SystemMessage.class::isInstance);
+
+		List<Message> finalMessages = new ArrayList<>();
+		if (hasNewSystemMessage) {
+			memoryMessages.stream().filter(msg -> !(msg instanceof SystemMessage)).forEach(finalMessages::add);
+			finalMessages.addAll(uniqueNewMessages);
+		}
+		else {
+			finalMessages.addAll(memoryMessages);
+			finalMessages.addAll(uniqueNewMessages);
+		}
 
-		Set<Message> memoryMessagesSet = new HashSet<>(memoryMessages);
-		boolean hasNewSystemMessage = newMessages.stream()
-			.filter(SystemMessage.class::isInstance)
-			.anyMatch(message -> !memoryMessagesSet.contains(message));
+		if (finalMessages.size() > this.maxMessages) {
+			List<Message> trimmedMessages = new ArrayList<>();
+			int messagesToRemove = finalMessages.size() - this.maxMessages;
+			int removed = 0;
+			for (Message message : finalMessages) {
+				if (message instanceof SystemMessage || removed >= messagesToRemove) {
+					trimmedMessages.add(message);
+				}
+				else {
+					removed++;
+				}
+			}
+			finalMessages = trimmedMessages;
+		}
 
-		memoryMessages.stream()
-			.filter(message -> !(hasNewSystemMessage && message instanceof SystemMessage))
-			.forEach(processedMessages::add);
+		Set<Message> finalMessageSet = new LinkedHashSet<>(finalMessages);
 
-		processedMessages.addAll(newMessages);
+		List<Message> toDelete = originalMessageSet.stream().filter(m -> !finalMessageSet.contains(m)).toList();
 
-		if (processedMessages.size() <= this.maxMessages) {
-			return processedMessages;
-		}
+		List<Message> toAdd = finalMessageSet.stream().filter(m -> !originalMessageSet.contains(m)).toList();
 
-		int messagesToRemove = processedMessages.size() - this.maxMessages;
+		return new MessageChanges(toDelete, toAdd);
+	}
 
-		List<Message> trimmedMessages = new ArrayList<>();
-		int removed = 0;
-		for (Message message : processedMessages) {
-			if (message instanceof SystemMessage || removed >= messagesToRemove) {
-				trimmedMessages.add(message);
-			}
-			else {
-				removed++;
-			}
+	private static class MessageChanges {
+
+		final List<Message> toDelete;
+
+		final List<Message> toAdd;
+
+		MessageChanges(List<Message> toDelete, List<Message> toAdd) {
+			this.toDelete = toDelete;
+			this.toAdd = toAdd;
 		}
 
-		return trimmedMessages;
 	}
 
 	public static Builder builder() {