From 5e053076ed302ec7ff90dcc921e83bb53f7517bb Mon Sep 17 00:00:00 2001 From: "copybara-service[bot]" <56741989+copybara-service[bot]@users.noreply.github.com> Date: Mon, 8 Apr 2024 10:26:12 -0700 Subject: [PATCH] chore: [vertexai] eliminate null values in VertexAI credentialsProvider (#10654) PiperOrigin-RevId: 621574129 Co-authored-by: Jaycee Li --- .../com/google/cloud/vertexai/VertexAI.java | 68 ++++++++----------- .../google/cloud/vertexai/VertexAITest.java | 12 ---- 2 files changed, 29 insertions(+), 51 deletions(-) diff --git a/java-vertexai/google-cloud-vertexai/src/main/java/com/google/cloud/vertexai/VertexAI.java b/java-vertexai/google-cloud-vertexai/src/main/java/com/google/cloud/vertexai/VertexAI.java index f69f1b38a0e6..cb67d7bd2ff3 100644 --- a/java-vertexai/google-cloud-vertexai/src/main/java/com/google/cloud/vertexai/VertexAI.java +++ b/java-vertexai/google-cloud-vertexai/src/main/java/com/google/cloud/vertexai/VertexAI.java @@ -35,6 +35,7 @@ import com.google.common.collect.ImmutableList; import java.io.IOException; import java.util.List; +import java.util.Optional; import java.util.concurrent.locks.ReentrantLock; import java.util.logging.Level; import java.util.logging.Logger; @@ -59,7 +60,6 @@ public class VertexAI implements AutoCloseable { private final String location; private final String apiEndpoint; private final Transport transport; - // Will be null if the user doesn't provide Credentials or scopes private final CredentialsProvider credentialsProvider; // The clients will be instantiated lazily private PredictionServiceClient predictionServiceClient = null; @@ -73,64 +73,63 @@ public class VertexAI implements AutoCloseable { * @param location the default location to use when making API calls */ public VertexAI(String projectId, String location) { - this.projectId = projectId; - this.location = location; - this.apiEndpoint = String.format("%s-aiplatform.googleapis.com", this.location); - this.transport = Transport.GRPC; - this.credentialsProvider = null; + this( + projectId, + location, + Transport.GRPC, + ImmutableList.of(), + Optional.empty(), + Optional.empty()); } private VertexAI( String projectId, String location, - String apiEndpoint, Transport transport, - Credentials credentials, - List scopes) { - if (!scopes.isEmpty() && credentials != null) { + List scopes, + Optional credentials, + Optional apiEndpoint) { + if (!scopes.isEmpty() && credentials.isPresent()) { throw new IllegalArgumentException( "At most one of Credentials and scopes should be specified."); } checkArgument(!Strings.isNullOrEmpty(projectId), "projectId can't be null or empty"); checkArgument(!Strings.isNullOrEmpty(location), "location can't be null or empty"); - checkArgument(!Strings.isNullOrEmpty(apiEndpoint), "apiEndpoint can't be null or empty"); checkNotNull(transport, "transport can't be null"); this.projectId = projectId; this.location = location; - this.apiEndpoint = apiEndpoint; this.transport = transport; - if (credentials != null) { - this.credentialsProvider = FixedCredentialsProvider.create(credentials); + + if (credentials.isPresent()) { + this.credentialsProvider = FixedCredentialsProvider.create(credentials.get()); } else { this.credentialsProvider = - scopes.size() == 0 - ? null + scopes.isEmpty() + ? PredictionServiceSettings.defaultCredentialsProviderBuilder().build() : GoogleCredentialsProvider.newBuilder() .setScopesToApply(scopes) .setUseJwtAccessWithScope(true) .build(); } + + this.apiEndpoint = apiEndpoint.orElse(String.format("%s-aiplatform.googleapis.com", location)); } /** Builder for {@link VertexAI}. */ public static class Builder { private String projectId; private String location; - private Credentials credentials; - private String apiEndpoint; private Transport transport = Transport.GRPC; private ImmutableList scopes = ImmutableList.of(); + private Optional credentials = Optional.empty(); + private Optional apiEndpoint = Optional.empty(); public VertexAI build() { checkNotNull(projectId, "projectId must be set."); checkNotNull(location, "location must be set."); - // Default ApiEndpoint is set here as we need to make sure location is set. - if (apiEndpoint == null) { - apiEndpoint = String.format("%s-aiplatform.googleapis.com", location); - } - return new VertexAI(projectId, location, apiEndpoint, transport, credentials, scopes); + return new VertexAI(projectId, location, transport, scopes, credentials, apiEndpoint); } public Builder setProjectId(String projectId) { @@ -150,7 +149,7 @@ public Builder setLocation(String location) { public Builder setApiEndpoint(String apiEndpoint) { checkArgument(!Strings.isNullOrEmpty(apiEndpoint), "apiEndpoint can't be null or empty"); - this.apiEndpoint = apiEndpoint; + this.apiEndpoint = Optional.of(apiEndpoint); return this; } @@ -164,7 +163,7 @@ public Builder setTransport(Transport transport) { public Builder setCredentials(Credentials credentials) { checkNotNull(credentials, "credentials can't be null"); - this.credentials = credentials; + this.credentials = Optional.of(credentials); return this; } @@ -215,16 +214,9 @@ public String getApiEndpoint() { /** * Returns the default credentials to use when making API calls. * - * @return {@link Credentials} if the user has provided either scopes or credentials to the - * VertexAI object. + * @return {@link Credentials} to use when making API calls. */ public Credentials getCredentials() throws IOException { - // TODO(b/330780087): support getCredentials() when default credentials (no user provided - // credentials or scopes) are used. - if (credentialsProvider == null) { - throw new IllegalStateException( - "Either Credentials or scopes needs to be provided while instantiating VertexAI."); - } return credentialsProvider.getCredentials(); } @@ -266,9 +258,8 @@ private PredictionServiceSettings getPredictionServiceSettings() throws IOExcept builder = PredictionServiceSettings.newBuilder(); } builder.setEndpoint(String.format("%s:443", this.apiEndpoint)); - if (this.credentialsProvider != null) { - builder.setCredentialsProvider(this.credentialsProvider); - } + builder.setCredentialsProvider(this.credentialsProvider); + HeaderProvider headerProvider = FixedHeaderProvider.create( "user-agent", @@ -318,9 +309,8 @@ private LlmUtilityServiceSettings getLlmUtilityServiceClientSettings() throws IO settingsBuilder = LlmUtilityServiceSettings.newBuilder(); } settingsBuilder.setEndpoint(String.format("%s:443", this.apiEndpoint)); - if (this.credentialsProvider != null) { - settingsBuilder.setCredentialsProvider(this.credentialsProvider); - } + settingsBuilder.setCredentialsProvider(this.credentialsProvider); + HeaderProvider headerProvider = FixedHeaderProvider.create( "user-agent", diff --git a/java-vertexai/google-cloud-vertexai/src/test/java/com/google/cloud/vertexai/VertexAITest.java b/java-vertexai/google-cloud-vertexai/src/test/java/com/google/cloud/vertexai/VertexAITest.java index 2d9b38e57a44..9ee2a14c7f79 100644 --- a/java-vertexai/google-cloud-vertexai/src/test/java/com/google/cloud/vertexai/VertexAITest.java +++ b/java-vertexai/google-cloud-vertexai/src/test/java/com/google/cloud/vertexai/VertexAITest.java @@ -60,12 +60,6 @@ public void testInstantiateVertexAI_usingConstructor_shouldContainRightFields() assertThat(vertexAi.getLocation()).isEqualTo(TEST_LOCATION); assertThat(vertexAi.getTransport()).isEqualTo(Transport.GRPC); assertThat(vertexAi.getApiEndpoint()).isEqualTo(TEST_DEFAULT_ENDPOINT); - IllegalStateException thrown = - assertThrows(IllegalStateException.class, () -> vertexAi.getCredentials()); - assertThat(thrown) - .hasMessageThat() - .isEqualTo( - "Either Credentials or scopes needs to be provided while instantiating VertexAI."); } @Test @@ -143,11 +137,5 @@ public void testInstantiateVertexAI_builderWithTransport_shouldContainRightField assertThat(vertexAi.getLocation()).isEqualTo(TEST_LOCATION); assertThat(vertexAi.getTransport()).isEqualTo(Transport.REST); assertThat(vertexAi.getApiEndpoint()).isEqualTo(TEST_DEFAULT_ENDPOINT); - IllegalStateException thrown = - assertThrows(IllegalStateException.class, () -> vertexAi.getCredentials()); - assertThat(thrown) - .hasMessageThat() - .isEqualTo( - "Either Credentials or scopes needs to be provided while instantiating VertexAI."); } }