记一次给 Spring AI 修 Bug 的全过程

1
@RestController
2
public class ChatController {
3

4
    private final ChatClient chatClient;
5

6
    public ChatController(ChatClient.Builder chatClientBuilder) {
7
        this.chatClient = chatClientBuilder.build();
8
    }
9

10
    @GetMapping("/chat")
11
    public Flux<String> chat(@RequestParam(value = "message") String message) {
12
        return chatClient
13
                .prompt()
14
                .user(message)
15
                .stream()
16
                .content()
17
                .doOnNext(System.out::println);
18
    }
19
}

1
<properties>
2
    <java.version>21</java.version>
3
    <spring-ai.version>2.0.0-M7</spring-ai.version>
4
</properties>
5
<dependencies>
6
    <dependency>
7
        <groupId>org.springframework.boot</groupId>
8
        <artifactId>spring-boot-starter-webmvc</artifactId>
9
    </dependency>
10

11
    <dependency>
12
        <groupId>org.springframework.ai</groupId>
13
        <artifactId>spring-ai-starter-model-openai</artifactId>
14
    </dependency>
15
</dependencies>

1
ChatClient.prompt().user(msg).stream().content()      ← Fluent API，返回 Flux<String>
2
        │
3
        ▼
4
DefaultStreamResponseSpec.content()                   ← 只是把 Flux<ChatResponse> 映射成文本
5
        │
6
        ▼
7
advisorChain.nextStream(chatClientRequest)            ← 进入顾问链（Reactive）
8
        │
9
        ▼
10
ChatModelStreamAdvisor.adviseStream(...)              ← 链路末端的「终结 Advisor」
11
        │   this.chatModel.stream(prompt)  ★ 真正调用模型的唯一入口
12
        ▼
13
OpenAiChatModel.stream(Prompt)                        ← StreamingChatModel 接口实现
14
        │
15
        ▼
16
OpenAiChatModel.internalStream(Prompt, prevResponse)  ← 真正订阅 OpenAI SDK 的地方
17
        │
18
        ▼
19
openAiClientAsync.chat().completions().createStreaming(request)

1
Flux<ChatResponse> chatResponses = Flux.<ChatResponse>create(sink -> {
2
    this.openAiClientAsync.chat().completions().createStreaming(request).subscribe(chunk -> {
3
        try {
4
            ChatCompletion chatCompletion = chunkToChatCompletion(chunk);
5

6
            List<Generation> generations = chatCompletion.choices().stream()
7
                    .map(choice -> buildGeneration(choice, metadata, request))
8
                    .toList();
9

10
            sink.next(new ChatResponse(generations, from(chatCompletion, accumulated)));
11
        }
12
        catch (Exception e) {
13
            sink.error(e);
14
        }
15
    }).onCompleteFuture().whenComplete((unused, throwable) -> {
16
        if (throwable != null) {
17
            sink.error(throwable);
18
        }
19
        else {
20
            sink.complete();
21
        }
22
    });
23
});

1
return flux.collectList().flatMapMany(list -> {
2
    if (list.isEmpty()) {
3
        return Flux.empty();
4
    }
5

6
    boolean hasToolCalls = list.stream()
7
        .map(this::safeAssistantMessage)
8
        .filter(Objects::nonNull)
9
        .anyMatch(am -> !CollectionUtils.isEmpty(am.getToolCalls()));
10

11
    if (!hasToolCalls) {
12
        if (list.size() > 2) {
13
            ChatResponse penultimateResponse = list.get(list.size() - 2);
14
            ChatResponse lastResponse = list.get(list.size() - 1);
15
            Usage usage = lastResponse.getMetadata().getUsage();
16
            observationContext.setResponse(new ChatResponse(
17
                    penultimateResponse.getResults(),
18
                    from(penultimateResponse.getMetadata(), usage)));
19
        }
20
        return Flux.fromIterable(list);
21
    }
22

23
    // tool call 的聚合和执行逻辑
24
});

1
Flux<ChatResponse> flux = chatResponses
2
    .contextWrite(ctx -> ctx.put(ObservationThreadLocalAccessor.KEY, observation));
3

4
// 旁路容器：边流式输出边收集一份，供结束时聚合用
5
List<ChatResponse> streamedResponses = Collections.synchronizedList(new ArrayList<>());
6
AtomicReference<ChatResponse> aggregatedResponse = new AtomicReference<>();
7

8
return flux
9
    .doOnNext(streamedResponses::add)        // 1. 旁路收集（不阻塞主流）
10
    .filter(this::shouldStreamResponse)      // 2. 普通内容 chunk 立刻放行
11
    .doOnComplete(() -> {                    // 3. 流结束时：聚合 → Observation
12
        ChatResponse aggregated = aggregateStreamingResponses(streamedResponses);
13
        if (aggregated != null) {
14
            aggregatedResponse.set(aggregated);
15
            observationContext.setResponse(aggregated);
16
        }
17
    })
18
    .concatWith(Flux.defer(() -> {           // ④ 流结束后追加：处理工具调用
19
        ChatResponse aggregated = aggregatedResponse.get();
20
        if (aggregated == null) {
21
            return Flux.empty();
22
        }
23
        return handleStreamingToolExecution(prompt, aggregated);
24
    }))
25
    .doOnError(observation::error)
26
    .doFinally(s -> observation.stop());

