Spring AI工具调用回调与流式前端展示实现方案

1. Spring AI 工具调用回调与流式前端展示的完整落地方案

最近在重构一个零代码生成项目时，遇到了一个棘手的问题：如何在 Spring AI 中实现工具调用的回调机制，并将这些调用信息实时展示给前端。这让我不得不深入研究了 Spring AI 的工具调用机制，并最终找到了一套完整的解决方案。

如果你也在使用 Spring AI 开发 AI 应用，特别是需要在前端展示工具调用过程的应用，这篇文章将为你提供一套可直接落地的方案。我会从问题背景、设计思路到具体实现，一步步带你了解整个过程。

2. 问题背景与需求分析

2.1 为什么需要工具调用回调

在 AI 应用中，工具调用（Tool Calling）是一个非常重要的功能。它允许 AI 模型在执行过程中调用外部工具或函数来完成特定任务。比如：

• 读写文件
• 查询数据库
• 调用外部 API
• 执行特定业务逻辑

然而，Spring AI 目前（截至 2024 年）的工具调用机制相比 Langchain4j 还比较基础，缺少一些关键功能：

1. 缺少回调机制：无法在工具调用前后执行自定义逻辑
2. 缺少流式通知：无法将工具调用过程实时通知给前端
3. 缺少上下文管理：难以获取和管理会话 ID（conversationId）

2.2 具体业务需求

在我们的零代码生成项目中，我们需要：

1. 区分不同应用的生成目录：通过 conversationId 隔离每个应用的生成路径
2. 记录工具调用次数：用于后续分析和优化
3. 实时展示工具调用过程：让用户看到 AI 正在执行哪些操作

这些需求在 Langchain4j 中可以通过 StreamingChatResponseHandler 和 @ToolMemoryId 轻松实现，但在 Spring AI 中需要我们自己搭建这套机制。

3. 整体架构设计

3.1 核心思路

我们的解决方案基于以下核心思想：

1. AOP 切面拦截：通过 Spring AOP 在工具调用前后插入自定义逻辑
2. 事件发布机制：使用 Reactor 的 Sinks 实现事件发布/订阅
3. 流合并：将 AI 响应流和工具调用事件流合并返回给前端

3.2 架构流程图

plaintext复制用户发请求 → Ai2ChatClient 接收 → SpringAI 处理 → 切面拦截工具调用 → 事件发布 → 实时推给前端

3.3 技术选型

1. Spring AOP：用于拦截工具方法调用
2. Project Reactor：实现事件流和响应流
3. Caffeine Cache：用于临时存储工具调用状态

4. 核心实现细节

4.1 工具类实现示例

我们先来看一个实际的工具类实现 - TodoList 工具：

java复制@Component
public class TodolistTools extends BaseTools {
    private static final Cache<String, String> TODOLIST_CACHE = Caffeine.newBuilder()
            .maximumSize(10_00)
            .expireAfterWrite(Duration.ofMinutes(30))
            .build();

    @Tool(description = "Write or update the todo list for current task.")
    public String todoWrite(
        @ToolParam(description = "The todo list content to save.")
        String todoContent,
        ToolContext toolContext
    ) {
        String conversationId = ConversationIdUtils.getConversationId(toolContext);
        if (StringUtils.isBlank(todoContent)) {
            TODOLIST_CACHE.invalidate(conversationId);
            return "Todo list cleared.";
        }
        TODOLIST_CACHE.put(conversationId, todoContent);
        return "Todo list saved successfully.";
    }

    @Tool(description = "Read the current todo list for this conversation.")
    public String todoRead(ToolContext toolContext) {
        String conversationId = ConversationIdUtils.getConversationId(toolContext);
        String todoContent = TODOLIST_CACHE.getIfPresent(conversationId);
        if (StringUtils.isBlank(todoContent)) {
            return "No todo list for this conversation.";
        }
        return "Current todo list:\n" + todoContent;
    }

    @Override
    String getToolName() { return "Todo List Tool"; }
    
    @Override
    String getToolDes() { return "Read and write task todo lists"; }
}

关键点：

1. 每个工具方法都接收 ToolContext 参数，用于获取会话信息
2. 使用 Caffeine Cache 存储临时数据，按 conversationId 隔离
3. 通过 @Tool 和 @ToolParam 注解声明工具及其参数

4.2 AOP 切面实现

