VertexAiGeminiChatModel support InternalToolExecutionMaxAttempts

Signed-off-by: lambochen <[email protected]>

Author: lambochen

models/spring-ai-vertex-ai-gemini/src/main/java/org/springframework/ai/vertexai/gemini/VertexAiGeminiChatModel.java

@@ -136,10 +136,12 @@
  * @author Jihoon Kim
  * @author Alexandros Pappas
  * @author Ilayaperumal Gopinathan
+ * @author lambochen
  * @since 0.8.1
  * @see VertexAiGeminiChatOptions
  * @see ToolCallingManager
  * @see ChatModel
+ * @see ToolCallingChatOptions
  */
 public class VertexAiGeminiChatModel implements ChatModel, DisposableBean {
 
@@ -389,10 +391,10 @@ private static Schema jsonToSchema(String json) {
 	@Override
 	public ChatResponse call(Prompt prompt) {
 		var requestPrompt = this.buildRequestPrompt(prompt);
-		return this.internalCall(requestPrompt, null);
+		return this.internalCall(requestPrompt, null, 1);
 	}
 
-	private ChatResponse internalCall(Prompt prompt, ChatResponse previousChatResponse) {
+	private ChatResponse internalCall(Prompt prompt, ChatResponse previousChatResponse, int attempts) {
 
 		ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
 			.prompt(prompt)
@@ -425,7 +427,7 @@ private ChatResponse internalCall(Prompt prompt, ChatResponse previousChatRespon
 				return chatResponse;
 			}));
 
-		if (this.toolExecutionEligibilityPredicate.isToolExecutionRequired(prompt.getOptions(), response)) {
+		if (this.toolExecutionEligibilityPredicate.isToolExecutionRequired(prompt.getOptions(), response, attempts)) {
 			var toolExecutionResult = this.toolCallingManager.executeToolCalls(prompt, response);
 			if (toolExecutionResult.returnDirect()) {
 				// Return tool execution result directly to the client.
@@ -437,7 +439,7 @@ private ChatResponse internalCall(Prompt prompt, ChatResponse previousChatRespon
 			else {
 				// Send the tool execution result back to the model.
 				return this.internalCall(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()),
-						response);
+						response, attempts + 1);
 			}
 		}
 
@@ -469,6 +471,11 @@ Prompt buildRequestPrompt(Prompt prompt) {
 			requestOptions.setInternalToolExecutionEnabled(
 					ModelOptionsUtils.mergeOption(runtimeOptions.getInternalToolExecutionEnabled(),
 							this.defaultOptions.getInternalToolExecutionEnabled()));
+			requestOptions.setInternalToolExecutionMaxAttempts(
+					ModelOptionsUtils.mergeOption(
+							runtimeOptions.getInternalToolExecutionMaxAttempts(),
+							this.defaultOptions.getInternalToolExecutionMaxAttempts())
+			);
 			requestOptions.setToolNames(ToolCallingChatOptions.mergeToolNames(runtimeOptions.getToolNames(),
 					this.defaultOptions.getToolNames()));
 			requestOptions.setToolCallbacks(ToolCallingChatOptions.mergeToolCallbacks(runtimeOptions.getToolCallbacks(),
@@ -483,6 +490,7 @@ Prompt buildRequestPrompt(Prompt prompt) {
 		}
 		else {
 			requestOptions.setInternalToolExecutionEnabled(this.defaultOptions.getInternalToolExecutionEnabled());
+			requestOptions.setInternalToolExecutionMaxAttempts(this.defaultOptions.getInternalToolExecutionMaxAttempts());
 			requestOptions.setToolNames(this.defaultOptions.getToolNames());
 			requestOptions.setToolCallbacks(this.defaultOptions.getToolCallbacks());
 			requestOptions.setToolContext(this.defaultOptions.getToolContext());
@@ -499,10 +507,10 @@ Prompt buildRequestPrompt(Prompt prompt) {
 	@Override
 	public Flux<ChatResponse> stream(Prompt prompt) {
 		var requestPrompt = this.buildRequestPrompt(prompt);
-		return this.internalStream(requestPrompt, null);
+		return this.internalStream(requestPrompt, null, 1);
 	}
 
-	public Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse previousChatResponse) {
+	public Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse previousChatResponse, int attempts) {
 		return Flux.deferContextual(contextView -> {
 
 			ChatModelObservationContext observationContext = ChatModelObservationContext.builder()
@@ -538,7 +546,7 @@ public Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse previousCha
 
 				// @formatter:off
 				Flux<ChatResponse> flux = chatResponseFlux.flatMap(response -> {
-					if (this.toolExecutionEligibilityPredicate.isToolExecutionRequired(prompt.getOptions(), response)) {
+					if (this.toolExecutionEligibilityPredicate.isToolExecutionRequired(prompt.getOptions(), response, attempts)) {
 						// FIXME: bounded elastic needs to be used since tool calling
 						// is currently only synchronous
 						return Flux.defer(() -> {
@@ -551,7 +559,10 @@ public Flux<ChatResponse> internalStream(Prompt prompt, ChatResponse previousCha
 							}
 							else {
 								// Send the tool execution result back to the model.
-								return this.internalStream(new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()), response);
+								return this.internalStream(
+										new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()),
+										response,
+										attempts + 1);
 							}
 						}).subscribeOn(Schedulers.boundedElastic());
 					}