From 4c13333db16e0086b0837e7d7c860275fb617c7c Mon Sep 17 00:00:00 2001 From: 014-code <161032298+014-code@users.noreply.github.com> Date: Sat, 17 Oct 2026 12:29:34 +0000 Subject: [PATCH] feat(core): expose Gemini retry options --- .../java/com/google/adk/models/Gemini.java | 121 ++++++++---- .../google/adk/models/GeminiBuilderTest.java | 180 ++++++++++++++++++ .../com/google/adk/models/GeminiTest.java | 26 +++ 3 files changed, 291 insertions(+), 36 deletions(-) create mode 100644 core/src/test/java/com/google/adk/models/GeminiBuilderTest.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 7bd524500..4cfd26955 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -31,6 +31,7 @@ import com.google.genai.types.GenerateContentConfig; import com.google.genai.types.GenerateContentResponse; import com.google.genai.types.HttpOptions; +import com.google.genai.types.HttpRetryOptions; import com.google.genai.types.LiveConnectConfig; import com.google.genai.types.Part; import com.google.genai.types.PartialArg; @@ -72,35 +73,69 @@ private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executor } // Every client ADK builds starts here, so each sends ADK's tracking tag. - private static Client.Builder adkClientBuilder(@Nullable ExecutorService executorService) { + private static Client.Builder adkClientBuilder( + @Nullable ExecutorService executorService, @Nullable HttpRetryOptions retryOptions) { return Client.builder() - .httpOptions(HttpOptions.builder().headers(TrackingHeaders.GENAI_SDK_HEADERS).build()) + .httpOptions(buildHttpOptions(retryOptions)) .clientOptions( ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build()); } private static Client buildApiKeyClient( - String apiKey, @Nullable ExecutorService executorService) { - return adkClientBuilder(executorService).apiKey(apiKey).build(); + String apiKey, + @Nullable ExecutorService executorService, + @Nullable HttpRetryOptions retryOptions) { + return adkClientBuilder(executorService, retryOptions).apiKey(apiKey).build(); } private static Client buildVertexClient( - VertexCredentials vertexCredentials, @Nullable ExecutorService executorService) { - Client.Builder apiClientBuilder = adkClientBuilder(executorService); + VertexCredentials vertexCredentials, + @Nullable ExecutorService executorService, + @Nullable HttpRetryOptions retryOptions) { + Client.Builder apiClientBuilder = adkClientBuilder(executorService, retryOptions); vertexCredentials.project().ifPresent(apiClientBuilder::project); vertexCredentials.location().ifPresent(apiClientBuilder::location); vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials); return apiClientBuilder.build(); } - private static Client buildDefaultClient(@Nullable ExecutorService executorService) { - return adkClientBuilder(executorService).build(); + private static Client buildDefaultClient( + @Nullable ExecutorService executorService, @Nullable HttpRetryOptions retryOptions) { + return adkClientBuilder(executorService, retryOptions).build(); + } + + @VisibleForTesting + static HttpOptions buildHttpOptions(@Nullable HttpRetryOptions retryOptions) { + HttpOptions.Builder httpOptionsBuilder = + HttpOptions.builder().headers(TrackingHeaders.GENAI_SDK_HEADERS); + if (retryOptions != null) { + httpOptionsBuilder.retryOptions(retryOptions); + } + return httpOptionsBuilder.build(); } private final Client apiClient; // True when apiClient already sends ADK's tracking tag. private final boolean clientSendsTrackingHeaders; + private Gemini(Builder builder) { + super(builder.modelName); + if (builder.apiClient != null) { + this.apiClient = builder.apiClient; + } else if (builder.apiKey != null) { + this.apiClient = + buildApiKeyClient(builder.apiKey, builder.httpExecutorService, builder.retryOptions); + } else if (builder.vertexCredentials != null) { + this.apiClient = + buildVertexClient( + builder.vertexCredentials, builder.httpExecutorService, builder.retryOptions); + } else { + this.apiClient = buildDefaultClient(builder.httpExecutorService, builder.retryOptions); + } + this.clientSendsTrackingHeaders = + builder.apiClient == null || builder.clientSendsTrackingHeaders; + } + /** * Constructs a new Gemini instance. Each generate request sent through {@code apiClient} carries * ADK's tracking tag, which the genai SDK appends to the client's tracking headers. If the @@ -111,13 +146,18 @@ private static Client buildDefaultClient(@Nullable ExecutorService executorServi * @param apiClient The genai {@link com.google.genai.Client} instance for making API calls. */ public Gemini(String modelName, Client apiClient) { - this(modelName, apiClient, /* clientSendsTrackingHeaders= */ false); + this( + builder() + .modelName(modelName) + .apiClient(Objects.requireNonNull(apiClient, "apiClient cannot be null"))); } Gemini(String modelName, Client apiClient, boolean clientSendsTrackingHeaders) { - super(modelName); - this.apiClient = Objects.requireNonNull(apiClient, "apiClient cannot be null"); - this.clientSendsTrackingHeaders = clientSendsTrackingHeaders; + this( + builder() + .modelName(modelName) + .apiClient(Objects.requireNonNull(apiClient, "apiClient cannot be null")) + .clientSendsTrackingHeaders(clientSendsTrackingHeaders)); } /** @@ -128,29 +168,28 @@ public Gemini(String modelName, Client apiClient) { */ public Gemini(String modelName, String apiKey) { this( - modelName, - buildApiKeyClient(Objects.requireNonNull(apiKey, "apiKey cannot be null"), null), - /* clientSendsTrackingHeaders= */ true); + builder() + .modelName(modelName) + .apiKey(Objects.requireNonNull(apiKey, "apiKey cannot be null"))); } /** - * Constructs a new Gemini instance with a Google Gemini API key. + * Constructs a new Gemini instance with Vertex AI credentials. * * @param modelName The name of the Gemini model to use (e.g., "gemini-2.0-flash"). * @param vertexCredentials The Vertex AI credentials to access the Gemini model. */ public Gemini(String modelName, VertexCredentials vertexCredentials) { this( - modelName, - buildVertexClient( - Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null"), null), - /* clientSendsTrackingHeaders= */ true); + builder() + .modelName(modelName) + .vertexCredentials( + Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null"))); } /** - * Returns a new Builder instance for constructing Gemini objects. Note that when building a - * Gemini object, at least one of apiKey, vertexCredentials, or an explicit apiClient must be set. - * If multiple are set, the explicit apiClient will take precedence. + * Returns a builder for configuring Gemini. An explicit client takes precedence over an API key + * or Vertex credentials. If none are supplied, the client uses its default configuration. * * @return A new {@link Builder}. */ @@ -165,9 +204,16 @@ public static class Builder { private String apiKey; private VertexCredentials vertexCredentials; private ExecutorService httpExecutorService; + private HttpRetryOptions retryOptions; + private boolean clientSendsTrackingHeaders; private Builder() {} + private Builder clientSendsTrackingHeaders(boolean clientSendsTrackingHeaders) { + this.clientSendsTrackingHeaders = clientSendsTrackingHeaders; + return this; + } + /** * Sets the name of the Gemini model to use. * @@ -236,6 +282,21 @@ public Builder httpExecutorService(ExecutorService httpExecutorService) { return this; } + /** + * Sets options for retrying transient HTTP failures through the Gemini client. + * + *