Spring AI 没有提供内置的回调机制，我们通过 AOP 自己实现：

java复制@Aspect
@Component
@Slf4j
public class ToolContextAspect {
    private final ToolEventPublisher toolEventPublisher;
    
    public ToolContextAspect(ToolEventPublisher toolEventPublisher) {
        this.toolEventPublisher = toolEventPublisher;
    }
    
    @Pointcut("execution(* com.leikooo.codemother.ai.tools..*.*(..)) && @annotation(org.springframework.ai.tool.annotation.Tool)")
    public void anyToolExecution() {}
    
    @Before("anyToolExecution()")
    public void beforeToolCall(JoinPoint joinPoint) {
        ToolContext toolContext = getToolContext(joinPoint);
        if (toolContext == null) return;
        handleToolContext(toolContext, joinPoint, null, true);
    }
    
    @AfterReturning(pointcut = "anyToolExecution()", returning = "result")
    public void afterToolCall(JoinPoint joinPoint, Object result) {
        ToolContext toolContext = getToolContext(joinPoint);
        if (toolContext != null) {
            handleToolContext(toolContext, joinPoint, result, false);
        }
    }
    
    private void handleToolContext(ToolContext context, JoinPoint joinPoint, Object result, boolean isBefore) {
        String className = joinPoint.getTarget().getClass().getSimpleName();
        String methodName = joinPoint.getSignature().getName();
        Message message = context.getToolCallHistory().getLast();
        AssistantMessage.ToolCall toolCallInfo = ((AssistantMessage) message).getToolCalls().getLast();
        String toolCallId = toolCallInfo.id();
        String sessionId = ConversationIdUtils.getConversationId(context);
        
        if (isBefore) {
            toolEventPublisher.publishToolCall(sessionId, className, methodName, toolCallId);
        } else {
            toolEventPublisher.publishToolResult(sessionId, className, methodName, toolCallId, result);
        }
    }
    
    private ToolContext getToolContext(JoinPoint joinPoint) {
        for (Object arg : joinPoint.getArgs()) {
            if (arg instanceof ToolContext) 
                return (ToolContext) arg;
        }
        return null;
    }
}

关键点：

1. 使用 @Before 和 @AfterReturning 分别在工具调用前后触发事件
2. 从 ToolContext 中提取会话 ID 和工具调用信息
3. 通过 ToolEventPublisher 发布事件

4.3 事件发布机制

我们使用 Reactor 的 Sinks 实现事件发布：

java复制@Component
public class ToolEventPublisher {
    private final Map<String, Sinks.Many<ToolEvent>> sinks = new ConcurrentHashMap<>();
    
    private Sinks.Many<ToolEvent> getSink(String sessionId) {
        return sinks.computeIfAbsent(sessionId, k -> Sinks.many().multicast().onBackpressureBuffer());
    }
    
    public void publishToolCall(String sessionId, String toolName, String methodName, String toolCallId) {
        getSink(sessionId).tryEmitNext(new ToolEvent(sessionId, "tool_call", toolName, methodName, toolCallId, null));
    }
    
    public void publishToolResult(String sessionId, String toolName, String methodName, String toolCallId, Object result) {
        getSink(sessionId).tryEmitNext(new ToolEvent(sessionId, "tool_result", toolName, methodName, toolCallId, result));
    }
    
    public Flux<ToolEvent> events(String sessionId) {
        return getSink(sessionId).asFlux();
    }
    
    public void complete(String sessionId) {
        Sinks.Many<ToolEvent> sink = sinks.remove(sessionId);
        if (sink != null) sink.tryEmitComplete();
    }
    
    public record ToolEvent(String sessionId, String type, String toolName, String methodName, String toolCallId, Object result) {}
}

关键点：

1. 每个会话（sessionId）对应一个独立的 Sink
2. 支持发布工具调用开始和结束事件
3. 提供事件流的订阅接口

5. 流式响应合并方案

5.1 方案一：手动合并流

java复制@Component
public class Ai2ChatClient {
    private final ChatClient chatClient;
    private final ToolEventPublisher toolEventPublisher;
    
    public Ai2ChatClient(ChatModel openAiChatModel, FileTools fileTools,
                         ToolEventPublisher toolEventPublisher) {
        this.toolEventPublisher = toolEventPublisher;
        this.chatClient = ChatClient.builder(openAiChatModel)
                .defaultTools(fileTools)
                .build();
    }
    
