Skip to content

Commit cfaf88f

Browse files
014-code014-code
committed
feat(core): expose Gemini retry options
1 parent 97f22a2 commit cfaf88f

3 files changed

Lines changed: 276 additions & 28 deletions

File tree

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

Lines changed: 70 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@
3434
import com.google.genai.types.GenerateContentConfig;
3535
import com.google.genai.types.GenerateContentResponse;
3636
import com.google.genai.types.HttpOptions;
37+
import com.google.genai.types.HttpRetryOptions;
3738
import com.google.genai.types.LiveConnectConfig;
3839
import com.google.genai.types.Part;
3940
import com.google.genai.types.PartialArg;
@@ -76,20 +77,24 @@ private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executor
7677
}
7778

7879
private static Client buildApiKeyClient(
79-
String apiKey, @Nullable ExecutorService executorService) {
80+
String apiKey,
81+
@Nullable ExecutorService executorService,
82+
@Nullable HttpRetryOptions retryOptions) {
8083
return Client.builder()
8184
.apiKey(apiKey)
82-
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
85+
.httpOptions(buildHttpOptions(retryOptions))
8386
.clientOptions(
8487
ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build())
8588
.build();
8689
}
8790

8891
private static Client buildVertexClient(
89-
VertexCredentials vertexCredentials, @Nullable ExecutorService executorService) {
92+
VertexCredentials vertexCredentials,
93+
@Nullable ExecutorService executorService,
94+
@Nullable HttpRetryOptions retryOptions) {
9095
Client.Builder apiClientBuilder =
9196
Client.builder()
92-
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
97+
.httpOptions(buildHttpOptions(retryOptions))
9398
.clientOptions(
9499
ClientOptions.builder()
95100
.customHttpClient(prepareHttpClient(executorService))
@@ -100,14 +105,24 @@ private static Client buildVertexClient(
100105
return apiClientBuilder.build();
101106
}
102107

103-
private static Client buildDefaultClient(@Nullable ExecutorService executorService) {
108+
private static Client buildDefaultClient(
109+
@Nullable ExecutorService executorService, @Nullable HttpRetryOptions retryOptions) {
104110
return Client.builder()
105-
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
111+
.httpOptions(buildHttpOptions(retryOptions))
106112
.clientOptions(
107113
ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build())
108114
.build();
109115
}
110116

117+
@VisibleForTesting
118+
static HttpOptions buildHttpOptions(@Nullable HttpRetryOptions retryOptions) {
119+
HttpOptions.Builder httpOptionsBuilder = HttpOptions.builder().headers(TRACKING_HEADERS);
120+
if (retryOptions != null) {
121+
httpOptionsBuilder.retryOptions(retryOptions);
122+
}
123+
return httpOptionsBuilder.build();
124+
}
125+
111126
static {
112127
String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION;
113128
String languageLabel = "gl-java/" + JAVA_VERSION.value();
@@ -121,15 +136,33 @@ private static Client buildDefaultClient(@Nullable ExecutorService executorServi
121136

122137
private final Client apiClient;
123138

139+
private Gemini(Builder builder) {
140+
super(builder.modelName);
141+
if (builder.apiClient != null) {
142+
this.apiClient = builder.apiClient;
143+
} else if (builder.apiKey != null) {
144+
this.apiClient =
145+
buildApiKeyClient(builder.apiKey, builder.httpExecutorService, builder.retryOptions);
146+
} else if (builder.vertexCredentials != null) {
147+
this.apiClient =
148+
buildVertexClient(
149+
builder.vertexCredentials, builder.httpExecutorService, builder.retryOptions);
150+
} else {
151+
this.apiClient = buildDefaultClient(builder.httpExecutorService, builder.retryOptions);
152+
}
153+
}
154+
124155
/**
125156
* Constructs a new Gemini instance.
126157
*
127158
* @param modelName The name of the Gemini model to use (e.g., "gemini-2.0-flash").
128159
* @param apiClient The genai {@link com.google.genai.Client} instance for making API calls.
129160
*/
130161
public Gemini(String modelName, Client apiClient) {
131-
super(modelName);
132-
this.apiClient = Objects.requireNonNull(apiClient, "apiClient cannot be null");
162+
this(
163+
builder()
164+
.modelName(modelName)
165+
.apiClient(Objects.requireNonNull(apiClient, "apiClient cannot be null")));
133166
}
134167

135168
/**
@@ -139,27 +172,29 @@ public Gemini(String modelName, Client apiClient) {
139172
* @param apiKey The Google Gemini API key.
140173
*/
141174
public Gemini(String modelName, String apiKey) {
142-
super(modelName);
143-
Objects.requireNonNull(apiKey, "apiKey cannot be null");
144-
this.apiClient = buildApiKeyClient(apiKey, null);
175+
this(
176+
builder()
177+
.modelName(modelName)
178+
.apiKey(Objects.requireNonNull(apiKey, "apiKey cannot be null")));
145179
}
146180

147181
/**
148-
* Constructs a new Gemini instance with a Google Gemini API key.
182+
* Constructs a new Gemini instance with Vertex AI credentials.
149183
*
150184
* @param modelName The name of the Gemini model to use (e.g., "gemini-2.0-flash").
151185
* @param vertexCredentials The Vertex AI credentials to access the Gemini model.
152186
*/
153187
public Gemini(String modelName, VertexCredentials vertexCredentials) {
154-
super(modelName);
155-
Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null");
156-
this.apiClient = buildVertexClient(vertexCredentials, null);
188+
this(
189+
builder()
190+
.modelName(modelName)
191+
.vertexCredentials(
192+
Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null")));
157193
}
158194

159195
/**
160-
* Returns a new Builder instance for constructing Gemini objects. Note that when building a
161-
* Gemini object, at least one of apiKey, vertexCredentials, or an explicit apiClient must be set.
162-
* If multiple are set, the explicit apiClient will take precedence.
196+
* Returns a builder for configuring Gemini. An explicit client takes precedence over an API key
197+
* or Vertex credentials. If none are supplied, the client uses its default configuration.
163198
*
164199
* @return A new {@link Builder}.
165200
*/
@@ -174,6 +209,7 @@ public static class Builder {
174209
private String apiKey;
175210
private VertexCredentials vertexCredentials;
176211
private ExecutorService httpExecutorService;
212+
private HttpRetryOptions retryOptions;
177213

178214
private Builder() {}
179215

@@ -245,6 +281,21 @@ public Builder httpExecutorService(ExecutorService httpExecutorService) {
245281
return this;
246282
}
247283

284+
/**
285+
* Sets options for retrying transient HTTP failures through the Gemini client.
286+
*
287+
* <p>If an explicit {@link #apiClient(Client)} is supplied, its existing HTTP options take
288+
* precedence and these retry options are ignored.
289+
*
290+
* @param retryOptions The retry options to pass to the Google Gen AI client.
291+
* @return This builder.
292+
*/
293+
@CanIgnoreReturnValue
294+
public Builder retryOptions(HttpRetryOptions retryOptions) {
295+
this.retryOptions = retryOptions;
296+
return this;
297+
}
298+
248299
/**
249300
* Builds the {@link Gemini} instance.
250301
*
@@ -253,16 +304,7 @@ public Builder httpExecutorService(ExecutorService httpExecutorService) {
253304
*/
254305
public Gemini build() {
255306
Objects.requireNonNull(modelName, "modelName must be set.");
256-
257-
if (apiClient != null) {
258-
return new Gemini(modelName, apiClient);
259-
} else if (apiKey != null) {
260-
return new Gemini(modelName, buildApiKeyClient(apiKey, httpExecutorService));
261-
} else if (vertexCredentials != null) {
262-
return new Gemini(modelName, buildVertexClient(vertexCredentials, httpExecutorService));
263-
} else {
264-
return new Gemini(modelName, buildDefaultClient(httpExecutorService));
265-
}
307+
return new Gemini(this);
266308
}
267309
}
268310

‎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)