diff --git a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java index ecc2bb412..08d620211 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/Functions.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/Functions.java @@ -586,8 +586,7 @@ private static Maybe> maybeInvokeAfterToolCall( private static Maybe> callTool( BaseTool tool, Map args, ToolContext toolContext, Context parentContext) { - return tool.runAsync(args, toolContext) - .toMaybe() + return tool.runMaybeAsync(args, toolContext) .doOnSubscribe( d -> Tracing.traceToolCall( diff --git a/core/src/main/java/com/google/adk/tools/BaseTool.java b/core/src/main/java/com/google/adk/tools/BaseTool.java index 1ea2808a1..10c740fd5 100644 --- a/core/src/main/java/com/google/adk/tools/BaseTool.java +++ b/core/src/main/java/com/google/adk/tools/BaseTool.java @@ -34,6 +34,7 @@ import com.google.genai.types.LiveConnectConfig; import com.google.genai.types.Tool; import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; import java.util.HashMap; import java.util.Map; @@ -48,6 +49,8 @@ public abstract class BaseTool { private final String name; private final String description; private final boolean isLongRunning; + private final boolean overridesRunAsync; + private final boolean overridesRunMaybeAsync; private final HashMap customMetadata; protected BaseTool(@Nonnull String name, @Nonnull String description) { @@ -58,6 +61,8 @@ protected BaseTool(@Nonnull String name, @Nonnull String description, boolean is this.name = name; this.description = description; this.isLongRunning = isLongRunning; + overridesRunAsync = overridesMethod("runAsync"); + overridesRunMaybeAsync = overridesMethod("runMaybeAsync"); customMetadata = new HashMap<>(); } @@ -90,6 +95,24 @@ public void setCustomMetadata(String key, Object value) { /** Calls a tool. */ public Single> runAsync(Map args, ToolContext toolContext) { + if (overridesRunMaybeAsync) { + return runMaybeAsync(args, toolContext).defaultIfEmpty(ImmutableMap.of()); + } + throw new UnsupportedOperationException("This method is not implemented."); + } + + /** + * Calls a tool and optionally returns a function response. + * + *

Override this method for long-running tools that may end the current invocation without + * emitting a function response event. This default implementation delegates to {@link + * #runAsync(Map, ToolContext)} for backwards compatibility. + */ + public Maybe> runMaybeAsync( + Map args, ToolContext toolContext) { + if (overridesRunAsync && !overridesRunMaybeAsync) { + return runAsync(args, toolContext).toMaybe(); + } throw new UnsupportedOperationException("This method is not implemented."); } @@ -180,6 +203,15 @@ private static ImmutableList findToolsWithoutFunctionDeclarations(LlmReque .orElse(ImmutableList.of()); } + private boolean overridesMethod(String methodName) { + try { + return getClass().getMethod(methodName, Map.class, ToolContext.class).getDeclaringClass() + != BaseTool.class; + } catch (NoSuchMethodException e) { + throw new IllegalStateException("Missing tool method: " + methodName, e); + } + } + /** * Creates a tool instance from a config. * diff --git a/core/src/main/java/com/google/adk/tools/FunctionTool.java b/core/src/main/java/com/google/adk/tools/FunctionTool.java index 4323b4569..05a78561a 100644 --- a/core/src/main/java/com/google/adk/tools/FunctionTool.java +++ b/core/src/main/java/com/google/adk/tools/FunctionTool.java @@ -245,6 +245,12 @@ public boolean isStreaming() { @Override public Single> runAsync(Map args, ToolContext toolContext) { + return runMaybeAsync(args, toolContext).defaultIfEmpty(ImmutableMap.of()); + } + + @Override + public Maybe> runMaybeAsync( + Map args, ToolContext toolContext) { try { if (requireConfirmation) { if (toolContext.toolConfirmation().isEmpty()) { @@ -253,17 +259,20 @@ public Single> runAsync(Map args, ToolContex "Please approve or reject the tool call %s() by responding with a" + " FunctionResponse with an expected ToolConfirmation payload.", name())); - return Single.just( + return Maybe.just( ImmutableMap.of( "error", "This tool call requires confirmation, please approve or reject.")); } else if (!toolContext.toolConfirmation().get().confirmed()) { - return Single.just(ImmutableMap.of("error", "This tool call is rejected.")); + return Maybe.just(ImmutableMap.of("error", "This tool call is rejected.")); } } - return this.call(args, toolContext).defaultIfEmpty(ImmutableMap.of()); + Maybe> functionResult = this.call(args, toolContext); + return longRunning() + ? functionResult + : functionResult.switchIfEmpty(Maybe.just(ImmutableMap.of())); } catch (Exception e) { logger.error("Exception occurred while calling function tool: " + func.getName(), e); - return Single.just( + return Maybe.just( ImmutableMap.of("status", "error", "message", "An internal error occurred.")); } } diff --git a/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java b/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java index d5db4d4b3..551b8d5a7 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/FunctionsTest.java @@ -27,12 +27,16 @@ import com.google.adk.agents.RunConfig.ToolExecutionMode; import com.google.adk.events.Event; import com.google.adk.testing.TestUtils; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.ToolContext; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Maybe; +import java.util.Map; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -146,6 +150,39 @@ public void handleFunctionCalls_singleFunctionCall() { .build()); } + @Test + public void handleFunctionCalls_longRunningToolWithEmptyResponse() { + InvocationContext invocationContext = createInvocationContext(createRootAgent()); + Event event = + createEvent("event").toBuilder() + .content( + Content.fromParts( + Part.fromText("..."), + Part.builder() + .functionCall( + FunctionCall.builder() + .id("function_call_id") + .name("empty_tool") + .args(ImmutableMap.of()) + .build()) + .build())) + .build(); + BaseTool tool = + new BaseTool("empty_tool", "Long-running tool without an immediate response", true) { + @Override + public Maybe> runMaybeAsync( + Map args, ToolContext toolContext) { + return Maybe.empty(); + } + }; + + Event functionResponseEvent = + Functions.handleFunctionCalls(invocationContext, event, ImmutableMap.of("empty_tool", tool)) + .blockingGet(); + + assertThat(functionResponseEvent).isNull(); + } + @Test public void handleFunctionCalls_multipleFunctionCalls_parallel() { InvocationContext invocationContext = diff --git a/core/src/test/java/com/google/adk/tools/BaseToolTest.java b/core/src/test/java/com/google/adk/tools/BaseToolTest.java index 2a07e7a44..ba265dbbc 100644 --- a/core/src/test/java/com/google/adk/tools/BaseToolTest.java +++ b/core/src/test/java/com/google/adk/tools/BaseToolTest.java @@ -1,6 +1,7 @@ package com.google.adk.tools; import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LlmAgent; @@ -16,6 +17,7 @@ import com.google.genai.types.ToolCodeExecution; import com.google.genai.types.UrlContext; import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; import java.util.Map; import java.util.Optional; @@ -117,6 +119,42 @@ public Single> runAsync( .build()); } + @Test + public void runAsync_withOnlyRunMaybeAsyncOverride_returnsEmptyMap() { + BaseTool tool = + new BaseTool("test_tool", "test_description", /* isLongRunning= */ true) { + @Override + public Maybe> runMaybeAsync( + Map args, ToolContext toolContext) { + return Maybe.empty(); + } + }; + + assertThat(tool.runAsync(Map.of(), /* toolContext= */ null).blockingGet()).isEmpty(); + } + + @Test + public void runAsync_withRunAsyncAndRunMaybeAsyncOverridesCallingSuper_throwsException() { + BaseTool tool = + new BaseTool("test_tool", "test_description", /* isLongRunning= */ true) { + @Override + public Single> runAsync( + Map args, ToolContext toolContext) { + return super.runAsync(args, toolContext); + } + + @Override + public Maybe> runMaybeAsync( + Map args, ToolContext toolContext) { + return super.runMaybeAsync(args, toolContext); + } + }; + + assertThrows( + UnsupportedOperationException.class, + () -> tool.runAsync(Map.of(), /* toolContext= */ null)); + } + @Test public void processLlmRequestWithGoogleSearchToolAddsToolToConfig() { FunctionDeclaration functionDeclaration = diff --git a/core/src/test/java/com/google/adk/tools/FunctionToolTest.java b/core/src/test/java/com/google/adk/tools/FunctionToolTest.java index 0939c6506..fdb9c2b33 100644 --- a/core/src/test/java/com/google/adk/tools/FunctionToolTest.java +++ b/core/src/test/java/com/google/adk/tools/FunctionToolTest.java @@ -551,6 +551,14 @@ public void call_withMaybeMapReturnType() throws Exception { assertThat(result).containsExactly("key", "value"); } + @Test + public void runMaybeAsync_longRunningWithEmptyMaybeReturnType_returnsEmpty() throws Exception { + Method method = Functions.class.getMethod("returnsEmptyMaybeMap"); + FunctionTool tool = new FunctionTool(null, method, /* isLongRunning= */ true); + + assertThat(tool.runMaybeAsync(new HashMap<>(), null).blockingGet()).isNull(); + } + @Test public void create_withSingleMapReturnType() { FunctionTool tool = FunctionTool.create(Functions.class, "returnsSingleMap"); @@ -804,6 +812,10 @@ public static Maybe> returnsMaybeMap() { return Maybe.just(ImmutableMap.of("key", "value")); } + public static Maybe> returnsEmptyMaybeMap() { + return Maybe.empty(); + } + public static Maybe returnsMaybeString() { return Maybe.just("not supported"); }