Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
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
79 changes: 68 additions & 11 deletions core/src/main/java/com/google/adk/models/Gemini.java
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,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 @@ -76,20 +77,24 @@ private static OkHttpClient prepareHttpClient(@Nullable ExecutorService executor
}

private static Client buildApiKeyClient(
String apiKey, @Nullable ExecutorService executorService) {
String apiKey,
@Nullable ExecutorService executorService,
@Nullable HttpRetryOptions retryOptions) {
return Client.builder()
.apiKey(apiKey)
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
.httpOptions(buildHttpOptions(retryOptions))
.clientOptions(
ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build())
.build();
}

private static Client buildVertexClient(
VertexCredentials vertexCredentials, @Nullable ExecutorService executorService) {
VertexCredentials vertexCredentials,
@Nullable ExecutorService executorService,
@Nullable HttpRetryOptions retryOptions) {
Client.Builder apiClientBuilder =
Client.builder()
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
.httpOptions(buildHttpOptions(retryOptions))
.clientOptions(
ClientOptions.builder()
.customHttpClient(prepareHttpClient(executorService))
Expand All @@ -100,14 +105,24 @@ private static Client buildVertexClient(
return apiClientBuilder.build();
}

private static Client buildDefaultClient(@Nullable ExecutorService executorService) {
private static Client buildDefaultClient(
@Nullable ExecutorService executorService, @Nullable HttpRetryOptions retryOptions) {
return Client.builder()
.httpOptions(HttpOptions.builder().headers(TRACKING_HEADERS).build())
.httpOptions(buildHttpOptions(retryOptions))
.clientOptions(
ClientOptions.builder().customHttpClient(prepareHttpClient(executorService)).build())
.build();
}

@VisibleForTesting
static HttpOptions buildHttpOptions(@Nullable HttpRetryOptions retryOptions) {
HttpOptions.Builder httpOptionsBuilder = HttpOptions.builder().headers(TRACKING_HEADERS);
if (retryOptions != null) {
httpOptionsBuilder.retryOptions(retryOptions);
}
return httpOptionsBuilder.build();
}

static {
String frameworkLabel = "google-adk/" + Version.JAVA_ADK_VERSION;
String languageLabel = "gl-java/" + JAVA_VERSION.value();
Expand Down Expand Up @@ -139,9 +154,20 @@ public Gemini(String modelName, Client apiClient) {
* @param apiKey The Google Gemini API key.
*/
public Gemini(String modelName, String apiKey) {
this(modelName, apiKey, /* retryOptions= */ null);
}

/**
* Constructs a new Gemini instance with a Google Gemini API key and optional retry options.
*
* @param modelName The name of the Gemini model to use (e.g., "gemini-2.0-flash").
* @param apiKey The Google Gemini API key.
* @param retryOptions Options for retrying transient HTTP failures.
*/
public Gemini(String modelName, String apiKey, @Nullable HttpRetryOptions retryOptions) {
super(modelName);
Objects.requireNonNull(apiKey, "apiKey cannot be null");
this.apiClient = buildApiKeyClient(apiKey, null);
this.apiClient = buildApiKeyClient(apiKey, null, retryOptions);
}

/**
Expand All @@ -151,9 +177,23 @@ public Gemini(String modelName, String apiKey) {
* @param vertexCredentials The Vertex AI credentials to access the Gemini model.
*/
public Gemini(String modelName, VertexCredentials vertexCredentials) {
this(modelName, vertexCredentials, /* retryOptions= */ null);
}

/**
* Constructs a new Gemini instance with Vertex AI credentials and optional retry options.
*
* @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.
* @param retryOptions Options for retrying transient HTTP failures.
*/
public Gemini(
String modelName,
VertexCredentials vertexCredentials,
@Nullable HttpRetryOptions retryOptions) {
super(modelName);
Objects.requireNonNull(vertexCredentials, "vertexCredentials cannot be null");
this.apiClient = buildVertexClient(vertexCredentials, null);
this.apiClient = buildVertexClient(vertexCredentials, null, retryOptions);
}

/**
Expand All @@ -174,6 +214,7 @@ public static class Builder {
private String apiKey;
private VertexCredentials vertexCredentials;
private ExecutorService httpExecutorService;
private HttpRetryOptions retryOptions;

private Builder() {}

Expand Down Expand Up @@ -245,6 +286,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 @@ -257,11 +313,12 @@ public Gemini build() {
if (apiClient != null) {
return new Gemini(modelName, apiClient);
} else if (apiKey != null) {
return new Gemini(modelName, buildApiKeyClient(apiKey, httpExecutorService));
return new Gemini(modelName, buildApiKeyClient(apiKey, httpExecutorService, retryOptions));
} else if (vertexCredentials != null) {
return new Gemini(modelName, buildVertexClient(vertexCredentials, httpExecutorService));
return new Gemini(
modelName, buildVertexClient(vertexCredentials, httpExecutorService, retryOptions));
} else {
return new Gemini(modelName, buildDefaultClient(httpExecutorService));
return new Gemini(modelName, buildDefaultClient(httpExecutorService, retryOptions));
}
}
}
Expand Down
26 changes: 26 additions & 0 deletions core/src/test/java/com/google/adk/models/GeminiTest.java
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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_leavesRetriesDisabled() {
HttpOptions httpOptions = Gemini.buildHttpOptions(null);

assertThat(httpOptions.retryOptions()).isEmpty();
}

// Test cases for processRawResponses static method
@Test
public void processRawResponses_withTextChunks_emitsPartialResponses() {
Expand Down
Loading