Skip to content

refactor: use genai SDK continuation #1633

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Oct 14, 2026
Merged

refactor: use genai SDK continuation #1633

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
199 changes: 44 additions & 155 deletions core/src/main/java/com/google/adk/models/Gemini.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -68,10 +62,6 @@ public class Gemini extends BaseLlm {
private static final Logger logger = LoggerFactory.getLogger(Gemini.class);
private static final ImmutableMap<String, String> 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
Expand Down Expand Up @@ -297,31 +287,31 @@ private Flowable<LlmResponse> 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<Content> contents = llmRequest.contents();
if (stream) {
logger.debug("Sending streaming generateContent request to model {}", effectiveModelName);
CompletableFuture<ResponseStream<GenerateContentResponse>> 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<GenerateContentResponse> 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));
}
}
Expand All @@ -335,131 +325,8 @@ private static LlmResponse withCacheMetadata(

@VisibleForTesting
static Flowable<LlmResponse> processRawResponses(Flowable<GenerateContentResponse> 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<LlmResponse> generateResumingPauses(
String modelName,
GeminiContinuation continuation,
CompletableFuture<GenerateContentResponse> firstResponse) {
Deque<CompletableFuture<GenerateContentResponse>> 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<LlmResponse> completeOrResume(
String modelName,
GeminiContinuation continuation,
GenerateContentResponse response,
Deque<CompletableFuture<GenerateContentResponse>> responses) {
LlmResponse llmResponse = LlmResponse.create(response);
Optional<ResumeRequest> 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<LlmResponse> streamResumingPauses(
String modelName,
GeminiContinuation continuation,
CompletableFuture<ResponseStream<GenerateContentResponse>> firstStream,
@Nullable CacheMetadata cacheMetadata) {
Deque<CompletableFuture<ResponseStream<GenerateContentResponse>>> 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<GenerateContentResponse> streamRequest(
String modelName,
GeminiContinuation continuation,
Deque<CompletableFuture<ResponseStream<GenerateContentResponse>>> 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<CompletableFuture<ResponseStream<GenerateContentResponse>>> 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
Expand Down Expand Up @@ -487,7 +354,7 @@ protected CompletableFuture<GeminiLiveTransport> connectLiveTransport(
return apiClient.async.live.connect(modelName, config).thenApply(GenAiLiveTransport::new);
}

static final class StreamingResponseAggregator {
private static final class StreamingResponseAggregator {
private final List<Part> accumulatedSequence = new ArrayList<>();
private final StringBuilder currentTextBuffer = new StringBuilder();
// Always reassigned in accumulateParts() before it is read; the initializer is never observed.
Expand All @@ -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;
}
Expand All @@ -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<LlmResponse> process(Flowable<GenerateContentResponse> rawResponses) {
private Flowable<LlmResponse> process(
Flowable<GenerateContentResponse> rawResponses, @Nullable CacheMetadata cacheMetadata) {
return rawResponses
.concatMap(this::processRawResponse)
.concatWith(Flowable.defer(this::processFinalResponse));
.concatWith(
Flowable.defer(
() ->
processFinalResponse()
.map(response -> withCacheMetadata(response, cacheMetadata))));
}

/**
Expand Down Expand Up @@ -558,8 +431,10 @@ private Flowable<LlmResponse> 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());
Expand Down Expand Up @@ -846,6 +721,20 @@ private Flowable<LlmResponse> 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
Expand Down
Loading
Loading