From 265feeef0280bc5c7b6678f0e2dc06835ef1b75a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mi=C5=82osz=20Sobczyk?= Date: Thu, 29 Oct 2026 08:42:53 -0700 Subject: [PATCH] feat: run a `SINGLE_TURN` `LlmAgent` node on its node input PiperOrigin-RevId: 1009716409 --- .../adk/models/langchain4j/LangChain4j.java | 36 +- .../LangChain4jIntegrationTest.java | 22 +- .../models/langchain4j/LangChain4jTest.java | 92 +++++ .../java/com/google/adk/agents/AgentMode.java | 29 ++ .../java/com/google/adk/agents/BaseAgent.java | 8 +- .../google/adk/agents/InvocationContext.java | 31 ++ .../java/com/google/adk/agents/LlmAgent.java | 115 ++++++ .../adk/flows/llmflows/AgentTransfer.java | 9 +- .../google/adk/flows/llmflows/Identity.java | 10 +- .../adk/tools/SetModelResponseTool.java | 28 +- .../adk/tools/SetModelResponseToolTest.java | 91 +++++ .../adk/workflow/LlmAgentAsNodeTest.java | 364 ++++++++++++++++++ 12 files changed, 799 insertions(+), 36 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/agents/AgentMode.java b/core/src/main/java/com/google/adk/agents/AgentMode.java new file mode 100644 index 000000000..c6c847118 --- /dev/null +++ b/core/src/main/java/com/google/adk/agents/AgentMode.java @@ -0,0 +1,29 @@ +/* + * 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. The node does not yet + * give its response as its output. + */ + 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..7fda5e29f 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) @@ -802,6 +819,98 @@ 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; 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 -> run, + injected -> injected.ifPresent(event -> removeEvent(invocationContext.session(), event)), + /* eager= */ true); + } + + /** + * 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); @@ -908,6 +1017,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/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..ce620d8bb --- /dev/null +++ b/core/src/test/java/com/google/adk/workflow/LlmAgentAsNodeTest.java @@ -0,0 +1,364 @@ +/* + * 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.createFunctionCallLlmResponse; +import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createTextLlmResponse; +import static com.google.common.collect.Iterables.getLast; +import static com.google.common.collect.Iterables.getOnlyElement; +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.common.collect.ImmutableList; +import com.google.common.collect.ImmutableMap; +import com.google.genai.types.Content; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Flowable; +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 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(); + } + + 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"); + } + + private static void 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(); + runner.runAsync(USER_ID, session.id(), userTurn(message)).ignoreElements().blockingAwait(); + } + + 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()); + } +}