If an explicit {@link #apiClient(Client)} is supplied, its existing HTTP options take + * precedence and these retry options are ignored. + * + * @param retryOptions The retry options to pass to the Google Gen AI client. + * @return This builder. + */ + @CanIgnoreReturnValue + public Builder retryOptions(HttpRetryOptions retryOptions) { + this.retryOptions = retryOptions; + return this; + } + /** * Builds the {@link Gemini} instance. * @@ -244,19 +305,7 @@ public Builder httpExecutorService(ExecutorService httpExecutorService) { */ public Gemini build() { Objects.requireNonNull(modelName, "modelName must be set."); - - if (apiClient != null) { - return new Gemini(modelName, apiClient); - } - Client client; - if (apiKey != null) { - client = buildApiKeyClient(apiKey, httpExecutorService); - } else if (vertexCredentials != null) { - client = buildVertexClient(vertexCredentials, httpExecutorService); - } else { - client = buildDefaultClient(httpExecutorService); - } - return new Gemini(modelName, client, /* clientSendsTrackingHeaders= */ true); + return new Gemini(this); } } diff --git a/core/src/test/java/com/google/adk/models/GeminiBuilderTest.java b/core/src/test/java/com/google/adk/models/GeminiBuilderTest.java new file mode 100644 index 000000000..31f28a422 --- /dev/null +++ b/core/src/test/java/com/google/adk/models/GeminiBuilderTest.java @@ -0,0 +1,180 @@ +/* + * 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 static com.google.common.truth.Truth.assertThat; +import static org.junit.Assert.assertThrows; +import static org.mockito.Mockito.RETURNS_SELF; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.verifyNoInteractions; +import static org.mockito.Mockito.when; + +import com.google.genai.Client; +import com.google.genai.types.HttpOptions; +import com.google.genai.types.HttpRetryOptions; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; +import org.mockito.ArgumentCaptor; +import org.mockito.MockedStatic; + +/** Tests the public Gemini construction paths without contacting a provider. */ +@RunWith(JUnit4.class) +public final class GeminiBuilderTest { + + private static final String MODEL = "gemini-test-model"; + private static final HttpRetryOptions RETRY_OPTIONS = + HttpRetryOptions.builder().attempts(3).httpStatusCodes(429, 503).build(); + private static final VertexCredentials VERTEX_CREDENTIALS = + VertexCredentials.builder().project("test-project").location("us-central1").build(); + + private Client client; + private Client.Builder clientBuilder; + private MockedStatic clientFactory; + + @Before + public void setUp() { + client = mock(Client.class); + clientBuilder = mock(Client.Builder.class, RETURNS_SELF); + when(clientBuilder.build()).thenReturn(client); + clientFactory = mockStatic(Client.class); + clientFactory.when(Client::builder).thenReturn(clientBuilder); + } + + @After + public void tearDown() { + clientFactory.close(); + } + + @Test + public void builder_withApiKey_passesRetryOptionsAndPrefersApiKeyOverVertexCredentials() { + Gemini gemini = + Gemini.builder() + .modelName(MODEL) + .apiKey("test-key") + .vertexCredentials(VERTEX_CREDENTIALS) + .retryOptions(RETRY_OPTIONS) + .build(); + + assertThat(gemini.model()).isEqualTo(MODEL); + verify(clientBuilder).apiKey("test-key"); + verify(clientBuilder, never()).project("test-project"); + verify(clientBuilder, never()).location("us-central1"); + assertThat(capturedHttpOptions().retryOptions()).hasValue(RETRY_OPTIONS); + } + + @Test + public void builder_withVertexCredentials_passesRetryOptionsAndCredentials() { + Gemini gemini = + Gemini.builder() + .modelName(MODEL) + .vertexCredentials(VERTEX_CREDENTIALS) + .retryOptions(RETRY_OPTIONS) + .build(); + + assertThat(gemini.model()).isEqualTo(MODEL); + verify(clientBuilder).project("test-project"); + verify(clientBuilder).location("us-central1"); + assertThat(capturedHttpOptions().retryOptions()).hasValue(RETRY_OPTIONS); + } + + @Test + public void builder_withDefaultClient_passesRetryOptions() { + Gemini gemini = Gemini.builder().modelName(MODEL).retryOptions(RETRY_OPTIONS).build(); + + assertThat(gemini.model()).isEqualTo(MODEL); + assertThat(capturedHttpOptions().retryOptions()).hasValue(RETRY_OPTIONS); + } + + @Test + public void builder_withoutRetryOptions_leavesSdkRetryOptionsUnset() { + Gemini.builder().modelName(MODEL).apiKey("test-key").build(); + + HttpOptions httpOptions = capturedHttpOptions(); + assertThat(httpOptions.retryOptions()).isEmpty(); + assertThat(httpOptions.headers().get()).containsKey("x-goog-api-client"); + assertThat(httpOptions.headers().get()).containsKey("user-agent"); + } + + @Test + public void builder_withExplicitClient_doesNotCreateOrReconfigureClient() { + Gemini gemini = + Gemini.builder() + .modelName(MODEL) + .apiClient(client) + .apiKey("test-key") + .vertexCredentials(VERTEX_CREDENTIALS) + .retryOptions(RETRY_OPTIONS) + .build(); + + assertThat(gemini.model()).isEqualTo(MODEL); + clientFactory.verifyNoInteractions(); + verifyNoInteractions(client); + } + + @Test + public void constructor_withApiKey_preservesClientConfiguration() { + Gemini gemini = new Gemini(MODEL, "test-key"); + + assertThat(gemini.model()).isEqualTo(MODEL); + verify(clientBuilder).apiKey("test-key"); + assertThat(capturedHttpOptions().retryOptions()).isEmpty(); + } + + @Test + public void constructor_withVertexCredentials_preservesClientConfiguration() { + Gemini gemini = new Gemini(MODEL, VERTEX_CREDENTIALS); + + assertThat(gemini.model()).isEqualTo(MODEL); + verify(clientBuilder).project("test-project"); + verify(clientBuilder).location("us-central1"); + assertThat(capturedHttpOptions().retryOptions()).isEmpty(); + } + + @Test + public void constructor_withExplicitClient_doesNotCreateAnotherClient() { + Gemini gemini = new Gemini(MODEL, client); + + assertThat(gemini.model()).isEqualTo(MODEL); + clientFactory.verifyNoInteractions(); + verifyNoInteractions(client); + } + + @Test + public void constructors_withNullArguments_rejectInsteadOfUsingDefaultClient() { + assertThrows(NullPointerException.class, () -> new Gemini(MODEL, (Client) null)); + assertThrows(NullPointerException.class, () -> new Gemini(MODEL, (String) null)); + assertThrows(NullPointerException.class, () -> new Gemini(MODEL, (VertexCredentials) null)); + clientFactory.verifyNoInteractions(); + } + + @Test + public void builder_withoutModelName_rejectsBeforeCreatingClient() { + assertThrows(NullPointerException.class, () -> Gemini.builder().apiKey("test-key").build()); + clientFactory.verifyNoInteractions(); + } + + private HttpOptions capturedHttpOptions() { + ArgumentCaptor httpOptions = ArgumentCaptor.forClass(HttpOptions.class); + verify(clientBuilder).httpOptions(httpOptions.capture()); + return httpOptions.getValue(); + } +} diff --git a/core/src/test/java/com/google/adk/models/GeminiTest.java b/core/src/test/java/com/google/adk/models/GeminiTest.java index a56628493..74a9c717f 100644 --- a/core/src/test/java/com/google/adk/models/GeminiTest.java +++ b/core/src/test/java/com/google/adk/models/GeminiTest.java @@ -28,6 +28,8 @@ import com.google.genai.types.FunctionCall; import com.google.genai.types.GenerateContentResponse; import com.google.genai.types.GenerateContentResponseUsageMetadata; +import com.google.genai.types.HttpOptions; +import com.google.genai.types.HttpRetryOptions; import com.google.genai.types.Part; import com.google.genai.types.PartialArg; import com.google.genai.types.ToolCall; @@ -42,6 +44,30 @@ @RunWith(JUnit4.class) public final class GeminiTest { + @Test + public void buildHttpOptions_withRetryOptions_setsRetryOptions() { + HttpRetryOptions retryOptions = + HttpRetryOptions.builder() + .attempts(3) + .initialDelay(1.0) + .maxDelay(8.0) + .expBase(2.0) + .jitter(0.2) + .httpStatusCodes(429, 500, 503) + .build(); + + HttpOptions httpOptions = Gemini.buildHttpOptions(retryOptions); + + assertThat(httpOptions.retryOptions()).hasValue(retryOptions); + } + + @Test + public void buildHttpOptions_withoutRetryOptions_leavesSdkRetryOptionsUnset() { + HttpOptions httpOptions = Gemini.buildHttpOptions(null); + + assertThat(httpOptions.retryOptions()).isEmpty(); + } + // Test cases for processRawResponses static method @Test public void processRawResponses_withTextChunks_emitsPartialResponses() {