聊聊Spring AI Alibaba的RedisChatMemory

序

本文主要研究一下Spring AI Alibaba的RedisChatMemory

RedisChatMemory

community/memories/spring-ai-alibaba-redis-memory/src/main/java/com/alibaba/cloud/ai/memory/redis/RedisChatMemory.java

public class RedisChatMemory implements ChatMemory, AutoCloseable {

  private static final Logger logger = LoggerFactory.getLogger(RedisChatMemory.class);

  private static final String DEFAULT_KEY_PREFIX = "spring_ai_alibaba_chat_memory";

  private static final String DEFAULT_HOST = "127.0.0.1";

  private static final int DEFAULT_PORT = 6379;

  private static final String DEFAULT_PASSWORD = null;

  private final JedisPool jedisPool;

  private final Jedis jedis;

  private final ObjectMapper objectMapper;

  public RedisChatMemory() {

    this(DEFAULT_HOST, DEFAULT_PORT, DEFAULT_PASSWORD);
  }

  public RedisChatMemory(String host, int port, String password) {

    JedisPoolConfig poolConfig = new JedisPoolConfig();

    this.jedisPool = new JedisPool(poolConfig, host, port, 2000, password);
    this.jedis = jedisPool.getResource();
    this.objectMapper = new ObjectMapper();
    SimpleModule module = new SimpleModule();
    module.addDeserializer(Message.class, new MessageDeserializer());
    this.objectMapper.registerModule(module);

    logger.info("Connected to Redis at {}:{}", host, port);
  }

  @Override
  public void add(String conversationId, List<Message> messages) {

    String key = DEFAULT_KEY_PREFIX + conversationId;

    for (Message message : messages) {
      try {
        String messageJson = objectMapper.writeValueAsString(message);
        jedis.rpush(key, messageJson);
      }
      catch (JsonProcessingException e) {
        throw new RuntimeException("Error serializing message", e);
      }
    }

    logger.info("Added messages to conversationId: {}", conversationId);
  }

  @Override
  public List<Message> get(String conversationId, int lastN) {

    String key = DEFAULT_KEY_PREFIX + conversationId;

    List<String> messageStrings = jedis.lrange(key, -lastN, -1);
    List<Message> messages = new ArrayList<>();

    for (String messageString : messageStrings) {
      try {
        Message message = objectMapper.readValue(messageString, Message.class);
        messages.add(message);
      }
      catch (JsonProcessingException e) {
        throw new RuntimeException("Error deserializing message", e);
      }
    }

    logger.info("Retrieved {} messages for conversationId: {}", messages.size(), conversationId);

    return messages;
  }

  @Override
  public void clear(String conversationId) {

    String key = DEFAULT_KEY_PREFIX + conversationId;

    jedis.del(key);
    logger.info("Cleared messages for conversationId: {}", conversationId);
  }

  @Override
  public void close() {

    if (jedis != null) {

      jedis.close();

      logger.info("Redis connection closed.");
    }
    if (jedisPool != null) {

      jedisPool.close();

      logger.info("Jedis pool closed.");
    }
  }

  public void clearOverLimit(String conversationId, int maxLimit, int deleteSize) {
    try {
      String key = DEFAULT_KEY_PREFIX + conversationId;

      List<String> all = jedis.lrange(key, 0, -1);

      if (all.size() >= maxLimit) {
        all = all.stream().skip(Math.max(0, deleteSize)).toList();
      }
      this.clear(conversationId);
      for (String message : all) {
        jedis.rpush(key, message);
      }
    }
    catch (Exception e) {
      logger.error("Error clearing messages from Redis chat memory", e);
      throw new RuntimeException(e);
    }
  }

  public void updateMessageById(String conversationId, String messages) {
    String key = "spring_ai_alibaba_chat_memory:" + conversationId;
    try {
      this.jedis.del(key);
      this.jedis.rpush(key, new String[] { messages });
    }
    catch (Exception var6) {
      logger.error("Error updating messages from Redis chat memory", var6);
      throw new RuntimeException(var6);
    }
  }

}

RedisChatMemory的构造器初始化了JedisPool并给ObjectMapper注册了org.springframework.ai.chat.messages.Message类型的MessageDeserializer；其add方法遍历messages挨个序列化为json然后rpush到spring_ai_alibaba_chat_memory{conversationId}中；其get方法通过lrange取出最近n条记录，再反序列化为message对象；其clear方法直接删除该key；close方法则先关闭jedis再关闭jedisPool

MessageDeserializer

community/memories/spring-ai-alibaba-redis-memory/src/main/java/com/alibaba/cloud/ai/memory/redis/serializer/MessageDeserializer.java

public class MessageDeserializer extends JsonDeserializer<Message> {

  private static final Logger logger = LoggerFactory.getLogger(MessageDeserializer.class);

  public Message deserialize(JsonParser p, DeserializationContext ctxt) {
    ObjectMapper mapper = (ObjectMapper) p.getCodec();
    JsonNode node = null;
    Message message = null;
    try {
      node = mapper.readTree(p);
      String messageType = node.get("messageType").asText();
      switch (messageType) {
        case "USER" -> message = new UserMessage(node.get("text").asText(),
            mapper.convertValue(node.get("media"), new TypeReference<Collection<Media>>() {
            }), mapper.convertValue(node.get("metadata"), new TypeReference<Map<String, Object>>() {
            }));
        case "ASSISTANT" -> message = new AssistantMessage(node.get("text").asText(),
            mapper.convertValue(node.get("metadata"), new TypeReference<Map<String, Object>>() {
            }), (List<AssistantMessage.ToolCall>) mapper.convertValue(node.get("toolCalls"),
                new TypeReference<Collection<AssistantMessage.ToolCall>>() {
                }),
            (List<Media>) mapper.convertValue(node.get("media"), new TypeReference<Collection<Media>>() {
            }));
        default -> throw new IllegalArgumentException("Unknown message type: " + messageType);
      }
      ;
    }
    catch (IOException e) {
      logger.error("Error deserializing message", e);
    }
    return message;
  }

}

MessageDeserializer继承了JsonDeserializer，它读取messageType字段，然后对于USER类型创建UserMessage、对于ASSISTANT类型创建AssistantMessage

小结

spring-ai-alibaba-redis-memory提供了ChatMemory的redis实现，它通过jedis使用rpush添加message，通过lrange取出最近N条，通过del删除指定会话的消息。

doc

• java2ai