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();