Multiple Roundtrips for functions calls are failing in Azure OpenAI Chat Chatcompletion

models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java

@@ -44,7 +44,9 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.ai.azure.openai.metadata.AzureOpenAiChatResponseMetadata;
+import org.springframework.ai.chat.messages.AssistantMessage;
 import org.springframework.ai.chat.messages.Message;
+import org.springframework.ai.chat.messages.ToolResponseMessage;
 import org.springframework.ai.chat.metadata.ChatGenerationMetadata;
 import org.springframework.ai.chat.metadata.PromptMetadata;
 import org.springframework.ai.chat.metadata.PromptMetadata.PromptFilterMetadata;
@@ -59,13 +61,9 @@
 import org.springframework.util.Assert;
 import org.springframework.util.CollectionUtils;
 import reactor.core.publisher.Flux;
+import reactor.core.publisher.Mono;
 
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashSet;
-import java.util.List;
-import java.util.Optional;
-import java.util.Set;
+import java.util.*;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -160,8 +158,20 @@ public ChatResponse call(Prompt prompt) {
 
 	@Override
 	public Flux<ChatResponse> stream(Prompt prompt) {
-
 		ChatCompletionsOptions options = toAzureChatCompletionsOptions(prompt);
+
+		// we have to map with a custom function to handle the tool call requests
+		// due to the existing bugs in the azure api (see comments in streamWithAzureApi)
+		// we have to recursively call this specific method for tool calls instead of
+		// using the one from the AbstractFunctionCallSupport
+		return streamWithAzureOpenAi(options).flatMapIterable(ChatCompletions::getChoices).map(choice -> {
+			var content = Optional.ofNullable(choice.getMessage()).orElse(choice.getDelta()).getContent();
+			var generation = new Generation(content).withGenerationMetadata(generateChoiceMetadata(choice));
+			return new ChatResponse(List.of(generation));
+		});
+	}
+
+	private Flux<ChatCompletions> streamWithAzureOpenAi(ChatCompletionsOptions options) {
 		options.setStream(true);
 
 		IterableStream<ChatCompletions> chatCompletionsStream = this.openAIClient
@@ -193,16 +203,36 @@ public Flux<ChatResponse> stream(Prompt prompt) {
 				return List.of(reduce);
 			})
 			.flatMap(mono -> mono);
-		return accessibleChatCompletionsFlux
-			.switchMap(accessibleChatCompletions -> handleFunctionCallOrReturnStream(options,
-					Flux.just(accessibleChatCompletions)))
-			.flatMapIterable(ChatCompletions::getChoices)
-			.map(choice -> {
-				var content = Optional.ofNullable(choice.getMessage()).orElse(choice.getDelta()).getContent();
-				var generation = new Generation(content).withGenerationMetadata(generateChoiceMetadata(choice));
-				return new ChatResponse(List.of(generation));
-			});
+		return accessibleChatCompletionsFlux.switchMap(
+				accessibleChatCompletions -> handleToolCallRequests(options, Flux.just(accessibleChatCompletions)));
+
+	}
+
+	private Flux<ChatCompletions> handleToolCallRequests(ChatCompletionsOptions request,
+			Flux<ChatCompletions> response) {
+		return response.switchMap(resp -> {
+			if (!this.isToolFunctionCall(resp)) {
+				return Mono.just(resp);
+			}
+
+			// The chat completion tool call requires the complete conversation
+			// history. Including the initial user message.
+			List<ChatRequestMessage> conversationHistory = new ArrayList<>();
+
+			conversationHistory.addAll(this.doGetUserMessages(request));
+
+			ChatRequestMessage responseMessage = this.doGetToolResponseMessage(resp);
 
+			// Add the assistant response to the message conversation history.
+			conversationHistory.add(responseMessage);
+
+			ChatCompletionsOptions newRequest = this.doCreateToolResponseRequest(request, responseMessage,
+					conversationHistory);
+
+			// recursively go backwards and call our stream again (including all bug fixes
+			// / workarounds for the azure api)
+			return this.streamWithAzureOpenAi(newRequest);
+		});
 	}
 
 	/**
@@ -486,6 +516,9 @@ private ChatCompletionsOptions merge(ChatCompletionsOptions fromOptions, ChatCom
 		if (fromOptions.getResponseFormat() != null) {
 			mergedOptions.setResponseFormat(fromOptions.getResponseFormat());
 		}
+		if (fromOptions.getTools() != null) {
+			mergedOptions.setTools(fromOptions.getTools());
+		}
 
 		return mergedOptions;
 	}

models/spring-ai-azure-openai/src/test/java/org/springframework/ai/azure/openai/function/AzureOpenAiChatModelFunctionCallIT.java

@@ -17,6 +17,7 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
@@ -85,6 +86,70 @@ void functionCallTest() {
 		assertThat(response.getResult().getOutput().getContent()).containsAnyOf("15.0", "15");
 	}
 
+	@Test
+	void functionCallSequentialTest() {
+
+		UserMessage userMessage = new UserMessage(
+				"What's the weather like in San Francisco? If the weather is above 25 degrees, please check the weather in Tokyo and Paris.");
+
+		List<Message> messages = new ArrayList<>(List.of(userMessage));
+
+		var promptOptions = AzureOpenAiChatOptions.builder()
+			.withDeploymentName(selectedModel)
+			.withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService())
+				.withName("getCurrentWeather")
+				.withDescription("Get the current weather in a given location")
+				.withResponseConverter((response) -> "" + response.temp() + response.unit())
+				.build()))
+			.build();
+
+		ChatResponse response = chatModel.call(new Prompt(messages, promptOptions));
+
+		logger.info("Response: {}", response);
+
+		assertThat(response.getResult().getOutput().getContent()).containsAnyOf("30.0", "30");
+		assertThat(response.getResult().getOutput().getContent()).containsAnyOf("10.0", "10");
+		assertThat(response.getResult().getOutput().getContent()).containsAnyOf("15.0", "15");
+	}
+
+	@Test
+	void functionCallSequentialAndStreamTest() {
+
+		UserMessage userMessage = new UserMessage(
+				"What's the weather like in San Francisco? If the weather is above 25 degrees, please check the weather in Tokyo and Paris.");
+
+		List<Message> messages = new ArrayList<>(List.of(userMessage));
+
+		var promptOptions = AzureOpenAiChatOptions.builder()
+			.withDeploymentName(selectedModel)
+			.withFunctionCallbacks(List.of(FunctionCallbackWrapper.builder(new MockWeatherService())
+				.withName("getCurrentWeather")
+				.withDescription("Get the current weather in a given location")
+				.withResponseConverter((response) -> "" + response.temp() + response.unit())
+				.build()))
+			.build();
+
+		var response = chatModel.stream(new Prompt(messages, promptOptions));
+
+		final var counter = new AtomicInteger();
+		String content = response.doOnEach(listSignal -> counter.getAndIncrement())
+			.collectList()
+			.block()
+			.stream()
+			.map(ChatResponse::getResults)
+			.flatMap(List::stream)
+			.map(Generation::getOutput)
+			.map(AssistantMessage::getContent)
+			.filter(Objects::nonNull)
+			.collect(Collectors.joining());
+
+		logger.info("Response: {}", response);
+
+		assertThat(content).containsAnyOf("30.0", "30");
+		assertThat(content).containsAnyOf("10.0", "10");
+		assertThat(content).containsAnyOf("15.0", "15");
+	}
+
 	@Test
 	void streamFunctionCallTest() {
 		UserMessage userMessage = new UserMessage("What's the weather like in San Francisco, Tokyo, and Paris?");