AI 超级智能体 - 基于mysql的消息存储

▼
java
复制代码
package com.cestbon.aiagent.chatmemory;


import cn.hutool.json.JSONUtil;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.cestbon.aiagent.entity.ChatConversation;
import com.cestbon.aiagent.entity.ChatMessage;
import com.cestbon.aiagent.mapper.ChatConversationMapper;
import com.cestbon.aiagent.mapper.ChatMessageMapper;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.memory.ChatMemoryRepository;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.util.Assert;

import java.util.Date;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;

@Component
public class InMySqlChatMemoryRepository implements ChatMemoryRepository {


    @Resource
    private ChatConversationMapper chatConversationMapper;

    @Resource
    private ChatMessageMapper chatMessageMapper;


    @Override
    public List<String> findConversationIds() {
        return chatConversationMapper.selectList(null)
                .stream()
                .map(ChatConversation::getConversationId)
                .collect(Collectors.toList());
    }

    @Override
    public List<Message> findByConversationId(String conversationId) {
        Assert.hasText(conversationId, "conversationId must not be empty");
        QueryWrapper<ChatMessage> queryWrapper = new QueryWrapper<>();
        queryWrapper.eq("conversation_id", conversationId).orderByAsc("message_index")
                .in("role", "USER", "ASSISTANT");
        List<ChatMessage> records = chatMessageMapper.selectList(queryWrapper);
        if (records == null || records.isEmpty()) {
            return List.of();
        }
        return records.stream().map(this::convertToMessage)
                .filter(Objects::nonNull)
                .collect(Collectors.toList());
    }

    private Message convertToMessage(ChatMessage chatMessage) {
        String role = chatMessage.getRole();
        String content = chatMessage.getContent();
        if ("USER".equals(role)) {
            return new UserMessage(content);
        } else if ("ASSISTANT".equals(role)) {
            return new AssistantMessage(content);
        }
        return null;
    }

    @Override
    @Transactional
    public void saveAll(String conversationId, List<Message> messages) {
        Assert.hasText(conversationId, "conversationId must not be empty");
        Assert.notNull(messages, "messages must not be null");

        // 查询当前对话，检查是否已存在
        QueryWrapper<ChatConversation> conversationWrapper = new QueryWrapper<>();
        conversationWrapper.eq("conversation_id", conversationId);
        ChatConversation existingConversation = chatConversationMapper.selectOne(conversationWrapper);

        // 如果对话不存在，则新建
        if (existingConversation == null) {
            ChatConversation newConversation = new ChatConversation();
            newConversation.setConversationId(conversationId);
            newConversation.setTitle(messages.getFirst().getText());
            chatConversationMapper.insert(newConversation);
        } else {
            // 更新对话的updateTime
            existingConversation.setUpdateTime(new Date());
            chatConversationMapper.updateById(existingConversation);
        }
        // TODO 当会话消息超过当前窗口时，应该调用llm总结当前会话的内容为一条消息，然后再维护当前窗口，而不是直接删除清空窗口
        chatMessageMapper.delete(new QueryWrapper<ChatMessage>().eq("conversation_id", conversationId));
        for (int i = 0; i < messages.size(); i++) {
            Message message = messages.get(i);
            ChatMessage chatMessage = new ChatMessage();
            chatMessage.setConversationId(conversationId);
            chatMessage.setMessageIndex(i);
            chatMessage.setContent(message.getText());
            chatMessage.setRole(message.getMessageType().name());
            chatMessage.setMetadata(JSONUtil.toJsonStr(message.getMetadata()));
            chatMessage.setCreateTime(new Date());
            chatMessageMapper.insert(chatMessage);
        }
    }

    @Override
    public void deleteByConversationId(String conversationId) {
        Assert.hasText(conversationId, "conversationId must not be empty");
        chatMessageMapper.delete(new QueryWrapper<ChatMessage>().eq("conversation_id", conversationId));
        chatConversationMapper.delete(new QueryWrapper<ChatConversation>().eq("conversation_id", conversationId));
    }
}