    public Flux<String> chat2Ai(String msg, String appId) {
        Flux<String> mainFlux = chatClient.prompt()
                .system("""
                        You are a helpful, precise, and reliable AI assistant.
                        Respond clearly and concisely.
                        Prioritize correctness, safety, and practicality.
                        If information is uncertain, state the uncertainty explicitly.
                        """)
                .user(msg)
                .advisors(spec -> spec.param(CONVERSATION_ID, appId))
                .toolContext(Map.of(CONVERSATION_ID, appId))
                .stream().content()
                .doFinally(s -> toolEventPublisher.complete(appId));
        
        Flux<String> toolEventFlux = toolEventPublisher.events(appId)
                .map(event -> {
                    Object result = Optional.ofNullable(event.result()).orElse("");
                    String message = switch (event.type()) {
                        case "tool_call" -> String.format("正在进行工具调用: %s", event.methodName());
                        case "tool_result" -> String.format("工具调用完成: %s", event.methodName());
                        default -> "";
                    };
                    return String.format("\n\n[选择工具] %s \n\n", message);
                });
        
        return Flux.merge(mainFlux, toolEventFlux);
    }
}

注意事项：

1. 必须在 mainFlux 的 doFinally 中调用 complete，否则事件流会一直挂起
2. 两个流的合并顺序会影响前端展示的顺序
3. 需要处理事件类型，生成友好的提示信息

5.2 方案二：使用 StreamAdvisor

更优雅的方案是实现 StreamAdvisor：

java复制@Slf4j
@Component
public class ToolAdvisor implements CallAdvisor, StreamAdvisor {
    private final ToolEventPublisher toolEventPublisher;
    
    @Override
    public ChatClientResponse adviseCall(ChatClientRequest request, CallAdvisorChain chain) {
        return chain.nextCall(request);
    }
    
    @Override
    public Flux<ChatClientResponse> adviseStream(ChatClientRequest request, StreamAdvisorChain chain) {
        String appId = ConversationIdUtils.getConversationId(request.context());
        Flux<ChatClientResponse> mainFlux = chain.nextStream(request)
                .doFinally(s -> toolEventPublisher.complete(appId));
        Flux<ChatClientResponse> toolEventFlux = getToolEventFlux(appId);
        return Flux.merge(mainFlux, toolEventFlux);
    }
    
    private Flux<ChatClientResponse> getToolEventFlux(String sessionId) {
        return toolEventPublisher.events(sessionId)
                .map(event -> {
                    String message = switch (event.type()) {
                        case "tool_call" -> String.format("正在进行工具调用: %s", event.methodName());
                        case "tool_result" -> String.format("工具调用完成: %s", event.methodName());
                        default -> "";
                    };
                    AssistantMessage msg = new AssistantMessage(String.format("\n\n[选择工具] %s \n\n", message));
                    return ChatClientResponse.builder()
                            .chatResponse(ChatResponse.builder().generations(List.of(new Generation(msg))).build())
                            .build();
                });
    }
    
    @Override
    public String getName() { return "ToolAdvisor"; }
    
    @Override
    public int getOrder() { return Integer.MIN_VALUE + 100; }
}

然后在 ChatClient 中注册：

java复制@Component
public class Ai2ChatClient {
    private final ChatClient chatClient;
    
    public Ai2ChatClient(ChatModel openAiChatModel, FileTools fileTools, ToolAdvisor toolAdvisor) {
        this.chatClient = ChatClient.builder(openAiChatModel)
                .defaultTools(fileTools)
                .defaultAdvisors(toolAdvisor)
                .build();
    }
    
    public Flux<String> chat(String msg, String appId) {
        return chatClient.prompt()
                .system("""
                        You are a helpful, precise, and reliable AI assistant.
                        Respond clearly and concisely.
                        Prioritize correctness, safety, and practicality.
                        If information is uncertain, state the uncertainty explicitly.
                        """)
                .user(msg)
                .advisors(spec -> spec.param(CONVERSATION_ID, appId))
                .toolContext(Map.of(CONVERSATION_ID, appId))
                .stream().content();
    }
}

5.3 两种方案对比

推荐使用 StreamAdvisor 方案，它更符合 Spring AI 的设计理念，也更易于维护和扩展。

6. 关键问题解决方案

6.1 如何获取 conversationId

