From a5bda238aed2f3ccd668999e4026be0bf5ffeb54 Mon Sep 17 00:00:00 2001 From: Mateusz Krawiec Date: Tue, 10 Nov 2026 09:14:37 -0800 Subject: [PATCH] fix: build lists from streamed function call argument paths PiperOrigin-RevId: 1018280171 --- .../adk/models/langchain4j/LangChain4j.java | 80 ++++++--- .../models/langchain4j/LangChain4jTest.java | 125 ++++++++++++++ .../main/java/com/google/adk/SchemaUtils.java | 56 +++--- .../java/com/google/adk/models/Gemini.java | 103 +++++++++-- .../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 ++++++ .../com/google/adk/models/GeminiTest.java | 163 +++++++++++++++++- .../adk/tools/SetModelResponseToolTest.java | 19 ++ 10 files changed, 704 insertions(+), 73 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/models/Gemini.java b/core/src/main/java/com/google/adk/models/Gemini.java index cf1291f00..e34b3715f 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -46,6 +46,8 @@ import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.regex.Matcher; +import java.util.regex.Pattern; import okhttp3.OkHttpClient; import org.jspecify.annotations.Nullable; import org.slf4j.Logger; @@ -440,6 +442,16 @@ protected CompletableFuture connectLiveTransport( } private static final class StreamingResponseAggregator { + /** The largest list index a streamed argument path may set, as in ADK Python. */ + private static final int MAX_ARRAY_INDEX = 10_000; + + /** A JSONPath segment: a list index such as {@code [0]}, or a map key with an optional dot. */ + private static final Pattern JSON_PATH_SEGMENT = + Pattern.compile("\\[(\\d+)]|\\.?([^.\\[\\]]+)"); + + /** A whole number whose magnitude is at most 2^53 - 1 has an exact double; others may not. */ + private static final double MAX_EXACT_DOUBLE_INTEGER = 9007199254740991d; + private final List accumulatedSequence = new ArrayList<>(); private final StringBuilder currentTextBuffer = new StringBuilder(); // Always reassigned in accumulateParts() before it is read; the initializer is never observed. @@ -674,7 +686,7 @@ private void applyPartialArg(PartialArg partialArg, String jsonPath) { String chunk = partialArg.stringValue().get(); setValueByJsonPath(jsonPath, existing instanceof String s ? s + chunk : chunk); } else if (partialArg.numberValue().isPresent()) { - setValueByJsonPath(jsonPath, partialArg.numberValue().get()); + setValueByJsonPath(jsonPath, wholeNumberAsInteger(partialArg.numberValue().get())); } else if (partialArg.boolValue().isPresent()) { setValueByJsonPath(jsonPath, partialArg.boolValue().get()); } else if (partialArg.nullValue().isPresent()) { @@ -682,14 +694,36 @@ private void applyPartialArg(PartialArg partialArg, String jsonPath) { } } + /** + * Returns a whole number whose magnitude is at most 2^53 - 1 as the Integer or Long a + * non-streamed call carries, because schema validation rejects a Double for an INTEGER field. + * Any other value stays a Double. + */ + private static Number wholeNumberAsInteger(double value) { + if (value != Math.rint(value) || Math.abs(value) > MAX_EXACT_DOUBLE_INTEGER) { + return value; + } + long whole = (long) value; + if (whole == (int) whole) { + return (int) whole; + } + return whole; + } + /** * Returns the value currently stored at {@code jsonPath} in {@link #currentFcArgs}, or null. */ private @Nullable Object getValueByJsonPath(String jsonPath) { Object current = currentFcArgs; - for (String key : splitJsonPath(jsonPath)) { - if (current instanceof Map map && map.containsKey(key)) { + for (Object segment : parseJsonPath(jsonPath)) { + if (segment instanceof String key + && current instanceof Map map + && map.containsKey(key)) { current = map.get(key); + } else if (segment instanceof Integer index + && current instanceof List list + && index < list.size()) { + current = list.get(index); } else { return null; } @@ -698,27 +732,60 @@ private void applyPartialArg(PartialArg partialArg, String jsonPath) { } /** - * Sets {@code value} at {@code jsonPath} in {@link #currentFcArgs}, creating maps as needed. + * Sets {@code value} at {@code jsonPath} in {@link #currentFcArgs}, creating maps and lists as + * needed and padding lists with nulls, as ADK Python's {@code _set_value_by_json_path} does. A + * path that runs into a value of the wrong type or an index above 10,000 is ignored. */ @SuppressWarnings("unchecked") - private void setValueByJsonPath(String jsonPath, Object value) { - String[] keys = splitJsonPath(jsonPath); - Map current = currentFcArgs; - for (int i = 0; i < keys.length - 1; i++) { - Object next = current.get(keys[i]); - if (!(next instanceof Map)) { - next = new LinkedHashMap<>(); - current.put(keys[i], next); + private void setValueByJsonPath(String jsonPath, @Nullable Object value) { + List segments = parseJsonPath(jsonPath); + Object current = currentFcArgs; + for (int i = 0; i < segments.size(); i++) { + boolean last = i == segments.size() - 1; + Object child = + last + ? value + : segments.get(i + 1) instanceof Integer + ? new ArrayList<>() + : new LinkedHashMap<>(); + if (segments.get(i) instanceof String key && current instanceof Map) { + Map map = (Map) current; + if (last || !map.containsKey(key)) { + map.put(key, child); + } + current = map.get(key); + } else if (segments.get(i) instanceof Integer index + && current instanceof List + && index <= MAX_ARRAY_INDEX) { + List list = (List) current; + while (list.size() <= index) { + list.add(null); + } + if (last || list.get(index) == null) { + list.set(index, child); + } + current = list.get(index); + } else { + return; } - current = (Map) next; } - current.put(keys[keys.length - 1], value); } - /** Splits a JSONPath such as {@code "$.location.city"} into its component keys. */ - private static String[] splitJsonPath(String jsonPath) { - String path = jsonPath.startsWith("$.") ? jsonPath.substring(2) : jsonPath; - return path.split("\\."); + /** Splits a JSONPath such as {@code "$.items[0].name"} into map keys and list indices. */ + private static List parseJsonPath(String jsonPath) { + Matcher matcher = + JSON_PATH_SEGMENT.matcher(jsonPath.startsWith("$") ? jsonPath.substring(1) : jsonPath); + List segments = new ArrayList<>(); + while (matcher.find()) { + String index = matcher.group(1); + if (index == null) { + segments.add(matcher.group(2)); + } else { + // Anything longer than five digits is above MAX_ARRAY_INDEX. + segments.add(index.length() > 5 ? Integer.MAX_VALUE : Integer.parseInt(index)); + } + } + return segments; } /** 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/models/GeminiTest.java b/core/src/test/java/com/google/adk/models/GeminiTest.java index 07cd7d7da..bea5f884e 100644 --- a/core/src/test/java/com/google/adk/models/GeminiTest.java +++ b/core/src/test/java/com/google/adk/models/GeminiTest.java @@ -19,6 +19,12 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static java.util.concurrent.TimeUnit.SECONDS; +import com.google.adk.agents.LlmAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.events.Event; +import com.google.adk.runner.InMemoryRunner; +import com.google.adk.sessions.Session; +import com.google.adk.tools.FunctionTool; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.common.collect.Iterables; @@ -31,13 +37,16 @@ import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.Part; import com.google.genai.types.PartialArg; +import com.google.genai.types.Schema; import com.google.genai.types.ToolCall; import com.google.genai.types.ToolResponse; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.functions.Predicate; import io.reactivex.rxjava3.schedulers.Schedulers; import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.util.Arrays; import java.util.Iterator; +import java.util.Map; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; @@ -329,7 +338,159 @@ public void processRawResponses_streamingFunctionCallArgs_buildsNestedArgs() { finalResponse.content().get().parts().get().get(0).functionCall().get(); assertThat(finalCall.name()).hasValue("book"); assertThat(finalCall.args().get()) - .containsExactly("location", ImmutableMap.of("city", "Paris"), "guests", 2.0); + .containsExactly("location", ImmutableMap.of("city", "Paris"), "guests", 2); + } + + @Test + public void processRawResponses_streamedArrayPaths_buildLists() { + ImmutableList partialArgs = + ImmutableList.of( + PartialArg.builder().jsonPath("$.tags[0]").stringValue("red").build(), + PartialArg.builder().jsonPath("$.tags[1]").stringValue("bl").build(), + PartialArg.builder().jsonPath("$.tags[1]").stringValue("ue").build(), + PartialArg.builder().jsonPath("$.items[0].name").stringValue("pen").build(), + PartialArg.builder().jsonPath("$.items[0].count").numberValue(2.0).build(), + PartialArg.builder().jsonPath("$.matrix[0][1]").numberValue(7.0).build()); + + Map args = streamedArgs(partialArgs); + + assertThat(args) + .containsExactly( + "tags", + ImmutableList.of("red", "blue"), + "items", + ImmutableList.of(ImmutableMap.of("name", "pen", "count", 2)), + "matrix", + ImmutableList.of(Arrays.asList(null, 7))); + } + + @Test + public void processRawResponses_streamedArrayPathsOutOfOrder_padWithNullsLikeAdkPython() { + ImmutableList partialArgs = + ImmutableList.of( + PartialArg.builder().jsonPath("$.tags[2]").stringValue("c").build(), + PartialArg.builder().jsonPath("$.tags[0]").stringValue("a").build(), + PartialArg.builder().jsonPath("$.tags[20000]").stringValue("ignored").build(), + PartialArg.builder().jsonPath("$.tags.key").stringValue("ignored").build()); + + Map args = streamedArgs(partialArgs); + + assertThat(args).containsExactly("tags", Arrays.asList("a", null, "c")); + } + + @Test + public void processRawResponses_streamedNumbers_storesExactWholeNumbersAsIntegers() { + ImmutableMap args = + streamedArgs( + ImmutableMap.of( + "fraction", + 2.5, + "beyondInt", + 3_000_000_000d, + "largestExact", + 9007199254740991d, + "beyondExact", + 9007199254740992d)); + + assertThat(args.get("fraction")).isEqualTo(2.5); + assertThat(args.get("beyondInt")).isEqualTo(3_000_000_000L); + assertThat(args.get("largestExact")).isEqualTo(9007199254740991L); + assertThat(args.get("beyondExact")).isEqualTo(9007199254740992d); + } + + // A streamed whole number must validate against an INTEGER field of the output schema. + @Test + public void runner_streamedSetModelResponseWithAWholeNumber_finishesWithTheAnswer() { + BaseLlm streamingGemini2 = + new BaseLlm("gemini-2.5-flash") { + @Override + public Flowable generateContent(LlmRequest llmRequest, boolean stream) { + return Gemini.processRawResponses( + streamedCall( + "set_model_response", + ImmutableList.of( + PartialArg.builder().jsonPath("$.guests").numberValue(2.0).build()))); + } + + @Override + public BaseLlmConnection connect(LlmRequest llmRequest) { + throw new UnsupportedOperationException(); + } + }; + LlmAgent agent = + LlmAgent.builder() + .name("agent") + .model(streamingGemini2) + .tools(FunctionTool.create(GeminiTest.class, "lookup")) + .outputSchema( + Schema.builder() + .type("OBJECT") + .properties(ImmutableMap.of("guests", Schema.builder().type("INTEGER").build())) + .build()) + .build(); + InMemoryRunner runner = new InMemoryRunner(agent, "test-app"); + Session session = runner.sessionService().createSession("test-app", "user").blockingGet(); + + ImmutableList events = + ImmutableList.copyOf( + runner + .runAsync( + "user", + session.id(), + Content.fromParts(Part.fromText("hello")), + RunConfig.builder().streamingMode(RunConfig.StreamingMode.SSE).build()) + .blockingIterable()); + + assertThat(Iterables.getLast(events).content().get().text()).isEqualTo("{\"guests\":2}"); + } + + /** Streams one call whose arguments are {@code numbers}, one partial arg each. */ + private ImmutableMap streamedArgs(ImmutableMap numbers) { + ImmutableList.Builder partialArgs = ImmutableList.builder(); + for (Map.Entry number : numbers.entrySet()) { + partialArgs.add( + PartialArg.builder() + .jsonPath("$." + number.getKey()) + .numberValue(number.getValue()) + .build()); + } + return ImmutableMap.copyOf(streamedArgs(partialArgs.build())); + } + + /** Streams one call through {@code partialArgs} and returns the arguments it ends with. */ + private Map streamedArgs(ImmutableList partialArgs) { + LlmResponse last = + Iterables.getLast( + Gemini.processRawResponses(streamedCall("f", partialArgs)).blockingIterable()); + return last.content().get().parts().get().get(0).functionCall().get().args().get(); + } + + /** Returns the raw chunks of one call to {@code name} streamed through {@code partialArgs}. */ + private Flowable streamedCall( + String name, ImmutableList partialArgs) { + GenerateContentResponse first = + toResponse( + functionCallPart( + FunctionCall.builder() + .name(name) + .partialArgs(partialArgs) + .willContinue(true) + .build())); + GenerateContentResponse last = + toResponse( + Candidate.builder() + .content( + Content.builder() + .parts(functionCallPart(FunctionCall.builder().willContinue(false).build())) + .build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + return Flowable.just(first, last); + } + + /** A tool for the runner tests; never called, since the fake model answers without it. */ + public static ImmutableMap lookup() { + return ImmutableMap.of(); } // Two streamed function calls back-to-back must not bleed arguments into each other: a completed 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 =