diff --git a/core/src/main/java/com/google/adk/events/EventActions.java b/core/src/main/java/com/google/adk/events/EventActions.java index cde23c10e..677ee2dde 100644 --- a/core/src/main/java/com/google/adk/events/EventActions.java +++ b/core/src/main/java/com/google/adk/events/EventActions.java @@ -45,6 +45,7 @@ public class EventActions extends JsonBaseModel { private ConcurrentMap requestedToolConfirmations; private boolean endOfAgent; private @Nullable EventCompaction compaction; + private @Nullable Object setModelResponse; /** Default constructor for Jackson. */ public EventActions() { @@ -67,6 +68,7 @@ private EventActions(Builder builder) { this.requestedToolConfirmations = builder.requestedToolConfirmations; this.endOfAgent = builder.endOfAgent; this.compaction = builder.compaction; + this.setModelResponse = builder.setModelResponse; } @JsonProperty("skipSummarization") @@ -201,6 +203,19 @@ public void setCompaction(@Nullable EventCompaction compaction) { this.compaction = compaction; } + /** + * The successfully validated structured response set by the {@code set_model_response} tool. + * Empty when the tool was not called or its arguments failed output-schema validation. + */ + @JsonProperty("setModelResponse") + public Optional setModelResponse() { + return Optional.ofNullable(setModelResponse); + } + + public void setSetModelResponse(@Nullable Object setModelResponse) { + this.setModelResponse = setModelResponse; + } + public static Builder builder() { return new Builder(); } @@ -226,7 +241,8 @@ public boolean equals(Object o) { && Objects.equals(requestedAuthConfigs, that.requestedAuthConfigs) && Objects.equals(requestedToolConfirmations, that.requestedToolConfirmations) && (endOfAgent == that.endOfAgent) - && Objects.equals(compaction, that.compaction); + && Objects.equals(compaction, that.compaction) + && Objects.equals(setModelResponse, that.setModelResponse); } @Override @@ -241,7 +257,8 @@ public int hashCode() { requestedAuthConfigs, requestedToolConfirmations, endOfAgent, - compaction); + compaction, + setModelResponse); } /** Builder for {@link EventActions}. */ @@ -256,6 +273,7 @@ public static class Builder { private ConcurrentMap requestedToolConfirmations; private boolean endOfAgent = false; private @Nullable EventCompaction compaction; + private @Nullable Object setModelResponse; public Builder() { this.stateDelta = new ConcurrentHashMap<>(); @@ -277,6 +295,7 @@ private Builder(EventActions eventActions) { new ConcurrentHashMap<>(eventActions.requestedToolConfirmations()); this.endOfAgent = eventActions.endOfAgent; this.compaction = eventActions.compaction; + this.setModelResponse = eventActions.setModelResponse; } @CanIgnoreReturnValue @@ -383,6 +402,13 @@ public Builder compaction(@Nullable EventCompaction value) { return this; } + @CanIgnoreReturnValue + @JsonProperty("setModelResponse") + public Builder setModelResponse(@Nullable Object value) { + this.setModelResponse = value; + return this; + } + @CanIgnoreReturnValue public Builder merge(EventActions other) { other.skipSummarization().ifPresent(this::skipSummarization); @@ -395,6 +421,7 @@ public Builder merge(EventActions other) { this.requestedToolConfirmations.putAll(other.requestedToolConfirmations()); this.endOfAgent = this.endOfAgent || other.endOfAgent(); other.compaction().ifPresent(this::compaction); + other.setModelResponse().ifPresent(this::setModelResponse); return this; } diff --git a/core/src/main/java/com/google/adk/flows/llmflows/OutputSchema.java b/core/src/main/java/com/google/adk/flows/llmflows/OutputSchema.java index d1f322f18..d1d214551 100644 --- a/core/src/main/java/com/google/adk/flows/llmflows/OutputSchema.java +++ b/core/src/main/java/com/google/adk/flows/llmflows/OutputSchema.java @@ -78,18 +78,24 @@ public Single processRequest( } /** - * Check if function response contains set_model_response and extract JSON. + * Extracts a successfully validated {@code set_model_response} result as JSON. + * + *

Only a result that passed output-schema validation (recorded on the event actions by {@link + * SetModelResponseTool}) is returned. Validation feedback sent back to the model is never + * promoted to the final structured response. * * @param functionResponseEvent The function response event to check. - * @return JSON response string if set_model_response was called, Optional.empty() otherwise. + * @return JSON response string if set_model_response succeeded, Optional.empty() otherwise. */ public static Optional getStructuredModelResponse(Event functionResponseEvent) { for (FunctionResponse funcResponse : functionResponseEvent.functionResponses()) { if (Objects.equals(funcResponse.name().orElse(""), SetModelResponseTool.NAME)) { - Object response = funcResponse.response(); - // The tool returns the args map directly. + Optional validatedResponse = functionResponseEvent.actions().setModelResponse(); + if (validatedResponse.isEmpty()) { + return Optional.empty(); + } try { - return Optional.of(JsonBaseModel.getMapper().writeValueAsString(response)); + return Optional.of(JsonBaseModel.getMapper().writeValueAsString(validatedResponse.get())); } catch (JsonProcessingException e) { logger.error("Failed to serialize set_model_response result", e); return Optional.empty(); 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 94569dd7d..487290528 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.ImmutableMap; import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.Schema; import io.reactivex.rxjava3.core.Single; @@ -33,6 +34,12 @@ public class SetModelResponseTool extends BaseTool { public static final String NAME = "set_model_response"; + // Prefix of the SchemaUtils validation message after which the full schema is appended. Used to + // strip the schema dump from feedback on a best-effort basis; if SchemaUtils changes its wording + // the feedback simply stays unstripped. runAsync_unknownArg_feedbackOmitsSchemaDump pins the + // current format. + private static final String OUTPUT_SCHEMA_DUMP_MARKER = " does not match agent output schema: "; + private final Schema outputSchema; public SetModelResponseTool(Schema outputSchema) { @@ -56,12 +63,35 @@ public Optional declaration() { @Override public Single> runAsync(Map args, ToolContext toolContext) { - // This tool is a marker for the final response, it doesn't do anything but return its arguments - // which will be captured as the final result. + // Record validated responses on the event actions; return validation feedback so the model can + // retry. return Single.fromCallable( () -> { - SchemaUtils.validateMapOnSchema(args, outputSchema, /* isInput= */ false); + try { + SchemaUtils.validateMapOnSchema(args, outputSchema, /* isInput= */ false); + } catch (IllegalArgumentException e) { + return ImmutableMap.of( + "error", + "Validation Error found:\n" + + sanitizeValidationMessage(e.getMessage()) + + "\nRecall the set_model_response function correctly, fix the errors, and" + + " call it again with all required fields using the correct types."); + } + toolContext.actions().setSetModelResponse(args); return args; }); } + + private static String sanitizeValidationMessage(String message) { + if (message == null) { + return "Arguments do not match the output schema."; + } + // The model already knows the schema from the tool declaration, so the appended schema dump is + // redundant in feedback. + int schemaDumpIndex = message.indexOf(OUTPUT_SCHEMA_DUMP_MARKER); + if (schemaDumpIndex >= 0) { + message = message.substring(0, schemaDumpIndex) + " does not match agent output schema."; + } + return message; + } } diff --git a/core/src/test/java/com/google/adk/events/EventActionsTest.java b/core/src/test/java/com/google/adk/events/EventActionsTest.java index c5949caf7..bced8086a 100644 --- a/core/src/test/java/com/google/adk/events/EventActionsTest.java +++ b/core/src/test/java/com/google/adk/events/EventActionsTest.java @@ -89,6 +89,7 @@ public void merge_mergesAllFields() { .requestedToolConfirmations( new ConcurrentHashMap<>(ImmutableMap.of("tool2", TOOL_CONFIRMATION))) .endOfAgent(true) + .setModelResponse(ImmutableMap.of("field1", "value1")) .build(); EventActions merged = eventActions1.toBuilder().merge(eventActions2).build(); @@ -109,6 +110,7 @@ public void merge_mergesAllFields() { .containsExactly("tool1", TOOL_CONFIRMATION, "tool2", TOOL_CONFIRMATION); assertThat(merged.endOfAgent()).isTrue(); assertThat(merged.compaction()).hasValue(COMPACTION); + assertThat(merged.setModelResponse()).hasValue(ImmutableMap.of("field1", "value1")); } @Test @@ -177,6 +179,7 @@ public void jsonSerialization_works() throws Exception { EventActions.builder() .deletedArtifactIds(ImmutableSet.of("d1", "d2")) .stateDelta(new ConcurrentHashMap<>(ImmutableMap.of("k", "v"))) + .setModelResponse(ImmutableMap.of("field1", "value1")) .build(); String json = eventActions.toJson(); @@ -184,6 +187,7 @@ public void jsonSerialization_works() throws Exception { assertThat(deserialized).isEqualTo(eventActions); assertThat(deserialized.deletedArtifactIds()).containsExactly("d1", "d2"); + assertThat(deserialized.setModelResponse()).hasValue(ImmutableMap.of("field1", "value1")); } @Test diff --git a/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java b/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java index ffd56de6c..98c5e9390 100644 --- a/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java +++ b/core/src/test/java/com/google/adk/flows/llmflows/OutputSchemaTest.java @@ -17,15 +17,22 @@ package com.google.adk.flows.llmflows; import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createLlmResponse; +import static com.google.adk.testing.TestUtils.createTestAgent; import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.common.truth.Truth.assertThat; import com.google.adk.agents.InvocationContext; import com.google.adk.agents.LlmAgent; import com.google.adk.events.Event; +import com.google.adk.events.EventActions; import com.google.adk.flows.llmflows.RequestProcessor.RequestProcessingResult; +import com.google.adk.models.BaseLlm; +import com.google.adk.models.BaseLlmConnection; import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; +import com.google.adk.runner.InMemoryRunner; +import com.google.adk.sessions.Session; import com.google.adk.testing.TestLlm; import com.google.adk.tools.BaseTool; import com.google.adk.tools.SetModelResponseTool; @@ -36,8 +43,12 @@ import com.google.genai.types.FunctionResponse; import com.google.genai.types.Part; import com.google.genai.types.Schema; +import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Single; +import java.util.List; import java.util.Map; +import java.util.Objects; +import java.util.Optional; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -143,10 +154,14 @@ public void getStructuredModelResponse_withSetModelResponse_returnsJson() { FunctionResponse fr = FunctionResponse.builder() .name(SetModelResponseTool.NAME) - .response(ImmutableMap.of("field1", "value1")) + .response(ImmutableMap.of("field1", "rawResponse")) .build(); Event event = Event.builder() + .actions( + EventActions.builder() + .setModelResponse(ImmutableMap.of("field1", "validatedValue")) + .build()) .content( Content.builder() .parts(Part.builder().functionResponse(fr).build()) @@ -154,7 +169,32 @@ public void getStructuredModelResponse_withSetModelResponse_returnsJson() { .build()) .build(); - assertThat(OutputSchema.getStructuredModelResponse(event)).hasValue("{\"field1\":\"value1\"}"); + // The result must come from the validated response on the event actions, not from the + // function response content. + assertThat(OutputSchema.getStructuredModelResponse(event)) + .hasValue("{\"field1\":\"validatedValue\"}"); + } + + @Test + public void getStructuredModelResponse_withValidationFeedback_returnsEmpty() { + FunctionResponse fr = + FunctionResponse.builder() + .name(SetModelResponseTool.NAME) + .response( + ImmutableMap.of( + "error", + "Validation Error found: field1 is required. Fix the errors and call it again.")) + .build(); + Event event = + Event.builder() + .content( + Content.builder() + .parts(Part.builder().functionResponse(fr).build()) + .role("user") + .build()) + .build(); + + assertThat(OutputSchema.getStructuredModelResponse(event)).isEmpty(); } @Test @@ -189,4 +229,123 @@ public void createFinalModelResponseEvent_createsModelResponseEvent() { assertThat(event.content().get().role()).hasValue("model"); assertThat(event.content().get().parts().get()).containsExactly(Part.fromText(jsonResponse)); } + + @Test + public void run_invalidThenValidSetModelResponse_emitsFeedbackThenFinalResponse() { + Content invalidCall = + Content.fromParts(Part.fromFunctionCall(SetModelResponseTool.NAME, ImmutableMap.of())); + Content validCall = + Content.fromParts( + Part.fromFunctionCall(SetModelResponseTool.NAME, ImmutableMap.of("field1", "value1"))); + TestLlm testLlm = createTestLlm(createLlmResponse(invalidCall), createLlmResponse(validCall)); + InvocationContext invocationContext = createInvocationContext(createTestAgent(testLlm)); + // Registers set_model_response on the request the same way OutputSchema.processRequest does, + // without the model-name gating which is covered separately above. + RequestProcessor injectSetModelResponseTool = + (context, request) -> { + LlmRequest.Builder builder = request.toBuilder(); + return new SetModelResponseTool(TEST_OUTPUT_SCHEMA) + .processLlmRequest(builder, ToolContext.builder(context).build()) + .andThen( + Single.fromCallable( + () -> RequestProcessingResult.create(builder.build(), ImmutableList.of()))); + }; + BaseLlmFlow flow = + new BaseLlmFlow( + ImmutableList.of(injectSetModelResponseTool), ImmutableList.of(), Optional.empty()) {}; + + List events = flow.run(invocationContext).toList().blockingGet(); + + // The invalid call must trigger a second LLM call (the retry); the tool must be declared on + // both requests. History assembly for the retry request is the Contents processor's own + // responsibility, covered by ContentsTest. + assertThat(testLlm.getRequests()).hasSize(2); + assertThat(testLlm.getRequests().get(0).tools()).containsKey(SetModelResponseTool.NAME); + assertThat(testLlm.getRequests().get(1).tools()).containsKey(SetModelResponseTool.NAME); + + // Invalid call produces only the feedback function response (no promoted final event), the + // corrected call produces the validated function response plus the final model response. + assertThat(events).hasSize(5); + + Event feedbackEvent = events.get(1); + Map feedback = feedbackEvent.functionResponses().get(0).response().get(); + assertThat(feedback).containsKey("error"); + assertThat((String) feedback.get("error")).contains("field1"); + assertThat(feedbackEvent.actions().setModelResponse()).isEmpty(); + + Event validatedEvent = events.get(3); + assertThat(validatedEvent.functionResponses().get(0).response().get()) + .containsExactly("field1", "value1"); + assertThat(validatedEvent.actions().setModelResponse()) + .hasValue(ImmutableMap.of("field1", "value1")); + + Event finalEvent = events.get(4); + assertThat(finalEvent.functionCalls()).isEmpty(); + assertThat(finalEvent.functionResponses()).isEmpty(); + assertThat(finalEvent.content().get().role()).hasValue("model"); + assertThat(finalEvent.content().get().parts().get().get(0).text()) + .hasValue("{\"field1\":\"value1\"}"); + } + + @Test + public void runner_invalidSetModelResponse_feedbackIsSentBackToModel() { + Content invalidCall = + Content.fromParts(Part.fromFunctionCall(SetModelResponseTool.NAME, ImmutableMap.of())); + Content validCall = + Content.fromParts( + Part.fromFunctionCall(SetModelResponseTool.NAME, ImmutableMap.of("field1", "value1"))); + TestLlm scriptedLlm = + createTestLlm(createLlmResponse(invalidCall), createLlmResponse(validCall)); + // The output-schema workaround only activates for models that cannot combine tools with an + // output schema, so the scripted TestLlm is wrapped under a gemini-2 model name. + BaseLlm gemini2NamedLlm = + new BaseLlm("gemini-2.0-flash") { + @Override + public Flowable generateContent(LlmRequest llmRequest, boolean stream) { + return scriptedLlm.generateContent(llmRequest, stream); + } + + @Override + public BaseLlmConnection connect(LlmRequest llmRequest) { + return scriptedLlm.connect(llmRequest); + } + }; + LlmAgent agent = + LlmAgent.builder() + .name("agent") + .model(gemini2NamedLlm) + .outputSchema(TEST_OUTPUT_SCHEMA) + .tools(ImmutableList.of(new TestTool())) + .build(); + InMemoryRunner runner = new InMemoryRunner(agent, "test-app"); + Session session = runner.sessionService().createSession("test-app", "user").blockingGet(); + + List events = + runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("hello"))) + .toList() + .blockingGet(); + + // The retry request must carry the validation feedback from the first call back to the model: + // the feedback names the missing required field, tying it to the first call's failure. + assertThat(scriptedLlm.getRequests()).hasSize(2); + boolean feedbackSentBack = + scriptedLlm.getRequests().get(1).contents().stream() + .flatMap(content -> content.parts().orElse(ImmutableList.of()).stream()) + .map(Part::functionResponse) + .flatMap(Optional::stream) + .anyMatch( + fr -> + Objects.equals(fr.name().orElse(""), SetModelResponseTool.NAME) + && fr.response().orElse(ImmutableMap.of()).get("error") + instanceof String error + && error.contains("Validation Error found") + && error.contains("field1")); + assertThat(feedbackSentBack).isTrue(); + + // Only the corrected, validated response becomes the final structured output. + Event finalEvent = events.get(events.size() - 1); + assertThat(finalEvent.content().get().parts().get().get(0).text()) + .hasValue("{\"field1\":\"value1\"}"); + } } 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 64b600af9..7ade26f31 100644 --- a/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java +++ b/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java @@ -16,9 +16,12 @@ package com.google.adk.tools; +import static com.google.adk.testing.TestUtils.createInvocationContext; +import static com.google.adk.testing.TestUtils.createTestAgent; +import static com.google.adk.testing.TestUtils.createTestLlm; import static com.google.common.truth.Truth.assertThat; -import static org.junit.Assert.assertThrows; +import com.google.adk.models.LlmResponse; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.FunctionDeclaration; @@ -31,6 +34,12 @@ @RunWith(JUnit4.class) public final class SetModelResponseToolTest { + private static ToolContext createToolContext() { + return ToolContext.builder( + createInvocationContext(createTestAgent(createTestLlm(LlmResponse.builder().build())))) + .build(); + } + @Test public void declaration_returnsCorrectFunctionDeclaration() { Schema outputSchema = @@ -50,7 +59,7 @@ public void declaration_returnsCorrectFunctionDeclaration() { } @Test - public void runAsync_returnsArgs() { + public void runAsync_returnsArgsAndRecordsValidatedResponse() { Schema outputSchema = Schema.builder() .type("OBJECT") @@ -58,15 +67,37 @@ public void runAsync_returnsArgs() { .build(); SetModelResponseTool tool = new SetModelResponseTool(outputSchema); + ToolContext toolContext = createToolContext(); Map args = ImmutableMap.of("field1", "value1"); - Map result = tool.runAsync(args, null).blockingGet(); + Map result = tool.runAsync(args, toolContext).blockingGet(); assertThat(result).isEqualTo(args); + assertThat(toolContext.actions().setModelResponse()).hasValue(args); + } + + @Test + public void runAsync_invalidArgs_returnsValidationFeedback() { + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("field1", Schema.builder().type("STRING").build())) + .required(ImmutableList.of("field1")) + .build(); + + SetModelResponseTool tool = new SetModelResponseTool(outputSchema); + ToolContext toolContext = createToolContext(); + Map invalidArgs = ImmutableMap.of(); + + Map result = tool.runAsync(invalidArgs, toolContext).blockingGet(); + + assertThat(result).containsKey("error"); + assertThat((String) result.get("error")).contains("field1"); + assertThat(toolContext.actions().setModelResponse()).isEmpty(); } @Test - public void runAsync_validatesArgs() { + public void runAsync_unknownArg_feedbackOmitsSchemaDump() { Schema outputSchema = Schema.builder() .type("OBJECT") @@ -77,12 +108,12 @@ public void runAsync_validatesArgs() { SetModelResponseTool tool = new SetModelResponseTool(outputSchema); Map invalidArgs = ImmutableMap.of("field2", "value2"); - // Should throw validation error - IllegalArgumentException exception = - assertThrows( - IllegalArgumentException.class, () -> tool.runAsync(invalidArgs, null).blockingGet()); + Map result = tool.runAsync(invalidArgs, createToolContext()).blockingGet(); - assertThat(exception).hasMessageThat().contains("does not match agent output schema"); + String error = (String) result.get("error"); + assertThat(error).contains("field2"); + assertThat(error).contains("does not match agent output schema"); + assertThat(error).doesNotContain(outputSchema.toString()); } @Test @@ -108,16 +139,18 @@ public void runAsync_validatesComplexArgs() { .build(); SetModelResponseTool tool = new SetModelResponseTool(complexSchema); + ToolContext toolContext = createToolContext(); Map complexArgs = ImmutableMap.of( "id", 123, "tags", ImmutableList.of("tag1", "tag2"), "metadata", ImmutableMap.of("key", "value")); - Map result = tool.runAsync(complexArgs, null).blockingGet(); + Map result = tool.runAsync(complexArgs, toolContext).blockingGet(); assertThat(result).containsEntry("id", 123); assertThat(result).containsEntry("tags", ImmutableList.of("tag1", "tag2")); assertThat(result).containsEntry("metadata", ImmutableMap.of("key", "value")); + assertThat(toolContext.actions().setModelResponse()).hasValue(complexArgs); } }