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 09402c2b5..e866a8c42 100644 --- a/core/src/main/java/com/google/adk/models/Gemini.java +++ b/core/src/main/java/com/google/adk/models/Gemini.java @@ -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; @@ -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)) @@ -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(); @@ -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); } /** @@ -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); } /** @@ -174,6 +214,7 @@ public static class Builder { private String apiKey; private VertexCredentials vertexCredentials; private ExecutorService httpExecutorService; + private HttpRetryOptions retryOptions; private Builder() {} @@ -245,6 +286,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. * @@ -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)); } } } 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..f9596c48f 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_leavesRetriesDisabled() { + HttpOptions httpOptions = Gemini.buildHttpOptions(null); + + assertThat(httpOptions.retryOptions()).isEmpty(); + } + // Test cases for processRawResponses static method @Test public void processRawResponses_withTextChunks_emitsPartialResponses() {