chat-memory-jdbc : Fix message order when retrieving + also save order for batch inserts

Signed-off-by: linarkou <[email protected]>

Author: linarkou

auto-configurations/models/chat/memory/spring-ai-autoconfigure-model-chat-memory-jdbc/src/test/java/org/springframework/ai/model/chat/memory/jdbc/autoconfigure/JdbcChatMemoryAutoConfigurationIT.java

@@ -82,6 +82,14 @@ void addGetAndClear_shouldAllExecute() {
 			assertThat(chatMemory.get(conversationId, Integer.MAX_VALUE)).hasSize(1);
 			assertThat(chatMemory.get(conversationId, Integer.MAX_VALUE)).isEqualTo(List.of(userMessage));
 
+			var assistantMessage = new AssistantMessage("Message from the assistant");
+
+			chatMemory.add(conversationId, assistantMessage);
+
+			assertThat(chatMemory.get(conversationId, Integer.MAX_VALUE)).hasSize(2);
+			assertThat(chatMemory.get(conversationId, Integer.MAX_VALUE))
+				.isEqualTo(List.of(userMessage, assistantMessage));
+
 			chatMemory.clear(conversationId);
 
 			assertThat(chatMemory.get(conversationId, Integer.MAX_VALUE)).isEmpty();

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

@@ -19,6 +19,9 @@
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Timestamp;
+import java.time.Instant;
+import java.util.Collections;
 import java.util.List;
 
 import org.springframework.ai.chat.memory.ChatMemory;
@@ -42,7 +45,7 @@
 public class JdbcChatMemory implements ChatMemory {
 
 	private static final String QUERY_ADD = """
-			INSERT INTO ai_chat_memory (conversation_id, content, type) VALUES (?, ?, ?)""";
+			INSERT INTO ai_chat_memory (conversation_id, content, type, "timestamp") VALUES (?, ?, ?, ?)""";
 
 	private static final String QUERY_GET = """
 			SELECT content, type FROM ai_chat_memory WHERE conversation_id = ? ORDER BY "timestamp" DESC LIMIT ?""";
@@ -66,7 +69,9 @@ public void add(String conversationId, List<Message> messages) {
 
 	@Override
 	public List<Message> get(String conversationId, int lastN) {
-		return this.jdbcTemplate.query(QUERY_GET, new MessageRowMapper(), conversationId, lastN);
+		List<Message> messages = this.jdbcTemplate.query(QUERY_GET, new MessageRowMapper(), conversationId, lastN);
+		Collections.reverse(messages);
+		return messages;
 	}
 
 	@Override
@@ -83,6 +88,7 @@ public void setValues(PreparedStatement ps, int i) throws SQLException {
 			ps.setString(1, this.conversationId);
 			ps.setString(2, message.getText());
 			ps.setString(3, message.getMessageType().name());
+			ps.setTimestamp(4, Timestamp.from(Instant.now()));
 		}
 
 		@Override

memory/spring-ai-model-chat-memory-jdbc/src/test/java/org/springframework/ai/chat/memory/jdbc/JdbcChatMemoryIT.java

@@ -16,27 +16,12 @@
 
 package org.springframework.ai.chat.memory.jdbc;
 
-import java.sql.Timestamp;
-import java.util.List;
-import java.util.UUID;
-
-import javax.sql.DataSource;
-
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.CsvSource;
-import org.testcontainers.containers.PostgreSQLContainer;
-import org.testcontainers.junit.jupiter.Container;
-import org.testcontainers.junit.jupiter.Testcontainers;
-import org.testcontainers.utility.MountableFile;
-
 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.messages.*;
 import org.springframework.boot.SpringBootConfiguration;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
 import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
@@ -46,6 +31,15 @@
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Primary;
 import org.springframework.jdbc.core.JdbcTemplate;
+import org.testcontainers.containers.PostgreSQLContainer;
+import org.testcontainers.junit.jupiter.Container;
+import org.testcontainers.junit.jupiter.Testcontainers;
+import org.testcontainers.utility.MountableFile;
+
+import javax.sql.DataSource;
+import java.sql.Timestamp;
+import java.util.List;
+import java.util.UUID;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -147,10 +141,11 @@ void get_shouldReturnMessages() {
 		this.contextRunner.run(context -> {
 			var chatMemory = context.getBean(ChatMemory.class);
 			var conversationId = UUID.randomUUID().toString();
-			var messages = List.<Message>of(new AssistantMessage("Message from assistant 1 - " + conversationId),
-					new AssistantMessage("Message from assistant 2 - " + conversationId),
-					new UserMessage("Message from user - " + conversationId),
-					new SystemMessage("Message from system - " + conversationId));
+			var messages = List.<Message>of(new SystemMessage("Message from system - " + conversationId),
+					new UserMessage("Message from user 1 - " + conversationId),
+					new AssistantMessage("Message from assistant 1 - " + conversationId),
+					new UserMessage("Message from user 2 - " + conversationId),
+					new AssistantMessage("Message from assistant 2 - " + conversationId));
 
 			chatMemory.add(conversationId, messages);
 
@@ -161,6 +156,24 @@ void get_shouldReturnMessages() {
 		});
 	}
 
+	@Test
+	void get_afterMultipleAdds_shouldReturnMessagesInSameOrder() {
+		this.contextRunner.run(context -> {
+			var chatMemory = context.getBean(ChatMemory.class);
+			var conversationId = UUID.randomUUID().toString();
+			var userMessage = new UserMessage("Message from user - " + conversationId);
+			var assistantMessage = new AssistantMessage("Message from assistant - " + conversationId);
+
+			chatMemory.add(conversationId, userMessage);
+			chatMemory.add(conversationId, assistantMessage);
+
+			var results = chatMemory.get(conversationId, Integer.MAX_VALUE);
+
+			assertThat(results.size()).isEqualTo(2);
+			assertThat(results).isEqualTo(List.of(userMessage, assistantMessage));
+		});
+	}
+
 	@Test
 	void clear_shouldDeleteMessages() {
 		this.contextRunner.run(context -> {