From 8bdb589a9b84f762860936d6bbe4398839638ef2 Mon Sep 17 00:00:00 2001
From: Matthew Tang <tangmatthew@google.com>
Date: Mon, 20 Jul 2026 13:45:50 -0700
Subject: [PATCH] feat: Allow api key + proj/location for enterprise mode

PiperOrigin-RevId: 951031340
---
 google/genai/_api_client.py                   | 34 ++++++++++-------
 .../client/test_client_initialization.py      | 38 ++++++++++++++-----
 2 files changed, 49 insertions(+), 23 deletions(-)

diff --git a/google/genai/_api_client.py b/google/genai/_api_client.py
index c8e821ea2..e71e63216 100644
--- a/google/genai/_api_client.py
+++ b/google/genai/_api_client.py
@@ -641,13 +641,10 @@ def __init__(
         self.vertexai = env_vertexai
 
     # Validate explicitly set initializer values.
-    if (project or location) and api_key:
-      # API cannot consume both project/location and api_key.
-      raise ValueError(
-          'Project/location and API key are mutually exclusive in the client'
-          ' initializer.'
-      )
-    elif credentials and api_key:
+    if (project or location) and not self.vertexai:
+      raise ValueError('Gemini API does not support project/location.')
+
+    if credentials and api_key:
       # API cannot consume both credentials and api_key.
       raise ValueError(
           'Credentials and API key are mutually exclusive in the client'
@@ -699,7 +696,12 @@ def __init__(
             + ' over the API key from the environment variable.'
         )
         self.api_key = None
-      elif (env_location or env_project) and api_key:
+      elif (
+          api_key
+          and not project
+          and not location
+          and (env_project or env_location)
+      ):
         # Explicit api_key takes precedence over implicit project/location.
         logger.info(
             'The user provided Vertex AI API key will take precedence over the'
@@ -707,14 +709,20 @@ def __init__(
         )
         self.project = None
         self.location = None
-      elif (project or location) and env_api_key:
+      elif (project or location) and not api_key and env_api_key:
         # Explicit project/location takes precedence over implicit api_key.
         logger.info(
             'The user provided project/location will take precedence over the'
             + ' Vertex AI API key from the environment variable.'
         )
         self.api_key = None
-      elif (env_location or env_project) and env_api_key:
+      elif (
+          not project
+          and not location
+          and not api_key
+          and (env_project or env_location)
+          and env_api_key
+      ):
         # Implicit project/location takes precedence over implicit api_key.
         logger.info(
             'The project/location from the environment variables will take'
@@ -755,7 +763,7 @@ def __init__(
             'Project or API key must be set when using the Vertex AI API.'
         )
       if (
-          self.api_key or self.location == 'global'
+          (self.api_key and not self.location) or self.location == 'global'
       ) and not self.custom_base_url:
         self._http_options.base_url = f'https://aiplatform.googleapis.com/'
       elif (
@@ -1412,7 +1420,7 @@ def _request_once(
   ) -> HttpResponse:
     data: Optional[Union[str, bytes]] = None
     # If using proj/location, fetch ADC
-    if self.vertexai and (self.project or self.location):
+    if self.vertexai and (self.project or self.location) and not self.api_key:
       http_request.headers['Authorization'] = f'Bearer {self._access_token()}'
       if self._credentials and self._credentials.quota_project_id:
         http_request.headers['x-goog-user-project'] = (
@@ -1492,7 +1500,7 @@ async def _async_request_once(
     data: Optional[bytes] = None
 
     # If using proj/location, fetch ADC
-    if self.vertexai and (self.project or self.location):
+    if self.vertexai and (self.project or self.location) and not self.api_key:
       http_request.headers['Authorization'] = (
           f'Bearer {await self._async_access_token()}'
       )
diff --git a/google/genai/tests/client/test_client_initialization.py b/google/genai/tests/client/test_client_initialization.py
index 9a99ddc37..659d1ca4b 100644
--- a/google/genai/tests/client/test_client_initialization.py
+++ b/google/genai/tests/client/test_client_initialization.py
@@ -931,19 +931,37 @@ def test_vertexai_apikey_from_env_both_api_keys(monkeypatch, caplog):
   )
 
 
-def test_vertexai_apikey_invalid_constructor1():
-  # Vertex AI Express mode uses API key on Vertex AI.
+def test_gemini_project_location_invalid(monkeypatch):
+  monkeypatch.setenv("GOOGLE_API_KEY", "test_key")
+  with pytest.raises(
+      ValueError, match="Gemini API does not support project/location."
+  ):
+    Client(project="fake_project_id", vertexai=False)
+
+
+def test_vertexai_apikey_with_project_and_location(monkeypatch):
   api_key = "vertexai_api_key"
   project_id = "fake_project_id"
-  location = "fake-location"
+  location = "us-central1"
+  monkeypatch.setenv("GOOGLE_CLOUD_PROJECT", "")
+  monkeypatch.setenv("GOOGLE_CLOUD_LOCATION", "")
+  monkeypatch.setenv("GOOGLE_API_KEY", "")
 
-  with pytest.raises(ValueError):
-    Client(
-        api_key=api_key,
-        project=project_id,
-        location=location,
-        vertexai=True,
-    )
+  client = Client(
+      api_key=api_key,
+      project=project_id,
+      location=location,
+      vertexai=True,
+  )
+
+  assert client.models._api_client.vertexai
+  assert client.models._api_client.api_key == api_key
+  assert client.models._api_client.project == project_id
+  assert client.models._api_client.location == location
+  assert client._api_client._http_options.base_url == (
+      "https://us-central1-aiplatform.googleapis.com/"
+  )
+  assert isinstance(client.models._api_client, api_client.BaseApiClient)
 
 
 def test_vertexai_apikey_combo1(monkeypatch):