From 28bceb3ac5de0c031133767ec1b010784bc1a18d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Dariusz=20J=C4=99drzejczyk?= Date: Mon, 10 Mar 2025 12:36:04 +0100 Subject: [PATCH] Prevent blocking event loop calling tools MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When streaming happens on an event loop, due to the fact that currently tools can only be synchronous and blocking, every tool call needs to be offloaded to a blocking-safe Scheduler. Resolves #2341 Signed-off-by: Dariusz Jędrzejczyk --- .../ai/anthropic/AnthropicChatModel.java | 27 ++++++++------- .../ai/azure/openai/AzureOpenAiChatModel.java | 33 +++++++++++-------- .../converse/BedrockProxyChatModel.java | 33 +++++++++++-------- .../ai/minimax/MiniMaxChatModel.java | 13 +++++--- .../ai/mistralai/MistralAiChatModel.java | 27 ++++++++------- .../ai/moonshot/MoonshotChatModel.java | 13 +++++--- .../ai/ollama/OllamaChatModel.java | 27 ++++++++------- .../ai/openai/OpenAiChatModel.java | 28 +++++++++------- .../gemini/VertexAiGeminiChatModel.java | 27 ++++++++------- .../ai/zhipuai/ZhiPuAiChatModel.java | 13 +++++--- 10 files changed, 145 insertions(+), 96 deletions(-) diff --git a/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java b/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java index 5218829fc..96ffcbaf6 100644 --- a/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java +++ b/models/spring-ai-anthropic/src/main/java/org/springframework/ai/anthropic/AnthropicChatModel.java @@ -39,6 +39,7 @@ import org.springframework.ai.util.json.JsonParser; import org.springframework.lang.Nullable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; import org.springframework.ai.anthropic.api.AnthropicApi; import org.springframework.ai.anthropic.api.AnthropicApi.AnthropicMessage; @@ -351,17 +352,21 @@ public class AnthropicChatModel extends AbstractToolCallSupport implements ChatM ChatResponse chatResponse = toChatResponse(chatCompletionResponse, accumulatedUsage); if (ToolCallingChatOptions.isInternalToolExecutionEnabled(prompt.getOptions()) && chatResponse.hasToolCalls() && chatResponse.hasFinishReasons(Set.of("tool_use"))) { - var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, chatResponse); - if (toolExecutionResult.returnDirect()) { - // Return tool execution result directly to the client. - return Flux.just(ChatResponse.builder().from(chatResponse) - .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) - .build()); - } else { - // Send the tool execution result back to the model. - return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), - chatResponse); - } + // FIXME: bounded elastic needs to be used since tool calling + // is currently only synchronous + return Flux.defer(() -> { + var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, chatResponse); + if (toolExecutionResult.returnDirect()) { + // Return tool execution result directly to the client. + return Flux.just(ChatResponse.builder().from(chatResponse) + .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) + .build()); + } else { + // Send the tool execution result back to the model. + return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), + chatResponse); + } + }).subscribeOn(Schedulers.boundedElastic()); } return Mono.just(chatResponse); diff --git a/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java b/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java index 7222d6b9f..0ac363b2a 100644 --- a/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java +++ b/models/spring-ai-azure-openai/src/main/java/org/springframework/ai/azure/openai/AzureOpenAiChatModel.java @@ -59,6 +59,7 @@ import io.micrometer.observation.contextpropagation.ObservationThreadLocalAccess import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; +import reactor.core.scheduler.Schedulers; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.Message; @@ -404,20 +405,24 @@ public class AzureOpenAiChatModel extends AbstractToolCallSupport implements Cha return chatResponseFlux.flatMap(chatResponse -> { if (ToolCallingChatOptions.isInternalToolExecutionEnabled(prompt.getOptions()) && chatResponse.hasToolCalls()) { - var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, chatResponse); - if (toolExecutionResult.returnDirect()) { - // Return tool execution result directly to the client. - return Flux.just(ChatResponse.builder() - .from(chatResponse) - .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) - .build()); - } - else { - // Send the tool execution result back to the model. - return this.internalStream( - new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), - chatResponse); - } + // FIXME: bounded elastic needs to be used since tool calling + // is currently only synchronous + return Flux.defer(() -> { + var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, chatResponse); + if (toolExecutionResult.returnDirect()) { + // Return tool execution result directly to the client. + return Flux.just(ChatResponse.builder() + .from(chatResponse) + .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) + .build()); + } + else { + // Send the tool execution result back to the model. + return this.internalStream( + new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), + chatResponse); + } + }).subscribeOn(Schedulers.boundedElastic()); } Flux flux = Flux.just(chatResponse) diff --git a/models/spring-ai-bedrock-converse/src/main/java/org/springframework/ai/bedrock/converse/BedrockProxyChatModel.java b/models/spring-ai-bedrock-converse/src/main/java/org/springframework/ai/bedrock/converse/BedrockProxyChatModel.java index f847bb98a..3098f4632 100644 --- a/models/spring-ai-bedrock-converse/src/main/java/org/springframework/ai/bedrock/converse/BedrockProxyChatModel.java +++ b/models/spring-ai-bedrock-converse/src/main/java/org/springframework/ai/bedrock/converse/BedrockProxyChatModel.java @@ -36,6 +36,7 @@ import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Sinks; import reactor.core.publisher.Sinks.EmitFailureHandler; +import reactor.core.scheduler.Schedulers; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; import software.amazon.awssdk.core.SdkBytes; import software.amazon.awssdk.core.document.Document; @@ -703,21 +704,25 @@ public class BedrockProxyChatModel extends AbstractToolCallSupport implements Ch && chatResponse.hasToolCalls() && chatResponse.hasFinishReasons(Set.of(StopReason.TOOL_USE.toString()))) { - var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, chatResponse); + // FIXME: bounded elastic needs to be used since tool calling + // is currently only synchronous + return Flux.defer(() -> { + var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, chatResponse); - if (toolExecutionResult.returnDirect()) { - // Return tool execution result directly to the client. - return Flux.just(ChatResponse.builder() - .from(chatResponse) - .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) - .build()); - } - else { - // Send the tool execution result back to the model. - return this.internalStream( - new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), - chatResponse); - } + if (toolExecutionResult.returnDirect()) { + // Return tool execution result directly to the client. + return Flux.just(ChatResponse.builder() + .from(chatResponse) + .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) + .build()); + } + else { + // Send the tool execution result back to the model. + return this.internalStream( + new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), + chatResponse); + } + }).subscribeOn(Schedulers.boundedElastic()); } else { return Flux.just(chatResponse); diff --git a/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxChatModel.java b/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxChatModel.java index e1660d18a..fe09f8cb1 100644 --- a/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxChatModel.java +++ b/models/spring-ai-minimax/src/main/java/org/springframework/ai/minimax/MiniMaxChatModel.java @@ -30,6 +30,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.MessageType; @@ -337,10 +338,14 @@ public class MiniMaxChatModel extends AbstractToolCallSupport implements ChatMod Flux flux = chatResponse.flatMap(response -> { if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(response, Set.of(ChatCompletionFinishReason.TOOL_CALLS.name(), ChatCompletionFinishReason.STOP.name()))) { - var toolCallConversation = handleToolCalls(prompt, response); - // Recursively call the stream method with the tool call message - // conversation that contains the call responses. - return this.stream(new Prompt(toolCallConversation, prompt.getOptions())); + // FIXME: bounded elastic needs to be used since tool calling + // is currently only synchronous + return Flux.defer(() -> { + var toolCallConversation = handleToolCalls(prompt, response); + // Recursively call the stream method with the tool call message + // conversation that contains the call responses. + return this.stream(new Prompt(toolCallConversation, prompt.getOptions())); + }).subscribeOn(Schedulers.boundedElastic()); } return Flux.just(response); }) diff --git a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatModel.java b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatModel.java index c14f8ae0f..18d8d3b70 100644 --- a/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatModel.java +++ b/models/spring-ai-mistral-ai/src/main/java/org/springframework/ai/mistralai/MistralAiChatModel.java @@ -35,6 +35,7 @@ import org.springframework.ai.tool.definition.ToolDefinition; import org.springframework.lang.Nullable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.SystemMessage; @@ -369,17 +370,21 @@ public class MistralAiChatModel extends AbstractToolCallSupport implements ChatM // @formatter:off Flux chatResponseFlux = chatResponse.flatMap(response -> { if (ToolCallingChatOptions.isInternalToolExecutionEnabled(prompt.getOptions()) && response.hasToolCalls()) { - var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response); - if (toolExecutionResult.returnDirect()) { - // Return tool execution result directly to the client. - return Flux.just(ChatResponse.builder().from(response) - .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) - .build()); - } else { - // Send the tool execution result back to the model. - return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), - response); - } + // FIXME: bounded elastic needs to be used since tool calling + // is currently only synchronous + return Flux.defer(() -> { + var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response); + if (toolExecutionResult.returnDirect()) { + // Return tool execution result directly to the client. + return Flux.just(ChatResponse.builder().from(response) + .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) + .build()); + } else { + // Send the tool execution result back to the model. + return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), + response); + } + }).subscribeOn(Schedulers.boundedElastic()); } else { return Flux.just(response); diff --git a/models/spring-ai-moonshot/src/main/java/org/springframework/ai/moonshot/MoonshotChatModel.java b/models/spring-ai-moonshot/src/main/java/org/springframework/ai/moonshot/MoonshotChatModel.java index 403c15a46..049b21fd7 100644 --- a/models/spring-ai-moonshot/src/main/java/org/springframework/ai/moonshot/MoonshotChatModel.java +++ b/models/spring-ai-moonshot/src/main/java/org/springframework/ai/moonshot/MoonshotChatModel.java @@ -29,6 +29,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.MessageType; @@ -321,10 +322,14 @@ public class MoonshotChatModel extends AbstractToolCallSupport implements ChatMo Flux flux = chatResponse.flatMap(response -> { if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(response, Set.of(ChatCompletionFinishReason.TOOL_CALLS.name(), ChatCompletionFinishReason.STOP.name()))) { - var toolCallConversation = handleToolCalls(prompt, response); - // Recursively call the stream method with the tool call message - // conversation that contains the call responses. - return this.internalStream(new Prompt(toolCallConversation, prompt.getOptions()), response); + // FIXME: bounded elastic needs to be used since tool calling + // is currently only synchronous + return Flux.defer(() -> { + var toolCallConversation = handleToolCalls(prompt, response); + // Recursively call the stream method with the tool call message + // conversation that contains the call responses. + return this.internalStream(new Prompt(toolCallConversation, prompt.getOptions()), response); + }).subscribeOn(Schedulers.boundedElastic()); } return Flux.just(response); }) diff --git a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatModel.java b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatModel.java index ab8bfe273..cf823fc97 100644 --- a/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatModel.java +++ b/models/spring-ai-ollama/src/main/java/org/springframework/ai/ollama/OllamaChatModel.java @@ -36,6 +36,7 @@ import org.springframework.ai.tool.definition.ToolDefinition; import org.springframework.ai.util.json.JsonParser; import org.springframework.lang.Nullable; import reactor.core.publisher.Flux; +import reactor.core.scheduler.Schedulers; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.SystemMessage; @@ -341,17 +342,21 @@ public class OllamaChatModel extends AbstractToolCallSupport implements ChatMode // @formatter:off Flux chatResponseFlux = chatResponse.flatMap(response -> { if (ToolCallingChatOptions.isInternalToolExecutionEnabled(prompt.getOptions()) && response.hasToolCalls()) { - var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response); - if (toolExecutionResult.returnDirect()) { - // Return tool execution result directly to the client. - return Flux.just(ChatResponse.builder().from(response) - .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) - .build()); - } else { - // Send the tool execution result back to the model. - return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), - response); - } + // FIXME: bounded elastic needs to be used since tool calling + // is currently only synchronous + return Flux.defer(() -> { + var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response); + if (toolExecutionResult.returnDirect()) { + // Return tool execution result directly to the client. + return Flux.just(ChatResponse.builder().from(response) + .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) + .build()); + } else { + // Send the tool execution result back to the model. + return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), + response); + } + }).subscribeOn(Schedulers.boundedElastic()); } else { return Flux.just(response); 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 8065f0b8b..c410f31dc 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 @@ -37,6 +37,7 @@ import org.springframework.ai.tool.definition.ToolDefinition; import org.springframework.lang.Nullable; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.MessageType; @@ -437,19 +438,22 @@ public class OpenAiChatModel extends AbstractToolCallSupport implements ChatMode // @formatter:off Flux flux = chatResponse.flatMap(response -> { - if (ToolCallingChatOptions.isInternalToolExecutionEnabled(prompt.getOptions()) && response.hasToolCalls()) { - var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response); - if (toolExecutionResult.returnDirect()) { - // Return tool execution result directly to the client. - return Flux.just(ChatResponse.builder().from(response) - .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) - .build()); - } else { - // Send the tool execution result back to the model. - return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), - response); - } + return Flux.defer(() -> { + // FIXME: bounded elastic needs to be used since tool calling + // is currently only synchronous + var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response); + if (toolExecutionResult.returnDirect()) { + // Return tool execution result directly to the client. + return Flux.just(ChatResponse.builder().from(response) + .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) + .build()); + } else { + // Send the tool execution result back to the model. + return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), + response); + } + }).subscribeOn(Schedulers.boundedElastic()); } else { return Flux.just(response); diff --git a/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java b/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java index 2378d7bc9..3c80d837b 100644 --- a/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java +++ b/models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java @@ -49,6 +49,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.Message; @@ -545,18 +546,22 @@ public class VertexAiGeminiChatModel extends AbstractToolCallSupport implements })); // @formatter:off - Flux chatResponseFlux = chatResponse1.flatMap(response -> { + Flux chatResponseFlux = chatResponse1.flatMap(response -> { if (ToolCallingChatOptions.isInternalToolExecutionEnabled(prompt.getOptions()) && response.hasToolCalls()) { - var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response); - if (toolExecutionResult.returnDirect()) { - // Return tool execution result directly to the client. - return Flux.just(ChatResponse.builder().from(response) - .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) - .build()); - } else { - // Send the tool execution result back to the model. - return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions())); - } + // FIXME: bounded elastic needs to be used since tool calling + // is currently only synchronous + return Flux.defer(() -> { + var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response); + if (toolExecutionResult.returnDirect()) { + // Return tool execution result directly to the client. + return Flux.just(ChatResponse.builder().from(response) + .generations(ToolExecutionResult.buildGenerations(toolExecutionResult)) + .build()); + } else { + // Send the tool execution result back to the model. + return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions())); + } + }).subscribeOn(Schedulers.boundedElastic()); } else { return Flux.just(response); diff --git a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java index 17d562852..3ea4fe7a0 100644 --- a/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java +++ b/models/spring-ai-zhipuai/src/main/java/org/springframework/ai/zhipuai/ZhiPuAiChatModel.java @@ -31,6 +31,7 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import reactor.core.publisher.Flux; import reactor.core.publisher.Mono; +import reactor.core.scheduler.Schedulers; import org.springframework.ai.chat.messages.AssistantMessage; import org.springframework.ai.chat.messages.MessageType; @@ -306,10 +307,14 @@ public class ZhiPuAiChatModel extends AbstractToolCallSupport implements ChatMod // @formatter:off Flux flux = chatResponse.flatMap(response -> { if (!isProxyToolCalls(prompt, this.defaultOptions) && isToolCall(response, Set.of(ChatCompletionFinishReason.TOOL_CALLS.name(), ChatCompletionFinishReason.STOP.name()))) { - var toolCallConversation = handleToolCalls(prompt, response); - // Recursively call the stream method with the tool call message - // conversation that contains the call responses. - return this.stream(new Prompt(toolCallConversation, prompt.getOptions())); + // FIXME: bounded elastic needs to be used since tool calling + // is currently only synchronous + return Flux.defer(() -> { + var toolCallConversation = handleToolCalls(prompt, response); + // Recursively call the stream method with the tool call message + // conversation that contains the call responses. + return this.stream(new Prompt(toolCallConversation, prompt.getOptions())); + }).subscribeOn(Schedulers.boundedElastic()); } return Flux.just(response); })