Skip to content

Fix OpenAI streaming tool call chunk merging #6525

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
liushenling2001 wants to merge 2 commits into from
+225 −34
Closed

Fix OpenAI streaming tool call chunk merging #6525

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
cd73809
Fix OpenAI stream empty tool calls
liushenling Jun 24, 2026
0b83a31
Fix OpenAI streaming tool call chunk merging
liushenling Jun 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
93 changes: 59 additions & 34 deletions models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1034,7 +1034,8 @@ private Prompt buildRequestPrompt(Prompt prompt) {
private static final class ChunkMerger {

static boolean hasToolCall(ChatCompletionChunk chunk) {
return !chunk.choices().isEmpty() && chunk.choices().get(0).delta().toolCalls().isPresent();
return !chunk.choices().isEmpty()
&& chunk.choices().get(0).delta().toolCalls().filter(toolCalls -> !toolCalls.isEmpty()).isPresent();
}

static boolean toolCallsDone(ChatCompletionChunk chunk) {
Expand Down Expand Up @@ -1068,40 +1069,59 @@ private static Choice mergeChoices(Choice c1, Choice c2) {
}

private static Delta mergeDeltas(Delta left, Delta right) {
var tcs = Stream.of(left.toolCalls(), right.toolCalls()).flatMap(Optional::stream).reduce((tcs1, tcs2) -> {
if (tcs2.isEmpty()) {
return tcs1;
}
Assert.isTrue(tcs2.size() == 1, "no more than one tool call per message currently supported");
ToolCall toolCall = tcs2.get(0);
if (toolCall.id().isPresent()) {
List<ToolCall> result = new ArrayList<>(tcs1);
result.add(toolCall);
return result;
}
else {
ToolCall lastFromTc1 = tcs1.get(tcs1.size() - 1);
Function lastFromTc1F = lastFromTc1.function().get();

var concatenatedArgs = Stream
.of(lastFromTc1F.arguments(), toolCall.function().flatMap(Function::arguments))
.flatMap(Optional::stream)
.reduce((args1, args2) -> args1 + args2)
.orElse("");

List<ToolCall> result = new ArrayList<>(tcs1);
result.set(tcs1.size() - 1,
lastFromTc1.toBuilder()
.putAllAdditionalProperties(toolCall._additionalProperties())
.function(lastFromTc1F.toBuilder().arguments(concatenatedArgs).build())
.build());
return result;
}
}).orElse(List.of());
List<ToolCall> tcs = new ArrayList<>(left.toolCalls().orElse(List.of()));
right.toolCalls().ifPresent(toolCalls -> toolCalls.forEach(toolCall -> mergeToolCall(tcs, toolCall)));

return left.toBuilder().toolCalls(tcs).build();
}

private static void mergeToolCall(List<ToolCall> existingToolCalls, ToolCall toolCall) {
Optional<Integer> existingIndex = findExistingToolCallIndex(existingToolCalls, toolCall);
if (existingIndex.isEmpty()) {
existingToolCalls.add(toolCall);
return;
}

ToolCall existingToolCall = existingToolCalls.get(existingIndex.get());
ToolCall.Builder builder = existingToolCall.toBuilder()
.putAllAdditionalProperties(toolCall._additionalProperties())
.function(mergeToolCallFunctions(existingToolCall.function(), toolCall.function()));
if (existingToolCall.id().isEmpty()) {
toolCall.id().ifPresent(builder::id);
}
existingToolCalls.set(existingIndex.get(), builder.build());
}

private static Optional<Integer> findExistingToolCallIndex(List<ToolCall> existingToolCalls,
ToolCall toolCall) {
for (int i = existingToolCalls.size() - 1; i >= 0; i--) {
ToolCall existingToolCall = existingToolCalls.get(i);
if (existingToolCall.index() == toolCall.index() || existingToolCall.id()
.filter(id -> toolCall.id().filter(id::equals).isPresent())
.isPresent()) {
return Optional.of(i);
}
}
return Optional.empty();
}

private static Function mergeToolCallFunctions(Optional<Function> left, Optional<Function> right) {
Function.Builder builder = left.map(Function::toBuilder).orElseGet(Function::builder);

Stream.of(left.flatMap(Function::name), right.flatMap(Function::name))
.flatMap(Optional::stream)
.findFirst()
.ifPresent(builder::name);

var concatenatedArgs = Stream.of(left.flatMap(Function::arguments), right.flatMap(Function::arguments))
.flatMap(Optional::stream)
.reduce((args1, args2) -> args1 + args2)
.orElse("");
builder.arguments(concatenatedArgs);

return builder.build();
}

/**
* Convert a ChatCompletionChunk into a ChatCompletion.
*/
Expand Down Expand Up @@ -1135,11 +1155,16 @@ static ChatCompletion chunkToChatCompletion(ChatCompletionChunk chunk) {
msgBuilder.toolCalls(ccctcs.stream().map(tc -> {
ChatCompletionMessageFunctionToolCall.Builder toolCallBuilder = ChatCompletionMessageFunctionToolCall
.builder();
Function function = tc.function()
.orElseThrow(() -> new IllegalStateException("Tool call function is missing"));
toolCallBuilder.putAllAdditionalProperties(tc._additionalProperties());
toolCallBuilder.id(tc.id().get());
toolCallBuilder
.id(tc.id().orElseThrow(() -> new IllegalStateException("Tool call id is missing")));
toolCallBuilder.function(ChatCompletionMessageFunctionToolCall.Function.builder()
.name(tc.function().get().name().get())
.arguments(tc.function().get().arguments().get())
.name(function.name()
.filter(StringUtils::hasText)
.orElseThrow(() -> new IllegalStateException("Tool call function name is missing")))
.arguments(function.arguments().orElse(""))
.build());
return ChatCompletionMessageToolCall.ofFunction(toolCallBuilder.build());
}).toList());
Expand Down
166 changes: 166 additions & 0 deletions ...ls/spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiChatModelTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -263,6 +263,172 @@ void preserveUnmappedToolCallAdditionalPropertiesFromStream() {
Map.of("google", Map.of("thought_signature", "signature-123")));
}

@Test
void mergeStreamToolCallWhenIdNameAndArgumentsArriveInSeparateChunks() {
ChatServiceAsync chatServiceAsync = mock(ChatServiceAsync.class);
ChatCompletionServiceAsync chatCompletionServiceAsync = mock(ChatCompletionServiceAsync.class);
when(this.openAiClientAsync.chat()).thenReturn(chatServiceAsync);
when(chatServiceAsync.completions()).thenReturn(chatCompletionServiceAsync);
when(chatCompletionServiceAsync.createStreaming(any(ChatCompletionCreateParams.class)))
.thenReturn(asyncStreamResponse(ChatCompletionChunk.builder()
.id("chatcmpl-stream-test")
.created(1777799928)
.model("vllm-openai-compatible")
.addChoice(ChatCompletionChunk.Choice.builder()
.index(0)
.finishReason(Optional.empty())
.delta(ChatCompletionChunk.Choice.Delta.builder()
.addToolCall(ChatCompletionChunk.Choice.Delta.ToolCall.builder().index(0).id("call_1").build())
.build())
.build())
.build(),
ChatCompletionChunk.builder()
.id("chatcmpl-stream-test")
.created(1777799928)
.model("vllm-openai-compatible")
.addChoice(ChatCompletionChunk.Choice.builder()
.index(0)
.finishReason(Optional.empty())
.delta(ChatCompletionChunk.Choice.Delta.builder()
.addToolCall(ChatCompletionChunk.Choice.Delta.ToolCall.builder()
.index(0)
.function(ChatCompletionChunk.Choice.Delta.ToolCall.Function.builder()
.name("get_current_weather")
.build())
.build())
.build())
.build())
.build(),
ChatCompletionChunk.builder()
.id("chatcmpl-stream-test")
.created(1777799928)
.model("vllm-openai-compatible")
.addChoice(ChatCompletionChunk.Choice.builder()
.index(0)
.finishReason(Optional.empty())
.delta(ChatCompletionChunk.Choice.Delta.builder()
.addToolCall(ChatCompletionChunk.Choice.Delta.ToolCall.builder()
.index(0)
.function(ChatCompletionChunk.Choice.Delta.ToolCall.Function.builder()
.arguments("{\"location\":\"Seoul\"}")
.build())
.build())
.build())
.build())
.build(),
ChatCompletionChunk.builder()
.id("chatcmpl-stream-test")
.created(1777799928)
.model("vllm-openai-compatible")
.addChoice(ChatCompletionChunk.Choice.builder()
.index(0)
.finishReason(ChatCompletionChunk.Choice.FinishReason.TOOL_CALLS)
.delta(ChatCompletionChunk.Choice.Delta.builder().build())
.build())
.build()));

OpenAiChatOptions options = OpenAiChatOptions.builder().model("vllm-openai-compatible").build();
OpenAiChatModel chatModel = OpenAiChatModel.builder()
.openAiClient(this.openAiClient)
.openAiClientAsync(this.openAiClientAsync)
.options(options)
.build();

ChatResponse response = chatModel.stream(new Prompt("hi", options)).blockLast();

assertThat(response).isNotNull();
AssistantMessage assistantMessage = response.getResult().getOutput();
assertThat(assistantMessage.getToolCalls()).singleElement().satisfies(toolCall -> {
assertThat(toolCall.id()).isEqualTo("call_1");
assertThat(toolCall.name()).isEqualTo("get_current_weather");
assertThat(toolCall.arguments()).isEqualTo("{\"location\":\"Seoul\"}");
});
}

@Test
void mergeStreamToolCallWhenProviderRepeatsIdAcrossChunks() {
ChatServiceAsync chatServiceAsync = mock(ChatServiceAsync.class);
ChatCompletionServiceAsync chatCompletionServiceAsync = mock(ChatCompletionServiceAsync.class);
when(this.openAiClientAsync.chat()).thenReturn(chatServiceAsync);
when(chatServiceAsync.completions()).thenReturn(chatCompletionServiceAsync);
when(chatCompletionServiceAsync.createStreaming(any(ChatCompletionCreateParams.class)))
.thenReturn(asyncStreamResponse(ChatCompletionChunk.builder()
.id("chatcmpl-stream-test")
.created(1777799928)
.model("vllm-openai-compatible")
.addChoice(ChatCompletionChunk.Choice.builder()
.index(0)
.finishReason(Optional.empty())
.delta(ChatCompletionChunk.Choice.Delta.builder()
.addToolCall(ChatCompletionChunk.Choice.Delta.ToolCall.builder().index(0).id("call_1").build())
.build())
.build())
.build(),
ChatCompletionChunk.builder()
.id("chatcmpl-stream-test")
.created(1777799928)
.model("vllm-openai-compatible")
.addChoice(ChatCompletionChunk.Choice.builder()
.index(0)
.finishReason(Optional.empty())
.delta(ChatCompletionChunk.Choice.Delta.builder()
.addToolCall(ChatCompletionChunk.Choice.Delta.ToolCall.builder()
.index(0)
.id("call_1")
.function(ChatCompletionChunk.Choice.Delta.ToolCall.Function.builder()
.name("get_current_weather")
.build())
.build())
.build())
.build())
.build(),
ChatCompletionChunk.builder()
.id("chatcmpl-stream-test")
.created(1777799928)
.model("vllm-openai-compatible")
.addChoice(ChatCompletionChunk.Choice.builder()
.index(0)
.finishReason(Optional.empty())
.delta(ChatCompletionChunk.Choice.Delta.builder()
.addToolCall(ChatCompletionChunk.Choice.Delta.ToolCall.builder()
.index(0)
.id("call_1")
.function(ChatCompletionChunk.Choice.Delta.ToolCall.Function.builder()
.arguments("{\"location\":\"Seoul\"}")
.build())
.build())
.build())
.build())
.build(),
ChatCompletionChunk.builder()
.id("chatcmpl-stream-test")
.created(1777799928)
.model("vllm-openai-compatible")
.addChoice(ChatCompletionChunk.Choice.builder()
.index(0)
.finishReason(ChatCompletionChunk.Choice.FinishReason.TOOL_CALLS)
.delta(ChatCompletionChunk.Choice.Delta.builder().build())
.build())
.build()));

OpenAiChatOptions options = OpenAiChatOptions.builder().model("vllm-openai-compatible").build();
OpenAiChatModel chatModel = OpenAiChatModel.builder()
.openAiClient(this.openAiClient)
.openAiClientAsync(this.openAiClientAsync)
.options(options)
.build();

ChatResponse response = chatModel.stream(new Prompt("hi", options)).blockLast();

assertThat(response).isNotNull();
AssistantMessage assistantMessage = response.getResult().getOutput();
assertThat(assistantMessage.getToolCalls()).singleElement().satisfies(toolCall -> {
assertThat(toolCall.id()).isEqualTo("call_1");
assertThat(toolCall.name()).isEqualTo("get_current_weather");
assertThat(toolCall.arguments()).isEqualTo("{\"location\":\"Seoul\"}");
});
}

@Test
void restoreUnmappedToolCallAdditionalProperties() {
OpenAiChatOptions options = OpenAiChatOptions.builder().model("gemini-3.5-flash").build();
Expand Down
Loading