Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -53,6 +55,8 @@ 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");
Expand Down Expand Up @@ -104,11 +108,13 @@ 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
Expand All @@ -124,15 +130,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);
};
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down