AI 超级智能体 - 历史消息存储Advisor

在实现自定义mysql存储消息时发现，ChatMemory和MessageWindowChatMemory都是只能维持当前窗口的消息，即maxMessage窗口内的消息。一但消息超过maxMessage，就会进行一次覆盖。若是在自定义mysql存储消息实行增量存储而不覆盖的话，llm的上下文记忆又会出错，因此我在自定义mysql存储消息时也和内存和文件存储一样使用了覆盖并实现了这个历史消息存储Advisor，在消息发送和llm输出时实现拦截并将消息存储到mysql中，存储整个会话的所有消息。

▼
text
复制代码
@Resource
private HistoryMessageMapper historyMessageMapper;

private int order = 1;

public MessageToDBAdvisor() {
}

@Override
public ChatClientResponse adviseCall(ChatClientRequest chatClientRequest, CallAdvisorChain callAdvisorChain) {
    HistoryMessage userHistoryMessage = before(chatClientRequest);
    ChatClientResponse chatClientResponse = callAdvisorChain.nextCall(chatClientRequest);
    after(chatClientResponse, chatClientRequest, userHistoryMessage);
    return chatClientResponse;
}

@Override
public Flux<ChatClientResponse> adviseStream(ChatClientRequest chatClientRequest, StreamAdvisorChain streamAdvisorChain) {
    HistoryMessage userHistoryMessage = before(chatClientRequest);
    Flux<ChatClientResponse> chatClientResponseFlux = streamAdvisorChain.nextStream(chatClientRequest);
    return afterFlux(chatClientResponseFlux, chatClientRequest, userHistoryMessage);
}

private Flux<ChatClientResponse> afterFlux(Flux<ChatClientResponse> chatClientResponseFlux, ChatClientRequest chatClientRequest, HistoryMessage userHistoryMessage) {
    String conversationId = chatClientRequest.context().get("chat_memory_conversation_id").toString();
    String model = chatClientRequest.prompt().getOptions().getModel();
    StringBuilder contentBuilder = new StringBuilder();
    AtomicInteger promptTokens = new AtomicInteger(0);
    AtomicInteger completionTokens = new AtomicInteger(0);
    Flux<ChatClientResponse> responseFlux = chatClientResponseFlux.doOnNext(chatClientResponse -> {
        String chunk = chatClientResponse.chatResponse().getResult().getOutput().getText();
        if (StrUtil.isNotBlank(chunk)) {
            contentBuilder.append(chunk);
        }
        promptTokens.set(chatClientResponse.chatResponse().getMetadata().getUsage().getPromptTokens());
        completionTokens.set(chatClientResponse.chatResponse().getMetadata().getUsage().getCompletionTokens());
    }).doOnComplete(() -> {
        String content = contentBuilder.toString();
        HistoryMessage assistantHistoryMessage = new HistoryMessage();
        assistantHistoryMessage.setConversationId(conversationId);
        assistantHistoryMessage.setRole(MessageType.ASSISTANT.getValue());
        assistantHistoryMessage.setModel(model);
        assistantHistoryMessage.setContent(content);
        assistantHistoryMessage.setTokenCount(completionTokens.get());
        historyMessageMapper.insert(assistantHistoryMessage);
        userHistoryMessage.setTokenCount(promptTokens.get());
        historyMessageMapper.updateById(userHistoryMessage);
    });
    return responseFlux;
}

public HistoryMessage before(ChatClientRequest chatClientRequest) {
    String conversationId = chatClientRequest.context().get("chat_memory_conversation_id").toString();
    String model = chatClientRequest.prompt().getOptions().getModel();
    String content = chatClientRequest.prompt().getUserMessage().getText();
    HistoryMessage userHistoryMessage = new HistoryMessage();
    userHistoryMessage.setConversationId(conversationId);
    userHistoryMessage.setModel(model);
    userHistoryMessage.setRole(MessageType.USER.getValue());
    userHistoryMessage.setContent(content);
    historyMessageMapper.insert(userHistoryMessage);
    return userHistoryMessage;
}

public void after(ChatClientResponse chatClientResponse, ChatClientRequest chatClientRequest, HistoryMessage userHistoryMessage) {
    String conversationId = chatClientRequest.context().get("chat_memory_conversation_id").toString();
    String model = chatClientRequest.prompt().getOptions().getModel();
    String content = chatClientResponse.chatResponse().getResult().getOutput().getText();
    Integer promptTokens = chatClientResponse.chatResponse().getMetadata().getUsage().getPromptTokens();
    Integer completionTokens = chatClientResponse.chatResponse().getMetadata().getUsage().getCompletionTokens();
    HistoryMessage assistantHistoryMessage = new HistoryMessage();
    assistantHistoryMessage.setConversationId(conversationId);
    assistantHistoryMessage.setRole(MessageType.ASSISTANT.getValue());
    assistantHistoryMessage.setModel(model);
    assistantHistoryMessage.setContent(content);
    assistantHistoryMessage.setTokenCount(completionTokens);
    userHistoryMessage.setTokenCount(promptTokens);
    historyMessageMapper.updateById(userHistoryMessage);
    historyMessageMapper.insert(assistantHistoryMessage);
}


@Override
public int getOrder() {
    return this.order;
}

@Override
public String getName() {
    return this.getClass().getSimpleName();
}