Spring AI 没有提供类似 Langchain4j 的 @ToolMemoryId 注解，我们需要自己实现：

java复制public class ConversationIdUtils {
    public static String getConversationId(ToolContext toolContext) {
        return toolContext.getContext().get(ChatMemory.CONVERSATION_ID).toString();
    }
}

在调用 ChatClient 时设置 conversationId：

java复制public Flux<String> chat2AiAdvisor(String msg, String appId) {
    return chatClient.prompt()
            .system("你是有用的小助手")
            .user(msg)
            .advisors(advisorSpec -> advisorSpec.param(CONVERSATION_ID, appId))
            .toolContext(Map.of(CONVERSATION_ID, appId))
            .stream().content();
}

6.2 工具调用监控

通过 AOP 切面，我们可以轻松实现工具调用监控：

1. 记录每个工具调用的开始和结束时间
2. 统计工具调用次数
3. 监控工具调用成功率
4. 收集工具调用性能数据

这些数据可以存储到数据库或推送到监控系统，用于后续分析和优化。

7. 测试与验证

7.1 测试用例

java复制@SpringBootTest
class Ai2ChatClientTest {
    @Resource
    private Ai2ChatClient ai2ChatClient;
    
    @Test
    void chat2Ai() throws InterruptedException {
        Flux<String> flux = ai2ChatClient.chat("帮我生成一个企业级别的后端，帮我生成 todolist", "12345");
        flux.doOnNext(System.out::println).subscribe();
        Thread.sleep(10000);
    }
}

7.2 预期输出

plaintext复制正在为您生成企业级后端代码...

[选择工具] 正在进行工具调用: todoWrite 

[选择工具] 工具调用完成: todoWrite 

生成完成！以下是您需要的企业级后端代码框架：
1. Spring Boot 基础结构
2. 数据库访问层
3. RESTful API 接口
...

7.3 实际效果验证

通过测试，我们可以验证：

1. 工具调用事件被正确捕获和发布
2. 事件流和主响应流正确合并
3. 前端能够实时显示工具调用状态
4. 会话 ID 正确传递和使用

8. 性能优化与注意事项

8.1 性能考量

1. 事件发布性能：Sinks 的并发性能足够好，但要注意背压处理
2. 内存占用：每个会话都会创建一个 Sink，需要及时清理
3. 网络传输：工具调用信息会增加响应体积，考虑压缩或精简

8.2 注意事项

1. 资源清理：务必在流结束时调用 complete() 清理资源
2. 错误处理：处理好工具调用失败的情况
3. 会话隔离：确保不同会话的事件不会互相干扰
4. 线程安全：工具类需要是线程安全的

8.3 常见问题排查

1. 工具调用事件未触发：

   • 检查 AOP 是否生效
   • 检查工具方法是否有 @Tool 注解
   • 检查 ToolContext 是否正确传递

2. 前端未收到事件通知：

   • 检查事件发布逻辑
   • 检查流合并是否正确
   • 检查前端订阅逻辑

3. 内存泄漏：

   • 检查是否调用了 complete()
   • 监控 sinks Map 的大小

9. 扩展与改进方向

9.1 功能扩展

1. 工具调用审批：在前端拦截特定工具调用，等待用户确认
2. 调用链追踪：记录完整的工具调用链，便于调试
3. 性能监控：收集并展示工具调用性能指标

9.2 性能改进

1. 批量事件发布：合并多个工具调用事件一起发布
2. 事件压缩：减少传输数据量
3. 智能流控：根据网络状况调整事件发布频率

9.3 架构优化

1. 分布式事件总线：支持多实例部署
2. 持久化事件存储：支持断线重连和历史查询
3. 更细粒度的权限控制：基于工具调用的权限管理

10. 总结与个人实践心得

这套方案在实际项目中运行良好，解决了 Spring AI 工具调用回调的痛点。在实现过程中，有几个关键点值得分享：

1. AOP 是强大的扩展点：通过切面可以在不修改原有代码的情况下添加新功能
2. Reactor 流处理要小心资源泄漏：务必处理好流的生命周期
3. 会话隔离很重要：特别是在多租户场景下
4. Advisor 模式更优雅：比手动合并流更符合 Spring AI 的设计哲学

在实际使用中，建议从简单方案开始，根据需要逐步扩展。同时，要做好监控和日志，便于排查问题。