From 7dd1f6a33e49ae8c7d218d02ab47c2cdae5521d8 Mon Sep 17 00:00:00 2001 From: Matthew Tang Date: Mon, 20 Jul 2026 13:45:50 -0700 Subject: [PATCH] feat: Allow api key + proj/location for enterprise mode PiperOrigin-RevId: 951031340 --- src/main/java/com/google/genai/ApiClient.java | 32 ++++----- src/main/java/com/google/genai/Client.java | 16 ++--- .../java/com/google/genai/ClientTest.java | 20 ++++-- .../com/google/genai/HttpApiClientTest.java | 69 ++++++++++++++----- 4 files changed, 85 insertions(+), 52 deletions(-) diff --git a/src/main/java/com/google/genai/ApiClient.java b/src/main/java/com/google/genai/ApiClient.java index c43de67cc9d..6910d9e8857 100644 --- a/src/main/java/com/google/genai/ApiClient.java +++ b/src/main/java/com/google/genai/ApiClient.java @@ -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" @@ -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" @@ -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())) { @@ -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())); @@ -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/")) { diff --git a/src/main/java/com/google/genai/Client.java b/src/main/java/com/google/genai/Client.java index 763d04507f8..7401a4424b1 100644 --- a/src/main/java/com/google/genai/Client.java +++ b/src/main/java/com/google/genai/Client.java @@ -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; @@ -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) { @@ -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()); diff --git a/src/test/java/com/google/genai/ClientTest.java b/src/test/java/com/google/genai/ClientTest.java index c6a170ee673..457105b64ac 100644 --- a/src/test/java/com/google/genai/ClientTest.java +++ b/src/test/java/com/google/genai/ClientTest.java @@ -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 @@ -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 diff --git a/src/test/java/com/google/genai/HttpApiClientTest.java b/src/test/java/com/google/genai/HttpApiClientTest.java index a58520c9153..ca97a0eeb33 100644 --- a/src/test/java/com/google/genai/HttpApiClientTest.java +++ b/src/test/java/com/google/genai/HttpApiClientTest.java @@ -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 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(); @@ -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