From aa74564c3a53035dbbdf0e2a323fcb34d864c2ca Mon Sep 17 00:00:00 2001 From: mumu-1029521 <186442148+mumu-1029521@users.noreply.github.com> Date: Mon, 31 Aug 2026 18:22:59 +0800 Subject: [PATCH] fix(spring-ai): make tool execution configurable --- contrib/spring-ai/README.md | 42 ++- .../springai/AdkToolContextResolver.java | 44 +++ .../adk/models/springai/MessageConverter.java | 8 +- .../google/adk/models/springai/SpringAI.java | 288 +++++++++++++----- .../adk/models/springai/ToolConverter.java | 229 +++++++------- .../models/springai/ToolExecutionMode.java | 25 ++ .../SpringAIAutoConfiguration.java | 38 ++- .../properties/SpringAIProperties.java | 26 ++ ...itional-spring-configuration-metadata.json | 8 +- .../springai/SpringAIIntegrationTest.java | 71 ++++- .../adk/models/springai/SpringAITest.java | 150 +++++++++ .../ToolConverterArgumentProcessingTest.java | 203 ++++-------- .../models/springai/ToolConverterTest.java | 134 ++++++++ .../SpringAIAutoConfigurationTest.java | 48 +++ 14 files changed, 975 insertions(+), 339 deletions(-) create mode 100644 contrib/spring-ai/src/main/java/com/google/adk/models/springai/AdkToolContextResolver.java create mode 100644 contrib/spring-ai/src/main/java/com/google/adk/models/springai/ToolExecutionMode.java diff --git a/contrib/spring-ai/README.md b/contrib/spring-ai/README.md index 0ce7de4fe..adf8de942 100644 --- a/contrib/spring-ai/README.md +++ b/contrib/spring-ai/README.md @@ -327,6 +327,9 @@ adk: observability: enabled: true metrics-enabled: true + tool-execution: + # Default: ADK executes tool calls with its InvocationContext-backed ToolContext. + mode: ADK_MANAGED ``` ## Architecture @@ -399,15 +402,16 @@ Converts between ADK tools and Spring AI function calling format. **Key Features:** - Converts ADK `BaseTool` to Spring AI `ToolCallback` - Schema conversion from ADK format to Spring AI JSON schema -- Intelligent argument processing for different provider formats - **Function Schema Registration:** Properly registers JSON schemas with Spring AI using `inputSchema()` method +- Configurable tool execution ownership: `ADK_MANAGED` (default) or `SPRING_AI_MANAGED` +- Never executes an ADK tool with a hard-coded `null` `ToolContext` - Debug logging for troubleshooting function calling issues **Function Calling Flow:** 1. ADK `FunctionDeclaration` → Spring AI `FunctionToolCallback` 2. ADK schema → JSON schema string -3. Runtime argument conversion and validation -4. Tool execution and result serialization +3. In `ADK_MANAGED` mode, Spring AI returns the function call and ADK executes it with the current invocation context +4. In `SPRING_AI_MANAGED` mode, Spring AI executes the callback using an application-provided `AdkToolContextResolver` #### 4. SpringAIEmbedding (SpringAIEmbedding.java) @@ -482,7 +486,9 @@ Flowable stream = springAI.generateContent(llmRequest, true); ### Function Calling -The library supports function calling through ADK tools: +The library supports function calling through ADK tools. Tool execution is owned by ADK by +default, preserving ADK's tool callbacks, confirmation flow, state changes, and +`InvocationContext`-backed `ToolContext`: ```java // Create agent with tools @@ -495,6 +501,25 @@ LlmAgent agent = LlmAgent.builder() // Tools are automatically converted to Spring AI format ``` +Applications that intentionally want Spring AI to own the tool-calling loop can opt in explicitly: + +```java +AdkToolContextResolver resolver = (tool, arguments, springAiContext) -> + resolveToolContextForTheCurrentRequest(); + +SpringAI springAI = new SpringAI( + chatModel, + "gpt-4o-mini", + ToolExecutionMode.SPRING_AI_MANAGED, + resolver); +``` + +The resolver must return a non-null ADK `ToolContext` for every call. Spring Boot applications can +select the same mode with `adk.spring-ai.tool-execution.mode=SPRING_AI_MANAGED` and must provide a +single `AdkToolContextResolver` bean. Configuration fails fast when that bean is missing. In this +mode, each streaming model turn is buffered so tool calls can be detected safely; chunks from the +final model turn are emitted after that turn completes. + ### Embedding Generation ```java @@ -594,6 +619,8 @@ adk: enabled: true metrics-enabled: true include-content: false + tool-execution: + mode: ADK_MANAGED ``` ### Auto-Configuration Beans @@ -679,8 +706,9 @@ The library provides comprehensive error handling through `SpringAIErrorMapper`: ### Function Calling 1. Ensure function schemas are properly defined in ADK tools 2. Test function calling with each provider separately -3. Handle provider-specific argument format differences -4. Use debug logging to troubleshoot function calling issues +3. Keep `ADK_MANAGED` unless the application deliberately supplies ADK context for Spring AI-managed execution +4. Never return `null` from an `AdkToolContextResolver` +5. Use debug logging to troubleshoot function calling issues ### Performance 1. Use streaming for long responses @@ -745,4 +773,4 @@ The library provides comprehensive error handling through `SpringAIErrorMapper`: - Java: 17+ - ADK: 0.3.1+ -This library provides a robust foundation for integrating Spring AI models with the ADK framework, offering enterprise-grade features like observability, error handling, and multi-provider support while maintaining the flexibility and power of both frameworks. \ No newline at end of file +This library provides a robust foundation for integrating Spring AI models with the ADK framework, offering enterprise-grade features like observability, error handling, and multi-provider support while maintaining the flexibility and power of both frameworks. diff --git a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/AdkToolContextResolver.java b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/AdkToolContextResolver.java new file mode 100644 index 000000000..85954833e --- /dev/null +++ b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/AdkToolContextResolver.java @@ -0,0 +1,44 @@ +/* + * Copyright 2026 Google LLC + * + * 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 + * + * http://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 com.google.adk.models.springai; + +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; +import java.util.Map; + +/** + * Resolves the ADK {@link ToolContext} used when Spring AI owns tool execution. + * + *

