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 36267551c..fe4ce99a5 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -321,6 +321,9 @@ private static final class StreamingResponseAggregator { private final StringBuilder currentTextBuffer = new StringBuilder(); // Always reassigned in accumulateParts() before it is read; the initializer is never observed. private boolean currentTextIsThought = false; + // Signature of the buffered text run, kept apart from the streamed function call's slot below + // so an interleaved chunk cannot flush one part carrying the other's signature. + private byte[] currentTextThoughtSignature = null; private byte[] currentThoughtSignature = null; private GenerateContentResponse lastRawResponse = null; @@ -421,18 +424,17 @@ private boolean accumulateParts(List parts) { String text = part.text().orElse(""); if (!text.isEmpty()) { hasContent = true; - // The signature belongs to this text; capture it so flushTextBufferToSequence attaches - // it. - part.thoughtSignature().ifPresent(sig -> currentThoughtSignature = sig); boolean isThought = part.thought().orElse(false); - // Immediately flush the active text buffer to preserve the exact interleaved blocks of - // text/thoughts. + // Flush before capturing this chunk's signature below, or the signature of the run + // starting here lands on the run being flushed. if (!currentTextBuffer.isEmpty() && isThought != currentTextIsThought) { flushTextBufferToSequence(); } if (currentTextBuffer.isEmpty()) { currentTextIsThought = isThought; } + // The signature rides on the merged part that flushTextBufferToSequence builds. + part.thoughtSignature().ifPresent(sig -> currentTextThoughtSignature = sig); currentTextBuffer.append(text); } else if (part.functionCall().isPresent()) { hasContent = true; @@ -443,16 +445,16 @@ private boolean accumulateParts(List parts) { // future part types) rather than an allowlist that silently drops unlisted types. Flush // buffered text first so parts keep their order, then append the part verbatim keeping // any - // thoughtSignature it carries. The signature is intentionally not captured into - // currentThoughtSignature, which would leak it onto the preceding part. + // thoughtSignature it carries. The signature is intentionally not captured, which would + // leak it onto the preceding part. hasContent = true; flushTextBufferToSequence(); accumulatedSequence.add(part); } else { // Standalone thought/thought-signature part with no renderable content: not emitted on - // its - // own; capture its signature to re-attach to the last real part in processFinalResponse. - part.thoughtSignature().ifPresent(sig -> currentThoughtSignature = sig); + // its own; its signature rides on the text run it sits in, and overrides what that run's + // own chunks carried, because a signature-only part is an explicit carrier. + part.thoughtSignature().ifPresent(sig -> currentTextThoughtSignature = sig); } } return hasContent; @@ -605,9 +607,9 @@ private void flushTextBufferToSequence() { if (!currentTextBuffer.isEmpty()) { Part.Builder partBuilder = Part.builder().text(currentTextBuffer.toString()).thought(currentTextIsThought); - if (currentThoughtSignature != null) { - partBuilder.thoughtSignature(currentThoughtSignature); - currentThoughtSignature = null; + if (currentTextThoughtSignature != null) { + partBuilder.thoughtSignature(currentTextThoughtSignature); + currentTextThoughtSignature = null; } accumulatedSequence.add(partBuilder.build()); currentTextBuffer.setLength(0); 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 a965e5b68..80efe3a24 100644 --- a/core/src/test/java/com/google/adk/models/GeminiTest.java +++ b/core/src/test/java/com/google/adk/models/GeminiTest.java @@ -15,6 +15,7 @@ */ package com.google.adk.models; +import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.truth.Truth.assertThat; import static java.nio.charset.StandardCharsets.UTF_8; @@ -35,6 +36,7 @@ import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.functions.Predicate; import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.util.Optional; import org.junit.Test; import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @@ -1068,6 +1070,170 @@ public void processRawResponses_emptyPartsThenSignature_doesNotThrowException() isFinalThoughtResponseWithUsageMetadataAndSignature("", metadata, "sig")); } + // Consecutive text chunks are merged into a single part the aggregator builds from scratch, so a + // thought signature the chunks carried is lost unless it is copied across. The model expects its + // signature back verbatim; without it, it redoes the reasoning the signature stood for. Mirrors + // ADK Python's TestStreamingThoughtSignature. + @Test + public void processRawResponses_signatureOnMergedText_isPreserved() { + GenerateContentResponse chunk1 = toResponseWithTextAndSignature("At minute 5 ", "text-sig"); + GenerateContentResponse chunk2 = + toResponseWithText("the presenter speaks.", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(1); + assertThat(parts.get(0).text()).hasValue("At minute 5 the presenter speaks."); + assertThat(parts.get(0).thoughtSignature()).hasValue("text-sig".getBytes(UTF_8)); + } + + // The signature can land on any chunk of the run, not just the first. + @Test + public void processRawResponses_signatureOnLaterTextChunk_isPreserved() { + GenerateContentResponse chunk1 = toResponseWithText("At minute 5 "); + GenerateContentResponse chunk2 = toResponseWithTextAndSignature("the presenter ", "late-sig"); + GenerateContentResponse chunk3 = toResponseWithText("speaks.", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(1); + assertThat(parts.get(0).text()).hasValue("At minute 5 the presenter speaks."); + assertThat(parts.get(0).thoughtSignature()).hasValue("late-sig".getBytes(UTF_8)); + } + + // A merged part carries one signature. This class keeps the last of the run rather than the + // first, because the part it flushes can still be rewritten by the final chunk's signature. + @Test + public void processRawResponses_multipleSignaturesInOneRun_keepsTheLast() { + GenerateContentResponse chunk1 = toResponseWithTextAndSignature("At minute 5 ", "first-sig"); + GenerateContentResponse chunk2 = toResponseWithTextAndSignature("the presenter ", "second-sig"); + GenerateContentResponse chunk3 = + toResponse( + Candidate.builder() + .content( + Content.builder() + .parts(Part.fromFunctionCall("done", ImmutableMap.of())) + .build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).thoughtSignature()).hasValue("second-sig".getBytes(UTF_8)); + } + + // A thought run and an answer run flush separately and must not swap signatures: the answer's + // signature arrives on the chunk that triggers the flush of the thought. + @Test + public void processRawResponses_thoughtAndAnswerRuns_keepTheirOwnSignatures() { + GenerateContentResponse chunk1 = + toResponse( + Part.builder() + .text("Let me check.") + .thought(true) + .thoughtSignature("thought-sig".getBytes(UTF_8)) + .build()); + GenerateContentResponse chunk2 = + toResponseWithTextAndSignature("It is a dog.", "answer-sig", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).thought()).hasValue(true); + assertThat(parts.get(0).thoughtSignature()).hasValue("thought-sig".getBytes(UTF_8)); + assertThat(parts.get(1).thoughtSignature()).hasValue("answer-sig".getBytes(UTF_8)); + } + + // A signature-only thought part is not emitted on its own, so its signature has to ride on the + // text run it sits in. It is dropped if it lands in the function call's slot instead. + @Test + public void processRawResponses_standaloneSignatureMidTextRun_ridesOnTheMergedText() { + GenerateContentResponse chunk1 = toResponseWithText("At minute 5 "); + GenerateContentResponse chunk2 = + toResponse( + Part.builder().thought(true).thoughtSignature("carried-sig".getBytes(UTF_8)).build()); + GenerateContentResponse chunk3 = + toResponseWithText("the presenter speaks.", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(1); + assertThat(parts.get(0).text()).hasValue("At minute 5 the presenter speaks."); + assertThat(parts.get(0).thoughtSignature()).hasValue("carried-sig".getBytes(UTF_8)); + } + + // A text chunk arriving mid-stream of a function call must not take the call's signature with it: + // the two runs flush together and each keeps its own. + @Test + public void processRawResponses_textInterleavedWithStreamedCall_keepsBothSignatures() { + GenerateContentResponse chunk1 = + toResponse( + Part.builder() + .functionCall( + FunctionCall.builder() + .name("search") + .partialArgs( + PartialArg.builder().jsonPath("$.q").stringValue("hel").build()) + .willContinue(true) + .build()) + .thoughtSignature("fc-sig".getBytes(UTF_8)) + .build()); + GenerateContentResponse chunk2 = toResponseWithTextAndSignature("Working on it.", "text-sig"); + GenerateContentResponse chunk3 = + toResponse( + Candidate.builder() + .content( + Content.builder() + .parts( + functionCallPart( + FunctionCall.builder() + .partialArgs( + PartialArg.builder() + .jsonPath("$.q") + .stringValue("lo") + .build()) + .willContinue(false) + .build())) + .build()) + .finishReason(new FinishReason(FinishReason.Known.STOP)) + .build()); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList parts = ImmutableList.copyOf(finalResponse.content().get().parts().get()); + assertThat(parts).hasSize(2); + assertThat(parts.get(0).text()).hasValue("Working on it."); + assertThat(parts.get(0).thoughtSignature()).hasValue("text-sig".getBytes(UTF_8)); + assertThat(parts.get(1).functionCall().get().name()).hasValue("search"); + assertThat(parts.get(1).thoughtSignature()).hasValue("fc-sig".getBytes(UTF_8)); + } + + // Server-side media tools return signatures on parts holding nothing else. Such a part must + // survive as its own part rather than being folded into the surrounding text. + @Test + public void processRawResponses_contentFreeSignaturePart_isKept() { + GenerateContentResponse chunk1 = toResponseWithText("At minute 5 the presenter speaks."); + GenerateContentResponse chunk2 = + toResponse(Part.builder().thoughtSignature("call-context".getBytes(UTF_8)).build()); + GenerateContentResponse chunk3 = toResponseWithText("", FinishReason.Known.STOP); + + LlmResponse finalResponse = aggregateFinalResponse(chunk1, chunk2, chunk3); + + ImmutableList signatures = + finalResponse.content().get().parts().get().stream() + .map(Part::thoughtSignature) + .flatMap(Optional::stream) + .map(signature -> new String(signature, UTF_8)) + .collect(toImmutableList()); + assertThat(signatures).containsExactly("call-context"); + } + @Test public void functionCallThenEmptyTextWithStop_emitsPartialThenFinalAggregatedFunctionCall() { Flowable rawResponses = @@ -1477,6 +1643,28 @@ private GenerateContentResponse toResponseWithText( .build(); } + private GenerateContentResponse toResponseWithTextAndSignature(String text, String signature) { + return toResponse( + Part.builder().text(text).thoughtSignature(signature.getBytes(UTF_8)).build()); + } + + private GenerateContentResponse toResponseWithTextAndSignature( + String text, String signature, FinishReason.Known finishReason) { + Part part = Part.builder().text(text).thoughtSignature(signature.getBytes(UTF_8)).build(); + return toResponse( + Candidate.builder() + .content(Content.builder().parts(part).build()) + .finishReason(new FinishReason(finishReason)) + .build()); + } + + /** Runs the chunks through the aggregator and returns the final (non-partial) response. */ + private static LlmResponse aggregateFinalResponse(GenerateContentResponse... chunks) { + return Iterables.getLast( + ImmutableList.copyOf( + Gemini.processRawResponses(Flowable.fromArray(chunks)).blockingIterable())); + } + private static Part functionCallPart(FunctionCall functionCall) { return Part.builder().functionCall(functionCall).build(); }