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..937ef61fa 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -20,6 +20,7 @@ import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.common.primitives.Ints; import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.Client; import com.google.genai.ResponseStream; @@ -46,6 +47,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 +443,28 @@ 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, as in ADK Python: a list index such as {@code [0]}, a quoted key such as + * {@code ['a.b']} or {@code ["a.b"]}, or a key with an optional dot. Keys may hold escapes. The + * key loops are possessive: a greedy one recurses per character and overflows the stack. + */ + private static final Pattern JSON_PATH_SEGMENT = + Pattern.compile( + "\\[\\s*(\\d+)\\s*]" + + "|\\['((?:[^'\\\\]|\\\\.)*+)'\\]" + + "|\\[\"((?:[^\"\\\\]|\\\\.)*+)\"\\]" + + "|\\.?((?:[^.\\[\\]\\\\]|\\\\.)++)"); + + /** A backslash escape in a JSONPath key: {@code \\uXXXX} or one escaped character. */ + private static final Pattern JSON_PATH_ESCAPE = + Pattern.compile("\\\\u([0-9a-fA-F]{4})|\\\\(.)"); + + /** 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 +699,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 +707,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 +745,92 @@ 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) { + String path = + jsonPath.startsWith("$.") + ? jsonPath.substring(2) + : jsonPath.startsWith("$") ? jsonPath.substring(1) : jsonPath; + Matcher matcher = JSON_PATH_SEGMENT.matcher(path); + List segments = new ArrayList<>(); + while (matcher.find()) { + if (matcher.group(1) != null) { + Integer index = Ints.tryParse(matcher.group(1)); + // An index too large for an int is above MAX_ARRAY_INDEX anyway. + segments.add(index == null ? Integer.MAX_VALUE : index); + } else { + String key = + matcher.group(2) != null + ? matcher.group(2) + : matcher.group(3) != null ? matcher.group(3) : matcher.group(4); + segments.add(unescapeJsonPathKey(key)); + } + } + return segments; + } + + /** Resolves backslash escapes in a JSONPath key the way ADK Python does. */ + private static String unescapeJsonPathKey(String key) { + Matcher matcher = JSON_PATH_ESCAPE.matcher(key); + StringBuilder result = new StringBuilder(); + while (matcher.find()) { + String replacement; + if (matcher.group(1) != null) { + replacement = String.valueOf((char) Integer.parseInt(matcher.group(1), 16)); + } else { + replacement = + switch (matcher.group(2)) { + case "n" -> "\n"; + case "r" -> "\r"; + case "t" -> "\t"; + case "b" -> "\b"; + case "f" -> "\f"; + default -> matcher.group(2); + }; + } + matcher.appendReplacement(result, Matcher.quoteReplacement(replacement)); + } + matcher.appendTail(result); + return result.toString(); } /** 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 af5368859..19857d22b 100644 --- a/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java +++ b/core/src/main/java/com/google/adk/tools/SetModelResponseTool.java @@ -21,6 +21,7 @@ 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; @@ -55,11 +56,16 @@ public SetModelResponseTool(Schema outputSchema) { @Override public Optional declaration() { + // Function parameters must be an object, so an untyped object schema gets that type. + Schema parameters = + outputSchema.type().isPresent() + ? outputSchema + : outputSchema.toBuilder().type(Type.Known.OBJECT).build(); return Optional.of( FunctionDeclaration.builder() .name(name()) .description(description()) - .parameters(outputSchema) + .parameters(parameters) .build()); } 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 c08946d22..78014f07f 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 @@ -41,11 +41,14 @@ import com.google.adk.tools.ToolContext; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; +import com.google.errorprone.annotations.CanIgnoreReturnValue; import com.google.genai.types.Content; +import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.FunctionResponse; import com.google.genai.types.GenerateContentConfig; import com.google.genai.types.Part; import com.google.genai.types.Schema; +import com.google.genai.types.Type; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Single; import java.time.Instant; @@ -71,6 +74,9 @@ public final class OutputSchemaTest { .required(ImmutableList.of("field1")) .build(); + /** A Gemini 2 model name, which answers through set_model_response. */ + private static final String GEMINI_2_MODEL = "gemini-2.5-flash"; + private OutputSchema outputSchemaProcessor; private TestLlm testLlm; private LlmRequest initialRequest; @@ -87,6 +93,12 @@ public TestTool() { super("test_tool", "test description"); } + @Override + public Optional declaration() { + return Optional.of( + FunctionDeclaration.builder().name(name()).description(description()).build()); + } + @Override public Single> runAsync(Map args, ToolContext toolContext) { return Single.just(ImmutableMap.of()); @@ -563,6 +575,25 @@ public void processRequest_modelWithNullCapabilities_followsTheRequestName() { assertThat(otherRequest.tools()).doesNotContainKey(SetModelResponseTool.NAME); } + @Test + public void runner_untypedSchemaWithProperties_getsSetModelResponseWithObjectParameters() { + Schema untypedSchema = + Schema.builder() + .properties(ImmutableMap.of("field1", Schema.builder().type("STRING").build())) + .build(); + + LlmRequest request = firstRequest(GEMINI_2_MODEL, untypedSchema, new TestTool()); + + assertUsesSetModelResponse(request); + assertThat( + request + .tools() + .get(SetModelResponseTool.NAME) + .declaration() + .flatMap(FunctionDeclaration::parameters)) + .hasValue(untypedSchema.toBuilder().type(Type.Known.OBJECT).build()); + } + /** * Forwards generateContent and connect to {@code delegate} under the model name {@code name}. It * does not forward capabilities(), which follows {@code name}. @@ -667,4 +698,42 @@ public BaseLlmConnection connect(LlmRequest llmRequest) { assertThat(finalEvent.content().get().parts().get().get(0).text()).hasValue("{\"extra\":null}"); assertThat(finalEvent.actions().stateDelta()).containsEntry("answer", args); } + + private static LlmAgent.Builder agentBuilder( + BaseLlm model, Schema outputSchema, BaseTool... tools) { + return LlmAgent.builder() + .name("agent") + .model(model) + .outputSchema(outputSchema) + .tools(ImmutableList.copyOf(tools)); + } + + /** Runs {@code agent} once in an {@link InMemoryRunner} and returns its events. */ + @CanIgnoreReturnValue + private static List runAgent(LlmAgent agent) { + InMemoryRunner runner = new InMemoryRunner(agent, "test-app"); + Session session = runner.sessionService().createSession("test-app", "user").blockingGet(); + return runner + .runAsync("user", session.id(), Content.fromParts(Part.fromText("hello"))) + .toList() + .blockingGet(); + } + + /** + * Returns the first request an agent with {@code outputSchema} and {@code tools} sends to a model + * named {@code modelName}. + */ + private static LlmRequest firstRequest(String modelName, Schema outputSchema, BaseTool... tools) { + TestLlm scriptedLlm = createTestLlm(createLlmResponse(Content.fromParts(Part.fromText("{}")))); + runAgent(agentBuilder(new RenamedLlm(modelName, scriptedLlm), outputSchema, tools).build()); + return scriptedLlm.getRequests().get(0); + } + + private static void assertUsesSetModelResponse(LlmRequest request) { + assertThat(request.tools()).containsKey(SetModelResponseTool.NAME); + assertThat(request.tools()).containsKey("test_tool"); + assertThat(String.join("\n", request.getSystemInstructions())) + .contains("you must provide your final response using the set_model_response tool"); + assertThat(request.config().flatMap(GenerateContentConfig::responseSchema)).isEmpty(); + } } 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..f12a3f5a8 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,209 @@ 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_streamedQuotedKeys_parseLikeAdkPython() { + ImmutableList partialArgs = + ImmutableList.of( + PartialArg.builder().jsonPath("$['a.b']").stringValue("dotted").build(), + PartialArg.builder().jsonPath("$[\"x\"]").stringValue("double").build(), + PartialArg.builder().jsonPath("$['it\\'s']").stringValue("escaped").build(), + PartialArg.builder().jsonPath("$.foo['bar.baz']").stringValue("nested").build()); + + Map args = streamedArgs(partialArgs); + + assertThat(args) + .containsExactly( + "a.b", "dotted", + "x", "double", + "it's", "escaped", + "foo", ImmutableMap.of("bar.baz", "nested")); + } + + @Test + public void processRawResponses_streamedLongKeys_parseWithoutOverflowingTheStack() { + String plainKey = "k".repeat(100_000); + String quotedKey = "q.".repeat(50_000); + ImmutableList partialArgs = + ImmutableList.of( + PartialArg.builder().jsonPath("$." + plainKey).stringValue("plain").build(), + PartialArg.builder().jsonPath("$['" + quotedKey + "']").stringValue("quoted").build()); + + Map args = streamedArgs(partialArgs); + + assertThat(args).containsExactly(plainKey, "plain", quotedKey, "quoted"); + } + + @Test + public void processRawResponses_streamedIndicesWithSpacesOrZeros_parseLikeAdkPython() { + ImmutableList partialArgs = + ImmutableList.of( + PartialArg.builder().jsonPath("$.tags[ 0 ]").stringValue("a").build(), + PartialArg.builder().jsonPath("$.tags[000001]").stringValue("b").build()); + + Map args = streamedArgs(partialArgs); + + assertThat(args).containsExactly("tags", ImmutableList.of("a", "b")); + } + + @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[99999999999]").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")), + // A rejected answer would make the model retry; stop instead of looping. + RunConfig.builder() + .streamingMode(RunConfig.StreamingMode.SSE) + .maxLlmCalls(2) + .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