Skip to content

feat: Allow api key + proj/location for enterprise mode #1125

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
copybara-service wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: Allow api key + proj/location for enterprise mode #1125

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
32 changes: 12 additions & 20 deletions src/main/java/com/google/genai/ApiClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,18 +163,6 @@ protected ApiClient(
customHttpOptions.flatMap(HttpOptions::baseUrl).map(url -> url.replaceAll("/$", ""));

// Validate constructor arguments combinations.
if (hasProject && hasApiKey) {
throw new IllegalArgumentException(
"For Vertex AI APIs, project and API key are mutually exclusive in the client"
+ " initializer. Please provide only one of them.");
}

if (hasLocation && hasApiKey) {
throw new IllegalArgumentException(
"For Vertex AI APIs, location and API key are mutually exclusive in the client"
+ " initializer. Please provide only one of them.");
}

if (hasCredentials && hasApiKey) {
throw new IllegalArgumentException(
"For Vertex AI APIs, API key cannot be set together with credentials. Please provide"
Expand All @@ -189,20 +177,24 @@ protected ApiClient(
+ " key from the environment variable.");
apiKeyValue = null;
}
if (hasApiKey && (hasEnvProjectValue || hasEnvLocationValue)) {
if (hasApiKey && !hasProject && !hasLocation && (hasEnvProjectValue || hasEnvLocationValue)) {
// Explicit API key takes precedence over implicit project/location.
logger.warning(
"Warning: The user provided Vertex AI API key will take precedence over the"
+ " project/location from the environment variables.");
projectValue = null;
locationValue = null;
} else if ((hasProject || hasLocation) && hasEnvApiKeyValue) {
} else if ((hasProject || hasLocation) && !hasApiKey && hasEnvApiKeyValue) {
// Explicit project/location takes precedence over implicit API key.
logger.warning(
"Warning: The user provided project/location will take precedence over the API key from"
+ " the environment variable.");
apiKeyValue = null;
} else if ((hasEnvProjectValue || hasEnvLocationValue) && hasEnvApiKeyValue) {
} else if (!hasProject
&& !hasLocation
&& !hasApiKey
&& (hasEnvProjectValue || hasEnvLocationValue)
&& hasEnvApiKeyValue) {
// Implicit project/location takes precedence over implicit API key.
logger.warning(
"Warning: The project/location from the environment variables will take precedence over"
Expand Down Expand Up @@ -233,7 +225,7 @@ protected ApiClient(
initHttpOptionsBuilder.baseUrl(customBaseUrl.get());
projectValue = null;
locationValue = null;
} else if (apiKeyValue != null
} else if ((apiKeyValue != null && locationValue == null && !customBaseUrl.isPresent())
|| (locationValue != null
&& locationValue.equals("global")
&& !customBaseUrl.isPresent())) {
Expand All @@ -253,9 +245,9 @@ protected ApiClient(
this.location = Optional.ofNullable(locationValue);
this.customBaseUrl = customBaseUrl;

// Only set credentials if using project/location.
// Only set credentials if using project/location and no API key is provided.
this.credentials =
!this.project.isPresent()
(!this.project.isPresent() || this.apiKey.isPresent())
? Optional.empty()
: Optional.of(credentials.orElseGet(() -> defaultCredentials()));

Expand Down Expand Up @@ -794,10 +786,10 @@ boolean shouldPrependVertexProjectPath(
&& httpOptions.baseUrlResourceScope().get().knownEnum() == Known.COLLECTION) {
return false;
}
if (this.apiKey.isPresent()) {
if (!this.vertexAI) {
return false;
}
if (!this.vertexAI) {
if (!this.project.isPresent() || !this.location.isPresent()) {
return false;
}
if (path.startsWith("projects/")) {
Expand Down
16 changes: 6 additions & 10 deletions src/main/java/com/google/genai/Client.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -227,7 +227,8 @@ private Client(

if (enterprise.isPresent() && vertexAI.isPresent() && !enterprise.get().equals(vertexAI.get())) {
throw new IllegalArgumentException(
"enterprise and vertexAI flags have conflicting values, please set enterprise value only.");
"enterprise and vertexAI flags have conflicting values, please set enterprise value"
+ " only.");
}

boolean useVertexAI;
Expand All @@ -245,7 +246,8 @@ private Client(

if (enterpriseEnvPresent && vertexEnvPresent && !enterpriseEnv.equalsIgnoreCase(vertexEnv)) {
logger.warning(
"Warning: Both GOOGLE_GENAI_USE_ENTERPRISE and GOOGLE_GENAI_USE_VERTEXAI are set with conflicting values. The value of GOOGLE_GENAI_USE_ENTERPRISE will be used.");
"Warning: Both GOOGLE_GENAI_USE_ENTERPRISE and GOOGLE_GENAI_USE_VERTEXAI are set with"
+ " conflicting values. The value of GOOGLE_GENAI_USE_ENTERPRISE will be used.");
}

if (enterpriseEnvPresent) {
Expand All @@ -257,14 +259,8 @@ private Client(
}
}

if (project.isPresent() || location.isPresent()) {
if (apiKey.isPresent()) {
throw new IllegalArgumentException(
"Project/location and API key are mutually exclusive in the client initializer.");
}
if (!useVertexAI) {
throw new IllegalArgumentException("Gemini API do not support project/location.");
}
if ((project.isPresent() || location.isPresent()) && !useVertexAI) {
throw new IllegalArgumentException("Gemini API does not support project/location.");
}

this.debugConfig = debugConfig.orElse(new DebugConfig());
Expand Down
20 changes: 16 additions & 4 deletions src/test/java/com/google/genai/ClientTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -150,7 +150,7 @@ public void testInitClientFromBuilder_setProjectInMldev_throwsException() {
() -> Client.builder().vertexAI(false).project(PROJECT).build());

// Assert
assertEquals("Gemini API do not support project/location.", exception.getMessage());
assertEquals("Gemini API does not support project/location.", exception.getMessage());
}

@Test
Expand All @@ -162,9 +162,21 @@ public void testInitClientFromBuilder_setBothApiKeyAndProject_throwsException()
() -> Client.builder().apiKey(API_KEY).project(PROJECT).build());

// Assert
assertEquals(
"Project/location and API key are mutually exclusive in the client initializer.",
exception.getMessage());
assertEquals("Gemini API does not support project/location.", exception.getMessage());
}

@Test
public void testInitClientFromBuilder_setApiKeyAndProjectAndLocationInVertex_success() {
// Act
Client client =
Client.builder().apiKey(API_KEY).project(PROJECT).location(LOCATION).vertexAI(true).build();

// Assert
assertEquals(API_KEY, client.apiKey());
assertEquals(PROJECT, client.project());
assertEquals(LOCATION, client.location());
assertTrue(client.vertexAI());
assertEquals("https://location-aiplatform.googleapis.com", client.baseUrl().orElse(null));
}

@Test
Expand Down
69 changes: 51 additions & 18 deletions src/test/java/com/google/genai/HttpApiClientTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,44 @@ public void testRequestPostMethodWithVertexAI() throws Exception {
assertEquals("Bearer", capturedRequest.header("Authorization"));
assertNull(capturedRequest.header("x-goog-api-key"));

RequestBody body = capturedRequest.body();
assertNotNull(body);
final Buffer buffer = new Buffer();
body.writeTo(buffer);
assertEquals("application/json; charset=utf-8", body.contentType().toString());
}

@Test
public void testRequestPostMethodWithVertexExpressMode() throws Exception {
// Arrange
HttpApiClient client =
new HttpApiClient(
Optional.of(API_KEY),
Optional.of(PROJECT),
Optional.of(LOCATION),
Optional.empty(),
Optional.empty(),
Optional.empty());
setMockClient(client);

// Act
client.request("POST", TEST_PATH, TEST_REQUEST_JSON, Optional.empty());

// Assert
ArgumentCaptor<Request> requestCaptor = ArgumentCaptor.forClass(Request.class);
verify(mockHttpClient).newCall(requestCaptor.capture());
Request capturedRequest = requestCaptor.getValue();

assertEquals("POST", capturedRequest.method());
assertEquals(
String.format(
"https://%s-aiplatform.googleapis.com/v1beta1/projects/%s/locations/%s/%s",
LOCATION, PROJECT, LOCATION, TEST_PATH),
capturedRequest.url().toString());
assertNotNull(capturedRequest.header("x-goog-api-key"));
assertEquals(API_KEY, capturedRequest.header("x-goog-api-key"));
assertNull(capturedRequest.header("Authorization"));

RequestBody body = capturedRequest.body();
assertNotNull(body);
final Buffer buffer = new Buffer();
Expand Down Expand Up @@ -797,24 +835,19 @@ public void testInitHttpClientWithCustomCredentialsAndApiKey_throwsException() t
}

@Test
public void testInitHttpClientVertexWithProjectLocationAndApiKey_throwsException()
throws Exception {
// Explicit proj/location and API key are not allowed.
IllegalArgumentException exception =
assertThrows(
IllegalArgumentException.class,
() ->
new HttpApiClient(
Optional.of(API_KEY),
Optional.of(PROJECT),
Optional.of(LOCATION),
Optional.empty(),
Optional.empty(),
Optional.empty()));
assertEquals(
"For Vertex AI APIs, project and API key are mutually exclusive in the client initializer."
+ " Please provide only one of them.",
exception.getMessage());
public void testInitHttpClientVertexWithProjectLocationAndApiKey_success() throws Exception {
HttpApiClient client =
new HttpApiClient(
Optional.of(API_KEY),
Optional.of(PROJECT),
Optional.of(LOCATION),
Optional.empty(),
Optional.empty(),
Optional.empty());
assertEquals(API_KEY, client.apiKey());
assertEquals(PROJECT, client.project());
assertEquals(LOCATION, client.location());
assertTrue(client.vertexAI());
}

@Test
Expand Down
Loading