聊聊Spring AI Alibaba的JdbcChatMemory

序

本文主要研究一下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);
    }
  }

  @Override
  public 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);
    }
  }

  @Override
  public List<Message> get(String conversationId, int lastN) {
    return this.selectMessageById(conversationId, lastN);
  }

  @Override
  public 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);
    }

  }

  @Override
  public 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 first
        String 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 approach
          all = all.stream().skip(Math.max(0, deleteSize)).toList();
          // Clear all messages and reinsert the remaining ones
          clear(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);
  }

  @Override
  protected String jdbcType() {
    return JDBC_TYPE;
  }

  @Override
  protected String hasTableSql(String tableName) {
    return String.format("SHOW TABLES LIKE '%s'", tableName);
  }

  @Override
  protected 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);
  }

  @Override
  protected String jdbcType() {
    return JDBC_TYPE;
  }

  @Override
  protected String hasTableSql(String tableName) {
    return String.format("SELECT table_name FROM user_tables WHERE table_name = '%s'", tableName.toUpperCase());
  }

  @Override
  protected 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);
  }

  @Override
  protected 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);
  }

  @Override
  protected String jdbcType() {
    return JDBC_TYPE;
  }

  @Override
  protected 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);
  }

  @Override
  protected 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);
  }

  @Override
  protected String jdbcType() {
    return JDBC_TYPE;
  }

  @Override
  protected String hasTableSql(String tableName) {
    return String.format("SELECT name FROM sqlite_master WHERE type = 'table' AND name LIKE '%s'", tableName);
  }

  @Override
  protected 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);
  }

  @Override
  protected String jdbcType() {
    return JDBC_TYPE;
  }

  @Override
  protected String hasTableSql(String tableName) {
    return String.format("SELECT name FROM sys.tables WHERE name LIKE '%s';", tableName);
  }

  @Override
  protected 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);
  }

  @Override
  protected 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 {

  @Test
  public void test() {
    MysqlChatMemory mock = mock(MysqlChatMemory.class);
    Assertions.assertNotNull(mock);
  }

  // @Test
  public 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