fix: Fixed the issue in JdbcChatMemoryRepository where executing saveAll would update the timestamps of all historical messages to the current timestamp.

Signed-off-by: Sun Yuhan <[email protected]>
Signed-off-by: Sun Yuhan <[email protected]>

Author: sunyuhan1998

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

@@ -23,7 +23,7 @@ public class HsqldbChatMemoryRepositoryDialect implements JdbcChatMemoryReposito
 
 	@Override
 	public String getSelectMessagesSql() {
-		return "SELECT content, type FROM SPRING_AI_CHAT_MEMORY WHERE conversation_id = ? ORDER BY timestamp ASC";
+		return "SELECT content, type, timestamp FROM SPRING_AI_CHAT_MEMORY WHERE conversation_id = ? ORDER BY timestamp ASC";
 	}
 
 	@Override

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

@@ -23,6 +23,8 @@
 import java.time.Instant;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.atomic.AtomicLong;
 
 import org.springframework.ai.chat.memory.ChatMemoryRepository;
@@ -53,6 +55,9 @@ public class JdbcChatMemoryRepository implements ChatMemoryRepository {
 
 	private final JdbcChatMemoryRepositoryDialect dialect;
 
+	public static final String JDBC_MESSAGE_TS = JdbcChatMemoryRepository.class.getSimpleName()
+			+ "_message_timestamp";
+
 	private JdbcChatMemoryRepository(JdbcTemplate jdbcTemplate, JdbcChatMemoryRepositoryDialect dialect) {
 		Assert.notNull(jdbcTemplate, "jdbcTemplate cannot be null");
 		Assert.notNull(dialect, "dialect cannot be null");
@@ -104,11 +109,14 @@ private AddBatchPreparedStatement(String conversationId, List<Message> messages)
 		@Override
 		public void setValues(PreparedStatement ps, int i) throws SQLException {
 			var message = this.messages.get(i);
+			Long timestamp = Optional.ofNullable(message.getMetadata())
+					.map(metadata -> (long) metadata.getOrDefault(JDBC_MESSAGE_TS, instantSeq.get()))
+					.orElse(instantSeq.get());
 
 			ps.setString(1, this.conversationId);
 			ps.setString(2, message.getText());
 			ps.setString(3, message.getMessageType().name());
-			ps.setTimestamp(4, new Timestamp(instantSeq.getAndIncrement()));
+			ps.setTimestamp(4, new Timestamp(timestamp));
 		}
 
 		@Override
@@ -124,15 +132,18 @@ private static class MessageRowMapper implements RowMapper<Message> {
 		public Message mapRow(ResultSet rs, int i) throws SQLException {
 			var content = rs.getString(1);
 			var type = MessageType.valueOf(rs.getString(2));
+			long timestamp = rs.getLong(3);
+
+			Map<String, Object> metadata = Map.of(JDBC_MESSAGE_TS, timestamp);
 
 			return switch (type) {
-				case USER -> new UserMessage(content);
-				case ASSISTANT -> new AssistantMessage(content);
-				case SYSTEM -> new SystemMessage(content);
+				case USER -> UserMessage.builder().text(content).metadata(metadata).build();
+				case ASSISTANT -> new AssistantMessage(content, metadata);
+				case SYSTEM -> SystemMessage.builder().text(content).metadata(metadata).build();
 				// The content is always stored empty for ToolResponseMessages.
 				// If we want to capture the actual content, we need to extend
 				// AddBatchPreparedStatement to support it.
-				case TOOL -> new ToolResponseMessage(List.of());
+				case TOOL -> new ToolResponseMessage(List.of(), metadata);
 			};
 		}
 

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

@@ -26,7 +26,7 @@ public class MysqlChatMemoryRepositoryDialect implements JdbcChatMemoryRepositor
 
 	@Override
 	public String getSelectMessagesSql() {
-		return "SELECT content, type FROM SPRING_AI_CHAT_MEMORY WHERE conversation_id = ? ORDER BY `timestamp` DESC LIMIT ?";
+		return "SELECT content, type, `timestamp` FROM SPRING_AI_CHAT_MEMORY WHERE conversation_id = ? ORDER BY `timestamp` DESC LIMIT ?";
 	}
 
 	@Override

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

@@ -26,7 +26,7 @@ public class PostgresChatMemoryRepositoryDialect implements JdbcChatMemoryReposi
 
 	@Override
 	public String getSelectMessagesSql() {
-		return "SELECT content, type FROM SPRING_AI_CHAT_MEMORY WHERE conversation_id = ? ORDER BY \"timestamp\"";
+		return "SELECT content, type, \"timestamp\" FROM SPRING_AI_CHAT_MEMORY WHERE conversation_id = ? ORDER BY \"timestamp\"";
 	}
 
 	@Override

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

@@ -26,7 +26,7 @@ public class SqlServerChatMemoryRepositoryDialect implements JdbcChatMemoryRepos
 
 	@Override
 	public String getSelectMessagesSql() {
-		return "SELECT content, type FROM SPRING_AI_CHAT_MEMORY WHERE conversation_id = ? ORDER BY [timestamp] ASC";
+		return "SELECT content, type, [timestamp] FROM SPRING_AI_CHAT_MEMORY WHERE conversation_id = ? ORDER BY [timestamp] ASC";
 	}
 
 	@Override