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..97ffe90f6 100644 --- a/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java +++ b/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java @@ -17,9 +17,14 @@ 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 com.google.genai.types.Type; import io.reactivex.rxjava3.core.Single; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; import java.util.Map; import java.util.Optional; @@ -33,6 +38,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 +67,66 @@ 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); - return args; + 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."); + } + // Match Python's model_dump(exclude_none=True) for Java's map-shaped response. + Map validatedResponse = excludeNullFields(args, outputSchema); + toolContext.actions().setSetModelResponse(validatedResponse); + return validatedResponse; }); } + + private static Map excludeNullFields(Map values, Schema schema) { + Map properties = schema.properties().get(); + Map result = new LinkedHashMap<>(); + for (Map.Entry entry : values.entrySet()) { + Object value = entry.getValue(); + if (value != null) { + result.put(entry.getKey(), excludeNullFields(value, properties.get(entry.getKey()))); + } + } + return result; + } + + @SuppressWarnings("unchecked") + private static Object excludeNullFields(Object value, Schema schema) { + Type.Known type = schema.type().get().knownEnum(); + if (type == Type.Known.OBJECT) { + return excludeNullFields((Map) value, schema); + } + if (type == Type.Known.ARRAY) { + Schema itemSchema = schema.items().get(); + List result = new ArrayList<>(); + for (Object item : (List) value) { + result.add(item == null ? null : excludeNullFields(item, itemSchema)); + } + return result; + } + return value; + } + + 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..273a5caf7 100644 --- a/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java +++ b/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java @@ -16,13 +16,17 @@ 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; import com.google.genai.types.Schema; +import java.util.LinkedHashMap; import java.util.Map; import org.junit.Test; import org.junit.runner.RunWith; @@ -31,6 +35,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 +60,7 @@ public void declaration_returnsCorrectFunctionDeclaration() { } @Test - public void runAsync_returnsArgs() { + public void runAsync_returnsArgsAndRecordsValidatedResponse() { Schema outputSchema = Schema.builder() .type("OBJECT") @@ -58,15 +68,135 @@ 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_nullableNull_excludesNullFromReturnedAndRecordedResponse() { + 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")) + .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"); + assertThat(toolContext.actions().setModelResponse()).hasValue(ImmutableMap.of("a", "x")); + } + + @Test + public void runAsync_nestedNullableNull_excludesNullFromReturnedAndRecordedResponse() { + Schema addressSchema = + 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")) + .build(); + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "name", Schema.builder().type("STRING").build(), "addr", addressSchema)) + .required(ImmutableList.of("name", "addr")) + .build(); + Map address = new LinkedHashMap<>(); + address.put("city", "NYC"); + address.put("zip", null); + + Map args = ImmutableMap.of("name", "x", "addr", address); + + ToolContext toolContext = createToolContext(); + Map result = + new SetModelResponseTool(outputSchema).runAsync(args, toolContext).blockingGet(); + + Map expected = + ImmutableMap.of("name", "x", "addr", ImmutableMap.of("city", "NYC")); + assertThat(result).isEqualTo(expected); + assertThat(toolContext.actions().setModelResponse()).hasValue(expected); + } + + @Test + public void runAsync_arrayOfObjects_excludesNestedNullFromReturnedAndRecordedResponse() { + Schema addressSchema = + 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")) + .build(); + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "addresses", Schema.builder().type("ARRAY").items(addressSchema).build())) + .required(ImmutableList.of("addresses")) + .build(); + Map address = new LinkedHashMap<>(); + address.put("city", "NYC"); + address.put("zip", null); + Map args = ImmutableMap.of("addresses", ImmutableList.of(address)); + + ToolContext toolContext = createToolContext(); + Map result = + new SetModelResponseTool(outputSchema).runAsync(args, toolContext).blockingGet(); + + Map expected = + ImmutableMap.of("addresses", ImmutableList.of(ImmutableMap.of("city", "NYC"))); + assertThat(result).isEqualTo(expected); + assertThat(toolContext.actions().setModelResponse()).hasValue(expected); } @Test - public void runAsync_validatesArgs() { + public void runAsync_unknownArg_feedbackOmitsSchemaDump() { Schema outputSchema = Schema.builder() .type("OBJECT") @@ -77,12 +207,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 +238,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); } }