chat-memory-jdbc:

* set timestamp manually to prevent saving equal timestamps while batch insert
* return DESC order into get query (in `JdbcChatMemory`)

Signed-off-by: Linar Abzaltdinov <[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/JdbcChatMemoryPostgresqlAutoConfigurationIT.java

@@ -87,6 +87,13 @@ 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, List.of(assistantMessage));
+
+			assertThat(chatMemory.get(conversationId)).hasSize(2);
+			assertThat(chatMemory.get(conversationId)).isEqualTo(List.of(userMessage, assistantMessage));
+
 			chatMemory.clear(conversationId);
 
 			assertThat(chatMemory.get(conversationId, Integer.MAX_VALUE)).isEmpty();
@@ -142,6 +149,13 @@ void useAutoConfiguredChatMemoryWithJdbc() {
 			assertThat(chatMemory.get(conversationId)).hasSize(1);
 			assertThat(chatMemory.get(conversationId)).isEqualTo(List.of(userMessage));
 
+			var assistantMessage = new AssistantMessage("Message from the assistant");
+
+			chatMemory.add(conversationId, List.of(assistantMessage));
+
+			assertThat(chatMemory.get(conversationId)).hasSize(2);
+			assertThat(chatMemory.get(conversationId)).isEqualTo(List.of(userMessage, assistantMessage));
+
 			chatMemory.clear(conversationId);
 
 			assertThat(chatMemory.get(conversationId)).isEmpty();

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

@@ -23,6 +23,7 @@
 import java.time.Instant;
 import java.util.Collections;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicLong;
 
 import org.springframework.ai.chat.memory.ChatMemory;
 import org.springframework.ai.chat.memory.MessageWindowChatMemory;
@@ -52,7 +53,7 @@ public class JdbcChatMemory implements ChatMemory {
 			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" LIMIT ?""";
+			SELECT content, type FROM ai_chat_memory WHERE conversation_id = ? ORDER BY "timestamp" DESC LIMIT ?""";
 
 	private static final String QUERY_CLEAR = "DELETE FROM ai_chat_memory WHERE conversation_id = ?";
 
@@ -83,16 +84,21 @@ public void clear(String conversationId) {
 		this.jdbcTemplate.update(QUERY_CLEAR, conversationId);
 	}
 
-	private record AddBatchPreparedStatement(String conversationId,
-			List<Message> messages) implements BatchPreparedStatementSetter {
+	private record AddBatchPreparedStatement(String conversationId, List<Message> messages,
+											 AtomicLong instantSeq) implements BatchPreparedStatementSetter {
+
+		private AddBatchPreparedStatement(String conversationId, List<Message> messages) {
+			this(conversationId, messages, new AtomicLong(Instant.now().toEpochMilli()));
+		}
+
 		@Override
 		public void setValues(PreparedStatement ps, int i) throws SQLException {
 			var message = this.messages.get(i);
 
 			ps.setString(1, this.conversationId);
 			ps.setString(2, message.getText());
 			ps.setString(3, message.getMessageType().name());
-			ps.setTimestamp(4, Timestamp.from(Instant.now()));
+			ps.setTimestamp(4, new Timestamp(instantSeq.getAndIncrement()));
 		}
 
 		@Override

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

@@ -27,8 +27,11 @@
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Timestamp;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicLong;
 
 /**
  * An implementation of {@link ChatMemoryRepository} for JDBC.
@@ -44,7 +47,7 @@ public class JdbcChatMemoryRepository implements ChatMemoryRepository {
 			""";
 
 	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 = """
@@ -93,15 +96,21 @@ public void deleteByConversationId(String conversationId) {
 		this.jdbcTemplate.update(QUERY_CLEAR, conversationId);
 	}
 
-	private record AddBatchPreparedStatement(String conversationId,
-			List<Message> messages) implements BatchPreparedStatementSetter {
+	private record AddBatchPreparedStatement(String conversationId, List<Message> messages,
+			AtomicLong instantSeq) implements BatchPreparedStatementSetter {
+
+		private AddBatchPreparedStatement(String conversationId, List<Message> messages) {
+			this(conversationId, messages, new AtomicLong(Instant.now().toEpochMilli()));
+		}
+
 		@Override
 		public void setValues(PreparedStatement ps, int i) throws SQLException {
 			var message = this.messages.get(i);
 
 			ps.setString(1, this.conversationId);
 			ps.setString(2, message.getText());
 			ps.setString(3, message.getMessageType().name());
+			ps.setTimestamp(4, new Timestamp(instantSeq.getAndIncrement()));
 		}
 
 		@Override