隔离工具调用记录，解决 AI 偷懒不执行工具调用的问题

leikooo

2025-09-10 09:23

阅读1067

leikooo

2025-09-10 09:23

阅读 1067

背景

项目：AI 代码生成平台

场景：当 Redis 里面的记录过期之后，从 db 里面获取到数据，AI 学习到之前的工具调用的对话记录导致偷懒，没有真正调用工具只输出修改成功

解决方案

我个人的解决方案是把 tool_executed 工具执行结果，单独设置一个 messageType 比如 tool ，这样在 AiService 在没有 Redis 数据的时候获取 db 记录排除一下类型为 tool 的记录。

优化后的思路：

后端新添加一个类型 FRONTED 这个类型的数据是专门前端显示并且包含工具调用的记录

同时记录一个类型为 AI 的，这个类型的数据不包含工具调用的记录，此类型不返回给前端

工具调用记录结果类似（罪魁祸首）：

▼
c
复制代码
[工具调用] 读取目录 根目录

[工具调用] 读取目录 src

[工具调用] 读取文件 src/Home.vue

[工具调用] 读取目录 src/pages

[工具调用] 读取文件 src/pages/Home.vue

[工具调用] 修改文件 src/pages/Home.vue

替换前：

  <h1 class="hero-title">展示创意，连接世界</h1>


替换后：

  <h1 class="hero-title">连接世界</h1>

代码修改

最终修改后的完整的代码：https://github.com/lieeew/yu-ai-code-mother/tree/dev/leikooo

下面是具体修改的细节

1、修改 ChatHistoryMessageTypeEnum

▼
java
复制代码
/**
 * 对话历史消息枚举类型
 */
@Getter
public enum ChatHistoryMessageTypeEnum {

    USER("用户", "user"),
    // 方便前端显示, 包含工具调用的流程
    FRONTEND("前端显示", "frontend"),
    // 不包含工具调用的流程
    AI("AI", "ai");

    private final String text;

    private final String value;

    ChatHistoryMessageTypeEnum(String text, String value) {
        this.text = text;
        this.value = value;
    }

    /**
     * 根据 value 获取枚举
     *
     * @param value 枚举值的value
     * @return 枚举值
     */
    public static ChatHistoryMessageTypeEnum getEnumByValue(String value) {
        if (ObjUtil.isEmpty(value)) {
            return null;
        }
        for (ChatHistoryMessageTypeEnum anEnum : ChatHistoryMessageTypeEnum.values()) {
            if (anEnum.value.equals(value)) {
                return anEnum;
            }
        }
        return null;
    }
}

2、定义了一个内部类

方便在 originFlux 里面的 map 解析 JSON 类型，拿到对应的 content 信息和 type 类型，方便分类收集信息

▼
java
复制代码
@Data
@Builder
@NoArgsConstructor
@AllArgsConstructor
public static class AiResponse {
    private String content;

    /**
     * @see StreamMessageTypeEnum
     */
    private String type;
}

3、handleJsonMessageChunk 方法里面返回的是 JSON 字符串数据，方便调用方进行解析

▼
java
复制代码
/**
 * JSON 消息流处理器
 * 处理 VUE_PROJECT 类型的复杂流式响应，包含工具调用信息
 */
@Slf4j
@Component
public class JsonMessageStreamHandler {

    @Resource
    private ToolManager toolManager;

