diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java index b17d16a68b..c19195e99b 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java @@ -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) { @@ -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 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 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 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 existingToolCalls, ToolCall toolCall) { + Optional 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 findExistingToolCallIndex(List 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 left, Optional 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. */ @@ -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()); diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiChatModelTests.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiChatModelTests.java index 82e8349993..c0054e0a65 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiChatModelTests.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/OpenAiChatModelTests.java @@ -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();