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..3c36cd866 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -440,6 +440,9 @@ protected CompletableFuture connectLiveTransport( } private static final class StreamingResponseAggregator { + /** 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 +677,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,6 +685,22 @@ 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. */ 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..0428ac112 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,6 +37,7 @@ 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; @@ -38,6 +45,7 @@ import io.reactivex.rxjava3.schedulers.Schedulers; import io.reactivex.rxjava3.subscribers.TestSubscriber; 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 +337,122 @@ 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_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()); + } + LlmResponse last = + Iterables.getLast( + Gemini.processRawResponses(streamedCall("f", partialArgs.build())).blockingIterable()); + return ImmutableMap.copyOf( + 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