From e516b797ec2f8b501db3473238a9b4b04f3ad3f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Sobczyk?= Date: Fri, 30 Oct 2026 03:11:44 -0700 Subject: [PATCH] feat: make a `SINGLE_TURN` `LlmAgent` node's answer its output PiperOrigin-RevId: 1010374021 --- .../adk/models/langchain4j/LangChain4j.java | 36 +- .../LangChain4jIntegrationTest.java | 22 +- .../models/langchain4j/LangChain4jTest.java | 92 +++ .../main/java/com/google/adk/SchemaUtils.java | 153 +++- .../java/com/google/adk/agents/AgentMode.java | 26 + .../java/com/google/adk/agents/BaseAgent.java | 8 +- .../google/adk/agents/InvocationContext.java | 31 + .../java/com/google/adk/agents/LlmAgent.java | 214 +++++- .../adk/flows/llmflows/AgentTransfer.java | 9 +- .../google/adk/flows/llmflows/Identity.java | 10 +- .../adk/tools/SetModelResponseTool.java | 28 +- .../google/adk/workflow/ParallelWorkers.java | 22 +- .../com/google/adk/workflow/ResumeScan.java | 110 +-- .../com/google/adk/workflow/ToolNodes.java | 10 +- .../java/com/google/adk/SchemaUtilsTest.java | 168 +++++ .../adk/tools/SetModelResponseToolTest.java | 91 +++ .../adk/workflow/LlmAgentAsNodeTest.java | 708 ++++++++++++++++++ 17 files changed, 1560 insertions(+), 178 deletions(-) create mode 100644 core/src/main/java/com/google/adk/agents/AgentMode.java create mode 100644 core/src/test/java/com/google/adk/workflow/LlmAgentAsNodeTest.java diff --git a/contrib/langchain4j/src/main/java/com/google/adk/models/langchain4j/LangChain4j.java b/contrib/langchain4j/src/main/java/com/google/adk/models/langchain4j/LangChain4j.java index ae2544e6e..34e730a15 100644 --- a/contrib/langchain4j/src/main/java/com/google/adk/models/langchain4j/LangChain4j.java +++ b/contrib/langchain4j/src/main/java/com/google/adk/models/langchain4j/LangChain4j.java @@ -217,30 +217,28 @@ public Flowable generateContent(LlmRequest llmRequest, boolean stre public void onPartialResponse(String s) { emitter.onNext( LlmResponse.builder() - .content(Content.fromParts(Part.fromText(s))) + .content( + Content.builder().role("model").parts(Part.fromText(s)).build()) + .partial(true) .build()); } @Override public void onCompleteResponse(ChatResponse chatResponse) { - if (chatResponse.aiMessage().hasToolExecutionRequests()) { - AiMessage aiMessage = chatResponse.aiMessage(); - toParts(aiMessage).stream() - .map(Part::functionCall) - .forEach( - functionCall -> { - functionCall.ifPresent( - function -> { - emitter.onNext( - LlmResponse.builder() - .content( - Content.fromParts( - Part.fromFunctionCall( - function.name().orElse(""), - function.args().orElse(Map.of())))) - .build()); - }); - }); + // The chunks were partial, so one response with the whole turn ends it. + AiMessage aiMessage = chatResponse.aiMessage(); + List parts = new ArrayList<>(); + if (aiMessage.text() != null && !aiMessage.text().isEmpty()) { + parts.add(Part.fromText(aiMessage.text())); + } + if (aiMessage.hasToolExecutionRequests()) { + parts.addAll(toParts(aiMessage)); + } + if (!parts.isEmpty()) { + emitter.onNext( + LlmResponse.builder() + .content(Content.builder().role("model").parts(parts).build()) + .build()); } emitter.onComplete(); } diff --git a/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java b/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java index 4a11cd886..a3dbf87ac 100644 --- a/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java +++ b/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jIntegrationTest.java @@ -375,6 +375,7 @@ void testSimpleStreamingResponse() { "", responses .blockingStream() + .filter(llmResponse -> !llmResponse.partial().orElse(false)) .map(llmResponse -> llmResponse.content().get().text()) .toList()); @@ -430,15 +431,30 @@ void testStreamingRunConfig() { // when List eventsHi = askAgentStreaming(agent, "Hi"); String responseToHi = - String.join("", eventsHi.stream().map(event -> event.content().get().text()).toList()); + String.join( + "", + eventsHi.stream() + .filter(event -> !event.partial().orElse(false)) + .map(event -> event.content().get().text()) + .toList()); List eventsQubit = askAgentStreaming(agent, "Tell me about qubits"); String responseToQubit = - String.join("", eventsQubit.stream().map(event -> event.content().get().text()).toList()); + String.join( + "", + eventsQubit.stream() + .filter(event -> !event.partial().orElse(false)) + .map(event -> event.content().get().text()) + .toList()); List eventsWeather = askAgentStreaming(agent, "What's the weather in Paris?"); String responseToWeather = - String.join("", eventsWeather.stream().map(Event::stringifyContent).toList()); + String.join( + "", + eventsWeather.stream() + .filter(event -> !event.partial().orElse(false)) + .map(Event::stringifyContent) + .toList()); // then diff --git a/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jTest.java b/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jTest.java index 9c56a0743..ff6af9614 100644 --- a/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jTest.java +++ b/contrib/langchain4j/src/test/java/com/google/adk/models/langchain4j/LangChain4jTest.java @@ -350,6 +350,17 @@ void testGenerateContentWithStreamingChatModel() { assertThat(responses).hasSize(5); // 4 partial responses + 1 function call // Verify the partial responses + assertThat(responses.stream().map(response -> response.content().orElseThrow().role())) + .containsExactly( + Optional.of("model"), + Optional.of("model"), + Optional.of("model"), + Optional.of("model"), + Optional.of("model")); + assertThat(responses.subList(0, 4).stream().map(LlmResponse::partial)) + .containsExactly( + Optional.of(true), Optional.of(true), Optional.of(true), Optional.of(true)); + assertThat(responses.get(4).partial()).isEmpty(); assertThat(responses.get(0).content().orElseThrow().text()).isEqualTo("Hello"); assertThat(responses.get(1).content().orElseThrow().text()).isEqualTo(", how"); assertThat(responses.get(2).content().orElseThrow().text()).isEqualTo(" can I help"); @@ -369,12 +380,93 @@ void testGenerateContentWithStreamingChatModel() { .get(0) .functionCall() .orElseThrow(); + assertThat(functionCall.id()).hasValue("123"); assertThat(functionCall.name()).isEqualTo(Optional.of("getWeather")); assertThat(functionCall.args().orElseThrow()).containsEntry("city", "Paris"); disposable.dispose(); } + @Test + @DisplayName("Should end a streamed text turn with its full text") + void testGenerateContentWithStreamingChatModel_endsWithTheFullText() { + final LlmRequest llmRequest = + LlmRequest.builder().contents(List.of(Content.fromParts(Part.fromText("Hello")))).build(); + final List responses = new ArrayList<>(); + final StreamingChatResponseHandler[] handlerRef = new StreamingChatResponseHandler[1]; + doAnswer( + invocation -> { + handlerRef[0] = invocation.getArgument(1); + return null; + }) + .when(streamingChatModel) + .chat(any(ChatRequest.class), any(StreamingChatResponseHandler.class)); + final var disposable = + streamingLangChain4j.generateContent(llmRequest, true).subscribe(responses::add); + + handlerRef[0].onPartialResponse("Par"); + handlerRef[0].onPartialResponse("is"); + final ChatResponse chatResponse = mock(ChatResponse.class); + when(chatResponse.aiMessage()).thenReturn(AiMessage.from("Paris")); + handlerRef[0].onCompleteResponse(chatResponse); + + assertThat(responses).hasSize(3); + final LlmResponse last = responses.get(2); + assertThat(last.partial()).isEmpty(); + assertThat(last.content().orElseThrow().role()).hasValue("model"); + assertThat(last.content().orElseThrow().text()).isEqualTo("Paris"); + disposable.dispose(); + } + + @Test + @DisplayName("Should end a streamed turn with its text and calls in one response") + void testGenerateContentWithStreamingChatModel_endsWithTextAndCallsInOneResponse() { + final LlmRequest llmRequest = + LlmRequest.builder().contents(List.of(Content.fromParts(Part.fromText("Hello")))).build(); + final List responses = new ArrayList<>(); + final StreamingChatResponseHandler[] handlerRef = new StreamingChatResponseHandler[1]; + doAnswer( + invocation -> { + handlerRef[0] = invocation.getArgument(1); + return null; + }) + .when(streamingChatModel) + .chat(any(ChatRequest.class), any(StreamingChatResponseHandler.class)); + final var disposable = + streamingLangChain4j.generateContent(llmRequest, true).subscribe(responses::add); + + handlerRef[0].onPartialResponse("Let me check."); + final AiMessage aiMessage = + AiMessage.builder() + .text("Let me check.") + .toolExecutionRequests( + List.of( + ToolExecutionRequest.builder() + .id("c1") + .name("getWeather") + .arguments("{\"city\":\"Paris\"}") + .build(), + ToolExecutionRequest.builder() + .id("c2") + .name("getWeather") + .arguments("{\"city\":\"Rome\"}") + .build())) + .build(); + final ChatResponse chatResponse = mock(ChatResponse.class); + when(chatResponse.aiMessage()).thenReturn(aiMessage); + handlerRef[0].onCompleteResponse(chatResponse); + + assertThat(responses).hasSize(2); + final LlmResponse last = responses.get(1); + assertThat(last.partial()).isEmpty(); + final List parts = last.content().orElseThrow().parts().orElseThrow(); + assertThat(parts).hasSize(3); + assertThat(parts.get(0).text()).hasValue("Let me check."); + assertThat(parts.subList(1, 3).stream().map(part -> part.functionCall().orElseThrow().id())) + .containsExactly(Optional.of("c1"), Optional.of("c2")); + disposable.dispose(); + } + @Test @DisplayName("Should pass configuration options to LangChain4j") void testGenerateContentWithConfigOptions() { diff --git a/core/src/main/java/com/google/adk/SchemaUtils.java b/core/src/main/java/com/google/adk/SchemaUtils.java index 0df7ee775..53eb05fed 100644 --- a/core/src/main/java/com/google/adk/SchemaUtils.java +++ b/core/src/main/java/com/google/adk/SchemaUtils.java @@ -17,17 +17,29 @@ package com.google.adk; import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.ObjectReader; import com.google.common.base.Preconditions; +import com.google.common.collect.ImmutableList; import com.google.genai.types.Schema; import com.google.genai.types.Type; +import java.math.BigInteger; +import java.util.Collection; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Map.Entry; +import java.util.Optional; +import org.jspecify.annotations.Nullable; -/** Utility class for validating schemas. */ +/** Utility class for reading JSON and validating values against schemas. */ public final class SchemaUtils { + private static final ObjectReader STRICT_JSON = + JsonBaseModel.getMapper() + .readerFor(Object.class) + .with(DeserializationFeature.FAIL_ON_TRAILING_TOKENS); + private SchemaUtils() {} // Private constructor for utility class /** @@ -128,6 +140,68 @@ public static void validateMapOnSchema(Map args, Schema schema, } } + /** + * Parses model text against {@code schema} as one JSON value, or returns the text itself when + * that fails and the schema accepts a string. Blank text, and a JSON null the schema allows, give + * {@code null}. + * + *

This method is only for use by Agent Development Kit. + * + * @throws IllegalArgumentException if the text does not match the schema + */ + public static @Nullable Object parseModelOutput(String output, Schema schema) { + if (output.isBlank()) { + return null; + } + String reason; + try { + Object value = readStrictJson(output); + reason = mismatch(value, schema); + if (reason == null) { + return value; + } + } catch (JsonProcessingException e) { + reason = "it is not JSON"; + } + if (acceptsString(schema)) { + return output; + } + // The reason names only parts of the schema, never the model's output. + throw new IllegalArgumentException( + "validation error: the output does not match its schema; " + reason + "."); + } + + /** + * Returns whether {@code schema} accepts a string, directly or through an {@code anyOf}. + * + *

This method is only for use by Agent Development Kit. + */ + public static boolean acceptsString(Schema schema) { + return typeOf(schema) == Type.Known.STRING + || schema.anyOf().orElse(ImmutableList.of()).stream().anyMatch(SchemaUtils::acceptsString); + } + + /** + * Reads {@code text} as one JSON value, failing on trailing text. + * + *

This method is only for use by Agent Development Kit. + * + * @throws JsonProcessingException if the text is not one JSON value + */ + public static @Nullable Object readStrictJson(String text) throws JsonProcessingException { + return STRICT_JSON.readValue(text); + } + + /** + * Returns why {@code value} does not match {@code schema}, or empty if it does. The reason names + * parts of the schema but never the value, which comes from the user. + * + *

This method is only for use by Agent Development Kit. + */ + public static Optional findMismatch(@Nullable Object value, Schema schema) { + return Optional.ofNullable(mismatch(value, schema)); + } + /** * Validates an output string against a schema. * @@ -144,4 +218,81 @@ public static Map validateOutputSchema(String output, Schema sch validateMapOnSchema(outputMap, schema, false); return outputMap; } + + private static @Nullable String mismatch(@Nullable Object value, Schema schema) { + if (value == null && schema.nullable().orElse(false)) { + return null; + } + // Check maps before anyOf so each member is validated as an object with its required keys. + if (value instanceof Map map) { + return mapMismatch(map, schema); + } + List anyOf = schema.anyOf().orElse(ImmutableList.of()); + if (!anyOf.isEmpty() && anyOf.stream().allMatch(member -> mismatch(value, member) != null)) { + return "it matches none of the schema's anyOf alternatives"; + } + Type.Known type = typeOf(schema); + if (type == null || type == Type.Known.TYPE_UNSPECIFIED) { + return null; + } + boolean matches = + switch (type) { + case STRING -> value instanceof String; + case INTEGER -> + value instanceof Integer + || value instanceof Long + || value instanceof Short + || value instanceof Byte + || value instanceof BigInteger; + case NUMBER -> value instanceof Number; + case BOOLEAN -> value instanceof Boolean; + case NULL -> value == null; + case ARRAY -> { + if (!(value instanceof Collection items)) { + yield false; + } + Optional itemSchema = schema.items(); + yield itemSchema.isEmpty() + || items.stream().allMatch(item -> mismatch(item, itemSchema.get()) == null); + } + case OBJECT -> false; + default -> true; + }; + return matches ? null : "it is not of the schema's type " + type; + } + + /** Returns why {@code map} does not match {@code schema}, or {@code null} if it does. */ + private static @Nullable String mapMismatch(Map map, Schema schema) { + Type.Known type = typeOf(schema); + if (type != null && type != Type.Known.OBJECT && type != Type.Known.TYPE_UNSPECIFIED) { + return "the schema's type is " + type + ", not an object"; + } + List anyOf = schema.anyOf().orElse(ImmutableList.of()); + if (!anyOf.isEmpty() && anyOf.stream().allMatch(member -> mapMismatch(map, member) != null)) { + return "it matches none of the schema's anyOf alternatives"; + } + Optional> properties = schema.properties(); + if (properties.isPresent()) { + for (Map.Entry entry : map.entrySet()) { + Schema property = properties.get().get(String.valueOf(entry.getKey())); + // An undeclared key comes from the user, so the reason does not name it. + if (property == null) { + return "it has a key the schema does not declare"; + } + if (mismatch(entry.getValue(), property) != null) { + return "property '" + entry.getKey() + "' does not match its schema"; + } + } + } + for (String required : schema.required().orElse(ImmutableList.of())) { + if (!map.containsKey(required)) { + return "required property '" + required + "' is missing"; + } + } + return null; + } + + private static Type.@Nullable Known typeOf(Schema schema) { + return schema.type().map(Type::knownEnum).orElse(null); + } } diff --git a/core/src/main/java/com/google/adk/agents/AgentMode.java b/core/src/main/java/com/google/adk/agents/AgentMode.java new file mode 100644 index 000000000..204f13c9c --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/AgentMode.java @@ -0,0 +1,26 @@ +/* + * 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.agents; + +import com.google.adk.annotations.Experimental; + +/** How an {@link LlmAgent} converses when it runs as a workflow node. */ +@Experimental +public enum AgentMode { + /** Answers one input with one response, without conversing with the user. */ + SINGLE_TURN, +} diff --git a/core/src/main/java/com/google/adk/agents/BaseAgent.java b/core/src/main/java/com/google/adk/agents/BaseAgent.java index da43622b1..03715a341 100644 --- a/core/src/main/java/com/google/adk/agents/BaseAgent.java +++ b/core/src/main/java/com/google/adk/agents/BaseAgent.java @@ -509,8 +509,14 @@ private Event checkpointEvent(InvocationContext context, EventActions actions) { @Override public Flowable runNode( com.google.adk.agents.Context context, @Nullable Object nodeInput) { + return runAsNode(context, context.invocationContext()); + } + + /** Runs this agent as {@link #runNode} does, but on {@code invocationContext}. */ + final Flowable runAsNode( + com.google.adk.agents.Context context, InvocationContext invocationContext) { NodeExecutionState nodeState = context.requireNodeState(); - return runWithCallbacks(context.invocationContext(), /* keepBranch= */ true, this::runAsyncImpl) + return runWithCallbacks(invocationContext, /* keepBranch= */ true, this::runAsyncImpl) .doOnNext( event -> { if (!isNullOrEmpty(event.author())) { diff --git a/core/src/main/java/com/google/adk/agents/InvocationContext.java b/core/src/main/java/com/google/adk/agents/InvocationContext.java index 338eaa00f..04c81a1dd 100644 --- a/core/src/main/java/com/google/adk/agents/InvocationContext.java +++ b/core/src/main/java/com/google/adk/agents/InvocationContext.java @@ -86,6 +86,7 @@ public class InvocationContext { @Nullable private String branch; private BaseAgent agent; private final @Nullable Node node; + private final @Nullable AgentModeOverride agentModeOverride; private boolean endInvocation; protected InvocationContext(Builder builder) { @@ -100,6 +101,7 @@ protected InvocationContext(Builder builder) { this.invocationId = builder.resolveInvocationId(); this.agent = builder.agent; this.node = builder.node; + this.agentModeOverride = builder.agentModeOverride; this.session = builder.session; this.userContent = builder.userContent; this.runConfig = builder.runConfig; @@ -194,6 +196,19 @@ public Optional node() { return Optional.ofNullable(node); } + /** + * Returns the mode {@code agent} runs in as a workflow node in this context, or empty when it + * runs without one. + * + *

This method is only for use by Agent Development Kit. + */ + @Experimental + public Optional agentMode(BaseAgent agent) { + return agentModeOverride != null && agentModeOverride.agent().equals(agent) + ? Optional.of(agentModeOverride.mode()) + : Optional.empty(); + } + /** Returns the session associated with this invocation. */ public Session session() { return session; @@ -650,6 +665,11 @@ public int hashCode() { } } + /** + * The mode one agent runs in, keyed to that agent so the sub-agents it reaches run without it. + */ + private record AgentModeOverride(BaseAgent agent, AgentMode mode) {} + /** Builder for {@link InvocationContext}. */ public static class Builder { @@ -667,6 +687,7 @@ private Builder(InvocationContext context) { this.invocationIdSet = true; this.agent = context.agent; this.node = context.node; + this.agentModeOverride = context.agentModeOverride; this.session = context.session; this.userContent = context.userContent; this.runConfig = context.runConfig; @@ -699,6 +720,7 @@ private Builder(InvocationContext context) { private boolean invocationIdSet; private BaseAgent agent; private @Nullable Node node; + private @Nullable AgentModeOverride agentModeOverride; private Session session; @Nullable private Content userContent = null; private RunConfig runConfig = RunConfig.builder().build(); @@ -822,6 +844,13 @@ public Builder node(@Nullable Node node) { return this; } + /** Sets the mode {@code agent} runs in as a workflow node, replacing any earlier one. */ + @CanIgnoreReturnValue + Builder agentMode(BaseAgent agent, AgentMode mode) { + this.agentModeOverride = new AgentModeOverride(agent, mode); + return this; + } + /** * Sets the session associated with this invocation. * @@ -1024,6 +1053,7 @@ public boolean equals(Object o) { && Objects.equals(invocationId, that.invocationId) && Objects.equals(agent, that.agent) && Objects.equals(node, that.node) + && Objects.equals(agentModeOverride, that.agentModeOverride) && Objects.equals(session, that.session) && Objects.equals(userContent, that.userContent) && Objects.equals(runConfig, that.runConfig) @@ -1050,6 +1080,7 @@ public int hashCode() { invocationId, agent, node, + agentModeOverride, session, userContent, runConfig, diff --git a/core/src/main/java/com/google/adk/agents/LlmAgent.java b/core/src/main/java/com/google/adk/agents/LlmAgent.java index 737058039..fdf2b0314 100644 --- a/core/src/main/java/com/google/adk/agents/LlmAgent.java +++ b/core/src/main/java/com/google/adk/agents/LlmAgent.java @@ -21,6 +21,7 @@ import static java.util.stream.Collectors.joining; import com.fasterxml.jackson.core.JsonProcessingException; +import com.google.adk.JsonBaseModel; import com.google.adk.SchemaUtils; import com.google.adk.agents.Callbacks.AfterAgentCallbackSync; import com.google.adk.agents.Callbacks.AfterModelCallback; @@ -52,8 +53,10 @@ import com.google.adk.models.BaseLlm; import com.google.adk.models.LlmRegistry; import com.google.adk.models.Model; +import com.google.adk.sessions.Session; import com.google.adk.tools.BaseTool; import com.google.adk.tools.BaseToolset; +import com.google.adk.workflow.NodeInfo; import com.google.common.base.Preconditions; import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; @@ -118,6 +121,7 @@ public enum IncludeContents { private final Optional executor; private final Optional outputKey; private final Optional codeExecutor; + private final @Nullable AgentMode mode; private volatile Model resolvedModel; private final BaseLlmFlow llmFlow; @@ -153,6 +157,7 @@ protected LlmAgent(Builder builder) { this.toolsUnion = requireNonNullElse(builder.toolsUnion, ImmutableList.of()); this.toolsets = extractToolsets(this.toolsUnion); this.codeExecutor = Optional.ofNullable(builder.codeExecutor); + this.mode = builder.mode; this.llmFlow = determineLlmFlow(); @@ -197,6 +202,7 @@ public static class Builder extends BaseAgent.Builder { private Executor executor; private String outputKey; private BaseCodeExecutor codeExecutor; + private AgentMode mode; @CanIgnoreReturnValue public Builder model(String model) { @@ -577,6 +583,17 @@ public Builder codeExecutor(BaseCodeExecutor codeExecutor) { return this; } + /** + * Sets how this agent converses when it runs as a workflow node. An agent without a mode runs + * as it does outside a workflow. + */ + @CanIgnoreReturnValue + @Experimental + public Builder mode(AgentMode mode) { + this.mode = mode; + return this; + } + private static @Nullable ImmutableList convertCallbacks( @Nullable List callbacks, Function converter, String callbackType) { return Optional.ofNullable(callbacks) @@ -621,10 +638,14 @@ protected BaseLlmFlow determineLlmFlow() { } } - private void maybeSaveOutputToState(Event event) { + private void maybeSaveOutputToState(Event event, InvocationContext invocationContext) { if (outputKey().isEmpty() || !event.finalResponse() || event.content().isEmpty()) { return; } + // A single-turn node's outputKey holds its own answer, not that of an agent it transferred to. + if (runsSingleTurn(invocationContext) && !name().equals(event.author())) { + return; + } List parts = event.content().flatMap(Content::parts).orElseGet(ImmutableList::of); // Skip events with no non-thought text part (e.g. a function-call-only long-running call, or a @@ -636,16 +657,21 @@ private void maybeSaveOutputToState(Event event) { return; } - // Concatenate text from all parts, excluding thoughts. - String rawResult = - parts.stream() - .filter(part -> !isThought(part)) - .map(part -> part.text().orElse("")) - .collect(joining()); + String rawResult = visibleText(event.content().get()); Object output = rawResult; Optional outputSchema = outputSchema(); - if (outputSchema.isPresent()) { + if (outputSchema.isPresent() && runsSingleTurn(invocationContext)) { + try { + output = SchemaUtils.parseModelOutput(rawResult, outputSchema.get()); + } catch (IllegalArgumentException e) { + // Like the node output, the saved value must match the schema. + return; + } + if (output == null) { + return; + } + } else if (outputSchema.isPresent()) { try { output = SchemaUtils.validateOutputSchema(rawResult, outputSchema.get()); } catch (JsonProcessingException e) { @@ -669,10 +695,24 @@ private static boolean isThought(Part part) { return part.thought().isPresent() && part.thought().get(); } + /** Returns the text of the parts of {@code content} that are not thoughts. */ + private static String visibleText(Content content) { + return content.parts().orElseGet(ImmutableList::of).stream() + .filter(part -> !isThought(part)) + .map(part -> part.text().orElse("")) + .collect(joining()); + } + + private boolean runsSingleTurn(InvocationContext invocationContext) { + return invocationContext.agentMode(this).filter(AgentMode.SINGLE_TURN::equals).isPresent(); + } + @Override protected Flowable runAsyncImpl(InvocationContext invocationContext) { if (!invocationContext.isResumable()) { - return llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState); + return llmFlow + .run(invocationContext) + .doOnNext(event -> maybeSaveOutputToState(event, invocationContext)); } return Flowable.defer( () -> { @@ -702,7 +742,9 @@ protected Flowable runAsyncImpl(InvocationContext invocationContext) { } // Normal path: emit an end-of-agent checkpoint so a later run can skip this agent. Flowable events = - llmFlow.run(invocationContext).doOnNext(this::maybeSaveOutputToState); + llmFlow + .run(invocationContext) + .doOnNext(event -> maybeSaveOutputToState(event, invocationContext)); AtomicBoolean paused = new AtomicBoolean(false); return events .doOnNext( @@ -802,9 +844,153 @@ private Optional findSubAgentToResume(InvocationContext context) { return Optional.empty(); } + /** + * Runs this agent as a workflow node. A {@link AgentMode#SINGLE_TURN} agent gets the node input + * as its user turn, and its first final model turn is the node output. An agent without a mode + * runs as {@link BaseAgent#runNode} runs it. + */ + @Experimental + @Override + public Flowable runNode(Context context, @Nullable Object nodeInput) { + if (mode != AgentMode.SINGLE_TURN) { + return super.runNode(context, nodeInput); + } + InvocationContext invocationContext = context.invocationContext(); + Flowable run = + runAsNode( + context, invocationContext.toBuilder().agentMode(this, AgentMode.SINGLE_TURN).build()); + // Eager disposal removes the input before the run ends, so a successor never sees it. + return Flowable.using( + () -> injectNodeInput(context, nodeInput), + unused -> { + // The first final turn with an output is the node's output; later content is not. + AtomicBoolean outputMarked = new AtomicBoolean(); + return run.map( + event -> { + if (outputMarked.get()) { + return event; + } + Optional marked = markOutputIfFinalTurn(event); + outputMarked.set(marked.isPresent()); + return marked.orElse(event); + }); + }, + injected -> injected.ifPresent(event -> removeEvent(invocationContext.session(), event)), + /* eager= */ true); + } + + /** + * Returns {@code event} carrying the node output when it is this agent's final model turn with + * text, or else empty. Without an output schema, the event's message is the output. + * + * @throws IllegalArgumentException if the text does not match the output schema + */ + private Optional markOutputIfFinalTurn(Event event) { + Optional content = + event.content().filter(c -> c.role().filter("model"::equals).isPresent()); + if (!name().equals(event.author()) + || !event.finalResponse() + || !event.functionCalls().isEmpty() + || content.isEmpty()) { + return Optional.empty(); + } + String text = visibleText(content.get()); + Object output = + outputSchema.isPresent() + ? SchemaUtils.parseModelOutput(text, outputSchema.get()) + : (text.isBlank() ? null : text.trim()); + if (output == null) { + return Optional.empty(); + } + // The output field gives the successor the trimmed text; the stored event drops it. + Event.Builder marked = event.toBuilder().output(output); + if (outputSchema.isEmpty()) { + marked.nodeInfo( + event + .nodeInfo() + .map(NodeInfo::toBuilder) + .orElseGet(NodeInfo::builder) + .messageAsOutput(true) + .build()); + } + return Optional.of(marked.build()); + } + + /** + * Adds the node input to the session as the agent's user turn and returns it. A resuming node + * gets none, and the user message the runner already recorded is not added twice. + */ + private static Optional injectNodeInput(Context context, @Nullable Object nodeInput) { + InvocationContext invocationContext = context.invocationContext(); + if (nodeInput == null + || !context.resumeInputs().isEmpty() + || isRecordedUserMessage(invocationContext, nodeInput)) { + return Optional.empty(); + } + Event event = + Event.builder() + .id(invocationContext.newUuid()) + .invocationId(invocationContext.invocationId()) + .author(Role.USER) + .content(nodeInputToUserContent(nodeInput)) + .branch(invocationContext.branch().orElse(null)) + .nodeInfo(NodeInfo.builder().path(context.requireNodeState().nodePath()).build()) + .timestamp(invocationContext.now().toEpochMilli()) + .build(); + invocationContext.session().addEvent(event); + return Optional.of(event); + } + + private static boolean isRecordedUserMessage( + InvocationContext invocationContext, Object nodeInput) { + return invocationContext.userContent().filter(nodeInput::equals).isPresent() + && invocationContext.session().immutableEvents().stream() + .anyMatch( + event -> + invocationContext.invocationId().equals(event.invocationId()) + && Objects.equals(event.author(), Role.USER) + && event.nodeInfo().map(NodeInfo::path).orElse("").isEmpty() + && event.content().filter(nodeInput::equals).isPresent()); + } + + /** + * Converts a node input into a user turn, writing a map or list as JSON text. + * + * @throws IllegalArgumentException if a map or list holds a value that JSON cannot represent + */ + @SuppressWarnings("UnusedException") // Jackson's message can quote the input, so it is dropped. + private static Content nodeInputToUserContent(Object nodeInput) { + if (nodeInput instanceof Content content) { + return content.toBuilder().role(Role.USER).build(); + } + String text; + if (nodeInput instanceof Map || nodeInput instanceof List) { + try { + text = JsonBaseModel.getMapper().writeValueAsString(nodeInput); + } catch (JsonProcessingException e) { + throw new IllegalArgumentException( + "A node input map or list must hold only values that JSON can represent."); + } + } else { + text = String.valueOf(nodeInput); + } + return Content.fromParts(Part.fromText(text)); + } + + // Session has no remove helper, so this locks the event list as Session.addEvent does. + @SuppressWarnings("deprecation") + private static void removeEvent(Session session, Event event) { + List events = session.events(); + synchronized (events) { + events.removeIf(event::equals); + } + } + @Override protected Flowable runLiveImpl(InvocationContext invocationContext) { - return llmFlow.runLive(invocationContext).doOnNext(this::maybeSaveOutputToState); + return llmFlow + .runLive(invocationContext) + .doOnNext(event -> maybeSaveOutputToState(event, invocationContext)); } /** @@ -908,6 +1094,12 @@ public IncludeContents includeContents() { return includeContents; } + /** Returns how this agent converses when it runs as a workflow node, or empty without a mode. */ + @Experimental + public Optional mode() { + return Optional.ofNullable(mode); + } + public Single> tools() { return canonicalTools().toList(); } diff --git a/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java b/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java index 62ff664b5..14c9fcce9 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/AgentTransfer.java @@ -16,6 +16,7 @@ package com.google.adk.flows.llmflows; +import com.google.adk.agents.AgentMode; import com.google.adk.agents.BaseAgent; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LlmAgent; @@ -50,9 +51,11 @@ public Single processRequest( RequestProcessor.RequestProcessingResult.create(request, ImmutableList.of())); } - LlmRequest.Builder builder = - request.toBuilder() - .appendInstructions(ImmutableList.of(buildTargetAgentsInstructions(agent, targets))); + LlmRequest.Builder builder = request.toBuilder(); + // As in Python ADK, a single-turn node can transfer but is not told to. + if (context.agentMode(agent).filter(AgentMode.SINGLE_TURN::equals).isEmpty()) { + builder.appendInstructions(ImmutableList.of(buildTargetAgentsInstructions(agent, targets))); + } FunctionTool agentTransferTool = createTransferToAgentTool(); agentTransferTool.processLlmRequest(builder, ToolContext.builder(context).build()); diff --git a/core/src/main/java/com/google/adk/flows/llmflows/Identity.java b/core/src/main/java/com/google/adk/flows/llmflows/Identity.java index 54aeb3b3a..8ddfa58a7 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/Identity.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/Identity.java @@ -16,6 +16,7 @@ package com.google.adk.flows.llmflows; +import com.google.adk.agents.AgentMode; import com.google.adk.agents.BaseAgent; import com.google.adk.agents.InvocationContext; import com.google.adk.models.LlmRequest; @@ -23,7 +24,10 @@ import com.google.common.collect.ImmutableList; import io.reactivex.rxjava3.core.Single; -/** {@link RequestProcessor} that gives the agent identity from the framework */ +/** + * {@link RequestProcessor} that gives the agent identity from the framework. An agent that runs as + * a single-turn workflow node gets none. + */ public final class Identity implements RequestProcessor { public Identity() {} @@ -32,6 +36,10 @@ public Identity() {} public Single processRequest( InvocationContext context, LlmRequest request) { BaseAgent agent = context.agent(); + if (context.agentMode(agent).filter(AgentMode.SINGLE_TURN::equals).isPresent()) { + return Single.just( + RequestProcessor.RequestProcessingResult.create(request, ImmutableList.of())); + } StringBuilder builder = new StringBuilder() .append("You are an agent. Your internal name is ") diff --git a/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java b/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java index 4f7022914..88f52cf02 100644 --- a/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java +++ b/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java @@ -17,6 +17,7 @@ package com.google.adk.tools; import com.google.adk.SchemaUtils; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.Schema; @@ -77,34 +78,43 @@ public Single> runAsync(Map args, ToolContex + "\nRecall the set_model_response function correctly, fix the errors, and" + " call it again with all required fields using the correct types."); } - // Match Python's model_dump(exclude_none=True) for Java's map-shaped response. - Map validatedResponse = excludeNullFields(args); + Map validatedResponse = excludeNullFields(args, outputSchema); toolContext.actions().setSetModelResponse(validatedResponse); return validatedResponse; }); } - private static Map excludeNullFields(Map values) { + /** + * Drops nulls like Python's {@code model_dump(exclude_none=True)}, but keeps the null of a + * property that {@code schema} marks required and nullable, which the schema demands. + */ + private static Map excludeNullFields(Map values, Schema schema) { + // Validation passed, so each key has a property schema and each non-empty list an item schema. + Map properties = schema.properties().orElseThrow(); + List required = schema.required().orElse(ImmutableList.of()); Map result = new LinkedHashMap<>(); for (Map.Entry entry : values.entrySet()) { + Schema property = properties.get(entry.getKey()); Object value = entry.getValue(); if (value != null) { - result.put(entry.getKey(), excludeNullsFromValue(value)); + result.put(entry.getKey(), excludeNullsFromValue(value, property)); + } else if (property.nullable().orElse(false) && required.contains(entry.getKey())) { + result.put(entry.getKey(), null); } } return result; } - private static Object excludeNullsFromValue(Object value) { + private static Object excludeNullsFromValue(Object value, Schema schema) { if (value instanceof Map) { @SuppressWarnings("unchecked") // Function-call args always have String keys. Map map = (Map) value; - return excludeNullFields(map); + return excludeNullFields(map, schema); } - if (value instanceof List) { + if (value instanceof List list) { List result = new ArrayList<>(); - for (Object item : (List) value) { - result.add(excludeNullsFromValue(item)); + for (Object item : list) { + result.add(excludeNullsFromValue(item, schema.items().orElseThrow())); } return result; } diff --git a/core/src/main/java/com/google/adk/workflow/ParallelWorkers.java b/core/src/main/java/com/google/adk/workflow/ParallelWorkers.java index 129d4f5c6..c6d3b0423 100644 --- a/core/src/main/java/com/google/adk/workflow/ParallelWorkers.java +++ b/core/src/main/java/com/google/adk/workflow/ParallelWorkers.java @@ -17,16 +17,12 @@ package com.google.adk.workflow; import com.fasterxml.jackson.core.JsonProcessingException; -import com.fasterxml.jackson.databind.DeserializationFeature; -import com.fasterxml.jackson.databind.ObjectReader; -import com.google.adk.JsonBaseModel; +import com.google.adk.SchemaUtils; import com.google.adk.agents.Context; import com.google.adk.annotations.Experimental; import com.google.common.collect.ImmutableList; import com.google.genai.types.Content; import com.google.genai.types.Part; -import com.google.genai.types.Schema; -import com.google.genai.types.Type; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Single; import java.util.ArrayList; @@ -80,10 +76,6 @@ public static Node parallelNode(Node node, int maxParallelWorkers) { private static final class ParallelWorkerNode extends BaseNode { - private static final ObjectReader JSON_READER = - JsonBaseModel.getMapper() - .readerFor(Object.class) - .with(DeserializationFeature.FAIL_ON_TRAILING_TOKENS); // No run ID: items subscribe in input order, so each gets the same generated ID on resume. private static final RunNodeConfig ITEM_CONFIG = RunNodeConfig.builder().useSubBranch(true).build(); @@ -132,22 +124,14 @@ public Flowable runNode(Context context, @Nullable Object nodeInput) { String text = content.text(); Object json; try { - json = JSON_READER.readValue(text); + json = SchemaUtils.readStrictJson(text); } catch (JsonProcessingException e) { return text; } if (json instanceof List list) { return list; } - return worker.inputSchema().map(ParallelWorkerNode::acceptsString).orElse(false) - ? text - : json; - } - - private static boolean acceptsString(Schema schema) { - return schema.type().map(type -> type.knownEnum() == Type.Known.STRING).orElse(false) - || schema.anyOf().orElse(ImmutableList.of()).stream() - .anyMatch(ParallelWorkerNode::acceptsString); + return worker.inputSchema().map(SchemaUtils::acceptsString).orElse(false) ? text : json; } private final class FanOut { diff --git a/core/src/main/java/com/google/adk/workflow/ResumeScan.java b/core/src/main/java/com/google/adk/workflow/ResumeScan.java index b1791981d..ef20168fb 100644 --- a/core/src/main/java/com/google/adk/workflow/ResumeScan.java +++ b/core/src/main/java/com/google/adk/workflow/ResumeScan.java @@ -19,9 +19,8 @@ import static com.google.common.base.Strings.isNullOrEmpty; import com.fasterxml.jackson.core.JsonProcessingException; -import com.fasterxml.jackson.databind.DeserializationFeature; -import com.fasterxml.jackson.databind.ObjectReader; import com.google.adk.JsonBaseModel; +import com.google.adk.SchemaUtils; import com.google.adk.agents.Context; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.Role; @@ -37,9 +36,6 @@ import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; import com.google.genai.types.Schema; -import com.google.genai.types.Type; -import java.math.BigInteger; -import java.util.Collection; import java.util.Collections; import java.util.HashMap; import java.util.HashSet; @@ -65,10 +61,6 @@ final class ResumeScan { private static final String RESULT_KEY = "result"; // Stands for text that is not JSON, since the JSON text null reads as a null value. private static final Object NOT_JSON = new Object(); - private static final ObjectReader STRICT_JSON = - JsonBaseModel.getMapper() - .readerFor(Object.class) - .with(DeserializationFeature.FAIL_ON_TRAILING_TOKENS); private ResumeScan() {} @@ -206,7 +198,7 @@ static Context replayContext( } Optional schema = node.outputSchema(); // Model text is raw rather than JSON-quoted, so it stays text when the schema accepts strings. - if (schema.isEmpty() || acceptsString(schema.get())) { + if (schema.isEmpty() || SchemaUtils.acceptsString(schema.get())) { return text; } Object value = readJson(text); @@ -230,10 +222,10 @@ static Context replayContext( if (schema == null) { return answer; } - String mismatch = valueMismatch(answer, schema); - if (mismatch != null) { + Optional mismatch = SchemaUtils.findMismatch(answer, schema); + if (mismatch.isPresent()) { throw new IllegalArgumentException( - "Validation failed for interrupt " + interruptId + ": " + mismatch); + "Validation failed for interrupt " + interruptId + ": " + mismatch.get()); } return answer; } @@ -243,110 +235,22 @@ static Context replayContext( return response; } Object value = map.get(RESULT_KEY); - if (!(value instanceof String text) || (schema != null && acceptsString(schema))) { + if (!(value instanceof String text) || (schema != null && SchemaUtils.acceptsString(schema))) { return value; } Object json = readJson(text); return json == NOT_JSON ? text : json; } - /** Returns whether {@code schema} accepts a string, directly or through an {@code anyOf}. */ - private static boolean acceptsString(Schema schema) { - return typeOf(schema) == Type.Known.STRING - || schema.anyOf().orElse(ImmutableList.of()).stream().anyMatch(ResumeScan::acceptsString); - } - /** Reads {@code text} as strict JSON, or returns {@link #NOT_JSON} if it is not JSON. */ private static @Nullable Object readJson(String text) { try { - return STRICT_JSON.readValue(text); + return SchemaUtils.readStrictJson(text); } catch (JsonProcessingException e) { return NOT_JSON; } } - /** - * Returns why {@code value} does not match {@code schema}, or null if it does. The reason names - * parts of the schema but never the value, which comes from the user. - */ - private static @Nullable String valueMismatch(@Nullable Object value, Schema schema) { - if (value == null && schema.nullable().orElse(false)) { - return null; - } - // Check maps before anyOf so each member is validated as an object with its required keys. - if (value instanceof Map map) { - return mapMismatch(map, schema); - } - List anyOf = schema.anyOf().orElse(ImmutableList.of()); - if (!anyOf.isEmpty() - && anyOf.stream().allMatch(member -> valueMismatch(value, member) != null)) { - return "it matches none of the schema's anyOf alternatives"; - } - Type.Known type = typeOf(schema); - if (type == null || type == Type.Known.TYPE_UNSPECIFIED) { - return null; - } - boolean matches = - switch (type) { - case STRING -> value instanceof String; - case INTEGER -> - value instanceof Integer - || value instanceof Long - || value instanceof Short - || value instanceof Byte - || value instanceof BigInteger; - case NUMBER -> value instanceof Number; - case BOOLEAN -> value instanceof Boolean; - case NULL -> value == null; - case ARRAY -> { - if (!(value instanceof Collection items)) { - yield false; - } - Optional itemSchema = schema.items(); - yield itemSchema.isEmpty() - || items.stream().allMatch(item -> valueMismatch(item, itemSchema.get()) == null); - } - case OBJECT -> false; - default -> true; - }; - return matches ? null : "it is not of the schema's type " + type; - } - - /** Returns why {@code map} does not match {@code schema}, or null if it does. */ - private static @Nullable String mapMismatch(Map map, Schema schema) { - Type.Known type = typeOf(schema); - if (type != null && type != Type.Known.OBJECT && type != Type.Known.TYPE_UNSPECIFIED) { - return "the schema's type is " + type + ", not an object"; - } - List anyOf = schema.anyOf().orElse(ImmutableList.of()); - if (!anyOf.isEmpty() && anyOf.stream().allMatch(member -> mapMismatch(map, member) != null)) { - return "it matches none of the schema's anyOf alternatives"; - } - Optional> properties = schema.properties(); - if (properties.isPresent()) { - for (Map.Entry entry : map.entrySet()) { - Schema property = properties.get().get(String.valueOf(entry.getKey())); - // An undeclared key comes from the user, so the reason does not name it. - if (property == null) { - return "it has a key the schema does not declare"; - } - if (valueMismatch(entry.getValue(), property) != null) { - return "property '" + entry.getKey() + "' does not match its schema"; - } - } - } - for (String required : schema.required().orElse(ImmutableList.of())) { - if (!map.containsKey(required)) { - return "required property '" + required + "' is missing"; - } - } - return null; - } - - private static Type.@Nullable Known typeOf(Schema schema) { - return schema.type().map(Type::knownEnum).orElse(null); - } - /** Returns the text of the non-thought parts of {@code content}. */ private static String textOf(Content content) { StringBuilder text = new StringBuilder(); diff --git a/core/src/main/java/com/google/adk/workflow/ToolNodes.java b/core/src/main/java/com/google/adk/workflow/ToolNodes.java index cbc220f58..b4f3654e6 100644 --- a/core/src/main/java/com/google/adk/workflow/ToolNodes.java +++ b/core/src/main/java/com/google/adk/workflow/ToolNodes.java @@ -20,10 +20,9 @@ import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.core.type.TypeReference; -import com.fasterxml.jackson.databind.DeserializationFeature; import com.fasterxml.jackson.databind.JsonNode; -import com.fasterxml.jackson.databind.ObjectReader; import com.google.adk.JsonBaseModel; +import com.google.adk.SchemaUtils; import com.google.adk.agents.Context; import com.google.adk.agents.InvocationContext; import com.google.adk.annotations.Experimental; @@ -85,11 +84,6 @@ public static Node asNode(BaseTool tool, String name, NodeConfig 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) { @@ -226,7 +220,7 @@ private static boolean declaresNoResult(BaseTool tool) { return null; } try { - Object json = JSON_READER.readValue(text); + Object json = SchemaUtils.readStrictJson(text); return json instanceof Map ? json : text; } catch (JsonProcessingException e) { return text; diff --git a/core/src/test/java/com/google/adk/SchemaUtilsTest.java b/core/src/test/java/com/google/adk/SchemaUtilsTest.java index ac5df39a8..fd2fd44f8 100644 --- a/core/src/test/java/com/google/adk/SchemaUtilsTest.java +++ b/core/src/test/java/com/google/adk/SchemaUtilsTest.java @@ -16,12 +16,17 @@ package com.google.adk; +import static com.google.common.truth.Truth.assertThat; import static org.junit.Assert.assertThrows; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Schema; +import java.io.IOException; +import java.math.BigInteger; import java.util.HashMap; import java.util.Map; +import java.util.Optional; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -30,6 +35,9 @@ @RunWith(JUnit4.class) public final class SchemaUtilsTest { + private static final Schema INTEGER = Schema.builder().type("INTEGER").build(); + private static final Schema STRING = Schema.builder().type("STRING").build(); + @Test public void validateMapOnSchema_nullableField_allowsNull() { Schema schema = @@ -98,4 +106,164 @@ public void validateMapOnSchema_integerField_allowsIntegerAndLong() { args.put("intField", 123); SchemaUtils.validateMapOnSchema(args, schema, /* isInput= */ true); } + + @Test + public void readStrictJson_oneValue_readsIt() throws Exception { + Object value = SchemaUtils.readStrictJson("[1, \"a\"]"); + + assertThat(value).isEqualTo(ImmutableList.of(1, "a")); + } + + @Test + public void readStrictJson_trailingText_throws() { + // JsonProcessingException is an IOException; Jackson itself is not visible to this test. + assertThrows(IOException.class, () -> SchemaUtils.readStrictJson("7 days")); + } + + @Test + public void findMismatch_valueMatchingOneAnyOfAlternative_matches() { + Schema schema = Schema.builder().anyOf(ImmutableList.of(INTEGER, STRING)).build(); + + Optional mismatch = SchemaUtils.findMismatch(7, schema); + + assertThat(mismatch).isEmpty(); + } + + @Test + public void findMismatch_valueMatchingNoAnyOfAlternative_explainsIt() { + Schema schema = Schema.builder().anyOf(ImmutableList.of(INTEGER, STRING)).build(); + + Optional mismatch = SchemaUtils.findMismatch(true, schema); + + assertThat(mismatch.orElseThrow()).contains("anyOf"); + } + + @Test + public void findMismatch_objectSchemaWithoutProperties_acceptsAnyMap() { + Schema schema = Schema.builder().type("OBJECT").build(); + + Optional mismatch = SchemaUtils.findMismatch(ImmutableMap.of("a", 1), schema); + + assertThat(mismatch).isEmpty(); + } + + @Test + public void findMismatch_undeclaredKey_explainsItWithoutNamingTheKey() { + Schema schema = + Schema.builder().type("OBJECT").properties(ImmutableMap.of("a", INTEGER)).build(); + + Optional mismatch = SchemaUtils.findMismatch(ImmutableMap.of("secret_key", 1), schema); + + assertThat(mismatch.orElseThrow()).doesNotContain("secret_key"); + } + + @Test + public void acceptsString_anyOfWithAString_isTrue() { + Schema schema = Schema.builder().anyOf(ImmutableList.of(INTEGER, STRING)).build(); + + assertThat(SchemaUtils.acceptsString(schema)).isTrue(); + } + + @Test + public void acceptsString_integerSchema_isFalse() { + assertThat(SchemaUtils.acceptsString(INTEGER)).isFalse(); + } + + @Test + public void findMismatch_nullForANullableSchema_matches() { + Schema schema = Schema.builder().type("INTEGER").nullable(true).build(); + + assertThat(SchemaUtils.findMismatch(null, schema)).isEmpty(); + } + + @Test + public void findMismatch_bigIntegerForAnIntegerSchema_matches() { + assertThat(SchemaUtils.findMismatch(new BigInteger("123456789012345678901"), INTEGER)) + .isEmpty(); + } + + @Test + public void findMismatch_arrayItemOfTheWrongType_explainsIt() { + Schema schema = Schema.builder().type("ARRAY").items(INTEGER).build(); + + assertThat(SchemaUtils.findMismatch(ImmutableList.of(1, "two"), schema)).isPresent(); + } + + @Test + public void findMismatch_missingRequiredProperty_namesIt() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("a", INTEGER)) + .required(ImmutableList.of("a")) + .build(); + + Optional mismatch = SchemaUtils.findMismatch(ImmutableMap.of(), schema); + + assertThat(mismatch.orElseThrow()).contains("'a'"); + } + + @Test + public void findMismatch_mapForANonObjectSchema_explainsIt() { + assertThat(SchemaUtils.findMismatch(ImmutableMap.of("a", 1), INTEGER)).isPresent(); + } + + @Test + public void parseModelOutput_quotedTextWithAStringSchema_decodesIt() { + Object output = SchemaUtils.parseModelOutput("\"Paris\"", STRING); + + assertThat(output).isEqualTo("Paris"); + } + + @Test + public void parseModelOutput_numberWithAStringSchema_keepsTheText() { + Object output = SchemaUtils.parseModelOutput("42", STRING); + + assertThat(output).isEqualTo("42"); + } + + @Test + public void parseModelOutput_blankText_givesNull() { + Object output = SchemaUtils.parseModelOutput(" ", STRING); + + assertThat(output).isNull(); + } + + @Test + public void parseModelOutput_valueMatchingAnAnyOfAlternative_givesIt() { + Schema schema = Schema.builder().anyOf(ImmutableList.of(INTEGER, STRING)).build(); + + Object output = SchemaUtils.parseModelOutput("7", schema); + + assertThat(output).isEqualTo(7); + } + + @Test + public void parseModelOutput_stringForAnIntegerProperty_throwsWithoutQuotingIt() { + Schema schema = + Schema.builder().type("OBJECT").properties(ImmutableMap.of("n", INTEGER)).build(); + + var thrown = + assertThrows( + IllegalArgumentException.class, + () -> SchemaUtils.parseModelOutput("{\"n\": \"seven\"}", schema)); + + assertThat(thrown).hasMessageThat().contains("validation error"); + assertThat(thrown).hasMessageThat().doesNotContain("seven"); + } + + @Test + public void parseModelOutput_jsonNullWithANullableSchema_givesNull() { + Schema schema = Schema.builder().type("INTEGER").nullable(true).build(); + + Object output = SchemaUtils.parseModelOutput("null", schema); + + assertThat(output).isNull(); + } + + @Test + public void parseModelOutput_jsonNullWithANonNullableSchema_throwsIllegalArgumentException() { + assertThrows( + IllegalArgumentException.class, () -> SchemaUtils.parseModelOutput("null", INTEGER)); + } } diff --git a/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java b/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java index 1964daa98..ea35f2f85 100644 --- a/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java +++ b/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java @@ -21,6 +21,7 @@ import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.common.truth.Truth.assertThat; +import com.google.adk.SchemaUtils; import com.google.adk.models.LlmResponse; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -123,6 +124,96 @@ public void runAsync_nullableNull_excludesNullFromReturnedAndRecordedResponse() assertThat(toolContext.actions().setModelResponse()).hasValue(ImmutableMap.of("a", "x")); } + @Test + public void runAsync_requiredNullableNull_keepsTheNull() { + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "a", + Schema.builder().type("STRING").build(), + "b", + Schema.builder().type("STRING").nullable(true).build())) + .required(ImmutableList.of("a", "b")) + .build(); + Map args = new LinkedHashMap<>(); + args.put("a", "x"); + args.put("b", null); + + ToolContext toolContext = createToolContext(); + Map result = + new SetModelResponseTool(outputSchema).runAsync(args, toolContext).blockingGet(); + + assertThat(result).containsExactly("a", "x", "b", null).inOrder(); + assertThat(toolContext.actions().setModelResponse()).hasValue(result); + // Throws if the result does not match the schema, as dropping the null did. + SchemaUtils.validateMapOnSchema(result, outputSchema, /* isInput= */ false); + } + + @Test + public void runAsync_nestedRequiredNullableNull_keepsTheNull() { + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("addr", requiredNullableZipAddressSchema())) + .required(ImmutableList.of("addr")) + .build(); + Map address = new LinkedHashMap<>(); + address.put("city", "NYC"); + address.put("zip", null); + + ToolContext toolContext = createToolContext(); + Map result = + new SetModelResponseTool(outputSchema) + .runAsync(ImmutableMap.of("addr", address), toolContext) + .blockingGet(); + + assertThat(result).isEqualTo(ImmutableMap.of("addr", address)); + assertThat(toolContext.actions().setModelResponse()).hasValue(result); + } + + @Test + public void runAsync_arrayItemRequiredNullableNull_keepsTheNull() { + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "addresses", + Schema.builder() + .type("ARRAY") + .items(requiredNullableZipAddressSchema()) + .build())) + .required(ImmutableList.of("addresses")) + .build(); + Map address = new LinkedHashMap<>(); + address.put("city", "NYC"); + address.put("zip", null); + + ToolContext toolContext = createToolContext(); + Map result = + new SetModelResponseTool(outputSchema) + .runAsync(ImmutableMap.of("addresses", ImmutableList.of(address)), toolContext) + .blockingGet(); + + assertThat(result).isEqualTo(ImmutableMap.of("addresses", ImmutableList.of(address))); + assertThat(toolContext.actions().setModelResponse()).hasValue(result); + } + + private static Schema requiredNullableZipAddressSchema() { + return Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "city", + Schema.builder().type("STRING").build(), + "zip", + Schema.builder().type("STRING").nullable(true).build())) + .required(ImmutableList.of("city", "zip")) + .build(); + } + @Test public void runAsync_nestedNullableNull_excludesNullFromReturnedAndRecordedResponse() { Schema addressSchema = diff --git a/core/src/test/java/com/google/adk/workflow/LlmAgentAsNodeTest.java b/core/src/test/java/com/google/adk/workflow/LlmAgentAsNodeTest.java new file mode 100644 index 000000000..e9853aed4 --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/LlmAgentAsNodeTest.java @@ -0,0 +1,708 @@ +/* + * 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.ResumabilityTestUtils.pendingFunctionTool; +import static com.google.adk.testing.TestUtils.createFunctionCallLlmResponse; +import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createLlmResponse; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.adk.workflow.WorkflowResumeTesting.answer; +import static com.google.adk.workflow.WorkflowResumeTesting.go; +import static com.google.adk.workflow.WorkflowResumeTesting.longRunningAsk; +import static com.google.common.collect.ImmutableList.toImmutableList; +import static com.google.common.collect.Iterables.getLast; +import static com.google.common.collect.Iterables.getOnlyElement; +import static com.google.common.collect.MoreCollectors.onlyElement; +import static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; + +import com.google.adk.agents.AgentMode; +import com.google.adk.agents.Context; +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.LlmAgent; +import com.google.adk.apps.App; +import com.google.adk.events.Event; +import com.google.adk.models.LlmRequest; +import com.google.adk.runner.Runner; +import com.google.adk.sessions.InMemorySessionService; +import com.google.adk.sessions.Session; +import com.google.adk.testing.TestLlm; +import com.google.adk.testing.TestNode; +import com.google.adk.testing.TestUtils.EchoTool; +import com.google.adk.workflow.WorkflowResumeTesting.WorkflowSession; +import com.google.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.errorprone.annotations.CanIgnoreReturnValue; +import com.google.genai.types.Content; +import com.google.genai.types.FunctionCall; +import com.google.genai.types.Part; +import com.google.genai.types.Schema; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import org.jspecify.annotations.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for running an {@link LlmAgent} as a workflow node. */ +@RunWith(JUnit4.class) +public final class LlmAgentAsNodeTest { + + private static final String APP_NAME = "app"; + private static final String USER_ID = "user"; + private static final NodeRunner.EventSink SINK = (event, nodeFailure) -> Completable.complete(); + + @Test + public void runAsync_singleTurnNode_getsItsNodeInputAsItsUserTurn() { + TestLlm model = TestLlm.create("done"); + + runAfterSource("translate this", singleTurnAgent("worker", model)); + + assertThat(getLast(model.getLastRequest().contents())).isEqualTo(userTurn("translate this")); + } + + @Test + public void runAsync_singleTurnNode_getsAContentInputAsAUserTurn() { + TestLlm model = TestLlm.create("done"); + Content input = Content.builder().role("model").parts(Part.fromText("from upstream")).build(); + + runAfterSource(input, singleTurnAgent("worker", model)); + + assertThat(getLast(model.getLastRequest().contents())).isEqualTo(userTurn("from upstream")); + } + + @Test + public void runAsync_singleTurnNode_getsAMapInputAsJson() { + TestLlm model = TestLlm.create("done"); + + runAfterSource(ImmutableMap.of("city", "Kraków", "days", 3), singleTurnAgent("worker", model)); + + assertThat(getLast(model.getLastRequest().contents())) + .isEqualTo(userTurn("{\"city\":\"Kraków\",\"days\":3}")); + } + + @Test + public void runAsync_singleTurnNode_getsAListInputAsJson() { + TestLlm model = TestLlm.create("done"); + + runAfterSource( + ImmutableList.of(1, ImmutableMap.of("tags", ImmutableList.of("a", "b"))), + singleTurnAgent("worker", model)); + + assertThat(getLast(model.getLastRequest().contents())) + .isEqualTo(userTurn("[1,{\"tags\":[\"a\",\"b\"]}]")); + } + + @Test + public void runAsync_firstSingleTurnNode_seesTheUserMessageOnce() { + TestLlm model = TestLlm.create("done"); + LlmAgent agent = singleTurnAgent("worker", model); + + run(Workflow.builder().name("wf").edges(new Edge(Start.INSTANCE, agent)).build(), "go"); + + assertThat(model.getLastRequest().contents()).containsExactly(userTurn("go")); + } + + @Test + public void runAsync_singleTurnNodeWithoutInput_addsNoUserTurn() { + TestLlm model = TestLlm.create("done"); + Node silent = new TestNode("source", (context, input) -> Flowable.empty()); + LlmAgent agent = singleTurnAgent("worker", model); + + run( + Workflow.builder() + .name("wf") + .edges(new Edge(Start.INSTANCE, silent), new Edge(silent, agent)) + .build(), + "go"); + + assertThat(model.getLastRequest().contents()).containsExactly(userTurn("go")); + } + + @Test + public void runAsync_nodeAfterASingleTurnNode_doesNotSeeItsInput() { + TestLlm laterModel = TestLlm.create("later"); + Node source = new TestNode("source", (context, input) -> Flowable.just("one-off input")); + LlmAgent first = singleTurnAgent("first", TestLlm.create("first reply")); + LlmAgent later = LlmAgent.builder().name("later").model(laterModel).build(); + + run( + Workflow.builder() + .name("wf") + .edges( + new Edge(Start.INSTANCE, source), new Edge(source, first), new Edge(first, later)) + .build(), + "go"); + + assertThat(laterModel.getLastRequest().contents()).doesNotContain(userTurn("one-off input")); + } + + @Test + public void runAsync_agentWithoutAMode_doesNotGetItsNodeInput() { + TestLlm model = TestLlm.create("done"); + + runAfterSource("current task", LlmAgent.builder().name("worker").model(model).build()); + + assertThat(model.getLastRequest().contents()).doesNotContain(userTurn("current task")); + } + + @Test + public void runAsync_agentWithoutAMode_keepsItsIdentity() { + TestLlm model = TestLlm.create("done"); + + runAfterSource("go", LlmAgent.builder().name("worker").model(model).build()); + + assertThat(String.join("\n", model.getLastRequest().getSystemInstructions())) + .contains("Your internal name is \"worker\""); + } + + @Test + public void runAsync_singleTurnNode_getsNoIdentity() { + TestLlm model = TestLlm.create("done"); + + runAfterSource("go", singleTurnAgent("worker", model)); + + assertThat(String.join("\n", model.getLastRequest().getSystemInstructions())) + .doesNotContain("Your internal name is"); + } + + @Test + public void runAsync_singleTurnNodeWithASubAgent_getsTheTransferToolWithoutInstructions() { + TestLlm model = TestLlm.create("done"); + LlmAgent helper = LlmAgent.builder().name("helper").model(TestLlm.create("help")).build(); + LlmAgent router = + LlmAgent.builder() + .name("router") + .model(model) + .mode(AgentMode.SINGLE_TURN) + .subAgents(helper) + .build(); + + runAfterSource("go", router); + + LlmRequest request = model.getLastRequest(); + assertThat(request.tools()).containsKey("transfer_to_agent"); + assertThat(String.join("\n", request.getSystemInstructions())) + .doesNotContain("transfer_to_agent"); + } + + @Test + public void runAsync_subAgentReachedFromASingleTurnNode_keepsItsIdentity() { + TestLlm helperModel = TestLlm.create("help"); + LlmAgent helper = LlmAgent.builder().name("helper").model(helperModel).build(); + TestLlm routerModel = + TestLlm.create( + createFunctionCallLlmResponse( + "call_1", "transfer_to_agent", ImmutableMap.of("agent_name", "helper"))); + LlmAgent router = + LlmAgent.builder() + .name("router") + .model(routerModel) + .mode(AgentMode.SINGLE_TURN) + .subAgents(helper) + .build(); + + runAfterSource("go", router); + + assertThat(String.join("\n", helperModel.getLastRequest().getSystemInstructions())) + .contains("Your internal name is \"helper\""); + } + + @Test + public void runAsync_secondSingleTurnNode_getsTheFirstNodesAnswerAsItsUserTurn() { + LlmAgent first = singleTurnAgent("first", modelAnswering("Paris")); + TestLlm secondModel = TestLlm.create("done"); + LlmAgent second = singleTurnAgent("second", secondModel); + + run( + Workflow.builder() + .name("wf") + .edges(new Edge(Start.INSTANCE, first), new Edge(first, second)) + .build(), + "Capital of France?"); + + assertThat(getLast(secondModel.getLastRequest().contents())).isEqualTo(userTurn("Paris")); + } + + @Test + public void runAsync_singleTurnNode_givesItsTrimmedAnswerOnTheAnswerEventOnly() { + NodeRun nodeRun = runWithSuccessor(singleTurnAgent("worker", modelAnswering(" the answer\n"))); + + assertThat(nodeRun.output()).isEqualTo("the answer"); + Event outputEvent = + nodeRun.events().stream() + .filter(event -> !outputFor(event).isEmpty()) + .collect(onlyElement()); + assertThat(outputEvent.nodeInfo().map(NodeInfo::messageAsOutput)).hasValue(true); + assertThat(outputEvent.output()).isEmpty(); + } + + @Test + public void runAsync_singleTurnNode_leavesThoughtsOutOfItsOutputAndOutputKey() { + Content answer = + Content.builder() + .role("model") + .parts( + Part.builder().text("internal reasoning").thought(true).build(), + Part.fromText("visible answer")) + .build(); + LlmAgent agent = + LlmAgent.builder() + .name("worker") + .model(TestLlm.create(createLlmResponse(answer))) + .mode(AgentMode.SINGLE_TURN) + .outputKey("summary") + .build(); + + NodeRun nodeRun = runWithSuccessor(agent); + + assertThat(nodeRun.output()).isEqualTo("visible answer"); + assertThat(stateDeltaValues(nodeRun.events(), "summary")).containsExactly("visible answer"); + } + + @Test + public void runAsync_singleTurnNodeThatCallsATool_givesItsAnswerTurnAsItsOutput() { + TestLlm model = + TestLlm.create( + createFunctionCallLlmResponse("call_1", "echo_tool", ImmutableMap.of()), + createTextLlmResponse("It is 72")); + LlmAgent agent = + LlmAgent.builder() + .name("worker") + .model(model) + .mode(AgentMode.SINGLE_TURN) + .tools(new EchoTool()) + .build(); + + NodeRun nodeRun = runWithSuccessor(agent); + + assertThat(nodeRun.output()).isEqualTo("It is 72"); + assertThat(nodeRun.events().stream().filter(event -> !outputFor(event).isEmpty())).hasSize(1); + } + + @Test + public void runAsync_textBesideALongRunningCall_isSavedButIsNotTheOutput() { + Content turn = + Content.builder() + .role("model") + .parts( + Part.fromText("Checking with the user."), + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call_1") + .name("pendingTool") + .args(ImmutableMap.of("message", "hello"))) + .build()) + .build(); + LlmAgent agent = + LlmAgent.builder() + .name("worker") + .model(TestLlm.create(createLlmResponse(turn))) + .mode(AgentMode.SINGLE_TURN) + .tools(pendingFunctionTool()) + .outputKey("answer") + .build(); + + List events = + run(Workflow.builder().name("wf").edges(new Edge(Start.INSTANCE, agent)).build(), "go"); + + assertThat(stateDeltaValues(events, "answer")).containsExactly("Checking with the user."); + assertThat(events.stream().filter(event -> !outputFor(event).isEmpty())).isEmpty(); + } + + @Test + public void runAsync_contentAfterTheAnswer_isNotASecondOutput() { + LlmAgent agent = + LlmAgent.builder() + .name("worker") + .model(modelAnswering("answer")) + .mode(AgentMode.SINGLE_TURN) + .afterAgentCallback( + context -> + Maybe.just( + Content.builder().role("model").parts(Part.fromText("note")).build())) + .build(); + + NodeRun nodeRun = runWithSuccessor(agent); + + assertThat(nodeRun.output()).isEqualTo("answer"); + } + + @Test + public void runAsync_thoughtOnlyAnswer_givesNoOutput() { + Content thoughtOnly = + Content.builder() + .role("model") + .parts(Part.builder().text("internal reasoning").thought(true).build()) + .build(); + + NodeRun nodeRun = + runWithSuccessor(singleTurnAgent("worker", TestLlm.create(createLlmResponse(thoughtOnly)))); + + assertThat(nodeRun.output()).isNull(); + } + + @Test + public void runAsync_answerWithoutTheModelRole_givesNoOutput() { + // As in Python ADK, only a model turn is the output; the model adapters all set the role. + NodeRun nodeRun = runWithSuccessor(singleTurnAgent("worker", TestLlm.create("done"))); + + assertThat(nodeRun.output()).isNull(); + } + + @Test + public void runAsync_singleTurnNodeWithAnOutputSchema_givesAndSavesTheParsedValue() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "title", + Schema.builder().type("STRING").build(), + "count", + Schema.builder().type("INTEGER").build())) + .build(); + LlmAgent agent = + schemaAgent(schema, "{\"title\": \"Story\", \"count\": 3}").outputKey("meta").build(); + + NodeRun nodeRun = runWithSuccessor(agent); + + ImmutableMap expected = ImmutableMap.of("title", "Story", "count", 3); + assertThat(nodeRun.output()).isEqualTo(expected); + assertThat(stateDeltaValues(nodeRun.events(), "meta")).containsExactly(expected); + Event outputEvent = + nodeRun.events().stream() + .filter(event -> !outputFor(event).isEmpty()) + .collect(onlyElement()); + assertThat(outputEvent.output()).hasValue(expected); + assertThat(outputEvent.nodeInfo().map(NodeInfo::messageAsOutput)).hasValue(false); + } + + @Test + public void runAsync_singleTurnNodeWithAStringSchema_givesItsPlainText() { + LlmAgent agent = + schemaAgent(Schema.builder().type("STRING").build(), "The answer is 42").build(); + + NodeRun nodeRun = runWithSuccessor(agent); + + assertThat(nodeRun.output()).isEqualTo("The answer is 42"); + } + + @Test + public void runAsync_singleTurnNodeWithAnArraySchema_givesAndSavesAList() { + Schema schema = + Schema.builder().type("ARRAY").items(Schema.builder().type("INTEGER").build()).build(); + + NodeRun nodeRun = + runWithSuccessor(schemaAgent(schema, "[1, 2, 3]").outputKey("numbers").build()); + + assertThat(nodeRun.output()).isEqualTo(ImmutableList.of(1, 2, 3)); + assertThat(stateDeltaValues(nodeRun.events(), "numbers")) + .containsExactly(ImmutableList.of(1, 2, 3)); + } + + @Test + public void runAsync_blankAnswerWithAnOutputSchema_givesNoOutput() { + Schema schema = Schema.builder().type("INTEGER").build(); + + NodeRun nodeRun = runWithSuccessor(schemaAgent(schema, " ").build()); + + assertThat(nodeRun.output()).isNull(); + } + + @Test + public void runAsync_answerNotMatchingTheOutputSchema_failsWithoutQuotingIt() { + Schema schema = Schema.builder().type("INTEGER").build(); + LlmAgent agent = schemaAgent(schema, "sorry, I do not know").build(); + + var thrown = assertThrows(IllegalArgumentException.class, () -> runWithSuccessor(agent)); + + assertThat(thrown).hasMessageThat().contains("validation error"); + assertThat(thrown).hasMessageThat().doesNotContain("sorry"); + } + + @Test + public void runAsync_singleTurnNodeThatTransfers_givesNoOutputAndSavesNoReply() { + LlmAgent helper = LlmAgent.builder().name("helper").model(modelAnswering("help")).build(); + TestLlm routerModel = + TestLlm.create( + createFunctionCallLlmResponse( + "call_1", "transfer_to_agent", ImmutableMap.of("agent_name", "helper"))); + LlmAgent router = + LlmAgent.builder() + .name("router") + .model(routerModel) + .mode(AgentMode.SINGLE_TURN) + .subAgents(helper) + .outputKey("route") + .build(); + + NodeRun nodeRun = runWithSuccessor(router); + + assertThat(nodeRun.output()).isNull(); + assertThat(stateDeltaValues(nodeRun.events(), "route")).isEmpty(); + } + + @Test + public void runAsync_agentWithoutAModeThatTransfers_stillSavesTheReplyToItsOutputKey() { + LlmAgent helper = LlmAgent.builder().name("helper").model(modelAnswering("help")).build(); + TestLlm routerModel = + TestLlm.create( + createFunctionCallLlmResponse( + "call_1", "transfer_to_agent", ImmutableMap.of("agent_name", "helper"))); + LlmAgent router = + LlmAgent.builder() + .name("router") + .model(routerModel) + .subAgents(helper) + .outputKey("route") + .build(); + + NodeRun nodeRun = runWithSuccessor(router); + + assertThat(stateDeltaValues(nodeRun.events(), "route")).contains("help"); + } + + @Test + public void runAsync_agentWithoutAMode_givesNoOutput() { + LlmAgent agent = LlmAgent.builder().name("worker").model(modelAnswering("done")).build(); + + NodeRun nodeRun = runWithSuccessor(agent); + + assertThat(nodeRun.output()).isNull(); + } + + @Test + public void runAsync_singleTurnNodeBeforeAPause_isNotRunAgainOnResume() { + TestLlm model = modelAnswering("summary"); + LlmAgent summarizer = singleTurnAgent("summarizer", model); + Node ask = longRunningAsk("ask", "ask_1"); + WorkflowSession session = + new WorkflowSession( + Workflow.builder() + .name("wf") + .edges(new Edge(Start.INSTANCE, summarizer), new Edge(summarizer, ask)) + .build()); + session.advance(go()); + + ImmutableList resumed = session.turn(answer("ask_1", ImmutableMap.of("ok", true))); + + assertThat(model.getRequests()).hasSize(1); + assertThat(resumed.stream().map(Event::author)).doesNotContain("summarizer"); + } + + @Test + public void runNode_singleTurnNode_givesItsInputTheNodeBranchAndPath() { + InMemorySessionService sessionService = new InMemorySessionService(); + Session session = sessionService.createSession(APP_NAME, USER_ID).blockingGet(); + List seenInputs = new ArrayList<>(); + TestLlm model = + new TestLlm( + () -> { + seenInputs.add(getLast(session.immutableEvents())); + return Flowable.just(createTextLlmResponse("done")); + }); + LlmAgent agent = singleTurnAgent("worker", model); + InvocationContext invocationContext = + createInvocationContext(agent, sessionService, session).toBuilder() + .branch("wf@1.a@1") + .build(); + + agent + .runNode(nodeContext(agent, invocationContext), "scoped input") + .ignoreElements() + .blockingAwait(); + + Event input = getLast(seenInputs); + assertThat(input.content()).hasValue(userTurn("scoped input")); + assertThat(input.branch()).hasValue("wf@1.a@1"); + assertThat(input.nodeInfo().map(NodeInfo::path)).hasValue("worker@1"); + } + + @Test + public void runNode_singleTurnNode_removesItsInputBeforeItCompletes() { + // A successor can start on completion, so the input must already be gone. + LlmAgent agent = singleTurnAgent("worker", TestLlm.create("done")); + InvocationContext invocationContext = createInvocationContext(agent); + List> eventsAtCompletion = new ArrayList<>(); + + agent + .runNode(nodeContext(agent, invocationContext), "input") + .doOnComplete(() -> eventsAtCompletion.add(invocationContext.session().immutableEvents())) + .ignoreElements() + .blockingAwait(); + + assertThat(getOnlyElement(eventsAtCompletion)).isEmpty(); + } + + @Test + public void runNode_singleTurnNodeWhoseModelFails_stillRemovesItsInput() { + IllegalStateException failure = new IllegalStateException("model failed"); + LlmAgent agent = singleTurnAgent("worker", TestLlm.create(ImmutableList.of(), failure)); + InvocationContext invocationContext = createInvocationContext(agent); + Single> run = + agent.runNode(nodeContext(agent, invocationContext), "input").toList(); + + var thrown = assertThrows(IllegalStateException.class, run::blockingGet); + + assertThat(thrown).isSameInstanceAs(failure); + + assertThat(invocationContext.session().immutableEvents()).isEmpty(); + } + + @Test + public void runNode_canceledSingleTurnNode_removesItsInput() { + LlmAgent agent = singleTurnAgent("worker", new TestLlm(Flowable::never)); + InvocationContext invocationContext = createInvocationContext(agent); + TestSubscriber subscriber = + agent.runNode(nodeContext(agent, invocationContext), "input").test(); + assertThat(invocationContext.session().immutableEvents()).hasSize(1); + + subscriber.cancel(); + + assertThat(invocationContext.session().immutableEvents()).isEmpty(); + } + + @Test + public void runNode_resumingSingleTurnNode_addsNoUserTurn() { + TestLlm model = TestLlm.create("done"); + LlmAgent agent = singleTurnAgent("worker", model); + Context context = + nodeContext(agent, createInvocationContext(agent), ImmutableMap.of("ask_1", "yes")); + + agent.runNode(context, "input").ignoreElements().blockingAwait(); + + assertThat(model.getLastRequest().contents()).doesNotContain(userTurn("input")); + } + + @Test + public void runNode_mapInputHoldingAValueJsonCannotRepresent_failsWithoutQuotingIt() { + LlmAgent agent = singleTurnAgent("worker", TestLlm.create("done")); + Context context = nodeContext(agent, createInvocationContext(agent)); + // Jackson has no JSON form for an object without properties. + Object secret = + new Object() { + @Override + public String toString() { + return "secret-day"; + } + }; + Single> run = agent.runNode(context, ImmutableMap.of("slot", secret)).toList(); + + var thrown = assertThrows(IllegalArgumentException.class, run::blockingGet); + + assertThat(thrown).hasMessageThat().contains("node input"); + assertThat(thrown).hasMessageThat().doesNotContain("secret-day"); + } + + private static LlmAgent singleTurnAgent(String name, TestLlm model) { + return LlmAgent.builder().name(name).model(model).mode(AgentMode.SINGLE_TURN).build(); + } + + /** Returns a model that answers once with {@code text} in a model turn, as Gemini does. */ + private static TestLlm modelAnswering(String text) { + return TestLlm.create(createTextLlmResponse(text)); + } + + private static Content userTurn(String text) { + return Content.fromParts(Part.fromText(text)); + } + + /** Runs {@code agent} after a node whose output is {@code input}. */ + private static void runAfterSource(Object input, Node agent) { + Node source = new TestNode("source", (context, unused) -> Flowable.just(input)); + run( + Workflow.builder() + .name("wf") + .edges(new Edge(Start.INSTANCE, source), new Edge(source, agent)) + .build(), + "go"); + } + + @CanIgnoreReturnValue + private static List run(Node root, String message) { + Runner runner = + Runner.builder().app(App.builder().name(APP_NAME).rootNode(root).build()).build(); + Session session = runner.sessionService().createSession(APP_NAME, USER_ID).blockingGet(); + return runner.runAsync(USER_ID, session.id(), userTurn(message)).toList().blockingGet(); + } + + /** The events of a run and the input the agent's successor got, which is the agent's output. */ + private record NodeRun(ImmutableList events, @Nullable Object output) {} + + private static NodeRun runWithSuccessor(Node agent) { + List<@Nullable Object> inputs = new ArrayList<>(); + Node successor = + new TestNode( + "successor", + (context, input) -> { + inputs.add(input); + return Flowable.empty(); + }); + List events = + run( + Workflow.builder() + .name("wf") + .edges(new Edge(Start.INSTANCE, agent), new Edge(agent, successor)) + .build(), + "go"); + return new NodeRun(ImmutableList.copyOf(events), getOnlyElement(inputs)); + } + + private static LlmAgent.Builder schemaAgent(Schema schema, String answer) { + return LlmAgent.builder() + .name("worker") + .model(modelAnswering(answer)) + .mode(AgentMode.SINGLE_TURN) + .outputSchema(schema); + } + + private static ImmutableList outputFor(Event event) { + return event.nodeInfo().map(NodeInfo::outputFor).orElse(ImmutableList.of()); + } + + private static ImmutableList stateDeltaValues(List events, String key) { + return events.stream() + .filter(event -> event.actions().stateDelta().containsKey(key)) + .map(event -> event.actions().stateDelta().get(key)) + .collect(toImmutableList()); + } + + private static Context nodeContext(LlmAgent agent, InvocationContext invocationContext) { + return nodeContext(agent, invocationContext, ImmutableMap.of()); + } + + private static Context nodeContext( + LlmAgent agent, InvocationContext invocationContext, Map answers) { + return new Context( + invocationContext, + NodeExecutionState.builder() + .node(agent) + .eventSink(SINK) + .runId("1") + .resumeInputs(answers) + .build()); + } +}