Skip to content

Commit 5bed284

Browse files
copybara-service[bot]jaycee-li
copybara-service[bot]
and
jaycee-li
authored
feat: [vertexai] add custom headers support in VertexAI (#11085)
PiperOrigin-RevId: 665451706 Co-authored-by: Jaycee Li <jayceeli@google.com>
1 parent 2ba5930 commit 5bed284

2 files changed

Lines changed: 93 additions & 0 deletions

File tree

‎java-vertexai/google-cloud-vertexai/src/main/java/com/google/cloud/vertexai/VertexAI.java‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,12 @@
3636
import com.google.common.base.Supplier;
3737
import com.google.common.base.Suppliers;
3838
import com.google.common.collect.ImmutableList;
39+
import com.google.common.collect.ImmutableMap;
3940
import com.google.errorprone.annotations.CanIgnoreReturnValue;
4041
import java.io.IOException;
42+
import java.util.HashMap;
4143
import java.util.List;
44+
import java.util.Map;
4245
import java.util.Optional;
4346
import java.util.logging.Level;
4447
import java.util.logging.Logger;
@@ -63,6 +66,7 @@ public class VertexAI implements AutoCloseable {
6366
private final String location;
6467
private final String apiEndpoint;
6568
private final Transport transport;
69+
private final HeaderProvider headerProvider;
6670
private final CredentialsProvider credentialsProvider;
6771

6872
private final transient Supplier<PredictionServiceClient> predictionClientSupplier;
@@ -85,6 +89,7 @@ public VertexAI(String projectId, String location) {
8589
location,
8690
Transport.GRPC,
8791
ImmutableList.of(),
92+
/* customHeaders= */ ImmutableMap.of(),
8893
/* credentials= */ Optional.empty(),
8994
/* apiEndpoint= */ Optional.empty(),
9095
/* predictionClientSupplierOpt= */ Optional.empty(),
@@ -108,6 +113,7 @@ public VertexAI() {
108113
null,
109114
Transport.GRPC,
110115
ImmutableList.of(),
116+
/* customHeaders= */ ImmutableMap.of(),
111117
/* credentials= */ Optional.empty(),
112118
/* apiEndpoint= */ Optional.empty(),
113119
/* predictionClientSupplierOpt= */ Optional.empty(),
@@ -119,6 +125,7 @@ private VertexAI(
119125
String location,
120126
Transport transport,
121127
List<String> scopes,
128+
Map<String, String> customHeaders,
122129
Optional<Credentials> credentials,
123130
Optional<String> apiEndpoint,
124131
Optional<Supplier<PredictionServiceClient>> predictionClientSupplierOpt,
@@ -131,6 +138,15 @@ private VertexAI(
131138
this.location = Strings.isNullOrEmpty(location) ? inferLocation() : location;
132139
this.transport = transport;
133140

141+
String sdkHeader =
142+
String.format(
143+
"%s/%s",
144+
Constants.USER_AGENT_HEADER,
145+
GaxProperties.getLibraryVersion(PredictionServiceSettings.class));
146+
Map<String, String> headers = new HashMap<>(customHeaders);
147+
headers.compute("user-agent", (k, v) -> v == null ? sdkHeader : sdkHeader + " " + v);
148+
this.headerProvider = FixedHeaderProvider.create(headers);
149+
134150
if (credentials.isPresent()) {
135151
this.credentialsProvider = FixedCredentialsProvider.create(credentials.get());
136152
} else {
@@ -160,6 +176,7 @@ public static class Builder {
160176
private String location;
161177
private Transport transport = Transport.GRPC;
162178
private ImmutableList<String> scopes = ImmutableList.of();
179+
private ImmutableMap<String, String> customHeaders = ImmutableMap.of();
163180
private Optional<Credentials> credentials = Optional.empty();
164181
private Optional<String> apiEndpoint = Optional.empty();
165182

@@ -174,6 +191,7 @@ public VertexAI build() {
174191
location,
175192
transport,
176193
scopes,
194+
customHeaders,
177195
credentials,
178196
apiEndpoint,
179197
Optional.ofNullable(predictionClientSupplier),
@@ -240,6 +258,14 @@ public Builder setScopes(List<String> scopes) {
240258
this.scopes = ImmutableList.copyOf(scopes);
241259
return this;
242260
}
261+
262+
@CanIgnoreReturnValue
263+
public Builder setCustomHeaders(Map<String, String> customHeaders) {
264+
checkNotNull(customHeaders, "customHeaders can't be null");
265+
266+
this.customHeaders = ImmutableMap.copyOf(customHeaders);
267+
return this;
268+
}
243269
}
244270

245271
/**
@@ -278,6 +304,15 @@ public String getApiEndpoint() {
278304
return apiEndpoint;
279305
}
280306

307+
/**
308+
* Returns the headers to use when making API calls.
309+
*
310+
* @return a map of headers to use when making API calls.
311+
*/
312+
public Map<String, String> getHeaders() {
313+
return headerProvider.getHeaders();
314+
}
315+
281316
/**
282317
* Returns the default credentials to use when making API calls.
283318
*

‎java-vertexai/google-cloud-vertexai/src/test/java/com/google/cloud/vertexai/VertexAITest.java‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,12 +22,15 @@
2222
import static org.mockito.Mockito.mockStatic;
2323
import static org.mockito.Mockito.when;
2424

25+
import com.google.api.gax.core.GaxProperties;
2526
import com.google.api.gax.core.GoogleCredentialsProvider;
2627
import com.google.auth.oauth2.GoogleCredentials;
2728
import com.google.cloud.vertexai.api.PredictionServiceClient;
2829
import com.google.cloud.vertexai.api.PredictionServiceSettings;
2930
import com.google.common.collect.ImmutableList;
3031
import java.io.IOException;
32+
import java.util.HashMap;
33+
import java.util.Map;
3134
import java.util.Optional;
3235
import org.junit.Rule;
3336
import org.junit.Test;
@@ -397,4 +400,59 @@ public void testInstantiateVertexAI_builderWithTransport_shouldContainRightField
397400
assertThat(vertexAi.getTransport()).isEqualTo(Transport.REST);
398401
assertThat(vertexAi.getApiEndpoint()).isEqualTo(TEST_DEFAULT_ENDPOINT);
399402
}
403+
404+
@Test
405+
public void testInstantiateVertexAI_builderWithCustomHeaders_shouldContainRightFields()
406+
throws IOException {
407+
Map<String, String> customHeaders = new HashMap<>();
408+
customHeaders.put("test_key", "test_value");
409+
410+
vertexAi =
411+
new VertexAI.Builder()
412+
.setProjectId(TEST_PROJECT)
413+
.setLocation(TEST_LOCATION)
414+
.setCustomHeaders(customHeaders)
415+
.build();
416+
417+
assertThat(vertexAi.getProjectId()).isEqualTo(TEST_PROJECT);
418+
assertThat(vertexAi.getLocation()).isEqualTo(TEST_LOCATION);
419+
// headers should include both the sdk header and the custom headers
420+
Map<String, String> expectedHeaders = new HashMap<>(customHeaders);
421+
expectedHeaders.put(
422+
"user-agent",
423+
String.format(
424+
"%s/%s",
425+
Constants.USER_AGENT_HEADER,
426+
GaxProperties.getLibraryVersion(PredictionServiceSettings.class)));
427+
assertThat(vertexAi.getHeaders()).isEqualTo(expectedHeaders);
428+
}
429+
430+
@Test
431+
public void
432+
testInstantiateVertexAI_builderWithCustomHeadersWithSdkReservedKey_shouldContainRightFields()
433+
throws IOException {
434+
Map<String, String> customHeadersWithSdkReservedKey = new HashMap<>();
435+
customHeadersWithSdkReservedKey.put("user-agent", "test_value");
436+
437+
vertexAi =
438+
new VertexAI.Builder()
439+
.setProjectId(TEST_PROJECT)
440+
.setLocation(TEST_LOCATION)
441+
.setCustomHeaders(customHeadersWithSdkReservedKey)
442+
.build();
443+
444+
assertThat(vertexAi.getProjectId()).isEqualTo(TEST_PROJECT);
445+
assertThat(vertexAi.getLocation()).isEqualTo(TEST_LOCATION);
446+
// headers should include sdk reserved key with value of both the sdk header and the custom
447+
// headers
448+
Map<String, String> expectedHeaders = new HashMap<>();
449+
expectedHeaders.put(
450+
"user-agent",
451+
String.format(
452+
"%s/%s %s",
453+
Constants.USER_AGENT_HEADER,
454+
GaxProperties.getLibraryVersion(PredictionServiceSettings.class),
455+
"test_value"));
456+
assertThat(vertexAi.getHeaders()).isEqualTo(expectedHeaders);
457+
}
400458
}

0 commit comments

Comments
 (0)