From dbba7556d3c8a00c46fdfd2d0d57bd6cc1bf9943 Mon Sep 17 00:00:00 2001 From: Mateusz Krawiec Date: Tue, 10 Nov 2026 05:14:58 -0800 Subject: [PATCH] fix(langchain4j): accept free-form schemas in tool declarations PiperOrigin-RevId: 1018130280 --- .../adk/models/langchain4j/LangChain4j.java | 80 ++++++---- .../models/langchain4j/LangChain4jTest.java | 125 ++++++++++++++++ .../main/java/com/google/adk/SchemaUtils.java | 56 ++++--- .../adk/tools/SetModelResponseTool.java | 8 +- .../java/com/google/adk/SchemaUtilsTest.java | 138 ++++++++++++++++++ .../com/google/adk/agents/LlmAgentTest.java | 37 +++++ .../adk/flows/llmflows/OutputSchemaTest.java | 48 ++++++ .../adk/tools/SetModelResponseToolTest.java | 19 +++ 8 files changed, 457 insertions(+), 54 deletions(-) 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 fde730dee..f9161ad42 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 @@ -59,6 +59,7 @@ import dev.langchain4j.model.chat.request.ChatRequest; import dev.langchain4j.model.chat.request.ChatRequestParameters; import dev.langchain4j.model.chat.request.ToolChoice; +import dev.langchain4j.model.chat.request.json.JsonAnyOfSchema; import dev.langchain4j.model.chat.request.json.JsonArraySchema; import dev.langchain4j.model.chat.request.json.JsonBooleanSchema; import dev.langchain4j.model.chat.request.json.JsonIntegerSchema; @@ -564,7 +565,11 @@ private List toToolSpecifications(LlmRequest llmRequest) { } private JsonObjectSchema toParameters(Schema schema) { - if (schema.type().isPresent() && Type.Known.OBJECT.equals(schema.type().get().knownEnum())) { + // An untyped schema with properties describes an object. + if (schema + .type() + .map(t -> t.knownEnum() == Type.Known.OBJECT) + .orElse(schema.properties().isPresent())) { return JsonObjectSchema.builder() .addProperties(toProperties(schema)) .required(schema.required().orElse(List.of())) @@ -582,31 +587,56 @@ private Map toProperties(Schema schema) { return result; } - private JsonSchemaElement toJsonSchemaElement(Schema schema) { - if (schema != null && schema.type().isPresent()) { - Type type = schema.type().get(); - return switch (type.knownEnum()) { - case STRING -> - JsonStringSchema.builder().description(schema.description().orElse(null)).build(); - case NUMBER -> - JsonNumberSchema.builder().description(schema.description().orElse(null)).build(); - case INTEGER -> - JsonIntegerSchema.builder().description(schema.description().orElse(null)).build(); - case BOOLEAN -> - JsonBooleanSchema.builder().description(schema.description().orElse(null)).build(); - case ARRAY -> - JsonArraySchema.builder() - .description(schema.description().orElse(null)) - .items(toJsonSchemaElement(schema.items().orElseThrow())) - .build(); - case OBJECT -> toParameters(schema); - default -> - throw new UnsupportedFeatureException( - "LangChain4jLlm does not support schema of type: " + type); - }; - } else { - throw new IllegalArgumentException("Schema type cannot be null or absent"); + private JsonSchemaElement toJsonSchemaElement(@Nullable Schema schema) { + if (schema == null) { + return anyValue(null); + } + // An untyped schema with properties is an object; without them, it accepts any value. + if (schema.type().isEmpty()) { + return schema.properties().isPresent() + ? toParameters(schema) + : anyValue(schema.description().orElse(null)); } + Type type = schema.type().get(); + return switch (type.knownEnum()) { + case STRING -> + JsonStringSchema.builder().description(schema.description().orElse(null)).build(); + case NUMBER -> + JsonNumberSchema.builder().description(schema.description().orElse(null)).build(); + case INTEGER -> + JsonIntegerSchema.builder().description(schema.description().orElse(null)).build(); + case BOOLEAN -> + JsonBooleanSchema.builder().description(schema.description().orElse(null)).build(); + case ARRAY -> + JsonArraySchema.builder() + .description(schema.description().orElse(null)) + .items(schema.items().map(this::toJsonSchemaElement).orElseGet(() -> anyValue(null))) + .build(); + case OBJECT -> toParameters(schema); + default -> + throw new UnsupportedFeatureException( + "LangChain4jLlm does not support schema of type: " + type); + }; + } + + /** + * Returns an element that accepts a string, number, boolean, object, or an array of those, for a + * schema without a type. Not a {@code JsonRawSchema}, which the Gemini provider rejects. The + * array branch has items, which Gemini requires. + */ + private static JsonSchemaElement anyValue(@Nullable String description) { + List nonArrays = + List.of( + new JsonStringSchema(), + new JsonNumberSchema(), + new JsonBooleanSchema(), + JsonObjectSchema.builder().build()); + List anyOf = new ArrayList<>(nonArrays); + anyOf.add( + JsonArraySchema.builder() + .items(JsonAnyOfSchema.builder().anyOf(nonArrays).build()) + .build()); + return JsonAnyOfSchema.builder().description(description).anyOf(anyOf).build(); } private LlmResponse toLlmResponse(ChatResponse chatResponse, ChatRequest chatRequest) { 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 7ae5fc061..77b565644 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 @@ -15,11 +15,13 @@ */ package com.google.adk.models.langchain4j; +import static java.nio.charset.StandardCharsets.UTF_8; import static org.assertj.core.api.Assertions.*; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.*; +import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.adk.agents.LlmAgent; import com.google.adk.agents.RunConfig; @@ -28,7 +30,9 @@ import com.google.adk.runner.InMemoryRunner; import com.google.adk.sessions.Session; import com.google.adk.tools.FunctionTool; +import com.google.adk.tools.SetModelResponseTool; import com.google.genai.types.*; +import com.sun.net.httpserver.HttpServer; import dev.langchain4j.agent.tool.ToolExecutionRequest; import dev.langchain4j.data.message.AiMessage; import dev.langchain4j.data.message.UserMessage; @@ -37,17 +41,25 @@ import dev.langchain4j.model.chat.StreamingChatModel; import dev.langchain4j.model.chat.request.ChatRequest; import dev.langchain4j.model.chat.request.ChatRequestParameters; +import dev.langchain4j.model.chat.request.json.JsonAnyOfSchema; +import dev.langchain4j.model.chat.request.json.JsonArraySchema; import dev.langchain4j.model.chat.request.json.JsonObjectSchema; +import dev.langchain4j.model.chat.request.json.JsonSchemaElement; import dev.langchain4j.model.chat.request.json.JsonStringSchema; import dev.langchain4j.model.chat.response.ChatResponse; import dev.langchain4j.model.chat.response.StreamingChatResponseHandler; +import dev.langchain4j.model.googleai.GoogleAiGeminiChatModel; import dev.langchain4j.model.output.TokenUsage; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.io.OutputStream; +import java.net.InetAddress; +import java.net.InetSocketAddress; import java.util.ArrayList; import java.util.List; import java.util.Map; import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.DisplayName; import org.junit.jupiter.api.Test; @@ -57,6 +69,25 @@ class LangChain4jTest { private static final String MODEL_NAME = "test-model"; + /** + * An output schema with an untyped property, an array without items, and an untyped object with + * properties. + */ + private static final Schema FREE_FORM_SCHEMA = + Schema.builder() + .type("OBJECT") + .properties( + Map.of( + "extra", + Schema.builder().description("Anything").build(), + "tags", + Schema.builder().type("ARRAY").build(), + "nested", + Schema.builder() + .properties(Map.of("a", Schema.builder().type("STRING").build())) + .build())) + .build(); + private ChatModel chatModel; private StreamingChatModel streamingChatModel; private LangChain4j langChain4j; @@ -977,6 +1008,100 @@ void testGenerateContentWithMcpToolParametersJsonSchemaAsSchema() { assertThat(capturedRequest.toolSpecifications().get(0).description()).isEqualTo("An MCP tool"); } + @Test + @DisplayName("Should convert free-form schemas in tool parameters") + void testGenerateContentWithFreeFormToolParameters() { + // Given + final ChatResponse chatResponse = mock(ChatResponse.class); + when(chatResponse.aiMessage()).thenReturn(AiMessage.from("Done")); + when(chatModel.chat(any(ChatRequest.class))).thenReturn(chatResponse); + + // When + final LlmResponse response = + langChain4j.generateContent(freeFormSetModelResponseRequest(), false).blockingFirst(); + + // Then + assertThat(response.content().get().text()).isEqualTo("Done"); + final ArgumentCaptor requestCaptor = ArgumentCaptor.forClass(ChatRequest.class); + verify(chatModel).chat(requestCaptor.capture()); + final Map properties = + requestCaptor.getValue().toolSpecifications().get(0).parameters().properties(); + final JsonAnyOfSchema extra = (JsonAnyOfSchema) properties.get("extra"); + assertThat(extra.description()).isEqualTo("Anything"); + assertThat(extra.anyOf()).hasSize(5); + assertThat(((JsonArraySchema) properties.get("tags")).items()) + .isInstanceOf(JsonAnyOfSchema.class); + assertThat(((JsonObjectSchema) properties.get("nested")).properties().get("a")) + .isInstanceOf(JsonStringSchema.class); + } + + @Test + @DisplayName("Should send free-form tool parameters through the Gemini provider") + void testGenerateContentWithFreeFormToolParameters_geminiProvider() throws Exception { + // Given + final AtomicReference requestBody = new AtomicReference<>(); + final HttpServer server = + HttpServer.create(new InetSocketAddress(InetAddress.getLoopbackAddress(), 0), 0); + server.createContext( + "/", + exchange -> { + requestBody.set(new String(exchange.getRequestBody().readAllBytes(), UTF_8)); + final byte[] reply = + ("{\"candidates\":[{\"content\":{\"role\":\"model\",\"parts\":[{\"text\":\"Done\"}]}," + + "\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":1," + + "\"candidatesTokenCount\":1,\"totalTokenCount\":2}}") + .getBytes(UTF_8); + exchange.getResponseHeaders().add("Content-Type", "application/json"); + exchange.sendResponseHeaders(200, reply.length); + try (OutputStream body = exchange.getResponseBody()) { + body.write(reply); + } + }); + + // When + final LlmResponse response; + server.start(); + try { + final LangChain4j gemini = + LangChain4j.builder() + .chatModel( + GoogleAiGeminiChatModel.builder() + .apiKey("test-key") + .modelName("gemini-2.5-flash") + .baseUrl("http://127.0.0.1:" + server.getAddress().getPort()) + .maxRetries(0) + .build()) + .build(); + response = gemini.generateContent(freeFormSetModelResponseRequest(), false).blockingFirst(); + } finally { + server.stop(0); + } + + // Then + assertThat(response.content().get().text()).isEqualTo("Done"); + // LangChain4j 1.12 sends "tools" as one object, later versions as a list of them. + final JsonNode tools = new ObjectMapper().readTree(requestBody.get()).path("tools"); + final JsonNode properties = + (tools.isArray() ? tools.path(0) : tools) + .at("/functionDeclarations/0/parameters/properties"); + assertThat(properties.isObject()).isTrue(); + assertThat(properties.at("/extra/anyOf").size()).isEqualTo(5); + assertThat(properties.at("/nested/properties/a/type").asText()).isEqualToIgnoringCase("string"); + // Gemini rejects an array schema without items. + assertThat( + properties.findParents("type").stream() + .filter(node -> node.get("type").asText().equalsIgnoreCase("array"))) + .isNotEmpty() + .allMatch(node -> node.has("items")); + } + + private static LlmRequest freeFormSetModelResponseRequest() { + return LlmRequest.builder() + .contents(List.of(Content.fromParts(Part.fromText("Answer")))) + .tools(Map.of(SetModelResponseTool.NAME, new SetModelResponseTool(FREE_FORM_SCHEMA))) + .build(); + } + @Test @DisplayName( "Should use TokenCountEstimator to estimate token usage when TokenUsage is not available") diff --git a/core/src/main/java/com/google/adk/SchemaUtils.java b/core/src/main/java/com/google/adk/SchemaUtils.java index 0df7ee775..bcdad7d7f 100644 --- a/core/src/main/java/com/google/adk/SchemaUtils.java +++ b/core/src/main/java/com/google/adk/SchemaUtils.java @@ -41,11 +41,17 @@ private SchemaUtils() {} // Private constructor for utility class */ @SuppressWarnings("unchecked") // For tool parameter type casting. private static Boolean matchType(Object value, Schema schema, Boolean isInput) { + // A schema with no type and no properties accepts any value, and null unless nullable=false. + boolean freeForm = schema.type().isEmpty() && schema.properties().isEmpty(); if (value == null) { - return schema.nullable().orElse(false); + return schema.nullable().orElse(freeForm); + } + if (freeForm) { + return true; } // Based on types from https://cloud.google.com/vertex-ai/docs/reference/rest/v1/Schema - Type.Known type = schema.type().get().knownEnum(); + // An untyped schema with properties describes an object. + Type.Known type = schema.type().map(Type::knownEnum).orElse(Type.Known.OBJECT); switch (type) { case STRING: return value instanceof String; @@ -57,6 +63,9 @@ private static Boolean matchType(Object value, Schema schema, Boolean isInput) { return value instanceof Number; case ARRAY: if (value instanceof List) { + if (schema.items().isEmpty()) { + return true; + } for (Object element : (List) value) { if (!matchType(element, schema.items().get(), isInput)) { return false; @@ -91,29 +100,14 @@ private static Boolean matchType(Object value, Schema schema, Boolean isInput) { */ public static void validateMapOnSchema(Map args, Schema schema, Boolean isInput) { Preconditions.checkNotNull(isInput, "IsInput cannot be null"); - Map properties = schema.properties().get(); - for (Entry arg : args.entrySet()) { - // Check if the argument is in the schema. - if (!properties.containsKey(arg.getKey())) { - if (isInput) { - throw new IllegalArgumentException( - "Input arg: " + arg.getKey() + " does not match agent input schema: " + schema); - } else { - throw new IllegalArgumentException( - "Output arg: " + arg.getKey() + " does not match agent output schema: " + schema); - } - } - // Check if the argument type matches the schema type. - if (!matchType(arg.getValue(), properties.get(arg.getKey()), isInput)) { - if (isInput) { - throw new IllegalArgumentException( - "Input arg: " + arg.getKey() + " does not match agent input schema: " + schema); - } else { - throw new IllegalArgumentException( - "Output arg: " + arg.getKey() + " does not match agent output schema: " + schema); - } - } + if (schema.type().map(Type::knownEnum).orElse(Type.Known.OBJECT) != Type.Known.OBJECT) { + throw new IllegalArgumentException( + (isInput ? "Input" : "Output") + " schema does not describe an object: " + schema); } + // A schema without properties leaves the keys open, as for a free-form map. + schema + .properties() + .ifPresent(properties -> validateProperties(args, properties, schema, isInput)); // Check if all required arguments are present. if (schema.required().isPresent()) { for (String required : schema.required().get()) { @@ -128,6 +122,20 @@ public static void validateMapOnSchema(Map args, Schema schema, } } + /** Checks that every key in {@code args} is one of {@code properties} and matches its schema. */ + private static void validateProperties( + Map args, Map properties, Schema schema, boolean isInput) { + for (Entry arg : args.entrySet()) { + if (!properties.containsKey(arg.getKey()) + || !matchType(arg.getValue(), properties.get(arg.getKey()), isInput)) { + throw new IllegalArgumentException( + isInput + ? "Input arg: " + arg.getKey() + " does not match agent input schema: " + schema + : "Output arg: " + arg.getKey() + " does not match agent output schema: " + schema); + } + } + } + /** * Validates an output string against a schema. * 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 9cab1749a..af5368859 100644 --- a/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java +++ b/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java @@ -87,8 +87,8 @@ public Single> runAsync(Map args, ToolContex /** * 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. A value - * without a schema loses all its nulls. + * property that {@code schema} marks required, which validation already accepted. A value without + * a schema loses all its nulls. */ private static Map excludeNullFields( Map values, @Nullable Schema schema) { @@ -102,9 +102,7 @@ private static Map excludeNullFields( Object value = entry.getValue(); if (value != null) { result.put(entry.getKey(), excludeNullsFromValue(value, property)); - } else if (property != null - && property.nullable().orElse(false) - && required.contains(entry.getKey())) { + } else if (required.contains(entry.getKey())) { result.put(entry.getKey(), null); } } diff --git a/core/src/test/java/com/google/adk/SchemaUtilsTest.java b/core/src/test/java/com/google/adk/SchemaUtilsTest.java index ac5df39a8..5bb275bb4 100644 --- a/core/src/test/java/com/google/adk/SchemaUtilsTest.java +++ b/core/src/test/java/com/google/adk/SchemaUtilsTest.java @@ -18,6 +18,7 @@ import static org.junit.Assert.assertThrows; +import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Schema; import java.util.HashMap; @@ -98,4 +99,141 @@ public void validateMapOnSchema_integerField_allowsIntegerAndLong() { args.put("intField", 123); SchemaUtils.validateMapOnSchema(args, schema, /* isInput= */ true); } + + @Test + public void validateMapOnSchema_untypedProperty_acceptsAnyValue() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("anything", Schema.builder().build())) + .build(); + + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("anything", ImmutableMap.of("nested", 1)), schema, /* isInput= */ false); + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("anything", "text"), schema, /* isInput= */ false); + } + + @Test + public void validateMapOnSchema_objectWithoutProperties_acceptsAnyKeys() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("metadata", Schema.builder().type("OBJECT").build())) + .build(); + + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("metadata", ImmutableMap.of("source", "web", "score", 0.5)), + schema, + /* isInput= */ false); + } + + @Test + public void validateMapOnSchema_topLevelObjectWithoutProperties_stillEnforcesRequired() { + Schema schema = Schema.builder().type("OBJECT").required(ImmutableList.of("id")).build(); + + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("id", "1", "extra", true), schema, /* isInput= */ true); + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("extra", true), schema, /* isInput= */ true)); + } + + @Test + public void validateMapOnSchema_arrayWithoutItems_acceptsAnyElements() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("tags", Schema.builder().type("ARRAY").build())) + .build(); + + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("tags", ImmutableList.of("a", 1, ImmutableMap.of())), + schema, + /* isInput= */ false); + } + + @Test + public void validateMapOnSchema_arrayWithoutItems_stillRejectsANonList() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("tags", Schema.builder().type("ARRAY").build())) + .build(); + + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("tags", "a"), schema, /* isInput= */ false)); + } + + @Test + public void validateMapOnSchema_nonObjectSchema_throwsException() { + Schema schema = Schema.builder().type("ARRAY").build(); + + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema(ImmutableMap.of("a", 1), schema, /* isInput= */ false)); + } + + @Test + public void validateMapOnSchema_untypedProperty_acceptsNull() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("anything", Schema.builder().build())) + .build(); + Map args = new HashMap<>(); + args.put("anything", null); + + SchemaUtils.validateMapOnSchema(args, schema, /* isInput= */ false); + } + + @Test + public void validateMapOnSchema_untypedNonNullableProperty_throwsException() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("anything", Schema.builder().nullable(false).build())) + .build(); + Map args = new HashMap<>(); + args.put("anything", null); + + assertThrows( + IllegalArgumentException.class, + () -> SchemaUtils.validateMapOnSchema(args, schema, /* isInput= */ false)); + } + + @Test + public void validateMapOnSchema_untypedNestedSchemaWithProperties_validatesAsObject() { + Schema nested = + Schema.builder() + .properties(ImmutableMap.of("a", Schema.builder().type("STRING").build())) + .required(ImmutableList.of("a")) + .build(); + Schema schema = + Schema.builder().type("OBJECT").properties(ImmutableMap.of("nested", nested)).build(); + + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("nested", ImmutableMap.of("a", "x")), schema, /* isInput= */ false); + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("nested", 5), schema, /* isInput= */ false)); + assertThrows( + IllegalArgumentException.class, + () -> + SchemaUtils.validateMapOnSchema( + ImmutableMap.of("nested", ImmutableMap.of()), schema, /* isInput= */ false)); + Map nullNested = new HashMap<>(); + nullNested.put("nested", null); + assertThrows( + IllegalArgumentException.class, + () -> SchemaUtils.validateMapOnSchema(nullNested, schema, /* isInput= */ false)); + } } diff --git a/core/src/test/java/com/google/adk/agents/LlmAgentTest.java b/core/src/test/java/com/google/adk/agents/LlmAgentTest.java index c32eafc4d..2f10651ba 100644 --- a/core/src/test/java/com/google/adk/agents/LlmAgentTest.java +++ b/core/src/test/java/com/google/adk/agents/LlmAgentTest.java @@ -71,7 +71,9 @@ import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Maybe; import io.reactivex.rxjava3.core.Single; +import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.atomic.AtomicBoolean; import org.junit.After; @@ -142,6 +144,41 @@ public void testRun_withOutputKey_savesState() { assertThat(events.get(0).actions().stateDelta()).containsEntry("myOutput", "Saved output"); } + @Test + public void testRun_withOutputKeyAndFreeFormOutputSchema_savesParsedState() { + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties( + ImmutableMap.of( + "city", + Schema.builder().type("STRING").build(), + "details", + Schema.builder().type("OBJECT").build(), + "tags", + Schema.builder().type("ARRAY").build(), + "extra", + Schema.builder().build())) + .build(); + TestLlm testLlm = + createTestLlm( + createTextLlmResponse( + "{\"city\":\"Paris\",\"details\":{\"population\":2100000}," + + "\"tags\":[\"a\",1],\"extra\":null}")); + LlmAgent agent = + createTestAgentBuilder(testLlm).outputSchema(outputSchema).outputKey("answer").build(); + InvocationContext invocationContext = createInvocationContext(agent); + + List events = agent.runAsync(invocationContext).toList().blockingGet(); + + Map expected = new HashMap<>(); + expected.put("city", "Paris"); + expected.put("details", ImmutableMap.of("population", 2100000)); + expected.put("tags", ImmutableList.of("a", 1)); + expected.put("extra", null); + assertThat(getOnlyElement(events).actions().stateDelta()).containsEntry("answer", expected); + } + @Test public void testRun_withOutputKey_savesMultiPartState() { Content modelContent = Content.fromParts(Part.fromText("Part 1."), Part.fromText(" Part 2.")); 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 8b9a55263..5e44be92a 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 @@ -613,4 +613,52 @@ private LlmRequest basicThenOutputSchema(BaseLlm model) { .blockingGet() .updatedRequest(); } + + @Test + public void setModelResponseWithARequiredUntypedNull_keepsItInTheOutputKey() { + Schema schema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("extra", Schema.builder().build())) + .required(ImmutableList.of("extra")) + .build(); + Map args = new LinkedHashMap<>(); + args.put("extra", null); + TestLlm scriptedLlm = + createTestLlm( + createLlmResponse( + Content.fromParts(Part.fromFunctionCall(SetModelResponseTool.NAME, args)))); + 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(schema) + .tools(ImmutableList.of(new TestTool())) + .outputKey("answer") + .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(); + + Event finalEvent = events.get(events.size() - 1); + assertThat(finalEvent.content().get().parts().get().get(0).text()).hasValue("{\"extra\":null}"); + assertThat(finalEvent.actions().stateDelta()).containsEntry("answer", args); + } } 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 959c1c6dc..9335fe475 100644 --- a/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java +++ b/core/src/test/java/com/google/adk/tools/SetModelResponseToolTest.java @@ -151,6 +151,25 @@ public void runAsync_requiredNullableNull_keepsTheNull() { SchemaUtils.validateMapOnSchema(result, outputSchema, /* isInput= */ false); } + @Test + public void runAsync_requiredUntypedNull_keepsTheNull() { + Schema outputSchema = + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("extra", Schema.builder().build())) + .required(ImmutableList.of("extra")) + .build(); + Map args = new LinkedHashMap<>(); + args.put("extra", null); + + ToolContext toolContext = createToolContext(); + Map result = + new SetModelResponseTool(outputSchema).runAsync(args, toolContext).blockingGet(); + + assertThat(result).isEqualTo(args); + assertThat(toolContext.actions().setModelResponse()).hasValue(args); + } + @Test public void runAsync_requiredAndOptionalNullableNulls_keepsOnlyTheRequiredOne() { Schema outputSchema =