Spring AI's tool context does not contain ADK's invocation state. Applications opting into + * {@link ToolExecutionMode#SPRING_AI_MANAGED} must therefore supply this resolver explicitly. + */ +@FunctionalInterface +public interface AdkToolContextResolver { + + /** + * Resolves a non-null ADK tool context for one tool invocation. + * + * @param tool the ADK tool being called + * @param arguments the decoded tool arguments + * @param springAiToolContext the context supplied by Spring AI; it may be {@code null} when a + * callback is invoked directly without a context + * @return the ADK context to pass to the tool; must not be {@code null} + */ + ToolContext resolve( + BaseTool tool, + Map arguments, + org.springframework.ai.chat.model.ToolContext springAiToolContext); +} diff --git a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/MessageConverter.java b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/MessageConverter.java index 815ca2b7b..e91482f80 100644 --- a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/MessageConverter.java +++ b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/MessageConverter.java @@ -77,8 +77,14 @@ public class MessageConverter { private final ConfigMapper configMapper; public MessageConverter(ObjectMapper objectMapper) { + this( + objectMapper, + new ToolConverter(objectMapper, ToolExecutionMode.ADK_MANAGED, /* resolver= */ null)); + } + + public MessageConverter(ObjectMapper objectMapper, ToolConverter toolConverter) { this.objectMapper = objectMapper; - this.toolConverter = new ToolConverter(); + this.toolConverter = toolConverter; this.configMapper = new ConfigMapper(); } diff --git a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/SpringAI.java b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/SpringAI.java index ece9b1e2f..e60247856 100644 --- a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/SpringAI.java +++ b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/SpringAI.java @@ -26,11 +26,15 @@ import io.reactivex.rxjava3.core.BackpressureStrategy; import io.reactivex.rxjava3.core.Flowable; import java.util.Objects; +import java.util.concurrent.atomic.AtomicReference; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.MessageAggregator; import org.springframework.ai.chat.model.StreamingChatModel; import org.springframework.ai.chat.prompt.ChatOptions; import org.springframework.ai.chat.prompt.Prompt; +import org.springframework.ai.model.tool.ToolCallingManager; +import org.springframework.ai.model.tool.ToolExecutionResult; import reactor.core.publisher.Flux; /** @@ -41,67 +45,86 @@ */ public class SpringAI extends BaseLlm { + private static final int MAX_TOOL_CALL_ITERATIONS = 100; + private final ChatModel chatModel; private final StreamingChatModel streamingChatModel; private final ObjectMapper objectMapper; private final MessageConverter messageConverter; private final SpringAIObservabilityHandler observabilityHandler; + private final ToolExecutionMode toolExecutionMode; + private final ToolCallingManager toolCallingManager; + private final int maxToolCallIterations; public SpringAI(ChatModel chatModel) { - super(extractModelName(chatModel)); - this.chatModel = Objects.requireNonNull(chatModel, "chatModel cannot be null"); - this.streamingChatModel = - (chatModel instanceof StreamingChatModel) ? (StreamingChatModel) chatModel : null; - this.objectMapper = new ObjectMapper(); - this.messageConverter = new MessageConverter(objectMapper); - this.observabilityHandler = - new SpringAIObservabilityHandler(createDefaultObservabilityConfig()); + this(chatModel, extractModelName(chatModel)); } public SpringAI(ChatModel chatModel, String modelName) { - super(Objects.requireNonNull(modelName, "model name cannot be null")); - this.chatModel = Objects.requireNonNull(chatModel, "chatModel cannot be null"); - this.streamingChatModel = - (chatModel instanceof StreamingChatModel) ? (StreamingChatModel) chatModel : null; - this.objectMapper = new ObjectMapper(); - this.messageConverter = new MessageConverter(objectMapper); - this.observabilityHandler = - new SpringAIObservabilityHandler(createDefaultObservabilityConfig()); + this( + chatModel, + chatModel instanceof StreamingChatModel ? (StreamingChatModel) chatModel : null, + modelName, + createDefaultObservabilityConfig(), + ToolExecutionMode.ADK_MANAGED, + null, + MAX_TOOL_CALL_ITERATIONS); + } + + public SpringAI( + ChatModel chatModel, + String modelName, + ToolExecutionMode toolExecutionMode, + AdkToolContextResolver toolContextResolver) { + this( + chatModel, + chatModel instanceof StreamingChatModel ? (StreamingChatModel) chatModel : null, + modelName, + createDefaultObservabilityConfig(), + toolExecutionMode, + toolContextResolver, + MAX_TOOL_CALL_ITERATIONS); } public SpringAI(StreamingChatModel streamingChatModel) { - super(extractModelName(streamingChatModel)); - this.chatModel = - (streamingChatModel instanceof ChatModel) ? (ChatModel) streamingChatModel : null; - this.streamingChatModel = - Objects.requireNonNull(streamingChatModel, "streamingChatModel cannot be null"); - this.objectMapper = new ObjectMapper(); - this.messageConverter = new MessageConverter(objectMapper); - this.observabilityHandler = - new SpringAIObservabilityHandler(createDefaultObservabilityConfig()); + this(streamingChatModel, extractModelName(streamingChatModel)); } public SpringAI(StreamingChatModel streamingChatModel, String modelName) { - super(Objects.requireNonNull(modelName, "model name cannot be null")); - this.chatModel = - (streamingChatModel instanceof ChatModel) ? (ChatModel) streamingChatModel : null; - this.streamingChatModel = - Objects.requireNonNull(streamingChatModel, "streamingChatModel cannot be null"); - this.objectMapper = new ObjectMapper(); - this.messageConverter = new MessageConverter(objectMapper); - this.observabilityHandler = - new SpringAIObservabilityHandler(createDefaultObservabilityConfig()); + this( + streamingChatModel instanceof ChatModel ? (ChatModel) streamingChatModel : null, + streamingChatModel, + modelName, + createDefaultObservabilityConfig(), + ToolExecutionMode.ADK_MANAGED, + null, + MAX_TOOL_CALL_ITERATIONS); + } + + public SpringAI( + StreamingChatModel streamingChatModel, + String modelName, + ToolExecutionMode toolExecutionMode, + AdkToolContextResolver toolContextResolver) { + this( + streamingChatModel instanceof ChatModel ? (ChatModel) streamingChatModel : null, + streamingChatModel, + modelName, + createDefaultObservabilityConfig(), + toolExecutionMode, + toolContextResolver, + MAX_TOOL_CALL_ITERATIONS); } public SpringAI(ChatModel chatModel, StreamingChatModel streamingChatModel, String modelName) { - super(Objects.requireNonNull(modelName, "model name cannot be null")); - this.chatModel = Objects.requireNonNull(chatModel, "chatModel cannot be null"); - this.streamingChatModel = - Objects.requireNonNull(streamingChatModel, "streamingChatModel cannot be null"); - this.objectMapper = new ObjectMapper(); - this.messageConverter = new MessageConverter(objectMapper); - this.observabilityHandler = - new SpringAIObservabilityHandler(createDefaultObservabilityConfig()); + this( + chatModel, + streamingChatModel, + modelName, + createDefaultObservabilityConfig(), + ToolExecutionMode.ADK_MANAGED, + null, + MAX_TOOL_CALL_ITERATIONS); } public SpringAI( @@ -109,44 +132,116 @@ public SpringAI( StreamingChatModel streamingChatModel, String modelName, SpringAIProperties.Observability observabilityConfig) { - super(Objects.requireNonNull(modelName, "model name cannot be null")); - this.chatModel = Objects.requireNonNull(chatModel, "chatModel cannot be null"); - this.streamingChatModel = - Objects.requireNonNull(streamingChatModel, "streamingChatModel cannot be null"); - this.objectMapper = new ObjectMapper(); - this.messageConverter = new MessageConverter(objectMapper); - this.observabilityHandler = - new SpringAIObservabilityHandler( - Objects.requireNonNull(observabilityConfig, "observabilityConfig cannot be null")); + this( + chatModel, + streamingChatModel, + modelName, + observabilityConfig, + ToolExecutionMode.ADK_MANAGED, + null, + MAX_TOOL_CALL_ITERATIONS); + } + + public SpringAI( + ChatModel chatModel, + StreamingChatModel streamingChatModel, + String modelName, + SpringAIProperties.Observability observabilityConfig, + ToolExecutionMode toolExecutionMode, + AdkToolContextResolver toolContextResolver) { + this( + chatModel, + streamingChatModel, + modelName, + observabilityConfig, + toolExecutionMode, + toolContextResolver, + MAX_TOOL_CALL_ITERATIONS); } public SpringAI( ChatModel chatModel, String modelName, SpringAIProperties.Observability observabilityConfig) { - super(Objects.requireNonNull(modelName, "model name cannot be null")); - this.chatModel = Objects.requireNonNull(chatModel, "chatModel cannot be null"); - this.streamingChatModel = - (chatModel instanceof StreamingChatModel) ? (StreamingChatModel) chatModel : null; - this.objectMapper = new ObjectMapper(); - this.messageConverter = new MessageConverter(objectMapper); - this.observabilityHandler = - new SpringAIObservabilityHandler( - Objects.requireNonNull(observabilityConfig, "observabilityConfig cannot be null")); + this( + chatModel, + chatModel instanceof StreamingChatModel ? (StreamingChatModel) chatModel : null, + modelName, + observabilityConfig, + ToolExecutionMode.ADK_MANAGED, + null, + MAX_TOOL_CALL_ITERATIONS); + } + + public SpringAI( + ChatModel chatModel, + String modelName, + SpringAIProperties.Observability observabilityConfig, + ToolExecutionMode toolExecutionMode, + AdkToolContextResolver toolContextResolver) { + this( + chatModel, + chatModel instanceof StreamingChatModel ? (StreamingChatModel) chatModel : null, + modelName, + observabilityConfig, + toolExecutionMode, + toolContextResolver, + MAX_TOOL_CALL_ITERATIONS); } public SpringAI( StreamingChatModel streamingChatModel, String modelName, SpringAIProperties.Observability observabilityConfig) { + this( + streamingChatModel instanceof ChatModel ? (ChatModel) streamingChatModel : null, + streamingChatModel, + modelName, + observabilityConfig, + ToolExecutionMode.ADK_MANAGED, + null, + MAX_TOOL_CALL_ITERATIONS); + } + + public SpringAI( + StreamingChatModel streamingChatModel, + String modelName, + SpringAIProperties.Observability observabilityConfig, + ToolExecutionMode toolExecutionMode, + AdkToolContextResolver toolContextResolver) { + this( + streamingChatModel instanceof ChatModel ? (ChatModel) streamingChatModel : null, + streamingChatModel, + modelName, + observabilityConfig, + toolExecutionMode, + toolContextResolver, + MAX_TOOL_CALL_ITERATIONS); + } + + private SpringAI( + ChatModel chatModel, + StreamingChatModel streamingChatModel, + String modelName, + SpringAIProperties.Observability observabilityConfig, + ToolExecutionMode toolExecutionMode, + AdkToolContextResolver toolContextResolver, + int maxToolCallIterations) { super(Objects.requireNonNull(modelName, "model name cannot be null")); - this.chatModel = - (streamingChatModel instanceof ChatModel) ? (ChatModel) streamingChatModel : null; - this.streamingChatModel = - Objects.requireNonNull(streamingChatModel, "streamingChatModel cannot be null"); + if (chatModel == null && streamingChatModel == null) { + throw new NullPointerException("At least one chat model must be configured"); + } + this.chatModel = chatModel; + this.streamingChatModel = streamingChatModel; this.objectMapper = new ObjectMapper(); - this.messageConverter = new MessageConverter(objectMapper); + ToolConverter toolConverter = + new ToolConverter(objectMapper, toolExecutionMode, toolContextResolver); + this.messageConverter = new MessageConverter(objectMapper, toolConverter); this.observabilityHandler = new SpringAIObservabilityHandler( Objects.requireNonNull(observabilityConfig, "observabilityConfig cannot be null")); + this.toolExecutionMode = toolExecutionMode; + this.maxToolCallIterations = maxToolCallIterations; + this.toolCallingManager = + ToolCallingManager.builder().maxTotalToolCalls(maxToolCallIterations).build(); } @Override @@ -174,7 +269,7 @@ private Flowable generateContent(LlmRequest llmRequest) { Prompt prompt = messageConverter.toLlmPrompt(llmRequest, resolveDefaultOptions()); observabilityHandler.logRequest(prompt.toString(), model()); - ChatResponse chatResponse = chatModel.call(prompt); + ChatResponse chatResponse = callChatModel(prompt); LlmResponse llmResponse = messageConverter.toLlmResponse(chatResponse); observabilityHandler.logResponse(extractTextFromResponse(llmResponse), model()); @@ -204,7 +299,7 @@ private Flowable generateStreamingContent(LlmRequest llmRequest) { Prompt prompt = messageConverter.toLlmPrompt(llmRequest, resolveDefaultOptions()); observabilityHandler.logRequest(prompt.toString(), model()); - Flux responseFlux = streamingChatModel.stream(prompt); + Flux responseFlux = streamChatModel(prompt, 0); responseFlux .doOnError( @@ -251,6 +346,63 @@ private Flowable generateStreamingContent(LlmRequest llmRequest) { BackpressureStrategy.BUFFER); } + private ChatResponse callChatModel(Prompt initialPrompt) { + Prompt prompt = initialPrompt; + ChatResponse chatResponse = chatModel.call(prompt); + if (toolExecutionMode == ToolExecutionMode.ADK_MANAGED) { + return chatResponse; + } + + int iteration = 0; + while (chatResponse.hasToolCalls()) { + if (iteration++ >= maxToolCallIterations) { + throw new IllegalStateException( + "Spring AI tool execution exceeded " + maxToolCallIterations + " iterations"); + } + ToolExecutionResult toolExecutionResult = + toolCallingManager.executeToolCalls(prompt, chatResponse); + if (toolExecutionResult.returnDirect()) { + return new ChatResponse(ToolExecutionResult.buildGenerations(toolExecutionResult)); + } + prompt = new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()); + chatResponse = chatModel.call(prompt); + } + return chatResponse; + } + + private Flux streamChatModel(Prompt prompt, int iteration) { + if (toolExecutionMode == ToolExecutionMode.ADK_MANAGED) { + return streamingChatModel.stream(prompt); + } + if (iteration >= maxToolCallIterations) { + return Flux.error( + new IllegalStateException( + "Spring AI tool execution exceeded " + maxToolCallIterations + " iterations")); + } + + AtomicReference aggregatedResponse = new AtomicReference<>(); + return new MessageAggregator() + .aggregate(streamingChatModel.stream(prompt), aggregatedResponse::set) + .collectList() + .flatMapMany( + chunks -> { + ChatResponse chatResponse = aggregatedResponse.get(); + if (chatResponse == null || !chatResponse.hasToolCalls()) { + return Flux.fromIterable(chunks); + } + + ToolExecutionResult toolExecutionResult = + toolCallingManager.executeToolCalls(prompt, chatResponse); + if (toolExecutionResult.returnDirect()) { + return Flux.just( + new ChatResponse(ToolExecutionResult.buildGenerations(toolExecutionResult))); + } + Prompt nextPrompt = + new Prompt(toolExecutionResult.conversationHistory(), prompt.getOptions()); + return streamChatModel(nextPrompt, iteration + 1); + }); + } + @Override public BaseLlmConnection connect(LlmRequest llmRequest) { throw new UnsupportedOperationException( @@ -283,7 +435,7 @@ private static String extractModelName(Object model) { return className.toLowerCase().replace("chatmodel", "").replace("model", ""); } - private SpringAIProperties.Observability createDefaultObservabilityConfig() { + private static SpringAIProperties.Observability createDefaultObservabilityConfig() { SpringAIProperties.Observability config = new SpringAIProperties.Observability(); config.setEnabled(true); config.setMetricsEnabled(true); diff --git a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/ToolConverter.java b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/ToolConverter.java index 4012ee5d6..d539cc56b 100644 --- a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/ToolConverter.java +++ b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/ToolConverter.java @@ -15,7 +15,10 @@ */ package com.google.adk.models.springai; +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.ObjectMapper; import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.Schema; import com.google.genai.types.Type; @@ -23,7 +26,10 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Set; +import java.util.function.BiFunction; +import java.util.function.Function; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import org.springframework.ai.tool.ToolCallback; @@ -40,6 +46,34 @@ public class ToolConverter { private static final Logger logger = LoggerFactory.getLogger(ToolConverter.class); + private final ObjectMapper objectMapper; + private final ToolExecutionMode toolExecutionMode; + private final AdkToolContextResolver toolContextResolver; + + /** Creates a converter that exposes tool definitions while ADK owns tool execution. */ + public ToolConverter() { + this(new ObjectMapper(), ToolExecutionMode.ADK_MANAGED, null); + } + + /** Creates a converter with an explicit tool execution owner and context resolver. */ + public ToolConverter( + ObjectMapper objectMapper, + ToolExecutionMode toolExecutionMode, + AdkToolContextResolver toolContextResolver) { + this.objectMapper = Objects.requireNonNull(objectMapper, "objectMapper must not be null"); + this.toolExecutionMode = + Objects.requireNonNull(toolExecutionMode, "toolExecutionMode must not be null"); + if (toolExecutionMode == ToolExecutionMode.SPRING_AI_MANAGED && toolContextResolver == null) { + throw new IllegalArgumentException( + "An AdkToolContextResolver is required when tool execution is SPRING_AI_MANAGED"); + } + this.toolContextResolver = toolContextResolver; + } + + public ToolExecutionMode getToolExecutionMode() { + return toolExecutionMode; + } + /** * Creates a tool registry from ADK tools for internal tracking. * @@ -120,143 +154,128 @@ public List convertToSpringAiTools(Map tools) { if (tool.declaration().isPresent()) { FunctionDeclaration declaration = tool.declaration().get(); - // Create a ToolCallback that wraps the ADK tool - // Create a Function that takes Map input and calls the ADK tool - java.util.function.Function, String> toolFunction = - args -> { - try { - logger.debug("Spring AI calling tool '{}'", tool.name()); - logger.debug("Raw args from Spring AI: {}", args); - logger.debug("Args type: {}", args.getClass().getName()); - logger.debug("Args keys: {}", args.keySet()); - for (Map.Entry entry : args.entrySet()) { - logger.debug( - " {} -> {} ({})", - entry.getKey(), - entry.getValue(), - entry.getValue().getClass().getName()); - } - - // Handle different argument formats that Spring AI might pass - Map processedArgs = processArguments(args, declaration); - logger.debug("Processed args for ADK: {}", processedArgs); - - // Call the ADK tool and wait for the result - Map result = tool.runAsync(processedArgs, null).blockingGet(); - // Convert result back to JSON string - return new com.fasterxml.jackson.databind.ObjectMapper().writeValueAsString(result); - } catch (Exception e) { - throw new RuntimeException("Tool execution failed: " + e.getMessage(), e); - } - }; - - FunctionToolCallback.Builder callbackBuilder = - FunctionToolCallback.builder(tool.name(), toolFunction).description(tool.description()); - - // Convert ADK schema to Spring AI schema if available - if (declaration.parameters().isPresent()) { - // Use Map.class to indicate the input is an object/map - callbackBuilder.inputType(Map.class); - - // Convert ADK schema to Spring AI JSON schema format - Map springAiSchema = - convertSchemaToSpringAi(declaration.parameters().get()); - logger.debug("Generated Spring AI schema for {}: {}", tool.name(), springAiSchema); - - // Provide the schema as JSON string using inputSchema method - try { - String schemaJson = - new com.fasterxml.jackson.databind.ObjectMapper() - .writeValueAsString(springAiSchema); - callbackBuilder.inputSchema(schemaJson); - logger.debug("Set input schema JSON: {}", schemaJson); - } catch (Exception e) { - logger.error("Error serializing schema to JSON: {}", e.getMessage(), e); - } - } else if (declaration.parametersJsonSchema().isPresent()) { - callbackBuilder.inputType(Map.class); - try { - String schemaJson = - new com.fasterxml.jackson.databind.ObjectMapper() - .writeValueAsString(declaration.parametersJsonSchema().get()); - callbackBuilder.inputSchema(schemaJson); - logger.debug("Set input schema JSON from parametersJsonSchema: {}", schemaJson); - } catch (Exception e) { - logger.error("Error serializing parametersJsonSchema to JSON: {}", e.getMessage(), e); - } + if (toolExecutionMode == ToolExecutionMode.ADK_MANAGED) { + // Spring AI still requires callbacks to expose tool definitions to the model. The + // callback intentionally has no side effect: ADK will execute the returned function call + // later with the InvocationContext-backed ToolContext. + Function, String> definitionOnlyCallback = arguments -> ""; + toolCallbacks.add( + configureCallback( + tool, + declaration, + FunctionToolCallback.builder(tool.name(), definitionOnlyCallback)) + .build()); + } else { + BiFunction< + Map, + org.springframework.ai.chat.model.ToolContext, + Map> + executableCallback = + (arguments, springAiToolContext) -> { + Map processedArguments = + processArguments(arguments, declaration); + ToolContext adkToolContext = + Objects.requireNonNull( + toolContextResolver.resolve( + tool, processedArguments, springAiToolContext), + "AdkToolContextResolver returned null for tool " + tool.name()); + return tool.runAsync(processedArguments, adkToolContext).blockingGet(); + }; + toolCallbacks.add( + configureCallback( + tool, + declaration, + FunctionToolCallback.builder(tool.name(), executableCallback)) + .build()); } - - toolCallbacks.add(callbackBuilder.build()); } } return toolCallbacks; } - /** - * Process arguments from Spring AI format to ADK format. Spring AI might pass arguments in - * different formats depending on the provider. - */ + private FunctionToolCallback.Builder, O> configureCallback( + BaseTool tool, + FunctionDeclaration declaration, + FunctionToolCallback.Builder, O> callbackBuilder) { + callbackBuilder.description(tool.description()).inputType(Map.class); + + if (declaration.parameters().isPresent()) { + Map springAiSchema = convertSchemaToSpringAi(declaration.parameters().get()); + logger.debug("Generated Spring AI schema for {}: {}", tool.name(), springAiSchema); + configureInputSchema(callbackBuilder, springAiSchema, tool.name()); + } else if (declaration.parametersJsonSchema().isPresent()) { + configureInputSchema(callbackBuilder, declaration.parametersJsonSchema().get(), tool.name()); + } + + return callbackBuilder; + } + + private void configureInputSchema( + FunctionToolCallback.Builder, O> callbackBuilder, + Object schema, + String toolName) { + try { + String schemaJson = objectMapper.writeValueAsString(schema); + callbackBuilder.inputSchema(schemaJson); + logger.debug("Set input schema JSON for {}: {}", toolName, schemaJson); + } catch (JsonProcessingException e) { + throw new IllegalArgumentException( + "Unable to serialize input schema for tool " + toolName, e); + } + } + + /** Normalizes provider-specific argument wrappers to the ADK tool's declared parameters. */ private Map processArguments( - Map args, FunctionDeclaration declaration) { + Map arguments, FunctionDeclaration declaration) { if (declaration.parameters().isPresent()) { - var schema = declaration.parameters().get(); + Schema schema = declaration.parameters().get(); if (schema.properties().isPresent()) { - return normalizeArguments(args, schema.properties().get().keySet()); + return normalizeArguments(arguments, schema.properties().get().keySet()); } } else if (declaration.parametersJsonSchema().isPresent()) { try { @SuppressWarnings("unchecked") - Map schemaMap = - new com.fasterxml.jackson.databind.ObjectMapper() - .convertValue(declaration.parametersJsonSchema().get(), Map.class); - Object propertiesObj = schemaMap.get("properties"); - if (propertiesObj instanceof Map) { - @SuppressWarnings("unchecked") - Set expectedParams = ((Map) propertiesObj).keySet(); - return normalizeArguments(args, expectedParams); + Map schema = + objectMapper.convertValue(declaration.parametersJsonSchema().get(), Map.class); + Object properties = schema.get("properties"); + if (properties instanceof Map propertiesMap) { + return normalizeArguments(arguments, propertiesMap.keySet()); } - } catch (Exception e) { + } catch (IllegalArgumentException e) { logger.warn( "Error processing parametersJsonSchema for argument mapping: {}", e.getMessage()); } } - // If no processing worked, return original args and let ADK handle the error - return args; + return arguments; } private Map normalizeArguments( - Map args, Set expectedParams) { - // Check if all expected parameters are present at the top level - boolean allParamsPresent = expectedParams.stream().allMatch(args::containsKey); - if (allParamsPresent) { - return args; + Map arguments, Set expectedParameters) { + if (expectedParameters.stream().allMatch(arguments::containsKey)) { + return arguments; } - // Check if arguments are nested under a single key (common pattern) - if (args.size() == 1) { - var singleValue = args.values().iterator().next(); - if (singleValue instanceof Map) { + if (arguments.size() == 1) { + Object singleValue = arguments.values().iterator().next(); + if (singleValue instanceof Map nestedArguments + && expectedParameters.stream().allMatch(nestedArguments::containsKey)) { @SuppressWarnings("unchecked") - Map nestedArgs = (Map) singleValue; - boolean allNestedParamsPresent = expectedParams.stream().allMatch(nestedArgs::containsKey); - if (allNestedParamsPresent) { - return nestedArgs; - } + Map normalizedArguments = (Map) nestedArguments; + return normalizedArguments; } } - // Check if we have a single parameter function and got a direct value - if (expectedParams.size() == 1) { - String expectedParam = expectedParams.iterator().next(); - if (args.size() == 1 && !args.containsKey(expectedParam)) { - Object singleValue = args.values().iterator().next(); - return Map.of(expectedParam, singleValue); + if (expectedParameters.size() == 1 && arguments.size() == 1) { + Object expectedParameter = expectedParameters.iterator().next(); + if (expectedParameter instanceof String parameterName + && !arguments.containsKey(parameterName)) { + return Map.of(parameterName, arguments.values().iterator().next()); } } - return args; + return arguments; } /** Simple metadata holder for tool information. */ diff --git a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/ToolExecutionMode.java b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/ToolExecutionMode.java new file mode 100644 index 000000000..14d2d278c --- /dev/null +++ b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/ToolExecutionMode.java @@ -0,0 +1,25 @@ +/* + * Copyright 2026 Google LLC + * + * 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 + * + * http://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 com.google.adk.models.springai; + +/** Selects which framework owns the tool-calling lifecycle. */ +public enum ToolExecutionMode { + /** ADK receives model tool calls and executes them with its current invocation context. */ + ADK_MANAGED, + + /** Spring AI executes tool calls before returning the final model response to ADK. */ + SPRING_AI_MANAGED +} diff --git a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/autoconfigure/SpringAIAutoConfiguration.java b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/autoconfigure/SpringAIAutoConfiguration.java index 7a312ca88..853367c84 100644 --- a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/autoconfigure/SpringAIAutoConfiguration.java +++ b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/autoconfigure/SpringAIAutoConfiguration.java @@ -15,6 +15,7 @@ */ package com.google.adk.models.springai.autoconfigure; +import com.google.adk.models.springai.AdkToolContextResolver; import com.google.adk.models.springai.SpringAI; import com.google.adk.models.springai.SpringAIEmbedding; import com.google.adk.models.springai.properties.SpringAIProperties; @@ -23,6 +24,7 @@ import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.model.StreamingChatModel; import org.springframework.ai.embedding.EmbeddingModel; +import org.springframework.beans.factory.ObjectProvider; import org.springframework.boot.autoconfigure.AutoConfiguration; import org.springframework.boot.autoconfigure.condition.ConditionalOnBean; import org.springframework.boot.autoconfigure.condition.ConditionalOnClass; @@ -83,7 +85,10 @@ public class SpringAIAutoConfiguration { @ConditionalOnMissingBean(SpringAI.class) @ConditionalOnBean({ChatModel.class, StreamingChatModel.class}) public SpringAI springAIWithBothModels( - ChatModel chatModel, StreamingChatModel streamingChatModel, SpringAIProperties properties) { + ChatModel chatModel, + StreamingChatModel streamingChatModel, + SpringAIProperties properties, + ObjectProvider toolContextResolverProvider) { String modelName = determineModelName(chatModel, properties); logger.info( @@ -91,7 +96,13 @@ public SpringAI springAIWithBothModels( modelName); validateConfiguration(properties); - return new SpringAI(chatModel, streamingChatModel, modelName, properties.getObservability()); + return new SpringAI( + chatModel, + streamingChatModel, + modelName, + properties.getObservability(), + properties.getToolExecution().getMode(), + toolContextResolverProvider.getIfAvailable()); } /** @@ -104,13 +115,21 @@ public SpringAI springAIWithBothModels( @Bean @ConditionalOnMissingBean(SpringAI.class) @ConditionalOnBean(ChatModel.class) - public SpringAI springAIWithChatModel(ChatModel chatModel, SpringAIProperties properties) { + public SpringAI springAIWithChatModel( + ChatModel chatModel, + SpringAIProperties properties, + ObjectProvider toolContextResolverProvider) { String modelName = determineModelName(chatModel, properties); logger.info("Auto-configuring SpringAI with ChatModel only. Model: {}", modelName); validateConfiguration(properties); - return new SpringAI(chatModel, modelName, properties.getObservability()); + return new SpringAI( + chatModel, + modelName, + properties.getObservability(), + properties.getToolExecution().getMode(), + toolContextResolverProvider.getIfAvailable()); } /** @@ -124,13 +143,20 @@ public SpringAI springAIWithChatModel(ChatModel chatModel, SpringAIProperties pr @ConditionalOnMissingBean({SpringAI.class, ChatModel.class}) @ConditionalOnBean(StreamingChatModel.class) public SpringAI springAIWithStreamingModel( - StreamingChatModel streamingChatModel, SpringAIProperties properties) { + StreamingChatModel streamingChatModel, + SpringAIProperties properties, + ObjectProvider toolContextResolverProvider) { String modelName = determineModelName(streamingChatModel, properties); logger.info("Auto-configuring SpringAI with StreamingChatModel only. Model: {}", modelName); validateConfiguration(properties); - return new SpringAI(streamingChatModel, modelName, properties.getObservability()); + return new SpringAI( + streamingChatModel, + modelName, + properties.getObservability(), + properties.getToolExecution().getMode(), + toolContextResolverProvider.getIfAvailable()); } /** diff --git a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/properties/SpringAIProperties.java b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/properties/SpringAIProperties.java index 5e972ebcd..403910bc4 100644 --- a/contrib/spring-ai/src/main/java/com/google/adk/models/springai/properties/SpringAIProperties.java +++ b/contrib/spring-ai/src/main/java/com/google/adk/models/springai/properties/SpringAIProperties.java @@ -15,6 +15,7 @@ */ package com.google.adk.models.springai.properties; +import com.google.adk.models.springai.ToolExecutionMode; import jakarta.annotation.Nullable; import jakarta.validation.constraints.DecimalMax; import jakarta.validation.constraints.DecimalMin; @@ -67,6 +68,9 @@ public class SpringAIProperties { /** Observability settings. */ private Observability observability = new Observability(); + /** Tool execution ownership settings. */ + private ToolExecution toolExecution = new ToolExecution(); + public String getModel() { return model; } @@ -123,6 +127,14 @@ public void setObservability(Observability observability) { this.observability = observability; } + public ToolExecution getToolExecution() { + return toolExecution; + } + + public void setToolExecution(ToolExecution toolExecution) { + this.toolExecution = toolExecution; + } + /** Configuration validation settings. */ public static class Validation { /** Whether to enable strict validation of configuration parameters. */ @@ -183,4 +195,18 @@ public void setMetricsEnabled(boolean metricsEnabled) { this.metricsEnabled = metricsEnabled; } } + + /** Tool execution ownership configuration. */ + public static class ToolExecution { + /** Framework responsible for executing model tool calls. */ + private ToolExecutionMode mode = ToolExecutionMode.ADK_MANAGED; + + public ToolExecutionMode getMode() { + return mode; + } + + public void setMode(ToolExecutionMode mode) { + this.mode = mode; + } + } } diff --git a/contrib/spring-ai/src/main/resources/META-INF/additional-spring-configuration-metadata.json b/contrib/spring-ai/src/main/resources/META-INF/additional-spring-configuration-metadata.json index 29f3a7e8e..398bd5be0 100644 --- a/contrib/spring-ai/src/main/resources/META-INF/additional-spring-configuration-metadata.json +++ b/contrib/spring-ai/src/main/resources/META-INF/additional-spring-configuration-metadata.json @@ -53,6 +53,12 @@ "description": "Whether to collect metrics.", "defaultValue": true }, + { + "name": "adk.spring-ai.tool-execution.mode", + "type": "com.google.adk.models.springai.ToolExecutionMode", + "description": "Framework responsible for executing model tool calls. SPRING_AI_MANAGED requires an AdkToolContextResolver bean.", + "defaultValue": "ADK_MANAGED" + }, { "name": "adk.spring-ai.auto-configuration.enabled", "type": "java.lang.Boolean", @@ -60,4 +66,4 @@ "defaultValue": true } ] -} \ No newline at end of file +} diff --git a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAIIntegrationTest.java b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAIIntegrationTest.java index 328df0415..39964cf09 100644 --- a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAIIntegrationTest.java +++ b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAIIntegrationTest.java @@ -20,16 +20,23 @@ import com.google.adk.agents.LlmAgent; import com.google.adk.agents.RunConfig; import com.google.adk.events.Event; -import com.google.adk.models.springai.integrations.tools.WeatherTool; import com.google.adk.runner.InMemoryRunner; import com.google.adk.runner.Runner; import com.google.adk.sessions.Session; -import com.google.adk.tools.FunctionTool; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; import com.google.genai.types.Content; +import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Single; import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.ToolResponseMessage; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.model.Generation; @@ -108,7 +115,39 @@ public ChatResponse call(Prompt prompt) { @Test void testAgentWithToolsUsingDummyModel() { - // given - Create a dummy ChatModel that simulates tool calling + AtomicInteger toolExecutions = new AtomicInteger(); + AtomicReference receivedToolContext = new AtomicReference<>(); + FunctionDeclaration declaration = + FunctionDeclaration.builder() + .name("get_weather") + .description("Get the weather for a city") + .parametersJsonSchema( + Map.of( + "type", + "object", + "properties", + Map.of("city", Map.of("type", "string")), + "required", + List.of("city"))) + .build(); + BaseTool weatherTool = + new BaseTool("get_weather", "Get the weather for a city") { + @Override + public Optional declaration() { + return Optional.of(declaration); + } + + @Override + public Single> runAsync( + Map args, ToolContext toolContext) { + toolExecutions.incrementAndGet(); + receivedToolContext.set(toolContext); + return Single.just(Map.of("forecast", "Sunny in " + args.get("city"))); + } + }; + + // The model returns a real tool call on its first turn and consumes ADK's tool response on its + // second turn. ChatModel dummyChatModel = new ChatModel() { private int callCount = 0; @@ -119,14 +158,22 @@ public ChatResponse call(Prompt prompt) { AssistantMessage message; if (callCount == 1) { - // First call - simulate asking for weather - message = new AssistantMessage("I need to check the weather for Paris."); - } else { - // Subsequent calls - provide final answer message = - new AssistantMessage( - "The weather in Paris is beautiful and sunny with temperatures from 10°C in" - + " the morning up to 24°C in the afternoon."); + AssistantMessage.builder() + .content("") + .toolCalls( + List.of( + new AssistantMessage.ToolCall( + "weather-call-1", + "function", + "get_weather", + "{\"city\":\"Paris\"}"))) + .build(); + } else { + assertTrue( + prompt.getInstructions().stream() + .anyMatch(instruction -> instruction instanceof ToolResponseMessage)); + message = new AssistantMessage("The weather in Paris is sunny."); } Generation generation = new Generation(message); @@ -146,7 +193,7 @@ public ChatResponse call(Prompt prompt) { If asked about the weather forecast for a city, you MUST call the `getWeather` function. """) - .tools(FunctionTool.create(WeatherTool.class, "getWeather")) + .tools(weatherTool) .build(); // when @@ -186,6 +233,8 @@ public ChatResponse call(Prompt prompt) { && event.content().get().text().toLowerCase().contains("paris")); assertTrue(hasParisResponse, "Should have a response mentioning Paris"); + assertEquals(1, toolExecutions.get(), "ADK should execute the tool exactly once"); + assertNotNull(receivedToolContext.get(), "ADK must pass a non-null ToolContext"); } @Test diff --git a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAITest.java b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAITest.java index 51fe60abb..4be2a63dc 100644 --- a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAITest.java +++ b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/SpringAITest.java @@ -23,15 +23,24 @@ import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; import com.google.genai.types.Content; +import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Single; import io.reactivex.rxjava3.subscribers.TestSubscriber; import java.time.Duration; +import java.util.ArrayList; import java.util.List; +import java.util.Map; +import java.util.Optional; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.messages.ToolResponseMessage; import org.springframework.ai.chat.model.ChatModel; import org.springframework.ai.chat.model.ChatResponse; import org.springframework.ai.chat.model.Generation; @@ -256,6 +265,121 @@ void testGenerateContentWithEmptyResponse() { assertThat(response.content()).isEmpty(); } + @Test + void adkManagedModeReturnsToolCallWithoutExecutingTool() { + AtomicInteger executions = new AtomicInteger(); + BaseTool tool = contextTool(executions, null); + AssistantMessage toolCallMessage = + AssistantMessage.builder() + .content("") + .toolCalls( + List.of( + new AssistantMessage.ToolCall( + "call-1", "function", tool.name(), "{\"value\":\"hello\"}"))) + .build(); + when(mockChatModel.call(any(Prompt.class))) + .thenReturn(new ChatResponse(List.of(new Generation(toolCallMessage)))); + LlmRequest request = + LlmRequest.builder() + .contents(testRequest.contents()) + .tools(Map.of(tool.name(), tool)) + .build(); + + LlmResponse response = + new SpringAI(mockChatModel).generateContent(request, false).blockingFirst(); + + assertThat(response.content()).isPresent(); + assertThat(response.content().get().parts().get().get(0).functionCall()).isPresent(); + assertThat(executions).hasValue(0); + } + + @Test + void springAiManagedModeExecutesToolOnceAndReturnsFinalModelResponse() { + AtomicInteger executions = new AtomicInteger(); + ToolContext resolvedContext = mock(ToolContext.class); + BaseTool tool = contextTool(executions, resolvedContext); + AssistantMessage toolCallMessage = + AssistantMessage.builder() + .content("") + .toolCalls( + List.of( + new AssistantMessage.ToolCall( + "call-1", "function", tool.name(), "{\"value\":\"hello\"}"))) + .build(); + List prompts = new ArrayList<>(); + ChatModel twoTurnModel = + prompt -> { + prompts.add(prompt); + if (prompts.size() == 1) { + return new ChatResponse(List.of(new Generation(toolCallMessage))); + } + return new ChatResponse(List.of(new Generation(new AssistantMessage("final answer")))); + }; + LlmRequest request = + LlmRequest.builder() + .contents(testRequest.contents()) + .tools(Map.of(tool.name(), tool)) + .build(); + SpringAI springAI = + new SpringAI( + twoTurnModel, + "test-model", + ToolExecutionMode.SPRING_AI_MANAGED, + (baseTool, arguments, springAiContext) -> resolvedContext); + + LlmResponse response = springAI.generateContent(request, false).blockingFirst(); + + assertThat(response.content().orElseThrow().text()).contains("final answer"); + assertThat(executions).hasValue(1); + assertThat(prompts).hasSize(2); + assertThat(prompts.get(1).getInstructions()) + .anyMatch(message -> message instanceof ToolResponseMessage); + } + + @Test + void springAiManagedModeExecutesToolOnceWhenStreaming() { + AtomicInteger executions = new AtomicInteger(); + AtomicInteger modelCalls = new AtomicInteger(); + ToolContext resolvedContext = mock(ToolContext.class); + BaseTool tool = contextTool(executions, resolvedContext); + AssistantMessage toolCallMessage = + AssistantMessage.builder() + .content("") + .toolCalls( + List.of( + new AssistantMessage.ToolCall( + "call-1", "function", tool.name(), "{\"value\":\"hello\"}"))) + .build(); + StreamingChatModel twoTurnStreamingModel = + prompt -> { + if (modelCalls.incrementAndGet() == 1) { + return Flux.just(new ChatResponse(List.of(new Generation(toolCallMessage)))); + } + return Flux.just( + createStreamingChatResponse("final "), createStreamingChatResponse("answer")); + }; + LlmRequest request = + LlmRequest.builder() + .contents(testRequest.contents()) + .tools(Map.of(tool.name(), tool)) + .build(); + SpringAI springAI = + new SpringAI( + twoTurnStreamingModel, + "test-model", + ToolExecutionMode.SPRING_AI_MANAGED, + (baseTool, arguments, springAiContext) -> resolvedContext); + + List responses = springAI.generateContent(request, true).toList().blockingGet(); + + assertThat(responses).hasSize(2); + assertThat(responses) + .extracting(response -> response.content().orElseThrow().text()) + .containsExactly("final ", "answer"); + assertThat(executions).hasValue(1); + assertThat(modelCalls).hasValue(2); + } + @Test void testGenerateContentStreamingBackpressure() { // Create a large number of streaming responses to test backpressure @@ -281,4 +405,30 @@ private ChatResponse createStreamingChatResponse(String text) { Generation generation = new Generation(assistantMessage); return new ChatResponse(List.of(generation)); } + + private BaseTool contextTool(AtomicInteger executions, ToolContext expectedContext) { + FunctionDeclaration declaration = + FunctionDeclaration.builder() + .name("context_tool") + .description("context tool") + .parametersJsonSchema( + Map.of("type", "object", "properties", Map.of("value", Map.of("type", "string")))) + .build(); + return new BaseTool("context_tool", "context tool") { + @Override + public Optional declaration() { + return Optional.of(declaration); + } + + @Override + public Single> runAsync( + Map args, ToolContext toolContext) { + if (expectedContext != null) { + assertThat(toolContext).isSameAs(expectedContext); + } + executions.incrementAndGet(); + return Single.just(Map.of("echo", args.get("value"))); + } + }; + } } diff --git a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/ToolConverterArgumentProcessingTest.java b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/ToolConverterArgumentProcessingTest.java index 77b988837..2eb5495a0 100644 --- a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/ToolConverterArgumentProcessingTest.java +++ b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/ToolConverterArgumentProcessingTest.java @@ -18,195 +18,118 @@ import static org.assertj.core.api.Assertions.assertThat; import com.google.adk.tools.FunctionTool; +import com.google.genai.types.FunctionDeclaration; import java.lang.reflect.Method; -import java.util.List; import java.util.Map; import org.junit.jupiter.api.Test; -import org.springframework.ai.tool.ToolCallback; -/** Test argument processing logic in ToolConverter. */ +/** Tests provider-specific argument normalization in {@link ToolConverter}. */ class ToolConverterArgumentProcessingTest { @Test - void testArgumentProcessingWithCorrectFormat() throws Exception { - // Create tool converter and tool + void declaredSchemaLeavesCorrectArgumentsUnchanged() throws Exception { ToolConverter converter = new ToolConverter(); FunctionTool tool = FunctionTool.create(WeatherTools.class, "getWeatherInfo"); - Map tools = Map.of("getWeatherInfo", tool); + Map arguments = Map.of("location", "San Francisco"); - // Convert to Spring AI format - List toolCallbacks = converter.convertToSpringAiTools(tools); - assertThat(toolCallbacks).hasSize(1); - - // Test with correct argument format - ToolCallback callback = toolCallbacks.get(0); - Method processArguments = getProcessArgumentsMethod(converter); - - Map correctArgs = Map.of("location", "San Francisco"); - Map processedArgs = - invokeProcessArguments(processArguments, converter, correctArgs, tool.declaration().get()); - - assertThat(processedArgs).isEqualTo(correctArgs); + assertThat(process(converter, arguments, tool.declaration().orElseThrow())) + .isEqualTo(arguments); } @Test - void testArgumentProcessingWithNestedFormat() throws Exception { + void declaredSchemaUnwrapsNestedArguments() throws Exception { ToolConverter converter = new ToolConverter(); FunctionTool tool = FunctionTool.create(WeatherTools.class, "getWeatherInfo"); - Method processArguments = getProcessArgumentsMethod(converter); - - // Test with nested arguments - Map nestedArgs = Map.of("args", Map.of("location", "San Francisco")); - Map processedArgs = - invokeProcessArguments(processArguments, converter, nestedArgs, tool.declaration().get()); - - assertThat(processedArgs).containsEntry("location", "San Francisco"); + assertThat( + process( + converter, + Map.of("args", Map.of("location", "San Francisco")), + tool.declaration().orElseThrow())) + .containsEntry("location", "San Francisco"); } @Test - void testArgumentProcessingWithDirectValue() throws Exception { + void declaredSchemaMapsDirectValueToSingleParameter() throws Exception { ToolConverter converter = new ToolConverter(); FunctionTool tool = FunctionTool.create(WeatherTools.class, "getWeatherInfo"); - Method processArguments = getProcessArgumentsMethod(converter); - - // Test with single direct value (wrong key name) - Map directValueArgs = Map.of("value", "San Francisco"); - Map processedArgs = - invokeProcessArguments( - processArguments, converter, directValueArgs, tool.declaration().get()); - - // Should map the single value to the expected parameter name - assertThat(processedArgs).containsEntry("location", "San Francisco"); + assertThat( + process(converter, Map.of("value", "San Francisco"), tool.declaration().orElseThrow())) + .containsEntry("location", "San Francisco"); } @Test - void testArgumentProcessingWithNoMatch() throws Exception { + void declaredSchemaLeavesUnmatchedArgumentsUnchanged() throws Exception { ToolConverter converter = new ToolConverter(); FunctionTool tool = FunctionTool.create(WeatherTools.class, "getWeatherInfo"); + Map arguments = Map.of("city", "San Francisco", "country", "USA"); - Method processArguments = getProcessArgumentsMethod(converter); - - // Test with completely wrong format - Map wrongArgs = Map.of("city", "San Francisco", "country", "USA"); - Map processedArgs = - invokeProcessArguments(processArguments, converter, wrongArgs, tool.declaration().get()); - - // Should return original args when no processing applies - assertThat(processedArgs).isEqualTo(wrongArgs); + assertThat(process(converter, arguments, tool.declaration().orElseThrow())) + .isEqualTo(arguments); } - private Method getProcessArgumentsMethod(ToolConverter converter) throws Exception { - Method method = - ToolConverter.class.getDeclaredMethod( - "processArguments", Map.class, com.google.genai.types.FunctionDeclaration.class); - method.setAccessible(true); - return method; - } + @Test + void jsonSchemaLeavesCorrectArgumentsUnchanged() throws Exception { + ToolConverter converter = new ToolConverter(); + Map arguments = Map.of("location", "San Francisco"); - @SuppressWarnings("unchecked") - private Map invokeProcessArguments( - Method method, - ToolConverter converter, - Map args, - com.google.genai.types.FunctionDeclaration declaration) - throws Exception { - return (Map) method.invoke(converter, args, declaration); + assertThat(process(converter, arguments, jsonSchemaDeclaration())).isEqualTo(arguments); } @Test - void testArgumentProcessingWithParametersJsonSchema_correctFormat() throws Exception { + void jsonSchemaUnwrapsNestedArguments() throws Exception { ToolConverter converter = new ToolConverter(); - Method processArguments = getProcessArgumentsMethod(converter); - - com.google.genai.types.FunctionDeclaration declaration = - com.google.genai.types.FunctionDeclaration.builder() - .name("getWeatherInfo") - .description("Get weather information") - .parametersJsonSchema( - Map.of( - "type", "object", "properties", Map.of("location", Map.of("type", "string")))) - .build(); - - Map correctArgs = Map.of("location", "San Francisco"); - Map processedArgs = - invokeProcessArguments(processArguments, converter, correctArgs, declaration); - - assertThat(processedArgs).isEqualTo(correctArgs); + + assertThat( + process( + converter, + Map.of("args", Map.of("location", "San Francisco")), + jsonSchemaDeclaration())) + .containsEntry("location", "San Francisco"); } @Test - void testArgumentProcessingWithParametersJsonSchema_nestedFormat() throws Exception { + void jsonSchemaMapsDirectValueToSingleParameter() throws Exception { ToolConverter converter = new ToolConverter(); - Method processArguments = getProcessArgumentsMethod(converter); - - com.google.genai.types.FunctionDeclaration declaration = - com.google.genai.types.FunctionDeclaration.builder() - .name("getWeatherInfo") - .description("Get weather information") - .parametersJsonSchema( - Map.of( - "type", "object", "properties", Map.of("location", Map.of("type", "string")))) - .build(); - - Map nestedArgs = Map.of("args", Map.of("location", "San Francisco")); - Map processedArgs = - invokeProcessArguments(processArguments, converter, nestedArgs, declaration); - - assertThat(processedArgs).containsEntry("location", "San Francisco"); + + assertThat(process(converter, Map.of("value", "San Francisco"), jsonSchemaDeclaration())) + .containsEntry("location", "San Francisco"); } @Test - void testArgumentProcessingWithParametersJsonSchema_directValue() throws Exception { + void jsonSchemaLeavesUnmatchedArgumentsUnchanged() throws Exception { ToolConverter converter = new ToolConverter(); - Method processArguments = getProcessArgumentsMethod(converter); - - com.google.genai.types.FunctionDeclaration declaration = - com.google.genai.types.FunctionDeclaration.builder() - .name("getWeatherInfo") - .description("Get weather information") - .parametersJsonSchema( - Map.of( - "type", "object", "properties", Map.of("location", Map.of("type", "string")))) - .build(); - - Map directValueArgs = Map.of("value", "San Francisco"); - Map processedArgs = - invokeProcessArguments(processArguments, converter, directValueArgs, declaration); - - assertThat(processedArgs).containsEntry("location", "San Francisco"); + Map arguments = Map.of("city", "San Francisco", "country", "USA"); + + assertThat(process(converter, arguments, jsonSchemaDeclaration())).isEqualTo(arguments); } - @Test - void testArgumentProcessingWithParametersJsonSchema_noMatch() throws Exception { - ToolConverter converter = new ToolConverter(); - Method processArguments = getProcessArgumentsMethod(converter); - - com.google.genai.types.FunctionDeclaration declaration = - com.google.genai.types.FunctionDeclaration.builder() - .name("getWeatherInfo") - .description("Get weather information") - .parametersJsonSchema( - Map.of( - "type", "object", "properties", Map.of("location", Map.of("type", "string")))) - .build(); - - Map wrongArgs = Map.of("city", "San Francisco", "country", "USA"); - Map processedArgs = - invokeProcessArguments(processArguments, converter, wrongArgs, declaration); - - assertThat(processedArgs).isEqualTo(wrongArgs); + private Map process( + ToolConverter converter, Map arguments, FunctionDeclaration declaration) + throws Exception { + Method method = + ToolConverter.class.getDeclaredMethod( + "processArguments", Map.class, FunctionDeclaration.class); + method.setAccessible(true); + @SuppressWarnings("unchecked") + Map processed = + (Map) method.invoke(converter, arguments, declaration); + return processed; + } + + private FunctionDeclaration jsonSchemaDeclaration() { + return FunctionDeclaration.builder() + .name("getWeatherInfo") + .description("Get weather information") + .parametersJsonSchema( + Map.of("type", "object", "properties", Map.of("location", Map.of("type", "string")))) + .build(); } public static class WeatherTools { public static Map getWeatherInfo(String location) { - return Map.of( - "location", location, - "temperature", "72°F", - "condition", "sunny and clear", - "humidity", "45%", - "forecast", "Perfect weather for outdoor activities!"); + return Map.of("location", location); } } } diff --git a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/ToolConverterTest.java b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/ToolConverterTest.java index 1f3044159..ebd69c053 100644 --- a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/ToolConverterTest.java +++ b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/ToolConverterTest.java @@ -16,14 +16,20 @@ package com.google.adk.models.springai; import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import com.fasterxml.jackson.databind.ObjectMapper; import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.Schema; +import io.reactivex.rxjava3.core.Single; +import java.util.ArrayList; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.springframework.ai.tool.ToolCallback; @@ -211,5 +217,133 @@ public Optional declaration() { assertThat(toolCallbacks).hasSize(1); assertThat(toolCallbacks.get(0).getToolDefinition().name()).isEqualTo("get_weather"); + assertThat(toolCallbacks.get(0).getToolDefinition().inputSchema()) + .contains("\"location\"") + .contains("\"required\""); + } + + @Test + void convertedToolCallbackOnlyProvidesDefinitionAndDoesNotExecuteAdkTool() { + FunctionDeclaration function = + FunctionDeclaration.builder() + .name("context_tool") + .description("A tool that requires ADK execution context") + .parametersJsonSchema( + Map.of( + "type", + "object", + "properties", + Map.of("value", Map.of("type", "string")), + "required", + List.of("value"))) + .build(); + List receivedContexts = new ArrayList<>(); + BaseTool contextTool = + new BaseTool("context_tool", "A tool that requires ADK execution context") { + @Override + public Optional declaration() { + return Optional.of(function); + } + + @Override + public Single> runAsync( + Map args, ToolContext toolContext) { + receivedContexts.add(toolContext); + return Single.just(Map.of("value", args.get("value"))); + } + }; + + ToolCallback callback = + toolConverter.convertToSpringAiTools(Map.of(contextTool.name(), contextTool)).get(0); + + String resultWithoutSpringContext = callback.call("{\"value\":\"test\"}"); + String resultWithSpringContext = + callback.call( + "{\"value\":\"test\"}", + new org.springframework.ai.chat.model.ToolContext(Map.of("requestId", "test"))); + + assertThat(resultWithoutSpringContext).isEmpty(); + assertThat(resultWithSpringContext).isEmpty(); + assertThat(receivedContexts).isEmpty(); + } + + @Test + void springAiManagedModeRequiresToolContextResolver() { + assertThatThrownBy( + () -> new ToolConverter(new ObjectMapper(), ToolExecutionMode.SPRING_AI_MANAGED, null)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("AdkToolContextResolver"); + } + + @Test + void springAiManagedCallbackExecutesToolWithResolvedAdkContext() { + FunctionDeclaration function = + FunctionDeclaration.builder() + .name("context_tool") + .description("A tool that requires ADK execution context") + .parametersJsonSchema( + Map.of("type", "object", "properties", Map.of("value", Map.of("type", "string")))) + .build(); + AtomicReference receivedAdkContext = new AtomicReference<>(); + AtomicReference receivedSpringAiContext = + new AtomicReference<>(); + AtomicReference> receivedArguments = new AtomicReference<>(); + ToolContext resolvedContext = org.mockito.Mockito.mock(ToolContext.class); + BaseTool contextTool = + new BaseTool("context_tool", "A tool that requires ADK execution context") { + @Override + public Optional declaration() { + return Optional.of(function); + } + + @Override + public Single> runAsync( + Map args, ToolContext toolContext) { + receivedAdkContext.set(toolContext); + return Single.just(Map.of("echo", args.get("value"))); + } + }; + ToolConverter executableConverter = + new ToolConverter( + new ObjectMapper(), + ToolExecutionMode.SPRING_AI_MANAGED, + (tool, arguments, springAiContext) -> { + receivedArguments.set(arguments); + receivedSpringAiContext.set(springAiContext); + return resolvedContext; + }); + ToolCallback callback = + executableConverter.convertToSpringAiTools(Map.of(contextTool.name(), contextTool)).get(0); + org.springframework.ai.chat.model.ToolContext springAiContext = + new org.springframework.ai.chat.model.ToolContext(Map.of("requestId", "test")); + + String result = callback.call("{\"wrapper\":{\"value\":\"hello\"}}", springAiContext); + + assertThat(result).contains("\"echo\":\"hello\""); + assertThat(receivedArguments.get()).containsEntry("value", "hello"); + assertThat(receivedSpringAiContext.get()).isSameAs(springAiContext); + assertThat(receivedAdkContext.get()).isSameAs(resolvedContext); + } + + @Test + void springAiManagedCallbackRejectsNullResolvedContext() { + FunctionDeclaration function = + FunctionDeclaration.builder().name("context_tool").description("context tool").build(); + BaseTool contextTool = + new BaseTool("context_tool", "context tool") { + @Override + public Optional declaration() { + return Optional.of(function); + } + }; + ToolConverter executableConverter = + new ToolConverter( + new ObjectMapper(), ToolExecutionMode.SPRING_AI_MANAGED, (tool, args, context) -> null); + ToolCallback callback = + executableConverter.convertToSpringAiTools(Map.of(contextTool.name(), contextTool)).get(0); + + assertThatThrownBy(() -> callback.call("{}")) + .hasRootCauseInstanceOf(NullPointerException.class) + .hasRootCauseMessage("AdkToolContextResolver returned null for tool context_tool"); } } diff --git a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/autoconfigure/SpringAIAutoConfigurationTest.java b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/autoconfigure/SpringAIAutoConfigurationTest.java index 7c55b8d68..a23106015 100644 --- a/contrib/spring-ai/src/test/java/com/google/adk/models/springai/autoconfigure/SpringAIAutoConfigurationTest.java +++ b/contrib/spring-ai/src/test/java/com/google/adk/models/springai/autoconfigure/SpringAIAutoConfigurationTest.java @@ -17,7 +17,9 @@ import static org.assertj.core.api.Assertions.assertThat; +import com.google.adk.models.springai.AdkToolContextResolver; import com.google.adk.models.springai.SpringAI; +import com.google.adk.models.springai.ToolExecutionMode; import com.google.adk.models.springai.properties.SpringAIProperties; import org.junit.jupiter.api.Test; import org.springframework.ai.chat.messages.AssistantMessage; @@ -161,6 +163,37 @@ void testDefaultConfiguration() { assertThat(properties.getObservability().isEnabled()).isTrue(); assertThat(properties.getObservability().isMetricsEnabled()).isTrue(); assertThat(properties.getObservability().isIncludeContent()).isFalse(); + assertThat(properties.getToolExecution().getMode()) + .isEqualTo(ToolExecutionMode.ADK_MANAGED); + }); + } + + @Test + void springAiManagedModeFailsFastWithoutContextResolver() { + contextRunner + .withUserConfiguration(TestConfigurationWithChatModel.class) + .withPropertyValues("adk.spring-ai.tool-execution.mode=spring-ai-managed") + .run( + context -> { + assertThat(context).hasFailed(); + assertThat(context.getStartupFailure()) + .hasRootCauseInstanceOf(IllegalArgumentException.class) + .hasRootCauseMessage( + "An AdkToolContextResolver is required when tool execution is SPRING_AI_MANAGED"); + }); + } + + @Test + void springAiManagedModeStartsWithContextResolverBean() { + contextRunner + .withUserConfiguration(TestConfigurationWithChatModelAndResolver.class) + .withPropertyValues("adk.spring-ai.tool-execution.mode=spring-ai-managed") + .run( + context -> { + assertThat(context).hasNotFailed(); + assertThat(context).hasSingleBean(SpringAI.class); + assertThat(context.getBean(SpringAIProperties.class).getToolExecution().getMode()) + .isEqualTo(ToolExecutionMode.SPRING_AI_MANAGED); }); } @@ -200,4 +233,19 @@ public StreamingChatModel streamingChatModel() { java.util.List.of(new Generation(new AssistantMessage("streaming"))))); } } + + @Configuration + static class TestConfigurationWithChatModelAndResolver { + @Bean + public ChatModel chatModel() { + return prompt -> + new ChatResponse(java.util.List.of(new Generation(new AssistantMessage("response")))); + } + + @Bean + public AdkToolContextResolver adkToolContextResolver() { + return (tool, arguments, context) -> + org.mockito.Mockito.mock(com.google.adk.tools.ToolContext.class); + } + } }