From 923d602e51aea8be4811676afb57d6264f4fe877 Mon Sep 17 00:00:00 2001 From: Damian Momot Date: Tue, 13 Oct 2026 20:31:00 -0700 Subject: [PATCH] refactor: use genai SDK continuation PiperOrigin-RevId: 999074260 --- .../java/com/google/adk/models/Gemini.java | 199 ++---- .../google/adk/models/GeminiContinuation.java | 328 --------- .../adk/models/GeminiContinuationTest.java | 670 ++++-------------- pom.xml | 2 +- 4 files changed, 195 insertions(+), 1004 deletions(-) delete mode 100644 core/src/main/java/com/google/adk/models/GeminiContinuation.java 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 63ba27bde..09402c2b5 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -20,8 +20,6 @@ import com.google.adk.Version; import com.google.adk.internal.http.HttpClientFactory; -import com.google.adk.models.GeminiContinuation.ResumeRequest; -import com.google.adk.models.GeminiContinuation.StreamedOutput; import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableMap; @@ -40,16 +38,12 @@ import com.google.genai.types.Part; import com.google.genai.types.PartialArg; import io.reactivex.rxjava3.core.Flowable; -import java.lang.reflect.Method; import java.time.Duration; -import java.util.ArrayDeque; import java.util.ArrayList; -import java.util.Deque; import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.Optional; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import okhttp3.OkHttpClient; @@ -68,10 +62,6 @@ public class Gemini extends BaseLlm { private static final Logger logger = LoggerFactory.getLogger(Gemini.class); private static final ImmutableMap TRACKING_HEADERS; - /** The google-genai setter for automatic continuation, or null on releases without it. */ - private static final @Nullable Method automaticContinuationSetter = - findAutomaticContinuationSetter(); - private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executorService) { OkHttpClient client = executorService == null @@ -297,31 +287,31 @@ private Flowable sendRequest( LlmRequest llmRequest, boolean stream, @Nullable CacheMetadata cacheMetadata) { - GenerateContentConfig config = withoutAutomaticContinuation(llmRequest.config().orElse(null)); + GenerateContentConfig config = llmRequest.config().orElse(null); logger.trace("Request Contents: {}", llmRequest.contents()); logger.trace("Request Config: {}", config); - List contents = llmRequest.contents(); if (stream) { logger.debug("Sending streaming generateContent request to model {}", effectiveModelName); CompletableFuture> streamFuture = - apiClient.async.models.generateContentStream(effectiveModelName, contents, config); + apiClient.async.models.generateContentStream( + effectiveModelName, llmRequest.contents(), config); + return Flowable.defer( () -> - streamResumingPauses( - effectiveModelName, - new GeminiContinuation(contents, config), - streamFuture, - cacheMetadata)); + new StreamingResponseAggregator() + .process( + Flowable.fromFuture(streamFuture).flatMapIterable(iterable -> iterable), + cacheMetadata)); } else { logger.debug("Sending generateContent request to model {}", effectiveModelName); - CompletableFuture responseFuture = - apiClient.async.models.generateContent(effectiveModelName, contents, config); - return Flowable.defer( - () -> - generateResumingPauses( - effectiveModelName, new GeminiContinuation(contents, config), responseFuture)) + return Flowable.fromFuture( + apiClient + .async + .models + .generateContent(effectiveModelName, llmRequest.contents(), config) + .thenApplyAsync(LlmResponse::create)) .map(response -> withCacheMetadata(response, cacheMetadata)); } } @@ -335,131 +325,8 @@ private static LlmResponse withCacheMetadata( @VisibleForTesting static Flowable processRawResponses(Flowable rawResponses) { - return Flowable.defer(() -> new StreamingResponseAggregator().process(rawResponses)); - } - - /** - * Returns {@code config} with google-genai's automatic continuation turned off, because ADK - * resumes paused generations itself. Calls the setter by reflection because ADK also builds - * against google-genai releases without it. - */ - private static @Nullable GenerateContentConfig withoutAutomaticContinuation( - @Nullable GenerateContentConfig config) { - if (automaticContinuationSetter == null) { - return config; - } - GenerateContentConfig.Builder builder = - config == null ? GenerateContentConfig.builder() : config.toBuilder(); - try { - automaticContinuationSetter.invoke(builder, false); - } catch (ReflectiveOperationException e) { - // The setter is a public builder method, so it is accessible and does not throw. - throw new AssertionError(e); - } - return builder.build(); - } - - private static @Nullable Method findAutomaticContinuationSetter() { - try { - return GenerateContentConfig.Builder.class.getMethod("automaticContinuation", boolean.class); - } catch (NoSuchMethodException e) { - return null; - } - } - - /** - * Returns the final response of a generation, sending another request each time the model pauses. - * Repeats instead of recursing, so a generation resumed many times keeps a flat stack. - */ - private Flowable generateResumingPauses( - String modelName, - GeminiContinuation continuation, - CompletableFuture firstResponse) { - Deque> responses = new ArrayDeque<>(); - responses.add(firstResponse); - return Flowable.defer(() -> Flowable.fromFuture(responses.remove())) - .mapOptional(response -> completeOrResume(modelName, continuation, response, responses)) - .repeatUntil(responses::isEmpty); - } - - /** - * Returns the final response of the generation, or empty after queueing the request to resume it. - */ - private Optional completeOrResume( - String modelName, - GeminiContinuation continuation, - GenerateContentResponse response, - Deque> responses) { - LlmResponse llmResponse = LlmResponse.create(response); - Optional next = - continuation.advance( - GeminiContinuation.resumeToken(response), - llmResponse.content().flatMap(Content::parts).orElse(ImmutableList.of()), - llmResponse.usageMetadata().orElse(null)); - if (next.isEmpty()) { - return Optional.of(continuation.complete(llmResponse)); - } - responses.add( - apiClient.async.models.generateContent( - modelName, next.get().contents(), next.get().config())); - return Optional.empty(); - } - - /** - * Streams a generation, sending another request each time the model pauses, and sums the usage of - * all requests on the final response. Repeats instead of recursing, so a generation resumed many - * times keeps a flat stack. - */ - private Flowable streamResumingPauses( - String modelName, - GeminiContinuation continuation, - CompletableFuture> firstStream, - @Nullable CacheMetadata cacheMetadata) { - Deque>> streams = new ArrayDeque<>(); - streams.add(firstStream); - StreamingResponseAggregator aggregator = new StreamingResponseAggregator(); - return Flowable.defer(() -> streamRequest(modelName, continuation, streams)) - .repeatUntil(streams::isEmpty) - .concatMap(aggregator::processRawResponse) - .concatWith( - Flowable.defer( - () -> - aggregator - .processFinalResponse() - .map(continuation::withSummedUsage) - .map(response -> withCacheMetadata(response, cacheMetadata)))); - } - - /** - * Streams the chunks of one request, then queues the request that resumes the generation, if any. - */ - private Flowable streamRequest( - String modelName, - GeminiContinuation continuation, - Deque>> streams) { - StreamedOutput output = new StreamedOutput(continuation); - return Flowable.fromFuture(streams.remove()) - .flatMapIterable(stream -> stream) - .mapOptional(output::record) - .doOnComplete(() -> queueResume(modelName, continuation, output, streams)); - } - - /** - * Records {@code output} in {@code continuation} and queues the request that resumes the - * generation, if any. - */ - private void queueResume( - String modelName, - GeminiContinuation continuation, - StreamedOutput output, - Deque>> streams) { - continuation - .advance(output.token(), output.parts(), output.usage()) - .ifPresent( - next -> - streams.add( - apiClient.async.models.generateContentStream( - modelName, next.contents(), next.config()))); + return Flowable.defer( + () -> new StreamingResponseAggregator().process(rawResponses, /* cacheMetadata= */ null)); } @Override @@ -487,7 +354,7 @@ protected CompletableFuture connectLiveTransport( return apiClient.async.live.connect(modelName, config).thenApply(GenAiLiveTransport::new); } - static final class StreamingResponseAggregator { + private static final class StreamingResponseAggregator { 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. @@ -498,7 +365,7 @@ static final class StreamingResponseAggregator { * and nothing else worth keeping. Compared by rebuilding rather than against a single literal, * so a terminator that also carries an explicit {@code thought=false} is still recognised. */ - static boolean isStreamTerminator(Part part) { + private static boolean isStreamTerminator(Part part) { if (!part.text().map(String::isEmpty).orElse(false)) { return false; } @@ -521,12 +388,18 @@ static boolean isStreamTerminator(Part part) { private String currentFcId = null; /** - * Processes a stream of raw responses, emitting partial and aggregated {@link LlmResponse}s. + * Processes a stream of raw responses, emitting partial and aggregated {@link LlmResponse}s, + * and puts {@code cacheMetadata}, if any, on the aggregated response. */ - private Flowable process(Flowable rawResponses) { + private Flowable process( + Flowable rawResponses, @Nullable CacheMetadata cacheMetadata) { return rawResponses .concatMap(this::processRawResponse) - .concatWith(Flowable.defer(this::processFinalResponse)); + .concatWith( + Flowable.defer( + () -> + processFinalResponse() + .map(response -> withCacheMetadata(response, cacheMetadata)))); } /** @@ -558,8 +431,10 @@ private Flowable processRawResponse(GenerateContentResponse rawResp // If the chunk has no text or function calls (e.g. metadata-only or empty), we suppress it // during streaming so it doesn't emit an empty partial response. // Exception: If this is a standalone empty chunk in an otherwise completely empty stream - // (and not a STOP chunk), we emit it directly as a non-partial empty response. + // (and not a STOP chunk), we emit it directly as a non-partial empty response, unless more + // of the stream follows it. if (!isStop(currentProcessedLlmResponse) + && !continuesStream(rawResponse) && accumulatedSequence.isEmpty() && currentTextBuffer.isEmpty()) { return Flowable.just(currentProcessedLlmResponse.toBuilder().partial(null).build()); @@ -846,6 +721,20 @@ private Flowable processFinalResponse() { .build()); } + /** + * Returns whether more of the stream follows this chunk: it carries a continuation token but no + * finish reason, like a pause whose finish reason the SDK clears before resuming it. + */ + private static boolean continuesStream(GenerateContentResponse response) { + return response + .candidates() + .flatMap(candidates -> candidates.stream().findFirst()) + .filter(candidate -> candidate.finishReason().isEmpty()) + .flatMap(Candidate::continuationToken) + .filter(token -> token.length > 0) + .isPresent(); + } + /** Checks whether the response finish reason indicates the stream has finished with STOP. */ private static boolean isStop(LlmResponse response) { return response diff --git a/core/src/main/java/com/google/adk/models/GeminiContinuation.java b/core/src/main/java/com/google/adk/models/GeminiContinuation.java deleted file mode 100644 index ec113566e..000000000 --- a/core/src/main/java/com/google/adk/models/GeminiContinuation.java +++ /dev/null @@ -1,328 +0,0 @@ -/* - * Copyright 2026 Google LLC - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.google.adk.models; - -import com.google.common.annotations.VisibleForTesting; -import com.google.common.collect.ImmutableList; -import com.google.genai.types.Candidate; -import com.google.genai.types.Content; -import com.google.genai.types.FinishReason; -import com.google.genai.types.GenerateContentConfig; -import com.google.genai.types.GenerateContentResponse; -import com.google.genai.types.GenerateContentResponseUsageMetadata; -import com.google.genai.types.MediaModality; -import com.google.genai.types.ModalityTokenCount; -import com.google.genai.types.Part; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.LinkedHashMap; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import org.jspecify.annotations.Nullable; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** One generation, carried across the requests that resume it after the model pauses it. */ -final class GeminiContinuation { - @VisibleForTesting static final int MAX_RESUMES = 256; - - private static final Logger logger = LoggerFactory.getLogger(GeminiContinuation.class); - - private final ImmutableList contents; - private final @Nullable GenerateContentConfig config; - private final List parts = new ArrayList<>(); - private @Nullable GenerateContentResponseUsageMetadata usage = null; - private byte @Nullable [] token = null; - private int resumes = 0; - - GeminiContinuation(List contents, @Nullable GenerateContentConfig config) { - this.contents = ImmutableList.copyOf(contents); - this.config = config; - } - - /** - * Returns the token that resumes {@code response}, or null if it did not pause or has no token. - */ - static byte @Nullable [] resumeToken(GenerateContentResponse response) { - Optional paused = - response - .candidates() - .flatMap(candidates -> candidates.stream().findFirst()) - .filter( - candidate -> - candidate - .finishReason() - .map(reason -> reason.knownEnum() == FinishReason.Known.CONTINUATION) - .orElse(false)); - if (paused.isEmpty()) { - return null; - } - byte[] token = paused.get().continuationToken().filter(bytes -> bytes.length > 0).orElse(null); - if (token == null) { - logger.warn( - "The model paused the generation for continuation, but the response carries no" - + " continuation token, so the partial output is returned."); - } - return token; - } - - /** - * Records one request's output and returns the next request to resume generation. Returns empty - * if generation finished, paused without a new token, or already resumed {@link #MAX_RESUMES} - * times. - */ - Optional advance( - byte @Nullable [] nextToken, - List newParts, - @Nullable GenerateContentResponseUsageMetadata newUsage) { - usage = addUsage(usage, newUsage); - if (nextToken == null && !resumed()) { - return Optional.empty(); - } - appendParts(newParts); - if (nextToken == null) { - return Optional.empty(); - } - if (!willResume(nextToken)) { - if (Arrays.equals(nextToken, token)) { - // A token that did not change means the model made no progress. - logger.warn( - "The model returned the same continuation token twice; returning the output" - + " generated so far."); - } else { - logger.warn( - "The model paused the generation {} times; returning the output generated so far.", - MAX_RESUMES); - } - return Optional.empty(); - } - token = nextToken; - resumes++; - logger.info("The model paused the generation; resuming it."); - return Optional.of(nextRequest(nextToken)); - } - - /** Returns {@code response} with the content and usage of all requests, if it was resumed. */ - LlmResponse complete(LlmResponse response) { - if (!resumed()) { - return response; - } - return response.toBuilder() - .content(parts.isEmpty() ? response.content().orElse(null) : content()) - .usageMetadata(usage) - .build(); - } - - /** - * Returns the final streamed {@code response}, with usage summed across all requests if resumed. - */ - LlmResponse withSummedUsage(LlmResponse response) { - return resumed() ? response.toBuilder().usageMetadata(usage).build() : response; - } - - /** Returns whether the generation took more than one request. */ - private boolean resumed() { - return token != null; - } - - /** Returns whether a request that ended with {@code nextToken} is resumed. */ - private boolean willResume(byte[] nextToken) { - return !Arrays.equals(nextToken, token) && resumes < MAX_RESUMES; - } - - private Content content() { - return Content.builder().role("model").parts(ImmutableList.copyOf(parts)).build(); - } - - private ResumeRequest nextRequest(byte[] nextToken) { - ImmutableList.Builder nextContents = ImmutableList.builder().addAll(contents); - if (!parts.isEmpty()) { - nextContents.add(content()); - } - GenerateContentConfig baseConfig = - config != null ? config : GenerateContentConfig.builder().build(); - return new ResumeRequest( - nextContents.build(), baseConfig.toBuilder().continuationToken(nextToken).build()); - } - - /** Appends {@code newParts}, joining text as the streaming aggregator does. */ - private void appendParts(List newParts) { - for (Part part : newParts) { - int last = parts.size() - 1; - Part previous = last >= 0 ? parts.get(last) : null; - if (previous != null && canJoin(previous, part)) { - Part.Builder joined = - Part.builder().text(previous.text().orElseThrow() + part.text().orElseThrow()); - previous.thought().ifPresent(joined::thought); - previous - .thoughtSignature() - .filter(signature -> signature.length > 0) - .or(part::thoughtSignature) - .ifPresent(joined::thoughtSignature); - parts.set(last, joined.build()); - } else { - parts.add(part); - } - } - } - - /** Returns whether two parts are non-empty text of the same kind. */ - private static boolean canJoin(Part first, Part second) { - return isText(first) - && isText(second) - && first.thought().orElse(false).equals(second.thought().orElse(false)); - } - - /** Returns whether the part is non-empty text with at most a thought flag and a signature. */ - private static boolean isText(Part part) { - if (part.text().filter(text -> !text.isEmpty()).isEmpty()) { - return false; - } - Part.Builder text = Part.builder().text(part.text().orElseThrow()); - part.thought().ifPresent(text::thought); - part.thoughtSignature().ifPresent(text::thoughtSignature); - return text.build().equals(part); - } - - /** Returns the combined token usage of two requests. */ - private static @Nullable GenerateContentResponseUsageMetadata addUsage( - @Nullable GenerateContentResponseUsageMetadata total, - @Nullable GenerateContentResponseUsageMetadata next) { - if (total == null || next == null) { - return total == null ? next : total; - } - GenerateContentResponseUsageMetadata.Builder combined = next.toBuilder(); - addCounts(total.promptTokenCount(), next.promptTokenCount()) - .ifPresent(combined::promptTokenCount); - addCounts(total.candidatesTokenCount(), next.candidatesTokenCount()) - .ifPresent(combined::candidatesTokenCount); - addCounts(total.totalTokenCount(), next.totalTokenCount()).ifPresent(combined::totalTokenCount); - addCounts(total.cachedContentTokenCount(), next.cachedContentTokenCount()) - .ifPresent(combined::cachedContentTokenCount); - addCounts(total.thoughtsTokenCount(), next.thoughtsTokenCount()) - .ifPresent(combined::thoughtsTokenCount); - addCounts(total.toolUsePromptTokenCount(), next.toolUsePromptTokenCount()) - .ifPresent(combined::toolUsePromptTokenCount); - addModalityCounts(total.promptTokensDetails(), next.promptTokensDetails()) - .ifPresent(combined::promptTokensDetails); - addModalityCounts(total.candidatesTokensDetails(), next.candidatesTokensDetails()) - .ifPresent(combined::candidatesTokensDetails); - addModalityCounts(total.cacheTokensDetails(), next.cacheTokensDetails()) - .ifPresent(combined::cacheTokensDetails); - addModalityCounts(total.toolUsePromptTokensDetails(), next.toolUsePromptTokensDetails()) - .ifPresent(combined::toolUsePromptTokensDetails); - return combined.build(); - } - - private static Optional addCounts(Optional first, Optional second) { - if (first.isEmpty() && second.isEmpty()) { - return Optional.empty(); - } - return Optional.of(first.orElse(0) + second.orElse(0)); - } - - /** Sums token counts that share a modality, keeping first-seen order. */ - private static Optional> addModalityCounts( - Optional> first, Optional> second) { - if (first.isEmpty() && second.isEmpty()) { - return Optional.empty(); - } - List counts = new ArrayList<>(first.orElse(ImmutableList.of())); - counts.addAll(second.orElse(ImmutableList.of())); - Map<@Nullable MediaModality, Integer> totals = new LinkedHashMap<>(); - for (ModalityTokenCount count : counts) { - totals.merge(count.modality().orElse(null), count.tokenCount().orElse(0), Integer::sum); - } - ImmutableList.Builder summed = ImmutableList.builder(); - totals.forEach( - (modality, tokenCount) -> { - ModalityTokenCount.Builder count = ModalityTokenCount.builder().tokenCount(tokenCount); - if (modality != null) { - count.modality(modality); - } - summed.add(count.build()); - }); - return Optional.of(summed.build()); - } - - /** The contents and config of a request that resumes a paused generation. */ - record ResumeRequest(ImmutableList contents, GenerateContentConfig config) {} - - /** What one streamed request generated, recorded to resume the generation if it pauses. */ - static final class StreamedOutput { - private final GeminiContinuation continuation; - private final List parts = new ArrayList<>(); - private byte @Nullable [] token = null; - private @Nullable GenerateContentResponseUsageMetadata usage = null; - - StreamedOutput(GeminiContinuation continuation) { - this.continuation = continuation; - } - - /** Returns the parts recorded so far, excluding stream terminators. */ - ImmutableList parts() { - return ImmutableList.copyOf(parts); - } - - /** Returns the token that resumes the generation if this request paused, or null. */ - byte @Nullable [] token() { - return token; - } - - /** Returns the request's usage metadata, taken from the last chunk that carried it. */ - @Nullable GenerateContentResponseUsageMetadata usage() { - return usage; - } - - /** - * Records a chunk and returns what to aggregate: the chunk, with the finish reason of a resumed - * pause cleared since that pause does not end the generation, or nothing for a resumed pause - * without output. - */ - Optional record(GenerateContentResponse chunk) { - chunk.usageMetadata().ifPresent(chunkUsage -> usage = chunkUsage); - List candidates = chunk.candidates().orElse(ImmutableList.of()); - if (candidates.isEmpty()) { - return Optional.of(chunk); - } - Candidate candidate = candidates.get(0); - // A stream terminator carries nothing, and resending it would add an empty text part. - ImmutableList chunkParts = - candidate.content().flatMap(Content::parts).orElse(ImmutableList.of()).stream() - .filter(part -> !Gemini.StreamingResponseAggregator.isStreamTerminator(part)) - .collect(ImmutableList.toImmutableList()); - parts.addAll(chunkParts); - byte[] chunkToken = resumeToken(chunk); - if (chunkToken == null) { - return Optional.of(chunk); - } - token = chunkToken; - if (!continuation.willResume(chunkToken)) { - // A pause that is not resumed ends the generation, so it keeps its finish reason. - return Optional.of(chunk); - } - if (chunkParts.isEmpty()) { - // Drop an empty pause chunk: as the first chunk, it would emit a complete, empty response. - return Optional.empty(); - } - List resumable = new ArrayList<>(candidates); - resumable.set(0, candidate.toBuilder().clearFinishReason().build()); - return Optional.of(chunk.toBuilder().candidates(resumable).build()); - } - } -} diff --git a/core/src/test/java/com/google/adk/models/GeminiContinuationTest.java b/core/src/test/java/com/google/adk/models/GeminiContinuationTest.java index be3a182b5..f1b9cd5eb 100644 --- a/core/src/test/java/com/google/adk/models/GeminiContinuationTest.java +++ b/core/src/test/java/com/google/adk/models/GeminiContinuationTest.java @@ -16,7 +16,6 @@ package com.google.adk.models; -import static com.google.adk.models.GeminiContinuation.MAX_RESUMES; import static com.google.common.collect.ImmutableList.toImmutableList; import static com.google.common.collect.Iterables.getLast; import static com.google.common.collect.Iterables.getOnlyElement; @@ -26,24 +25,16 @@ import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.adk.agents.ContextCacheConfig; import com.google.common.collect.ImmutableList; -import com.google.common.collect.ImmutableMap; import com.google.genai.Client; -import com.google.genai.errors.ClientException; import com.google.genai.types.Candidate; import com.google.genai.types.ClientOptions; import com.google.genai.types.Content; import com.google.genai.types.FinishReason; -import com.google.genai.types.FunctionCall; -import com.google.genai.types.GenerateContentConfig; import com.google.genai.types.GenerateContentResponse; import com.google.genai.types.GenerateContentResponseUsageMetadata; -import com.google.genai.types.HttpOptions; -import com.google.genai.types.MediaModality; -import com.google.genai.types.ModalityTokenCount; import com.google.genai.types.Part; -import io.reactivex.rxjava3.core.Flowable; -import io.reactivex.rxjava3.subscribers.TestSubscriber; import java.io.IOException; import java.util.ArrayDeque; import java.util.ArrayList; @@ -51,8 +42,6 @@ import java.util.Base64; import java.util.Deque; import java.util.List; -import java.util.Optional; -import java.util.concurrent.atomic.AtomicReference; import okhttp3.Interceptor; import okhttp3.MediaType; import okhttp3.OkHttpClient; @@ -66,576 +55,228 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; -/** Tests for {@link Gemini} resuming generations paused by the model with a continuation token. */ +/** Tests that {@link Gemini} leaves resuming paused generations to the GenAI SDK. */ @RunWith(JUnit4.class) public final class GeminiContinuationTest { private static final ObjectMapper objectMapper = new ObjectMapper(); private static final Content QUESTION = Content.builder().role("user").parts(Part.fromText("What is the answer?")).build(); - private static final GenerateContentConfig CONFIG = - GenerateContentConfig.builder().temperature(0.5f).build(); + private static final Content EARLIER_ANSWER = + Content.builder().role("model").parts(Part.fromText("It depends.")).build(); + private static final Content FOLLOW_UP = + Content.builder().role("user").parts(Part.fromText("On what?")).build(); + private static final String CACHE_NAME = "cachedContents/answers"; - private final FakeGeminiApi api = new FakeGeminiApi(); + private final ScriptedGeminiApi api = new ScriptedGeminiApi(); private final Gemini gemini = new Gemini("gemini-test-model", api.client()); @Test - public void generateContent_pausedGeneration_resumesAndReturnsWholeGeneration() { + public void generateContent_pausedGeneration_isResumedWithSummedUsage() { api.respond( - withUsage( - paused("\0state", Part.fromText("The answer is")), - GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(10) - .candidatesTokenCount(5) - .totalTokenCount(15) - .thoughtsTokenCount(2) - .toolUsePromptTokenCount(1) - .cachedContentTokenCount(4) - .promptTokensDetails( - ImmutableList.of( - tokens(MediaModality.Known.TEXT, 8), tokens(MediaModality.Known.IMAGE, 2))) - .candidatesTokensDetails(ImmutableList.of(tokens(MediaModality.Known.TEXT, 5))) - .cacheTokensDetails(ImmutableList.of(tokens(MediaModality.Known.TEXT, 4))) - .toolUsePromptTokensDetails(ImmutableList.of(tokens(/* modality= */ null, 1))) - .build()), - withUsage( - finished(Part.fromText(" 42.")), - GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(15) - .candidatesTokenCount(3) - .totalTokenCount(18) - .thoughtsTokenCount(3) - .toolUsePromptTokenCount(2) - .promptTokensDetails( - ImmutableList.of( - tokens(MediaModality.Known.IMAGE, 1), tokens(MediaModality.Known.TEXT, 14))) - .candidatesTokensDetails(ImmutableList.of(tokens(MediaModality.Known.TEXT, 3))) - .toolUsePromptTokensDetails(ImmutableList.of(tokens(/* modality= */ null, 2))) - .build())); - - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); - - assertThat(response.content()).hasValue(modelText("The answer is 42.")); - assertThat(response.finishReason().map(FinishReason::knownEnum)) - .hasValue(FinishReason.Known.STOP); - assertThat(response.usageMetadata()) - .hasValue( - GenerateContentResponseUsageMetadata.builder() - .promptTokenCount(25) - .candidatesTokenCount(8) - .totalTokenCount(33) - .thoughtsTokenCount(5) - .toolUsePromptTokenCount(3) - .cachedContentTokenCount(4) - .promptTokensDetails( - ImmutableList.of( - tokens(MediaModality.Known.TEXT, 22), tokens(MediaModality.Known.IMAGE, 3))) - .candidatesTokensDetails(ImmutableList.of(tokens(MediaModality.Known.TEXT, 8))) - .cacheTokensDetails(ImmutableList.of(tokens(MediaModality.Known.TEXT, 4))) - .toolUsePromptTokensDetails(ImmutableList.of(tokens(/* modality= */ null, 3))) - .build()); - assertThat(api.requests()).hasSize(2); - JsonNode resumed = api.requests().get(1); - assertThat(contents(resumed)).containsExactly(QUESTION, modelText("The answer is")).inOrder(); - assertThat(resumed.get("continuationToken").asText()).isEqualTo(base64("\0state")); - } - - @Test - public void generateContent_streamPausedGeneration_resumesInOneAggregatedStream() { - api.respondStream( - chunk(Part.fromText("The answer")), - withUsage( - paused("state", Part.fromText(" is")), - GenerateContentResponseUsageMetadata.builder().candidatesTokenCount(5).build())); - api.respondStream( - withUsage( - finished(Part.fromText(" 42.")), - GenerateContentResponseUsageMetadata.builder().candidatesTokenCount(3).build())); + response("CONTINUATION", "state", usage(10, 5, 15), Part.fromText("The answer is")), + response("STOP", /* token= */ null, usage(15, 3, 18), Part.fromText(" 42."))); - List responses = generate(CONFIG, /* stream= */ true); + LlmResponse response = getOnlyElement(generate(/* stream= */ false)); - assertThat(responses.stream().map(r -> r.partial().orElse(false)).collect(toImmutableList())) - .containsExactly(true, true, true, false) - .inOrder(); - assertThat(responses.stream().filter(r -> r.errorCode().isPresent())).isEmpty(); - // A pause is not the end of the generation, so no response reports it. - assertThat(responses.stream().map(r -> r.finishReason().map(FinishReason::knownEnum))) - .doesNotContain(Optional.of(FinishReason.Known.CONTINUATION)); - LlmResponse last = getLast(responses); - assertThat(firstText(last)).isEqualTo("The answer is 42."); - assertThat(last.finishReason().map(FinishReason::knownEnum)).hasValue(FinishReason.Known.STOP); - assertThat( - last.usageMetadata() - .flatMap(GenerateContentResponseUsageMetadata::candidatesTokenCount)) - .hasValue(8); - assertThat(api.requests()).hasSize(2); - JsonNode resumed = api.requests().get(1); - assertThat(getLast(contents(resumed))).isEqualTo(modelText("The answer is")); - assertThat(resumed.get("continuationToken").asText()).isEqualTo(base64("state")); + assertThat(texts(response)).containsExactly("The answer is", " 42.").inOrder(); + assertThat(finishReason(response)).isEqualTo(FinishReason.Known.STOP); + assertUsage(response, /* prompt= */ 10 + 15, /* candidates= */ 5 + 3, /* total= */ 15 + 18); + ImmutableList requests = api.requests(); + assertThat(requests).hasSize(2); + assertThat(requests.get(0).has("continuationToken")).isFalse(); + assertThat(requests.get(1).get("continuationToken").asText()).isEqualTo(base64("state")); + assertThat(requests.get(1).get("contents")).isEqualTo(requests.get(0).get("contents")); } @Test - public void generateContent_pausedTwice_resumesUntilComplete() { - api.respond( - paused("first", Part.fromText("a")), - paused("second", Part.fromText("b")), - finished(Part.fromText("c"))); - - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); - - assertThat(firstText(response)).isEqualTo("abc"); - assertThat(api.requests()).hasSize(3); - JsonNode secondResume = api.requests().get(2); - assertThat(contents(secondResume)).containsExactly(QUESTION, modelText("ab")).inOrder(); - assertThat(secondResume.get("continuationToken").asText()).isEqualTo(base64("second")); - } - - @Test - public void generateContent_streamPausedTwice_resumesUntilComplete() { - api.respondStream(paused("first", Part.fromText("a"))); - api.respondStream(paused("second", Part.fromText("b"))); - api.respondStream(finished(Part.fromText("c"))); - - List responses = generate(CONFIG, /* stream= */ true); - - assertThat(firstText(getLast(responses))).isEqualTo("abc"); - assertThat(api.requests()).hasSize(3); - JsonNode secondResume = api.requests().get(2); - assertThat(contents(secondResume)).containsExactly(QUESTION, modelText("ab")).inOrder(); - assertThat(secondResume.get("continuationToken").asText()).isEqualTo(base64("second")); - } - - @Test - public void generateContent_pauseWithoutOutput_resumesWithOriginalContents() { - api.respond(paused("state"), finished(Part.fromText("Done."))); + public void generateContent_pausedStream_isResumedWithSummedUsage() { + api.respondStream( + response( + /* finishReason= */ null, /* token= */ null, /* usage= */ null, Part.fromText("The")), + response("CONTINUATION", "state", usage(10, 5, 15), Part.fromText(" answer is"))); + api.respondStream(response("STOP", /* token= */ null, usage(15, 3, 18), Part.fromText(" 42."))); - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); + List responses = generate(/* stream= */ true); - assertThat(firstText(response)).isEqualTo("Done."); + LlmResponse last = getLast(responses); + assertThat(last.partial()).isEmpty(); + assertThat(texts(last)).containsExactly("The answer is 42."); + assertThat(finishReason(last)).isEqualTo(FinishReason.Known.STOP); + assertUsage(last, /* prompt= */ 10 + 15, /* candidates= */ 5 + 3, /* total= */ 15 + 18); + assertThat(responses.stream().map(GeminiContinuationTest::finishReason)) + .doesNotContain(FinishReason.Known.CONTINUATION); assertThat(api.requests()).hasSize(2); - assertThat(contents(api.requests().get(1))).containsExactly(QUESTION); } @Test public void generateContent_streamPauseWithoutOutput_emitsNoEmptyResponse() { - api.respondStream(paused("state")); - api.respondStream(finished(Part.fromText("Done."))); + api.respondStream(response("CONTINUATION", "state", usage(10, 0, 10))); + api.respondStream( + response("STOP", /* token= */ null, usage(15, 3, 18), Part.fromText("Done."))); - List responses = generate(CONFIG, /* stream= */ true); + List responses = generate(/* stream= */ true); - assertThat(responses.stream().map(r -> r.partial().orElse(false)).collect(toImmutableList())) + assertThat(responses.stream().map(response -> response.partial().orElse(false))) .containsExactly(true, false) .inOrder(); - assertThat(firstText(getLast(responses))).isEqualTo("Done."); + assertThat(texts(getLast(responses))).containsExactly("Done."); assertThat(api.requests()).hasSize(2); - assertThat(contents(api.requests().get(1))).containsExactly(QUESTION); } @Test public void generateContent_streamPauseWithOnlyStreamTerminator_emitsNoEmptyResponse() { - api.respondStream(paused("state", Part.fromText(""))); - api.respondStream(finished(Part.fromText("Done."))); + api.respondStream(response("CONTINUATION", "state", usage(10, 0, 10), Part.fromText(""))); + api.respondStream( + response("STOP", /* token= */ null, usage(15, 3, 18), Part.fromText("Done."))); - List responses = generate(CONFIG, /* stream= */ true); + List responses = generate(/* stream= */ true); - assertThat(responses.stream().map(r -> r.partial().orElse(false)).collect(toImmutableList())) + assertThat(responses.stream().map(response -> response.partial().orElse(false))) .containsExactly(true, false) .inOrder(); - assertThat(firstText(getLast(responses))).isEqualTo("Done."); - assertThat(api.requests()).hasSize(2); - assertThat(contents(api.requests().get(1))).containsExactly(QUESTION); - } - - @Test - public void generateContent_resumed_keepsRequestConfig() { - GenerateContentConfig config = - CONFIG.toBuilder() - .httpOptions( - HttpOptions.builder().extraBody(ImmutableMap.of("custom", "value")).build()) - .build(); - api.respond(paused("state", Part.fromText("a")), finished(Part.fromText("b"))); - - List unused = generate(config, /* stream= */ false); - - assertThat(api.requests()).hasSize(2); - JsonNode first = api.requests().get(0); - JsonNode resumed = api.requests().get(1); - assertThat(first.has("continuationToken")).isFalse(); - assertThat(resumed.get("custom").asText()).isEqualTo("value"); - assertThat(resumed.get("continuationToken").asText()).isEqualTo(base64("state")); - assertThat(resumed.get("generationConfig")).isEqualTo(first.get("generationConfig")); - } - - @Test - public void generateContent_configWithToken_sendsItFirstAndReplacesItOnResume() { - api.respond(paused("next", Part.fromText("a")), finished(Part.fromText("b"))); - - List unused = - generate( - CONFIG.toBuilder().continuationToken("caller".getBytes(UTF_8)).build(), - /* stream= */ false); - - assertThat(api.requests().stream().map(request -> request.get("continuationToken").asText())) - .containsExactly(base64("caller"), base64("next")) - .inOrder(); - } - - @Test - public void generateContent_resumed_joinsSignedTextKeepingFirstSignature() { - byte[] first = "first".getBytes(UTF_8); - byte[] second = "second".getBytes(UTF_8); - api.respond( - paused("state", Part.builder().text("The answer is").thoughtSignature(first).build()), - finished(Part.builder().text(" 42.").thoughtSignature(second).build())); - - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); - - Part part = getOnlyElement(response.content().flatMap(Content::parts).get()); - assertThat(part.text()).hasValue("The answer is 42."); - assertThat(part.thoughtSignature().get()).isEqualTo(first); - } - - @Test - public void generateContent_resumed_keepsThoughtApartFromAnswer() { - Part thought = Part.builder().text("Thinking.").thought(true).build(); - api.respond(paused("state", thought), finished(Part.fromText("The answer is 42."))); - - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); - - assertThat(response.content().flatMap(Content::parts).get()) - .containsExactly(thought, Part.fromText("The answer is 42.")) - .inOrder(); - assertThat(api.requests()).hasSize(2); - assertThat(getLast(contents(api.requests().get(1)))) - .isEqualTo(Content.builder().role("model").parts(thought).build()); - } - - @Test - public void generateContent_resumed_keepsFunctionCallApartFromText() { - Part call = Part.fromFunctionCall("lookup", ImmutableMap.of("query", "answer")); - api.respond( - paused("state", Part.fromText("Let me check."), call), finished(Part.fromText("42."))); - - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); - - assertThat(response.content().flatMap(Content::parts).get()) - .containsExactly(Part.fromText("Let me check."), call, Part.fromText("42.")) - .inOrder(); - assertThat(api.requests()).hasSize(2); - assertThat(getLast(contents(api.requests().get(1)))) - .isEqualTo( - Content.builder().role("model").parts(Part.fromText("Let me check."), call).build()); - } - - @Test - public void generateContent_streamResumed_resendsFunctionCallWithoutClientId() { - Part call = Part.fromFunctionCall("lookup", ImmutableMap.of("q", "x")); - api.respondStream(chunk(call), paused("state")); - api.respondStream(finished(Part.fromText("Done."))); - - List responses = generate(CONFIG, /* stream= */ true); - - FunctionCall finalCall = - getLast(responses).content().flatMap(Content::parts).get().stream() - .flatMap(part -> part.functionCall().stream()) - .findFirst() - .orElseThrow(); - assertThat(finalCall.id().orElse("")).startsWith("adk-"); - assertThat(api.requests()).hasSize(2); - assertThat(getLast(contents(api.requests().get(1)))) - .isEqualTo(Content.builder().role("model").parts(call).build()); - } - - @Test - public void generateContent_resumed_joinsThoughtSplitByPause() { - api.respond( - paused("state", Part.builder().text("Think").thought(true).build()), - finished(Part.builder().text("ing.").thought(true).build(), Part.fromText("42."))); - - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); - - assertThat(response.content().flatMap(Content::parts).get()) - .containsExactly( - Part.builder().text("Thinking.").thought(true).build(), Part.fromText("42.")) - .inOrder(); - } - - @Test - public void generateContent_resumed_joinsTextKeepingResumedSignature() { - byte[] signature = "sig".getBytes(UTF_8); - api.respond( - paused("state", Part.fromText("The answer is")), - finished(Part.builder().text(" 42.").thoughtSignature(signature).build())); - - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); - - Part part = getOnlyElement(response.content().flatMap(Content::parts).get()); - assertThat(part.text()).hasValue("The answer is 42."); - assertThat(part.thoughtSignature().get()).isEqualTo(signature); - } - - @Test - public void generateContent_resumed_keepsEmptyTextApart() { - api.respond( - paused("state", Part.fromText("a"), Part.fromText("")), finished(Part.fromText("b"))); - - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); - - assertThat(response.content().flatMap(Content::parts).get()) - .containsExactly(Part.fromText("a"), Part.fromText(""), Part.fromText("b")) - .inOrder(); - } - - @Test - public void generateContent_pauseWithoutToken_returnsOutput() { - api.respond(response("CONTINUATION", /* token= */ null, Part.fromText("a"))); - - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); - - assertThat(api.requests()).hasSize(1); - assertThat(firstText(response)).isEqualTo("a"); - assertThat(response.finishReason().map(FinishReason::knownEnum)) - .hasValue(FinishReason.Known.CONTINUATION); - } - - @Test - public void generateContent_streamPauseWithoutToken_returnsOutput() { - api.respondStream(response("CONTINUATION", /* token= */ null, Part.fromText("a"))); - - List responses = generate(CONFIG, /* stream= */ true); - - assertThat(api.requests()).hasSize(1); - LlmResponse last = getLast(responses); - assertThat(firstText(last)).isEqualTo("a"); - assertThat(last.finishReason().map(FinishReason::knownEnum)) - .hasValue(FinishReason.Known.CONTINUATION); - assertThat(last.errorCode().map(FinishReason::knownEnum)) - .hasValue(FinishReason.Known.CONTINUATION); - } - - @Test - public void generateContent_emptyToken_returnsOutput() { - api.respond(paused("", Part.fromText("a"))); - - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); - - assertThat(api.requests()).hasSize(1); - assertThat(firstText(response)).isEqualTo("a"); - } - - @Test - public void generateContent_repeatedToken_stopsResuming() { - api.respond(paused("state", Part.fromText("a")), paused("state", Part.fromText("b"))); - - LlmResponse response = getOnlyElement(generate(CONFIG, /* stream= */ false)); - + assertThat(texts(getLast(responses))).containsExactly("Done."); assertThat(api.requests()).hasSize(2); - assertThat(firstText(response)).isEqualTo("ab"); } @Test - public void generateContent_streamRepeatedToken_endsWithContinuation() { - api.respondStream(paused("state", Part.fromText("a"))); - api.respondStream(paused("state", Part.fromText("b"))); - - List responses = generate(CONFIG, /* stream= */ true); + public void generateContent_maxTokensWithToken_isNotResumed() { + api.respond(response("MAX_TOKENS", "state", usage(10, 5, 15), Part.fromText("The answer is"))); - assertThat(api.requests()).hasSize(2); - LlmResponse last = getLast(responses); - assertThat(firstText(last)).isEqualTo("ab"); - assertThat(last.errorCode().map(FinishReason::knownEnum)) - .hasValue(FinishReason.Known.CONTINUATION); - } - - @Test - public void generateContent_pausedOnEveryRequest_stopsAfterMaxResumesOnSmallStack() - throws InterruptedException { - for (int i = 0; i <= MAX_RESUMES; i++) { - api.respond(paused("token" + i, Part.fromText("a"))); - } - - LlmResponse response = getOnlyElement(generateOnSmallStack(/* stream= */ false)); - - assertThat(api.requests()).hasSize(MAX_RESUMES + 1); - assertThat(firstText(response)).isEqualTo("a".repeat(MAX_RESUMES + 1)); - assertThat(response.finishReason().map(FinishReason::knownEnum)) - .hasValue(FinishReason.Known.CONTINUATION); - } - - @Test - public void generateContent_streamPausedOnEveryRequest_stopsAfterMaxResumesOnSmallStack() - throws InterruptedException { - for (int i = 0; i <= MAX_RESUMES; i++) { - api.respondStream(paused("token" + i, Part.fromText("a"))); - } - - List responses = generateOnSmallStack(/* stream= */ true); - - assertThat(api.requests()).hasSize(MAX_RESUMES + 1); - LlmResponse last = getLast(responses); - assertThat(firstText(last)).isEqualTo("a".repeat(MAX_RESUMES + 1)); - assertThat(last.errorCode().map(FinishReason::knownEnum)) - .hasValue(FinishReason.Known.CONTINUATION); - } - - @Test - public void generateContent_cancelledBeforeResume_sendsNoFurtherRequest() { - api.respond(paused("state", Part.fromText("a")), finished(Part.fromText("b"))); - TestSubscriber subscriber = new TestSubscriber<>(); - api.onRequest(subscriber::cancel); - - gemini.generateContent(request(CONFIG), /* stream= */ false).subscribe(subscriber); + LlmResponse response = getOnlyElement(generate(/* stream= */ false)); + assertThat(finishReason(response)).isEqualTo(FinishReason.Known.MAX_TOKENS); assertThat(api.requests()).hasSize(1); - subscriber.assertNoValues(); } @Test - public void generateContent_streamCancelledBeforeResume_sendsNoFurtherRequest() { - api.respondStream(paused("state", Part.fromText("a"))); - api.respondStream(finished(Part.fromText("b"))); - - List responses = - gemini.generateContent(request(CONFIG), /* stream= */ true).take(1).toList().blockingGet(); - - assertThat(responses).hasSize(1); - assertThat(api.requests()).hasSize(1); - } - - @Test - public void generateContent_subscribedTwice_resumesEachTime() { + public void generateContent_pausedGenerationWithCache_isResumedWithTheCache() { + LlmRequest request = requestThatCreatesCache(); api.respond( - paused("state", Part.fromText("a")), - finished(Part.fromText("b")), - finished(Part.fromText("b"))); - Flowable responses = gemini.generateContent(request(CONFIG), /* stream= */ false); + response("CONTINUATION", "state", usage(10, 5, 15), Part.fromText("The answer is")), + response("STOP", /* token= */ null, usage(15, 3, 18), Part.fromText(" 42."))); - LlmResponse first = responses.blockingSingle(); - LlmResponse second = responses.blockingSingle(); + LlmResponse response = + getOnlyElement(gemini.generateContent(request, /* stream= */ false).toList().blockingGet()); - assertThat(firstText(first)).isEqualTo("ab"); - assertThat(firstText(second)).isEqualTo("ab"); + assertResumedWithCache(api.requests()); + assertThat(response.cacheMetadata().flatMap(CacheMetadata::cacheName)).hasValue(CACHE_NAME); + assertUsage(response, /* prompt= */ 10 + 15, /* candidates= */ 5 + 3, /* total= */ 15 + 18); } @Test - public void generateContent_resumeFails_returnsError() { - api.respond(paused("state", Part.fromText("a"))); - api.respondWithError(400); - - TestSubscriber subscriber = - gemini.generateContent(request(CONFIG), /* stream= */ false).test(); - - assertThat(api.requests()).hasSize(2); - subscriber.assertNoValues(); - subscriber.assertError(error -> error.getCause() instanceof ClientException); - } - - @Test - public void generateContent_streamCompleteGeneration_sendsOneRequest() { - GenerateContentResponseUsageMetadata usage = - GenerateContentResponseUsageMetadata.builder().totalTokenCount(7).build(); + public void generateContent_pausedStreamWithCache_isResumedWithTheCache() { + LlmRequest request = requestThatCreatesCache(); api.respondStream( - chunk(Part.fromText("Hello")), withUsage(finished(Part.fromText(" world")), usage)); + response( + /* finishReason= */ null, /* token= */ null, /* usage= */ null, Part.fromText("The")), + response("CONTINUATION", "state", usage(10, 5, 15), Part.fromText(" answer is"))); + api.respondStream(response("STOP", /* token= */ null, usage(15, 3, 18), Part.fromText(" 42."))); - List responses = generate(CONFIG, /* stream= */ true); + List responses = + gemini.generateContent(request, /* stream= */ true).toList().blockingGet(); - assertThat(api.requests()).hasSize(1); - assertThat(api.requests().get(0).has("continuationToken")).isFalse(); + assertResumedWithCache(api.requests()); LlmResponse last = getLast(responses); - assertThat(firstText(last)).isEqualTo("Hello world"); - assertThat(last.usageMetadata()).hasValue(usage); - } - - private List generate(GenerateContentConfig config, boolean stream) { - return gemini.generateContent(request(config), stream).toList().blockingGet(); - } - - /** Generates on a thread with a 256 KB stack, which overflows if every resume adds frames. */ - private List generateOnSmallStack(boolean stream) throws InterruptedException { - AtomicReference> responses = new AtomicReference<>(); - AtomicReference failure = new AtomicReference<>(); - Thread thread = - new Thread( - /* group= */ null, - () -> responses.set(generate(CONFIG, stream)), - "small-stack", - /* stackSize= */ 256 * 1024); - thread.setUncaughtExceptionHandler((unused, error) -> failure.set(error)); - thread.start(); - thread.join(); - assertThat(failure.get()).isNull(); - return responses.get(); - } - - private static LlmRequest request(GenerateContentConfig config) { - return LlmRequest.builder().contents(ImmutableList.of(QUESTION)).config(config).build(); - } - - private static GenerateContentResponse paused(String token, Part... parts) { - return response("CONTINUATION", token.getBytes(UTF_8), parts); - } - - private static GenerateContentResponse finished(Part... parts) { - return response("STOP", /* token= */ null, parts); + assertThat(responses.stream().filter(response -> response.cacheMetadata().isPresent())) + .containsExactly(last); + assertThat(last.cacheMetadata().flatMap(CacheMetadata::cacheName)).hasValue(CACHE_NAME); + assertUsage(last, /* prompt= */ 10 + 15, /* candidates= */ 5 + 3, /* total= */ 15 + 18); + } + + private List generate(boolean stream) { + LlmRequest request = LlmRequest.builder().contents(ImmutableList.of(QUESTION)).build(); + return gemini.generateContent(request, stream).toList().blockingGet(); + } + + /** + * Returns a request whose next call creates a cache: a first call records the fingerprint of the + * cacheable contents, and the fake answers the cache creation with {@link #CACHE_NAME}. + */ + private LlmRequest requestThatCreatesCache() { + LlmRequest request = + LlmRequest.builder() + .contents(ImmutableList.of(QUESTION, EARLIER_ANSWER, FOLLOW_UP)) + .cacheConfig(new ContextCacheConfig()) + .build(); + api.respond(response("STOP", /* token= */ null, usage(5000, 1, 5001), Part.fromText("OK"))); + CacheMetadata fingerprint = + getOnlyElement(gemini.generateContent(request, /* stream= */ false).toList().blockingGet()) + .cacheMetadata() + .orElseThrow(); + api.respondWithCache(CACHE_NAME); + return request.toBuilder().cacheMetadata(fingerprint).cacheableContentsTokenCount(5000).build(); } - private static GenerateContentResponse chunk(Part... parts) { - return response(/* finishReason= */ null, /* token= */ null, parts); + /** + * Asserts that, after the fingerprint call and the cache creation, the paused generation and its + * resume both name the cache and send the same contents. + */ + private static void assertResumedWithCache(ImmutableList requests) { + assertThat(requests).hasSize(4); + JsonNode paused = requests.get(2); + JsonNode resume = requests.get(3); + assertThat(paused.path("cachedContent").asText()).isEqualTo(CACHE_NAME); + assertThat(resume.path("cachedContent").asText()).isEqualTo(CACHE_NAME); + assertThat(resume.get("contents")).isEqualTo(paused.get("contents")); + assertThat(resume.path("continuationToken").asText()).isEqualTo(base64("state")); } private static GenerateContentResponse response( - @Nullable String finishReason, byte @Nullable [] token, Part... parts) { - Candidate.Builder candidate = Candidate.builder(); - if (parts.length > 0) { - candidate.content(Content.builder().role("model").parts(parts).build()); - } + @Nullable String finishReason, + @Nullable String token, + @Nullable GenerateContentResponseUsageMetadata usage, + Part... parts) { + Candidate.Builder candidate = + Candidate.builder().content(Content.builder().role("model").parts(parts).build()); if (finishReason != null) { candidate.finishReason(finishReason); } if (token != null) { - candidate.continuationToken(token); + candidate.continuationToken(token.getBytes(UTF_8)); } - return GenerateContentResponse.builder() - .candidates(ImmutableList.of(candidate.build())) - .build(); + GenerateContentResponse.Builder response = + GenerateContentResponse.builder().candidates(ImmutableList.of(candidate.build())); + if (usage != null) { + response.usageMetadata(usage); + } + return response.build(); } - private static GenerateContentResponse withUsage( - GenerateContentResponse response, GenerateContentResponseUsageMetadata usage) { - return response.toBuilder().usageMetadata(usage).build(); + private static GenerateContentResponseUsageMetadata usage( + int promptTokens, int candidatesTokens, int totalTokens) { + return GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(promptTokens) + .candidatesTokenCount(candidatesTokens) + .totalTokenCount(totalTokens) + .build(); } - private static ModalityTokenCount tokens(MediaModality.@Nullable Known modality, int tokenCount) { - ModalityTokenCount.Builder count = ModalityTokenCount.builder().tokenCount(tokenCount); - if (modality != null) { - count.modality(modality); - } - return count.build(); + private static void assertUsage(LlmResponse response, int prompt, int candidates, int total) { + GenerateContentResponseUsageMetadata usage = response.usageMetadata().orElseThrow(); + assertThat(usage.promptTokenCount()).hasValue(prompt); + assertThat(usage.candidatesTokenCount()).hasValue(candidates); + assertThat(usage.totalTokenCount()).hasValue(total); } - private static Content modelText(String text) { - return Content.builder().role("model").parts(Part.fromText(text)).build(); + private static ImmutableList texts(LlmResponse response) { + return response.content().flatMap(Content::parts).orElse(ImmutableList.of()).stream() + .map(part -> part.text().orElse("")) + .collect(toImmutableList()); } - private static String firstText(LlmResponse response) { - return response.content().flatMap(Content::parts).get().get(0).text().get(); + private static FinishReason.@Nullable Known finishReason(LlmResponse response) { + return response.finishReason().map(FinishReason::knownEnum).orElse(null); } private static String base64(String token) { return Base64.getEncoder().encodeToString(token.getBytes(UTF_8)); } - private static ImmutableList contents(JsonNode request) { - ImmutableList.Builder contents = ImmutableList.builder(); - request.get("contents").forEach(content -> contents.add(Content.fromJson(content.toString()))); - return contents.build(); - } - /** A fake Gemini API that records each request and answers it with the next canned response. */ - private static final class FakeGeminiApi implements Interceptor { - private final Deque responses = new ArrayDeque<>(); + private static final class ScriptedGeminiApi implements Interceptor { + private final Deque responses = new ArrayDeque<>(); private final List requests = new ArrayList<>(); - private Runnable onRequest = () -> {}; Client client() { OkHttpClient httpClient = new OkHttpClient.Builder().addInterceptor(this).build(); @@ -648,28 +289,21 @@ Client client() { synchronized void respond(GenerateContentResponse... responses) { for (GenerateContentResponse response : responses) { - this.responses.add(new CannedResponse(200, response.toJson())); + this.responses.add(response.toJson()); } } /** Answers the next request with a server-sent event stream of {@code chunks}. */ synchronized void respondStream(GenerateContentResponse... chunks) { responses.add( - new CannedResponse( - 200, - Arrays.stream(chunks) - .map(chunk -> "data: " + chunk.toJson() + "\n\n") - .collect(joining()))); + Arrays.stream(chunks) + .map(chunk -> "data: " + chunk.toJson() + "\n\n") + .collect(joining())); } - synchronized void respondWithError(int code) { - responses.add( - new CannedResponse(code, "{\"error\": {\"code\": " + code + ", \"message\": \"fake\"}}")); - } - - /** Runs {@code action} when a request arrives, before answering it. */ - synchronized void onRequest(Runnable action) { - onRequest = action; + /** Answers the next request, which creates a cache, with a cache named {@code name}. */ + synchronized void respondWithCache(String name) { + responses.add("{\"name\":\"" + name + "\",\"expireTime\":\"2099-01-01T00:00:00Z\"}"); } synchronized ImmutableList requests() { @@ -682,17 +316,13 @@ public synchronized Response intercept(Chain chain) throws IOException { Buffer body = new Buffer(); request.body().writeTo(body); requests.add(objectMapper.readTree(body.readUtf8())); - onRequest.run(); - CannedResponse response = responses.remove(); return new Response.Builder() .request(request) .protocol(Protocol.HTTP_1_1) - .code(response.code()) + .code(200) .message("fake") - .body(ResponseBody.create(response.body(), MediaType.get("application/json"))) + .body(ResponseBody.create(responses.remove(), MediaType.get("application/json"))) .build(); } - - private record CannedResponse(int code, String body) {} } } diff --git a/pom.xml b/pom.xml index 4e076d448..e18a91c75 100644 --- a/pom.xml +++ b/pom.xml @@ -56,7 +56,7 @@ 1.51.0 2.0.0 2.47.0 - 1.75.0 + 1.77.0 4.33.5 5.11.4 5.23.0