Skip to content

Commit

Permalink
chore: [vertexai] eliminate null values in VertexAI credentialsProvid…
Browse files Browse the repository at this point in the history 
…er (#10654)

PiperOrigin-RevId: 621574129

Co-authored-by: Jaycee Li <[email protected]>
  • Loading branch information
@copybara-service @jaycee-li
copybara-service[bot] and jaycee-li authored Apr 8, 2024
1 parent 2cf8711 commit 5e05307
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 51 deletions.
68 changes: 29 additions & 39 deletions java-vertexai/google-cloud-vertexai/src/main/java/com/google/cloud/vertexai/VertexAI.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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<String> scopes) {
if (!scopes.isEmpty() && credentials != null) {
List<String> scopes,
Optional<Credentials> credentials,
Optional<String> 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<String> scopes = ImmutableList.of();
private Optional<Credentials> credentials = Optional.empty();
private Optional<String> 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) {
Expand All @@ -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;
}

Expand All @@ -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;
}

Expand Down Expand Up @@ -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();
}

Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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",
Expand Down
12 changes: 0 additions & 12 deletions ...-vertexai/google-cloud-vertexai/src/test/java/com/google/cloud/vertexai/VertexAITest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.");
}
}

0 comments on commit 5e05307

Please sign in to comment.