Skip to content

feat(core): expose Gemini retry options #1413

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

Open
014-code wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(core): expose Gemini retry options #1413

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
121 changes: 85 additions & 36 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 @@ -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;
Expand Down Expand Up @@ -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
Expand All @@ -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));
}

/**
Expand All @@ -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}.
*/
Expand All @@ -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.
*
Expand Down Expand Up @@ -236,6 +282,21 @@ public Builder httpExecutorService(ExecutorService httpExecutorService) {
return this;
}

/**
* Sets options for retrying transient HTTP failures through the Gemini client.
*
* <p>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.
*
Expand All @@ -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);
}
}

Expand Down
180 changes: 180 additions & 0 deletions core/src/test/java/com/google/adk/models/GeminiBuilderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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<Client> 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> httpOptions = ArgumentCaptor.forClass(HttpOptions.class);
verify(clientBuilder).httpOptions(httpOptions.capture());
return httpOptions.getValue();
}
}
Loading
Loading