From e44bf9115994b05c80f25b68f5b967e3257eab75 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Sobczyk?= Date: Tue, 20 Oct 2026 01:40:03 -0700 Subject: [PATCH] feat: run a tool as a workflow node PiperOrigin-RevId: 1002922637 --- .../com/google/adk/workflow/ToolNodes.java | 246 ++++++ .../google/adk/workflow/ToolNodesTest.java | 718 ++++++++++++++++++ 2 files changed, 964 insertions(+) create mode 100644 core/src/main/java/com/google/adk/workflow/ToolNodes.java create mode 100644 core/src/test/java/com/google/adk/workflow/ToolNodesTest.java diff --git a/core/src/main/java/com/google/adk/workflow/ToolNodes.java b/core/src/main/java/com/google/adk/workflow/ToolNodes.java new file mode 100644 index 000000000..774cfe9ab --- /dev/null +++ b/core/src/main/java/com/google/adk/workflow/ToolNodes.java @@ -0,0 +1,246 @@ +/* + * 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.workflow; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.ObjectReader; +import com.google.adk.JsonBaseModel; +import com.google.adk.agents.Context; +import com.google.adk.annotations.Experimental; +import com.google.adk.events.EventActions; +import com.google.adk.plugins.Plugin; +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.FunctionDeclaration; +import com.google.genai.types.Schema; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Single; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.jspecify.annotations.Nullable; + +/** + * Static factories for workflow nodes that run a {@link BaseTool}, such as an MCP tool or an agent + * wrapped as a tool. A node calls its tool between the plugins' tool callbacks and keeps only the + * tool's state and artifact changes. Workflows cannot pause for input yet, so a tool that asks for + * confirmation outputs its placeholder error. + */ +@Experimental +public final class ToolNodes { + + /** + * Returns a node named after {@code tool}, with the default {@link NodeConfig}. Otherwise it + * behaves as {@link #asNode(BaseTool, String, NodeConfig)} describes. + */ + public static Node asNode(BaseTool tool) { + return asNode(tool, tool.name()); + } + + /** + * Returns a node named {@code name}, with the default {@link NodeConfig}. Otherwise it behaves as + * {@link #asNode(BaseTool, String, NodeConfig)} describes. + */ + public static Node asNode(BaseTool tool, String name) { + return asNode(tool, name, NodeConfig.builder().build()); + } + + /** + * Returns a node named {@code name} that runs {@code tool} with {@code config}; session state + * supplies any required parameter that the input lacks. The input may be a map of arguments, a + * JSON object as text or as a {@link Content}, or null or blank text for none. A tool with one + * parameter, or one required parameter, takes any other input as that parameter; other tools + * fail. + */ + public static Node asNode(BaseTool tool, String name, NodeConfig config) { + return new ToolNode(tool, name, config); + } + + private static final class ToolNode extends BaseNode { + + private static final ObjectReader JSON_READER = + JsonBaseModel.getMapper() + .readerFor(Object.class) + .with(DeserializationFeature.FAIL_ON_TRAILING_TOKENS); + + private final BaseTool tool; + + ToolNode(BaseTool tool, String name, NodeConfig config) { + super( + name, + tool.description(), + /* rerunOnResume= */ false, + /* waitForOutput= */ false, + config, + /* inputSchema= */ null, + /* outputSchema= */ null, + /* stateSchema= */ null); + this.tool = tool; + } + + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + return Single.defer(() -> runTool(context, nodeInput)).toFlowable(); + } + + private Single> runTool(Context context, @Nullable Object nodeInput) { + Parameters parameters = Parameters.of(tool); + Map args = argumentsFrom(nodeInput, parameters); + for (String param : parameters.required()) { + if (!args.containsKey(param) && context.state().containsKey(param)) { + args.put(param, context.state().get(param)); + } + } + ToolContext toolContext = + ToolContext.builder(context.invocationContext()) + .functionCallId(context.invocationContext().newUuid()) + .build(); + return callTool(context.invocationContext().pluginManager(), args, toolContext) + // Kept on failure too, so the recorded deltas match the live session state. + .doOnEvent((unusedResult, unusedError) -> keepDeltas(context, toolContext)); + } + + /** Calls the tool between the plugins' tool callbacks, as a model's tool call does. */ + private Single> callTool( + Plugin plugins, Map args, ToolContext toolContext) { + Single> toolCall = + Single.defer(() -> tool.runAsync(args, toolContext)) + .onErrorResumeNext( + error -> + plugins + .onToolErrorCallback(tool, args, toolContext, error) + .switchIfEmpty(Single.error(error))); + return plugins + .beforeToolCallback(tool, args, toolContext) + .switchIfEmpty(toolCall) + .flatMap( + result -> + plugins + .afterToolCallback(tool, args, toolContext, result) + .defaultIfEmpty(result)); + } + + private Map argumentsFrom(@Nullable Object nodeInput, Parameters parameters) { + Object input = nodeInput instanceof Content content ? content.text() : nodeInput; + if (input instanceof String text) { + input = parseText(text.strip()); + } + Map args = new LinkedHashMap<>(); + if (input == null) { + return args; + } + if (input instanceof Map map) { + map.forEach((key, value) -> args.put(String.valueOf(key), value)); + return args; + } + String param = parameters.forNonMapInput(); + if (param == null) { + throw new IllegalArgumentException( + String.format( + "The input to tool node '%s' must be a dictionary of tool arguments, a JSON object" + + " as text, or null, but got %s.", + name(), + input instanceof String + ? "text that is not a JSON object" + : input.getClass().getSimpleName())); + } + args.put(param, input); + return args; + } + + private static @Nullable Object parseText(String text) { + if (text.isEmpty()) { + return null; + } + try { + Object json = JSON_READER.readValue(text); + return json instanceof Map ? json : text; + } catch (JsonProcessingException e) { + return text; + } + } + + private static void keepDeltas(Context context, ToolContext toolContext) { + EventActions toolActions = toolContext.eventActions(); + context.eventActions().stateDelta().putAll(toolActions.stateDelta()); + context.eventActions().artifactDelta().putAll(toolActions.artifactDelta()); + } + } + + /** + * The declared parameters that a node fills itself. + * + * @param required the parameters read from state when the input lacks them + * @param forNonMapInput the parameter that a non-map input fills, as in Python ADK: the only + * parameter, or else the only required one + */ + private record Parameters(ImmutableList required, @Nullable String forNonMapInput) { + private static final Parameters NONE = new Parameters(ImmutableList.of(), null); + private static final TypeReference> MAP_TYPE = new TypeReference<>() {}; + + /** Reads {@code tool}'s raw JSON schema, which MCP tools set, or else its typed schema. */ + static Parameters of(BaseTool tool) { + FunctionDeclaration declaration = tool.declaration().orElse(null); + if (declaration == null) { + return NONE; + } + Map rawSchema = + declaration + .parametersJsonSchema() + .map(schema -> JsonBaseModel.getMapper().convertValue(schema, MAP_TYPE)) + .orElse(ImmutableMap.of()); + ImmutableList all; + ImmutableList required; + if (!rawSchema.isEmpty()) { + all = + rawSchema.get("properties") instanceof Map properties + ? names(properties.keySet()) + : ImmutableList.of(); + required = + rawSchema.get("required") instanceof List list ? names(list) : ImmutableList.of(); + } else { + Schema schema = declaration.parameters().orElse(null); + if (schema == null) { + return NONE; + } + all = + schema + .properties() + .map(properties -> ImmutableList.copyOf(properties.keySet())) + .orElse(ImmutableList.of()); + required = schema.required().map(ImmutableList::copyOf).orElse(ImmutableList.of()); + } + if (all.size() == 1) { + return new Parameters(required, all.get(0)); + } + return new Parameters(required, required.size() == 1 ? required.get(0) : null); + } + + private static ImmutableList names(Iterable values) { + ImmutableList.Builder names = ImmutableList.builder(); + values.forEach(value -> names.add(String.valueOf(value))); + return names.build(); + } + } + + private ToolNodes() {} +} diff --git a/core/src/test/java/com/google/adk/workflow/ToolNodesTest.java b/core/src/test/java/com/google/adk/workflow/ToolNodesTest.java new file mode 100644 index 000000000..ff67ab262 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/ToolNodesTest.java @@ -0,0 +1,718 @@ +/* + * 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.workflow; + +import static com.google.adk.testing.TestUtils.createTestAgentBuilder; +import static com.google.adk.testing.TestUtils.createTestLlm; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.common.collect.ImmutableMap.toImmutableMap; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.adk.agents.CallbackContext; +import com.google.adk.agents.Context; +import com.google.adk.agents.InvocationContext; +import com.google.adk.events.EventActions; +import com.google.adk.plugins.BasePlugin; +import com.google.adk.plugins.PluginManager; +import com.google.adk.testing.TestLlm; +import com.google.adk.testing.TestUtils; +import com.google.adk.tools.AgentTool; +import com.google.adk.tools.Annotations; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.FunctionTool; +import com.google.adk.tools.ToolContext; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.common.collect.Iterables; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionDeclaration; +import com.google.genai.types.Part; +import com.google.genai.types.Schema; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.function.BiFunction; +import org.jspecify.annotations.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +@RunWith(JUnit4.class) +public final class ToolNodesTest { + + @Test + public void asNode_withOnlyATool_usesItsNameAndDescriptionAndDefaults() { + Node node = ToolNodes.asNode(FakeTool.echo("lookup")); + + assertThat(node.name()).isEqualTo("lookup"); + assertThat(node.description()).isEqualTo("A fake tool."); + assertThat(node.config()).isEqualTo(NodeConfig.builder().build()); + assertThat(node.inputSchema()).isEmpty(); + assertThat(node.outputSchema()).isEmpty(); + } + + @Test + public void asNode_withAName_overridesTheToolsName() { + Node node = ToolNodes.asNode(FakeTool.echo("lookup"), "custom"); + + assertThat(node.name()).isEqualTo("custom"); + assertThat(node.description()).isEqualTo("A fake tool."); + } + + @Test + public void asNode_withANameAndConfig_usesThem() { + NodeConfig config = NodeConfig.builder().timeout(Duration.ofSeconds(5)).build(); + + Node node = ToolNodes.asNode(FakeTool.echo("lookup"), "custom", config); + + assertThat(node.name()).isEqualTo("custom"); + assertThat(node.config()).isEqualTo(config); + } + + @Test + public void runNode_mapInput_passesItAsTheToolsArguments() { + ImmutableMap args = ImmutableMap.of("param_a", 1, "param_b", "value"); + + assertThat(run(FakeTool.echo("mock_tool"), args)).isEqualTo(args); + } + + @Test + public void runNode_jsonObjectText_parsesItIntoArguments() { + Object output = run(FakeTool.echo("mock_tool"), " {\"param_a\": 1, \"param_b\": \"value\"} "); + + assertThat(output).isEqualTo(ImmutableMap.of("param_a", 1, "param_b", "value")); + } + + @Test + public void runNode_nullInput_callsTheToolWithNoArguments() { + assertThat(run(FakeTool.echo("mock_tool"), null)).isEqualTo(ImmutableMap.of()); + } + + @Test + public void runNode_blankText_callsTheToolWithNoArguments() { + assertThat(run(FakeTool.echo("mock_tool"), " \n\t")).isEqualTo(ImmutableMap.of()); + } + + @Test + public void runNode_jsonObjectContent_readsItsTextWithoutThoughts() { + Content content = + Content.builder() + .role("model") + .parts( + Part.builder().text("reasoning").thought(true).build(), Part.fromText("{\"x\": 1}")) + .build(); + + assertThat(run(FakeTool.echo("mock_tool"), content)).isEqualTo(ImmutableMap.of("x", 1)); + } + + @Test + public void runNode_contentWithOnlyThoughts_callsTheToolWithNoArguments() { + Content content = + Content.builder() + .role("model") + .parts(Part.builder().text("{\"x\": 1}").thought(true).build()) + .build(); + + assertThat(run(FakeTool.echo("mock_tool"), content)).isEqualTo(ImmutableMap.of()); + } + + @Test + public void runNode_textThatIsNotJson_failsTheNode() { + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, () -> run(FakeTool.echo("mock_tool"), "not a json")); + + assertThat(thrown) + .hasMessageThat() + .contains("The input to tool node 'mock_tool' must be a dictionary of tool arguments"); + assertThat(thrown).hasMessageThat().contains("but got text that is not a JSON object"); + } + + @Test + public void runNode_jsonArrayText_failsTheNode() { + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, () -> run(FakeTool.echo("mock_tool"), "[1, 2, 3]")); + + assertThat(thrown).hasMessageThat().contains("but got text that is not a JSON object"); + } + + @Test + public void runNode_jsonNullText_failsTheNode() { + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, () -> run(FakeTool.echo("mock_tool"), " null ")); + + assertThat(thrown).hasMessageThat().contains("but got text that is not a JSON object"); + } + + @Test + public void runNode_objectTextWithUnquotedValue_failsTheNode() { + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> run(FakeTool.echo("mock_tool"), "{\"city\": Paris}")); + + assertThat(thrown).hasMessageThat().contains("but got text that is not a JSON object"); + } + + @Test + public void runNode_objectTextWithTrailingText_failsTheNode() { + IllegalArgumentException thrown = + assertThrows( + IllegalArgumentException.class, + () -> run(FakeTool.echo("mock_tool"), "{\"city\": \"Paris\"} and more")); + + assertThat(thrown).hasMessageThat().contains("but got text that is not a JSON object"); + } + + @Test + public void runNode_nonMapValue_failsTheNode() { + IllegalArgumentException thrown = + assertThrows(IllegalArgumentException.class, () -> run(FakeTool.echo("mock_tool"), 42)); + + assertThat(thrown).hasMessageThat().contains("but got Integer"); + } + + @Test + public void runNode_toolWithOneParameter_takesStrippedTextAsIt() { + Object output = run(FakeTool.echo("reader", "path"), " /tmp/a.txt "); + + assertThat(output).isEqualTo(ImmutableMap.of("path", "/tmp/a.txt")); + } + + @Test + public void runNode_toolWithOneParameter_takesANonMapValueAsIt() { + Object output = run(FakeTool.echo("counter", "count"), 42); + + assertThat(output).isEqualTo(ImmutableMap.of("count", 42)); + } + + @Test + public void runNode_toolWithOneParameter_takesJsonNullTextAsText() { + Object output = run(FakeTool.echo("writer", "text"), "null"); + + assertThat(output).isEqualTo(ImmutableMap.of("text", "null")); + } + + @Test + public void runNode_toolWithOneParameter_takesContentTextAsIt() { + Object output = + run(FakeTool.echo("search", "query"), Content.fromParts(Part.fromText("hello"))); + + assertThat(output).isEqualTo(ImmutableMap.of("query", "hello")); + } + + @Test + public void runNode_toolWithOneRequiredParameter_takesANonMapValueAsIt() { + FakeTool tool = + FakeTool.echo("weather", ImmutableList.of("city", "units"), ImmutableList.of("city")); + + assertThat(run(tool, "Paris")).isEqualTo(ImmutableMap.of("city", "Paris")); + } + + @Test + public void runNode_toolWithTwoParameters_rejectsANonMapValue() { + FakeTool tool = FakeTool.echo("lookup", ImmutableList.of("a", "b"), ImmutableList.of()); + + IllegalArgumentException thrown = + assertThrows(IllegalArgumentException.class, () -> run(tool, "scalar")); + + assertThat(thrown).hasMessageThat().contains("must be a dictionary"); + } + + @Test + public void runNode_toolWithRawJsonSchema_readsItsParameters() { + // MCP tools declare their parameters as a raw JSON schema. + FakeTool tool = + FakeTool.withRawSchema( + "weather", + ImmutableMap.of( + "type", + "object", + "properties", + ImmutableMap.of( + "city", ImmutableMap.of("type", "string"), + "units", ImmutableMap.of("type", "string")), + "required", + ImmutableList.of("city"))); + + assertThat(run(tool, "Paris")).isEqualTo(ImmutableMap.of("city", "Paris")); + } + + @Test + public void runNode_toolWithRawAndTypedSchemas_readsTheRawOne() { + Schema typed = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "a", Schema.builder().type("STRING").build(), + "b", Schema.builder().type("STRING").build())) + .build(); + FunctionDeclaration declaration = + FunctionDeclaration.builder() + .name("weather") + .parameters(typed) + .parametersJsonSchema( + ImmutableMap.of( + "type", + "object", + "properties", + ImmutableMap.of("city", ImmutableMap.of("type", "string")))) + .build(); + FakeTool tool = new FakeTool("weather", declaration, (args, toolContext) -> args); + + assertThat(run(tool, "Paris")).isEqualTo(ImmutableMap.of("city", "Paris")); + } + + @Test + public void runNode_requiredParametersMissing_readsOnlyThoseFromState() { + FakeTool tool = + FakeTool.echo("weather", ImmutableList.of("city", "units"), ImmutableList.of("city")); + Context context = context(); + context.invocationContext().session().state().put("city", "Paris"); + context.invocationContext().session().state().put("units", "fahrenheit"); + + assertThat(run(tool, context, null)).isEqualTo(ImmutableMap.of("city", "Paris")); + } + + @Test + public void runNode_requiredParameterInNeitherInputNorState_isLeftOut() { + FakeTool tool = FakeTool.echo("weather", "city"); + + assertThat(run(tool, ImmutableMap.of())).isEqualTo(ImmutableMap.of()); + } + + @Test + public void runNode_argumentInInputAndState_usesTheInput() { + FakeTool tool = FakeTool.echo("weather", "city"); + Context context = context(); + context.invocationContext().session().state().put("city", "Paris"); + + assertThat(run(tool, context, ImmutableMap.of("city", "Rome"))) + .isEqualTo(ImmutableMap.of("city", "Rome")); + } + + @Test + public void runNode_functionTool_readsAMissingRequiredArgumentFromState() { + FunctionTool tool = FunctionTool.create(Functions.class, "forecast"); + Context context = context(); + context.invocationContext().session().state().put("city", "Paris"); + + assertThat(run(tool, context, null)).isEqualTo(ImmutableMap.of("forecast", "Sunny in Paris")); + } + + @Test + public void runNode_functionToolWithOneParameter_takesTextAsIt() { + FunctionTool tool = FunctionTool.create(Functions.class, "forecast"); + + assertThat(run(tool, "Rome")).isEqualTo(ImmutableMap.of("forecast", "Sunny in Rome")); + } + + @Test + public void runNode_functionToolRequiringConfirmation_outputsItsPlaceholderWithoutRunning() { + FunctionTool tool = + FunctionTool.create(Functions.class, "forecast", /* requireConfirmation= */ true); + Context context = context(); + + Object output = run(tool, context, "Rome"); + + assertThat(output) + .isEqualTo( + ImmutableMap.of( + "error", "This tool call requires confirmation, please approve or reject.")); + assertThat(context.eventActions().requestedToolConfirmations()).isEmpty(); + } + + @Test + public void runNode_toolWritesState_putsTheStateDeltaOnTheNodeContextBeforeTheOutput() { + FakeTool tool = + FakeTool.running( + "writer", + (args, toolContext) -> { + toolContext.state().put("color", "blue"); + return ImmutableMap.of("status", "ok"); + }); + Context context = context(); + + // The runtime attaches the context's deltas to the output event as it emits it. + ImmutableMap deltaAtOutput = + ToolNodes.asNode(tool) + .runNode(context, null) + .map(unusedOutput -> ImmutableMap.copyOf(context.eventActions().stateDelta())) + .blockingSingle(); + + assertThat(deltaAtOutput).containsExactly("color", "blue"); + } + + @Test + public void runNode_toolSavesArtifact_keepsTheArtifactDeltaOnTheNodeContext() { + FakeTool tool = + FakeTool.running( + "saver", + (args, toolContext) -> { + toolContext.saveArtifact("report.txt", Part.fromText("v1")).blockingAwait(); + return ImmutableMap.of("status", "saved"); + }); + Context context = context(); + + Object output = run(tool, context, null); + + assertThat(output).isEqualTo(ImmutableMap.of("status", "saved")); + assertThat(context.eventActions().artifactDelta()).containsKey("report.txt"); + } + + @Test + public void runNode_toolWritesStateThenFails_stillKeepsTheStateDelta() { + FakeTool tool = + FakeTool.running( + "writer", + (args, toolContext) -> { + toolContext.state().put("color", "blue"); + throw new IllegalStateException("transient failure"); + }); + Context context = context(); + + assertThrows(IllegalStateException.class, () -> run(tool, context, null)); + + assertThat(context.eventActions().stateDelta()).containsExactly("color", "blue"); + } + + @Test + public void runNode_toolRequestsModelOnlyActions_dropsThem() { + ImmutableMap placeholder = ImmutableMap.of("error", "Needs confirmation."); + FakeTool tool = + FakeTool.running( + "handoff", + (args, toolContext) -> { + toolContext.requestConfirmation("Approve?"); + toolContext.eventActions().setTransferToAgent("elsewhere"); + toolContext.eventActions().setSkipSummarization(true); + return placeholder; + }); + Context context = context(); + + Object output = run(tool, context, null); + + assertThat(output).isEqualTo(placeholder); + EventActions actions = context.eventActions(); + assertThat(actions.requestedToolConfirmations()).isEmpty(); + assertThat(actions.transferToAgent()).isEmpty(); + assertThat(actions.skipSummarization()).isEmpty(); + } + + @Test + public void runNode_twoRuns_giveTheToolAFreshFunctionCallIdEachTime() { + List seenIds = new ArrayList<>(); + FakeTool tool = + FakeTool.running( + "record", + (args, toolContext) -> { + seenIds.add(toolContext.functionCallId().orElse("")); + return ImmutableMap.of(); + }); + Node node = ToolNodes.asNode(tool); + Context context = context(); + + Object unusedFirst = node.runNode(context, null).blockingSingle(); + Object unusedSecond = node.runNode(context, null).blockingSingle(); + + assertThat(seenIds).hasSize(2); + assertThat(seenIds.get(0)).isNotEmpty(); + assertThat(seenIds.get(1)).isNotEqualTo(seenIds.get(0)); + } + + @Test + public void runNode_toolReturnsErrorMap_outputsIt() { + ImmutableMap errorMap = + ImmutableMap.of("error", "Missing required parameter city"); + FakeTool tool = FakeTool.running("lookup", (args, toolContext) -> errorMap); + + // No model reads the error map, so it is ordinary output for the next node. + assertThat(run(tool, null)).isEqualTo(errorMap); + } + + @Test + public void runNode_toolThrows_failsTheNode() { + FakeTool tool = + FakeTool.running( + "flaky", + (args, toolContext) -> { + throw new IllegalStateException("transient failure"); + }); + + IllegalStateException thrown = assertThrows(IllegalStateException.class, () -> run(tool, null)); + + assertThat(thrown).hasMessageThat().isEqualTo("transient failure"); + } + + @Test + public void runNode_agentTool_runsTheAgentOnTheRequestAndOutputsItsReply() { + TestLlm llm = createTestLlm(createTextLlmResponse("Paris is sunny.")); + AgentTool tool = AgentTool.create(createTestAgentBuilder(llm).name("forecaster").build()); + + Object output = run(tool, "Weather in Paris?"); + + assertThat(output).isEqualTo(ImmutableMap.of("result", "Paris is sunny.")); + assertThat(Iterables.getLast(llm.getLastRequest().contents()).text()) + .isEqualTo("Weather in Paris?"); + } + + @Test + public void runNode_runsThePluginsToolCallbacksAroundTheCall() { + RecordingPlugin plugin = RecordingPlugin.silent(); + + Object output = run(FakeTool.echo("lookup"), context(plugin), ImmutableMap.of("city", "Paris")); + + assertThat(output).isEqualTo(ImmutableMap.of("city", "Paris")); + assertThat(plugin.calls) + .containsExactly("before lookup {city=Paris}", "after lookup {city=Paris}") + .inOrder(); + } + + @Test + public void runNode_pluginBeforeAnswer_replacesTheToolCall() { + List> toolCalls = new ArrayList<>(); + FakeTool tool = + FakeTool.running( + "lookup", + (args, toolContext) -> { + toolCalls.add(args); + return args; + }); + RecordingPlugin plugin = RecordingPlugin.answeringBefore(ImmutableMap.of("cached", true)); + + Object output = run(tool, context(plugin), ImmutableMap.of("city", "Paris")); + + assertThat(output).isEqualTo(ImmutableMap.of("cached", true)); + assertThat(toolCalls).isEmpty(); + assertThat(plugin.calls).contains("after lookup {cached=true}"); + } + + @Test + public void runNode_pluginAfterAnswer_replacesTheOutput() { + RecordingPlugin plugin = RecordingPlugin.answeringAfter(ImmutableMap.of("redacted", true)); + + Object output = run(FakeTool.echo("lookup"), context(plugin), ImmutableMap.of("city", "Paris")); + + assertThat(output).isEqualTo(ImmutableMap.of("redacted", true)); + } + + @Test + public void runNode_pluginErrorAnswer_becomesTheOutput() { + FakeTool tool = + FakeTool.running( + "lookup", + (args, toolContext) -> { + throw new IllegalStateException("no data"); + }); + RecordingPlugin plugin = RecordingPlugin.answeringError(ImmutableMap.of("error", "handled")); + + Object output = run(tool, context(plugin), ImmutableMap.of("city", "Paris")); + + assertThat(output).isEqualTo(ImmutableMap.of("error", "handled")); + assertThat(plugin.calls) + .containsExactly( + "before lookup {city=Paris}", "error lookup no data", "after lookup {error=handled}") + .inOrder(); + } + + @Test + public void runNode_toolFailsAndNoPluginAnswers_failsTheNode() { + FakeTool tool = + FakeTool.running( + "lookup", + (args, toolContext) -> { + throw new IllegalStateException("no data"); + }); + RecordingPlugin plugin = RecordingPlugin.silent(); + + assertThrows( + IllegalStateException.class, + () -> run(tool, context(plugin), ImmutableMap.of("city", "Paris"))); + + assertThat(plugin.calls).contains("error lookup no data"); + } + + private static Object run(BaseTool tool, @Nullable Object nodeInput) { + return run(tool, context(), nodeInput); + } + + private static Object run(BaseTool tool, Context context, @Nullable Object nodeInput) { + return ToolNodes.asNode(tool).runNode(context, nodeInput).blockingSingle(); + } + + private static Context context() { + return context(new PluginManager()); + } + + private static Context context(BasePlugin plugin) { + return context(new PluginManager(ImmutableList.of(plugin))); + } + + private static Context context(PluginManager pluginManager) { + InvocationContext invocationContext = + TestUtils.createInvocationContext(TestUtils.createRootAgent()).toBuilder() + .pluginManager(pluginManager) + .build(); + return new CallbackContext(invocationContext, /* eventActions= */ null); + } + + /** Functions for the {@link FunctionTool} tests. */ + public static final class Functions { + public static ImmutableMap forecast( + @Annotations.Schema(name = "city") String city) { + return ImmutableMap.of("forecast", "Sunny in " + city); + } + + private Functions() {} + } + + /** A tool with an optional declaration that runs {@code body}. */ + private static final class FakeTool extends BaseTool { + private final @Nullable FunctionDeclaration declaration; + private final BiFunction, ToolContext, Map> body; + + private FakeTool( + String name, + @Nullable FunctionDeclaration declaration, + BiFunction, ToolContext, Map> body) { + super(name, "A fake tool."); + this.declaration = declaration; + this.body = body; + } + + /** Returns a tool without a declaration that outputs its arguments. */ + static FakeTool echo(String name) { + return running(name, (args, toolContext) -> args); + } + + /** Returns a tool with one required string parameter that outputs its arguments. */ + static FakeTool echo(String name, String param) { + return echo(name, ImmutableList.of(param), ImmutableList.of(param)); + } + + /** Returns a tool that declares string {@code params} and outputs its arguments. */ + static FakeTool echo( + String name, ImmutableList params, ImmutableList required) { + Schema parameters = + Schema.builder() + .type("OBJECT") + .properties( + params.stream() + .collect( + toImmutableMap( + param -> param, param -> Schema.builder().type("STRING").build()))) + .required(required) + .build(); + return new FakeTool( + name, + FunctionDeclaration.builder().name(name).parameters(parameters).build(), + (args, toolContext) -> args); + } + + /** + * Returns a tool that declares {@code schema} as a raw JSON schema and outputs its arguments. + */ + static FakeTool withRawSchema(String name, Map schema) { + return new FakeTool( + name, + FunctionDeclaration.builder().name(name).parametersJsonSchema(schema).build(), + (args, toolContext) -> args); + } + + /** Returns a tool without a declaration that runs {@code body}. */ + static FakeTool running( + String name, BiFunction, ToolContext, Map> body) { + return new FakeTool(name, /* declaration= */ null, body); + } + + @Override + public Optional declaration() { + return Optional.ofNullable(declaration); + } + + @Override + public Single> runAsync(Map args, ToolContext toolContext) { + return Single.fromCallable(() -> body.apply(args, toolContext)); + } + } + + /** Records the plugins' tool callbacks, answering at most one of them with a map. */ + private static final class RecordingPlugin extends BasePlugin { + final List calls = new ArrayList<>(); + private final @Nullable Map beforeAnswer; + private final @Nullable Map afterAnswer; + private final @Nullable Map errorAnswer; + + private RecordingPlugin( + @Nullable Map beforeAnswer, + @Nullable Map afterAnswer, + @Nullable Map errorAnswer) { + super("recording"); + this.beforeAnswer = beforeAnswer; + this.afterAnswer = afterAnswer; + this.errorAnswer = errorAnswer; + } + + static RecordingPlugin silent() { + return new RecordingPlugin(null, null, null); + } + + static RecordingPlugin answeringBefore(Map answer) { + return new RecordingPlugin(answer, null, null); + } + + static RecordingPlugin answeringAfter(Map answer) { + return new RecordingPlugin(null, answer, null); + } + + static RecordingPlugin answeringError(Map answer) { + return new RecordingPlugin(null, null, answer); + } + + @Override + public Maybe> beforeToolCallback( + BaseTool tool, Map toolArgs, ToolContext toolContext) { + calls.add("before " + tool.name() + " " + toolArgs); + return beforeAnswer == null ? Maybe.empty() : Maybe.just(beforeAnswer); + } + + @Override + public Maybe> afterToolCallback( + BaseTool tool, + Map toolArgs, + ToolContext toolContext, + Map result) { + calls.add("after " + tool.name() + " " + result); + return afterAnswer == null ? Maybe.empty() : Maybe.just(afterAnswer); + } + + @Override + public Maybe> onToolErrorCallback( + BaseTool tool, Map toolArgs, ToolContext toolContext, Throwable error) { + calls.add("error " + tool.name() + " " + error.getMessage()); + return errorAnswer == null ? Maybe.empty() : Maybe.just(errorAnswer); + } + } +}