    /**
     * 处理 TokenStream（VUE_PROJECT）
     * 解析 JSON 消息并重组为完整的响应格式
     *
     * @param originFlux         原始流
     * @param chatHistoryService 聊天历史服务
     * @param appId              应用ID
     * @param loginUser          登录用户
     * @return 处理后的流
     */
     public Flux<String> handle(Flux<String> originFlux,
                                   ChatHistoryService chatHistoryService,
                                   long appId, User loginUser) {
            // 收集数据用于生成后端记忆格式
            StringBuilder aiHistoryStringBuilder = new StringBuilder();
            StringBuilder gatherHistoryStringBuilder = new StringBuilder();
            // 用于跟踪已经见过的工具ID，判断是否是第一次调用
            Set<String> seenToolIds = new HashSet<>();
            return originFlux
                .map(chunk -> {
                    // 解析每个 JSON 消息块
                    AiResponse aiResponse = JSONUtil.toBean(handleJsonMessageChunk(chunk, seenToolIds), AiResponse.class);
                    String content = aiResponse.content;
                    // 收集完整信息方便前端进行显示
                    gatherHistoryStringBuilder.append(content);
                    // 收集数据用于生成后端记忆格式
                    if (aiResponse.type.equals(StreamMessageTypeEnum.AI_RESPONSE.getValue())) {
                        aiHistoryStringBuilder.append(content);
                    }
                    return content;
                })
                // 过滤空字串
                .filter(StrUtil::isNotEmpty)
                .doOnComplete(() -> {
                    // ====== 不包含 tool 调用信息 ======
                    String aiResponse = aiHistoryStringBuilder.toString();
                    if (StringUtils.isNoneBlank(aiResponse)) {
                        chatHistoryService.addChatMessage(appId, aiResponse, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                    }
                    // ====== 前端显示包含 tool 调用信息 ======
                    String frontendHistory = gatherHistoryStringBuilder.toString();
                    if (StringUtils.isNoneBlank(frontendHistory)) {
                        chatHistoryService.addChatMessage(appId, frontendHistory, ChatHistoryMessageTypeEnum.FRONTEND.getValue(), loginUser.getId());
                    }                })
                .doOnError(error -> {
                    // 如果AI回复失败，也要记录错误消息
                    String errorMessage = "AI回复失败: " + error.getMessage();
                    chatHistoryService.addChatMessage(appId, errorMessage, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                });
    }
    /**
     * 解析并收集 TokenStream 数据
     */
    private String handleJsonMessageChunk(String chunk, Set<String> seenToolIds) {
        // 解析 JSON
        StreamMessage streamMessage = JSONUtil.toBean(chunk, StreamMessage.class);
        StreamMessageTypeEnum typeEnum = StreamMessageTypeEnum.getEnumByValue(streamMessage.getType());
        switch (typeEnum) {
            case AI_RESPONSE -> {
                AiResponseMessage aiMessage = JSONUtil.toBean(chunk, AiResponseMessage.class);
                String data = aiMessage.getData();
                return getAiMsgJson(data, typeEnum);
            }
            case TOOL_REQUEST -> {
                ToolRequestMessage toolRequestMessage = JSONUtil.toBean(chunk, ToolRequestMessage.class);
                String toolId = toolRequestMessage.getId();
                String toolName = toolRequestMessage.getName();
                // 检查是否是第一次看到这个工具 ID
                if (toolId != null && !seenToolIds.contains(toolId)) {
                    // 第一次调用这个工具，记录 ID 并完整返回工具信息
                    seenToolIds.add(toolId);
                    // 根据工具名称获取工具实例
                    BaseTool tool = toolManager.getTool(toolName);
                    // 返回格式化的工具调用信息
                    return getAiMsgJson(tool.generateToolRequestResponse(), typeEnum);
                } else {
                    // 不是第一次调用这个工具，直接返回空
                    return getAiMsgJson("", typeEnum);
                }
            }
            case TOOL_EXECUTED -> {
                ToolExecutedMessage toolExecutedMessage = JSONUtil.toBean(chunk, ToolExecutedMessage.class);
                JSONObject jsonObject = JSONUtil.parseObj(toolExecutedMessage.getArguments());
                // 根据工具名称获取工具实例
                String toolName = toolExecutedMessage.getName();
                BaseTool tool = toolManager.getTool(toolName);
                String result = tool.generateToolExecutedResult(jsonObject);
                // 输出前端和要持久化的内容
                String output = String.format("\n\n%s\n\n", result);
                return getAiMsgJson(output, typeEnum);
            }
            default -> {
                log.error("不支持的消息类型: {}", typeEnum);
                return getAiMsgJson("", typeEnum);
            }
        }
    }

    private String getAiMsgJson(String content, StreamMessageTypeEnum typeEnum) {
        AiResponse aiResponse = AiResponse.builder().content(content).type(typeEnum.getValue()).build();
        return JSONUtil.toJsonStr(aiResponse);
    }

    @Data
    @Builder
    @NoArgsConstructor
    @AllArgsConstructor
    public static class AiResponse {
        /**
         * 内容
         */
        private String content;

        /**
         * @see StreamMessageTypeEnum
         */
        private String type;
    }
}

4、修改 historyService 方法

FRONTED 类型包含 tool 调用的信息，前端显示这个类型的数据
USER 类型是用户发给前端的信息，前端 + AiService 使用这个类型的数据
AI 类型不包含 tool 调用的信息，AiService 加载会话使用这个类型的数据

ChatHistoryServiceImpl#listAppChatHistoryByPage

主要修改位置是查询的 queryWapper 属于 frontend、user 这两个类型

▼
java
复制代码
@Override
public Page<ChatHistory> listAppChatHistoryByPage(Long appId, int pageSize,
                                                  LocalDateTime lastCreateTime,
                                                  User loginUser) {
    ThrowUtils.throwIf(appId == null || appId <= 0, ErrorCode.PARAMS_ERROR, "应用ID不能为空");
    ThrowUtils.throwIf(pageSize <= 0 || pageSize > 50, ErrorCode.PARAMS_ERROR, "页面大小必须在1-50之间");
    ThrowUtils.throwIf(loginUser == null, ErrorCode.NOT_LOGIN_ERROR);
    // 验证权限：只有应用创建者和管理员可以查看
    App app = appService.getById(appId);
    ThrowUtils.throwIf(app == null, ErrorCode.NOT_FOUND_ERROR, "应用不存在");
    boolean isAdmin = UserConstant.ADMIN_ROLE.equals(loginUser.getUserRole());
    boolean isCreator = app.getUserId().equals(loginUser.getId());
    ThrowUtils.throwIf(!isAdmin && !isCreator, ErrorCode.NO_AUTH_ERROR, "无权查看该应用的对话历史");
    // 构建查询条件
    ChatHistoryQueryRequest queryRequest = new ChatHistoryQueryRequest();
    queryRequest.setAppId(appId);
    queryRequest.setLastCreateTime(lastCreateTime);
    QueryWrapper queryWrapper = this.getQueryWrapper(queryRequest);
    queryWrapper.in(ChatHistory::getMessageType,
            ChatHistoryMessageTypeEnum.FRONTEND.getValue(),
            ChatHistoryMessageTypeEnum.USER.getValue()
    );
    // 查询数据
    return this.page(Page.of(1, pageSize), queryWrapper);
}

5、ChatHistoryServiceImpl#loadChatHistoryToMemory

修改加载到 AiService 的记录，包含的 messageType 属于 ai、user 这两个类型

▼
java
复制代码
@Override
public int loadChatHistoryToMemory(Long appId, MessageWindowChatMemory chatMemory, int maxCount) {
    try {
        QueryWrapper queryWrapper = QueryWrapper.create()
                .eq(ChatHistory::getAppId, appId)
                .in(ChatHistory::getMessageType, 
                    ChatHistoryMessageTypeEnum.AI, 
                    ChatHistoryMessageTypeEnum.USER
                   )
                .orderBy(ChatHistory::getCreateTime, false)
                .limit(1, maxCount);
        List<ChatHistory> historyList = this.list(queryWrapper);
        if (CollUtil.isEmpty(historyList)) {
            return 0;
        }
        // 反转列表，确保按照时间正序（老的在前，新的在后）
        historyList = historyList.reversed();
        // 按照时间顺序将消息添加到记忆中
        int loadedCount = 0;
        // 先清理历史缓存，防止重复加载
        chatMemory.clear();
        for (ChatHistory history : historyList) {
            if (ChatHistoryMessageTypeEnum.USER.getValue().equals(history.getMessageType())) {
                chatMemory.add(UserMessage.from(history.getMessage()));
            } else if (ChatHistoryMessageTypeEnum.AI.getValue().equals(history.getMessageType())) {
                chatMemory.add(AiMessage.from(history.getMessage()));
            }
            loadedCount++;
        }
        log.info("成功为 appId: {} 加载 {} 条历史消息", appId, loadedCount);
        return loadedCount;
    } catch (Exception e) {
        log.error("加载历史对话失败，appId: {}, error: {}", appId, e.getMessage(), e);
        // 加载失败不影响系统运行，只是没有历史上下文
        return 0;
    }
}

6、SimpleTextStreamHandler 存储的格式也需要进行修改
需要保存两份，虽然内容一眼但是效果不一样

1）fronted 类型：前端显示

2）ai 类型：加载 Service 时的历史纪录

▼
java
复制代码
@Slf4j
public class SimpleTextStreamHandler {


    /**
     * 处理传统流（HTML, MULTI_FILE）
     * 直接收集完整的文本响应
     *
     * @param originFlux         原始流
     * @param chatHistoryService 聊天历史服务
     * @param appId              应用ID
     * @param loginUser          登录用户
     * @return 处理后的流
     */
     public Flux<String> handle(Flux<String> originFlux,
                               ChatHistoryService chatHistoryService,
                               long appId, User loginUser) {
        StringBuilder aiResponseBuilder = new StringBuilder();
        return originFlux
                .map(chunk -> {
                    // 收集AI响应内容
                    aiResponseBuilder.append(chunk);
                    return chunk;
                })
                .doOnComplete(() -> {
                    // 流式响应完成后，添加AI消息到对话历史
                    String aiResponse = aiResponseBuilder.toString();
                    chatHistoryService.addChatMessage(appId, aiResponse, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                    chatHistoryService.addChatMessage(appId, aiResponse, ChatHistoryMessageTypeEnum.FRONTEND.getValue(), loginUser.getId());
                })
                .doOnError(error -> {
                    // 如果AI回复失败，也要记录错误消息
                    String errorMessage = "AI回复失败: " + error.getMessage();
                    chatHistoryService.addChatMessage(appId, errorMessage, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                    chatHistoryService.addChatMessage(appId, errorMessage, ChatHistoryMessageTypeEnum.FRONTEND.getValue(), loginUser.getId());
                });
    }
}

最终测试

最终测试正常调用工具，并且前端显示也是正常

总结

通过把 tool 工具调用的记录隔离出来，避免 AI 学习导致不实际上调用工具，最终解决这个 BUG，并且前端不需要进行任何修改