diff --git a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatClient.java b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatClient.java index ee212817d..9f506e02d 100644 --- a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatClient.java +++ b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatClient.java @@ -18,6 +18,7 @@ package org.springframework.ai.mistralai; import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; @@ -34,6 +35,8 @@ import org.springframework.ai.chat.prompt.ChatOptions; import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.mistralai.api.MistralAiApi; import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletion; +import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletion.Choice; +import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionChunk; import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionMessage; import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionMessage.ToolCall; import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionRequest; @@ -131,13 +134,21 @@ public class MistralAiChatClient extends // The rest of the chunks with same ID share the same role. ConcurrentHashMap roleMap = new ConcurrentHashMap<>(); - return completionChunks.map(chunk -> { - String chunkId = chunk.id(); - List generations = chunk.choices().stream().map(choice -> { - if (choice.delta().role() != null) { - roleMap.putIfAbsent(chunkId, choice.delta().role().name()); + return completionChunks.map(chunk -> toChatCompletion(chunk)).map(chatCompletion -> { + + chatCompletion = handleFunctionCallOrReturn(request, ResponseEntity.of(Optional.of(chatCompletion))) + .getBody(); + + @SuppressWarnings("null") + String id = chatCompletion.id(); + + List generations = chatCompletion.choices().stream().map(choice -> { + if (choice.message().role() != null) { + roleMap.putIfAbsent(id, choice.message().role().name()); } - var generation = new Generation(choice.delta().content(), Map.of("role", roleMap.get(chunkId))); + String finish = (choice.finishReason() != null ? choice.finishReason().name() : ""); + var generation = new Generation(choice.message().content(), + Map.of("id", id, "role", roleMap.get(id), "finishReason", finish)); if (choice.finishReason() != null) { generation = generation .withGenerationMetadata(ChatGenerationMetadata.from(choice.finishReason().name(), null)); @@ -149,6 +160,15 @@ public class MistralAiChatClient extends }); } + private ChatCompletion toChatCompletion(ChatCompletionChunk chunk) { + List choices = chunk.choices() + .stream() + .map(cc -> new Choice(cc.index(), cc.delta(), cc.finishReason())) + .toList(); + + return new ChatCompletion(chunk.id(), "chat.completion", chunk.created(), chunk.model(), choices, null); + } + /** * Accessible for testing. */ @@ -194,10 +214,6 @@ public class MistralAiChatClient extends // Add the enabled functions definitions to the request's tools parameter. if (!CollectionUtils.isEmpty(functionsForThisRequest)) { - if (stream) { - throw new IllegalArgumentException("Currently tool functions are not supported in streaming mode"); - } - request = ModelOptionsUtils.merge( MistralAiChatOptions.builder().withTools(this.getFunctionTools(functionsForThisRequest)).build(), request, ChatCompletionRequest.class); @@ -241,7 +257,7 @@ public class MistralAiChatClient extends // Recursively call chatCompletionWithTools until the model doesn't call a // functions anymore. - ChatCompletionRequest newRequest = new ChatCompletionRequest(conversationHistory, previousRequest.stream()); + ChatCompletionRequest newRequest = new ChatCompletionRequest(conversationHistory, false); newRequest = ModelOptionsUtils.merge(newRequest, previousRequest, ChatCompletionRequest.class); return newRequest; @@ -252,6 +268,7 @@ public class MistralAiChatClient extends return request.messages(); } + @SuppressWarnings("null") @Override protected ChatCompletionMessage doGetToolResponseMessage(ResponseEntity chatCompletion) { return chatCompletion.getBody().choices().iterator().next().message(); diff --git a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MIstralAiStreamFunctionCallingHelper.java b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MIstralAiStreamFunctionCallingHelper.java new file mode 100644 index 000000000..6165ce900 --- /dev/null +++ b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MIstralAiStreamFunctionCallingHelper.java @@ -0,0 +1,198 @@ +/* + * Copyright 2023 - 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.mistralai.api; + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.UUID; + +import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionChunk; +import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionChunk.ChunkChoice; +import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionFinishReason; +import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionMessage; +import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionMessage.ChatCompletionFunction; +import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionMessage.Role; +import org.springframework.ai.mistralai.api.MistralAiApi.ChatCompletionMessage.ToolCall; +import org.springframework.util.CollectionUtils; + +/** + * Helper class to support Streaming function calling. + * + * It can merge the streamed ChatCompletionChunk in case of function calling message. + * + * @author Christian Tzolov + * @since 0.8.1 + */ +public class MIstralAiStreamFunctionCallingHelper { + + /** + * Merge the previous and current ChatCompletionChunk into a single one. + * @param previous the previous ChatCompletionChunk + * @param current the current ChatCompletionChunk + * @return the merged ChatCompletionChunk + */ + public ChatCompletionChunk merge(ChatCompletionChunk previous, ChatCompletionChunk current) { + + if (previous == null) { + return current; + } + + String id = (current.id() != null ? current.id() : previous.id()); + Long created = (current.created() != null ? current.created() : previous.created()); + String model = (current.model() != null ? current.model() : previous.model()); + String object = (current.object() != null ? current.object() : previous.object()); + + ChunkChoice previousChoice0 = (CollectionUtils.isEmpty(previous.choices()) ? null : previous.choices().get(0)); + ChunkChoice currentChoice0 = (CollectionUtils.isEmpty(current.choices()) ? null : current.choices().get(0)); + + ChunkChoice choice = merge(previousChoice0, currentChoice0); + + return new ChatCompletionChunk(id, object, created, model, List.of(choice)); + } + + private ChunkChoice merge(ChunkChoice previous, ChunkChoice current) { + if (previous == null) { + if (current.delta() != null && current.delta().toolCalls() != null) { + Optional id = current.delta() + .toolCalls() + .stream() + .filter(tool -> tool.id() != null) + .map(tool -> tool.id()) + .findFirst(); + if (!id.isPresent()) { + var newId = UUID.randomUUID().toString(); + + var toolCallsWithID = current.delta() + .toolCalls() + .stream() + .map(toolCall -> new ToolCall(newId, "function", toolCall.function())) + .toList(); + + var role = current.delta().role() != null ? current.delta().role() : Role.ASSISTANT; + current = new ChunkChoice(current.index(), new ChatCompletionMessage(current.delta().content(), + role, current.delta().name(), toolCallsWithID), current.finishReason()); + } + } + return current; + } + + ChatCompletionFinishReason finishReason = (current.finishReason() != null ? current.finishReason() + : previous.finishReason()); + Integer index = (current.index() != null ? current.index() : previous.index()); + + ChatCompletionMessage message = merge(previous.delta(), current.delta()); + + return new ChunkChoice(index, message, finishReason); + } + + private ChatCompletionMessage merge(ChatCompletionMessage previous, ChatCompletionMessage current) { + String content = (current.content() != null ? current.content() + : "" + ((previous.content() != null) ? previous.content() : "")); + Role role = (current.role() != null ? current.role() : previous.role()); + role = (role != null ? role : Role.ASSISTANT); // default to ASSISTANT (if null + String name = (current.name() != null ? current.name() : previous.name()); + + List toolCalls = new ArrayList<>(); + ToolCall lastPreviousTooCall = null; + if (previous.toolCalls() != null) { + lastPreviousTooCall = previous.toolCalls().get(previous.toolCalls().size() - 1); + if (previous.toolCalls().size() > 1) { + toolCalls.addAll(previous.toolCalls().subList(0, previous.toolCalls().size() - 1)); + } + } + if (current.toolCalls() != null) { + if (current.toolCalls().size() > 1) { + throw new IllegalStateException("Currently only one tool call is supported per message!"); + } + var currentToolCall = current.toolCalls().iterator().next(); + if (currentToolCall.id() != null) { + if (lastPreviousTooCall != null) { + toolCalls.add(lastPreviousTooCall); + } + toolCalls.add(currentToolCall); + } + else { + toolCalls.add(merge(lastPreviousTooCall, currentToolCall)); + } + } + else { + if (lastPreviousTooCall != null) { + toolCalls.add(lastPreviousTooCall); + } + } + return new ChatCompletionMessage(content, role, name, toolCalls); + } + + private ToolCall merge(ToolCall previous, ToolCall current) { + if (previous == null) { + return current; + } + String id = (current.id() != null ? current.id() : previous.id()); + String type = (current.type() != null ? current.type() : previous.type()); + ChatCompletionFunction function = merge(previous.function(), current.function()); + return new ToolCall(id, type, function); + } + + private ChatCompletionFunction merge(ChatCompletionFunction previous, ChatCompletionFunction current) { + if (previous == null) { + return current; + } + String name = (current.name() != null ? current.name() : previous.name()); + StringBuilder arguments = new StringBuilder(); + if (previous.arguments() != null) { + arguments.append(previous.arguments()); + } + if (current.arguments() != null) { + arguments.append(current.arguments()); + } + return new ChatCompletionFunction(name, arguments.toString()); + } + + /** + * @param chatCompletion the ChatCompletionChunk to check + * @return true if the ChatCompletionChunk is a streaming tool function call. + */ + public boolean isStreamingToolFunctionCall(ChatCompletionChunk chatCompletion) { + + var choices = chatCompletion.choices(); + if (CollectionUtils.isEmpty(choices)) { + return false; + } + + var choice = choices.get(0); + return !CollectionUtils.isEmpty(choice.delta().toolCalls()); + } + + /** + * @param chatCompletion the ChatCompletionChunk to check + * @return true if the ChatCompletionChunk is a streaming tool function call and it is + * the last one. + */ + public boolean isStreamingToolFunctionCallFinish(ChatCompletionChunk chatCompletion) { + + var choices = chatCompletion.choices(); + if (CollectionUtils.isEmpty(choices)) { + return false; + } + + var choice = choices.get(0); + return choice.finishReason() == ChatCompletionFinishReason.TOOL_CALL + || choice.finishReason() == ChatCompletionFinishReason.TOOL_CALLS; + } + +} +// --- \ No newline at end of file diff --git a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java index d2bd75c4a..bdf6efea8 100644 --- a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java +++ b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/api/MistralAiApi.java @@ -17,6 +17,7 @@ package org.springframework.ai.mistralai.api; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Consumer; import java.util.function.Predicate; @@ -704,6 +705,8 @@ public class MistralAiApi { .toEntity(ChatCompletion.class); } + private MIstralAiStreamFunctionCallingHelper chunkMerger = new MIstralAiStreamFunctionCallingHelper(); + /** * Creates a streaming chat response for the given chat conversation. * @param chatRequest The chat completion request. Must have the stream property set @@ -715,6 +718,8 @@ public class MistralAiApi { Assert.notNull(chatRequest, "The request body can not be null."); Assert.isTrue(chatRequest.stream(), "Request must set the steam property to true."); + AtomicBoolean isInsideTool = new AtomicBoolean(false); + return this.webClient.post() .uri("/v1/chat/completions") .body(Mono.just(chatRequest), ChatCompletionRequest.class) @@ -722,7 +727,26 @@ public class MistralAiApi { .bodyToFlux(String.class) .takeUntil(SSE_DONE_PREDICATE) .filter(SSE_DONE_PREDICATE.negate()) - .map(content -> ModelOptionsUtils.jsonToObject(content, ChatCompletionChunk.class)); + .map(content -> ModelOptionsUtils.jsonToObject(content, ChatCompletionChunk.class)) + .map(chunk -> { + if (this.chunkMerger.isStreamingToolFunctionCall(chunk)) { + isInsideTool.set(true); + } + return chunk; + }) + .windowUntil(chunk -> { + if (isInsideTool.get() && this.chunkMerger.isStreamingToolFunctionCallFinish(chunk)) { + isInsideTool.set(false); + return true; + } + return !isInsideTool.get(); + }) + .concatMapIterable(window -> { + Mono mono1 = window.reduce(new ChatCompletionChunk(null, null, null, null, null), + (previous, current) -> this.chunkMerger.merge(previous, current)); + return List.of(mono1); + }) + .flatMap(mono -> mono); } } diff --git a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatClientIT.java b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatClientIT.java index bdcc03b48..00cbf6973 100644 --- a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatClientIT.java +++ b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MistralAiChatClientIT.java @@ -15,6 +15,7 @@ */ package org.springframework.ai.mistralai; +import java.util.ArrayList; import java.util.Arrays; import java.util.List; import java.util.Map; @@ -24,6 +25,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; import org.springframework.ai.chat.ChatClient; import org.springframework.ai.chat.ChatResponse; @@ -35,6 +37,8 @@ import org.springframework.ai.chat.messages.UserMessage; import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.chat.prompt.PromptTemplate; import org.springframework.ai.chat.prompt.SystemPromptTemplate; +import org.springframework.ai.mistralai.api.MistralAiApi; +import org.springframework.ai.model.function.FunctionCallbackWrapper; import org.springframework.ai.parser.BeanOutputParser; import org.springframework.ai.parser.ListOutputParser; import org.springframework.ai.parser.MapOutputParser; @@ -181,4 +185,58 @@ class MistralAiChatClientIT { assertThat(actorsFilms.movies()).hasSize(5); } + @Test + void functionCallTest() { + + UserMessage userMessage = new UserMessage("What's the weather like in San Francisco?"); + + List messages = new ArrayList<>(List.of(userMessage)); + + var promptOptions = MistralAiChatOptions.builder() + .withModel(MistralAiApi.ChatModel.SMALL.getValue()) + .withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService()) + .withName("getCurrentWeather") + .withDescription("Get the weather in location") + .withResponseConverter((response) -> "" + response.temp() + response.unit()) + .build())) + .build(); + + ChatResponse response = chatClient.call(new Prompt(messages, promptOptions)); + + logger.info("Response: {}", response); + + assertThat(response.getResult().getOutput().getContent()).containsAnyOf("30.0", "30"); + } + + @Test + void streamFunctionCallTest() { + + UserMessage userMessage = new UserMessage("What's the weather like in Tokyo, Japan?"); + + List messages = new ArrayList<>(List.of(userMessage)); + + var promptOptions = MistralAiChatOptions.builder() + .withModel(MistralAiApi.ChatModel.SMALL.getValue()) + .withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService()) + .withName("getCurrentWeather") + .withDescription("Get the weather in location") + .withResponseConverter((response) -> "" + response.temp() + response.unit()) + .build())) + .build(); + + Flux response = streamingChatClient.stream(new Prompt(messages, promptOptions)); + + String content = response.collectList() + .block() + .stream() + .map(ChatResponse::getResults) + .flatMap(List::stream) + .map(Generation::getOutput) + .map(AssistantMessage::getContent) + .collect(Collectors.joining()); + logger.info("Response: {}", content); + + assertThat(content).containsAnyOf("10.0", "10"); + } + } diff --git a/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MockWeatherService.java b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MockWeatherService.java new file mode 100644 index 000000000..4a6f594a9 --- /dev/null +++ b/models/spring-ai-mistral-ai/src/test/java/org/springframework/ai/mistralai/MockWeatherService.java @@ -0,0 +1,90 @@ +/* + * Copyright 2023 - 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.mistralai; + +import java.util.function.Function; + +import com.fasterxml.jackson.annotation.JsonClassDescription; +import com.fasterxml.jackson.annotation.JsonInclude; +import com.fasterxml.jackson.annotation.JsonInclude.Include; +import com.fasterxml.jackson.annotation.JsonProperty; +import com.fasterxml.jackson.annotation.JsonPropertyDescription; + +/** + * @author Christian Tzolov + */ +public class MockWeatherService implements Function { + + /** + * Weather Function request. + */ + @JsonInclude(Include.NON_NULL) + @JsonClassDescription("Weather API request") + public record Request(@JsonProperty(required = true, + value = "location") @JsonPropertyDescription("The city and state e.g. San Francisco, CA") String location, + @JsonProperty(required = true, value = "unit") @JsonPropertyDescription("Temperature unit") Unit unit) { + } + + /** + * Temperature units. + */ + public enum Unit { + + /** + * Celsius. + */ + C("metric"), + /** + * Fahrenheit. + */ + F("imperial"); + + /** + * Human readable unit name. + */ + public final String unitName; + + private Unit(String text) { + this.unitName = text; + } + + } + + /** + * Weather Function response. + */ + public record Response(double temp, double feels_like, double temp_min, double temp_max, int pressure, int humidity, + Unit unit) { + } + + @Override + public Response apply(Request request) { + + double temperature = 0; + if (request.location().contains("Paris")) { + temperature = 15; + } + else if (request.location().contains("Tokyo")) { + temperature = 10; + } + else if (request.location().contains("San Francisco")) { + temperature = 30; + } + + return new Response(temperature, 15, 20, 2, 53, 45, Unit.C); + } + +} \ No newline at end of file diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatClient.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatClient.java index 14ef4e783..65410d40f 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatClient.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatClient.java @@ -19,6 +19,7 @@ import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; @@ -39,6 +40,8 @@ import org.springframework.ai.model.function.AbstractFunctionCallSupport; import org.springframework.ai.model.function.FunctionCallbackContext; import org.springframework.ai.openai.api.OpenAiApi; import org.springframework.ai.openai.api.OpenAiApi.ChatCompletion; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletion.Choice; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionFinishReason; import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage; import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.Role; import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.ToolCall; @@ -155,8 +158,10 @@ public class OpenAiChatClient extends @Override public Flux stream(Prompt prompt) { + + ChatCompletionRequest request = createRequest(prompt, true); + return this.retryTemplate.execute(ctx -> { - ChatCompletionRequest request = createRequest(prompt, true); Flux completionChunks = this.openAiApi.chatCompletionStream(request); @@ -164,24 +169,56 @@ public class OpenAiChatClient extends // The rest of the chunks with same ID share the same role. ConcurrentHashMap roleMap = new ConcurrentHashMap<>(); - return completionChunks.map(chunk -> { - String chunkId = chunk.id(); - List generations = chunk.choices().stream().map(choice -> { - if (choice.delta().role() != null) { - roleMap.putIfAbsent(chunkId, choice.delta().role().name()); - } - var generation = new Generation(choice.delta().content(), Map.of("role", roleMap.get(chunkId))); - if (choice.finishReason() != null) { - generation = generation - .withGenerationMetadata(ChatGenerationMetadata.from(choice.finishReason().name(), null)); - } - return generation; - }).toList(); - return new ChatResponse(generations); + // Convert the ChatCompletionChunk into a ChatCompletion to be able to reuse + // the function call handling logic. + return completionChunks.map(chunk -> chunkToChatCompletion(chunk)).map(chatCompletion -> { + try { + chatCompletion = handleFunctionCallOrReturn(request, ResponseEntity.of(Optional.of(chatCompletion))) + .getBody(); + + @SuppressWarnings("null") + String id = chatCompletion.id(); + + List generations = chatCompletion.choices().stream().map(choice -> { + if (choice.message().role() != null) { + roleMap.putIfAbsent(id, choice.message().role().name()); + } + String finish = (choice.finishReason() != null ? choice.finishReason().name() : ""); + var generation = new Generation(choice.message().content(), + Map.of("id", id, "role", roleMap.get(id), "finishReason", finish)); + if (choice.finishReason() != null) { + generation = generation.withGenerationMetadata( + ChatGenerationMetadata.from(choice.finishReason().name(), null)); + } + return generation; + }).toList(); + + return new ChatResponse(generations); + } + catch (Exception e) { + logger.error("Error processing chat completion", e); + return new ChatResponse(List.of()); + } + }); }); } + /** + * Convert the ChatCompletionChunk into a ChatCompletion. The Usage is set to null. + * @param chunk the ChatCompletionChunk to convert + * @return the ChatCompletion + */ + private OpenAiApi.ChatCompletion chunkToChatCompletion(OpenAiApi.ChatCompletionChunk chunk) { + List choices = chunk.choices() + .stream() + .map(cc -> new Choice(cc.finishReason(), cc.index(), cc.delta(), cc.logprobs())) + .toList(); + + return new OpenAiApi.ChatCompletion(chunk.id(), choices, chunk.created(), chunk.model(), + chunk.systemFingerprint(), "chat.completion", null); + } + /** * Accessible for testing. */ @@ -227,10 +264,6 @@ public class OpenAiChatClient extends // Add the enabled functions definitions to the request's tools parameter. if (!CollectionUtils.isEmpty(functionsForThisRequest)) { - if (stream) { - throw new IllegalArgumentException("Currently tool functions are not supported in streaming mode"); - } - request = ModelOptionsUtils.merge( OpenAiChatOptions.builder().withTools(this.getFunctionTools(functionsForThisRequest)).build(), request, ChatCompletionRequest.class); @@ -250,16 +283,6 @@ public class OpenAiChatClient extends private Map toMap(ChatCompletionMessage message) { Map map = new HashMap<>(); - // The tool_calls and tool_call_id are not used by the OpenAiChatClient functions - // call support! Useful only for users that want to use the tool_calls and - // tool_call_id in their applications. - if (message.toolCalls() != null) { - map.put("tool_calls", message.toolCalls()); - } - if (message.toolCallId() != null) { - map.put("tool_call_id", message.toolCallId()); - } - if (message.role() != null) { map.put("role", message.role().name()); } @@ -290,7 +313,7 @@ public class OpenAiChatClient extends // Recursively call chatCompletionWithTools until the model doesn't call a // functions anymore. - ChatCompletionRequest newRequest = new ChatCompletionRequest(conversationHistory, previousRequest.stream()); + ChatCompletionRequest newRequest = new ChatCompletionRequest(conversationHistory, false); newRequest = ModelOptionsUtils.merge(newRequest, previousRequest, ChatCompletionRequest.class); return newRequest; @@ -323,7 +346,9 @@ public class OpenAiChatClient extends return false; } - return !CollectionUtils.isEmpty(choices.get(0).message().toolCalls()); + var choice = choices.get(0); + return !CollectionUtils.isEmpty(choice.message().toolCalls()) + && choice.finishReason() == ChatCompletionFinishReason.TOOL_CALLS; } } diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java index b105d8824..b07c7579c 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiApi.java @@ -17,6 +17,7 @@ package org.springframework.ai.openai.api; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.function.Predicate; import com.fasterxml.jackson.annotation.JsonInclude; @@ -673,6 +674,8 @@ public class OpenAiApi { .toEntity(ChatCompletion.class); } + private OpenAiStreamFunctionCallingHelper chunkMerger = new OpenAiStreamFunctionCallingHelper(); + /** * Creates a streaming chat response for the given chat conversation. * @@ -684,6 +687,8 @@ public class OpenAiApi { Assert.notNull(chatRequest, "The request body can not be null."); Assert.isTrue(chatRequest.stream(), "Request must set the steam property to true."); + AtomicBoolean isInsideTool = new AtomicBoolean(false); + return this.webClient.post() .uri("/v1/chat/completions") .body(Mono.just(chatRequest), ChatCompletionRequest.class) @@ -693,7 +698,34 @@ public class OpenAiApi { .takeUntil(SSE_DONE_PREDICATE) // filters out the "[DONE]" message. .filter(SSE_DONE_PREDICATE.negate()) - .map(content -> ModelOptionsUtils.jsonToObject(content, ChatCompletionChunk.class)); + .map(content -> ModelOptionsUtils.jsonToObject(content, ChatCompletionChunk.class)) + // Detect is the chunk is part of a streaming function call. + .map(chunk -> { + if (this.chunkMerger.isStreamingToolFunctionCall(chunk)) { + isInsideTool.set(true); + } + return chunk; + }) + // Group all chunks belonging to the same function call. + // Flux -> Flux> + .windowUntil(chunk -> { + if (isInsideTool.get() && this.chunkMerger.isStreamingToolFunctionCallFinish(chunk)) { + isInsideTool.set(false); + return true; + } + return !isInsideTool.get(); + }) + // Merging the window chunks into a single chunk. + // Reduce the inner Flux window into a single Mono, + // Flux> -> Flux> + .concatMapIterable(window -> { + Mono monoChunk = window.reduce( + new ChatCompletionChunk(null, null, null, null, null, null), + (previous, current) -> this.chunkMerger.merge(previous, current)); + return List.of(monoChunk); + }) + // Flux> -> Flux + .flatMap(mono -> mono); } // Embeddings API @@ -851,86 +883,5 @@ public class OpenAiApi { }); } - // Transcription API - - // @JsonInclude(Include.NON_NULL) - // public record Transcription( - // @JsonProperty("text") String text) { - // } - - // /** - // * - // * @param model ID of the model to use. - // * @param language The language of the input audio. Supplying the input language in ISO-639-1 format will improve accuracy and latency. - // * @param prompt An optional text to guide the model's style or continue a previous audio segment. The prompt should match the audio language. - // * @param responseFormat An object specifying the format that the model must output. - // * @param temperature What sampling temperature to use, between 0 and 1. Higher values like 0.8 will make the output - // * more random, while lower values like 0.2 will make it more focused and deterministic. */ - // @JsonInclude(Include.NON_NULL) - // public record TranscriptionRequest ( - // @JsonProperty("model") String model, - // @JsonProperty("language") String language, - // @JsonProperty("prompt") String prompt, - // @JsonProperty("response_format") ResponseFormat responseFormat, - // @JsonProperty("temperature") Float temperature) { - - // /** - // * Shortcut constructor for a transcription request with the given model and temperature - // * - // * @param model ID of the model to use. - // * @param temperature What sampling temperature to use, between 0 and 1. - // */ - // public TranscriptionRequest(String model, Float temperature) { - // this(model, null, null, null, temperature); - // } - - // public TranscriptionRequest() { - // this(null, null, null, null, null); - // } - - // /** - // * An object specifying the format that the model must output. - // * @param type Must be one of 'text' or 'json_object'. - // */ - // @JsonInclude(Include.NON_NULL) - // public record ResponseFormat( - // @JsonProperty("type") String type) { - // } - // } - - // /** - // * Creates a model response for the given transcription. - // * - // * @param transcriptionRequest The transcription request. - // * @return Entity response with {@link Transcription} as a body and HTTP status code and headers. - // */ - // public ResponseEntity transcriptionEntityJson(MultiValueMap transcriptionRequest) { - - // Assert.notNull(transcriptionRequest, "The request body can not be null."); - - // return this.multipartRestClient.post() - // .uri("/v1/audio/transcriptions") - // .body(transcriptionRequest) - // .retrieve() - // .toEntity(Transcription.class); - // } - - // /** - // * Creates a model response for the given transcription. - // * - // * @param transcriptionRequest The transcription request. - // * @return Entity response with {@link String} as a body and HTTP status code and headers. - // */ - // public ResponseEntity transcriptionEntityText(MultiValueMap transcriptionRequest) { - - // Assert.notNull(transcriptionRequest, "The request body can not be null."); - - // return this.multipartRestClient.post() - // .uri("/v1/audio/transcriptions") - // .body(transcriptionRequest) - // .accept(MediaType.TEXT_PLAIN) - // .retrieve() - // .toEntity(String.class); - // } } // @formatter:on diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiStreamFunctionCallingHelper.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiStreamFunctionCallingHelper.java new file mode 100644 index 000000000..5c4132077 --- /dev/null +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/api/OpenAiStreamFunctionCallingHelper.java @@ -0,0 +1,201 @@ +/* + * Copyright 2023 - 2024 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.springframework.ai.openai.api; + +import java.util.ArrayList; +import java.util.List; + +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletion; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletion.Choice; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionChunk; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionFinishReason; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage; +import org.springframework.ai.openai.api.OpenAiApi.LogProbs; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionChunk.ChunkChoice; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.ChatCompletionFunction; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.Role; +import org.springframework.ai.openai.api.OpenAiApi.ChatCompletionMessage.ToolCall; +import org.springframework.util.CollectionUtils; + +/** + * Helper class to support Streaming function calling. + * + * It can merge the streamed ChatCompletionChunk in case of function calling message. + * + * @author Christian Tzolov + * @since 0.8.1 + */ +public class OpenAiStreamFunctionCallingHelper { + + /** + * Merge the previous and current ChatCompletionChunk into a single one. + * @param previous the previous ChatCompletionChunk + * @param current the current ChatCompletionChunk + * @return the merged ChatCompletionChunk + */ + public ChatCompletionChunk merge(ChatCompletionChunk previous, ChatCompletionChunk current) { + + if (previous == null) { + return current; + } + + String id = (current.id() != null ? current.id() : previous.id()); + Long created = (current.created() != null ? current.created() : previous.created()); + String model = (current.model() != null ? current.model() : previous.model()); + String systemFingerprint = (current.systemFingerprint() != null ? current.systemFingerprint() + : previous.systemFingerprint()); + String object = (current.object() != null ? current.object() : previous.object()); + + ChunkChoice previousChoice0 = (CollectionUtils.isEmpty(previous.choices()) ? null : previous.choices().get(0)); + ChunkChoice currentChoice0 = (CollectionUtils.isEmpty(current.choices()) ? null : current.choices().get(0)); + + ChunkChoice choice = merge(previousChoice0, currentChoice0); + + return new ChatCompletionChunk(id, List.of(choice), created, model, systemFingerprint, object); + } + + private ChunkChoice merge(ChunkChoice previous, ChunkChoice current) { + if (previous == null) { + return current; + } + + ChatCompletionFinishReason finishReason = (current.finishReason() != null ? current.finishReason() + : previous.finishReason()); + Integer index = (current.index() != null ? current.index() : previous.index()); + + ChatCompletionMessage message = merge(previous.delta(), current.delta()); + + LogProbs logprobs = (current.logprobs() != null ? current.logprobs() : previous.logprobs()); + return new ChunkChoice(finishReason, index, message, logprobs); + } + + private ChatCompletionMessage merge(ChatCompletionMessage previous, ChatCompletionMessage current) { + String content = (current.content() != null ? current.content() + : "" + ((previous.content() != null) ? previous.content() : "")); + Role role = (current.role() != null ? current.role() : previous.role()); + role = (role != null ? role : Role.ASSISTANT); // default to ASSISTANT (if null + String name = (current.name() != null ? current.name() : previous.name()); + String toolCallId = (current.toolCallId() != null ? current.toolCallId() : previous.toolCallId()); + + List toolCalls = new ArrayList<>(); + ToolCall lastPreviousTooCall = null; + if (previous.toolCalls() != null) { + lastPreviousTooCall = previous.toolCalls().get(previous.toolCalls().size() - 1); + if (previous.toolCalls().size() > 1) { + toolCalls.addAll(previous.toolCalls().subList(0, previous.toolCalls().size() - 1)); + } + } + if (current.toolCalls() != null) { + if (current.toolCalls().size() > 1) { + throw new IllegalStateException("Currently only one tool call is supported per message!"); + } + var currentToolCall = current.toolCalls().iterator().next(); + if (currentToolCall.id() != null) { + if (lastPreviousTooCall != null) { + toolCalls.add(lastPreviousTooCall); + } + toolCalls.add(currentToolCall); + } + else { + toolCalls.add(merge(lastPreviousTooCall, currentToolCall)); + } + } + else { + if (lastPreviousTooCall != null) { + toolCalls.add(lastPreviousTooCall); + } + } + return new ChatCompletionMessage(content, role, name, toolCallId, toolCalls); + } + + private ToolCall merge(ToolCall previous, ToolCall current) { + if (previous == null) { + return current; + } + String id = (current.id() != null ? current.id() : previous.id()); + String type = (current.type() != null ? current.type() : previous.type()); + ChatCompletionFunction function = merge(previous.function(), current.function()); + return new ToolCall(id, type, function); + } + + private ChatCompletionFunction merge(ChatCompletionFunction previous, ChatCompletionFunction current) { + if (previous == null) { + return current; + } + String name = (current.name() != null ? current.name() : previous.name()); + StringBuilder arguments = new StringBuilder(); + if (previous.arguments() != null) { + arguments.append(previous.arguments()); + } + if (current.arguments() != null) { + arguments.append(current.arguments()); + } + return new ChatCompletionFunction(name, arguments.toString()); + } + + /** + * @param chatCompletion the ChatCompletionChunk to check + * @return true if the ChatCompletionChunk is a streaming tool function call. + */ + public boolean isStreamingToolFunctionCall(ChatCompletionChunk chatCompletion) { + + if (chatCompletion == null || CollectionUtils.isEmpty(chatCompletion.choices())) { + return false; + } + + var choice = chatCompletion.choices().get(0); + if (choice == null || choice.delta() == null) { + return false; + } + return !CollectionUtils.isEmpty(choice.delta().toolCalls()); + } + + /** + * @param chatCompletion the ChatCompletionChunk to check + * @return true if the ChatCompletionChunk is a streaming tool function call and it is + * the last one. + */ + public boolean isStreamingToolFunctionCallFinish(ChatCompletionChunk chatCompletion) { + + if (chatCompletion == null || CollectionUtils.isEmpty(chatCompletion.choices())) { + return false; + } + + var choice = chatCompletion.choices().get(0); + if (choice == null || choice.delta() == null) { + return false; + } + return choice.finishReason() == ChatCompletionFinishReason.TOOL_CALLS; + } + + /** + * Convert the ChatCompletionChunk into a ChatCompletion. The Usage is set to null. + * @param chunk the ChatCompletionChunk to convert + * @return the ChatCompletion + */ + public ChatCompletion chunkToChatCompletion(ChatCompletionChunk chunk) { + List choices = chunk.choices() + .stream() + .map(chunkChoice -> new Choice(chunkChoice.finishReason(), chunkChoice.index(), chunkChoice.delta(), + chunkChoice.logprobs())) + .toList(); + + return new OpenAiApi.ChatCompletion(chunk.id(), choices, chunk.created(), chunk.model(), + chunk.systemFingerprint(), "chat.completion", null); + } + +} +// --- \ No newline at end of file diff --git a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatClientIT.java b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatClientIT.java index 579502daf..b23af540b 100644 --- a/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatClientIT.java +++ b/models/spring-ai-openai/src/test/java/org/springframework/ai/openai/chat/OpenAiChatClientIT.java @@ -25,6 +25,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; import org.springframework.ai.chat.ChatResponse; import org.springframework.ai.chat.Generation; @@ -37,6 +38,7 @@ import org.springframework.ai.chat.prompt.SystemPromptTemplate; import org.springframework.ai.model.function.FunctionCallbackWrapper; import org.springframework.ai.openai.OpenAiChatOptions; import org.springframework.ai.openai.OpenAiTestConfiguration; +import org.springframework.ai.openai.api.OpenAiApi; import org.springframework.ai.openai.api.tool.MockWeatherService; import org.springframework.ai.openai.testutils.AbstractIT; import org.springframework.ai.parser.BeanOutputParser; @@ -187,7 +189,7 @@ class OpenAiChatClientIT extends AbstractIT { List messages = new ArrayList<>(List.of(userMessage)); var promptOptions = OpenAiChatOptions.builder() - .withModel("gpt-4-turbo-preview") + .withModel(OpenAiApi.ChatModel.GPT_4_TURBO_PREVIEW.getValue()) .withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService()) .withName("getCurrentWeather") .withDescription("Get the weather in location") @@ -204,4 +206,37 @@ class OpenAiChatClientIT extends AbstractIT { assertThat(response.getResult().getOutput().getContent()).containsAnyOf("15.0", "15"); } -} + @Test + void streamFunctionCallTest() { + + UserMessage userMessage = new UserMessage("What's the weather like in San Francisco, Tokyo, and Paris?"); + + List messages = new ArrayList<>(List.of(userMessage)); + + var promptOptions = OpenAiChatOptions.builder() + // .withModel(OpenAiApi.ChatModel.GPT_4_TURBO_PREVIEW.getValue()) + .withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService()) + .withName("getCurrentWeather") + .withDescription("Get the weather in location") + .withResponseConverter((response) -> "" + response.temp() + response.unit()) + .build())) + .build(); + + Flux response = openStreamingChatClient.stream(new Prompt(messages, promptOptions)); + + String content = response.collectList() + .block() + .stream() + .map(ChatResponse::getResults) + .flatMap(List::stream) + .map(Generation::getOutput) + .map(AssistantMessage::getContent) + .collect(Collectors.joining()); + logger.info("Response: {}", content); + + assertThat(content).containsAnyOf("30.0", "30"); + assertThat(content).containsAnyOf("10.0", "10"); + assertThat(content).containsAnyOf("15.0", "15"); + } + +} \ No newline at end of file diff --git a/spring-ai-core/pom.xml b/spring-ai-core/pom.xml index 4e52278b1..15d79f0af 100644 --- a/spring-ai-core/pom.xml +++ b/spring-ai-core/pom.xml @@ -93,6 +93,18 @@ ${jsonschema.version} + + com.fasterxml.jackson.core + jackson-databind + ${jackson.version} + + + + com.fasterxml.jackson.datatype + jackson-datatype-jsr310 + ${jackson.version} + + org.springframework.boot diff --git a/spring-ai-core/src/main/java/org/springframework/ai/model/ModelOptionsUtils.java b/spring-ai-core/src/main/java/org/springframework/ai/model/ModelOptionsUtils.java index b9b3deb82..458445fc9 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/model/ModelOptionsUtils.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/model/ModelOptionsUtils.java @@ -35,6 +35,7 @@ import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.SerializationFeature; import com.fasterxml.jackson.databind.node.ArrayNode; import com.fasterxml.jackson.databind.node.ObjectNode; +import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule; import com.github.victools.jsonschema.generator.Option; import com.github.victools.jsonschema.generator.OptionPreset; import com.github.victools.jsonschema.generator.SchemaGenerator; @@ -60,7 +61,8 @@ public final class ModelOptionsUtils { private final static ObjectMapper OBJECT_MAPPER = new ObjectMapper() .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES) - .disable(SerializationFeature.FAIL_ON_EMPTY_BEANS); + .disable(SerializationFeature.FAIL_ON_EMPTY_BEANS) + .registerModule(new JavaTimeModule()); private final static List BEAN_MERGE_FIELD_EXCISIONS = List.of("class"); diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackInPromptIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackInPromptIT.java index 09e17abe5..67770f6af 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackInPromptIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackInPromptIT.java @@ -16,15 +16,19 @@ package org.springframework.ai.autoconfigure.openai.tool; import java.util.List; +import java.util.stream.Collectors; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration; import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration; import org.springframework.ai.chat.ChatResponse; +import org.springframework.ai.chat.Generation; +import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.UserMessage; import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.model.function.FunctionCallbackWrapper; @@ -70,4 +74,39 @@ public class FunctionCallbackInPromptIT { }); } + @Test + void streamingFunctionCallTest() { + + contextRunner.withPropertyValues("spring.ai.openai.chat.options.model=gpt-4-turbo-preview").run(context -> { + + OpenAiChatClient chatClient = context.getBean(OpenAiChatClient.class); + + UserMessage userMessage = new UserMessage("What's the weather like in San Francisco, Tokyo, and Paris?"); + + var promptOptions = OpenAiChatOptions.builder() + .withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService()) + .withName("CurrentWeatherService") + .withDescription("Get the weather in location") + .withResponseConverter((response) -> "" + response.temp() + response.unit()) + .build())) + .build(); + + Flux response = chatClient.stream(new Prompt(List.of(userMessage), promptOptions)); + + String content = response.collectList() + .block() + .stream() + .map(ChatResponse::getResults) + .flatMap(List::stream) + .map(Generation::getOutput) + .map(AssistantMessage::getContent) + .collect(Collectors.joining()); + logger.info("Response: {}", content); + + assertThat(content).containsAnyOf("30.0", "30"); + assertThat(content).containsAnyOf("10.0", "10"); + assertThat(content).containsAnyOf("15.0", "15"); + }); + } + } \ No newline at end of file diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackWithPlainFunctionBeanIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackWithPlainFunctionBeanIT.java index 51bd8721d..d5437a62f 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackWithPlainFunctionBeanIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackWithPlainFunctionBeanIT.java @@ -17,15 +17,19 @@ package org.springframework.ai.autoconfigure.openai.tool; import java.util.List; import java.util.function.Function; +import java.util.stream.Collectors; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration; import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration; import org.springframework.ai.chat.ChatResponse; +import org.springframework.ai.chat.Generation; +import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.UserMessage; import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.model.function.FunctionCallingOptions; @@ -98,6 +102,52 @@ class FunctionCallbackWithPlainFunctionBeanIT { }); } + @Test + void streamFunctionCallTest() { + contextRunner.withPropertyValues("spring.ai.openai.chat.options.model=gpt-4-turbo-preview").run(context -> { + + OpenAiChatClient chatClient = context.getBean(OpenAiChatClient.class); + + // Test weatherFunction + UserMessage userMessage = new UserMessage("What's the weather like in San Francisco, Tokyo, and Paris?"); + + Flux response = chatClient.stream(new Prompt(List.of(userMessage), + OpenAiChatOptions.builder().withFunction("weatherFunction").build())); + + String content = response.collectList() + .block() + .stream() + .map(ChatResponse::getResults) + .flatMap(List::stream) + .map(Generation::getOutput) + .map(AssistantMessage::getContent) + .collect(Collectors.joining()); + logger.info("Response: {}", content); + + assertThat(content).containsAnyOf("30.0", "30"); + assertThat(content).containsAnyOf("10.0", "10"); + assertThat(content).containsAnyOf("15.0", "15"); + + // Test weatherFunctionTwo + response = chatClient.stream(new Prompt(List.of(userMessage), + OpenAiChatOptions.builder().withFunction("weatherFunctionTwo").build())); + + content = response.collectList() + .block() + .stream() + .map(ChatResponse::getResults) + .flatMap(List::stream) + .map(Generation::getOutput) + .map(AssistantMessage::getContent) + .collect(Collectors.joining()); + logger.info("Response: {}", content); + + assertThat(content).containsAnyOf("30.0", "30"); + assertThat(content).containsAnyOf("10.0", "10"); + assertThat(content).containsAnyOf("15.0", "15"); + }); + } + @Configuration static class Config { diff --git a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackWrapperIT.java b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackWrapperIT.java index 9a4e70072..44e2c3af7 100644 --- a/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackWrapperIT.java +++ b/spring-ai-spring-boot-autoconfigure/src/test/java/org/springframework/ai/autoconfigure/openai/tool/FunctionCallbackWrapperIT.java @@ -16,15 +16,19 @@ package org.springframework.ai.autoconfigure.openai.tool; import java.util.List; +import java.util.stream.Collectors; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import reactor.core.publisher.Flux; import org.springframework.ai.autoconfigure.openai.OpenAiAutoConfiguration; import org.springframework.ai.autoconfigure.retry.SpringAiRetryAutoConfiguration; import org.springframework.ai.chat.ChatResponse; +import org.springframework.ai.chat.Generation; +import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.UserMessage; import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.model.function.FunctionCallbackWrapper; @@ -68,6 +72,34 @@ public class FunctionCallbackWrapperIT { }); } + @Test + void streamFunctionCallTest() { + contextRunner.withPropertyValues("spring.ai.openai.chat.options.model=gpt-4-turbo-preview").run(context -> { + + OpenAiChatClient chatClient = context.getBean(OpenAiChatClient.class); + + UserMessage userMessage = new UserMessage("What's the weather like in San Francisco, Tokyo, and Paris?"); + + Flux response = chatClient.stream( + new Prompt(List.of(userMessage), OpenAiChatOptions.builder().withFunction("WeatherInfo").build())); + + String content = response.collectList() + .block() + .stream() + .map(ChatResponse::getResults) + .flatMap(List::stream) + .map(Generation::getOutput) + .map(AssistantMessage::getContent) + .collect(Collectors.joining()); + logger.info("Response: {}", content); + + assertThat(content).containsAnyOf("30.0", "30"); + assertThat(content).containsAnyOf("10.0", "10"); + assertThat(content).containsAnyOf("15.0", "15"); + + }); + } + @Configuration static class Config {