From 94353776710149cdb2d88922df4f79dccbf92165 Mon Sep 17 00:00:00 2001 From: Jaehee Shin Date: Tue, 13 Oct 2026 15:08:57 +0900 Subject: [PATCH] fix: aggregate live text chunks and skip partial events in runLive GeminiLlmConnection turned each live server message into responses without aggregating text, and marked a response partial only when its message carried turnComplete: false. Runner.runLive passed every event to appendEvent, while runAsync skips partial events. A live turn's text was therefore stored as one model event per chunk, or, with partial chunks, dropped by InMemorySessionService and stored chunk by chunk by FirestoreSessionService and VertexAiSessionService. Follow ADK Python's live connection and runner: - Only responses that stream text are partial. The connection also collects the text and emits it as one non-partial response before the end of the turn, an interruption or a tool call, and when thought and other text alternate or a part that is neither text nor inline data arrives. Text that arrives with the end of the turn or an interruption is emitted only in full. - The full response reports an interruption that ends the text, and no other response is emitted just to report it. - runLive skips partial events when persisting, as runAsync does. Responses without text are no longer partial when their message carries turnComplete: false. Transcriptions are not aggregated yet. --- .../adk/models/GeminiLlmConnection.java | 143 ++++++- .../java/com/google/adk/runner/Runner.java | 6 +- .../adk/models/GeminiLlmConnectionTest.java | 380 ++++++++++++++++-- .../com/google/adk/runner/RunnerTest.java | 37 ++ 4 files changed, 528 insertions(+), 38 deletions(-) diff --git a/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java b/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java index 422fe99a7..782f063e1 100644 --- a/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java +++ b/core/src/main/java/com/google/adk/models/GeminiLlmConnection.java @@ -34,9 +34,11 @@ import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Flowable; import io.reactivex.rxjava3.core.Observable; +import io.reactivex.rxjava3.core.ObservableEmitter; import io.reactivex.rxjava3.disposables.CompositeDisposable; import io.reactivex.rxjava3.processors.PublishProcessor; import java.net.SocketException; +import java.util.ArrayList; import java.util.List; import java.util.Objects; import java.util.Optional; @@ -63,6 +65,8 @@ public final class GeminiLlmConnection implements BaseLlmConnection { private final Flowable responseFlowable = responseProcessor.serialize(); private final CompositeDisposable disposables = new CompositeDisposable(); private final AtomicBoolean closed = new AtomicBoolean(false); + // Used only by handleServerMessage, which genai's WebSocket listener calls one message at a time. + private final PendingText pendingText = new PendingText(); /** * Establishes a new connection over the given live transport. @@ -117,7 +121,7 @@ private void handleServerMessage(LiveServerMessage message) { logger.debug("Received server message: {}", message.toJson()); - Observable llmResponse = convertToServerResponse(message); + Observable llmResponse = convertToServerResponse(message, pendingText); if (!disposables.add( llmResponse.subscribe(responseProcessor::onNext, responseProcessor::onError))) { logger.warn( @@ -125,8 +129,13 @@ private void handleServerMessage(LiveServerMessage message) { } } - /** Converts a server message into the standardized LlmResponse format. */ - static Observable convertToServerResponse(LiveServerMessage message) { + /** + * Converts a server message into the standardized LlmResponse format. + * + * @param pendingText the text streamed by earlier messages and not yet emitted in full + */ + static Observable convertToServerResponse( + LiveServerMessage message, PendingText pendingText) { return Observable.create( emitter -> { // AtomicBoolean is used to modify state from within lambdas, which @@ -136,13 +145,14 @@ static Observable convertToServerResponse(LiveServerMessage message .serverContent() .ifPresent( serverContent -> { - emitter.onNext(createServerContentResponse(serverContent)); + emitServerContentResponses(serverContent, pendingText, emitter); handled.set(true); }); message .toolCall() .ifPresent( toolCall -> { + pendingText.flush().ifPresent(emitter::onNext); emitter.onNext(createToolCallResponse(toolCall)); handled.set(true); }); @@ -178,16 +188,74 @@ static Observable convertToServerResponse(LiveServerMessage message }); } - private static LlmResponse createServerContentResponse(LiveServerContent serverContent) { + /** + * Emits the responses for a server content message. Text parts are streamed in partial responses + * and collected in {@code pendingText}, which is emitted as one non-partial response when other + * content ends it and before the end of the turn or an interruption. Text that arrives with the + * end of the turn or an interruption is emitted only in full. The full text reports an + * interruption, and no response is emitted just to report it again. + */ + private static void emitServerContentResponses( + LiveServerContent serverContent, + PendingText pendingText, + ObservableEmitter emitter) { + List parts = serverContent.modelTurn().flatMap(Content::parts).orElse(ImmutableList.of()); + List streamedParts = new ArrayList<>(); + for (Part part : parts) { + if (pendingText.isEndedBy(part)) { + pendingText.flush().ifPresent(emitter::onNext); + // The full response now carries the text parts of this message so far. + streamedParts.removeIf(GeminiLlmConnection::hasText); + } + if (hasText(part)) { + pendingText.append(part); + } + streamedParts.add(part); + } + boolean turnComplete = serverContent.turnComplete().orElse(false); + boolean interrupted = serverContent.interrupted().orElse(false); + boolean interruptionInFullText = false; + if (turnComplete || interrupted) { + streamedParts.removeIf(GeminiLlmConnection::hasText); + Optional fullText = pendingText.flush(interrupted); + fullText.ifPresent(emitter::onNext); + interruptionInFullText = + interrupted && fullText.isPresent() && !turnComplete && streamedParts.isEmpty(); + } + boolean hasTranscription = + serverContent.inputTranscription().isPresent() + || serverContent.outputTranscription().isPresent(); + if (interruptionInFullText && !hasTranscription) { + return; + } + LlmResponse.Builder builder = LlmResponse.builder(); - serverContent.modelTurn().ifPresent(builder::content); - return builder - .partial(serverContent.turnComplete().map(completed -> !completed).orElse(false)) - .turnComplete(serverContent.turnComplete().orElse(false)) - .interrupted(serverContent.interrupted().orElse(null)) - .inputTranscription(serverContent.inputTranscription().orElse(null)) - .outputTranscription(serverContent.outputTranscription().orElse(null)) - .build(); + if (streamedParts.size() == parts.size()) { + serverContent.modelTurn().ifPresent(builder::content); + } else if (!streamedParts.isEmpty()) { + builder.content(serverContent.modelTurn().get().toBuilder().parts(streamedParts).build()); + } + // Only text is partial: it is emitted again in full. + boolean streamsText = streamedParts.stream().anyMatch(GeminiLlmConnection::hasText); + builder + .partial(streamsText) + .turnComplete(turnComplete) + .interrupted(interruptionInFullText ? null : serverContent.interrupted().orElse(null)); + if (streamsText && hasTranscription) { + // A partial response is not persisted, so the transcriptions get a response of their own. + emitter.onNext(builder.build()); + builder = LlmResponse.builder().partial(false).turnComplete(turnComplete); + } + emitter.onNext( + builder + .inputTranscription(serverContent.inputTranscription().orElse(null)) + .outputTranscription(serverContent.outputTranscription().orElse(null)) + .build()); + } + + /** Returns whether {@code part} has text; an empty string does not count. */ + private static boolean hasText(Part part) { + return part.text().filter(text -> !text.isEmpty()).isPresent(); } private static LlmResponse createToolCallResponse(LiveServerToolCall toolCall) { @@ -356,4 +424,53 @@ private void closeTransportIgnoringErrors(GeminiLiveTransport transport) { }); } } + + /** Model text that has been streamed in partial responses but not yet emitted in full. */ + static final class PendingText { + private final StringBuilder text = new StringBuilder(); + private boolean thought; + + /** + * Returns whether {@code part} ends the pending text: text of the other thought kind, or a part + * that is neither text nor inline data. + */ + private boolean isEndedBy(Part part) { + if (text.length() == 0) { + return false; + } + return hasText(part) ? part.thought().orElse(false) != thought : part.inlineData().isEmpty(); + } + + private void append(Part part) { + text.append(part.text().get()); + thought = part.thought().orElse(false); + } + + /** Returns the pending text as one non-partial response and clears it, or empty if none. */ + private Optional flush() { + return flush(/* interrupted= */ false); + } + + /** + * Returns the pending text as one non-partial response and clears it, or empty if none. + * + * @param interrupted whether an interruption ended the text, which the response then reports + */ + private Optional flush(boolean interrupted) { + if (text.length() == 0) { + return Optional.empty(); + } + Part.Builder part = Part.builder().text(text.toString()); + if (thought) { + part.thought(true); + } + text.setLength(0); + return Optional.of( + LlmResponse.builder() + .content(Content.builder().role("model").parts(part.build()).build()) + .partial(false) + .interrupted(interrupted ? true : null) + .build()); + } + } } diff --git a/core/src/main/java/com/google/adk/runner/Runner.java b/core/src/main/java/com/google/adk/runner/Runner.java index e25ab3a7d..0e8db15df 100644 --- a/core/src/main/java/com/google/adk/runner/Runner.java +++ b/core/src/main/java/com/google/adk/runner/Runner.java @@ -1609,8 +1609,12 @@ protected Flowable runLiveImpl( updatedInvocationContext .agent() .runLive(updatedInvocationContext) + // As in runAsync, partial events are streamed but not persisted. .concatMapSingle( - event -> this.sessionService.appendEvent(session, event))) + event -> + event.partial().orElse(false) + ? Single.just(event) + : this.sessionService.appendEvent(session, event))) .doOnError( throwable -> { Span span = Span.current(); diff --git a/core/src/test/java/com/google/adk/models/GeminiLlmConnectionTest.java b/core/src/test/java/com/google/adk/models/GeminiLlmConnectionTest.java index 587a32c93..d218d6025 100644 --- a/core/src/test/java/com/google/adk/models/GeminiLlmConnectionTest.java +++ b/core/src/test/java/com/google/adk/models/GeminiLlmConnectionTest.java @@ -16,14 +16,18 @@ package com.google.adk.models; +import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.Iterables.getOnlyElement; import static com.google.common.truth.Truth.assertThat; +import com.google.adk.models.GeminiLlmConnection.PendingText; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; +import com.google.genai.types.ExecutableCode; import com.google.genai.types.FunctionCall; import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.Language; import com.google.genai.types.LiveSendClientContentParameters; import com.google.genai.types.LiveSendRealtimeInputParameters; import com.google.genai.types.LiveSendToolResponseParameters; @@ -33,9 +37,11 @@ import com.google.genai.types.LiveServerToolCall; import com.google.genai.types.LiveServerToolCallCancellation; import com.google.genai.types.Part; +import com.google.genai.types.Transcription; import com.google.genai.types.UsageMetadata; import io.reactivex.rxjava3.observers.TestObserver; import io.reactivex.rxjava3.subscribers.TestSubscriber; +import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.CopyOnWriteArrayList; @@ -149,16 +155,15 @@ public void convertToServerResponse_withInterruptedTrue_mapsInterruptedField() { LiveServerMessage message = LiveServerMessage.builder().serverContent(serverContent).build(); TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); + // The text of the interrupted message is emitted in full and reports the interruption. testObserver.assertValueCount(1); testObserver.assertComplete(); LlmResponse response = testObserver.values().get(0); - assertThat(response.content()).isPresent(); assertThat(response.content().get().text()).isEqualTo("Model response"); - assertThat(response.partial()).hasValue(true); - assertThat(response.turnComplete()).hasValue(false); + assertThat(response.partial()).hasValue(false); assertThat(response.interrupted()).hasValue(true); } @@ -175,7 +180,7 @@ public void convertToServerResponse_withInterruptedFalse_mapsInterruptedField() TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); testObserver.assertValueCount(1); testObserver.assertComplete(); @@ -196,11 +201,11 @@ public void convertToServerResponse_withoutInterruptedField_mapsEmptyOptional() TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); - testObserver.assertValueCount(1); + testObserver.assertValueCount(2); testObserver.assertComplete(); - LlmResponse response = testObserver.values().get(0); + LlmResponse response = testObserver.values().get(1); assertThat(response.interrupted()).isEmpty(); assertThat(response.turnComplete()).hasValue(true); } @@ -217,11 +222,16 @@ public void convertToServerResponse_withTurnCompleteTrue_mapsPartialFalse() { TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); - testObserver.assertValueCount(1); + // The text is emitted in full before the response that completes the turn. + testObserver.assertValueCount(2); testObserver.assertComplete(); - LlmResponse response = testObserver.values().get(0); + LlmResponse textResponse = testObserver.values().get(0); + assertThat(textResponse.content().get().text()).isEqualTo("Final response"); + assertThat(textResponse.partial()).hasValue(false); + LlmResponse response = testObserver.values().get(1); + assertThat(response.content()).isEmpty(); assertThat(response.partial()).hasValue(false); assertThat(response.turnComplete()).hasValue(true); } @@ -238,7 +248,7 @@ public void convertToServerResponse_withTurnCompleteFalse_mapsPartialTrue() { TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); testObserver.assertValueCount(1); testObserver.assertComplete(); @@ -257,7 +267,7 @@ public void convertToServerResponse_withToolCall_mapsContentWithFunctionCall() { TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); testObserver.assertValueCount(1); testObserver.assertComplete(); @@ -289,7 +299,7 @@ public void convertToServerResponse_withMultipleFunctionCalls_preservesAllCallsA TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); testObserver.assertValueCount(1); testObserver.assertComplete(); @@ -318,7 +328,7 @@ public void convertToServerResponse_withUsageMetadata_mapsGenerateResponseUsageM TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); testObserver.assertValueCount(1); testObserver.assertComplete(); LlmResponse response = testObserver.values().get(0); @@ -341,7 +351,7 @@ public void convertToServerResponse_withToolCallCancellation_returnsNoValues() { TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); testObserver.assertNoValues(); testObserver.assertComplete(); } @@ -355,7 +365,7 @@ public void convertToServerResponse_withSetupComplete_returnsNoValues() { TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); testObserver.assertNoValues(); testObserver.assertComplete(); @@ -367,7 +377,7 @@ public void convertToServerResponse_withUnknownMessage_returnsErrorResponse() { TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); testObserver.assertValueCount(1); testObserver.assertComplete(); @@ -399,21 +409,22 @@ public void convertToServerResponse_withContentAndUsageMetadata_emitsMultiple() TestObserver testObserver = new TestObserver<>(); - GeminiLlmConnection.convertToServerResponse(message).subscribe(testObserver); + GeminiLlmConnection.convertToServerResponse(message, new PendingText()).subscribe(testObserver); - testObserver.assertValueCount(2); + testObserver.assertValueCount(3); testObserver.assertComplete(); List responses = testObserver.values(); - // Check for ServerContent response + // Check for ServerContent response: the text in full, then the turn completion LlmResponse contentResponse = responses.get(0); assertThat(contentResponse.content()).isPresent(); assertThat(contentResponse.content().get().text()).isEqualTo("Model response"); assertThat(contentResponse.usageMetadata()).isEmpty(); + assertThat(responses.get(1).turnComplete()).hasValue(true); // Check for UsageMetadata response - LlmResponse usageResponse = responses.get(1); + LlmResponse usageResponse = responses.get(2); assertThat(usageResponse.content()).isEmpty(); assertThat(usageResponse.usageMetadata()).isPresent(); GenerateContentResponseUsageMetadata expectedUsageMetadata = @@ -425,6 +436,250 @@ public void convertToServerResponse_withContentAndUsageMetadata_emitsMultiple() assertThat(usageResponse.usageMetadata()).hasValue(expectedUsageMetadata); } + @Test + public void convertToServerResponse_textArrivingWithTurnComplete_isEmittedOnlyInFull() { + // Text that arrives with the end of the turn is not streamed on its own. + assertThat( + convertAll( + modelTurn(Part.fromText("Hel")), + serverContent( + LiveServerContent.builder() + .modelTurn(modelContent(Part.fromText("lo"))) + .turnComplete(true)))) + .containsExactly("partial text:Hel", "text:Hello", "turnComplete") + .inOrder(); + } + + @Test + public void convertToServerResponse_interrupted_fullTextReportsIt() { + assertThat( + convertAll( + modelTurn(Part.fromText("Hel")), + serverContent(LiveServerContent.builder().interrupted(true)))) + .containsExactly("partial text:Hel", "text:Hel interrupted") + .inOrder(); + } + + @Test + public void convertToServerResponse_textArrivingWithInterrupted_isEmittedOnlyInFull() { + assertThat( + convertAll( + modelTurn(Part.fromText("Hel")), + serverContent( + LiveServerContent.builder() + .modelTurn(modelContent(Part.fromText("lo"))) + .interrupted(true)))) + .containsExactly("partial text:Hel", "text:Hello interrupted") + .inOrder(); + } + + @Test + public void convertToServerResponse_interruptedWithTurnComplete_bothResponsesReportIt() { + assertThat( + convertAll( + modelTurn(Part.fromText("Hel")), + serverContent(LiveServerContent.builder().interrupted(true).turnComplete(true)))) + .containsExactly("partial text:Hel", "text:Hel interrupted", "turnComplete interrupted") + .inOrder(); + } + + @Test + public void convertToServerResponse_interruptedWithInlineData_bothResponsesReportIt() { + assertThat( + convertAll( + modelTurn(Part.fromText("Hel")), + serverContent( + LiveServerContent.builder() + .modelTurn(modelContent(Part.fromBytes(new byte[] {1, 2}, "audio/pcm"))) + .interrupted(true)))) + .containsExactly("partial text:Hel", "text:Hel interrupted", "inlineData interrupted") + .inOrder(); + } + + @Test + public void convertToServerResponse_interruptedWithTranscription_fullTextReportsIt() { + assertThat( + convertAll( + modelTurn(Part.fromText("Hel")), + serverContent( + LiveServerContent.builder() + .interrupted(true) + .outputTranscription(Transcription.builder().text("Hel").build())))) + .containsExactly("partial text:Hel", "text:Hel interrupted", "output:Hel") + .inOrder(); + } + + @Test + public void convertToServerResponse_interruptedWithoutText_isReportedOnMessageResponse() { + assertThat(convertAll(serverContent(LiveServerContent.builder().interrupted(true)))) + .containsExactly("interrupted"); + assertThat( + convertAll( + serverContent( + LiveServerContent.builder() + .interrupted(true) + .outputTranscription(Transcription.builder().text("Hel").build())))) + .containsExactly("output:Hel interrupted"); + assertThat( + convertAll( + serverContent( + LiveServerContent.builder() + .modelTurn(modelContent(Part.fromBytes(new byte[] {1, 2}, "audio/pcm"))) + .interrupted(true)))) + .containsExactly("inlineData interrupted"); + assertThat( + convertAll( + serverContent(LiveServerContent.builder().interrupted(true).turnComplete(true)))) + .containsExactly("turnComplete interrupted"); + } + + @Test + public void convertToServerResponse_textAndOtherPartInOneMessage_keepsOtherPartAfterFullText() { + Part code = + Part.builder() + .executableCode( + ExecutableCode.builder().code("print(1)").language(Language.Known.PYTHON)) + .build(); + LiveServerMessage message = + serverContent( + LiveServerContent.builder() + .modelTurn( + Content.builder().role("model").parts(Part.fromText("Run"), code).build())); + + List responses = + GeminiLlmConnection.convertToServerResponse(message, new PendingText()) + .toList() + .blockingGet(); + + assertThat(describe(responses)).containsExactly("text:Run", "executableCode").inOrder(); + assertThat(responses.get(1).content().get().role()).hasValue("model"); + } + + @Test + public void convertToServerResponse_textAndInlineDataWithTurnComplete_keepsInlineData() { + assertThat( + convertAll( + serverContent( + LiveServerContent.builder() + .modelTurn( + Content.builder() + .role("model") + .parts( + Part.fromText("Hello"), + Part.fromBytes(new byte[] {1, 2}, "audio/pcm")) + .build()) + .turnComplete(true)))) + .containsExactly("text:Hello", "inlineData turnComplete") + .inOrder(); + } + + @Test + public void convertToServerResponse_toolCall_emitsPendingTextInFullFirst() { + LiveServerMessage toolCall = + LiveServerMessage.builder() + .toolCall( + LiveServerToolCall.builder() + .functionCalls(ImmutableList.of(FunctionCall.builder().name("tool").build())) + .build()) + .build(); + + assertThat(convertAll(modelTurn(Part.fromText("Checking")), toolCall)) + .containsExactly("partial text:Checking", "text:Checking", "call:tool") + .inOrder(); + } + + @Test + public void convertToServerResponse_thoughtThenText_emitsEachInFull() { + assertThat( + convertAll( + modelTurn(Part.builder().text("Plan").thought(true).build()), + modelTurn(Part.fromText("Answer")), + serverContent(LiveServerContent.builder().turnComplete(true)))) + .containsExactly( + "partial thought:Plan", + "thought:Plan", + "partial text:Answer", + "text:Answer", + "turnComplete") + .inOrder(); + } + + @Test + public void convertToServerResponse_inlineData_isNotPartialAndDoesNotEndPendingText() { + // Only text is partial, even in a message that sets turnComplete to false. + assertThat( + convertAll( + modelTurn(Part.fromText("Hel")), + serverContent( + LiveServerContent.builder() + .modelTurn(modelContent(Part.fromBytes(new byte[] {1, 2}, "audio/pcm"))) + .turnComplete(false)), + modelTurn(Part.fromText("lo")), + serverContent(LiveServerContent.builder().turnComplete(true)))) + .containsExactly( + "partial text:Hel", "inlineData", "partial text:lo", "text:Hello", "turnComplete") + .inOrder(); + } + + @Test + public void convertToServerResponse_textWithTranscription_emitsTranscriptionSeparately() { + // The partial text response is not persisted, so the transcription is not put on it. + assertThat( + convertAll( + serverContent( + LiveServerContent.builder() + .modelTurn(modelContent(Part.fromText("Hel"))) + .inputTranscription(Transcription.builder().text("Hi").build())))) + .containsExactly("partial text:Hel", "input:Hi") + .inOrder(); + } + + @Test + public void convertToServerResponse_emptyText_endsPendingTextLikeOtherParts() { + // A part with empty text, such as a bare thought signature, is not text. + assertThat(convertAll(modelTurn(Part.fromText("Hel")), modelTurn(Part.fromText("")))) + .containsExactly("partial text:Hel", "text:Hel", "text:") + .inOrder(); + } + + @Test + public void convertToServerResponse_otherPart_endsPendingText() { + Part code = + Part.builder() + .executableCode( + ExecutableCode.builder().code("print(1)").language(Language.Known.PYTHON)) + .build(); + + assertThat(convertAll(modelTurn(Part.fromText("Running")), modelTurn(code))) + .containsExactly("partial text:Running", "text:Running", "executableCode") + .inOrder(); + } + + @Test + public void receive_textChunks_streamedAsPartialAndEmittedInFullAtTurnComplete() { + // The Live API may leave turnComplete out of a message that does not end the turn. + CompletableFuture transportFuture = new CompletableFuture<>(); + GeminiLlmConnection connection = new GeminiLlmConnection(transportFuture); + TestSubscriber subscriber = connection.receive().test(); + + transportFuture.complete( + new FakeTransport( + /* endStreamOnReceive= */ true, + ImmutableList.of( + modelTurn(Part.fromText("Hel")), + serverContent( + LiveServerContent.builder() + .modelTurn(modelContent(Part.fromText("lo"))) + .turnComplete(false)), + serverContent(LiveServerContent.builder().turnComplete(true))))); + subscriber.awaitDone(5, TimeUnit.SECONDS); + + subscriber.assertComplete(); + assertThat(describe(subscriber.values())) + .containsExactly("partial text:Hel", "partial text:lo", "text:Hello", "turnComplete") + .inOrder(); + } + @Test public void receive_completesWhenTransportSignalsStreamEnd() { // A transport that ends its own receive stream completes the connection's response stream, so a @@ -457,18 +712,94 @@ public void close_afterTransportStreamEnd_isNoOp() { assertThat(transport.closeCount.get()).isEqualTo(0); } + /** Converts {@code messages} in order with one {@link PendingText}, as a connection does. */ + private static ImmutableList convertAll(LiveServerMessage... messages) { + PendingText pendingText = new PendingText(); + List responses = new ArrayList<>(); + for (LiveServerMessage message : messages) { + responses.addAll( + GeminiLlmConnection.convertToServerResponse(message, pendingText).toList().blockingGet()); + } + return describe(responses); + } + + private static LiveServerMessage modelTurn(Part part) { + return serverContent(LiveServerContent.builder().modelTurn(modelContent(part))); + } + + private static LiveServerMessage serverContent(LiveServerContent.Builder serverContent) { + return LiveServerMessage.builder().serverContent(serverContent.build()).build(); + } + + private static Content modelContent(Part part) { + return Content.builder().role("model").parts(part).build(); + } + /** - * A minimal in-process {@link GeminiLiveTransport} that records sends and closes and can end its - * stream. + * Describes each response by its partial flag, parts and turn signals, such as {@code partial + * text:Hel} or {@code turnComplete}. + */ + private static ImmutableList describe(List responses) { + return responses.stream() + .map( + response -> { + List words = new ArrayList<>(); + if (response.partial().orElse(false)) { + words.add("partial"); + } + response.content().flatMap(Content::parts).orElse(ImmutableList.of()).stream() + .map(GeminiLlmConnectionTest::describePart) + .forEach(words::add); + response + .inputTranscription() + .flatMap(Transcription::text) + .ifPresent(text -> words.add("input:" + text)); + response + .outputTranscription() + .flatMap(Transcription::text) + .ifPresent(text -> words.add("output:" + text)); + if (response.turnComplete().orElse(false)) { + words.add("turnComplete"); + } + if (response.interrupted().orElse(false)) { + words.add("interrupted"); + } + return String.join(" ", words); + }) + .collect(toImmutableList()); + } + + private static String describePart(Part part) { + if (part.text().isPresent()) { + return (part.thought().orElse(false) ? "thought:" : "text:") + part.text().get(); + } + if (part.functionCall().isPresent()) { + return "call:" + part.functionCall().get().name().orElse(""); + } + if (part.inlineData().isPresent()) { + return "inlineData"; + } + return part.executableCode().isPresent() ? "executableCode" : "other"; + } + + /** + * A minimal in-process {@link GeminiLiveTransport} that delivers the given messages, records + * sends and closes and can end its stream. */ private static final class FakeTransport implements GeminiLiveTransport { private final boolean endStreamOnReceive; + private final ImmutableList messages; final AtomicInteger closeCount = new AtomicInteger(); final List clientContents = new CopyOnWriteArrayList<>(); final List toolResponses = new CopyOnWriteArrayList<>(); FakeTransport(boolean endStreamOnReceive) { + this(endStreamOnReceive, ImmutableList.of()); + } + + FakeTransport(boolean endStreamOnReceive, ImmutableList messages) { this.endStreamOnReceive = endStreamOnReceive; + this.messages = messages; } @Override @@ -491,6 +822,7 @@ public CompletableFuture sendToolResponse(LiveSendToolResponseParameters p @Override public CompletableFuture receive( Consumer onMessage, Runnable onStreamEnd) { + messages.forEach(onMessage); if (endStreamOnReceive) { onStreamEnd.run(); } diff --git a/core/src/test/java/com/google/adk/runner/RunnerTest.java b/core/src/test/java/com/google/adk/runner/RunnerTest.java index 2b9951121..a64c781a3 100644 --- a/core/src/test/java/com/google/adk/runner/RunnerTest.java +++ b/core/src/test/java/com/google/adk/runner/RunnerTest.java @@ -3551,6 +3551,43 @@ public void runAsync_partialEvent_streamedButNotPassedToSessionService() { .isFalse(); } + @Test + public void runLive_partialEvent_streamedButNotPassedToSessionService() throws Exception { + // The live model streams a partial text chunk followed by the full text. + LlmResponse partialResponse = + LlmResponse.builder() + .content(Content.builder().role("model").parts(Part.fromText("partial")).build()) + .partial(true) + .build(); + LlmResponse finalResponse = + LlmResponse.builder() + .content(Content.builder().role("model").parts(Part.fromText("final")).build()) + .build(); + TestLlm testLlm = new TestLlm(() -> Flowable.just(partialResponse, finalResponse)); + RecordingSessionService sessionService = new RecordingSessionService(); + Runner runner = + Runner.builder() + .app(App.builder().name("test").rootAgent(createTestAgent(testLlm)).build()) + .sessionService(sessionService) + .build(); + Session session = sessionService.createSession("test", "user").blockingGet(); + LiveRequestQueue liveRequestQueue = new LiveRequestQueue(); + + TestSubscriber testSubscriber = + runner.runLive(session, liveRequestQueue, RunConfig.builder().build()).test(); + liveRequestQueue.content(createContent("hi")); + liveRequestQueue.close(); + testSubscriber.await(); + + testSubscriber.assertComplete(); + // The partial event is still streamed to the caller. + assertThat(testSubscriber.values().stream().anyMatch(event -> event.partial().orElse(false))) + .isTrue(); + // As in runAsync, only the final event is handed to the session service. + assertThat(sessionService.appendedEvents.stream().map(e -> e.content().get().text()).toList()) + .containsExactly("final"); + } + /** A session service that records every event passed to {@code appendEvent} for assertions. */ private static final class RecordingSessionService implements BaseSessionService { private final InMemorySessionService delegate = new InMemorySessionService();