1
private @Nullable ChatResponse aggregateStreamingResponses(List<ChatResponse> responses) {
2
    if (responses.isEmpty()) {
3
        return null;
4
    }
5

6
    Map<String, ToolCallBuilder> builders = new HashMap<>();
7
    StringBuilder text = new StringBuilder();
8

9
    // ... 遍历所有 chunk：拼接 text、merge 工具调用碎片、收集 metadata ...
10
    ChatGenerationMetadata finalGenMetadata = ChatGenerationMetadata.NULL;
11
    Map<String, Object> props = new HashMap<>();
12
    ChatResponse lastResponseWithResult = null;
13

14
    for (ChatResponse chatResponse : responses) {
15
        AssistantMessage assistantMessage = safeAssistantMessage(chatResponse);
16
        if (assistantMessage == null) {
17
            continue;
18
        }
19

20
        lastResponseWithResult = chatResponse;
21

22
        if (assistantMessage.getText() != null) {
23
            text.append(assistantMessage.getText());
24
        }
25
        if (assistantMessage.getMetadata() != null) {
26
            props.putAll(assistantMessage.getMetadata());
27
        }
28
        mergeToolCalls(builders, assistantMessage);
29

30
        Generation generation = chatResponse.getResult();
31
        if (generation != null && generation.getMetadata() != ChatGenerationMetadata.NULL) {
32
            finalGenMetadata = generation.getMetadata();
33
        }
34
    }
35

36
    List<AssistantMessage.ToolCall> mergedToolCalls = builders.values()
37
        .stream()
38
        .map(ToolCallBuilder::build)
39
        .filter(toolCall -> StringUtils.hasText(toolCall.name()))
40
        .toList();
41

42
    AssistantMessage.Builder assistantMessageBuilder = AssistantMessage.builder()
43
        .content(text.toString())
44
        .properties(props);
45

46
    if (!mergedToolCalls.isEmpty()) {
47
        assistantMessageBuilder.toolCalls(mergedToolCalls);
48
    }
49

50
    ChatResponseMetadata finalMetadata = aggregateStreamingMetadata(responses, lastResponseWithResult);
51
    return new ChatResponse(List.of(new Generation(assistantMessageBuilder.build(), finalGenMetadata)),
52
            finalMetadata);
53
}

1
private void mergeToolCalls(Map<String, ToolCallBuilder> builders, AssistantMessage assistantMessage) {
2
    if (CollectionUtils.isEmpty(assistantMessage.getToolCalls())) {
3
        return;
4
    }
5

6
    Object chunkChoiceObject = assistantMessage.getMetadata().get("chunkChoice");
7
    if (chunkChoiceObject instanceof ChatCompletionChunk.Choice chunkChoice
8
            && chunkChoice.delta().toolCalls().isPresent()) {
9
        List<ChatCompletionChunk.Choice.Delta.ToolCall> deltaToolCalls = chunkChoice.delta().toolCalls().get();
10
        for (int i = 0; i < assistantMessage.getToolCalls().size() && i < deltaToolCalls.size(); i++) {
11
            AssistantMessage.ToolCall toolCall = assistantMessage.getToolCalls().get(i);
12
            ChatCompletionChunk.Choice.Delta.ToolCall deltaToolCall = deltaToolCalls.get(i);
13
            String key = chunkChoice.index() + "-" + deltaToolCall.index();
14
            builders.computeIfAbsent(key, keyValue -> new ToolCallBuilder()).merge(toolCall);
15
        }
16
        return;
17
    }
18

19
    for (AssistantMessage.ToolCall toolCall : assistantMessage.getToolCalls()) {
20
        builders.computeIfAbsent(toolCall.id(), key -> new ToolCallBuilder()).merge(toolCall);
21
    }
22
}

1
private Flux<ChatResponse> handleStreamingToolExecution(Prompt prompt, ChatResponse aggregated) {
2
    ChatOptions options = prompt.getOptions();
3
    Assert.state(options != null, "ChatOptions must not be null");
4

5
    // 不需要执行工具：是工具调用就发合并后的那一个、否则什么都不发
6
    if (!this.toolExecutionEligibilityPredicate.isToolExecutionRequired(options, aggregated)) {
7
        return aggregated.hasToolCalls() ? Flux.just(aggregated) : Flux.empty();
8
    }
9

10
    // 需要执行工具：执行后用历史重新发起一次 internalStream（递归流式）
11
    if (this.internalToolExecutionWarned.compareAndSet(false, true)) {
12
        logger.warn(
13
            "Internal tool execution in OpenAiChatModel is deprecated since 2.0.0 and will be removed in 3.0.0. "
14
                + "Use ChatClient with ToolCallAdvisor instead.");
15
    }
16

17
    return Flux.deferContextual(ctx -> {
18
        ToolExecutionResult toolExecutionResult;
19
        try {
20
            ToolCallReactiveContextHolder.setContext(ctx);
21
            toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, aggregated);
22
        }
23
        finally {
24
            ToolCallReactiveContextHolder.clearContext();
25
        }
26

27
        if (toolExecutionResult.returnDirect()) {
28
            return Flux.just(ChatResponse.builder()
29
                .from(aggregated)
30
                .generations(ToolExecutionResult.buildGenerations(toolExecutionResult))
31
                .build());
32
        }
33

34
        return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), options), aggregated);
35
    }).subscribeOn(Schedulers.boundedElastic());
36
}