序
本文主要研究一下Spring AI Alibaba的JdbcChatMemory
JdbcChatMemory
community/memories/spring-ai-alibaba-jdbc-memory/src/main/java/com/alibaba/cloud/ai/memory/jdbc/JdbcChatMemory.java
public abstract class JdbcChatMemory implements ChatMemory, AutoCloseable {private static final Logger logger = LoggerFactory.getLogger(JdbcChatMemory.class);private static final String DEFAULT_TABLE_NAME = "chat_memory";private final Connection connection;private final String tableName;protected JdbcChatMemory(String username, String password, String jdbcUrl) {this(username, password, jdbcUrl, DEFAULT_TABLE_NAME);}protected JdbcChatMemory(String username, String password, String jdbcUrl, String tableName) {this.tableName = tableName;try {this.connection = DriverManager.getConnection(jdbcUrl, username, password);checkAndCreateTable();}catch (SQLException e) {throw new RuntimeException("Error connecting to the database", e);}}protected JdbcChatMemory(Connection connection) {this(connection, DEFAULT_TABLE_NAME);}protected JdbcChatMemory(Connection connection, String tableName) {this.connection = connection;this.tableName = tableName;try {checkAndCreateTable();}catch (SQLException e) {throw new RuntimeException("Error checking the database table", e);}}protected abstract String jdbcType();protected abstract String hasTableSql(String tableName);protected abstract String createTableSql(String tableName);/*** Generate paginated query SQL based on the database type. Default implementation* uses LIMIT clause. Subclasses can override this method if their database uses* different pagination syntax.* @param tableName The name of the table* @param lastN Number of records to return, if greater than 0* @return SQL query string with pagination*/protected String generatePaginatedQuerySql(String tableName, int lastN) {StringBuilder sqlBuilder = new StringBuilder("SELECT messages,type FROM ").append(tableName).append(" WHERE conversation_id = ?");if (lastN > 0) {sqlBuilder.append(" LIMIT ?");}return sqlBuilder.toString();}private void checkAndCreateTable() throws SQLException {String checkTableQuery = hasTableSql(tableName);try (Statement stmt = connection.createStatement(); ResultSet rs = stmt.executeQuery(checkTableQuery)) {if (rs.next()) {logger.info("Table {} exists.", tableName);}else {logger.info("Table {} does not exist. Creating table...", tableName);createTable();}}}private void createTable() {try (Statement stmt = connection.createStatement()) {stmt.execute(createTableSql(tableName));logger.info("Table {} created successfully.", tableName);}catch (Exception e) {throw new RuntimeException("Error creating table " + tableName + " ", e);}}@Overridepublic void add(String conversationId, List<Message> messages) {try {for (Message message : messages) {String sql = "INSERT INTO " + tableName + " (messages, conversation_id, type) VALUES (?, ?, ?)";try (PreparedStatement stmt = this.connection.prepareStatement(sql)) {stmt.setString(1, message.getText());stmt.setString(2, conversationId);stmt.setString(3, message.getMessageType().name());stmt.executeUpdate();}}}catch (Exception e) {logger.error("Error adding messages to {} chat memory", jdbcType(), e);throw new RuntimeException(e);}}@Overridepublic List<Message> get(String conversationId, int lastN) {return this.selectMessageById(conversationId, lastN);}@Overridepublic void clear(String conversationId) {String sql = "DELETE FROM " + tableName + " WHERE conversation_id = ?";try (PreparedStatement stmt = connection.prepareStatement(sql)) {stmt.setString(1, conversationId);stmt.executeUpdate();}catch (Exception e) {logger.error("Error clearing messages from {} chat memory", jdbcType(), e);throw new RuntimeException("Error executing delete ", e);}}@Overridepublic void close() throws Exception {if (connection != null) {connection.close();}}public void clearOverLimit(String conversationId, int maxLimit, int deleteSize) {try {List<Message> all = this.selectMessageById(conversationId);if (all.size() >= maxLimit) {// Delete oldest messages firstString sql = "DELETE FROM " + tableName + " WHERE conversation_id = ? ORDER BY ROWID LIMIT ?";try (PreparedStatement stmt = this.connection.prepareStatement(sql)) {stmt.setString(1, conversationId);stmt.setInt(2, deleteSize);stmt.executeUpdate();}catch (SQLException e) {// If the database doesn't support ORDER BY in DELETE, fallback to// alternative approachall = all.stream().skip(Math.max(0, deleteSize)).toList();// Clear all messages and reinsert the remaining onesclear(conversationId);for (Message message : all) {add(conversationId, List.of(message));}}}}catch (Exception e) {logger.error("Error clearing messages from {} chat memory", jdbcType(), e);throw new RuntimeException(e);}}public List<Message> selectMessageById(String conversationId, int lastN) {List<Message> totalMessage = new ArrayList<>();String sql = generatePaginatedQuerySql(tableName, lastN);try (PreparedStatement stmt = connection.prepareStatement(sql)) {// Set conversation ID parameters.stmt.setString(1, conversationId);// If there is a limit, set the limit parameter.if (lastN > 0) {stmt.setInt(2, lastN);}ResultSet resultSet = stmt.executeQuery();while (resultSet.next()) {var content = resultSet.getString("messages");var type = MessageType.valueOf(resultSet.getString("type"));var message = switch (type) {case USER -> new UserMessage(content);case ASSISTANT -> new AssistantMessage(content);case SYSTEM -> new SystemMessage(content);default -> null;};totalMessage.add(message);}}catch (SQLException e) {logger.error("select message by {} error,sql:{}", jdbcType(), sql, e);throw new RuntimeException(e);}return totalMessage;}public List<Message> selectMessageById(String conversationId) {return this.selectMessageById(conversationId, 0);}}
JdbcChatMemory构造器要求输入String username, String password, String jdbcUrl, String tableName,默认表名为
chat_memory
,它定义了jdbcType、hasTableSql、createTableSql三个抽象方法;其add方法遍历messages,挨个使用PreparedStatement插入人数据;其get方法则查询指定conversation_id的最近几条记录;其clear方法则删除指定conversation_id的所有记录。
MysqlChatMemory
community/memories/spring-ai-alibaba-jdbc-memory/src/main/java/com/alibaba/cloud/ai/memory/jdbc/MysqlChatMemory.java
public class MysqlChatMemory extends JdbcChatMemory {private static final String JDBC_TYPE = "mysql";public MysqlChatMemory(String username, String password, String jdbcUrl) {super(username, password, jdbcUrl);}public MysqlChatMemory(String username, String password, String jdbcUrl, String tableName) {super(username, password, jdbcUrl, tableName);}public MysqlChatMemory(Connection connection) {super(connection);}public MysqlChatMemory(Connection connection, String tableName) {super(connection, tableName);}@Overrideprotected String jdbcType() {return JDBC_TYPE;}@Overrideprotected String hasTableSql(String tableName) {return String.format("SHOW TABLES LIKE '%s'", tableName);}@Overrideprotected String createTableSql(String tableName) {return String.format("CREATE TABLE %s( id BIGINT AUTO_INCREMENT PRIMARY KEY,conversation_id VARCHAR(256) NULL,messages TEXT NULL,type varchar(100) NULL) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;",tableName);}}
MysqlChatMemory继承了JdbcChatMemory,其jdbcType为mysql
OracleChatMemory
community/memories/spring-ai-alibaba-jdbc-memory/src/main/java/com/alibaba/cloud/ai/memory/jdbc/OracleChatMemory.java
public class OracleChatMemory extends JdbcChatMemory {private static final String JDBC_TYPE = "oracle";public OracleChatMemory(String username, String password, String url) {super(username, password, url);}public OracleChatMemory(String username, String password, String url, String tableName) {super(username, password, url, tableName);}public OracleChatMemory(Connection connection) {super(connection);}public OracleChatMemory(Connection connection, String tableName) {super(connection, tableName);}@Overrideprotected String jdbcType() {return JDBC_TYPE;}@Overrideprotected String hasTableSql(String tableName) {return String.format("SELECT table_name FROM user_tables WHERE table_name = '%s'", tableName.toUpperCase());}@Overrideprotected String createTableSql(String tableName) {return String.format("CREATE TABLE %s ( id NUMBER GENERATED BY DEFAULT AS IDENTITY PRIMARY KEY, conversation_id VARCHAR2(256), messages CLOB, type VARCHAR2(100));",tableName);}@Overrideprotected String generatePaginatedQuerySql(String tableName, int lastN) {StringBuilder sqlBuilder = new StringBuilder().append("SELECT * FROM (").append("SELECT messages,type FROM ").append(tableName).append(" WHERE conversation_id = ?");if (lastN > 0) {sqlBuilder.append(" AND ROWNUM <= ?");}sqlBuilder.append(")");return sqlBuilder.toString();}}
OracleChatMemory继承了JdbcChatMemory,其jdbcType为oracle,它还额外覆盖了generatePaginatedQuerySql方法
PostgresChatMemory
public class PostgresChatMemory extends JdbcChatMemory {private static final String JDBC_TYPE = "postgresql";public PostgresChatMemory(String username, String password, String url) {super(username, password, url);}public PostgresChatMemory(String username, String password, String jdbcUrl, String tableName) {super(username, password, jdbcUrl, tableName);}public PostgresChatMemory(Connection connection) {super(connection);}public PostgresChatMemory(Connection connection, String tableName) {super(connection, tableName);}@Overrideprotected String jdbcType() {return JDBC_TYPE;}@Overrideprotected String hasTableSql(String tableName) {return String.format("SELECT table_name FROM information_schema.tables WHERE table_schema = 'public' AND table_type = 'BASE TABLE' AND table_name LIKE '%s'",tableName);}@Overrideprotected String createTableSql(String tableName) {return String.format("CREATE TABLE %s ( id BIGSERIAL PRIMARY KEY, conversation_id VARCHAR(256), messages TEXT, type VARCHAR(100));",tableName);}}
PostgresChatMemory继承了JdbcChatMemory,其jdbcType为postgresql
SQLiteChatMemory
community/memories/spring-ai-alibaba-jdbc-memory/src/main/java/com/alibaba/cloud/ai/memory/jdbc/SQLiteChatMemory.java
public class SQLiteChatMemory extends JdbcChatMemory {private static final String JDBC_TYPE = "sqlite";public SQLiteChatMemory(String username, String password, String jdbcUrl) {super(username, password, jdbcUrl);}public SQLiteChatMemory(String username, String password, String jdbcUrl, String tableName) {super(username, password, jdbcUrl, tableName);}public SQLiteChatMemory(Connection connection) {super(connection);}public SQLiteChatMemory(Connection connection, String tableName) {super(connection, tableName);}@Overrideprotected String jdbcType() {return JDBC_TYPE;}@Overrideprotected String hasTableSql(String tableName) {return String.format("SELECT name FROM sqlite_master WHERE type = 'table' AND name LIKE '%s'", tableName);}@Overrideprotected String createTableSql(String tableName) {return String.format("CREATE TABLE IF NOT EXISTS %s ( id INTEGER PRIMARY KEY AUTOINCREMENT, conversation_id TEXT, messages TEXT, type TEXT);",tableName);}}
SQLiteChatMemory继承了JdbcChatMemory,其jdbcType为sqlite
SqlServerChatMemory
community/memories/spring-ai-alibaba-jdbc-memory/src/main/java/com/alibaba/cloud/ai/memory/jdbc/SqlServerChatMemory.java
public class SqlServerChatMemory extends JdbcChatMemory {private static final String JDBC_TYPE = "sqlserver";public SqlServerChatMemory(String username, String password, String url) {super(username, password, url);}public SqlServerChatMemory(String username, String password, String jdbcUrl, String tableName) {super(username, password, jdbcUrl, tableName);}public SqlServerChatMemory(Connection connection) {super(connection);}public SqlServerChatMemory(Connection connection, String tableName) {super(connection, tableName);}@Overrideprotected String jdbcType() {return JDBC_TYPE;}@Overrideprotected String hasTableSql(String tableName) {return String.format("SELECT name FROM sys.tables WHERE name LIKE '%s';", tableName);}@Overrideprotected String createTableSql(String tableName) {return String.format("CREATE TABLE %s (id BIGINT IDENTITY(1,1) PRIMARY KEY, conversation_id NVARCHAR(256), messages NVARCHAR(MAX), type VARCHAR(100));",tableName);}@Overrideprotected String generatePaginatedQuerySql(String tableName, int lastN) {StringBuilder sqlBuilder = new StringBuilder("SELECT ");if (lastN > 0) {sqlBuilder.append("TOP ").append(lastN).append(" ");}sqlBuilder.append("messages,type FROM ").append(tableName).append(" WHERE conversation_id = ?");return sqlBuilder.toString();}}
SqlServerChatMemory继承了JdbcChatMemory,其jdbcType为sqlserver,它还额外覆盖了generatePaginatedQuerySql方法
示例
community/memories/spring-ai-alibaba-jdbc-memory/src/test/java/com/alibaba/cloud/ai/memory/jdbc/MysqlChatMemoryTest.java
class MysqlChatMemoryTest {@Testpublic void test() {MysqlChatMemory mock = mock(MysqlChatMemory.class);Assertions.assertNotNull(mock);}// @Testpublic void mysql() {MysqlChatMemory chatMemory = new MysqlChatMemory("root", "123456","jdbc:mysql://127.0.0.1:3306/spring_ai_alibaba_chat_memory");String apiKey = System.getenv().getOrDefault("AI_DASHSCOPE_API_KEY", "test-api-key");ChatClient chatClient = ChatClient.create(new DashScopeChatModel(new DashScopeApi(apiKey)));String content1 = chatClient.prompt().advisors(new MessageChatMemoryAdvisor(chatMemory)).system("你是一个AI聊天小助手,给人提供情绪价值。").user("我是张三😄").call().content();System.out.println(content1);String content2 = chatClient.prompt().advisors(new MessageChatMemoryAdvisor(chatMemory)).user("我是谁").call().content();System.out.println(content2);Assertions.assertTrue(content2.contains("张三"));}
小结
spring-ai-alibaba-jdbc-memory提供了ChatMemory的jdbc实现,具体有MysqlChatMemory、OracleChatMemory、PostgresChatMemory、SQLiteChatMemory、SqlServerChatMemory这几种。
doc
- java2ai