Skip to content

Commit 4c13333

Browse files
014-code014-code
committed
feat(core): expose Gemini retry options
1 parent 84bae34 commit 4c13333

3 files changed

Lines changed: 291 additions & 36 deletions

File tree

‎core/src/main/java/com/google/adk/models/Gemini.java‎

Lines changed: 85 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import com.google.genai.types.GenerateContentConfig;
3232
import com.google.genai.types.GenerateContentResponse;
3333
import com.google.genai.types.HttpOptions;
34+
import com.google.genai.types.HttpRetryOptions;
3435
import com.google.genai.types.LiveConnectConfig;
3536
import com.google.genai.types.Part;
3637
import com.google.genai.types.PartialArg;
@@ -72,35 +73,69 @@ private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executor
7273
}
7374

7475
// Every client ADK builds starts here, so each sends ADK's tracking tag.
75-
private static Client.Builder adkClientBuilder(@Nullable ExecutorService executorService) {
76+
private static Client.Builder adkClientBuilder(
77+
@Nullable ExecutorService executorService, @Nullable HttpRetryOptions retryOptions) {
7678
return Client.builder()
77-
.httpOptions(HttpOptions.builder().headers(TrackingHeaders.GENAI_SDK_HEADERS).build())
79+
.httpOptions(buildHttpOptions(retryOptions))
7880
.clientOptions(
7981
ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build());
8082
}
8183

8284
private static Client buildApiKeyClient(
83-
String apiKey, @Nullable ExecutorService executorService) {
84-
return adkClientBuilder(executorService).apiKey(apiKey).build();
85+
String apiKey,
86+
@Nullable ExecutorService executorService,
87+
@Nullable HttpRetryOptions retryOptions) {
88+
return adkClientBuilder(executorService, retryOptions).apiKey(apiKey).build();
8589
}
8690

8791
private static Client buildVertexClient(
88-
VertexCredentials vertexCredentials, @Nullable ExecutorService executorService) {
89-
Client.Builder apiClientBuilder = adkClientBuilder(executorService);
92+
VertexCredentials vertexCredentials,
93+
@Nullable ExecutorService executorService,
94+
@Nullable HttpRetryOptions retryOptions) {
95+
Client.Builder apiClientBuilder = adkClientBuilder(executorService, retryOptions);
9096
vertexCredentials.project().ifPresent(apiClientBuilder::project);
9197
vertexCredentials.location().ifPresent(apiClientBuilder::location);
9298
vertexCredentials.credentials().ifPresent(apiClientBuilder::credentials);
9399
return apiClientBuilder.build();
94100
}
95101

96-
private static Client buildDefaultClient(@Nullable ExecutorService executorService) {
97-
return adkClientBuilder(executorService).build();
102+
private static Client buildDefaultClient(
103+
@Nullable ExecutorService executorService, @Nullable HttpRetryOptions retryOptions) {
104+
return adkClientBuilder(executorService, retryOptions).build();
105+
}
106+
107+
@VisibleForTesting
108+
static HttpOptions buildHttpOptions(@Nullable HttpRetryOptions retryOptions) {
109+
HttpOptions.Builder httpOptionsBuilder =
110+
HttpOptions.builder().headers(TrackingHeaders.GENAI_SDK_HEADERS);
111+
if (retryOptions != null) {
112+
httpOptionsBuilder.retryOptions(retryOptions);
113+
}
114+
return httpOptionsBuilder.build();
98115
}
99116

100117
private final Client apiClient;
101118
// True when apiClient already sends ADK's tracking tag.
102119
private final boolean clientSendsTrackingHeaders;
103120

121+
private Gemini(Builder builder) {
122+
super(builder.modelName);
123+
if (builder.apiClient != null) {
124+
this.apiClient = builder.apiClient;
125+
} else if (builder.apiKey != null) {
126+
this.apiClient =
127+
buildApiKeyClient(builder.apiKey, builder.httpExecutorService, builder.retryOptions);
128+
} else if (builder.vertexCredentials != null) {
129+
this.apiClient =
130+
buildVertexClient(
131+
builder.vertexCredentials, builder.httpExecutorService, builder.retryOptions);
132+
} else {
133+
this.apiClient = buildDefaultClient(builder.httpExecutorService, builder.retryOptions);
134+
}
135+
this.clientSendsTrackingHeaders =
136+
builder.apiClient == null || builder.clientSendsTrackingHeaders;
137+
}
138+
104139
/**
105140
* Constructs a new Gemini instance. Each generate request sent through {@code apiClient} carries
106141
* 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
111146
* @param apiClient The genai {@link com.google.genai.Client} instance for making API calls.
112147
*/
113148
public Gemini(String modelName, Client apiClient) {
114-
this(modelName, apiClient, /* clientSendsTrackingHeaders= */ false);
149+
this(
150+
builder()
151+
.modelName(modelName)
152+
.apiClient(Objects.requireNonNull(apiClient, "apiClient cannot be null")));
115153
}
116154

117155
Gemini(String modelName, Client apiClient, boolean clientSendsTrackingHeaders) {
118-
super(modelName);
119-
this.apiClient = Objects.requireNonNull(apiClient, "apiClient cannot be null");
120-
this.clientSendsTrackingHeaders = clientSendsTrackingHeaders;
156+
this(
157+
builder()
158+
.modelName(modelName)
159+
.apiClient(Objects.requireNonNull(apiClient, "apiClient cannot be null"))
160+
.clientSendsTrackingHeaders(clientSendsTrackingHeaders));
121161
}
122162

123163
/**
@@ -128,29 +168,28 @@ public Gemini(String modelName, Client apiClient) {
128168
*/
129169
public Gemini(String modelName, String apiKey) {
130170
this(
131-
modelName,
132-
buildApiKeyClient(Objects.requireNonNull(apiKey, "apiKey cannot be null"), null),
133-
/* clientSendsTrackingHeaders= */ true);
171+
builder()
172+
.modelName(modelName)
173+
.apiKey(Objects.requireNonNull(apiKey, "apiKey cannot be null")));
134174
}
135175

136176
/**
137-
* Constructs a new Gemini instance with a Google Gemini API key.
177+
* Constructs a new Gemini instance with Vertex AI credentials.
138178
*
139179
* @param modelName The name of the Gemini model to use (e.g., "gemini-2.0-flash").
140180
* @param vertexCredentials The Vertex AI credentials to access the Gemini model.
141181
*/
142182
public Gemini(String modelName, VertexCredentials vertexCredentials) {
143183
this(
144-
modelName,
145-
buildVertexClient(
146-
Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null"), null),
147-
/* clientSendsTrackingHeaders= */ true);
184+
builder()
185+
.modelName(modelName)
186+
.vertexCredentials(
187+
Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null")));
148188
}
149189

150190
/**
151-
* Returns a new Builder instance for constructing Gemini objects. Note that when building a
152-
* Gemini object, at least one of apiKey, vertexCredentials, or an explicit apiClient must be set.
153-
* If multiple are set, the explicit apiClient will take precedence.
191+
* Returns a builder for configuring Gemini. An explicit client takes precedence over an API key
192+
* or Vertex credentials. If none are supplied, the client uses its default configuration.
154193
*
155194
* @return A new {@link Builder}.
156195
*/
@@ -165,9 +204,16 @@ public static class Builder {
165204
private String apiKey;
166205
private VertexCredentials vertexCredentials;
167206
private ExecutorService httpExecutorService;
207+
private HttpRetryOptions retryOptions;
208+
private boolean clientSendsTrackingHeaders;
168209

169210
private Builder() {}
170211

212+
private Builder clientSendsTrackingHeaders(boolean clientSendsTrackingHeaders) {
213+
this.clientSendsTrackingHeaders = clientSendsTrackingHeaders;
214+
return this;
215+
}
216+
171217
/**
172218
* Sets the name of the Gemini model to use.
173219
*
@@ -236,6 +282,21 @@ public Builder httpExecutorService(ExecutorService httpExecutorService) {
236282
return this;
237283
}
238284

285+
/**
286+
* Sets options for retrying transient HTTP failures through the Gemini client.
287+
*
288+
* <p>If an explicit {@link #apiClient(Client)} is supplied, its existing HTTP options take
289+
* precedence and these retry options are ignored.
290+
*
291+
* @param retryOptions The retry options to pass to the Google Gen AI client.
292+
* @return This builder.
293+
*/
294+
@CanIgnoreReturnValue
295+
public Builder retryOptions(HttpRetryOptions retryOptions) {
296+
this.retryOptions = retryOptions;
297+
return this;
298+
}
299+
239300
/**
240301
* Builds the {@link Gemini} instance.
241302
*
@@ -244,19 +305,7 @@ public Builder httpExecutorService(ExecutorService httpExecutorService) {
244305
*/
245306
public Gemini build() {
246307
Objects.requireNonNull(modelName, "modelName must be set.");
247-
248-
if (apiClient != null) {
249-
return new Gemini(modelName, apiClient);
250-
}
251-
Client client;
252-
if (apiKey != null) {
253-
client = buildApiKeyClient(apiKey, httpExecutorService);
254-
} else if (vertexCredentials != null) {
255-
client = buildVertexClient(vertexCredentials, httpExecutorService);
256-
} else {
257-
client = buildDefaultClient(httpExecutorService);
258-
}
259-
return new Gemini(modelName, client, /* clientSendsTrackingHeaders= */ true);
308+
return new Gemini(this);
260309
}
261310
}
262311

‎core/src/test/java/com/google/adk/models/GeminiBuilderTest.java‎

Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
1+
/*
2+
* Copyright 2026 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.adk.models;
17+
18+
import static com.google.common.truth.Truth.assertThat;
19+
import static org.junit.Assert.assertThrows;
20+
import static org.mockito.Mockito.RETURNS_SELF;
21+
import static org.mockito.Mockito.mock;
22+
import static org.mockito.Mockito.mockStatic;
23+
import static org.mockito.Mockito.never;
24+
import static org.mockito.Mockito.verify;
25+
import static org.mockito.Mockito.verifyNoInteractions;
26+
import static org.mockito.Mockito.when;
27+
28+
import com.google.genai.Client;
29+
import com.google.genai.types.HttpOptions;
30+
import com.google.genai.types.HttpRetryOptions;
31+
import org.junit.After;
32+
import org.junit.Before;
33+
import org.junit.Test;
34+
import org.junit.runner.RunWith;
35+
import org.junit.runners.JUnit4;
36+
import org.mockito.ArgumentCaptor;
37+
import org.mockito.MockedStatic;
38+
39+
/** Tests the public Gemini construction paths without contacting a provider. */
40+
@RunWith(JUnit4.class)
41+
public final class GeminiBuilderTest {
42+
43+
private static final String MODEL = "gemini-test-model";
44+
private static final HttpRetryOptions RETRY_OPTIONS =
45+
HttpRetryOptions.builder().attempts(3).httpStatusCodes(429, 503).build();
46+
private static final VertexCredentials VERTEX_CREDENTIALS =
47+
VertexCredentials.builder().project("test-project").location("us-central1").build();
48+
49+
private Client client;
50+
private Client.Builder clientBuilder;
51+
private MockedStatic<Client> clientFactory;
52+
53+
@Before
54+
public void setUp() {
55+
client = mock(Client.class);
56+
clientBuilder = mock(Client.Builder.class, RETURNS_SELF);
57+
when(clientBuilder.build()).thenReturn(client);
58+
clientFactory = mockStatic(Client.class);
59+
clientFactory.when(Client::builder).thenReturn(clientBuilder);
60+
}
61+
62+
@After
63+
public void tearDown() {
64+
clientFactory.close();
65+
}
66+
67+
@Test
68+
public void builder_withApiKey_passesRetryOptionsAndPrefersApiKeyOverVertexCredentials() {
69+
Gemini gemini =
70+
Gemini.builder()
71+
.modelName(MODEL)
72+
.apiKey("test-key")
73+
.vertexCredentials(VERTEX_CREDENTIALS)
74+
.retryOptions(RETRY_OPTIONS)
75+
.build();
76+
77+
assertThat(gemini.model()).isEqualTo(MODEL);
78+
verify(clientBuilder).apiKey("test-key");
79+
verify(clientBuilder, never()).project("test-project");
80+
verify(clientBuilder, never()).location("us-central1");
81+
assertThat(capturedHttpOptions().retryOptions()).hasValue(RETRY_OPTIONS);
82+
}
83+
84+
@Test
85+
public void builder_withVertexCredentials_passesRetryOptionsAndCredentials() {
86+
Gemini gemini =
87+
Gemini.builder()
88+
.modelName(MODEL)
89+
.vertexCredentials(VERTEX_CREDENTIALS)
90+
.retryOptions(RETRY_OPTIONS)
91+
.build();
92+
93+
assertThat(gemini.model()).isEqualTo(MODEL);
94+
verify(clientBuilder).project("test-project");
95+
verify(clientBuilder).location("us-central1");
96+
assertThat(capturedHttpOptions().retryOptions()).hasValue(RETRY_OPTIONS);
97+
}
98+
99+
@Test
100+
public void builder_withDefaultClient_passesRetryOptions() {
101+
Gemini gemini = Gemini.builder().modelName(MODEL).retryOptions(RETRY_OPTIONS).build();
102+
103+
assertThat(gemini.model()).isEqualTo(MODEL);
104+
assertThat(capturedHttpOptions().retryOptions()).hasValue(RETRY_OPTIONS);
105+
}
106+
107+
@Test
108+
public void builder_withoutRetryOptions_leavesSdkRetryOptionsUnset() {
109+
Gemini.builder().modelName(MODEL).apiKey("test-key").build();
110+
111+
HttpOptions httpOptions = capturedHttpOptions();
112+
assertThat(httpOptions.retryOptions()).isEmpty();
113+
assertThat(httpOptions.headers().get()).containsKey("x-goog-api-client");
114+
assertThat(httpOptions.headers().get()).containsKey("user-agent");
115+
}
116+
117+
@Test
118+
public void builder_withExplicitClient_doesNotCreateOrReconfigureClient() {
119+
Gemini gemini =
120+
Gemini.builder()
121+
.modelName(MODEL)
122+
.apiClient(client)
123+
.apiKey("test-key")
124+
.vertexCredentials(VERTEX_CREDENTIALS)
125+
.retryOptions(RETRY_OPTIONS)
126+
.build();
127+
128+
assertThat(gemini.model()).isEqualTo(MODEL);
129+
clientFactory.verifyNoInteractions();
130+
verifyNoInteractions(client);
131+
}
132+
133+
@Test
134+
public void constructor_withApiKey_preservesClientConfiguration() {
135+
Gemini gemini = new Gemini(MODEL, "test-key");
136+
137+
assertThat(gemini.model()).isEqualTo(MODEL);
138+
verify(clientBuilder).apiKey("test-key");
139+
assertThat(capturedHttpOptions().retryOptions()).isEmpty();
140+
}
141+
142+
@Test
143+
public void constructor_withVertexCredentials_preservesClientConfiguration() {
144+
Gemini gemini = new Gemini(MODEL, VERTEX_CREDENTIALS);
145+
146+
assertThat(gemini.model()).isEqualTo(MODEL);
147+
verify(clientBuilder).project("test-project");
148+
verify(clientBuilder).location("us-central1");
149+
assertThat(capturedHttpOptions().retryOptions()).isEmpty();
150+
}
151+
152+
@Test
153+
public void constructor_withExplicitClient_doesNotCreateAnotherClient() {
154+
Gemini gemini = new Gemini(MODEL, client);
155+
156+
assertThat(gemini.model()).isEqualTo(MODEL);
157+
clientFactory.verifyNoInteractions();
158+
verifyNoInteractions(client);
159+
}
160+
161+
@Test
162+
public void constructors_withNullArguments_rejectInsteadOfUsingDefaultClient() {
163+
assertThrows(NullPointerException.class, () -> new Gemini(MODEL, (Client) null));
164+
assertThrows(NullPointerException.class, () -> new Gemini(MODEL, (String) null));
165+
assertThrows(NullPointerException.class, () -> new Gemini(MODEL, (VertexCredentials) null));
166+
clientFactory.verifyNoInteractions();
167+
}
168+
169+
@Test
170+
public void builder_withoutModelName_rejectsBeforeCreatingClient() {
171+
assertThrows(NullPointerException.class, () -> Gemini.builder().apiKey("test-key").build());
172+
clientFactory.verifyNoInteractions();
173+
}
174+
175+
private HttpOptions capturedHttpOptions() {
176+
ArgumentCaptor<HttpOptions> httpOptions = ArgumentCaptor.forClass(HttpOptions.class);
177+
verify(clientBuilder).httpOptions(httpOptions.capture());
178+
return httpOptions.getValue();
179+
}
180+
}

0 commit comments

Comments
 (0)